Upload 4 files

data_utils.py

@@ -1,14 +1,13 @@
-import time
 import os
 import random
+
 import numpy as np
 import torch
 import torch.utils.data
 
-import modules.commons as commons
 import utils
-from modules.mel_processing import spectrogram_torch, spec_to_mel_torch, spectrogram_torch
-from utils import load_wav_to_torch, load_filepaths_and_text
+from modules.mel_processing import spectrogram_torch
+from utils import load_filepaths_and_text, load_wav_to_torch
 
 # import h5py
 
@@ -31,6 +30,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         self.filter_length = hparams.data.filter_length
         self.hop_length = hparams.data.hop_length
         self.win_length = hparams.data.win_length
+        self.unit_interpolate_mode = hparams.data.unit_interpolate_mode
         self.sampling_rate = hparams.data.sampling_rate
         self.use_sr = hparams.train.use_sr
         self.spec_len = hparams.train.max_speclen
@@ -73,7 +73,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         uv = torch.FloatTensor(np.array(uv,dtype=float))
 
         c = torch.load(filename+ ".soft.pt")
-        c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[0])
+        c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[0], mode=self.unit_interpolate_mode)
         if self.vol_emb:
             volume_path = filename + ".vol.npy"
             volume = np.load(volume_path)
@@ -86,7 +86,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         assert abs(audio_norm.shape[1]-lmin * self.hop_length) < 3 * self.hop_length
         spec, c, f0, uv = spec[:, :lmin], c[:, :lmin], f0[:lmin], uv[:lmin]
         audio_norm = audio_norm[:, :lmin * self.hop_length]
-        if volume!= None:
+        if volume is not None:
             volume = volume[:lmin]
         return c, f0, spec, audio_norm, spk, uv, volume
 
@@ -95,7 +95,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
         #     print("skip too short audio:", filename)
         #     return None
 
-        if random.choice([True, False]) and self.vol_aug and volume!=None:
+        if random.choice([True, False]) and self.vol_aug and volume is not None:
             max_amp = float(torch.max(torch.abs(audio_norm))) + 1e-5
             max_shift = min(1, np.log10(1/max_amp))
             log10_vol_shift = random.uniform(-1, max_shift)
@@ -113,7 +113,7 @@ class TextAudioSpeakerLoader(torch.utils.data.Dataset):
             end = start + 790
             spec, c, f0, uv = spec[:, start:end], c[:, start:end], f0[start:end], uv[start:end]
             audio_norm = audio_norm[:, start * self.hop_length : end * self.hop_length]
-            if volume !=None:
+            if volume is not None:
                 volume = volume[start:end]
         return c, f0, spec, audio_norm, spk, uv,volume
 
@@ -177,7 +177,7 @@ class TextAudioCollate:
             uv = row[5]
             uv_padded[i, :uv.size(0)] = uv
             volume = row[6]
-            if volume != None:
+            if volume is not None:
                 volume_padded[i, :volume.size(0)] = volume
             else :
                 volume_padded = None

inference_main.py

@@ -1,16 +1,10 @@
-import io
 import logging
-import time
-from pathlib import Path
 
-import librosa
-import matplotlib.pyplot as plt
-import numpy as np
 import soundfile
 
 from inference import infer_tool
-from inference import slicer
 from inference.infer_tool import Svc
+from spkmix import spk_mix_map
 
 logging.getLogger('numba').setLevel(logging.WARNING)
 chunks_dict = infer_tool.read_temp("inference/chunks_temp.json")
@@ -23,21 +17,33 @@ def main():
     parser = argparse.ArgumentParser(description='sovits4 inference')
 
     # 一定要设置的部分
