"""NFA Track R — Fun depth ControlNet ZeroGPU (VideoX-Fun / ALIMAMA).

REAL Fun ControlNet Union depth — NOT soft Flux2 image=depth (banned forever).

Host-RAM fix (2026-07-21):
  Prior hang: bf16 from_pretrained DiT+TE (~112GB) then post-hoc qfloat8.
  Now: stream DiT shards straight into float8 (never full bf16 materialize),
  local quantized TE (bnb-4bit / fp8-class — NO HF remote TE),
  ZeroGPU size=xlarge (96GB). Abort early if Fun CN not loaded in time.
"""

from __future__ import annotations

import gc
import glob
import json
import os
import shutil
import subprocess
import sys
import time
import traceback
from pathlib import Path
from typing import Optional

import gradio as gr
import spaces
import torch
from huggingface_hub import hf_hub_download, login, snapshot_download
from omegaconf import OmegaConf
from PIL import Image

APP_DIR = Path(__file__).resolve().parent
_VX = APP_DIR / "vendor" / "VideoX-Fun"
_VX_CACHE = Path.home() / "VideoX-Fun"

# Exclude from float8 (VideoX Fun CN + embedding / timestep stability).
_FP8_EXCLUDE = ("img_in", "txt_in", "timestep", "control_img_in", "embed")


def _patch_videox_inits(root: Path) -> None:
    models_init = root / "videox_fun" / "models" / "__init__.py"
    if models_init.is_file():
        text = models_init.read_text(encoding="utf-8", errors="replace")
        if not (
            "Flux2ControlTransformer2DModel" in text
            and "fantasytalking" not in text
            and "AutoProcessor" in text
        ):
            models_init.write_text(
                "from transformers import (\n"
                "    AutoProcessor,\n"
                "    Mistral3ForConditionalGeneration,\n"
                "    PixtralProcessor,\n"
                ")\n"
                "from .flux2_image_processor import Flux2ImageProcessor\n"
                "from .flux2_transformer2d import Flux2Transformer2DModel\n"
                "from .flux2_transformer2d_control import Flux2ControlTransformer2DModel\n"
                "from .flux2_vae import AutoencoderKLFlux2\n"
                "__all__ = [\n"
                "    'AutoProcessor',\n"
                "    'AutoencoderKLFlux2',\n"
                "    'Flux2ControlTransformer2DModel',\n"
                "    'Flux2ImageProcessor',\n"
                "    'Flux2Transformer2DModel',\n"
                "    'Mistral3ForConditionalGeneration',\n"
                "    'PixtralProcessor',\n"
                "]\n",
                encoding="utf-8",
            )
            print("[nfa-fun-cn] patched videox_fun.models.__init__", flush=True)

    pipe_init = root / "videox_fun" / "pipeline" / "__init__.py"
    if pipe_init.is_file():
        text = pipe_init.read_text(encoding="utf-8", errors="replace")
        if "pipeline_cogvideox" in text or "Flux2ControlPipeline" not in text:
            pipe_init.write_text(
                "from .pipeline_flux2_control import Flux2ControlPipeline\n"
                "__all__ = ['Flux2ControlPipeline']\n",
                encoding="utf-8",
            )
            print("[nfa-fun-cn] patched videox_fun.pipeline.__init__", flush=True)


def _ensure_videox_on_path() -> None:
    for candidate in (_VX, _VX_CACHE):
        if (candidate / "videox_fun" / "models").is_dir():
            _patch_videox_inits(candidate)
            p = str(candidate)
            if p not in sys.path:
                sys.path.insert(0, p)
            return
    print("[nfa-fun-cn] cloning VideoX-Fun for Fun CN runtime…", flush=True)
    if _VX_CACHE.exists():
        shutil.rmtree(_VX_CACHE, ignore_errors=True)
    subprocess.check_call(
        [
            "git",
            "clone",
            "--depth",
            "1",
            "https://github.com/aigc-apps/VideoX-Fun.git",
            str(_VX_CACHE),
        ]
    )
    _patch_videox_inits(_VX_CACHE)
    sys.path.insert(0, str(_VX_CACHE))


