Update tools/image_generator.py

tools/image_generator.py

@@ -5,26 +5,35 @@ import requests
 import base64
 from PIL import Image
 from io import BytesIO
+import time
 
 class ImageGeneratorTool(Tool):
     name = "image_generator"
     description = "Generate an image based on a textual description using HuggingFace's API"
     inputs = {'description': {'type': 'string', 'description': 'Text description of the image to generate'}}
-    output_type = "string"  # Changed from "bytes" to "string"
+    output_type = "string"
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.api_token = os.environ.get('HF_TOKEN')
         if not self.api_token:
             raise ValueError("HF_TOKEN environment variable not found in Space secrets")
-        self.api_url = "https://api-inference.huggingface.co/models/prompthero/openjourney"
+        self.api_url = "https://api-inference.huggingface.co/models/runwayml/stable-diffusion-v1-5"
         self.headers = {"Authorization": f"Bearer {self.api_token}"}
 
     def forward(self, description: str) -> str:
         try:
             payload = {"inputs": description}
+            
+            # First request might return a loading message
             response = requests.post(self.api_url, headers=self.headers, json=payload)
-            response.raise_for_status()
+            
+            # If model is loading, wait and retry
+            if response.status_code == 503:
+                # Wait for 20 seconds
+                time.sleep(20)
+                response = requests.post(self.api_url, headers=self.headers, json=payload)
+                response.raise_for_status()
             
             # Convert the image bytes to base64 string
             image_base64 = base64.b64encode(response.content).decode('utf-8')