exp4: decode timers + H3_PROFILE profiler hook

app.py

@@ -139,6 +139,8 @@ VSA_BLOCKS = 0
 VSA_GATES = 0
 # Wall seconds of the last request's transformer forwards (reset by every `_generate` call).
 FORWARD_TIMES: list[float] = []
+# Wall seconds of the last request's two VAE decodes, appended by the wrappers `_arm_decode_timers` arms.
+DECODE_TIMES: dict[str, float] = {}
 
 
 def status() -> str:
@@ -203,13 +205,8 @@ def load_models() -> str | None:
         else:
             pipe.transformer.set_attention_backend(ATTENTION)
 
-        # The video VAE ships fp32 (`_keep_in_fp32_modules` over every module); decode it in bf16 instead —
-        # half the decode FLOPs and traffic on Blackwell. The audio VAE keeps its fp32 modules: a bf16
-        # audio VAE decodes the soundtrack roughly 20 dB too quiet (see the module docstring).
-        pipe.vae.to(torch.bfloat16)
-        print(f"[gen] video VAE cast to {next(pipe.vae.parameters()).dtype}", flush=True)
-
         _install_forward_timing(pipe)
+        _arm_decode_timers(pipe)
 
         if PLACEMENT == "offload":
             manager.enable_auto_cpu_offload(device="cuda")
@@ -227,6 +224,27 @@ def load_models() -> str | None:
     return LOAD_ERROR
 
 
+def _arm_decode_timers(pipe) -> None:
+    """Time each VAE's `decode` call, so a request can report the video/audio decode split.
+
+    Instance-attribute wrapper, like `_arm_decode_hooks`; only meaningful with `PLACEMENT=lazy` (no offload hook in
+    front). The video VAE decode runs tiled under fp16 autocast over fp32 weights inside the pipeline's decode
+    block; this just measures it.
+    """
+    for name in ("vae", "audio_vae"):
+        module = getattr(pipe, name)
+        inner = module.decode
+
+        def timed(*args, _name=name, _decode=inner, **kwargs):
+            started = time.perf_counter()
+            try:
+                return _decode(*args, **kwargs)
+            finally:
+                DECODE_TIMES[_name] = DECODE_TIMES.get(_name, 0.0) + time.perf_counter() - started
+
+        module.decode = timed
+
+
 def _install_forward_timing(pipe) -> None:
     """Record the wall time of every transformer forward, so each request can report where its seconds went.
 
@@ -335,21 +353,49 @@ def _generate(prompt_embeds, text_token_tags, height: int, width: int, num_frame
 
     torch.cuda.reset_peak_memory_stats()
     FORWARD_TIMES.clear()
+    DECODE_TIMES.clear()
     started = time.perf_counter()
-    state = PIPE(
-        prompt_embeds=prompt_embeds.to("cuda"),
-        text_token_tags=text_token_tags,
-        height=height,
-        width=width,
-        num_frames=num_frames,
-        num_inference_steps=SIGMA_GRID_POINTS,
-        generator=torch.Generator("cpu").manual_seed(int(seed)),
-    )
+    if os.environ.get("H3_PROFILE") == "1":
+        from torch.profiler import ProfilerActivity, profile
+
+        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
+            state = PIPE(
+                prompt_embeds=prompt_embeds.to("cuda"),
+                text_token_tags=text_token_tags,
+                height=height,
+                width=width,
+                num_frames=num_frames,
+                num_inference_steps=SIGMA_GRID_POINTS,
+                generator=torch.Generator("cpu").manual_seed(int(seed)),
+            )
+        print(
+            prof.key_averages().table(sort_by="cuda_time_total", row_limit=40, max_name_column_width=60),
+            flush=True,
+        )
+    else:
+        state = PIPE(
+            prompt_embeds=prompt_embeds.to("cuda"),
+            text_token_tags=text_token_tags,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            num_inference_steps=SIGMA_GRID_POINTS,
+            generator=torch.Generator("cpu").manual_seed(int(seed)),
+        )
     gen_s = time.perf_counter() - started
     peak_gib = torch.cuda.max_memory_allocated() / 1024**3
     forwards = [round(t, 2) for t in FORWARD_TIMES]
-    print(f"[gen] pipeline {gen_s:.2f}s · forwards {forwards} · peak {peak_gib:.2f} GiB", flush=True)
-    return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate"), gen_s, peak_gib, forwards
+    decodes = {name: round(t, 2) for name, t in DECODE_TIMES.items()}
+    print(f"[gen] pipeline {gen_s:.2f}s · forwards {forwards} · decodes {decodes} · peak {peak_gib:.2f} GiB", flush=True)
+    return (
+        state.get("videos")[0],
+        state.get("audio")[0].cpu(),
+        state.get("sampling_rate"),
+        gen_s,
+        peak_gib,
+        forwards,
+        decodes,
+    )
 
 
 @spaces.GPU(duration=120, size=GPU_SIZE)
@@ -462,7 +508,7 @@ def generate(
     refined = plan.get("refined_prompt") or ""
 
     progress(0.2, desc=f"{NUM_FORWARDS} transformer forwards at {width}x{height}, {num_frames} frames ...")
-    frames, audio, sampling_rate, gen_s, peak_gib, forwards = _generate(
+    frames, audio, sampling_rate, gen_s, peak_gib, forwards, decodes = _generate(
         prompt_embeds, text_token_tags, height, width, num_frames, seed
     )
 
@@ -474,10 +520,11 @@ def generate(
 
     forward_total = sum(forwards) if forwards else float("nan")
     forward_detail = "+".join(f"{t:.1f}" for t in forwards) if forwards else "n/a"
+    decode_detail = " + ".join(f"{k} {v:.1f}s" for k, v in decodes.items()) if decodes else "n/a"
     report = (
         f"`{width}x{height}`, {num_frames} frames @ {FPS} fps, {NUM_FORWARDS} transformer forwards · "
-        f"pipeline **{gen_s:.1f}s** (forwards {forward_detail} = {forward_total:.1f}s, decode+other "
-        f"{gen_s - forward_total:.1f}s) · peak {peak_gib:.1f} GiB · "
+        f"pipeline **{gen_s:.1f}s** (forwards {forward_detail} = {forward_total:.1f}s, {decode_detail}, "
+        f"other {gen_s - forward_total - sum(decodes.values()):.1f}s) · peak {peak_gib:.1f} GiB · "
         f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
         f"{', rewritten' if refined else ''}) · seed {int(seed)} · req {request_id}"
     )