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"]