Update app.py

app.py

@@ -4,6 +4,7 @@ import random
 import os
 import base64
 import requests
+import time
 import io
 from PIL import Image, ImageOps
 import pillow_heif  # For HEIF/AVIF support
@@ -23,57 +24,115 @@ def get_headers():
         "X-HF-Bill-To": "huggingface"
     }
 
-def query_api(payload):
+def query_api(payload, progress_callback=None):
     """Send request to the API and return response"""
     headers = get_headers()
+    
+    # Submit the job
     response = requests.post(API_URL, headers=headers, json=payload)
     
     if response.status_code != 200:
         raise gr.Error(f"API request failed with status {response.status_code}: {response.text}")
     
-    # 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()
+    # Parse the initial response
+    try:
+        json_response = response.json()
+        print(f"Initial response: {json_response}")
+    except:
+        raise gr.Error("Failed to parse initial API response as JSON")
     
-    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 job was queued
+    if json_response.get("status") == "IN_QUEUE":
+        status_url = json_response.get("status_url")
+        if not status_url:
+            raise gr.Error("No status URL provided in queue response")
+        
+        # Poll for completion
+        max_attempts = 60  # Wait up to 5 minutes (60 * 5 seconds)
+        attempt = 0
+        
+        while attempt < max_attempts:
+            if progress_callback:
+                progress_callback(0.1 + (attempt / max_attempts) * 0.8, f"Processing... (attempt {attempt + 1}/60)")
             
-            # 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)}")
+            time.sleep(5)  # Wait 5 seconds between polls
+            
+            # Check status
+            status_response = requests.get(status_url, headers=headers)
+            
+            if status_response.status_code != 200:
+                raise gr.Error(f"Status check failed: {status_response.status_code}")
+            
+            try:
+                status_data = status_response.json()
+                print(f"Status check {attempt + 1}: {status_data}")
+                
+                if status_data.get("status") == "COMPLETED":
+                    # Job completed, get the result
+                    response_url = json_response.get("response_url")
+                    if not response_url:
+                        raise gr.Error("No response URL provided")
+                    
+                    result_response = requests.get(response_url, headers=headers)
+                    if result_response.status_code != 200:
+                        raise gr.Error(f"Failed to get result: {result_response.status_code}")
+                    
+                    # Check if result is JSON with image data
+                    try:
+                        result_data = result_response.json()
+                        print(f"Result data: {result_data}")
+                        
+                        # Look for image in various possible fields
+                        if 'images' in result_data and len(result_data['images']) > 0:
+                            # Images array with URLs or base64
+                            image_data = result_data['images'][0]
+                            if isinstance(image_data, dict) and 'url' in image_data:
+                                # Image URL - fetch it
+                                img_response = requests.get(image_data['url'])
+                                return img_response.content
+                            elif isinstance(image_data, str):
+                                # Assume base64
+                                return base64.b64decode(image_data)
+                        elif 'image' in result_data:
+                            # Single image field
+                            if isinstance(result_data['image'], str):
+                                return base64.b64decode(result_data['image'])
+                        elif 'url' in result_data:
+                            # Direct URL
+                            img_response = requests.get(result_data['url'])
+                            return img_response.content
+                        else:
+                            raise gr.Error(f"No image found in result: {result_data}")
+                            
+                    except requests.exceptions.JSONDecodeError:
+                        # Result might be direct image bytes
+                        return result_response.content
+                        
+                elif status_data.get("status") == "FAILED":
+                    error_msg = status_data.get("error", "Unknown error")
+                    raise gr.Error(f"Job failed: {error_msg}")
+                
+                # Still processing, continue polling
+                attempt += 1
+                
+            except requests.exceptions.JSONDecodeError:
+                raise gr.Error("Failed to parse status response")
+    
+        raise gr.Error("Job timed out after 5 minutes")
     
-    elif 'image/' in content_type:
-        # Response is direct image bytes
-        return response.content
+    elif json_response.get("status") == "COMPLETED":
+        # Job completed immediately
+        if 'images' in json_response and len(json_response['images']) > 0:
+            image_data = json_response['images'][0]
+            if isinstance(image_data, str):
+                return base64.b64decode(image_data)
+        elif 'image' in json_response:
+            return base64.b64decode(json_response['image'])
+        else:
+            raise gr.Error(f"No image found in immediate response: {json_response}")
     
     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
+        raise gr.Error(f"Unexpected response status: {json_response.get('status', 'unknown')}")
 
 # --- Core Inference Function for ChatInterface ---
 def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress()):