app.py

"""HuggingFace Spaces (Gradio) demo for Foveated Diffusion.

Mirrors the behavior of `webgui/server.py` from the project repo, adapted to
Gradio + an HF Spaces deployment:

  - LoRA checkpoints (no_fov, fov_random, fov_saliency, fov_bbox) are fetched
    from `bchao1/foveated_diffusion` via `hf_hub_download` on startup. Same
    fuse/unfuse machinery as the Flask server — switching is fast and bounded
    in VRAM.
  - The painted mask + "preset circle" mask spec are quantized to LR blocks
    under the same any-touch-upgrades-the-block rule used in `webgui/`.
  - Generation runs under the same fixed settings (1024x1024, 50 steps, soft
    foveation blend on, decode_mode=merge, prediction_type=clean).

Layout (mirrors webgui/static/index.html):
  Left column   = prompt, seed, LoRA, mask mode, draw/preset controls, Generate.
  Right column  = side-by-side tokenization mask preview + generated image.
"""

from __future__ import annotations

import io
import os
import sys
import threading
import time
from types import SimpleNamespace

# Match webgui/server.py: disable sageattention before importing torch-deps.
sys.modules["sageattention"] = None

# ZeroGPU runtime. MUST be imported before torch so that `spaces` can intercept
# CUDA initialization during module load. On non-Spaces envs `spaces` isn't
# installed, so we install a no-op shim that makes `@spaces.GPU(...)` a
# transparent decorator. This keeps `python app.py` working locally.
try:
    import spaces  # type: ignore
except ImportError:
    class _SpacesShim:
        @staticmethod
        def GPU(*args, **kwargs):
            def decorator(fn):
                return fn
            # @spaces.GPU (no parens): the first arg is the function itself.
            if args and callable(args[0]) and not kwargs:
                return args[0]
            return decorator
    spaces = _SpacesShim()  # type: ignore[assignment]

import numpy as np
import torch
import torch.nn.functional as F
import gradio as gr
from huggingface_hub import hf_hub_download
from PIL import Image

# Make the bundled `src/` package importable.
#
# On Spaces, `src/` is vendored into the Space root by `upload_space.py`, so
# it lives next to app.py at `_HERE/src`. For local dev (running app.py from
# the project repo before upload), `src/` only exists at the repo root, so we
# also add `_HERE/../..` to sys.path. The second insert is a no-op on Spaces
# (the path doesn't exist there).
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)
_REPO_ROOT = os.path.normpath(os.path.join(_HERE, "..", ".."))
if os.path.isdir(os.path.join(_REPO_ROOT, "src")):
    sys.path.insert(0, _REPO_ROOT)

from diffsynth.core import load_state_dict  # noqa: E402
from src.inference import load_pipeline  # noqa: E402
from src.inference.visualize import create_tokenization_mask_vis  # noqa: E402
from src.masks import create_foveation_mask_full_res  # noqa: E402

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

# Fixed generation settings (per webgui spec).
MODEL_ID = "black-forest-labs/FLUX.2-klein-base-4B"
HEIGHT = 1024
WIDTH = 1024
NUM_INFERENCE_STEPS = 50
GUIDANCE_SCALE = 4.0
DECODE_MODE = "merge"
PREDICTION_TYPE = "clean"
LR_DOWNSAMPLE_FACTOR = 2
SOFT_FOVEATION_BLEND = True

# Masks are tiny (1024x1024 fp32 ~ 4 MB) and the only thing they need to do
# outside a `@spaces.GPU` call is feed the tokenization preview. Building them
# on CPU keeps `refresh_tokenization` (fires on every slider move) off the GPU
# entirely; we move them to `_pipe.device` once, inside `generate`, just before
# the actual forward.
MASK_DEVICE = "cpu"

# LoRA registry. Files come from the bchao1/foveated_diffusion model repo.
LORA_REPO = "bchao1/foveated_diffusion"
LORA_FILES = {
    "no_fov":   "image/no_fov.safetensors",
    "random":   "image/fov_random.safetensors",
    "saliency": "image/fov_saliency.safetensors",
    # "bbox":   "image/fov_bbox.safetensors",   # disabled in this UI build
}
DEFAULT_LORA = "random"

