Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import gradio as gr
 import requests
+import base64
 
 # Get API key from environment variable
 api_key = os.environ.get("NVCF_API_KEY")
@@ -35,15 +36,30 @@ def generate_image(prompt, negative_prompt, sampler, seed, guidance_scale, infer
         fetch_url = fetch_url_format + request_id
         response = session.get(fetch_url, headers=headers)
 
+
     response.raise_for_status()
     response_body = response.json()
+
     # Print the API response for debugging
     print("API Response:", response_body)
 
-    if response_body.get("output") is None:
-        return "Error: API response does not contain 'output' key."
-    # Extract image URL from response
-    image_url = response_body.get("output").get("image_url")
+    if response_body.get("b64_json") is None:
+        return "Error: API response does not contain 'b64_json' key."
+
+    # 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