Update app.py

app.py

@@ -21,7 +21,7 @@ def generate_image(prompt):
             image = image_model(prompt).images[0]
         return image
     except Exception as e:
-        return f"Error generating image: {str(e)}"
+        return None, f"Error generating image: {str(e)}"
 
 # Face swap function placeholder (replace with actual implementation)
 def swap_faces(image1, image2):
@@ -31,54 +31,62 @@ def swap_faces(image1, image2):
 def face_swap(image1, image2):
     try:
         if image1 is None or image2 is None:
-            return None
+            return None, "Images for face swap are required"
         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)
+        return cv2.cvtColor(swapped_image, cv2.COLOR_BGR2RGB), None
     except Exception as e:
-        return f"Error during face swap: {str(e)}"
+        return None, f"Error during face swap: {str(e)}"
 
 # Upscaling function
 def upscale_image(image, scale_factor=2):
     try:
+        if image is None:
+            return None, "Image for upscaling is required"
         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
+        return upscaled_image, None
     except Exception as e:
-        return f"Error during upscaling: {str(e)}"
+        return None, f"Error during upscaling: {str(e)}"
 
 # Gradio interface function
 def process_image(prompt, image1=None, image2=None, scale_factor=2):
     try:
         if prompt:
             # Generate image from prompt
-            generated_image = generate_image(prompt)
-            return generated_image, None, None
+            generated_image, error = generate_image(prompt)
+            if error:
+                return None, error
+            return generated_image, None
         elif image1 and image2:
             # Perform faceswap if two images are provided
-            swapped_image = face_swap(image1, image2)
-            upscaled_image = upscale_image(swapped_image, scale_factor)
-            return None, swapped_image, upscaled_image
+            swapped_image, error = face_swap(image1, image2)
+            if error:
+                return None, error
+            upscaled_image, error = upscale_image(swapped_image, scale_factor)
+            if error:
+                return None, error
+            return upscaled_image, None
         else:
-            return None, None, None
+            return None, "Either a prompt or two images must be provided"
     except Exception as e:
-        return f"Error in process_image function: {str(e)}", None, None
+        return None, f"Error in process_image function: {str(e)}"
 
 # Gradio interface setup
 iface = gr.Interface(
     fn=process_image,
     inputs=[
         gr.Textbox(label="Enter your prompt", placeholder="Type your prompt here..."),
-        gr.Image(label="Image 1 (for faceswap)", type="pil", optional=True),
-        gr.Image(label="Image 2 (for faceswap)", type="pil", optional=True),
+        gr.Image(label="Image 1 (for faceswap)", type="pil"),
+        gr.Image(label="Image 2 (for faceswap)", type="pil"),
         gr.Slider(label="Upscale Factor", minimum=1, maximum=4, step=1, value=2)
     ],
     outputs=[
         gr.Image(label="Output Image"),
-        gr.Textbox(label="Error Message", placeholder="Error details will appear here...", visible=True)
+        gr.Textbox(label="Error Message", placeholder="Error details will appear here...")
     ],
     title="Fooocus Image Processing",
     description="Generate images from prompts, swap faces, and upscale images."