renamed utils to helpers

This commit is contained in:
SayanoAI
2025-06-23 10:44:15 -04:00
parent e7583b1601
commit b28fcaeae8
20 changed files with 106 additions and 83 deletions
+2 -2
View File
@@ -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"
+2 -2
View File
@@ -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
+2 -2
View File
@@ -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()
+2 -2
View File
@@ -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
+2 -2
View File
@@ -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"
+2 -2
View File
@@ -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
+4 -1
View File
@@ -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
+1 -1
View File
@@ -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
View File
+1 -1
View File
@@ -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:
+1 -1
View File
@@ -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")
+1 -1
View File
@@ -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
+15 -14
View File
@@ -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
+1 -1
View File
@@ -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):
+1 -1
View File
@@ -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
+1 -1
View File
@@ -1,5 +1,5 @@
librosa~=0.10.2
numpy
numpy<2
faiss-cpu==1.7.4
ffmpeg-python
tqdm
+66 -47
View File
@@ -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<least_loss:
least_loss = loss_gen_all
logger.info(f"\t>>>[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)
+1 -1
View File
@@ -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")
+1 -1
View File
@@ -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")