[feat] v2 VideoGenerator: A/V convenience path (generate_video -> T2VS -> mp4 + 24kHz wav)

Makes the 'Full A/V' LTX-2.3 deliverable reachable from the user-facing entrypoint, not just the engine. A model advertising TEXT_TO_VIDEO_SOUND (LTX-2.3) now auto-issues a T2VS request, so generate() / generate_video() return BOTH modalities in one joint pass: - VideoGenerator stores the resident instance + a supports_av flag (from card.capabilities); generate() gains want_audio (None=auto-by-capability, True/False to force) and routes T2V vs T2VS+{video,audio}. - _result saves the stereo waveform as a sibling .wav at the vocoder's REAL rate (24000) — read off the built audio_vae adapter (TorchLTX2AudioVAE.sample_rate = Vocoder.output_sample_rate), since the AudioArtifact default rate is a placeholder. Populates GenerationResult.audio/.audio_sample_rate and extra['audio_path']. scipy IEEE-float WAV; [channels,samples] auto-transposed. - Rewrote v2_basic_ltx2_3_distilled.py: registry routes to build_ltx2_3_card (its own joint T2VS A/V card, not the LTX-2 base/2-stage card); the example prints both the mp4 and the wav. - GPU-verified via the convenience API: ev.mp4 + ev.wav (24000 Hz, stereo 61920x2, nonzero, std 0.043). CPU mini green (202 passed, 2 skipped); engine/program/toy paths untouched (test_ltx2_av pins 44100).
2026-07-04 17:15:34 +00:00
parent 6096b00aeb
commit c6d2976fc2
4 changed files with 99 additions and 32 deletions
@@ -37,11 +37,13 @@ rebuild the venv for the current arch — system python `/usr/bin/python3.12`, r
 | `basic_wan2_2_ti2v.py` (T2V branch) | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | wan2.2-ti2v | `v2_basic_wan2_2_ti2v.py` |
 | `basic_wan2_2.py` | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` (MoE) | wan2.2-a14b + **CPU expert offload** (one 14B expert resident at a time, swapped at the boundary → 60GB peak) | `v2_basic_wan2_2.py` |
 | `basic_ltx2.py` | `Davids048/LTX2-Base-Diffusers` | ltx2 base (single-stage, request-driven many-step) | `v2_basic_ltx2.py` |
-| `basic_ltx2_3_distilled.py` | `FastVideo/LTX-2.3-Distilled-Diffusers` (18.99B) | ltx2 base card (single-stage; few-step) | `v2_basic_ltx2_3_distilled.py` |
+| `basic_ltx2_3_distilled.py` | `FastVideo/LTX-2.3-Distilled-Diffusers` (18.99B) | ltx2.3-distilled card — single-stage **joint T2VS** (video **+** audio): SEPARATE video/audio text connectors + gated attention, joint DiT forward (video↔audio cross-attention), then video VAE + AudioDecoder→Vocoder | `v2_basic_ltx2_3_distilled.py` |

-All seven rows GPU-verified to generate real video (LTX-2 base / 2.3 show inter-frame motion 4.5 / 6.6
-on the rebuilt x86 stack). \* `basic_mps.py` (Apple MPS) and `basic_ray.py` (ray executor) are
-device/executor variants of the same Wan2.1 model — covered by wan21; those runtime modes are out of scope.
+All seven rows GPU-verified to generate real video on the rebuilt x86 stack. LTX-2.3 is the only **A/V**
+model: GPU-verified T2VS produces video (3,33,256,384) **and** stereo audio (2×61920 @ 24 kHz, nonzero,
+std 0.059) in one joint pass — the dual-connector / joint-cross-attention path is exercised end-to-end.
+\* `basic_mps.py` (Apple MPS) and `basic_ray.py` (ray executor) are device/executor variants of the same
+Wan2.1 model — covered by wan21; those runtime modes are out of scope.

 ## Needs per-model work (architecture/sampler matches partially)
 | Official example(s) | Model | What v2 needs |
@@ -1,25 +1,34 @@
-"""v2 port of basic_ltx2_3_distilled.py — LTX-2.3 Distilled (single-stage) through the v2 VideoGenerator.
+"""v2 port of basic_ltx2_3_distilled.py — LTX-2.3 Distilled (single-stage, joint A/V) through the v2
+VideoGenerator.

-LTX-2.3-Distilled is single-stage (no spatial upsampler, unlike LTX-2.0 distilled's two-stage), so the
-architecture-driven dispatch (v2/video_generator.py:_select_builders, via has_spatial_upsampler=False)
-routes it to the SAME single-stage card as LTX-2 base (build_ltx2_base_card). Being distilled, it wants
-FEW steps — pass a small num_inference_steps for the few-step schedule.
+Unlike LTX-2.0 distilled (two-stage, video-only), LTX-2.3 is a single-stage *audio+video* model. The
+shared registry (v2/registry.py) maps ``FastVideo/LTX-2.3-Distilled-Diffusers`` to its OWN card,
+``build_ltx2_3_card`` — distinct from the LTX-2 base/2-stage cards — which wires the 2.3-specific path:
+  * SEPARATE video + audio text connectors (the Gemma encoder projects the prompt to two embeddings,
+    2048-dim for audio, 4096-dim for video) plus gated attention;
+  * a JOINT DiT forward where video and audio latents cross-attend in a single denoise per step;
+  * a video VAE decode + an AudioDecoder→Vocoder decode → video frames AND a stereo waveform @24kHz.

-NOTE: the v2 single-stage loop currently builds a request-driven *linspace* flow-match schedule
-(base_flow_sigmas); the distilled-tuned few-step schedule is a follow-up. GPU re-verify of this example
-is pending the env (the box was rescheduled aarch64->x86 mid-session — see V2_PORTING_STATUS.md).
+Because the model advertises TEXT_TO_VIDEO_SOUND, the VideoGenerator issues a T2VS request by default,
+so ``generate_video`` returns BOTH modalities: the mp4 plus a sibling ``.wav`` (and ``result.audio`` /
+``result.audio_sample_rate`` in memory). Being distilled, it wants FEW steps (8). GPU-verified on the
+rebuilt x86 stack: video (3,33,256,384) + stereo audio (2×61920 @ 24kHz).
 """
 from v2 import VideoGenerator

-PROMPT = "a hummingbird hovering beside a red flower, wings blurred, macro, cinematic, highly detailed"
+PROMPT = "ocean waves crashing on rocks at sunset, seagulls calling in the distance, cinematic, highly detailed"


 def main() -> None:
    generator = VideoGenerator.from_pretrained("FastVideo/LTX-2.3-Distilled-Diffusers", num_gpus=1)
-    video = generator.generate_video(
-        prompt=PROMPT, output_path="v2_video_samples_ltx2_3", output_video_name="ltx2_3_hummingbird",
-        save_video=True, num_frames=9, height=512, width=768, num_inference_steps=8)
-    print(f"Output: {video.video_path}")
+    # audio=None auto-enables sound for this A/V model (pass audio=False to force video-only).
+    result = generator.generate_video(
+        prompt=PROMPT, output_path="v2_video_samples_ltx2_3", output_video_name="ltx2_3_ocean",
+        save_video=True, num_frames=33, height=512, width=768, num_inference_steps=8, seed=1)
+    print(f"Video: {result.video_path}")
+    audio_path = result.extra.get("audio_path")
+    if audio_path:
+        print(f"Audio: {audio_path}  ({result.audio_sample_rate} Hz)")
    generator.shutdown()


@@ -166,6 +166,9 @@ class TorchLTX2AudioVAE:
        self.module = module.to(device=device, dtype=dtype).eval()      # AudioDecoder
        self.vocoder = vocoder                                          # TorchLTX2Vocoder | None
        self.device, self.dtype = device, dtype
+        # Final waveform rate = the Vocoder's output_sample_rate (24000 for LTX-2; AudioDecoder's own
+        # sample_rate is the internal mel rate). Surfaced so a saver writes the wav at the right rate.
+        self.sample_rate = int(getattr(getattr(vocoder, "module", None), "output_sample_rate", 24000))

    @torch.no_grad()
    def decode(self, audio_latent):
@@ -27,9 +27,11 @@ from typing import Any
 class VideoGenerator:
    """Typed t2v entrypoint over a single resident v2 ``ModelInstance`` (mirrors fastvideo's)."""

-    def __init__(self, engine: Any, model_id: str) -> None:
+    def __init__(self, engine: Any, model_id: str, *, instance: Any = None, supports_av: bool = False) -> None:
        self._engine = engine
        self._model_id = model_id
+        self._inst = instance              # the resident ModelInstance (to read e.g. the audio sample rate)
+        self._supports_av = supports_av    # model emits joint video+sound (LTX-2.3 T2VS) -> auto-request audio

    # --------------------------------------------------------------------- #
    @classmethod
@@ -84,25 +86,34 @@ class VideoGenerator:
        from .models.wan21 import stamp_wan21_checkpoints   # transformer/vae/text_encoder subfolder layout
        stamp_wan21_checkpoints(card, root)

+        from ._enums import Capability
        from .card import load_card
        from .cache import CacheManager
        from .runtime import Engine
        inst = load_card(card, cache_manager=CacheManager.from_card(card))   # platform auto-detect -> cuda
        eng = Engine()
        eng.register(card.model_id, inst, program)
-        return cls(eng, card.model_id)
+        # A model that advertises TEXT_TO_VIDEO_SOUND (LTX-2.3) generates joint video+audio; the
+        # entrypoint then issues a T2VS request by default so a plain generate() yields both modalities.
+        supports_av = card.capabilities.has(Capability.TEXT_TO_VIDEO_SOUND)
+        return cls(eng, card.model_id, instance=inst, supports_av=supports_av)

    def shutdown(self) -> None:
        """API parity with ``fastvideo.VideoGenerator.shutdown()``. The v2 bring-up holds a single
        resident instance that is freed at process exit; there is nothing to tear down explicitly."""

    # --------------------------------------------------------------------- #
-    def generate(self, request: Any) -> Any:
+    def generate(self, request: Any, *, want_audio: bool | None = None) -> Any:
        """``request``: ``fastvideo.api.GenerationRequest``. Returns a ``GenerationResult`` (or a list
-        for multi-prompt) with ``.video_path`` (saved mp4) and optionally ``.frames``."""
-        from .request import DiffusionParams, TaskType, make_request
+        for multi-prompt) with ``.video_path`` (saved mp4) and optionally ``.frames``.
+
+        ``want_audio``: ``None`` (default) auto-enables sound for an A/V model (LTX-2.3 → T2VS, both
+        video+audio in one joint pass); ``True``/``False`` force it on/off. For a video-only model the
+        auto resolves to video-only (T2V), unchanged."""
+        from .request import DiffusionParams, OutputSpec, TaskType, make_request
        s = request.sampling
        prompts = request.prompt if isinstance(request.prompt, list) else [request.prompt]
+        audio = self._supports_av if want_audio is None else bool(want_audio)
        results = []
        for idx, prompt in enumerate(prompts):
            diff = DiffusionParams(
@@ -110,7 +121,11 @@ class VideoGenerator:
                height=s.height, width=s.width, guidance_scale=s.guidance_scale,
                negative_prompt=request.negative_prompt or "",
                sigmas=tuple(s.sigmas) if getattr(s, "sigmas", None) else None)
-            req = make_request(TaskType.T2V, self._model_id, prompt, diffusion=diff)
+            if audio:   # joint text->video+sound: ask for both modalities so the audio branch runs
+                req = make_request(TaskType.T2VS, self._model_id, prompt, diffusion=diff,
+                                   outputs=OutputSpec(modalities=frozenset({"video", "audio"})))
+            else:
+                req = make_request(TaskType.T2V, self._model_id, prompt, diffusion=diff)
            out = self._engine.run(req)
            results.append(self._result(out, request.output, s.fps, idx))
        return results[0] if len(results) == 1 else results
@@ -121,8 +136,10 @@ class VideoGenerator:
        ``GenerationRequest`` from loose kwargs (+ an optional ``SamplingParam`` override) and calls
        ``generate()``. kwargs: num_inference_steps / seed / num_frames / height / width /
        guidance_scale / fps / sigmas / negative_prompt / output_path / output_video_name /
-        save_video / return_frames."""
+        save_video / return_frames / audio (None=auto for an A/V model, True/False to force; the audio
+        is saved as a sibling ``.wav`` next to the mp4)."""
        from fastvideo.api import GenerationRequest, OutputConfig, SamplingConfig
+        audio = kwargs.pop("audio", None)      # None -> auto (by model capability), True/False -> force
        sp = sampling_param

        def pick(key: str, default: Any) -> Any:
@@ -148,7 +165,7 @@ class VideoGenerator:
            return_frames=bool(kwargs.get("return_frames", False)))
        req = GenerationRequest(prompt=prompt, negative_prompt=pick("negative_prompt", None),
                               sampling=sampling, output=output)
-        return self.generate(req)
+        return self.generate(req, want_audio=audio)

    # --------------------------------------------------------------------- #
    def _result(self, out: Any, output: Any, fps: int, idx: int) -> Any:
@@ -158,13 +175,49 @@ class VideoGenerator:
        # -> [T,H,W,C] uint8 (the on-disk / return convention)
        vid = ((np.clip(frames, -1.0, 1.0) + 1.0) / 2.0 * 255.0).astype("uint8")
        vid = vid.transpose(1, 2, 3, 0) if vid.ndim == 4 else vid
+        save = bool(getattr(output, "save_video", True))
+        out_dir = getattr(output, "output_path", "outputs/")
+        stem = (output.output_video_name or f"v2_{self._model_id}_{idx}")
+        stem = stem[:-4] if stem.endswith(".mp4") else stem
        video_path = None
-        if getattr(output, "save_video", True):
-            os.makedirs(output.output_path, exist_ok=True)
-            name = output.output_video_name or f"v2_{self._model_id}_{idx}"
-            name = name if name.endswith(".mp4") else f"{name}.mp4"
-            video_path = os.path.join(output.output_path, name)
+        if save:
+            os.makedirs(out_dir, exist_ok=True)
+            video_path = os.path.join(out_dir, f"{stem}.mp4")
            import imageio.v2 as imageio
            imageio.mimsave(video_path, list(vid), fps=int(fps), format="mp4")
-        return GenerationResult(frames=(vid if getattr(output, "return_frames", True) else None),
-                                video_path=video_path)
+        # Audio (T2VS): the AudioArtifact carries the raw stereo waveform; save a sibling .wav at the
+        # vocoder's real rate (read from the built audio_vae adapter — the artifact's default rate is a
+        # placeholder). A plain T2V request produces no audio artifact, so this is skipped.
+        audio = audio_sr = audio_path = None
+        art = out.artifacts.get("audio") if hasattr(out.artifacts, "get") else None
+        wav = getattr(art, "samples", None) if art is not None else None
+        if wav is not None:
+            audio = np.asarray(wav, dtype="float32")
+            audio_sr = self._audio_sample_rate(default=int(getattr(art, "sample_rate", 24000) or 24000))
+            if save:
+                audio_path = os.path.join(out_dir, f"{stem}.wav")
+                self._write_wav(audio_path, audio, audio_sr)
+        res = GenerationResult(frames=(vid if getattr(output, "return_frames", True) else None),
+                               video_path=video_path, audio=audio, audio_sample_rate=audio_sr)
+        if audio_path:
+            res.extra["audio_path"] = audio_path
+        return res
+
+    def _audio_sample_rate(self, default: int = 24000) -> int:
+        """Read the true output rate off the built audio_vae adapter (LTX-2 vocoder = 24000); fall back
+        to ``default`` if the model has no audio component or the adapter doesn't expose it."""
+        try:
+            return int(getattr(self._inst.component("audio_vae"), "sample_rate", default))
+        except Exception:
+            return int(default)
+
+    @staticmethod
+    def _write_wav(path: str, wav: Any, sample_rate: int) -> None:
+        """Write ``wav`` ([samples] or [channels, samples] / [samples, channels], float in [-1,1]) to a
+        WAV at ``sample_rate`` (scipy: IEEE-float WAV, no extra deps beyond the installed stack)."""
+        import numpy as np
+        from scipy.io import wavfile
+        a = np.asarray(wav, dtype="float32")
+        if a.ndim == 2 and a.shape[0] in (1, 2) and a.shape[0] < a.shape[1]:
+            a = a.T                       # [channels, samples] -> [samples, channels] (scipy convention)
+        wavfile.write(path, int(sample_rate), np.clip(a, -1.0, 1.0))