Calibrate duration from three timed runs; set verified 44k token ceiling; document memory findings

README.md

@@ -46,6 +46,7 @@ the transformer's own self-attention. Reference tokens ride along at `conditioni
 | **Code** | `diffusers` main @ `d57cecde` — `LTX2InContextPipeline`, subclassed in `msr.py` |
 | **Hardware** | ZeroGPU `xlarge` (96 GB). The 22B DiT plus the Gemma-4 encoder packs to ~71 GB, over the 48 GB slice |
 | **Output** | 1280×704 @ 24 fps by default (1664×960 available), with LTX-2.5's native audio |
+| **Memory** | `expandable_segments:True`. Stage 2 cycles ~1 GB activation buffers against a nearly full card; the default arena fragments and OOMs at 43k tokens where expandable segments complete in 87 s. Gemma and the connectors are released once each has produced its output |
 
 ## Faithfulness to the reference implementation
 
@@ -69,8 +70,13 @@ The author's reference is a ComfyUI graph:
 ## Deviations
 
 * **Resolution.** The workflow renders 1664×960 × 361 frames (15 s). The default here is
-  1280×704 × 97 frames (4 s) so a run fits comfortably in a visitor's ZeroGPU quota; both the
-  workflow geometry and 8 s are selectable under **Advanced**.
+  1280×704 × 97 frames (4 s), which runs in **53 s** so it fits comfortably in a visitor's
+  ZeroGPU quota; both the workflow geometry (1664×960 × 97, 87 s) and 8 s at 1280×704 (73 s)
+  are selectable under **Advanced**.
+* **Length ceiling.** 76.2 GB stays resident during denoising and stage-2 activations cost
+  ~0.33 MB per token, so the 96 GB slice tops out near 48k tokens. Combinations above 44,000
+  (`(latent_frames + slots × 5) × H/32 × W/32`) are refused up front instead of dying in an
+  OOM traceback — 15 s at the workflow resolution is simply out of reach here.
 * **Tiled VAE decode.** The final decode is tiled (untiled it wants >100 GB). Reference
   encoding is *not* tiled, matching the node's `use_tiled_encode=False`, with a tiled fallback
   only if it would OOM.

app.py

@@ -69,10 +69,11 @@ DEFAULT_REFERENCE_FRAMES = 33  # the workflow's `reference_frames`
 
 # Stage-2 sequence length is the memory driver, and here it is user-driven twice over
 # (resolution x length, plus up to five reference slots each contributing a full clip of
-# tokens). Measured on the 96 GB ZeroGPU slice: ~1.0 MiB of activations per token on top of
-# ~45 GiB of resident weights. Refuse past the ceiling rather than let the request die in an
-# OOM traceback.
-MAX_STAGE_2_TOKENS = 46_000
+# tokens). Measured on the 96 GB slice: 76.2 GiB stays resident during denoising and peak
+# activations run at 0.33 MiB/token (87.4 GiB peak at 35,200 tokens, 90.1 GiB at 43,680),
+# which puts the wall just under 48k. 44,000 is the last value verified to complete; refuse
+# past it rather than let the request die in an OOM traceback.
+MAX_STAGE_2_TOKENS = 44_000
 
 print("[msr] downloading LTX-2.5 (diffusers)...", flush=True)
 # `from_pretrained` derives allow-patterns from every model-like file in the repo, so it would
@@ -164,11 +165,6 @@ def _plan(resolution: str, seconds: float, reference_frames: int, num_refs: int)
     return width, height, num_frames, ref_keep, tokens
 
 
-# Reference point for both the cost model and the token ceiling: the default geometry with
-# three reference slots, timed on the live Space.
-REF_TOKENS = 24_640
-REF_PIXELS = 1280 * 704 * 97
-
 
 def _count_refs(*images) -> int:
     return max(1, sum(1 for image in images if image is not None))
