Update app.py

app.py

@@ -21,9 +21,12 @@ import gradio as gr
 
 MODEL_REPO = os.environ.get("H3_MODEL_REPO", "MiniMaxAI/MiniMax-H3")
 CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "dagloop5/qwen3vl-conditioner")
-# `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`. Startup placement is not an option here — see `load_models`.
-PLACEMENT = os.environ.get("H3_PLACEMENT", "lazy").lower()
+# `pack` moves only `transformer_ref` onto the card at startup, the same way fl2va scopes this to its own
+# `transformer` partition: packing the full ~72.16 GiB pipe (plus `spaces`' on-disk pack copy) busts the 150 GB
+# storage quota, but the ~61.7 GiB `transformer_ref` alone fits. The ~10 GB of fp32 VAEs move on the first GPU
+# call instead. `lazy` moves everything on the first GPU call rather than packing anything; `offload` hands
+# placement to `ComponentsManager.enable_auto_cpu_offload`.
+PLACEMENT = os.environ.get("H3_PLACEMENT", "pack").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()
@@ -80,9 +83,9 @@ MIN_STEPS = 4
 # Seconds of GPU one request needs, from the packed sequence it is about to denoise: linear in the rows for the
 # matmuls, quadratic for the attention, against the AoTI block package this Space runs.
 STEP_LINEAR, STEP_QUADRATIC, SAFETY = 1.1745e-4, 3.8396e-9, 1.3
-# The lazy 72.16 GiB `PIPE.to("cuda")` a cold worker pays inside its first GPU call; every request carries it, because
-# nothing here knows whether the worker it lands on is cold.
-PLACEMENT_ALLOWANCE = int(os.environ.get("H3_PLACEMENT_ALLOWANCE", "90"))
+# `pack` mode: only the ~10 GB of fp32 VAEs move, and only on a cold worker — matches fl2va's own allowance for the
+# identical move. Every request still carries it, because nothing here knows whether the worker it lands on is cold.
+PLACEMENT_ALLOWANCE = int(os.environ.get("H3_PLACEMENT_ALLOWANCE", "8"))
 AUDIO_LATENTS_PER_SECOND, AUDIO_CHANNELS = 40, 2
 REFERENCE_IMAGE_SHORT_EDGE, CANVAS_MULTIPLE = 2048, 32
 DECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124
@@ -183,16 +186,15 @@ LOAD_ERROR: str | None = None
 
 
 def load_models() -> str | None:
-    """Load the denoising half at startup, but *not* onto the card.
-
+    """Load the denoising half at startup, packing `transformer_ref` onto the card under `pack` placement.
     `MiniMaxH3Ref2VAGeneratorBlocks` declares `transformer_ref`, `vae`, `audio_vae`, the two schedulers and
     `video_processor`, so `load_components` fetches exactly those subfolders — `text_encoder/` and the `transformer/`
     partition are never touched. Both autoencoders carry `_keep_in_fp32_modules` over every module and stay float32: a
     bfloat16 audio VAE decodes the soundtrack roughly 20 dB too quiet.
-
-    Nothing moves onto the card here, for storage rather than memory: `spaces`' startup `torch.pack()` writes every
-    startup-resident CUDA tensor to a second copy on disk, and 77.3 GB of weights plus its pack busts the 150 GB quota
-    (`OSError: [Errno 28] No space left on device` out of `os.posix_fallocate`, mid-pack).
+    Only `transformer_ref` moves onto the card here, for storage rather than memory: `spaces`' startup `torch.pack()`
+    writes every startup-resident CUDA tensor to a second copy on disk, and 77.3 GB of weights plus its pack busts
+    the 150 GB quota (`OSError: [Errno 28] No space left on device` out of `os.posix_fallocate`, mid-pack); the
+    ~61.7 GB `transformer_ref` alone fits, same as fl2va's `transformer`.
     """
     global PIPE, MANAGER, LOAD_ERROR
 
@@ -231,6 +233,10 @@ def load_models() -> str | None:
         if PLACEMENT == "offload":
             manager.enable_auto_cpu_offload(device="cuda")
             _arm_decode_hooks(pipe)
+        elif PLACEMENT == "pack":
+            # Scoped to `transformer_ref`, exactly as fl2va scopes this to its `transformer` partition — see the
+            # docstring above for the quota math. The ~10 GB of fp32 VAEs move on the first GPU call instead.
+            pipe.transformer_ref.to("cuda")
 
         PIPE, MANAGER = pipe, manager
         print(f"[ref2va] ready in {time.time() - started:.0f}s", flush=True)
@@ -430,6 +436,9 @@ def _generate(prompt_embeds, text_token_tags, references, height, width, num_fra
     import torch
     if PLACEMENT == "lazy":
         PIPE.to("cuda")
+    elif PLACEMENT == "pack":
+        PIPE.vae.to("cuda")
+        PIPE.audio_vae.to("cuda")
     state = PIPE(
         prompt_embeds=prompt_embeds.to("cuda"),
         text_token_tags=text_token_tags,