Update model.py

model.py

@@ -2,18 +2,26 @@ from diffusers import StableDiffusionPipeline
 import torch
 
 def load_model(model_name):
-    if model_name == "Stable Diffusion":
-        # Check if GPU is available, else use CPU
-        if torch.cuda.is_available():
-            pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
-            pipe = pipe.to("cuda")  # Use GPU
-        else:
-            pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
-            pipe = pipe.to("cpu")  # Use CPU if no GPU is found
-        return pipe
+    try:
+        if model_name == "Stable Diffusion":
+            # Check if GPU is available, else use CPU
+            if torch.cuda.is_available():
+                pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
+                pipe = pipe.to("cuda")  # Use GPU
+            else:
+                pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
+                pipe = pipe.to("cpu")  # Use CPU if no GPU is found
+            return pipe
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None
 
 def generate_image(model, prompt, model_name):
-    if model_name == "Stable Diffusion":
-        result = model(prompt, num_inference_steps=20, height=512, width=512)  # Reduced resolution for faster generation
-        image = result.images[0]
-        return image
+    try:
+        if model_name == "Stable Diffusion":
+            result = model(prompt, num_inference_steps=10, height=256, width=256)  # Reduced resolution for faster generation
+            image = result.images[0]
+            return image
+    except Exception as e:
+        print(f"Error generating image: {e}")
+        return None