-    parser.add_argument('-m', '--model_path', type=str, default="logs/44k/G_0.pth", help='模型路径')
-    parser.add_argument('-c', '--config_path', type=str, default="configs/config.json", help='配置文件路径')
+    parser.add_argument('-m', '--model_path', type=str, default="logs/44k/G_37600.pth", help='模型路径')
+    parser.add_argument('-c', '--config_path', type=str, default="logs/44k/config.json", help='配置文件路径')
     parser.add_argument('-cl', '--clip', type=float, default=0, help='音频强制切片，默认0为自动切片，单位为秒/s')
     parser.add_argument('-n', '--clean_names', type=str, nargs='+', default=["君の知らない物語-src.wav"], help='wav文件名列表，放在raw文件夹下')
     parser.add_argument('-t', '--trans', type=int, nargs='+', default=[0], help='音高调整，支持正负（半音）')
-    parser.add_argument('-s', '--spk_list', type=str, nargs='+', default=['nen'], help='合成目标说话人名称')
-
+    parser.add_argument('-s', '--spk_list', type=str, nargs='+', default=['buyizi'], help='合成目标说话人名称')
+    
     # 可选项部分
-    parser.add_argument('-a', '--auto_predict_f0', action='store_true', default=False,help='语音转换自动预测音高，转换歌声时不要打开这个会严重跑调')
-    parser.add_argument('-cm', '--cluster_model_path', type=str, default="logs/44k/kmeans_10000.pt", help='聚类模型路径，如果没有训练聚类则随便填')
-    parser.add_argument('-cr', '--cluster_infer_ratio', type=float, default=0, help='聚类方案占比，范围0-1，若没有训练聚类模型则默认0即可')
+    parser.add_argument('-a', '--auto_predict_f0', action='store_true', default=False, help='语音转换自动预测音高，转换歌声时不要打开这个会严重跑调')
+    parser.add_argument('-cm', '--cluster_model_path', type=str, default="", help='聚类模型或特征检索索引路径，留空则自动设为各方案模型的默认路径，如果没有训练聚类或特征检索则随便填')
+    parser.add_argument('-cr', '--cluster_infer_ratio', type=float, default=0, help='聚类方案或特征检索占比，范围0-1，若没有训练聚类模型或特征检索则默认0即可')
     parser.add_argument('-lg', '--linear_gradient', type=float, default=0, help='两段音频切片的交叉淡入长度，如果强制切片后出现人声不连贯可调整该数值，如果连贯建议采用默认值0，单位为秒')
-    parser.add_argument('-fmp', '--f0_mean_pooling', type=bool, default=False, help='是否对F0使用均值滤波器(池化)，对部分哑音有改善。注意，启动该选项会导致推理速度下降，默认关闭')
-    parser.add_argument('-eh', '--enhance', type=bool, default=False, help='是否使用NSF_HIFIGAN增强器,该选项对部分训练集少的模型有一定的音质增强效果，但是对训练好的模型有反面效果，默认关闭')
- 
+    parser.add_argument('-f0p', '--f0_predictor', type=str, default="pm", help='选择F0预测器,可选择crepe,pm,dio,harvest,rmvpe,默认为pm(注意：crepe为原F0使用均值滤波器)')
+    parser.add_argument('-eh', '--enhance', action='store_true', default=False, help='是否使用NSF_HIFIGAN增强器,该选项对部分训练集少的模型有一定的音质增强效果，但是对训练好的模型有反面效果，默认关闭')
+    parser.add_argument('-shd', '--shallow_diffusion', action='store_true', default=False, help='是否使用浅层扩散，使用后可解决一部分电音问题，默认关闭，该选项打开时，NSF_HIFIGAN增强器将会被禁止')
+    parser.add_argument('-usm', '--use_spk_mix', action='store_true', default=False, help='是否使用角色融合')
+    parser.add_argument('-lea', '--loudness_envelope_adjustment', type=float, default=1, help='输入源响度包络替换输出响度包络融合比例，越靠近1越使用输出响度包络')
+    parser.add_argument('-fr', '--feature_retrieval', action='store_true', default=False, help='是否使用特征检索，如果使用聚类模型将被禁用，且cm与cr参数将会变成特征检索的索引路径与混合比例')
+
+    # 浅扩散设置
+    parser.add_argument('-dm', '--diffusion_model_path', type=str, default="logs/44k/diffusion/model_0.pt", help='扩散模型路径')
+    parser.add_argument('-dc', '--diffusion_config_path', type=str, default="logs/44k/diffusion/config.yaml", help='扩散模型配置文件路径')
+    parser.add_argument('-ks', '--k_step', type=int, default=100, help='扩散步数，越大越接近扩散模型的结果，默认100')
+    parser.add_argument('-se', '--second_encoding', action='store_true', default=False, help='二次编码，浅扩散前会对原始音频进行二次编码，玄学选项，有时候效果好，有时候效果差')
+    parser.add_argument('-od', '--only_diffusion', action='store_true', default=False, help='纯扩散模式，该模式不会加载sovits模型，以扩散模型推理')
+    
+
     # 不用动的部分
     parser.add_argument('-sd', '--slice_db', type=int, default=-40, help='默认-40，嘈杂的音频可以-30，干声保留呼吸可以-50')
     parser.add_argument('-d', '--device', type=str, default=None, help='推理设备，None则为自动选择cpu和gpu')
