Updated for ZeroGPU

app.py

@@ -1,28 +1,38 @@
-from diffusers import StableDiffusionPipeline
 import gradio as gr
-import torch
+from transformers import pipeline
+import spaces
 
-# Load the Stable Diffusion model
-pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5"
-)
+# Load the pipeline (lazy-load the model to save resources)
+@spaces.GPU
+def load_model():
+    return pipeline(
+        "text-to-image",
+        model="stabilityai/stable-diffusion-2-1",
+        torch_dtype="float16"  # Ensure compatibility with ZeroGPU
+    )
 
-pipe.to("cuda" if torch.cuda.is_available() else "cpu")
+# Initialize the pipeline
+model = load_model()
 
-# Define the image generation function
-def generate_image(prompt):
-    image = pipe(prompt).images[0]
-    return image
+# Function to generate images
+@spaces.GPU
+def generate_image(prompt, guidance_scale=7.5):
+    print(f"Generating image for prompt: {prompt}")
+    images = model(prompt, guidance_scale=guidance_scale)
+    return images[0]
 
-# Build the Gradio interface
+# Gradio interface
 interface = gr.Interface(
     fn=generate_image,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your creative prompt here..."),
-    outputs="image",
+    inputs=[
+        gr.Textbox(label="Prompt", placeholder="Describe your image..."),
+        gr.Slider(1, 20, value=7.5, label="Guidance Scale")
+    ],
+    outputs=gr.Image(label="Generated Image"),
     title="Text-to-Image Generator",
-    description="Enter a creative prompt and see it turned into an image!"
+    description="Generate images from text prompts using Stable Diffusion."
 )
 
 # Launch the app
-interface.launch()
-
+if __name__ == "__main__":
+    interface.launch(server_name="0.0.0.0", server_port=7860)