Update app.py

app.py

@@ -1,38 +1,58 @@
 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
-    )
-
-# Initialize the pipeline
-model = load_model()
-
-# 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]
-
-# 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."
-)
-
-# Launch the app
-if __name__ == "__main__":
-    interface.launch(server_name="0.0.0.0", server_port=7860)
+    try:
+        # Load the model with explicit variant for half-precision weights
+        pipeline = StableDiffusionPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-2-1",
+            torch_dtype=torch.float16,
+            variant="fp16",  # Updated from 'revision' to 'variant'
+            safety_checker=None  # Disable safety checker for faster inference
+        )
+    except Exception as e:
+        print(f"Error loading the model: {e}")
+        raise
+
+    # Configure the scheduler for faster generation
+    pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
+
+    # Move to CPU if no GPU is available
+    try:
+        pipeline = pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
+    except Exception as e:
+        print(f"Error moving the model to device: {e}")
+        raise
+
+    return pipeline
+
+# Initialize the model
+try:
+    model = load_model()
+except Exception as e:
+    print(f"Error initializing the model: {e}")
+
+# Define Gradio interface
+def generate(prompt, guidance_scale=7.5, num_inference_steps=50):
+    try:
+        # Generate image from the prompt
+        images = model(prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps).images
+        return images[0]
+    except Exception as e:
+        return f"Error generating image: {e}"
+
+# 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()
+
+    submit.click(generate, inputs=[prompt, guidance_scale, steps], outputs=output)
+
+demo.launch()