From 37eb3b50da8fdcb24f8b420b5f91246d85d4f5b5 Mon Sep 17 00:00:00 2001 From: unknown Date: Tue, 29 Oct 2024 14:20:50 +0200 Subject: [PATCH 1/6] add tensorboard and add export sample for mel and audio --- src/f5_tts/model/trainer.py | 161 +++++++++++++++++++++++++--- src/f5_tts/train/finetune_cli.py | 11 ++ src/f5_tts/train/finetune_gradio.py | 79 ++++++++++++++ 3 files changed, 237 insertions(+), 14 deletions(-) diff --git a/src/f5_tts/model/trainer.py b/src/f5_tts/model/trainer.py index 3e64112..88652f1 100644 --- a/src/f5_tts/model/trainer.py +++ b/src/f5_tts/model/trainer.py @@ -3,7 +3,11 @@ from __future__ import annotations import os import gc from tqdm import tqdm -import wandb + +try: + from torch.utils.tensorboard import SummaryWriter +except ImportError: + print("TensorBoard is not installed") import torch from torch.optim import AdamW @@ -19,9 +23,26 @@ from f5_tts.model import CFM from f5_tts.model.utils import exists, default from f5_tts.model.dataset import DynamicBatchSampler, collate_fn - +import numpy as np +import matplotlib.pyplot as plt # trainer +# audio imports +import torchaudio +import soundfile as sf +from vocos import Vocos +import warnings + +warnings.filterwarnings("ignore", category=FutureWarning) + +# ----------------------------------------- +target_sample_rate = 24000 +hop_length = 256 +nfe_step = 16 +cfg_strength = 2.0 +sway_sampling_coef = -1.0 +# ----------------------------------------- + class Trainer: def __init__( @@ -39,6 +60,8 @@ class Trainer: max_grad_norm=1.0, noise_scheduler: str | None = None, duration_predictor: torch.nn.Module | None = None, + logger: str = "wandb", # Add logger parameter wandb,tensorboard , none + log_dir: str = "logs", # Add log directory parameter wandb_project="test_e2-tts", wandb_run_name="test_run", wandb_resume_id: str = None, @@ -46,24 +69,24 @@ class Trainer: accelerate_kwargs: dict = dict(), ema_kwargs: dict = dict(), bnb_optimizer: bool = False, + export_samples=False, ): ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) - logger = "wandb" if wandb.api.api_key else None - print(f"Using logger: {logger}") + self.logger = logger + if self.logger == "wandb": + self.accelerator = Accelerator( + log_with="wandb", + kwargs_handlers=[ddp_kwargs], + gradient_accumulation_steps=grad_accumulation_steps, + **accelerate_kwargs, + ) - self.accelerator = Accelerator( - log_with=logger, - kwargs_handlers=[ddp_kwargs], - gradient_accumulation_steps=grad_accumulation_steps, - **accelerate_kwargs, - ) - - if logger == "wandb": if exists(wandb_resume_id): init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name, "id": wandb_resume_id}} else: init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name}} + self.accelerator.init_trackers( project_name=wandb_project, init_kwargs=init_kwargs, @@ -80,12 +103,37 @@ class Trainer: "noise_scheduler": noise_scheduler, }, ) + elif self.logger == "tensorboard": + self.accelerator = Accelerator( + kwargs_handlers=[ddp_kwargs], + gradient_accumulation_steps=grad_accumulation_steps, + **accelerate_kwargs, + ) + if self.is_main: + path_log_dir = os.path.join(log_dir, wandb_project) + os.makedirs(path_log_dir, exist_ok=True) + existing_folders = [folder for folder in os.listdir(path_log_dir) if folder.startswith("exp")] + next_number = len(existing_folders) + 2 + folder_name = f"exp{next_number}" + folder_path = os.path.join(path_log_dir, folder_name) + os.makedirs(folder_path, exist_ok=True) + + self.writer = SummaryWriter(log_dir=folder_path) + + # export audio and mel + self.export_samples = export_samples + if self.export_samples: + self.path_ckpts_project = checkpoint_path + + self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz") + self.vocos.to("cpu") + self.file_path_samples = os.path.join(self.path_ckpts_project, "samples") + os.makedirs(self.file_path_samples, exist_ok=True) self.model = model if self.is_main: self.ema_model = EMA(model, include_online_model=False, **ema_kwargs) - self.ema_model.to(self.accelerator.device) self.epochs = epochs @@ -175,6 +223,82 @@ class Trainer: gc.collect() return step + def log(self, metrics, step): + """Unified logging method for both WandB and TensorBoard""" + if self.logger == "none": + return + if self.logger == "wandb": + self.accelerator.log(metrics, step=step) + elif self.is_main: + for key, value in metrics.items(): + self.writer.add_scalar(key, value, step) + + def export_add_log(self, global_step, mel_org, text_inputs): + try: + generated_wave_org = self.vocos.decode(mel_org.unsqueeze(0).cpu()) + generated_wave_org = generated_wave_org.squeeze().cpu().numpy() + file_wav_org = os.path.join(self.file_path_samples, f"step_{global_step}_org.wav") + sf.write(file_wav_org, generated_wave_org, target_sample_rate) + + audio, sr = torchaudio.load(file_wav_org) + audio = audio.to("cuda") + + ref_audio_len = audio.shape[-1] // hop_length + text = [text_inputs[0] + [" . "] + text_inputs[0]] + duration = int((audio.shape[1] / 256) * 2.0) + + with torch.inference_mode(): + generated_gen, _ = self.model.sample( + cond=audio, + text=text, + duration=duration, + steps=nfe_step, + cfg_strength=cfg_strength, + sway_sampling_coef=sway_sampling_coef, + ) + + generated_gen = generated_gen.to(torch.float32) + generated_gen = generated_gen[:, ref_audio_len:, :] + generated_mel_spec_gen = generated_gen.permute(0, 2, 1) + generated_wave_gen = self.vocos.decode(generated_mel_spec_gen.cpu()) + generated_wave_gen = generated_wave_gen.squeeze().cpu().numpy() + file_wav_gen = os.path.join(self.file_path_samples, f"step_{global_step}_gen.wav") + sf.write(file_wav_gen, generated_wave_gen, target_sample_rate) + + if self.logger == "tensorboard": + self.writer.add_audio("Audio/original", generated_wave_org, global_step, sample_rate=target_sample_rate) + + self.writer.add_audio("Audio/generate", generated_wave_gen, global_step, sample_rate=target_sample_rate) + + mel_org = mel_org + mel_min, mel_max = mel_org.min(), mel_org.max() + mel_norm = (mel_org - mel_min) / (mel_max - mel_min + 1e-8) + mel_colored = plt.get_cmap("viridis")(mel_norm.detach().cpu().numpy())[:, :, :3] + mel_colored = np.transpose(mel_colored, (2, 0, 1)) + + if self.logger == "tensorboard": + self.writer.add_image("Mel/oginal", mel_colored, global_step, dataformats="CHW") + + mel_colored_hwc = np.transpose(mel_colored, (1, 2, 0)) + file_gen_org = os.path.join(self.file_path_samples, f"step_{global_step}_org.png") + plt.imsave(file_gen_org, mel_colored_hwc) + + mel_gen = generated_mel_spec_gen[0] + mel_min, mel_max = mel_gen.min(), mel_gen.max() + mel_norm = (mel_gen - mel_min) / (mel_max - mel_min + 1e-8) + mel_colored = plt.get_cmap("viridis")(mel_norm.detach().cpu().numpy())[:, :, :3] + mel_colored = np.transpose(mel_colored, (2, 0, 1)) + + if self.logger == "tensorboard": + self.writer.add_image("Mel/generate", mel_colored, global_step, dataformats="CHW") + + mel_colored_hwc = np.transpose(mel_colored, (1, 2, 0)) + file_gen_gen = os.path.join(self.file_path_samples, f"step_{global_step}_gen.png") + plt.imsave(file_gen_gen, mel_colored_hwc) + + except Exception as e: + print("An error occurred:", e) + def train(self, train_dataset: Dataset, num_workers=16, resumable_with_seed: int = None): if exists(resumable_with_seed): generator = torch.Generator() @@ -270,6 +394,15 @@ class Trainer: loss, cond, pred = self.model( mel_spec, text=text_inputs, lens=mel_lengths, noise_scheduler=self.noise_scheduler ) + + # save 4 audio per save step + if ( + self.accelerator.is_local_main_process + and self.export_samples + and global_step % (int(self.save_per_updates * 0.25) * self.grad_accumulation_steps) == 0 + ): + self.export_add_log(global_step, batch["mel"][0], text_inputs) + self.accelerator.backward(loss) if self.max_grad_norm > 0 and self.accelerator.sync_gradients: @@ -285,7 +418,7 @@ class Trainer: global_step += 1 if self.accelerator.is_local_main_process: - self.accelerator.log({"loss": loss.item(), "lr": self.scheduler.get_last_lr()[0]}, step=global_step) + self.log({"loss": loss.item(), "lr": self.scheduler.get_last_lr()[0]}, step=global_step) progress_bar.set_postfix(step=str(global_step), loss=loss.item()) diff --git a/src/f5_tts/train/finetune_cli.py b/src/f5_tts/train/finetune_cli.py index 9a95647..c09ae7d 100644 --- a/src/f5_tts/train/finetune_cli.py +++ b/src/f5_tts/train/finetune_cli.py @@ -56,6 +56,14 @@ def parse_args(): help="Path to custom tokenizer vocab file (only used if tokenizer = 'custom')", ) + parser.add_argument( + "--export_samples", + type=bool, + default=False, + help="Export 4 audio and spect samples for the checkpoint audio, per step.", + ) + parser.add_argument("--logger", type=str, default="wandb", choices=["none", "wandb", "tensorboard"], help="logger") + return parser.parse_args() @@ -64,6 +72,7 @@ def parse_args(): def main(): args = parse_args() + checkpoint_path = str(files("f5_tts").joinpath(f"../../ckpts/{args.dataset_name}")) # Model parameters based on experiment name @@ -136,6 +145,8 @@ def main(): wandb_run_name=args.exp_name, wandb_resume_id=wandb_resume_id, last_per_steps=args.last_per_steps, + logger=args.logger, + export_samples=args.export_samples, ) train_dataset = load_dataset(args.dataset_name, tokenizer, mel_spec_kwargs=mel_spec_kwargs) diff --git a/src/f5_tts/train/finetune_gradio.py b/src/f5_tts/train/finetune_gradio.py index 0637b38..3693f2f 100644 --- a/src/f5_tts/train/finetune_gradio.py +++ b/src/f5_tts/train/finetune_gradio.py @@ -447,6 +447,8 @@ def start_training( cmd += f" --tokenizer {tokenizer_type} " + cmd += " --export_samples True --logger wandb " + print(cmd) save_settings( @@ -1223,6 +1225,27 @@ def get_checkpoints_project(project_name, is_gradio=True): return files_checkpoints, selelect_checkpoint +def get_audio_project(project_name, is_gradio=True): + if project_name is None: + return [], "" + project_name = project_name.replace("_pinyin", "").replace("_char", "") + + if os.path.isdir(path_project_ckpts): + files_audios = glob(os.path.join(path_project_ckpts, project_name, "samples", "*.wav")) + files_audios = sorted(files_audios, key=lambda x: int(os.path.basename(x).split("_")[1].split(".")[0])) + + files_audios = [item.replace("_gen.wav", "") for item in files_audios if item.endswith("_gen.wav")] + else: + files_audios = [] + + selelect_checkpoint = None if not files_audios else files_audios[0] + + if is_gradio: + return gr.update(choices=files_audios, value=selelect_checkpoint) + + return files_audios, selelect_checkpoint + + def get_gpu_stats(): gpu_stats = "" @@ -1290,6 +1313,21 @@ def get_combined_stats(): return combined_stats +def get_audio_select(file_sample): + select_audio_org = file_sample + select_audio_gen = file_sample + select_image_org = file_sample + select_image_gen = file_sample + + if file_sample is not None: + select_audio_org += "_org.wav" + select_audio_gen += "_gen.wav" + select_image_org += "_org.png" + select_image_gen += "_gen.png" + + return select_audio_org, select_audio_gen, select_image_org, select_image_gen + + with gr.Blocks() as app: gr.Markdown( """ @@ -1511,6 +1549,47 @@ If you encounter a memory error, try reducing the batch size per GPU to a smalle ch_stream = gr.Checkbox(label="stream output experiment.", value=True) txt_info_train = gr.Text(label="info", value="") + + list_audios, select_audio = get_audio_project(projects_selelect, False) + + select_audio_org = select_audio + select_audio_gen = select_audio + select_image_org = select_audio + select_image_gen = select_audio + + if select_audio is not None: + select_audio_org += "_org.wav" + select_audio_gen += "_gen.wav" + select_image_org += "_org.png" + select_image_gen += "_gen.png" + + with gr.Row(): + ch_list_audio = gr.Dropdown( + choices=list_audios, + value=select_audio, + label="audios", + allow_custom_value=True, + scale=6, + interactive=True, + ) + bt_stream_audio = gr.Button("refresh", scale=1) + bt_stream_audio.click(fn=get_audio_project, inputs=[cm_project], outputs=[ch_list_audio]) + cm_project.change(fn=get_audio_project, inputs=[cm_project], outputs=[ch_list_audio]) + + with gr.Row(): + audio_org_stream = gr.Audio(label="original", type="filepath", value=select_audio_org) + mel_org_stream = gr.Image(label="original", type="filepath", value=select_image_org) + + with gr.Row(): + audio_gen_stream = gr.Audio(label="generate", type="filepath", value=select_audio_gen) + mel_gen_stream = gr.Image(label="generate", type="filepath", value=select_image_gen) + + ch_list_audio.change( + fn=get_audio_select, + inputs=[ch_list_audio], + outputs=[audio_org_stream, audio_gen_stream, mel_org_stream, mel_gen_stream], + ) + start_button.click( fn=start_training, inputs=[ From 3409192662062f297e9c252fb7890988c57963d7 Mon Sep 17 00:00:00 2001 From: unknown Date: Tue, 29 Oct 2024 16:44:23 +0200 Subject: [PATCH 2/6] update --- src/f5_tts/model/trainer.py | 120 +++++++++--------------------------- src/f5_tts/model/utils.py | 74 ++++++++++++++++++++++ 2 files changed, 102 insertions(+), 92 deletions(-) diff --git a/src/f5_tts/model/trainer.py b/src/f5_tts/model/trainer.py index 88652f1..c45b6b1 100644 --- a/src/f5_tts/model/trainer.py +++ b/src/f5_tts/model/trainer.py @@ -4,10 +4,6 @@ import os import gc from tqdm import tqdm -try: - from torch.utils.tensorboard import SummaryWriter -except ImportError: - print("TensorBoard is not installed") import torch from torch.optim import AdamW @@ -22,27 +18,11 @@ from ema_pytorch import EMA from f5_tts.model import CFM from f5_tts.model.utils import exists, default from f5_tts.model.dataset import DynamicBatchSampler, collate_fn +from f5_tts.infer.utils_infer import target_sample_rate, hop_length, nfe_step, cfg_strength, sway_sampling_coef, vocos +from f5_tts.model.utils import gen_sample -import numpy as np -import matplotlib.pyplot as plt # trainer -# audio imports -import torchaudio -import soundfile as sf -from vocos import Vocos -import warnings - -warnings.filterwarnings("ignore", category=FutureWarning) - -# ----------------------------------------- -target_sample_rate = 24000 -hop_length = 256 -nfe_step = 16 -cfg_strength = 2.0 -sway_sampling_coef = -1.0 -# ----------------------------------------- - class Trainer: def __init__( @@ -104,6 +84,8 @@ class Trainer: }, ) elif self.logger == "tensorboard": + from torch.utils.tensorboard import SummaryWriter + self.accelerator = Accelerator( kwargs_handlers=[ddp_kwargs], gradient_accumulation_steps=grad_accumulation_steps, @@ -124,9 +106,6 @@ class Trainer: self.export_samples = export_samples if self.export_samples: self.path_ckpts_project = checkpoint_path - - self.vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz") - self.vocos.to("cpu") self.file_path_samples = os.path.join(self.path_ckpts_project, "samples") os.makedirs(self.file_path_samples, exist_ok=True) @@ -233,72 +212,6 @@ class Trainer: for key, value in metrics.items(): self.writer.add_scalar(key, value, step) - def export_add_log(self, global_step, mel_org, text_inputs): - try: - generated_wave_org = self.vocos.decode(mel_org.unsqueeze(0).cpu()) - generated_wave_org = generated_wave_org.squeeze().cpu().numpy() - file_wav_org = os.path.join(self.file_path_samples, f"step_{global_step}_org.wav") - sf.write(file_wav_org, generated_wave_org, target_sample_rate) - - audio, sr = torchaudio.load(file_wav_org) - audio = audio.to("cuda") - - ref_audio_len = audio.shape[-1] // hop_length - text = [text_inputs[0] + [" . "] + text_inputs[0]] - duration = int((audio.shape[1] / 256) * 2.0) - - with torch.inference_mode(): - generated_gen, _ = self.model.sample( - cond=audio, - text=text, - duration=duration, - steps=nfe_step, - cfg_strength=cfg_strength, - sway_sampling_coef=sway_sampling_coef, - ) - - generated_gen = generated_gen.to(torch.float32) - generated_gen = generated_gen[:, ref_audio_len:, :] - generated_mel_spec_gen = generated_gen.permute(0, 2, 1) - generated_wave_gen = self.vocos.decode(generated_mel_spec_gen.cpu()) - generated_wave_gen = generated_wave_gen.squeeze().cpu().numpy() - file_wav_gen = os.path.join(self.file_path_samples, f"step_{global_step}_gen.wav") - sf.write(file_wav_gen, generated_wave_gen, target_sample_rate) - - if self.logger == "tensorboard": - self.writer.add_audio("Audio/original", generated_wave_org, global_step, sample_rate=target_sample_rate) - - self.writer.add_audio("Audio/generate", generated_wave_gen, global_step, sample_rate=target_sample_rate) - - mel_org = mel_org - mel_min, mel_max = mel_org.min(), mel_org.max() - mel_norm = (mel_org - mel_min) / (mel_max - mel_min + 1e-8) - mel_colored = plt.get_cmap("viridis")(mel_norm.detach().cpu().numpy())[:, :, :3] - mel_colored = np.transpose(mel_colored, (2, 0, 1)) - - if self.logger == "tensorboard": - self.writer.add_image("Mel/oginal", mel_colored, global_step, dataformats="CHW") - - mel_colored_hwc = np.transpose(mel_colored, (1, 2, 0)) - file_gen_org = os.path.join(self.file_path_samples, f"step_{global_step}_org.png") - plt.imsave(file_gen_org, mel_colored_hwc) - - mel_gen = generated_mel_spec_gen[0] - mel_min, mel_max = mel_gen.min(), mel_gen.max() - mel_norm = (mel_gen - mel_min) / (mel_max - mel_min + 1e-8) - mel_colored = plt.get_cmap("viridis")(mel_norm.detach().cpu().numpy())[:, :, :3] - mel_colored = np.transpose(mel_colored, (2, 0, 1)) - - if self.logger == "tensorboard": - self.writer.add_image("Mel/generate", mel_colored, global_step, dataformats="CHW") - - mel_colored_hwc = np.transpose(mel_colored, (1, 2, 0)) - file_gen_gen = os.path.join(self.file_path_samples, f"step_{global_step}_gen.png") - plt.imsave(file_gen_gen, mel_colored_hwc) - - except Exception as e: - print("An error occurred:", e) - def train(self, train_dataset: Dataset, num_workers=16, resumable_with_seed: int = None): if exists(resumable_with_seed): generator = torch.Generator() @@ -383,6 +296,7 @@ class Trainer: for batch in progress_bar: with self.accelerator.accumulate(self.model): text_inputs = batch["text"] + mel_spec = batch["mel"].permute(0, 2, 1) mel_lengths = batch["mel_lengths"] @@ -401,7 +315,29 @@ class Trainer: and self.export_samples and global_step % (int(self.save_per_updates * 0.25) * self.grad_accumulation_steps) == 0 ): - self.export_add_log(global_step, batch["mel"][0], text_inputs) + wave_org, wave_gen, mel_org, mel_gen = gen_sample( + vocos, + self.model, + self.file_path_samples, + global_step, + batch["mel"][0], + text_inputs, + target_sample_rate, + hop_length, + nfe_step, + cfg_strength, + sway_sampling_coef, + ) + + if self.logger == "tensorboard": + self.writer.add_audio( + "Audio/original", wave_org, global_step, sample_rate=target_sample_rate + ) + self.writer.add_audio( + "Audio/generate", wave_gen, global_step, sample_rate=target_sample_rate + ) + self.writer.add_image("Mel/original", mel_org, global_step, dataformats="CHW") + self.writer.add_image("Mel/generate", mel_gen, global_step, dataformats="CHW") self.accelerator.backward(loss) diff --git a/src/f5_tts/model/utils.py b/src/f5_tts/model/utils.py index 76cfa4d..8877a7e 100644 --- a/src/f5_tts/model/utils.py +++ b/src/f5_tts/model/utils.py @@ -11,6 +11,10 @@ from torch.nn.utils.rnn import pad_sequence import jieba from pypinyin import lazy_pinyin, Style +import numpy as np +import matplotlib.pyplot as plt +import soundfile as sf +import torchaudio # seed everything @@ -183,3 +187,73 @@ def repetition_found(text, length=2, tolerance=10): if count > tolerance: return True return False + + +def normalize_and_colorize_spectrogram(mel_org): + mel_min, mel_max = mel_org.min(), mel_org.max() + mel_norm = (mel_org - mel_min) / (mel_max - mel_min + 1e-8) + mel_colored = plt.get_cmap("viridis")(mel_norm.detach().cpu().numpy())[:, :, :3] + mel_colored = np.transpose(mel_colored, (2, 0, 1)) + return mel_colored + + +def export_audio(file_out, wav, target_sample_rate): + sf.write(file_out, wav, samplerate=target_sample_rate) + + +def export_mel(mel_colored_hwc, file_out): + plt.imsave(file_out, mel_colored_hwc) + + +def get_sample(model, vocos, file_wav_org, text_inputs, hop_length, nfe_step, cfg_strength, sway_sampling_coef): + audio, sr = torchaudio.load(file_wav_org) + audio = audio.to("cuda") + ref_audio_len = audio.shape[-1] // hop_length + text = [text_inputs[0] + [" . "] + text_inputs[0]] + duration = int((audio.shape[1] / 256) * 2.0) + with torch.inference_mode(): + generated_gen, _ = model.sample( + cond=audio, + text=text, + duration=duration, + steps=nfe_step, + cfg_strength=cfg_strength, + sway_sampling_coef=sway_sampling_coef, + ) + generated_gen = generated_gen.to(torch.float32) + generated_gen = generated_gen[:, ref_audio_len:, :] + generated_mel_spec_gen = generated_gen.permute(0, 2, 1) + generated_wave_gen = vocos.decode(generated_mel_spec_gen.cpu()) + generated_wave_gen = generated_wave_gen.squeeze().cpu().numpy() + return generated_wave_gen, generated_mel_spec_gen + + +def gen_sample( + vocos, + model, + file_path_samples, + global_step, + mel_org, + text_inputs, + target_sample_rate, + hop_length, + nfe_step, + cfg_strength, + sway_sampling_coef, +): + generated_wave_org = vocos.decode(mel_org.unsqueeze(0).cpu()) + generated_wave_org = generated_wave_org.squeeze().cpu().numpy() + file_wav_org = os.path.join(file_path_samples, f"step_{global_step}_org.wav") + export_audio(file_wav_org, generated_wave_org, target_sample_rate) + generated_wave_gen, generated_mel_spec_gen = get_sample( + model, vocos, file_wav_org, text_inputs, hop_length, nfe_step, cfg_strength, sway_sampling_coef + ) + file_wav_gen = os.path.join(file_path_samples, f"step_{global_step}_gen.wav") + export_audio(file_wav_gen, generated_wave_gen, target_sample_rate) + mel_org = normalize_and_colorize_spectrogram(mel_org) + mel_gen = normalize_and_colorize_spectrogram(generated_mel_spec_gen[0]) + file_gen_org = os.path.join(file_path_samples, f"step_{global_step}_org.png") + export_mel(np.transpose(mel_org, (1, 2, 0)), file_gen_org) + file_gen_gen = os.path.join(file_path_samples, f"step_{global_step}_gen.png") + export_mel(np.transpose(mel_gen, (1, 2, 0)), file_gen_gen) + return generated_wave_org, generated_wave_gen, mel_org, mel_gen From 2ca1fb7c2563f35d0fff48e00e4a690c206ced81 Mon Sep 17 00:00:00 2001 From: unknown Date: Tue, 29 Oct 2024 16:46:38 +0200 Subject: [PATCH 3/6] update --- src/f5_tts/model/trainer.py | 4 ++-- src/f5_tts/model/utils.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/f5_tts/model/trainer.py b/src/f5_tts/model/trainer.py index c45b6b1..e9ed5d4 100644 --- a/src/f5_tts/model/trainer.py +++ b/src/f5_tts/model/trainer.py @@ -19,7 +19,7 @@ from f5_tts.model import CFM from f5_tts.model.utils import exists, default from f5_tts.model.dataset import DynamicBatchSampler, collate_fn from f5_tts.infer.utils_infer import target_sample_rate, hop_length, nfe_step, cfg_strength, sway_sampling_coef, vocos -from f5_tts.model.utils import gen_sample +from f5_tts.model.utils import get_sample # trainer @@ -315,7 +315,7 @@ class Trainer: and self.export_samples and global_step % (int(self.save_per_updates * 0.25) * self.grad_accumulation_steps) == 0 ): - wave_org, wave_gen, mel_org, mel_gen = gen_sample( + wave_org, wave_gen, mel_org, mel_gen = get_sample( vocos, self.model, self.file_path_samples, diff --git a/src/f5_tts/model/utils.py b/src/f5_tts/model/utils.py index 8877a7e..0314d1f 100644 --- a/src/f5_tts/model/utils.py +++ b/src/f5_tts/model/utils.py @@ -205,7 +205,7 @@ def export_mel(mel_colored_hwc, file_out): plt.imsave(file_out, mel_colored_hwc) -def get_sample(model, vocos, file_wav_org, text_inputs, hop_length, nfe_step, cfg_strength, sway_sampling_coef): +def gen_sample(model, vocos, file_wav_org, text_inputs, hop_length, nfe_step, cfg_strength, sway_sampling_coef): audio, sr = torchaudio.load(file_wav_org) audio = audio.to("cuda") ref_audio_len = audio.shape[-1] // hop_length @@ -228,7 +228,7 @@ def get_sample(model, vocos, file_wav_org, text_inputs, hop_length, nfe_step, cf return generated_wave_gen, generated_mel_spec_gen -def gen_sample( +def get_sample( vocos, model, file_path_samples, @@ -245,7 +245,7 @@ def gen_sample( generated_wave_org = generated_wave_org.squeeze().cpu().numpy() file_wav_org = os.path.join(file_path_samples, f"step_{global_step}_org.wav") export_audio(file_wav_org, generated_wave_org, target_sample_rate) - generated_wave_gen, generated_mel_spec_gen = get_sample( + generated_wave_gen, generated_mel_spec_gen = gen_sample( model, vocos, file_wav_org, text_inputs, hop_length, nfe_step, cfg_strength, sway_sampling_coef ) file_wav_gen = os.path.join(file_path_samples, f"step_{global_step}_gen.wav") From 886500ac971d7d467c84b7a2377e8061f3702477 Mon Sep 17 00:00:00 2001 From: unknown Date: Tue, 29 Oct 2024 17:02:21 +0200 Subject: [PATCH 4/6] update --- src/f5_tts/model/trainer.py | 29 ++++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/src/f5_tts/model/trainer.py b/src/f5_tts/model/trainer.py index e9ed5d4..420089a 100644 --- a/src/f5_tts/model/trainer.py +++ b/src/f5_tts/model/trainer.py @@ -18,8 +18,6 @@ from ema_pytorch import EMA from f5_tts.model import CFM from f5_tts.model.utils import exists, default from f5_tts.model.dataset import DynamicBatchSampler, collate_fn -from f5_tts.infer.utils_infer import target_sample_rate, hop_length, nfe_step, cfg_strength, sway_sampling_coef, vocos -from f5_tts.model.utils import get_sample # trainer @@ -51,6 +49,11 @@ class Trainer: bnb_optimizer: bool = False, export_samples=False, ): + # export audio and mel + self.export_samples = export_samples + if export_samples: + self.path_ckpts_project = checkpoint_path + ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) self.logger = logger @@ -102,13 +105,6 @@ class Trainer: self.writer = SummaryWriter(log_dir=folder_path) - # export audio and mel - self.export_samples = export_samples - if self.export_samples: - self.path_ckpts_project = checkpoint_path - self.file_path_samples = os.path.join(self.path_ckpts_project, "samples") - os.makedirs(self.file_path_samples, exist_ok=True) - self.model = model if self.is_main: @@ -213,6 +209,21 @@ class Trainer: self.writer.add_scalar(key, value, step) def train(self, train_dataset: Dataset, num_workers=16, resumable_with_seed: int = None): + # import only when export_sample True + if self.export_samples: + from f5_tts.infer.utils_infer import ( + target_sample_rate, + hop_length, + nfe_step, + cfg_strength, + sway_sampling_coef, + vocos, + ) + from f5_tts.model.utils import get_sample + + self.file_path_samples = os.path.join(self.path_ckpts_project, "samples") + os.makedirs(self.file_path_samples, exist_ok=True) + if exists(resumable_with_seed): generator = torch.Generator() generator.manual_seed(resumable_with_seed) From 8a5041ef9f4b2d217a7709e6d86a3f3fb37c30d2 Mon Sep 17 00:00:00 2001 From: unknown Date: Tue, 29 Oct 2024 17:14:03 +0200 Subject: [PATCH 5/6] update --- src/f5_tts/model/trainer.py | 47 ++++++++++++++++++++----------------- 1 file changed, 25 insertions(+), 22 deletions(-) diff --git a/src/f5_tts/model/trainer.py b/src/f5_tts/model/trainer.py index 420089a..10aa275 100644 --- a/src/f5_tts/model/trainer.py +++ b/src/f5_tts/model/trainer.py @@ -326,29 +326,32 @@ class Trainer: and self.export_samples and global_step % (int(self.save_per_updates * 0.25) * self.grad_accumulation_steps) == 0 ): - wave_org, wave_gen, mel_org, mel_gen = get_sample( - vocos, - self.model, - self.file_path_samples, - global_step, - batch["mel"][0], - text_inputs, - target_sample_rate, - hop_length, - nfe_step, - cfg_strength, - sway_sampling_coef, - ) + try: + wave_org, wave_gen, mel_org, mel_gen = get_sample( + vocos, + self.model, + self.file_path_samples, + global_step, + batch["mel"][0], + text_inputs, + target_sample_rate, + hop_length, + nfe_step, + cfg_strength, + sway_sampling_coef, + ) - if self.logger == "tensorboard": - self.writer.add_audio( - "Audio/original", wave_org, global_step, sample_rate=target_sample_rate - ) - self.writer.add_audio( - "Audio/generate", wave_gen, global_step, sample_rate=target_sample_rate - ) - self.writer.add_image("Mel/original", mel_org, global_step, dataformats="CHW") - self.writer.add_image("Mel/generate", mel_gen, global_step, dataformats="CHW") + if self.logger == "tensorboard": + self.writer.add_audio( + "Audio/original", wave_org, global_step, sample_rate=target_sample_rate + ) + self.writer.add_audio( + "Audio/generate", wave_gen, global_step, sample_rate=target_sample_rate + ) + self.writer.add_image("Mel/original", mel_org, global_step, dataformats="CHW") + self.writer.add_image("Mel/generate", mel_gen, global_step, dataformats="CHW") + except Exception as e: + print("An error occurred:", e) self.accelerator.backward(loss) From d5c307b56aa85bbc9b210ebe459ea4a11317053f Mon Sep 17 00:00:00 2001 From: unknown Date: Tue, 29 Oct 2024 17:30:28 +0200 Subject: [PATCH 6/6] add logger --- src/f5_tts/train/finetune_gradio.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/f5_tts/train/finetune_gradio.py b/src/f5_tts/train/finetune_gradio.py index 3693f2f..b5a45c7 100644 --- a/src/f5_tts/train/finetune_gradio.py +++ b/src/f5_tts/train/finetune_gradio.py @@ -69,6 +69,7 @@ def save_settings( tokenizer_type, tokenizer_file, mixed_precision, + logger, ): path_project = os.path.join(path_project_ckpts, project_name) os.makedirs(path_project, exist_ok=True) @@ -91,6 +92,7 @@ def save_settings( "tokenizer_type": tokenizer_type, "tokenizer_file": tokenizer_file, "mixed_precision": mixed_precision, + "logger": logger, } with open(file_setting, "w") as f: json.dump(settings, f, indent=4) @@ -121,6 +123,7 @@ def load_settings(project_name): "tokenizer_type": "pinyin", "tokenizer_file": "", "mixed_precision": "none", + "logger": "wandb", } return ( settings["exp_name"], @@ -139,6 +142,7 @@ def load_settings(project_name): settings["tokenizer_type"], settings["tokenizer_file"], settings["mixed_precision"], + settings["logger"], ) with open(file_setting, "r") as f: @@ -160,6 +164,7 @@ def load_settings(project_name): settings["tokenizer_type"], settings["tokenizer_file"], settings["mixed_precision"], + settings["logger"], ) @@ -374,6 +379,7 @@ def start_training( tokenizer_file="", mixed_precision="fp16", stream=False, + logger="wandb", ): global training_process, tts_api, stop_signal @@ -447,7 +453,7 @@ def start_training( cmd += f" --tokenizer {tokenizer_type} " - cmd += " --export_samples True --logger wandb " + cmd += f" --export_samples True --logger {logger} " print(cmd) @@ -469,6 +475,7 @@ def start_training( tokenizer_type, tokenizer_file, mixed_precision, + logger, ) try: @@ -1508,6 +1515,7 @@ If you encounter a memory error, try reducing the batch size per GPU to a smalle with gr.Row(): mixed_precision = gr.Radio(label="mixed_precision", choices=["none", "fp16", "fpb16"], value="none") + cd_logger = gr.Radio(label="logger", choices=["none", "wandb", "tensorboard"], value="wandb") start_button = gr.Button("Start Training") stop_button = gr.Button("Stop Training", interactive=False) @@ -1529,6 +1537,7 @@ If you encounter a memory error, try reducing the batch size per GPU to a smalle tokenizer_typev, tokenizer_filev, mixed_precisionv, + cd_loggerv, ) = load_settings(projects_selelect) exp_name.value = exp_namev learning_rate.value = learning_ratev @@ -1546,6 +1555,7 @@ If you encounter a memory error, try reducing the batch size per GPU to a smalle tokenizer_type.value = tokenizer_typev tokenizer_file.value = tokenizer_filev mixed_precision.value = mixed_precisionv + cd_logger.value = cd_loggerv ch_stream = gr.Checkbox(label="stream output experiment.", value=True) txt_info_train = gr.Text(label="info", value="") @@ -1611,6 +1621,7 @@ If you encounter a memory error, try reducing the batch size per GPU to a smalle tokenizer_file, mixed_precision, ch_stream, + cd_logger, ], outputs=[txt_info_train, start_button, stop_button], ) @@ -1662,6 +1673,7 @@ If you encounter a memory error, try reducing the batch size per GPU to a smalle tokenizer_type, tokenizer_file, mixed_precision, + cd_logger, ] return output_components