Nick088's picture
Update app.py
827c867 verified
raw
history blame
1.87 kB
import torch
from diffusers import StableDiffusion3Pipeline
import gradio as gr
import os
import spaces
HF_TOKEN = os.environ.get('HF_TOKEN')
if torch.cuda.is_available():
device = "cuda"
print("Using GPU")
else:
device = "cpu"
print("Using CPU")
# Initialize the pipeline and download the model
pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium", torch_dtype=torch.float16, use_auth_token=HF_TOKEN)
pipe.to(device)
# Define the image generation function
@spaces.GPU(duration=60)
def generate_image(prompt):
image = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
num_inference_steps=num_inference_steps,
height=height,
width=width,
guidance_scale=guidance_scale,
).images[0]
return image
# Create the Gradio interface
prompt = gr.Textbox(label="Prompt", info="Describe the image you want", placeholder="A cat...")
negative_prompt = gr.Textbox(label="Negative Prompt", info="Describe what you don't want in the image", placeholder="Ugly, bad anatomy...")
num_inference_steps = gr.Number(label="Number of Inference Steps", precision=0, value=25)
height = gr.Number(label="Number of Inference Steps", precision=0, value=1024)
width = gr.Number(label="Number of Inference Steps", precision=0, value=1024)
guidance_scale = gr.Number(minimum=0.1, value=7.5, label="Guidance Scale", info="The number of denoising steps of the image. More denoising steps usually lead to a higher quality image at the cost of slower inference")
interface = gr.Interface(
fn=generate_image,
inputs=[prompt, negative_prompt, num_inference_steps, height, width, guidance_scale]
outputs="image",
title="Stable Diffusion 3 Medium",
description="Made by [Nick088](https://linktr.ee/Nick088"
)
# Launch the interface
interface.launch(share=False)