Update app.py

app.py

@@ -2,44 +2,44 @@ import gradio as gr
 import numpy as np
 import random
 import os
-import tempfile
+import base64
+import requests
+import io
 from PIL import Image, ImageOps
 import pillow_heif  # For HEIF/AVIF support
-from huggingface_hub import InferenceClient
-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"
 
-# --- Global client variable ---
-client = None
+def get_headers():
+    """Get headers for 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.")
+    
+    return {
+        "Authorization": f"Bearer {hf_token}",
+        "X-HF-Bill-To": "huggingface"
+    }
 
-def load_client():
-    """Initialize the Inference Client"""
-    global client
-    if client is None:
-        # 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 hf_token:
-            client = InferenceClient(
-                provider="fal-ai",
-                api_key=hf_token,
-                bill_to="huggingface",
-            )
-        else:
-            raise gr.Error("HF_TOKEN environment variable not found. Please add your Hugging Face token to the Space settings.")
-    return client
+def query_api(payload):
+    """Send request to the API and return response"""
+    headers = get_headers()
+    response = requests.post(API_URL, headers=headers, json=payload)
+    
+    if response.status_code != 200:
+        raise gr.Error(f"API request failed with status {response.status_code}: {response.text}")
+    
+    return response.content
 
 # --- 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.
     """
-    # Load client
-    client = load_client()
+    # Register HEIF opener with PIL for AVIF/HEIF support
+    pillow_heif.register_heif_opener()
     
     prompt = message["text"]
     files = message["files"]
@@ -50,7 +50,16 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    input_image = None
+    # Prepare the payload
+    payload = {
+        "parameters": {
+            "prompt": prompt,
+            "seed": seed,
+            "guidance_scale": guidance_scale,
+            "num_inference_steps": steps
+        }
+    }
+
     if files:
         print(f"Received image: {files[0]}")
         try:
@@ -62,38 +71,36 @@ 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 bytes for the API
+            # Convert PIL image to base64 for the API
             img_byte_arr = io.BytesIO()
             input_image.save(img_byte_arr, format='PNG')
-            input_image_bytes = img_byte_arr.getvalue()
+            img_byte_arr.seek(0)
+            image_base64 = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
+            
+            # Add image to payload for image-to-image
+            payload["inputs"] = image_base64
             
         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).")
             
-        # Use image_to_image for editing
         progress(0.1, desc="Processing image...")
-        image = client.image_to_image(
-            input_image_bytes,
-            prompt=prompt,
-            model="black-forest-labs/FLUX.1-Kontext-dev",
-            # Note: guidance_scale and steps might not be supported by the API
-            # Check the API documentation for available parameters
-        )
-        progress(1.0, desc="Complete!")
     else:
         print(f"Received prompt for text-to-image: {prompt}")
-        # Use text_to_image for generation
+        # For text-to-image, we don't need the inputs field
         progress(0.1, desc="Generating image...")
-        image = client.text_to_image(
-            prompt=prompt,
-            model="black-forest-labs/FLUX.1-Kontext-dev",
-            # Note: guidance_scale and steps might not be supported by the API
-            # Check the API documentation for available parameters
-        )
+
+    try:
+        # Make API request
+        image_bytes = query_api(payload)
+        
+        # Convert response bytes to PIL Image
+        image = Image.open(io.BytesIO(image_bytes))
+        
         progress(1.0, desc="Complete!")
-    
-    # The client returns a PIL Image object
-    return gr.Image(value=image)
+        return gr.Image(value=image)
+        
+    except Exception as e:
+        raise gr.Error(f"Failed to generate image: {str(e)}")
 
 # --- UI Definition using gr.ChatInterface ---
 
@@ -102,14 +109,11 @@ randomize_checkbox = gr.Checkbox(label="Randomize seed", value=False)
 guidance_slider = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=2.5)
 steps_slider = gr.Slider(label="Steps", minimum=1, maximum=30, value=28, step=1)
 
-# Note: The Inference Client API may not support all parameters like guidance_scale and steps
-# Check the API documentation for supported parameters
-
 demo = gr.ChatInterface(
     fn=chat_fn,
-    title="FLUX.1 Kontext [dev] - Inference Client",
+    title="FLUX.1 Kontext [dev] - Direct API",
     description="""<p style='text-align: center;'>
-    A simple chat UI for the <b>FLUX.1 Kontext</b> model using Hugging Face Inference Client with fal-ai provider.
+    A simple chat UI for the <b>FLUX.1 Kontext</b> model using direct API calls with requests.
     <br>
     To edit an image, upload it and type your instructions (e.g., "Add a hat").
     <br>