Update app.py

app.py

@@ -57,10 +57,15 @@ def query_api(payload, progress_callback=None):
             
             time.sleep(5)  # Wait 5 seconds between polls
             
-            # Check status
-            status_response = requests.get(status_url, headers=headers)
+            # Check status - try without auth headers first, then with auth headers
+            status_response = requests.get(status_url)
+            
+            # If unauthorized, try with headers
+            if status_response.status_code == 401:
+                status_response = requests.get(status_url, headers=headers)
             
             if status_response.status_code != 200:
+                print(f"Status response: {status_response.status_code} - {status_response.text}")
                 raise gr.Error(f"Status check failed: {status_response.status_code}")
             
             try:
@@ -73,8 +78,13 @@ def query_api(payload, progress_callback=None):
                     if not response_url:
                         raise gr.Error("No response URL provided")
                     
-                    result_response = requests.get(response_url, headers=headers)
+                    # Try to get result without auth first, then with auth
+                    result_response = requests.get(response_url)
+                    if result_response.status_code == 401:
+                        result_response = requests.get(response_url, headers=headers)
+                    
                     if result_response.status_code != 200:
+                        print(f"Result response: {result_response.status_code} - {result_response.text}")
                         raise gr.Error(f"Failed to get result: {result_response.status_code}")
                     
                     # Check if result is JSON with image data
@@ -191,25 +201,22 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
         progress(0.1, desc="Generating image...")
 
     try:
-        # Make API request
-        image_bytes = query_api(payload)
+        # Make API request with progress callback
+        image_bytes = query_api(payload, progress_callback=progress)
         
-        # Try to convert response bytes to PIL Image with better error handling
+        # Try to convert response bytes to PIL Image
         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")
+            print(f"Failed to open image: {img_error}")
+            print(f"Image bytes type: {type(image_bytes)}, length: {len(image_bytes) if hasattr(image_bytes, '__len__') else 'unknown'}")
             
             # 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'}")
+                raise gr.Error(f"Could not process API response as image. Response length: {len(image_bytes) if hasattr(image_bytes, '__len__') else 'unknown'}")
         
         progress(1.0, desc="Complete!")
         return gr.Image(value=image)