@@ -189,20 +185,24 @@ def gpu_duration(
 ):
     """Scale the ZeroGPU reservation with the work actually requested.
 
-    Calibrated against the reference run, which took 53.1 s split as 9.0 s stage 1 + 0.5 s
-    upsample + 19.0 s stage 2 (sampling, scales super-linearly with sequence length) and
-    24.6 s of Gemma pass, text-encoder eviction and tiled VAE decode (scales with output
-    pixels). 15% margin on top and nothing more: `duration` is charged against every
-    visitor's daily quota, and a fat reservation also costs queue priority.
+    Least-squares fit through three timed runs on this Space — 53.1 s at 24,640 tokens /
+    87 Mpx, 72.8 s at 35,200 / 174 Mpx, 86.8 s at 43,680 / 155 Mpx — which land on a plain
+    linear model to within 0.2 s:
+
+        seconds = 9.5 + 0.0206 * megapixels + 1.697 * kilotokens
+
+    Megapixels cover the tiled VAE decode and the vocoder, kilotokens cover both sampling
+    stages, and the constant is the Gemma pass plus the mp4 mux. 15% margin on top and
+    nothing more: `duration` is charged against every visitor's daily quota, and a fat
+    reservation also costs queue priority.
     """
     num_refs = _count_refs(subject_1, subject_2, subject_3, subject_4, scene_image)
     width, height, num_frames, _ref_keep, tokens = _plan(
         resolution, seconds, reference_frames, num_refs
     )
-    fixed = 12.0  # Gemma pass, 24 GB eviction, mp4 mux
-    decode = 15.6 * (width * height * num_frames) / REF_PIXELS
-    sample = 28.5 * (tokens / REF_TOKENS) ** 1.25
-    return int(min(200, math.ceil(1.15 * (fixed + decode + sample))))
+    megapixels = width * height * num_frames / 1e6
+    seconds_est = 9.5 + 0.0206 * megapixels + 1.697 * (tokens / 1000)
+    return int(min(140, math.ceil(1.15 * seconds_est)))
 
 
 def _write_video(frames, path: str, audio=None) -> None:
@@ -231,13 +231,16 @@ def _vram(label: str) -> None:
 
 
 def _encode_once(prompt: str):
-    """Run Gemma once and then take its 24 GB off the GPU for the rest of the call.
+    """Run Gemma once, then send it back to the host for the rest of the call.
+
+    Not CPU offloading: the text encoder is genuinely finished after this, and nothing moves
+    inside a sampling loop. Both `__call__`s below are handed the resulting `prompt_embeds`,
+    so `encode_prompt` short-circuits and never touches it again.
 
-    Not CPU offloading: the text encoder is genuinely finished after this. Both `__call__`s
-    below are handed the resulting `prompt_embeds`, so `encode_prompt` short-circuits and never
-    touches it again, while `connectors` (which does run per stage) stays resident. Without
-    this the 22B DiT plus the encoder leaves ~20 GiB for activations and stage 2 OOMs at the
-    default geometry. `DiffusionPipeline.device` deliberately prefers a non-CPU component, so
+    Measured effect is only ~1.9 GiB — ZeroGPU packs many parameters into shared CUDA
+    storages, so releasing one module's references frees far less than its nominal 24 GB —
+    but it is free and it lowers the peak the Gemma pass itself reaches.
+    `DiffusionPipeline.device` deliberately prefers a non-CPU component, so
     `_execution_device` still reports cuda afterwards.
     """
     pipe.text_encoder.to("cuda")  # no-op on a fresh worker; restores it if one is reused
@@ -251,10 +254,8 @@ def _encode_once(prompt: str):
             device=torch.device("cuda"),
         )
     finally:
-        _vram("before text-encoder eviction")
         pipe.text_encoder.to("cpu")
         _free()
-        _vram(f"after text-encoder eviction ({next(pipe.text_encoder.parameters()).device})")
     return prompt_embeds, prompt_mask
 
 
@@ -340,7 +341,6 @@ def generate(
     # One Gemma pass, reused by both stages, then evicted (see `_encode_once`).
     torch.cuda.reset_peak_memory_stats()
     prompt_embeds, prompt_mask = _encode_once(prompt)
-    _vram("weights resident")
 
     shared = dict(
         prompt=None,