| |
|
|
| |
|
|
| """Anti-aliased 1D resampling layers for LTX-2 audio (BigVGAN-v2 style). |
| |
| Device equivalents of the reference's ``LowPassFilter1d`` / ``UpSample1d`` / |
| ``DownSample1d`` / ``Activation1d``. The transposed-conv upsample is expressed as |
| zero-stuff + zero-pad + a depthwise tap filter (bit-equivalent to |
| ``F.conv_transpose1d`` for the symmetric sinc kernel). ``UpSample1d`` is shared |
| by the vocoder (kaiser window, T-shard/halo aware) and the BWE resampler (hann |
| window, single-device). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch |
|
|
| import ttnn |
|
|
| from ..parallel.config import ParallelFactor |
| from ..parallel.manager import CCLManager |
| from .audio_ops import ( |
| _make_kaiser_sinc_kernel_1d, |
| _replicate_pad_t, |
| _t_neighbor_pad, |
| _zero_pad_t, |
| _zero_stuff_t, |
| depthwise_tap_filter, |
| ) |
| from .module import Module |
|
|
|
|
| def _make_hann_sinc_kernel_1d(*, ratio: int) -> tuple[torch.Tensor, int, int, int, int]: |
| """Return ``(kernel, kernel_size, pad, pad_left_crop, pad_right_crop)`` for the |
| Hann-window sinc resampler (torchaudio-equivalent), used by the BWE skip path.""" |
| rolloff = 0.99 |
| lowpass_filter_width = 6 |
| width = math.ceil(lowpass_filter_width / rolloff) |
| kernel_size = 2 * width * ratio + 1 |
| pad = width |
| pad_left = 2 * width * ratio |
| pad_right = kernel_size - ratio |
|
|
| time_axis = (torch.arange(kernel_size, dtype=torch.float64) / ratio - width) * rolloff |
| time_clamped = time_axis.clamp(-lowpass_filter_width, lowpass_filter_width) |
| window = torch.cos(time_clamped * math.pi / lowpass_filter_width / 2) ** 2 |
| sinc_filter = torch.sinc(time_axis) * window * rolloff / ratio |
| return sinc_filter.float().reshape(kernel_size), kernel_size, pad, pad_left, pad_right |
|
|
|
|
| class LowPassFilter1d(Module): |
| """Depthwise low-pass conv1d with a fixed kaiser-sinc kernel. |
| |
| The kernel is constant (baked at __init__), so ``_prepare_torch_state`` only |
| absorbs a checkpoint-provided kernel if present (BigVGAN convention). |
| """ |
|
|
| def __init__( |
| self, |
| *, |
| cutoff: float = 0.5, |
| half_width: float = 0.6, |
| stride: int = 1, |
| kernel_size: int = 12, |
| padding: bool = True, |
| padding_mode: str = "replicate", |
| mesh_device: ttnn.MeshDevice, |
| dtype: ttnn.DataType = ttnn.float32, |
| parallel_config: ParallelFactor | None = None, |
| ccl_manager: CCLManager | None = None, |
| ) -> None: |
| super().__init__() |
| if cutoff < 0.0 or cutoff > 0.5: |
| raise ValueError("cutoff must be in [0, 0.5]") |
| if padding_mode not in ("replicate", "zeros"): |
| raise ValueError(f"padding_mode must be replicate or zeros, got {padding_mode!r}") |
| sharded = parallel_config is not None and parallel_config.factor > 1 |
| if sharded: |
| assert ccl_manager is not None, "T-sharding requires ccl_manager" |
| self.kernel_size = kernel_size |
| self.even = kernel_size % 2 == 0 |
| self.pad_left = kernel_size // 2 - int(self.even) |
| self.pad_right = kernel_size // 2 |
| self.stride = stride |
| self.padding = padding |
| self.padding_mode = padding_mode |
| self.mesh_device = mesh_device |
| self.dtype = dtype |
| self.parallel_config = parallel_config |
| self.ccl_manager = ccl_manager |
|
|
| kernel = _make_kaiser_sinc_kernel_1d(cutoff, half_width, kernel_size) |
| self._taps_cpu = kernel.tolist() |
| self._conv1d_cache: dict = {} |
|
|
| def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None: |
| if "filter" in state: |
| t = state.pop("filter") |
| assert tuple(t.shape) == (1, 1, self.kernel_size) |
| self._taps_cpu = t.reshape(self.kernel_size).float().tolist() |
|
|
| def forward(self, x_BTC: ttnn.Tensor) -> ttnn.Tensor: |
| """``x_BTC``: ``(B, T, C)`` ROW_MAJOR. Returns ``(B, T_out, C)``. |
| |
| When ``parallel_config.factor > 1``, ``T`` is the *per-device* extent and |
| the replicate/zero pad becomes a halo exchange via ``_t_neighbor_pad``. |
| """ |
| assert x_BTC.layout == ttnn.ROW_MAJOR_LAYOUT |
| sharded = self.parallel_config is not None and self.parallel_config.factor > 1 |
|
|
| if self.padding: |
| if sharded: |
| x = _t_neighbor_pad( |
| x_BTC, |
| pad_left=self.pad_left, |
| pad_right=self.pad_right, |
| parallel_config=self.parallel_config, |
| ccl_manager=self.ccl_manager, |
| padding_mode=self.padding_mode, |
| ) |
| elif self.padding_mode == "replicate": |
| x = _replicate_pad_t(x_BTC, self.pad_left, self.pad_right, self.mesh_device) |
| else: |
| x = _zero_pad_t(x_BTC, self.pad_left, self.pad_right, self.mesh_device) |
| else: |
| x = x_BTC |
|
|
| return depthwise_tap_filter( |
| x, |
| self._taps_cpu, |
| self.stride, |
| mesh_device=self.mesh_device, |
| dtype=self.dtype, |
| cache=self._conv1d_cache, |
| ) |
|
|
|
|
| class UpSample1d(Module): |
| """Anti-aliased sinc upsampler (zero-stuff + depthwise lowpass). |
| |
| ``window="kaiser"`` is the BigVGAN anti-alias upsampler (T-shard/halo aware); |
| ``window="hann"`` is the torchaudio-equivalent resampler for the BWE skip path |
| (single-device). The depthwise Conv1d formulation is bit-equivalent to |
| ``F.conv_transpose1d`` for the symmetric sinc kernel. |
| """ |
|
|
| def __init__( |
| self, |
| *, |
| ratio: int = 2, |
| window: str = "kaiser", |
| kernel_size: int | None = None, |
| mesh_device: ttnn.MeshDevice, |
| dtype: ttnn.DataType = ttnn.float32, |
| parallel_config: ParallelFactor | None = None, |
| ccl_manager: CCLManager | None = None, |
| ) -> None: |
| super().__init__() |
| sharded = parallel_config is not None and parallel_config.factor > 1 |
| if sharded: |
| assert ccl_manager is not None, "T-sharding requires ccl_manager" |
| self.ratio = ratio |
| self.stride = ratio |
| self.window = window |
| self.mesh_device = mesh_device |
| self.dtype = dtype |
| self.parallel_config = parallel_config |
| self.ccl_manager = ccl_manager |
|
|
| if window == "hann": |
| kernel, self.kernel_size, self.pad, self.pad_left_crop, self.pad_right_crop = _make_hann_sinc_kernel_1d( |
| ratio=ratio |
| ) |
| self._taps_cpu = kernel.tolist() |
| elif window == "kaiser": |
| self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size |
| self.pad = self.kernel_size // ratio - 1 |
| self.pad_left_crop = self.pad * self.stride + (self.kernel_size - self.stride) // 2 |
| self.pad_right_crop = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 |
| kernel = _make_kaiser_sinc_kernel_1d( |
| cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size |
| ) |
| self._taps_cpu = kernel.tolist() |
| else: |
| raise ValueError(f"window must be kaiser or hann, got {window!r}") |
| self._conv1d_cache: dict = {} |
| self._use_polyphase = (self.kernel_size % ratio) == 0 |
| if self._use_polyphase: |
| self._poly_K_sub = self.kernel_size // ratio |
|
|
| def _prepare_torch_state(self, state: dict[str, torch.Tensor]) -> None: |
| if "filter" in state: |
| t = state.pop("filter") |
| assert tuple(t.shape) == (1, 1, self.kernel_size) |
| self._taps_cpu = t.reshape(self.kernel_size).float().tolist() |
|
|
| def forward(self, x_BTC: ttnn.Tensor) -> ttnn.Tensor: |
| assert x_BTC.layout == ttnn.ROW_MAJOR_LAYOUT |
| B, T, C = x_BTC.shape |
| sharded = self.parallel_config is not None and self.parallel_config.factor > 1 |
|
|
| poly2 = self._use_polyphase and self.ratio == 2 |
| |
| crop = 2 if (poly2 and self.pad >= 2) else 0 |
| eff_pad = self.pad - crop |
|
|
| if sharded and eff_pad > 0: |
| x_pad = _t_neighbor_pad( |
| x_BTC, |
| pad_left=eff_pad, |
| pad_right=eff_pad, |
| parallel_config=self.parallel_config, |
| ccl_manager=self.ccl_manager, |
| padding_mode="replicate", |
| ) |
| else: |
| x_pad = _replicate_pad_t(x_BTC, eff_pad, eff_pad, self.mesh_device) |
|
|
| if poly2: |
| B_, T_pad, C_ = x_pad.shape |
| |
| if crop: |
| base = x_pad |
| else: |
| base = ttnn.slice(x_pad, [0, 2, 0], [B_, T_pad - 2, C_]) |
| scaled_taps = [t * self.ratio for t in self._taps_cpu] |
| sub0 = [scaled_taps[2 * j + 0] for j in range(self._poly_K_sub)] + [0.0] |
| sub1 = [0.0] + [scaled_taps[2 * j + 1] for j in range(self._poly_K_sub)] |
| ph0 = depthwise_tap_filter( |
| base, |
| sub0, |
| 1, |
| mesh_device=self.mesh_device, |
| dtype=self.dtype, |
| cache=self._conv1d_cache, |
| ) |
| ph1 = depthwise_tap_filter( |
| base, |
| sub1, |
| 1, |
| mesh_device=self.mesh_device, |
| dtype=self.dtype, |
| cache=self._conv1d_cache, |
| ) |
| if base is not x_pad: |
| ttnn.deallocate(base) |
| T_out = ph0.shape[1] |
| ph0_b = ttnn.reshape(ph0, (B_, T_out, 1, C_)) |
| ph1_b = ttnn.reshape(ph1, (B_, T_out, 1, C_)) |
| stacked = ttnn.concat([ph0_b, ph1_b], dim=2) |
| return ttnn.reshape(stacked, (B_, T_out * 2, C_)) |
|
|
| x_zs = _zero_stuff_t(x_pad, stride=self.stride, mesh_device=self.mesh_device) |
| x_padded = _zero_pad_t(x_zs, self.kernel_size - 1, self.kernel_size - 1, self.mesh_device) |
|
|
| y = depthwise_tap_filter( |
| x_padded, |
| [t * self.ratio for t in self._taps_cpu], |
| 1, |
| mesh_device=self.mesh_device, |
| dtype=self.dtype, |
| cache=self._conv1d_cache, |
| ) |
|
|
| T_y = y.shape[1] |
| y_cropped = ttnn.slice(y, [0, self.pad_left_crop, 0], [B, T_y - self.pad_right_crop, C]) |
| ttnn.deallocate(y) |
| return y_cropped |
|
|
|
|
| class DownSample1d(Module): |
| """Strided kaiser-sinc lowpass downsampler wrapping ``LowPassFilter1d``.""" |
|
|
| def __init__( |
| self, |
| *, |
| ratio: int = 2, |
| kernel_size: int | None = None, |
| mesh_device: ttnn.MeshDevice, |
| dtype: ttnn.DataType = ttnn.float32, |
| parallel_config: ParallelFactor | None = None, |
| ccl_manager: CCLManager | None = None, |
| ) -> None: |
| super().__init__() |
| self.ratio = ratio |
| self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size |
| self.lowpass = LowPassFilter1d( |
| cutoff=0.5 / ratio, |
| half_width=0.6 / ratio, |
| stride=ratio, |
| kernel_size=self.kernel_size, |
| padding=True, |
| padding_mode="replicate", |
| mesh_device=mesh_device, |
| dtype=dtype, |
| parallel_config=parallel_config, |
| ccl_manager=ccl_manager, |
| ) |
|
|
| def forward(self, x_BTC: ttnn.Tensor) -> ttnn.Tensor: |
| return self.lowpass(x_BTC) |
|
|
|
|
| class Activation1d(Module): |
| """Anti-aliased activation: ``UpSample1d(2x) → activation → DownSample1d(2x)``.""" |
|
|
| def __init__( |
| self, |
| *, |
| channels: int, |
| activation: Module, |
| up_ratio: int = 2, |
| down_ratio: int = 2, |
| up_kernel_size: int = 12, |
| down_kernel_size: int = 12, |
| mesh_device: ttnn.MeshDevice, |
| dtype: ttnn.DataType = ttnn.float32, |
| parallel_config: ParallelFactor | None = None, |
| ccl_manager: CCLManager | None = None, |
| ) -> None: |
| super().__init__() |
| self.channels = channels |
| self.act = activation |
| self.upsample = UpSample1d( |
| ratio=up_ratio, |
| window="kaiser", |
| kernel_size=up_kernel_size, |
| mesh_device=mesh_device, |
| dtype=dtype, |
| parallel_config=parallel_config, |
| ccl_manager=ccl_manager, |
| ) |
| self.downsample = DownSample1d( |
| ratio=down_ratio, |
| kernel_size=down_kernel_size, |
| mesh_device=mesh_device, |
| dtype=dtype, |
| parallel_config=parallel_config, |
| ccl_manager=ccl_manager, |
| ) |
|
|
| def forward(self, x_BTC: ttnn.Tensor) -> ttnn.Tensor: |
| y = self.upsample(x_BTC) |
| y = self.act(y) |
| if y.layout != ttnn.ROW_MAJOR_LAYOUT: |
| y = ttnn.to_layout(y, ttnn.ROW_MAJOR_LAYOUT) |
| y = self.downsample(y) |
| return y |
|
|