@@ -46,7 +52,9 @@ def main():
     parser.add_argument('-wf', '--wav_format', type=str, default='flac', help='音频输出格式')
     parser.add_argument('-lgr', '--linear_gradient_retain', type=float, default=0.75, help='自动音频切片后，需要舍弃每段切片的头尾。该参数设置交叉长度保留的比例，范围0-1,左开右闭')
     parser.add_argument('-eak', '--enhancer_adaptive_key', type=int, default=0, help='使增强器适应更高的音域(单位为半音数)|默认为0')
-    
+    parser.add_argument('-ft', '--f0_filter_threshold', type=float, default=0.05,help='F0过滤阈值，只有使用crepe时有效. 数值范围从0-1. 降低该值可减少跑调概率，但会增加哑音')
+
+
     args = parser.parse_args()
 
     clean_names = args.clean_names
@@ -61,75 +69,85 @@ def main():
     clip = args.clip
     lg = args.linear_gradient
     lgr = args.linear_gradient_retain
-    F0_mean_pooling = args.f0_mean_pooling
+    f0p = args.f0_predictor
     enhance = args.enhance
     enhancer_adaptive_key = args.enhancer_adaptive_key
+    cr_threshold = args.f0_filter_threshold
+    diffusion_model_path = args.diffusion_model_path
+    diffusion_config_path = args.diffusion_config_path
+    k_step = args.k_step
+    only_diffusion = args.only_diffusion
+    shallow_diffusion = args.shallow_diffusion
+    use_spk_mix = args.use_spk_mix
+    second_encoding = args.second_encoding
+    loudness_envelope_adjustment = args.loudness_envelope_adjustment
 
-    svc_model = Svc(args.model_path, args.config_path, args.device, args.cluster_model_path,enhance)
-    infer_tool.mkdir(["raw", "results"])
+    if cluster_infer_ratio != 0:
+        if args.cluster_model_path == "":
+            if args.feature_retrieval:  # 若指定了占比但没有指定模型路径，则按是否使用特征检索分配默认的模型路径
+                args.cluster_model_path = "logs/44k/feature_and_index.pkl"
+            else:
+                args.cluster_model_path = "logs/44k/kmeans_10000.pt"
+    else:  # 若未指定占比，则无论是否指定模型路径，都将其置空以避免之后的模型加载
+        args.cluster_model_path = ""
 
