Update app.py

app.py

@@ -1,38 +1,47 @@
 import gradio as gr
-from transformers import pipeline
-import spaces
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+import torch
 
-# 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
+    # Specify the Stable Diffusion pipeline with an appropriate model type
+    pipeline = StableDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-2-1",
+        torch_dtype=torch.float16,
+        revision="fp16",
+        safety_checker=None  # Disable safety checker if necessary
     )
+    
+    # Set the scheduler (optional but recommended)
+    pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
+    
+    # Move pipeline to GPU or ZeroGPU
+    pipeline = pipeline.to("cuda")  # or ZeroGPU-specific setup
+    
+    return pipeline
 
-# Initialize the pipeline
-model = load_model()
+# Initialize the model
+try:
+    model = load_model()
+except Exception as e:
+    print(f"Error loading the model: {e}")
 
-# 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)
+# Define Gradio interface
+def generate(prompt, guidance_scale=7.5, num_inference_steps=50):
+    # Generate the image
+    images = model(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images
     return images[0]
 
-# Gradio interface
-interface = gr.Interface(
-    fn=generate_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="Generate images from text prompts using Stable Diffusion."
-)
+# Gradio Interface
+with gr.Blocks() as demo:
+    with gr.Row():
+        prompt = gr.Textbox(label="Enter your prompt")
+        guidance_scale = gr.Slider(1.0, 10.0, value=7.5, label="Guidance Scale")
+        steps = gr.Slider(10, 100, value=50, label="Number of Inference Steps")
+    with gr.Row():
+        submit = gr.Button("Generate")
+    with gr.Row():
+        output = gr.Image()
 
-# Launch the app
-if __name__ == "__main__":
-    interface.launch(server_name="0.0.0.0", server_port=7860)
+    submit.click(generate, inputs=[prompt, guidance_scale, steps], outputs=output)
+
+demo.launch()