_ensure_videox_on_path()

HF_TOKEN = (
    os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") or ""
).strip()
if HF_TOKEN:
    try:
        login(token=HF_TOKEN, add_to_git_credential=False)
    except Exception as exc:  # noqa: BLE001
        print(f"[nfa-fun-cn] HF login warning: {exc}", flush=True)

BASE_MODEL = os.environ.get(
    "NFA_FLUX2_MODEL_ID", "black-forest-labs/FLUX.2-dev"
).strip()
# Pre-quantized local TE (bnb-4bit). Remote TE is banned (HF endpoint broken).
TE_MODEL = os.environ.get(
    "NFA_FLUX2_TE_MODEL_ID", "diffusers/FLUX.2-dev-bnb-4bit"
).strip()
CN_REPO = os.environ.get(
    "NFA_FUN_CN_REPO", "alibaba-pai/FLUX.2-dev-Fun-Controlnet-Union"
).strip()
CN_FILE = os.environ.get(
    "NFA_FUN_CN_FILE", "FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors"
).strip()
# Director lock: 96GB xlarge + Q8/fp8 painter (not soft; not bf16 host dump).
GPU_SIZE = (os.environ.get("NFA_FUN_CN_GPU_SIZE") or "xlarge").strip().lower()
if GPU_SIZE not in ("large", "xlarge"):
    GPU_SIZE = "xlarge"
GPU_DURATION = int(os.environ.get("NFA_FUN_CN_GPU_DURATION") or "300")
WEIGHT_DTYPE = torch.bfloat16
MEM_MODE = (
    os.environ.get("NFA_FUN_CN_MEM_MODE") or "model_cpu_offload_and_qfloat8"
).strip()
# Abort CPU load if Fun CN not ready — do not thrash 35+ min again.
LOAD_DEADLINE_SEC = int(os.environ.get("NFA_FUN_CN_LOAD_DEADLINE_SEC") or "480")

CONFIG_PATH = APP_DIR / "config" / "flux2_control.yaml"
MODEL_DIR = Path(os.environ.get("NFA_FLUX2_MOUNT") or "/data/FLUX.2-dev")
CN_MOUNT_DIR = Path(os.environ.get("NFA_FUN_CN_MOUNT") or "/data/Fun-CN")
TE_MOUNT_DIR = Path(os.environ.get("NFA_FLUX2_TE_MOUNT") or "/data/FLUX.2-TE-bnb4")
CACHE_ROOT = Path(
    os.environ.get("NFA_FUN_CN_CACHE") or (Path.home() / ".cache" / "nfa_fun_cn")
)

_PIPE = None
_PIPE_OFFLOAD_READY = False
_CN_FILE_PATH: Path | None = None
_GET_IMAGE_LATENT = None
_FUN_CN_LOADED = False
_LOAD_T0: float | None = None


def _rss_gb() -> float:
    try:
        import resource

        # Linux: ru_maxrss is KB
        return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / (1024 * 1024)
    except Exception:  # noqa: BLE001
        return -1.0


def _check_load_deadline(stage: str) -> None:
    if _LOAD_T0 is None:
        return
    elapsed = time.time() - _LOAD_T0
    if elapsed > LOAD_DEADLINE_SEC and not _FUN_CN_LOADED:
        raise RuntimeError(
            f"FUN_CN_LOAD_ABORT: stage={stage} elapsed={elapsed:.0f}s "
            f"> deadline={LOAD_DEADLINE_SEC}s (never reached Fun CN loaded). "
            "Refusing to thrash ZeroGPU host RAM like the prior bf16 hang."
        )


