lt-tensor 0.0.1a0__py3-none-any.whl

lt_tensor/model_zoo/bsc.py

@@ -0,0 +1,210 @@
+__all__ = [
+    "FeedForward",
+    "MLP",
+    "TimestepEmbedder",
+    "GRUEncoder",
+    "ConvBlock1D",
+    "TemporalPredictor",
+    "StyleEncoder",
+    "PatchEmbed1D",
+    "MultiScaleEncoder1D",
+]
+
+from .._torch_commons import *
+from .._basics import Model
+from ..transform import get_sinusoidal_embedding
+
+
+class FeedForward(Model):
+    def __init__(
+        self,
+        d_model: int,
+        ff_dim: int,
+        dropout: float = 0.01,
+        activation: nn.Module = nn.LeakyReLU(0.1),
+        normalizer: nn.Module = nn.Identity(),
+    ):
+        """Creates a Feed-Forward Layer, with the chosen activation function and the normalizer."""
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(d_model, ff_dim),
+            activation,
+            nn.Dropout(dropout),
+            nn.Linear(ff_dim, d_model),
+            normalizer,
+        )
+
+    def forward(self, x: Tensor):
+        return self.net(x)
+
+
+class MLP(Model):
+    def __init__(
+        self,
+        d_model: int,
+        ff_dim: int,
+        n_classes: int,
+        dropout: float = 0.01,
+        activation: nn.Module = nn.LeakyReLU(0.1),
+    ):
+        """Creates a MLP block, with the chosen activation function and the normalizer."""
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(d_model, ff_dim),
+            activation,
+            nn.Dropout(dropout),
+            nn.Linear(ff_dim, n_classes),
+        )
+
+    def forward(self, x: Tensor):
+        return self.net(x)
+
+
+class TimestepEmbedder(Model):
+    def __init__(self, dim_emb: int, proj_dim: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim_emb, proj_dim),
+            nn.SiLU(),
+            nn.Linear(proj_dim, proj_dim),
+        )
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        # t: [B] (long)
+        emb = get_sinusoidal_embedding(t, self.net[0].in_features)  # [B, dim_emb]
+        return self.net(emb)  # [B, proj_dim]
+
+
+class GRUEncoder(Model):
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        num_layers: int = 1,
+        bidirectional: bool = False,
+    ):
+        super().__init__()
+        self.gru = nn.GRU(
+            input_dim,
+            hidden_dim,
+            num_layers=num_layers,
+            batch_first=True,
+            bidirectional=bidirectional,
+        )
+        self.output_dim = hidden_dim * (2 if bidirectional else 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, T, input_dim]
+        output, _ = self.gru(x)  # output: [B, T, hidden_dim*D]
+        return output
+
+
+class ConvBlock1D(Model):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        norm: bool = True,
+        residual: bool = False,
+    ):
+        super().__init__()
+        padding = (kernel_size - 1) // 2
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding)
+        self.norm = nn.BatchNorm1d(out_channels) if norm else nn.Identity()
+        self.act = nn.LeakyReLU(0.1)
+        self.residual = residual and in_channels == out_channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.act(self.norm(self.conv(x)))
+        return x + y if self.residual else y
+
+
+class TemporalPredictor(Model):
+    def __init__(
+        self,
+        d_model: int,
+        hidden_dim: int = 128,
+        n_layers: int = 2,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        layers = []
+        for _ in range(n_layers):
+            layers.append(nn.Conv1d(d_model, hidden_dim, kernel_size=3, padding=1))
+            layers.append(nn.ReLU())
+            layers.append(nn.LayerNorm(hidden_dim))
+            layers.append(nn.Dropout(dropout))
+            d_model = hidden_dim
+        self.network = nn.Sequential(*layers)
+        self.proj = nn.Linear(hidden_dim, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, T, D]
+        x = x.transpose(1, 2)  # [B, D, T]
+        x = self.network(x)  # [B, H, T]
+        x = x.transpose(1, 2)  # [B, T, H]
+        return self.proj(x).squeeze(-1)  # [B, T]
+
+
+class StyleEncoder(Model):
+    def __init__(self, in_channels: int = 80, hidden: int = 128, out_dim: int = 256):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv1d(in_channels, hidden, kernel_size=3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv1d(hidden, hidden, kernel_size=3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.AdaptiveAvgPool1d(1),
+        )
+        self.linear = nn.Linear(hidden, out_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, Mels, T]
+        x = self.net(x).squeeze(-1)  # [B, hidden]
+        return self.linear(x)  # [B, out_dim]
+
+
+class PatchEmbed1D(Model):
+    def __init__(self, in_channels: int, patch_size: int, embed_dim: int):
+        """
+        Args:
+            in_channels: number of input channels (e.g., mel bins)
+            patch_size: number of time-steps per patch
+            embed_dim: dimension of the patch embedding
+        """
+        super().__init__()
+        self.patch_size = patch_size
+        self.proj = nn.Conv1d(
+            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, C, T]
+        x = self.proj(x)  # [B, embed_dim, T//patch_size]
+        return x.transpose(1, 2)  # [B, T_patches, embed_dim]
+
+
+class MultiScaleEncoder1D(Model):
+    def __init__(
+        self, in_channels: int, hidden: int, num_layers: int = 4, kernel_size: int = 3
+    ):
+        super().__init__()
+        layers = []
+        for i in range(num_layers):
+            layers.append(
+                nn.Conv1d(
+                    in_channels if i == 0 else hidden,
+                    hidden,
+                    kernel_size=kernel_size,
+                    dilation=2**i,
+                    padding=(kernel_size - 1) * (2**i) // 2,
+                )
+            )
+            layers.append(nn.GELU())
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, C, T]
+        return self.net(x)  # [B, hidden, T]

lt_tensor/model_zoo/dfs.py

@@ -0,0 +1,181 @@
+__all__ = [
+    "Downsample1D",
+    "Upsample1D",
+    "DiffusionUNet",
+    "UNetConvBlock1D",
+    "UNetUpBlock1D",
+    "NoisePredictor1D",
+]
+
+from .._torch_commons import *
+from .._basics import Model
+from .rsd import ResBlock1D
+from ..misc_utils import log_tensor
+
+
+class Downsample1D(Model):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+    ):
+        super().__init__()
+        self.pool = nn.Conv1d(in_channels, out_channels, 4, stride=2, padding=1)
+
+    def forward(self, x):
+        return self.pool(x)
+
+
+class Upsample1D(Model):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        activation=nn.ReLU(inplace=True),
+    ):
+        super().__init__()
+        self.up = nn.Sequential(
+            nn.ConvTranspose1d(
+                in_channels, out_channels, kernel_size=4, stride=2, padding=1
+            ),
+            nn.BatchNorm1d(out_channels),
+            activation,
+        )
+
+    def forward(self, x):
+        return self.up(x)
+
+
+class DiffusionUNet(Model):
+    def __init__(self, in_channels=1, base_channels=64, out_channels=1, depth=4):
+        super().__init__()
+
+        self.depth = depth
+        self.encoder_blocks = nn.ModuleList()
+        self.downsamples = nn.ModuleList()
+        self.upsamples = nn.ModuleList()
+        self.decoder_blocks = nn.ModuleList()
+        # Keep track of channel sizes per layer for skip connections
+        self.channels = [in_channels]  # starting input channel
+        for i in range(depth):
+            enc_in = self.channels[-1]
+            enc_out = base_channels * (2**i)
+            # Encoder block and downsample
+            self.encoder_blocks.append(ResBlock1D(enc_in, enc_out))
+            self.downsamples.append(
+                Downsample1D(enc_out, enc_out)
+            )  # halve time, keep channels
+            self.channels.append(enc_out)
+        # Bottleneck
+        bottleneck_ch = self.channels[-1]
+        self.bottleneck = ResBlock1D(bottleneck_ch, bottleneck_ch)
+        # Decoder blocks (reverse channel flow)
+        for i in reversed(range(depth)):
+            skip_ch = self.channels[i + 1]  # from encoder
+            dec_out = self.channels[i]  # match earlier stage's output
+            self.upsamples.append(Upsample1D(skip_ch, skip_ch))
+            self.decoder_blocks.append(ResBlock1D(skip_ch * 2, dec_out))
+        # Final output projection (out_channels)
+        self.final = nn.Conv1d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x: Tensor):
+        skips = []
+
+        # Encoder
+        for enc, down in zip(self.encoder_blocks, self.downsamples):
+            # log_tensor(x, "before enc")
+            x = enc(x)
+            skips.append(x)
+            x = down(x)
+
+        # Bottleneck
+        x = self.bottleneck(x)
+
+        # Decoder
+        for up, dec, skip in zip(self.upsamples, self.decoder_blocks, reversed(skips)):
+            x = up(x)
+
+            # Match lengths via trimming or padding
+            if x.shape[-1] > skip.shape[-1]:
+                x = x[..., : skip.shape[-1]]
+            elif x.shape[-1] < skip.shape[-1]:
+                diff = skip.shape[-1] - x.shape[-1]
+                x = F.pad(x, (0, diff))
+
+            x = torch.cat([x, skip], dim=1)  # concat on channels
+            x = dec(x)
+
+        # Final 1x1 conv
+        return self.final(x)
+
+
+class UNetConvBlock1D(Model):
+    def __init__(self, in_channels: int, out_channels: int, down: bool = True):
+        super().__init__()
+        self.down = down
+        self.conv = nn.Sequential(
+            nn.Conv1d(
+                in_channels,
+                out_channels,
+                kernel_size=3,
+                stride=2 if down else 1,
+                padding=1,
+            ),
+            nn.BatchNorm1d(out_channels),
+            nn.LeakyReLU(0.2),
+            nn.Conv1d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm1d(out_channels),
+            nn.LeakyReLU(0.2),
+        )
+        self.downsample = (
+            nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=2 if down else 1)
+            if in_channels != out_channels
+            else nn.Identity()
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, C, T]
+        residual = self.downsample(x)
+        return self.conv(x) + residual
+
+
+class UNetUpBlock1D(Model):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv1d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm1d(out_channels),
+            nn.LeakyReLU(0.2),
+            nn.Conv1d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm1d(out_channels),
+            nn.LeakyReLU(0.2),
+        )
+        self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
+
+    def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
+        x = self.upsample(x)
+        x = torch.cat([x, skip], dim=1)  # skip connection
+        return self.conv(x)
+
+
+class NoisePredictor1D(Model):
+    def __init__(self, in_channels: int, cond_dim: int = 0, hidden: int = 128):
+        """
+        Args:
+            in_channels: channels of the noisy input [B, C, T]
+            cond_dim: optional condition vector [B, cond_dim]
+        """
+        super().__init__()
+        self.proj = nn.Linear(cond_dim, hidden) if cond_dim > 0 else None
+        self.net = nn.Sequential(
+            nn.Conv1d(in_channels, hidden, kernel_size=3, padding=1),
+            nn.SiLU(),
+            nn.Conv1d(hidden, in_channels, kernel_size=3, padding=1),
+        )
+
+    def forward(self, x: torch.Tensor, cond: Optional[torch.Tensor] = None):
+        # x: [B, C, T], cond: [B, cond_dim]
+        if cond is not None:
+            cond_proj = self.proj(cond).unsqueeze(-1)  # [B, hidden, 1]
+            x = x + cond_proj  # simple conditioning
+        return self.net(x)  # [B, C, T]

