Update ImageGenerator.py

ImageGenerator.py

@@ -18,17 +18,72 @@ import requests
 from PIL import Image, UnidentifiedImageError, ImageDraw, ImageFont
 from io import BytesIO
 
+# def generate_image(prompt):
+#     """
+#     Generate an image from a prompt via an API.
+#     Returns a PIL.Image object. If the API fails, returns a placeholder image.
+#     """
+#     try:
+#         api_url="https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5"
+#         headers = {"Authorization": f"Bearer {HF_API_KEY}"}
+#         response = requests.post(api_url, headers=headers, json={"inputs": prompt.split('-')[-1]})
+
+#         # Check if the response is OK
+#         if response.status_code == 200:
+#             try:
+#                 img = Image.open(BytesIO(response.content))
+#                 return img
+#             except UnidentifiedImageError:
+#                 print("API response is not a valid image.")
+#         else:
+#             print(f"API returned status {response.status_code}")
+#             print("Response content:", response.text[:200])
+
+#     except requests.RequestException as e:
+#         print(f"Request failed: {e}")
+
+#     # Fallback placeholder image
+#     print("Using placeholder image.")
+#     placeholder = Image.new("RGB", (512, 512), color=(200, 200, 200))
+#     draw = ImageDraw.Draw(placeholder)
+#     text = "Image\nNot\nAvailable"
+    
+#     # Optional: Use a basic font
+#     try:
+#         font = ImageFont.load_default()
+#         w, h = draw.multiline_textsize(text, font=font)
+#         draw.multiline_text(
+#             ((512 - w) / 2, (512 - h) / 2), text, fill=(50, 50, 50), font=font, align="center"
+#         )
+#     except Exception:
+#         draw.text((50, 200), text, fill=(50, 50, 50))
+    
+#     return placeholder
+
+import os
+import requests
+from PIL import Image, UnidentifiedImageError, ImageDraw, ImageFont
+from io import BytesIO
+
+HF_API_KEY = os.getenv("HF_API_KEY")
+
 def generate_image(prompt):
     """
-    Generate an image from a prompt via an API.
+    Generate an image from the part of the prompt after the dash (-) via Hugging Face API.
     Returns a PIL.Image object. If the API fails, returns a placeholder image.
     """
+    api_url = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-v1-5"
+    headers = {"Authorization": f"Bearer {HF_API_KEY}"}
+
+    # Get the part after the dash
+    if "-" in prompt:
+        prompt_to_send = prompt.split("-", 1)[-1].strip()
+    else:
+        prompt_to_send = prompt.strip()
+
     try:
-        api_url="https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5"
-        headers = {"Authorization": f"Bearer {HF_API_KEY}"}
-        response = requests.post(api_url, headers=headers, json={"inputs": prompt.split('-')[-1]})
+        response = requests.post(api_url, headers=headers, json={"inputs": prompt_to_send})
 
-        # Check if the response is OK
         if response.status_code == 200:
             try:
                 img = Image.open(BytesIO(response.content))
@@ -48,7 +103,6 @@ def generate_image(prompt):
     draw = ImageDraw.Draw(placeholder)
     text = "Image\nNot\nAvailable"
     
-    # Optional: Use a basic font
     try:
         font = ImageFont.load_default()
         w, h = draw.multiline_textsize(text, font=font)
@@ -59,3 +113,4 @@ def generate_image(prompt):
         draw.text((50, 200), text, fill=(50, 50, 50))
     
     return placeholder
+