Update app.py

app.py

@@ -1,4 +1,4 @@
-from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
+from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool, Tool
 import datetime
 import requests
 import pytz
@@ -9,6 +9,8 @@ from Gradio_UI import GradioUI
 import fitz  # PyMuPDF
 from sentence_transformers import SentenceTransformer, util
 from transformers import pipeline
+from PIL import Image
+import io
 
 # API Key for weather
 API_KEY = os.getenv("Weather_Token")
@@ -135,14 +137,42 @@ def document_qna_tool(pdf_path: str, question: str) -> str:
         return f"[EXCEPTION] {type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
 
 # -------------------- TOOL 4: Image Generation --------------------
+@tool
+def generate_image(prompt: str) -> Image.Image:
+    """
+    A tool that generates an image from a text prompt using a Hugging Face Space.
+    Args:
+        prompt (str): The text description of the image to generate.
+    Returns:
+        Image.Image: The generated image as a PIL Image object.
+    """
+    try:
+        # Use the hosted image generation model from Hugging Face Spaces
+        image_generator = Tool.from_space(
+            "black-forest-labs/FLUX.1-schnell",
+            name="image_generator",
+            description="Generate an image from a prompt"
+        )
 
-from smolagents import Tool
+        # Call the model with the prompt
+        result = image_generator(prompt=prompt)
 
-image_generation_tool = Tool.from_space(
-    "black-forest-labs/FLUX.1-schnell",
-    name="image_generator",  # You can name it whatever makes sense for your agent
-    description="Generate an image from a prompt"
-)
+        # If the result is bytes → convert to PIL.Image
+        if isinstance(result, bytes):
+            return Image.open(io.BytesIO(result))
+
+        # If the result is already a PIL Image
+        if isinstance(result, Image.Image):
+            return result
+
+        # If the model gave back a path, open it
+        if isinstance(result, str):
+            return Image.open(result)
+
+        raise ValueError("Unexpected output format from image generator.")
+
+    except Exception as e:
+        raise RuntimeError(f"Error generating image: {str(e)}")
 
 # -------------------- Other Components --------------------
 final_answer = FinalAnswerTool()