Update app.py

app.py

@@ -6,12 +6,14 @@ import tempfile
 from PIL import Image, ImageOps
 import pillow_heif  # For HEIF/AVIF support
 import io
+import fal_client
+import base64
 
 # --- Constants ---
 MAX_SEED = np.iinfo(np.int32).max
 
 def load_client():
-    """Initialize the Inference Client"""
+    """Initialize the FAL Client through HF"""
     # Register HEIF opener with PIL for AVIF/HEIF support
     pillow_heif.register_heif_opener()
     
@@ -20,121 +22,125 @@ def load_client():
     if not hf_token:
         raise gr.Error("HF_TOKEN environment variable not found. Please add your Hugging Face token to the Space settings.")
     
-    return hf_token
+    # Set the HF token for fal_client to use HF routing
+    os.environ["FAL_KEY"] = hf_token
+    return True
 
 def query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callback=None):
-    """Send request to the API using HF Router for fal.ai provider"""
-    import requests
-    import json
-    import base64
+    """Send request using fal_client"""
     
-    hf_token = load_client()
+    load_client()
     
     if progress_callback:
         progress_callback(0.1, "Submitting request...")
     
-    # 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}",
-        "X-HF-Bill-To": "huggingface",
-        "Content-Type": "application/json"
-    }
-    
-    # Convert image to base64
+    # Convert image bytes to base64
     image_base64 = base64.b64encode(image_bytes).decode('utf-8')
     
-    # Fixed payload structure - prompt should be at the top level
-    payload = {
-        "prompt": prompt,
-        "inputs": image_base64,
-        "seed": seed,
-        "guidance_scale": guidance_scale,
-        "num_inference_steps": steps
-    }
+    # Create a temporary file for the image
+    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file:
+        temp_file.write(image_bytes)
+        temp_file_path = temp_file.name
     
-    if progress_callback:
-        progress_callback(0.3, "Processing request...")
+    def on_queue_update(update):
+        if isinstance(update, fal_client.InProgress):
+            for log in update.logs:
+                print(f"FAL Log: {log['message']}")
+                if progress_callback:
+                    progress_callback(0.5, f"Processing: {log['message'][:50]}...")
     
     try:
-        response = requests.post(url, headers=headers, json=payload, timeout=300)
+        if progress_callback:
+            progress_callback(0.3, "Connecting to FAL API...")
         
-        if response.status_code != 200:
-            raise gr.Error(f"API request failed with status {response.status_code}: {response.text}")
+        # Use fal_client.subscribe following the pattern you provided
+        result = fal_client.subscribe(
+            "fal-ai/flux-kontext/dev",
+            arguments={
+                "prompt": prompt,
+                "image_url": f"data:image/png;base64,{image_base64}",
+                "seed": seed,
+                "guidance_scale": guidance_scale,
+                "num_inference_steps": steps,
+            },
+            with_logs=True,
+            on_queue_update=on_queue_update,
+        )
         
-        # 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)}")
+        print(f"FAL Result: {result}")
         
-        if 'image/' in content_type:
-            # Direct image response
-            if progress_callback:
-                progress_callback(1.0, "Complete!")
-            return response.content
-        elif 'application/json' in content_type:
-            # JSON response - might be queue status or result
-            try:
-                json_response = response.json()
-                print(f"JSON response: {json_response}")
-                
-                # 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:
-                        # Download image from URL
+        if progress_callback:
+            progress_callback(0.9, "Processing result...")
+        
+        # Handle the result
+        if isinstance(result, dict):
+            if 'images' in result and len(result['images']) > 0:
+                # Get the first image
+                image_info = result['images'][0]
+                if isinstance(image_info, dict) and 'url' in image_info:
+                    # Download image from URL
+                    import requests
+                    img_response = requests.get(image_info['url'])
+                    if img_response.status_code == 200:
                         if progress_callback:
-                            progress_callback(0.9, "Downloading result...")
-                        img_response = requests.get(image_info['url'])
+                            progress_callback(1.0, "Complete!")
+                        return img_response.content
+                    else:
+                        raise gr.Error(f"Failed to download result image: {img_response.status_code}")
+                elif isinstance(image_info, str):
+                    # Direct URL
+                    import requests
+                    img_response = requests.get(image_info)
+                    if img_response.status_code == 200:
+                        if progress_callback:
+                            progress_callback(1.0, "Complete!")
+                        return img_response.content
+            elif 'image' in result:
+                # Single image field
+                if isinstance(result['image'], dict) and 'url' in result['image']:
+                    import requests
+                    img_response = requests.get(result['image']['url'])
+                    if img_response.status_code == 200:
+                        if progress_callback:
+                            progress_callback(1.0, "Complete!")
+                        return img_response.content
+                elif isinstance(result['image'], str):
+                    # Could be URL or base64
+                    if result['image'].startswith('http'):
+                        import requests
+                        img_response = requests.get(result['image'])
                         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:
-                            progress_callback(1.0, "Complete!")
-                        return base64.b64decode(image_info)
-                elif 'image' in json_response:
-                    # Single image field
+                    else:
+                        # Assume base64
+                        try:
+                            if progress_callback:
+                                progress_callback(1.0, "Complete!")
+                            return base64.b64decode(result['image'])
+                        except:
+                            pass
+            elif 'url' in result:
+                # Direct URL in result
+                import requests
+                img_response = requests.get(result['url'])
+                if img_response.status_code == 200:
                     if progress_callback:
                         progress_callback(1.0, "Complete!")
-                    return base64.b64decode(json_response['image'])
-                else:
-                    raise gr.Error(f"Unexpected JSON response format: {json_response}")
-            except json.JSONDecodeError as e:
-                raise gr.Error(f"Failed to parse JSON response: {str(e)}")
-        else:
-            # Try to treat as image bytes
-            if len(response.content) > 1000:  # Likely an image
-                if progress_callback:
-                    progress_callback(1.0, "Complete!")
-                return response.content
-            else:
-                # 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
+                    return img_response.content
+        
+        # If we get here, the result format is unexpected
+        raise gr.Error(f"Unexpected result format from FAL API: {result}")
+        
     except Exception as e:
-        raise gr.Error(f"Unexpected error: {str(e)}")
+        raise gr.Error(f"FAL API error: {str(e)}")
+    finally:
+        # Clean up temporary file
+        try:
+            os.unlink(temp_file_path)
+        except:
+            pass
 
 # --- Core Inference Function for ChatInterface ---
 def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress()):
@@ -213,15 +219,17 @@ 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] - HF Inference Client",
+    title="FLUX.1 Kontext [dev] - FAL Client",
     description="""<p style='text-align: center;'>
-    A simple chat UI for the <b>FLUX.1 Kontext [dev]</b> model using Hugging Face Inference Client approach.
+    A simple chat UI for the <b>FLUX.1 Kontext [dev]</b> model using FAL AI client through Hugging Face.
     <br>
     <b>Upload an image</b> and type your editing instructions (e.g., "Turn the cat into a tiger", "Add a hat").
     <br>
     This model specializes in understanding context and making precise edits to your images.
     <br>
     Find the model on <a href='https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev' target='_blank'>Hugging Face</a>.
+    <br>
+    <b>Note:</b> Uses HF_TOKEN environment variable through HF inference providers.
     </p>""",
     multimodal=True,
     textbox=gr.MultimodalTextbox(