Update app.py

app.py

@@ -97,72 +97,30 @@ model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',#
 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.
+def image_generator(prompt: str) -> bytes:
+    """Generate an image based on a text description.
     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)
+    Returns:
+        The image data as bytes that can be displayed
     """
-    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
-       
-        
-        # 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, 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}"
+    import requests
+    
+    API_URL = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
+    
+    # Make API request without authentication headers
+    response = requests.post(
+        API_URL,
+        json={"inputs": prompt}
+    )
+    
+    # Check if the request was successful
+    if response.status_code != 200:
+        raise Exception(f"Error generating image: {response.text}")
+    
+    # Return the raw image bytes
+    return response.content
     
-    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.
-#    Args:
-#        promt: description of image
-#    """
-#   try:
-#        image = image_generation_tool(promt)
-#        return image
-        
-#    except Exception as e:
-#        return f'We have some problem {e}'
 # Import tool from Hub
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
 search_tool = DuckDuckGoSearchTool()
@@ -171,7 +129,7 @@ with open("prompts.yaml", 'r') as stream:
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer,enhanced_translate,image_generation_tool,search_tool,generate_image_from_prompt, ], ## add your tools here (don't remove final answer)
+    tools=[final_answer,enhanced_translate,image_generator,search_tool ], ## add your tools here (don't remove final answer)
     max_steps=6,
     verbosity_level=1,
     grammar=None,