Format code

Author: neoncloud

models/networks.py

@@ -116,15 +116,17 @@ def get_target_tensor(self, input, target_is_real):
             create_label = ((self.real_label_var is None) or
                             (self.real_label_var.shape != input.shape))
             if create_label:
-                real_tensor = torch.Tensor(input.size()).fill_(self.real_label).to(self.device)
+                real_tensor = torch.Tensor(input.size()).fill_(
+                    self.real_label).to(self.device)
                 self.real_label_var = Variable(
                     real_tensor, requires_grad=False)
             target_tensor = self.real_label_var
         else:
             create_label = ((self.fake_label_var is None) or
                             (self.fake_label_var.shape != input.shape))
             if create_label:
-                fake_tensor = torch.Tensor(input.size()).fill_(self.fake_label).to(self.device)
+                fake_tensor = torch.Tensor(input.size()).fill_(
+                    self.fake_label).to(self.device)
                 self.fake_label_var = Variable(
                     fake_tensor, requires_grad=False)
             target_tensor = self.fake_label_var
@@ -213,41 +215,41 @@ def __init__(self, input_nc, output_nc, ngf=32, n_downsample_global=3, n_blocks_
         model_downsample = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0),
                             norm_layer(ngf_global), nn.ReLU(True),
                             downsample_layer(ngf_global, ngf_global * 2,
-                                                kernel_size=3, stride=2, padding=1),
+                                             kernel_size=3, stride=2, padding=1),
                             norm_layer(ngf_global * 2), nn.ReLU(True)]
         # residual blocks
         model_upsample = []
         for i in range(n_blocks_local):
             model_upsample += [ResnetBlock(ngf_global * 2,
-                                            padding_type=padding_type, norm_layer=norm_layer)]
+                                           padding_type=padding_type, norm_layer=norm_layer)]
         # attention bottleneck
         if n_attn_l > 0:
             middle = n_blocks_local//2
             # 8x downsample
             down = [downsample_layer(ngf_global * 2, ngf_global,
-                                        kernel_size=3, stride=2, padding=1),
+                                     kernel_size=3, stride=2, padding=1),
                     norm_layer(ngf_global), nn.ReLU(True)]
             down += [downsample_layer(ngf_global, ngf_global,
-                                        kernel_size=3, stride=2, padding=1),
-                        norm_layer(ngf_global), nn.ReLU(True)]*2
+                                      kernel_size=3, stride=2, padding=1),
+                     norm_layer(ngf_global), nn.ReLU(True)]*2
             down = nn.Sequential(*down)
             model_upsample.insert(middle, down)
 
             middle += 1
             input_size = tuple(map(lambda x: x//16, input_size))
             from bottleneck_transformer_pytorch import BottleStack
             attn_block = BottleStack(dim=ngf_global, fmap_size=input_size, dim_out=ngf_global*2, num_layers=n_attn_l, proj_factor=proj_factor_l,
-                                        downsample=False, heads=heads_l, dim_head=dim_head_l, activation=nn.ReLU(True), rel_pos_emb=False)
+                                     downsample=False, heads=heads_l, dim_head=dim_head_l, activation=nn.ReLU(True), rel_pos_emb=False)
             model_upsample.insert(middle, attn_block)
             model_upsample += [upsample_layer(in_channels=ngf_global*2, out_channels=ngf_global*2, kernel_size=3, stride=2, padding=1, output_padding=1),
-                                norm_layer(ngf_global), nn.ReLU(True)]*3
+                               norm_layer(ngf_global), nn.ReLU(True)]*3
 
         model_upsample += [upsample_layer(in_channels=ngf_global*2, out_channels=ngf_global, kernel_size=3, stride=2, padding=1, output_padding=1),
-                            norm_layer(ngf_global), nn.ReLU(True)]
+                           norm_layer(ngf_global), nn.ReLU(True)]
 
         # final convolution
         model_upsample += [nn.ReflectionPad2d(3), nn.Conv2d(
-                ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
+            ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
 
         self.model1_1 = nn.Sequential(*model_downsample)
         self.model1_2 = nn.Sequential(*model_upsample)

models/pix2pixHD_model.py

@@ -10,9 +10,10 @@
 #from dct.dct_native import DCT_2N_native, IDCT_2N_native
 import torchaudio.functional as aF
 
+
 class Audio2Spectro(torch.nn.Module):
     def __init__(self, opt) -> None:
-        super(Audio2Spectro,self).__init__()
+        super(Audio2Spectro, self).__init__()
         opt_dict = vars(opt)
         for k, v in opt_dict.items():
             setattr(self, k, v)
@@ -27,7 +28,7 @@ def __init__(self, opt) -> None:
         self._imdct = IMDCT4(n_fft=self.n_fft, hop_length=self.hop_length,
                              win_length=self.win_length, window=self.window, device=self.device)
 
-    def to_spectro(self, audio:torch.Tensor, mask:bool=False, mask_size:int=-1):
+    def to_spectro(self, audio: torch.Tensor, mask: bool = False, mask_size: int = -1):
         # Forward Transformation (MDCT)
         spectro, frames = self._mdct(audio.to(self.device), True)
         spectro = spectro.unsqueeze(1)
@@ -59,12 +60,14 @@ def to_spectro(self, audio:torch.Tensor, mask:bool=False, mask_size:int=-1):
                 mask_size = int(size[3]*(1-1/self.up_ratio))
 
             # fill the blank mask with noise
-            _noise = torch.randn(size[0], size[1], size[2], mask_size, device=self.device)
+            _noise = torch.randn(
+                size[0], size[1], size[2], mask_size, device=self.device)
             _noise_min = _noise.min()
             _noise_max = _noise.max()
 
             if self.fit_residual:
-                _noise = torch.zeros(size[0], size[1], size[2], mask_size, device=self.device)
+                _noise = torch.zeros(
+                    size[0], size[1], size[2], mask_size, device=self.device)
             else:
                 # fill empty with randn noise, single peak, centered at 0
                 _noise = _noise/(_noise_max - _noise_min)
@@ -108,16 +111,18 @@ def normalize(self, spectro):
             audio_min = log_spectro.flatten(-2).min(dim=-
                                                     1).values[:, :, None, None].float()
         else:
-            audio_min = torch.tensor([self.src_range[0]])[None,None,None,:].to(self.device)
-            audio_max = torch.tensor([self.src_range[1]])[None,None,None,:].to(self.device)
+            audio_min = torch.tensor([self.src_range[0]])[
+                None, None, None, :].to(self.device)
+            audio_max = torch.tensor([self.src_range[1]])[
+                None, None, None, :].to(self.device)
         log_spectro = (log_spectro-audio_min)/(audio_max-audio_min)
         log_spectro = log_spectro * \
             (self.norm_range[1]-self.norm_range[0]
              )+self.norm_range[0]
 
         return log_spectro, audio_max, audio_min, mean, std
 
-    def denormalize(self, log_spectro:torch.Tensor, min:torch.Tensor, max:torch.Tensor):
+    def denormalize(self, log_spectro: torch.Tensor, min: torch.Tensor, max: torch.Tensor):
         log_spectro = (
             log_spectro.to(torch.float64)-self.norm_range[0])/(self.norm_range[1]-self.norm_range[0])
         log_spectro = log_spectro*(max-min)+min
@@ -127,8 +132,9 @@ def denormalize(self, log_spectro:torch.Tensor, min:torch.Tensor, max:torch.Tens
         else:
             return aF.DB_to_amplitude(log_spectro.to(self.device), 10.0, 0.5)-self.min_value
 
-    def to_audio(self, log_spectro:torch.Tensor, norm_param:Dict[str,torch.Tensor], pha:torch.Tensor):
-        spectro = self.denormalize(log_spectro, norm_param['min'], norm_param['max'])
+    def to_audio(self, log_spectro: torch.Tensor, norm_param: Dict[str, torch.Tensor], pha: torch.Tensor):
+        spectro = self.denormalize(
+            log_spectro, norm_param['min'], norm_param['max'])
         if self.explicit_encoding:
             spectro = (spectro[..., 0, :, :] -
                        spectro[..., 1, :, :])/(2*self.alpha-1)
@@ -151,7 +157,8 @@ def to_audio(self, log_spectro:torch.Tensor, norm_param:Dict[str,torch.Tensor],
         return audio
 
     def to_frames(self, log_spectro, norm_param):
-        spectro = self.denormalize(log_spectro, norm_param['min'],norm_param['max'])
+        spectro = self.denormalize(
+            log_spectro, norm_param['min'], norm_param['max'])
         if self.explicit_encoding:
             spectro = (spectro[..., 0, :, :] -
                        spectro[..., 1, :, :])/(2*self.alpha-1)
@@ -165,21 +172,22 @@ def norm_frames(self, frames):
         frames = frames / frames.max()
         return frames * (self.norm_range[1]-self.norm_range[0]) + self.norm_range[0]
 
-    def forward(self, lr_audio:torch.Tensor):
+    def forward(self, lr_audio: torch.Tensor):
         # low-res audio for training
         with torch.no_grad():
             lr_spectro, lr_pha, lr_norm_param = self.to_spectro(
                 lr_audio, mask=self.mask)
         return lr_spectro, lr_pha, lr_norm_param
 
-    def hr_forward(self, hr_audio:torch.Tensor):
+    def hr_forward(self, hr_audio: torch.Tensor):
         # high-res audio for training
         with torch.no_grad():
             hr_spectro, hr_pha, hr_norm_param = self.to_spectro(hr_audio, mask=self.mask_hr, mask_size=int(
                 self.n_fft*(1-self.sr_sampling_rate/self.hr_sampling_rate)//2))
 
         return hr_spectro, hr_pha, hr_norm_param
 
+
 class Pix2PixHDModel(BaseModel):
     def name(self):
         return 'Pix2PixHDModel'
@@ -376,7 +384,8 @@ def discriminate_hifi(self, input, norm_param=None, pha=None, is_spectro=True):
     def forward(self, lr_audio, hr_audio):
         # Encode Inputs
         lr_spectro, lr_pha, lr_norm_param = self.preprocess.forward(lr_audio)
-        hr_spectro, hr_pha, hr_norm_param = self.preprocess.hr_forward(hr_audio)
+        hr_spectro, hr_pha, hr_norm_param = self.preprocess.hr_forward(
+            hr_audio)
         #### G Forward ####
         if self.abs_spectro and self.arcsinh_transform:
             lr_input = lr_spectro.abs()*2+self.norm_range[0]
@@ -395,11 +404,14 @@ def forward(self, lr_audio, hr_audio):
         return sr_spectro, sr_pha, hr_spectro, hr_pha, hr_norm_param, lr_spectro, lr_pha, lr_norm_param
 
     def _forward(self, lr_audio, hr_audio, infer=False):
-        sr_spectro, sr_pha, hr_spectro, hr_pha, hr_norm_param, lr_spectro, lr_pha, lr_norm_param = self.forward(lr_audio, hr_audio)
+        sr_spectro, sr_pha, hr_spectro, hr_pha, hr_norm_param, lr_spectro, lr_pha, lr_norm_param = self.forward(
+            lr_audio, hr_audio)
         # Fake Detection and Loss
         if self.abs_spectro and self.arcsinh_transform:
-            sr_input = torch.cat((sr_spectro, sr_spectro.abs()*2+self.norm_range[0]), dim=1)
-            hr_input = torch.cat((hr_spectro, hr_spectro.abs()*2+self.norm_range[0]), dim=1)
+            sr_input = torch.cat(
+                (sr_spectro, sr_spectro.abs()*2+self.norm_range[0]), dim=1)
+            hr_input = torch.cat(
+                (hr_spectro, hr_spectro.abs()*2+self.norm_range[0]), dim=1)
         else:
             sr_input = sr_spectro
             hr_input = hr_spectro
@@ -584,7 +596,7 @@ def inference(self, lr_audio):
         # Encode Inputs
         with torch.no_grad():
             lr_spectro, lr_pha, lr_norm_param = self.preprocess.forward(
-            lr_audio)
+                lr_audio)
 
             if self.abs_spectro and self.arcsinh_transform:
                 lr_input = lr_spectro.abs()*2+self.norm_range[0]
@@ -673,4 +685,4 @@ def get_current_visuals(self):
 
 class InferenceModel(Pix2PixHDModel):
     def forward(self, lr_audio):
-        return self.inference(lr_audio)
+        return self.inference(lr_audio)

run_script.py

@@ -4,13 +4,10 @@
 
 from options.train_options import TrainOptions
 from data.data_loader import CreateDataLoader
-from util.visualizer import Visualizer
-from util.spectro_img import compute_visuals
 from util.util import compute_matrics
 
 # Initilize the setup
 opt = TrainOptions().parse()
-visualizer = Visualizer(opt)
 data_loader = CreateDataLoader(opt)
 dataset = data_loader.load_data()
 dataset_size = len(data_loader)
@@ -47,9 +44,6 @@
 print('MSE: %.4f' % _mse)
 print('SNR_SR: %.4f' % _snr_sr)
 print('SNR_LR: %.4f' % _snr_lr)
-#print('SSNR_SR: %.4f' % _ssnr_sr)
-#print('SSNR_LR: %.4f' % _ssnr_lr)
-#print('PESQ: %.4f' % _pesq)
 print('LSD: %.4f' % _lsd)
 with open(os.path.join(opt.checkpoints_dir, opt.name, 'metric.txt'), 'w') as f:
     f.write('MSE,SNR_SR,LSD\n')

save_model.py

@@ -9,4 +9,5 @@
 opt.isTrain = False
 model = create_model(opt)
 model_scripted = torch.jit.script(model)
-torch.jit.save(model_scripted,os.path.join(opt.checkpoints_dir, opt.name, 'model_scripted.pt'))
+torch.jit.save(model_scripted,os.path.join(opt.checkpoints_dir, opt.name, 'model_scripted.pt'))
+torch.save(opt,os.path.join(opt.checkpoints_dir, opt.name, 'opt.pt'))