Update app.py

app.py

@@ -23,25 +23,25 @@ def load_client():
     return hf_token
 
 def query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callback=None):
-    """Send request to the API using HF Inference Client approach"""
+    """Send request to the API using HF Router for fal.ai provider"""
     import requests
     import json
+    import base64
     
     hf_token = load_client()
     
     if progress_callback:
         progress_callback(0.1, "Submitting request...")
     
-    # Prepare the request data similar to HF Inference Client
-    url = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-Kontext-dev"
+    # Use the HF router to access fal.ai provider
+    url = "https://router.huggingface.co/fal-ai/fal-ai/flux-kontext/dev"
     headers = {
         "Authorization": f"Bearer {hf_token}",
-        "Content-Type": "application/json",
-        "X-Use-Billing": "huggingface"
+        "X-HF-Bill-To": "huggingface",
+        "Content-Type": "application/json"
     }
     
     # Convert image to base64
-    import base64
     image_base64 = base64.b64encode(image_bytes).decode('utf-8')
     
     payload = {
@@ -51,9 +51,6 @@ def query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callbac
             "seed": seed,
             "guidance_scale": guidance_scale,
             "num_inference_steps": steps
-        },
-        "options": {
-            "wait_for_model": True
         }
     }
     
@@ -63,35 +60,13 @@ def query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callbac
     try:
         response = requests.post(url, headers=headers, json=payload, timeout=300)
         
-        if response.status_code != 200:
-            # Try alternative approach with fal-ai provider routing
-            alt_url = "https://router.huggingface.co/fal-ai/black-forest-labs/FLUX.1-Kontext-dev"
-            alt_headers = {
-                "Authorization": f"Bearer {hf_token}",
-                "X-HF-Bill-To": "huggingface",
-                "Content-Type": "application/json"
-            }
-            
-            alt_payload = {
-                "inputs": image_base64,
-                "parameters": {
-                    "prompt": prompt,
-                    "seed": seed,
-                    "guidance_scale": guidance_scale,
-                    "num_inference_steps": steps
-                }
-            }
-            
-            if progress_callback:
-                progress_callback(0.5, "Trying alternative routing...")
-            
-            response = requests.post(alt_url, headers=alt_headers, json=alt_payload, timeout=300)
-        
         if response.status_code != 200:
             raise gr.Error(f"API request failed with status {response.status_code}: {response.text}")
         
         # Check if response is image bytes or JSON
         content_type = response.headers.get('content-type', '').lower()
+        print(f"Response content type: {content_type}")
+        print(f"Response length: {len(response.content)}")
         
         if 'image/' in content_type:
             # Direct image response
@@ -99,12 +74,18 @@ def query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callbac
                 progress_callback(1.0, "Complete!")
             return response.content
         elif 'application/json' in content_type:
-            # JSON response, might contain image URL or base64
+            # JSON response - might be queue status or result
             try:
                 json_response = response.json()
                 print(f"JSON response: {json_response}")
                 
-                # Handle different response formats
+                # Check if it's a queue response
+                if json_response.get("status") == "IN_QUEUE":
+                    if progress_callback:
+                        progress_callback(0.4, "Request queued, please wait...")
+                    raise gr.Error("Request is being processed. Please try again in a few moments.")
+                
+                # Handle immediate completion or result
                 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:
@@ -112,9 +93,12 @@ def query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callbac
                         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
+                        if img_response.status_code == 200:
+                            if progress_callback:
+                                progress_callback(1.0, "Complete!")
+                            return img_response.content
+                        else:
+                            raise gr.Error(f"Failed to download image: {img_response.status_code}")
                     elif isinstance(image_info, str):
                         # Base64 encoded image
                         if progress_callback:
@@ -127,7 +111,7 @@ def query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callbac
                     return base64.b64decode(json_response['image'])
                 else:
                     raise gr.Error(f"Unexpected JSON response format: {json_response}")
-            except Exception as e:
+            except json.JSONDecodeError as e:
                 raise gr.Error(f"Failed to parse JSON response: {str(e)}")
         else:
             # Try to treat as image bytes
@@ -136,12 +120,22 @@ def query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callbac
                     progress_callback(1.0, "Complete!")
                 return response.content
             else:
-                raise gr.Error(f"Unexpected response format. Content: {response.content[:500]}")
+                # Small response, probably an error
+                try:
+                    error_text = response.content.decode('utf-8')
+                    raise gr.Error(f"Unexpected response: {error_text[:500]}")
+                except:
+                    raise gr.Error(f"Unexpected response format. Content length: {len(response.content)}")
                 
     except requests.exceptions.Timeout:
         raise gr.Error("Request timed out. Please try again.")
     except requests.exceptions.RequestException as e:
         raise gr.Error(f"Request failed: {str(e)}")
+    except gr.Error:
+        # Re-raise Gradio errors as-is
+        raise
+    except Exception as e:
+        raise gr.Error(f"Unexpected error: {str(e)}")
 
 # --- Core Inference Function for ChatInterface ---
 def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress()):