Update app.py

app.py

@@ -31,7 +31,49 @@ def query_api(payload):
     if response.status_code != 200:
         raise gr.Error(f"API request failed with status {response.status_code}: {response.text}")
     
-    return response.content
+    # Debug: Check response content type and first few bytes
+    print(f"Response status: {response.status_code}")
+    print(f"Response headers: {dict(response.headers)}")
+    print(f"Response content type: {response.headers.get('content-type', 'unknown')}")
+    print(f"Response content length: {len(response.content)}")
+    print(f"First 200 chars of response: {response.content[:200]}")
+    
+    # Check if response is JSON (error case) or binary (image case)
+    content_type = response.headers.get('content-type', '').lower()
+    
+    if 'application/json' in content_type:
+        # Response is JSON, might contain base64 image or error
+        try:
+            json_response = response.json()
+            print(f"JSON response: {json_response}")
+            
+            # Check if there's a base64 image in the response
+            if 'image' in json_response:
+                # Decode base64 image
+                image_data = base64.b64decode(json_response['image'])
+                return image_data
+            elif 'images' in json_response and len(json_response['images']) > 0:
+                # Multiple images, take the first one
+                image_data = base64.b64decode(json_response['images'][0])
+                return image_data
+            else:
+                raise gr.Error(f"Unexpected JSON response format: {json_response}")
+        except Exception as e:
+            raise gr.Error(f"Failed to parse JSON response: {str(e)}")
+    
+    elif 'image/' in content_type:
+        # Response is direct image bytes
+        return response.content
+    
+    else:
+        # Try to decode as base64 first, then as direct bytes
+        try:
+            # Maybe the entire response is base64 encoded
+            image_data = base64.b64decode(response.content)
+            return image_data
+        except:
+            # Return as-is and let PIL try to handle it
+            return response.content
 
 # --- Core Inference Function for ChatInterface ---
 def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress()):
@@ -93,12 +135,29 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
         # Make API request
         image_bytes = query_api(payload)
         
-        # Convert response bytes to PIL Image
-        image = Image.open(io.BytesIO(image_bytes))
+        # Try to convert response bytes to PIL Image with better error handling
+        try:
+            image = Image.open(io.BytesIO(image_bytes))
+        except Exception as img_error:
+            print(f"Failed to open image directly: {img_error}")
+            # Maybe it's a different format, try to save and examine
+            with open('/tmp/debug_response.bin', 'wb') as f:
+                f.write(image_bytes)
+            print(f"Saved response to /tmp/debug_response.bin for debugging")
+            
+            # Try to decode as base64 if direct opening failed
+            try:
+                decoded_bytes = base64.b64decode(image_bytes)
+                image = Image.open(io.BytesIO(decoded_bytes))
+            except:
+                raise gr.Error(f"Could not process API response as image. Response type: {type(image_bytes)}, Length: {len(image_bytes) if isinstance(image_bytes, (bytes, str)) else 'unknown'}")
         
         progress(1.0, desc="Complete!")
         return gr.Image(value=image)
         
+    except gr.Error:
+        # Re-raise gradio errors as-is
+        raise
     except Exception as e:
         raise gr.Error(f"Failed to generate image: {str(e)}")