app.py

import gradio as gr
import numpy as np
import spaces
import torch
import random
import os
from PIL import Image

# Import the pipeline from diffusers
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
try:
    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}")

# --- Core Inference Function for ChatInterface ---

@spaces.GPU
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.
    """
    if pipe is None:
        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"]

    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)

    generator = torch.Generator(device="cuda").manual_seed(int(seed))

    input_image = None
    if files:
        print(f"Received image: {files[0]}")
        input_image = Image.open(files[0]).convert("RGB")
        image = pipe(
            image=input_image,
            prompt=prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=steps,
            generator=generator,
        ).images[0]
    else:
        print(f"Received prompt for text-to-image: {prompt}")
        image = pipe(
            prompt=prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=steps,
            generator=generator,
        ).images[0]
        
    return image

# --- UI Definition using gr.ChatInterface ---

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)

# --- 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]",
    description="""<p style='text-align: center;'>
    A simple chat UI for the <b>FLUX.1 Kontext</b> model.
    <br>
    To edit an image, upload it and type your instructions (e.g., "Add a hat").
    <br>
    To generate an image, just type a prompt (e.g., "A photo of an astronaut on a horse").
    <br>
    Find the model on <a href='https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev' target='_blank'>Hugging Face</a>.
    </p>""",
    textbox=gr.MultimodalTextbox(
        file_types=["image"],
        placeholder="Type a prompt and/or upload an image...",
        render=False
    ),
    additional_inputs=[
        seed_slider,
        randomize_checkbox,
        guidance_slider,
        steps_slider
    ],
    examples=examples, # Use the correctly formatted list
    theme="soft"
)

if __name__ == "__main__":
    demo.launch()