lt-tensor 0.0.1.dev0__py3-none-any.whl

lt_tensor/model_zoo/diffusion/models.py

@@ -0,0 +1,114 @@
+__all__ = [
+    "ResidualBlock1D_B",
+    "Downsample1D",
+    "Upsample1D",
+    "DiffusionUNet",
+    "DiffusionUNetT",
+]
+
+from ..._torch_commons import *
+from ..._basics import Model
+from ..residual import ResBlock1D, ResBlock1DT
+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)
+
+
+
+
+

lt_tensor/model_zoo/residual.py

@@ -0,0 +1,236 @@
+from .._torch_commons import *
+from .._basics import Model
+import math
+from ..misc_utils import log_tensor
+
+
+def initialize__weights(model: nn.Module, method: str = "xavier"):
+    """Initialize model weights using specified method."""
+    for name, param in model.named_parameters():
+        if "weight" in name:
+            if method == "xavier":
+                nn.init.xavier_uniform_(param)
+            elif method == "kaiming":
+                nn.init.kaiming_uniform_(param, nonlinearity="relu")
+        elif "bias" in name:
+            nn.init.constant_(param, 0)
+
+
+def spectral_norm_select(module: Module, enabled: bool):
+    if enabled:
+        return spectral_norm(module)
+    return module
+
+
+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 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(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
+
+
+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 ResBlock1DT(Model):
+    """For time based residual layers"""
+
+    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,
+        time_emb_dim: int = 1,
+        batched: bool = True,
+    ):
+        super().__init__()
+        self.conv = nn.ModuleList()
+        if isinstance(dilation, int):
+            dilation = [dilation]
+        self.time_proj = nn.Linear(time_emb_dim, out_channels)
+        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(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, t_embed: Optional[Tensor] = None):
+        if t_embed is not None:
+            t_emb = self.time_proj(t_embed).unsqueeze(-1)  # [B, C, 1]
+
+        for i, layer in enumerate(self.conv):
+            if t_embed is not None:
+                xt = layer["net"](x) + t_emb
+            else:
+                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

lt_tensor/model_zoo/transformer_models/__init__.py

@@ -0,0 +1,6 @@
+from .models import (
+    TransformerEncoderLayer,
+    TransformerDecoderLayer,
+    TransformerEncoder,
+    TransformerDecoder,
+)

lt_tensor/model_zoo/transformer_models/models.py

@@ -0,0 +1,132 @@
+import math
+from ..._torch_commons import *
+from ..._basics import Model
+from lt_utils.misc_utils import default
+
+from .positional_encoders import *
+from ..basic 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(nn.Module):
+    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(nn.Module):
+    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/model_zoo/transformer_models/positional_encoders.py

@@ -0,0 +1,95 @@
+__all__ = ["RotaryEmbedding", "PositionalEncoding"]
+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

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)