Update app.py

app.py

@@ -3,6 +3,7 @@ import numpy as np
 import spaces
 import torch
 import random
+import os
 from PIL import Image
 
 # Import the pipeline from diffusers
@@ -11,13 +12,24 @@ from diffusers import FluxKontextPipeline
 # --- Constants and Model Loading ---
 MAX_SEED = np.iinfo(np.int32).max
 
+# --- FIX 1: Handle Hugging Face Authentication ---
+# This is a gated model. You must have access on Hugging Face and provide a token.
+# 1. Visit https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev and accept the terms.
+# 2. Get an access token from https://huggingface.co/settings/tokens
+# 3. Add the token below or set it as an environment variable `HF_TOKEN`.
+HF_TOKEN = os.getenv("HF_TOKEN", "YOUR_HUGGING_FACE_TOKEN") # Replace with your token
+
 # Load the pretrained model
-# Note: This requires a CUDA-enabled GPU. Error handling is added for environments without it.
 try:
-    pipe = FluxKontextPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-Kontext-dev", 
-        torch_dtype=torch.bfloat16
-    ).to("cuda")
+    if HF_TOKEN == "YOUR_HUGGING_FACE_TOKEN":
+        pipe = None
+        print("Warning: Hugging Face token not provided. Please replace 'YOUR_HUGGING_FACE_TOKEN' with your actual token.")
+    else:
+        pipe = FluxKontextPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-Kontext-dev", 
+            torch_dtype=torch.bfloat16,
+            token=HF_TOKEN,  # Use the token for authentication
+        ).to("cuda")
 except Exception as e:
     pipe = None
     print(f"Warning: Could not load the model on CUDA. GPU is required. Error: {e}")
@@ -28,40 +40,23 @@ except Exception as e:
 def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps, progress=gr.Progress(track_tqdm=True)):
     """
     Performs image generation or editing based on user input from the chat interface.
-
-    Args:
-        message (dict): A dictionary from gr.MultimodalTextbox, containing:
-                        - "text" (str): The user's text prompt.
-                        - "files" (list): A list of paths to uploaded files.
-        chat_history (list): The history of the conversation (managed by ChatInterface).
-        seed (int): The random seed for generation.
-        randomize_seed (bool): If True, a random seed is used.
-        guidance_scale (float): Controls adherence to the prompt.
-        steps (int): Number of inference steps.
-        progress (gr.Progress): Gradio progress tracker.
-
-    Returns:
-        PIL.Image.Image: The generated or edited image to be displayed in the chat.
     """
     if pipe is None:
-        raise gr.Error("Model could not be loaded. A CUDA-enabled GPU is required to run this application.")
+        raise gr.Error("Model could not be loaded. This could be due to a missing Hugging Face token, no access to the model, or no CUDA-enabled GPU.")
 
     prompt = message["text"]
     files = message["files"]
 
-    # Input validation
     if not prompt and not files:
         raise gr.Error("Please provide a prompt and/or upload an image.")
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    # Set up a PyTorch generator for reproducible results
-    generator = torch.Generator(device="cuda").manual_seed(seed)
+    generator = torch.Generator(device="cuda").manual_seed(int(seed))
 
     input_image = None
     if files:
-        # User has uploaded an image for editing (image-to-image)
         print(f"Received image: {files[0]}")
         input_image = Image.open(files[0]).convert("RGB")
         image = pipe(
@@ -72,7 +67,6 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
             generator=generator,
         ).images[0]
     else:
-        # No image uploaded, perform text-to-image generation
         print(f"Received prompt for text-to-image: {prompt}")
         image = pipe(
             prompt=prompt,
@@ -81,20 +75,32 @@ def chat_fn(message, chat_history, seed, randomize_seed, guidance_scale, steps,
             generator=generator,
         ).images[0]
         
-    # To also inform the user of the seed, you could optionally return a tuple,
-    # but for a clean image output, we just return the image.
-    # For example: return (image, f"Seed: {seed}")
     return image
 
 # --- UI Definition using gr.ChatInterface ---
 
-# Define the components for "Advanced Settings" that will be passed to `additional_inputs`
 seed_slider = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
 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)
 
-# Create the ChatInterface
+# --- FIX 2: Correctly format the examples as a list of lists ---
+# Format: [ [message_dict, seed, randomize, guidance, steps], ... ]
+examples = [
+    [
+        {"text": "A cute robot reading a book", "files": []}, 
+        42, False, 2.5, 28
+    ],
+    [
+        {"text": "change his shirt to a hawaiian shirt", "files": ["https://gradio-builds.s3.amazonaws.com/demo-files/chewbacca.png"]},
+        12345, False, 3.0, 25
+    ],
+    [
+        {"text": "make it a wooden house, add a chimney", "files": ["https://gradio-builds.s3.amazonaws.com/demo-files/house.png"]},
+        54321, False, 2.0, 30
+    ],
+]
+
 demo = gr.ChatInterface(
     fn=chat_fn,
     title="FLUX.1 Kontext [dev]",
@@ -107,11 +113,10 @@ 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>""",
-    # Use a multimodal textbox to allow both text and image uploads
     textbox=gr.MultimodalTextbox(
         file_types=["image"],
         placeholder="Type a prompt and/or upload an image...",
-        render=False  # Important: Let ChatInterface render the textbox
+        render=False
     ),
     additional_inputs=[
         seed_slider,
@@ -119,14 +124,9 @@ demo = gr.ChatInterface(
         guidance_slider,
         steps_slider
     ],
-    examples=[
-        {"text": "A cute robot reading a book", "files": []},
-        {"text": "change his shirt to a hawaiian shirt", "files": ["https://gradio-builds.s3.amazonaws.com/demo-files/chewbacca.png"]},
-        {"text": "make it a wooden house", "files": ["https://gradio-builds.s3.amazonaws.com/demo-files/house.png"]},
-    ],
+    examples=examples, # Use the correctly formatted list
     theme="soft"
 )
 
-# Launch the application
 if __name__ == "__main__":
     demo.launch()