anupamdas's picture
Update tools/image_tool.py
d74de07 verified
Raw
History Blame Contribute Delete
2.53 kB
import os
import tempfile
import requests
from phi.tools import Toolkit
from phi.utils.log import logger
class ImageGenerationTool(Toolkit):
def __init__(self, api_key: str, style_guide: str = "photographic, high detail"):
super().__init__(name="image_generation_tool")
self.api_key = api_key
self.style_guide = style_guide
self.api_url = "https://clipdrop-api.co/text-to-image/v1"
logger.info(f"Clipdrop Stable Diffusion ImageGenerationTool initialized with endpoint: {self.api_url}")
self.register(self.generate_image)
def generate_image(self, prompt: str) -> str:
"""
Generates an image using Stable Diffusion via the Clipdrop API,
strictly following the official documentation.
"""
if not self.api_key:
return "Error: Image Generation Tool not initialized. Check Clipdrop API Key."
# Enhance the prompt with our style guide for better results
full_prompt = f"{prompt}, {self.style_guide}"
logger.info(f"Generating image with Clipdrop prompt: '{full_prompt}'")
payload = {
'prompt': (None, full_prompt, 'text/plain')
}
headers = {
'x-api-key': self.api_key
}
try:
# Send the POST request using `files=` as specified in the docs.
response = requests.post(self.api_url, files=payload, headers=headers)
response.raise_for_status() # This will raise an HTTPError for non-2xx responses
image_bytes = response.content
# --- FIX: Explicitly define the save directory as /tmp ---
save_dir = "/tmp"
# --- END OF FIX ---
filename = f"clipdrop_img_{hash(prompt)}.png"
save_path = os.path.join(save_dir, filename)
with open(save_path, "wb") as f:
f.write(image_bytes)
logger.info(f"Image successfully saved to: {save_path}")
return save_path
except requests.exceptions.HTTPError as e:
# Provide more detailed error information
error_message = f"Clipdrop API error - Status Code: {e.response.status_code}, Response: {e.response.text}"
logger.error(error_message)
return f"Error: {error_message}"
except Exception as e:
logger.error(f"Clipdrop image generation failed with an exception: {e}")
return f"Error: Image generation failed. Details: {e}"