Update app.py

app.py

@@ -1,64 +1,57 @@
 import gradio as gr
 import requests
 import base64
-import logging
 from PIL import Image
-from io import BytesIO
+import io
+import json
 
-# Set up logging to display debug information
-logging.basicConfig(level=logging.DEBUG)
+def call_api_and_generate_image(api_key, negative_prompt, positive_prompt, controlnet_image):
+    if not api_key:
+        return None, "Error: API key is required. Please provide a valid API key."
 
-def call_api(negative_prompt, positive_prompt, api_key):
-    # API URL
-    url = "https://model-5qe9kjp3.api.baseten.co/development/predict"
-    
-    # Headers including the provided API Key
-    headers = {"Authorization": f"Api-Key {api_key}"}
-    
-    # JSON payload with the prompts
-    payload = {
-        'workflow_values': {
-            'negative_prompt': negative_prompt,
-            'positive_prompt': positive_prompt
+    resp = requests.post(
+        "https://model-4w561emq.api.baseten.co/development/predict",
+        headers={"Authorization": f"Api-Key {api_key}"},
+        json={
+            'workflow_values': {
+                'negative_prompt': negative_prompt,
+                'positive_prompt': positive_prompt,
+                'controlnet_image': controlnet_image
+            }
         }
-    }
-    
-    # POST request to the API
-    response = requests.post(url, headers=headers, json=payload)
-    
-    # Check if the request was successful
-    if response.status_code == 200:
-        data = response.json()['result'][0]['data']
-        
-        # Decode the base64 string to bytes
-        image_data = base64.b64decode(data)
-        
-        # Convert bytes data to an image
-        image = Image.open(BytesIO(image_data))
-        
-        # Save the image as PNG to a buffer
-        buffer = BytesIO()
-        image.save(buffer, format="PNG")
-        buffer.seek(0)
-        
-        # Return the buffer content to display it directly in the interface
-        return buffer
-    else:
-        return f"Failed to fetch the image. Status Code: {response.status_code}, Error: {response.text}"
+    )
 
-# Define Gradio interface
-interface = gr.Interface(
-    fn=call_api,
+    if resp.status_code != 200:
+        return None, f"Error: API call failed with status code {resp.status_code}, message: {resp.text}"
+
+    result = resp.json()
+    base64_data = result['result'][0]['data']
+    image_data = base64.b64decode(base64_data)
+    image = Image.open(io.BytesIO(image_data))
+    return image, json.dumps(result, indent=2)
+
+def generate_image(api_key, negative_prompt, positive_prompt, controlnet_image):
+    try:
+        return call_api_and_generate_image(api_key, negative_prompt, positive_prompt, controlnet_image)
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=generate_image,
     inputs=[
-        gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt here..."),
-        gr.Textbox(label="Positive Prompt", placeholder="Enter positive prompt here..."),
-        gr.Textbox(label="API Key", placeholder="Enter your API key here..."),
+        gr.Textbox(label="API Key", placeholder="Enter your API key here"),
+        gr.Textbox(label="Negative Prompt", value="blurry, text, low quality", placeholder="Enter negative prompt here..."),
+        gr.Textbox(label="Positive Prompt", value="An igloo on a snowy day, 4k, hd", placeholder="Enter positive prompt here..."),
+        gr.Textbox(label="ControlNet Image URL", value="https://storage.googleapis.com/logos-bucket-01/baseten_logo.png", placeholder="Enter ControlNet image URL here...")
     ],
-    outputs="image",
-    title="Image Generation API Interface",
-    description="Enter the negative and positive prompts to generate an image."
+    outputs=[
+        gr.Image(type="pil", label="Generated Image"),
+        gr.Textbox(label="API Response JSON", lines=10)
+    ],
+    title="Image Generation with API",
+    description="Enter your API key and prompts to generate an image using the provided API."
 )
 
 # Launch the interface
-interface.launch()
-
+iface.launch()