Spaces:
Runtime error
Runtime error
import torch | |
from diffusers import StableDiffusion3Pipeline | |
import gradio as gr | |
import os | |
import spaces | |
from huggingface_hub import snapshot_download | |
HF_TOKEN = os.getenv("HF_TOKEN") | |
model_path = snapshot_download( | |
repo_id="stabilityai/stable-diffusion-3-medium", | |
revision="refs/pr/26", | |
repo_type="model", | |
ignore_patterns=["*.md", "*..gitattributes"], | |
local_dir="stable-diffusion-3-medium", | |
token=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(model_path, torch_dtype=torch.float16) | |
pipe.to(device) | |
# Define the image generation function | |
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) |