Update app.py

app.py

@@ -34,6 +34,9 @@ def query_api(payload, progress_callback=None):
         del payload["image_bytes"]
     
     # Submit the job
+    if progress_callback:
+        progress_callback(0.1, "Submitting request...")
+    
     response = requests.post(API_URL, headers=headers, json=payload)
     
     if response.status_code != 200:
@@ -41,10 +44,8 @@ def query_api(payload, progress_callback=None):
     
     # Debug the response
     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 500 chars of response: {response.content[:500]}")
     
     # Check if response is JSON (queue status) or binary (direct image)
     content_type = response.headers.get('content-type', '').lower()
@@ -57,103 +58,36 @@ def query_api(payload, progress_callback=None):
             
             # Check if job was queued
             if json_response.get("status") == "IN_QUEUE":
-                request_id = json_response.get("request_id")
-                if not request_id:
-                    raise gr.Error("No request_id provided in queue response")
+                # For HF router, we should wait and let it handle the queue
+                # The router should eventually return the result automatically
+                if progress_callback:
+                    progress_callback(0.5, "Processing in queue...")
                 
-                # Poll for completion using the proper HF router endpoints
-                max_attempts = 60  # Wait up to 5 minutes
-                attempt = 0
+                # Wait a bit and try to get the response again
+                import time
+                time.sleep(10)  # Wait 10 seconds for processing
                 
-                while attempt < max_attempts:
-                    if progress_callback:
-                        progress_callback(0.1 + (attempt / max_attempts) * 0.8, f"Processing... (attempt {attempt + 1}/60)")
-                    
-                    time.sleep(5)  # Wait 5 seconds between polls
-                    
-                    # Check status using HF router format
-                    status_url = f"https://router.huggingface.co/fal-ai/fal-ai/flux-kontext/dev/requests/{request_id}/status"
-                    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}")
-                        # Continue polling even if status check fails temporarily
-                        attempt += 1
-                        continue
-                    
-                    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
-                            result_url = f"https://router.huggingface.co/fal-ai/fal-ai/flux-kontext/dev/requests/{request_id}"
-                            result_response = requests.get(result_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 direct image bytes or JSON
-                            result_content_type = result_response.headers.get('content-type', '').lower()
-                            if 'image/' in result_content_type:
-                                # Direct image bytes
-                                return result_response.content
-                            else:
-                                # Try to parse as JSON for image URL or base64
-                                try:
-                                    result_data = result_response.json()
-                                    print(f"Result data: {result_data}")
-                                    
-                                    # Look for images in various formats
-                                    if 'images' in result_data and len(result_data['images']) > 0:
-                                        image_info = result_data['images'][0]
-                                        if isinstance(image_info, dict) and 'url' in image_info:
-                                            # Download the image
-                                            img_response = requests.get(image_info['url'])
-                                            return img_response.content
-                                        elif isinstance(image_info, str):
-                                            # Base64 encoded
-                                            return base64.b64decode(image_info)
-                                    elif 'image' in result_data:
-                                        # Single image field
-                                        if isinstance(result_data['image'], str):
-                                            return base64.b64decode(result_data['image'])
-                                    else:
-                                        # Maybe it's direct image bytes
-                                        return result_response.content
-                                        
-                                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:
-                        print("Failed to parse status response, continuing...")
-                        attempt += 1
-                        continue
-            
-                raise gr.Error("Job timed out after 5 minutes")
+                # Make another request to check if it's ready
+                # But based on the error, we shouldn't poll manually
+                # Instead, let's return an error asking user to try again
+                raise gr.Error("Request is queued for processing. Please try again in a few moments.")
             
             elif json_response.get("status") == "COMPLETED":
                 # Job completed immediately
                 if 'images' in json_response and len(json_response['images']) > 0:
                     image_info = json_response['images'][0]
                     if isinstance(image_info, dict) and 'url' in image_info:
+                        if progress_callback:
+                            progress_callback(0.9, "Downloading result...")
                         img_response = requests.get(image_info['url'])
                         return img_response.content
                     elif isinstance(image_info, str):
+                        # Base64 encoded image
                         return base64.b64decode(image_info)
                 elif 'image' in json_response:
                     return base64.b64decode(json_response['image'])
                 else:
-                    raise gr.Error(f"No images found in immediate response: {json_response}")
+                    raise gr.Error(f"No images found in response: {json_response}")
             
             else:
                 raise gr.Error(f"Unexpected response status: {json_response.get('status', 'unknown')}")
@@ -163,11 +97,24 @@ def query_api(payload, progress_callback=None):
     
     elif 'image/' in content_type:
         # Response is direct image bytes
+        if progress_callback:
+            progress_callback(1.0, "Complete!")
         return response.content
     
     else:
-        # Unknown content type, try to handle as bytes
-        return response.content
+        # Unknown content type, but try to handle as image bytes
+        # This might be the case where the router returns the image directly
+        if len(response.content) > 1000:  # Likely an image if it's substantial
+            if progress_callback:
+                progress_callback(1.0, "Complete!")
+            return response.content
+        else:
+            # Small response, probably an error
+            try:
+                error_response = response.json()
+                raise gr.Error(f"API Error: {error_response}")
+            except:
+                raise gr.Error(f"Unexpected response: {response.content.decode()[:500]}")
 
 def upload_image_to_fal(image_bytes):
     """Upload image to fal.ai and return the URL"""