Update app.py

app.py

@@ -3,142 +3,156 @@ import numpy as np
 import random
 import os
 import tempfile
+import subprocess
+import json
 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 FAL Client through HF"""
-    # Register HEIF opener with PIL for AVIF/HEIF support
-    pillow_heif.register_heif_opener()
+def setup_node_environment():
+    """Setup Node.js environment and install required packages"""
+    try:
+        # Check if node is available
+        result = subprocess.run(['node', '--version'], capture_output=True, text=True)
+        if result.returncode != 0:
+            raise gr.Error("Node.js is not installed. Please install Node.js to use this feature.")
+        
+        # Check if @huggingface/inference is installed, if not install it
+        package_check = subprocess.run(['npm', 'list', '@huggingface/inference'], capture_output=True, text=True)
+        if package_check.returncode != 0:
+            print("Installing @huggingface/inference package...")
+            install_result = subprocess.run(['npm', 'install', '@huggingface/inference'], capture_output=True, text=True)
+            if install_result.returncode != 0:
+                raise gr.Error(f"Failed to install @huggingface/inference: {install_result.stderr}")
+        
+        return True
+    except FileNotFoundError:
+        raise gr.Error("Node.js or npm not found. Please install Node.js and npm.")
+
+def create_js_inference_script(image_path, prompt, hf_token):
+    """Create JavaScript inference script"""
+    js_code = f"""
+const {{ InferenceClient }} = require("@huggingface/inference");
+const fs = require("fs");
+
+async function runInference() {{
+    try {{
+        const client = new InferenceClient("{hf_token}");
+        const data = fs.readFileSync("{image_path}");
+        
+        const image = await client.imageToImage({{
+            provider: "replicate",
+            model: "black-forest-labs/FLUX.1-Kontext-dev",
+            inputs: data,
+            parameters: {{ prompt: "{prompt}" }},
+        }}, {{
+            billTo: "huggingface",
+        }});
+        
+        // Convert blob to buffer
+        const arrayBuffer = await image.arrayBuffer();
+        const buffer = Buffer.from(arrayBuffer);
+        
+        // Output as base64 for Python to read
+        const base64 = buffer.toString('base64');
+        console.log(JSON.stringify({{
+            success: true,
+            image_base64: base64,
+            content_type: image.type || 'image/jpeg'
+        }}));
+        
+    }} catch (error) {{
+        console.log(JSON.stringify({{
+            success: false,
+            error: error.message
+        }}));
+        process.exit(1);
+    }}
+}}
+
+runInference();
+"""
+    return js_code
+
+def query_api_js(image_bytes, prompt, seed, guidance_scale, steps, progress_callback=None):
+    """Send request using JavaScript HF Inference Client"""
     
     # 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.")
+        raise gr.Error("HF_TOKEN environment variable not found. Please add your Hugging Face token to the environment.")
     
-    # 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 using fal_client"""
+    if progress_callback:
+        progress_callback(0.1, "Setting up Node.js environment...")
     
-    load_client()
+    # Setup Node.js environment
+    setup_node_environment()
     
     if progress_callback:
-        progress_callback(0.1, "Submitting request...")
-    
-    # Convert image bytes to base64
-    image_base64 = base64.b64encode(image_bytes).decode('utf-8')
+        progress_callback(0.2, "Preparing image...")
     
     # 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
+        temp_image_path = temp_file.name
     
-    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]}...")
+    # Create temporary JavaScript file
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.js', delete=False) as js_file:
+        js_code = create_js_inference_script(temp_image_path, prompt.replace('"', '\\"'), hf_token)
+        js_file.write(js_code)
+        js_file_path = js_file.name
     
     try:
         if progress_callback:
-            progress_callback(0.3, "Connecting to FAL API...")
-        
-        # 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,
+            progress_callback(0.3, "Running JavaScript inference...")
+        
+        # Run the JavaScript code
+        result = subprocess.run(
+            ['node', js_file_path],
+            capture_output=True,
+            text=True,
+            timeout=300  # 5 minute timeout
         )
         
-        print(f"FAL Result: {result}")
+        if progress_callback:
+            progress_callback(0.8, "Processing result...")
+        
+        if result.returncode != 0:
+            raise gr.Error(f"JavaScript inference failed: {result.stderr}")
+        
+        # Parse the JSON output
+        try:
+            output = json.loads(result.stdout.strip())
+        except json.JSONDecodeError:
+            raise gr.Error(f"Failed to parse JavaScript output: {result.stdout}")
+        
+        if not output.get('success'):
+            raise gr.Error(f"Inference error: {output.get('error', 'Unknown error')}")
         
         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(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:
-                        # 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 img_response.content
-        
-        # If we get here, the result format is unexpected
-        raise gr.Error(f"Unexpected result format from FAL API: {result}")
+            progress_callback(0.9, "Decoding image...")
         
+        # Decode base64 image
+        import base64
+        image_data = base64.b64decode(output['image_base64'])
+        
+        if progress_callback:
+            progress_callback(1.0, "Complete!")
+        
+        return image_data
+        
+    except subprocess.TimeoutExpired:
+        raise gr.Error("Inference timed out. Please try again.")
     except Exception as e:
-        raise gr.Error(f"FAL API error: {str(e)}")
+        raise gr.Error(f"Error running JavaScript inference: {str(e)}")
     finally:
-        # Clean up temporary file
+        # Clean up temporary files
         try:
-            os.unlink(temp_file_path)
+            os.unlink(temp_image_path)
+            os.unlink(js_file_path)
         except:
             pass
 
@@ -159,6 +173,9 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
     if files:
         print(f"Received image: {files[0]}")
         try:
+            # Register HEIF opener with PIL for AVIF/HEIF support
+            pillow_heif.register_heif_opener()
+            
             # Try to open and convert the image
             input_image = Image.open(files[0])
             # Convert to RGB if needed (handles RGBA, P, etc.)
@@ -183,8 +200,8 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
         raise gr.Error("This model (FLUX.1 Kontext) requires an input image. Please upload an image to edit.")
 
     try:
-        # Make API request
-        result_bytes = query_api(image_bytes, prompt, seed, guidance_scale, steps, progress_callback=progress)
+        # Make API request using JavaScript
+        result_bytes = query_api_js(image_bytes, prompt, seed, guidance_scale, steps, progress_callback=progress)
         
         # Try to convert response bytes to PIL Image
         try:
@@ -192,14 +209,7 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
         except Exception as img_error:
             print(f"Failed to open image: {img_error}")
             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:
-                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(result_bytes) if hasattr(result_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)
@@ -219,9 +229,9 @@ 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] - FAL Client",
+    title="FLUX.1 Kontext [dev] - HF Inference Client (JS)",
     description="""<p style='text-align: center;'>
-    A simple chat UI for the <b>FLUX.1 Kontext [dev]</b> model using FAL AI client through Hugging Face.
+    A simple chat UI for the <b>FLUX.1 Kontext [dev]</b> model using Hugging Face Inference Client via JavaScript.
     <br>
     <b>Upload an image</b> and type your editing instructions (e.g., "Turn the cat into a tiger", "Add a hat").
     <br>
@@ -229,7 +239,7 @@ demo = gr.ChatInterface(
     <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.
+    <b>Requirements:</b> Node.js and npm must be installed. Uses HF_TOKEN environment variable.
     </p>""",
     multimodal=True,
     textbox=gr.MultimodalTextbox(