Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+RingFormer/Utils/JDC/__pycache__/bst.t7 filter=lfs diff=lfs merge=lfs -text
+RingFormer/Utils/JDC/bst.t7 filter=lfs diff=lfs merge=lfs -text

RingFormer/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Aaron (Yinghao) Li
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

RingFormer/README.md

@@ -0,0 +1,38 @@
+# HiFTNet: A Fast High-Quality Neural Vocoder with Harmonic-plus-Noise Filter and Inverse Short Time Fourier Transform
+
+### Yinghao Aaron Li, Cong Han, Xilin Jiang, Nima Mesgarani
+
+> Recent advancements in speech synthesis have leveraged GAN-based networks like HiFi-GAN and BigVGAN to produce high-fidelity waveforms from mel-spectrograms. However, these networks are computationally expensive and parameter-heavy. iSTFTNet addresses these limitations by integrating inverse short-time Fourier transform (iSTFT) into the network, achieving both speed and parameter efficiency. In this paper, we introduce an extension to iSTFTNet, termed HiFTNet, which incorporates a harmonic-plus-noise source filter in the time-frequency domain that uses a sinusoidal source from the fundamental frequency (F0) inferred via a pre-trained F0 estimation network for fast inference speed. Subjective evaluations on LJSpeech show that our model significantly outperforms both iSTFTNet and HiFi-GAN, achieving ground-truth-level performance. HiFTNet also outperforms BigVGAN-base on LibriTTS for unseen speakers and achieves comparable performance to BigVGAN while being four times faster with only 1/6 of the parameters. Our work sets a new benchmark for efficient, high-quality neural vocoding, paving the way for real-time applications that demand high quality speech synthesis.
+
+Paper: [https://arxiv.org/abs/2309.09493](https://arxiv.org/abs/2309.09493)
+
+Audio samples: [https://hiftnet.github.io/](https://hiftnet.github.io/)
+
+**Check our TTS work that uses HiFTNet as speech decoder for human-level speech synthesis here: https://github.com/yl4579/StyleTTS2**
+
+## Pre-requisites
+1. Python >= 3.7
+2. Clone this repository:
+```bash
+git clone https://github.com/yl4579/HiFTNet.git
+cd HiFTNet
+```
+3. Install python requirements: 
+```bash
+pip install -r requirements.txt
+```
+
+## Training
+```bash
+python train.py --config config_v1.json --[args]
+```
+For the F0 model training, please refer to [yl4579/PitchExtractor](https://github.com/yl4579/PitchExtractor). This repo includes a pre-trained F0 model on LibriTTS. Still, you may want to train your own F0 model for the best performance, particularly for noisy or non-speech data, as we found that F0 estimation accuracy is essential for the vocoder performance. 
+
+## Inference
+Please refer to the notebook [inference.ipynb](https://github.com/yl4579/HiFTNet/blob/main/inference.ipynb) for details.
+### Pre-Trained Models
+You can download the pre-trained LJSpeech model [here](https://huggingface.co/yl4579/HiFTNet/blob/main/LJSpeech/cp_hifigan.zip) and the pre-trained LibriTTS model [here](https://huggingface.co/yl4579/HiFTNet/blob/main/LibriTTS/cp_hifigan.zip). The pre-trained models contain parameters of the optimizers and discriminators that can be used for fine-tuning.  
+
+## References
+- [rishikksh20/iSTFTNet-pytorch](https://github.com/rishikksh20/iSTFTNet-pytorch)
+- [nii-yamagishilab/project-NN-Pytorch-scripts/project/01-nsf](https://github.com/nii-yamagishilab/project-NN-Pytorch-scripts/tree/master/project/01-nsf)

RingFormer/Utils/JDC/__init__.py

@@ -0,0 +1 @@
+

RingFormer/Utils/JDC/__pycache__/bst.t7

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

RingFormer/Utils/JDC/bst.t7

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

RingFormer/Utils/JDC/model.py

@@ -0,0 +1,192 @@
+"""
+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, 64, 31, 128) from conv_block, out = (b, 64, 31, 4)
+        self.maxpool1 = nn.MaxPool2d(kernel_size=(1, 32))  # 128 / 32 = 4
+        # in = (b, 128, 31, 64) from res_block1, out = (b, 128, 31, 4)
+        self.maxpool2 = nn.MaxPool2d(kernel_size=(1, 16))  # 64 / 16 = 4
+        # in = (b, 192, 31, 32) from res_block2, out = (b, 192, 31, 4)
+        self.maxpool3 = nn.MaxPool2d(kernel_size=(1, 8))  # 32 / 8 = 4
+
+
+        # 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=1024, 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=1024, 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, 1024))
+        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

RingFormer/Utils/__init__.py

@@ -0,0 +1 @@
+

RingFormer/config_v1.json

@@ -0,0 +1,42 @@
+{
+    "F0_path": "/home/ubuntu/Darya_Speech/Utils/JDC/epoch_00079.pth",
+    
+    "resblock": "1",
+    "num_gpus": 1,
+    "batch_size": 18,
+    "learning_rate": 0.0002,
+    "adam_b1": 0.8,
+    "adam_b2": 0.99,
+    "lr_decay": 0.999,
+    "seed": 1234,
+
+
+    "upsample_rates": [16,8],
+    "upsample_kernel_sizes": [32, 16],
+    "upsample_initial_channel": 512,
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+    "gen_istft_n_fft": 32,
+    "gen_istft_hop_size": 4,
+
+
+    "segment_size": 65536,
+    "num_mels": 128,
+    "n_fft": 2048,
+    "hop_size": 512,
+    "win_size": 2048,
+
+    "sampling_rate": 44100,
+
+    "fmin": 0,
+    "fmax": null,
+    "fmax_for_loss": null,
+
+    "num_workers": 8,
+
+    "dist_config": {
+        "dist_backend": "nccl",
+        "dist_url": "tcp://localhost:54321",
+        "world_size": 1
+    }
+}

RingFormer/conformer.py

@@ -0,0 +1,228 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from einops.layers.torch import Rearrange
+from ring_attention_pytorch import RingAttention
+
+# helper functions
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+def calc_same_padding(kernel_size):
+    pad = kernel_size // 2
+    return (pad, pad - (kernel_size + 1) % 2)
+
+
+# helper classes
+
+
+class Swish(nn.Module):
+    def forward(self, x):
+        return x * x.sigmoid()
+
+
+class GLU(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        out, gate = x.chunk(2, dim=self.dim)
+        return out * gate.sigmoid()
+
+
+class DepthWiseConv1d(nn.Module):
+    def __init__(self, chan_in, chan_out, kernel_size, padding):
+        super().__init__()
+        self.padding = padding
+        self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, groups=chan_in)
+
+    def forward(self, x):
+        x = F.pad(x, self.padding)
+        return self.conv(x)
+
+
+# attention, feedforward, and conv module
+
+
+class Scale(nn.Module):
+    def __init__(self, scale, fn):
+        super().__init__()
+        self.fn = fn
+        self.scale = scale
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) * self.scale
+
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x, **kwargs):
+
+        x = self.norm(x.to(x.device))
+
+        out = self.fn(x.to(x.device), **kwargs)
+   
+        return out
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, mult=4, dropout=0.0):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, dim * mult),
+            Swish(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * mult, dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class ConformerConvModule(nn.Module):
+    def __init__(
+        self, dim, causal=False, expansion_factor=2, kernel_size=31, dropout=0.0
+    ):
+        super().__init__()
+
+        inner_dim = dim * expansion_factor
+        padding = calc_same_padding(kernel_size) if not causal else (kernel_size - 1, 0)
+
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            Rearrange("b n c -> b c n"),
+            nn.Conv1d(dim, inner_dim * 2, 1),
+            GLU(dim=1),
+            DepthWiseConv1d(
+                inner_dim, inner_dim, kernel_size=kernel_size, padding=padding
+            ),
+            nn.BatchNorm1d(inner_dim) if not causal else nn.Identity(),
+            Swish(),
+            nn.Conv1d(inner_dim, dim, 1),
+            Rearrange("b c n -> b n c"),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# Conformer Block
+
+
+class ConformerBlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        dim_head=64,
+        heads=8,
+        ff_mult=4,
+        conv_expansion_factor=2,
+        conv_kernel_size=31,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        conv_dropout=0.0,
+        conv_causal=False
+    ):
+        super().__init__()
+        self.ff1 = FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)
+        self.attn = RingAttention(
+            dim=dim,
+            dim_head=dim_head,
+            heads=heads,
+            causal=True,
+            auto_shard_seq=False, # doesn't work on multi-gpu setup for some reason
+            ring_attn=True,
+            ring_seq_size=512,
+        )
+        self.self_attn_dropout = torch.nn.Dropout(attn_dropout)
+        self.conv = ConformerConvModule(
+            dim=dim,
+            causal=conv_causal,
+            expansion_factor=conv_expansion_factor,
+            kernel_size=conv_kernel_size,
+            dropout=conv_dropout,
+        )
+        self.ff2 = FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)
+
+        self.attn = PreNorm(dim, self.attn)
+        self.ff1 = Scale(0.5, PreNorm(dim, self.ff1))
+        self.ff2 = Scale(0.5, PreNorm(dim, self.ff2))
+
+        self.post_norm = nn.LayerNorm(dim)
+
+
+    def forward(self, x, mask=None):
+        x_ff1 = self.ff1(x) + x
+        
+        x = self.attn(x, mask=mask)
+        x = self.self_attn_dropout(x)
+        x = x + x_ff1
+        x = self.conv(x) + x
+        x = self.ff2(x) + x
+        x = self.post_norm(x)
+        return x
+
+
+
+# Conformer
+
+
+class Conformer(nn.Module):
+    def __init__(
+        self,
+
+        dim,
+        *,
+        depth,
+        dim_head=64,
+        heads=8,
+        ff_mult=4,
+        conv_expansion_factor=2,
+        conv_kernel_size=31,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        conv_dropout=0.0,
+        conv_causal=False
+    ):
+        super().__init__()
+        self.dim = dim
+
+        self.layers = nn.ModuleList([])
+
+        for _ in range(depth):
+            self.layers.append(
+                ConformerBlock(
+                    dim=dim,
+                    dim_head=dim_head,
+                    heads=heads,
+                    ff_mult=ff_mult,
+                    conv_expansion_factor=conv_expansion_factor,
+                    conv_kernel_size=conv_kernel_size,
+                    conv_causal=conv_causal,
+                )
+            )
+
+
+    def forward(self, x):
+
+        for block in self.layers:
+            
+            x = block(x)
+
+        return x

RingFormer/env.py

@@ -0,0 +1,15 @@
+import os
+import shutil
+
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self
+
+
+def build_env(config, config_name, path):
+    t_path = os.path.join(path, config_name)
+    if config != t_path:
+        os.makedirs(path, exist_ok=True)
+        shutil.copyfile(config, os.path.join(path, config_name))

RingFormer/inference.ipynb

@@ -0,0 +1,292 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "7b82eb58",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/ubuntu/miniconda3/envs/respair/lib/python3.11/site-packages/IPython/core/magics/osm.py:417: UserWarning: This is now an optional IPython functionality, setting dhist requires you to install the `pickleshare` library.\n",
+      "  self.shell.db['dhist'] = compress_dhist(dhist)[-100:]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/home/ubuntu/RINGFORMER\n"
+     ]
+    }
+   ],
+   "source": [
+    "%cd /home/ubuntu/RINGFORMER\n",
+    "\n",
+    "from __future__ import absolute_import, division, print_function, unicode_literals\n",
+    "\n",
+    "import glob\n",
+    "import os\n",
+    "import argparse\n",
+    "import json\n",
+    "import torch\n",
+    "from scipy.io.wavfile import write\n",
+    "from env import AttrDict\n",
+    "from meldataset import mel_spectrogram, MAX_WAV_VALUE, load_wav\n",
+    "from models import Generator\n",
+    "from stft import TorchSTFT\n",
+    "\n",
+    "from Utils.JDC.model import JDCNet"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "3ee13ffd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "h = None\n",
+    "device = None\n",
+    "\n",
+    "\n",
+    "def load_checkpoint(filepath, device):\n",
+    "    assert os.path.isfile(filepath)\n",
+    "    print(\"Loading '{}'\".format(filepath))\n",
+    "    checkpoint_dict = torch.load(filepath, map_location=device)\n",
+    "    print(\"Complete.\")\n",
+    "    return checkpoint_dict\n",
+    "\n",
+    "\n",
+    "def get_mel(x):\n",
+    "    return mel_spectrogram(x, h.n_fft, h.num_mels, h.sampling_rate, h.hop_size, h.win_size, h.fmin, h.fmax)\n",
+    "\n",
+    "\n",
+    "def scan_checkpoint(cp_dir, prefix):\n",
+    "    pattern = os.path.join(cp_dir, prefix + '*')\n",
+    "    cp_list = glob.glob(pattern)\n",
+    "    if len(cp_list) == 0:\n",
+    "        return ''\n",
+    "    return sorted(cp_list)[-1]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "003b1249",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "F0_model = JDCNet(num_class=1, seq_len=192)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "321eb3b5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cp_path = \"/home/ubuntu/RINGFORMER/cp_ringformer_44.1khz\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "3dcc2764",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'F0_path': '/home/ubuntu/Darya_Speech/Utils/JDC/epoch_00079.pth',\n",
+       " 'resblock': '1',\n",
+       " 'num_gpus': 1,\n",
+       " 'batch_size': 18,\n",
+       " 'learning_rate': 0.0002,\n",
+       " 'adam_b1': 0.8,\n",
+       " 'adam_b2': 0.99,\n",
+       " 'lr_decay': 0.999,\n",
+       " 'seed': 1234,\n",
+       " 'upsample_rates': [16, 8],\n",
+       " 'upsample_kernel_sizes': [32, 16],\n",
+       " 'upsample_initial_channel': 512,\n",
+       " 'resblock_kernel_sizes': [3, 7, 11],\n",
+       " 'resblock_dilation_sizes': [[1, 3, 5], [1, 3, 5], [1, 3, 5]],\n",
+       " 'gen_istft_n_fft': 32,\n",
+       " 'gen_istft_hop_size': 4,\n",
+       " 'segment_size': 65536,\n",
+       " 'num_mels': 128,\n",
+       " 'n_fft': 2048,\n",
+       " 'hop_size': 512,\n",
+       " 'win_size': 2048,\n",
+       " 'sampling_rate': 44100,\n",
+       " 'fmin': 0,\n",
+       " 'fmax': None,\n",
+       " 'fmax_for_loss': None,\n",
+       " 'num_workers': 8,\n",
+       " 'dist_config': {'dist_backend': 'nccl',\n",
+       "  'dist_url': 'tcp://localhost:54321',\n",
+       "  'world_size': 1}}"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "with open(cp_path + \"/config.json\") as f:\n",
+    "    data = f.read()\n",
+    "\n",
+    "json_config = json.loads(data)\n",
+    "h = AttrDict(json_config)\n",
+    "h"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "a4c78cb6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# device = torch.device('cuda:{:d}'.format(0))\n",
+    "\n",
+    "device = 'cuda:0'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "5a782adb",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/ubuntu/miniconda3/envs/respair/lib/python3.11/site-packages/torch/nn/utils/weight_norm.py:143: FutureWarning: `torch.nn.utils.weight_norm` is deprecated in favor of `torch.nn.utils.parametrizations.weight_norm`.\n",
+      "  WeightNorm.apply(module, name, dim)\n"
+     ]
+    }
+   ],
+   "source": [
+    "generator = Generator(h, F0_model).to(device)\n",
+    "stft = TorchSTFT(filter_length=h.gen_istft_n_fft, hop_length=h.gen_istft_hop_size, win_length=h.gen_istft_n_fft).to(device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "6f0a7c64",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Loading '/home/ubuntu/RINGFORMER/cp_ringformer_44.1khz/g_00017000'\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_3972638/3295719764.py:8: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  checkpoint_dict = torch.load(filepath, map_location=device)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Complete.\n",
+      "Removing weight norm...\n"
+     ]
+    }
+   ],
+   "source": [
+    "cp_g = scan_checkpoint(cp_path, 'g_')\n",
+    "state_dict_g = load_checkpoint(cp_g, device)\n",
+    "generator.load_state_dict(state_dict_g['generator'])\n",
+    "generator.remove_weight_norm()\n",
+    "_ = generator.eval()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a115a967",
+   "metadata": {},
+   "source": [
+    "### Resynthesis"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cbeee500",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torchaudio\n",
+    "from librosa.filters import mel as librosa_mel_fn\n",
+    "from IPython.display import Audio\n",
+    "import librosa\n",
+    "\n",
+    "to_mel = torchaudio.transforms.MelSpectrogram(\n",
+    "    n_mels=128, n_fft=2048, win_length=2048, hop_length=512, sample_rate=44100, power=2.5)\n",
+    "mean, std = -4, 4\n",
+    "\n",
+    "def preprocess(wave):\n",
+    "    wave_tensor = torch.FloatTensor(wav)\n",
+    "    mel_tensor = to_mel(wave_tensor)\n",
+    "    mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - mean) / std\n",
+    "    return mel_tensor.to('cuda:0')\n",
+    "\n",
+    "\n",
+    "wav = librosa.load(\"/your.wav\", sr=44100)[0]\n",
+    "\n",
+    "x = preprocess(wav)\n",
+    "print(x.shape)\n",
+    "\n",
+    "n = 1\n",
+    "xxx = torch.load(\"/home/ubuntu/RINGFORMER/gt.pt\").to('cuda:0')[n:n+1,:,:]\n",
+    "with torch.no_grad():\n",
+    "\n",
+    "    spec, phase = generator(xxx)\n",
+    "    y_g_hat = stft.inverse(spec, phase)\n",
+    "    audio = y_g_hat.squeeze()\n",
+    "    # audio = audio * MAX_WAV_VALUE\n",
+    "    audio = audio.cpu().numpy()\n",
+    "\n",
+    "\n",
+    "print('Synthesized:')\n",
+    "display(Audio(audio, rate=44100))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "respair",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

RingFormer/meldataset.py

@@ -0,0 +1,203 @@
+import math
+import os
+import random
+import torch
+import torch.utils.data
+import numpy as np
+from librosa.util import normalize
+from scipy.io.wavfile import read
+import torchaudio
+import librosa 
+from librosa.filters import mel as librosa_mel_fn
+
+MAX_WAV_VALUE = 32768.0
+import soundfile as sf
+
+
+def normalize_audio(wav):
+    return wav / torch.max(torch.abs(torch.from_numpy(wav)))  # Correct peak normalization
+
+def load_wav(full_path):
+    data, sampling_rate = librosa.load(full_path, sr=None)
+    return data, sampling_rate
+
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+    return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
+
+
+def dynamic_range_decompression(x, C=1):
+    return np.exp(x) / C
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression_torch(x, C=1):
+    return torch.exp(x) / C
+
+
+def spectral_normalize_torch(magnitudes):
+    output = dynamic_range_compression_torch(magnitudes)
+    return output
+
+
+def spectral_de_normalize_torch(magnitudes):
+    output = dynamic_range_decompression_torch(magnitudes)
+    return output
+
+
+mel_basis = {}
+hann_window = {}
+
+
+def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
+    
+    y = torch.clamp(y, min=-1, max=1)
+    
+    if torch.min(y) < -1.:
+        print('min value is ', torch.min(y))
+    if torch.max(y) > 1.:
+        print('max value is ', torch.max(y))
+
+    global mel_basis, hann_window
+    if fmax not in mel_basis:
+        mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
+        mel_basis[str(fmax)+'_'+str(y.device)] = torch.from_numpy(mel).float().to(y.device)
+        hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device)
+
+    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
+    y = y.squeeze(1)
+
+    # complex tensor as default, then use view_as_real for future pytorch compatibility
+    spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[str(y.device)],
+                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=True)
+    spec = torch.view_as_real(spec)
+    spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9))
+
+    spec = torch.matmul(mel_basis[str(fmax)+'_'+str(y.device)], spec)
+    spec = spectral_normalize_torch(spec)
+
+    return spec
+
+
+
+to_mel = torchaudio.transforms.MelSpectrogram(
+   sample_rate=44_100, n_mels=128, n_fft=2048, win_length=2048, hop_length=512)
+
+
+
+# to_mel = torchaudio.transforms.MelSpectrogram(
+#    sample_rate=24000, n_mels=80, n_fft=2048, win_length=1200, hop_length=300, center='center')
+
+mean, std = -4, 4
+
+def preproces(wave,to_mel=to_mel, device='cpu'):
+    
+    to_mel = to_mel.to(device)
+    # wave_tensor = torch.from_numpy(wave).float()
+    mel_tensor = to_mel(wave)
+    mel_tensor = (torch.log(1e-5 + mel_tensor) - mean) / std
+    return mel_tensor
+
+
+def get_dataset_filelist(a):
+    with open(a.input_training_file, 'r', encoding='utf-8') as fi:
+        training_files = [os.path.join(a.input_wavs_dir, x.split('|')[0] + ('' if '' not in x else ''))
+                          for x in fi.read().split('\n') if len(x) > 0]
+
+    with open(a.input_validation_file, 'r', encoding='utf-8') as fi:
+        validation_files = [os.path.join(a.input_wavs_dir, x.split('|')[0] + ('' if '' not in x else ''))
+                            for x in fi.read().split('\n') if len(x) > 0]
+    return training_files, validation_files
+
+
+class MelDataset(torch.utils.data.Dataset):
+    def __init__(self, training_files, segment_size, n_fft, num_mels,
+                 hop_size, win_size, sampling_rate,  fmin, fmax, split=True, shuffle=True, n_cache_reuse=1,
+                 device=None, fmax_loss=None, fine_tuning=False, base_mels_path=None):
+        self.audio_files = training_files
+        random.seed(1234)
+        if shuffle:
+            random.shuffle(self.audio_files)
+        self.segment_size = segment_size
+        self.sampling_rate = sampling_rate
+        self.split = split
+        self.n_fft = n_fft
+        self.num_mels = num_mels
+        self.hop_size = hop_size
+        self.win_size = win_size
+        self.fmin = fmin
+        self.fmax = fmax
+        self.fmax_loss = fmax_loss
+        self.cached_wav = None
+        self.n_cache_reuse = n_cache_reuse
+        self._cache_ref_count = 0
+        self.device = device
+        self.fine_tuning = fine_tuning
+        self.base_mels_path = base_mels_path
+
+    def __getitem__(self, index):
+        filename = self.audio_files[index]
+        if self._cache_ref_count == 0:
+            audio, sampling_rate = load_wav(filename)
+          
+
+            self.cached_wav = audio
+            if sampling_rate != self.sampling_rate:
+                raise ValueError("{} SR doesn't match target {} SR".format(
+                    sampling_rate, self.sampling_rate))
+            self._cache_ref_count = self.n_cache_reuse
+        else:
+            audio = self.cached_wav
+            self._cache_ref_count -= 1
+
+        audio = torch.FloatTensor(audio)
+        audio = audio.unsqueeze(0)
+
+        if not self.fine_tuning:
+            if self.split:
+                if audio.size(1) >= self.segment_size:
+                    max_audio_start = audio.size(1) - self.segment_size
+                    audio_start = random.randint(0, max_audio_start)
+                    audio = audio[:, audio_start:audio_start+self.segment_size]
+                else:
+                    audio = torch.nn.functional.pad(audio, (0, self.segment_size - audio.size(1)), 'constant')
+
+            # mel = mel_spectrogram(audio, self.n_fft, self.num_mels,
+            #                       self.sampling_rate, self.hop_size, self.win_size, self.fmin, self.fmax,
+            #                       center=False)
+            
+            mel = preproces(audio)
+        else:
+            mel = np.load(
+                os.path.join(self.base_mels_path, os.path.splitext(os.path.split(filename)[-1])[0] + '.npy'))
+            mel = torch.from_numpy(mel)
+
+            if len(mel.shape) < 3:
+                mel = mel.unsqueeze(0)
+
+            if self.split:
+                frames_per_seg = math.ceil(self.segment_size / self.hop_size)
+
+                if audio.size(1) >= self.segment_size:
+                    mel_start = random.randint(0, mel.size(2) - frames_per_seg - 1)
+                    mel = mel[:, :, mel_start:mel_start + frames_per_seg]
+                    audio = audio[:, mel_start * self.hop_size:(mel_start + frames_per_seg) * self.hop_size]
+                else:
+                    mel = torch.nn.functional.pad(mel, (0, frames_per_seg - mel.size(2)), 'constant')
+                    audio = torch.nn.functional.pad(audio, (0, self.segment_size - audio.size(1)), 'constant')
+
+        mel_loss = mel_spectrogram(audio, self.n_fft, self.num_mels,
+                                   self.sampling_rate, self.hop_size, self.win_size, self.fmin, self.fmax_loss,
+                                   center=False)
+        
+        # mel_loss = mel_spectrogram(audio)
+        if mel.shape[-1] != mel_loss.shape[-1]:
+            mel = mel[..., :mel_loss.shape[-1]]
+            
+        return (mel.squeeze(), audio.squeeze(0), filename, mel_loss.squeeze())
+
+    def __len__(self):
+        return len(self.audio_files)

RingFormer/models.py

@@ -0,0 +1,943 @@
+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 numpy as np 
+from stft import TorchSTFT
+import torchaudio
+from nnAudio import features
+from einops import rearrange
+from norm2d import NormConv2d
+from utils import get_padding
+from munch import Munch
+from conformer import Conformer
+
+LRELU_SLOPE = 0.1
+
+
+def get_2d_padding(kernel_size, dilation=(1, 1)):
+    return (
+        ((kernel_size[0] - 1) * dilation[0]) // 2,
+        ((kernel_size[1] - 1) * dilation[1]) // 2,
+    )
+
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.h = h
+        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.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):
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, self.alpha1, self.alpha2):
+            xt = x + (1 / a1) * (torch.sin(a1 * x) ** 2)  # Snake1D
+            xt = c1(xt)
+            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 ResBlock1_old(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.h = h
+        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)
+
+    def forward(self, x):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            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 ResBlock2(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)):
+        super(ResBlock2, self).__init__()
+        self.h = h
+        self.convs = 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])))
+        ])
+        self.convs.apply(init_weights)
+
+    def forward(self, x):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            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, h, F0_model):
+        super(Generator, self).__init__()
+        self.h = h
+        self.num_kernels = len(h.resblock_kernel_sizes)
+        self.num_upsamples = len(h.upsample_rates)
+        self.conv_pre = weight_norm(Conv1d(128, h.upsample_initial_channel, 7, 1, padding=3))
+        
+        
+
+        resblock = ResBlock1 if h.resblock == '1' else ResBlock2
+
+        self.m_source = SourceModuleHnNSF(
+                    sampling_rate=h.sampling_rate,
+                    upsample_scale=np.prod(h.upsample_rates) * h.gen_istft_hop_size,
+                    harmonic_num=8, voiced_threshod=10)
+        
+        self.f0_upsamp = torch.nn.Upsample(
+            scale_factor=np.prod(h.upsample_rates) * h.gen_istft_hop_size)
+        self.noise_convs = nn.ModuleList()
+        self.noise_res = nn.ModuleList()
+        
+        self.F0_model = F0_model
+        
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+            self.ups.append(weight_norm(
+                ConvTranspose1d(h.upsample_initial_channel//(2**i),
+                                h.upsample_initial_channel//(2**(i+1)),
+                                k,
+                                u,
+                                padding=(k-u)//2)))
+
+            c_cur = h.upsample_initial_channel // (2 ** (i + 1))
+            
+            if i + 1 < len(h.upsample_rates):  #
+                stride_f0 = np.prod(h.upsample_rates[i + 1:])
+                self.noise_convs.append(Conv1d(
+                    h.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(h, c_cur, 7, [1,3,5]))
+            else:
+                self.noise_convs.append(Conv1d(h.gen_istft_n_fft + 2, c_cur, kernel_size=1))
+                self.noise_res.append(resblock(h, c_cur, 11, [1,3,5]))
+                
+        self.alphas = nn.ParameterList()
+        self.alphas.append(nn.Parameter(torch.ones(1, h.upsample_initial_channel, 1)))
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = h.upsample_initial_channel//(2**(i+1))
+            self.alphas.append(nn.Parameter(torch.ones(1, ch, 1)))
+            for j, (k, d) in enumerate(
+                zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
+                self.resblocks.append(resblock(h, ch, k, d))
+
+        
+        self.conformers = nn.ModuleList()
+        self.post_n_fft = h.gen_istft_n_fft
+        self.conv_post = weight_norm(Conv1d(128, self.post_n_fft + 2, 7, 1, padding=3))
+        
+        for i in range(len(self.ups)):
+            ch = h.upsample_initial_channel // (2**i)
+            self.conformers.append(
+                Conformer(
+                    dim=ch,
+                    depth=2,
+                    dim_head=64,
+                    heads=8,
+                    ff_mult=4,
+                    conv_expansion_factor=2,
+                    conv_kernel_size=31,
+                    attn_dropout=0.1,
+                    ff_dropout=0.1,
+                    conv_dropout=0.1,
+                    # device=self.device
+                )
+            )
+            
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+        self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
+        self.stft = TorchSTFT(filter_length=h.gen_istft_n_fft,
+                              hop_length=h.gen_istft_hop_size,
+                              win_length=h.gen_istft_n_fft)
+
+
+
+    def forward(self, x):
+        
+      
+        
+        f0, _, _ = self.F0_model(x.unsqueeze(1))
+        if len(f0.shape) == 1:
+            f0 = f0.unsqueeze(0)
+        
+        f0 = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
+
+        har_source, _, _ = 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)
+        
+
+        x = self.conv_pre(x)
+   
+        for i in range(self.num_upsamples):
+
+            x = x + (1 / self.alphas[i]) * (torch.sin(self.alphas[i] * x) ** 2)
+            x = rearrange(x, "b f t -> b t f")
+        
+            x = self.conformers[i](x)
+
+            x = rearrange(x, "b t f -> b f t")
+            
+            # x = F.leaky_relu(x, LRELU_SLOPE)
+            x_source = self.noise_convs[i](har)
+            x_source = self.noise_res[i](x_source)
+            
+            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)
+                else:
+                    xs += self.resblocks[i*self.num_kernels+j](x)
+            x = xs / self.num_kernels
+        # x = F.leaky_relu(x)
+        
+
+        x = x + (1 / self.alphas[i + 1]) * (torch.sin(self.alphas[i + 1] * x) ** 2)
+        
+        x = self.conv_post(x)
+        spec = torch.exp(x[:,:self.post_n_fft // 2 + 1, :]).to(x.device)
+        phase = torch.sin(x[:, self.post_n_fft // 2 + 1:, :]).to(x.device)
+
+        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)
+
+
+
+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]
+
+    # NOTE(kan-bayashi): clamp is needed to avoid nan or inf
+    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 DiscriminatorS(torch.nn.Module):
+    def __init__(self, use_spectral_norm=False):
+        super(DiscriminatorS, self).__init__()
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.convs = nn.ModuleList([
+            norm_f(Conv1d(1, 128, 15, 1, padding=7)),
+            norm_f(Conv1d(128, 128, 41, 2, groups=4, padding=20)),
+            norm_f(Conv1d(128, 256, 41, 2, groups=16, padding=20)),
+            norm_f(Conv1d(256, 512, 41, 4, groups=16, padding=20)),
+            norm_f(Conv1d(512, 1024, 41, 4, groups=16, padding=20)),
+            norm_f(Conv1d(1024, 1024, 41, 1, groups=16, padding=20)),
+            norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
+        ])
+        self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, 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)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiScaleDiscriminator(torch.nn.Module):
+    def __init__(self):
+        super(MultiScaleDiscriminator, self).__init__()
+        self.discriminators = nn.ModuleList([
+            DiscriminatorS(use_spectral_norm=True),
+            DiscriminatorS(),
+            DiscriminatorS(),
+        ])
+        self.meanpools = nn.ModuleList([
+            AvgPool1d(4, 2, padding=2),
+            AvgPool1d(4, 2, padding=2)
+        ])
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            if i != 0:
+                y = self.meanpools[i-1](y)
+                y_hat = self.meanpools[i-1](y_hat)
+            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
+
+
+
+
+
+########################### from ringformer
+
+multiscale_subband_cfg = {
+    "hop_lengths": [1024, 512, 512],  # Doubled to maintain similar time resolution
+    "sampling_rate": 44100,           # New sampling rate
+    "filters": 32,                    # Kept same as it controls initial feature dimension
+    "max_filters": 1024,              # Kept same as it's a maximum limit
+    "filters_scale": 1,               # Kept same as it's a scaling factor
+    "dilations": [1, 2, 4],          # Kept same as they control receptive field growth
+    "in_channels": 1,                 # Kept same (mono audio)
+    "out_channels": 1,                # Kept same (mono audio)
+    "n_octaves": [10, 10, 10],       # Increased by 1 to handle higher frequency range
+    "bins_per_octaves": [24, 36, 48], # Kept same as they control frequency resolution
+}
+
+
+
+# multiscale_subband_cfg = {
+#     "hop_lengths": [512, 256, 256],
+#     "sampling_rate": 24000,
+#     "filters": 32,
+#     "max_filters": 1024,
+#     "filters_scale": 1,
+#     "dilations": [1, 2, 4],
+#     "in_channels": 1,
+#     "out_channels": 1,
+#     "n_octaves": [9, 9, 9],
+#     "bins_per_octaves": [24, 36, 48],
+# }
+
+class DiscriminatorCQT(nn.Module): 
+    def __init__(self, cfg, hop_length, n_octaves, bins_per_octave):
+        super(DiscriminatorCQT, self).__init__()
+        self.cfg = cfg
+
+        self.filters = cfg.filters
+        self.max_filters = cfg.max_filters
+        self.filters_scale = cfg.filters_scale
+        self.kernel_size = (3, 9)
+        self.dilations = cfg.dilations
+        self.stride = (1, 2)
+
+        self.in_channels = cfg.in_channels
+        self.out_channels = cfg.out_channels
+        self.fs = cfg.sampling_rate
+        self.hop_length = hop_length
+        self.n_octaves = n_octaves
+        self.bins_per_octave = bins_per_octave
+
+        self.cqt_transform = features.cqt.CQT2010v2(
+            sr=self.fs * 2,
+            hop_length=self.hop_length,
+            n_bins=self.bins_per_octave * self.n_octaves,
+            bins_per_octave=self.bins_per_octave,
+            output_format="Complex",
+            pad_mode="constant",
+        )
+
+        self.conv_pres = nn.ModuleList()
+        for i in range(self.n_octaves):
+            self.conv_pres.append(
+                NormConv2d(
+                    self.in_channels * 2,
+                    self.in_channels * 2,
+                    kernel_size=self.kernel_size,
+                    padding=get_2d_padding(self.kernel_size),
+                )
+            )
+
+        self.convs = nn.ModuleList()
+
+        self.convs.append(
+            NormConv2d(
+                self.in_channels * 2,
+                self.filters,
+                kernel_size=self.kernel_size,
+                padding=get_2d_padding(self.kernel_size),
+            )
+        )
+
+        in_chs = min(self.filters_scale * self.filters, self.max_filters)
+        for i, dilation in enumerate(self.dilations):
+            out_chs = min(
+                (self.filters_scale ** (i + 1)) * self.filters, self.max_filters
+            )
+            self.convs.append(
+                NormConv2d(
+                    in_chs,
+                    out_chs,
+                    kernel_size=self.kernel_size,
+                    stride=self.stride,
+                    dilation=(dilation, 1),
+                    padding=get_2d_padding(self.kernel_size, (dilation, 1)),
+                    norm="weight_norm",
+                )
+            )
+            in_chs = out_chs
+        out_chs = min(
+            (self.filters_scale ** (len(self.dilations) + 1)) * self.filters,
+            self.max_filters,
+        )
+        self.convs.append(
+            NormConv2d(
+                in_chs,
+                out_chs,
+                kernel_size=(self.kernel_size[0], self.kernel_size[0]),
+                padding=get_2d_padding((self.kernel_size[0], self.kernel_size[0])),
+                norm="weight_norm",
+            )
+        )
+
+        self.conv_post = NormConv2d(
+            out_chs,
+            self.out_channels,
+            kernel_size=(self.kernel_size[0], self.kernel_size[0]),
+            padding=get_2d_padding((self.kernel_size[0], self.kernel_size[0])),
+            norm="weight_norm",
+        )
+
+        self.activation = torch.nn.LeakyReLU(negative_slope=LRELU_SLOPE)
+        self.resample = torchaudio.transforms.Resample(
+            orig_freq=self.fs, new_freq=self.fs * 2
+        )
+
+    def forward(self, x):
+        fmap = []
+
+        x = self.resample(x)
+
+        z = self.cqt_transform(x)
+
+        z_amplitude = z[:, :, :, 0].unsqueeze(1)
+        z_phase = z[:, :, :, 1].unsqueeze(1)
+
+        z = torch.cat([z_amplitude, z_phase], dim=1)
+        z = rearrange(z, "b c w t -> b c t w")
+
+        latent_z = []
+        for i in range(self.n_octaves):
+            latent_z.append(
+                self.conv_pres[i](
+                    z[
+                        :,
+                        :,
+                        :,
+                        i * self.bins_per_octave : (i + 1) * self.bins_per_octave,
+                    ]
+                )
+            )
+        latent_z = torch.cat(latent_z, dim=-1)
+
+        for i, l in enumerate(self.convs):
+            latent_z = l(latent_z)
+
+            latent_z = self.activation(latent_z)
+            fmap.append(latent_z)
+
+        latent_z = self.conv_post(latent_z)
+
+        return latent_z, fmap
+    
+    
+    
+class MultiScaleSubbandCQTDiscriminator(nn.Module): # replacing "MultiResSpecDiscriminator"
+    def __init__(self):
+        super(MultiScaleSubbandCQTDiscriminator, self).__init__()
+        cfg = Munch(multiscale_subband_cfg)
+        self.cfg = cfg
+        self.discriminators = nn.ModuleList(
+            [
+                DiscriminatorCQT(
+                    cfg,
+                    hop_length=cfg.hop_lengths[i],
+                    n_octaves=cfg.n_octaves[i],
+                    bins_per_octave=cfg.bins_per_octaves[i],
+                )
+                for i in range(len(cfg.hop_lengths))
+            ]
+        )
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+
+        for disc in self.discriminators:
+            y_d_r, fmap_r = disc(y)
+            y_d_g, fmap_g = disc(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
+
+    
+    
+    #############################
+
+
+
+def feature_loss(fmap_r, fmap_g):
+    loss = 0
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            loss += torch.mean(torch.abs(rl - gl))
+
+    return loss*2
+
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+    loss = 0
+    r_losses = []
+    g_losses = []
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        r_loss = torch.mean((1-dr)**2)
+        g_loss = torch.mean(dg**2)
+        loss += (r_loss + g_loss)
+        r_losses.append(r_loss.item())
+        g_losses.append(g_loss.item())
+
+    return loss, r_losses, g_losses
+
+
+def generator_loss(disc_outputs):
+    loss = 0
+    gen_losses = []
+    for dg in disc_outputs:
+        l = torch.mean((1-dg)**2)
+        gen_losses.append(l)
+        loss += l
+
+    return loss, gen_losses
+
+def discriminator_TPRLS_loss(disc_real_outputs, disc_generated_outputs):
+    loss = 0
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        tau = 0.04
+        m_DG = torch.median((dr-dg))
+        L_rel = torch.mean((((dr - dg) - m_DG)**2)[dr < dg + m_DG])
+        loss += tau - F.relu(tau - L_rel)
+    return loss
+
+def generator_TPRLS_loss(disc_real_outputs, disc_generated_outputs):
+    loss = 0
+    for dg, dr in zip(disc_real_outputs, disc_generated_outputs):
+        tau = 0.04
+        m_DG = torch.median((dr-dg))
+        L_rel = torch.mean((((dr - dg) - m_DG)**2)[dr < dg + m_DG])
+        loss += tau - F.relu(tau - L_rel)
+    return loss

RingFormer/norm2d.py

@@ -0,0 +1,92 @@
+# Copyright (c) 2023 Amphion.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+#################### Norm2D for Discriminators ####################
+
+import torch
+import torch.nn as nn
+import einops
+from torch.nn.utils import spectral_norm, weight_norm
+
+CONV_NORMALIZATIONS = frozenset(
+    [
+        "none",
+        "weight_norm",
+        "spectral_norm",
+        "time_layer_norm",
+        "layer_norm",
+        "time_group_norm",
+    ]
+)
+
+
+class ConvLayerNorm(nn.LayerNorm):
+    """
+    Convolution-friendly LayerNorm that moves channels to last dimensions
+    before running the normalization and moves them back to original position right after.
+    """
+
+    def __init__(self, normalized_shape, **kwargs):
+        super().__init__(normalized_shape, **kwargs)
+
+    def forward(self, x):
+        x = einops.rearrange(x, "b ... t -> b t ...")
+        x = super().forward(x)
+        x = einops.rearrange(x, "b t ... -> b ... t")
+        return
+
+
+def apply_parametrization_norm(module: nn.Module, norm: str = "none") -> nn.Module:
+    assert norm in CONV_NORMALIZATIONS
+    if norm == "weight_norm":
+        return weight_norm(module)
+    elif norm == "spectral_norm":
+        return spectral_norm(module)
+    else:
+        # We already check was in CONV_NORMALIZATION, so any other choice
+        # doesn't need reparametrization.
+        return module
+
+
+def get_norm_module(
+    module: nn.Module, causal: bool = False, norm: str = "none", **norm_kwargs
+) -> nn.Module:
+    """Return the proper normalization module. If causal is True, this will ensure the returned
+    module is causal, or return an error if the normalization doesn't support causal evaluation.
+    """
+    assert norm in CONV_NORMALIZATIONS
+    if norm == "layer_norm":
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return ConvLayerNorm(module.out_channels, **norm_kwargs)
+    elif norm == "time_group_norm":
+        if causal:
+            raise ValueError("GroupNorm doesn't support causal evaluation.")
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+    else:
+        return nn.Identity()
+
+
+class NormConv2d(nn.Module):
+    """Wrapper around Conv2d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+
+    def __init__(
+        self,
+        *args,
+        norm: str = "none",
+        norm_kwargs={},
+        **kwargs,
+    ):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x

RingFormer/requirements.txt

@@ -0,0 +1,10 @@
+torch
+numpy
+librosa
+scipy==1.4.1
+tensorboard==2.0
+soundfile
+matplotlib==3.1.3
+nnaudio
+ring-attention-pytorch
+--extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/triton-nightly

RingFormer/stft.py

@@ -0,0 +1,254 @@
+"""
+BSD 3-Clause License
+Copyright (c) 2017, Prem Seetharaman
+All rights reserved.
+* Redistribution and use in source and binary forms, with or without
+  modification, are permitted provided that the following conditions are met:
+* Redistributions of source code must retain the above copyright notice,
+  this list of conditions and the following disclaimer.
+* Redistributions in binary form must reproduce the above copyright notice, this
+  list of conditions and the following disclaimer in the
+  documentation and/or other materials provided with the distribution.
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from this
+  software without specific prior written permission.
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+
+import torch
+import numpy as np
+import torch.nn.functional as F
+from torch.autograd import Variable
+from scipy.signal import get_window
+from librosa.util import pad_center, tiny
+import librosa.util as librosa_util
+
+
+def window_sumsquare(
+    window,
+    n_frames,
+    hop_length=200,
+    win_length=800,
+    n_fft=800,
+    dtype=np.float32,
+    norm=None,
+):
+    """
+    # from librosa 0.6
+    Compute the sum-square envelope of a window function at a given hop length.
+    This is used to estimate modulation effects induced by windowing
+    observations in short-time fourier transforms.
+    Parameters
+    ----------
+    window : string, tuple, number, callable, or list-like
+        Window specification, as in `get_window`
+    n_frames : int > 0
+        The number of analysis frames
+    hop_length : int > 0
+        The number of samples to advance between frames
+    win_length : [optional]
+        The length of the window function.  By default, this matches `n_fft`.
+    n_fft : int > 0
+        The length of each analysis frame.
+    dtype : np.dtype
+        The data type of the output
+    Returns
+    -------
+    wss : np.ndarray, shape=`(n_fft + hop_length * (n_frames - 1))`
+        The sum-squared envelope of the window function
+    """
+    if win_length is None:
+        win_length = n_fft
+
+    n = n_fft + hop_length * (n_frames - 1)
+    x = np.zeros(n, dtype=dtype)
+
+    # Compute the squared window at the desired length
+    win_sq = get_window(window, win_length, fftbins=True)
+    win_sq = librosa_util.normalize(win_sq, norm=norm) ** 2
+    win_sq = librosa_util.pad_center(win_sq, n_fft)
+
+    # Fill the envelope
+    for i in range(n_frames):
+        sample = i * hop_length
+        x[sample : min(n, sample + n_fft)] += win_sq[: max(0, min(n_fft, n - sample))]
+    return x
+
+
+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=False)
+    real = x_stft[..., 0]
+    imag = x_stft[..., 1]
+
+    # NOTE(kan-bayashi): clamp is needed to avoid nan or inf
+    return torch.sqrt(torch.clamp(real**2 + imag**2, min=1e-7)).transpose(2, 1)
+
+
+class STFT(torch.nn.Module):
+    """adapted from Prem Seetharaman's https://github.com/pseeth/pytorch-stft"""
+
+    def __init__(
+        self, filter_length=800, hop_length=200, win_length=800, window="hann"
+    ):
+        super(STFT, self).__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.window = window
+        self.forward_transform = None
+        scale = self.filter_length / self.hop_length
+        fourier_basis = np.fft.fft(np.eye(self.filter_length))
+
+        cutoff = int((self.filter_length / 2 + 1))
+        fourier_basis = np.vstack(
+            [np.real(fourier_basis[:cutoff, :]), np.imag(fourier_basis[:cutoff, :])]
+        )
+
+        forward_basis = torch.FloatTensor(fourier_basis[:, None, :])
+        inverse_basis = torch.FloatTensor(
+            np.linalg.pinv(scale * fourier_basis).T[:, None, :]
+        )
+
+        if window is not None:
+            assert filter_length >= win_length
+            # get window and zero center pad it to filter_length
+            fft_window = get_window(window, win_length, fftbins=True)
+            fft_window = pad_center(fft_window, filter_length)
+            fft_window = torch.from_numpy(fft_window).float()
+
+            # window the bases
+            forward_basis *= fft_window
+            inverse_basis *= fft_window
+
+        self.register_buffer("forward_basis", forward_basis.float())
+        self.register_buffer("inverse_basis", inverse_basis.float())
+
+    def transform(self, input_data):
+        num_batches = input_data.size(0)
+        num_samples = input_data.size(1)
+
+        self.num_samples = num_samples
+
+        # similar to librosa, reflect-pad the input
+        input_data = input_data.view(num_batches, 1, num_samples)
+        input_data = F.pad(
+            input_data.unsqueeze(1),
+            (int(self.filter_length / 2), int(self.filter_length / 2), 0, 0),
+            mode="reflect",
+        )
+        input_data = input_data.squeeze(1)
+
+        forward_transform = F.conv1d(
+            input_data,
+            Variable(self.forward_basis, requires_grad=False),
+            stride=self.hop_length,
+            padding=0,
+        )
+
+        cutoff = int((self.filter_length / 2) + 1)
+        real_part = forward_transform[:, :cutoff, :]
+        imag_part = forward_transform[:, cutoff:, :]
+
+        magnitude = torch.sqrt(real_part**2 + imag_part**2)
+        phase = torch.autograd.Variable(torch.atan2(imag_part.data, real_part.data))
+
+        return magnitude, phase
+
+    def inverse(self, magnitude, phase):
+        recombine_magnitude_phase = torch.cat(
+            [magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1
+        )
+
+        inverse_transform = F.conv_transpose1d(
+            recombine_magnitude_phase,
+            Variable(self.inverse_basis, requires_grad=False),
+            stride=self.hop_length,
+            padding=0,
+        )
+
+        if self.window is not None:
+            window_sum = window_sumsquare(
+                self.window,
+                magnitude.size(-1),
+                hop_length=self.hop_length,
+                win_length=self.win_length,
+                n_fft=self.filter_length,
+                dtype=np.float32,
+            )
+            # remove modulation effects
+            approx_nonzero_indices = torch.from_numpy(
+                np.where(window_sum > tiny(window_sum))[0]
+            )
+            window_sum = torch.autograd.Variable(
+                torch.from_numpy(window_sum), requires_grad=False
+            )
+            window_sum = (
+                window_sum.to(inverse_transform.device())
+                if magnitude.is_cuda
+                else window_sum
+            )
+            inverse_transform[:, :, approx_nonzero_indices] /= window_sum[
+                approx_nonzero_indices
+            ]
+
+            # scale by hop ratio
+            inverse_transform *= float(self.filter_length) / self.hop_length
+
+        inverse_transform = inverse_transform[:, :, int(self.filter_length / 2) :]
+        inverse_transform = inverse_transform[:, :, : -int(self.filter_length / 2) :]
+
+        return inverse_transform
+
+    def forward(self, input_data):
+        self.magnitude, self.phase = self.transform(input_data)
+        reconstruction = self.inverse(self.magnitude, self.phase)
+        return reconstruction
+
+
+
+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

RingFormer/train.py

@@ -0,0 +1,671 @@
+# import warnings
+# warnings.simplefilter(action='ignore', category=FutureWarning)
+# import itertools
+# import os
+# import time
+# import argparse
+# import json
+# import torch
+# import torch.nn.functional as F
+# from torch.utils.tensorboard import SummaryWriter
+# from torch.utils.data import DistributedSampler, DataLoader
+# import torch.multiprocessing as mp
+# from torch.distributed import init_process_group
+# from torch.nn.parallel import DistributedDataParallel
+# from env import AttrDict, build_env
+# from meldataset import MelDataset, mel_spectrogram, get_dataset_filelist
+# from models import Generator, MultiPeriodDiscriminator, feature_loss, generator_loss,\
+#     discriminator_loss, discriminator_TPRLS_loss, generator_TPRLS_loss, MultiScaleSubbandCQTDiscriminator
+# from utils import plot_spectrogram, scan_checkpoint, load_checkpoint, save_checkpoint
+# from stft import TorchSTFT
+# from Utils.JDC.model import JDCNet
+
+# torch.backends.cudnn.benchmark = True
+
+
+# def train(rank, a, h):
+#     if h.num_gpus > 1:
+#         init_process_group(backend=h.dist_config['dist_backend'], init_method=h.dist_config['dist_url'],
+#                            world_size=h.dist_config['world_size'] * h.num_gpus, rank=rank)
+
+#     torch.cuda.manual_seed(h.seed)
+#     device = torch.device('cuda:{:d}'.format(rank))
+
+#     F0_model = JDCNet(num_class=1, seq_len=192)
+#     params = torch.load(h.F0_path)['net']
+#     F0_model.load_state_dict(params)
+    
+#     generator = Generator(h, F0_model, device=device).to(device)
+#     mpd = MultiPeriodDiscriminator().to(device)
+#     msd = MultiScaleSubbandCQTDiscriminator().to(device)
+#     stft = TorchSTFT(filter_length=h.gen_istft_n_fft, hop_length=h.gen_istft_hop_size, win_length=h.gen_istft_n_fft).to(device)
+
+#     if rank == 0:
+#         print(generator)
+#         os.makedirs(a.checkpoint_path, exist_ok=True)
+#         print("checkpoints directory : ", a.checkpoint_path)
+
+#     if os.path.isdir(a.checkpoint_path):
+#         cp_g = scan_checkpoint(a.checkpoint_path, 'g_')
+#         cp_do = scan_checkpoint(a.checkpoint_path, 'do_')
+
+#     steps = 0
+#     if cp_g is None or cp_do is None:
+#         state_dict_do = None
+#         last_epoch = -1
+#     else:
+#         state_dict_g = load_checkpoint(cp_g, device)
+#         state_dict_do = load_checkpoint(cp_do, device)
+#         generator.load_state_dict(state_dict_g['generator'])
+#         mpd.load_state_dict(state_dict_do['mpd'])
+#         msd.load_state_dict(state_dict_do['msd'])
+#         steps = state_dict_do['steps'] + 1
+#         last_epoch = state_dict_do['epoch']
+
+#     if h.num_gpus > 1:
+#         generator = DistributedDataParallel(generator, device_ids=[rank], find_unused_parameters=True).to(device)
+#         mpd = DistributedDataParallel(mpd, device_ids=[rank]).to(device)
+#         msd = DistributedDataParallel(msd, device_ids=[rank]).to(device)
+
+#     optim_g = torch.optim.AdamW(generator.parameters(), h.learning_rate, betas=[h.adam_b1, h.adam_b2])
+#     optim_d = torch.optim.AdamW(itertools.chain(msd.parameters(), mpd.parameters()),
+#                                 h.learning_rate, betas=[h.adam_b1, h.adam_b2])
+
+#     if state_dict_do is not None:
+#         optim_g.load_state_dict(state_dict_do['optim_g'])
+#         optim_d.load_state_dict(state_dict_do['optim_d'])
+
+#     scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=h.lr_decay, last_epoch=last_epoch)
+#     scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=h.lr_decay, last_epoch=last_epoch)
+
+#     training_filelist, validation_filelist = get_dataset_filelist(a)
+
+#     trainset = MelDataset(training_filelist, h.segment_size, h.n_fft, h.num_mels,
+#                           h.hop_size, h.win_size, h.sampling_rate, h.fmin, h.fmax, n_cache_reuse=0,
+#                           shuffle=False if h.num_gpus > 1 else True, fmax_loss=h.fmax_for_loss, device=device,
+#                           fine_tuning=a.fine_tuning, base_mels_path=a.input_mels_dir)
+
+#     train_sampler = DistributedSampler(trainset) if h.num_gpus > 1 else None
+
+#     train_loader = DataLoader(trainset, num_workers=h.num_workers, shuffle=False,
+#                               sampler=train_sampler,
+#                               batch_size=h.batch_size,
+#                               pin_memory=True,
+#                               drop_last=True)
+    
+    
+    
+
+#     if rank == 0:
+#         validset = MelDataset(validation_filelist, h.segment_size, h.n_fft, h.num_mels,
+#                               h.hop_size, h.win_size, h.sampling_rate, h.fmin, h.fmax, False, False, n_cache_reuse=0,
+#                               fmax_loss=h.fmax_for_loss, device=device, fine_tuning=a.fine_tuning,
+#                               base_mels_path=a.input_mels_dir)
+#         validation_loader = DataLoader(validset, num_workers=1, shuffle=False,
+#                                        sampler=None,
+#                                        batch_size=1,
+#                                        pin_memory=True,
+#                                        drop_last=True)
+
+#         sw = SummaryWriter(os.path.join(a.checkpoint_path, 'logs'))
+
+#     generator.train()
+#     mpd.train()
+#     msd.train()
+#     for epoch in range(max(0, last_epoch), a.training_epochs):
+#         if rank == 0:
+#             start = time.time()
+#             print("Epoch: {}".format(epoch+1))
+
+#         if h.num_gpus > 1:
+#             train_sampler.set_epoch(epoch)
+
+#         for i, batch in enumerate(train_loader):
+#             if rank == 0:
+#                 start_b = time.time()
+#             x, y, _, y_mel = batch
+#             x = torch.autograd.Variable(x.to(device, non_blocking=True))
+#             y = torch.autograd.Variable(y.to(device, non_blocking=True))
+#             y_mel = torch.autograd.Variable(y_mel.to(device, non_blocking=True))
+#             y = y.unsqueeze(1)
+#             # y_g_hat = generator(x)
+#             spec, phase = generator(x)
+
+#             y_g_hat = stft.inverse(spec, phase)
+
+#             y_g_hat_mel = mel_spectrogram(y_g_hat.squeeze(1), h.n_fft, h.num_mels, h.sampling_rate, h.hop_size, h.win_size,
+#                                           h.fmin, h.fmax_for_loss)
+
+#             optim_d.zero_grad()
+
+#             # MPD
+#             y_df_hat_r, y_df_hat_g, _, _ = mpd(y, y_g_hat.detach())
+#             loss_disc_f, losses_disc_f_r, losses_disc_f_g = discriminator_loss(y_df_hat_r, y_df_hat_g)
+#             loss_disc_f += discriminator_TPRLS_loss(y_df_hat_r, y_df_hat_g)
+            
+#             # MSD
+#             y_ds_hat_r, y_ds_hat_g, _, _ = msd(y, y_g_hat.detach())
+#             loss_disc_s, losses_disc_s_r, losses_disc_s_g = discriminator_loss(y_ds_hat_r, y_ds_hat_g)
+#             loss_disc_s += discriminator_TPRLS_loss(y_ds_hat_r, y_ds_hat_g)
+            
+#             loss_disc_all = loss_disc_s + loss_disc_f
+
+#             loss_disc_all.backward()
+#             optim_d.step()
+
+#             # Generator
+#             optim_g.zero_grad()
+
+#             # L1 Mel-Spectrogram Loss
+#             loss_mel = F.l1_loss(y_mel, y_g_hat_mel) * 45
+
+#             y_df_hat_r, y_df_hat_g, fmap_f_r, fmap_f_g = mpd(y, y_g_hat)
+#             y_ds_hat_r, y_ds_hat_g, fmap_s_r, fmap_s_g = msd(y, y_g_hat)
+#             loss_fm_f = feature_loss(fmap_f_r, fmap_f_g)
+#             loss_fm_s = feature_loss(fmap_s_r, fmap_s_g)
+#             loss_gen_f, losses_gen_f = generator_loss(y_df_hat_g)
+#             loss_gen_s, losses_gen_s = generator_loss(y_ds_hat_g)
+            
+#             loss_gen_f += generator_TPRLS_loss(y_df_hat_r, y_df_hat_g)
+#             loss_gen_s += generator_TPRLS_loss(y_ds_hat_r, y_ds_hat_g)
+
+#             loss_gen_all = loss_gen_s + loss_gen_f + loss_fm_s + loss_fm_f + loss_mel
+
+#             loss_gen_all.backward()
+#             optim_g.step()
+
+#             if rank == 0:
+#                 # STDOUT logging
+#                 if steps % a.stdout_interval == 0:
+#                     with torch.no_grad():
+#                         mel_error = F.l1_loss(y_mel, y_g_hat_mel).item()
+
+#                     print('Steps : {:d}, Gen Loss Total : {:4.3f}, Mel-Spec. Error : {:4.3f}, s/b : {:4.3f}'.
+#                           format(steps, loss_gen_all, mel_error, time.time() - start_b))
+
+#                 # checkpointing
+#                 if steps % a.checkpoint_interval == 0 and steps != 0:
+#                     checkpoint_path = "{}/g_{:08d}".format(a.checkpoint_path, steps)
+#                     save_checkpoint(checkpoint_path,
+#                                     {'generator': (generator.module if h.num_gpus > 1 else generator).state_dict()})
+#                     checkpoint_path = "{}/do_{:08d}".format(a.checkpoint_path, steps)
+#                     save_checkpoint(checkpoint_path, 
+#                                     {'mpd': (mpd.module if h.num_gpus > 1
+#                                                          else mpd).state_dict(),
+#                                      'msd': (msd.module if h.num_gpus > 1
+#                                                          else msd).state_dict(),
+#                                      'optim_g': optim_g.state_dict(), 'optim_d': optim_d.state_dict(), 'steps': steps,
+#                                      'epoch': epoch})
+
+#                 # Tensorboard summary logging
+#                 if steps % a.summary_interval == 0:
+#                     sw.add_scalar("training/gen_loss_total", loss_gen_all, steps)
+#                     sw.add_scalar("training/mel_spec_error", mel_error, steps)
+
+#                 # Validation
+#                 if steps % a.validation_interval == 0:  # and steps != 0:
+#                     generator.eval()
+#                     torch.cuda.empty_cache()
+#                     val_err_tot = 0
+#                     with torch.no_grad():
+#                         for j, batch in enumerate(validation_loader):
+#                             x, y, _, y_mel = batch
+#                             # y_g_hat = generator(x.to(device))
+#                             spec, phase = generator(x.to(device))
+
+#                             y_g_hat = stft.inverse(spec, phase)
+
+#                             y_mel = torch.autograd.Variable(y_mel.to(device, non_blocking=True))
+#                             y_g_hat_mel = mel_spectrogram(y_g_hat.squeeze(1), h.n_fft, h.num_mels, h.sampling_rate,
+#                                                           h.hop_size, h.win_size,
+#                                                           h.fmin, h.fmax_for_loss)
+#                             val_err_tot += F.l1_loss(y_mel, y_g_hat_mel).item()
+
+#                             if j <= 4:
+#                                 if steps == 0:
+#                                     sw.add_audio('gt/y_{}'.format(j), y[0], steps, h.sampling_rate)
+#                                     sw.add_figure('gt/y_spec_{}'.format(j), plot_spectrogram(x[0]), steps)
+
+#                                 sw.add_audio('generated/y_hat_{}'.format(j), y_g_hat[0], steps, h.sampling_rate)
+#                                 y_hat_spec = mel_spectrogram(y_g_hat.squeeze(1), h.n_fft, h.num_mels,
+#                                                              h.sampling_rate, h.hop_size, h.win_size,
+#                                                              h.fmin, h.fmax)
+#                                 sw.add_figure('generated/y_hat_spec_{}'.format(j),
+#                                               plot_spectrogram(y_hat_spec.squeeze(0).cpu().numpy()), steps)
+
+#                         val_err = val_err_tot / (j+1)
+#                         sw.add_scalar("validation/mel_spec_error", val_err, steps)
+
+#                     generator.train()
+
+#             steps += 1
+
+#         scheduler_g.step()
+#         scheduler_d.step()
+        
+#         if rank == 0:
+#             print('Time taken for epoch {} is {} sec\n'.format(epoch + 1, int(time.time() - start)))
+
+
+# def main():
+#     print('Initializing Training Process..')
+
+#     parser = argparse.ArgumentParser()
+
+#     parser.add_argument('--group_name', default=None)
+#     parser.add_argument('--input_wavs_dir', default='')
+#     parser.add_argument('--input_mels_dir', default='ft_dataset')
+#     # parser.add_argument('--input_training_file', default='/home/ubuntu/RINGFORMER/LJSpeech-1.1/training.txt')
+#     parser.add_argument('--input_training_file', default='/home/ubuntu/RINGFORMER/LJSpeech-1.1/validation.txt')
+#     parser.add_argument('--input_validation_file', default='/home/ubuntu/RINGFORMER/LJSpeech-1.1/validation.txt')
+#     parser.add_argument('--checkpoint_path', default='cp_ringformer_24')
+#     parser.add_argument('--config', default='config_v1.json')
+#     parser.add_argument('--training_epochs', default=3100, type=int)
+#     parser.add_argument('--stdout_interval', default=5, type=int)
+#     parser.add_argument('--checkpoint_interval', default=5000, type=int)
+#     parser.add_argument('--summary_interval', default=100, type=int)
+#     parser.add_argument('--validation_interval', default=1000, type=int)
+#     parser.add_argument('--fine_tuning', default=False, type=bool)
+
+#     a = parser.parse_args()
+
+#     with open(a.config) as f:
+#         data = f.read()
+
+#     json_config = json.loads(data)
+#     h = AttrDict(json_config)
+#     build_env(a.config, 'config.json', a.checkpoint_path)
+
+#     torch.manual_seed(h.seed)
+#     if torch.cuda.is_available():
+#         torch.cuda.manual_seed(h.seed)
+#         h.num_gpus = torch.cuda.device_count()
+#         h.batch_size = int(h.batch_size / h.num_gpus)
+#         print('Batch size per GPU :', h.batch_size)
+#     else:
+#         pass
+
+#     if h.num_gpus > 1:
+#         mp.spawn(train, nprocs=h.num_gpus, args=(a, h,))
+#     else:
+#         train(0, a, h)
+
+
+# if __name__ == '__main__':
+#     main()
+
+
+import warnings
+warnings.simplefilter(action='ignore', category=FutureWarning)
+import itertools
+import os
+import time
+import argparse
+import json
+import torch
+import torch.nn.functional as F
+from torch.utils.tensorboard import SummaryWriter
+from torch.utils.data import DistributedSampler, DataLoader
+import torch.multiprocessing as mp
+from torch.distributed import init_process_group
+from torch.nn.parallel import DistributedDataParallel
+from env import AttrDict, build_env
+from meldataset import MelDataset, mel_spectrogram, get_dataset_filelist
+from models import Generator, MultiPeriodDiscriminator, feature_loss, generator_loss,\
+    discriminator_loss, discriminator_TPRLS_loss, generator_TPRLS_loss, MultiScaleSubbandCQTDiscriminator
+from utils import plot_spectrogram, scan_checkpoint, load_checkpoint, save_checkpoint
+from stft import TorchSTFT
+from Utils.JDC.model import JDCNet
+from accelerate import Accelerator
+from accelerate.utils import LoggerType
+from accelerate import DistributedDataParallelKwargs
+
+
+torch.backends.cudnn.benchmark = True
+
+
+def train(accelerator, a, h):
+    # if h.num_gpus > 1:
+    #     init_process_group(backend=h.dist_config['dist_backend'], init_method=h.dist_config['dist_url'],
+    #                        world_size=h.dist_config['world_size'] * h.num_gpus, rank=rank)
+        
+
+
+    torch.cuda.manual_seed(h.seed)
+    device = accelerator.device
+
+    F0_model = JDCNet(num_class=1, seq_len=192)
+    params = torch.load(h.F0_path)['model']
+    F0_model.load_state_dict(params)
+    
+    generator = Generator(h, F0_model).to(device)
+    mpd = MultiPeriodDiscriminator().to(device)
+    msd = MultiScaleSubbandCQTDiscriminator().to(device)
+    stft = TorchSTFT(filter_length=h.gen_istft_n_fft, hop_length=h.gen_istft_hop_size, win_length=h.gen_istft_n_fft).to(device)
+
+    with accelerator.main_process_first():
+        accelerator.print(generator)
+        os.makedirs(a.checkpoint_path, exist_ok=True)
+        accelerator.print("checkpoints directory : ", a.checkpoint_path)
+
+    if os.path.isdir(a.checkpoint_path):
+        cp_g = scan_checkpoint(a.checkpoint_path, 'g_')
+        cp_do = scan_checkpoint(a.checkpoint_path, 'do_')
+
+    steps = 0
+    if cp_g is None or cp_do is None:
+        state_dict_do = None
+        last_epoch = -1
+    else:
+        state_dict_g = load_checkpoint(cp_g, device)
+        state_dict_do = load_checkpoint(cp_do, device)
+        generator.load_state_dict(state_dict_g['generator'])
+        mpd.load_state_dict(state_dict_do['mpd'])
+        msd.load_state_dict(state_dict_do['msd'])
+        steps = state_dict_do['steps'] + 1
+        last_epoch = state_dict_do['epoch']
+
+    # if h.num_gpus > 1:
+    generator = accelerator.prepare(generator).to(device)
+    mpd = accelerator.prepare(mpd).to(device)
+    msd = accelerator.prepare(msd).to(device)
+
+    optim_g = accelerator.prepare(torch.optim.AdamW(generator.parameters(), h.learning_rate, betas=[h.adam_b1, h.adam_b2]))
+    optim_d = accelerator.prepare(torch.optim.AdamW(itertools.chain(msd.parameters(), mpd.parameters()),
+                                h.learning_rate, betas=[h.adam_b1, h.adam_b2]))
+    
+    
+
+    if state_dict_do is not None:
+        optim_g.load_state_dict(state_dict_do['optim_g'])
+        optim_d.load_state_dict(state_dict_do['optim_d'])
+
+    scheduler_g = accelerator.prepare(torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=h.lr_decay, last_epoch=last_epoch))
+    scheduler_d = accelerator.prepare(torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=h.lr_decay, last_epoch=last_epoch))
+
+    training_filelist, validation_filelist = get_dataset_filelist(a)
+
+    trainset = MelDataset(training_filelist, h.segment_size, h.n_fft, h.num_mels,
+                          h.hop_size, h.win_size, h.sampling_rate, h.fmin, h.fmax, n_cache_reuse=0,
+                          shuffle=False if h.num_gpus > 1 else True, fmax_loss=h.fmax_for_loss, device=device,
+                          fine_tuning=a.fine_tuning, base_mels_path=a.input_mels_dir)
+
+    # train_sampler = DistributedSampler(trainset) if h.num_gpus > 1 else None
+    
+    accelerator.print("bs", h.batch_size)
+
+    train_loader = DataLoader(trainset, num_workers=h.num_workers, shuffle=False,
+                            #   sampler=train_sampler,
+                              batch_size=h.batch_size,
+                              pin_memory=True,
+                              drop_last=True)
+    
+    
+    train_loader = accelerator.prepare(train_loader)
+    
+    
+
+    validset = MelDataset(validation_filelist, h.segment_size, h.n_fft, h.num_mels,
+                            h.hop_size, h.win_size, h.sampling_rate, h.fmin, h.fmax, False, False, n_cache_reuse=0,
+                            fmax_loss=h.fmax_for_loss, device=device, fine_tuning=a.fine_tuning,
+                            base_mels_path=a.input_mels_dir)
+    
+    
+    
+    validation_loader = DataLoader(validset, num_workers=1, shuffle=False,
+                                    sampler=None,
+                                    batch_size=1,
+                                    pin_memory=True,
+                                    drop_last=True)
+
+    sw = SummaryWriter(os.path.join(a.checkpoint_path, 'logs'))
+    
+    # validation_loader = accelerator.prepare(validation_loader)
+
+    generator.train()
+    mpd.train()
+    msd.train()
+    for epoch in range(max(0, last_epoch), a.training_epochs):
+        with accelerator.main_process_first():
+            start = time.time()
+            accelerator.print("Epoch: {}".format(epoch+1))
+
+        # if h.num_gpus > 1:
+        #     train_sampler.set_epoch(epoch)
+
+        for i, batch in enumerate(train_loader):
+            # with accelerator.main_process_first():
+            start_b = time.time()
+            x, y, _, y_mel = batch
+            x = torch.autograd.Variable(x.to(device, non_blocking=True))
+            y = torch.autograd.Variable(y.to(device, non_blocking=True))
+            y_mel = torch.autograd.Variable(y_mel.to(device, non_blocking=True))
+            y = y.unsqueeze(1)
+            # y_g_hat = generator(x)
+            spec, phase = generator(x)
+
+            y_g_hat = stft.inverse(spec, phase)
+
+            y_g_hat_mel = mel_spectrogram(y_g_hat.squeeze(1), h.n_fft, h.num_mels, h.sampling_rate, h.hop_size, h.win_size,
+                                          h.fmin, h.fmax_for_loss)
+
+            optim_d.zero_grad()
+
+            # MPD
+            y_df_hat_r, y_df_hat_g, _, _ = mpd(y, y_g_hat.detach())
+            loss_disc_f, losses_disc_f_r, losses_disc_f_g = discriminator_loss(y_df_hat_r, y_df_hat_g)
+            loss_disc_f += discriminator_TPRLS_loss(y_df_hat_r, y_df_hat_g)
+            
+            # MSD
+            y_ds_hat_r, y_ds_hat_g, _, _ = msd(y, y_g_hat.detach())
+            loss_disc_s, losses_disc_s_r, losses_disc_s_g = discriminator_loss(y_ds_hat_r, y_ds_hat_g)
+            loss_disc_s += discriminator_TPRLS_loss(y_ds_hat_r, y_ds_hat_g)
+            
+            loss_disc_all = loss_disc_s + loss_disc_f
+            
+            accelerator.backward(loss_disc_all)
+            # accelerator.clip_grad_norm_(mpd.parameters(), max_norm=10.0)
+            # accelerator.clip_grad_norm_(msd.parameters(), max_norm=10.0)
+            
+            # loss_disc_all.backward()
+            optim_d.step()
+
+            # Generator
+            optim_g.zero_grad()
+
+            # L1 Mel-Spectrogram Loss
+            loss_mel = F.l1_loss(y_mel, y_g_hat_mel) * 45
+
+            y_df_hat_r, y_df_hat_g, fmap_f_r, fmap_f_g = mpd(y, y_g_hat)
+            y_ds_hat_r, y_ds_hat_g, fmap_s_r, fmap_s_g = msd(y, y_g_hat)
+            loss_fm_f = feature_loss(fmap_f_r, fmap_f_g)
+            loss_fm_s = feature_loss(fmap_s_r, fmap_s_g)
+            loss_gen_f, losses_gen_f = generator_loss(y_df_hat_g)
+            loss_gen_s, losses_gen_s = generator_loss(y_ds_hat_g)
+            
+            loss_gen_f += generator_TPRLS_loss(y_df_hat_r, y_df_hat_g)
+            loss_gen_s += generator_TPRLS_loss(y_ds_hat_r, y_ds_hat_g)
+
+            loss_gen_all = loss_gen_s + loss_gen_f + loss_fm_s + loss_fm_f + loss_mel
+            
+            
+            accelerator.backward(loss_gen_all)
+            
+            # accelerator.clip_grad_norm_(generator.parameters(), max_norm=10.0)
+
+            # loss_gen_all.backward()
+            optim_g.step()
+            
+            # accelerator.print('done +',i)
+
+
+            # STDOUT logging
+            if steps % a.stdout_interval == 0:
+                
+                with accelerator.main_process_first():
+                    
+                    with torch.no_grad():
+                        mel_error = F.l1_loss(y_mel, y_g_hat_mel).item()
+
+                    accelerator.print('Steps : {:d}, Gen Loss Total : {:4.3f},  Disc Loss Total : {:4.3f}, Mel-Spec. Error : {:4.3f}, s/b : {:4.3f}'.
+                            format(steps, loss_gen_all, loss_disc_all, mel_error, time.time() - start_b))
+
+                # checkpointing
+                # In the main code:
+                    # if steps % a.checkpoint_interval == 0 and steps != 0:
+                    # # with accelerator.main_process_first():
+                    #     # Unwrap the generator
+                    #     unwrapped_generator = accelerator.unwrap_model(generator)
+                    #     checkpoint_path = "{}/g_{:08d}".format(a.checkpoint_path, steps)
+                    #     save_checkpoint(checkpoint_path,
+                    #                 {'generator': unwrapped_generator.state_dict()})
+                        
+                    #     # Unwrap discriminators
+                    #     unwrapped_mpd = accelerator.unwrap_model(mpd)
+                    #     unwrapped_msd = accelerator.unwrap_model(msd)
+                        
+                    #     checkpoint_path = "{}/do_{:08d}".format(a.checkpoint_path, steps)
+                    #     save_checkpoint(checkpoint_path,
+                    #                 {'mpd': unwrapped_mpd.state_dict(),
+                    #                     'msd': unwrapped_msd.state_dict(),
+                    #                     'optim_g': optim_g.state_dict(), 
+                    #                     'optim_d': optim_d.state_dict(), 
+                    #                     'steps': steps,
+                    #                     'epoch': epoch})
+                    
+                    
+                
+                    
+                if steps % a.checkpoint_interval == 0 and steps != 0:
+                    
+                    with accelerator.main_process_first():   
+                        checkpoint_path = "{}/g_{:08d}".format(a.checkpoint_path, steps)
+                        save_checkpoint(checkpoint_path,
+                                        {'generator': (generator.module if h.num_gpus > 1 else generator).state_dict()})
+                        checkpoint_path = "{}/do_{:08d}".format(a.checkpoint_path, steps)
+                        save_checkpoint(checkpoint_path, 
+                                        {'mpd': (mpd.module if h.num_gpus > 1
+                                                            else mpd).state_dict(),
+                                        'msd': (msd.module if h.num_gpus > 1
+                                                            else msd).state_dict(),
+                                        'optim_g': optim_g.state_dict(), 'optim_d': optim_d.state_dict(), 'steps': steps,
+                                        'epoch': epoch})
+
+                    # Tensorboard summary logging
+                    if steps % a.summary_interval == 0:
+                        sw.add_scalar("training/gen_loss_total", loss_gen_all, steps)
+                        sw.add_scalar("training/mel_spec_error", mel_error, steps)
+
+                # Validation
+                if steps % a.validation_interval == 0:  # and steps != 0:
+                    
+                    # with accelerator.main_process_first():
+                    accelerator.print('evalution...')
+                        
+                    
+                    
+                    generator.eval()
+                    torch.cuda.empty_cache()
+                    val_err_tot = 0
+                    with torch.no_grad():
+                        for j, batch in enumerate(validation_loader):
+                            x, y, _, y_mel = batch
+                            # y_g_hat = generator(x.to(device))
+                            spec, phase = generator(x.to('cuda'))
+
+                            y_g_hat = stft.inverse(spec, phase)
+                            
+                            
+
+                            y_mel = torch.autograd.Variable(y_mel.to('cuda', non_blocking=True))
+                            y_g_hat_mel = mel_spectrogram(y_g_hat.squeeze(1), h.n_fft, h.num_mels, h.sampling_rate,
+                                                        h.hop_size, h.win_size,
+                                                        h.fmin, h.fmax_for_loss)
+                            
+                            accelerator.print('done: ',j)
+                            
+                            
+                            val_err_tot += F.l1_loss(y_mel, y_g_hat_mel).item()
+
+                            if j <= 4:
+                                if steps == 0:
+                                    sw.add_audio('gt/y_{}'.format(j), y[0], steps, h.sampling_rate)
+                                    sw.add_figure('gt/y_spec_{}'.format(j), plot_spectrogram(x[0]), steps)
+                                    
+
+                                sw.add_audio('generated/y_hat_{}'.format(j), y_g_hat[0], steps, h.sampling_rate)
+                                y_hat_spec = mel_spectrogram(y_g_hat.squeeze(1), h.n_fft, h.num_mels,
+                                                            h.sampling_rate, h.hop_size, h.win_size,
+                                                            h.fmin, h.fmax)
+                                
+                                
+                                sw.add_figure('generated/y_hat_spec_{}'.format(j),
+                                            plot_spectrogram(y_hat_spec.squeeze(0).cpu().numpy()), steps)
+                                
+                                
+                        
+                        val_err = val_err_tot / (j+1)
+                        sw.add_scalar("validation/mel_spec_error", val_err, steps)
+                    
+
+                generator.train()
+
+            steps += 1
+
+        scheduler_g.step()
+        scheduler_d.step()
+        
+        # with accelerator.main_process_first():
+        accelerator.print('Time taken for epoch {} is {} sec\n'.format(epoch + 1, int(time.time() - start)))
+
+
+def main():
+    
+    ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+
+    
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--group_name', default=None)
+    parser.add_argument('--input_wavs_dir', default='')
+    parser.add_argument('--input_mels_dir', default='ft_dataset')
+    # parser.add_argument('--input_training_file', default='/home/ubuntu/RINGFORMER/LJSpeech-1.1/training.txt')
+    parser.add_argument('--input_training_file', default='/home/ubuntu/respair/audio_files_over_2sec_SUB.txt')
+    parser.add_argument('--input_validation_file', default='/home/ubuntu/RINGFORMER/LJSpeech-1.1/validation.txt')
+    parser.add_argument('--checkpoint_path', default='cp_ringformer_44.1khz_NUM2')
+    parser.add_argument('--config', default='config_v1.json')
+    parser.add_argument('--training_epochs', default=3100, type=int)
+    parser.add_argument('--stdout_interval', default=10, type=int)
+    parser.add_argument('--checkpoint_interval', default=1000, type=int)
+    parser.add_argument('--summary_interval', default=100, type=int)
+    parser.add_argument('--validation_interval', default=1000, type=int)
+    parser.add_argument('--fine_tuning', default=False, type=bool)
+
+    a = parser.parse_args()
+    
+    accelerator = Accelerator(project_dir=a.checkpoint_path, split_batches=False, 
+                            kwargs_handlers=[ddp_kwargs])
+  
+    accelerator.print('Initializing Training Process..')
+
+    with open(a.config) as f:
+        data = f.read()
+
+    json_config = json.loads(data)
+    h = AttrDict(json_config)
+    build_env(a.config, 'config.json', a.checkpoint_path)
+
+    torch.manual_seed(h.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(h.seed)
+        h.num_gpus = torch.cuda.device_count()
+        # h.batch_size = int(h.batch_size / h.num_gpus)
+        # print('Batch size per GPU :', h.batch_size)
+    else:
+        pass
+  # Pass accelerator to train function instead of rank
+    train(accelerator, a, h)
+
+if __name__ == '__main__':
+    main()

RingFormer/utils.py

@@ -0,0 +1,58 @@
+import glob
+import os
+import matplotlib
+import torch
+from torch.nn.utils import weight_norm
+matplotlib.use("Agg")
+import matplotlib.pylab as plt
+
+
+def plot_spectrogram(spectrogram):
+    fig, ax = plt.subplots(figsize=(10, 2))
+    im = ax.imshow(spectrogram, aspect="auto", origin="lower",
+                   interpolation='none')
+    plt.colorbar(im, ax=ax)
+
+    fig.canvas.draw()
+    plt.close()
+
+    return fig
+
+
+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)
+
+
+def load_checkpoint(filepath, device):
+    assert os.path.isfile(filepath)
+    print("Loading '{}'".format(filepath))
+    checkpoint_dict = torch.load(filepath, map_location=device)
+    print("Complete.")
+    return checkpoint_dict
+
+
+def save_checkpoint(filepath, obj):
+    print("Saving checkpoint to {}".format(filepath))
+    torch.save(obj, filepath)
+    print("Complete.")
+
+
+def scan_checkpoint(cp_dir, prefix):
+    pattern = os.path.join(cp_dir, prefix + '????????')
+    cp_list = glob.glob(pattern)
+    if len(cp_list) == 0:
+        return None
+    return sorted(cp_list)[-1]
+