Fun CN: stream fp8 DiT on CUDA inside @spaces.GPU xlarge (no host bf16)

app.py

@@ -135,7 +135,7 @@ CN_FILE = os.environ.get(
 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")
+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"
@@ -281,32 +281,38 @@ def _fp8_dtype_for_key(key: str) -> torch.dtype:
     return torch.float8_e4m3fn
 
 
-def _set_tensor(model: torch.nn.Module, key: str, tensor: torch.Tensor) -> None:
-    """Materialize on CPU as bf16, then in-place cast to float8 when allowed.
-
-    Direct float8 via set_module_tensor_to_device hung on ZeroGPU CPU; bf16
-    assign + ``param.data = param.data.to(float8)`` is the stable path.
-    """
+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)
-    set_module_tensor_to_device(
-        model, key, device="cpu", value=value, 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
-    if _fp8_dtype_for_key(key) != torch.float8_e4m3fn:
-        return
-    # Walk to the Parameter and cast storage in-place (frees bf16 bits).
-    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)
-    elif isinstance(param, torch.Tensor):
-        setattr(mod, leaf, param.to(torch.float8_e4m3fn))
 
 
 def _stream_shards_into_model(
@@ -314,39 +320,46 @@ def _stream_shards_into_model(
     shard_paths: list[str],
     *,
     label: str,
+    device: str,
 ) -> None:
-    """Load one safetensors shard at a time → fp8; never accumulate full bf16."""
-    from safetensors.torch import load_file
+    """Stream shards key-by-key (no full-shard RAM spike) into float8 on device."""
+    from safetensors import safe_open
 
-    model_sd = model.state_dict()
+    # 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"rss_max≈{_rss_gb():.1f}GB path={Path(path).name}",
+            f"device={device} rss_max≈{_rss_gb():.1f}GB path={Path(path).name}",
             flush=True,
         )
-        sd = load_file(path, device="cpu")
-        n_keys = len(sd)
-        for j, (key, tensor) in enumerate(sd.items()):
-            if key not in model_sd:
-                skipped += 1
-                continue
-            if tuple(tensor.shape) != tuple(model_sd[key].shape):
-                skipped += 1
-                continue
-            _set_tensor(model, key, tensor)
-            loaded += 1
-            if (j + 1) % 200 == 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,
-                )
-        del sd
+        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",
@@ -354,24 +367,16 @@ def _stream_shards_into_model(
     )
 
 
-def _init_missing_control_params(model: torch.nn.Module) -> None:
+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()}
-    # 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)
+    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
@@ -381,65 +386,72 @@ def _init_missing_control_params(model: torch.nn.Module) -> None:
         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:
+            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="cpu", value=value, 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) -> tuple[int, int]:
+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
 
-    model_sd = model.state_dict()
+    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} "
+        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"]:
-            # 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):
+                if key not in shape_map or tuple(tensor.shape) != shape_map[key]:
                     skipped += 1
                     continue
-                _set_tensor(model, key, tensor)
+                _set_tensor(model, key, tensor, device=device)
                 loaded += 1
             del wrapped, inner
             gc.collect()
         else:
-            for key in keys:
-                if key not in model_sd:
+            for j, key in enumerate(keys):
+                if key not in shape_map:
                     skipped += 1
                     continue
                 tensor = f.get_tensor(key)
-                if tuple(tensor.shape) != tuple(model_sd[key].shape):
+                if tuple(tensor.shape) != shape_map[key]:
                     skipped += 1
+                    del tensor
                     continue
-                _set_tensor(model, key, tensor)
+                _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):
-    """Q8-class painter: empty meta → stream bf16 shards as float8 → Fun CN."""
+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 (
@@ -456,8 +468,8 @@ def _load_control_transformer_fp8(model_name: str, cn_file: str):
     ])
 
     print(
-        "[nfa-fun-cn] CPU-load Flux2Control as float8 stream "
-        f"(NOT full bf16) mem={MEM_MODE} gpu_size={GPU_SIZE}",
+        "[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():
@@ -467,15 +479,19 @@ def _load_control_transformer_fp8(model_name: str, cn_file: str):
     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")
+    _stream_shards_into_model(
+        transformer, shards, label="DiT-fp8", device=device
+    )
     _check_load_deadline("after_dit_stream")
-    _init_missing_control_params(transformer)
+    _init_missing_control_params(transformer, device=device)
     _check_load_deadline("after_control_init")
-    loaded, skipped = _overlay_fun_cn(transformer, Path(cn_file))
+    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)",
+        f"rss_max≈{_rss_gb():.1f}GB (REAL Fun CN, fp8 painter, {device})",
         flush=True,
     )
     return transformer
@@ -508,9 +524,17 @@ def _load_local_quantized_te(te_root: Path):
 
 
 def get_pipe(*, prepare_gpu_offload: bool = False):
-    """CPU-load Fun CN pipeline (fp8 DiT stream); arm GPU offload inside @spaces.GPU."""
+    """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()
@@ -525,14 +549,21 @@ def get_pipe(*, prepare_gpu_offload: bool = False):
         _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)
+        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 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(
@@ -547,7 +578,7 @@ def get_pipe(*, prepare_gpu_offload: bool = False):
         )
         elapsed = time.time() - _LOAD_T0
         print(
-            f"[nfa-fun-cn] Flux2ControlPipeline CPU-ready "
+            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,
@@ -558,12 +589,9 @@ def get_pipe(*, prepare_gpu_offload: bool = False):
 
         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)
+        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"):
@@ -602,7 +630,10 @@ def _generate_still_gpu(
     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
@@ -651,12 +682,7 @@ def generate_still(
                 "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"
-            )
+        # Mounts only on host; DiT+Fun CN load is inside @spaces.GPU.
         return _generate_still_gpu(
             positive,
             negative or "",