+    svc_model = Svc(args.model_path,
+                    args.config_path,
+                    args.device,
+                    args.cluster_model_path,
+                    enhance,
+                    diffusion_model_path,
+                    diffusion_config_path,
+                    shallow_diffusion,
+                    only_diffusion,
+                    use_spk_mix,
+                    args.feature_retrieval)
+    
+    infer_tool.mkdir(["raw", "results"])
+    
+    if len(spk_mix_map)<=1:
+        use_spk_mix = False
+    if use_spk_mix:
+        spk_list = [spk_mix_map]
+    
     infer_tool.fill_a_to_b(trans, clean_names)
     for clean_name, tran in zip(clean_names, trans):
         raw_audio_path = f"raw/{clean_name}"
         if "." not in raw_audio_path:
             raw_audio_path += ".wav"
         infer_tool.format_wav(raw_audio_path)
-        wav_path = Path(raw_audio_path).with_suffix('.wav')
-        chunks = slicer.cut(wav_path, db_thresh=slice_db)
-        audio_data, audio_sr = slicer.chunks2audio(wav_path, chunks)
-        per_size = int(clip*audio_sr)
-        lg_size = int(lg*audio_sr)
-        lg_size_r = int(lg_size*lgr)
-        lg_size_c_l = (lg_size-lg_size_r)//2
-        lg_size_c_r = lg_size-lg_size_r-lg_size_c_l
-        lg = np.linspace(0,1,lg_size_r) if lg_size!=0 else 0
-
         for spk in spk_list:
-            audio = []
-            for (slice_tag, data) in audio_data:
-                print(f'#=====segment start, {round(len(data) / audio_sr, 3)}s======')
-                
-                length = int(np.ceil(len(data) / audio_sr * svc_model.target_sample))
-                if slice_tag:
-                    print('jump empty segment')
-                    _audio = np.zeros(length)
-                    audio.extend(list(infer_tool.pad_array(_audio, length)))
-                    continue
-                if per_size != 0:
-                    datas = infer_tool.split_list_by_n(data, per_size,lg_size)
-                else:
-                    datas = [data]
-                for k,dat in enumerate(datas):
-                    per_length = int(np.ceil(len(dat) / audio_sr * svc_model.target_sample)) if clip!=0 else length
-                    if clip!=0: print(f'###=====segment clip start, {round(len(dat) / audio_sr, 3)}s======')
-                    # padd
-                    pad_len = int(audio_sr * pad_seconds)
-                    dat = np.concatenate([np.zeros([pad_len]), dat, np.zeros([pad_len])])
-                    raw_path = io.BytesIO()
-                    soundfile.write(raw_path, dat, audio_sr, format="wav")
-                    raw_path.seek(0)
-                    out_audio, out_sr = svc_model.infer(spk, tran, raw_path,
-                                                        cluster_infer_ratio=cluster_infer_ratio,
-                                                        auto_predict_f0=auto_predict_f0,
-                                                        noice_scale=noice_scale,
-                                                        F0_mean_pooling = F0_mean_pooling,
-                                                        enhancer_adaptive_key = enhancer_adaptive_key
-                                                        )
-                    _audio = out_audio.cpu().numpy()
-                    pad_len = int(svc_model.target_sample * pad_seconds)
-                    _audio = _audio[pad_len:-pad_len]
-                    _audio = infer_tool.pad_array(_audio, per_length)
-                    if lg_size!=0 and k!=0:
-                        lg1 = audio[-(lg_size_r+lg_size_c_r):-lg_size_c_r] if lgr != 1 else audio[-lg_size:]
-                        lg2 = _audio[lg_size_c_l:lg_size_c_l+lg_size_r]  if lgr != 1 else _audio[0:lg_size]
-                        lg_pre = lg1*(1-lg)+lg2*lg
-                        audio = audio[0:-(lg_size_r+lg_size_c_r)] if lgr != 1 else audio[0:-lg_size]
-                        audio.extend(lg_pre)
-                        _audio = _audio[lg_size_c_l+lg_size_r:] if lgr != 1 else _audio[lg_size:]
-                    audio.extend(list(_audio))
+            kwarg = {
+                "raw_audio_path" : raw_audio_path,
+                "spk" : spk,
+                "tran" : tran,
+                "slice_db" : slice_db,
+                "cluster_infer_ratio" : cluster_infer_ratio,
+                "auto_predict_f0" : auto_predict_f0,
+                "noice_scale" : noice_scale,
+                "pad_seconds" : pad_seconds,
+                "clip_seconds" : clip,
+                "lg_num": lg,
+                "lgr_num" : lgr,
+                "f0_predictor" : f0p,
+                "enhancer_adaptive_key" : enhancer_adaptive_key,
+                "cr_threshold" : cr_threshold,
+                "k_step":k_step,
+                "use_spk_mix":use_spk_mix,
+                "second_encoding":second_encoding,
+                "loudness_envelope_adjustment":loudness_envelope_adjustment
+            }
+            audio = svc_model.slice_inference(**kwarg)
             key = "auto" if auto_predict_f0 else f"{tran}key"
             cluster_name = "" if cluster_infer_ratio == 0 else f"_{cluster_infer_ratio}"
