[feat] LTX-2.3 audio: BWE vocoder path - #1398
Conversation
Port the bandwidth-extension (BWE) vocoder stack from the internal FastVideo repo into the upstream worktree. The change is additive: the non-BWE vocoder path is unchanged; BWE only activates when the vocoder config dict contains a "bwe" key. New symbols in fastvideo/models/audio/ltx2_audio_vae.py: - get_padding, _sinc, kaiser_sinc_filter1d - LowPassFilter1d, UpSample1d, DownSample1d, Activation1d - Snake, SnakeBeta, AMPBlock1 - _STFTFn, MelSTFT, VocoderWithBWE Vocoder extended with activation / use_tanh_at_final / apply_final_activation / use_bias_at_final params and AMP1 resblock support to match the internal implementation. VocoderConfigurator.from_config gains an "if bwe in vocoder_cfg" branch that builds the 16 kHz base Vocoder + BWE generator + MelSTFT and wraps them in VocoderWithBWE (48 kHz output). fastvideo/pipelines/basic/ltx2/stages/ltx2_audio_decoding.py: replace hardcoded DEFAULT_LTX2_VOCODER_OUTPUT_SAMPLE_RATE with a _resolve_audio_sample_rate() helper so BWE's 48 kHz output is propagated correctly while the default (24 000 Hz) is preserved.
mergify
Bot
added
type: feat
Merge ProtectionsYour pull request matches the following merge protections and will not be merged until they are valid. 🔴 PR merge requirementsWaiting for
This rule is failing.
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces BigVGAN v2 and Bandwidth Extension (BWE) helper modules, including low-pass filters, upsampling/downsampling, Snake activations, and a VocoderWithBWE wrapper. It also updates the audio decoding pipeline to dynamically resolve the output sample rate. The code review highlights critical dtype mismatch bugs in LowPassFilter1d and _STFTFn when processing half-precision tensors, potential silent failures due to initializing STFT and Mel basis buffers with zeros, redundant custom sinc functions that should be replaced with torch.sinc, and a potential sample rate mismatch in the configurator's default fallback values.
| def forward(self, x: torch.Tensor) -> torch.Tensor: | ||
| _, n_channels, _ = x.shape | ||
| if self.padding: | ||
| x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) | ||
| return F.conv1d( | ||
| x, | ||
| self.filter.expand(n_channels, -1, -1), | ||
| stride=self.stride, | ||
| groups=n_channels, | ||
| ) |
There was a problem hiding this comment.
The self.filter buffer is initialized as float32. If the input x is in half precision (e.g., bfloat16 or float16), F.conv1d will crash due to a dtype mismatch. Cast self.filter to x.dtype and x.device before calling F.conv1d, similar to how it is done in UpSample1d.
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| _, n_channels, _ = x.shape | |
| if self.padding: | |
| x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) | |
| return F.conv1d( | |
| x, | |
| self.filter.expand(n_channels, -1, -1), | |
| stride=self.stride, | |
| groups=n_channels, | |
| ) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| _, n_channels, _ = x.shape | |
| if self.padding: | |
| x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) | |
| filt = self.filter.to(dtype=x.dtype, device=x.device).expand( | |
| n_channels, -1, -1 | |
| ) | |
| return F.conv1d( | |
| x, | |
| filt, | |
| stride=self.stride, | |
| groups=n_channels, | |
| ) |
| def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | ||
| if y.dim() == 2: | ||
| y = y.unsqueeze(1) | ||
| left_pad = max(0, self.win_length - self.hop_length) | ||
| y = F.pad(y, (left_pad, 0)) | ||
| spec = F.conv1d(y, self.forward_basis, stride=self.hop_length, padding=0) | ||
| n_freqs = spec.shape[1] // 2 | ||
| real, imag = spec[:, :n_freqs], spec[:, n_freqs:] | ||
| magnitude = torch.sqrt(real**2 + imag**2) | ||
| phase = torch.atan2(imag.float(), real.float()).to(real.dtype) | ||
| return magnitude, phase |
There was a problem hiding this comment.
The self.forward_basis buffer is initialized as float32. If the input y is in half precision (e.g., bfloat16 or float16), F.conv1d will crash due to a dtype mismatch. Cast self.forward_basis to y.dtype and y.device before calling F.conv1d.
| def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | |
| if y.dim() == 2: | |
| y = y.unsqueeze(1) | |
| left_pad = max(0, self.win_length - self.hop_length) | |
| y = F.pad(y, (left_pad, 0)) | |
| spec = F.conv1d(y, self.forward_basis, stride=self.hop_length, padding=0) | |
| n_freqs = spec.shape[1] // 2 | |
| real, imag = spec[:, :n_freqs], spec[:, n_freqs:] | |
| magnitude = torch.sqrt(real**2 + imag**2) | |
| phase = torch.atan2(imag.float(), real.float()).to(real.dtype) | |
| return magnitude, phase | |
| def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | |
| if y.dim() == 2: | |
| y = y.unsqueeze(1) | |
| left_pad = max(0, self.win_length - self.hop_length) | |
| y = F.pad(y, (left_pad, 0)) | |
| basis = self.forward_basis.to(dtype=y.dtype, device=y.device) | |
| spec = F.conv1d(y, basis, stride=self.hop_length, padding=0) | |
| n_freqs = spec.shape[1] // 2 | |
| real, imag = spec[:, :n_freqs], spec[:, n_freqs:] | |
| magnitude = torch.sqrt(real**2 + imag**2) | |
| phase = torch.atan2(imag.float(), real.float()).to(real.dtype) | |
| return magnitude, phase |
| n_freqs = filter_length // 2 + 1 | ||
| self.register_buffer("forward_basis", torch.zeros(n_freqs * 2, 1, filter_length)) | ||
| self.register_buffer("inverse_basis", torch.zeros(n_freqs * 2, 1, filter_length)) |
There was a problem hiding this comment.
Initializing forward_basis and inverse_basis to all zeros makes the model entirely dependent on the checkpoint containing these exact buffers. If the checkpoint does not contain them, or if the model is initialized without a checkpoint (e.g., in unit tests), it will silently produce zero outputs. Consider initializing them with the actual Fourier basis functions, or at least raise a warning if they remain all zeros after loading the state dict.
| n_freqs = filter_length // 2 + 1 | ||
| self.register_buffer("mel_basis", torch.zeros(n_mel_channels, n_freqs)) |
There was a problem hiding this comment.
Initializing mel_basis to all zeros makes the model entirely dependent on the checkpoint containing this exact buffer. If the checkpoint does not contain it, or if the model is initialized without a checkpoint (e.g., in unit tests), it will silently produce zero outputs. Consider initializing it with the actual Mel scale filterbank (e.g., using librosa or torchaudio), or at least raise a warning if it remains all zeros after loading the state dict.
| def _sinc(x: torch.Tensor) -> torch.Tensor: | ||
| return torch.where( | ||
| x == 0, | ||
| torch.tensor(1.0, device=x.device, dtype=x.dtype), | ||
| torch.sin(math.pi * x) / math.pi / x, | ||
| ) |
There was a problem hiding this comment.
The custom _sinc helper function can be completely replaced by PyTorch's native torch.sinc function, which is highly optimized, fully differentiable, and handles the
| def _sinc(x: torch.Tensor) -> torch.Tensor: | |
| return torch.where( | |
| x == 0, | |
| torch.tensor(1.0, device=x.device, dtype=x.dtype), | |
| torch.sin(math.pi * x) / math.pi / x, | |
| ) | |
| # Removed in favor of torch.sinc |
| if cutoff == 0: | ||
| filter_ = torch.zeros_like(time) | ||
| else: | ||
| filter_ = 2 * cutoff * window * _sinc(2 * cutoff * time) |
There was a problem hiding this comment.
Use torch.sinc directly instead of the custom _sinc helper.
| filter_ = 2 * cutoff * window * _sinc(2 * cutoff * time) | |
| filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time) |
| vocoder=base_vocoder, | ||
| bwe_generator=bwe_generator, | ||
| mel_stft=mel_stft, | ||
| input_sampling_rate=bwe_cfg.get("input_sampling_rate", 16000), |
There was a problem hiding this comment.
There is a potential sample rate mismatch if bwe_cfg does not contain "input_sampling_rate". In that case, base_vocoder's output_sample_rate (lines 1751-1754) defaults to nested_vocoder_cfg.get("output_sampling_rate", 24000), while VocoderWithBWE's input_sampling_rate defaults to 16000. This mismatch will cause incorrect resampling ratios and shape mismatches. Consider resolving the base vocoder sample rate once and passing it to both places.
|
Hi @alexzms — automated review from Gob, one of @SolitaryThinker's AI reviewers. Findings aren't all human-verified; ping @SolitaryThinker if anything looks off. Thanks for the focused BWE vocoder patch. I found two S1 issues that should be fixed before merge:
— Gob (@SolitaryThinker's AI reviewer). Full review (including S3 items) is archived locally. |
Addresses Gob's S1 finding on hao-ai-lab#1398: ``_STFTFn.forward_basis`` / ``_STFTFn.inverse_basis`` and ``MelSTFT.mel_basis`` were registered as ``torch.zeros(...)`` placeholders and used immediately by ``F.conv1d`` and ``torch.matmul``. Because ``VocoderLoader`` loads with ``strict=False`` (``fastvideo/models/loader/component_loader.py:829``), a converted vocoder checkpoint that omits these deterministic signal-processing buffers no longer raises — the buffers stay at zero, ``magnitude = sqrt(0 + 0) = 0``, ``log_mel = log(clamp(0, 1e-5))`` is a constant ≈ -11.51, and the BWE generator runs on a constant feature, silently producing no high-band synthesis. Fix: add ``_build_stft_basis`` (Hann-windowed FFT basis as a Conv1d kernel) and ``_build_mel_basis`` (Slaney-normalised triangular mel filterbank, librosa-compatible defaults) and use them in the constructors: - ``_STFTFn.__init__`` initialises ``forward_basis`` and ``inverse_basis`` from the windowed FFT basis instead of zeros. - ``MelSTFT.__init__`` gains a required ``sampling_rate`` argument and initialises ``mel_basis`` from the mel filterbank instead of zeros. - The construction site in the vocoder factory now passes ``bwe_cfg["input_sampling_rate"]`` (default 16000) into ``MelSTFT``. A loaded checkpoint still overrides the deterministic init exactly as before; the only behavioural change is that an *omitted* buffer no longer silently zeros the signal-processing path. Adds ``fastvideo/tests/audio/test_ltx2_bwe.py`` (11 CPU-only regression tests) covering forward/mel basis non-zero invariants, shape, sampling-rate dependence, mel filterbank non-negativity, end-to-end log-mel of a sine tone not constant, and the standalone helpers.
|
Pushed S1 #1 (zero-init buffers). Added S1 #2 (no CI coverage). Added |
2 participants
Summary
Adds the LTX-2.3 audio BWE (bandwidth-extension) vocoder path to
ltx2_audio_vae.py. Additive + backward-compatible: it only activates when a vocoder config contains a"bwe"key; the existing non-BWEVocoderpath is behavior-identical.VocoderWithBWE,MelSTFT/_STFTFn,Snake/SnakeBeta,AMPBlock1, anti-alias resampling (LowPassFilter1d/UpSample1d/DownSample1d/Activation1d,kaiser_sinc_filter1d).VocoderConfigurator.from_configgains theif "bwe" in vocoder_cfgbranch (16 kHz base + BWE generator + mel re-analysis → 48 kHz)._resolve_audio_sample_rate(handles 48 kHz BWE vs the 24 kHz default).Backward compatibility
Non-BWE
Vocoderis unchanged: the in-loopleaky_relu(x, LRELU_SLOPE)is gatedif not is_amp, and the final path staysleaky_relu() → conv_post → tanh(verified against the original forward). BWE classes are inert unless the checkpoint config carries abweblock.Test plan
py_compile+ import (VocoderWithBWEpresent)test_ltx2_continuation+test_ltx2_stage_overrides(29)Independent of the DiT-2.3 PR; can be reviewed/merged separately.