Update app.py

app.py

@@ -58,19 +58,109 @@ def query_api(payload, progress_callback=None):
             
             # Check if job was queued
             if json_response.get("status") == "IN_QUEUE":
-                # 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...")
+                # Get the status and response URLs provided by the API
+                status_url = json_response.get("status_url")
+                response_url = json_response.get("response_url")
                 
-                # Wait a bit and try to get the response again
-                import time
-                time.sleep(10)  # Wait 10 seconds for processing
+                if not status_url or not response_url:
+                    raise gr.Error("Missing status_url or response_url in queue response")
                 
-                # 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.")
+                # For fal.ai queue, we need to use direct authentication
+                fal_headers = {
+                    "Authorization": f"Key {os.getenv('HF_TOKEN')}",
+                    "Content-Type": "application/json"
+                }
+                
+                # Poll for completion using the provided URLs
+                max_attempts = 60  # Wait up to 5 minutes
+                attempt = 0
+                
+                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 the provided status_url
+                    status_response = requests.get(status_url, headers=fal_headers)
+                    
+                    if status_response.status_code != 200:
+                        print(f"Status response: {status_response.status_code} - {status_response.text}")
+                        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 using response_url
+                            result_response = requests.get(response_url, headers=fal_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
+                                if progress_callback:
+                                    progress_callback(1.0, "Complete!")
+                                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
+                                            if progress_callback:
+                                                progress_callback(0.9, "Downloading result...")
+                                            img_response = requests.get(image_info['url'])
+                                            if progress_callback:
+                                                progress_callback(1.0, "Complete!")
+                                            return img_response.content
+                                        elif isinstance(image_info, str):
+                                            # Base64 encoded
+                                            if progress_callback:
+                                                progress_callback(1.0, "Complete!")
+                                            return base64.b64decode(image_info)
+                                    elif 'image' in result_data:
+                                        # Single image field
+                                        if isinstance(result_data['image'], str):
+                                            if progress_callback:
+                                                progress_callback(1.0, "Complete!")
+                                            return base64.b64decode(result_data['image'])
+                                    else:
+                                        # Maybe it's direct image bytes
+                                        if progress_callback:
+                                            progress_callback(1.0, "Complete!")
+                                        return result_response.content
+                                        
+                                except requests.exceptions.JSONDecodeError:
+                                    # Result might be direct image bytes
+                                    if progress_callback:
+                                        progress_callback(1.0, "Complete!")
+                                    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")
             
             elif json_response.get("status") == "COMPLETED":
                 # Job completed immediately
@@ -80,11 +170,17 @@ def query_api(payload, progress_callback=None):
                         if progress_callback:
                             progress_callback(0.9, "Downloading result...")
                         img_response = requests.get(image_info['url'])
+                        if progress_callback:
+                            progress_callback(1.0, "Complete!")
                         return img_response.content
                     elif isinstance(image_info, str):
                         # Base64 encoded image
+                        if progress_callback:
+                            progress_callback(1.0, "Complete!")
                         return base64.b64decode(image_info)
                 elif 'image' in json_response:
+                    if progress_callback:
+                        progress_callback(1.0, "Complete!")
                     return base64.b64decode(json_response['image'])
                 else:
                     raise gr.Error(f"No images found in response: {json_response}")