-            res_path = f'./results/{clean_name}_{key}_{spk}{cluster_name}.{wav_format}'
+            isdiffusion = "sovits"
+            if shallow_diffusion :
+                isdiffusion = "sovdiff"
+            if only_diffusion :
+                isdiffusion = "diff"
+            if use_spk_mix:
+                spk = "spk_mix"
+            res_path = f'results/{clean_name}_{key}_{spk}{cluster_name}_{isdiffusion}_{f0p}.{wav_format}'
             soundfile.write(res_path, audio, svc_model.target_sample, format=wav_format)
             svc_model.clear_empty()
             

models.py

@@ -1,21 +1,17 @@
-import copy
-import math
 import torch
 from torch import nn
+from torch.nn import Conv1d, Conv2d
 from torch.nn import functional as F
+from torch.nn.utils import spectral_norm, weight_norm
 
 import modules.attentions as attentions
 import modules.commons as commons
 import modules.modules as modules
-
-from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
-from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
-
 import utils
-from modules.commons import init_weights, get_padding
-from vdecoder.hifigan.models import Generator
+from modules.commons import get_padding
 from utils import f0_to_coarse
 
+
 class ResidualCouplingBlock(nn.Module):
     def __init__(self,
                  channels,
@@ -126,7 +122,7 @@ class DiscriminatorP(torch.nn.Module):
         super(DiscriminatorP, self).__init__()
         self.period = period
         self.use_spectral_norm = use_spectral_norm
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
         self.convs = nn.ModuleList([
             norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
             norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(kernel_size, 1), 0))),
@@ -161,7 +157,7 @@ class DiscriminatorP(torch.nn.Module):
 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
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
         self.convs = nn.ModuleList([
             norm_f(Conv1d(1, 16, 15, 1, padding=7)),
             norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
@@ -321,6 +317,10 @@ class SynthesizerTrn(nn.Module):
                  n_speakers,
                  sampling_rate=44100,
                  vol_embedding=False,
+                 vocoder_name = "nsf-hifigan",
+                 use_depthwise_conv = False,
+                 use_automatic_f0_prediction = True,
+                 n_flow_layer = 4,
                  **kwargs):
 
         super().__init__()
@@ -343,6 +343,8 @@ class SynthesizerTrn(nn.Module):
         self.ssl_dim = ssl_dim
         self.vol_embedding = vol_embedding
         self.emb_g = nn.Embedding(n_speakers, gin_channels)
+        self.use_depthwise_conv = use_depthwise_conv
+        self.use_automatic_f0_prediction = use_automatic_f0_prediction
         if vol_embedding:
            self.emb_vol = nn.Linear(1, hidden_channels)
 
@@ -367,20 +369,35 @@ class SynthesizerTrn(nn.Module):
             "upsample_initial_channel": upsample_initial_channel,
             "upsample_kernel_sizes": upsample_kernel_sizes,
             "gin_channels": gin_channels,
+            "use_depthwise_conv":use_depthwise_conv
         }
