From da1b40968ac98937dfeb9af72e62c69cdcb6e5ab Mon Sep 17 00:00:00 2001 From: SWivid Date: Wed, 30 Oct 2024 02:38:58 +0800 Subject: [PATCH] basic structure --- src/f5_tts/model/trainer.py | 134 +++++++++------------------- src/f5_tts/model/utils.py | 75 ---------------- src/f5_tts/train/finetune_cli.py | 10 +-- src/f5_tts/train/finetune_gradio.py | 30 ++----- 4 files changed, 58 insertions(+), 191 deletions(-) diff --git a/src/f5_tts/model/trainer.py b/src/f5_tts/model/trainer.py index 10aa275..53fa8cd 100644 --- a/src/f5_tts/model/trainer.py +++ b/src/f5_tts/model/trainer.py @@ -3,9 +3,10 @@ from __future__ import annotations import os import gc from tqdm import tqdm - +import wandb import torch +import torchaudio from torch.optim import AdamW from torch.utils.data import DataLoader, Dataset, SequentialSampler from torch.optim.lr_scheduler import LinearLR, SequentialLR @@ -19,6 +20,7 @@ from f5_tts.model import CFM from f5_tts.model.utils import exists, default from f5_tts.model.dataset import DynamicBatchSampler, collate_fn + # trainer @@ -38,33 +40,32 @@ 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 + logger: str | None = "wandb", # "wandb" | "tensorboard" | None wandb_project="test_e2-tts", wandb_run_name="test_run", wandb_resume_id: str = None, + log_samples: bool = False, last_per_steps=None, accelerate_kwargs: dict = dict(), ema_kwargs: dict = dict(), 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) + if logger == "wandb" and not wandb.api.api_key: + logger = None + print(f"Using logger: {logger}") + self.log_samples = log_samples + + self.accelerator = Accelerator( + log_with=logger if logger == "wandb" else None, + kwargs_handlers=[ddp_kwargs], + gradient_accumulation_steps=grad_accumulation_steps, + **accelerate_kwargs, + ) + 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, - ) - if exists(wandb_resume_id): init_kwargs = {"wandb": {"resume": "allow", "name": wandb_run_name, "id": wandb_resume_id}} else: @@ -86,24 +87,11 @@ class Trainer: "noise_scheduler": noise_scheduler, }, ) + elif self.logger == "tensorboard": from torch.utils.tensorboard import SummaryWriter - 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) + self.writer = SummaryWriter(log_dir=f"runs/{wandb_run_name}") self.model = model @@ -198,31 +186,13 @@ 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 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 + if self.log_samples: + from f5_tts.infer.utils_infer import vocos, nfe_step, cfg_strength, sway_sampling_coef - self.file_path_samples = os.path.join(self.path_ckpts_project, "samples") - os.makedirs(self.file_path_samples, exist_ok=True) + target_sample_rate = self.model.mel_spec.mel_stft.sample_rate + log_samples_path = f"{self.checkpoint_path}/samples" + os.makedirs(log_samples_path, exist_ok=True) if exists(resumable_with_seed): generator = torch.Generator() @@ -307,7 +277,6 @@ 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"] @@ -319,40 +288,6 @@ 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 - ): - 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") - except Exception as e: - print("An error occurred:", e) - self.accelerator.backward(loss) if self.max_grad_norm > 0 and self.accelerator.sync_gradients: @@ -368,13 +303,32 @@ class Trainer: global_step += 1 if self.accelerator.is_local_main_process: - self.log({"loss": loss.item(), "lr": self.scheduler.get_last_lr()[0]}, step=global_step) + self.accelerator.log({"loss": loss.item(), "lr": self.scheduler.get_last_lr()[0]}, step=global_step) + if self.logger == "tensorboard": + self.writer.add_scalar("loss", loss.item(), global_step) + self.writer.add_scalar("lr", self.scheduler.get_last_lr()[0], global_step) progress_bar.set_postfix(step=str(global_step), loss=loss.item()) if global_step % (self.save_per_updates * self.grad_accumulation_steps) == 0: self.save_checkpoint(global_step) + if self.log_samples: + ref_audio, ref_audio_len = vocos.decode([batch["mel"][0]].cpu()), mel_lengths[0] + torchaudio.save(f"{log_samples_path}/step_{global_step}_ref.wav", ref_audio, target_sample_rate) + with torch.inference_mode(): + generated, _ = self.model.sample( + cond=[mel_spec[0][:ref_audio_len]], + text=[text_inputs[0] + [" "] + text_inputs[0]], + duration=ref_audio_len * 2, + steps=nfe_step, + cfg_strength=cfg_strength, + sway_sampling_coef=sway_sampling_coef, + ) + generated = generated.to(torch.float32) + gen_audio = vocos.decode(generated[:, ref_audio_len:, :].permute(0, 2, 1).cpu()) + torchaudio.save(f"{log_samples_path}/step_{global_step}_gen.wav", gen_audio, target_sample_rate) + if global_step % self.last_per_steps == 0: self.save_checkpoint(global_step, last=True) diff --git a/src/f5_tts/model/utils.py b/src/f5_tts/model/utils.py index 7ac4aaa..76cfa4d 100644 --- a/src/f5_tts/model/utils.py +++ b/src/f5_tts/model/utils.py @@ -11,10 +11,6 @@ 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 @@ -187,74 +183,3 @@ 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)) - mel_colored = np.flip(mel_colored, axis=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 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 - 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 get_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 = 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") - 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 diff --git a/src/f5_tts/train/finetune_cli.py b/src/f5_tts/train/finetune_cli.py index c09ae7d..43d766c 100644 --- a/src/f5_tts/train/finetune_cli.py +++ b/src/f5_tts/train/finetune_cli.py @@ -57,12 +57,12 @@ def parse_args(): ) parser.add_argument( - "--export_samples", + "--log_samples", type=bool, default=False, - help="Export 4 audio and spect samples for the checkpoint audio, per step.", + help="Log inferenced samples per ckpt save steps", ) - parser.add_argument("--logger", type=str, default="wandb", choices=["none", "wandb", "tensorboard"], help="logger") + parser.add_argument("--logger", type=str, default=None, choices=["wandb", "tensorboard"], help="logger") return parser.parse_args() @@ -141,12 +141,12 @@ def main(): max_samples=args.max_samples, grad_accumulation_steps=args.grad_accumulation_steps, max_grad_norm=args.max_grad_norm, + logger=args.logger, wandb_project=args.dataset_name, wandb_run_name=args.exp_name, wandb_resume_id=wandb_resume_id, + log_samples=args.log_samples, 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 b5a45c7..725320d 100644 --- a/src/f5_tts/train/finetune_gradio.py +++ b/src/f5_tts/train/finetune_gradio.py @@ -453,7 +453,7 @@ def start_training( cmd += f" --tokenizer {tokenizer_type} " - cmd += f" --export_samples True --logger {logger} " + cmd += f" --log_samples True --logger {logger} " print(cmd) @@ -1321,18 +1321,14 @@ def get_combined_stats(): def get_audio_select(file_sample): - select_audio_org = file_sample + select_audio_ref = 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_ref += "_ref.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 + return select_audio_ref, select_audio_gen with gr.Blocks() as app: @@ -1515,7 +1511,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") + cd_logger = gr.Radio(label="logger", choices=["wandb", "tensorboard"], value="wandb") start_button = gr.Button("Start Training") stop_button = gr.Button("Stop Training", interactive=False) @@ -1562,16 +1558,12 @@ If you encounter a memory error, try reducing the batch size per GPU to a smalle list_audios, select_audio = get_audio_project(projects_selelect, False) - select_audio_org = select_audio + select_audio_ref = 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_ref += "_ref.wav" select_audio_gen += "_gen.wav" - select_image_org += "_org.png" - select_image_gen += "_gen.png" with gr.Row(): ch_list_audio = gr.Dropdown( @@ -1587,17 +1579,13 @@ If you encounter a memory error, try reducing the batch size per GPU to a smalle 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_ref_stream = gr.Audio(label="original", type="filepath", value=select_audio_ref) 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], + outputs=[audio_ref_stream, audio_gen_stream], ) start_button.click(