basic structure

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)
 

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

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)

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(