"""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 "600")
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 "900")

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,
    *,
    device: str,
) -> None:
    """Assign one weight; prefer CUDA float8 (CPU float8 hung on ZeroGPU host)."""
    from accelerate.utils import set_module_tensor_to_device

    target = _fp8_dtype_for_key(key)
    # Always stage through bf16 for numeric stability, then target dtype on device.
    value = tensor.detach().to(dtype=WEIGHT_DTYPE)
    if target == torch.float8_e4m3fn and device.startswith("cuda"):
        value = value.to(dtype=target)
        set_module_tensor_to_device(
            model, key, device=device, value=value, dtype=target
        )
    else:
        set_module_tensor_to_device(
            model, key, device=device, value=value, dtype=WEIGHT_DTYPE
        )
        if target == torch.float8_e4m3fn:
            mod: torch.nn.Module = model
            parts = key.split(".")
            for p in parts[:-1]:
                mod = getattr(mod, p)
            leaf = parts[-1]
            param = getattr(mod, leaf)
            if isinstance(param, torch.nn.Parameter):
                param.data = param.data.to(torch.float8_e4m3fn)
    del value


def _stream_shards_into_model(
    model: torch.nn.Module,
    shard_paths: list[str],
    *,
    label: str,
    device: str,
) -> None:
    """Stream shards key-by-key (no full-shard RAM spike) into float8 on device."""
    from safetensors import safe_open

    # Shapes only — avoid pulling meta tensors repeatedly.
    shape_map = {k: tuple(v.shape) for k, v in model.state_dict().items()}
    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"device={device} rss_max≈{_rss_gb():.1f}GB path={Path(path).name}",
            flush=True,
        )
        with safe_open(path, framework="pt", device="cpu") as f:
            keys = list(f.keys())
            n_keys = len(keys)
            for j, key in enumerate(keys):
                if key not in shape_map:
                    skipped += 1
                    continue
                tensor = f.get_tensor(key)
                if tuple(tensor.shape) != shape_map[key]:
                    skipped += 1
                    del tensor
                    continue
                _set_tensor(model, key, tensor, device=device)
                loaded += 1
                del tensor
                if (j + 1) % 50 == 0 or (j + 1) == n_keys:
                    print(
                        f"[nfa-fun-cn] {label} shard {i+1} keys {j+1}/{n_keys} "
                        f"loaded={loaded} rss_max≈{_rss_gb():.1f}GB",
                        flush=True,
                    )
        gc.collect()
        if device.startswith("cuda"):
            torch.cuda.empty_cache()
    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, *, device: str) -> 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()}
    meta_sd = model.state_dict()
    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)
            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 or "bias" in key:
                value = torch.zeros(shape, dtype=dtype)
            else:
                value = torch.zeros(shape, dtype=dtype)
            set_module_tensor_to_device(
                model, key, device=device, value=value, dtype=dtype
            )


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

    shape_map = {k: tuple(v.shape) for k, v in model.state_dict().items()}
    loaded = 0
    skipped = 0
    print(
        f"[nfa-fun-cn] Fun CN overlay {cn_path.name} device={device} "
        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"]:
            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 shape_map or tuple(tensor.shape) != shape_map[key]:
                    skipped += 1
                    continue
                _set_tensor(model, key, tensor, device=device)
                loaded += 1
            del wrapped, inner
            gc.collect()
        else:
            for j, key in enumerate(keys):
                if key not in shape_map:
                    skipped += 1
                    continue
                tensor = f.get_tensor(key)
                if tuple(tensor.shape) != shape_map[key]:
                    skipped += 1
                    del tensor
                    continue
                _set_tensor(model, key, tensor, device=device)
                loaded += 1
                del tensor
                if (j + 1) % 50 == 0:
                    print(
                        f"[nfa-fun-cn] Fun CN overlay keys {j+1}/{len(keys)} "
                        f"loaded={loaded}",
                        flush=True,
                    )
    gc.collect()
    return loaded, skipped


def _load_control_transformer_fp8(model_name: str, cn_file: str, *, device: str):
    """Q8-class painter: empty meta → stream to device 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] load Flux2Control float8 stream "
        f"device={device} (NOT full bf16 host) 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", device=device
    )
    _check_load_deadline("after_dit_stream")
    _init_missing_control_params(transformer, device=device)
    _check_load_deadline("after_control_init")
    loaded, skipped = _overlay_fun_cn(
        transformer, Path(cn_file), device=device
    )
    _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, {device})",
        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):
    """Build Fun CN pipeline on CUDA when available (avoids host-RAM bf16 hang)."""
    global _PIPE, _PIPE_OFFLOAD_READY, _GET_IMAGE_LATENT, _LOAD_T0, _FUN_CN_LOADED
    if _PIPE is None:
        if not prepare_gpu_offload:
            # Do not CPU-preload DiT — that was the prior 35min hang.
            print(
                "[nfa-fun-cn] defer DiT+Fun CN load until @spaces.GPU "
                f"(size={GPU_SIZE})",
                flush=True,
            )
            return 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())
        device = "cuda" if torch.cuda.is_available() else "cpu"
        if device != "cuda":
            raise RuntimeError(
                "FUN_CN_REQUIRES_CUDA: fp8 Fun CN load must run inside "
                "@spaces.GPU (CPU path hangs / thrash)."
            )

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

        vae = AutoencoderKLFlux2.from_pretrained(model_name, subfolder="vae").to(
            WEIGHT_DTYPE
        )
        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 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
        convert_weight_dtype_wrapper(transformer, WEIGHT_DTYPE)
        # DiT already resident on CUDA as fp8 — keep full GPU load (xlarge 96GB).
        # TE stays CPU via pipeline hooks when using cpu_offload for VAE/TE only.
        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)
    # Load DiT+Fun CN HERE (on GPU) — never full bf16 on host.
    pipe = get_pipe(prepare_gpu_offload=True)
    if pipe is None or not _FUN_CN_LOADED:
        raise RuntimeError("FUN_CN_LOAD_ABORT: Fun CN not loaded on GPU")
    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()
        # Mounts only on host; DiT+Fun CN load is inside @spaces.GPU.
        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()