Update app.py

app.py

@@ -1,7 +1,38 @@
 from fastapi import FastAPI
+import requests
+import io
+from PIL import Image
+from fastapi.responses import StreamingResponse
+import os
 
+# Initialize the FastAPI app
 app = FastAPI()
 
-@app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+# Get the Hugging Face API key from environment variables
+API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+
+headers = {"Authorization": f"Bearer {HUGGINGFACE_TOKEN}"}
+
+# Function to query Hugging Face model API
+def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.content
+
+# FastAPI route to return the image
+@app.get("/generate-image")
+def generate_image():
+    # Send a request to the Hugging Face model
+    image_bytes = query({"inputs": "cat with dog"})
+    
+    # Open the image
+    image = Image.open(io.BytesIO(image_bytes))
+    
+    # Create a StreamingResponse to return the image as a response
+    img_byte_arr = io.BytesIO()
+    image.save(img_byte_arr, format='PNG')
+    img_byte_arr.seek(0)
+    
+    return StreamingResponse(img_byte_arr, media_type="image/png")
+
+