Sync the split MiniMax-H3 Spaces

README.md

@@ -31,11 +31,15 @@ reason alone. Cut the `MiniMaxH3Blocks` sequence at its `text_encoder` step and
 | Space | Subfolders | Download | Resident |
 |---|---|---|---|
 | [`minimax-h3-conditioner`](https://huggingface.co/spaces/diffusers-internal-dev/minimax-h3-conditioner) | `text_encoder/` + `tokenizer/` + `processor/` | 66.7 GB | 62.15 GiB bf16 |
-| this one | `transformer/` + `vae/` + `audio_vae/` | 77.3 GB | 61.73 GiB bf16 + ~20.5 GiB float32 |
+| this one | `transformer/` + `vae/` + `audio_vae/` | 77.3 GB | 61.73 GiB bf16 + 10.43 GiB float32 |
 
 Besides the quality argument, unquantized weights are the ones AoTI can export; an NVFP4 checkpoint cannot be
 exported at all.
 
+At 72.16 GiB of weights on a 95.0 GiB card there is no offloading in the request path at all, which is why an
+*unquantized* MiniMax-H3 is also the fastest one measured here: **10.6 s/step** against the 19–21 s/step the 4 bit
+Space pays, where the whole cost is the traffic auto-offload has to move.
+
 ## How the split is expressed
 
 `MiniMaxH3Blocks` is a `SequentialPipelineBlocks` of eight steps:
@@ -71,17 +75,47 @@ before any GPU is attached, and these are plain bfloat16 tensors, so ZeroGPU's s
 torchao `Float8Tensor` that cannot be packed, and there is none here). The conditioner round trip is a network call
 on this Space's CPU. A `@spaces.GPU` call is therefore only the denoise loop and the two decoders.
 
+### The 150 GB quota, not the 95 GiB card, is what limits startup placement
+
+`spaces`' startup `torch.pack()` writes every startup-resident CUDA tensor to a **second copy on disk**, and only
+deletes the downloaded originals afterwards (`Cleaned 62.13GB of tensor files ... after packing`). Moving all
+77.3 GB onto the card therefore needs 154.6 GB at once, and the Space is evicted mid-pack:
+
+```
+ZeroGPU tensors packing:   0%|          | 0.00/77.3G
+OSError: [Errno 28] No space left on device      # os.posix_fallocate in spaces/zero/torch/packing.py
+```
+
+So only the transformer is preplaced — a 66.3 GB pack — and `app.py` unlinks the downloaded shards itself *before*
+`launch()` rather than after, which `.to("cuda")` has already made unreferenced. The autoencoders stay on the host
+and cross PCIe on the first request, 10.43 GiB and a no-op on every request after it.
+
 ## Generation constraints
 
 Fixed by the checkpoint: 24 fps, a 768 pixel short edge, 5 to 15 s, `num_frames` snapped up to the next `17 * n + 5`,
 no CFG and no negative prompt (it is guidance-distilled, so every step is one forward pass).
 
+## Measured
+
+An `rtx-pro-6000` Job — the same silicon as the ZeroGPU pool (RTX PRO 6000 Blackwell, sm120, 95.0 GiB) — running
+exactly this blockset over the wire format, 1344x768, 124 frames, 30 steps, bfloat16, cuDNN attention, everything
+resident:
+
+| | |
+|---|---|
+| `load_components` (77.3 GB, warm Xet) | 43 s |
+| `.to("cuda")` | 10 s |
+| resident weights | 72.16 GiB |
+| denoise + decode | 317 s, **10.58 s/step** |
+| peak allocated / reserved | 78.54 / 85.37 GiB |
+| output | h264 1344x768 @ 24 fps, 5.167 s + stereo AAC @ 32 kHz |
+
 ## Space variables
 
 | Variable | Default | Meaning |
 |---|---|---|
 | `H3_CONDITIONER` | `diffusers-internal-dev/minimax-h3-conditioner` | The Space this one asks for embeddings. |
