Create image_generation.py

application/utils/image_generation.py

@@ -0,0 +1,36 @@
+# application/utils/image_generation.py - NEW FILE
+from huggingface_hub import InferenceClient
+import os
+import base64
+from io import BytesIO
+
+def generate_image(prompt: str) -> str:
+    """
+    Generates an image based on the given prompt using the Hugging Face Inference API.
+
+    Args:
+        prompt: The text prompt to generate the image from.
+
+    Returns:
+        The base64 encoded image data, or None if an error occurred.
+    """
+    try:
+        api_key = os.environ.get('auth')
+        if not api_key:
+            raise ValueError("Hugging Face API key ('auth') not found in environment variables.")
+
+        client = InferenceClient(api_key=api_key)
+        image = client.text_to_image(
+            prompt,
+            model="black-forest-labs/FLUX.1-schnell"
+        )
+
+        # Convert PIL Image to base64 string
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")  # Or "JPEG" if appropriate
+        img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+        return img_str
+
+    except Exception as e:
+        print(f"Error in image generation: {e}")
+        return None