Sync the split MiniMax-H3 Spaces

README.md

@@ -70,25 +70,29 @@ it resolved so this Space pins the same canvas rather than re-deriving it.
 
 ## Nothing is paid for with GPU time
 
-Download, load and the move onto the card happen at **startup**: `import spaces` at module top patches `torch.cuda`
-before any GPU is attached, and these are plain bfloat16 tensors, so ZeroGPU's startup packing handles them (it is a
-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 77.3 GB download and the load happen at **startup**: `import spaces` at module top patches `torch.cuda` before
+any GPU is attached, so nothing about the load needs a card. The conditioner round trip is a network call on this
+Space's CPU. A `@spaces.GPU` call is therefore only the placement (once) and the denoise loop and the two decoders.
 
-### The 150 GB quota, not the 95 GiB card, is what limits startup placement
+### The 150 GB quota, not the 95 GiB card, is what rules out 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:
+One thing does *not* happen at startup: the move onto the card. `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`, which is what keeps the conditioner half comfortable at
+66.7 GB). Packing 77.3 GB needs 154.6 GB at once, and this 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
+OSError: [Errno 28] No space left on device      # os.posix_fallocate, 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.
+Unlinking the shards first does not rescue it. `.to("cuda")` under the startup patch does not release the
+memory-mapped safetensors, so nothing is freed — and the pack's own cleanup walks those still-open mappings and
+`lstat`s them, so a deleted blob becomes `FileNotFoundError: .../blobs/3d449... (deleted)`.
+
+Placement therefore happens on the **first GPU call**, `PIPE.to("cuda")` at the top of the `@spaces.GPU` function:
+about 10 s of PCIe once, then a no-op walk, and the denoise loop runs with everything resident and no offloading at
+all. It is the same trick the 4 bit Space uses, for the same reason.
 
 ## Generation constraints
 
@@ -104,7 +108,7 @@ resident:
 | | |
 |---|---|
 | `load_components` (77.3 GB, warm Xet) | 43 s |
-| `.to("cuda")` | 10 s |
+| `.to("cuda")`, once | 10 s |
 | resident weights | 72.16 GiB |
 | denoise + decode | 317 s, **10.58 s/step** |
 | peak allocated / reserved | 78.54 / 85.37 GiB |
@@ -115,7 +119,7 @@ resident:
 | Variable | Default | Meaning |
 |---|---|---|
 | `H3_CONDITIONER` | `diffusers-internal-dev/minimax-h3-conditioner` | The Space this one asks for embeddings. |
-| `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_PLACEMENT` | `lazy` | `lazy` moves all 72.16 GiB onto the card on the first GPU call and leaves it there; `offload` hands placement to `ComponentsManager.enable_auto_cpu_offload` instead. |
 | `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,11 +30,10 @@ 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")
-# `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()
+# `lazy` moves all 72.16 GiB onto the card on the first GPU call and leaves it there; `offload` hands placement to
+# `ComponentsManager.enable_auto_cpu_offload` instead. Neither puts anything on the card at *startup*, which is
+# deliberate — see `load_models`: the 150 GB storage quota, not the 95 GiB card, is what rules that out here.
+PLACEMENT = os.environ.get("H3_PLACEMENT", "lazy").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()
@@ -82,14 +81,23 @@ def status() -> str:
 
 
 def load_models() -> str | None:
-    """Load the denoising half. At **startup**.
+    """Load the denoising half. At **startup**, but *not* onto the card.
 
     `MiniMaxH3GeneratorBlocks` declares `transformer`, `vae`, `audio_vae`, `scheduler`, `audio_scheduler` and
     `video_processor`, so `load_components` fetches exactly those subfolders out of the shared
     `modular_model_index.json` — `text_encoder/` and `transformer_ref/` are never touched.
 
     Both autoencoders carry `_keep_in_fp32_modules` over every module, so the `dtype` below is refused for them and
-    they load float32 (~20.5 GiB rather than 10.26): a bfloat16 audio VAE decodes the soundtrack ~20 dB too quiet.
+    they stay float32: a bfloat16 audio VAE decodes the soundtrack roughly 20 dB too quiet.
+
+    Nothing is moved onto the card here, which is the one place this Space departs from the ZeroGPU idiom, and the
+    reason is storage rather than memory. `spaces`' startup `torch.pack()` writes every startup-resident CUDA tensor
+    to a **second copy on disk** and only deletes the downloaded originals afterwards; 77.3 GB of weights plus a
+    77.3 GB pack is 154.6 GB against a 150 GB quota, and the Space is evicted mid-pack with `OSError: [Errno 28] No
+    space left on device` out of `os.posix_fallocate`. Deleting the shards first does not help either: the pack's own
+    cleanup walks the still-open mappings and `lstat`s them, so an unlinked blob turns into `FileNotFoundError:
+    ... (deleted)`. Placement therefore happens on the first GPU call, where it costs about 10 s of PCIe and then
+    persists across every later request in the same worker.
     """
     global PIPE, MANAGER, LOAD_ERROR, LOADED_IN
 
@@ -115,13 +123,7 @@ def load_models() -> str | None:
         pipe.load_components(dtype=torch.bfloat16, token=token)
         pipe.transformer.set_attention_backend(ATTENTION)
 
-        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. 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:
+        if PLACEMENT == "offload":
             manager.enable_auto_cpu_offload(device="cuda")
             _arm_decode_hooks(pipe)
 
@@ -134,46 +136,6 @@ 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.
 
@@ -231,11 +193,10 @@ def _generate(prompt_embeds, text_token_tags, image, last_image, height, width,
     """
     import torch
 
-    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")
+    if PLACEMENT == "lazy":
+        # 72.16 GiB across PCIe on the first request of a worker, a no-op walk on every one after it. Startup
+        # placement is not an option here — see `load_models` — and this is what buys the offload-free denoise loop.
+        PIPE.to("cuda")
 
     state = PIPE(
         prompt_embeds=prompt_embeds.to("cuda"),