Update app.py

app.py

@@ -2,171 +2,113 @@ import gradio as gr
 import numpy as np
 import random
 import os
-import base64
-import requests
-import time
-import io
+import tempfile
 from PIL import Image, ImageOps
 import pillow_heif  # For HEIF/AVIF support
+import io
 
 # --- Constants ---
 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 Hugging Face router API requests"""
+def load_client():
+    """Initialize the Inference Client"""
+    # Register HEIF opener with PIL for AVIF/HEIF support
+    pillow_heif.register_heif_opener()
+    
+    # Get token from environment variable
     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.")
     
-    return {
-        "Authorization": f"Bearer {hf_token}",
-        "X-HF-Bill-To": "huggingface"
-    }
+    return hf_token
 
-def query_api(payload, progress_callback=None):
-    """Send request to the API and return response"""
-    headers = get_headers()
+def query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callback=None):
+    """Send request to the API using HF Inference Client approach"""
+    import requests
+    import json
     
-    # Convert image to base64 if it's bytes
-    if "image_bytes" in payload:
-        payload["inputs"] = base64.b64encode(payload["image_bytes"]).decode("utf-8")
-        del payload["image_bytes"]
+    hf_token = load_client()
     
-    # Submit the job
     if progress_callback:
         progress_callback(0.1, "Submitting request...")
     
-    response = requests.post(API_URL, headers=headers, json=payload)
+    # Prepare the request data similar to HF Inference Client
+    url = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-Kontext-dev"
+    headers = {
+        "Authorization": f"Bearer {hf_token}",
+        "Content-Type": "application/json",
+        "X-Use-Billing": "huggingface"
+    }
     
-    if response.status_code != 200:
-        raise gr.Error(f"API request failed with status {response.status_code}: {response.text}")
+    # Convert image to base64
+    import base64
+    image_base64 = base64.b64encode(image_bytes).decode('utf-8')
     
-    # Debug the response
-    print(f"Response status: {response.status_code}")
-    print(f"Response content type: {response.headers.get('content-type', 'unknown')}")
-    print(f"Response content length: {len(response.content)}")
+    payload = {
+        "inputs": image_base64,
+        "parameters": {
+            "prompt": prompt,
+            "seed": seed,
+            "guidance_scale": guidance_scale,
+            "num_inference_steps": steps
+        },
+        "options": {
+            "wait_for_model": True
+        }
+    }
     
-    # Check if response is JSON (queue status) or binary (direct image)
-    content_type = response.headers.get('content-type', '').lower()
+    if progress_callback:
+        progress_callback(0.3, "Processing request...")
     
-    if 'application/json' in content_type:
-        # Response is JSON, likely queue status
-        try:
-            json_response = response.json()
-            print(f"JSON response: {json_response}")
+    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"
+            }
             
-            # Check if job was queued
-            if json_response.get("status") == "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")
-                
-                if not status_url or not response_url:
-                    raise gr.Error("Missing status_url or response_url in queue response")
-                
-                # For fal.ai queue, we need to use direct authentication
-                fal_headers = {
-                    "Authorization": f"Key {os.getenv('HF_TOKEN')}",
-                    "Content-Type": "application/json"
+            alt_payload = {
+                "inputs": image_base64,
+                "parameters": {
+                    "prompt": prompt,
+                    "seed": seed,
+                    "guidance_scale": guidance_scale,
+                    "num_inference_steps": steps
                 }
-                
-                # 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")
+            if progress_callback:
+                progress_callback(0.5, "Trying alternative routing...")
             
-            elif json_response.get("status") == "COMPLETED":
-                # Job completed immediately
+            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()
+        
+        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 contain image URL or base64
+            try:
+                json_response = response.json()
+                print(f"JSON response: {json_response}")
+                
+                # Handle different response formats
                 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, "Downloading result...")
                         img_response = requests.get(image_info['url'])
@@ -179,177 +121,33 @@ def query_api(payload, progress_callback=None):
                             progress_callback(1.0, "Complete!")
                         return base64.b64decode(image_info)
                 elif 'image' in json_response:
+                    # Single image field
                     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}")
-            
-            else:
-                raise gr.Error(f"Unexpected response status: {json_response.get('status', 'unknown')}")
-                
-        except requests.exceptions.JSONDecodeError as e:
-            raise gr.Error(f"Failed to parse JSON response: {str(e)}")
-    
-    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, 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
+                    raise gr.Error(f"Unexpected JSON response format: {json_response}")
+            except Exception as e:
+                raise gr.Error(f"Failed to parse JSON response: {str(e)}")
         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"""