-| `H3_PLACEMENT` | `resident` | `resident` keeps all 82.3 GiB of weights on the 95.0 GiB card; `offload` hands placement to `ComponentsManager.enable_auto_cpu_offload`. |
+| `H3_PLACEMENT` | `preplace` | `preplace` puts the transformer on the card at startup and the autoencoders on the first request; `offload` hands placement to `ComponentsManager.enable_auto_cpu_offload`. |
 | `H3_ATTENTION` | `_native_cudnn` | cuDNN's fused kernel, 10–20% faster than the SDPA default and needs nothing installed. flash-attention 3 is sm90-only and this pool is sm120. |
 | `H3_GPU_DURATION` | `900` | Seconds per request; the pool applies a 1.5 duration factor. |
 | `H3_GPU_SIZE` | `xlarge` | ZeroGPU allocation size. `large` does not fit. |

app.py

@@ -30,9 +30,11 @@ import gradio as gr
 
 MODEL_REPO = os.environ.get("H3_MODEL_REPO", "diffusers-internal-dev/MiniMax-H3")
 CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "diffusers-internal-dev/minimax-h3-conditioner")
-# `resident` keeps the 61.73 GiB transformer and the ~20.5 GiB of float32 VAEs on the card at once (82.3 of 95.0 GiB,
-# leaving ~12.7 GiB for activations); `offload` hands placement to `ComponentsManager.enable_auto_cpu_offload`.
-PLACEMENT = os.environ.get("H3_PLACEMENT", "resident").lower()
+# `preplace` moves the 61.73 GiB transformer onto the card at startup and leaves the 10.43 GiB of float32
+# autoencoders for the first GPU call; `offload` hands the whole thing to
+# `ComponentsManager.enable_auto_cpu_offload` instead. See `_drop_cached_weight_blobs` for why the *transformer
+# alone* rather than everything: the 150 GB storage quota, not the 95 GiB card, is the binding constraint at startup.
+PLACEMENT = os.environ.get("H3_PLACEMENT", "preplace").lower()
 # cuDNN's fused attention is 10-20% faster than the SDPA default on this pool and needs nothing installed.
 # flash-attention 3 is sm90-only and this card is sm120 (the `zero-a10g` flavour name is legacy).
 ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
@@ -113,10 +115,12 @@ def load_models() -> str | None:
         pipe.load_components(dtype=torch.bfloat16, token=token)
         pipe.transformer.set_attention_backend(ATTENTION)
 
-        if PLACEMENT == "resident":
+        if PLACEMENT == "preplace":
             # Plain bfloat16 tensors, so ZeroGPU's startup packing handles them — the thing that cannot be moved at
-            # startup is a torchao `Float8Tensor`, whose `aten.empty_like(pin_memory=True)` is unimplemented.
-            pipe.to("cuda")
+            # startup is a torchao `Float8Tensor`, whose `aten.empty_like(pin_memory=True)` is unimplemented. The
+            # autoencoders deliberately stay behind; `_generate` moves them on the first request.
+            pipe.transformer.to("cuda")
+            _log_disk(f"freed {_drop_cached_weight_blobs():.1f} GiB of downloaded weights")
         else:
             manager.enable_auto_cpu_offload(device="cuda")
             _arm_decode_hooks(pipe)
@@ -130,6 +134,46 @@ def load_models() -> str | None:
     return LOAD_ERROR
 
 
+def _log_disk(label: str) -> None:
+    import shutil
+
+    from huggingface_hub.constants import HF_HUB_CACHE
+
+    probe = HF_HUB_CACHE if os.path.isdir(HF_HUB_CACHE) else os.path.expanduser("~")
+    _total, used, free = shutil.disk_usage(probe)
+    print(f"[gen] {label} · disk {used / 2**30:.1f} used, {free / 2**30:.1f} free", flush=True)
+
+
+def _drop_cached_weight_blobs() -> float:
+    """Delete the downloaded weight files once their tensors are copies in (fake-)CUDA memory.
+
+    A ZeroGPU Space is evicted at **150 GB of storage**, and `spaces`' startup `torch.pack()` writes every
+    startup-resident CUDA tensor to a *second* copy on disk before deleting the originals itself. This half
+    downloads 77.3 GB, so packing all of it needs 154.6 GB and the Space dies with `OSError: [Errno 28] No space
+    left on device` in `os.posix_fallocate` — which is exactly how its first bring-up died.
+
+    Two things keep it under the cap. Only the transformer is moved onto the card, so the pack is 66.3 GB rather
+    than 77.3 GB; and its shards are dropped here, before `launch()`, rather than after the pack. `.to("cuda")`
+    copies each parameter out of the memory-mapped shard, so those blobs are unreferenced by then and unlinking
+    them frees the space for real. The autoencoders are still mapped, so their ~11 GB survives the unlink — which
+    is fine, 66.3 + 11 is nowhere near the cap, and it is the arithmetic this Space is designed around:
+
+        worst case, nothing frees   77.3 (download) + 66.3 (pack) = 143.6 GB
+        expected                    11 (still-mapped VAEs) + 66.3 (pack) = 77.3 GB
+    """
+    import glob
+
+    from huggingface_hub.constants import HF_HUB_CACHE
+
+    blobs = os.path.join(HF_HUB_CACHE, "models--" + MODEL_REPO.replace("/", "--"), "blobs", "*")
+    freed = 0
+    for blob in glob.glob(blobs):
+        size = os.path.getsize(blob)
+        os.unlink(blob)
+        freed += size
+    return freed / 2**30
+
+
 def _arm_decode_hooks(pipe):
     """Make the offload hooks fire for the two VAEs.
 
@@ -180,10 +224,20 @@ def encode_remote(prompt, image_path, last_image_path, canvas, num_frames):
 
 @spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE)
 def _generate(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed):
-    """The only thing on GPU time: the packed-sequence denoise loop and the two decoders."""
+    """The only thing on GPU time: the packed-sequence denoise loop and the two decoders.
+
+    Only the three generated outputs come back. A `@spaces.GPU` return crosses a process boundary by pickling, and
+    the full `PipelineState` still holds the packed latents, the rotary grid and the row indices on the card.
+    """
     import torch
 
-    return PIPE(
+    if PLACEMENT == "preplace":
+        # The two autoencoders were left on the host so that ZeroGPU's startup pack stayed inside the storage quota;
+        # 10.43 GiB across PCIe on the first request, and a no-op walk on every one after it.
+        PIPE.vae.to("cuda")
+        PIPE.audio_vae.to("cuda")
+
+    state = PIPE(
         prompt_embeds=prompt_embeds.to("cuda"),
         text_token_tags=text_token_tags,
         image=image,
@@ -194,6 +248,7 @@ def _generate(prompt_embeds, text_token_tags, image, last_image, height, width,
         num_inference_steps=int(steps),
         generator=torch.Generator("cpu").manual_seed(int(seed)),
     )
+    return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")
 
 
 def generate(prompt, image_path, last_image_path, canvas, duration, steps, seed, progress=gr.Progress()):
@@ -220,7 +275,7 @@ def generate(prompt, image_path, last_image_path, canvas, duration, steps, seed,
 
     progress(0.1, desc=f"Denoising {steps} steps at {width}x{height}, {num_frames} frames ...")
     started = time.time()
-    state = _generate(
+    frames, audio, sampling_rate = _generate(
         prompt_embeds,
         text_token_tags,
         Image.open(image_path) if image_path else None,
@@ -236,13 +291,7 @@ def generate(prompt, image_path, last_image_path, canvas, duration, steps, seed,
     directory = os.path.join(tempfile.gettempdir(), "h3-outputs")
     os.makedirs(directory, exist_ok=True)
     path = os.path.join(directory, f"h3-{int(time.time() * 1000)}.mp4")
-    encode_video(
-        state.get("videos")[0],
-        fps=FPS,
-        output_path=path,
-        audio=state.get("audio")[0],
-        audio_sample_rate=state.get("sampling_rate"),
-    )
+    encode_video(frames, fps=FPS, output_path=path, audio=audio, audio_sample_rate=sampling_rate)
 
     report = (
         f"`{width}x{height}`, {num_frames} frames ({num_frames / FPS:.3f} s), {int(steps)} steps · "