Delete CHATTS/model/dvae.py

CHATTS/model/dvae.py

@@ -1,155 +0,0 @@
-import math
-from einops import rearrange
-from vector_quantize_pytorch import GroupedResidualFSQ
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-class ConvNeXtBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        intermediate_dim: int,
-        kernel, dilation,
-        layer_scale_init_value: float = 1e-6,
-    ):
-        # ConvNeXt Block copied from Vocos.
-        super().__init__()
-        self.dwconv = nn.Conv1d(dim, dim, 
-                                kernel_size=kernel, padding=dilation*(kernel//2), 
-                                dilation=dilation, groups=dim
-                            )  # depthwise conv
-        
-        self.norm = nn.LayerNorm(dim, eps=1e-6)
-        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
-        self.act = nn.GELU()
-        self.pwconv2 = nn.Linear(intermediate_dim, dim)
-        self.gamma = (
-            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
-            if layer_scale_init_value > 0
-            else None
-        )
-
-    def forward(self, x: torch.Tensor, cond = None) -> torch.Tensor:
-        residual = x
-        x = self.dwconv(x)
-        x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
-        x = self.norm(x)
-        x = self.pwconv1(x)
-        x = self.act(x)
-        x = self.pwconv2(x)
-        if self.gamma is not None:
-            x = self.gamma * x
-        x = x.transpose(1, 2)  # (B, T, C) -> (B, C, T)
-
-        x = residual + x
-        return x
-    
-
-
-class GFSQ(nn.Module):
-
-    def __init__(self, 
-            dim, levels, G, R, eps=1e-5, transpose = True
-        ):
-        super(GFSQ, self).__init__()
-        self.quantizer = GroupedResidualFSQ(
-            dim=dim,
-            levels=levels,
-            num_quantizers=R,
-            groups=G,
-        )
-        self.n_ind = math.prod(levels)
-        self.eps = eps
-        self.transpose = transpose
-        self.G = G
-        self.R = R
-        
-    def _embed(self, x):
-        if self.transpose:
-            x = x.transpose(1,2)
-        x = rearrange(
-            x, "b t (g r) -> g b t r", g = self.G, r = self.R,
-        )  
-        feat = self.quantizer.get_output_from_indices(x)
-        return feat.transpose(1,2) if self.transpose else feat
-        
-    def forward(self, x,):
-        if self.transpose:
-            x = x.transpose(1,2)
-        feat, ind = self.quantizer(x)
-        ind = rearrange(
-            ind, "g b t r ->b t (g r)",
-        )  
-        embed_onehot = F.one_hot(ind.long(), self.n_ind).to(x.dtype)
-        e_mean = torch.mean(embed_onehot, dim=[0,1])
-        e_mean = e_mean / (e_mean.sum(dim=1) + self.eps).unsqueeze(1)
-        perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + self.eps), dim=1))
-        
-        return (
-            torch.zeros(perplexity.shape, dtype=x.dtype, device=x.device),
-            feat.transpose(1,2) if self.transpose else feat,
-            perplexity,
-            None,
-            ind.transpose(1,2) if self.transpose else ind,
-        )
-        
-class DVAEDecoder(nn.Module):
-    def __init__(self, idim, odim,
-                 n_layer = 12, bn_dim = 64, hidden = 256, 
-                 kernel = 7, dilation = 2, up = False
-                ):
-        super().__init__()
-        self.up = up
-        self.conv_in = nn.Sequential(
-            nn.Conv1d(idim, bn_dim, 3, 1, 1), nn.GELU(),
-            nn.Conv1d(bn_dim, hidden, 3, 1, 1)
-        )
-        self.decoder_block = nn.ModuleList([
-            ConvNeXtBlock(hidden, hidden* 4, kernel, dilation,)
-            for _ in range(n_layer)])
-        self.conv_out = nn.Conv1d(hidden, odim, kernel_size=1, bias=False)
-
-    def forward(self, input, conditioning=None):
-        # B, T, C
-        x = input.transpose(1, 2)
-        x = self.conv_in(x)
-        for f in self.decoder_block:
-            x = f(x, conditioning)
-        
-        x = self.conv_out(x)
-        return x.transpose(1, 2)
-    
-
-class DVAE(nn.Module):
-    def __init__(
-        self, decoder_config, vq_config, dim=512
-    ):
-        super().__init__()
-        self.register_buffer('coef', torch.randn(1, 100, 1))
-
-        self.decoder = DVAEDecoder(**decoder_config)
-        self.out_conv = nn.Conv1d(dim, 100, 3, 1, 1, bias=False)
-        if vq_config is not None:
-            self.vq_layer = GFSQ(**vq_config)
-        else:
-            self.vq_layer = None
-
-    def forward(self, inp):
-
-        if self.vq_layer is not None:
-            vq_feats = self.vq_layer._embed(inp)
-        else:
-            vq_feats = inp.detach().clone()
-            
-        temp = torch.chunk(vq_feats, 2, dim=1) # flatten trick :)
-        temp = torch.stack(temp, -1)
-        vq_feats = temp.reshape(*temp.shape[:2], -1)
-
-        vq_feats = vq_feats.transpose(1, 2)
-        dec_out = self.decoder(input=vq_feats)
-        dec_out = self.out_conv(dec_out.transpose(1, 2))
-        mel = dec_out * self.coef
-
-        return mel