def _resolve_cn_path() -> Path:
    global _CN_FILE_PATH
    if _CN_FILE_PATH is not None and _CN_FILE_PATH.is_file():
        return _CN_FILE_PATH
    candidates = [CN_MOUNT_DIR / CN_FILE, CACHE_ROOT / CN_FILE]
    if CN_MOUNT_DIR.is_dir():
        candidates.extend(CN_MOUNT_DIR.rglob(CN_FILE))
    if CACHE_ROOT.is_dir():
        candidates.extend(CACHE_ROOT.rglob(CN_FILE))
    for c in candidates:
        if c.is_file():
            _CN_FILE_PATH = c
            return c
    raise FileNotFoundError(
        f"Fun CN weights missing: {CN_FILE}. "
        "Mount alibaba-pai/FLUX.2-dev-Fun-Controlnet-Union at /data/Fun-CN."
    )


def _ensure_weights() -> None:
    global MODEL_DIR, _CN_FILE_PATH
    if MODEL_DIR.is_dir() and (MODEL_DIR / "model_index.json").is_file():
        try:
            cn = _resolve_cn_path()
            print(f"[nfa-fun-cn] using mounts model={MODEL_DIR} cn={cn}", flush=True)
            return
        except FileNotFoundError:
            pass
    print("[nfa-fun-cn] WARN mounts missing; selective download fallback", flush=True)
    cache_model = CACHE_ROOT / "FLUX.2-dev"
    CACHE_ROOT.mkdir(parents=True, exist_ok=True)
    token = HF_TOKEN or None
    snapshot_download(
        repo_id=BASE_MODEL,
        local_dir=str(cache_model),
        token=token,
        allow_patterns=[
            "model_index.json",
            "transformer/*",
            "vae/*",
            "tokenizer/*",
            "scheduler/*",
        ],
    )
    path = hf_hub_download(
        repo_id=CN_REPO, filename=CN_FILE, local_dir=str(CACHE_ROOT), token=token
    )
    MODEL_DIR = cache_model
    _CN_FILE_PATH = Path(path)


def _ensure_te_weights() -> Path:
    """Local quantized TE only — never remote TE."""
    if TE_MOUNT_DIR.is_dir() and (
        (TE_MOUNT_DIR / "text_encoder").is_dir()
        or (TE_MOUNT_DIR / "config.json").is_file()
    ):
        print(f"[nfa-fun-cn] TE mount={TE_MOUNT_DIR}", flush=True)
        return TE_MOUNT_DIR
    cache_te = CACHE_ROOT / "FLUX.2-TE-bnb4"
    if (cache_te / "text_encoder").is_dir() or (cache_te / "model_index.json").is_file():
        print(f"[nfa-fun-cn] TE cache={cache_te}", flush=True)
        return cache_te
    print(
        f"[nfa-fun-cn] downloading local quantized TE from {TE_MODEL} "
        "(NO remote TE)",
        flush=True,
    )
    CACHE_ROOT.mkdir(parents=True, exist_ok=True)
    snapshot_download(
        repo_id=TE_MODEL,
        local_dir=str(cache_te),
        token=HF_TOKEN or None,
        allow_patterns=[
            "model_index.json",
            "text_encoder/*",
            "tokenizer/*",
        ],
    )
    return cache_te


def _prep_depth(depth_image: Image.Image, width: int, height: int) -> Image.Image:
    img = depth_image.convert("RGB")
    if img.size != (width, height):
        img = img.resize((width, height), Image.Resampling.LANCZOS)
    return img


def _compose_prompt(positive: str, negative: str) -> tuple[str, str]:
    return (positive or "").strip(), ((negative or "").strip() or " ")


def _fp8_dtype_for_key(key: str) -> torch.dtype:
    for ex in _FP8_EXCLUDE:
        if ex in key:
            return WEIGHT_DTYPE
    return torch.float8_e4m3fn


def _set_tensor(model: torch.nn.Module, key: str, tensor: torch.Tensor) -> None:
    from accelerate.utils import set_module_tensor_to_device

    target = _fp8_dtype_for_key(key)
    # float8 cast must go through a float dtype first on some builds
    if target == torch.float8_e4m3fn and tensor.dtype not in (
        torch.float8_e4m3fn,
        torch.float8_e5m2,
    ):
        value = tensor.detach().to(dtype=torch.bfloat16).to(dtype=target)
    else:
        value = tensor.detach().to(dtype=target)
    set_module_tensor_to_device(model, key, device="cpu", value=value, dtype=target)


