Update app.py

app.py

@@ -14,7 +14,7 @@ MAX_SEED = np.iinfo(np.int32).max
 API_URL = "https://router.huggingface.co/fal-ai/fal-ai/flux-kontext/dev?_subdomain=queue"
 
 def get_headers():
-    """Get headers for API requests"""
+    """Get headers for Hugging Face router API requests"""
     hf_token = os.getenv("HF_TOKEN")
     if not hf_token:
         raise gr.Error("HF_TOKEN environment variable not found. Please add your Hugging Face token to the Space settings.")
@@ -24,12 +24,25 @@ def get_headers():
         "X-HF-Bill-To": "huggingface"
     }
 
-def query_api(payload, progress_callback=None):
+def upload_image_to_fal(image_bytes):
+    """Upload image to fal.ai and return the URL"""
+    # For now, we'll use base64 data URI as mentioned in the docs
+    # fal.ai supports base64 data URIs for image_url
+    image_base64 = base64.b64encode(image_bytes).decode('utf-8')
+    # Detect image format
+    try:
+        img = Image.open(io.BytesIO(image_bytes))
+        format_map = {'JPEG': 'jpeg', 'PNG': 'png', 'WEBP': 'webp'}
+        img_format = format_map.get(img.format, 'jpeg')
+    except:
+        img_format = 'jpeg'
+    
+    return f"data:image/{img_format};base64,{image_base64}"
     """Send request to the API and return response"""
-    headers = get_headers()
+    hf_headers = get_headers()
     
     # Submit the job
-    response = requests.post(API_URL, headers=headers, json=payload)
+    response = requests.post(API_URL, headers=hf_headers, json=payload)
     
     if response.status_code != 200:
         raise gr.Error(f"API request failed with status {response.status_code}: {response.text}")
@@ -47,6 +60,9 @@ def query_api(payload, progress_callback=None):
         if not status_url:
             raise gr.Error("No status URL provided in queue response")
         
+        # For fal.ai endpoints, we need different headers
+        fal_headers = get_fal_headers()
+        
         # Poll for completion
         max_attempts = 60  # Wait up to 5 minutes (60 * 5 seconds)
         attempt = 0
@@ -57,12 +73,8 @@ def query_api(payload, progress_callback=None):
             
             time.sleep(5)  # Wait 5 seconds between polls
             
-            # 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)
+            # Check status with fal.ai headers
+            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}")
@@ -78,10 +90,8 @@ def query_api(payload, progress_callback=None):
                     if not response_url:
                         raise gr.Error("No response URL provided")
                     
-                    # 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)
+                    # Get result with fal.ai headers
+                    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}")
@@ -161,7 +171,7 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    # Prepare the payload
+    # Prepare the payload for Hugging Face router
     payload = {
         "parameters": {
             "prompt": prompt,
@@ -182,14 +192,14 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
             # Auto-orient the image based on EXIF data
             input_image = ImageOps.exif_transpose(input_image)
             
-            # Convert PIL image to base64 for the API
+            # Convert PIL image to bytes
             img_byte_arr = io.BytesIO()
             input_image.save(img_byte_arr, format='PNG')
             img_byte_arr.seek(0)
-            image_base64 = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
+            image_bytes = img_byte_arr.getvalue()
             
-            # Add image to payload for image-to-image
-            payload["inputs"] = image_base64
+            # Add image bytes to payload - will be converted to base64 in query_api
+            payload["image_bytes"] = image_bytes
             
         except Exception as e:
             raise gr.Error(f"Could not process the uploaded image: {str(e)}. Please try uploading a different image format (JPEG, PNG, WebP).")
@@ -197,12 +207,12 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
         progress(0.1, desc="Processing image...")
     else:
         print(f"Received prompt for text-to-image: {prompt}")
-        # For text-to-image, we don't need the inputs field
+        # For text-to-image, we don't need an input image
         progress(0.1, desc="Generating image...")
 
     try:
-        # Make API request with progress callback
-        image_bytes = query_api(payload, progress_callback=progress)
+        # Make API request
+        image_bytes = query_api(payload)
         
         # Try to convert response bytes to PIL Image
         try:
@@ -236,11 +246,11 @@ steps_slider = gr.Slider(label="Steps", minimum=1, maximum=30, value=28, step=1)
 
 demo = gr.ChatInterface(
     fn=chat_fn,
-    title="FLUX.1 Kontext [dev] - Direct API",
+    title="FLUX.1 Kontext [dev] - Hugging Face Router",
     description="""<p style='text-align: center;'>
-    A simple chat UI for the <b>FLUX.1 Kontext</b> model using direct API calls with requests.
+    A simple chat UI for the <b>FLUX.1 Kontext [dev]</b> model using Hugging Face router.
     <br>
-    To edit an image, upload it and type your instructions (e.g., "Add a hat").
+    To edit an image, upload it and type your instructions (e.g., "Add a hat", "Turn the cat into a tiger").
     <br>
     To generate an image, just type a prompt (e.g., "A photo of an astronaut on a horse").
     <br>