Update app.py

app.py

@@ -3,33 +3,44 @@ from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 import torch
 
 def load_model():
-    # 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)
+    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 pipeline to GPU or ZeroGPU
-    pipeline = pipeline.to("cuda")  # or ZeroGPU-specific setup
-    
+
+    # 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 loading the model: {e}")
+    print(f"Error initializing the model: {e}")
 
 # 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]
+    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: