Update app.py

app.py

@@ -16,9 +16,12 @@ def load_image_model():
 image_model = load_image_model()
 
 def generate_image(prompt):
-    with torch.no_grad():
-        image = image_model(prompt).images[0]
-    return image
+    try:
+        with torch.no_grad():
+            image = image_model(prompt).images[0]
+        return image
+    except Exception as e:
+        return f"Error generating image: {str(e)}"
 
 # Face swap function placeholder (replace with actual implementation)
 def swap_faces(image1, image2):
@@ -26,37 +29,50 @@ def swap_faces(image1, image2):
     return image1
 
 def face_swap(image1, image2):
-    if image1 is None or image2 is None:
-        return None
-    image1 = cv2.cvtColor(np.array(image1), cv2.COLOR_RGB2BGR)
-    image2 = cv2.cvtColor(np.array(image2), cv2.COLOR_RGB2BGR)
-    swapped_image = swap_faces(image1, image2)
-    return cv2.cvtColor(swapped_image, cv2.COLOR_BGR2RGB)
+    try:
+        if image1 is None or image2 is None:
+            return None
+        image1 = cv2.cvtColor(np.array(image1), cv2.COLOR_RGB2BGR)
+        image2 = cv2.cvtColor(np.array(image2), cv2.COLOR_RGB2BGR)
+        swapped_image = swap_faces(image1, image2)
+        return cv2.cvtColor(swapped_image, cv2.COLOR_BGR2RGB)
+    except Exception as e:
+        return f"Error during face swap: {str(e)}"
 
 # Upscaling function
 def upscale_image(image, scale_factor=2):
-    image = Image.fromarray(np.array(image))
-    width, height = image.size
-    new_size = (int(width * scale_factor), int(height * scale_factor))
-    upscaled_image = image.resize(new_size, Image.LANCZOS)
-    return upscaled_image
+    try:
+        image = Image.fromarray(np.array(image))
+        width, height = image.size
+        new_size = (int(width * scale_factor), int(height * scale_factor))
+        upscaled_image = image.resize(new_size, Image.LANCZOS)
+        return upscaled_image
+    except Exception as e:
+        return f"Error during upscaling: {str(e)}"
 
 # Gradio interface function
 def process_image(prompt, image1=None, image2=None, scale_factor=2):
-    # Generate image from prompt
-    generated_image = generate_image(prompt)
-    
-    # Perform faceswap if two images are provided
-    swapped_image = face_swap(image1, image2)
-    
-    # Perform upscaling if scale factor is specified
-    upscaled_image = upscale_image(generated_image, scale_factor)
+    try:
+        # Generate image from prompt
+        generated_image = generate_image(prompt)
+        
+        # Perform faceswap if two images are provided
+        swapped_image = face_swap(image1, image2)
+        
+        # Perform upscaling if scale factor is specified
+        upscaled_image = upscale_image(generated_image, scale_factor)
 
-    return {
-        "Generated Image": generated_image,
-        "Swapped Image": swapped_image,
-        "Upscaled Image": upscaled_image
-    }
+        return {
+            "Generated Image": generated_image,
+            "Swapped Image": swapped_image,
+            "Upscaled Image": upscaled_image
+        }
+    except Exception as e:
+        return {
+            "Generated Image": f"Error in process_image function: {str(e)}",
+            "Swapped Image": None,
+            "Upscaled Image": None
+        }
 
 # Gradio interface setup
 iface = gr.Interface(