"""ICTone Hugging Face Spaces demo optimized for ZeroGPU.

Space setup:
1. Select ZeroGPU hardware in the Space settings.
2. Add a Space secret named HF_TOKEN. The token owner must have accepted
   the access conditions for black-forest-labs/FLUX.1-Fill-dev.
3. Keep inference.py in the same directory as this file.
"""

from __future__ import annotations

import os
import random

import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import FluxFillPipeline
from PIL import Image

from inference import (
    DEFAULT_INSTANCE_PROMPT,
    apply_lut,
    estimate_lut,
    run_one,
)


MAX_SEED = np.iinfo(np.int32).max

FLUX_PATH = os.getenv(
    "FLUX_PATH",
    "black-forest-labs/FLUX.1-Fill-dev",
)
LORA_PATH = os.getenv(
    "LORA_PATH",
    "ToneStyle/ICTone-Fill-LoRA",
)
IMAGE_SIZE = int(os.getenv("IMAGE_SIZE", "512"))
LUT_SIZE = 33
HF_TOKEN = os.getenv("HF_TOKEN")


def load_pipeline() -> FluxFillPipeline:
    """Load FluxFill + ICTone LoRA once at Space startup.

    ZeroGPU recommends placing the model on CUDA at module scope. During Space
    startup this uses ZeroGPU's CUDA emulation; a real GPU is attached only
    while a @spaces.GPU function is running.
    """
    print(f"[ICTone] Loading base model: {FLUX_PATH}")
    print(f"[ICTone] Loading LoRA: {LORA_PATH}")

    load_kwargs = {
        "torch_dtype": torch.bfloat16,
    }
    if HF_TOKEN:
        load_kwargs["token"] = HF_TOKEN

    pipe = FluxFillPipeline.from_pretrained(
        FLUX_PATH,
        **load_kwargs,
    )
    pipe.load_lora_weights(LORA_PATH)

    # Required placement pattern for ZeroGPU. Do not lazy-load/move the model
    # inside infer().
    pipe.to("cuda")

    print("[ICTone] Pipeline ready.")
    return pipe


# Load once at module scope for efficient ZeroGPU model placement.
pipe = load_pipeline()


@spaces.GPU(duration=60)
def infer(
    content: Image.Image,
    reference: Image.Image,
    seed: int,
    randomize_seed: bool,
    guidance_scale: float,
    num_inference_steps: int,
    progress=gr.Progress(track_tqdm=True),
):
    """Run ICTone and reconstruct the result at the original content resolution."""
    if content is None or reference is None:
        raise gr.Error("Please upload both a content image and a reference image.")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    seed = int(seed)
    guidance_scale = float(guidance_scale)
    num_inference_steps = int(num_inference_steps)

    content_rgb = content.convert("RGB")
    reference_rgb = reference.convert("RGB")

    with torch.inference_mode():
        pred, panel, _, _ = run_one(
            pipe,
            content_rgb,
            reference_rgb,
            size=IMAGE_SIZE,
            prompt=DEFAULT_INSTANCE_PROMPT,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            seed=seed,
            generator_device="cuda",
        )

        # Lift the low-resolution Flux prediction back to the original content
        # resolution using ICTone's fitted 3D LUT.
        before = np.asarray(content_rgb)
        after = np.asarray(
            pred.resize(content_rgb.size, Image.Resampling.BILINEAR)
        )

        before_flat = before.reshape(-1, 3)
        after_flat = after.reshape(-1, 3)

        # Bound LUT fitting cost for very large uploaded images.
        max_samples = 500_000
        if len(before_flat) > max_samples:
            rng = np.random.default_rng(0)
            selected = rng.choice(
                len(before_flat),
                max_samples,
                replace=False,
            )
            before_flat = before_flat[selected]
            after_flat = after_flat[selected]

        lut = estimate_lut(
            before_flat,
            after_flat,
            size=LUT_SIZE,
            device="cuda",
        )
        output = Image.fromarray(
            apply_lut(
                before,
                lut,
                device="cuda",
            )
        )

    return output, panel, seed


with gr.Blocks(title="ICTone · In-Context Tone Style Transfer") as demo:
    gr.Markdown(
        """
# ICTone

**In-Context Tone Style Transfer**

Upload a **content image** and a **reference image**. ICTone transfers the
reference color, contrast, and photographic tone while preserving the content
of the source image.

The demo uses **FLUX.1-Fill-dev** with the **ICTone LoRA** and runs on
Hugging Face **ZeroGPU**. A short queue may appear when shared GPUs are busy.
"""
    )

    with gr.Row():
        content = gr.Image(
            label="Content image",
            type="pil",
        )
        reference = gr.Image(
            label="Reference image",
            type="pil",
        )

    with gr.Accordion("Generation settings", open=False):
        with gr.Row():
            seed = gr.Number(
                label="Seed",
                value=666,
                precision=0,
            )
            randomize_seed = gr.Checkbox(
                label="Randomize seed",
                value=False,
            )
        with gr.Row():
            guidance = gr.Slider(
                label="Guidance scale",
                minimum=1,
                maximum=100,
                value=50,
                step=1,
            )
            steps = gr.Slider(
                label="Inference steps",
                minimum=1,
                maximum=28,
                value=4,
                step=1,
            )

    run = gr.Button(
        "Transfer tone",
        variant="primary",
    )

    with gr.Row():
        output = gr.Image(
            label="Result",
            type="pil",
        )
        preview = gr.Image(
            label="Content | Reference | Result",
            type="pil",
        )

    used_seed = gr.Number(
        label="Used seed",
        precision=0,
    )

    run.click(
        fn=infer,
        inputs=[
            content,
            reference,
            seed,
            randomize_seed,
            guidance,
            steps,
        ],
        outputs=[
            output,
            preview,
            used_seed,
        ],
        show_progress="full",
    )

    gr.Markdown(
        """
**Models:** `black-forest-labs/FLUX.1-Fill-dev` +
`ToneStyle/ICTone-Fill-LoRA`

FLUX.1-Fill-dev is subject to the FLUX.1 [dev] license and access conditions.
"""
    )


if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
    )