From c17d33bf3392bcafa9b75913ddf59bdbc67b9b40 Mon Sep 17 00:00:00 2001 From: Mook <68294499+Godmook@users.noreply.github.com> Date: Tue, 5 May 2026 03:36:58 -0700 Subject: [PATCH] [ci] Replace flaky LTX-2 pixel SSIM with latent-slice cosine regression (#1253) Co-authored-by: SolitaryThinker --- .agents/skills/seed-ssim-references/SKILL.md | 155 +++++- fastvideo/entrypoints/video_generator.py | 56 +- .../basic/stable_audio/stages/decoding.py | 9 + fastvideo/tests/modal/ssim_test.py | 4 + .../tests/ssim/inference_similarity_utils.py | 12 +- .../tests/ssim/latent_similarity_utils.py | 494 ++++++++++++++++++ fastvideo/tests/ssim/reference_videos_cli.py | 25 +- fastvideo/tests/ssim/test_ltx2_similarity.py | 40 +- .../ssim/test_stable_audio_similarity.py | 111 ++++ 9 files changed, 845 insertions(+), 61 deletions(-) create mode 100644 fastvideo/tests/ssim/latent_similarity_utils.py create mode 100644 fastvideo/tests/ssim/test_stable_audio_similarity.py diff --git a/.agents/skills/seed-ssim-references/SKILL.md b/.agents/skills/seed-ssim-references/SKILL.md index e494660f2..527b76ef9 100644 --- a/.agents/skills/seed-ssim-references/SKILL.md +++ b/.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 `` 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 `//*.mp4` +downloaded tree; `copy-local` walks all `//*.{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///.mp4`. +Result for pixel: `fastvideo/tests/ssim/reference_videos/default/L40S_reference_videos///.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/`) + 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/`) 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). | diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 7aced5f34..033ad1502 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/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"), } diff --git a/fastvideo/pipelines/basic/stable_audio/stages/decoding.py b/fastvideo/pipelines/basic/stable_audio/stages/decoding.py index 63cbd1bd6..3b3fe12b9 100644 --- a/fastvideo/pipelines/basic/stable_audio/stages/decoding.py +++ b/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()) diff --git a/fastvideo/tests/modal/ssim_test.py b/fastvideo/tests/modal/ssim_test.py index f37325481..106aa304a 100644 --- a/fastvideo/tests/modal/ssim_test.py +++ b/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" """ diff --git a/fastvideo/tests/ssim/inference_similarity_utils.py b/fastvideo/tests/ssim/inference_similarity_utils.py index 57f7d6ee2..e2946d28e 100644 --- a/fastvideo/tests/ssim/inference_similarity_utils.py +++ b/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, diff --git a/fastvideo/tests/ssim/latent_similarity_utils.py b/fastvideo/tests/ssim/latent_similarity_utils.py new file mode 100644 index 000000000..e38033e36 --- /dev/null +++ b/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 + ``_reference_videos///.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, + ) diff --git a/fastvideo/tests/ssim/reference_videos_cli.py b/fastvideo/tests/ssim/reference_videos_cli.py index 526d5f82b..eef18baf2 100644 --- a/fastvideo/tests/ssim/reference_videos_cli.py +++ b/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 diff --git a/fastvideo/tests/ssim/test_ltx2_similarity.py b/fastvideo/tests/ssim/test_ltx2_similarity.py index acda07d56..7bd38d108 100644 --- a/fastvideo/tests/ssim/test_ltx2_similarity.py +++ b/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, ) diff --git a/fastvideo/tests/ssim/test_stable_audio_similarity.py b/fastvideo/tests/ssim/test_stable_audio_similarity.py new file mode 100644 index 000000000..acf9ae0f6 --- /dev/null +++ b/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}, + )