-        self.dec = Generator(h=hps)
+        
+        modules.set_Conv1dModel(self.use_depthwise_conv)
+
+        if vocoder_name == "nsf-hifigan":
+            from vdecoder.hifigan.models import Generator
+            self.dec = Generator(h=hps)
+        elif vocoder_name == "nsf-snake-hifigan":
+            from vdecoder.hifiganwithsnake.models import Generator
+            self.dec = Generator(h=hps)
+        else:
+            print("[?] Unkown vocoder: use default(nsf-hifigan)")
+            from vdecoder.hifigan.models import Generator
+            self.dec = Generator(h=hps)
+
         self.enc_q = Encoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)
-        self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)
-        self.f0_decoder = F0Decoder(
-            1,
-            hidden_channels,
-            filter_channels,
-            n_heads,
-            n_layers,
-            kernel_size,
-            p_dropout,
-            spk_channels=gin_channels
-        )
+        self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, n_flow_layer, gin_channels=gin_channels)
+        if self.use_automatic_f0_prediction:
+            self.f0_decoder = F0Decoder(
+                1,
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+                spk_channels=gin_channels
+            )
         self.emb_uv = nn.Embedding(2, hidden_channels)
         self.character_mix = False
 
@@ -395,17 +412,21 @@ class SynthesizerTrn(nn.Module):
         g = self.emb_g(g).transpose(1,2)
 
         # vol proj
-        vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol!=None and self.vol_embedding else 0
+        vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol is not None and self.vol_embedding else 0
 
         # ssl prenet
         x_mask = torch.unsqueeze(commons.sequence_mask(c_lengths, c.size(2)), 1).to(c.dtype)
         x = self.pre(c) * x_mask + self.emb_uv(uv.long()).transpose(1,2) + vol
-
+        
         # f0 predict
-        lf0 = 2595. * torch.log10(1. + f0.unsqueeze(1) / 700.) / 500
-        norm_lf0 = utils.normalize_f0(lf0, x_mask, uv)
-        pred_lf0 = self.f0_decoder(x, norm_lf0, x_mask, spk_emb=g)
-
+        if self.use_automatic_f0_prediction:
+            lf0 = 2595. * torch.log10(1. + f0.unsqueeze(1) / 700.) / 500
+            norm_lf0 = utils.normalize_f0(lf0, x_mask, uv)
+            pred_lf0 = self.f0_decoder(x, norm_lf0, x_mask, spk_emb=g)
+        else:
+            lf0 = 0
+            norm_lf0 = 0
+            pred_lf0 = 0
         # encoder
         z_ptemp, m_p, logs_p, _ = self.enc_p(x, x_mask, f0=f0_to_coarse(f0))
         z, m_q, logs_q, spec_mask = self.enc_q(spec, spec_lengths, g=g)
@@ -419,6 +440,7 @@ class SynthesizerTrn(nn.Module):
 
         return o, ids_slice, spec_mask, (z, z_p, m_p, logs_p, m_q, logs_q), pred_lf0, norm_lf0, lf0
 
+    @torch.no_grad()
     def infer(self, c, f0, uv, g=None, noice_scale=0.35, seed=52468, predict_f0=False, vol = None):
 
         if c.device == torch.device("cuda"):
@@ -440,11 +462,13 @@ class SynthesizerTrn(nn.Module):
         
         x_mask = torch.unsqueeze(commons.sequence_mask(c_lengths, c.size(2)), 1).to(c.dtype)
         # vol proj
-        vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol!=None and self.vol_embedding else 0
-           
-        x = self.pre(c) * x_mask + self.emb_uv(uv.long()).transpose(1,2) + vol
         
