nodes/SayanoAI-Comfy-RVC
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactored training params to node

Browse Source
This commit is contained in:
SayanoAI
2024-08-30 23:04:28 -04:00
parent b52bae0be0
commit 5ecd6a84f2
4 changed files with 74 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
custom_nodes/rvc_nodes.py
+40 -22
View File
@@ -8,7 +8,7 @@ import torch
from . import BASE_DIR
from ..lib.train.utils import HParams
from ..lib.train.utils import HParams, DEFAULT_TRAINING_PARAMS
from ..training_cli import train_model
@@ -153,13 +153,9 @@ class LoadRVCModelNode:
class RVCNode:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"audio": (MultipleTypeProxy('AUDIO,VHS_AUDIO'),),
@@ -213,9 +209,6 @@ class RVCNode:
class RVCProcessDatasetNode:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
@@ -358,10 +351,36 @@ class RVCProcessDatasetNode:
),)
class RVCTrainModelNode:
class RVCTrainParamsNode:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"optional": dict(
batch_size=("INT",dict(default=4,min=1,max=64,step=1)),
c_mel=("FLOAT",dict(default=45.,min=0.,max=100.,step=.1)),
c_kl=("FLOAT",dict(default=1.,min=0.,max=100.,step=.1)),
c_fm=("FLOAT",dict(default=2.,min=0.,max=100.,step=.1)),
c_mfcc=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
c_lfcc=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
c_hd=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
c_sts=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
c_gp=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
)
}
RETURN_TYPES = ('RVC_TRAINING_PARAMS', )
RETURN_NAMES = ("rvc_training_params", )
FUNCTION = "init"
CATEGORY = CATEGORY
def init(self,batch_size=4,c_mel=45,c_kl=1.,c_fm=1.,c_mfcc=0.,c_lfcc=0.,c_hd=0.,c_sts=0.,c_gp=0.,):
return (dict(batch_size=batch_size,c_mel=c_mel,c_kl=c_kl,c_fm=c_fm,c_mfcc=c_mfcc,c_gp=c_gp,c_lfcc=c_lfcc,c_hd=c_hd,c_sts=c_sts),)
class RVCTrainModelNode:
@classmethod
def INPUT_TYPES(cls):
@@ -374,10 +393,10 @@ class RVCTrainModelNode:
return {
"required": {
"rvc_dataset_pipe": ("RVC_DATASET_PIPE",),
"rvc_training_params": ("RVC_TRAINING_PARAMS",dict(default=DEFAULT_TRAINING_PARAMS)),
},
"optional": dict(
gpu=(DEVICES, {"default": DEVICES[0]}),
batch_size=("INT",dict(default=4,min=1,max=64,step=1)),
total_epoch=("INT",dict(default=100,min=10,max=1000,step=10)),
save_every_epoch=("INT",dict(default=0,min=0,max=100)),
pretrained_G=(PRETRAINED_G,{"default": PRETRAINED_G[0]}),
@@ -388,9 +407,8 @@ class RVCTrainModelNode:
train_index=("BOOLEAN",{"default": True}),
retrain=("BOOLEAN",{"default": False}),
save_best_model=("BOOLEAN",{"default": True}),
best_model_threshold=("INT",dict(default=50,min=20,max=80)),
log_every_epoch=("FLOAT",dict(default=1.,min=0.,max=2.,step=.5)),
gradient_lambda=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
timbre_lambda=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
num_workers=("INT",dict(default=1,min=1,max=16))
)
}
@@ -406,8 +424,8 @@ class RVCTrainModelNode:
def train_model(self,
rvc_dataset_pipe,
rvc_training_params,
gpu="0",
batch_size=4,
total_epoch=100,
save_every_epoch=0,
pretrained_G="",
@@ -418,15 +436,14 @@ class RVCTrainModelNode:
train_index=True,
retrain=False,
save_best_model=True,
best_model_threshold=50,
log_every_epoch=1.,
gradient_lambda=0.,
timbre_lambda=0.,
num_workers=1):
sample_rate = rvc_dataset_pipe["sample_rate"]
name = rvc_dataset_pipe["name"]
dataset_dir = rvc_dataset_pipe["dataset_dir"]
cache_name = get_hash(batch_size,pretrained_G,pretrained_D,gradient_lambda,timbre_lambda)
cache_name = get_hash(pretrained_G,pretrained_D,*rvc_training_params.items())
model_dir = os.path.join(output_path,"logs",cache_name)
if_f0 = rvc_dataset_pipe["if_f0"]
@@ -444,7 +461,6 @@ class RVCTrainModelNode:
hparams.pretrainD = model_downloader(pretrained_D)
hparams.version = "v2"
hparams.gpus = gpu
hparams.train.batch_size = max(batch_size,1)
hparams.sample_rate = sample_rate
hparams.if_f0 = if_f0
hparams.if_latest = if_save_latest
@@ -452,10 +468,10 @@ class RVCTrainModelNode:
hparams.if_cache_data_in_gpu = if_cache_gpu
hparams.data.training_files = rvc_dataset_pipe["training_files"]
hparams.save_best_model = save_best_model
hparams.best_model_threshold = best_model_threshold
hparams.log_every_epoch = log_every_epoch
hparams.train.gradient_lambda = gradient_lambda
hparams.train.num_workers = num_workers
hparams.train.timbre_lambda = timbre_lambda
hparams.train.update(rvc_training_params)
file_index = self.train_index(dataset_dir, sample_rate, name) if train_index else None
model_path = os.path.join(BASE_MODELS_DIR,"RVC",f"{name}_{sample_rate}.pth")
@@ -531,7 +547,8 @@ NODE_CLASS_MAPPINGS = {
"LoadHubertModel": LoadHubertModel,
"LoadPitchExtractionParams": LoadPitchExtractionParams,
"RVCProcessDatasetNode": RVCProcessDatasetNode,
"RVCTrainModelNode": RVCTrainModelNode
"RVCTrainModelNode": RVCTrainModelNode,
"RVCTrainParamsNode": RVCTrainParamsNode
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -542,4 +559,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"LoadPitchExtractionParams": "🌺Load Pitch Extraction Params",
"RVCProcessDatasetNode": "🌺Process RVC Dataset",
"RVCTrainModelNode": "🌺Train RVC Model",
"RVCTrainParamsNode": "🌺RVC Training Params"
}
lib/train/losses.py
+1 -1
View File
@@ -35,7 +35,7 @@ def feature_loss(fmap_r, fmap_g):
gl = gl.float()
loss += torch.mean(torch.abs(rl - gl))
return loss * 2
return loss
def discriminator_loss(disc_real_outputs, disc_generated_outputs):
lib/train/utils.py
+23 -37
View File
@@ -1,4 +1,4 @@
import os, traceback
import os
import glob
import sys
import argparse
@@ -10,6 +10,17 @@ from scipy.io.wavfile import read
import torch
MATPLOTLIB_FLAG = False
DEFAULT_TRAINING_PARAMS = dict(
batch_size=4,
c_mel=45,
c_kl=1.,
c_fm=2.,
c_gp=0.,
c_mfcc=0.,
c_lfcc=0.,
c_hd=0.,
c_sts=0.
)
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
logger = logging
@@ -63,36 +74,6 @@ def load_checkpoint_d(checkpoint_path, combd, sbd, optimizer=None, load_opt=1):
logger.info("Loaded checkpoint '{}' (epoch {})".format(checkpoint_path, iteration))
return model, optimizer, learning_rate, iteration
# def load_checkpoint(checkpoint_path, model, optimizer=None):
# assert os.path.isfile(checkpoint_path)
# checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
# iteration = checkpoint_dict['iteration']
# learning_rate = checkpoint_dict['learning_rate']
# if optimizer is not None:
# optimizer.load_state_dict(checkpoint_dict['optimizer'])
# # print(1111)
# saved_state_dict = checkpoint_dict['model']
# # print(1111)
#
# if hasattr(model, 'module'):
# state_dict = model.module.state_dict()
# else:
# state_dict = model.state_dict()
# new_state_dict= {}
# for k, v in state_dict.items():
# try:
# new_state_dict[k] = saved_state_dict[k]
# except:
# logger.info("%s is not in the checkpoint" % k)
# new_state_dict[k] = v
# if hasattr(model, 'module'):
# model.module.load_state_dict(new_state_dict)
# else:
# model.load_state_dict(new_state_dict)
# logger.info("Loaded checkpoint '{}' (epoch {})" .format(
# checkpoint_path, iteration))
# return model, optimizer, learning_rate, iteration
def load_checkpoint(checkpoint_path, model, optimizer=None, load_opt=1):
assert os.path.isfile(checkpoint_path)
checkpoint_dict = torch.load(checkpoint_path, map_location="cpu")
@@ -382,10 +363,9 @@ def get_hparams(init=True):
hparams.if_cache_data_in_gpu = args.if_cache_data_in_gpu
hparams.data.training_files = os.path.join(experiment_dir,"filelist.txt")
hparams.save_best_model = True
hparams.train.gradient_lambda = 0.
hparams.train.num_workers = 4
hparams.train.timbre_lambda = 0.
hparams.model_path = None
hparams.best_model_threshold = 50
return hparams
@@ -449,11 +429,10 @@ def get_logger(model_dir, filename="train.log"):
return logger
class HParams:
class HParams():
def __init__(self, **kwargs):
for k, v in kwargs.items():
if isinstance(v, dict):
v = HParams(**v)
if isinstance(v, dict): v = HParams(**v)
self[k] = v
def keys(self):
@@ -483,7 +462,14 @@ class HParams:
def get(self, key, default=None):
try: return self[key]
except: return default
def update(self, data={}, **kwargs):
data.update(kwargs)
for k, v in data.items():
if isinstance(v, dict): v = HParams(**v)
if k in self and hasattr(self[k], "update") and isinstance(v, HParams): self[k].update(v)
else: self[k] = v
def sync_log_interval(self, dataset_length):
log_interval = getattr(self,"log_every_epoch",1)
log_interval = getattr(self,"log_every_epoch",1.)
self.train.log_interval = int(dataset_length*log_interval)
training_cli.py
+10 -10
View File
@@ -127,10 +127,10 @@ def run(rank, n_gpus, hps, device):
if os.path.isfile(loss_file):
with open(loss_file,"r") as f:
data: dict = json.load(f)
least_loss = data.get("least_loss",50)
least_loss = data.get("least_loss",hps.best_model_threshold)
best_model_name = data.get("best_model_name","")
else:
least_loss = 50
least_loss = hps.best_model_threshold
best_model_name = ""
if hps.version == "v1":
@@ -481,7 +481,7 @@ def train_and_evaluate(
with autocast(enabled=False):
gradient_penalty = torch.Tensor([0.]).to(wave.device)
if hps.train.get("gradient_lambda",0.)>0:
if hps.train.get("c_gp",0.)>0:
# Compute the gradient penalty
# Randomly interpolate between real and generated data
size = [1]*wave.ndim
@@ -502,10 +502,10 @@ def train_and_evaluate(
)[0]
gradient = gradient.view(gradient.size(0), -1)
gradient_penalty += torch.log(gradient.norm(2,dim=-1)).abs().mean() # force gradnorm=1
gradient_penalty *= hps.train.gradient_lambda/len(disc_interpolated_output)
gradient_penalty *= hps.train.c_gp/len(disc_interpolated_output)
loss_mfcc = torch.Tensor([0.]).to(wave.device)
if hps.train.get("timbre_lambda",0)>0:
if hps.train.get("c_mfcc",0)>0:
loss_mfcc = mfcc_loss(wave, gen_wave, SR_MAP[hps.sample_rate],
n_fft=hps.data.filter_length,
hop_length=hps.data.hop_length,
@@ -514,7 +514,7 @@ def train_and_evaluate(
norm="slaney",
f_min=hps.data.mel_fmin,
f_max=hps.data.mel_fmax
) * hps.train.timbre_lambda
) * hps.train.c_mfcc
loss_mfcc = loss_mfcc.squeeze()
gradient_penalty = gradient_penalty.squeeze()
@@ -531,9 +531,9 @@ def train_and_evaluate(
# Generator
y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(wave, y_hat)
with autocast(enabled=False):
loss_mel = F.smooth_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_mel = F.smooth_l1_loss(y_mel, y_hat_mel) * hps.train.get("c_mel",45.)
loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * hps.train.get("c_kl",1.)
loss_fm = feature_loss(fmap_r, fmap_g) * hps.train.get("c_fm",2.)
loss_gen, losses_gen = generator_loss(y_d_hat_g)
loss_gen_all = loss_gen + loss_fm + loss_mel + loss_kl
@@ -556,7 +556,7 @@ def train_and_evaluate(
else:
ckpt = net_g.state_dict()
best_model_name = f"{hps.name}_loss{least_loss:2.0f}" if hps.if_save_latest else f"{hps.name}_e{epoch}_s{global_step}_loss{least_loss:2.0f}"
best_model_name = f"{hps.name}_e{epoch}_s{global_step}_loss{least_loss:2.0f}" if hps.save_every_weights else f"{hps.name}_loss{least_loss:2.0f}"
status = save_checkpoint(ckpt,best_model_name,epoch,hps)
logger.info(f"=== saving best model {best_model_name}: {status=} ===")