Create image_modals.py

image_modals.py

@@ -0,0 +1,30 @@
+import requests
+import io
+from PIL import Image
+
+
+headers = {"Authorization": f"Bearer {os.getenv('ImagiGen_HF_secret')}"}
+def query(API_URL, payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.content
+
+def flux(prompt):
+    image = ""
+    API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
+    image_bytes = query(
+        API_URL,
+        {
+            "inputs": prompt,
+        },
+    )
+    return image_bytes
+
+def stable_diffusion(prompt):
+    API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3-medium-diffusers"
+    image_bytes = query(
+        API_URL,
+        {
+            "inputs": prompt,
+        },
+    )
+    return image_bytes