Update app.py

app.py

@@ -1,8 +1,11 @@
 import gradio as gr
 import requests
+import base64
+from PIL import Image
+from io import BytesIO
 
-def query_huggingface_api(api_key, model, prompt):
-    API_URL = f"https://api-inference.huggingface.co/models/{model}"
+def query_hf_image_generation(api_key, prompt):
+    API_URL = f"https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
     headers = {
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json"
@@ -13,28 +16,25 @@ def query_huggingface_api(api_key, model, prompt):
 
     response = requests.post(API_URL, headers=headers, json=data)
     result = response.json()
+    
+    # Check if the API response contains an error.
+    if 'error' in result:
+        return f"Error: {result['error']}", None
 
-    # Check for errors or handle the response appropriately.
-    if isinstance(result, dict) and 'error' in result:
-        return f"Error: {result['error']}"
-
-    return result['generated_text'] if isinstance(result, dict) and 'generated_text' in result else result
-
-# List of some available models
-model_list = [
-    "stabilityai/stable-diffusion-xl-base-1.0",
-]
+    # Assuming the API returns an image in base64 format.
+    image_data = result['data'][0]  # You might need to adjust this path according to the actual API response
+    image = Image.open(BytesIO(base64.b64decode(image_data)))
+    return image
 
 iface = gr.Interface(
-    fn=query_huggingface_api,
+    fn=query_hf_image_generation,
     inputs=[
         gr.Textbox(label="Hugging Face API Key", placeholder="Enter your Hugging Face API Key here..."),
-        gr.Dropdown(choices=model_list, label="Model"),
-        gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt")
+        gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt"),
     ],
-    outputs="text",
-    title="Hugging Face Model API Explorer",
-    description="Enter your API Key, select a model, and enter a prompt to generate text using the Hugging Face Inference API."
+    outputs=gr.outputs.Image(label="Generated Image"),
+    title="Stable Diffusion XL Image Generator",
+    description="Enter your API Key and a prompt to generate an image using the Stable Diffusion XL model from Hugging Face."
 )
 
 iface.launch()