Update app.py

app.py

@@ -1,25 +1,46 @@
-# This is a Gradio app that integrates a chat interface with a text-to-image and image editing model.
 import gradio as gr
 import numpy as np
 import random
 import os
+import tempfile
+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
 
-# --- Initialize Inference Client ---
-client = InferenceClient(
-    provider="fal-ai",
-    api_key=os.environ["HF_TOKEN"],
-    bill_to="huggingface",
-)
+# --- Global client variable ---
+client = None
+
+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
 
 # --- Core Inference Function for ChatInterface ---
-def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps):
+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()
+    
     prompt = message["text"]
     files = message["files"]
 
@@ -34,32 +55,44 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps):
         print(f"Received image: {files[0]}")
         try:
             # Try to open and convert the image
-            with open(files[0], "rb") as image_file:
-                input_image = image_file.read()
+            input_image = Image.open(files[0])
+            # Convert to RGB if needed (handles RGBA, P, etc.)
+            if input_image.mode != "RGB":
+                input_image = input_image.convert("RGB")
+            # Auto-orient the image based on EXIF data
+            input_image = ImageOps.exif_transpose(input_image)
+            
+            # Convert PIL image to bytes for the API
+            img_byte_arr = io.BytesIO()
+            input_image.save(img_byte_arr, format='PNG')
+            input_image_bytes = img_byte_arr.getvalue()
+            
         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).")
-
-    if input_image:
-        print(f"Received prompt for image editing: {prompt}")
+            
+        # Use image_to_image for editing
+        progress(0.1, desc="Processing image...")
         image = client.image_to_image(
-            input_image,
+            input_image_bytes,
             prompt=prompt,
             model="black-forest-labs/FLUX.1-Kontext-dev",
-            guidance_scale=guidance_scale,
-            num_inference_steps=steps,
-            seed=seed
+            # 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
+        progress(0.1, desc="Generating image...")
         image = client.text_to_image(
             prompt=prompt,
             model="black-forest-labs/FLUX.1-Kontext-dev",
-            guidance_scale=guidance_scale,
-            num_inference_steps=steps,
-            seed=seed
+            # Note: guidance_scale and steps might not be supported by the API
+            # Check the API documentation for available parameters
         )
-
-    # Return the PIL Image as a Gradio Image component
+        progress(1.0, desc="Complete!")
+    
+    # The client returns a PIL Image object
     return gr.Image(value=image)
 
 # --- UI Definition using gr.ChatInterface ---
@@ -69,15 +102,14 @@ 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)
 
-# --- Examples without external URLs ---
-# Remove examples temporarily to avoid format issues
-examples = None
+# 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]",
+    title="FLUX.1 Kontext [dev] - Inference Client",
     description="""<p style='text-align: center;'>
-    A simple chat UI for the <b>FLUX.1 Kontext</b> model running on ZeroGPU.
+    A simple chat UI for the <b>FLUX.1 Kontext</b> model using Hugging Face Inference Client with fal-ai provider.
     <br>
     To edit an image, upload it and type your instructions (e.g., "Add a hat").
     <br>
@@ -85,7 +117,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>.
     </p>""",
-    multimodal=True,  # This is important for MultimodalTextbox to work
+    multimodal=True,
     textbox=gr.MultimodalTextbox(
         file_types=["image"],
         placeholder="Type a prompt and/or upload an image...",
@@ -97,9 +129,8 @@ demo = gr.ChatInterface(
         guidance_slider,
         steps_slider
     ],
-    examples=examples,
     theme="soft"
 )
 
 if __name__ == "__main__":
-    demo.launch(show_error=True)
+    demo.launch()