-    # 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"""
-    hf_headers = get_headers()
-    
-    # Submit the job
-    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}")
-    
-    # Parse the initial response
-    try:
-        json_response = response.json()
-        print(f"Initial response: {json_response}")
-    except:
-        raise gr.Error("Failed to parse initial API response as JSON")
-    
-    # Check if job was queued
-    if json_response.get("status") == "IN_QUEUE":
-        status_url = json_response.get("status_url")
-        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
-        
-        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 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}")
-                raise gr.Error(f"Status check failed: {status_response.status_code}")
-            
-            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
-                    response_url = json_response.get("response_url")
-                    if not response_url:
-                        raise gr.Error("No response URL provided")
-                    
-                    # 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}")
-                        raise gr.Error(f"Failed to get result: {result_response.status_code}")
-                    
-                    # Check if result is JSON with image data
-                    try:
-                        result_data = result_response.json()
-                        print(f"Result data: {result_data}")
-                        
-                        # Look for image in various possible fields
-                        if 'images' in result_data and len(result_data['images']) > 0:
-                            # Images array with URLs or base64
-                            image_data = result_data['images'][0]
-                            if isinstance(image_data, dict) and 'url' in image_data:
-                                # Image URL - fetch it
-                                img_response = requests.get(image_data['url'])
-                                return img_response.content
-                            elif isinstance(image_data, str):
-                                # Assume base64
-                                return base64.b64decode(image_data)
-                        elif 'image' in result_data:
-                            # Single image field
-                            if isinstance(result_data['image'], str):
-                                return base64.b64decode(result_data['image'])
-                        elif 'url' in result_data:
-                            # Direct URL
-                            img_response = requests.get(result_data['url'])
-                            return img_response.content
-                        else:
-                            raise gr.Error(f"No image found in result: {result_data}")
-                            
-                    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
+            # 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:
+                raise gr.Error(f"Unexpected response format. Content: {response.content[:500]}")
                 
-            except requests.exceptions.JSONDecodeError:
-                raise gr.Error("Failed to parse status response")
-    
-        raise gr.Error("Job timed out after 5 minutes")
-    
-    elif json_response.get("status") == "COMPLETED":
-        # Job completed immediately
-        if 'images' in json_response and len(json_response['images']) > 0:
-            image_data = json_response['images'][0]
-            if isinstance(image_data, str):
-                return base64.b64decode(image_data)
-        elif 'image' in json_response:
-            return base64.b64decode(json_response['image'])
-        else:
-            raise gr.Error(f"No image found in immediate response: {json_response}")
-    
-    else:
-        raise gr.Error(f"Unexpected response status: {json_response.get('status', 'unknown')}")
+    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)}")
 
 # --- Core Inference Function for ChatInterface ---
 def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress()):
     """
     Performs image generation or editing based on user input from the chat interface.
     """
-    # Register HEIF opener with PIL for AVIF/HEIF support
-    pillow_heif.register_heif_opener()
-    
     prompt = message["text"]
     files = message["files"]
 
@@ -359,16 +157,6 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    # Prepare the payload for Hugging Face router
-    payload = {
-        "parameters": {
-            "prompt": prompt,
-            "seed": seed,
-            "guidance_scale": guidance_scale,
-            "num_inference_steps": steps
-        }
-    }
-
     if files:
         print(f"Received image: {files[0]}")
         try:
@@ -386,35 +174,33 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
             img_byte_arr.seek(0)
             image_bytes = img_byte_arr.getvalue()
             
-            # 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).")
             
         progress(0.1, desc="Processing image...")
     else:
-        print(f"Received prompt for text-to-image: {prompt}")
-        # For text-to-image, we don't need an input image
-        progress(0.1, desc="Generating image...")
+        # For text-to-image, we need a placeholder image or handle differently
+        # FLUX.1 Kontext is primarily an image-to-image model
+        raise gr.Error("This model (FLUX.1 Kontext) requires an input image. Please upload an image to edit.")
 
     try:
-        # Make API request with progress callback
-        image_bytes = query_api(payload, progress_callback=progress)
+        # Make API request
+        result_bytes = query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callback=progress)
         
         # Try to convert response bytes to PIL Image
         try:
-            image = Image.open(io.BytesIO(image_bytes))
+            image = Image.open(io.BytesIO(result_bytes))
         except Exception as img_error:
             print(f"Failed to open image: {img_error}")
-            print(f"Image bytes type: {type(image_bytes)}, length: {len(image_bytes) if hasattr(image_bytes, '__len__') else 'unknown'}")
+            print(f"Image bytes type: {type(result_bytes)}, length: {len(result_bytes) if hasattr(result_bytes, '__len__') else 'unknown'}")
             
             # Try to decode as base64 if direct opening failed
             try:
-                decoded_bytes = base64.b64decode(image_bytes)
+                import base64
+                decoded_bytes = base64.b64decode(result_bytes)
                 image = Image.open(io.BytesIO(decoded_bytes))
             except:
-                raise gr.Error(f"Could not process API response as image. Response length: {len(image_bytes) if hasattr(image_bytes, '__len__') else 'unknown'}")
+                raise gr.Error(f"Could not process API response as image. Response length: {len(result_bytes) if hasattr(result_bytes, '__len__') else 'unknown'}")
         
         progress(1.0, desc="Complete!")
         return gr.Image(value=image)
@@ -434,20 +220,20 @@ 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] - Hugging Face Router",
+    title="FLUX.1 Kontext [dev] - HF Inference Client",
     description="""<p style='text-align: center;'>
-    A simple chat UI for the <b>FLUX.1 Kontext [dev]</b> model using Hugging Face router.
+    A simple chat UI for the <b>FLUX.1 Kontext [dev]</b> model using Hugging Face Inference Client approach.
     <br>
-    To edit an image, upload it and type your instructions (e.g., "Add a hat", "Turn the cat into a tiger").
+    <b>Upload an image</b> and type your editing instructions (e.g., "Turn the cat into a tiger", "Add a hat").
     <br>
-    To generate an image, just type a prompt (e.g., "A photo of an astronaut on a horse").
+    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>.
     </p>""",
     multimodal=True,
     textbox=gr.MultimodalTextbox(
         file_types=["image"],
-        placeholder="Type a prompt and/or upload an image...",
+        placeholder="Upload an image and type your editing instructions...",
         render=False
     ),
     additional_inputs=[