Update app.py

app.py

@@ -43,24 +43,18 @@ def generate_image(prompt, negative_prompt, sampler, seed, guidance_scale, infer
     # Print the API response for debugging
     print("API Response:", response_body)
 
-    if response_body.get("b64_json") is None:
+    # Decode the base64-encoded image data
+    b64_image_data = response_body.get("b64_json")
+    if b64_image_data is None:
         return "Error: API response does not contain 'b64_json' key."
+    
+    image_data = base64.b64decode(b64_image_data)
+    
+    # Since Gradio's Image component can accept a PIL Image, convert the binary data to a PIL Image
+    image = Image.open(BytesIO(image_data))
+    
+    return image
 
-    # Decode the base64-encoded JSON string
-    b64_json_string = response_body.get("b64_json")
-    decoded_json_bytes = base64.b64decode(b64_json_string)
-    decoded_json_string = decoded_json_bytes.decode("utf-8")
-
-    # Load the decoded JSON string into a Python dictionary
-    output_dict = json.loads(decoded_json_string)
-
-    # Extract the image URL from the output dictionary
-    image_url = output_dict.get("output").get("image_url")
-
-    if image_url is None:
-        return "Error: API response does not contain 'image_url' key."
-
-    return image_url
 
 # Create Gradio interface
 iface = gr.Interface(
@@ -73,7 +67,7 @@ iface = gr.Interface(
         gr.Slider(label="Guidance Scale", minimum=0, maximum=20, value=5),
         gr.Slider(label="Inference Steps", minimum=1, maximum=50, value=25)
     ],
-    outputs=gr.Image(label="Generated Image")
+    outputs=gr.Image(label="Generated Image"),
 )
 
 # Launch the Gradio app