Remove files

Modules/__init__.py

@@ -1 +0,0 @@
-

Modules/diffusion/__init__.py

@@ -1 +0,0 @@
-

Modules/diffusion/diffusion.py

@@ -1,94 +0,0 @@
-from math import pi
-from random import randint
-from typing import Any, Optional, Sequence, Tuple, Union
-
-import torch
-from einops import rearrange
-from torch import Tensor, nn
-from tqdm import tqdm
-
-from .utils import *
-from .sampler import *
-
-"""
-Diffusion Classes (generic for 1d data)
-"""
-
-
-class Model1d(nn.Module):
-    def __init__(self, unet_type: str = "base", **kwargs):
-        super().__init__()
-        diffusion_kwargs, kwargs = groupby("diffusion_", kwargs)
-        self.unet = None
-        self.diffusion = None
-
-    def forward(self, x: Tensor, **kwargs) -> Tensor:
-        return self.diffusion(x, **kwargs)
-
-    def sample(self, *args, **kwargs) -> Tensor:
-        return self.diffusion.sample(*args, **kwargs)
-
-
-"""
-Audio Diffusion Classes (specific for 1d audio data)
-"""
-
-
-def get_default_model_kwargs():
-    return dict(
-        channels=128,
-        patch_size=16,
-        multipliers=[1, 2, 4, 4, 4, 4, 4],
-        factors=[4, 4, 4, 2, 2, 2],
-        num_blocks=[2, 2, 2, 2, 2, 2],
-        attentions=[0, 0, 0, 1, 1, 1, 1],
-        attention_heads=8,
-        attention_features=64,
-        attention_multiplier=2,
-        attention_use_rel_pos=False,
-        diffusion_type="v",
-        diffusion_sigma_distribution=UniformDistribution(),
-    )
-
-
-def get_default_sampling_kwargs():
-    return dict(sigma_schedule=LinearSchedule(), sampler=VSampler(), clamp=True)
-
-
-class AudioDiffusionModel(Model1d):
-    def __init__(self, **kwargs):
-        super().__init__(**{**get_default_model_kwargs(), **kwargs})
-
-    def sample(self, *args, **kwargs):
-        return super().sample(*args, **{**get_default_sampling_kwargs(), **kwargs})
-
-
-class AudioDiffusionConditional(Model1d):
-    def __init__(
-        self,
-        embedding_features: int,
-        embedding_max_length: int,
-        embedding_mask_proba: float = 0.1,
-        **kwargs,
-    ):
-        self.embedding_mask_proba = embedding_mask_proba
-        default_kwargs = dict(
-            **get_default_model_kwargs(),
-            unet_type="cfg",
-            context_embedding_features=embedding_features,
-            context_embedding_max_length=embedding_max_length,
-        )
-        super().__init__(**{**default_kwargs, **kwargs})
-
-    def forward(self, *args, **kwargs):
-        default_kwargs = dict(embedding_mask_proba=self.embedding_mask_proba)
-        return super().forward(*args, **{**default_kwargs, **kwargs})
-
-    def sample(self, *args, **kwargs):
-        default_kwargs = dict(
-            **get_default_sampling_kwargs(),
-            embedding_scale=5.0,
-        )
-        return super().sample(*args, **{**default_kwargs, **kwargs})
-
-

Modules/diffusion/modules.py