_pipe = None

_lora_registry: dict = {}    # name -> {"path": str, "state_dict": dict or None}
_lora_order: list = []
# LoRA state. The dropdown (`lora_dd`) is the actual source of truth — its
# value is plumbed into every handler (and `generate`) as an explicit Gradio
# input, so it crosses the @spaces.GPU process boundary safely under either
# fork or spawn. The two globals below are kept only for:
#   _selected_lora — parent-only book-keeping for `switch_lora`'s rollback on
#                    an invalid dropdown value (defensive; not load-bearing).
#   _current_lora  — worker-local. Tracks what's fused in *this* worker's
#                    model copy so `_apply_selected_lora` can skip work when
#                    the dropdown choice already matches.
_selected_lora: str = DEFAULT_LORA
_current_lora: str | None = None
_lora_lock = threading.Lock()
_job_lock = threading.Lock()  # serializes generation + LoRA switching


# ---------------------------------------------------------------------------
# LoRA fuse / unfuse (ported verbatim from webgui/server.py)
# ---------------------------------------------------------------------------

def _convert_lora_sd(state_dict):
    loader = _pipe.lora_loader(torch_dtype=_pipe.torch_dtype, device="cpu")
    return loader.convert_state_dict(state_dict)


def _apply_lora_delta(state_dict, sign):
    lora_layer_names = {
        k[: -len(".lora_B.weight")]
        for k in state_dict if k.endswith(".lora_B.weight")
    }
    updated = 0
    for name, module in _pipe.dit.named_modules():
        if name not in lora_layer_names:
            continue
        wb = state_dict[name + ".lora_B.weight"].to(device=_pipe.device, dtype=_pipe.torch_dtype)
        wa = state_dict[name + ".lora_A.weight"].to(device=_pipe.device, dtype=_pipe.torch_dtype)
        if wb.dim() == 4:
            wb = wb.squeeze(3).squeeze(2)
            wa = wa.squeeze(3).squeeze(2)
            delta = torch.mm(wb, wa).unsqueeze(2).unsqueeze(3)
        else:
            delta = torch.mm(wb, wa)
        base = module.weight.data
        base.add_(delta.to(device=base.device, dtype=base.dtype), alpha=float(sign))
        updated += 1
    return updated


def _switch_lora_unlocked(target_name: str):
    global _current_lora
    if target_name == _current_lora:
        return
    if target_name not in _lora_registry:
        raise ValueError(f"unknown lora: {target_name}")

    if _current_lora is not None:
        cur = _lora_registry[_current_lora]
        n = _apply_lora_delta(cur["state_dict"], sign=-1)
        print(f"[lora] unfused {_current_lora} ({n} tensors)")

    nxt = _lora_registry[target_name]
    if nxt["state_dict"] is None:
        raw = load_state_dict(nxt["path"], torch_dtype=_pipe.torch_dtype)
        nxt["state_dict"] = _convert_lora_sd(raw)
    n = _apply_lora_delta(nxt["state_dict"], sign=+1)
    print(f"[lora] fused {target_name} ({n} tensors)")
    _current_lora = target_name


def _fetch_lora_files() -> dict[str, str]:
    """hf_hub_download each LoRA from LORA_REPO; return {name: local_path}."""
    out = {}
    for name, fn in LORA_FILES.items():
        print(f"[lora] fetching {LORA_REPO}/{fn} …")
        out[name] = hf_hub_download(repo_id=LORA_REPO, filename=fn)
    return out


