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

README.md

@@ -1,13 +1,36 @@
 ---
-title: NFA Track R FLUX2 Fun CN ZeroGPU
-emoji: 🏆
-colorFrom: indigo
-colorTo: gray
+title: NFA Track R FLUX.2 Fun CN ZeroGPU
+emoji: 🎬
+colorFrom: green
+colorTo: blue
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.13'
+sdk_version: "5.49.1"
 app_file: app.py
 pinned: false
+license: other
+short_description: Real Fun depth ControlNet (NOT soft image=depth)
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Track R — FLUX.2 Fun depth ControlNet (ZeroGPU)
+
+**Hard rule:** This Space runs **real** ALIMAMA / VideoX-Fun Fun ControlNet Union depth.
+Soft Flux2 `image=depth` is **banned forever** on this path.
+
+## Stack (VRAM choice)
+
+| Piece | Value |
+|---|---|
+| Painter + CN | VideoX-Fun `Flux2ControlPipeline` + `FLUX.2-dev-Fun-Controlnet-Union-2602` |
+| Base | `black-forest-labs/FLUX.2-dev` |
+| GPU | `@spaces.GPU(duration=300, size="large")` = **48GB** @ **1×** Pro minutes |
+| Memory mode | VideoX-Fun **`model_cpu_offload_and_qfloat8`** (official low-VRAM Fun CN path) |
+| Escalation | If OOM → redeploy with `size="xlarge"` + `model_cpu_offload` (2× quota) |
+
+Requires Space secret **`HF_TOKEN`** (gated FLUX.2-dev license accepted on the account).
+
+## API
+
+Same client contract as Plan A soft Space: `/generate_still`  
+(`positive`, `negative`, `depth_image`, `seed`, `width`, `height`, `steps`, `guidance`, `cn_strength`)
+
+Depth is **required**. Soft depth is never used.

app.py

@@ -0,0 +1,372 @@
+"""NFA Track R — Fun depth ControlNet ZeroGPU (VideoX-Fun / ALIMAMA).
+
+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.
+  Escalation if OOM: size=\"xlarge\" + model_cpu_offload (2× quota).
+"""
+
+from __future__ import annotations
+
+import os
+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
+
+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()
+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()
+# 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 (
+        "model_cpu_offload"
+        if GPU_SIZE == "xlarge"
+        else "model_cpu_offload_and_qfloat8"
+    )
+).strip()
+
+APP_DIR = Path(__file__).resolve().parent
+CONFIG_PATH = APP_DIR / "config" / "flux2_control.yaml"
+CACHE_ROOT = Path(
+    os.environ.get("NFA_FUN_CN_CACHE")
+    or (Path.home() / ".cache" / "nfa_fun_cn")
+)
+MODEL_DIR = CACHE_ROOT / "FLUX.2-dev"
+
+_PIPE = None
+_CN_FILE_PATH: Path | None = None
+_WEIGHTS_READY = False
+
+
+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
+    direct = CACHE_ROOT / CN_FILE
+    if direct.is_file():
+        _CN_FILE_PATH = direct
+        return direct
+    nested = list(CACHE_ROOT.rglob(CN_FILE))
+    if nested:
+        _CN_FILE_PATH = nested[0]
+        return nested[0]
+    raise FileNotFoundError(f"Fun CN weights missing: {CN_FILE}")
+
+
+def _ensure_weights() -> None:
+    """Download on CPU (must run outside @spaces.GPU so quota is not burned)."""
+    global _WEIGHTS_READY, _CN_FILE_PATH
+    if _WEIGHTS_READY and MODEL_DIR.is_dir():
+        try:
+            _resolve_cn_path()
+            return
+        except FileNotFoundError:
+            pass
+    CACHE_ROOT.mkdir(parents=True, exist_ok=True)
+    token = HF_TOKEN or None
+    print(f"[nfa-fun-cn] snapshot {BASE_MODEL} -> {MODEL_DIR}", flush=True)
+    snapshot_download(
+        repo_id=BASE_MODEL,
+        local_dir=str(MODEL_DIR),
+        local_dir_use_symlinks=False,
+        token=token,
+    )
+    print(f"[nfa-fun-cn] download {CN_REPO}/{CN_FILE}", flush=True)
+    path = hf_hub_download(
+        repo_id=CN_REPO,
+        filename=CN_FILE,
+        local_dir=str(CACHE_ROOT),
+        local_dir_use_symlinks=False,
+        token=token,
+    )
+    _CN_FILE_PATH = Path(path)
+    _WEIGHTS_READY = True
+    print(f"[nfa-fun-cn] weights ready cn={_CN_FILE_PATH}", flush=True)
+
+
+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]:
+    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()
+    from diffusers import FlowMatchEulerDiscreteScheduler
+    from safetensors.torch import load_file
+    from videox_fun.models import (
+        AutoencoderKLFlux2,
+        Flux2ControlTransformer2DModel,
+        Mistral3ForConditionalGeneration,
+        PixtralProcessor,
+    )
+    from videox_fun.pipeline 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
+
+    # 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,
+    )
+
+    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,
+        )
+        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)
+    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:
+    """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}",
+            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()
+    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]
+    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 = 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",
+        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:
+    """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()
+        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"- Base: `{BASE_MODEL}`\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)."
+    )
+    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()

config/flux2_control.yaml

@@ -0,0 +1,5 @@
+format: diffusers
+pipeline: flux2
+transformer_additional_kwargs:
+    control_layers: [0, 2, 4, 6]
+    control_in_dim: 260

requirements.txt

@@ -0,0 +1,19 @@
+torch
+torchvision
+torchaudio
+accelerate
+safetensors
+sentencepiece
+protobuf
+pillow
+numpy
+omegaconf
+einops
+spaces
+gradio>=5.0.0
+huggingface_hub
+requests
+# Official Fun CN stack (ALIMAMA / VideoX-Fun) — NOT Diffusers soft image=
+videox-fun @ git+https://github.com/aigc-apps/VideoX-Fun.git
+diffusers>=0.36.0
+transformers>=4.46.2