diff --git a/custom_nodes/audio_nodes.py b/custom_nodes/audio_nodes.py index d99cc30..61818f4 100644 --- a/custom_nodes/audio_nodes.py +++ b/custom_nodes/audio_nodes.py @@ -5,9 +5,9 @@ import numpy as np import yt_dlp import torch from .settings import MERGE_OPTIONS -from .utils import MultipleTypeProxy, increment_filename_no_overwrite +from .helper import MultipleTypeProxy, increment_filename_no_overwrite from ..lib.audio import MAX_INT16, SUPPORTED_AUDIO, AudioProcessor, audio_to_bytes, get_audio, load_input_audio, pad_audio, remix_audio, save_input_audio -from ..lib.utils import get_filenames, get_hash, get_merge_func +from ..lib.helper import get_filenames, get_hash, get_merge_func import folder_paths CATEGORY = "🌺RVC-Studio/audio" diff --git a/custom_nodes/utils.py b/custom_nodes/helper.py similarity index 100% rename from custom_nodes/utils.py rename to custom_nodes/helper.py diff --git a/custom_nodes/musetalk_nodes.py b/custom_nodes/musetalk_nodes.py index aaa5f10..d6bcb7d 100644 --- a/custom_nodes/musetalk_nodes.py +++ b/custom_nodes/musetalk_nodes.py @@ -8,8 +8,8 @@ from tqdm import tqdm from ..lib.audio import get_audio, remix_audio from ..lib import BASE_CACHE_DIR -from ..lib.utils import gc_collect, get_hash, get_optimal_torch_device -from .utils import MultipleTypeProxy, model_downloader +from ..lib.helper import gc_collect, get_hash, get_optimal_torch_device +from .helper import MultipleTypeProxy, model_downloader from ..lib.musetalk.models.unet import PositionalEncoding, UNet from ..lib.musetalk.models.vae import VAE from ..lib.musetalk.whisper.audio2feature import Audio2Feature diff --git a/custom_nodes/rvc_nodes.py b/custom_nodes/rvc_nodes.py index 920acb0..1869ebf 100644 --- a/custom_nodes/rvc_nodes.py +++ b/custom_nodes/rvc_nodes.py @@ -19,7 +19,7 @@ from ..config import config from ..lib.model_utils import load_hubert -from .utils import MultipleTypeProxy, increment_filename_no_overwrite, model_downloader +from .helper import MultipleTypeProxy, increment_filename_no_overwrite, model_downloader from .settings import PITCH_EXTRACTION_OPTIONS from .settings.downloader import PRETRAINED_MODELS_D, PRETRAINED_MODELS_G, RVC_DOWNLOAD_LINK, RVC_INDEX, RVC_MODELS, download_file, extract_zip_without_structure @@ -27,7 +27,7 @@ from ..lib.audio import SUPPORTED_AUDIO, audio_to_bytes, load_input_audio, save_ from ..vc_infer_pipeline import get_vc, vc_single import folder_paths -from ..lib.utils import get_filenames, get_hash, get_optimal_threads, get_optimal_torch_device +from ..lib.helper import get_filenames, get_hash, get_optimal_threads, get_optimal_torch_device from ..lib import BASE_CACHE_DIR, BASE_MODELS_DIR input_path = folder_paths.get_input_directory() diff --git a/custom_nodes/stt.py b/custom_nodes/stt.py index f58a4fa..a922cf8 100644 --- a/custom_nodes/stt.py +++ b/custom_nodes/stt.py @@ -10,8 +10,8 @@ import torch from .settings import SUPPORTED_LANGUAGES from ..lib import BASE_CACHE_DIR, BASE_MODELS_DIR -from .utils import increment_filename_no_overwrite -from ..lib.utils import get_hash, get_optimal_torch_device +from .helper import increment_filename_no_overwrite +from ..lib.helper import get_hash, get_optimal_torch_device import folder_paths from ..lib.audio import bytes_to_audio, remix_audio import spacy diff --git a/custom_nodes/utility_nodes.py b/custom_nodes/utility_nodes.py index 3399de0..74eb8b0 100644 --- a/custom_nodes/utility_nodes.py +++ b/custom_nodes/utility_nodes.py @@ -2,9 +2,9 @@ import itertools import math import os import torch -from .utils import AlwaysEqualProxy, MultipleTypeProxy +from .helper import AlwaysEqualProxy, MultipleTypeProxy import folder_paths -from ..lib.utils import gc_collect +from ..lib.helper import gc_collect import numpy as np CATEGORY = "🌺RVC-Studio/utils" diff --git a/custom_nodes/uvr.py b/custom_nodes/uvr.py index 26e423f..0d53969 100644 --- a/custom_nodes/uvr.py +++ b/custom_nodes/uvr.py @@ -1,10 +1,10 @@ import os import audio_separator.separator as uvr -from .utils import MultipleTypeProxy +from .helper import MultipleTypeProxy from ..lib.audio import audio_to_bytes, save_input_audio, load_input_audio, get_audio import folder_paths -from ..lib.utils import get_filenames, get_hash, get_optimal_torch_device +from ..lib.helper import get_filenames, get_hash, get_optimal_torch_device from ..lib import BASE_CACHE_DIR, BASE_MODELS_DIR, karafan from .settings.downloader import KARAFAN_MODELS, MDX_MODELS, RVC_DOWNLOAD_LINK, VR_MODELS, download_file diff --git a/lib/__init__.py b/lib/__init__.py index 42c32e2..a13a006 100644 --- a/lib/__init__.py +++ b/lib/__init__.py @@ -7,7 +7,10 @@ import weakref @lru_cache def get_cwd(): - CWD = os.getcwd() + import folder_paths + # CWD = os.getcwd() + CWD = folder_paths.base_path + print(f"{CWD=}") if CWD not in sys.path: sys.path.append(CWD) return CWD diff --git a/lib/audio.py b/lib/audio.py index eba5138..8797ffb 100644 --- a/lib/audio.py +++ b/lib/audio.py @@ -2,7 +2,7 @@ import base64 import io import os import zlib -from .utils import get_hash, get_merge_func +from .helper import get_hash, get_merge_func import numpy as np import librosa import soundfile as sf diff --git a/lib/utils.py b/lib/helper.py similarity index 100% rename from lib/utils.py rename to lib/helper.py diff --git a/lib/karafan/inference.py b/lib/karafan/inference.py index 311a105..63be1a6 100644 --- a/lib/karafan/inference.py +++ b/lib/karafan/inference.py @@ -17,7 +17,7 @@ import numpy as np, onnxruntime as ort import yaml from .. import BASE_CACHE_DIR, BASE_DIR, BASE_MODELS_DIR, ObjectNamespace as ConfigDict from ..audio import remix_audio -from ..utils import gc_collect +from ..helper import gc_collect # Maybe not available on PC try: diff --git a/lib/mdx.py b/lib/mdx.py index 0090c7c..e0884c6 100644 --- a/lib/mdx.py +++ b/lib/mdx.py @@ -7,7 +7,7 @@ import torch from tqdm import tqdm from .model_utils import get_hash -from .utils import gc_collect +from .helper import gc_collect import onnxruntime as ort warnings.filterwarnings("ignore") diff --git a/lib/separators.py b/lib/separators.py index 8f70cf6..82a3e2e 100644 --- a/lib/separators.py +++ b/lib/separators.py @@ -2,7 +2,7 @@ import hashlib import json from multiprocessing.pool import ThreadPool import os -from .utils import ObjectNamespace +from .helper import ObjectNamespace import numpy as np import torch from tqdm import tqdm diff --git a/lib/train/losses.py b/lib/train/losses.py index 1c3e738..810e0a3 100644 --- a/lib/train/losses.py +++ b/lib/train/losses.py @@ -8,7 +8,7 @@ import torch.nn.functional as F from .mel_processing import mel_spectrogram_torch, spectral_de_normalize_torch -from ..utils import gc_collect +from ..helper import gc_collect from ..infer_pack.commons import compute_correlation, median_pool1d, minmax_scale class LossBalancer: @@ -71,13 +71,13 @@ class LossBalancer: """Calculates the slope of the loss using the current loss and its EMA.""" ema_loss = self.historical_losses.get(key, current_loss)+self.epsilon slope = (current_loss-ema_loss)/ema_loss # relative loss change - return slope.abs() + return slope.abs().requires_grad_() def calculate_gradients(self, key: str, current_loss: torch.Tensor, input: torch.Tensor): """Calculates the gradient norm of the current loss wrt the model params.""" self.model.zero_grad() # Clear previous gradients - input.requires_grad_(True) - current_loss.requires_grad_(True) + input.requires_grad_() + current_loss.requires_grad_() # Backward pass to output layer output_params = torch.autograd.grad(current_loss, [input], retain_graph=True, allow_unused=True, materialize_grads=True)[0].nan_to_num(self.epsilon) @@ -89,7 +89,7 @@ class LossBalancer: if grad_norm<=self.epsilon: #use slope as fallback grad_norm = self.calculate_loss_slope(key, current_loss) - return grad_norm + return grad_norm.requires_grad_() def pareto_normalizer(self, loss_dict: dict, weight=.8): """ @@ -196,7 +196,7 @@ class LossBalancer: gradients[key] = max(loss_slope.item(), self.epsilon) valid_losses[key] = loss.nan_to_num(self.epsilon) - if not valid_losses or not gradients: return torch.tensor(0.0) # If all losses are skipped + if not valid_losses or not gradients: return torch.tensor(0.,device=input.device).requires_grad_() # If all losses are skipped # update historical losses self.update_historical_losses({k: loss.item() for k, loss in valid_losses.items()}) @@ -209,11 +209,12 @@ class LossBalancer: normalized_weights = self.update_ema_weights(normalized_weights) # Balance losses - balanced_loss = 0 + balanced_loss = torch.tensor(0.,device=input.device).requires_grad_() for k, loss in valid_losses.items(): - balanced_loss += normalized_weights.get(k, 1.) * loss + balanced_loss += normalized_weights.get(k, 1.) * loss.requires_grad_() + print(balanced_loss.requires_grad,loss.requires_grad) - return balanced_loss + return balanced_loss.requires_grad_() def on_epoch_end(self, weights_decay=None, loss_decay=None): """ @@ -380,23 +381,23 @@ def combined_aux_loss( # Define loss terms harmonic_loss = F.l1_loss(generated_harmonics, original_harmonics) harmonic_loss += F.l1_loss(generated_percussives, original_percussives) - else: harmonic_loss = 0 + else: harmonic_loss = torch.tensor(0.,device=original_audio.device).requires_grad_() # temporal invariant phase if c_tsi>0: freq_tsi = compute_tsi_loss(org_mag,gen_mag,dim=-1, eps=eps) temp_tsi = compute_tsi_loss(org_mag,gen_mag,dim=-2, eps=eps) tsi_loss = (freq_tsi+temp_tsi) - else: tsi_loss = 0 + else: tsi_loss = torch.tensor(0.,device=original_audio.device).requires_grad_() # Temperol Envelope and Fine Structure Loss if c_tefs>0: gen_te, gen_tfs = compute_tefs(generated_audio, eps=eps) org_te, org_tfs = compute_tefs(original_audio, eps=eps) tefs_loss = F.l1_loss(gen_te, org_te) + F.l1_loss(gen_tfs, org_tfs) - else: tefs_loss = 0 + else: tefs_loss = torch.tensor(0.,device=original_audio.device).requires_grad_() - return harmonic_loss, tefs_loss, tsi_loss + return harmonic_loss.requires_grad_(), tefs_loss.requires_grad_(), tsi_loss.requires_grad_() def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tensor, net_d: torch.nn.Module, eps=1e-8): # Compute the gradient penalty @@ -423,7 +424,7 @@ def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tens if gradients.ndim<=1: gradients = gradients.unsqueeze(0) grad_norm = gradients.view(gradients.size(0), -1).square().sum(-1).sqrt() loss = ((grad_norm - 1) ** 2).mean() - return loss + return loss.requires_grad_() # Adapted from https://github.com/NVIDIA/BigVGAN/blob/main/loss.py # LICENSE: https://github.com/NVIDIA/BigVGAN/blob/main/LICENSE diff --git a/pitch_extraction.py b/pitch_extraction.py index 05729c6..6c08ebc 100644 --- a/pitch_extraction.py +++ b/pitch_extraction.py @@ -8,7 +8,7 @@ import torch from .lib.rmvpe import RMVPE from .lib.audio import autotune_f0, pad_audio, hz_to_mel from .lib import BASE_MODELS_DIR -from .lib.utils import gc_collect, get_merge_func, get_optimal_threads, get_optimal_torch_device +from .lib.helper import gc_collect, get_merge_func, get_optimal_threads, get_optimal_torch_device class FeatureExtractor: def __init__(self, tgt_sr, config, onnx=False): diff --git a/preprocessing_utils.py b/preprocessing_utils.py index 48ba084..3013830 100644 --- a/preprocessing_utils.py +++ b/preprocessing_utils.py @@ -6,7 +6,7 @@ import traceback from scipy.io import wavfile from .pitch_extraction import FeatureExtractor from .lib.audio import hz_to_mel, load_input_audio, remix_audio, AudioProcessor -from .lib.utils import gc_collect +from .lib.helper import gc_collect from .config import config import torch diff --git a/requirements.txt b/requirements.txt index fef6df5..c531691 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ librosa~=0.10.2 -numpy +numpy<2 faiss-cpu==1.7.4 ffmpeg-python tqdm diff --git a/training_cli.py b/training_cli.py index 54fc342..de7bf6b 100644 --- a/training_cli.py +++ b/training_cli.py @@ -85,13 +85,8 @@ class EpochRecorder: current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") return f"[{current_time}] | ({elapsed_time_str})" -def train_model(hps: "utils.HParams"): - os.environ["CUDA_VISIBLE_DEVICES"] = hps.gpus.replace("-", ",") - os.environ["NCCL_P2P_DISABLE"] = "1" - os.environ["MASTER_ADDR"] = "localhost" - os.environ["MASTER_PORT"] = str(randint(8189, 8205+hps.train.num_workers**2)) - os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128,garbage_collection_threshold:0.8" - +def train_model(hps: "utils.HParams"): + print(hps) n_gpus = len(hps.gpus.split("-")) if hps.gpus else torch.cuda.device_count() if not torch.cuda.is_available() and torch.backends.mps.is_available(): @@ -102,23 +97,35 @@ def train_model(hps: "utils.HParams"): n_gpus = 1 gpu_devices = hps.gpus.split("-") if hps.gpus else range(n_gpus) + print(f"{gpu_devices=} {n_gpus=}") - children = {} - for i, device in enumerate(gpu_devices): - subproc = mp.Process( - target=run, - args=( - i, - n_gpus, - hps, - device - ), - ) - children[i]=subproc - subproc.start() + if n_gpus==1: + run(0,1,hps,"0") + else: + import sys + sys.path.insert(0,os.getcwd()) + mp.set_start_method("spawn") + os.environ["CUDA_VISIBLE_DEVICES"] = hps.gpus.replace("-", ",") + os.environ["NCCL_P2P_DISABLE"] = "1" + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(randint(8189, 8205+hps.train.num_workers**2)) + os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128,garbage_collection_threshold:0.8" + children = {} + for i, device in enumerate(gpu_devices): + subproc = mp.Process( + target=run, + args=( + i, + n_gpus, + hps, + device + ), + ) + children[i]=subproc + subproc.start() - for i in children: - children[i].join() + for i in children: + children[i].join() def run(rank, n_gpus, hps, device): print(f"{__name__=}") @@ -195,13 +202,13 @@ def run(rank, n_gpus, hps, device): collate_fn = TextAudioCollate() train_loader = DataLoader( train_dataset, - num_workers=hps.train.num_workers, + # num_workers=hps.train.num_workers, shuffle=False, - pin_memory=True, + # pin_memory=True, collate_fn=collate_fn, batch_sampler=train_sampler, - persistent_workers=True, - prefetch_factor=8, + # persistent_workers=True, + # prefetch_factor=8, ) hps.sync_log_interval(len(train_loader)) @@ -386,6 +393,7 @@ def train_and_evaluate( net_d.train() balancer_g, balancer_d = balancer gradient_clip_value = commons.sigmoid_value(global_step, total_steps=10000, start_value=1, end_value=500, midpoint=.2) + print(balancer) # Prepare data iterator if hps.if_cache_data_in_gpu: @@ -470,7 +478,10 @@ def train_and_evaluate( data_iterator = enumerate(train_loader) # Run steps + print(data_iterator) epoch_recorder = EpochRecorder() + print(epoch_recorder) + for batch_idx, info in tqdm(data_iterator,desc=f"[Epoch {epoch}]: "): # Data ## Unpack @@ -526,6 +537,7 @@ def train_and_evaluate( hps.data.mel_fmin, hps.data.mel_fmax, ) + y_hat.requires_grad_() y_mel = commons.slice_segments( mel, ids_slice, hps.train.segment_size // hps.data.hop_length ) @@ -545,22 +557,29 @@ def train_and_evaluate( wave = commons.slice_segments(wave, ids_slice * hps.data.hop_length, hps.train.segment_size) # slice # Discriminator - gen_wave = y_hat.detach() + gen_wave = y_hat.clone().detach().requires_grad_() y_d_hat_r, y_d_hat_g, _, _ = net_d(wave, gen_wave) with autocast(enabled=False): - gradient_penalty = gradient_norm_loss(wave,gen_wave, net_d, eps=hps.train.eps)*hps.train.c_gp if hps.train.get("c_gp",0.)>0 else 0 + if hps.train.get("c_gp",0.)>0: + gradient_penalty = gradient_norm_loss(wave,gen_wave, net_d, eps=hps.train.eps)*hps.train.c_gp + else: + gradient_penalty = torch.tensor(0., device=wave.device) loss_disc, losses_disc = discriminator_loss(y_d_hat_r, y_d_hat_g) loss_disc_all = balancer_d.on_train_batch_start(dict( loss_disc=loss_disc, gradient_penalty=gradient_penalty ),input=y_hat) - optim_d.zero_grad() - scaler.scale(loss_disc_all).backward() - scaler.unscale_(optim_d) + try: + optim_d.zero_grad() + scaler.scale(loss_disc_all.requires_grad_()).backward() + scaler.unscale_(optim_d) + scaler.step(optim_d) + # scaler.update() + except Exception as e: + print(f"error updating discriminator: {e}") grad_norm_d = commons.clip_grad_value_(net_d.parameters(), gradient_clip_value, batch_size=hps.train.batch_size) - scaler.step(optim_d) with autocast(enabled=hps.train.fp16_run): # Generator @@ -589,17 +608,18 @@ def train_and_evaluate( loss_fm=loss_fm, loss_mel=loss_mel, loss_kl=loss_kl, - harmonic_loss=harmonic_loss, - tsi_loss=tsi_loss, - tefs_loss=tefs_loss, + aux_loss=aux_loss, ),input=y_hat) - - optim_g.zero_grad() - scaler.scale(loss_gen_all).backward() - scaler.unscale_(optim_g) + print(f"{loss_gen.requires_grad=} {loss_gen_all.requires_grad=}") + try: + optim_g.zero_grad() + scaler.scale(loss_gen_all.requires_grad_()).backward() + scaler.unscale_(optim_g) + scaler.step(optim_g) + scaler.update() + except Exception as e: + print(f"error updating generator: {e}") grad_norm_g = commons.clip_grad_value_(net_g.parameters(), gradient_clip_value, batch_size=hps.train.batch_size) - scaler.step(optim_g) - scaler.update() if rank == 0: if hps.train.log_interval>0 and global_step % hps.train.log_interval == 0: #tensorboard logging @@ -693,9 +713,9 @@ def train_and_evaluate( if rank == 0: total_loss = balancer_g.weighted_ema_loss + balancer_d.weighted_ema_loss logger.info(f"====> Epoch {epoch} ({total_loss=:.3f}): {global_step=} {lr=:.2E} {epoch_recorder.record()}") - logger.info(f"|| {loss_disc_all=:.3f}: {loss_disc=:.3f}, {gradient_penalty=:.3f}") - logger.info(f"|| {loss_gen_all=:.3f}: {loss_gen=:.3f}, {loss_fm=:.3f}, {loss_mel=:.3f}, {loss_kl=:.3f}") - logger.info(f"|| {aux_loss=:.3f}: {harmonic_loss=:.3f}, {tefs_loss=:.3f}, {tsi_loss=:.3f}") + logger.info(f"|| {loss_disc_all.item()=:.3f}: {loss_disc.item()=:.3f}, {gradient_penalty.item()=:.3f}") + logger.info(f"|| {loss_gen_all.item()=:.3f}: {loss_gen.item()=:.3f}, {loss_fm.item()=:.3f}, {loss_mel.item()=:.3f}, {loss_kl.item()=:.3f}") + logger.info(f"|| {aux_loss.item()=:.3f}: {harmonic_loss.item()=:.3f}, {tefs_loss.item()=:.3f}, {tsi_loss.item()=:.3f}") #sigmoid scaling of ema weights_decay = commons.sigmoid_value(global_step,total_steps=10000,start_value=.5, end_value=.999, midpoint=.2) @@ -704,13 +724,13 @@ def train_and_evaluate( if loss_gen_all>>[lowest loss]: {least_loss:.3f}<<<") + logger.info(f"\t>>>[lowest loss]: {least_loss.item():.3f}<<<") if hps.save_best_model: if hasattr(net_g, "module"): ckpt = net_g.module.state_dict() else: ckpt = net_g.state_dict() - best_model_name = f"{hps.name}_e{epoch}_s{global_step}_loss{least_loss:.0f}" if hps.save_every_weights else f"{hps.name}_loss{least_loss:2.0f}" + best_model_name = f"{hps.name}_e{epoch}_s{global_step}_loss{least_loss.item():.0f}" if hps.save_every_weights else f"{hps.name}_loss{least_loss.item():2.0f}" status = save_checkpoint(ckpt,best_model_name,epoch,hps) logger.info(f"=== saving best model {best_model_name}: {status=} ===") @@ -755,6 +775,5 @@ def train_and_evaluate( os._exit(0) if __name__ == "__main__": - torch.multiprocessing.set_start_method("spawn") hps = utils.get_hparams() train_model(hps) \ No newline at end of file diff --git a/uvr5_cli.py b/uvr5_cli.py index c634254..2188586 100644 --- a/uvr5_cli.py +++ b/uvr5_cli.py @@ -4,7 +4,7 @@ import os, torch, warnings from .lib.separators import MDXNet, UVR5Base, UVR5New from .lib import BASE_CACHE_DIR, karafan from .lib.audio import load_input_audio, pad_audio, remix_audio, save_input_audio -from .lib.utils import gc_collect, get_optimal_threads, get_merge_func +from .lib.helper import gc_collect, get_optimal_threads, get_merge_func CACHED_SONGS_DIR = os.path.join(BASE_CACHE_DIR,"songs") diff --git a/vc_infer_pipeline.py b/vc_infer_pipeline.py index d6598c0..4a8d0b2 100644 --- a/vc_infer_pipeline.py +++ b/vc_infer_pipeline.py @@ -12,7 +12,7 @@ from .pitch_extraction import FeatureExtractor from .lib.audio import MAX_INT16, load_input_audio, remix_audio from .config import config -from .lib.utils import gc_collect +from .lib.helper import gc_collect # torchcrepe = lazyload("torchcrepe") # Fork Feature. Crepe algo for training and preprocess # torch = lazyload("torch")