Update app.py

app.py

@@ -1,44 +1,48 @@
 import gradio as gr
 import requests
-import json
+from PIL import Image
+from io import BytesIO
 import os
 
-# Hugging Face Inference API setup
-API_URL = "https://api-inference.huggingface.co/models/Benevolent/PonyDiffusionV10"
-HF_API_TOKEN = os.getenv("HF_API_TOKEN")  # Get token from environment variable
-headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
 
-def query(payload):
-    response = requests.post(API_URL, headers=headers, json=payload)
-    return json.loads(response.content)
 
-# Function to call the HF Inference API
-def generate_image(prompt):
-    payload = {
-        "inputs": prompt,
-    }
-    response = query(payload)
-    if 'error' in response:
-        return f"Error: {response['error']}"
-    # Check for the generated image
-    if isinstance(response, list) and len(response) > 0 and "generated_image" in response[0]:
-        return response[0]["generated_image"]  # Base64 encoded image
-    return "No image returned from model"
+# Load API Token from environment variable
+API_TOKEN = os.getenv("HF_API_TOKEN")  # Ensure you've set this environment variable
 
-# Gradio Blocks Web UI
-with gr.Blocks() as demo:
-    with gr.Row():
-        gr.Markdown("# PonyDiffusion V10 Text-to-Image Generator")
-    
-    with gr.Row():
-        prompt = gr.Textbox(label="Enter a prompt for the image")
+# Hugging Face Inference API URL
+API_URL = "https://api-inference.huggingface.co/models/Benevolent/PonyDiffusionV10"
 
-    with gr.Row():
-        generate_button = gr.Button("Generate Image")
-        output_image = gr.Image()
+# Function to call Hugging Face API and get the generated image
+def generate_image(prompt):
+    headers = {"Authorization": f"Bearer {API_TOKEN}"}
+    data = {"inputs": prompt}
+    
+    response = requests.post(API_URL, headers=headers, json=data)
+    
+    if response.status_code == 200:
+        image_bytes = BytesIO(response.content)
+        image = Image.open(image_bytes)
+        return image
+    else:
+        return f"Error: {response.status_code}, {response.text}"
 
-    # When generate_button is clicked, call generate_image
-    generate_button.click(fn=generate_image, inputs=prompt, outputs=output_image)
+# Create Gradio interface
+def create_ui():
+    with gr.Blocks() as ui:
+        gr.Markdown("## PonyDiffusionV10 - Text to Image Generator")
+        
+        with gr.Row():
+            prompt_input = gr.Textbox(label="Enter a Prompt", placeholder="Describe the image you want to generate", lines=3)
+            generate_button = gr.Button("Generate Image")
+        
+        with gr.Row():
+            output_image = gr.Image(label="Generated Image")
+        
+        # Link the button to the function
+        generate_button.click(fn=generate_image, inputs=prompt_input, outputs=output_image)
+    
+    return ui
 
-# Launch the app
-demo.launch()
+# Run the interface
+if __name__ == "__main__":
+    create_ui().launch()