[skill] add-model: fold in Cosmos3 omni-port lessons

Revisions to the /add-model skill from the Cosmos3-Nano omni port retrospective
(tests/local_tests/cosmos3/port_feedback.md). Each addition maps to a pitfall
that silently shipped wrong output or wasted significant time.

common_rules.md (Production Boundary):
- The official reference framework is a TEST-ONLY dependency — no production
  import; production must run with it uninstalled.
- Stock library model classes the reference merely vendors (e.g. a Qwen-VL ViT)
  may be reused under the documented model-import exception, proven bit-exact,
  rather than re-ported.
- Mandate verifying runtime independence (grep + sys.modules block) before handoff.

SKILL.md:
- Phase 1: checkpoint per-component configs can be lossy (verify vs framework
  instantiation, not the JSON — the black-video scheduler trap); dump safetensors
  keys+shapes early; temporal-VAE conditioning ≠ "only frame 0" (static-repeat);
  pipeline knobs may be resolution/axis-keyed. New multi-modal/MoT note:
  per-modality packing + MRoPE offsets + flat order are the real work; text
  reasoning reuses the LM backbone + lm_head; meta-device non-persistent buffers
  + fp32/bf16 boundary casts.
- Phase 2: the oracle MUST be the framework object (not a 2nd copy of the unit —
  the green-suite scheduler bug); report max AND mean diff; make the oracle
  CPU-runnable (SDPA monkeypatch; mirror un-importable constants).
- Phase 3: a reused module can diverge on an untested config path (re-run parity
  on the new config — the odd-stride output_padding case); check for a stock
  library class or an existing in-repo module before porting a large sub-model.
- Pre-handoff checklist: max/mean reporting; reference-framework import ban +
  runtime-independence verification.

.agents/skills/add-model/SKILL.md

@@ -141,6 +141,54 @@ Arch config rule:
   `transformer/config.json` from the official Python model-config class, not
   from data/eval config classes.
 
+Reference/checkpoint traps to check explicitly (each has silently shipped wrong
+output before):
+
+- **The checkpoint's per-component config can be a lossy generic-format export
+  that does NOT describe the runtime.** A "diffusers-format" `scheduler/`,
+  `vae/`, etc. config may carry fields the framework never uses (and that flip
+  behavior). Verify each component against *how the framework actually
+  instantiates it*, not the JSON. (A Cosmos3 `scheduler_config.json` with
+  `use_karras_sigmas=true` selected diffusion sigmas instead of the framework's
+  flow-matching sigmas → NaN latents → black video, while the DiT was correct.)
+- **Dump `safetensors` keys + shapes for every sub-checkpoint NOW.** This reveals
+  the on-disk naming (often a third convention, distinct from both the framework
+  module tree and your FastVideo classes), whether a sub-checkpoint is partial
+  (e.g. decoder-only), and frequently that it matches an existing native module
+  you can reuse.
+- **Temporal-VAE conditioning is not "only frame 0 matters."** When a condition
+  image/clip feeds a temporal autoencoder, replicate the framework's exact frame
+  fill (e.g. static-repeat of the conditioning frame across the clip), not an
+  assumed zero-fill — the temporal receptive field makes the kept-clean latent
+  frame depend on several pixel frames.
+- **Pipeline knobs may be keyed by an axis you didn't expect** (resolution,
+  aspect, embodiment), not by task. Confirm where a knob's default comes from
+  (e.g. flow_shift was resolution-bucketed, not T2I-vs-video).
+
+Multi-modal / mixture-of-transformers (omni) models:
+
+- The backbone is often a shared transformer with per-modality
+  embed-in/embed-out heads and a packed sequence (e.g. a causal text split + a
+  full-attention generation split shared by vision/audio/action). Once the first
+  modality (usually vision) works, each additional modality is the same *shape*
+  of work — pack tokens, project in + add a modality embedding + scatter timestep
+  embeds, run the shared backbone, project out, unpack — so build the first
+  modality with clean "a modality" seams.
+- The real novelty per modality is the **packing layout + position-id (MRoPE)
+  offsets** and the **combined flat-latent order** for joint denoising (the
+  per-sample concat order must match the framework). Spend the care there.
+- A modality may be **domain-conditioned** (e.g. per-embodiment action via a
+  domain-indexed linear); thread a per-token domain id through pack + project.
+- **Text reasoning usually needs no new model code** — it is the causal pathway
+  + `embed_tokens` / final norm / `lm_head`, all already in the DiT. A text-only
+  forward + `lm_head` should be token-for-token identical to the reference
+  reasoner.
+- **Derived non-persistent buffers** (e.g. RoPE `inv_freq`) are absent from the
+  checkpoint, so under a meta-device FSDP load they land on the `meta` device;
+  provide a `materialize_non_persistent_buffers`-style hook to recompute them on
+  the real device. Also cast at fp32/bf16 boundaries (fp32 noise/VAE latents into
+  a bf16 model) — no-ops in fp32 parity tests, required at inference.
+
 ## Phase 2: Early Parity Scaffolding
 
 Create component parity tests before or alongside implementation. Use
@@ -159,6 +207,20 @@ This phase is early by design:
 - The scaffold must still contain real official loading, deterministic inputs,
   output extraction, and concrete tensor comparisons. No unconditional skips,
   no shape-only tests.
+- **The oracle side MUST be the official framework object** — never a second
+  instance of the unit under test. A test that compares two copies of the wrong
+  thing passes while the real bug ships. (A scheduler "parity" test that used the
+  same diffusers scheduler on both sides hid the black-video bug behind a green
+  suite.) After writing each test, ask: "if the framework were wrong here, would
+  this fail?"
+- **Report both max AND mean abs diff** for every numerical parity test (not just
+  pass/fail) — the trend across components is the signal, and a clean port should
+  read max=mean=0.0 (or fp32-epsilon, called out).
+- **Budget for making the oracle runnable in the test env.** The framework's
+  attention may be CUDA-only (flash/natten) — monkeypatch it to SDPA so parity
+  runs on CPU/float32. If a framework module won't import headless (heavy/private
+  deps), mirror the one constant you need in the test with a cited source rather
+  than importing it.
 
 Use subagents here: dispatch one parity-test subagent per required component,
 including components that may be reused. Their output becomes the red/skip
@@ -192,6 +254,22 @@ Reuse decision:
    FastVideo-native code through the bucket-specific skill.
 4. Reused components still require non-skip component parity against the exact
    official instantiation used by the target pipeline.
