Deploy NFA Track R FLUX.2 Fun CN ZeroGPU (real depth CN)

app.py

@@ -4,12 +4,16 @@ REAL Fun ControlNet Union depth — NOT soft Flux2 image=depth (banned forever).
 
 VRAM choice (documented):
   size=\"large\" (48GB, 1× Pro) + VideoX-Fun model_cpu_offload_and_qfloat8.
+  Weights via HF Mount volumes (not 178GB ephemeral download).
+  CPU-preload outside @spaces.GPU so ZeroGPU minutes are not burned on load.
   Escalation if OOM: size=\"xlarge\" + model_cpu_offload (2× quota).
 """
 
 from __future__ import annotations
 
 import os
+import shutil
+import subprocess
 import sys
 import traceback
 from pathlib import Path
@@ -23,12 +27,11 @@ from omegaconf import OmegaConf
 from PIL import Image
 
 APP_DIR = Path(__file__).resolve().parent
-# Prefer vendored tree; else clone VideoX-Fun on Space (pip package omits submodules).
 _VX = APP_DIR / "vendor" / "VideoX-Fun"
 _VX_CACHE = Path.home() / "VideoX-Fun"
 
 
-def _patch_videox_models_init(root: Path) -> None:
+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")
@@ -58,7 +61,7 @@ def _patch_videox_models_init(root: Path) -> None:
                 "]\n",
                 encoding="utf-8",
             )
-            print("[nfa-fun-cn] patched videox_fun.models.__init__ (Flux2-only)", flush=True)
+            print("[nfa-fun-cn] patched videox_fun.models.__init__", flush=True)
 
     pipe_init = root / "videox_fun" / "pipeline" / "__init__.py"
     if pipe_init.is_file():
@@ -69,24 +72,19 @@ def _patch_videox_models_init(root: Path) -> None:
                 "__all__ = ['Flux2ControlPipeline']\n",
                 encoding="utf-8",
             )
-            print("[nfa-fun-cn] patched videox_fun.pipeline.__init__ (Flux2-only)", flush=True)
+            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_models_init(candidate)
+            _patch_videox_inits(candidate)
             p = str(candidate)
             if p not in sys.path:
                 sys.path.insert(0, p)
             return
-    import subprocess
-
     print("[nfa-fun-cn] cloning VideoX-Fun for Fun CN runtime…", flush=True)
-    _VX_CACHE.parent.mkdir(parents=True, exist_ok=True)
     if _VX_CACHE.exists():
-        import shutil
-
         shutil.rmtree(_VX_CACHE, ignore_errors=True)
     subprocess.check_call(
         [
@@ -98,16 +96,14 @@ def _ensure_videox_on_path() -> None:
             str(_VX_CACHE),
         ]
     )
-    _patch_videox_models_init(_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 ""
+    os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") or ""
 ).strip()
 if HF_TOKEN:
     try:
@@ -124,13 +120,11 @@ CN_REPO = os.environ.get(
 CN_FILE = os.environ.get(
     "NFA_FUN_CN_FILE", "FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors"
 ).strip()
-# large=48GB 1×; set NFA_FUN_CN_GPU_SIZE=xlarge only after OOM on large
 GPU_SIZE = (os.environ.get("NFA_FUN_CN_GPU_SIZE") or "large").strip().lower()
 if GPU_SIZE not in ("large", "xlarge"):
     GPU_SIZE = "large"
 GPU_DURATION = int(os.environ.get("NFA_FUN_CN_GPU_DURATION") or "300")
 WEIGHT_DTYPE = torch.bfloat16
-# VideoX-Fun official low-VRAM Fun CN mode for 48GB; offload-only on xlarge
 MEM_MODE = (
     os.environ.get("NFA_FUN_CN_MEM_MODE")
     or (
@@ -141,63 +135,47 @@ MEM_MODE = (
 ).strip()
 
 CONFIG_PATH = APP_DIR / "config" / "flux2_control.yaml"
-# Prefer HF Mount volumes (no 178GB ephemeral download). Fallback: cache download.
-MODEL_DIR = Path(
-    os.environ.get("NFA_FLUX2_MOUNT") or "/data/FLUX.2-dev"
-)
+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")
 CACHE_ROOT = Path(
-    os.environ.get("NFA_FUN_CN_CACHE")
-    or (Path.home() / ".cache" / "nfa_fun_cn")
+    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
-_WEIGHTS_READY = False
+_GET_IMAGE_LATENT = None
 
 
 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,
-    ]
-    candidates.extend(CN_MOUNT_DIR.rglob(CN_FILE) if CN_MOUNT_DIR.is_dir() else [])
-    candidates.extend(CACHE_ROOT.rglob(CN_FILE) if CACHE_ROOT.is_dir() else [])
+    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 "
-        "or download into cache."
+        "Mount alibaba-pai/FLUX.2-dev-Fun-Controlnet-Union at /data/Fun-CN."
     )
 
 
 def _ensure_weights() -> None:
-    """Resolve mounted Hub volumes, or (last resort) selective download."""
-    global _WEIGHTS_READY, MODEL_DIR, _CN_FILE_PATH
+    global MODEL_DIR, _CN_FILE_PATH
     if MODEL_DIR.is_dir() and (MODEL_DIR / "model_index.json").is_file():
         try:
-            cn_path = _resolve_cn_path()
-            _WEIGHTS_READY = True
-            print(
-                f"[nfa-fun-cn] using mounts model={MODEL_DIR} cn={cn_path}",
-                flush=True,
-            )
+            cn = _resolve_cn_path()
+            print(f"[nfa-fun-cn] using mounts model={MODEL_DIR} cn={cn}", flush=True)
             return
         except FileNotFoundError:
             pass
-
-    # Fallback: selective download (may OOM disk on ZeroGPU — mounts preferred)
-    print(
-        "[nfa-fun-cn] WARN mounts missing; selective download fallback "
-        "(transformer+vae+tokenizer+text_encoder+scheduler only)",
-        flush=True,
-    )
+    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
@@ -215,15 +193,10 @@ def _ensure_weights() -> None:
         ],
     )
     path = hf_hub_download(
-        repo_id=CN_REPO,
-        filename=CN_FILE,
-        local_dir=str(CACHE_ROOT),
-        token=token,
+        repo_id=CN_REPO, filename=CN_FILE, local_dir=str(CACHE_ROOT), token=token
     )
     MODEL_DIR = cache_model
     _CN_FILE_PATH = Path(path)
-    _WEIGHTS_READY = True
-    print(f"[nfa-fun-cn] weights ready model={MODEL_DIR} cn={_CN_FILE_PATH}", flush=True)
 
 
 def _prep_depth(depth_image: Image.Image, width: int, height: int) -> Image.Image:
@@ -234,103 +207,95 @@ def _prep_depth(depth_image: Image.Image, width: int, height: int) -> Image.Imag
 
 
 def _compose_prompt(positive: str, negative: str) -> tuple[str, str]:
-    pos = (positive or "").strip()
-    neg = (negative or "").strip() or " "
-    return pos, neg
-
-
-def get_pipe():
-    """Build VideoX-Fun Flux2ControlPipeline once per warm process."""
-    global _PIPE
-    if _PIPE is not None:
-        return _PIPE
-
-    _ensure_weights()
-    _ensure_videox_on_path()
-    from diffusers import FlowMatchEulerDiscreteScheduler
-    from safetensors.torch import load_file
-    from transformers import Mistral3ForConditionalGeneration, PixtralProcessor
-
-    # Direct imports — avoid videox_fun.models.__init__ (pulls librosa/audio stacks).
-    from videox_fun.models.flux2_transformer2d_control import (
-        Flux2ControlTransformer2DModel,
-    )
-    from videox_fun.models.flux2_vae import AutoencoderKLFlux2
-    from videox_fun.pipeline.pipeline_flux2_control import Flux2ControlPipeline
-    from videox_fun.utils.fp8_optimization import (
-        convert_model_weight_to_float8,
-        convert_weight_dtype_wrapper,
-    )
-    from videox_fun.utils.utils import get_image_latent
+    return (positive or "").strip(), ((negative or "").strip() or " ")
 
-    # stash for generate
-    get_pipe._get_image_latent = get_image_latent  # type: ignore[attr-defined]
 
-    model_name = str(MODEL_DIR)
-    cn_file = str(_resolve_cn_path())
-    config = OmegaConf.load(str(CONFIG_PATH))
-    print(
-        f"[nfa-fun-cn] load Flux2ControlTransformer + Fun CN "
-        f"mem={MEM_MODE} size={GPU_SIZE}",
-        flush=True,
-    )
-    transformer = Flux2ControlTransformer2DModel.from_pretrained(
-        model_name,
-        subfolder="transformer",
-        low_cpu_mem_usage=True,
-        torch_dtype=WEIGHT_DTYPE,
-        transformer_additional_kwargs=OmegaConf.to_container(
-            config["transformer_additional_kwargs"]
-        ),
-    ).to(WEIGHT_DTYPE)
-
-    state_dict = load_file(cn_file)
-    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
-    missing, unexpected = transformer.load_state_dict(state_dict, strict=False)
-    print(
-        f"[nfa-fun-cn] Fun CN loaded missing={len(missing)} unexpected={len(unexpected)}",
-        flush=True,
-    )
-
-    vae = AutoencoderKLFlux2.from_pretrained(model_name, subfolder="vae").to(
-        WEIGHT_DTYPE
-    )
-    tokenizer = PixtralProcessor.from_pretrained(model_name, subfolder="tokenizer")
-    text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
-        model_name,
-        subfolder="text_encoder",
-        torch_dtype=WEIGHT_DTYPE,
-        low_cpu_mem_usage=True,
-    )
-    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-        model_name, subfolder="scheduler"
-    )
-    pipeline = Flux2ControlPipeline(
-        vae=vae,
-        tokenizer=tokenizer,
-        text_encoder=text_encoder,
-        transformer=transformer,
-        scheduler=scheduler,
-    )
+def get_pipe(*, prepare_gpu_offload: bool = False):
+    """CPU-load Fun CN pipeline; arm GPU offload only inside @spaces.GPU."""
+    global _PIPE, _PIPE_OFFLOAD_READY, _GET_IMAGE_LATENT
+    if _PIPE is None:
+        _ensure_weights()
+        _ensure_videox_on_path()
+        from diffusers import FlowMatchEulerDiscreteScheduler
+        from safetensors.torch import load_file
+        from transformers import Mistral3ForConditionalGeneration, PixtralProcessor
+        from videox_fun.models.flux2_transformer2d_control import (
+            Flux2ControlTransformer2DModel,
+        )
+        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())
+        config = OmegaConf.load(str(CONFIG_PATH))
+        print(
+            f"[nfa-fun-cn] CPU-load Flux2Control + Fun CN mem={MEM_MODE}",
+            flush=True,
+        )
+        transformer = Flux2ControlTransformer2DModel.from_pretrained(
+            model_name,
+            subfolder="transformer",
+            low_cpu_mem_usage=True,
+            torch_dtype=WEIGHT_DTYPE,
+            transformer_additional_kwargs=OmegaConf.to_container(
+                config["transformer_additional_kwargs"]
+            ),
+        ).to(WEIGHT_DTYPE)
+        state_dict = load_file(cn_file)
+        state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+        missing, unexpected = transformer.load_state_dict(state_dict, strict=False)
+        print(
+            f"[nfa-fun-cn] Fun CN loaded missing={len(missing)} unexpected={len(unexpected)}",
+            flush=True,
+        )
+        vae = AutoencoderKLFlux2.from_pretrained(model_name, subfolder="vae").to(
+            WEIGHT_DTYPE
+        )
+        tokenizer = PixtralProcessor.from_pretrained(model_name, subfolder="tokenizer")
+        text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
+            model_name,
+            subfolder="text_encoder",
+            torch_dtype=WEIGHT_DTYPE,
+            low_cpu_mem_usage=True,
+        )
+        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model_name, subfolder="scheduler"
+        )
+        _PIPE = Flux2ControlPipeline(
+            vae=vae,
+            tokenizer=tokenizer,
+            text_encoder=text_encoder,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        print("[nfa-fun-cn] Flux2ControlPipeline CPU-ready (REAL Fun CN)", flush=True)
 
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    if MEM_MODE == "model_cpu_offload_and_qfloat8":
-        convert_model_weight_to_float8(
-            transformer,
-            exclude_module_name=["img_in", "txt_in", "timestep"],
-            device=device,
+    if prepare_gpu_offload and not _PIPE_OFFLOAD_READY and torch.cuda.is_available():
+        from videox_fun.utils.fp8_optimization import (
+            convert_model_weight_to_float8,
+            convert_weight_dtype_wrapper,
         )
-        convert_weight_dtype_wrapper(transformer, WEIGHT_DTYPE)
-        pipeline.enable_model_cpu_offload(device=device)
-    elif MEM_MODE == "sequential_cpu_offload":
-        pipeline.enable_sequential_cpu_offload(device=device)
-    elif MEM_MODE == "model_cpu_offload":
-        pipeline.enable_model_cpu_offload(device=device)
-    else:
-        pipeline.to(device=device)
-
-    _PIPE = pipeline
-    print("[nfa-fun-cn] Flux2ControlPipeline ready (REAL Fun CN)", flush=True)
+
+        device = "cuda"
+        transformer = _PIPE.transformer
+        if MEM_MODE == "model_cpu_offload_and_qfloat8":
+            convert_model_weight_to_float8(
+                transformer,
+                exclude_module_name=["img_in", "txt_in", "timestep"],
+                device=device,
+            )
+            convert_weight_dtype_wrapper(transformer, WEIGHT_DTYPE)
+            _PIPE.enable_model_cpu_offload(device=device)
+        elif MEM_MODE == "sequential_cpu_offload":
+            _PIPE.enable_sequential_cpu_offload(device=device)
+        elif MEM_MODE == "model_cpu_offload":
+            _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}", flush=True)
     return _PIPE
 
 
@@ -346,15 +311,13 @@ def _generate_still_gpu(
     guidance: float,
     cn_strength: float,
 ) -> Image.Image:
-    """GPU-billed Fun CN infer only — weights must already be on disk."""
     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} mem={MEM_MODE}",
+            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
@@ -362,22 +325,13 @@ def _generate_still_gpu(
     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()
-    get_image_latent = get_pipe._get_image_latent  # type: ignore[attr-defined]
-
-    # VideoX-Fun control latents (NOT Flux2Pipeline soft image=)
-    control_latent = get_image_latent(depth, sample_size=[h, w])[:, :, 0]
+    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 strength <= 0:
-        strength = 0.75
-    # ALIMAMA recommended band 0.65–0.80; allow Track R packet values
+    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(
@@ -414,13 +368,14 @@ def generate_still(
     guidance: float = 4.0,
     cn_strength: float = 0.75,
 ) -> Image.Image:
-    """CPU download + GPU Fun CN. Soft image=depth is never used."""
     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)
         return _generate_still_gpu(
             positive,
             negative or "",
@@ -444,10 +399,10 @@ 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"- Base: `{BASE_MODEL}`\n"
+        f"- Base: `{BASE_MODEL}` (HF Mount `/data/FLUX.2-dev`)\n"
         f"- GPU: `size={GPU_SIZE}` duration={GPU_DURATION}s mem=`{MEM_MODE}`\n"
         "- Soft `image=depth` is **banned** on this Space.\n"
-        "- First GPU call loads Fun CN (slow once)."
+        "- First call CPU-loads weights (slow once), then GPU Fun CN infer."
     )
     with gr.Row():
         with gr.Column():