lt_tensor/model_zoo/fsn.py

@@ -0,0 +1,67 @@
+__all__ = [
+    "ConcatFusion",
+    "FiLMFusion",
+    "BilinearFusion",
+    "CrossAttentionFusion",
+    "GatedFusion",
+]
+
+from .._torch_commons import *
+from .._basics import Model
+
+
+class ConcatFusion(Model):
+    def __init__(self, in_dim_a: int, in_dim_b: int, out_dim: int):
+        super().__init__()
+        self.proj = nn.Linear(in_dim_a + in_dim_b, out_dim)
+
+    def forward(self, a: Tensor, b: Tensor) -> Tensor:
+        x = torch.cat([a, b], dim=-1)
+        return self.proj(x)
+
+
+class FiLMFusion(Model):
+    def __init__(self, cond_dim: int, feature_dim: int):
+        super().__init__()
+        self.modulator = nn.Linear(cond_dim, 2 * feature_dim)
+
+    def forward(self, x: Tensor, cond: Tensor) -> Tensor:
+        scale, shift = self.modulator(cond).chunk(2, dim=-1)
+        return x * scale + shift
+
+
+class BilinearFusion(Model):
+    def __init__(self, in_dim_a: int, in_dim_b: int, out_dim: int):
+        super().__init__()
+        self.bilinear = nn.Bilinear(in_dim_a, in_dim_b, out_dim)
+
+    def forward(self, a: Tensor, b: Tensor) -> Tensor:
+        return self.bilinear(a, b)
+
+
+class CrossAttentionFusion(nn.Module):
+    def __init__(self, q_dim: int, kv_dim: int, n_heads: int = 4, d_model: int = 256):
+        super().__init__()
+        self.q_proj = nn.Linear(q_dim, d_model)
+        self.k_proj = nn.Linear(kv_dim, d_model)
+        self.v_proj = nn.Linear(kv_dim, d_model)
+        self.attn = nn.MultiheadAttention(
+            embed_dim=d_model, num_heads=n_heads, batch_first=True
+        )
+
+    def forward(self, query: Tensor, context: Tensor, mask: Tensor = None) -> Tensor:
+        Q = self.q_proj(query)
+        K = self.k_proj(context)
+        V = self.v_proj(context)
+        output, _ = self.attn(Q, K, V, key_padding_mask=mask)
+        return output
+
+
+class GatedFusion(nn.Module):
+    def __init__(self, in_dim: int):
+        super().__init__()
+        self.gate = nn.Sequential(nn.Linear(in_dim * 2, in_dim), nn.Sigmoid())
+
+    def forward(self, a: Tensor, b: Tensor) -> Tensor:
+        gate = self.gate(torch.cat([a, b], dim=-1))
+        return gate * a + (1 - gate) * b