+5. **A "matching" native module can still diverge on a config path the original
+   user never exercised.** Re-run parity with the NEW model's exact config — do
+   not assume a reused module is correct because it passed for its first model.
+   (FastVideo's `OobleckVAE` decoder matched bit-exact except for `output_padding
+   = stride % 2`, a no-op for the original model's even strides but required for
+   the new model's odd stride — a 1-sample drift otherwise.) A safe one-line fix
+   that is provably a no-op for existing users is preferred over forking the
+   module.
+6. **Before porting a large sub-model, check two things first:** (a) is it
+   literally a stock library class (transformers/diffusers) that the reference
+   framework merely vendors a copy of? then prefer reusing the library class
+   under the documented model-import exception, proven bit-exact vs the reference;
+   (b) does an existing in-repo native module already implement this architecture
+   under a different name? (Both the audio decoder and the vision encoder were
+   "already there" — one in-repo, one in transformers — saving a from-scratch
+   ViT/codec port.)
 
 Porting subagent dispatch:
 
@@ -373,6 +451,7 @@ Pre-handoff checklist:
 [ ] Prep handoff is complete and committed nowhere with token values.
 [ ] Conversion was run if needed and output loads with real weights.
 [ ] Every required component, reused or newly ported, has a non-skip local parity PASS.
+[ ] Every numerical parity test reports max AND mean abs diff (clean port reads max=mean=0.0, or fp32-epsilon called out).
 [ ] `local_tests_readme` lists every component parity test, command, status, and blocker if any.
 [ ] `port_state_file` has every open question/issue either resolved or listed as an explicit blocker.
 [ ] Any `next_step=ask_user` has a matching `escape_hatch` block and `E###` row.
@@ -380,7 +459,8 @@ Pre-handoff checklist:
 [ ] Pipeline parity has a non-skip local PASS against the official reference.
 [ ] Basic example runs and writes a non-corrupt output.
 [ ] Video SSIM or audio-specific quality regression is added or explicitly deferred.
-[ ] Runtime production code has no diffusers/transformers model-class imports.
+[ ] Runtime production code has no diffusers/transformers model-class imports AND no official-reference-framework imports.
+[ ] Runtime independence verified: production modules import + the model builds with the reference framework blocked (sys.modules raise-on-import), or proven by grep.
 [ ] Production comments are WHY-focused; examples have user-story docstrings.
 [ ] Post-parity hot-path pass is complete.
 ```
@@ -408,6 +488,8 @@ matching `*secret*`.
   registry/example wiring, smoke tests, and pipeline parity-debug.
 - `fastvideo/layers/AGENTS.md` for native layer selection and state-dict surface
   guidance.
+- `tests/local_tests/cosmos3/port_feedback.md` for a worked omni-port (Cosmos3)
+  retrospective: the concrete pitfalls + fixes the Phase notes above distill.
 - `docs/contributing/coding_agents.md` for narrative context.
 - `docs/design/overview.md` for pipeline/config/registry architecture.
 - `fastvideo/pipelines/basic/wan/` for standard T2V/I2V/DMD/Causal variants.
@@ -438,3 +520,4 @@ matching `*secret*`.
 | 2026-04-30 | Extracted handoff schemas into `contracts/` for shared use across skills. |
 | 2026-04-30 | Added pipeline skill contract and Phase 7 component-parity gate. |
 | 2026-04-30 | Added escape-hatch contract for user decisions and `ask_user` handoffs. |
+| 2026-06-07 | Folded Cosmos3 omni-port lessons (`tests/local_tests/cosmos3/port_feedback.md`): lossy checkpoint configs / dump-weights-early / temporal-VAE conditioning / resolution-keyed knobs (Phase 1); oracle-must-be-framework + report max/mean + CPU-runnable oracle (Phase 2); reuse-diverges-on-untested-config + check-for-stock/in-repo-module (Phase 3); multi-modal/MoT modality-bookkeeping + reasoning/buffer/dtype notes (Phase 1); reference-framework production ban + runtime-independence verification (common_rules + checklist). |

.agents/skills/add-model/shared/common_rules.md

@@ -75,14 +75,30 @@ weaken these shared rules.
 
 - No runtime `from diffusers import <model class>` or
   `from transformers import <model class>` in `fastvideo/` production code.
+- **The official reference framework (the upstream repo used as the parity
+  oracle) is a TEST-ONLY dependency.** No `fastvideo/` production module may
+  import it. Production must run with the reference framework *uninstalled* — it
+  exists only to validate the native port, never to run it.
 - Components that own weights or numerical behavior must be FastVideo-native
   unless the user explicitly accepts a documented lazy-wrapper exception.
 - Allowed third-party runtime exceptions are tokenizers and pure data utilities
-  when they match existing project patterns.
-- Tests may import diffusers/transformers as parity references.
+  (e.g. a seeded-randn helper) when they match existing project patterns.
+- When the checkpoint component IS literally a stock library model class and the
+  reference framework itself only vendors a copy of that class (e.g. a stock
+  Qwen-VL vision encoder), reusing the library class — proven bit-exact vs the
+  reference — is usually better than re-porting a large stock model. It is still
+  a model-class import, so it takes the documented lazy-wrapper-exception path
+  (flag it for the user), not the default-native path.
+- Tests may import diffusers / transformers / the reference framework as parity
+  references.
+- **Verify runtime independence before handoff** (see the final-handoff
+  checklist): `grep` production code for reference/diffusers/transformers model
+  imports, AND block those top-level packages in `sys.modules` (raise on import)
+  and confirm every production module still imports and the model still builds.
 - Production comments explain why, not what or provenance. Avoid narrative
   comments like `vendored from`, `matches upstream`, `REVIEW`, or session-history
-  commentary.
+  commentary. (A one-line "mirrors `<ref path>`" pointer on a parity-critical
+  function is acceptable when it documents the numerical contract, not history.)
 
 ## Verification Semantics
 

docker/Dockerfile.python3.10

@@ -55,12 +55,14 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install FastVideo Unified Kernel.
+# This build machine has no GPU, so target Hopper (sm_90a) explicitly instead
+# of probing a live device for the arch (matches the released kernel wheel).
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
     cd fastvideo-kernel && \
     git submodule update --init --recursive && \
-    ./build.sh
+    TORCH_CUDA_ARCH_LIST=9.0a ./build.sh
 
 
 EXPOSE 22

docker/Dockerfile.python3.11

@@ -55,12 +55,14 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install FastVideo Unified Kernel.
+# This build machine has no GPU, so target Hopper (sm_90a) explicitly instead
+# of probing a live device for the arch (matches the released kernel wheel).
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
     cd fastvideo-kernel && \
     git submodule update --init --recursive && \
-    ./build.sh
+    TORCH_CUDA_ARCH_LIST=9.0a ./build.sh
 
 
 EXPOSE 22

docker/Dockerfile.python3.12

@@ -55,11 +55,13 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install FastVideo Unified Kernel.
+# This build machine has no GPU, so target Hopper (sm_90a) explicitly instead
+# of probing a live device for the arch (matches the released kernel wheel).
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
     cd fastvideo-kernel && \
     git submodule update --init --recursive && \
-    ./build.sh
+    TORCH_CUDA_ARCH_LIST=9.0a ./build.sh
 
 EXPOSE 22

docker/Dockerfile.python3.12.cuda12.9.1

@@ -55,12 +55,14 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install FastVideo Unified Kernel.
+# This build machine has no GPU, so target Hopper (sm_90a) explicitly instead
+# of probing a live device for the arch (matches the released kernel wheel).
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
     cd fastvideo-kernel && \
     git submodule update --init --recursive && \
-    ./build.sh
+    TORCH_CUDA_ARCH_LIST=9.0a ./build.sh
 
 
 EXPOSE 22

examples/inference/basic/basic_ltx2_3_distilled_i2v.py

@@ -0,0 +1,286 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LTX-2.3 distilled image-to-video with torch.compile + timing breakdown.
+
+This example runs the LTX-2.3 distilled student model on a single GPU with
+torch.compile fully enabled, then prints a per-stage timing breakdown so the
+user can see where wall-time goes. It is meant as the canonical entry point
+for trying out the LTX-2.3 i2v path on `hao-ai-lab/FastVideo:main`.
+
+Quick start
+-----------
+    export LTX23_I2V_IMAGE=/path/to/your/portrait_or_product.jpg
+    # optional overrides:
+    #   export LTX23_I2V_PROMPT="a fashion model walks toward camera..."
+    #   export LTX23_OUTPUT_DIR=outputs_video/ltx2_3_distilled_i2v
+    python examples/inference/basic/basic_ltx2_3_distilled_i2v.py
+
+What the script does
+--------------------
+1. Loads FastVideo/LTX-2.3-Distilled-Diffusers (8 denoise + 3 refine steps,
+   CFG=1, no refine LoRA — the distilled production recipe).
+2. Compiles the DiT, text encoder, and VAE (fullgraph, max-autotune-no-
+   cudagraphs).
+3. Runs 2 warmup calls (untimed) + 2 measured calls. Two warmups are needed
+   even though distilled has no refine LoRA, because Inductor's per-shape
+   autotune for the first call still leaves a few cold guards on call 2;
+   the second warmup settles them. Skipping to a single warmup typically
+   inflates the first measured run by tens of seconds.
+4. Prints a per-stage breakdown and an average over the measured runs.
+
+Hardware notes
+--------------
+- Single-GPU example; for multi-GPU sequence-parallel see the gradio demo
+  under `examples/inference/gradio/local/gradio_local_demo_ltx2_3/`.
+- First-time compile + autotune takes ~10-30 min on H100 / GB200 (cached
+  in `$TORCHINDUCTOR_CACHE_DIR` afterwards). Subsequent invocations only
+  pay the one-time process load + a few seconds of dynamo trace.
+- On GB200 / Blackwell, run with `env -u LD_LIBRARY_PATH ...` to avoid a
+  system-cuBLAS / torch-cuBLAS mismatch that fails every GEMM. The
+  `_inductor.shape_padding = False` line below also avoids a pad_mm
+  landmine on the same generation of cards.
+"""
+from __future__ import annotations
+
+import os
+import time
+from collections import OrderedDict
+from pathlib import Path
+
+import torch._inductor.config as _inductor
+
+from fastvideo import VideoGenerator
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.utils import maybe_download_model
+
+# Env knobs (set BEFORE importing fastvideo where possible — but
+# FASTVIDEO_ATTENTION_BACKEND is fine here because the worker reads it
+# on generator construction).
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
+os.environ.setdefault("FASTVIDEO_STAGE_LOGGING", "1")
+
+# Inductor knobs. The first one (shape_padding=False) is mandatory on
+# Blackwell to avoid a cuBLAS INVALID_VALUE crash inside pad_mm during
+# the refine path. The rest are autotune-friendliness flags.
+_inductor.shape_padding = False
+_inductor.conv_1x1_as_mm = True
+_inductor.coordinate_descent_tuning = True
+_inductor.coordinate_descent_check_all_directions = True
+_inductor.epilogue_fusion = False
+
+MODEL_ID = os.path.expandvars(
+    os.path.expanduser(
+        os.getenv("LTX23_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers")
+    )
+)
+OUTPUT_DIR = Path(
+    os.getenv("LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_distilled_i2v")
+)
+I2V_IMAGE = os.getenv("LTX23_I2V_IMAGE", "")
+DEFAULT_PROMPT = (
+    "A fashion model takes a slow step forward and shifts her weight, "
+    "the soft fabric of her clothing swaying and rippling with the "
+    "motion, her hair shifting gently, soft even studio lighting on a "
+    "clean light background, elegant slow-motion runway feel."
+)
+PROMPT = os.getenv("LTX23_I2V_PROMPT", DEFAULT_PROMPT)
+
+# Per-stage timing helpers --------------------------------------------------
+
+def _print_stage_breakdown(result: dict, label: str) -> float | None:
+    """Print stage execution times and return the sum, or None if missing."""
+    logging_info = result.get("logging_info")
+    stages = getattr(logging_info, "stages", None) if logging_info else None
+    if not stages:
+        print(f"  [{label}] stage breakdown unavailable")
+        return None
+    print(f"  [{label}] stage breakdown:")
+    total = 0.0
+    for name, metrics in stages.items():
+        exec_s = float(metrics.get("execution_time", 0.0))
+        total += exec_s
+        print(f"    - {name}: {exec_s:.3f}s")
+    print(f"    - stage_sum: {total:.3f}s")
+    return total
+
+
+def _collect_stage_times(
+    result: dict,
+    stage_times: dict[str, list[float]],
+    stage_order: OrderedDict[str, None],
+) -> None:
+    logging_info = result.get("logging_info")
+    stages = getattr(logging_info, "stages", None) if logging_info else None
+    if not stages:
+        return
+    for name, metrics in stages.items():
+        stage_order.setdefault(name, None)
+        stage_times.setdefault(name, []).append(
+            float(metrics.get("execution_time", 0.0))
+        )
+
+
+def _resolve_refine_upsampler(model_root: str) -> Path:
+    """LTX-2.3 distilled snapshots ship a `spatial_upscaler/` subdir."""
+    for name in ("spatial_upscaler", "spatial_upsampler"):
+        cand = Path(model_root) / name
+        if (cand / "config.json").is_file():
+            return cand
+    raise FileNotFoundError(
+        f"No refine upsampler directory under {model_root}. "
+        f"Expected `{model_root}/spatial_upscaler/config.json`."
+    )
+
+
+# Main ---------------------------------------------------------------------
+
+def main() -> None:
+    if not I2V_IMAGE:
+        raise SystemExit(
+            "LTX23_I2V_IMAGE is required for i2v. Example:\n"
+            "  export LTX23_I2V_IMAGE=/path/to/portrait_or_product.jpg\n"
+            "  python examples/inference/basic/basic_ltx2_3_distilled_i2v.py"
+        )
+    if not Path(I2V_IMAGE).is_file():
+        raise SystemExit(f"LTX23_I2V_IMAGE not found: {I2V_IMAGE}")
+
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    model_root = maybe_download_model(MODEL_ID)
+    refine_upsampler_path = _resolve_refine_upsampler(model_root)
+    print(f"Model:           {model_root}")
+    print(f"Refine upsampler: {refine_upsampler_path}")
+    print(f"i2v image:       {I2V_IMAGE}")
+    print(f"Output dir:      {OUTPUT_DIR.resolve()}")
+
+    torch_compile_kwargs = {
+        "backend": "inductor",
+        "fullgraph": True,
+        "mode": "max-autotune-no-cudagraphs",
+        "dynamic": False,
+    }
+
+    # Loading the pipeline config *with model_path* binds model-specific
+    # tuning (notably VAE precision/decoder defaults) into the config. Without
+    # this, the generic pipeline config gives a substantially slower VAE
+    # decode stage. `basic_ltx2_distilled_fast_profile.py` uses the same
+    # pattern.
+    pipeline_config = PipelineConfig.from_pretrained(model_root)
+    pipeline_config.dit_config.quant_config = None
+
+    generator = VideoGenerator.from_pretrained(
+        model_root,
+        num_gpus=1,
+        # LTX-2.3 distilled uses the two-stage refine pipeline; the refine
+        # LoRA is intentionally empty for the distilled student.
+        ltx2_refine_enabled=True,
+        ltx2_refine_upsampler_path=str(refine_upsampler_path),
+        ltx2_refine_lora_path="",
+        ltx2_refine_num_inference_steps=3,
+        ltx2_refine_guidance_scale=1.0,
+        ltx2_refine_add_noise=True,
+        pipeline_config=pipeline_config,
+        enable_torch_compile=True,
+        enable_torch_compile_text_encoder=True,
+        # Compile the VAE codec submodules (encoder / decoder) too. The
+        # `LTX2CausalVideoAutoencoder` declares `_compile_conditions` so
+        # `_compile_with_conditions` targets just those submodules and
+        # leaves the surrounding tiling control flow eager — needed for
+        # fullgraph + dynamic=False to succeed. VAE eager decode is
+        # ~1.0s; compiling it brings the stage to ~0.3s.
+        enable_torch_compile_vae=True,
+        torch_compile_kwargs=torch_compile_kwargs,
+        torch_compile_kwargs_vae=torch_compile_kwargs,
+        # Keep everything resident — no CPU offload for serving-style runs.
+        dit_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+        vae_cpu_offload=False,
+        ltx2_vae_tiling=False,
+    )
+
+    common_kwargs = dict(
+        prompt=PROMPT,
+        negative_prompt="",         # distilled is CFG-free; no negative needed
+        guidance_scale=1.0,         # CFG=1 for distilled
+        height=1280, width=832,     # portrait runway aspect
+        num_frames=121, fps=24,     # ~5s clip
+        num_inference_steps=8,      # distilled denoise steps
+        # i2v: anchor the input image at frame 0 with full strength.
+        # `ltx2_image_crf=0.0` skips an extra JPEG re-encode of an already
+        # JPEG conditioning image.
+        ltx2_images=[(I2V_IMAGE, 0, 1.0)],
+        ltx2_image_crf=0.0,
+        save_video=True,
+    )
+
+    warmup_runs = 2
+    measured_runs = 2
+    warmup_secs: list[float] = []
+    measured_secs: list[float] = []
+    stage_times: dict[str, list[float]] = {}
+    stage_order: OrderedDict[str, None] = OrderedDict()
+
+    try:
+        # Warmup: untimed (but we still wall-clock them so the first compile
+        # cost is visible to the reader).
+        for w in range(warmup_runs):
+            t0 = time.perf_counter()
+            print(f"\n[warmup {w + 1}/{warmup_runs}] compiling + generating…")
+            generator.generate_video(
+                output_path=str(OUTPUT_DIR / f"_warmup_{w + 1}.mp4"),
+                seed=7,
+                **common_kwargs,
+            )
+            dt = time.perf_counter() - t0
+            warmup_secs.append(dt)
+            print(f"[warmup {w + 1}/{warmup_runs}] wall={dt:.1f}s")
+
+        # Cleanup warmup artifacts so the user only sees measured outputs.
+        for w in range(warmup_runs):
+            (OUTPUT_DIR / f"_warmup_{w + 1}.mp4").unlink(missing_ok=True)
+
+        # Measured.
+        for m in range(measured_runs):
+            out_path = OUTPUT_DIR / f"output_ltx2_3_distilled_i2v_run_{m + 1}.mp4"
+            print(f"\n[measured {m + 1}/{measured_runs}] generating: {out_path}")
+            t0 = time.perf_counter()
+            result = generator.generate_video(
+                output_path=str(out_path),
+                seed=2002 + m,
+                **common_kwargs,
+            )
+            wall = time.perf_counter() - t0
+            e2e = (
+                result.get("e2e_latency")
+                if isinstance(result, dict) else None
+            ) or wall
+            measured_secs.append(e2e)
+            print(f"[measured {m + 1}/{measured_runs}] e2e={e2e:.2f}s wall={wall:.2f}s")
+            if isinstance(result, dict):
+                _print_stage_breakdown(result, f"measured {m + 1}")
+                _collect_stage_times(result, stage_times, stage_order)
+
+        # Summary.
+        print("\n=== summary ===")
+        print(f"warmup wall-times:      {[round(x, 1) for x in warmup_secs]}")
+        if measured_secs:
+            avg = sum(measured_secs) / len(measured_secs)
+            print(
+                f"measured e2e (n={len(measured_secs)}): "
+                f"{[round(x, 2) for x in measured_secs]} -> avg {avg:.2f}s"
+            )
+        if stage_times:
+            print(f"average stage times over {measured_runs} measured runs:")
+            avg_total = 0.0
+            for name in stage_order:
+                vals = stage_times.get(name) or []
+                if not vals:
+                    continue
+                avg_v = sum(vals) / len(vals)
+                avg_total += avg_v
+                print(f"  - {name}: {avg_v:.3f}s")
+            print(f"  - stage_sum_avg: {avg_total:.3f}s")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/eval/eval_third_person_separation.py

@@ -0,0 +1,91 @@
+"""Run ``judge.third_person_separation`` (needs ``.[eval-judge]`` + a Gemini key)
+over each baseline and print the candidate's win-rate table — from a ``--manifest``
+of pairs, or by pairing ``--candidate-dir`` against each ``--reference`` dir by
+filename stem.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from collections import defaultdict
+from pathlib import Path
+
+from fastvideo.eval import create_evaluator
+
+METRIC = "judge.third_person_separation"
+VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".webm"}
+IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}
+
+
+def _by_stem(directory: Path, exts: set[str]) -> dict[str, Path]:
+    """Map filename stem -> path for files with the given extensions."""
+    return {p.stem: p for p in sorted(directory.iterdir()) if p.suffix.lower() in exts}
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__,
+                                formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--candidate-dir", type=Path, default=None,
+                   help="Directory of candidate clips (directory mode).")
+    p.add_argument("--reference", action="append", default=[], metavar="NAME=DIR",
+                   help="Baseline directory, repeatable: 'name=dir' or bare 'dir'.")
+    p.add_argument("--image-dir", type=Path, default=None,
+                   help="Optional first-frame images, matched to clips by stem.")
+    p.add_argument("--prompts-json", type=Path, default=None,
+                   help="Optional {stem: control-signal text} JSON.")
+    p.add_argument("--actions-json", type=Path, default=None,
+                   help="Optional {stem: action-label} JSON for the per-action breakdown.")
+    p.add_argument("--manifest", type=Path, default=None,
+                   help="JSON list of {baseline, video_path, reference_path, ...} rows.")
+    p.add_argument("--output", type=Path, default=None)
+    args = p.parse_args()
+
+    # Group path-only samples per baseline: {baseline: [sample dict, ...]}.
+    by_baseline: dict[str, list[dict]] = defaultdict(list)
+    if args.manifest is not None:
+        for row in json.loads(args.manifest.read_text()):
+            by_baseline[row.get("baseline", "baseline")].append(
+                {k: v for k, v in row.items() if k != "baseline"})
+    elif args.candidate_dir is not None and args.reference:
+        cands = _by_stem(args.candidate_dir, VIDEO_EXTS)
+        images = _by_stem(args.image_dir, IMAGE_EXTS) if args.image_dir else {}
+        prompts = json.loads(args.prompts_json.read_text()) if args.prompts_json else {}
+        actions = json.loads(args.actions_json.read_text()) if args.actions_json else {}
+        for spec in args.reference:
+            name, sep, ref_dir = spec.partition("=")
+            if not sep:
+                name, ref_dir = Path(spec).name, spec
+            refs = _by_stem(Path(ref_dir), VIDEO_EXTS)
+            for stem in sorted(cands.keys() & refs.keys()):
+                sample = {"video_path": str(cands[stem]), "reference_path": str(refs[stem])}
+                if stem in images:
+                    sample["image_path"] = str(images[stem])
+                if stem in prompts:
+                    sample["text_prompt"] = prompts[stem]
+                if stem in actions:
+                    sample["action"] = actions[stem]
+                by_baseline[name].append(sample)
+    else:
+        p.error("provide either --manifest, or --candidate-dir with at least one --reference")
+
+    ev = create_evaluator(metrics=[METRIC], device="cpu")
+    print("\n| Baseline | Candidate win-rate (excl. ties) | W / L / T | n |")
+    print("|---|---|---|---|")
+    rows = {}
+    for baseline, samples in by_baseline.items():
+        res = ev.evaluate(samples=samples).corpus[METRIC]
+        rows[baseline] = res
+        d = res.details
+        if res.score is None:
+            print(f"| {baseline} | — | — | 0 |")
+        else:
+            print(f"| {baseline} | {100 * res.score:.1f}% | {d['wins']}/{d['losses']}/{d['ties']} | {d['n']} |")
+
+    if args.output is not None:
+        payload = {b: {"score": r.score, "details": r.details} for b, r in rows.items()}
+        args.output.write_text(json.dumps(payload, indent=2))
+        print(f"\nWrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()

fastvideo-kernel/build.sh

@@ -78,16 +78,26 @@ print(f'{mj}.{mn}')"
 }
 
 if [ "${GPU_BACKEND}" = "CUDA" ]; then
-    detected_cc="$(detect_with_torch)" || {
-        echo "ERROR: torch-based CUDA arch detection failed in uv environment." >&2
-        echo "       Ensure torch is installed and CUDA is available in the uv-selected Python." >&2
-        exit 1
-    }
-
-    cc_major="${detected_cc%%.*}"
-    cc_minor="${detected_cc##*.}"
+    # Compute capability drives the arch/TK defaults below. Prefer an explicit
+    # TORCH_CUDA_ARCH_LIST (works on GPU-less build machines such as CI/Docker);
+    # only probe a live GPU via torch when no arch was provided.
+    if [ -n "${TORCH_CUDA_ARCH_LIST:-}" ]; then
+        echo "Using TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} (skipping torch GPU probe)"
+        first_arch="${TORCH_CUDA_ARCH_LIST%%[;, ]*}"   # first entry, e.g. 9.0a
+        first_arch="${first_arch%[af]}"                # strip trailing a/f suffix
+        cc_major="${first_arch%%.*}"
+        cc_minor="${first_arch##*.}"
+    else
+        detected_cc="$(detect_with_torch)" || {
+            echo "ERROR: torch-based CUDA arch detection failed and TORCH_CUDA_ARCH_LIST is unset." >&2
+            echo "       Set TORCH_CUDA_ARCH_LIST (e.g. 9.0a) for GPU-less builds, or build where CUDA is available." >&2
+            exit 1
+        }
+        cc_major="${detected_cc%%.*}"
+        cc_minor="${detected_cc##*.}"
+        echo "Detected compute capability via torch: ${detected_cc}"
+    fi
     cmake_arch="${cc_major}${cc_minor}"
-    echo "Detected compute capability via torch: ${detected_cc} (sm_${cmake_arch})"
 
     # Respect explicit overrides.
     if [ -z "${TORCH_CUDA_ARCH_LIST:-}" ]; then

fastvideo/eval/README.md

@@ -2,8 +2,9 @@
 
 In-process evaluation suite for video generations. Includes pixel
 metrics (SSIM, PSNR, LPIPS), Fréchet Video Distance (FVD), optical-flow
-comparisons, the full VBench suite, Physics-IQ, audio metrics, and a
-VLM scorer behind a single registry-driven API.
+comparisons, the full VBench suite, Physics-IQ, audio metrics, an
+absolute VLM scorer (`videoscore2`), and a pairwise VLM judge
+(`judge.third_person_separation`) — all behind a single registry-driven API.
 
 ## Install
 
@@ -159,6 +160,7 @@ fastvideo/
 │       ├── audio/                 # clap_score, audiobox_aesthetics, kl_divergence,
 │       │                          # frechet_distance, wer, desync, imagebind_score
 │       ├── videoscore2/           # VideoScore-2 (Qwen2.5-VL)
+│       ├── judge/                 # pairwise VLM judges (third_person_separation)
 │       ├── physics_iq/            # PhysicsIQ + sub-metrics
 │       └── vbench/                # adapter: sys.path bootstrap + shims
 │           ├── __init__.py
@@ -313,6 +315,40 @@ to control read/write behavior. The example script
 `examples/inference/eval/eval_fvd.py` demonstrates the full
 two-directory workflow.
 
+## `judge.third_person_separation` — pairwise VLM judge
+
+A **preference** metric (a judge, not an absolute score), and the suite's first
+remote-API one. For each pair the judge (Gemini) sees the shared first frame and
+two rollouts — a candidate and a reference model under the same control signal —
+and picks the one that better separates the third-person CHARACTER (foreground)
+from the BACKGROUND. The corpus score is the candidate's win-rate, excluding
+ties. Set-vs-set, motion-first; it reads native mp4s, so samples carry path
+strings, not decoded tensors.
+
+```bash
+uv pip install -e .[eval-judge]      # opt-in: needs network + an API key
+export GEMINI_API_KEY=...            # or GOOGLE_API_KEY, or ~/.gemini_token
+```
+
+```python
+from fastvideo.eval import create_evaluator
+
+ev = create_evaluator(metrics=["judge.third_person_separation"], device="cpu")
+result = ev.evaluate(samples=[
+    {"video_path": "cand/000.mp4", "reference_path": "base/000.mp4",
+     "image_path": "frames/000.png", "text_prompt": "W: moves forward", "action": "W"},
+    # ... more pairs ...
+]).corpus["judge.third_person_separation"]
+result.score    # candidate win-rate excl. ties; result.details has the breakdown
+```
+
+Only `video_path`/`reference_path` are required; `image_path`/`text_prompt`/
+`action` are optional. Verdicts are cached under `${FASTVIDEO_EVAL_CACHE}/eval/judge/`.
+The judge separates best when the control yields genuine parallax (e.g.
+translation); rigid whole-frame motion (e.g. pure camera rotation) is harder. To
+sweep several baselines into a table, see
+`examples/inference/eval/eval_third_person_separation.py`.
+
 ## Out of scope (follow-up PRs)
 
 - **MIND** metrics. Depend on a separate `vipe` upstream submodule.

fastvideo/eval/metrics/judge/third_person_separation/metric.py

@@ -0,0 +1,306 @@
+"""Pairwise VLM judge (Gemini): of two rollouts under the same control, which
+better separates the third-person character (foreground) from the background;
+the score is the candidate's win-rate over a reference.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+from fastvideo.eval.metrics.base import BaseMetric
+from fastvideo.eval.models import get_cache_dir
+from fastvideo.eval.registry import register
+from fastvideo.eval.types import MetricResult, Video
+
+# Part of the on-disk cache key; bump to invalidate cached verdicts.
+RUBRIC_ID = "v1"
+DEFAULT_MODEL = "gemini-2.5-pro"
+DEFAULT_K = 3
+
+SYSTEM_PROMPT = ("You are a strict comparative evaluator of third-person video-game "
+                 "rollouts. Two videos were generated by two different models from the "
+                 "SAME first frame and the SAME control signal. You judge which video "
+                 "better demonstrates that the model SEPARATES the third-person CHARACTER "
+                 "(foreground) from the BACKGROUND SCENE — i.e. the control animates the "
+                 "character as an INDEPENDENT AGENT with its own trajectory while the "
+                 "background moves with the camera. You do NOT reward whichever video "
+                 "merely looks cleaner, sharper, or higher-res.")
+
+RUBRIC = """\
+You are watching TWO generated video rollouts (Video 1, Video 2) played in full,
+from the same first frame under the same control signal. Pick the better
+third-person world-model rollout. Judge THREE things together, in this order:
+
+(C) MOTION / ACTION EXECUTION FIRST. The rollout must actually CARRY OUT the
+    control signal with substantial motion (the scene/character clearly moves as
+    commanded). A clip that is near-static, barely drifts, or only twitches has
+    NOT demonstrated controllable separation — it FAILS, no matter how clean it
+    looks. CRITICAL: do NOT reward a clip for looking "smoother" or "more
+    stable" when that smoothness is really just the ABSENCE OF MOTION. Less
+    motion is NOT better. If one clip executes the action with clear motion and
+    the other is comparatively static, the MOVING one wins (unless it fails B).
+
+(A) TEMPORAL COHERENCE — among clips that actually move, penalize GENUINE
+    corruption: flicker/strobing, texture boiling/crawling, geometry swimming,
+    the character or scene morphing/warping into mush, colors pulsing, or
+    progressive degradation into noise. Do NOT confuse LEGITIMATE large motion
+    (camera sweeping, character running, scene flowing past) with instability —
+    fast correct motion is GOOD, not a defect. Only true frame-to-frame
+    INCOHERENCE counts against a clip.
+
+(B) FOREGROUND/BACKGROUND SEPARATION — the character stays a distinct, coherent
+    entity with its own trajectory while the background responds to the control;
+    it does not dissolve/smear into the bg, and the whole frame does not slide
+    as one rigid sheet.
+
+Decision: among clips that genuinely execute the motion (C), pick the one that
+is both temporally coherent (A) and shows cleaner separation (B). A static or
+barely-moving clip loses to a moving one. Real motion is not instability. Do not
+reward resolution or placidity. "tie" only if truly equivalent on all three."""
+
+USER_TASK = ("Output a JSON object with four fields:\n"
+             " - video_1_analysis: FIRST, how much does Video 1 actually move — does it "
+             "execute the control with clear motion, or is it near-static / barely "
+             "drifting? THEN: among its motion, is there GENUINE corruption (flicker, "
+             "boiling, morphing into mush) as opposed to legitimate fast motion? THEN: "
+             "is the character a distinct coherent entity vs rigid-slide / dissolve?\n"
+             " - video_2_analysis: the same three checks for Video 2.\n"
+             " - comparison: apply (C) motion-first, then (A) genuine-coherence, then "
+             "(B) separation. A near-static clip loses to a moving one; legitimate large "
+             "motion is NOT a defect.\n"
+             " - winner: \"video_1\", \"video_2\", or \"tie\".")
+
+RESPONSE_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "video_1_analysis": {
+            "type": "string"
+        },
+        "video_2_analysis": {
+            "type": "string"
+        },
+        "comparison": {
+            "type": "string"
+        },
+        "winner": {
+            "type": "string",
+            "enum": ["video_1", "video_2", "tie"]
+        },
+    },
+    "required": ["video_1_analysis", "video_2_analysis", "comparison", "winner"],
+    "propertyOrdering": ["video_1_analysis", "video_2_analysis", "comparison", "winner"],
+}
+
+
+def _path_of(sample: dict, key: str) -> str | None:
+    """Resolve a native-file path from a string key or a Video wrapper."""
+    p = sample.get(f"{key}_path")
+    if isinstance(p, str):
+        return p
+    v = sample.get(key)
+    if isinstance(v, Video) and isinstance(v.source, str):
+        return v.source
+    if isinstance(v, str):
+        return v
+    return None
+
+
+def _resolve_api_key() -> str:
+    for env in ("GEMINI_API_KEY", "GOOGLE_API_KEY"):
+        key = os.environ.get(env)
+        if key:
+            return key.strip()
+    token = Path("~/.gemini_token").expanduser()
+    if token.is_file():
+        return token.read_text().strip()
+    raise ValueError("judge.third_person_separation needs a Gemini API key. Set "
+                     "GEMINI_API_KEY (or GOOGLE_API_KEY), or write it to ~/.gemini_token.")
+
+
+@register("judge.third_person_separation")
+class ThirdPersonSeparationMetric(BaseMetric):
+    """Pairwise VLM judge of third-person fg/bg separation; corpus win-rate."""
+
+    name = "judge.third_person_separation"
+    requires_reference = True
+    higher_is_better = True
+    needs_gpu = False
+    is_set_metric = True
+    dependencies = ["google.genai"]
+
+    def __init__(self, model: str = DEFAULT_MODEL, k: int = DEFAULT_K) -> None:
+        super().__init__()
+        self.model = model
+        self.k = k
+        self._client: Any = None
+        self._files: dict[str, Any] = {}  # path -> uploaded Gemini file handle
+        self._records: list[dict] = []  # one per accumulated pair
+
+    # --- model / client -----------------------------------------------------
+    def setup(self) -> None:
+        if self._client is not None:
+            return
+        from google import genai
+        self._client = genai.Client(api_key=_resolve_api_key())
+
+    # --- set-vs-set protocol ------------------------------------------------
+    def reset(self) -> None:
+        self._records = []
+        self._files = {}
+
+    def accumulate(self, sample: dict) -> None:
+        cand = _path_of(sample, "video")
+        base = _path_of(sample, "reference")
+        if cand is None or base is None:
+            return  # nothing to compare
+
+        image = sample.get("image_path")
+        action_text = sample.get("text_prompt") or ""
+        action = sample.get("action")
+
+        rec = self._cached(cand, base, action_text)
+        if rec is None:
+            if self._client is None:
+                self.setup()
+            rec = self._judge_pair(cand, base, image, action_text)
+            self._write_cache(cand, base, action_text, rec)
+        rec = {**rec, "action": action}
+        self._records.append(rec)
+
+    def finalize(self) -> MetricResult:
+        recs = [r for r in self._records if r.get("verdict")]
+        if not recs:
+            return MetricResult(name=self.name, score=None, details={"skipped": "no pairs judged"})
+        wins = sum(r["verdict"] == "candidate" for r in recs)
+        losses = sum(r["verdict"] == "baseline" for r in recs)
+        ties = sum(r["verdict"] == "tie" for r in recs)
+        decided = wins + losses
+        score = wins / decided if decided else None
+
+        # Per-action win-rate, grouped by the raw label (no assumed control scheme).
+        per_action: dict[str, dict] = {}
+        labels: set[str] = {str(r["action"]) for r in recs if r.get("action")}
+        for action in sorted(labels):
+            gr = [r for r in recs if r.get("action") == action]
+            gw = sum(r["verdict"] == "candidate" for r in gr)
+            gl = sum(r["verdict"] == "baseline" for r in gr)
+            per_action[action] = {
+                "n": len(gr),
+                "wins": gw,
+                "losses": gl,
+                "ties": len(gr) - gw - gl,
+                "win_rate": (gw / (gw + gl)) if (gw + gl) else None,
+            }
+        return MetricResult(name=self.name,
+                            score=score,
+                            details={
+                                "wins": wins,
+                                "losses": losses,
+                                "ties": ties,
+                                "n": len(recs),
+                                "win_rate_excl_ties": score,
+                                "per_action": per_action,
+                            })
+
+    def merge_from(self, other: BaseMetric) -> None:
+        assert isinstance(other, ThirdPersonSeparationMetric)
+        self._records.extend(other._records)
+
+    # --- judging ------------------------------------------------------------
+    def _judge_pair(self, cand: str, base: str, image: str | None, action_text: str) -> dict:
+        """k counterbalanced comparisons → aggregated per-pair verdict."""
+        # Seed the A/B alternation from the pair itself so it is reproducible and
+        # independent of evaluation order or which subset is being run.
+        seed = int(hashlib.sha1(f"{cand}|{base}".encode()).hexdigest(), 16)
+        mapped: list[str] = []
+        for i in range(self.k):
+            cand_first = (seed + i) % 2 == 0
+            v1, v2 = (cand, base) if cand_first else (base, cand)
+            winner = self._one_call(image, v1, v2, action_text)
+            if winner == "tie":
+                mapped.append("tie")
+            elif (winner == "video_1") == cand_first:
+                mapped.append("candidate")
+            else:
+                mapped.append("baseline")
+        cand_w = mapped.count("candidate")
+        base_w = mapped.count("baseline")
+        verdict = ("candidate" if cand_w > base_w else "baseline" if base_w > cand_w else "tie")
+        return {
+            "verdict": verdict,
+            "candidate_wins": cand_w,
+            "baseline_wins": base_w,
+            "ties": mapped.count("tie"),
+            "k": self.k,
+            "rubric_id": RUBRIC_ID
+        }
+
+    def _one_call(self, image: str | None, vid1: str, vid2: str, action_text: str) -> str:
+        from google.genai import types
+        contents: list[Any] = [action_text or "Compare these two rollouts."]
+        if image is not None:
+            contents += ["\nFirst frame (input condition, shared by BOTH "
+                         "videos):", self._upload(image)]
+        contents += [
+            "\nVideo 1 (model A's full rollout — watch it in motion):",
+            self._upload(vid1),
+            "\nVideo 2 (model B's full rollout — watch it in motion):",
+            self._upload(vid2),
+            "\n" + RUBRIC + "\n\n" + USER_TASK,
+        ]
+        for attempt in range(6):
+            try:
+                resp = self._client.models.generate_content(model=self.model,
+                                                            contents=contents,
+                                                            config=types.GenerateContentConfig(
+                                                                system_instruction=SYSTEM_PROMPT,
+                                                                response_mime_type="application/json",
+                                                                response_schema=RESPONSE_SCHEMA,
+                                                                temperature=0.4))
+                return json.loads(resp.text).get("winner", "tie")
+            except Exception as exc:  # noqa: BLE001 - transient API errors
+                if attempt == 5:
+                    print(f"[judge] giving up after 6 attempts ({exc}); scoring this call a tie")
+                    break
+                is_429 = "429" in str(exc) or "RESOURCE_EXHAUSTED" in str(exc)
+                time.sleep(40 if is_429 else 2**attempt)
+        return "tie"
+
+    def _upload(self, path: str) -> Any:
+        f = self._files.get(path)
+        if f is not None:
+            return f
+        f = self._client.files.upload(file=path)
+        while f.state.name != "ACTIVE":
+            time.sleep(1)
+            f = self._client.files.get(name=f.name)
+            if f.state.name == "FAILED":
+                raise RuntimeError(f"Gemini upload failed for {path}")
+        self._files[path] = f
+        return f
+
+    # --- per-pair cache -----------------------------------------------------
+    def _cache_path(self, cand: str, base: str, action_text: str) -> Path:
+        # k is intentionally NOT in the key so a larger-k run can reuse an
+        # existing verdict with enough samples (see ``_cached``).
+        h = hashlib.sha1(f"{RUBRIC_ID}|{self.model}|{cand}|{base}|{action_text}".encode()).hexdigest()[:16]
+        return get_cache_dir() / "judge" / "third_person_separation" / f"{h}.json"
+
+    def _cached(self, cand: str, base: str, action_text: str) -> dict | None:
+        cp = self._cache_path(cand, base, action_text)
+        if not cp.is_file():
+            return None
+        try:
+            rec = json.loads(cp.read_text())
+        except Exception:
+            return None
+        return rec if rec.get("k", 0) >= self.k and "verdict" in rec else None
+
+    def _write_cache(self, cand: str, base: str, action_text: str, rec: dict) -> None:
+        cp = self._cache_path(cand, base, action_text)
+        cp.parent.mkdir(parents=True, exist_ok=True)
+        cp.write_text(json.dumps(rec, indent=2))

fastvideo/eval/registry.py

@@ -85,6 +85,8 @@ def _extra_for(metric_name: str) -> str:
         return "eval-physics-iq"
     if metric_name.startswith("audio."):
         return "eval-audio"
+    if metric_name.startswith("judge."):
+        return "eval-judge"
     return "eval"
 
 

fastvideo/training/distillation_pipeline.py

@@ -245,7 +245,9 @@ class DistillationPipeline(TrainingPipeline):
 
         self.generator_ema: EMA_FSDP | None = None
         self.generator_ema_2: EMA_FSDP | None = None
-        if (self.training_args.ema_decay is not None) and (self.training_args.ema_decay > 0.0):
+        ema_enabled = (self.training_args.ema_decay is not None) and (self.training_args.ema_decay > 0.0)
+        if ema_enabled and (self.training_args.ema_start_step <= 0):
+            # Only eager-construct from the cold init weights when averaging starts at step 0.
             self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay)
             logger.info("Initialized generator EMA with decay=%s", self.training_args.ema_decay)
 
@@ -253,6 +255,12 @@ class DistillationPipeline(TrainingPipeline):
             if self.transformer_2 is not None:
                 self.generator_ema_2 = EMA_FSDP(self.transformer_2, decay=self.training_args.ema_decay)
                 logger.info("Initialized generator EMA_2 with decay=%s", self.training_args.ema_decay)
+        elif ema_enabled:
+            # Defer construction to the lazy block in the train loop, which builds the EMA AT
+            # ema_start_step from the already-trained weights. Eager-constructing here would anchor
+            # the shadow to the cold init and leave it base-contaminated (blurry) on short runs.
+            logger.info("Generator EMA deferred: built lazily at ema_start_step=%s from trained weights",
+                        self.training_args.ema_start_step)
         else:
             logger.info("Generator EMA disabled (ema_decay <= 0.0)")
 
@@ -326,22 +334,6 @@ class DistillationPipeline(TrainingPipeline):
                 return model
         return model
 
-    def get_ema_model_copy(self) -> torch.nn.Module | None:
-        """Get a copy of the model with EMA weights applied."""
-        if self.generator_ema is not None:
-            ema_model = copy.deepcopy(self.transformer)
-            self.generator_ema.copy_to_unwrapped(ema_model)
-            return ema_model
-        return None
-
-    def get_ema_2_model_copy(self) -> torch.nn.Module | None:
-        """Get a copy of the transformer_2 model with EMA weights applied."""
-        if self.generator_ema_2 is not None and self.transformer_2 is not None:
-            ema_2_model = copy.deepcopy(self.transformer_2)
-            self.generator_ema_2.copy_to_unwrapped(ema_2_model)
-            return ema_2_model
-        return None
-
     def is_ema_ready(self, current_step: int | None = None):
         """Check if EMA is ready for use (after ema_start_step)."""
         if current_step is None:
@@ -361,68 +353,61 @@ class DistillationPipeline(TrainingPipeline):
         try:
             # Save main transformer EMA
             if self.generator_ema is not None:
-                ema_model = self.get_ema_model_copy()
-                if ema_model is None:
-                    logger.warning("Failed to create EMA model copy")
-                else:
-                    ema_save_dir = os.path.join(output_dir, f"ema_checkpoint-{step}")
-                    os.makedirs(ema_save_dir, exist_ok=True)
+                ema_save_dir = os.path.join(output_dir, f"ema_checkpoint-{step}")
+                os.makedirs(ema_save_dir, exist_ok=True)
 
-                    # save as diffusers format
-                    from safetensors.torch import save_file
+                # save as diffusers format
+                from safetensors.torch import save_file
 
-                    from fastvideo.training.training_utils import (custom_to_hf_state_dict,
-                                                                   gather_state_dict_on_cpu_rank0)
-                    cpu_state = gather_state_dict_on_cpu_rank0(ema_model, device=None)
+                from fastvideo.training.training_utils import (custom_to_hf_state_dict, gather_state_dict_on_cpu_rank0)
+                # Swap EMA weights into the live FSDP module in place (no deepcopy) and gather the
+                # full state dict within the context; weights are restored on exit.
+                with self.generator_ema.apply_to_model(self.transformer):
+                    cpu_state = gather_state_dict_on_cpu_rank0(self.transformer, device=None)
 
-                    if self.global_rank == 0:
-                        weight_path = os.path.join(ema_save_dir, "diffusion_pytorch_model.safetensors")
-                        diffusers_state_dict = custom_to_hf_state_dict(cpu_state, ema_model.reverse_param_names_mapping)
-                        save_file(diffusers_state_dict, weight_path)
+                if self.global_rank == 0:
+                    weight_path = os.path.join(ema_save_dir, "diffusion_pytorch_model.safetensors")
+                    diffusers_state_dict = custom_to_hf_state_dict(cpu_state,
+                                                                   self.transformer.reverse_param_names_mapping)
+                    save_file(diffusers_state_dict, weight_path)
 
-                        config_dict = ema_model.hf_config
-                        if "dtype" in config_dict:
-                            del config_dict["dtype"]
-                        config_path = os.path.join(ema_save_dir, "config.json")
-                        with open(config_path, "w") as f:
-                            json.dump(config_dict, f, indent=4)
+                    # deepcopy so deleting "dtype" doesn't mutate the live model's hf_config
+                    config_dict = copy.deepcopy(self.transformer.hf_config)
+                    if "dtype" in config_dict:
+                        del config_dict["dtype"]
+                    config_path = os.path.join(ema_save_dir, "config.json")
+                    with open(config_path, "w") as f:
+                        json.dump(config_dict, f, indent=4)
 
-                        logger.info("EMA weights saved to %s", weight_path)
-
-                    del ema_model
+                    logger.info("EMA weights saved to %s", weight_path)
 
             # Save transformer_2 EMA
-            if self.generator_ema_2 is not None:
-                ema_2_model = self.get_ema_2_model_copy()
-                if ema_2_model is None:
-                    logger.warning("Failed to create EMA_2 model copy")
-                else:
-                    ema_2_save_dir = os.path.join(output_dir, f"ema_2_checkpoint-{step}")
-                    os.makedirs(ema_2_save_dir, exist_ok=True)
+            if self.generator_ema_2 is not None and self.transformer_2 is not None:
+                ema_2_save_dir = os.path.join(output_dir, f"ema_2_checkpoint-{step}")
+                os.makedirs(ema_2_save_dir, exist_ok=True)
 
-                    # save as diffusers format
-                    from safetensors.torch import save_file
+                # save as diffusers format
+                from safetensors.torch import save_file
 
-                    from fastvideo.training.training_utils import (custom_to_hf_state_dict,
-                                                                   gather_state_dict_on_cpu_rank0)
-                    cpu_state_2 = gather_state_dict_on_cpu_rank0(ema_2_model, device=None)
+                from fastvideo.training.training_utils import (custom_to_hf_state_dict, gather_state_dict_on_cpu_rank0)
+                with self.generator_ema_2.apply_to_model(self.transformer_2):
+                    cpu_state_2 = gather_state_dict_on_cpu_rank0(self.transformer_2, device=None)
 
-                    if self.global_rank == 0:
-                        weight_path_2 = os.path.join(ema_2_save_dir, "diffusion_pytorch_model.safetensors")
-                        diffusers_state_dict_2 = custom_to_hf_state_dict(cpu_state_2,
-                                                                         ema_2_model.reverse_param_names_mapping)
-                        save_file(diffusers_state_dict_2, weight_path_2)
+                if self.global_rank == 0:
+                    weight_path_2 = os.path.join(ema_2_save_dir, "diffusion_pytorch_model.safetensors")
+                    diffusers_state_dict_2 = custom_to_hf_state_dict(cpu_state_2,
+                                                                     self.transformer_2.reverse_param_names_mapping)
+                    save_file(diffusers_state_dict_2, weight_path_2)
 
-                        config_dict_2 = ema_2_model.hf_config
-                        if "dtype" in config_dict_2:
-                            del config_dict_2["dtype"]
-                        config_path_2 = os.path.join(ema_2_save_dir, "config.json")
-                        with open(config_path_2, "w") as f:
-                            json.dump(config_dict_2, f, indent=4)
+                    # deepcopy so deleting "dtype" doesn't mutate the live model's hf_config
+                    config_dict_2 = copy.deepcopy(self.transformer_2.hf_config)
+                    if "dtype" in config_dict_2:
+                        del config_dict_2["dtype"]
+                    config_path_2 = os.path.join(ema_2_save_dir, "config.json")
+                    with open(config_path_2, "w") as f:
+                        json.dump(config_dict_2, f, indent=4)
 
-                        logger.info("EMA_2 weights saved to %s", weight_path_2)
-
-                    del ema_2_model
+                    logger.info("EMA_2 weights saved to %s", weight_path_2)
 
         except Exception as e:
             logger.error("Failed to save EMA weights: %s", str(e))

pyproject.toml

@@ -106,6 +106,7 @@ torchaudio = [
   { index = "pytorch-cu128", marker = "sys_platform == 'linux'" },
 ]
 imagebind = { git = "https://github.com/facebookresearch/ImageBind.git", rev = "53680b02d7e37b19b124fa37bae4b6c98c38f5be" }
+flash-attn-cute = { git = "https://github.com/XOR-op/flash-attention.git", branch = "fa4-compile", subdirectory = "flash_attn/cute" }
 
 [[tool.uv.index]]
 name = "pytorch-cpu"
@@ -142,10 +143,13 @@ test = [
 # separately as documented in the README. decord has no aarch64 wheels
 # (upstream effectively unmaintained); it's split into its own extra so
 # the main eval extras install on ARM systems too. eval/io/video.py has
-# a PyAV fallback that covers the video-decode path without decord.
+# a PyAV fallback that covers the video-decode path without decord. The
+# judge.* metrics (eval-judge) call a remote API and need a key, so they
+# are opt-in and intentionally kept out of the default [eval] set.
 eval-vbench     = ["openai-clip", "pyiqa", "easydict"]
 eval-physics-iq = []
 eval-audio = ["jiwer", "librosa", "pyloudnorm", "hear21passt", "audiobox_aesthetics", "imagebind", "pytorchvideo"]
+eval-judge = ["google-genai"]
 eval = [
     "lpips", "ptlflow", "qwen-vl-utils",
     "fastvideo[eval-vbench]", "fastvideo[eval-physics-iq]",
@@ -172,7 +176,10 @@ streaming = [
 dreamverse = [
     "uvicorn[standard]>=0.41.0",
     "cerebras-cloud-sdk",
-    "flash-attn-cute @ git+https://github.com/XOR-op/flash-attention.git@fa4-compile#subdirectory=flash_attn/cute",
+    # PyPI forbids direct URL deps in published metadata; pin the fork via
+    # [tool.uv.sources] above (same pattern as imagebind) so the wheel stays
+    # publishable while uv workspace installs still get the fork.
+    "flash-attn-cute",
     "flashinfer-python",
     "openai>=1.40",
 ]