[ci] Replace flaky LTX-2 pixel SSIM with latent-slice cosine regression (#1253)

Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>

.agents/skills/seed-ssim-references/SKILL.md

@@ -1,26 +1,41 @@
 ---
 name: seed-ssim-references
-description: Seed HF reference videos for a single newly-added SSIM test. Runs the test on Modal L40S, downloads the generated mp4s via `modal volume get`, pauses for the user to eyeball quality, then uploads only that test's files to `FastVideo/ssim-reference-videos`. Use when a new `fastvideo/tests/ssim/test_*_similarity.py` has just been added and has no references on HF yet.
+description: Seed HF reference artefacts for a single newly-added SSIM test (pixel `.mp4` for `run_text_to_video_similarity_test`-style tests, or latent `.pt` for `run_text_to_latent_similarity_test`-style tests). Runs the test on Modal L40S, downloads the generated artefacts via `modal volume get`, pauses for the user to verify (visual eyeball for mp4, numerics dump for pt), then uploads only that test's files to `FastVideo/ssim-reference-videos`. Use when a new `fastvideo/tests/ssim/test_*_similarity.py` has just been added and has no references on HF yet.
 ---
 
-# Seed SSIM Reference Videos
+# Seed SSIM Reference Artefacts (mp4 or pt)
 
 ## Purpose
 
 A brand-new SSIM test in `fastvideo/tests/ssim/` fails forever until its
-reference videos exist on the HF dataset (`FastVideo/ssim-reference-videos`).
+reference artefacts exist on the HF dataset
+(`FastVideo/ssim-reference-videos`). The dataset hosts two kinds of artefacts
+side-by-side per `(model_id, backend, prompt)`:
+
+- **`.mp4`** — pixel ground-truth for tests that call
+  `run_text_to_video_similarity_test` / `run_image_to_video_similarity_test`
+  in `inference_similarity_utils.py`. Compared via SSIM.
+- **`.pt`** — pre-VAE latent bundle (fp16 full latent + fp32 slice +
+  metadata + `slice_spec` + `format_version`) for tests that call
+  `run_text_to_latent_similarity_test` in `latent_similarity_utils.py`.
+  Compared via cosine distance on the slice and the full tensor.
+
 This skill:
 
-1. Runs the test on Modal's L40S pool to generate the videos.
-2. Downloads them to the local repo via `modal volume get`.
-3. Pauses so the user can eyeball the mp4s and confirm quality.
-4. Uploads only the new test's files to HF, with a guard that refuses to
+1. Detects which artefact type the test produces (pixel vs latent).
+2. Runs the test on Modal's L40S pool to generate the artefacts.
+3. Downloads them to the local repo via `modal volume get`.
+4. Pauses so the user can verify quality:
+   - **mp4**: visual eyeball in a video player.
+   - **pt**: numerics dump (shape, slice stats, NaN/Inf check, metadata).
+5. Uploads only the new test's files to HF, with a guard that refuses to
    overwrite anything already present.
 
 The skill is run **manually**, once per new test. Before invoking it, the user
 has already sanity-tested the new test locally — it launches `VideoGenerator`
-and writes an mp4 without crashing. The skill does not re-test locally; it
-goes straight to Modal L40S (which is what CI uses).
+and writes an artefact without crashing (the missing-reference assertion at
+the end is expected). The skill does not re-test locally; it goes straight
+to Modal L40S (which is what CI uses).
 
 ## When to use
 
@@ -69,7 +84,7 @@ Fail fast if the token env var is missing.
 
 ## Steps
 
-### 1. Ask for the test file
+### 1. Ask for the test file, then detect artefact type
 
 If the user didn't name one, ask: *"Which SSIM test file do you want to seed
 references for? (e.g. `fastvideo/tests/ssim/test_ltx2_similarity.py`)"*.
@@ -80,6 +95,22 @@ Validate:
 - File defines a `*_MODEL_TO_PARAMS` dict — grep it to extract the set of
   model ids. Those ids drive step 5.
 
+Detect artefact type by inspecting the file's imports / helper call:
+
+- **latent** (`.pt`) — file imports `run_text_to_latent_similarity_test`
+  from `fastvideo.tests.ssim.latent_similarity_utils` (or any other helper
+  that ends with `_latent_similarity_test`).
+- **pixel** (`.mp4`) — file imports
+  `run_text_to_video_similarity_test` / `run_image_to_video_similarity_test`
+  from `fastvideo.tests.ssim.inference_similarity_utils`, OR uses the
+  legacy custom-inline helper pattern (see `test_gamecraft`,
+  `test_longcat`, etc.). Default to pixel when both heuristics fail.
+
+Record `ARTEFACT_TYPE ∈ {pixel, latent}` for use in step 4. Steps 2, 3, 5,
+and 6 are artefact-type-agnostic — `_iter_reference_files`,
+`copy_generated_to_reference`, and `upload_reference_videos` already walk
+both `.mp4` and `.pt` (see `reference_videos_cli.py`).
+
 If either check fails, stop and tell the user what's wrong.
 
 ### 2. Run the test on Modal L40S
@@ -166,17 +197,59 @@ get` preserves that trailing `generated_videos/` segment.
 
 ### 4. PAUSE — user reviews quality
 
-Print the list of downloaded mp4s and their paths, then stop. Tell the user:
+Type-aware verification.
+
+**For `ARTEFACT_TYPE = pixel`** — list the downloaded mp4s and ask the user to
+open them in a video player:
 
 > "Generated videos downloaded to `./generated_videos_modal/default/generated_videos/L40S_reference_videos/`. Please open them and confirm the quality looks correct. Reply **`upload`** to continue, or anything else to abort."
 
+**For `ARTEFACT_TYPE = latent`** — `.pt` files are not human-watchable. Print
+a numerics dump for each `.pt` so the user can sanity-check shape, distribution,
+and metadata:
+
+```python
+import torch
+from pathlib import Path
+ROOT = Path("./generated_videos_modal/default/generated_videos/L40S_reference_videos")
+for p in sorted(ROOT.rglob("*.pt")):
+    d = torch.load(p, map_location="cpu", weights_only=False)
+    s = d["expected_slice"]
+    L = d["latent"].float()
+    print(f"=== {p.relative_to(ROOT)} ===")
+    print(f"  format_version: {d['format_version']}")
+    print(f"  shape:          {d['shape']}")
+    print(f"  dtype_original: {d['dtype_original']}")
+    print(f"  slice_spec:     {d['slice_spec']}")
+    print(f"  slice  shape={tuple(s.shape)}  mean={s.mean():+.4f}  std={s.std():.4f}  min={s.min():+.4f}  max={s.max():+.4f}")
+    print(f"  latent shape={tuple(L.shape)}  mean={L.mean():+.4f}  std={L.std():.4f}  min={L.min():+.4f}  max={L.max():+.4f}")
+    print(f"  finite: latent NaN={torch.isnan(L).any().item()} Inf={torch.isinf(L).any().item()}; "
+          f"slice NaN={torch.isnan(s).any().item()} Inf={torch.isinf(s).any().item()}")
+    print(f"  metadata: {d['metadata']}\n")
+```
+
+Sanity criteria:
+- `format_version == 1` (matches `LATENT_REFERENCE_FORMAT_VERSION`).
+- `shape` matches what the model produces (e.g. LTX-2 distilled =
+  `[1, 128, T_lat, H_lat, W_lat]`; Stable Audio Open 1.0 = `[1, 64, 1024]`).
+- `slice_spec.kind` matches a registered kind (`corner_3x3_first_frame`
+  for video, `audio_first_8_timesteps` for audio).
+- No `NaN`/`Inf`. `mean ≈ 0`, `std ≈ 1` (denoised latents stay close to
+  the initial Gaussian distribution; very wide deviations suggest
+  numerical drift).
+- `metadata.prompt` matches the test's prompt.
+
+Then ask:
+
+> "Numerics look right? Reply **`upload`** to continue, or anything else to abort."
+
 Do not proceed until the user explicitly says `upload`. If they abort, leave
 everything on disk so they can inspect further — no cleanup.
 
 ### 5. Copy into the local reference layout
 
-Scoped copy — only the new test's mp4s. Loop over each `<model_id>` extracted
-in step 1:
+Scoped copy — only the new test's artefacts. Single command works for both
+artefact types because `_iter_reference_files` walks `.mp4` and `.pt`:
 
 ```bash
 python fastvideo/tests/ssim/reference_videos_cli.py copy-local \
@@ -186,12 +259,13 @@ python fastvideo/tests/ssim/reference_videos_cli.py copy-local \
 ```
 
 (The `--generated-dir` points at the device-folder root inside the
-downloaded tree; `copy-local` walks all `<model>/<backend>/*.mp4`
+downloaded tree; `copy-local` walks all `<model>/<backend>/*.{mp4,pt}`
 underneath it. Since the Modal run was scoped to a single test file via
 `--test-files`, only that test's model(s) are present — so the copy is
 implicitly per-test.)
 
-Result: `fastvideo/tests/ssim/reference_videos/default/L40S_reference_videos/<model_id>/<backend>/<prompt>.mp4`.
+Result for pixel: `fastvideo/tests/ssim/reference_videos/default/L40S_reference_videos/<model_id>/<backend>/<prompt>.mp4`.
+Result for latent: same path with `.pt` extension.
 
 ### 6. Upload to HF — scoped per model_id, with overwrite guard
 
@@ -224,33 +298,54 @@ it will auto-download the refs they just uploaded.
 ## Failure modes and how to handle them
 
 - **`HF_API_KEY` unset.** Stop before step 2. The Modal run needs it (passed
-  via `--hf-api-key`), and step 6 needs it for upload.
-- **Modal run fails before generation.** No mp4s on the volume — nothing to
-  download. Fix the test locally (`pytest fastvideo/tests/ssim/<test_file>`)
+  via `--hf-api-key`), and step 6 needs it for upload. If the user
+  ran `hf auth login` instead of exporting an env var, read the cached
+  token via `huggingface_hub.get_token()` and forward it to Modal as
+  `--hf-api-key="$CACHED_TOKEN"`.
+- **Modal run fails before generation.** No artefacts on the volume — nothing
+  to download. Fix the test locally (`pytest fastvideo/tests/ssim/<test_file>`)
   and retry from step 2.
 - **`./generated_videos_modal/default/L40S_reference_videos/` missing after
-  `modal volume get`.** The run didn't produce videos (most likely the test
-  crashed before writing, or `REQUIRED_GPUS` exceeded the partition capacity
-  — see Modal logs).
+  `modal volume get`.** The run didn't produce artefacts (most likely the
+  test crashed before writing, or `REQUIRED_GPUS` exceeded the partition
+  capacity — see Modal logs).
+- **Latent test crashed with FSDP / inference_mode error
+  (`RuntimeError: Inference tensors do not track version counter`).** The
+  test must pass `init_kwargs_override={"use_fsdp_inference": False}` when
+  `sp_size == 1` — see `test_stable_audio_similarity.py` for the pattern.
+  Fix in the test, push, retry.
 - **Upload guard fires (files already exist).** The test name / model id
   collides with something already on HF. Verify the user actually wants to
   replace existing refs; if so, re-run the upload with `--force`. If not,
   rename the model id in `*_MODEL_TO_PARAMS` and re-seed.
-- **Quality looks wrong in step 4.** Abort. The mp4s stay on disk for
+- **Quality looks wrong in step 4.** Abort. The artefacts stay on disk for
   inspection. The fix is usually in the test's params (resolution, steps,
   seed) — edit the test, then re-run the skill.
+  - For latent: also check `slice_spec.kind` matches the latent rank
+    (`corner_3x3_first_frame` requires 5-D, `audio_first_8_timesteps`
+    requires 3-D); a rank/kind mismatch raises in `_extract_expected_slice`.
 
 ## Design notes (for future skill maintainers)
 
 - The skill deliberately runs on Modal, **not** locally, because the CI
   runner is L40S. Seeding from a different GPU SKU produces refs that CI's
-  L40S runs can't match (SSIM drifts across SKUs).
+  L40S runs can't match (pixel SSIM drifts across SKUs; latent cosine has
+  tighter cross-SKU bf16 drift but the configured tolerances assume
+  same-SKU seed → same-SKU verify).
 - The skill is default-tier only. `full_quality` refs are seeded by a
   separate, deliberate operation — they double runtime and aren't what CI
   gates on.
 - The overwrite guard in `reference_videos_cli.py upload` is default-on
   specifically because this skill exists. Re-seeding is a distinct operation
   that requires explicit `--force`.
+- Both artefact types share the same Modal flow: the orchestrator sets
+  `--skip-reference-download` + `--no-fail-fast`, runs pytest, the test's
+  helper writes the artefact (`.mp4` via `imageio` for pixel,
+  `save_latent_reference` → `torch.save` for latent) BEFORE the
+  missing-reference assertion raises. `_sync_generated_videos_to_volume` in
+  `ssim_test.py` does a `shutil.copytree` of the whole `generated_videos/`
+  tree, picking up `.mp4`, `.pt`, and the `*_ssim.json` / `*_latent.json`
+  metric files alongside.
 
 ## References
 
@@ -259,10 +354,17 @@ it will auto-download the refs they just uploaded.
   `--skip-reference-download`, `--no-fail-fast`.
 - `fastvideo/tests/ssim/reference_videos_cli.py` — `copy-local`, `upload`
   (with `--model-id`, `--force`), `download`, `ensure` subcommands.
+  Extension allowlist is `REFERENCE_EXTENSIONS = VIDEO_EXTENSIONS +
+  LATENT_EXTENSIONS` (`.pt`).
 - `fastvideo/tests/ssim/README.md` — reference layout, HF repo conventions.
-- `fastvideo/tests/ssim/inference_similarity_utils.py` —
-  `run_text_to_video_similarity_test` + `_build_init_kwargs`: what each test
-  config passes to `VideoGenerator.from_pretrained`.
+- `fastvideo/tests/ssim/inference_similarity_utils.py` — pixel helpers
+  (`run_text_to_video_similarity_test`,
+  `run_image_to_video_similarity_test`, `build_init_kwargs`).
+- `fastvideo/tests/ssim/latent_similarity_utils.py` — latent helper
+  (`run_text_to_latent_similarity_test`), slice spec dispatch
+  (`_extract_expected_slice`), reference schema
+  (`save_latent_reference` / `load_latent_reference`),
+  `LATENT_REFERENCE_FORMAT_VERSION`.
 
 ## Changelog
 
@@ -271,3 +373,4 @@ it will auto-download the refs they just uploaded.
 | 2026-04-17 | Initial version (Modal sync-to-volume flow). |
 | 2026-04-21 | Rewrite: single-test scope, explicit user-review pause, per-`model_id` upload, HF overwrite guard. Dropped `scripts/seed_ssim.sh`. |
 | 2026-04-21 | Post-first-run fixes: `modal volume get` needs `--force` when parent exists; download tree has an extra `generated_videos/` level so `--generated-dir` must reflect it. |
+| 2026-05-01 | Latent (`*.pt`) artefact support: artefact-type detection in step 1, type-aware verification (visual eyeball for mp4, numerics dump for pt) in step 4, FSDP+inference_mode failure-mode added, design notes for the unified Modal flow. Triggered by PR #1253 (LTX-2 latent migration + Stable Audio latent test). |

fastvideo/entrypoints/video_generator.py

@@ -65,6 +65,7 @@ _FROM_PRETRAINED_CONVENIENCE_KWARGS = frozenset({
     "pin_cpu_memory",
     "enable_torch_compile",
     "torch_compile_kwargs",
+    "output_type",
 })
 
 
@@ -601,10 +602,20 @@ class VideoGenerator:
         thread = threading.Thread(target=execute_forward_thread)
         thread.start()
         latent_batch_size = _infer_latent_batch_size(batch)
-        samples = torch.empty(
-            (latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width),
-            device='cpu',
-            pin_memory=fastvideo_args.pin_cpu_memory)
+        # When ``output_type == "latent"`` the forward output has latent
+        # shape (e.g. ``[B, C_latent, T_latent, H_latent, W_latent]``)
+        # rather than the pre-allocation's pixel shape. Skip the pinned
+        # ~50 MB buffer entirely; we always fall through to the
+        # ``samples = output_batch.output.cpu()`` branch below in that
+        # mode. ``skip_pixel_prealloc`` also gates the slow-path warning.
+        skip_pixel_prealloc = fastvideo_args.output_type == "latent"
+        if skip_pixel_prealloc:
+            samples = torch.empty(0, device='cpu')
+        else:
+            samples = torch.empty(
+                (latent_batch_size, 3, sampling_param.num_frames, sampling_param.height, sampling_param.width),
+                device='cpu',
+                pin_memory=fastvideo_args.pin_cpu_memory)
         thread.join()
 
         if thread_error["error"] is not None:
@@ -619,29 +630,44 @@ class VideoGenerator:
         if output_batch.output.shape == samples.shape:
             samples.copy_(output_batch.output)
         else:
-            logger.warning("Output shape %s does not match expected shape %s; use slow path", output_batch.output.shape,
-                           samples.shape)
+            if not skip_pixel_prealloc:
+                logger.warning("Output shape %s does not match expected shape %s; use slow path",
+                               output_batch.output.shape, samples.shape)
             samples = output_batch.output.cpu()
         logging_info = output_batch.logging_info
 
         gen_time = time.perf_counter() - start_time
         logger.info("Generated successfully in %.2f seconds", gen_time)
 
-        # Process outputs (skip the make_grid loop for audio-only, where
-        # `samples` is a 1×3×1×8×8 placeholder no caller will use).
+        # Three mutually-exclusive output modes determine whether (a) we
+        # build an RGB frame buffer and (b) what file we write to disk:
+        #
+        #   1. `output_type == "latent"` — VAE is bypassed in DecodingStage
+        #      and `samples` holds raw latents (arbitrary channel count).
+        #      The RGB grid / uint8 / mp4 / png pipeline below cannot
+        #      consume those, so we skip it entirely and let callers work
+        #      with the latent tensor directly via `result["samples"]`.
+        #   2. Audio-only workload — `samples` is a 1×3×1×8×8 placeholder
+        #      no caller will use; skip the grid loop and save a `.wav`.
+        #   3. Pixel video / image — the historical happy path.
+        is_latent_output = fastvideo_args.output_type == "latent"
         audio_only = bool(output_batch.extra.get("audio_only"))
-        frames: list[np.ndarray] = []
-        if not audio_only:
+
+        frames: list[np.ndarray] | None
+        if is_latent_output or audio_only:
+            frames = None if is_latent_output else []
+        else:
             videos = rearrange(samples, "b c t h w -> t b c h w")
+            frames = []
             for x in videos:
                 x = torchvision.utils.make_grid(x, nrow=6)
                 x = x.permute(1, 2, 0).squeeze(-1)
                 x = (x * 255).to(torch.uint8)
                 frames.append(x.cpu().numpy())
 
-        # Save output if requested
-        if batch.save_video:
-            if output_batch.extra.get("audio_only"):
+        save_to_disk = batch.save_video and not is_latent_output
+        if save_to_disk:
+            if audio_only:
                 # Audio-only workload: write a standalone .wav rather than
                 # muxing the audio into a placeholder mp4 (which forces
                 # ffmpeg to round 8x8 placeholder frames up to 16x16).
@@ -654,9 +680,11 @@ class VideoGenerator:
                 logger.info("Saved audio to %s", output_path)
             elif self._is_image_workload():
                 # Image workloads (t2i, i2i, …): save the first frame as PNG.
+                assert frames is not None  # implied by save_to_disk and not audio_only
                 imageio.imwrite(output_path, frames[0])
                 logger.info("Saved image to %s", output_path)
             else:
+                assert frames is not None  # implied by save_to_disk and not audio_only
                 imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4")
                 logger.info("Saved video to %s", output_path)
                 audio = output_batch.extra.get("audio")
@@ -680,7 +708,7 @@ class VideoGenerator:
             "trajectory": output_batch.trajectory_latents,
             "trajectory_timesteps": output_batch.trajectory_timesteps,
             "trajectory_decoded": output_batch.trajectory_decoded,
-            "video_path": output_path if batch.save_video else None,
+            "video_path": output_path if save_to_disk else None,
             "peak_memory_mb": output_batch.extra.get("peak_memory_mb"),
         }
 

fastvideo/pipelines/basic/stable_audio/stages/decoding.py

@@ -33,6 +33,15 @@ class StableAudioDecodingStage(PipelineStage):
         pc = fastvideo_args.pipeline_config
         latents = batch.latents
 
+        # Latent regression path: hand back the un-decoded denoised latent
+        # so the LatentSimilarityUtils harness can compare on pre-VAE
+        # numerics. Mirrors the bypass in `pipelines/stages/decoding.py`
+        # used by the video DiTs. Skips the `.to(device)` VAE move so we
+        # don't pay decoder load cost on this shortcut path.
+        if fastvideo_args.output_type == "latent":
+            batch.output = latents.detach().cpu()
+            return batch
+
         # VAE may be CPU-parked under `vae_cpu_offload=True`.
         from fastvideo.distributed.parallel_state import get_local_torch_device
         self.vae = self.vae.to(get_local_torch_device())

fastvideo/tests/modal/ssim_test.py

@@ -485,6 +485,10 @@ def _prepare_ssim_workspace(
     ./build.sh
     cd ..
     uv pip install git+https://github.com/microsoft/MoGe.git
+    # Stable Audio Open 1.0 inference deps (optional in basic install,
+    # required by `StableAudioDenoisingStage`; consumed by
+    # `test_stable_audio_similarity.py`).
+    uv pip install k_diffusion einops_exts alias_free_torch torchsde
     export HF_HOME='/root/data/.cache'
     hf auth login --token "$HF_API_KEY"
     """

fastvideo/tests/ssim/inference_similarity_utils.py

@@ -132,7 +132,7 @@ def _assert_similarity(
     )
 
 
-def _build_init_kwargs(
+def build_init_kwargs(
     base_params: dict[str, object],
 ) -> dict[str, object]:
     init_kwargs: dict[str, object] = {
@@ -167,7 +167,7 @@ def _build_init_kwargs(
     return init_kwargs
 
 
-def _build_generation_kwargs(
+def build_generation_kwargs(
     base_params: dict[str, object],
     num_inference_steps: int,
     output_dir: str,
@@ -222,11 +222,11 @@ def run_text_to_video_similarity_test(
         base_params = params_map[model_id]
         num_inference_steps = int(base_params["num_inference_steps"])
 
-        init_kwargs = _build_init_kwargs(base_params)
+        init_kwargs = build_init_kwargs(base_params)
         if init_kwargs_override:
             init_kwargs.update(init_kwargs_override)
 
-        generation_kwargs = _build_generation_kwargs(
+        generation_kwargs = build_generation_kwargs(
             base_params,
             num_inference_steps,
             output_dir,
@@ -297,11 +297,11 @@ def run_image_to_video_similarity_test(
         base_params = params_map[model_id]
         num_inference_steps = int(base_params["num_inference_steps"])
 
-        init_kwargs = _build_init_kwargs(base_params)
+        init_kwargs = build_init_kwargs(base_params)
         if init_kwargs_override:
             init_kwargs.update(init_kwargs_override)
 
-        generation_kwargs = _build_generation_kwargs(
+        generation_kwargs = build_generation_kwargs(
             base_params,
             num_inference_steps,
             output_dir,

fastvideo/tests/ssim/latent_similarity_utils.py

@@ -0,0 +1,494 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Latent-space regression helpers for numerically fragile SSIM tests.
+
+Motivation
+----------
+Pixel-space SSIM is a poor regression signal for distilled / few-step
+models (e.g. LTX-2 distilled): a single mis-rounded bf16 accumulator in
+the VAE decoder can drive mean SSIM from ~0.95 to ~0.50 without any real
+quality regression.
+
+Inspired by diffusers' "small signature slice + bounded full-tensor
+distance" testing philosophy, applied here to the **pre-VAE latent**
+rather than the decoded pixel/audio output:
+
+* `tests/pipelines/ltx2/test_ltx2.py` (diffusers) compares
+  ``output_type='pt'`` (pixel) slices via
+  ``torch.allclose(generated_slice, expected_slice, atol=1e-4)``;
+* `tests/pipelines/stable_audio/test_stable_audio.py` (diffusers)
+  compares decoded audio samples via
+  ``np.abs(expected - actual).max() < 1.5e-3``;
+* `tests/pipelines/cogvideo/test_cogvideox.py` (diffusers) compares
+  full pixel video tensors via
+  ``numpy_cosine_similarity_distance(...) < 1e-3``.
+
+Diffusers does *not* assert on latents directly — that is a FastVideo
+adaptation. Distilled few-step pipelines amplify per-step bf16 noise
+enough that VAE-decoded comparisons are unreliable across our
+heterogeneous CI pool, so we move the assertion upstream of the VAE.
+
+Design
+------
+* Inference is run with ``output_type='latent'`` so ``DecodingStage``
+  hands back the un-decoded latent on ``result["samples"]``.
+* The reference artefact is a ``.pt`` bundle (tensor + metadata) hosted
+  on the same HF dataset as the mp4 references, selected by
+  ``<GPU>_reference_videos/<model_id>/<backend>/<prompt>.pt``.
+* Two assertions are performed:
+    1. A small signature slice (default video: ``latent[0, :, 0, :3, :3]``;
+       audio: ``latent[0, :, :8]``) is compared via cosine distance with
+       a loose tolerance. Primary pass/fail gate.
+    2. The full latent is compared via cosine distance with a slightly
+       tighter tolerance, guarding against shape-correct but globally
+       drifted outputs.
+* Tolerances default to 5e-3 (slice) and 1e-2 (full). diffusers uses
+  ``1e-3`` against deterministic CPU dummy components; we relax for
+  cross-GPU-arch bf16 differences on the rented CI pool
+  (A40/L40S/H100/B200).
+
+The helper intentionally reuses ``build_init_kwargs`` /
+``build_generation_kwargs`` from :mod:`inference_similarity_utils` so
+model params (vae tiling, sp_size, flow shift, …) flow through a single
+source of truth.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from logging import Logger
+from typing import Any
+
+import torch
+from torch.nn.functional import cosine_similarity
+
+from fastvideo import VideoGenerator
+from fastvideo.tests.ssim.inference_similarity_utils import (
+    attention_backend,
+    build_generation_kwargs,
+    build_init_kwargs,
+    shutdown_executor,
+)
+from fastvideo.tests.ssim.reference_utils import (
+    build_generated_output_dir,
+    build_reference_folder_path,
+    select_ssim_params,
+)
+
+LATENT_REFERENCE_EXTENSION = ".pt"
+LATENT_REFERENCE_FORMAT_VERSION = 1
+
+DEFAULT_SLICE_SPEC: dict[str, Any] = {
+    "kind": "corner_3x3_first_frame",
+    "version": 1,
+}
+
+AUDIO_FIRST_8_TIMESTEPS_SPEC: dict[str, Any] = {
+    "kind": "audio_first_8_timesteps",
+    "version": 1,
+}
+
+
+def _extract_expected_slice(
+    latent: torch.Tensor,
+    spec: dict[str, Any],
+) -> torch.Tensor:
+    """Return a 1-D fp32 signature slice from ``latent`` per ``spec``.
+
+    Dispatch by ``spec["kind"]``:
+
+    * ``"corner_3x3_first_frame"`` — 5-D video latent ``[B, C, T, H, W]``;
+      returns ``latent[0, :, 0, :3, :3]`` flattened (length = C * 9).
+    * ``"audio_first_8_timesteps"`` — 3-D audio latent ``[B, C, T]``;
+      returns ``latent[0, :, :8]`` flattened (length = C * 8).
+    """
+    kind = spec.get("kind", "corner_3x3_first_frame")
+    if kind == "corner_3x3_first_frame":
+        if latent.dim() != 5:
+            raise ValueError(
+                f"corner_3x3_first_frame requires 5-D latent [B,C,T,H,W]; "
+                f"got shape {tuple(latent.shape)}")
+        _, _, t, h, w = latent.shape
+        if t < 1 or h < 3 or w < 3:
+            raise ValueError(
+                "corner_3x3_first_frame requires T>=1, H>=3, W>=3; got "
+                f"shape {tuple(latent.shape)}")
+        return (latent[0, :, 0, :3, :3]
+                .detach()
+                .to(torch.float32)
+                .reshape(-1)
+                .contiguous())
+    if kind == "audio_first_8_timesteps":
+        if latent.dim() != 3:
+            raise ValueError(
+                f"audio_first_8_timesteps requires 3-D latent [B,C,T]; "
+                f"got shape {tuple(latent.shape)}")
+        _, _, t = latent.shape
+        if t < 8:
+            raise ValueError(
+                "audio_first_8_timesteps requires T>=8; got "
+                f"shape {tuple(latent.shape)}")
+        return (latent[0, :, :8]
+                .detach()
+                .to(torch.float32)
+                .reshape(-1)
+                .contiguous())
+    raise ValueError(f"Unknown slice kind: {kind!r}")
+
+
+def _cosine_distance(a: torch.Tensor, b: torch.Tensor) -> float:
+    """Return ``1 - cos(a, b)`` as a Python float, operating in fp32."""
+    a32 = a.detach().to(torch.float32).reshape(-1)
+    b32 = b.detach().to(torch.float32).reshape(-1)
+    if a32.shape != b32.shape:
+        raise ValueError(
+            f"Cosine shape mismatch: {tuple(a32.shape)} vs {tuple(b32.shape)}"
+        )
+    sim = cosine_similarity(a32.unsqueeze(0), b32.unsqueeze(0), dim=1).item()
+    return 1.0 - float(sim)
+
+
+def save_latent_reference(
+    path: str,
+    latent: torch.Tensor,
+    *,
+    metadata: dict[str, Any],
+    slice_spec: dict[str, Any] | None = None,
+) -> None:
+    """Persist a latent bundle to ``path``.
+
+    Storage format (dict pickled via ``torch.save``):
+
+    * ``latent``: full latent as fp16 on cpu
+    * ``shape``: original shape (list)
+    * ``dtype_original``: str
+    * ``expected_slice``: fp32 1-D signature slice
+    * ``slice_spec``: dict describing how the slice was built
+    * ``metadata``: caller-provided context (prompt, backend, steps, …)
+    * ``format_version``: int
+
+    fp16 is lossy but bounded; it keeps ref artefacts small (~a few MB
+    per prompt) while preserving enough dynamic range for cosine-based
+    regression. Slice values stay fp32 because the primary assertion is
+    computed against them.
+    """
+    spec = slice_spec if slice_spec is not None else DEFAULT_SLICE_SPEC
+    parent = os.path.dirname(path)
+    if parent:
+        os.makedirs(parent, exist_ok=True)
+
+    latent_cpu = latent.detach().to("cpu")
+    payload: dict[str, Any] = {
+        "latent": latent_cpu.to(torch.float16),
+        "shape": list(latent_cpu.shape),
+        "dtype_original": str(latent_cpu.dtype),
+        "expected_slice": _extract_expected_slice(latent_cpu, spec),
+        "slice_spec": spec,
+        "metadata": metadata,
+        "format_version": LATENT_REFERENCE_FORMAT_VERSION,
+    }
+    torch.save(payload, path)
+
+
+def load_latent_reference(path: str) -> dict[str, Any]:
+    """Inverse of :func:`save_latent_reference` — always loads to cpu.
+
+    Enforces ``format_version == LATENT_REFERENCE_FORMAT_VERSION`` so a
+    schema change forces a deliberate reseed instead of silently
+    misinterpreting old artefacts.
+    """
+    # ``weights_only=False`` is required because the payload is a dict of
+    # tensors + plain-Python metadata (slice_spec, prompt, ...). The trust
+    # boundary is the controlled HF dataset configured via
+    # FASTVIDEO_SSIM_REFERENCE_HF_REPO (default
+    # FastVideo/ssim-reference-videos), which is org-write-gated.
+    payload = torch.load(path, map_location="cpu", weights_only=False)
+    fmt = payload.get("format_version") if isinstance(payload, dict) else None
+    if fmt != LATENT_REFERENCE_FORMAT_VERSION:
+        raise ValueError(
+            f"Latent reference at {path!r} has format_version={fmt!r}; "
+            f"expected {LATENT_REFERENCE_FORMAT_VERSION}. Re-seed via the "
+            "test that produces this artefact, then re-upload through "
+            "fastvideo/tests/ssim/reference_videos_cli.py upload.")
+    return payload
+
+
+def write_latent_similarity_results(
+    output_dir: str,
+    metrics: dict[str, float],
+    *,
+    reference_path: str,
+    generated_path: str,
+    num_inference_steps: int,
+    prompt: str,
+    model_id: str,
+    attention_backend_name: str,
+    slice_spec: dict[str, Any],
+    slice_cosine_threshold: float,
+    full_cosine_threshold: float,
+    passed: bool,
+) -> bool:
+    """Persist latent regression metrics next to the generated artefact.
+
+    Mirrors :func:`fastvideo.tests.utils.write_ssim_results` so downstream
+    CI tooling can scrape one schema for both pixel and latent runs.
+    The filename is ``steps{N}_{prompt[:100]}_latent.json``.
+    """
+    try:
+        os.makedirs(output_dir, exist_ok=True)
+        prompt_prefix = prompt[:100].strip()
+        filename = f"steps{num_inference_steps}_{prompt_prefix}_latent.json"
+        target = os.path.join(output_dir, filename)
+        payload = {
+            "metrics": metrics,
+            "reference_latent": reference_path,
+            "generated_latent": generated_path,
+            "model_id": model_id,
+            "attention_backend": attention_backend_name,
+            "slice_spec": slice_spec,
+            "thresholds": {
+                "slice_cosine": slice_cosine_threshold,
+                "full_cosine": full_cosine_threshold,
+            },
+            "passed": passed,
+            "parameters": {
+                "num_inference_steps": num_inference_steps,
+                "prompt": prompt,
+            },
+        }
+        with open(target, "w", encoding="utf-8") as handle:
+            json.dump(payload, handle, indent=2, sort_keys=True)
+        return True
+    except OSError:
+        return False
+
+
+def _assert_latent_similarity(
+    *,
+    logger: Logger,
+    gen_latent: torch.Tensor,
+    reference_path: str,
+    slice_cosine_threshold: float,
+    full_cosine_threshold: float,
+    model_id: str,
+    attention_backend_name: str,
+    output_dir: str,
+    generated_path: str,
+    num_inference_steps: int,
+    prompt: str,
+) -> dict[str, float]:
+    ref = load_latent_reference(reference_path)
+    expected_slice = ref["expected_slice"]
+    ref_full = ref["latent"].to(torch.float32)
+    ref_shape = tuple(ref.get("shape", ref_full.shape))
+    slice_spec = ref.get("slice_spec", DEFAULT_SLICE_SPEC)
+
+    if tuple(gen_latent.shape) != ref_shape:
+        raise AssertionError(
+            f"Generated latent shape {tuple(gen_latent.shape)} does not "
+            f"match reference shape {ref_shape} for {model_id} with "
+            f"backend {attention_backend_name}")
+
+    gen_slice = _extract_expected_slice(gen_latent, slice_spec)
+    slice_cos = _cosine_distance(gen_slice, expected_slice)
+    # The reference full tensor was fp16-quantized at seed time. Round-trip
+    # the generated tensor through fp16 so both sides share the same
+    # quantization floor and the cosine distance is symmetric. This does
+    # NOT affect ``slice_cos`` because the reference slice is persisted in
+    # fp32 (see :func:`save_latent_reference`).
+    gen_full_matched = gen_latent.to(torch.float16).to(torch.float32)
+    full_cos = _cosine_distance(gen_full_matched, ref_full)
+    max_abs_diff = float(
+        (gen_full_matched - ref_full).abs().max().item())
+
+    metrics: dict[str, float] = {
+        "slice_cosine_distance": slice_cos,
+        "full_cosine_distance": full_cos,
+        "max_abs_diff": max_abs_diff,
+    }
+    logger.info(
+        "Latent regression metrics for %s/%s: %s",
+        model_id,
+        attention_backend_name,
+        metrics,
+    )
+
+    failures: list[str] = []
+    if slice_cos > slice_cosine_threshold:
+        failures.append(
+            f"slice cosine {slice_cos:.6e} > threshold "
+            f"{slice_cosine_threshold:.6e}")
+    if full_cos > full_cosine_threshold:
+        failures.append(
+            f"full cosine {full_cos:.6e} > threshold "
+            f"{full_cosine_threshold:.6e}")
+
+    passed = not failures
+    write_latent_similarity_results(
+        output_dir,
+        metrics,
+        reference_path=reference_path,
+        generated_path=generated_path,
+        num_inference_steps=num_inference_steps,
+        prompt=prompt,
+        model_id=model_id,
+        attention_backend_name=attention_backend_name,
+        slice_spec=slice_spec,
+        slice_cosine_threshold=slice_cosine_threshold,
+        full_cosine_threshold=full_cosine_threshold,
+        passed=passed,
+    )
+
+    if failures:
+        raise AssertionError(
+            f"Latent regression exceeded tolerance for {model_id} with "
+            f"backend {attention_backend_name}: {'; '.join(failures)}. "
+            f"Full metrics: {metrics}")
+
+    return metrics
+
+
+def _extract_latent_from_result(result: Any) -> torch.Tensor:
+    """Pull a fp32 cpu latent tensor (5-D video or 3-D audio) from
+    ``generate_video`` output.
+    """
+    if not isinstance(result, dict):
+        raise RuntimeError(
+            "VideoGenerator.generate_video returned unexpected payload "
+            f"(type={type(result)!r}); expected dict with 'samples'.")
+    samples = result.get("samples")
+    if samples is None:
+        raise RuntimeError(
+            "VideoGenerator did not return latent samples. Ensure "
+            "output_type='latent' and return_frames=True for this call.")
+    if not isinstance(samples, torch.Tensor):
+        raise RuntimeError(
+            f"Expected torch.Tensor samples; got type={type(samples)!r}.")
+    gen_latent = samples.detach().to(torch.float32).cpu()
+    if gen_latent.dim() not in (3, 5):
+        raise RuntimeError(
+            "Expected 5-D video latent (B,C,T,H,W) or 3-D audio latent "
+            f"(B,C,T); got shape {tuple(gen_latent.shape)}")
+    return gen_latent
+
+
+def run_text_to_latent_similarity_test(
+    *,
+    logger: Logger,
+    script_dir: str,
+    device_reference_folder: str,
+    prompt: str,
+    attention_backend_name: str,
+    model_id: str,
+    default_params_map: dict[str, dict[str, object]],
+    full_quality_params_map: dict[str, dict[str, object]],
+    slice_cosine_threshold: float = 5e-3,
+    full_cosine_threshold: float = 1e-2,
+    init_kwargs_override: dict[str, object] | None = None,
+    generation_kwargs_override: dict[str, object] | None = None,
+    slice_spec: dict[str, Any] | None = None,
+) -> dict[str, float]:
+    """Run T2V (or T2A) inference with ``output_type='latent'`` and
+    compare to a reference latent.
+
+    Returns the computed metrics dict on success. Raises
+    ``AssertionError`` if any cosine tolerance is exceeded and
+    ``FileNotFoundError`` if the reference artefact is missing.
+    """
+    spec = slice_spec if slice_spec is not None else DEFAULT_SLICE_SPEC
+    with attention_backend(attention_backend_name):
+        output_dir = build_generated_output_dir(
+            script_dir,
+            device_reference_folder,
+            model_id,
+            attention_backend_name,
+        )
+        prompt_prefix = prompt[:100].strip()
+        output_latent_name = f"{prompt_prefix}{LATENT_REFERENCE_EXTENSION}"
+        os.makedirs(output_dir, exist_ok=True)
+
+        params_map = select_ssim_params(
+            default_params_map,
+            full_quality_params_map,
+        )
+        base_params = params_map[model_id]
+        num_inference_steps = int(base_params["num_inference_steps"])
+
+        init_kwargs = build_init_kwargs(base_params)
+        if init_kwargs_override:
+            init_kwargs.update(init_kwargs_override)
+        # Always wins: the helper exists specifically to compare on latents,
+        # so an override can never silently turn it back into a pixel run.
+        init_kwargs["output_type"] = "latent"
+
+        generation_kwargs = build_generation_kwargs(
+            base_params,
+            num_inference_steps,
+            output_dir,
+        )
+        # We serialize latents ourselves; skip the RGB encoder path.
+        generation_kwargs["save_video"] = False
+        generation_kwargs["return_frames"] = True
+        if generation_kwargs_override:
+            generation_kwargs.update(generation_kwargs_override)
+
+        generator: VideoGenerator | None = None
+        try:
+            generator = VideoGenerator.from_pretrained(
+                model_path=base_params["model_path"],
+                **init_kwargs,
+            )
+            result = generator.generate_video(prompt, **generation_kwargs)
+        finally:
+            shutdown_executor(generator)
+
+    gen_latent = _extract_latent_from_result(result)
+
+    generated_latent_path = os.path.join(output_dir, output_latent_name)
+    save_latent_reference(
+        generated_latent_path,
+        gen_latent,
+        metadata={
+            "prompt": prompt,
+            "model_id": model_id,
+            "attention_backend": attention_backend_name,
+            "num_inference_steps": num_inference_steps,
+        },
+        slice_spec=spec,
+    )
+    logger.info("Saved generated latent to %s", generated_latent_path)
+
+    reference_folder = build_reference_folder_path(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        attention_backend_name,
+    )
+    if not os.path.exists(reference_folder):
+        raise FileNotFoundError(
+            f"Reference folder does not exist: {reference_folder}\n"
+            f"To download references, run:\n"
+            f"  python fastvideo/tests/ssim/reference_videos_cli.py download")
+
+    reference_latent_path = os.path.join(
+        reference_folder,
+        output_latent_name,
+    )
+    if not os.path.exists(reference_latent_path):
+        raise FileNotFoundError(
+            "Reference latent missing for prompt/backend: "
+            f"{reference_latent_path}")
+
+    return _assert_latent_similarity(
+        logger=logger,
+        gen_latent=gen_latent,
+        reference_path=reference_latent_path,
+        slice_cosine_threshold=slice_cosine_threshold,
+        full_cosine_threshold=full_cosine_threshold,
+        model_id=model_id,
+        attention_backend_name=attention_backend_name,
+        output_dir=output_dir,
+        generated_path=generated_latent_path,
+        num_inference_steps=num_inference_steps,
+        prompt=prompt,
+    )

fastvideo/tests/ssim/reference_videos_cli.py

@@ -13,6 +13,12 @@ from pathlib import Path
 
 
 VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv")
+# Additional artefact types stored under the same reference folders. Latent
+# tensors (`.pt`) back the latent-slice regression tests used by flaky
+# pixel-space models (e.g. LTX-2 distilled). Keeping them in the same upload
+# flow means seeding a new test only requires one HF round-trip.
+LATENT_EXTENSIONS = (".pt",)
+REFERENCE_EXTENSIONS = VIDEO_EXTENSIONS + LATENT_EXTENSIONS
 HF_TOKEN_ENV_KEYS = ("HF_API_KEY", "HUGGINGFACE_HUB_TOKEN", "HF_TOKEN")
 HF_REPO_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO"
 HF_REPO_TYPE_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE"
@@ -42,9 +48,14 @@ def _default_repo_type() -> str:
     return os.environ.get(HF_REPO_TYPE_ENV_KEY, DEFAULT_REPO_TYPE)
 
 
-def _iter_video_files(root: Path) -> Iterable[Path]:
+def _iter_reference_files(root: Path) -> Iterable[Path]:
+    """Yield video and latent (.pt) references under `root`.
+
+    Used by copy-local and the "has local references" marker probe so that
+    latent-only tests (no mp4) still satisfy readiness checks.
+    """
     for path in root.rglob("*"):
-        if path.is_file() and path.suffix.lower() in VIDEO_EXTENSIONS:
+        if path.is_file() and path.suffix.lower() in REFERENCE_EXTENSIONS:
             yield path
 
 
@@ -139,7 +150,7 @@ def _has_local_reference_videos(base_dir: Path, quality_tier: str) -> bool:
         return False
     tier_root = _reference_tier_root(base_dir, quality_tier)
     for ref_dir in _discover_reference_dirs(tier_root):
-        for _ in _iter_video_files(ref_dir):
+        for _ in _iter_reference_files(ref_dir):
             return True
     return False
 
@@ -173,14 +184,14 @@ def copy_generated_to_reference(
         raise FileNotFoundError(f"Generated directory not found: {generated_dir}")
 
     copied = 0
-    for video_file in _iter_video_files(generated_dir):
-        rel = video_file.relative_to(generated_dir)
+    for ref_file in _iter_reference_files(generated_dir):
+        rel = ref_file.relative_to(generated_dir)
         dst = reference_dir / rel
         if dry_run:
-            print(f"[dry-run] {video_file} -> {dst}")
+            print(f"[dry-run] {ref_file} -> {dst}")
         else:
             dst.parent.mkdir(parents=True, exist_ok=True)
-            shutil.copy2(video_file, dst)
+            shutil.copy2(ref_file, dst)
             print(f"Copied: {rel}")
         copied += 1
     return copied

fastvideo/tests/ssim/test_ltx2_similarity.py

@@ -1,11 +1,25 @@
 # SPDX-License-Identifier: Apache-2.0
-"""SSIM-based similarity test for LTX-2 distilled text-to-video.
+"""Latent-slice regression test for LTX-2 distilled text-to-video.
 
-Parameters derived from examples/inference/basic/basic_ltx2_distilled.py,
-with resolution + num_inference_steps reduced to keep GPU CI runtime
-bounded. Full-quality variant (via ``--ssim-full-quality``) falls back
-to the ``ltx2_distilled`` preset defaults.
+Pixel-space SSIM is not a useful signal for this model: 4 distilled
+steps + bf16 attention + tiled VAE decode produce outputs that pass
+visual QA but occupy a very wide region in pixel space.
+
+Inspired by diffusers' slice-vs-full regression philosophy — see
+``diffusers/tests/pipelines/ltx2/test_ltx2.py`` (compares pixel slices
+via ``torch.allclose(..., atol=1e-4)``) and
+``diffusers/tests/pipelines/cogvideo/test_cogvideox.py`` (full pixel
+tensors via ``numpy_cosine_similarity_distance(...) < 1e-3``).
+Diffusers itself does NOT compare latents; we apply the same "small
+signature slice + bounded full-tensor distance" idea to the **pre-VAE
+latent** because distilled few-step pipelines amplify per-step bf16
+noise enough that VAE-decoded comparisons are unreliable.
+
+Parameters are kept identical to the original SSIM run so that
+reference artefacts generated on Modal L40S remain bit-compatible with
+production inference.
 """
+
 import os
 
 import pytest
@@ -14,7 +28,9 @@ from fastvideo.api.sampling_param import SamplingParam
 from fastvideo.logger import init_logger
 from fastvideo.tests.ssim.inference_similarity_utils import (
     resolve_inference_device_reference_folder,
-    run_text_to_video_similarity_test,
+)
+from fastvideo.tests.ssim.latent_similarity_utils import (
+    run_text_to_latent_similarity_test,
 )
 
 logger = init_logger(__name__)
@@ -71,6 +87,13 @@ LTX2_DISTILLED_TEST_PROMPTS = [
     "deadpan, absurd, and quietly tragic.",
 ]
 
+# Tolerances chosen on top of diffusers' ``1e-3`` defaults. LTX-2 distilled
+# amplifies per-step numerical noise, and FastVideo's CI pool spans L40S /
+# A40 / H100 so cross-architecture bf16 drift must be absorbed. Values can
+# be tightened after an initial stable window of reference refreshes.
+SLICE_COSINE_DISTANCE_THRESHOLD = 5e-3
+FULL_COSINE_DISTANCE_THRESHOLD = 1e-2
+
 
 @pytest.mark.parametrize("prompt", LTX2_DISTILLED_TEST_PROMPTS)
 @pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"])
@@ -80,7 +103,7 @@ def test_ltx2_distilled_inference_similarity(
     attention_backend_name: str,
     model_id: str,
 ) -> None:
-    run_text_to_video_similarity_test(
+    run_text_to_latent_similarity_test(
         logger=logger,
         script_dir=os.path.dirname(os.path.abspath(__file__)),
         device_reference_folder=device_reference_folder,
@@ -89,5 +112,6 @@ def test_ltx2_distilled_inference_similarity(
         model_id=model_id,
         default_params_map=LTX2_DISTILLED_MODEL_TO_PARAMS,
         full_quality_params_map=FULL_QUALITY_LTX2_DISTILLED_MODEL_TO_PARAMS,
-        min_acceptable_ssim=0.60,
+        slice_cosine_threshold=SLICE_COSINE_DISTANCE_THRESHOLD,
+        full_cosine_threshold=FULL_COSINE_DISTANCE_THRESHOLD,
     )

fastvideo/tests/ssim/test_stable_audio_similarity.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Latent-slice regression test for Stable Audio Open 1.0 text-to-audio.
+
+Companion to ``test_ltx2_similarity.py`` — applies the same latent
+cosine-distance philosophy to 3-D audio latents ``[B, 64, T_latent]``.
+
+Why latent-space and not waveform-space SSIM:
+- ``dpmpp-3m-sde`` (k-diffusion) accumulates per-step bf16 noise; the
+  Oobleck VAE then magnifies any residual drift into the time-domain
+  waveform. A few mis-rounded accumulators drive sample-wise diff well
+  past audible thresholds without indicating a real regression.
+- Diffusers' own
+  ``tests/pipelines/stable_audio/test_stable_audio.py`` compares
+  decoded audio samples via
+  ``np.abs(expected - actual).max() < 1.5e-3``; that bound holds for
+  CPU dummy components but does not survive cross-architecture bf16
+  on our CI pool (L40S/A40/H100/B200).
+- Comparing the **pre-VAE latent** moves the assertion upstream of
+  the dominant noise source.
+
+Slice spec: ``audio_first_8_timesteps`` returns
+``latent[0, :, :8]`` (= 64 channels × 8 latent timesteps = 512
+elements). The full latent ``[1, 64, 1024]`` for SA-1.0 is also
+compared via cosine distance.
+"""
+
+import os
+
+import pytest
+
+from fastvideo.api.sampling_param import SamplingParam
+from fastvideo.logger import init_logger
+from fastvideo.tests.ssim.inference_similarity_utils import (
+    resolve_inference_device_reference_folder,
+)
+from fastvideo.tests.ssim.latent_similarity_utils import (
+    AUDIO_FIRST_8_TIMESTEPS_SPEC,
+    run_text_to_latent_similarity_test,
+)
+
+logger = init_logger(__name__)
+
+REQUIRED_GPUS = 1
+
+device_reference_folder = resolve_inference_device_reference_folder(logger)
+
+STABLE_AUDIO_PARAMS = {
+    "num_gpus": 1,
+    "model_path": "FastVideo/stable-audio-open-1.0-Diffusers",
+    # 8/8/1 are the SA-1.0 SamplingParam defaults; the values are
+    # placeholders unused by the audio pipeline but must satisfy the
+    # shared InputValidationStage divisible-by-8 check.
+    "height": 8,
+    "width": 8,
+    "num_frames": 1,
+    "num_inference_steps": 25,
+    "guidance_scale": 7.0,
+    "seed": 1024,
+    "sp_size": 1,
+    "tp_size": 1,
+    "fps": 24,
+}
+
+_SA_FULL_DEFAULTS = SamplingParam.from_pretrained(STABLE_AUDIO_PARAMS["model_path"])
+STABLE_AUDIO_FULL_QUALITY_PARAMS = {
+    **STABLE_AUDIO_PARAMS,
+    "num_inference_steps": _SA_FULL_DEFAULTS.num_inference_steps,
+    "guidance_scale": _SA_FULL_DEFAULTS.guidance_scale,
+    "seed": _SA_FULL_DEFAULTS.seed,
+}
+
+STABLE_AUDIO_MODEL_TO_PARAMS = {
+    "stable-audio-open-1.0-Diffusers": STABLE_AUDIO_PARAMS,
+}
+FULL_QUALITY_STABLE_AUDIO_MODEL_TO_PARAMS = {
+    "stable-audio-open-1.0-Diffusers": STABLE_AUDIO_FULL_QUALITY_PARAMS,
+}
+
+STABLE_AUDIO_TEST_PROMPTS = [
+    "Lo-fi hip hop instrumental with vinyl crackle and gentle piano.",
+]
+
+SLICE_COSINE_DISTANCE_THRESHOLD = 5e-3
+FULL_COSINE_DISTANCE_THRESHOLD = 1e-2
+
+
+@pytest.mark.parametrize("prompt", STABLE_AUDIO_TEST_PROMPTS)
+@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
+@pytest.mark.parametrize("model_id", list(STABLE_AUDIO_MODEL_TO_PARAMS.keys()))
+def test_stable_audio_inference_similarity(
+    prompt: str,
+    attention_backend_name: str,
+    model_id: str,
+) -> None:
+    run_text_to_latent_similarity_test(
+        logger=logger,
+        script_dir=os.path.dirname(os.path.abspath(__file__)),
+        device_reference_folder=device_reference_folder,
+        prompt=prompt,
+        attention_backend_name=attention_backend_name,
+        model_id=model_id,
+        default_params_map=STABLE_AUDIO_MODEL_TO_PARAMS,
+        full_quality_params_map=FULL_QUALITY_STABLE_AUDIO_MODEL_TO_PARAMS,
+        slice_cosine_threshold=SLICE_COSINE_DISTANCE_THRESHOLD,
+        full_cosine_threshold=FULL_COSINE_DISTANCE_THRESHOLD,
+        slice_spec=AUDIO_FIRST_8_TIMESTEPS_SPEC,
+        # FSDP + @torch.inference_mode in StableAudioDenoisingStage hits
+        # "Inference tensors do not track version counter" on single-GPU
+        # unshard. SA-1.0 fits on one B200 anyway.
+        init_kwargs_override={"use_fsdp_inference": False},
+    )