lt_tensor/model_zoo/pos.py

@@ -0,0 +1,121 @@
+__all__ = [
+    "RotaryEmbedding",
+    "PositionalEncoding",
+    "LearnedPositionalEncoding",
+]
+
+import math
+from .._torch_commons import *
+from .._basics import Model
+
+
+class RotaryEmbedding(Module):
+    def __init__(self, dim: int, base: int = 10000):
+        """
+        Rotary Positional Embedding Module.
+        Args:
+            dim (int): The dimension of the rotary embedding (must be even).
+            base (int): The base frequency scale (default: 10000).
+        """
+        super().__init__()
+        assert dim % 2 == 0, "Rotary dimension must be even"
+        self.dim = dim
+        self.base = base
+
+    def forward(self, x, seq_len=None):
+        """
+        Apply rotary embeddings to input tensor.
+        Args:
+            x (torch.Tensor): Input tensor of shape [batch, seq_len, dim].
+            seq_len (int, optional): Override for sequence length.
+        Returns:
+            torch.Tensor: Tensor with rotary embeddings applied.
+        """
+        bsz, seq_len = x.shape[0], seq_len or x.shape[1]
+        device = x.device
+
+        pos = torch.arange(seq_len, dtype=torch.float32, device=device).unsqueeze(1)
+        freqs = torch.pow(
+            self.base, -torch.arange(0, self.dim, 2, device=device).float() / self.dim
+        )
+        angle = pos * freqs  # [seq_len, dim/2]
+
+        sin = torch.sin(angle)
+        cos = torch.cos(angle)
+
+        # Expand and interleave to [seq_len, dim]
+        sin = torch.stack((sin, sin), dim=-1).reshape(seq_len, self.dim)
+        cos = torch.stack((cos, cos), dim=-1).reshape(seq_len, self.dim)
+
+        sin = sin.unsqueeze(0).expand(bsz, -1, -1)  # [batch, seq_len, dim]
+        cos = cos.unsqueeze(0).expand(bsz, -1, -1)
+
+        return self.apply_rotary(x, sin, cos)
+
+    def _apply_rotary(self, x, sin, cos):
+        """This version may still be useful, but for now its the problem for the text model"""
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1)
+
+    def apply_rotary(self, x, sin, cos):
+        """x: [batch, seq_len, dim]  → assume dim is even"""
+        b, s, d = x.shape
+        x = x.view(b, s, d // 2, 2)  # [b, s, d//2, 2]
+        sin = sin.view(b, s, d // 2, 2)
+        cos = cos.view(b, s, d // 2, 2)
+
+        # Apply rotation: even, odd = x[..., 0], x[..., 1]
+        x_rotated = torch.stack(
+            [
+                x[..., 0] * cos[..., 0] - x[..., 1] * sin[..., 0],
+                x[..., 0] * sin[..., 0] + x[..., 1] * cos[..., 0],
+            ],
+            dim=-1,
+        )
+
+        return x_rotated.view(b, s, d)  # Back to [b, s, d]
+
+
+class PositionalEncoding(Module):
+    def __init__(self, d_model: int, max_len: int = 8192):
+        super().__init__()
+        # create a matrix of [seq_len, hidden_dim] representing positional encoding for each token in sequence
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(
+            1
+        )  # (max_len, 1)
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2, dtype=torch.float)
+            * (-math.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe, persistent=False)  # Shape: (1, max_len, d_model)
+
+    def forward(self, x: Tensor, seq_len: Optional[Tensor] = None):
+        # x shape: (batch_size, seq_len, d_model)
+        s_sz = seq_len or x.size(1)
+        x = x + self.pe[:, :s_sz]
+        return x
+
+
+class LearnedPositionalEncoding(Module):
+    def __init__(self, max_len: int, dim_model: int, dropout: float = 0.1):
+        super().__init__()
+        self.embedding = nn.Embedding(max_len, dim_model)
+        self.dropout = nn.Dropout(dropout)
+        self.max_len = max_len
+
+    def forward(self, x: torch.Tensor, offset: int = 0) -> torch.Tensor:
+        # x: [B, T, D] or [T, D]
+        seq_len = x.size(1 if x.dim() == 3 else 0)
+        if seq_len + offset > self.max_len:
+            raise ValueError(
+                f"Sequence length {seq_len + offset} exceeds max length {self.max_len}"
+            )
+        positions = torch.arange(offset, offset + seq_len, device=x.device)
+        pos_embed = self.embedding(positions)
+        if x.dim() == 3:
+            pos_embed = pos_embed.unsqueeze(0).expand(x.size(0), -1, -1)  # [B, T, D]
+        return self.dropout(x + pos_embed)