def _stream_shards_into_model(
    model: torch.nn.Module,
    shard_paths: list[str],
    *,
    label: str,
) -> None:
    """Load one safetensors shard at a time → fp8; never accumulate full bf16."""
    from safetensors import safe_open

    model_sd = model.state_dict()
    loaded = 0
    skipped = 0
    for i, path in enumerate(shard_paths):
        _check_load_deadline(f"{label}_shard_{i}")
        print(
            f"[nfa-fun-cn] {label} shard {i+1}/{len(shard_paths)} "
            f"rss_max≈{_rss_gb():.1f}GB path={Path(path).name}",
            flush=True,
        )
        with safe_open(path, framework="pt", device="cpu") as f:
            for key in f.keys():
                if key not in model_sd:
                    skipped += 1
                    continue
                tensor = f.get_tensor(key)
                if tuple(tensor.shape) != tuple(model_sd[key].shape):
                    print(
                        f"[nfa-fun-cn] skip size mismatch {key} "
                        f"{tuple(tensor.shape)} vs {tuple(model_sd[key].shape)}",
                        flush=True,
                    )
                    skipped += 1
                    continue
                _set_tensor(model, key, tensor)
                loaded += 1
                del tensor
        gc.collect()
    print(
        f"[nfa-fun-cn] {label} stream done loaded={loaded} skipped={skipped} "
        f"rss_max≈{_rss_gb():.1f}GB",
        flush=True,
    )


def _init_missing_control_params(model: torch.nn.Module) -> None:
    """Mirror VideoX missing-key init for control blocks (zeros / clones)."""
    from accelerate.utils import set_module_tensor_to_device

    sd = {k: v for k, v in model.named_parameters()}
    # named_parameters may still be meta; use state_dict meta shapes
    meta_sd = model.state_dict()
    missing = [
        k
        for k, v in meta_sd.items()
        if getattr(v, "is_meta", False) or (hasattr(v, "device") and v.device.type == "meta")
    ]
    if not missing:
        # Also detect uninitialized via device
        missing = []
        for name, param in model.named_parameters():
            if param.device.type == "meta":
                missing.append(name)
    if not missing:
        print("[nfa-fun-cn] no meta params left before Fun CN overlay", flush=True)
        return

    print(f"[nfa-fun-cn] init {len(missing)} missing/meta params", flush=True)
    with torch.no_grad():
        for key in missing:
            shape = tuple(meta_sd[key].shape)
            dtype = _fp8_dtype_for_key(key)
            # Prefer clone from non-control twin when VideoX does
            twin = key.replace("control_", "")
            if "control" in key and twin in sd and sd[twin].device.type != "meta":
                value = sd[twin].detach().to(dtype=torch.bfloat16).to(dtype=dtype)
            elif "after_proj" in key or "before_proj" in key:
                value = torch.zeros(shape, dtype=dtype)
            elif "bias" in key:
                value = torch.zeros(shape, dtype=dtype)
            else:
                value = torch.zeros(shape, dtype=dtype)
            set_module_tensor_to_device(model, key, device="cpu", value=value, dtype=dtype)


