lt-tensor 0.0.1a0__py3-none-any.whl

lt_tensor/model_zoo/rsd.py

@@ -0,0 +1,158 @@
+__all__ = [
+    "spectral_norm_select",
+    "ResBlock1D",
+    "ResBlock2D",
+    "ResBlock1D_S",
+]
+
+from .._torch_commons import *
+from .._basics import Model
+import math
+from ..misc_utils import log_tensor
+
+
+def spectral_norm_select(module: Module, enabled: bool):
+    if enabled:
+        return spectral_norm(module)
+    return module
+
+
+class ResBlock1D(Model):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        dilation: Union[Sequence[int], int] = (1, 3, 5),
+        activation: nn.Module = nn.LeakyReLU(0.1),
+        num_groups: int = 1,
+        batched: bool = True,
+    ):
+        super().__init__()
+        self.conv = nn.ModuleList()
+        if isinstance(dilation, int):
+            dilation = [dilation]
+
+        if batched:
+            layernorm_fn = lambda x: nn.GroupNorm(num_groups=num_groups, num_channels=x)
+        else:
+            layernorm_fn = lambda x: nn.LayerNorm(normalized_shape=x)
+        for i, dil in enumerate(dilation):
+
+            self.conv.append(
+                nn.ModuleDict(
+                    dict(
+                        net=nn.Sequential(
+                            self._get_conv_layer(
+                                in_channels, in_channels, kernel_size, dil
+                            ),
+                            activation,
+                            self._get_conv_layer(
+                                in_channels, in_channels, kernel_size, 1, True
+                            ),
+                            activation,
+                        ),
+                        l_norm=layernorm_fn(in_channels),
+                    )
+                )
+            )
+        self.final = nn.Sequential(
+            self._get_conv_layer(in_channels, out_channels, kernel_size, 1, True),
+            activation,
+        )
+        self.conv.apply(self.init_weights)
+
+    def _get_conv_layer(
+        self,
+        channels_in: int,
+        channels_out: int,
+        kernel_size: int,
+        dilation: int,
+        pad_gate: bool = False,
+    ):
+        return weight_norm(
+            nn.Conv1d(
+                in_channels=channels_in,
+                out_channels=channels_out,
+                kernel_size=kernel_size,
+                stride=1,
+                dilation=dilation,
+                padding=(
+                    int((kernel_size * dilation - dilation) / 2)
+                    if not pad_gate
+                    else int((kernel_size * 1 - 1) / 2)
+                ),
+            )
+        )
+
+    def forward(self, x: Tensor):
+        for i, layer in enumerate(self.conv):
+            xt = layer["net"](x)
+            x = xt + x
+            x = layer["l_norm"](x)
+        return self.final(x)
+
+    def remove_weight_norm(self):
+        for module in self.modules():
+            try:
+                remove_weight_norm(module)
+            except ValueError:
+                pass  # Not normed, skip
+
+    @staticmethod
+    def init_weights(m, mean=0.0, std=0.01):
+        classname = m.__class__.__name__
+        if "Conv" in classname:
+            m.weight.data.normal_(mean, std)
+
+
+class ResBlock2D(Model):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        downsample=False,
+        spec_norm: bool = False,
+    ):
+        super().__init__()
+        stride = 2 if downsample else 1
+
+        self.block = nn.Sequential(
+            spectral_norm_select(
+                nn.Conv2d(in_channels, out_channels, 3, stride, 1), spec_norm
+            ),
+            nn.LeakyReLU(0.2),
+            spectral_norm_select(
+                nn.Conv2d(out_channels, out_channels, 3, 1, 1), spec_norm
+            ),
+        )
+
+        self.skip = nn.Identity()
+        if downsample or in_channels != out_channels:
+            self.skip = spectral_norm_select(
+                nn.Conv2d(in_channels, out_channels, 1, stride), spec_norm
+            )
+        # on less to be handled every cicle
+        self.sqrt_2 = math.sqrt(2)
+
+    def forward(self, x):
+        return (self.block(x) + self.skip(x)) / self.sqrt_2
+
+
+class ResBlock1D_S(Model):
+    """Simplified version"""
+
+    def __init__(self, channels: int, kernel_size: int = 3, dilation: int = 1):
+        super().__init__()
+        padding = (kernel_size - 1) // 2 * dilation
+        self.net = nn.Sequential(
+            nn.Conv1d(
+                channels, channels, kernel_size, padding=padding, dilation=dilation
+            ),
+            nn.LeakyReLU(0.1),
+            nn.Conv1d(channels, channels, kernel_size, padding=padding, dilation=1),
+        )
+        self.activation = nn.LeakyReLU(0.1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.activation(x + self.net(x))

lt_tensor/model_zoo/tfr.py

@@ -0,0 +1,140 @@
+__all__ = [
+    "TransformerEncoderLayer",
+    "TransformerDecoderLayer",
+    "TransformerEncoder",
+    "TransformerDecoder",
+    "init_weights",
+]
+
+import math
+from .._torch_commons import *
+from .._basics import Model
+from lt_utils.misc_utils import default
+
+from .pos import *
+from .bsc import FeedForward
+
+
+def init_weights(module):
+    if isinstance(module, nn.Linear):
+        nn.init.xavier_uniform_(module.weight)
+        if module.bias is not None:
+            nn.init.constant_(module.bias, 0)
+    elif isinstance(module, nn.Embedding):
+        nn.init.normal_(module.weight, mean=0.0, std=0.02)
+    elif isinstance(module, nn.LayerNorm):
+        nn.init.constant_(module.bias, 0)
+        nn.init.constant_(module.weight, 1.0)
+
+
+class TransformerEncoderLayer(Model):
+    def __init__(self, d_model: int, n_heads: int, ff_size: int, dropout: float = 0.1):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(
+            d_model, n_heads, dropout=dropout, batch_first=True
+        )
+        self.norm1 = nn.LayerNorm(d_model)
+        self.ff = FeedForward(d_model, ff_size, dropout)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: Tensor, src_mask: Optional[Tensor] = None) -> Tensor:
+        attn_output, _ = self.self_attn(x, x, x, attn_mask=src_mask)
+        x = self.norm1(x + self.dropout(attn_output))
+        ff_output = self.ff(x)
+        x = self.norm2(x + self.dropout(ff_output))
+        return x
+
+
+class TransformerDecoderLayer(Model):
+    def __init__(self, d_model: int, n_heads: int, ff_size: int, dropout: float = 0.1):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(
+            d_model, n_heads, dropout=dropout, batch_first=True
+        )
+        self.norm1 = nn.LayerNorm(d_model)
+
+        self.cross_attn = nn.MultiheadAttention(
+            d_model, n_heads, dropout=dropout, batch_first=True
+        )
+        self.norm2 = nn.LayerNorm(d_model)
+
+        self.ff = FeedForward(d_model, ff_size, dropout)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x: Tensor,  # Decoder input [B, T, d_model]
+        encoder_out: Tensor,  # Encoder output [B, S, d_model]
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        # 1. Masked Self-Attention
+        attn_output, _ = self.self_attn(x, x, x, attn_mask=tgt_mask)
+        x = self.norm1(x + self.dropout(attn_output))
+
+        # 2. Cross-Attention
+        cross_output, _ = self.cross_attn(
+            x, encoder_out, encoder_out, attn_mask=memory_mask
+        )
+        x = self.norm2(x + self.dropout(cross_output))
+
+        # 3. FeedForward
+        ff_output = self.ff(x)
+        x = self.norm3(x + self.dropout(ff_output))
+        return x
+
+
+class TransformerEncoder(Model):
+    def __init__(
+        self,
+        d_model: int = 64,
+        n_heads: int = 4,
+        ff_size: int = 256,
+        num_layers: int = 2,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                TransformerEncoderLayer(d_model, n_heads, ff_size, dropout)
+                for _ in range(num_layers)
+            ]
+        )
+
+    def forward(self, x: Tensor, src_mask: Optional[Tensor] = None) -> Tensor:
+
+        for layer in self.layers:
+            x = layer(x, src_mask)
+        return x
+
+
+class TransformerDecoder(Model):
+    def __init__(
+        self,
+        d_model: int = 64,
+        n_heads: int = 2,
+        ff_size: int = 256,
+        num_layers: int = 2,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+
+        self.layers = nn.ModuleList(
+            [
+                TransformerDecoderLayer(d_model, n_heads, ff_size, dropout)
+                for _ in range(num_layers)
+            ]
+        )
+
+    def forward(
+        self,
+        x: Tensor,
+        encoder_out: Tensor,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        for layer in self.layers:
+            x = layer(x, encoder_out, tgt_mask, memory_mask)
+        return x

lt_tensor/monotonic_align.py

@@ -0,0 +1,70 @@
+from numba import njit, prange
+
+
+@njit()
+def maximum_path_each(path, value, t_x, t_y, max_neg_val):
+    index = t_x - 1
+    # Forward pass: Calculate max path sums
+    for y in range(t_y):
+        for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
+            v_cur = max_neg_val if x == y else value[x, y - 1]
+            v_prev = (
+                0.0
+                if (x == 0 and y == 0)
+                else (max_neg_val if x == 0 else value[x - 1, y - 1])
+            )
+            value[x, y] = max(v_cur, v_prev) + value[x, y]
+
+    # Backtrack to store the path
+    for y in range(t_y - 1, -1, -1):
+        path[index, y] = 1
+        if index != 0 and (index == y or value[index, y - 1] < value[index - 1, y - 1]):
+            index -= 1
+
+
+@njit()  # Took almost 10x the time while testing using "parallel=True".
+def maximum_path(paths, values, t_xs, t_ys, max_neg_val=-1e9):
+    """
+    Example:
+    ```python
+        paths = tc.randn((2, 3, 3)).numpy()
+        values = tc.randn((2, 3, 3)).numpy()
+        t_xs = tc.tensor([3, 3, 3]).numpy()
+        t_ys = tc.tensor([3, 3]).numpy()
+
+        # to display values (before) and paths:
+        print("=====================")
+        print("Paths:")
+        print(paths)
+        print("Original Values:")
+        print(values)
+
+        maximum_path(paths, values, t_xs, t_ys)
+
+        print("Updated Values:")
+        print(values)
+        print("=====================")
+
+    ```
+    Outputs:
+    ```md
+        =====================
+        Paths:
+        [[[ 2.310408   -1.9375949  -0.57884663]
+        [ 1.0308106   1.0793993   0.4461908 ]
+        [ 0.26789713  0.48924422  0.3409592 ]]]
+        Original Values:
+        [[[-0.48256454  0.51348686 -1.8236492 ]
+        [ 0.9949021  -0.6066166   0.18991096]
+        [ 1.2555764  -0.24222293 -0.78757876]]]
+        Updated Values:
+        [[[-0.48256454  0.51348686 -1.8236492 ]
+        [ 0.9949021  -1.0891812   0.18991096]
+        [ 1.2555764  -0.24222293 -1.87676   ]]]
+        =====================
+    ```
+    This may not be the standard, but may work for your project.
+    """
+    batch_size = values.shape[0]
+    for i in prange(batch_size):
+        maximum_path_each(paths[i], values[i], t_xs[i], t_ys[i], max_neg_val)

lt_tensor/transform.py

@@ -0,0 +1,349 @@
+__all__ = [
+    "to_mel_spectrogram",
+    "stft",
+    "istft",
+    "fft",
+    "ifft",
+    "to_log_mel_spectrogram",
+    "normalize",
+    "min_max_scale",
+    "mel_to_linear",
+    "add_noise",
+    "shift_time",
+    "stretch_tensor",
+    "pad_tensor",
+    "get_sinusoidal_embedding",
+    "pad_center",
+    "normalize",
+    "window_sumsquare",
+    "inverse_transform",
+    "stft_istft_rebuild",
+]
+
+from ._torch_commons import *
+import torchaudio
+import math
+from .misc_utils import log_tensor
+
+
+def to_mel_spectrogram(
+    waveform: torch.Tensor,
+    sample_rate: int = 22050,
+    n_fft: int = 1024,
+    hop_length: Optional[int] = None,
+    win_length: Optional[int] = None,
+    n_mels: int = 80,
+    f_min: float = 0.0,
+    f_max: Optional[float] = None,
+) -> torch.Tensor:
+    """Converts waveform to mel spectrogram."""
+    mel_spectrogram = torchaudio.transforms.MelSpectrogram(
+        sample_rate=sample_rate,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        n_mels=n_mels,
+        f_min=f_min,
+        f_max=f_max,
+    )
+    return mel_spectrogram(waveform)
+
+
+def stft(
+    waveform: Tensor,
+    n_fft: int = 512,
+    hop_length: Optional[int] = None,
+    win_length: Optional[int] = None,
+    window_fn: str = "hann",
+    center: bool = True,
+    return_complex: bool = True,
+) -> Tensor:
+    """Performs short-time Fourier transform using PyTorch."""
+    window = (
+        torch.hann_window(win_length or n_fft).to(waveform.device)
+        if window_fn == "hann"
+        else None
+    )
+    return torch.stft(
+        input=waveform,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=window,
+        center=center,
+        return_complex=return_complex,
+    )
+
+
+def istft(
+    stft_matrix: Tensor,
+    n_fft: int = 512,
+    hop_length: Optional[int] = None,
+    win_length: Optional[int] = None,
+    window_fn: str = "hann",
+    center: bool = True,
+    length: Optional[int] = None,
+) -> Tensor:
+    """Performs inverse short-time Fourier transform using PyTorch."""
+    window = (
+        torch.hann_window(win_length or n_fft).to(stft_matrix.device)
+        if window_fn == "hann"
+        else None
+    )
+    return torch.istft(
+        input=stft_matrix,
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=window,
+        center=center,
+        length=length,
+    )
+
+
+def fft(x: Tensor, norm: Optional[str] = "backward") -> Tensor:
+    """Returns the FFT of a real tensor."""
+    return torch.fft.fft(x, norm=norm)
+
+
+def ifft(x: Tensor, norm: Optional[str] = "backward") -> Tensor:
+    """Returns the inverse FFT of a complex tensor."""
+    return torch.fft.ifft(x, norm=norm)
+
+
+def to_log_mel_spectrogram(
+    waveform: torch.Tensor, sample_rate: int = 22050, eps: float = 1e-9, **kwargs
+) -> torch.Tensor:
+    """Converts waveform to log-mel spectrogram."""
+    mel = to_mel_spectrogram(waveform, sample_rate, **kwargs)
+    return torch.log(mel + eps)
+
+
+def normalize(
+    x: torch.Tensor,
+    mean: Optional[float] = None,
+    std: Optional[float] = None,
+    eps: float = 1e-9,
+) -> torch.Tensor:
+    """Normalizes tensor by mean and std."""
+    if mean is None:
+        mean = x.mean()
+    if std is None:
+        std = x.std()
+    return (x - mean) / (std + eps)
+
+
+def min_max_scale(
+    x: torch.Tensor, min_val: float = 0.0, max_val: float = 1.0
+) -> torch.Tensor:
+    """Scales tensor to [min_val, max_val] range."""
+    x_min, x_max = x.min(), x.max()
+    return (x - x_min) / (x_max - x_min + 1e-8) * (max_val - min_val) + min_val
+
+
+def mel_to_linear(
+    mel_spec: torch.Tensor, mel_fb: torch.Tensor, eps: float = 1e-8
+) -> torch.Tensor:
+    """Approximate inversion of mel to linear spectrogram using pseudo-inverse."""
+    mel_fb_inv = torch.pinverse(mel_fb)
+    return torch.matmul(mel_fb_inv, mel_spec + eps)
+
+
+def add_noise(x: torch.Tensor, noise_level: float = 0.01) -> torch.Tensor:
+    """Adds Gaussian noise to tensor."""
+    return x + noise_level * torch.randn_like(x)
+
+
+def shift_time(x: torch.Tensor, shift: int) -> torch.Tensor:
+    """Shifts tensor along time axis (last dim)."""
+    return torch.roll(x, shifts=shift, dims=-1)
+
+
+def stretch_tensor(x: torch.Tensor, rate: float, mode: str = "linear") -> torch.Tensor:
+    """Time-stretch tensor using interpolation."""
+    B, C, T = x.shape if x.ndim == 3 else (1, 1, x.shape[0])
+    new_T = int(T * rate)
+    x_reshaped = x.view(B * C, T).unsqueeze(1)
+    stretched = torch.nn.functional.interpolate(x_reshaped, size=new_T, mode=mode)
+    return (
+        stretched.squeeze(1).view(B, C, new_T) if x.ndim == 3 else stretched.squeeze()
+    )
+
+
+def pad_tensor(
+    x: torch.Tensor, target_len: int, pad_value: float = 0.0
+) -> torch.Tensor:
+    """Pads tensor to target length along last dimension."""
+    current_len = x.shape[-1]
+    if current_len >= target_len:
+        return x[..., :target_len]
+    padding = [0] * (2 * (x.ndim - 1)) + [0, target_len - current_len]
+    return F.pad(x, padding, value=pad_value)
+
+
+def get_sinusoidal_embedding(timesteps: torch.Tensor, dim: int) -> torch.Tensor:
+    # Expect shape [B] or [B, 1]
+    if timesteps.dim() > 1:
+        timesteps = timesteps.view(-1)  # flatten to [B]
+
+    device = timesteps.device
+    half_dim = dim // 2
+    emb = torch.exp(
+        torch.arange(half_dim, device=device) * -(math.log(10000.0) / half_dim)
+    )
+    emb = timesteps[:, None].float() * emb[None, :]  # [B, half_dim]
+    emb = torch.cat((emb.sin(), emb.cos()), dim=-1)  # [B, dim]
+    return emb
+
+
+def _generate_window(
+    M: int, alpha: float = 0.5, device: Optional[DeviceType] = None
+) -> Tensor:
+    if M < 1:
+        raise ValueError("Window length M must be >= 1.")
+    if M == 1:
+        return torch.ones(1, device=device)
+
+    n = torch.arange(M, dtype=torch.float32, device=device)
+    window = alpha - (1.0 - alpha) * torch.cos(2.0 * math.pi * n / (M - 1))
+    return window
+
+
+def pad_center(tensor: torch.Tensor, size: int, axis: int = -1) -> torch.Tensor:
+    n = tensor.shape[axis]
+    if size < n:
+        raise ValueError(f"Target size ({size}) must be at least input size ({n})")
+
+    lpad = (size - n) // 2
+    rpad = size - n - lpad
+
+    pad = [0] * (2 * tensor.ndim)
+    pad[2 * axis + 1] = rpad
+    pad[2 * axis] = lpad
+
+    return F.pad(tensor, pad, mode="constant", value=0)
+
+
+def normalize(
+    S: torch.Tensor,
+    norm: float = float("inf"),
+    axis: int = 0,
+    threshold: float = 1e-10,
+    fill: bool = False,
+) -> torch.Tensor:
+    mag = S.abs().float()
+
+    if norm is None:
+        return S
+
+    elif norm == float("inf"):
+        length = mag.max(dim=axis, keepdim=True).values
+
+    elif norm == float("-inf"):
+        length = mag.min(dim=axis, keepdim=True).values
+
+    elif norm == 0:
+        length = (mag > 0).sum(dim=axis, keepdim=True).float()
+
+    elif norm > 0:
+        length = (mag**norm).sum(dim=axis, keepdim=True) ** (1.0 / norm)
+
+    else:
+        raise ValueError(f"Unsupported norm: {norm}")
+
+    small_idx = length < threshold
+    length = length.clone()
+    if fill:
+        length[small_idx] = float("nan")
+        Snorm = S / length
+        Snorm[Snorm != Snorm] = 1.0  # replace nan with fill_norm (default 1.0)
+    else:
+        length[small_idx] = float("inf")
+        Snorm = S / length
+
+    return Snorm
+
+
+def window_sumsquare(
+    window_spec: Union[str, int, float, Callable, List[Any], Tuple[Any, ...]],
+    n_frames: int,
+    hop_length: int = 300,
+    win_length: int = 1200,
+    n_fft: int = 2048,
+    dtype: torch.dtype = torch.float32,
+    norm: Optional[Union[int, float]] = None,
+    device: Optional[torch.device] = "cpu",
+):
+    if win_length is None:
+        win_length = n_fft
+
+    total_length = n_fft + hop_length * (n_frames - 1)
+    x = torch.zeros(total_length, dtype=dtype, device=device)
+
+    # Get the window (from scipy for now)
+    win = _generate_window(window_spec, win_length, fftbins=True)
+    win = torch.tensor(win, dtype=dtype, device=device)
+
+    # Normalize and square
+    win_sq = normalize(win, norm=norm, axis=0) ** 2
+    win_sq = pad_center(win_sq, size=n_fft, axis=0)
+
+    # Accumulate squared windows
+    for i in range(n_frames):
+        sample = i * hop_length
+        end = min(total_length, sample + n_fft)
+        length = end - sample
+        x[sample:end] += win_sq[:length]
+
+    return x
+
+
+def inverse_transform(
+    spec: Tensor,
+    phase: Tensor,
+    window: Optional[Tensor] = None,
+    n_fft: int = 2048,
+    hop_length: int = 300,
+    win_length: int = 1200,
+    length: Optional[Any] = None,
+):
+    if window is None:
+        window = _generate_window(win_length)
+    return torch.istft(
+        spec * torch.exp(phase * 1j),
+        n_fft,
+        hop_length,
+        win_length,
+        window=window,
+        length=length,
+    )
+
+
+def stft_istft_rebuild(
+    input_data: Tensor,
+    window: Optional[Tensor] = None,
+    n_fft: int = 2048,
+    hop_length: int = 300,
+    win_length: int = 1200,
+):
+    """
+    Perform STFT followed by ISTFT reconstruction using magnitude and phase.
+    """
+    if window is None:
+        window = _generate_window(win_length)
+    st = torch.stft(
+        input_data,
+        n_fft,
+        hop_length,
+        win_length,
+        window=window,
+        return_complex=True,
+    )
+    return torch.istft(
+        torch.abs(st) * torch.exp(1j * torch.angle(st)),
+        n_fft,
+        hop_length,
+        win_length,
+        window=window,
+        length=input_data.shape[-1],
+    ).squeeze(0)