update test code

.gitignore

@@ -129,3 +129,4 @@ dmypy.json
 .pyre/
 
 /logs/
+*.wav

configs/16k.json

@@ -0,0 +1,54 @@
+{
+    "trainer": {
+      "max_epochs": 20000,
+      "limit_val_batches": 1,
+      "accumulate_grad_batches": 1,
+      "default_root_dir": "./logs",
+      "val_check_interval": 1000
+    },
+    "train": {
+      "log_interval": 200,
+      "eval_interval": 1000,
+      "seed": 1234,
+      "max_epochs": 20000,
+      "learning_rate": 2e-4,
+      "betas": [0.8, 0.99],
+      "eps": 1e-9,
+      "batch_size": 64,
+      "fp16_run": true,
+      "lr_decay": 0.999875,
+      "segment_size": 16384,
+      "init_lr_ratio": 1,
+      "warmup_epochs": 0,
+      "c_mel": 45
+    },
+    "data": {
+      "training_files": "filelists/example_audio_filelist_train.txt",
+      "validation_files": "filelists/example_audio_filelist_valid.txt",
+      "max_wav_value": 32768.0,
+      "sampling_rate": 16000,
+      "filter_length": 1024,
+      "hop_length": 256,
+      "win_length": 1024,
+      "n_mel_channels": 256,
+      "mel_fmin": 0.0,
+      "mel_fmax": null
+    },
+    "model": {
+      "inter_channels": 256,
+      "hidden_channels": 256,
+      "filter_channels": 768,
+      "kernel_size": 3,
+      "p_dropout": 0.1,
+      "resblock": "1",
+      "resblock_kernel_sizes": [3,7,11],
+      "resblock_dilation_sizes": [
+        [1,3,5],
+        [1,3,5],
+        [1,3,5]
+      ],
+      "upsample_rates": [8,8,4],
+      "upsample_initial_channel": 512,
+      "upsample_kernel_sizes": [16,16,8]
+    }
+  }

configs/base.json → configs/48k.json

@@ -14,7 +14,7 @@
       "learning_rate": 2e-4,
       "betas": [0.8, 0.99],
       "eps": 1e-9,
-      "batch_size": 32,
+      "batch_size": 64,
       "fp16_run": true,
       "lr_decay": 0.999875,
       "segment_size": 16384,
@@ -37,10 +37,7 @@
     "model": {
       "inter_channels": 256,
       "hidden_channels": 256,
-      "hubert_channels": 1280,
       "filter_channels": 768,
-      "n_heads": 4,
-      "n_layers": 6,
       "kernel_size": 3,
       "p_dropout": 0.1,
       "resblock": "1",
@@ -52,9 +49,6 @@
       ],
       "upsample_rates": [8,8,4,2],
       "upsample_initial_channel": 512,
-      "upsample_kernel_sizes": [16,16,4,4],
-      "n_layers_q": 3,
-      "use_spectral_norm": false,
-      "gin_channels": 256
+      "upsample_kernel_sizes": [16,16,4,4]
     }
   }

hifigan/data/dataset.py

@@ -9,6 +9,8 @@
 
 from ..utils import load_filepaths, load_wav_to_torch
 
+resamplers = {}
+
 def load_audio(filename: str, sr: Optional[int] = None):
     global resamplers
     audio, sampling_rate = load_wav_to_torch(filename)

hifigan/model/generators/generator.py

@@ -12,15 +12,15 @@
 from ..commons import init_weights, get_padding
 
 
-class Generator(torch.nn.Module):
+class HifiGANGenerator(torch.nn.Module):
     def __init__(self, initial_channel: int,
                     resblock: Union[str, ResBlock2],
                     resblock_kernel_sizes: List[int],
                     resblock_dilation_sizes: List[int],
                     upsample_rates: List[int],
                     upsample_initial_channel: int,
                     upsample_kernel_sizes: List[int]):
-        super(Generator, self).__init__()
+        super(HifiGANGenerator, self).__init__()
         self.num_kernels = len(resblock_kernel_sizes)
         self.num_upsamples = len(upsample_rates)
         self.conv_pre = Conv1d(in_channels=initial_channel, out_channels=upsample_initial_channel, kernel_size=7, stride=1, padding=3)

hifigan/model/hifigan.py

@@ -15,7 +15,7 @@
 
 from .discriminators.multi_scale_discriminator import MultiScaleDiscriminator
 from .discriminators.multi_period_discriminator import MultiPeriodDiscriminator
-from .generators.generator import Generator
+from .generators.generator import HifiGANGenerator
 
 from ..mel_processing import spec_to_mel_torch, mel_spectrogram_torch, spectrogram_torch, spectrogram_torch_audio
 from .losses import discriminator_loss, kl_loss,feature_loss, generator_loss
@@ -27,7 +27,7 @@ def __init__(self, **kwargs):
         super().__init__()
         self.save_hyperparameters(*[k for k in kwargs])
 
-        self.net_g = Generator(
+        self.net_g = HifiGANGenerator(
             self.hparams.model.inter_channels,
             self.hparams.model.resblock,
             self.hparams.model.resblock_kernel_sizes,
@@ -163,11 +163,8 @@ def validation_step(self, batch, batch_idx):
         y_spec_lengths = (y_wav_lengths / self.hparams.data.hop_length).long()
 
         # remove else
-        y_spec = y_spec[:1]
-        y_spec_lengths = y_spec_lengths[:1]
-
         y_hat = self.net_g(x_mel)
-        y_hat_lengths = y_spec_lengths
+        y_hat_lengths = torch.tensor([y_hat.shape[2]], dtype=torch.long)
 
         mel = spec_to_mel_torch(
             y_spec, 

hub_conf.py

@@ -0,0 +1,3 @@
+dependencies = ["torch", "torchaudio"]
+
+from hifigan.model.generators.generator import HifiGANGenerator

test.py

@@ -0,0 +1,15 @@
+import torch, torchaudio
+
+from hifigan.mel_processing import mel_spectrogram_torch
+
+# Load checkpoint (either hubert_soft or hubert_discrete)
+hifigan = torch.hub.load("vtuber-plan/hifi-gan:main", "HifiGANGenerator").cuda()
+
+# Load audio
+wav, sr = torchaudio.load("test.wav")
+assert sr == 48000
+wav = wav.unsqueeze(0).cuda()
+
+mel = mel_spectrogram_torch(wav, 2048, 256, 48000, 512, 2048, 0, None, False)
+
+units = hifigan(mel)

train.py

@@ -31,7 +31,7 @@ def get_hparams(config_path: str) -> HParams:
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument('-c', '--config', type=str, default="./configs/base.json", help='JSON file for configuration')
+    parser.add_argument('-c', '--config', type=str, default="./configs/48k.json", help='JSON file for configuration')
     parser.add_argument('-a', '--accelerator', type=str, default="gpu", help='training device')
     parser.add_argument('-d', '--device', type=str, default="0", help='training device ids')
     args = parser.parse_args()