Update app.py

app.py

@@ -41,38 +41,28 @@ from PIL import Image
 def image_generator(arg1: str) -> Image.Image:
     """
     A tool that generates an image based on a text prompt.
-    
+
     Args:
         arg1: A description or prompt for the image to be generated.
-    
+
     Returns:
-        A PIL.Image object which Gradio can display.
+        A PIL.Image object which Gradio and final_answer can display properly.
     """
     try:
         result = image_generation_tool(prompt=arg1)
 
-        # If result is a PIL image, return it directly
+        # If result is a PIL image, return a fully loaded version
         if isinstance(result, Image.Image):
-            return result
-
-        # If result is a URL string
-        elif isinstance(result, str) and result.startswith("http"):
-            response = requests.get(result)
-            image = Image.open(BytesIO(response.content))
-            return image
-
-        # If result is a dictionary with an image URL
+            return result.convert("RGB")  # Ensures it's fully loaded and in a common format
+        elif isinstance(result, str):
+            raise ValueError("Expected an image, but got a string URL. Cannot pass this to final_answer.")
         elif isinstance(result, dict) and "image_url" in result:
-            image_url = result["image_url"]
-            response = requests.get(image_url)
-            image = Image.open(BytesIO(response.content))
-            return image
-
-        return Image.new("RGB", (256, 256), color="gray")  # fallback dummy image
-    
+            raise ValueError("Expected an image, but got a dict with 'image_url'. Cannot pass this to final_answer.")
+        else:
+            raise ValueError(f"Unexpected result type: {type(result)}")
     except Exception as e:
-        # Show a red dummy image with error in metadata or raise it
-        return Image.new("RGB", (256, 256), color="red")
+        # Raise errors to let the agent or UI handle them
+        raise RuntimeError(f"Error generating image from prompt '{arg1}': {str(e)}")