Update app.py

app.py

@@ -60,30 +60,8 @@ def generate_image_from_text(prompt: str) -> str:
         # Save the PIL image to file
         pil_image.save(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
+        # Return only the filename - the display formatting is now handled by the final_answer tool
+        return filename
     except Exception as e:
         # Handle any errors that occur during image generation
         return f"Error generating image: {str(e)}"
@@ -111,10 +89,45 @@ def search_web(query: str) -> str:
         return f"Error searching the web: {str(e)}"
 
 #############################################################
-# FINAL ANSWER TOOL (REQUIRED)
+# FINAL ANSWER TOOL (REQUIRED & CUSTOMIZED)
 #############################################################
 # Initialize the final answer tool - required for the agent to provide final answers
-final_answer = FinalAnswerTool()
+final_answer_base = FinalAnswerTool()
+
+@tool
+def final_answer(answer) -> str:
+    """Provides a final answer to the given problem.
+    Args:
+        answer: The final answer to the problem
+    """
+    # Check if the answer is a list of image paths
+    if isinstance(answer, list) and all(isinstance(item, str) and 'uploads/images/' in item for item in answer):
+        # Build HTML to display all images
+        html_output = "<div style='display: flex; flex-direction: column; gap: 20px;'>"
+        
+        # Get the HuggingFace Space name from environment variables
+        space_name = os.environ.get("SPACE_ID", "unknown-space")
+        
+        for i, image_path in enumerate(answer):
+            # Construct a URL that works in Hugging Face Spaces
+            file_url = f"https://huggingface.co/spaces/{space_name}/resolve/main/{image_path}"
+            
+            html_output += f"""
+            <div style="margin-bottom: 20px; border: 1px solid #ddd; border-radius: 8px; padding: 15px; background-color: #f9f9f9;">
+              <h3 style="margin-top: 0;">Image {i+1}</h3>
+              <img src="{image_path}" alt="Generated image {i+1}" style="max-width:100%; height:auto; border-radius:4px;">
+              <div style="margin-top:10px;">
+                <a href="{file_url}" target="_blank" style="display:inline-block; padding:8px 16px; background-color:#4CAF50; color:white; text-decoration:none; border-radius:4px; margin-right:10px;">View Full Size</a>
+                <code style="display:block; margin-top:10px; padding:8px; background-color:#f0f0f0; border-radius:4px; font-size:12px; word-break:break-all;">{file_url}</code>
+              </div>
+            </div>
+            """
+        
+        html_output += "</div>"
+        return final_answer_base(html_output)
+    
+    # Otherwise, just pass the answer to the regular final_answer tool
+    return final_answer_base(answer)
 
 #############################################################
 # VISIT WEBPAGE TOOL - From tools directory