@@ -1,693 +0,0 @@
-from math import floor, log, pi
-from typing import Any, List, Optional, Sequence, Tuple, Union
-
-from .utils import *
-
-import torch
-import torch.nn as nn
-from einops import rearrange, reduce, repeat
-from einops.layers.torch import Rearrange
-from einops_exts import rearrange_many
-from torch import Tensor, einsum
-
-
-"""
-Utils
-"""
-
-class AdaLayerNorm(nn.Module):
-    def __init__(self, style_dim, channels, eps=1e-5):
-        super().__init__()
-        self.channels = channels
-        self.eps = eps
-
-        self.fc = nn.Linear(style_dim, channels*2)
-
-    def forward(self, x, s):
-        x = x.transpose(-1, -2)
-        x = x.transpose(1, -1)
-                
-        h = self.fc(s)
-        h = h.view(h.size(0), h.size(1), 1)
-        gamma, beta = torch.chunk(h, chunks=2, dim=1)
-        gamma, beta = gamma.transpose(1, -1), beta.transpose(1, -1)
-        
-        
-        x = F.layer_norm(x, (self.channels,), eps=self.eps)
-        x = (1 + gamma) * x + beta
-        return x.transpose(1, -1).transpose(-1, -2)
-
-class StyleTransformer1d(nn.Module):
-    def __init__(
-        self,
-        num_layers: int,
-        channels: int,
-        num_heads: int,
-        head_features: int,
-        multiplier: int,
-        use_context_time: bool = True,
-        use_rel_pos: bool = False,
-        context_features_multiplier: int = 1,
-        rel_pos_num_buckets: Optional[int] = None,
-        rel_pos_max_distance: Optional[int] = None,
-        context_features: Optional[int] = None,
-        context_embedding_features: Optional[int] = None,
-        embedding_max_length: int = 512,
-    ):
-        super().__init__()
-
-        self.blocks = nn.ModuleList(
-            [
-                StyleTransformerBlock(
-                    features=channels + context_embedding_features,
-                    head_features=head_features,
-                    num_heads=num_heads,
-                    multiplier=multiplier,
-                    style_dim=context_features,
-                    use_rel_pos=use_rel_pos,
-                    rel_pos_num_buckets=rel_pos_num_buckets,
-                    rel_pos_max_distance=rel_pos_max_distance,
-                )
-                for i in range(num_layers)
-            ]
-        )
-
-        self.to_out = nn.Sequential(
-            Rearrange("b t c -> b c t"),
-            nn.Conv1d(
-                in_channels=channels + context_embedding_features,
-                out_channels=channels,
-                kernel_size=1,
-            ),
-        )
-        
-        use_context_features = exists(context_features)
-        self.use_context_features = use_context_features
-        self.use_context_time = use_context_time
-
-        if use_context_time or use_context_features:
-            context_mapping_features = channels + context_embedding_features
-
-            self.to_mapping = nn.Sequential(
-                nn.Linear(context_mapping_features, context_mapping_features),
-                nn.GELU(),
-                nn.Linear(context_mapping_features, context_mapping_features),
-                nn.GELU(),
-            )
-        
-        if use_context_time:
-            assert exists(context_mapping_features)
-            self.to_time = nn.Sequential(
-                TimePositionalEmbedding(
-                    dim=channels, out_features=context_mapping_features
-                ),
-                nn.GELU(),
-            )
-
-        if use_context_features:
-            assert exists(context_features) and exists(context_mapping_features)
-            self.to_features = nn.Sequential(
-                nn.Linear(
-                    in_features=context_features, out_features=context_mapping_features
-                ),
-                nn.GELU(),
-            )
-            
-        self.fixed_embedding = FixedEmbedding(
-            max_length=embedding_max_length, features=context_embedding_features
-        )
-        
-
-    def get_mapping(
-        self, time: Optional[Tensor] = None, features: Optional[Tensor] = None
-    ) -> Optional[Tensor]:
-        """Combines context time features and features into mapping"""
-        items, mapping = [], None
-        # Compute time features
-        if self.use_context_time:
-            assert_message = "use_context_time=True but no time features provided"
-            assert exists(time), assert_message
-            items += [self.to_time(time)]
-        # Compute features
-        if self.use_context_features:
-            assert_message = "context_features exists but no features provided"
-            assert exists(features), assert_message
-            items += [self.to_features(features)]
-
-        # Compute joint mapping
-        if self.use_context_time or self.use_context_features:
-            mapping = reduce(torch.stack(items), "n b m -> b m", "sum")
-            mapping = self.to_mapping(mapping)
-
-        return mapping
-            
-    def run(self, x, time, embedding, features):
-        
-        mapping = self.get_mapping(time, features)
-        x = torch.cat([x.expand(-1, embedding.size(1), -1), embedding], axis=-1)
-        mapping = mapping.unsqueeze(1).expand(-1, embedding.size(1), -1)
-        
-        for block in self.blocks:
-            x = x + mapping
-            x = block(x, features)
-        
-        x = x.mean(axis=1).unsqueeze(1)
-        x = self.to_out(x)
-        x = x.transpose(-1, -2)
-        
-        return x
-        
-    def forward(self, x: Tensor, 
-                time: Tensor, 
-                embedding_mask_proba: float = 0.0,
-                embedding: Optional[Tensor] = None, 
-                features: Optional[Tensor] = None,
-               embedding_scale: float = 1.0) -> Tensor:
-        
-        b, device = embedding.shape[0], embedding.device
-        fixed_embedding = self.fixed_embedding(embedding)
-        if embedding_mask_proba > 0.0:
-            # Randomly mask embedding
-            batch_mask = rand_bool(
-                shape=(b, 1, 1), proba=embedding_mask_proba, device=device
-            )
-            embedding = torch.where(batch_mask, fixed_embedding, embedding)
-
-        if embedding_scale != 1.0:
-            # Compute both normal and fixed embedding outputs
-            out = self.run(x, time, embedding=embedding, features=features)
-            out_masked = self.run(x, time, embedding=fixed_embedding, features=features)
-            # Scale conditional output using classifier-free guidance
-            return out_masked + (out - out_masked) * embedding_scale
-        else:
-            return self.run(x, time, embedding=embedding, features=features)
-        
-        return x
-
-
-class StyleTransformerBlock(nn.Module):
-    def __init__(
-        self,
-        features: int,
-        num_heads: int,
-        head_features: int,
-        style_dim: int,
-        multiplier: int,
-        use_rel_pos: bool,
-        rel_pos_num_buckets: Optional[int] = None,
-        rel_pos_max_distance: Optional[int] = None,
-        context_features: Optional[int] = None,
-    ):
-        super().__init__()
-
-        self.use_cross_attention = exists(context_features) and context_features > 0
-
-        self.attention = StyleAttention(
-            features=features,
-            style_dim=style_dim,
-            num_heads=num_heads,
-            head_features=head_features,
-            use_rel_pos=use_rel_pos,
-            rel_pos_num_buckets=rel_pos_num_buckets,
-            rel_pos_max_distance=rel_pos_max_distance,
-        )
-
-        if self.use_cross_attention:
-            self.cross_attention = StyleAttention(
-                features=features,
-                style_dim=style_dim,
-                num_heads=num_heads,
-                head_features=head_features,
-                context_features=context_features,
-                use_rel_pos=use_rel_pos,
-                rel_pos_num_buckets=rel_pos_num_buckets,
-                rel_pos_max_distance=rel_pos_max_distance,
-            )
-
-        self.feed_forward = FeedForward(features=features, multiplier=multiplier)
-
-    def forward(self, x: Tensor, s: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
-        x = self.attention(x, s) + x
-        if self.use_cross_attention:
-            x = self.cross_attention(x, s, context=context) + x
-        x = self.feed_forward(x) + x
-        return x
-
-class StyleAttention(nn.Module):
-    def __init__(
-        self,
-        features: int,
-        *,
-        style_dim: int,
-        head_features: int,
-        num_heads: int,
-        context_features: Optional[int] = None,
-        use_rel_pos: bool,
-        rel_pos_num_buckets: Optional[int] = None,
-        rel_pos_max_distance: Optional[int] = None,
-    ):
-        super().__init__()
-        self.context_features = context_features
-        mid_features = head_features * num_heads
-        context_features = default(context_features, features)
-
-        self.norm = AdaLayerNorm(style_dim, features)
-        self.norm_context = AdaLayerNorm(style_dim, context_features)
-        self.to_q = nn.Linear(
-            in_features=features, out_features=mid_features, bias=False
-        )
-        self.to_kv = nn.Linear(
-            in_features=context_features, out_features=mid_features * 2, bias=False
-        )
-        self.attention = AttentionBase(
-            features,
-            num_heads=num_heads,
-            head_features=head_features,
-            use_rel_pos=use_rel_pos,
-            rel_pos_num_buckets=rel_pos_num_buckets,
-            rel_pos_max_distance=rel_pos_max_distance,
-        )
-
-    def forward(self, x: Tensor, s: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
-        assert_message = "You must provide a context when using context_features"
-        assert not self.context_features or exists(context), assert_message
-        # Use context if provided
-        context = default(context, x)
-        # Normalize then compute q from input and k,v from context
-        x, context = self.norm(x, s), self.norm_context(context, s)
-        
-        q, k, v = (self.to_q(x), *torch.chunk(self.to_kv(context), chunks=2, dim=-1))
-        # Compute and return attention
-        return self.attention(q, k, v)
-        
-class Transformer1d(nn.Module):
-    def __init__(
-        self,
-        num_layers: int,
-        channels: int,
-        num_heads: int,
-        head_features: int,
-        multiplier: int,
-        use_context_time: bool = True,
-        use_rel_pos: bool = False,
-        context_features_multiplier: int = 1,
-        rel_pos_num_buckets: Optional[int] = None,
-        rel_pos_max_distance: Optional[int] = None,
-        context_features: Optional[int] = None,
-        context_embedding_features: Optional[int] = None,
-        embedding_max_length: int = 512,
-    ):
-        super().__init__()
-
-        self.blocks = nn.ModuleList(
-            [
-                TransformerBlock(
-                    features=channels + context_embedding_features,
-                    head_features=head_features,
-                    num_heads=num_heads,
-                    multiplier=multiplier,
-                    use_rel_pos=use_rel_pos,
-                    rel_pos_num_buckets=rel_pos_num_buckets,
-                    rel_pos_max_distance=rel_pos_max_distance,
-                )
-                for i in range(num_layers)
-            ]
-        )
-
-        self.to_out = nn.Sequential(
-            Rearrange("b t c -> b c t"),
-            nn.Conv1d(
-                in_channels=channels + context_embedding_features,
-                out_channels=channels,
-                kernel_size=1,
-            ),
-        )
-        
-        use_context_features = exists(context_features)
-        self.use_context_features = use_context_features
-        self.use_context_time = use_context_time
-
-        if use_context_time or use_context_features:
-            context_mapping_features = channels + context_embedding_features
-
-            self.to_mapping = nn.Sequential(
-                nn.Linear(context_mapping_features, context_mapping_features),
-                nn.GELU(),
-                nn.Linear(context_mapping_features, context_mapping_features),
-                nn.GELU(),
-            )
-        
-        if use_context_time:
-            assert exists(context_mapping_features)
-            self.to_time = nn.Sequential(
-                TimePositionalEmbedding(
-                    dim=channels, out_features=context_mapping_features
-                ),
-                nn.GELU(),
-            )
-
-        if use_context_features:
-            assert exists(context_features) and exists(context_mapping_features)
-            self.to_features = nn.Sequential(
-                nn.Linear(
-                    in_features=context_features, out_features=context_mapping_features
-                ),
-                nn.GELU(),
-            )
-            
-        self.fixed_embedding = FixedEmbedding(
-            max_length=embedding_max_length, features=context_embedding_features
-        )
-        
-
-    def get_mapping(
-        self, time: Optional[Tensor] = None, features: Optional[Tensor] = None
-    ) -> Optional[Tensor]:
-        """Combines context time features and features into mapping"""
-        items, mapping = [], None
-        # Compute time features
-        if self.use_context_time:
-            assert_message = "use_context_time=True but no time features provided"
-            assert exists(time), assert_message
-            items += [self.to_time(time)]
-        # Compute features
-        if self.use_context_features:
-            assert_message = "context_features exists but no features provided"
-            assert exists(features), assert_message
-            items += [self.to_features(features)]
-
-        # Compute joint mapping
-        if self.use_context_time or self.use_context_features:
-            mapping = reduce(torch.stack(items), "n b m -> b m", "sum")
-            mapping = self.to_mapping(mapping)
-
-        return mapping
-            
-    def run(self, x, time, embedding, features):
-        
-        mapping = self.get_mapping(time, features)
-        x = torch.cat([x.expand(-1, embedding.size(1), -1), embedding], axis=-1)
-        mapping = mapping.unsqueeze(1).expand(-1, embedding.size(1), -1)
-        
-        for block in self.blocks:
-            x = x + mapping
-            x = block(x)
-        
-        x = x.mean(axis=1).unsqueeze(1)
-        x = self.to_out(x)
-        x = x.transpose(-1, -2)
-        
-        return x
-        
-    def forward(self, x: Tensor, 
-                time: Tensor, 
-                embedding_mask_proba: float = 0.0,
-                embedding: Optional[Tensor] = None, 
-                features: Optional[Tensor] = None,
-               embedding_scale: float = 1.0) -> Tensor:
-        
-        b, device = embedding.shape[0], embedding.device
-        fixed_embedding = self.fixed_embedding(embedding)
-        if embedding_mask_proba > 0.0:
-            # Randomly mask embedding
-            batch_mask = rand_bool(
-                shape=(b, 1, 1), proba=embedding_mask_proba, device=device
-            )
-            embedding = torch.where(batch_mask, fixed_embedding, embedding)
-
-        if embedding_scale != 1.0:
-            # Compute both normal and fixed embedding outputs
-            out = self.run(x, time, embedding=embedding, features=features)
-            out_masked = self.run(x, time, embedding=fixed_embedding, features=features)
-            # Scale conditional output using classifier-free guidance
-            return out_masked + (out - out_masked) * embedding_scale
-        else:
-            return self.run(x, time, embedding=embedding, features=features)
-        
-        return x
-
-
-"""
-Attention Components
-"""
-
-
-class RelativePositionBias(nn.Module):
-    def __init__(self, num_buckets: int, max_distance: int, num_heads: int):
-        super().__init__()
-        self.num_buckets = num_buckets
-        self.max_distance = max_distance
-        self.num_heads = num_heads
-        self.relative_attention_bias = nn.Embedding(num_buckets, num_heads)
-
-    @staticmethod
-    def _relative_position_bucket(
-        relative_position: Tensor, num_buckets: int, max_distance: int
-    ):
-        num_buckets //= 2
-        ret = (relative_position >= 0).to(torch.long) * num_buckets
-        n = torch.abs(relative_position)
-
-        max_exact = num_buckets // 2
-        is_small = n < max_exact
-
-        val_if_large = (
-            max_exact
-            + (
-                torch.log(n.float() / max_exact)
-                / log(max_distance / max_exact)
-                * (num_buckets - max_exact)
-            ).long()
-        )
-        val_if_large = torch.min(
-            val_if_large, torch.full_like(val_if_large, num_buckets - 1)
-        )
-
-        ret += torch.where(is_small, n, val_if_large)
-        return ret
-
-    def forward(self, num_queries: int, num_keys: int) -> Tensor:
-        i, j, device = num_queries, num_keys, self.relative_attention_bias.weight.device
-        q_pos = torch.arange(j - i, j, dtype=torch.long, device=device)
-        k_pos = torch.arange(j, dtype=torch.long, device=device)
-        rel_pos = rearrange(k_pos, "j -> 1 j") - rearrange(q_pos, "i -> i 1")
-
-        relative_position_bucket = self._relative_position_bucket(
-            rel_pos, num_buckets=self.num_buckets, max_distance=self.max_distance
-        )
-
-        bias = self.relative_attention_bias(relative_position_bucket)
-        bias = rearrange(bias, "m n h -> 1 h m n")
-        return bias
-
-
-def FeedForward(features: int, multiplier: int) -> nn.Module:
-    mid_features = features * multiplier
-    return nn.Sequential(
-        nn.Linear(in_features=features, out_features=mid_features),
-        nn.GELU(),
-        nn.Linear(in_features=mid_features, out_features=features),
-    )
-
-
-class AttentionBase(nn.Module):
-    def __init__(
-        self,
-        features: int,
-        *,
-        head_features: int,
-        num_heads: int,
-        use_rel_pos: bool,
-        out_features: Optional[int] = None,
-        rel_pos_num_buckets: Optional[int] = None,
-        rel_pos_max_distance: Optional[int] = None,
-    ):
-        super().__init__()
-        self.scale = head_features ** -0.5
-        self.num_heads = num_heads
-        self.use_rel_pos = use_rel_pos
-        mid_features = head_features * num_heads
-
-        if use_rel_pos:
-            assert exists(rel_pos_num_buckets) and exists(rel_pos_max_distance)
-            self.rel_pos = RelativePositionBias(
-                num_buckets=rel_pos_num_buckets,
-                max_distance=rel_pos_max_distance,
-                num_heads=num_heads,
-            )
-        if out_features is None:
-            out_features = features
-            
-        self.to_out = nn.Linear(in_features=mid_features, out_features=out_features)
-
-    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
-        # Split heads
-        q, k, v = rearrange_many((q, k, v), "b n (h d) -> b h n d", h=self.num_heads)
-        # Compute similarity matrix
-        sim = einsum("... n d, ... m d -> ... n m", q, k)
-        sim = (sim + self.rel_pos(*sim.shape[-2:])) if self.use_rel_pos else sim
-        sim = sim * self.scale
-        # Get attention matrix with softmax
-        attn = sim.softmax(dim=-1)
-        # Compute values
-        out = einsum("... n m, ... m d -> ... n d", attn, v)
-        out = rearrange(out, "b h n d -> b n (h d)")
-        return self.to_out(out)
-
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        features: int,
-        *,
-        head_features: int,
-        num_heads: int,
-        out_features: Optional[int] = None,
-        context_features: Optional[int] = None,
-        use_rel_pos: bool,
-        rel_pos_num_buckets: Optional[int] = None,
-        rel_pos_max_distance: Optional[int] = None,
-    ):
-        super().__init__()
-        self.context_features = context_features
-        mid_features = head_features * num_heads
-        context_features = default(context_features, features)
-
-        self.norm = nn.LayerNorm(features)
-        self.norm_context = nn.LayerNorm(context_features)
-        self.to_q = nn.Linear(
-            in_features=features, out_features=mid_features, bias=False
-        )
-        self.to_kv = nn.Linear(
-            in_features=context_features, out_features=mid_features * 2, bias=False
-        )
-
-        self.attention = AttentionBase(
-            features,
-            out_features=out_features,
-            num_heads=num_heads,
-            head_features=head_features,
-            use_rel_pos=use_rel_pos,
-            rel_pos_num_buckets=rel_pos_num_buckets,
-            rel_pos_max_distance=rel_pos_max_distance,
-        )
-
-    def forward(self, x: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
-        assert_message = "You must provide a context when using context_features"
-        assert not self.context_features or exists(context), assert_message
-        # Use context if provided
-        context = default(context, x)
-        # Normalize then compute q from input and k,v from context
-        x, context = self.norm(x), self.norm_context(context)
-        q, k, v = (self.to_q(x), *torch.chunk(self.to_kv(context), chunks=2, dim=-1))
-        # Compute and return attention
-        return self.attention(q, k, v)
-
-
-"""
-Transformer Blocks
-"""
-
-
-class TransformerBlock(nn.Module):
-    def __init__(
-        self,
-        features: int,
-        num_heads: int,
-        head_features: int,
-        multiplier: int,
-        use_rel_pos: bool,
-        rel_pos_num_buckets: Optional[int] = None,
-        rel_pos_max_distance: Optional[int] = None,
-        context_features: Optional[int] = None,
-    ):
-        super().__init__()
-
-        self.use_cross_attention = exists(context_features) and context_features > 0
-
-        self.attention = Attention(
-            features=features,
-            num_heads=num_heads,
-            head_features=head_features,
-            use_rel_pos=use_rel_pos,
-            rel_pos_num_buckets=rel_pos_num_buckets,
-            rel_pos_max_distance=rel_pos_max_distance,
-        )
-
-        if self.use_cross_attention:
-            self.cross_attention = Attention(
-                features=features,
-                num_heads=num_heads,
-                head_features=head_features,
-                context_features=context_features,
-                use_rel_pos=use_rel_pos,
-                rel_pos_num_buckets=rel_pos_num_buckets,
-                rel_pos_max_distance=rel_pos_max_distance,
-            )
-
-        self.feed_forward = FeedForward(features=features, multiplier=multiplier)
-
-    def forward(self, x: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
-        x = self.attention(x) + x
-        if self.use_cross_attention:
-            x = self.cross_attention(x, context=context) + x
-        x = self.feed_forward(x) + x
-        return x
-
-
-
-"""
-Time Embeddings
-"""
-
-
-class SinusoidalEmbedding(nn.Module):
-    def __init__(self, dim: int):
-        super().__init__()
-        self.dim = dim
-
-    def forward(self, x: Tensor) -> Tensor:
-        device, half_dim = x.device, self.dim // 2
-        emb = torch.tensor(log(10000) / (half_dim - 1), device=device)
-        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
-        emb = rearrange(x, "i -> i 1") * rearrange(emb, "j -> 1 j")
-        return torch.cat((emb.sin(), emb.cos()), dim=-1)
-
-
-class LearnedPositionalEmbedding(nn.Module):
-    """Used for continuous time"""
-
-    def __init__(self, dim: int):
-        super().__init__()
-        assert (dim % 2) == 0
-        half_dim = dim // 2
-        self.weights = nn.Parameter(torch.randn(half_dim))
-
-    def forward(self, x: Tensor) -> Tensor:
-        x = rearrange(x, "b -> b 1")
-        freqs = x * rearrange(self.weights, "d -> 1 d") * 2 * pi
-        fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1)
-        fouriered = torch.cat((x, fouriered), dim=-1)
-        return fouriered
-
-
-def TimePositionalEmbedding(dim: int, out_features: int) -> nn.Module:
-    return nn.Sequential(
-        LearnedPositionalEmbedding(dim),
-        nn.Linear(in_features=dim + 1, out_features=out_features),
-    )
-
-class FixedEmbedding(nn.Module):
-    def __init__(self, max_length: int, features: int):
-        super().__init__()
-        self.max_length = max_length
-        self.embedding = nn.Embedding(max_length, features)
-
-    def forward(self, x: Tensor) -> Tensor:
-        batch_size, length, device = *x.shape[0:2], x.device
-        assert_message = "Input sequence length must be <= max_length"
-        assert length <= self.max_length, assert_message
-        position = torch.arange(length, device=device)
-        fixed_embedding = self.embedding(position)
-        fixed_embedding = repeat(fixed_embedding, "n d -> b n d", b=batch_size)
-        return fixed_embedding

Modules/diffusion/sampler.py

@@ -1,691 +0,0 @@
-from math import atan, cos, pi, sin, sqrt
-from typing import Any, Callable, List, Optional, Tuple, Type
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange, reduce
-from torch import Tensor
-
-from .utils import *
-
-"""
-Diffusion Training
-"""
-
-""" Distributions """
-
-
-class Distribution:
-    def __call__(self, num_samples: int, device: torch.device):
-        raise NotImplementedError()
-
-
-class LogNormalDistribution(Distribution):
-    def __init__(self, mean: float, std: float):
-        self.mean = mean
-        self.std = std
-
-    def __call__(
-        self, num_samples: int, device: torch.device = torch.device("cpu")
-    ) -> Tensor:
-        normal = self.mean + self.std * torch.randn((num_samples,), device=device)
-        return normal.exp()
-
-
-class UniformDistribution(Distribution):
-    def __call__(self, num_samples: int, device: torch.device = torch.device("cpu")):
-        return torch.rand(num_samples, device=device)
-
-
-class VKDistribution(Distribution):
-    def __init__(
-        self,
-        min_value: float = 0.0,
-        max_value: float = float("inf"),
-        sigma_data: float = 1.0,
-    ):
-        self.min_value = min_value
-        self.max_value = max_value
-        self.sigma_data = sigma_data
-
-    def __call__(
-        self, num_samples: int, device: torch.device = torch.device("cpu")
-    ) -> Tensor:
-        sigma_data = self.sigma_data
-        min_cdf = atan(self.min_value / sigma_data) * 2 / pi
-        max_cdf = atan(self.max_value / sigma_data) * 2 / pi
-        u = (max_cdf - min_cdf) * torch.randn((num_samples,), device=device) + min_cdf
-        return torch.tan(u * pi / 2) * sigma_data
-
-
-""" Diffusion Classes """
-
-
-def pad_dims(x: Tensor, ndim: int) -> Tensor:
-    # Pads additional ndims to the right of the tensor
-    return x.view(*x.shape, *((1,) * ndim))
-
-
-def clip(x: Tensor, dynamic_threshold: float = 0.0):
-    if dynamic_threshold == 0.0:
-        return x.clamp(-1.0, 1.0)
-    else:
-        # Dynamic thresholding
-        # Find dynamic threshold quantile for each batch
-        x_flat = rearrange(x, "b ... -> b (...)")
-        scale = torch.quantile(x_flat.abs(), dynamic_threshold, dim=-1)
-        # Clamp to a min of 1.0
-        scale.clamp_(min=1.0)
-        # Clamp all values and scale
-        scale = pad_dims(scale, ndim=x.ndim - scale.ndim)
-        x = x.clamp(-scale, scale) / scale
-        return x
-
-
-def to_batch(
-    batch_size: int,
-    device: torch.device,
-    x: Optional[float] = None,
-    xs: Optional[Tensor] = None,
-) -> Tensor:
-    assert exists(x) ^ exists(xs), "Either x or xs must be provided"
-    # If x provided use the same for all batch items
-    if exists(x):
-        xs = torch.full(size=(batch_size,), fill_value=x).to(device)
-    assert exists(xs)
-    return xs
-
-
-class Diffusion(nn.Module):
-
-    alias: str = ""
-
-    """Base diffusion class"""
-
-    def denoise_fn(
-        self,
-        x_noisy: Tensor,
-        sigmas: Optional[Tensor] = None,
-        sigma: Optional[float] = None,
-        **kwargs,
-    ) -> Tensor:
-        raise NotImplementedError("Diffusion class missing denoise_fn")
-
-    def forward(self, x: Tensor, noise: Tensor = None, **kwargs) -> Tensor:
-        raise NotImplementedError("Diffusion class missing forward function")
-
-
-class VDiffusion(Diffusion):
-
-    alias = "v"
-
-    def __init__(self, net: nn.Module, *, sigma_distribution: Distribution):
-        super().__init__()
-        self.net = net
-        self.sigma_distribution = sigma_distribution
-
-    def get_alpha_beta(self, sigmas: Tensor) -> Tuple[Tensor, Tensor]:
-        angle = sigmas * pi / 2
-        alpha = torch.cos(angle)
-        beta = torch.sin(angle)
-        return alpha, beta
-
-    def denoise_fn(
-        self,
-        x_noisy: Tensor,
-        sigmas: Optional[Tensor] = None,
-        sigma: Optional[float] = None,
-        **kwargs,
-    ) -> Tensor:
-        batch_size, device = x_noisy.shape[0], x_noisy.device
-        sigmas = to_batch(x=sigma, xs=sigmas, batch_size=batch_size, device=device)
-        return self.net(x_noisy, sigmas, **kwargs)
-
-    def forward(self, x: Tensor, noise: Tensor = None, **kwargs) -> Tensor:
-        batch_size, device = x.shape[0], x.device
-
-        # Sample amount of noise to add for each batch element
-        sigmas = self.sigma_distribution(num_samples=batch_size, device=device)
-        sigmas_padded = rearrange(sigmas, "b -> b 1 1")
-
-        # Get noise
-        noise = default(noise, lambda: torch.randn_like(x))
-
-        # Combine input and noise weighted by half-circle
-        alpha, beta = self.get_alpha_beta(sigmas_padded)
-        x_noisy = x * alpha + noise * beta
-        x_target = noise * alpha - x * beta
-
-        # Denoise and return loss
-        x_denoised = self.denoise_fn(x_noisy, sigmas, **kwargs)
-        return F.mse_loss(x_denoised, x_target)
-
-
-class KDiffusion(Diffusion):
-    """Elucidated Diffusion (Karras et al. 2022): https://arxiv.org/abs/2206.00364"""
-
-    alias = "k"
-
-    def __init__(
-        self,
-        net: nn.Module,
-        *,
-        sigma_distribution: Distribution,
-        sigma_data: float,  # data distribution standard deviation
-        dynamic_threshold: float = 0.0,
-    ):
-        super().__init__()
-        self.net = net
-        self.sigma_data = sigma_data
-        self.sigma_distribution = sigma_distribution
-        self.dynamic_threshold = dynamic_threshold
-
-    def get_scale_weights(self, sigmas: Tensor) -> Tuple[Tensor, ...]:
-        sigma_data = self.sigma_data
-        c_noise = torch.log(sigmas) * 0.25
-        sigmas = rearrange(sigmas, "b -> b 1 1")
-        c_skip = (sigma_data ** 2) / (sigmas ** 2 + sigma_data ** 2)
-        c_out = sigmas * sigma_data * (sigma_data ** 2 + sigmas ** 2) ** -0.5
-        c_in = (sigmas ** 2 + sigma_data ** 2) ** -0.5
-        return c_skip, c_out, c_in, c_noise
-
-    def denoise_fn(
-        self,
-        x_noisy: Tensor,
-        sigmas: Optional[Tensor] = None,
-        sigma: Optional[float] = None,
-        **kwargs,
-    ) -> Tensor:
-        batch_size, device = x_noisy.shape[0], x_noisy.device
-        sigmas = to_batch(x=sigma, xs=sigmas, batch_size=batch_size, device=device)
-
-        # Predict network output and add skip connection
-        c_skip, c_out, c_in, c_noise = self.get_scale_weights(sigmas)
-        x_pred = self.net(c_in * x_noisy, c_noise, **kwargs)
-        x_denoised = c_skip * x_noisy + c_out * x_pred
-
-        return x_denoised
-
-    def loss_weight(self, sigmas: Tensor) -> Tensor:
-        # Computes weight depending on data distribution
-        return (sigmas ** 2 + self.sigma_data ** 2) * (sigmas * self.sigma_data) ** -2
-
-    def forward(self, x: Tensor, noise: Tensor = None, **kwargs) -> Tensor:
-        batch_size, device = x.shape[0], x.device
-        from einops import rearrange, reduce
-
-        # Sample amount of noise to add for each batch element
-        sigmas = self.sigma_distribution(num_samples=batch_size, device=device)
-        sigmas_padded = rearrange(sigmas, "b -> b 1 1")
-
-        # Add noise to input
-        noise = default(noise, lambda: torch.randn_like(x))
-        x_noisy = x + sigmas_padded * noise
-        
-        # Compute denoised values
-        x_denoised = self.denoise_fn(x_noisy, sigmas=sigmas, **kwargs)
-
-        # Compute weighted loss
-        losses = F.mse_loss(x_denoised, x, reduction="none")
-        losses = reduce(losses, "b ... -> b", "mean")
-        losses = losses * self.loss_weight(sigmas)
-        loss = losses.mean()
-        return loss
-
-
-class VKDiffusion(Diffusion):
-
-    alias = "vk"
-
-    def __init__(self, net: nn.Module, *, sigma_distribution: Distribution):
-        super().__init__()
-        self.net = net
-        self.sigma_distribution = sigma_distribution
-
-    def get_scale_weights(self, sigmas: Tensor) -> Tuple[Tensor, ...]:
-        sigma_data = 1.0
-        sigmas = rearrange(sigmas, "b -> b 1 1")
-        c_skip = (sigma_data ** 2) / (sigmas ** 2 + sigma_data ** 2)
-        c_out = -sigmas * sigma_data * (sigma_data ** 2 + sigmas ** 2) ** -0.5
-        c_in = (sigmas ** 2 + sigma_data ** 2) ** -0.5
-        return c_skip, c_out, c_in
-
-    def sigma_to_t(self, sigmas: Tensor) -> Tensor:
-        return sigmas.atan() / pi * 2
-
-    def t_to_sigma(self, t: Tensor) -> Tensor:
-        return (t * pi / 2).tan()
-
-    def denoise_fn(
-        self,
-        x_noisy: Tensor,
-        sigmas: Optional[Tensor] = None,
-        sigma: Optional[float] = None,
-        **kwargs,
-    ) -> Tensor:
-        batch_size, device = x_noisy.shape[0], x_noisy.device
-        sigmas = to_batch(x=sigma, xs=sigmas, batch_size=batch_size, device=device)
-
-        # Predict network output and add skip connection
-        c_skip, c_out, c_in = self.get_scale_weights(sigmas)
-        x_pred = self.net(c_in * x_noisy, self.sigma_to_t(sigmas), **kwargs)
-        x_denoised = c_skip * x_noisy + c_out * x_pred
-        return x_denoised
-
-    def forward(self, x: Tensor, noise: Tensor = None, **kwargs) -> Tensor:
-        batch_size, device = x.shape[0], x.device
-
-        # Sample amount of noise to add for each batch element
-        sigmas = self.sigma_distribution(num_samples=batch_size, device=device)
-        sigmas_padded = rearrange(sigmas, "b -> b 1 1")
-
-        # Add noise to input
-        noise = default(noise, lambda: torch.randn_like(x))
-        x_noisy = x + sigmas_padded * noise
-
-        # Compute model output
-        c_skip, c_out, c_in = self.get_scale_weights(sigmas)
-        x_pred = self.net(c_in * x_noisy, self.sigma_to_t(sigmas), **kwargs)
-
-        # Compute v-objective target
-        v_target = (x - c_skip * x_noisy) / (c_out + 1e-7)
-
-        # Compute loss
-        loss = F.mse_loss(x_pred, v_target)
-        return loss
-
-
-"""
-Diffusion Sampling
-"""
-
-""" Schedules """
-
-
-class Schedule(nn.Module):
-    """Interface used by different sampling schedules"""
-
-    def forward(self, num_steps: int, device: torch.device) -> Tensor:
-        raise NotImplementedError()
-
-
-class LinearSchedule(Schedule):
-    def forward(self, num_steps: int, device: Any) -> Tensor:
-        sigmas = torch.linspace(1, 0, num_steps + 1)[:-1]
-        return sigmas
-
-
-class KarrasSchedule(Schedule):
-    """https://arxiv.org/abs/2206.00364 equation 5"""
-
-    def __init__(self, sigma_min: float, sigma_max: float, rho: float = 7.0):
-        super().__init__()
-        self.sigma_min = sigma_min
-        self.sigma_max = sigma_max
-        self.rho = rho
-
-    def forward(self, num_steps: int, device: Any) -> Tensor:
-        rho_inv = 1.0 / self.rho
-        steps = torch.arange(num_steps, device=device, dtype=torch.float32)
-        sigmas = (
-            self.sigma_max ** rho_inv
-            + (steps / (num_steps - 1))
-            * (self.sigma_min ** rho_inv - self.sigma_max ** rho_inv)
-        ) ** self.rho
-        sigmas = F.pad(sigmas, pad=(0, 1), value=0.0)
-        return sigmas
-
-
-""" Samplers """
-
-
-class Sampler(nn.Module):
-
-    diffusion_types: List[Type[Diffusion]] = []
-
-    def forward(
-        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int
-    ) -> Tensor:
-        raise NotImplementedError()
-
-    def inpaint(
-        self,
-        source: Tensor,
-        mask: Tensor,
-        fn: Callable,
-        sigmas: Tensor,
-        num_steps: int,
-        num_resamples: int,
-    ) -> Tensor:
-        raise NotImplementedError("Inpainting not available with current sampler")
-
-
-class VSampler(Sampler):
-
-    diffusion_types = [VDiffusion]
-
-    def get_alpha_beta(self, sigma: float) -> Tuple[float, float]:
-        angle = sigma * pi / 2
-        alpha = cos(angle)
-        beta = sin(angle)
-        return alpha, beta
-
-    def forward(
-        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int
-    ) -> Tensor:
-        x = sigmas[0] * noise
-        alpha, beta = self.get_alpha_beta(sigmas[0].item())
-
-        for i in range(num_steps - 1):
-            is_last = i == num_steps - 1
-
-            x_denoised = fn(x, sigma=sigmas[i])
-            x_pred = x * alpha - x_denoised * beta
-            x_eps = x * beta + x_denoised * alpha
-
-            if not is_last:
-                alpha, beta = self.get_alpha_beta(sigmas[i + 1].item())
-                x = x_pred * alpha + x_eps * beta
-
-        return x_pred
-
-
-class KarrasSampler(Sampler):
-    """https://arxiv.org/abs/2206.00364 algorithm 1"""
-
-    diffusion_types = [KDiffusion, VKDiffusion]
-
-    def __init__(
-        self,
-        s_tmin: float = 0,
-        s_tmax: float = float("inf"),
-        s_churn: float = 0.0,
-        s_noise: float = 1.0,
-    ):
-        super().__init__()
-        self.s_tmin = s_tmin
-        self.s_tmax = s_tmax
-        self.s_noise = s_noise
-        self.s_churn = s_churn
-
-    def step(
-        self, x: Tensor, fn: Callable, sigma: float, sigma_next: float, gamma: float
-    ) -> Tensor:
-        """Algorithm 2 (step)"""
-        # Select temporarily increased noise level
-        sigma_hat = sigma + gamma * sigma
-        # Add noise to move from sigma to sigma_hat
-        epsilon = self.s_noise * torch.randn_like(x)
-        x_hat = x + sqrt(sigma_hat ** 2 - sigma ** 2) * epsilon
-        # Evaluate ∂x/∂sigma at sigma_hat
-        d = (x_hat - fn(x_hat, sigma=sigma_hat)) / sigma_hat
-        # Take euler step from sigma_hat to sigma_next
-        x_next = x_hat + (sigma_next - sigma_hat) * d
-        # Second order correction
-        if sigma_next != 0:
-            model_out_next = fn(x_next, sigma=sigma_next)
-            d_prime = (x_next - model_out_next) / sigma_next
-            x_next = x_hat + 0.5 * (sigma - sigma_hat) * (d + d_prime)
-        return x_next
-
-    def forward(
-        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int
-    ) -> Tensor:
-        x = sigmas[0] * noise
-        # Compute gammas
-        gammas = torch.where(
-            (sigmas >= self.s_tmin) & (sigmas <= self.s_tmax),
-            min(self.s_churn / num_steps, sqrt(2) - 1),
-            0.0,
-        )
-        # Denoise to sample
-        for i in range(num_steps - 1):
-            x = self.step(
-                x, fn=fn, sigma=sigmas[i], sigma_next=sigmas[i + 1], gamma=gammas[i]  # type: ignore # noqa
-            )
-
-        return x
-
-
-class AEulerSampler(Sampler):
-
-    diffusion_types = [KDiffusion, VKDiffusion]
-
-    def get_sigmas(self, sigma: float, sigma_next: float) -> Tuple[float, float]:
-        sigma_up = sqrt(sigma_next ** 2 * (sigma ** 2 - sigma_next ** 2) / sigma ** 2)
-        sigma_down = sqrt(sigma_next ** 2 - sigma_up ** 2)
-        return sigma_up, sigma_down
-
-    def step(self, x: Tensor, fn: Callable, sigma: float, sigma_next: float) -> Tensor:
-        # Sigma steps
-        sigma_up, sigma_down = self.get_sigmas(sigma, sigma_next)
-        # Derivative at sigma (∂x/∂sigma)
-        d = (x - fn(x, sigma=sigma)) / sigma
-        # Euler method
-        x_next = x + d * (sigma_down - sigma)
-        # Add randomness
-        x_next = x_next + torch.randn_like(x) * sigma_up
-        return x_next
-
-    def forward(
-        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int
-    ) -> Tensor:
-        x = sigmas[0] * noise
-        # Denoise to sample
-        for i in range(num_steps - 1):
-            x = self.step(x, fn=fn, sigma=sigmas[i], sigma_next=sigmas[i + 1])  # type: ignore # noqa
-        return x
-
-
-class ADPM2Sampler(Sampler):
-    """https://www.desmos.com/calculator/jbxjlqd9mb"""
-
-    diffusion_types = [KDiffusion, VKDiffusion]
-
-    def __init__(self, rho: float = 1.0):
-        super().__init__()
-        self.rho = rho
-
-    def get_sigmas(self, sigma: float, sigma_next: float) -> Tuple[float, float, float]:
-        r = self.rho
-        sigma_up = sqrt(sigma_next ** 2 * (sigma ** 2 - sigma_next ** 2) / sigma ** 2)
-        sigma_down = sqrt(sigma_next ** 2 - sigma_up ** 2)
-        sigma_mid = ((sigma ** (1 / r) + sigma_down ** (1 / r)) / 2) ** r
-        return sigma_up, sigma_down, sigma_mid
-
-    def step(self, x: Tensor, fn: Callable, sigma: float, sigma_next: float) -> Tensor:
-        # Sigma steps
-        sigma_up, sigma_down, sigma_mid = self.get_sigmas(sigma, sigma_next)
-        # Derivative at sigma (∂x/∂sigma)
-        d = (x - fn(x, sigma=sigma)) / sigma
-        # Denoise to midpoint
-        x_mid = x + d * (sigma_mid - sigma)
-        # Derivative at sigma_mid (∂x_mid/∂sigma_mid)
-        d_mid = (x_mid - fn(x_mid, sigma=sigma_mid)) / sigma_mid
-        # Denoise to next
-        x = x + d_mid * (sigma_down - sigma)
-        # Add randomness
-        x_next = x + torch.randn_like(x) * sigma_up
-        return x_next
-
-    def forward(
-        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int
-    ) -> Tensor:
-        x = sigmas[0] * noise
-        # Denoise to sample
-        for i in range(num_steps - 1):
-            x = self.step(x, fn=fn, sigma=sigmas[i], sigma_next=sigmas[i + 1])  # type: ignore # noqa
-        return x
-
-    def inpaint(
-        self,
-        source: Tensor,
-        mask: Tensor,
-        fn: Callable,
-        sigmas: Tensor,
-        num_steps: int,
-        num_resamples: int,
-    ) -> Tensor:
-        x = sigmas[0] * torch.randn_like(source)
-
-        for i in range(num_steps - 1):
-            # Noise source to current noise level
-            source_noisy = source + sigmas[i] * torch.randn_like(source)
-            for r in range(num_resamples):
-                # Merge noisy source and current then denoise
-                x = source_noisy * mask + x * ~mask
-                x = self.step(x, fn=fn, sigma=sigmas[i], sigma_next=sigmas[i + 1])  # type: ignore # noqa
-                # Renoise if not last resample step
-                if r < num_resamples - 1:
-                    sigma = sqrt(sigmas[i] ** 2 - sigmas[i + 1] ** 2)
-                    x = x + sigma * torch.randn_like(x)
-
-        return source * mask + x * ~mask
-
-
-""" Main Classes """
-
-
-class DiffusionSampler(nn.Module):
-    def __init__(
-        self,
-        diffusion: Diffusion,
-        *,
-        sampler: Sampler,
-        sigma_schedule: Schedule,
-        num_steps: Optional[int] = None,
-        clamp: bool = True,
-    ):
-        super().__init__()
-        self.denoise_fn = diffusion.denoise_fn
-        self.sampler = sampler
-        self.sigma_schedule = sigma_schedule
-        self.num_steps = num_steps
-        self.clamp = clamp
-
-        # Check sampler is compatible with diffusion type
-        sampler_class = sampler.__class__.__name__
-        diffusion_class = diffusion.__class__.__name__
-        message = f"{sampler_class} incompatible with {diffusion_class}"
-        assert diffusion.alias in [t.alias for t in sampler.diffusion_types], message
-
-    def forward(
-        self, noise: Tensor, num_steps: Optional[int] = None, **kwargs
-    ) -> Tensor:
-        device = noise.device
-        num_steps = default(num_steps, self.num_steps)  # type: ignore
-        assert exists(num_steps), "Parameter `num_steps` must be provided"
-        # Compute sigmas using schedule
-        sigmas = self.sigma_schedule(num_steps, device)
-        # Append additional kwargs to denoise function (used e.g. for conditional unet)
-        fn = lambda *a, **ka: self.denoise_fn(*a, **{**ka, **kwargs})  # noqa
-        # Sample using sampler
-        x = self.sampler(noise, fn=fn, sigmas=sigmas, num_steps=num_steps)
-        x = x.clamp(-1.0, 1.0) if self.clamp else x
-        return x
-
-
-class DiffusionInpainter(nn.Module):
-    def __init__(
-        self,
-        diffusion: Diffusion,
-        *,
-        num_steps: int,
-        num_resamples: int,
-        sampler: Sampler,
-        sigma_schedule: Schedule,
-    ):
-        super().__init__()
-        self.denoise_fn = diffusion.denoise_fn
-        self.num_steps = num_steps
-        self.num_resamples = num_resamples
-        self.inpaint_fn = sampler.inpaint
-        self.sigma_schedule = sigma_schedule
-
-    @torch.no_grad()
-    def forward(self, inpaint: Tensor, inpaint_mask: Tensor) -> Tensor:
-        x = self.inpaint_fn(
-            source=inpaint,
-            mask=inpaint_mask,
-            fn=self.denoise_fn,
-            sigmas=self.sigma_schedule(self.num_steps, inpaint.device),
-            num_steps=self.num_steps,
-            num_resamples=self.num_resamples,
-        )
-        return x
-
-
-def sequential_mask(like: Tensor, start: int) -> Tensor:
-    length, device = like.shape[2], like.device
-    mask = torch.ones_like(like, dtype=torch.bool)
-    mask[:, :, start:] = torch.zeros((length - start,), device=device)
-    return mask
-
-
-class SpanBySpanComposer(nn.Module):
-    def __init__(
-        self,
-        inpainter: DiffusionInpainter,
-        *,
-        num_spans: int,
-    ):
-        super().__init__()
-        self.inpainter = inpainter
-        self.num_spans = num_spans
-
-    def forward(self, start: Tensor, keep_start: bool = False) -> Tensor:
-        half_length = start.shape[2] // 2
-
-        spans = list(start.chunk(chunks=2, dim=-1)) if keep_start else []
-        # Inpaint second half from first half
-        inpaint = torch.zeros_like(start)
-        inpaint[:, :, :half_length] = start[:, :, half_length:]
-        inpaint_mask = sequential_mask(like=start, start=half_length)
-
-        for i in range(self.num_spans):
-            # Inpaint second half
-            span = self.inpainter(inpaint=inpaint, inpaint_mask=inpaint_mask)
-            # Replace first half with generated second half
-            second_half = span[:, :, half_length:]
-            inpaint[:, :, :half_length] = second_half
-            # Save generated span
-            spans.append(second_half)
-
-        return torch.cat(spans, dim=2)
-
-
-class XDiffusion(nn.Module):
-    def __init__(self, type: str, net: nn.Module, **kwargs):
-        super().__init__()
-
-        diffusion_classes = [VDiffusion, KDiffusion, VKDiffusion]
-        aliases = [t.alias for t in diffusion_classes]  # type: ignore
-        message = f"type='{type}' must be one of {*aliases,}"
-        assert type in aliases, message
-        self.net = net
-
-        for XDiffusion in diffusion_classes:
-            if XDiffusion.alias == type:  # type: ignore
-                self.diffusion = XDiffusion(net=net, **kwargs)
-
-    def forward(self, *args, **kwargs) -> Tensor:
-        return self.diffusion(*args, **kwargs)
-
-    def sample(
-        self,
-        noise: Tensor,
-        num_steps: int,
-        sigma_schedule: Schedule,
-        sampler: Sampler,
-        clamp: bool,
-        **kwargs,
-    ) -> Tensor:
-        diffusion_sampler = DiffusionSampler(
-            diffusion=self.diffusion,
-            sampler=sampler,
-            sigma_schedule=sigma_schedule,
-            num_steps=num_steps,
-            clamp=clamp,
-        )
-        return diffusion_sampler(noise, **kwargs)

Modules/diffusion/utils.py

@@ -1,82 +0,0 @@
-from functools import reduce
-from inspect import isfunction
-from math import ceil, floor, log2, pi
-from typing import Callable, Dict, List, Optional, Sequence, Tuple, TypeVar, Union
-
-import torch
-import torch.nn.functional as F
-from einops import rearrange
-from torch import Generator, Tensor
-from typing_extensions import TypeGuard
-
-T = TypeVar("T")
-
-
-def exists(val: Optional[T]) -> TypeGuard[T]:
-    return val is not None
-
-
-def iff(condition: bool, value: T) -> Optional[T]:
-    return value if condition else None
-
-
-def is_sequence(obj: T) -> TypeGuard[Union[list, tuple]]:
-    return isinstance(obj, list) or isinstance(obj, tuple)
-
-
-def default(val: Optional[T], d: Union[Callable[..., T], T]) -> T:
-    if exists(val):
-        return val
-    return d() if isfunction(d) else d
-
-
-def to_list(val: Union[T, Sequence[T]]) -> List[T]:
-    if isinstance(val, tuple):
-        return list(val)
-    if isinstance(val, list):
-        return val
-    return [val]  # type: ignore
-
-
-def prod(vals: Sequence[int]) -> int:
-    return reduce(lambda x, y: x * y, vals)
-
-
-def closest_power_2(x: float) -> int:
-    exponent = log2(x)
-    distance_fn = lambda z: abs(x - 2 ** z)  # noqa
-    exponent_closest = min((floor(exponent), ceil(exponent)), key=distance_fn)
-    return 2 ** int(exponent_closest)
-
-def rand_bool(shape, proba, device = None):
-    if proba == 1:
-        return torch.ones(shape, device=device, dtype=torch.bool)
-    elif proba == 0:
-        return torch.zeros(shape, device=device, dtype=torch.bool)
-    else:
-        return torch.bernoulli(torch.full(shape, proba, device=device)).to(torch.bool)
-
-
-"""
-Kwargs Utils
-"""
-
-
-def group_dict_by_prefix(prefix: str, d: Dict) -> Tuple[Dict, Dict]:
-    return_dicts: Tuple[Dict, Dict] = ({}, {})
-    for key in d.keys():
-        no_prefix = int(not key.startswith(prefix))
-        return_dicts[no_prefix][key] = d[key]
-    return return_dicts
-
-
-def groupby(prefix: str, d: Dict, keep_prefix: bool = False) -> Tuple[Dict, Dict]:
-    kwargs_with_prefix, kwargs = group_dict_by_prefix(prefix, d)
-    if keep_prefix:
-        return kwargs_with_prefix, kwargs
-    kwargs_no_prefix = {k[len(prefix) :]: v for k, v in kwargs_with_prefix.items()}
-    return kwargs_no_prefix, kwargs
-
-
-def prefix_dict(prefix: str, d: Dict) -> Dict:
-    return {prefix + str(k): v for k, v in d.items()}

Modules/discriminators.py

@@ -1,188 +0,0 @@
-import torch
-import torch.nn.functional as F
-import torch.nn as nn
-from torch.nn import Conv1d, AvgPool1d, Conv2d
-from torch.nn.utils import weight_norm, spectral_norm
-
-from .utils import get_padding
-
-LRELU_SLOPE = 0.1
-
-def stft(x, fft_size, hop_size, win_length, window):
-    """Perform STFT and convert to magnitude spectrogram.
-    Args:
-        x (Tensor): Input signal tensor (B, T).
-        fft_size (int): FFT size.
-        hop_size (int): Hop size.
-        win_length (int): Window length.
-        window (str): Window function type.
-    Returns:
-        Tensor: Magnitude spectrogram (B, #frames, fft_size // 2 + 1).
-    """
-    x_stft = torch.stft(x, fft_size, hop_size, win_length, window,
-            return_complex=True)
-    real = x_stft[..., 0]
-    imag = x_stft[..., 1]
-
-    return torch.abs(x_stft).transpose(2, 1)
-
-class SpecDiscriminator(nn.Module):
-    """docstring for Discriminator."""
-
-    def __init__(self, fft_size=1024, shift_size=120, win_length=600, window="hann_window", use_spectral_norm=False):
-        super(SpecDiscriminator, self).__init__()
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
-        self.fft_size = fft_size
-        self.shift_size = shift_size
-        self.win_length = win_length
-        self.window = getattr(torch, window)(win_length)
-        self.discriminators = nn.ModuleList([
-            norm_f(nn.Conv2d(1, 32, kernel_size=(3, 9), padding=(1, 4))),
-            norm_f(nn.Conv2d(32, 32, kernel_size=(3, 9), stride=(1,2), padding=(1, 4))),
-            norm_f(nn.Conv2d(32, 32, kernel_size=(3, 9), stride=(1,2), padding=(1, 4))),
-            norm_f(nn.Conv2d(32, 32, kernel_size=(3, 9), stride=(1,2), padding=(1, 4))),
-            norm_f(nn.Conv2d(32, 32, kernel_size=(3, 3), stride=(1,1), padding=(1, 1))),
-        ])
-
-        self.out = norm_f(nn.Conv2d(32, 1, 3, 1, 1))
-
-    def forward(self, y):
-
-        fmap = []
-        y = y.squeeze(1)
-        y = stft(y, self.fft_size, self.shift_size, self.win_length, self.window.to(y.get_device()))
-        y = y.unsqueeze(1)
-        for i, d in enumerate(self.discriminators):
-            y = d(y)
-            y = F.leaky_relu(y, LRELU_SLOPE)
-            fmap.append(y)
-
-        y = self.out(y)
-        fmap.append(y)
-
-        return torch.flatten(y, 1, -1), fmap
-
-class MultiResSpecDiscriminator(torch.nn.Module):
-
-    def __init__(self,
-                 fft_sizes=[1024, 2048, 512],
-                 hop_sizes=[120, 240, 50],
-                 win_lengths=[600, 1200, 240],
-                 window="hann_window"):
-
-        super(MultiResSpecDiscriminator, self).__init__()
-        self.discriminators = nn.ModuleList([
-            SpecDiscriminator(fft_sizes[0], hop_sizes[0], win_lengths[0], window),
-            SpecDiscriminator(fft_sizes[1], hop_sizes[1], win_lengths[1], window),
-            SpecDiscriminator(fft_sizes[2], hop_sizes[2], win_lengths[2], window)
-            ])
-
-    def forward(self, y, y_hat):
-        y_d_rs = []
-        y_d_gs = []
-        fmap_rs = []
-        fmap_gs = []
-        for i, d in enumerate(self.discriminators):
-            y_d_r, fmap_r = d(y)
-            y_d_g, fmap_g = d(y_hat)
-            y_d_rs.append(y_d_r)
-            fmap_rs.append(fmap_r)
-            y_d_gs.append(y_d_g)
-            fmap_gs.append(fmap_g)
-
-        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
-
-
-class DiscriminatorP(torch.nn.Module):
-    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
-        super(DiscriminatorP, self).__init__()
-        self.period = period
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
-        self.convs = nn.ModuleList([
-            norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
-            norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
-            norm_f(Conv2d(128, 512, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
-            norm_f(Conv2d(512, 1024, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
-            norm_f(Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))),
-        ])
-        self.conv_post = norm_f(Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
-
-    def forward(self, x):
-        fmap = []
-
-        # 1d to 2d
-        b, c, t = x.shape
-        if t % self.period != 0: # pad first
-            n_pad = self.period - (t % self.period)
-            x = F.pad(x, (0, n_pad), "reflect")
-            t = t + n_pad
-        x = x.view(b, c, t // self.period, self.period)
-
-        for l in self.convs:
-            x = l(x)
-            x = F.leaky_relu(x, LRELU_SLOPE)
-            fmap.append(x)
-        x = self.conv_post(x)
-        fmap.append(x)
-        x = torch.flatten(x, 1, -1)
-
-        return x, fmap
-
-
-class MultiPeriodDiscriminator(torch.nn.Module):
-    def __init__(self):
-        super(MultiPeriodDiscriminator, self).__init__()
-        self.discriminators = nn.ModuleList([
-            DiscriminatorP(2),
-            DiscriminatorP(3),
-            DiscriminatorP(5),
-            DiscriminatorP(7),
-            DiscriminatorP(11),
-        ])
-
-    def forward(self, y, y_hat):
-        y_d_rs = []
-        y_d_gs = []
-        fmap_rs = []
-        fmap_gs = []
-        for i, d in enumerate(self.discriminators):
-            y_d_r, fmap_r = d(y)
-            y_d_g, fmap_g = d(y_hat)
-            y_d_rs.append(y_d_r)
-            fmap_rs.append(fmap_r)
-            y_d_gs.append(y_d_g)
-            fmap_gs.append(fmap_g)
-
-        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
-    
-class WavLMDiscriminator(nn.Module):
-    """docstring for Discriminator."""
-
-    def __init__(self, slm_hidden=768, 
-                 slm_layers=13, 
-                 initial_channel=64, 
-                 use_spectral_norm=False):
-        super(WavLMDiscriminator, self).__init__()
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
-        self.pre = norm_f(Conv1d(slm_hidden * slm_layers, initial_channel, 1, 1, padding=0))
-        
-        self.convs = nn.ModuleList([
-            norm_f(nn.Conv1d(initial_channel, initial_channel * 2, kernel_size=5, padding=2)),
-            norm_f(nn.Conv1d(initial_channel * 2, initial_channel * 4, kernel_size=5, padding=2)),
-            norm_f(nn.Conv1d(initial_channel * 4, initial_channel * 4, 5, 1, padding=2)),
-        ])
-
-        self.conv_post = norm_f(Conv1d(initial_channel * 4, 1, 3, 1, padding=1))
-        
-    def forward(self, x):
-        x = self.pre(x)
-        
-        fmap = []
-        for l in self.convs:
-            x = l(x)
-            x = F.leaky_relu(x, LRELU_SLOPE)
-            fmap.append(x)
-        x = self.conv_post(x)
-        x = torch.flatten(x, 1, -1)
-
-        return x

Modules/hifigan.py

@@ -1,477 +0,0 @@
-import torch
-import torch.nn.functional as F
-import torch.nn as nn
-from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
-from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
-from .utils import init_weights, get_padding
-
-import math
-import random
-import numpy as np 
-
-LRELU_SLOPE = 0.1
-
-class AdaIN1d(nn.Module):
-    def __init__(self, style_dim, num_features):
-        super().__init__()
-        self.norm = nn.InstanceNorm1d(num_features, affine=False)
-        self.fc = nn.Linear(style_dim, num_features*2)
-
-    def forward(self, x, s):
-        h = self.fc(s)
-        h = h.view(h.size(0), h.size(1), 1)
-        gamma, beta = torch.chunk(h, chunks=2, dim=1)
-        return (1 + gamma) * self.norm(x) + beta
-
-class AdaINResBlock1(torch.nn.Module):
-    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), style_dim=64):
-        super(AdaINResBlock1, self).__init__()
-        self.convs1 = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
-                               padding=get_padding(kernel_size, dilation[0]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
-                               padding=get_padding(kernel_size, dilation[1]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
-                               padding=get_padding(kernel_size, dilation[2])))
-        ])
-        self.convs1.apply(init_weights)
-
-        self.convs2 = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
-                               padding=get_padding(kernel_size, 1))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
-                               padding=get_padding(kernel_size, 1))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
-                               padding=get_padding(kernel_size, 1)))
-        ])
-        self.convs2.apply(init_weights)
-        
-        self.adain1 = nn.ModuleList([
-            AdaIN1d(style_dim, channels),
-            AdaIN1d(style_dim, channels),
-            AdaIN1d(style_dim, channels),
-        ])
-        
-        self.adain2 = nn.ModuleList([
-            AdaIN1d(style_dim, channels),
-            AdaIN1d(style_dim, channels),
-            AdaIN1d(style_dim, channels),
-        ])
-        
-        self.alpha1 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs1))])
-        self.alpha2 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs2))])
-
-
-    def forward(self, x, s):
-        for c1, c2, n1, n2, a1, a2 in zip(self.convs1, self.convs2, self.adain1, self.adain2, self.alpha1, self.alpha2):
-            xt = n1(x, s)
-            xt = xt + (1 / a1) * (torch.sin(a1 * xt) ** 2)  # Snake1D
-            xt = c1(xt)
-            xt = n2(xt, s)
-            xt = xt + (1 / a2) * (torch.sin(a2 * xt) ** 2)  # Snake1D
-            xt = c2(xt)
-            x = xt + x
-        return x
-
-    def remove_weight_norm(self):
-        for l in self.convs1:
-            remove_weight_norm(l)
-        for l in self.convs2:
-            remove_weight_norm(l)
-    
-class SineGen(torch.nn.Module):
-    """ Definition of sine generator
-    SineGen(samp_rate, harmonic_num = 0,
-            sine_amp = 0.1, noise_std = 0.003,
-            voiced_threshold = 0,
-            flag_for_pulse=False)
-    samp_rate: sampling rate in Hz
-    harmonic_num: number of harmonic overtones (default 0)
-    sine_amp: amplitude of sine-wavefrom (default 0.1)
-    noise_std: std of Gaussian noise (default 0.003)
-    voiced_thoreshold: F0 threshold for U/V classification (default 0)
-    flag_for_pulse: this SinGen is used inside PulseGen (default False)
-    Note: when flag_for_pulse is True, the first time step of a voiced
-        segment is always sin(np.pi) or cos(0)
-    """
-
-    def __init__(self, samp_rate, upsample_scale, harmonic_num=0,
-                 sine_amp=0.1, noise_std=0.003,
-                 voiced_threshold=0,
-                 flag_for_pulse=False):
-        super(SineGen, self).__init__()
-        self.sine_amp = sine_amp
-        self.noise_std = noise_std
-        self.harmonic_num = harmonic_num
-        self.dim = self.harmonic_num + 1
-        self.sampling_rate = samp_rate
-        self.voiced_threshold = voiced_threshold
-        self.flag_for_pulse = flag_for_pulse
-        self.upsample_scale = upsample_scale
-
-    def _f02uv(self, f0):
-        # generate uv signal
-        uv = (f0 > self.voiced_threshold).type(torch.float32)
-        return uv
-
-    def _f02sine(self, f0_values):
-        """ f0_values: (batchsize, length, dim)
-            where dim indicates fundamental tone and overtones
-        """
-        # convert to F0 in rad. The interger part n can be ignored
-        # because 2 * np.pi * n doesn't affect phase
-        rad_values = (f0_values / self.sampling_rate) % 1
-
-        # initial phase noise (no noise for fundamental component)
-        rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], \
-                              device=f0_values.device)
-        rand_ini[:, 0] = 0
-        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
-
-        # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad)
-        if not self.flag_for_pulse:
-#             # for normal case
-
-#             # To prevent torch.cumsum numerical overflow,
-#             # it is necessary to add -1 whenever \sum_k=1^n rad_value_k > 1.
-#             # Buffer tmp_over_one_idx indicates the time step to add -1.
-#             # This will not change F0 of sine because (x-1) * 2*pi = x * 2*pi
-#             tmp_over_one = torch.cumsum(rad_values, 1) % 1
-#             tmp_over_one_idx = (padDiff(tmp_over_one)) < 0
-#             cumsum_shift = torch.zeros_like(rad_values)
-#             cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
-
-#             phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi
-            rad_values = torch.nn.functional.interpolate(rad_values.transpose(1, 2), 
-                                                         scale_factor=1/self.upsample_scale, 
-                                                         mode="linear").transpose(1, 2)
-    
-#             tmp_over_one = torch.cumsum(rad_values, 1) % 1
-#             tmp_over_one_idx = (padDiff(tmp_over_one)) < 0
-#             cumsum_shift = torch.zeros_like(rad_values)
-#             cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
-    
-            phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi
-            phase = torch.nn.functional.interpolate(phase.transpose(1, 2) * self.upsample_scale, 
-                                                    scale_factor=self.upsample_scale, mode="linear").transpose(1, 2)
-            sines = torch.sin(phase)
-            
-        else:
-            # If necessary, make sure that the first time step of every
-            # voiced segments is sin(pi) or cos(0)
-            # This is used for pulse-train generation
-
-            # identify the last time step in unvoiced segments
-            uv = self._f02uv(f0_values)
-            uv_1 = torch.roll(uv, shifts=-1, dims=1)
-            uv_1[:, -1, :] = 1
-            u_loc = (uv < 1) * (uv_1 > 0)
-
-            # get the instantanouse phase
-            tmp_cumsum = torch.cumsum(rad_values, dim=1)
-            # different batch needs to be processed differently
-            for idx in range(f0_values.shape[0]):
-                temp_sum = tmp_cumsum[idx, u_loc[idx, :, 0], :]
-                temp_sum[1:, :] = temp_sum[1:, :] - temp_sum[0:-1, :]
-                # stores the accumulation of i.phase within
-                # each voiced segments
-                tmp_cumsum[idx, :, :] = 0
-                tmp_cumsum[idx, u_loc[idx, :, 0], :] = temp_sum
-
-            # rad_values - tmp_cumsum: remove the accumulation of i.phase
-            # within the previous voiced segment.
-            i_phase = torch.cumsum(rad_values - tmp_cumsum, dim=1)
-
-            # get the sines
-            sines = torch.cos(i_phase * 2 * np.pi)
-        return sines
-
-    def forward(self, f0):
-        """ sine_tensor, uv = forward(f0)
-        input F0: tensor(batchsize=1, length, dim=1)
-                  f0 for unvoiced steps should be 0
-        output sine_tensor: tensor(batchsize=1, length, dim)
-        output uv: tensor(batchsize=1, length, 1)
-        """
-        f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim,
-                             device=f0.device)
-        # fundamental component
-        fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device))
-
-        # generate sine waveforms
-        sine_waves = self._f02sine(fn) * self.sine_amp
-
-        # generate uv signal
-        # uv = torch.ones(f0.shape)
-        # uv = uv * (f0 > self.voiced_threshold)
-        uv = self._f02uv(f0)
-
-        # noise: for unvoiced should be similar to sine_amp
-        #        std = self.sine_amp/3 -> max value ~ self.sine_amp
-        # .       for voiced regions is self.noise_std
-        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
-        noise = noise_amp * torch.randn_like(sine_waves)
-
-        # first: set the unvoiced part to 0 by uv
-        # then: additive noise
-        sine_waves = sine_waves * uv + noise
-        return sine_waves, uv, noise
-
-
-class SourceModuleHnNSF(torch.nn.Module):
-    """ SourceModule for hn-nsf
-    SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
-                 add_noise_std=0.003, voiced_threshod=0)
-    sampling_rate: sampling_rate in Hz
-    harmonic_num: number of harmonic above F0 (default: 0)
-    sine_amp: amplitude of sine source signal (default: 0.1)
-    add_noise_std: std of additive Gaussian noise (default: 0.003)
-        note that amplitude of noise in unvoiced is decided
-        by sine_amp
-    voiced_threshold: threhold to set U/V given F0 (default: 0)
-    Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
-    F0_sampled (batchsize, length, 1)
-    Sine_source (batchsize, length, 1)
-    noise_source (batchsize, length 1)
-    uv (batchsize, length, 1)
-    """
-
-    def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1,
-                 add_noise_std=0.003, voiced_threshod=0):
-        super(SourceModuleHnNSF, self).__init__()
-
-        self.sine_amp = sine_amp
-        self.noise_std = add_noise_std
-
-        # to produce sine waveforms
-        self.l_sin_gen = SineGen(sampling_rate, upsample_scale, harmonic_num,
-                                 sine_amp, add_noise_std, voiced_threshod)
-
-        # to merge source harmonics into a single excitation
-        self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
-        self.l_tanh = torch.nn.Tanh()
-
-    def forward(self, x):
-        """
-        Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
-        F0_sampled (batchsize, length, 1)
-        Sine_source (batchsize, length, 1)
-        noise_source (batchsize, length 1)
-        """
-        # source for harmonic branch
-        with torch.no_grad():
-            sine_wavs, uv, _ = self.l_sin_gen(x)
-        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
-
-        # source for noise branch, in the same shape as uv
-        noise = torch.randn_like(uv) * self.sine_amp / 3
-        return sine_merge, noise, uv
-def padDiff(x):
-    return F.pad(F.pad(x, (0,0,-1,1), 'constant', 0) - x, (0,0,0,-1), 'constant', 0)
-
-class Generator(torch.nn.Module):
-    def __init__(self, style_dim, resblock_kernel_sizes, upsample_rates, upsample_initial_channel, resblock_dilation_sizes, upsample_kernel_sizes):
-        super(Generator, self).__init__()
-        self.num_kernels = len(resblock_kernel_sizes)
-        self.num_upsamples = len(upsample_rates)
-        resblock = AdaINResBlock1
-
-        self.m_source = SourceModuleHnNSF(
-                    sampling_rate=24000,
-                    upsample_scale=np.prod(upsample_rates),
-                    harmonic_num=8, voiced_threshod=10)
-
-        self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates))
-        self.noise_convs = nn.ModuleList()
-        self.ups = nn.ModuleList()
-        self.noise_res = nn.ModuleList()
-
-        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
-            c_cur = upsample_initial_channel // (2 ** (i + 1))
-            
-            self.ups.append(weight_norm(ConvTranspose1d(upsample_initial_channel//(2**i), 
-                         upsample_initial_channel//(2**(i+1)),
-                         k, u, padding=(u//2 + u%2), output_padding=u%2)))
-            
-            if i + 1 < len(upsample_rates):  #
-                stride_f0 = np.prod(upsample_rates[i + 1:])
-                self.noise_convs.append(Conv1d(
-                    1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0+1) // 2))
-                self.noise_res.append(resblock(c_cur, 7, [1,3,5], style_dim))
-            else:
-                self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
-                self.noise_res.append(resblock(c_cur, 11, [1,3,5], style_dim))
-            
-        self.resblocks = nn.ModuleList()
-        
-        self.alphas = nn.ParameterList()
-        self.alphas.append(nn.Parameter(torch.ones(1, upsample_initial_channel, 1)))
-        
-        for i in range(len(self.ups)):
-            ch = upsample_initial_channel//(2**(i+1))
-            self.alphas.append(nn.Parameter(torch.ones(1, ch, 1)))
-            
-            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
-                self.resblocks.append(resblock(ch, k, d, style_dim))
-
-        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
-        self.ups.apply(init_weights)
-        self.conv_post.apply(init_weights)
-
-    def forward(self, x, s, f0):
-        
-        f0 = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
-
-        har_source, noi_source, uv = self.m_source(f0)
-        har_source = har_source.transpose(1, 2)
-        
-        for i in range(self.num_upsamples):
-            x = x + (1 / self.alphas[i]) * (torch.sin(self.alphas[i] * x) ** 2)
-            x_source = self.noise_convs[i](har_source)
-            x_source = self.noise_res[i](x_source, s)
-            
-            x = self.ups[i](x)
-            x = x + x_source
-            
-            xs = None
-            for j in range(self.num_kernels):
-                if xs is None:
-                    xs = self.resblocks[i*self.num_kernels+j](x, s)
-                else:
-                    xs += self.resblocks[i*self.num_kernels+j](x, s)
-            x = xs / self.num_kernels
-        x = x + (1 / self.alphas[i+1]) * (torch.sin(self.alphas[i+1] * x) ** 2)
-        x = self.conv_post(x)
-        x = torch.tanh(x)
-
-        return x
-
-    def remove_weight_norm(self):
-        print('Removing weight norm...')
-        for l in self.ups:
-            remove_weight_norm(l)
-        for l in self.resblocks:
-            l.remove_weight_norm()
-        remove_weight_norm(self.conv_pre)
-        remove_weight_norm(self.conv_post)
-
-        
-class AdainResBlk1d(nn.Module):
-    def __init__(self, dim_in, dim_out, style_dim=64, actv=nn.LeakyReLU(0.2),
-                 upsample='none', dropout_p=0.0):
-        super().__init__()
-        self.actv = actv
-        self.upsample_type = upsample
-        self.upsample = UpSample1d(upsample)
-        self.learned_sc = dim_in != dim_out
-        self._build_weights(dim_in, dim_out, style_dim)
-        self.dropout = nn.Dropout(dropout_p)
-        
-        if upsample == 'none':
-            self.pool = nn.Identity()
-        else:
-            self.pool = weight_norm(nn.ConvTranspose1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1, output_padding=1))
-        
-        
-    def _build_weights(self, dim_in, dim_out, style_dim):
-        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
-        self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1))
-        self.norm1 = AdaIN1d(style_dim, dim_in)
-        self.norm2 = AdaIN1d(style_dim, dim_out)
-        if self.learned_sc:
-            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
-
-    def _shortcut(self, x):
-        x = self.upsample(x)
-        if self.learned_sc:
-            x = self.conv1x1(x)
-        return x
-
-    def _residual(self, x, s):
-        x = self.norm1(x, s)
-        x = self.actv(x)
-        x = self.pool(x)
-        x = self.conv1(self.dropout(x))
-        x = self.norm2(x, s)
-        x = self.actv(x)
-        x = self.conv2(self.dropout(x))
-        return x
-
-    def forward(self, x, s):
-        out = self._residual(x, s)
-        out = (out + self._shortcut(x)) / math.sqrt(2)
-        return out
-    
-class UpSample1d(nn.Module):
-    def __init__(self, layer_type):
-        super().__init__()
-        self.layer_type = layer_type
-
-    def forward(self, x):
-        if self.layer_type == 'none':
-            return x
-        else:
-            return F.interpolate(x, scale_factor=2, mode='nearest')
-
-class Decoder(nn.Module):
-    def __init__(self, dim_in=512, F0_channel=512, style_dim=64, dim_out=80, 
-                resblock_kernel_sizes = [3,7,11],
-                upsample_rates = [10,5,3,2],
-                upsample_initial_channel=512,
-                resblock_dilation_sizes=[[1,3,5], [1,3,5], [1,3,5]],
-                upsample_kernel_sizes=[20,10,6,4]):
-        super().__init__()
-        
-        self.decode = nn.ModuleList()
-        
-        self.encode = AdainResBlk1d(dim_in + 2, 1024, style_dim)
-        
-        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
-        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
-        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
-        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 512, style_dim, upsample=True))
-
-        self.F0_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
-        
-        self.N_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
-        
-        self.asr_res = nn.Sequential(
-            weight_norm(nn.Conv1d(512, 64, kernel_size=1)),
-        )
-        
-        
-        self.generator = Generator(style_dim, resblock_kernel_sizes, upsample_rates, upsample_initial_channel, resblock_dilation_sizes, upsample_kernel_sizes)
-
-        
-    def forward(self, asr, F0_curve, N, s):
-        if self.training:
-            downlist = [0, 3, 7]
-            F0_down = downlist[random.randint(0, 2)]
-            downlist = [0, 3, 7, 15]
-            N_down = downlist[random.randint(0, 3)]
-            if F0_down:
-                F0_curve = nn.functional.conv1d(F0_curve.unsqueeze(1), torch.ones(1, 1, F0_down).to('cuda'), padding=F0_down//2).squeeze(1) / F0_down
-            if N_down:
-                N = nn.functional.conv1d(N.unsqueeze(1), torch.ones(1, 1, N_down).to('cuda'), padding=N_down//2).squeeze(1)  / N_down
-
-        
-        F0 = self.F0_conv(F0_curve.unsqueeze(1))
-        N = self.N_conv(N.unsqueeze(1))
-        
-        x = torch.cat([asr, F0, N], axis=1)
-        x = self.encode(x, s)
-        
-        asr_res = self.asr_res(asr)
-        
-        res = True
-        for block in self.decode:
-            if res:
-                x = torch.cat([x, asr_res, F0, N], axis=1)
-            x = block(x, s)
-            if block.upsample_type != "none":
-                res = False
-                
-        x = self.generator(x, s, F0_curve)
-        return x
-    
-    

Modules/istftnet.py

@@ -1,530 +0,0 @@
-import torch
-import torch.nn.functional as F
-import torch.nn as nn
-from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
-from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
-from .utils import init_weights, get_padding
-
-import math
-import random
-import numpy as np 
-from scipy.signal import get_window
-
-LRELU_SLOPE = 0.1
-
-class AdaIN1d(nn.Module):
-    def __init__(self, style_dim, num_features):
-        super().__init__()
-        self.norm = nn.InstanceNorm1d(num_features, affine=False)
-        self.fc = nn.Linear(style_dim, num_features*2)
-
-    def forward(self, x, s):
-        h = self.fc(s)
-        h = h.view(h.size(0), h.size(1), 1)
-        gamma, beta = torch.chunk(h, chunks=2, dim=1)
-        return (1 + gamma) * self.norm(x) + beta
-
-class AdaINResBlock1(torch.nn.Module):
-    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), style_dim=64):
-        super(AdaINResBlock1, self).__init__()
-        self.convs1 = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
-                               padding=get_padding(kernel_size, dilation[0]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
-                               padding=get_padding(kernel_size, dilation[1]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
-                               padding=get_padding(kernel_size, dilation[2])))
-        ])
-        self.convs1.apply(init_weights)
-
-        self.convs2 = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
-                               padding=get_padding(kernel_size, 1))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
-                               padding=get_padding(kernel_size, 1))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
-                               padding=get_padding(kernel_size, 1)))
-        ])
-        self.convs2.apply(init_weights)
-        
-        self.adain1 = nn.ModuleList([
-            AdaIN1d(style_dim, channels),
-            AdaIN1d(style_dim, channels),
-            AdaIN1d(style_dim, channels),
-        ])
-        
-        self.adain2 = nn.ModuleList([
-            AdaIN1d(style_dim, channels),
-            AdaIN1d(style_dim, channels),
-            AdaIN1d(style_dim, channels),
-        ])
-        
-        self.alpha1 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs1))])
-        self.alpha2 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs2))])
-
-
-    def forward(self, x, s):
-        for c1, c2, n1, n2, a1, a2 in zip(self.convs1, self.convs2, self.adain1, self.adain2, self.alpha1, self.alpha2):
-            xt = n1(x, s)
-            xt = xt + (1 / a1) * (torch.sin(a1 * xt) ** 2)  # Snake1D
-            xt = c1(xt)
-            xt = n2(xt, s)
-            xt = xt + (1 / a2) * (torch.sin(a2 * xt) ** 2)  # Snake1D
-            xt = c2(xt)
-            x = xt + x
-        return x
-
-    def remove_weight_norm(self):
-        for l in self.convs1:
-            remove_weight_norm(l)
-        for l in self.convs2:
-            remove_weight_norm(l)
-            
-class TorchSTFT(torch.nn.Module):
-    def __init__(self, filter_length=800, hop_length=200, win_length=800, window='hann'):
-        super().__init__()
-        self.filter_length = filter_length
-        self.hop_length = hop_length
-        self.win_length = win_length
-        self.window = torch.from_numpy(get_window(window, win_length, fftbins=True).astype(np.float32))
-
-    def transform(self, input_data):
-        forward_transform = torch.stft(
-            input_data,
-            self.filter_length, self.hop_length, self.win_length, window=self.window.to(input_data.device),
-            return_complex=True)
-
-        return torch.abs(forward_transform), torch.angle(forward_transform)
-
-    def inverse(self, magnitude, phase):
-        inverse_transform = torch.istft(
-            magnitude * torch.exp(phase * 1j),
-            self.filter_length, self.hop_length, self.win_length, window=self.window.to(magnitude.device))
-
-        return inverse_transform.unsqueeze(-2)  # unsqueeze to stay consistent with conv_transpose1d implementation
-
-    def forward(self, input_data):
-        self.magnitude, self.phase = self.transform(input_data)
-        reconstruction = self.inverse(self.magnitude, self.phase)
-        return reconstruction
-    
-class SineGen(torch.nn.Module):
-    """ Definition of sine generator
-    SineGen(samp_rate, harmonic_num = 0,
-            sine_amp = 0.1, noise_std = 0.003,
-            voiced_threshold = 0,
-            flag_for_pulse=False)
-    samp_rate: sampling rate in Hz
-    harmonic_num: number of harmonic overtones (default 0)
-    sine_amp: amplitude of sine-wavefrom (default 0.1)
-    noise_std: std of Gaussian noise (default 0.003)
-    voiced_thoreshold: F0 threshold for U/V classification (default 0)
-    flag_for_pulse: this SinGen is used inside PulseGen (default False)
-    Note: when flag_for_pulse is True, the first time step of a voiced
-        segment is always sin(np.pi) or cos(0)
-    """
-
-    def __init__(self, samp_rate, upsample_scale, harmonic_num=0,
-                 sine_amp=0.1, noise_std=0.003,
-                 voiced_threshold=0,
-                 flag_for_pulse=False):
-        super(SineGen, self).__init__()
-        self.sine_amp = sine_amp
-        self.noise_std = noise_std
-        self.harmonic_num = harmonic_num
-        self.dim = self.harmonic_num + 1
-        self.sampling_rate = samp_rate
-        self.voiced_threshold = voiced_threshold
-        self.flag_for_pulse = flag_for_pulse
-        self.upsample_scale = upsample_scale
-
-    def _f02uv(self, f0):
-        # generate uv signal
-        uv = (f0 > self.voiced_threshold).type(torch.float32)
-        return uv
-
-    def _f02sine(self, f0_values):
-        """ f0_values: (batchsize, length, dim)
-            where dim indicates fundamental tone and overtones
-        """
-        # convert to F0 in rad. The interger part n can be ignored
-        # because 2 * np.pi * n doesn't affect phase
-        rad_values = (f0_values / self.sampling_rate) % 1
-
-        # initial phase noise (no noise for fundamental component)
-        rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], \
-                              device=f0_values.device)
-        rand_ini[:, 0] = 0
-        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
-
-        # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad)
-        if not self.flag_for_pulse:
-#             # for normal case
-
-#             # To prevent torch.cumsum numerical overflow,
-#             # it is necessary to add -1 whenever \sum_k=1^n rad_value_k > 1.
-#             # Buffer tmp_over_one_idx indicates the time step to add -1.
-#             # This will not change F0 of sine because (x-1) * 2*pi = x * 2*pi
-#             tmp_over_one = torch.cumsum(rad_values, 1) % 1
-#             tmp_over_one_idx = (padDiff(tmp_over_one)) < 0
-#             cumsum_shift = torch.zeros_like(rad_values)
-#             cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
-
-#             phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi
-            rad_values = torch.nn.functional.interpolate(rad_values.transpose(1, 2), 
-                                                         scale_factor=1/self.upsample_scale, 
-                                                         mode="linear").transpose(1, 2)
-    
-#             tmp_over_one = torch.cumsum(rad_values, 1) % 1
-#             tmp_over_one_idx = (padDiff(tmp_over_one)) < 0
-#             cumsum_shift = torch.zeros_like(rad_values)
-#             cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
-    
-            phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi
-            phase = torch.nn.functional.interpolate(phase.transpose(1, 2) * self.upsample_scale, 
-                                                    scale_factor=self.upsample_scale, mode="linear").transpose(1, 2)
-            sines = torch.sin(phase)
-            
-        else:
-            # If necessary, make sure that the first time step of every
-            # voiced segments is sin(pi) or cos(0)
-            # This is used for pulse-train generation
-
-            # identify the last time step in unvoiced segments
-            uv = self._f02uv(f0_values)
-            uv_1 = torch.roll(uv, shifts=-1, dims=1)
-            uv_1[:, -1, :] = 1
-            u_loc = (uv < 1) * (uv_1 > 0)
-
-            # get the instantanouse phase
-            tmp_cumsum = torch.cumsum(rad_values, dim=1)
-            # different batch needs to be processed differently
-            for idx in range(f0_values.shape[0]):
-                temp_sum = tmp_cumsum[idx, u_loc[idx, :, 0], :]
-                temp_sum[1:, :] = temp_sum[1:, :] - temp_sum[0:-1, :]
-                # stores the accumulation of i.phase within
-                # each voiced segments
-                tmp_cumsum[idx, :, :] = 0
-                tmp_cumsum[idx, u_loc[idx, :, 0], :] = temp_sum
-
-            # rad_values - tmp_cumsum: remove the accumulation of i.phase
-            # within the previous voiced segment.
-            i_phase = torch.cumsum(rad_values - tmp_cumsum, dim=1)
-
-            # get the sines
-            sines = torch.cos(i_phase * 2 * np.pi)
-        return sines
-
-    def forward(self, f0):
-        """ sine_tensor, uv = forward(f0)
-        input F0: tensor(batchsize=1, length, dim=1)
-                  f0 for unvoiced steps should be 0
-        output sine_tensor: tensor(batchsize=1, length, dim)
-        output uv: tensor(batchsize=1, length, 1)
-        """
-        f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim,
-                             device=f0.device)
-        # fundamental component
-        fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device))
-
-        # generate sine waveforms
-        sine_waves = self._f02sine(fn) * self.sine_amp
-
-        # generate uv signal
-        # uv = torch.ones(f0.shape)
-        # uv = uv * (f0 > self.voiced_threshold)
-        uv = self._f02uv(f0)
-
-        # noise: for unvoiced should be similar to sine_amp
-        #        std = self.sine_amp/3 -> max value ~ self.sine_amp
-        # .       for voiced regions is self.noise_std
-        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
-        noise = noise_amp * torch.randn_like(sine_waves)
-
-        # first: set the unvoiced part to 0 by uv
-        # then: additive noise
-        sine_waves = sine_waves * uv + noise
-        return sine_waves, uv, noise
-
-
-class SourceModuleHnNSF(torch.nn.Module):
-    """ SourceModule for hn-nsf
-    SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
-                 add_noise_std=0.003, voiced_threshod=0)
-    sampling_rate: sampling_rate in Hz
-    harmonic_num: number of harmonic above F0 (default: 0)
-    sine_amp: amplitude of sine source signal (default: 0.1)
-    add_noise_std: std of additive Gaussian noise (default: 0.003)
-        note that amplitude of noise in unvoiced is decided
-        by sine_amp
-    voiced_threshold: threhold to set U/V given F0 (default: 0)
-    Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
-    F0_sampled (batchsize, length, 1)
-    Sine_source (batchsize, length, 1)
-    noise_source (batchsize, length 1)
-    uv (batchsize, length, 1)
-    """
-
-    def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1,
-                 add_noise_std=0.003, voiced_threshod=0):
-        super(SourceModuleHnNSF, self).__init__()
-
-        self.sine_amp = sine_amp
-        self.noise_std = add_noise_std
-
-        # to produce sine waveforms
-        self.l_sin_gen = SineGen(sampling_rate, upsample_scale, harmonic_num,
-                                 sine_amp, add_noise_std, voiced_threshod)
-
-        # to merge source harmonics into a single excitation
-        self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
-        self.l_tanh = torch.nn.Tanh()
-
-    def forward(self, x):
-        """
-        Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
-        F0_sampled (batchsize, length, 1)
-        Sine_source (batchsize, length, 1)
-        noise_source (batchsize, length 1)
-        """
-        # source for harmonic branch
-        with torch.no_grad():
-            sine_wavs, uv, _ = self.l_sin_gen(x)
-        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
-
-        # source for noise branch, in the same shape as uv
-        noise = torch.randn_like(uv) * self.sine_amp / 3
-        return sine_merge, noise, uv
-def padDiff(x):
-    return F.pad(F.pad(x, (0,0,-1,1), 'constant', 0) - x, (0,0,0,-1), 'constant', 0)
-
-    
-class Generator(torch.nn.Module):
-    def __init__(self, style_dim, resblock_kernel_sizes, upsample_rates, upsample_initial_channel, resblock_dilation_sizes, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size):
-        super(Generator, self).__init__()
-
-        self.num_kernels = len(resblock_kernel_sizes)
-        self.num_upsamples = len(upsample_rates)
-        resblock = AdaINResBlock1
-
-        self.m_source = SourceModuleHnNSF(
-                    sampling_rate=24000,
-                    upsample_scale=np.prod(upsample_rates) * gen_istft_hop_size,
-                    harmonic_num=8, voiced_threshod=10)
-        self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * gen_istft_hop_size)
-        self.noise_convs = nn.ModuleList()
-        self.noise_res = nn.ModuleList()
-        
-        self.ups = nn.ModuleList()
-        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
-            self.ups.append(weight_norm(
-                ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
-                                k, u, padding=(k-u)//2)))
-
-        self.resblocks = nn.ModuleList()
-        for i in range(len(self.ups)):
-            ch = upsample_initial_channel//(2**(i+1))
-            for j, (k, d) in enumerate(zip(resblock_kernel_sizes,resblock_dilation_sizes)):
-                self.resblocks.append(resblock(ch, k, d, style_dim))
-                
-            c_cur = upsample_initial_channel // (2 ** (i + 1))
-            
-            if i + 1 < len(upsample_rates):  #
-                stride_f0 = np.prod(upsample_rates[i + 1:])
-                self.noise_convs.append(Conv1d(
-                    gen_istft_n_fft + 2, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0+1) // 2))
-                self.noise_res.append(resblock(c_cur, 7, [1,3,5], style_dim))
-            else:
-                self.noise_convs.append(Conv1d(gen_istft_n_fft + 2, c_cur, kernel_size=1))
-                self.noise_res.append(resblock(c_cur, 11, [1,3,5], style_dim))
-                
-                
-        self.post_n_fft = gen_istft_n_fft
-        self.conv_post = weight_norm(Conv1d(ch, self.post_n_fft + 2, 7, 1, padding=3))
-        self.ups.apply(init_weights)
-        self.conv_post.apply(init_weights)
-        self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
-        self.stft = TorchSTFT(filter_length=gen_istft_n_fft, hop_length=gen_istft_hop_size, win_length=gen_istft_n_fft)
-        
-        
-    def forward(self, x, s, f0):
-        with torch.no_grad():
-            f0 = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
-
-            har_source, noi_source, uv = self.m_source(f0)
-            har_source = har_source.transpose(1, 2).squeeze(1)
-            har_spec, har_phase = self.stft.transform(har_source)
-            har = torch.cat([har_spec, har_phase], dim=1)
-        
-        for i in range(self.num_upsamples):
-            x = F.leaky_relu(x, LRELU_SLOPE)
-            x_source = self.noise_convs[i](har)
-            x_source = self.noise_res[i](x_source, s)
-
-            x = self.ups[i](x)
-            if i == self.num_upsamples - 1:
-                x = self.reflection_pad(x)
-
-            x = x + x_source
-            xs = None
-            for j in range(self.num_kernels):
-                if xs is None:
-                    xs = self.resblocks[i*self.num_kernels+j](x, s)
-                else:
-                    xs += self.resblocks[i*self.num_kernels+j](x, s)
-            x = xs / self.num_kernels
-        x = F.leaky_relu(x)
-        x = self.conv_post(x)
-        spec = torch.exp(x[:,:self.post_n_fft // 2 + 1, :])
-        phase = torch.sin(x[:, self.post_n_fft // 2 + 1:, :])
-        return self.stft.inverse(spec, phase)
-    
-    def fw_phase(self, x, s):
-        for i in range(self.num_upsamples):
-            x = F.leaky_relu(x, LRELU_SLOPE)
-            x = self.ups[i](x)
-            xs = None
-            for j in range(self.num_kernels):
-                if xs is None:
-                    xs = self.resblocks[i*self.num_kernels+j](x, s)
-                else:
-                    xs += self.resblocks[i*self.num_kernels+j](x, s)
-            x = xs / self.num_kernels
-        x = F.leaky_relu(x)
-        x = self.reflection_pad(x)
-        x = self.conv_post(x)
-        spec = torch.exp(x[:,:self.post_n_fft // 2 + 1, :])
-        phase = torch.sin(x[:, self.post_n_fft // 2 + 1:, :])
-        return spec, phase
-
-    def remove_weight_norm(self):
-        print('Removing weight norm...')
-        for l in self.ups:
-            remove_weight_norm(l)
-        for l in self.resblocks:
-            l.remove_weight_norm()
-        remove_weight_norm(self.conv_pre)
-        remove_weight_norm(self.conv_post)
-
-        
-class AdainResBlk1d(nn.Module):
-    def __init__(self, dim_in, dim_out, style_dim=64, actv=nn.LeakyReLU(0.2),
-                 upsample='none', dropout_p=0.0):
-        super().__init__()
-        self.actv = actv
-        self.upsample_type = upsample
-        self.upsample = UpSample1d(upsample)
-        self.learned_sc = dim_in != dim_out
-        self._build_weights(dim_in, dim_out, style_dim)
-        self.dropout = nn.Dropout(dropout_p)
-        
-        if upsample == 'none':
-            self.pool = nn.Identity()
-        else:
-            self.pool = weight_norm(nn.ConvTranspose1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1, output_padding=1))
-        
-        
-    def _build_weights(self, dim_in, dim_out, style_dim):
-        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
-        self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1))
-        self.norm1 = AdaIN1d(style_dim, dim_in)
-        self.norm2 = AdaIN1d(style_dim, dim_out)
-        if self.learned_sc:
-            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
-
-    def _shortcut(self, x):
-        x = self.upsample(x)
-        if self.learned_sc:
-            x = self.conv1x1(x)
-        return x
-
-    def _residual(self, x, s):
-        x = self.norm1(x, s)
-        x = self.actv(x)
-        x = self.pool(x)
-        x = self.conv1(self.dropout(x))
-        x = self.norm2(x, s)
-        x = self.actv(x)
-        x = self.conv2(self.dropout(x))
-        return x
-
-    def forward(self, x, s):
-        out = self._residual(x, s)
-        out = (out + self._shortcut(x)) / math.sqrt(2)
-        return out
-    
-class UpSample1d(nn.Module):
-    def __init__(self, layer_type):
-        super().__init__()
-        self.layer_type = layer_type
-
-    def forward(self, x):
-        if self.layer_type == 'none':
-            return x
-        else:
-            return F.interpolate(x, scale_factor=2, mode='nearest')
-
-class Decoder(nn.Module):
-    def __init__(self, dim_in=512, F0_channel=512, style_dim=64, dim_out=80, 
-                resblock_kernel_sizes = [3,7,11],
-                upsample_rates = [10, 6],
-                upsample_initial_channel=512,
-                resblock_dilation_sizes=[[1,3,5], [1,3,5], [1,3,5]],
-                upsample_kernel_sizes=[20, 12], 
-                gen_istft_n_fft=20, gen_istft_hop_size=5):
-        super().__init__()
-        
-        self.decode = nn.ModuleList()
-        
-        self.encode = AdainResBlk1d(dim_in + 2, 1024, style_dim)
-        
-        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
-        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
-        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
-        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 512, style_dim, upsample=True))
-
-        self.F0_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
-        
-        self.N_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
-        
-        self.asr_res = nn.Sequential(
-            weight_norm(nn.Conv1d(512, 64, kernel_size=1)),
-        )
-        
-        
-        self.generator = Generator(style_dim, resblock_kernel_sizes, upsample_rates, 
-                                   upsample_initial_channel, resblock_dilation_sizes, 
-                                   upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size)
-        
-    def forward(self, asr, F0_curve, N, s):
-        if self.training:
-            downlist = [0, 3, 7]
-            F0_down = downlist[random.randint(0, 2)]
-            downlist = [0, 3, 7, 15]
-            N_down = downlist[random.randint(0, 3)]
-            if F0_down:
-                F0_curve = nn.functional.conv1d(F0_curve.unsqueeze(1), torch.ones(1, 1, F0_down).to('cuda'), padding=F0_down//2).squeeze(1) / F0_down
-            if N_down:
-                N = nn.functional.conv1d(N.unsqueeze(1), torch.ones(1, 1, N_down).to('cuda'), padding=N_down//2).squeeze(1)  / N_down
-
-        
-        F0 = self.F0_conv(F0_curve.unsqueeze(1))
-        N = self.N_conv(N.unsqueeze(1))
-        
-        x = torch.cat([asr, F0, N], axis=1)
-        x = self.encode(x, s)
-        
-        asr_res = self.asr_res(asr)
-        
-        res = True
-        for block in self.decode:
-            if res:
-                x = torch.cat([x, asr_res, F0, N], axis=1)
-            x = block(x, s)
-            if block.upsample_type != "none":
-                res = False
-                
-        x = self.generator(x, s, F0_curve)
-        return x
-    
-    

Modules/slmadv.py

@@ -1,195 +0,0 @@
-import torch
-import numpy as np
-import torch.nn.functional as F
-
-class SLMAdversarialLoss(torch.nn.Module):
-
-    def __init__(self, model, wl, sampler, min_len, max_len, batch_percentage=0.5, skip_update=10, sig=1.5):
-        super(SLMAdversarialLoss, self).__init__()
-        self.model = model
-        self.wl = wl
-        self.sampler = sampler
-        
-        self.min_len = min_len
-        self.max_len = max_len
-        self.batch_percentage = batch_percentage
-        
-        self.sig = sig
-        self.skip_update = skip_update
-        
-    def forward(self, iters, y_rec_gt, y_rec_gt_pred, waves, mel_input_length, ref_text, ref_lengths, use_ind, s_trg, ref_s=None):
-        text_mask = length_to_mask(ref_lengths).to(ref_text.device)
-        bert_dur = self.model.bert(ref_text, attention_mask=(~text_mask).int())
-        d_en = self.model.bert_encoder(bert_dur).transpose(-1, -2) 
-        
-        if use_ind and np.random.rand() < 0.5:
-            s_preds = s_trg
-        else:
-            num_steps = np.random.randint(3, 5)
-            if ref_s is not None:
-                s_preds = self.sampler(noise = torch.randn_like(s_trg).unsqueeze(1).to(ref_text.device), 
-                      embedding=bert_dur,
-                      embedding_scale=1,
-                               features=ref_s, # reference from the same speaker as the embedding
-                         embedding_mask_proba=0.1,
-                         num_steps=num_steps).squeeze(1)
-            else:
-                s_preds = self.sampler(noise = torch.randn_like(s_trg).unsqueeze(1).to(ref_text.device), 
-                      embedding=bert_dur,
-                      embedding_scale=1,
-                         embedding_mask_proba=0.1,
-                         num_steps=num_steps).squeeze(1)
-            
-        s_dur = s_preds[:, 128:]
-        s = s_preds[:, :128]
-        
-        d, _ = self.model.predictor(d_en, s_dur, 
-                                                ref_lengths, 
-                                                torch.randn(ref_lengths.shape[0], ref_lengths.max(), 2).to(ref_text.device), 
-                                                text_mask)
-        
-        bib = 0
-
-        output_lengths = []
-        attn_preds = []
-        
-        # differentiable duration modeling
-        for _s2s_pred, _text_length in zip(d, ref_lengths):
-
-            _s2s_pred_org = _s2s_pred[:_text_length, :]
-
-            _s2s_pred = torch.sigmoid(_s2s_pred_org)
-            _dur_pred = _s2s_pred.sum(axis=-1)
-
-            l = int(torch.round(_s2s_pred.sum()).item())
-            t = torch.arange(0, l).expand(l)
-
-            t = torch.arange(0, l).unsqueeze(0).expand((len(_s2s_pred), l)).to(ref_text.device)
-            loc = torch.cumsum(_dur_pred, dim=0) - _dur_pred / 2
-
-            h = torch.exp(-0.5 * torch.square(t - (l - loc.unsqueeze(-1))) / (self.sig)**2)
-
-            out = torch.nn.functional.conv1d(_s2s_pred_org.unsqueeze(0), 
-                                         h.unsqueeze(1), 
-                                         padding=h.shape[-1] - 1, groups=int(_text_length))[..., :l]
-            attn_preds.append(F.softmax(out.squeeze(), dim=0))
-
-            output_lengths.append(l)
-
-        max_len = max(output_lengths)
-        
-        with torch.no_grad():
-            t_en = self.model.text_encoder(ref_text, ref_lengths, text_mask)
-            
-        s2s_attn = torch.zeros(len(ref_lengths), int(ref_lengths.max()), max_len).to(ref_text.device)
-        for bib in range(len(output_lengths)):
-            s2s_attn[bib, :ref_lengths[bib], :output_lengths[bib]] = attn_preds[bib]
-
-        asr_pred = t_en @ s2s_attn
-
-        _, p_pred = self.model.predictor(d_en, s_dur, 
-                                                ref_lengths, 
-                                                s2s_attn, 
-                                                text_mask)
-        
-        mel_len = max(int(min(output_lengths) / 2 - 1), self.min_len // 2)
-        mel_len = min(mel_len, self.max_len // 2)
-        
-        # get clips
-        
-        en = []
-        p_en = []
-        sp = []
-        
-        F0_fakes = []
-        N_fakes = []
-        
-        wav = []
-
-        for bib in range(len(output_lengths)):
-            mel_length_pred = output_lengths[bib]
-            mel_length_gt = int(mel_input_length[bib].item() / 2)
-            if mel_length_gt <= mel_len or mel_length_pred <= mel_len:
-                continue
-
-            sp.append(s_preds[bib])
-
-            random_start = np.random.randint(0, mel_length_pred - mel_len)
-            en.append(asr_pred[bib, :, random_start:random_start+mel_len])
-            p_en.append(p_pred[bib, :, random_start:random_start+mel_len])
-
-            # get ground truth clips
-            random_start = np.random.randint(0, mel_length_gt - mel_len)
-            y = waves[bib][(random_start * 2) * 300:((random_start+mel_len) * 2) * 300]
-            wav.append(torch.from_numpy(y).to(ref_text.device))
-            
-            if len(wav) >= self.batch_percentage * len(waves): # prevent OOM due to longer lengths
-                break
-
-        if len(sp) <= 1:
-            return None
-            
-        sp = torch.stack(sp)
-        wav = torch.stack(wav).float()
-        en = torch.stack(en)
-        p_en = torch.stack(p_en)
-        
-        F0_fake, N_fake = self.model.predictor.F0Ntrain(p_en, sp[:, 128:])
-        y_pred = self.model.decoder(en, F0_fake, N_fake, sp[:, :128])
-        
-        # discriminator loss
-        if (iters + 1) % self.skip_update == 0:
-            if np.random.randint(0, 2) == 0:
-                wav = y_rec_gt_pred
-                use_rec = True
-            else:
-                use_rec = False
-
-            crop_size = min(wav.size(-1), y_pred.size(-1))
-            if use_rec: # use reconstructed (shorter lengths), do length invariant regularization
-                if wav.size(-1) > y_pred.size(-1):
-                    real_GP = wav[:, : , :crop_size]
-                    out_crop = self.wl.discriminator_forward(real_GP.detach().squeeze())
-                    out_org = self.wl.discriminator_forward(wav.detach().squeeze())
-                    loss_reg = F.l1_loss(out_crop, out_org[..., :out_crop.size(-1)])
-
-                    if np.random.randint(0, 2) == 0:
-                        d_loss = self.wl.discriminator(real_GP.detach().squeeze(), y_pred.detach().squeeze()).mean()
-                    else:
-                        d_loss = self.wl.discriminator(wav.detach().squeeze(), y_pred.detach().squeeze()).mean()
-                else:
-                    real_GP = y_pred[:, : , :crop_size]
-                    out_crop = self.wl.discriminator_forward(real_GP.detach().squeeze())
-                    out_org = self.wl.discriminator_forward(y_pred.detach().squeeze())
-                    loss_reg = F.l1_loss(out_crop, out_org[..., :out_crop.size(-1)])
-
-                    if np.random.randint(0, 2) == 0:
-                        d_loss = self.wl.discriminator(wav.detach().squeeze(), real_GP.detach().squeeze()).mean()
-                    else:
-                        d_loss = self.wl.discriminator(wav.detach().squeeze(), y_pred.detach().squeeze()).mean()
-                
-                # regularization (ignore length variation)
-                d_loss += loss_reg
-
-                out_gt = self.wl.discriminator_forward(y_rec_gt.detach().squeeze())
-                out_rec = self.wl.discriminator_forward(y_rec_gt_pred.detach().squeeze())
-
-                # regularization (ignore reconstruction artifacts)
-                d_loss += F.l1_loss(out_gt, out_rec)
-
-            else:
-                d_loss = self.wl.discriminator(wav.detach().squeeze(), y_pred.detach().squeeze()).mean()
-        else:
-            d_loss = 0
-            
-        # generator loss
-        gen_loss = self.wl.generator(y_pred.squeeze())
-        
-        gen_loss = gen_loss.mean()
-        
-        return d_loss, gen_loss, y_pred.detach().cpu().numpy()
-    
-def length_to_mask(lengths):
-    mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
-    mask = torch.gt(mask+1, lengths.unsqueeze(1))
-    return mask

Modules/utils.py

@@ -1,14 +0,0 @@
-def init_weights(m, mean=0.0, std=0.01):
-    classname = m.__class__.__name__
-    if classname.find("Conv") != -1:
-        m.weight.data.normal_(mean, std)
-
-
-def apply_weight_norm(m):
-    classname = m.__class__.__name__
-    if classname.find("Conv") != -1:
-        weight_norm(m)
-
-
-def get_padding(kernel_size, dilation=1):
-    return int((kernel_size*dilation - dilation)/2)

Utils/ASR/__init__.py

@@ -1 +0,0 @@
-

Utils/ASR/config.yml

@@ -1,29 +0,0 @@
-log_dir: "logs/20201006"
-save_freq: 5
-device: "cuda"
-epochs: 180
-batch_size: 64
-pretrained_model: ""
-train_data: "ASRDataset/train_list.txt"
-val_data: "ASRDataset/val_list.txt"
-
-dataset_params:
-  data_augmentation: false
-
-preprocess_parasm:
-  sr: 24000
-  spect_params:
-    n_fft: 2048
-    win_length: 1200
-    hop_length: 300
-  mel_params:
-    n_mels: 80
-
-model_params:
-   input_dim: 80
-   hidden_dim: 256
-   n_token: 181
-   token_embedding_dim: 512
-
-optimizer_params:
-  lr: 0.0005

Utils/ASR/layers.py

@@ -1,354 +0,0 @@
-import math
-import torch
-from torch import nn
-from typing import Optional, Any
-from torch import Tensor
-import torch.nn.functional as F
-import torchaudio
-import torchaudio.functional as audio_F
-
-import random
-random.seed(0)
-
-
-def _get_activation_fn(activ):
-    if activ == 'relu':
-        return nn.ReLU()
-    elif activ == 'lrelu':
-        return nn.LeakyReLU(0.2)
-    elif activ == 'swish':
-        return lambda x: x*torch.sigmoid(x)
-    else:
-        raise RuntimeError('Unexpected activ type %s, expected [relu, lrelu, swish]' % activ)
-
-class LinearNorm(torch.nn.Module):
-    def __init__(self, in_dim, out_dim, bias=True, w_init_gain='linear'):
-        super(LinearNorm, self).__init__()
-        self.linear_layer = torch.nn.Linear(in_dim, out_dim, bias=bias)
-
-        torch.nn.init.xavier_uniform_(
-            self.linear_layer.weight,
-            gain=torch.nn.init.calculate_gain(w_init_gain))
-
-    def forward(self, x):
-        return self.linear_layer(x)
-
-
-class ConvNorm(torch.nn.Module):
-    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1,
-                 padding=None, dilation=1, bias=True, w_init_gain='linear', param=None):
-        super(ConvNorm, self).__init__()
-        if padding is None:
-            assert(kernel_size % 2 == 1)
-            padding = int(dilation * (kernel_size - 1) / 2)
-
-        self.conv = torch.nn.Conv1d(in_channels, out_channels,
-                                    kernel_size=kernel_size, stride=stride,
-                                    padding=padding, dilation=dilation,
-                                    bias=bias)
-
-        torch.nn.init.xavier_uniform_(
-            self.conv.weight, gain=torch.nn.init.calculate_gain(w_init_gain, param=param))
-
-    def forward(self, signal):
-        conv_signal = self.conv(signal)
-        return conv_signal
-
-class CausualConv(nn.Module):
-    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=1, dilation=1, bias=True, w_init_gain='linear', param=None):
-        super(CausualConv, self).__init__()
-        if padding is None:
-            assert(kernel_size % 2 == 1)
-            padding = int(dilation * (kernel_size - 1) / 2) * 2
-        else:
-            self.padding = padding * 2
-        self.conv = nn.Conv1d(in_channels, out_channels,
-                              kernel_size=kernel_size, stride=stride,
-                              padding=self.padding,
-                              dilation=dilation,
-                              bias=bias)
-
-        torch.nn.init.xavier_uniform_(
-            self.conv.weight, gain=torch.nn.init.calculate_gain(w_init_gain, param=param))
-
-    def forward(self, x):
-        x = self.conv(x)
-        x = x[:, :, :-self.padding]
-        return x
-
-class CausualBlock(nn.Module):
-    def __init__(self, hidden_dim, n_conv=3, dropout_p=0.2, activ='lrelu'):
-        super(CausualBlock, self).__init__()
-        self.blocks = nn.ModuleList([
-            self._get_conv(hidden_dim, dilation=3**i, activ=activ, dropout_p=dropout_p)
-            for i in range(n_conv)])
-
-    def forward(self, x):
-        for block in self.blocks:
-            res = x
-            x = block(x)
-            x += res
-        return x
-
-    def _get_conv(self, hidden_dim, dilation, activ='lrelu', dropout_p=0.2):
-        layers = [
-            CausualConv(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
-            _get_activation_fn(activ),
-            nn.BatchNorm1d(hidden_dim),
-            nn.Dropout(p=dropout_p),
-            CausualConv(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
-            _get_activation_fn(activ),
-            nn.Dropout(p=dropout_p)
-        ]
-        return nn.Sequential(*layers)
-
-class ConvBlock(nn.Module):
-    def __init__(self, hidden_dim, n_conv=3, dropout_p=0.2, activ='relu'):
-        super().__init__()
-        self._n_groups = 8
-        self.blocks = nn.ModuleList([
-            self._get_conv(hidden_dim, dilation=3**i, activ=activ, dropout_p=dropout_p)
-            for i in range(n_conv)])
-
-
-    def forward(self, x):
-        for block in self.blocks:
-            res = x
-            x = block(x)
-            x += res
-        return x
-
-    def _get_conv(self, hidden_dim, dilation, activ='relu', dropout_p=0.2):
-        layers = [
-            ConvNorm(hidden_dim, hidden_dim, kernel_size=3, padding=dilation, dilation=dilation),
-            _get_activation_fn(activ),
-            nn.GroupNorm(num_groups=self._n_groups, num_channels=hidden_dim),
-            nn.Dropout(p=dropout_p),
-            ConvNorm(hidden_dim, hidden_dim, kernel_size=3, padding=1, dilation=1),
-            _get_activation_fn(activ),
-            nn.Dropout(p=dropout_p)
-        ]
-        return nn.Sequential(*layers)
-
-class LocationLayer(nn.Module):
-    def __init__(self, attention_n_filters, attention_kernel_size,
-                 attention_dim):
-        super(LocationLayer, self).__init__()
-        padding = int((attention_kernel_size - 1) / 2)
-        self.location_conv = ConvNorm(2, attention_n_filters,
-                                      kernel_size=attention_kernel_size,
-                                      padding=padding, bias=False, stride=1,
-                                      dilation=1)
-        self.location_dense = LinearNorm(attention_n_filters, attention_dim,
-                                         bias=False, w_init_gain='tanh')
-
-    def forward(self, attention_weights_cat):
-        processed_attention = self.location_conv(attention_weights_cat)
-        processed_attention = processed_attention.transpose(1, 2)
-        processed_attention = self.location_dense(processed_attention)
-        return processed_attention
-
-
-class Attention(nn.Module):
-    def __init__(self, attention_rnn_dim, embedding_dim, attention_dim,
-                 attention_location_n_filters, attention_location_kernel_size):
-        super(Attention, self).__init__()
-        self.query_layer = LinearNorm(attention_rnn_dim, attention_dim,
-                                      bias=False, w_init_gain='tanh')
-        self.memory_layer = LinearNorm(embedding_dim, attention_dim, bias=False,
-                                       w_init_gain='tanh')
-        self.v = LinearNorm(attention_dim, 1, bias=False)
-        self.location_layer = LocationLayer(attention_location_n_filters,
-                                            attention_location_kernel_size,
-                                            attention_dim)
-        self.score_mask_value = -float("inf")
-
-    def get_alignment_energies(self, query, processed_memory,
-                               attention_weights_cat):
-        """
-        PARAMS
-        ------
-        query: decoder output (batch, n_mel_channels * n_frames_per_step)
-        processed_memory: processed encoder outputs (B, T_in, attention_dim)
-        attention_weights_cat: cumulative and prev. att weights (B, 2, max_time)
-        RETURNS
-        -------
-        alignment (batch, max_time)
-        """
-
-        processed_query = self.query_layer(query.unsqueeze(1))
-        processed_attention_weights = self.location_layer(attention_weights_cat)
-        energies = self.v(torch.tanh(
-            processed_query + processed_attention_weights + processed_memory))
-
-        energies = energies.squeeze(-1)
-        return energies
-
-    def forward(self, attention_hidden_state, memory, processed_memory,
-                attention_weights_cat, mask):
-        """
-        PARAMS
-        ------
-        attention_hidden_state: attention rnn last output
-        memory: encoder outputs
-        processed_memory: processed encoder outputs
-        attention_weights_cat: previous and cummulative attention weights
-        mask: binary mask for padded data
-        """
-        alignment = self.get_alignment_energies(
-            attention_hidden_state, processed_memory, attention_weights_cat)
-
-        if mask is not None:
-            alignment.data.masked_fill_(mask, self.score_mask_value)
-
-        attention_weights = F.softmax(alignment, dim=1)
-        attention_context = torch.bmm(attention_weights.unsqueeze(1), memory)
-        attention_context = attention_context.squeeze(1)
-
-        return attention_context, attention_weights
-
-
-class ForwardAttentionV2(nn.Module):
-    def __init__(self, attention_rnn_dim, embedding_dim, attention_dim,
-                 attention_location_n_filters, attention_location_kernel_size):
-        super(ForwardAttentionV2, self).__init__()
-        self.query_layer = LinearNorm(attention_rnn_dim, attention_dim,
-                                      bias=False, w_init_gain='tanh')
-        self.memory_layer = LinearNorm(embedding_dim, attention_dim, bias=False,
-                                       w_init_gain='tanh')
-        self.v = LinearNorm(attention_dim, 1, bias=False)
-        self.location_layer = LocationLayer(attention_location_n_filters,
-                                            attention_location_kernel_size,
-                                            attention_dim)
-        self.score_mask_value = -float(1e20)
-
-    def get_alignment_energies(self, query, processed_memory,
-                               attention_weights_cat):
-        """
-        PARAMS
-        ------
-        query: decoder output (batch, n_mel_channels * n_frames_per_step)
-        processed_memory: processed encoder outputs (B, T_in, attention_dim)
-        attention_weights_cat:  prev. and cumulative att weights (B, 2, max_time)
-        RETURNS
-        -------
-        alignment (batch, max_time)
-        """
-
-        processed_query = self.query_layer(query.unsqueeze(1))
-        processed_attention_weights = self.location_layer(attention_weights_cat)
-        energies = self.v(torch.tanh(
-            processed_query + processed_attention_weights + processed_memory))
-
-        energies = energies.squeeze(-1)
-        return energies
-
-    def forward(self, attention_hidden_state, memory, processed_memory,
-                attention_weights_cat, mask, log_alpha):
-        """
-        PARAMS
-        ------
-        attention_hidden_state: attention rnn last output
-        memory: encoder outputs
-        processed_memory: processed encoder outputs
-        attention_weights_cat: previous and cummulative attention weights
-        mask: binary mask for padded data
-        """
-        log_energy = self.get_alignment_energies(
-            attention_hidden_state, processed_memory, attention_weights_cat)
-
-        #log_energy =
-
-        if mask is not None:
-            log_energy.data.masked_fill_(mask, self.score_mask_value)
-
-        #attention_weights = F.softmax(alignment, dim=1)
-
-        #content_score = log_energy.unsqueeze(1) #[B, MAX_TIME] -> [B, 1, MAX_TIME]
-        #log_alpha = log_alpha.unsqueeze(2) #[B, MAX_TIME] -> [B, MAX_TIME, 1]
-
-        #log_total_score = log_alpha + content_score
-
-        #previous_attention_weights = attention_weights_cat[:,0,:]
-
-        log_alpha_shift_padded = []
-        max_time = log_energy.size(1)
-        for sft in range(2):
-            shifted = log_alpha[:,:max_time-sft]
-            shift_padded = F.pad(shifted, (sft,0), 'constant', self.score_mask_value)
-            log_alpha_shift_padded.append(shift_padded.unsqueeze(2))
-
-        biased = torch.logsumexp(torch.cat(log_alpha_shift_padded,2), 2)
-
-        log_alpha_new = biased +  log_energy
-
-        attention_weights =  F.softmax(log_alpha_new, dim=1)
-
-        attention_context = torch.bmm(attention_weights.unsqueeze(1), memory)
-        attention_context = attention_context.squeeze(1)
-
-        return attention_context, attention_weights, log_alpha_new
-
-
-class PhaseShuffle2d(nn.Module):
-    def __init__(self, n=2):
-        super(PhaseShuffle2d, self).__init__()
-        self.n = n
-        self.random = random.Random(1)
-
-    def forward(self, x, move=None):
-        # x.size = (B, C, M, L)
-        if move is None:
-            move = self.random.randint(-self.n, self.n)
-
-        if move == 0:
-            return x
-        else:
-            left = x[:, :, :, :move]
-            right = x[:, :, :, move:]
-            shuffled = torch.cat([right, left], dim=3)
-        return shuffled
-
-class PhaseShuffle1d(nn.Module):
-    def __init__(self, n=2):
-        super(PhaseShuffle1d, self).__init__()
-        self.n = n
-        self.random = random.Random(1)
-
-    def forward(self, x, move=None):
-        # x.size = (B, C, M, L)
-        if move is None:
-            move = self.random.randint(-self.n, self.n)
-
-        if move == 0:
-            return x
-        else:
-            left = x[:, :,  :move]
-            right = x[:, :, move:]
-            shuffled = torch.cat([right, left], dim=2)
-
-        return shuffled
-
-class MFCC(nn.Module):
-    def __init__(self, n_mfcc=40, n_mels=80):
-        super(MFCC, self).__init__()
-        self.n_mfcc = n_mfcc
-        self.n_mels = n_mels
-        self.norm = 'ortho'
-        dct_mat = audio_F.create_dct(self.n_mfcc, self.n_mels, self.norm)
-        self.register_buffer('dct_mat', dct_mat)
-
-    def forward(self, mel_specgram):
-        if len(mel_specgram.shape) == 2:
-            mel_specgram = mel_specgram.unsqueeze(0)
-            unsqueezed = True
-        else:
-            unsqueezed = False
-        # (channel, n_mels, time).tranpose(...) dot (n_mels, n_mfcc)
-        # -> (channel, time, n_mfcc).tranpose(...)
-        mfcc = torch.matmul(mel_specgram.transpose(1, 2), self.dct_mat).transpose(1, 2)
-
-        # unpack batch
-        if unsqueezed:
-            mfcc = mfcc.squeeze(0)
-        return mfcc

Utils/ASR/models.py

@@ -1,186 +0,0 @@
-import math
-import torch
-from torch import nn
-from torch.nn import TransformerEncoder
-import torch.nn.functional as F
-from .layers import MFCC, Attention, LinearNorm, ConvNorm, ConvBlock
-
-class ASRCNN(nn.Module):
-    def __init__(self,
-                 input_dim=80,
-                 hidden_dim=256,
-                 n_token=35,
-                 n_layers=6,
-                 token_embedding_dim=256,
-
-    ):
-        super().__init__()
-        self.n_token = n_token
-        self.n_down = 1
-        self.to_mfcc = MFCC()
-        self.init_cnn = ConvNorm(input_dim//2, hidden_dim, kernel_size=7, padding=3, stride=2)
-        self.cnns = nn.Sequential(
-            *[nn.Sequential(
-                ConvBlock(hidden_dim),
-                nn.GroupNorm(num_groups=1, num_channels=hidden_dim)
-            ) for n in range(n_layers)])
-        self.projection = ConvNorm(hidden_dim, hidden_dim // 2)
-        self.ctc_linear = nn.Sequential(
-            LinearNorm(hidden_dim//2, hidden_dim),
-            nn.ReLU(),
-            LinearNorm(hidden_dim, n_token))
-        self.asr_s2s = ASRS2S(
-            embedding_dim=token_embedding_dim,
-            hidden_dim=hidden_dim//2,
-            n_token=n_token)
-
-    def forward(self, x, src_key_padding_mask=None, text_input=None):
-        x = self.to_mfcc(x)
-        x = self.init_cnn(x)
-        x = self.cnns(x)
-        x = self.projection(x)
-        x = x.transpose(1, 2)
-        ctc_logit = self.ctc_linear(x)
-        if text_input is not None:
-            _, s2s_logit, s2s_attn = self.asr_s2s(x, src_key_padding_mask, text_input)
-            return ctc_logit, s2s_logit, s2s_attn
-        else:
-            return ctc_logit
-
-    def get_feature(self, x):
-        x = self.to_mfcc(x.squeeze(1))
-        x = self.init_cnn(x)
-        x = self.cnns(x)
-        x = self.projection(x)
-        return x
-
-    def length_to_mask(self, lengths):
-        mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
-        mask = torch.gt(mask+1, lengths.unsqueeze(1)).to(lengths.device)
-        return mask
-
-    def get_future_mask(self, out_length, unmask_future_steps=0):
-        """
-        Args:
-            out_length (int): returned mask shape is (out_length, out_length).
-            unmask_futre_steps (int): unmasking future step size.
-        Return:
-            mask (torch.BoolTensor): mask future timesteps mask[i, j] = True if i > j + unmask_future_steps else False
-        """
-        index_tensor = torch.arange(out_length).unsqueeze(0).expand(out_length, -1)
-        mask = torch.gt(index_tensor, index_tensor.T + unmask_future_steps)
-        return mask
-
-class ASRS2S(nn.Module):
-    def __init__(self,
-                 embedding_dim=256,
-                 hidden_dim=512,
-                 n_location_filters=32,
-                 location_kernel_size=63,
-                 n_token=40):
-        super(ASRS2S, self).__init__()
-        self.embedding = nn.Embedding(n_token, embedding_dim)
-        val_range = math.sqrt(6 / hidden_dim)
-        self.embedding.weight.data.uniform_(-val_range, val_range)
-
-        self.decoder_rnn_dim = hidden_dim
-        self.project_to_n_symbols = nn.Linear(self.decoder_rnn_dim, n_token)
-        self.attention_layer = Attention(
-            self.decoder_rnn_dim,
-            hidden_dim,
-            hidden_dim,
-            n_location_filters,
-            location_kernel_size
-        )
-        self.decoder_rnn = nn.LSTMCell(self.decoder_rnn_dim + embedding_dim, self.decoder_rnn_dim)
-        self.project_to_hidden = nn.Sequential(
-            LinearNorm(self.decoder_rnn_dim * 2, hidden_dim),
-            nn.Tanh())
-        self.sos = 1
-        self.eos = 2
-
-    def initialize_decoder_states(self, memory, mask):
-        """
-        moemory.shape = (B, L, H) = (Batchsize, Maxtimestep, Hiddendim)
-        """
-        B, L, H = memory.shape
-        self.decoder_hidden = torch.zeros((B, self.decoder_rnn_dim)).type_as(memory)
-        self.decoder_cell = torch.zeros((B, self.decoder_rnn_dim)).type_as(memory)
-        self.attention_weights = torch.zeros((B, L)).type_as(memory)
-        self.attention_weights_cum = torch.zeros((B, L)).type_as(memory)
-        self.attention_context = torch.zeros((B, H)).type_as(memory)
-        self.memory = memory
-        self.processed_memory = self.attention_layer.memory_layer(memory)
-        self.mask = mask
-        self.unk_index = 3
-        self.random_mask = 0.1
-
-    def forward(self, memory, memory_mask, text_input):
-        """
-        moemory.shape = (B, L, H) = (Batchsize, Maxtimestep, Hiddendim)
-        moemory_mask.shape = (B, L, )
-        texts_input.shape = (B, T)
-        """
-        self.initialize_decoder_states(memory, memory_mask)
-        # text random mask
-        random_mask = (torch.rand(text_input.shape) < self.random_mask).to(text_input.device)
-        _text_input = text_input.clone()
-        _text_input.masked_fill_(random_mask, self.unk_index)
-        decoder_inputs = self.embedding(_text_input).transpose(0, 1) # -> [T, B, channel]
-        start_embedding = self.embedding(
-            torch.LongTensor([self.sos]*decoder_inputs.size(1)).to(decoder_inputs.device))
-        decoder_inputs = torch.cat((start_embedding.unsqueeze(0), decoder_inputs), dim=0)
-
-        hidden_outputs, logit_outputs, alignments = [], [], []
-        while len(hidden_outputs) < decoder_inputs.size(0):
-
-            decoder_input = decoder_inputs[len(hidden_outputs)]
-            hidden, logit, attention_weights = self.decode(decoder_input)
-            hidden_outputs += [hidden]
-            logit_outputs += [logit]
-            alignments += [attention_weights]
-
-        hidden_outputs, logit_outputs, alignments = \
-            self.parse_decoder_outputs(
-                hidden_outputs, logit_outputs, alignments)
-
-        return hidden_outputs, logit_outputs, alignments
-
-
-    def decode(self, decoder_input):
-
-        cell_input = torch.cat((decoder_input, self.attention_context), -1)
-        self.decoder_hidden, self.decoder_cell = self.decoder_rnn(
-            cell_input,
-            (self.decoder_hidden, self.decoder_cell))
-
-        attention_weights_cat = torch.cat(
-            (self.attention_weights.unsqueeze(1),
-            self.attention_weights_cum.unsqueeze(1)),dim=1)
-
-        self.attention_context, self.attention_weights = self.attention_layer(
-            self.decoder_hidden,
-            self.memory,
-            self.processed_memory,
-            attention_weights_cat,
-            self.mask)
-
-        self.attention_weights_cum += self.attention_weights
-
-        hidden_and_context = torch.cat((self.decoder_hidden, self.attention_context), -1)
-        hidden = self.project_to_hidden(hidden_and_context)
-
-        # dropout to increasing g
-        logit = self.project_to_n_symbols(F.dropout(hidden, 0.5, self.training))
-
-        return hidden, logit, self.attention_weights
-
-    def parse_decoder_outputs(self, hidden, logit, alignments):
-
-        # -> [B, T_out + 1, max_time]
-        alignments = torch.stack(alignments).transpose(0,1)
-        # [T_out + 1, B, n_symbols] -> [B, T_out + 1,  n_symbols]
-        logit = torch.stack(logit).transpose(0, 1).contiguous()
-        hidden = torch.stack(hidden).transpose(0, 1).contiguous()
-
-        return hidden, logit, alignments

Utils/JDC/__init__.py

@@ -1 +0,0 @@
-

Utils/JDC/model.py

@@ -1,190 +0,0 @@
-"""
-Implementation of model from:
-Kum et al. - "Joint Detection and Classification of Singing Voice Melody Using
-Convolutional Recurrent Neural Networks" (2019)
-Link: https://www.semanticscholar.org/paper/Joint-Detection-and-Classification-of-Singing-Voice-Kum-Nam/60a2ad4c7db43bace75805054603747fcd062c0d
-"""
-import torch
-from torch import nn
-        
-class JDCNet(nn.Module):
-    """
-    Joint Detection and Classification Network model for singing voice melody.
-    """
-    def __init__(self, num_class=722, seq_len=31, leaky_relu_slope=0.01):
-        super().__init__()
-        self.num_class = num_class
-
-        # input = (b, 1, 31, 513), b = batch size
-        self.conv_block = nn.Sequential(
-            nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, padding=1, bias=False),  # out: (b, 64, 31, 513)
-            nn.BatchNorm2d(num_features=64),
-            nn.LeakyReLU(leaky_relu_slope, inplace=True),
-            nn.Conv2d(64, 64, 3, padding=1, bias=False),  # (b, 64, 31, 513)
-        )
-
-        # res blocks
-        self.res_block1 = ResBlock(in_channels=64, out_channels=128)  # (b, 128, 31, 128)
-        self.res_block2 = ResBlock(in_channels=128, out_channels=192)  # (b, 192, 31, 32)
-        self.res_block3 = ResBlock(in_channels=192, out_channels=256)  # (b, 256, 31, 8)
-
-        # pool block
-        self.pool_block = nn.Sequential(
-            nn.BatchNorm2d(num_features=256),
-            nn.LeakyReLU(leaky_relu_slope, inplace=True),
-            nn.MaxPool2d(kernel_size=(1, 4)),  # (b, 256, 31, 2)
-            nn.Dropout(p=0.2),
-        )
-
-        # maxpool layers (for auxiliary network inputs)
-        # in = (b, 128, 31, 513) from conv_block, out = (b, 128, 31, 2)
-        self.maxpool1 = nn.MaxPool2d(kernel_size=(1, 40))
-        # in = (b, 128, 31, 128) from res_block1, out = (b, 128, 31, 2)
-        self.maxpool2 = nn.MaxPool2d(kernel_size=(1, 20))
-        # in = (b, 128, 31, 32) from res_block2, out = (b, 128, 31, 2)
-        self.maxpool3 = nn.MaxPool2d(kernel_size=(1, 10))
-
-        # in = (b, 640, 31, 2), out = (b, 256, 31, 2)
-        self.detector_conv = nn.Sequential(
-            nn.Conv2d(640, 256, 1, bias=False),
-            nn.BatchNorm2d(256),
-            nn.LeakyReLU(leaky_relu_slope, inplace=True),
-            nn.Dropout(p=0.2),
-        )
-
-        # input: (b, 31, 512) - resized from (b, 256, 31, 2)
-        self.bilstm_classifier = nn.LSTM(
-            input_size=512, hidden_size=256,
-            batch_first=True, bidirectional=True)  # (b, 31, 512)
-
-        # input: (b, 31, 512) - resized from (b, 256, 31, 2)
-        self.bilstm_detector = nn.LSTM(
-            input_size=512, hidden_size=256,
-            batch_first=True, bidirectional=True)  # (b, 31, 512)
-
-        # input: (b * 31, 512)
-        self.classifier = nn.Linear(in_features=512, out_features=self.num_class)  # (b * 31, num_class)
-
-        # input: (b * 31, 512)
-        self.detector = nn.Linear(in_features=512, out_features=2)  # (b * 31, 2) - binary classifier
-
-        # initialize weights
-        self.apply(self.init_weights)
-
-    def get_feature_GAN(self, x):
-        seq_len = x.shape[-2]
-        x = x.float().transpose(-1, -2)
-        
-        convblock_out = self.conv_block(x)
-        
-        resblock1_out = self.res_block1(convblock_out)
-        resblock2_out = self.res_block2(resblock1_out)
-        resblock3_out = self.res_block3(resblock2_out)
-        poolblock_out = self.pool_block[0](resblock3_out)
-        poolblock_out = self.pool_block[1](poolblock_out)
-        
-        return poolblock_out.transpose(-1, -2)
-        
-    def get_feature(self, x):
-        seq_len = x.shape[-2]
-        x = x.float().transpose(-1, -2)
-        
-        convblock_out = self.conv_block(x)
-        
-        resblock1_out = self.res_block1(convblock_out)
-        resblock2_out = self.res_block2(resblock1_out)
-        resblock3_out = self.res_block3(resblock2_out)
-        poolblock_out = self.pool_block[0](resblock3_out)
-        poolblock_out = self.pool_block[1](poolblock_out)
-        
-        return self.pool_block[2](poolblock_out)
-        
-    def forward(self, x):
-        """
-        Returns:
-            classification_prediction, detection_prediction
-            sizes: (b, 31, 722), (b, 31, 2)
-        """
-        ###############################
-        # forward pass for classifier #
-        ###############################
-        seq_len = x.shape[-1]
-        x = x.float().transpose(-1, -2)
-        
-        convblock_out = self.conv_block(x)
-        
-        resblock1_out = self.res_block1(convblock_out)
-        resblock2_out = self.res_block2(resblock1_out)
-        resblock3_out = self.res_block3(resblock2_out)
-        
-        
-        poolblock_out = self.pool_block[0](resblock3_out)
-        poolblock_out = self.pool_block[1](poolblock_out)
-        GAN_feature = poolblock_out.transpose(-1, -2)
-        poolblock_out = self.pool_block[2](poolblock_out)
-        
-        # (b, 256, 31, 2) => (b, 31, 256, 2) => (b, 31, 512)
-        classifier_out = poolblock_out.permute(0, 2, 1, 3).contiguous().view((-1, seq_len, 512))
-        classifier_out, _ = self.bilstm_classifier(classifier_out)  # ignore the hidden states
-
-        classifier_out = classifier_out.contiguous().view((-1, 512))  # (b * 31, 512)
-        classifier_out = self.classifier(classifier_out)
-        classifier_out = classifier_out.view((-1, seq_len, self.num_class))  # (b, 31, num_class)
-        
-        # sizes: (b, 31, 722), (b, 31, 2)
-        # classifier output consists of predicted pitch classes per frame
-        # detector output consists of: (isvoice, notvoice) estimates per frame
-        return torch.abs(classifier_out.squeeze()), GAN_feature, poolblock_out
-
-    @staticmethod
-    def init_weights(m):
-        if isinstance(m, nn.Linear):
-            nn.init.kaiming_uniform_(m.weight)
-            if m.bias is not None:
-                nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.Conv2d):
-            nn.init.xavier_normal_(m.weight)
-        elif isinstance(m, nn.LSTM) or isinstance(m, nn.LSTMCell):
-            for p in m.parameters():
-                if p.data is None:
-                    continue
-
-                if len(p.shape) >= 2:
-                    nn.init.orthogonal_(p.data)
-                else:
-                    nn.init.normal_(p.data)
-                    
-
-class ResBlock(nn.Module):
-    def __init__(self, in_channels: int, out_channels: int, leaky_relu_slope=0.01):
-        super().__init__()
-        self.downsample = in_channels != out_channels
-
-        # BN / LReLU / MaxPool layer before the conv layer - see Figure 1b in the paper
-        self.pre_conv = nn.Sequential(
-            nn.BatchNorm2d(num_features=in_channels),
-            nn.LeakyReLU(leaky_relu_slope, inplace=True),
-            nn.MaxPool2d(kernel_size=(1, 2)),  # apply downsampling on the y axis only
-        )
-
-        # conv layers
-        self.conv = nn.Sequential(
-            nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
-                      kernel_size=3, padding=1, bias=False),
-            nn.BatchNorm2d(out_channels),
-            nn.LeakyReLU(leaky_relu_slope, inplace=True),
-            nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False),
-        )
-
-        # 1 x 1 convolution layer to match the feature dimensions
-        self.conv1by1 = None
-        if self.downsample:
-            self.conv1by1 = nn.Conv2d(in_channels, out_channels, 1, bias=False)
-
-    def forward(self, x):
-        x = self.pre_conv(x)
-        if self.downsample:
-            x = self.conv(x) + self.conv1by1(x)
-        else:
-            x = self.conv(x) + x
-        return x

Utils/PLBERT/config.yml

@@ -1,30 +0,0 @@
-log_dir: "Checkpoint"
-mixed_precision: "fp16"
-data_folder: "wikipedia_20220301.en.processed"
-batch_size: 192
-save_interval: 5000
-log_interval: 10
-num_process: 1 # number of GPUs
-num_steps: 1000000
-
-dataset_params:
-    tokenizer: "transfo-xl-wt103"
-    token_separator: " " # token used for phoneme separator (space)
-    token_mask: "M" # token used for phoneme mask (M)
-    word_separator: 3039 # token used for word separator (<formula>)
-    token_maps: "token_maps.pkl" # token map path
-    
-    max_mel_length: 512 # max phoneme length
-    
-    word_mask_prob: 0.15 # probability to mask the entire word
-    phoneme_mask_prob: 0.1 # probability to mask each phoneme
-    replace_prob: 0.2 # probablity to replace phonemes
-    
-model_params:
-    vocab_size: 198
-    hidden_size: 768
-    num_attention_heads: 12
-    intermediate_size: 2048
-    max_position_embeddings: 512
-    num_hidden_layers: 12
-    dropout: 0.1

Utils/PLBERT/util.py

@@ -1,24 +0,0 @@
-import os
-import yaml
-import torch
-from transformers import AlbertConfig, AlbertModel
-
-class CustomAlbert(AlbertModel):
-    def forward(self, *args, **kwargs):
-        # Call the original forward method
-        outputs = super().forward(*args, **kwargs)
-
-        # Only return the last_hidden_state
-        return outputs.last_hidden_state
-
-
-def load_plbert(wights_path, config_path):
-    plbert_config = yaml.safe_load(open(config_path))
-    
-    albert_base_configuration = AlbertConfig(**plbert_config['model_params'])
-    bert = CustomAlbert(albert_base_configuration)
-
-    state_dict = torch.load(wights_path, map_location='cpu')
-    bert.load_state_dict(state_dict, strict=False)
-    
-    return bert

Utils/__init__.py

@@ -1 +0,0 @@
-

models.py

@@ -1,717 +0,0 @@
-#coding:utf-8
-
-import os
-import os.path as osp
-
-import copy
-import math
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
-
-from Utils.ASR.models import ASRCNN
-from Utils.JDC.model import JDCNet
-
-from Modules.diffusion.sampler import KDiffusion, LogNormalDistribution, DiffusionSampler, ADPM2Sampler, KarrasSchedule
-from Modules.diffusion.modules import Transformer1d, StyleTransformer1d
-from Modules.diffusion.diffusion import AudioDiffusionConditional
-
-from Modules.discriminators import MultiPeriodDiscriminator, MultiResSpecDiscriminator, WavLMDiscriminator
-
-from munch import Munch
-import yaml
-
-class LearnedDownSample(nn.Module):
-    def __init__(self, layer_type, dim_in):
-        super().__init__()
-        self.layer_type = layer_type
-
-        if self.layer_type == 'none':
-            self.conv = nn.Identity()
-        elif self.layer_type == 'timepreserve':
-            self.conv = spectral_norm(nn.Conv2d(dim_in, dim_in, kernel_size=(3, 1), stride=(2, 1), groups=dim_in, padding=(1, 0)))
-        elif self.layer_type == 'half':
-            self.conv = spectral_norm(nn.Conv2d(dim_in, dim_in, kernel_size=(3, 3), stride=(2, 2), groups=dim_in, padding=1))
-        else:
-            raise RuntimeError('Got unexpected donwsampletype %s, expected is [none, timepreserve, half]' % self.layer_type)
-            
-    def forward(self, x):
-        return self.conv(x)
-
-class LearnedUpSample(nn.Module):
-    def __init__(self, layer_type, dim_in):
-        super().__init__()
-        self.layer_type = layer_type
-        
-        if self.layer_type == 'none':
-            self.conv = nn.Identity()
-        elif self.layer_type == 'timepreserve':
-            self.conv = nn.ConvTranspose2d(dim_in, dim_in, kernel_size=(3, 1), stride=(2, 1), groups=dim_in, output_padding=(1, 0), padding=(1, 0))
-        elif self.layer_type == 'half':
-            self.conv = nn.ConvTranspose2d(dim_in, dim_in, kernel_size=(3, 3), stride=(2, 2), groups=dim_in, output_padding=1, padding=1)
-        else:
-            raise RuntimeError('Got unexpected upsampletype %s, expected is [none, timepreserve, half]' % self.layer_type)
-
-
-    def forward(self, x):
-        return self.conv(x)
-
-class DownSample(nn.Module):
-    def __init__(self, layer_type):
-        super().__init__()
-        self.layer_type = layer_type
-
-    def forward(self, x):
-        if self.layer_type == 'none':
-            return x
-        elif self.layer_type == 'timepreserve':
-            return F.avg_pool2d(x, (2, 1))
-        elif self.layer_type == 'half':
-            if x.shape[-1] % 2 != 0:
-                x = torch.cat([x, x[..., -1].unsqueeze(-1)], dim=-1)
-            return F.avg_pool2d(x, 2)
-        else:
-            raise RuntimeError('Got unexpected donwsampletype %s, expected is [none, timepreserve, half]' % self.layer_type)
-
-
-class UpSample(nn.Module):
-    def __init__(self, layer_type):
-        super().__init__()
-        self.layer_type = layer_type
-
-    def forward(self, x):
-        if self.layer_type == 'none':
-            return x
-        elif self.layer_type == 'timepreserve':
-            return F.interpolate(x, scale_factor=(2, 1), mode='nearest')
-        elif self.layer_type == 'half':
-            return F.interpolate(x, scale_factor=2, mode='nearest')
-        else:
-            raise RuntimeError('Got unexpected upsampletype %s, expected is [none, timepreserve, half]' % self.layer_type)
-
-
-class ResBlk(nn.Module):
-    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
-                 normalize=False, downsample='none'):
-        super().__init__()
-        self.actv = actv
-        self.normalize = normalize
-        self.downsample = DownSample(downsample)
-        self.downsample_res = LearnedDownSample(downsample, dim_in)
-        self.learned_sc = dim_in != dim_out
-        self._build_weights(dim_in, dim_out)
-
-    def _build_weights(self, dim_in, dim_out):
-        self.conv1 = spectral_norm(nn.Conv2d(dim_in, dim_in, 3, 1, 1))
-        self.conv2 = spectral_norm(nn.Conv2d(dim_in, dim_out, 3, 1, 1))
-        if self.normalize:
-            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
-            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
-        if self.learned_sc:
-            self.conv1x1 = spectral_norm(nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False))
-
-    def _shortcut(self, x):
-        if self.learned_sc:
-            x = self.conv1x1(x)
-        if self.downsample:
-            x = self.downsample(x)
-        return x
-
-    def _residual(self, x):
-        if self.normalize:
-            x = self.norm1(x)
-        x = self.actv(x)
-        x = self.conv1(x)
-        x = self.downsample_res(x)
-        if self.normalize:
-            x = self.norm2(x)
-        x = self.actv(x)
-        x = self.conv2(x)
-        return x
-
-    def forward(self, x):
-        x = self._shortcut(x) + self._residual(x)
-        return x / math.sqrt(2)  # unit variance
-
-class StyleEncoder(nn.Module):
-    def __init__(self, dim_in=48, style_dim=48, max_conv_dim=384):
-        super().__init__()
-        blocks = []
-        blocks += [spectral_norm(nn.Conv2d(1, dim_in, 3, 1, 1))]
-
-        repeat_num = 4
-        for _ in range(repeat_num):
-            dim_out = min(dim_in*2, max_conv_dim)
-            blocks += [ResBlk(dim_in, dim_out, downsample='half')]
-            dim_in = dim_out
-
-        blocks += [nn.LeakyReLU(0.2)]
-        blocks += [spectral_norm(nn.Conv2d(dim_out, dim_out, 5, 1, 0))]
-        blocks += [nn.AdaptiveAvgPool2d(1)]
-        blocks += [nn.LeakyReLU(0.2)]
-        self.shared = nn.Sequential(*blocks)
-
-        self.unshared = nn.Linear(dim_out, style_dim)
-
-    def forward(self, x):
-        h = self.shared(x)
-        h = h.view(h.size(0), -1)
-        s = self.unshared(h)
-    
-        return s
-
-class LinearNorm(torch.nn.Module):
-    def __init__(self, in_dim, out_dim, bias=True, w_init_gain='linear'):
-        super(LinearNorm, self).__init__()
-        self.linear_layer = torch.nn.Linear(in_dim, out_dim, bias=bias)
-
-        torch.nn.init.xavier_uniform_(
-            self.linear_layer.weight,
-            gain=torch.nn.init.calculate_gain(w_init_gain))
-
-    def forward(self, x):
-        return self.linear_layer(x)
-
-class Discriminator2d(nn.Module):
-    def __init__(self, dim_in=48, num_domains=1, max_conv_dim=384, repeat_num=4):
-        super().__init__()
-        blocks = []
-        blocks += [spectral_norm(nn.Conv2d(1, dim_in, 3, 1, 1))]
-
-        for lid in range(repeat_num):
-            dim_out = min(dim_in*2, max_conv_dim)
-            blocks += [ResBlk(dim_in, dim_out, downsample='half')]
-            dim_in = dim_out
-
-        blocks += [nn.LeakyReLU(0.2)]
-        blocks += [spectral_norm(nn.Conv2d(dim_out, dim_out, 5, 1, 0))]
-        blocks += [nn.LeakyReLU(0.2)]
-        blocks += [nn.AdaptiveAvgPool2d(1)]
-        blocks += [spectral_norm(nn.Conv2d(dim_out, num_domains, 1, 1, 0))]
-        self.main = nn.Sequential(*blocks)
-
-    def get_feature(self, x):
-        features = []
-        for l in self.main:
-            x = l(x)
-            features.append(x) 
-        out = features[-1]
-        out = out.view(out.size(0), -1)  # (batch, num_domains)
-        return out, features
-
-    def forward(self, x):
-        out, features = self.get_feature(x)
-        out = out.squeeze()  # (batch)
-        return out, features
-
-class ResBlk1d(nn.Module):
-    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
-                 normalize=False, downsample='none', dropout_p=0.2):
-        super().__init__()
-        self.actv = actv
-        self.normalize = normalize
-        self.downsample_type = downsample
-        self.learned_sc = dim_in != dim_out
-        self._build_weights(dim_in, dim_out)
-        self.dropout_p = dropout_p
-        
-        if self.downsample_type == 'none':
-            self.pool = nn.Identity()
-        else:
-            self.pool = weight_norm(nn.Conv1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1))
-
-    def _build_weights(self, dim_in, dim_out):
-        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_in, 3, 1, 1))
-        self.conv2 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
-        if self.normalize:
-            self.norm1 = nn.InstanceNorm1d(dim_in, affine=True)
-            self.norm2 = nn.InstanceNorm1d(dim_in, affine=True)
-        if self.learned_sc:
-            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
-
-    def downsample(self, x):
-        if self.downsample_type == 'none':
-            return x
-        else:
-            if x.shape[-1] % 2 != 0:
-                x = torch.cat([x, x[..., -1].unsqueeze(-1)], dim=-1)
-            return F.avg_pool1d(x, 2)
-
-    def _shortcut(self, x):
-        if self.learned_sc:
-            x = self.conv1x1(x)
-        x = self.downsample(x)
-        return x
-
-    def _residual(self, x):
-        if self.normalize:
-            x = self.norm1(x)
-        x = self.actv(x)
-        x = F.dropout(x, p=self.dropout_p, training=self.training)
-        
-        x = self.conv1(x)
-        x = self.pool(x)
-        if self.normalize:
-            x = self.norm2(x)
-            
-        x = self.actv(x)
-        x = F.dropout(x, p=self.dropout_p, training=self.training)
-        
-        x = self.conv2(x)
-        return x
-
-    def forward(self, x):
-        x = self._shortcut(x) + self._residual(x)
-        return x / math.sqrt(2)  # unit variance
-
-class LayerNorm(nn.Module):
-    def __init__(self, channels, eps=1e-5):
-        super().__init__()
-        self.channels = channels
-        self.eps = eps
-
-        self.gamma = nn.Parameter(torch.ones(channels))
-        self.beta = nn.Parameter(torch.zeros(channels))
-
-    def forward(self, x):
-        x = x.transpose(1, -1)
-        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
-        return x.transpose(1, -1)
-    
-class TextEncoder(nn.Module):
-    def __init__(self, channels, kernel_size, depth, n_symbols, actv=nn.LeakyReLU(0.2)):
-        super().__init__()
-        self.embedding = nn.Embedding(n_symbols, channels)
-
-        padding = (kernel_size - 1) // 2
-        self.cnn = nn.ModuleList()
-        for _ in range(depth):
-            self.cnn.append(nn.Sequential(
-                weight_norm(nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding)),
-                LayerNorm(channels),
-                actv,
-                nn.Dropout(0.2),
-            ))
-        # self.cnn = nn.Sequential(*self.cnn)
-
-        self.lstm = nn.LSTM(channels, channels//2, 1, batch_first=True, bidirectional=True)
-
-    def forward(self, x, input_lengths, m):
-        x = self.embedding(x)  # [B, T, emb]
-        x = x.transpose(1, 2)  # [B, emb, T]
-        m = m.to(input_lengths.device).unsqueeze(1)
-        x.masked_fill_(m, 0.0)
-        
-        for c in self.cnn:
-            x = c(x)
-            x.masked_fill_(m, 0.0)
-            
-        x = x.transpose(1, 2)  # [B, T, chn]
-
-        input_lengths = input_lengths.cpu().numpy()
-        x = nn.utils.rnn.pack_padded_sequence(
-            x, input_lengths, batch_first=True, enforce_sorted=False)
-
-        self.lstm.flatten_parameters()
-        x, _ = self.lstm(x)
-        x, _ = nn.utils.rnn.pad_packed_sequence(
-            x, batch_first=True)
-                
-        x = x.transpose(-1, -2)
-        x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]])
-
-        x_pad[:, :, :x.shape[-1]] = x
-        x = x_pad.to(x.device)
-        
-        x.masked_fill_(m, 0.0)
-        
-        return x
-
-    def inference(self, x):
-        x = self.embedding(x)
-        x = x.transpose(1, 2)
-        x = self.cnn(x)
-        x = x.transpose(1, 2)
-        self.lstm.flatten_parameters()
-        x, _ = self.lstm(x)
-        return x
-    
-    def length_to_mask(self, lengths):
-        mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
-        mask = torch.gt(mask+1, lengths.unsqueeze(1))
-        return mask
-
-
-
-class AdaIN1d(nn.Module):
-    def __init__(self, style_dim, num_features):
-        super().__init__()
-        self.norm = nn.InstanceNorm1d(num_features, affine=False)
-        self.fc = nn.Linear(style_dim, num_features*2)
-
-    def forward(self, x, s):
-        h = self.fc(s)
-        h = h.view(h.size(0), h.size(1), 1)
-        gamma, beta = torch.chunk(h, chunks=2, dim=1)
-        return (1 + gamma) * self.norm(x) + beta
-
-class UpSample1d(nn.Module):
-    def __init__(self, layer_type):
-        super().__init__()
-        self.layer_type = layer_type
-
-    def forward(self, x):
-        if self.layer_type == 'none':
-            return x
-        else:
-            return F.interpolate(x, scale_factor=2, mode='nearest')
-
-class AdainResBlk1d(nn.Module):
-    def __init__(self, dim_in, dim_out, style_dim=64, actv=nn.LeakyReLU(0.2),
-                 upsample='none', dropout_p=0.0):
-        super().__init__()
-        self.actv = actv
-        self.upsample_type = upsample
-        self.upsample = UpSample1d(upsample)
-        self.learned_sc = dim_in != dim_out
-        self._build_weights(dim_in, dim_out, style_dim)
-        self.dropout = nn.Dropout(dropout_p)
-        
-        if upsample == 'none':
-            self.pool = nn.Identity()
-        else:
-            self.pool = weight_norm(nn.ConvTranspose1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1, output_padding=1))
-        
-        
-    def _build_weights(self, dim_in, dim_out, style_dim):
-        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
-        self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1))
-        self.norm1 = AdaIN1d(style_dim, dim_in)
-        self.norm2 = AdaIN1d(style_dim, dim_out)
-        if self.learned_sc:
-            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
-
-    def _shortcut(self, x):
-        x = self.upsample(x)
-        if self.learned_sc:
-            x = self.conv1x1(x)
-        return x
-
-    def _residual(self, x, s):
-        x = self.norm1(x, s)
-        x = self.actv(x)
-        x = self.pool(x)
-        x = self.conv1(self.dropout(x))
-        x = self.norm2(x, s)
-        x = self.actv(x)
-        x = self.conv2(self.dropout(x))
-        return x
-
-    def forward(self, x, s):
-        out = self._residual(x, s)
-        out = (out + self._shortcut(x)) / math.sqrt(2)
-        return out
-    
-class AdaLayerNorm(nn.Module):
-    def __init__(self, style_dim, channels, eps=1e-5):
-        super().__init__()
-        self.channels = channels
-        self.eps = eps
-
-        self.fc = nn.Linear(style_dim, channels*2)
-
-    def forward(self, x, s):
-        x = x.transpose(-1, -2)
-        x = x.transpose(1, -1)
-                
-        h = self.fc(s)
-        h = h.view(h.size(0), h.size(1), 1)
-        gamma, beta = torch.chunk(h, chunks=2, dim=1)
-        gamma, beta = gamma.transpose(1, -1), beta.transpose(1, -1)
-        
-        
-        x = F.layer_norm(x, (self.channels,), eps=self.eps)
-        x = (1 + gamma) * x + beta
-        return x.transpose(1, -1).transpose(-1, -2)
-
-class ProsodyPredictor(nn.Module):
-
-    def __init__(self, style_dim, d_hid, nlayers, max_dur=50, dropout=0.1):
-        super().__init__() 
-        
-        self.text_encoder = DurationEncoder(sty_dim=style_dim, 
-                                            d_model=d_hid,
-                                            nlayers=nlayers, 
-                                            dropout=dropout)
-
-        self.lstm = nn.LSTM(d_hid + style_dim, d_hid // 2, 1, batch_first=True, bidirectional=True)
-        self.duration_proj = LinearNorm(d_hid, max_dur)
-        
-        self.shared = nn.LSTM(d_hid + style_dim, d_hid // 2, 1, batch_first=True, bidirectional=True)
-        self.F0 = nn.ModuleList()
-        self.F0.append(AdainResBlk1d(d_hid, d_hid, style_dim, dropout_p=dropout))
-        self.F0.append(AdainResBlk1d(d_hid, d_hid // 2, style_dim, upsample=True, dropout_p=dropout))
-        self.F0.append(AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim, dropout_p=dropout))
-
-        self.N = nn.ModuleList()
-        self.N.append(AdainResBlk1d(d_hid, d_hid, style_dim, dropout_p=dropout))
-        self.N.append(AdainResBlk1d(d_hid, d_hid // 2, style_dim, upsample=True, dropout_p=dropout))
-        self.N.append(AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim, dropout_p=dropout))
-        
-        self.F0_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0)
-        self.N_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0)
-
-
-    def forward(self, texts, style, text_lengths, alignment, m):
-        d = self.text_encoder(texts, style, text_lengths, m)
-        
-        batch_size = d.shape[0]
-        text_size = d.shape[1]
-        
-        # predict duration
-        input_lengths = text_lengths.cpu().numpy()
-        x = nn.utils.rnn.pack_padded_sequence(
-            d, input_lengths, batch_first=True, enforce_sorted=False)
-        
-        m = m.to(text_lengths.device).unsqueeze(1)
-        
-        self.lstm.flatten_parameters()
-        x, _ = self.lstm(x)
-        x, _ = nn.utils.rnn.pad_packed_sequence(
-            x, batch_first=True)
-        
-        x_pad = torch.zeros([x.shape[0], m.shape[-1], x.shape[-1]])
-
-        x_pad[:, :x.shape[1], :] = x
-        x = x_pad.to(x.device)
-                
-        duration = self.duration_proj(nn.functional.dropout(x, 0.5, training=self.training))
-        
-        en = (d.transpose(-1, -2) @ alignment)
-
-        return duration.squeeze(-1), en
-    
-    def F0Ntrain(self, x, s):
-        x, _ = self.shared(x.transpose(-1, -2))
-        
-        F0 = x.transpose(-1, -2)
-        for block in self.F0:
-            F0 = block(F0, s)
-        F0 = self.F0_proj(F0)
-
-        N = x.transpose(-1, -2)
-        for block in self.N:
-            N = block(N, s)
-        N = self.N_proj(N)
-        
-        return F0.squeeze(1), N.squeeze(1)
-    
-    def length_to_mask(self, lengths):
-        mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
-        mask = torch.gt(mask+1, lengths.unsqueeze(1))
-        return mask
-    
-class DurationEncoder(nn.Module):
-
-    def __init__(self, sty_dim, d_model, nlayers, dropout=0.1):
-        super().__init__()
-        self.lstms = nn.ModuleList()
-        for _ in range(nlayers):
-            self.lstms.append(nn.LSTM(d_model + sty_dim, 
-                                 d_model // 2, 
-                                 num_layers=1, 
-                                 batch_first=True, 
-                                 bidirectional=True, 
-                                 dropout=dropout))
-            self.lstms.append(AdaLayerNorm(sty_dim, d_model))
-        
-        
-        self.dropout = dropout
-        self.d_model = d_model
-        self.sty_dim = sty_dim
-
-    def forward(self, x, style, text_lengths, m):
-        masks = m.to(text_lengths.device)
-        
-        x = x.permute(2, 0, 1)
-        s = style.expand(x.shape[0], x.shape[1], -1)
-        x = torch.cat([x, s], axis=-1)
-        x.masked_fill_(masks.unsqueeze(-1).transpose(0, 1), 0.0)
-                
-        x = x.transpose(0, 1)
-        input_lengths = text_lengths.cpu().numpy()
-        x = x.transpose(-1, -2)
-        
-        for block in self.lstms:
-            if isinstance(block, AdaLayerNorm):
-                x = block(x.transpose(-1, -2), style).transpose(-1, -2)
-                x = torch.cat([x, s.permute(1, -1, 0)], axis=1)
-                x.masked_fill_(masks.unsqueeze(-1).transpose(-1, -2), 0.0)
-            else:
-                x = x.transpose(-1, -2)
-                x = nn.utils.rnn.pack_padded_sequence(
-                    x, input_lengths, batch_first=True, enforce_sorted=False)
-                block.flatten_parameters()
-                x, _ = block(x)
-                x, _ = nn.utils.rnn.pad_packed_sequence(
-                    x, batch_first=True)
-                x = F.dropout(x, p=self.dropout, training=self.training)
-                x = x.transpose(-1, -2)
-                
-                x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]])
-
-                x_pad[:, :, :x.shape[-1]] = x
-                x = x_pad.to(x.device)
-        
-        return x.transpose(-1, -2)
-    
-    def inference(self, x, style):
-        x = self.embedding(x.transpose(-1, -2)) * math.sqrt(self.d_model)
-        style = style.expand(x.shape[0], x.shape[1], -1)
-        x = torch.cat([x, style], axis=-1)
-        src = self.pos_encoder(x)
-        output = self.transformer_encoder(src).transpose(0, 1)
-        return output
-    
-    def length_to_mask(self, lengths):
-        mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
-        mask = torch.gt(mask+1, lengths.unsqueeze(1))
-        return mask
-    
-def load_F0_models(path):
-    # load F0 model
-
-    F0_model = JDCNet(num_class=1, seq_len=192)
-    params = torch.load(path, map_location='cpu')
-    F0_model.load_state_dict(params)
-    _ = F0_model.train()
-    
-    return F0_model
-
-def load_ASR_models(ASR_MODEL_PATH, ASR_MODEL_CONFIG):
-    # load ASR model
-    def _load_config(path):
-        with open(path) as f:
-            config = yaml.safe_load(f)
-        model_config = config['model_params']
-        return model_config
-
-    def _load_model(model_config, model_path):
-        model = ASRCNN(**model_config)
-        params = torch.load(model_path, map_location='cpu')
-        model.load_state_dict(params)
-        return model
-
-    asr_model_config = _load_config(ASR_MODEL_CONFIG)
-    asr_model = _load_model(asr_model_config, ASR_MODEL_PATH)
-    _ = asr_model.train()
-
-    return asr_model
-
-def build_model(args, text_aligner, pitch_extractor, bert):
-    assert args.decoder.type in ['istftnet', 'hifigan'], 'Decoder type unknown'
-    
-    if args.decoder.type == "istftnet":
-        from Modules.istftnet import Decoder
-        decoder = Decoder(dim_in=args.hidden_dim, style_dim=args.style_dim, dim_out=args.n_mels,
-                resblock_kernel_sizes = args.decoder.resblock_kernel_sizes,
-                upsample_rates = args.decoder.upsample_rates,
-                upsample_initial_channel=args.decoder.upsample_initial_channel,
-                resblock_dilation_sizes=args.decoder.resblock_dilation_sizes,
-                upsample_kernel_sizes=args.decoder.upsample_kernel_sizes, 
-                gen_istft_n_fft=args.decoder.gen_istft_n_fft, gen_istft_hop_size=args.decoder.gen_istft_hop_size) 
-    else:
-        from Modules.hifigan import Decoder
-        decoder = Decoder(dim_in=args.hidden_dim, style_dim=args.style_dim, dim_out=args.n_mels,
-                resblock_kernel_sizes = args.decoder.resblock_kernel_sizes,
-                upsample_rates = args.decoder.upsample_rates,
-                upsample_initial_channel=args.decoder.upsample_initial_channel,
-                resblock_dilation_sizes=args.decoder.resblock_dilation_sizes,
-                upsample_kernel_sizes=args.decoder.upsample_kernel_sizes) 
-        
-    text_encoder = TextEncoder(channels=args.hidden_dim, kernel_size=5, depth=args.n_layer, n_symbols=args.n_token)
-    
-    predictor = ProsodyPredictor(style_dim=args.style_dim, d_hid=args.hidden_dim, nlayers=args.n_layer, max_dur=args.max_dur, dropout=args.dropout)
-    
-    style_encoder = StyleEncoder(dim_in=args.dim_in, style_dim=args.style_dim, max_conv_dim=args.hidden_dim) # acoustic style encoder
-    predictor_encoder = StyleEncoder(dim_in=args.dim_in, style_dim=args.style_dim, max_conv_dim=args.hidden_dim) # prosodic style encoder
-        
-    # define diffusion model
-    if args.multispeaker:
-        transformer = StyleTransformer1d(channels=args.style_dim*2, 
-                                    context_embedding_features=bert.config.hidden_size,
-                                    context_features=args.style_dim*2, 
-                                    **args.diffusion.transformer)
-    else:
-        transformer = Transformer1d(channels=args.style_dim*2, 
-                                    context_embedding_features=bert.config.hidden_size,
-                                    **args.diffusion.transformer)
-    
-    diffusion = AudioDiffusionConditional(
-        in_channels=1,
-        embedding_max_length=bert.config.max_position_embeddings,
-        embedding_features=bert.config.hidden_size,
-        embedding_mask_proba=args.diffusion.embedding_mask_proba, # Conditional dropout of batch elements,
-        channels=args.style_dim*2,
-        context_features=args.style_dim*2,
-    )
-    
-    diffusion.diffusion = KDiffusion(
-        net=diffusion.unet,
-        sigma_distribution=LogNormalDistribution(mean = args.diffusion.dist.mean, std = args.diffusion.dist.std),
-        sigma_data=args.diffusion.dist.sigma_data, # a placeholder, will be changed dynamically when start training diffusion model
-        dynamic_threshold=0.0 
-    )
-    diffusion.diffusion.net = transformer
-    diffusion.unet = transformer
-
-    
-    nets = Munch(
-            bert=bert,
-            bert_encoder=nn.Linear(bert.config.hidden_size, args.hidden_dim),
-
-            predictor=predictor,
-            decoder=decoder,
-            text_encoder=text_encoder,
-
-            predictor_encoder=predictor_encoder,
-            style_encoder=style_encoder,
-            diffusion=diffusion,
-
-            text_aligner = text_aligner,
-            pitch_extractor=pitch_extractor,
-
-            mpd = MultiPeriodDiscriminator(),
-            msd = MultiResSpecDiscriminator(),
-        
-            # slm discriminator head
-            wd = WavLMDiscriminator(args.slm.hidden, args.slm.nlayers, args.slm.initial_channel),
-            sampler = DiffusionSampler(diffusion.diffusion,
-                                       sampler=ADPM2Sampler(),
-                                       sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0),
-                                       clamp=False )
-       )
-    
-    return nets
-
-def load_checkpoint(model, optimizer, path, load_only_params=True, ignore_modules=[]):
-    state = torch.load(path, map_location='cpu')
-    params = state['net']
-    for key in model:
-        if key in params and key not in ignore_modules:
-            print('%s loaded' % key)
-            model[key].load_state_dict(params[key], strict=False)
-    _ = [model[key].eval() for key in model]
-    
-    if not load_only_params:
-        epoch = state["epoch"]
-        iters = state["iters"]
-        optimizer.load_state_dict(state["optimizer"])
-    else:
-        epoch = 0
-        iters = 0
-        
-    return model, optimizer, epoch, iters

weights/asr.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ee69a5f32c76aff88bafb09f31379dc625e51e3f71d842a4cf30ecd69cb92b56
-size 31529032

weights/jdc.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:87aabb4a35814b581d13f4cc6ee352e99e80ede7fb6e7e963ea85b1feee940ac
-size 21023821

weights/plbert.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:74284b3d86962b0c4b52f2d3fe8507fc1ad87ae1defd1e7ddf28677380a290ef
-size 25188799