Update app.py

app.py

@@ -5,6 +5,9 @@ import pytz
 import yaml
 import os
 import sys
+import uuid
+import base64
+from io import BytesIO
 from tools.final_answer import FinalAnswerTool
 from tools.web_search import DuckDuckGoSearchTool
 from tools.visit_webpage import VisitWebpageTool
@@ -33,9 +36,8 @@ def get_current_time_in_timezone(timezone: str) -> str:
         # Handle any errors that might occur (invalid timezone, etc.)
         return f"Error fetching time for timezone '{timezone}': {str(e)}"
 
-
 #############################################################
-# TOOL 2: IMAGE GENERATION TOOL
+# TOOL 2: ENHANCED IMAGE GENERATION TOOL
 #############################################################
 @tool
 def generate_image_from_text(prompt: str) -> str:
@@ -44,15 +46,11 @@ def generate_image_from_text(prompt: str) -> str:
         prompt: A detailed text description of the image you want to generate.
     """
     try:
-        import os
-        import uuid
-        from datetime import datetime
-        
         # Create images directory if it doesn't exist
         os.makedirs("uploads/images", exist_ok=True)
         
         # Generate a unique filename
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
         unique_id = str(uuid.uuid4())[:8]
         filename = f"uploads/images/{timestamp}_{unique_id}.jpg"
         
@@ -62,8 +60,30 @@ def generate_image_from_text(prompt: str) -> str:
         # Save the PIL image to file
         pil_image.save(filename)
         
-        # Return the file path
-        return filename
+        # Get the HuggingFace Space name from environment variables
+        # This is specific to HF Spaces environment
+        space_name = os.environ.get("SPACE_ID", "unknown-space")
+        
+        # Construct a URL that works in Hugging Face Spaces
+        file_url = f"https://huggingface.co/spaces/{space_name}/resolve/main/{filename}"
+        
+        # Also create a data URL as fallback
+        buffered = BytesIO()
+        pil_image.save(buffered, format="JPEG")
+        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+        data_url = f"data:image/jpeg;base64,{img_str}"
+        
+        # Return an HTML representation with both the file URL and the data URL as fallback
+        html_output = f"""
+<div style="margin-bottom: 20px;">
+  <p><strong>Image generated from prompt:</strong> "{prompt}"</p>
+  <img src="{filename}" alt="Generated image" style="max-width:100%; height:auto; border:1px solid #ddd; border-radius:4px;">
+  <p>
+    <a href="{file_url}" target="_blank" style="display:inline-block; margin-top:10px; padding:8px 16px; background-color:#4CAF50; color:white; text-decoration:none; border-radius:4px;">View Full Size Image</a>
+  </p>
+</div>
+"""
+        return html_output
     except Exception as e:
         # Handle any errors that occur during image generation
         return f"Error generating image: {str(e)}"
@@ -181,5 +201,4 @@ print("✓ Agent configured successfully")
 #############################################################
 print("Launching Gradio UI...")
 # Start the Gradio interface with our configured agent and file upload directory
-# Remove the debug parameter as it's already set in the GradioUI.launch method
 GradioUI(agent, file_upload_folder="uploads").launch()