def _overlay_fun_cn(model: torch.nn.Module, cn_path: Path) -> tuple[int, int]:
    """Apply Fun CN Union weights without loading a second full DiT."""
    from safetensors import safe_open

    model_sd = model.state_dict()
    loaded = 0
    skipped = 0
    print(
        f"[nfa-fun-cn] Fun CN overlay {cn_path.name} "
        f"rss_max≈{_rss_gb():.1f}GB",
        flush=True,
    )
    with safe_open(str(cn_path), framework="pt", device="cpu") as f:
        keys = list(f.keys())
        if keys == ["state_dict"]:
            # rare wrapper — fall back to full load of inner dict only
            from safetensors.torch import load_file

            wrapped = load_file(str(cn_path))
            inner = wrapped.get("state_dict", wrapped)
            for key, tensor in inner.items():
                if key not in model_sd or tuple(tensor.shape) != tuple(model_sd[key].shape):
                    skipped += 1
                    continue
                _set_tensor(model, key, tensor)
                loaded += 1
            del wrapped, inner
            gc.collect()
        else:
            for key in keys:
                if key not in model_sd:
                    skipped += 1
                    continue
                tensor = f.get_tensor(key)
                if tuple(tensor.shape) != tuple(model_sd[key].shape):
                    skipped += 1
                    continue
                _set_tensor(model, key, tensor)
                loaded += 1
                del tensor
    gc.collect()
    return loaded, skipped


def _load_control_transformer_fp8(model_name: str, cn_file: str):
    """Q8-class painter: empty meta → stream bf16 shards as float8 → Fun CN."""
    global _FUN_CN_LOADED
    import accelerate
    from videox_fun.models.flux2_transformer2d_control import (
        Flux2ControlTransformer2DModel,
    )

    config_file = os.path.join(model_name, "transformer", "config.json")
    if not os.path.isfile(config_file):
        raise FileNotFoundError(config_file)
    with open(config_file, "r", encoding="utf-8") as fh:
        config = json.load(fh)
    extra = OmegaConf.to_container(OmegaConf.load(str(CONFIG_PATH))[
        "transformer_additional_kwargs"
    ])

    print(
        "[nfa-fun-cn] CPU-load Flux2Control as float8 stream "
        f"(NOT full bf16) mem={MEM_MODE} gpu_size={GPU_SIZE}",
        flush=True,
    )
    with accelerate.init_empty_weights():
        transformer = Flux2ControlTransformer2DModel.from_config(config, **extra)

    shard_dir = os.path.join(model_name, "transformer")
    shards = sorted(glob.glob(os.path.join(shard_dir, "*.safetensors")))
    if not shards:
        raise FileNotFoundError(f"No transformer shards under {shard_dir}")
    _stream_shards_into_model(transformer, shards, label="DiT-fp8")
    _check_load_deadline("after_dit_stream")
    _init_missing_control_params(transformer)
    _check_load_deadline("after_control_init")
    loaded, skipped = _overlay_fun_cn(transformer, Path(cn_file))
    _FUN_CN_LOADED = True
    print(
        f"[nfa-fun-cn] Fun CN loaded overlay_ok={loaded} skipped={skipped} "
        f"rss_max≈{_rss_gb():.1f}GB (REAL Fun CN, fp8 painter)",
        flush=True,
    )
    return transformer


def _load_local_quantized_te(te_root: Path):
    """Local quantized Mistral TE — never HF remote TE."""
    from transformers import Mistral3ForConditionalGeneration

    te_path = te_root / "text_encoder"
    if not te_path.is_dir():
        te_path = te_root
    print(
        f"[nfa-fun-cn] loading LOCAL quantized TE from {te_path} "
        "(bnb-4bit / no remote TE)",
        flush=True,
    )
    _check_load_deadline("te_start")
    text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
        str(te_path),
        torch_dtype=WEIGHT_DTYPE,
        low_cpu_mem_usage=True,
        device_map="cpu",
    )
    print(
        f"[nfa-fun-cn] local TE ready rss_max≈{_rss_gb():.1f}GB",
        flush=True,
    )
    return text_encoder