-        if predict_f0:
+        vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol is not None and self.vol_embedding else 0
+
+        x = self.pre(c) * x_mask + self.emb_uv(uv.long()).transpose(1, 2) + vol
+
+        
+        if self.use_automatic_f0_prediction and predict_f0:
             lf0 = 2595. * torch.log10(1. + f0.unsqueeze(1) / 700.) / 500
             norm_lf0 = utils.normalize_f0(lf0, x_mask, uv, random_scale=False)
             pred_lf0 = self.f0_decoder(x, norm_lf0, x_mask, spk_emb=g)

train.py

@@ -1,39 +1,30 @@
 import logging
 import multiprocessing
+import os
 import time
 
-logging.getLogger('matplotlib').setLevel(logging.WARNING)
-logging.getLogger('numba').setLevel(logging.WARNING)
-
-import os
-import json
-import argparse
-import itertools
-import math
 import torch
-from torch import nn, optim
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from torch.cuda.amp import GradScaler, autocast
 from torch.nn import functional as F
+from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
-import torch.multiprocessing as mp
-import torch.distributed as dist
-from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.cuda.amp import autocast, GradScaler
 
 import modules.commons as commons
 import utils
-from data_utils import TextAudioSpeakerLoader, TextAudioCollate
+from data_utils import TextAudioCollate, TextAudioSpeakerLoader
 from models import (
-    SynthesizerTrn,
     MultiPeriodDiscriminator,
+    SynthesizerTrn,
 )
-from modules.losses import (
-    kl_loss,
-    generator_loss, discriminator_loss, feature_loss
-)
-
+from modules.losses import discriminator_loss, feature_loss, generator_loss, kl_loss
 from modules.mel_processing import mel_spectrogram_torch, spec_to_mel_torch
 
+logging.getLogger('matplotlib').setLevel(logging.WARNING)
+logging.getLogger('numba').setLevel(logging.WARNING)
+
 torch.backends.cudnn.benchmark = True
 global_step = 0
 start_time = time.time()
@@ -61,7 +52,7 @@ def run(rank, n_gpus, hps):
         utils.check_git_hash(hps.model_dir)
         writer = SummaryWriter(log_dir=hps.model_dir)
         writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
-
+    
     # for pytorch on win, backend use gloo    
     dist.init_process_group(backend=  'gloo' if os.name == 'nt' else 'nccl', init_method='env://', world_size=n_gpus, rank=rank)
     torch.manual_seed(hps.train.seed)
@@ -108,7 +99,7 @@ def run(rank, n_gpus, hps):
         name=utils.latest_checkpoint_path(hps.model_dir, "D_*.pth")
         global_step=int(name[name.rfind("_")+1:name.rfind(".")])+1
         #global_step = (epoch_str - 1) * len(train_loader)
-    except:
+    except Exception:
         print("load old checkpoint failed...")
         epoch_str = 1
         global_step = 0
@@ -148,6 +139,8 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
     train_loader, eval_loader = loaders
     if writers is not None:
         writer, writer_eval = writers
+    
+    half_type = torch.bfloat16 if hps.train.half_type=="bf16" else torch.float16
 
     # train_loader.batch_sampler.set_epoch(epoch)
     global global_step
@@ -169,8 +162,8 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
             hps.data.sampling_rate,
             hps.data.mel_fmin,
             hps.data.mel_fmax)
-
-        with autocast(enabled=hps.train.fp16_run):
+        
+        with autocast(enabled=hps.train.fp16_run, dtype=half_type):
             y_hat, ids_slice, z_mask, \
             (z, z_p, m_p, logs_p, m_q, logs_q), pred_lf0, norm_lf0, lf0 = net_g(c, f0, uv, spec, g=g, c_lengths=lengths,
                                                                                 spec_lengths=lengths,vol = volume)
@@ -191,25 +184,26 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
             # Discriminator
             y_d_hat_r, y_d_hat_g, _, _ = net_d(y, y_hat.detach())
 
