Update app.py

app.py

@@ -1,51 +1,25 @@
+# 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 spaces
-import torch
 import random
 import os
-import tempfile
-from PIL import Image, ImageOps
-import pillow_heif  # For HEIF/AVIF support
-
-# Import the pipeline from diffusers
-from diffusers import FluxKontextPipeline
+from huggingface_hub import InferenceClient
 
 # --- Constants ---
 MAX_SEED = np.iinfo(np.int32).max
 
-# --- Global pipeline variable ---
-pipe = None
-
-def load_model():
-    """Load the model on CPU first, then move to GPU when needed"""
-    global pipe
-    if pipe 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:
-            pipe = FluxKontextPipeline.from_pretrained(
-                "black-forest-labs/FLUX.1-Kontext-dev", 
-                torch_dtype=torch.bfloat16,
-                token=hf_token,
-            )
-        else:
-            raise gr.Error("HF_TOKEN environment variable not found. Please add your Hugging Face token to the Space settings.")
-    return pipe
+# --- Initialize Inference Client ---
+client = InferenceClient(
+    provider="fal-ai",
+    api_key=os.environ["HF_TOKEN"],
+    bill_to="huggingface",
+)
 
 # --- Core Inference Function for ChatInterface ---
-@spaces.GPU(duration=120)  # Set duration based on expected inference time
-def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress(track_tqdm=True)):
+def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps):
     """
     Performs image generation or editing based on user input from the chat interface.
     """
-    # Load and move model to GPU within the decorated function
-    pipe = load_model()
-    pipe = pipe.to("cuda")
-    
     prompt = message["text"]
     files = message["files"]
 
@@ -55,42 +29,36 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    generator = torch.Generator(device="cuda").manual_seed(int(seed))
-
     input_image = None
     if files:
         print(f"Received image: {files[0]}")
         try:
             # Try to open and convert the image
-            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)
+            with open(files[0], "rb") as image_file:
+                input_image = image_file.read()
         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).")
-            
-        image = pipe(
-            image=input_image,
+
+    if input_image:
+        print(f"Received prompt for image editing: {prompt}")
+        image = client.image_to_image(
+            input_image,
             prompt=prompt,
+            model="black-forest-labs/FLUX.1-Kontext-dev",
             guidance_scale=guidance_scale,
             num_inference_steps=steps,
-            generator=generator,
-        ).images[0]
+            seed=seed
+        )
     else:
         print(f"Received prompt for text-to-image: {prompt}")
-        image = pipe(
+        image = client.text_to_image(
             prompt=prompt,
+            model="black-forest-labs/FLUX.1-Kontext-dev",
             guidance_scale=guidance_scale,
             num_inference_steps=steps,
-            generator=generator,
-        ).images[0]
-    
-    # Move model back to CPU to free GPU memory
-    pipe = pipe.to("cpu")
-    torch.cuda.empty_cache()
-    
+            seed=seed
+        )
+
     # Return the PIL Image as a Gradio Image component
     return gr.Image(value=image)
 
@@ -134,4 +102,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(show_error=True)