image-generator / app.py
Dubberkuro's picture
Update app.py
ae6c33e verified
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
pipe = None
def load_pipeline():
global pipe
if pipe is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
revision = "fp16" if device == "cuda" else "main"
torch_dtype = torch.float16 if device == "cuda" else torch.float32
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
revision=revision,
torch_dtype=torch_dtype
).to(device)
return pipe
def generate_image(prompt, num_inference_steps=30, guidance_scale=7.5):
try:
pipeline = load_pipeline()
with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
result = pipeline(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)
image = result.images[0]
return image
except Exception as e:
return f"Error: {str(e)}"
iface = gr.Interface(
fn=generate_image,
inputs=[
gr.Textbox(label="Prompt", lines=2),
gr.Slider(1, 100, step=1, value=30, label="Inference Steps"),
gr.Slider(0.0, 15.0, step=0.5, value=7.5, label="Guidance Scale")
],
outputs=[
gr.Image(type="pil", label="Generated Image")
],
title="Stable Diffusion Image Generator",
description="Generate images from text prompts using Stable Diffusion."
)
if __name__ == "__main__":
iface.launch()