def get_pipe(*, prepare_gpu_offload: bool = False):
    """CPU-load Fun CN pipeline (fp8 DiT stream); arm GPU offload inside @spaces.GPU."""
    global _PIPE, _PIPE_OFFLOAD_READY, _GET_IMAGE_LATENT, _LOAD_T0, _FUN_CN_LOADED
    if _PIPE is None:
        _LOAD_T0 = time.time()
        _FUN_CN_LOADED = False
        _ensure_weights()
        te_root = _ensure_te_weights()
        _ensure_videox_on_path()
        from diffusers import FlowMatchEulerDiscreteScheduler
        from transformers import PixtralProcessor
        from videox_fun.models.flux2_vae import AutoencoderKLFlux2
        from videox_fun.pipeline.pipeline_flux2_control import Flux2ControlPipeline
        from videox_fun.utils.utils import get_image_latent

        _GET_IMAGE_LATENT = get_image_latent
        model_name = str(MODEL_DIR)
        cn_file = str(_resolve_cn_path())

        transformer = _load_control_transformer_fp8(model_name, cn_file)
        _check_load_deadline("post_fun_cn")

        vae = AutoencoderKLFlux2.from_pretrained(model_name, subfolder="vae").to(
            WEIGHT_DTYPE
        )
        # Tokenizer from base FLUX.2 mount (same vocab as TE).
        tokenizer = PixtralProcessor.from_pretrained(model_name, subfolder="tokenizer")
        text_encoder = _load_local_quantized_te(te_root)
        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
            model_name, subfolder="scheduler"
        )
        _PIPE = Flux2ControlPipeline(
            vae=vae,
            tokenizer=tokenizer,
            text_encoder=text_encoder,
            transformer=transformer,
            scheduler=scheduler,
        )
        elapsed = time.time() - _LOAD_T0
        print(
            f"[nfa-fun-cn] Flux2ControlPipeline CPU-ready "
            f"(REAL Fun CN + fp8 DiT + local TE) in {elapsed:.0f}s "
            f"rss_max≈{_rss_gb():.1f}GB",
            flush=True,
        )

    if prepare_gpu_offload and not _PIPE_OFFLOAD_READY and torch.cuda.is_available():
        from videox_fun.utils.fp8_optimization import convert_weight_dtype_wrapper

        device = "cuda"
        transformer = _PIPE.transformer
        # Weights already float8 from stream load — only wrap compute dtype.
        if MEM_MODE in (
            "model_cpu_offload_and_qfloat8",
            "model_full_load_and_qfloat8",
        ):
            convert_weight_dtype_wrapper(transformer, WEIGHT_DTYPE)
        if MEM_MODE == "sequential_cpu_offload":
            _PIPE.enable_sequential_cpu_offload(device=device)
        elif MEM_MODE in ("model_cpu_offload", "model_cpu_offload_and_qfloat8"):
            _PIPE.enable_model_cpu_offload(device=device)
        else:
            _PIPE.to(device=device)
        _PIPE_OFFLOAD_READY = True
        print(f"[nfa-fun-cn] GPU offload armed mem={MEM_MODE} size={GPU_SIZE}", flush=True)
    return _PIPE


@spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE)
def _generate_still_gpu(
    positive: str,
    negative: str,
    depth_image: Image.Image,
    seed: int,
    width: int,
    height: int,
    steps: int,
    guidance: float,
    cn_strength: float,
) -> Image.Image:
    if torch.cuda.is_available():
        free, total = torch.cuda.mem_get_info()
        print(
            f"[nfa-fun-cn] cuda free={free/1e9:.1f}G total={total/1e9:.1f}G "
            f"duration={GPU_DURATION} size={GPU_SIZE}",
            flush=True,
        )
    w = int(width) if width else 1216
    h = int(height) if height else 832
    w -= w % 16
    h -= h % 16
    prompt, neg = _compose_prompt(positive, negative)
    if not prompt:
        raise gr.Error("positive prompt is required")
    depth = _prep_depth(depth_image, w, h)
    pipe = get_pipe(prepare_gpu_offload=True)
    control_latent = _GET_IMAGE_LATENT(depth, sample_size=[h, w])[:, :, 0]
    inpaint_image = torch.zeros([1, 3, h, w])
    mask_image = torch.ones([1, 1, h, w]) * 255
    strength = float(cn_strength) if float(cn_strength) > 0 else 0.75
    strength = max(0.05, min(1.5, strength))
    device = "cuda" if torch.cuda.is_available() else "cpu"
    generator = torch.Generator(device=device).manual_seed(int(seed))
    print(
        f"[nfa-fun-cn] REAL Fun CN generate seed={seed} {w}x{h} steps={steps} "
        f"cn={strength} path=videox_fun_flux2_control painter=fp8 te=local_bnb4 "
        f"size={GPU_SIZE}",
        flush=True,
    )
    with torch.no_grad():
        out = pipe(
            prompt=prompt,
            negative_prompt=neg,
            height=h,
            width=w,
            generator=generator,
            guidance_scale=float(guidance),
            image=None,
            inpaint_image=inpaint_image,
            mask_image=mask_image,
            control_image=control_latent,
            num_inference_steps=int(steps),
            control_context_scale=strength,
        ).images
    return out[0]


