Update app.py

app.py

@@ -96,7 +96,63 @@ temperature=0.5,
 model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',# it is possible that this model may be overloaded
 custom_role_conversions=None,
 )
-
+@tool
+def generate_image_from_prompt(prompt: str, negative_prompt: str = "", height: int = 512, width: int = 512) -> str:
+    """Generate an image based on a text description using AI.
+    Args:
+        prompt: Detailed description of the image you want to generate
+        negative_prompt: Elements to avoid in the generated image (optional)
+        height: Image height in pixels (default: 512)
+        width: Image width in pixels (default: 512)
+    """
+    try:
+        import os
+        import requests
+        import io
+        import base64
+        from PIL import Image
+        from datetime import datetime
+        
+        API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
+        # You need to set your Hugging Face API token 
+        # Get one from: https://huggingface.co/settings/tokens
+        headers = {
+            "Authorization": f"Bearer {os.environ.get('HF_API_TOKEN', 'YOUR_HF_TOKEN')}"
+        }
+        
+        # Prepare the payload
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "negative_prompt": negative_prompt,
+                "height": height,
+                "width": width,
+                "num_inference_steps": 50
+            }
+        }
+        
+        # Make the API call
+        response = requests.post(API_URL, headers=headers, json=payload)
+        
+        if response.status_code != 200:
+            return f"Error generating image: {response.text}"
+        
+        # Create output directory if it doesn't exist
+        output_dir = "generated_images"
+        os.makedirs(output_dir, exist_ok=True)
+        
+        # Save the image
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        image_name = f"{output_dir}/image_{timestamp}.png"
+        
+        # Convert response to image
+        image = Image.open(io.BytesIO(response.content))
+        image.save(image_name)
+        
+        return f"Image generated successfully and saved as {image_name}"
+    
+    except Exception as e:
+        return f"Image generation failed: {str(e)}"
 #@tool 
 #def image_generation_tool(promt: str)->any:
 #    """A tool that generate image from user`s promt.
@@ -117,7 +173,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer,enhanced_translate,image_generation_tool,search_tool], ## add your tools here (don't remove final answer)
+    tools=[final_answer,enhanced_translate,image_generation_tool,search_tool,generate_image_from_prompt, ], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,