Update main.py

main.py

@@ -9,9 +9,25 @@ import string
 import time
 from queue import Queue
 from threading import Thread
+import requests
+import io
+from PIL import Image
 
 app = FastAPI()
 
+API_URL = "https://api-inference.huggingface.co/models/playgroundai/playground-v2-1024px-aesthetic"
+API_TOKEN = os.getenv("HF_READ_TOKEN") # it is free
+headers = {"Authorization": f"Bearer {API_TOKEN}"}
+def generate_image(prompt):
+    payload = {"inputs": prompt}
+    response = requests.post(API_URL, headers=headers, json=payload)
+    image_bytes = response.content
+    image = Image.open(io.BytesIO(image_bytes))
+    return image
+
+
+
+
 text_gen = gr.Interface.load("models/Gustavosta/MagicPrompt-Stable-Diffusion")
 proc1 = gr.Interface.load("models/playgroundai/playground-v2-1024px-aesthetic")
 
@@ -89,7 +105,9 @@ def generate_prompts(prompt_text: str):
 
 @app.get("/send_inputs")
 def send_inputs(inputs: str, noise_level: float):
-    return send_it1(inputs, noise_level)
+    generated_image = generate_image(inputs)
+    return generated_image
+
 
 app.mount("/", StaticFiles(directory="static", html=True), name="static")