-            with autocast(enabled=False):
+            with autocast(enabled=False, dtype=half_type):
                 loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(y_d_hat_r, y_d_hat_g)
                 loss_disc_all = loss_disc
-
+        
         optim_d.zero_grad()
         scaler.scale(loss_disc_all).backward()
         scaler.unscale_(optim_d)
         grad_norm_d = commons.clip_grad_value_(net_d.parameters(), None)
         scaler.step(optim_d)
+        
 
-        with autocast(enabled=hps.train.fp16_run):
+        with autocast(enabled=hps.train.fp16_run, dtype=half_type):
             # Generator
             y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(y, y_hat)
-            with autocast(enabled=False):
+            with autocast(enabled=False, dtype=half_type):
                 loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel
                 loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * hps.train.c_kl
                 loss_fm = feature_loss(fmap_r, fmap_g)
                 loss_gen, losses_gen = generator_loss(y_d_hat_g)
-                loss_lf0 = F.mse_loss(pred_lf0, lf0)
+                loss_lf0 = F.mse_loss(pred_lf0, lf0) if net_g.module.use_automatic_f0_prediction else 0
                 loss_gen_all = loss_gen + loss_fm + loss_mel + loss_kl + loss_lf0
         optim_g.zero_grad()
         scaler.scale(loss_gen_all).backward()
@@ -241,13 +235,17 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade
                 image_dict = {
                     "slice/mel_org": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()),
                     "slice/mel_gen": utils.plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().numpy()),
-                    "all/mel": utils.plot_spectrogram_to_numpy(mel[0].data.cpu().numpy()),
-                    "all/lf0": utils.plot_data_to_numpy(lf0[0, 0, :].cpu().numpy(),
-                                                          pred_lf0[0, 0, :].detach().cpu().numpy()),
-                    "all/norm_lf0": utils.plot_data_to_numpy(lf0[0, 0, :].cpu().numpy(),
-                                                               norm_lf0[0, 0, :].detach().cpu().numpy())
+                    "all/mel": utils.plot_spectrogram_to_numpy(mel[0].data.cpu().numpy())
                 }
 
+                if net_g.module.use_automatic_f0_prediction:
+                    image_dict.update({
+                        "all/lf0": utils.plot_data_to_numpy(lf0[0, 0, :].cpu().numpy(),
+                                                              pred_lf0[0, 0, :].detach().cpu().numpy()),
+                        "all/norm_lf0": utils.plot_data_to_numpy(lf0[0, 0, :].cpu().numpy(),
+                                                                   norm_lf0[0, 0, :].detach().cpu().numpy())
+                    })
+
                 utils.summarize(
                     writer=writer,
                     global_step=global_step,
@@ -287,7 +285,7 @@ def evaluate(hps, generator, eval_loader, writer_eval):
             c = c[:1].cuda(0)
             f0 = f0[:1].cuda(0)
             uv= uv[:1].cuda(0)
-            if volume!=None:
+            if volume is not None:
                 volume = volume[:1].cuda(0)
             mel = spec_to_mel_torch(
                 spec,
@@ -314,7 +312,7 @@ def evaluate(hps, generator, eval_loader, writer_eval):
                 f"gt/audio_{batch_idx}": y[0]
             })
         image_dict.update({
-            f"gen/mel": utils.plot_spectrogram_to_numpy(y_hat_mel[0].cpu().numpy()),
+            "gen/mel": utils.plot_spectrogram_to_numpy(y_hat_mel[0].cpu().numpy()),
             "gt/mel": utils.plot_spectrogram_to_numpy(mel[0].cpu().numpy())
         })
     utils.summarize(
@@ -328,4 +326,4 @@ def evaluate(hps, generator, eval_loader, writer_eval):
 
 
 if __name__ == "__main__":
-    main()
+    main()