def generate_still(
    positive: str,
    negative: str = "",
    depth_image: Optional[Image.Image] = None,
    seed: int = 42,
    width: int = 1216,
    height: int = 832,
    steps: int = 28,
    guidance: float = 4.0,
    cn_strength: float = 0.75,
) -> Image.Image:
    try:
        if depth_image is None:
            raise gr.Error(
                "FUN_CN_REQUIRES_DEPTH: depth_image is required for real Fun ControlNet."
            )
        _ensure_weights()
        # Wall-clock CPU load — does not burn ZeroGPU minutes.
        get_pipe(prepare_gpu_offload=False)
        if not _FUN_CN_LOADED:
            raise RuntimeError(
                "FUN_CN_LOAD_ABORT: pipeline built but Fun CN flag false"
            )
        return _generate_still_gpu(
            positive,
            negative or "",
            depth_image,
            int(seed),
            int(width),
            int(height),
            int(steps),
            float(guidance),
            float(cn_strength),
        )
    except gr.Error:
        raise
    except Exception as exc:  # noqa: BLE001
        tb = traceback.format_exc()
        print(tb, flush=True)
        raise gr.Error(f"{type(exc).__name__}: {exc}\n\n{tb[-2500:]}") from exc


with gr.Blocks(title="NFA Track R FLUX.2 Fun CN ZeroGPU") as demo:
    gr.Markdown(
        "## NFA Track R — **Real Fun depth ControlNet** (ZeroGPU)\n"
        f"- Stack: VideoX-Fun `Flux2ControlPipeline` + `{CN_FILE}`\n"
        f"- Painter: **float8 stream** from `{BASE_MODEL}` (no full bf16 host dump)\n"
        f"- TE: **local quantized** `{TE_MODEL}` (NO remote TE)\n"
        f"- GPU: `size={GPU_SIZE}` duration={GPU_DURATION}s mem=`{MEM_MODE}`\n"
        f"- Load deadline: {LOAD_DEADLINE_SEC}s to reach `Fun CN loaded`\n"
        "- Soft `image=depth` is **banned** on this Space."
    )
    with gr.Row():
        with gr.Column():
            positive = gr.Textbox(label="positive", lines=12)
            negative = gr.Textbox(label="negative", lines=3)
            depth_image = gr.Image(label="depth_image (required)", type="pil")
            seed = gr.Number(label="seed", value=42, precision=0)
            width = gr.Number(label="width", value=1216, precision=0)
            height = gr.Number(label="height", value=832, precision=0)
            steps = gr.Number(label="steps", value=28, precision=0)
            guidance = gr.Number(label="guidance", value=4.0)
            cn_strength = gr.Number(label="cn_strength", value=0.75)
            btn = gr.Button("Generate (Fun CN)", variant="primary")
        with gr.Column():
            still = gr.Image(label="still")
    btn.click(
        fn=generate_still,
        inputs=[
            positive,
            negative,
            depth_image,
            seed,
            width,
            height,
            steps,
            guidance,
            cn_strength,
        ],
        outputs=[still],
        api_name="generate_still",
    )

if __name__ == "__main__":
    demo.queue(max_size=4).launch()