def _load_pipe():
    global _pipe

    # On HF Spaces we MUST pull FLUX.2 from HuggingFace Hub (not ModelScope, the
    # diffsynth default). The `preload_from_hub` block in README.md caches these
    # files during the Space build, so this call is just a fast cache lookup.
    # Locally, $DIFFSYNTH_MODEL_BASE_PATH points at a ModelScope checkout, so we
    # keep the diffsynth default there.
    if "SPACE_ID" in os.environ or os.environ.get("DIFFSYNTH_USE_HF") == "1":
        from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
        from src.diffsynth_fov import Flux2FoveatedImagePipeline

        print("[pipe] loading FLUX2 from HuggingFace Hub …")
        _pipe = Flux2FoveatedImagePipeline.from_pretrained(
            torch_dtype=torch.bfloat16,
            device="cuda" if torch.cuda.is_available() else "cpu",
            model_configs=[
                ModelConfig(model_id=MODEL_ID, origin_file_pattern="transformer/*.safetensors", download_source="huggingface"),
                ModelConfig(model_id=MODEL_ID, origin_file_pattern="text_encoder/*.safetensors", download_source="huggingface"),
                ModelConfig(model_id=MODEL_ID, origin_file_pattern="vae/diffusion_pytorch_model.safetensors", download_source="huggingface"),
            ],
            tokenizer_config=ModelConfig(model_id=MODEL_ID, origin_file_pattern="tokenizer/", download_source="huggingface"),
        )
    else:
        loader_args = SimpleNamespace(
            model_id=MODEL_ID,
            lora_checkpoint=None,
            dit_checkpoint=None,
            experiment="ours",
        )
        _pipe = load_pipeline(loader_args)

    print(f"[pipe] loaded on {_pipe.device}")

    paths = _fetch_lora_files()
    for name in LORA_FILES:  # preserve declared order
        _lora_registry[name] = {"path": paths[name], "state_dict": None}
        _lora_order.append(name)

    # Preload + convert all LoRA state dicts on CPU at startup. They stay in
    # parent memory and are visible (copy-on-write) to every `@spaces.GPU`
    # worker, so a cold worker doesn't pay disk-read + convert cost inside its
    # GPU window. ~hundreds of MB per LoRA in CPU RAM — fine.
    print("[lora] preloading state dicts on CPU …")
    for name in LORA_FILES:
        raw = load_state_dict(_lora_registry[name]["path"], torch_dtype=_pipe.torch_dtype)
        _lora_registry[name]["state_dict"] = _convert_lora_sd(raw)
    print("[lora] preload done")

    # NOTE: the actual fuse (`_apply_lora_delta`, which does `torch.mm` on
    # CUDA) is deferred to the first `@spaces.GPU` `generate` call. On
    # ZeroGPU the GPU is only attached inside @spaces.GPU-decorated
    # functions, AND those run in forked workers — so any fuse done outside
    # them either fails (no GPU) or doesn't propagate back to the parent.


def _apply_selected_lora(lora_name: str):
    """Reconcile the worker's fused LoRA with `lora_name`.

    Must be called from inside an `@spaces.GPU` function. `lora_name` is
    passed in (from the dropdown via Gradio inputs) rather than read from a
    global — globals don't reliably cross the @spaces.GPU process boundary
    on ZeroGPU."""
    if _current_lora == lora_name:
        return
    with _lora_lock:
        if _current_lora == lora_name:  # double-checked
            return
        _switch_lora_unlocked(lora_name)


# ---------------------------------------------------------------------------
# Mask construction (ported from webgui/server.py)
# ---------------------------------------------------------------------------

