RustyMark/dots.tts

1import torch2import torch.nn as nn3import torch.nn.functional as F4from einops import rearrange5 6 7class Dropout(nn.Module):8    def __init__(9        self, p: float = 0.5, inplace: bool = False, force_drop: bool = False, **_kwargs10    ):11        super().__init__()12        if p < 0.0 or p > 1.0:13            raise ValueError(14                f"dropout probability has to be between 0 and 1, but got {p}"15            )16        self.p = p17        self.inplace = inplace18        self.force_drop = force_drop19 20    def forward(self, x, **_kwargs):21        return F.dropout(22            x,23            p=self.p,24            training=True if self.force_drop else self.training,25            inplace=self.inplace,26        )27 28 29class Conv1d(nn.Conv1d):30    def __init__(31        self,32        in_channels: int,33        out_channels: int,34        kernel_size: int = 1,35        stride: int = 1,36        dilation: int = 1,37        groups: int = 1,38        padding_mode: str = "zeros",39        bias: bool = True,40        padding=None,41        causal: bool = False,42        **_kwargs,43    ):44        self.causal = causal45        if padding is None:46            if causal:47                padding = 048                self.left_padding = dilation * (kernel_size - 1)49            else:50                padding = int((kernel_size * dilation - dilation) / 2)51 52        super().__init__(53            in_channels,54            out_channels,55            kernel_size,56            stride=stride,57            padding=padding,58            dilation=dilation,59            groups=groups,60            padding_mode=padding_mode,61            bias=bias,62        )63 64        self.in_channels = in_channels65 66    def forward(self, x):67        if self.causal:68            x = F.pad(x.unsqueeze(2), (self.left_padding, 0, 0, 0)).squeeze(2)69        return super().forward(x)70 71 72class ConvTranspose1d(nn.ConvTranspose1d):73    def __init__(74        self,75        in_channels: int,76        out_channels: int,77        kernel_size: int,78        stride: int = 1,79        output_padding: int = 0,80        groups: int = 1,81        bias: bool = True,82        dilation: int = 1,83        padding=None,84        padding_mode: str = "zeros",85        causal: bool = False,86        **_kwargs,87    ):88        if padding is None:89            padding = 0 if causal else (kernel_size - stride) // 290        if causal:91            assert padding == 0, "padding is not allowed in causal ConvTranspose1d."92            assert kernel_size == 2 * stride, (93                "kernel_size must be equal to 2*stride in Causal ConvTranspose1d."94            )95 96        super().__init__(97            in_channels,98            out_channels,99            kernel_size,100            stride=stride,101            padding=padding,102            output_padding=output_padding,103            groups=groups,104            bias=bias,105            dilation=dilation,106            padding_mode=padding_mode,107        )108 109        self.causal = causal110        self.stride = stride111 112    def forward(self, x):113        x = super().forward(x)114        if self.causal:115            x = x[:, :, : -self.stride]116        return x117 118 119class Mlp(nn.Module):120    def __init__(121        self,122        hidden_size,123        ffn_hidden_size=4096,124        act_layer=nn.GELU,125        dropout=0.0,126        **_kwargs,127    ):128        super().__init__()129        self.fc1 = nn.Linear(hidden_size, ffn_hidden_size)130        self.act = act_layer()131        self.fc2 = nn.Linear(ffn_hidden_size, hidden_size)132        self.drop = Dropout(dropout)133 134    def forward(self, x, _mask=None):135        x = self.fc1(x)136        x = self.act(x)137        x = self.drop(x)138        x = self.fc2(x)139        return self.drop(x)140 141 142def rotate_half(x):143    x1, x2 = x.chunk(2, dim=-1)144    return torch.cat((-x2, x1), dim=-1)145 146 147@torch.autocast(enabled=False, device_type="cuda")148def apply_rotary_pos_emb(pos, t):149    if pos.dim() == 3:150        pos = pos.unsqueeze(1)151    return t * pos.cos() + rotate_half(t) * pos.sin()152 153 154class RotaryEmbedding(nn.Module):155    def __init__(self, dim, theta=50000):156        super().__init__()157        self.register_buffer(158            "inv_freq",159            1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)),160            persistent=False,161        )162        self._theta = float(theta)163 164    def _apply(self, fn):165        inv_freq = self.inv_freq166        super()._apply(fn)167        self.inv_freq = inv_freq.to(device=self.inv_freq.device, dtype=torch.float32)168        return self169 170    @torch.autocast(enabled=False, device_type="cuda")171    def forward(self, t):172        inv_freq = self.inv_freq173        if inv_freq.device != t.device:174            raise RuntimeError(175                "RotaryEmbedding buffer device mismatch: "176                f"inv_freq={inv_freq.device} input={t.device}."177            )178        t = t.to(dtype=inv_freq.dtype)179        if t.dim() == 1:180            freqs = torch.einsum("i , j -> i j", t, inv_freq)181        else:182            freqs = torch.einsum("bi, j -> bij", t, inv_freq)183        return torch.cat((freqs, freqs), dim=-1)184 185 186class MultiHeadAttention(nn.Module):187    """Multi-head attention"""188 189    def __init__(190        self,191        hidden_size: int,192        num_heads: int = 8,193        qkv_bias: bool = False,194        qk_norm: bool = False,195        attn_drop: float = 0.0,196        dropout: float = 0.0,197        norm_layer: str = "LayerNorm",198        rotary_bias: bool = False,199        rotary_theta: float | None = 50000,200        **_kwargs,201    ):202        super().__init__()203        assert hidden_size % num_heads == 0, (204            "hidden_size should be divisible by num_heads"205        )206        self.num_heads = num_heads207        self.head_dim = hidden_size // num_heads208        self.scale = self.head_dim**-0.5209        self.rotary_bias = rotary_bias210 211        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)212        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)213        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)214 215        norm_layer = getattr(nn, norm_layer)216        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()217        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()218 219        self.attn_drop = Dropout(attn_drop)220        self.o_proj = nn.Linear(hidden_size, hidden_size)221        self.o_dropout = Dropout(dropout)222 223        if self.rotary_bias:224            self.rotary = RotaryEmbedding(self.head_dim, theta=rotary_theta)225 226    def forward(self, q, k=None, v=None, mask=None, pos_ids=None, **_kwargs):227        k = k or q228        v = v or q229        B, L, _ = q.shape230        _, S, _ = v.shape231        if mask is not None:232            if mask.ndim == 2:  # [B, L]233                assert L == S234                mask = rearrange(mask, "b j -> b 1 1 j")235                mask = mask.expand(-1, self.num_heads, L, -1)236            elif mask.ndim == 3:  # [B, L, S]237                assert mask.size(1) == L and mask.size(2) == S238                mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)239 240        q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)241        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)242        k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads)243        v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads)244        q, k = self.q_norm(q), self.k_norm(k)245 246        # Apply rotary247        if self.rotary_bias:248            if L == S:249                if pos_ids is None:250                    rotary_emb = self.rotary(torch.arange(L, device=q.device))251                else:252                    rotary_emb = self.rotary(pos_ids)253                q, k = (apply_rotary_pos_emb(rotary_emb, tensor) for tensor in (q, k))254            else:255                q_rotary_emb = self.rotary(torch.arange(L, device=q.device))256                k_rotary_emb = self.rotary(torch.arange(S, device=k.device))257                q = apply_rotary_pos_emb(q_rotary_emb, q)258                k = apply_rotary_pos_emb(k_rotary_emb, k)259 260        attn_bias = torch.zeros(B, self.num_heads, L, S, dtype=q.dtype, device=q.device)261 262        if mask is not None:263            attn_bias.masked_fill_(mask.logical_not(), float("-inf"))264 265        out = F.scaled_dot_product_attention(266            q,267            k,268            v,269            attn_mask=attn_bias,270            dropout_p=self.attn_drop.p if self.training else 0.0,271        )272 273        out = rearrange(out, "b h n d -> b n (h d)")274        return self.o_dropout(self.o_proj(out))275 276    def decode_step(self, x, *, cache, positions: torch.Tensor):277        if x.size(1) <= 0:278            raise ValueError("MultiHeadAttention.decode_step expects a non-empty input.")279        if positions.ndim != 1 or positions.size(0) != x.size(1):280            raise ValueError(281                "MultiHeadAttention.decode_step positions must match the decode block length."282            )283 284        q = self.q_proj(x)285        k = self.k_proj(x)286        v = self.v_proj(x)287 288        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)289        k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads)290        v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads)291        q, k = self.q_norm(q), self.k_norm(k)292        block_len = q.size(2)293 294        if self.rotary_bias:295            rotary_emb = self.rotary(positions)296            q = apply_rotary_pos_emb(rotary_emb, q)297            k = apply_rotary_pos_emb(rotary_emb, k)298 299        cached_k, cached_v = cache300        cached_k.index_copy_(2, positions, k)301        cached_v.index_copy_(2, positions, v)302 303        cache_capacity = cached_k.size(2)304        key_positions = torch.arange(305            cache_capacity,306            device=x.device,307            dtype=torch.long,308        ).unsqueeze(0)309        query_positions = positions.unsqueeze(1)310        causal_mask = key_positions <= query_positions311        valid_mask = key_positions <= positions[-1]312        attn_bias = torch.zeros(313            q.size(0),314            self.num_heads,315            block_len,316            cache_capacity,317            dtype=q.dtype,318            device=q.device,319        )320        attn_bias.masked_fill_(321            (causal_mask & valid_mask).unsqueeze(0).unsqueeze(0).logical_not(),322            float("-inf"),323        )324 325        out = F.scaled_dot_product_attention(326            q,327            cached_k,328            cached_v,329            attn_mask=attn_bias,330            dropout_p=self.attn_drop.p if self.training else 0.0,331        )332        out = rearrange(out, "b h n d -> b n (h d)")333        return self.o_dropout(self.o_proj(out)), cache334