From 6096b00aeb740d2e37f90329ee128d3b53747bc5 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Thu, 18 Jun 2026 08:44:07 +0000 Subject: [PATCH] [feat] v2 LTX-2.3 T2VS GPU-verified: audio VAE/vocoder wiring + dual-connector audio fix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The full joint text->video+audio LTX-2.3 path now generates on the real 18.99B model: - GPU audio components: build_torch_audio_vae (AudioDecoderLoader -> LTX2AudioDecoder, chains the vocoder) + build_torch_vocoder (VocoderLoader -> LTX2Vocoder); registered the 'audio_vae'/'vocoder' cuda component kinds; stamped their checkpoint subfolders (_WAN21_SUBFOLDERS). - Fix: TorchGemma.encode_av must pass output_hidden_states=True — the 2.3 connector's SEPARATE audio projection lives in hidden_states[0] only then (gemma.py:703); without it the audio text fell back to the video embedding (4096 vs 2048 -> audio cross-attn shape mismatch). - GPU-verified T2VS: video (3,33,256,384, std 0.68) + audio (stereo 2x61920 @24kHz, nonzero, std 0.059). - test_torch_backend: cuda component kinds now include audio_vae + vocoder. CPU suite green (202+2). --- v2/models/wan21/card.py | 3 ++- v2/platform/backends/torch_adapters.py | 19 +++++++++++++++++++ v2/platform/backends/torch_cuda.py | 14 ++++++++++++++ v2/platform/backends/torch_ltx2.py | 5 ++++- v2/tests/test_torch_backend.py | 3 ++- 5 files changed, 41 insertions(+), 3 deletions(-) diff --git a/v2/models/wan21/card.py b/v2/models/wan21/card.py index 58f8f4416..ce03fbfed 100644 --- a/v2/models/wan21/card.py +++ b/v2/models/wan21/card.py @@ -182,7 +182,8 @@ def build_wan22_a14b_card(model_id: str = "wan2.2-t2v-a14b", *, checkpoint_root: # (``transformer_2``), and LTX-2 (``spatial_upsampler``) alike. _WAN21_SUBFOLDERS = {"transformer": "transformer", "transformer_2": "transformer_2", "vae": "vae", "text_encoder": "text_encoder", - "spatial_upsampler": "spatial_upsampler"} + "spatial_upsampler": "spatial_upsampler", + "audio_vae": "audio_vae", "vocoder": "vocoder"} # LTX-2.3 T2VS audio branch def stamp_wan21_checkpoints(card: ModelCard, model_root: str) -> ModelCard: diff --git a/v2/platform/backends/torch_adapters.py b/v2/platform/backends/torch_adapters.py index 74794e312..6b050c2ea 100644 --- a/v2/platform/backends/torch_adapters.py +++ b/v2/platform/backends/torch_adapters.py @@ -363,3 +363,22 @@ def build_torch_upsampler(spec, instance, platform): stats_owner = getattr(vae_module, "decoder", None) or getattr(vae_module, "encoder", None) or vae_module from .torch_ltx2 import TorchLTX2Upsampler return TorchLTX2Upsampler(module, stats_owner, device=device, dtype=dtype) + + +def build_torch_audio_vae(spec, instance, platform): + ckpt = _require_checkpoint(spec) + _ensure_fastvideo_runtime() + module = _load_component("AudioDecoderLoader", ckpt, _fastvideo_args(spec)) # LTX2AudioDecoder + device, dtype = _device(platform), _native_dtype(module) + voc = instance.component("vocoder") # build/fetch the vocoder; TorchLTX2AudioVAE chains decoder->vocoder + from .torch_ltx2 import TorchLTX2AudioVAE + return TorchLTX2AudioVAE(module, voc, device=device, dtype=dtype) + + +def build_torch_vocoder(spec, instance, platform): + ckpt = _require_checkpoint(spec) + _ensure_fastvideo_runtime() + module = _load_component("VocoderLoader", ckpt, _fastvideo_args(spec)) # LTX2Vocoder + device, dtype = _device(platform), _native_dtype(module) + from .torch_ltx2 import TorchLTX2Vocoder + return TorchLTX2Vocoder(module, device=device, dtype=dtype) diff --git a/v2/platform/backends/torch_cuda.py b/v2/platform/backends/torch_cuda.py index b365adf24..558021a8c 100644 --- a/v2/platform/backends/torch_cuda.py +++ b/v2/platform/backends/torch_cuda.py @@ -60,6 +60,16 @@ def _build_upsampler(spec, instance, platform): return build_torch_upsampler(spec, instance, platform) +def _build_audio_vae(spec, instance, platform): + from .torch_adapters import build_torch_audio_vae + return build_torch_audio_vae(spec, instance, platform) + + +def _build_vocoder(spec, instance, platform): + from .torch_adapters import build_torch_vocoder + return build_torch_vocoder(spec, instance, platform) + + def _flow_match_cuda(*args, **kwargs): from .torch_kernels import flow_match_step return flow_match_step(*args, **kwargs) @@ -80,6 +90,10 @@ register_component("text_encoder", _build_text_encoder, device="cuda", available source=f"{_ADAPTERS}:TorchT5Encoder") register_component("upsampler", _build_upsampler, device="cuda", available=_cuda_available, source=f"{_ADAPTERS}:TorchLTX2Upsampler") +register_component("audio_vae", _build_audio_vae, device="cuda", available=_cuda_available, + source=f"{_ADAPTERS}:TorchLTX2AudioVAE") +register_component("vocoder", _build_vocoder, device="cuda", available=_cuda_available, + source=f"{_ADAPTERS}:TorchLTX2Vocoder") # --- solver ops: plain torch elementwise (no fused solver kernel exists in fastvideo-kernel) ---- # # arch="generic" (elementwise, arch-agnostic), workspace_bytes=0 (no static scratch). diff --git a/v2/platform/backends/torch_ltx2.py b/v2/platform/backends/torch_ltx2.py index 756ee86b3..da1556194 100644 --- a/v2/platform/backends/torch_ltx2.py +++ b/v2/platform/backends/torch_ltx2.py @@ -139,7 +139,10 @@ class TorchGemma: mask = toks.attention_mask.to(self.device) from fastvideo.forward_context import set_forward_context with set_forward_context(current_timestep=0, attn_metadata=None): - out = self.module(input_ids=ids, attention_mask=mask) + # output_hidden_states=True populates hidden_states=(audio_encoding,) — the SEPARATE audio + # connector projection (gemma.py:703); without it hidden_states is None and we'd wrongly + # fall back to the video embedding (different dim -> cross-attn shape mismatch). + out = self.module(input_ids=ids, attention_mask=mask, output_hidden_states=True) video = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0] hs = getattr(out, "hidden_states", None) audio = hs[0] if hs else video diff --git a/v2/tests/test_torch_backend.py b/v2/tests/test_torch_backend.py index dc9f4a379..0227b44ef 100644 --- a/v2/tests/test_torch_backend.py +++ b/v2/tests/test_torch_backend.py @@ -42,7 +42,8 @@ def _make_cuda_available(monkeypatch): # --------------------------------------------------------------------------- # def test_cuda_has_all_components_and_two_solver_ops(): comps = sorted(r["kind"] for r in component_matrix() if r["device"] == "cuda") - assert comps == ["dit", "text_encoder", "upsampler", "vae"] # + upsampler: LTX-2 spatial upsampler + # + upsampler (LTX-2 spatial), audio_vae + vocoder (LTX-2.3 T2VS audio branch) + assert comps == ["audio_vae", "dit", "text_encoder", "upsampler", "vae", "vocoder"] ops = sorted(r["op"] for r in kernel_matrix() if r["device"] == "cuda") assert ops == ["flow_match_step", "flow_sde_step"]