def _quantize_to_lr_blocks(raw_mask: torch.Tensor):
    if raw_mask.shape != (HEIGHT, WIDTH):
        raw_mask = F.interpolate(
            raw_mask.unsqueeze(0).unsqueeze(0), size=(HEIGHT, WIDTH), mode="nearest",
        ).squeeze(0).squeeze(0)

    lr_h = HEIGHT // 16 // LR_DOWNSAMPLE_FACTOR
    lr_w = WIDTH // 16 // LR_DOWNSAMPLE_FACTOR
    lr_mask = F.adaptive_max_pool2d(
        (raw_mask > 0).float().unsqueeze(0).unsqueeze(0), (lr_h, lr_w),
    )
    token_grid = F.interpolate(
        lr_mask, size=(HEIGHT // 16, WIDTH // 16), mode="nearest",
    ).squeeze(0).squeeze(0)
    full_res = F.interpolate(
        lr_mask, size=(HEIGHT, WIDTH), mode="nearest",
    ).squeeze(0).squeeze(0)
    return token_grid, full_res


def _drawn_mask_to_tensors(rgba: np.ndarray):
    """ImageEditor output -> (token_grid_mask, full_res_mask). `rgba` is HxWx4 uint8."""
    if rgba.ndim == 3 and rgba.shape[2] == 4:
        alpha = rgba[..., 3]
    elif rgba.ndim == 3:
        alpha = rgba.max(axis=2)
    else:
        alpha = rgba
    raw = torch.from_numpy((alpha > 0).astype(np.float32)).to(MASK_DEVICE)
    return _quantize_to_lr_blocks(raw)


def _alpha_union_from_layers(layers) -> np.ndarray | None:
    """OR together alpha channels of all layers — gives painted region in RGBA layers
    even when the composite is opaque."""
    if not layers:
        return None
    out = None
    for layer in layers:
        arr = np.asarray(layer)
        if arr.ndim != 3 or arr.shape[2] != 4 or arr.size == 0:
            continue
        a = arr[..., 3]
        out = a if out is None else np.maximum(out, a)
    return out


def _preset_mask_to_tensors(cx: float, cy: float, r: float):
    full_res = create_foveation_mask_full_res(
        HEIGHT, WIDTH, (cx, cy), r, "circular", MASK_DEVICE,
    )
    return _quantize_to_lr_blocks(full_res)


def _empty_mask_tensors():
    token = torch.zeros(HEIGHT // 16, WIDTH // 16, device=MASK_DEVICE, dtype=torch.float32)
    full = torch.zeros(HEIGHT, WIDTH, device=MASK_DEVICE, dtype=torch.float32)
    return token, full


def _full_hr_mask_tensors():
    """All-HR mask — every token is foveal. Used with the no-fov LoRA."""
    token = torch.ones(HEIGHT // 16, WIDTH // 16, device=MASK_DEVICE, dtype=torch.float32)
    full = torch.ones(HEIGHT, WIDTH, device=MASK_DEVICE, dtype=torch.float32)
    return token, full


def _build_mask(lora_name: str, mode: str, editor_value, cx: float, cy: float, r: float):
    """Return (token_grid_mask, full_res_mask, is_empty).

    `lora_name` is passed explicitly (from the Gradio dropdown component, not
    a global) so the call works across the @spaces.GPU process boundary even
    if ZeroGPU uses spawn semantics."""
    # The no-fov LoRA was trained without foveation conditioning — it expects
    # a full-HR token mask and ignores the circle controls entirely.
    if lora_name == "no_fov":
        t, f = _full_hr_mask_tensors()
        return t, f, False
    if mode == "Preset (circle)":
        if r <= 0:
            return (*_empty_mask_tensors(), True)
        t, f = _preset_mask_to_tensors(float(cx), float(cy), float(r))
        return t, f, bool(f.sum().item() == 0)

    # Drawn. ImageEditor returns {"background": ..., "layers": [...], "composite": ...}
    # We MUST use the layers (not the composite): Gradio's composite has alpha=255
    # everywhere (opaque background), so `alpha > 0` matches every pixel and the
    # whole image registers as HR. The layer RGBA arrays preserve true paint alpha.
    if editor_value is None:
        return (*_empty_mask_tensors(), True)
    if isinstance(editor_value, dict):
        layers = editor_value.get("layers") or []
        alpha = _alpha_union_from_layers(layers)
        if alpha is None:
            # No usable layers; fall back to composite alpha (rare).
            composite = editor_value.get("composite")
            if composite is None:
                return (*_empty_mask_tensors(), True)
            arr = np.asarray(composite)
            if arr.size == 0:
                return (*_empty_mask_tensors(), True)
            t, f = _drawn_mask_to_tensors(arr)
            return t, f, bool(f.sum().item() == 0)
        raw = torch.from_numpy((alpha > 0).astype(np.float32)).to(MASK_DEVICE)
        t, f = _quantize_to_lr_blocks(raw)
        return t, f, bool(f.sum().item() == 0)
    # Bare numpy array — treat as composite-like.
    arr = np.asarray(editor_value)
    if arr.size == 0:
        return (*_empty_mask_tensors(), True)
    t, f = _drawn_mask_to_tensors(arr)
    return t, f, bool(f.sum().item() == 0)


def _tokenization_vis_image(token_mask) -> np.ndarray:
    return create_tokenization_mask_vis(token_mask, HEIGHT, WIDTH, lr_factor=LR_DOWNSAMPLE_FACTOR)


# ---------------------------------------------------------------------------
# Gradio handlers
# ---------------------------------------------------------------------------

def refresh_tokenization(lora_name, mode, editor_value, cx, cy, r):
    print(f"[tok] enter  lora={lora_name!r}  mode={mode!r}  cx={cx} cy={cy} r={r}  editor={type(editor_value).__name__}", flush=True)
    if _pipe is None:
        print("[tok] pipe is None, returning", flush=True)
        return None
    try:
        t0 = time.time()
        token_mask, _, _ = _build_mask(lora_name, mode, editor_value, cx, cy, r)
        print(f"[tok] mask built ({time.time() - t0:.3f}s)", flush=True)
        t1 = time.time()
        vis = _tokenization_vis_image(token_mask)
        print(f"[tok] vis built ({time.time() - t1:.3f}s)  shape={getattr(vis, 'shape', None)}", flush=True)
        return vis
    except Exception as e:
        import traceback
        print(f"[tok] EXCEPTION: {e!r}", flush=True)
        traceback.print_exc()
        return None


def on_mode_change(lora_name, mode, editor_value, cx, cy, r):
    """Switch active mode — refresh tokenization + swap CSS active/inactive
    classes on the two group wrappers.

    We deliberately do NOT toggle visibility of the draw/preset groups: hiding
    a Group via `gr.update(visible=False)` in Gradio 6.14 unmounts its children,
    and re-showing it later does not re-bind their event listeners — slider
    .change/.release would silently stop firing. Both groups stay mounted; CSS
    dims the inactive one via the `mode-active` / `mode-inactive` classes.
    """
    is_draw = (mode == "Draw")
    draw_cls = ["mode-group", "mode-active" if is_draw else "mode-inactive"]
    preset_cls = ["mode-group", "mode-inactive" if is_draw else "mode-active"]
    tok = refresh_tokenization(lora_name, mode, editor_value, cx, cy, r)
    return gr.update(elem_classes=draw_cls), gr.update(elem_classes=preset_cls), tok


def switch_lora(name: str, cx: float, cy: float, r: float):
    """Record the user's LoRA selection and refresh dependent UI.

    NOT decorated with `@spaces.GPU`: the actual fuse can't happen here
    because @spaces.GPU runs in a forked worker and its mutations to the
    model + globals don't propagate back to the parent. Instead, we just
    update `_selected_lora` here and let the next `generate` worker fuse.

    UI updates: dropdown value, description, status, slider interactivity
    (no_fov disables circle controls), tokenization preview.
    """
    global _selected_lora
    if name not in LORA_FILES:
        return (
            gr.update(value=_selected_lora),
            LORA_DESCRIPTIONS.get(_selected_lora, ""),
            f"unknown LoRA: {name}",
            gr.update(),
            gr.update(),
            gr.update(),
            None,
        )
    _selected_lora = name

    is_no_fov = (name == "no_fov")
    # Pass `name` explicitly so the no_fov full-HR short-circuit triggers
    # for the preview the user sees right after switching.
    token_mask, _, _ = _build_mask(name, "Preset (circle)", None, cx, cy, r)
    tok_vis = _tokenization_vis_image(token_mask)
    # interactive=False + CSS `pointer-events: none` so the handles can't fire
    # `.change`/`.release` even if the user clicks them. (Belt-and-suspenders:
    # the `_build_mask` short-circuit above already forces full-HR for no_fov,
    # but this prevents the sliders from looking active.)
    slider_classes = ["slider-locked"] if is_no_fov else []
    def _slider_update():
        return gr.update(interactive=not is_no_fov, elem_classes=slider_classes)
    return (
        gr.update(value=name),
        LORA_DESCRIPTIONS.get(name, ""),
        f"Selected: {name} (fuses on next generate)",
        _slider_update(),
        _slider_update(),
        _slider_update(),
        tok_vis,
    )


def toggle_mode(mode):
    is_draw = (mode == "Draw")
    return gr.update(visible=is_draw), gr.update(visible=not is_draw)


def fill_all_editor():
    """Paint the entire canvas with opaque black so the whole image becomes HR."""
    layer = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
    layer[..., 0:3] = 17
    layer[..., 3] = 255
    return {"background": None, "layers": [layer], "composite": layer}


def clear_editor():
    return {"background": None, "layers": [], "composite": None}


def clear_preset():
    return 0.0, 0.0, 0.0


@spaces.GPU(duration=120)
def generate(prompt, seed, lora_name, mode, editor_value, cx, cy, r,
             progress=gr.Progress(track_tqdm=False)):
    print(f"[gen] CLICK REACHED  lora={lora_name!r}  mode={mode!r}  seed={seed}", flush=True)
    if _pipe is None:
        raise gr.Error("Pipeline not loaded yet — please wait for cold start.")
    if lora_name not in LORA_FILES:
        raise gr.Error(f"unknown LoRA: {lora_name}")
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Prompt is required.")
    seed = int(seed or 0)

    # `lora_name` came in as a Gradio input (the live dropdown value), so it
    # crosses the @spaces.GPU process boundary safely — no global-state
    # propagation required.
    _apply_selected_lora(lora_name)

    progress(0.0, desc="Building mask")
    t_mask = time.time()
    token_mask, full_res_mask, is_empty = _build_mask(lora_name, mode, editor_value, cx, cy, r)
    print(f"[gen] mask built ({time.time() - t_mask:.2f}s)  is_empty={is_empty}", flush=True)
    if is_empty:
        raise gr.Error("Mask is empty — paint a region or set a circle radius > 0.")

    tok_vis = _tokenization_vis_image(token_mask)
    print("[gen] tokenization vis ready -> yielding to UI", flush=True)
    # Yield tok_vis right away so the tokenization panel updates *before* denoise.
    # gr.update() keeps out_img unchanged at this stage.
    yield tok_vis, gr.update()

    print("[gen] waiting for _job_lock", flush=True)
    with _job_lock:
        print("[gen] _job_lock acquired; calling _pipe", flush=True)
        torch.cuda.empty_cache()

        # Move CPU-built masks onto the pipeline's device for the forward.
        token_mask = token_mask.to(_pipe.device)
        full_res_mask = full_res_mask.to(_pipe.device)

        def progress_wrap(iterable):
            iterable = list(iterable)
            total = len(iterable) or NUM_INFERENCE_STEPS
            t_prev = time.time()
            for i, item in enumerate(iterable):
                yield item
                t_now = time.time()
                print(f"[gen] step {i + 1}/{total}  {t_now - t_prev:.3f}s", flush=True)
                t_prev = t_now
                progress((i + 1) / total, desc=f"Denoising {i + 1}/{total}")

        t0 = time.time()
        image = _pipe(
            prompt=prompt,
            height=HEIGHT,
            width=WIDTH,
            seed=seed,
            rand_device="cuda" if torch.cuda.is_available() else "cpu",
            num_inference_steps=NUM_INFERENCE_STEPS,
            cfg_scale=GUIDANCE_SCALE,
            foveation_mask=token_mask,
            full_res_foveation_mask=full_res_mask,
            decode_mode=DECODE_MODE,
            prediction_type=PREDICTION_TYPE,
            soft_foveation_blend=SOFT_FOVEATION_BLEND,
            lr_downsample_factor=LR_DOWNSAMPLE_FACTOR,
            progress_bar_cmd=progress_wrap,
        )
        print(f"[gen] total pipeline call: {time.time() - t0:.2f}s", flush=True)

    yield tok_vis, image


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

# Short description shown under the LoRA dropdown.
LORA_DESCRIPTIONS = {
    "no_fov":   "Finetuned baseline — **no foveation conditioning**. Always generates a uniform full-resolution image (mask controls disabled).",
    "random":   "Finetuned with **random gaze locations** during training. Most general-purpose foveation LoRA.",
    "saliency": "Finetuned with **saliency-driven** foveation masks (model attends to visually salient regions).",
}


DEFAULT_PROMPT = (
    "Plitvice Lakes cascade through a series of sixteen terraced turquoise pools "
    "connected by hundreds of waterfalls in a forested Croatian valley. The "
    "travertine barriers between each pool, the moss and the mineral deposits, "
    "and the submerged fallen logs in the crystal-clear water are all equally "
    "sharp. The surrounding beech and fir forest is in full autumn colour—gold, "
    "russet, and dark green—and the waterfalls range from wide curtains to thin "
    "threads, all frozen sharp in the even overcast light. The water's colour "
    "shifts from pale aquamarine in the shallows to deep teal in the pools."
)

CSS = """
.tok-img img, .out-img img { background: #f4f4f4 !important; }
/* Belt-and-suspenders lock for cx/cy/r when no_fov is active: gr.update
   (interactive=False) sets the disabled attribute, this kills pointer events
   so a stray click on the handle can't fire .change/.release. */
.slider-locked, .slider-locked * { pointer-events: none !important; }
.slider-locked { opacity: 0.35 !important; }
"""


def build_ui() -> gr.Blocks:
    with gr.Blocks(title="Foveated Diffusion") as demo:
        gr.Markdown(
            "## Foveated Diffusion: Efficient Spatially Aware Image and Video Generation  \n"
            "Image generation demo for Foveated Diffusion. The base model is <a href=\"https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B\" target=\"_blank\" rel=\"noopener noreferrer\">FLUX.2-klein-base-4B</a>. \n"
            "   \n"
            '<a href="https://bchao1.github.io/foveated-diffusion/" target="_blank" rel="noopener noreferrer">Project website</a> · '
            '<a href="https://huggingface.co/bchao1/foveated_diffusion" target="_blank" rel="noopener noreferrer">Model weights</a> · '
            '<a href="https://arxiv.org/abs/2603.23491" target="_blank" rel="noopener noreferrer">Paper</a> · '
            '<a href="https://github.com/bchao1/foveated_diffusion" target="_blank" rel="noopener noreferrer">Code</a>'
        )

        with gr.Row():
            # ---------------- Left column: controls ----------------
            with gr.Column(scale=1, min_width=340):
                prompt = gr.Textbox(label="Prompt", lines=6, value=DEFAULT_PROMPT)
                seed = gr.Number(label="Seed", value=0, precision=0)

                lora_dd = gr.Dropdown(
                    label="LoRA model",
                    choices=list(LORA_FILES.keys()),
                    value=DEFAULT_LORA,
                    interactive=True,
                )
                lora_desc = gr.Markdown(
                    LORA_DESCRIPTIONS[DEFAULT_LORA],
                    elem_classes=["hint"],
                )
                lora_status = gr.Markdown(f"Active: {DEFAULT_LORA}")

                # --- Mask mode ---
                # The "Draw" path (paint a freeform HR region with gr.ImageEditor)
                # is implemented end-to-end in this file but intentionally NOT
                # surfaced in the UI for this Spaces demo. The radio and the
                # ImageEditor are kept alive as hidden components (visible=False)
                # so the wiring + handler logic compiles unchanged — re-enable
                # them by flipping visible=True if a future demo wants both modes.
                mode = gr.Radio(
                    choices=["Draw", "Preset (circle)"],
                    value="Preset (circle)",
                    visible=False,
                )

                with gr.Group(visible=False) as draw_group:
                    editor = gr.ImageEditor(
                        image_mode="RGBA",
                        type="numpy",
                        sources=(),
                        canvas_size=(1024, 1024),
                        fixed_canvas=True,
                        layers=False,
                        transforms=(),
                        brush=gr.Brush(
                            default_size=80,
                            colors=["#111111"],
                            color_mode="fixed",
                            default_color="#111111",
                        ),
                        eraser=gr.Eraser(default_size=80),
                        height=420,
                    )
                    with gr.Row():
                        clear_btn = gr.Button("Clear", variant="secondary")
                        fill_btn = gr.Button("Fill all", variant="secondary")

                # Preset (circle) is the only mode exposed in the UI.
                with gr.Group() as preset_group:
                    gr.Markdown(
                        "**Foveal region (circular)** — place a circular HR region "
                        "on the image. `cx, cy ∈ [-0.5, 0.5]` (0 = image center); "
                        "`r` is relative to half the image diagonal.",
                        elem_classes=["hint"],
                    )
                    cx = gr.Slider(-0.5, 0.5, value=0.0, step=0.01, label="Center X")
                    cy = gr.Slider(-0.5, 0.5, value=0.0, step=0.01, label="Center Y")
                    r = gr.Slider(0.0, 1.0, value=0.3, step=0.01, label="Radius")
                    preset_clear_btn = gr.Button("Clear", variant="secondary")

                generate_btn = gr.Button("Generate", variant="primary")

            # ---------------- Right column: stages ----------------
            with gr.Column(scale=2):
                with gr.Row():
                    tok_img = gr.Image(
                        label="Tokenization mask",
                        height=520,
                        interactive=False,
                        elem_classes=["tok-img"],
                    )
                    out_img = gr.Image(
                        label="Generated image",
                        height=520,
                        interactive=False,
                        elem_classes=["out-img"],
                    )
                gr.Markdown(
                    "**Tokenization mask** — white = HR (foveal) tokens, gray = LR "
                    "(peripheral). Refreshes as you move the circle sliders. "
                    "**Generated image** — appears once denoising finishes."
                )

        # ---------------- Wiring ----------------
        # lora_dd is FIRST so every handler that reads `refresh_in` (and
        # `generate`) gets the live dropdown value as an explicit argument —
        # the only safe way to ferry it into @spaces.GPU workers.
        refresh_in = [lora_dd, mode, editor, cx, cy, r]

        # Mode change: visibility toggle + tokenization refresh in ONE handler.
        # Registering two separate listeners on the same event in Gradio 6.x
        # was unreliable; a single combined handler is the safe pattern.
        mode.change(
            on_mode_change, refresh_in, [draw_group, preset_group, tok_img], queue=False,
        )

        # LoRA switch: also disables circle sliders and refreshes tok preview
        # so no_fov immediately reflects "full-HR, controls disabled".
        lora_dd.change(
            switch_lora,
            [lora_dd, cx, cy, r],
            [lora_dd, lora_desc, lora_status, cx, cy, r, tok_img],
        )

        # Live tokenization preview on edits / slider moves.
        # queue=False so brush/slider activity isn't serialized behind a running
        # generation in the queue worker. We wire BOTH .release and .change on
        # sliders — .release alone can miss events on some Gradio 6.x builds
        # (especially for sliders inside a Group that was initially hidden).
        editor.apply(refresh_tokenization, refresh_in, tok_img, queue=False)
        for trig in (cx.release, cy.release, r.release,
                     cx.change, cy.change, r.change):
            trig(refresh_tokenization, refresh_in, tok_img, queue=False)

        # Draw helpers. queue=False — pure value-replacement, no GPU.
        clear_btn.click(clear_editor, None, editor, queue=False)
        fill_btn.click(fill_all_editor, None, editor, queue=False)
        preset_clear_btn.click(clear_preset, None, [cx, cy, r], queue=False)

        # Generate.
        # show_progress_on=[out_img] keeps the denoise progress bar from
        # overlaying the tokenization panel.
        generate_btn.click(
            generate,
            inputs=[prompt, seed, lora_dd, mode, editor, cx, cy, r],
            outputs=[tok_img, out_img],
            show_progress_on=[out_img],
        )

        # On page load: render the initial tokenization grid for the default
        # circle (cx=0, cy=0, r=0.3). Mode is fixed to "Preset (circle)" — the
        # Draw mode is kept in code (visible=False) but not wired into load.
        demo.load(refresh_tokenization, refresh_in, tok_img, queue=False)

    return demo


if __name__ == "__main__":
    print("Loading pipeline and LoRAs (cold start can take several minutes)…")
    _load_pipe()
    demo = build_ui()
    # Spaces sets PORT=7860 by convention.
    demo.queue(max_size=8).launch(
        server_name=os.environ.get("HOST", "0.0.0.0"),
        server_port=int(os.environ.get("PORT", 7860)),
        theme=gr.themes.Default(primary_hue="slate"),
        css=CSS,
        # SHARE=1 -> create a *.gradio.live tunnel. Use this to bypass flaky
        # local port-forwarders (e.g. VS Code) that mangle Gradio's SSE/WS.
        share=os.environ.get("SHARE", "").lower() in ("1", "true", "yes"),
    )