Sync from GitHub repo

This Space is synced from the GitHub repo: https://github.com/SWivid/F5-TTS. Please submit contributions to the Space there

app.py

@@ -37,7 +37,7 @@ from f5_tts.infer.utils_infer import (
     save_spectrogram,
 )
 
-vocos = load_vocoder()
+vocoder = load_vocoder()
 
 
 # load models
@@ -94,6 +94,7 @@ def infer(
         ref_text,
         gen_text,
         ema_model,
+        vocoder,
         cross_fade_duration=cross_fade_duration,
         speed=speed,
         show_info=show_info,

src/f5_tts/api.py

@@ -47,7 +47,7 @@ class F5TTS:
         self.load_ema_model(model_type, ckpt_file, vocab_file, ode_method, use_ema)
 
     def load_vocoder_model(self, local_path):
-        self.vocos = load_vocoder(local_path is not None, local_path, self.device)
+        self.vocoder = load_vocoder(local_path is not None, local_path, self.device)
 
     def load_ema_model(self, model_type, ckpt_file, vocab_file, ode_method, use_ema):
         if model_type == "F5-TTS":
@@ -102,6 +102,7 @@ class F5TTS:
             ref_text,
             gen_text,
             self.ema_model,
+            self.vocoder,
             show_info=show_info,
             progress=progress,
             target_rms=target_rms,

src/f5_tts/infer/infer_cli.py

@@ -113,7 +113,7 @@ wave_path = Path(output_dir) / "infer_cli_out.wav"
 # spectrogram_path = Path(output_dir) / "infer_cli_out.png"
 vocos_local_path = "../checkpoints/charactr/vocos-mel-24khz"
 
-vocos = load_vocoder(is_local=args.load_vocoder_from_local, local_path=vocos_local_path)
+vocoder = load_vocoder(is_local=args.load_vocoder_from_local, local_path=vocos_local_path)
 
 
 # load models
@@ -175,7 +175,9 @@ def main_process(ref_audio, ref_text, text_gen, model_obj, remove_silence, speed
         ref_audio = voices[voice]["ref_audio"]
         ref_text = voices[voice]["ref_text"]
         print(f"Voice: {voice}")
-        audio, final_sample_rate, spectragram = infer_process(ref_audio, ref_text, gen_text, model_obj, speed=speed)
+        audio, final_sample_rate, spectragram = infer_process(
+            ref_audio, ref_text, gen_text, model_obj, vocoder, speed=speed
+        )
         generated_audio_segments.append(audio)
 
     if generated_audio_segments:

src/f5_tts/infer/utils_infer.py

@@ -29,9 +29,6 @@ _ref_audio_cache = {}
 
 device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 
-vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
-
-
 # -----------------------------------------
 
 target_sample_rate = 24000
@@ -263,6 +260,7 @@ def infer_process(
     ref_text,
     gen_text,
     model_obj,
+    vocoder,
     show_info=print,
     progress=tqdm,
     target_rms=target_rms,
@@ -287,6 +285,7 @@ def infer_process(
         ref_text,
         gen_text_batches,
         model_obj,
+        vocoder,
         progress=progress,
         target_rms=target_rms,
         cross_fade_duration=cross_fade_duration,
@@ -307,6 +306,7 @@ def infer_batch_process(
     ref_text,
     gen_text_batches,
     model_obj,
+    vocoder,
     progress=tqdm,
     target_rms=0.1,
     cross_fade_duration=0.15,
@@ -362,7 +362,7 @@ def infer_batch_process(
         generated = generated.to(torch.float32)
         generated = generated[:, ref_audio_len:, :]
         generated_mel_spec = generated.permute(0, 2, 1)
-        generated_wave = vocos.decode(generated_mel_spec.cpu())
+        generated_wave = vocoder.decode(generated_mel_spec.cpu())
         if rms < target_rms:
             generated_wave = generated_wave * rms / target_rms
 

src/f5_tts/model/trainer.py

@@ -6,6 +6,7 @@ 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
@@ -39,9 +40,11 @@ class Trainer:
         max_grad_norm=1.0,
         noise_scheduler: str | None = None,
         duration_predictor: torch.nn.Module | None = None,
+        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(),
@@ -49,21 +52,25 @@ class Trainer:
     ):
         ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
 
-        logger = "wandb" if wandb.api.api_key else None
+        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,
+            log_with=logger if logger == "wandb" else None,
             kwargs_handlers=[ddp_kwargs],
             gradient_accumulation_steps=grad_accumulation_steps,
             **accelerate_kwargs,
         )
 
-        if logger == "wandb":
+        self.logger = logger
+        if self.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,
@@ -81,11 +88,15 @@ class Trainer:
                 },
             )
 
+        elif self.logger == "tensorboard":
+            from torch.utils.tensorboard import SummaryWriter
+
+            self.writer = SummaryWriter(log_dir=f"runs/{wandb_run_name}")
+
         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
@@ -176,6 +187,14 @@ class Trainer:
         return step
 
     def train(self, train_dataset: Dataset, num_workers=16, resumable_with_seed: int = None):
+        if self.log_samples:
+            from f5_tts.infer.utils_infer import load_vocoder, nfe_step, cfg_strength, sway_sampling_coef
+
+            vocoder = load_vocoder()
+            target_sample_rate = self.accelerator.unwrap_model(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()
             generator.manual_seed(resumable_with_seed)
@@ -286,12 +305,31 @@ class Trainer:
 
                 if self.accelerator.is_local_main_process:
                     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 and self.accelerator.is_local_main_process:
+                        ref_audio, ref_audio_len = vocoder.decode(batch["mel"][0].unsqueeze(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.accelerator.unwrap_model(self.model).sample(
+                                cond=mel_spec[0][:ref_audio_len].unsqueeze(0),
+                                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 = vocoder.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/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(
+        "--log_samples",
+        type=bool,
+        default=False,
+        help="Log inferenced samples per ckpt save steps",
+    )
+    parser.add_argument("--logger", type=str, default=None, choices=["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
@@ -132,9 +141,11 @@ 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,
     )
 

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,6 +453,8 @@ def start_training(
 
     cmd += f" --tokenizer {tokenizer_type} "
 
+    cmd += f" --log_samples True --logger {logger} "
+
     print(cmd)
 
     save_settings(
@@ -467,6 +475,7 @@ def start_training(
         tokenizer_type,
         tokenizer_file,
         mixed_precision,
+        logger,
     )
 
     try:
@@ -1223,6 +1232,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 +1320,17 @@ def get_combined_stats():
     return combined_stats
 
 
+def get_audio_select(file_sample):
+    select_audio_ref = file_sample
+    select_audio_gen = file_sample
+
+    if file_sample is not None:
+        select_audio_ref += "_ref.wav"
+        select_audio_gen += "_gen.wav"
+
+    return select_audio_ref, select_audio_gen
+
+
 with gr.Blocks() as app:
     gr.Markdown(
         """
@@ -1470,6 +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=["wandb", "tensorboard"], value="wandb")
                 start_button = gr.Button("Start Training")
                 stop_button = gr.Button("Stop Training", interactive=False)
 
@@ -1491,6 +1533,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
@@ -1508,9 +1551,43 @@ 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="")
+
+            list_audios, select_audio = get_audio_project(projects_selelect, False)
+
+            select_audio_ref = select_audio
+            select_audio_gen = select_audio
+
+            if select_audio is not None:
+                select_audio_ref += "_ref.wav"
+                select_audio_gen += "_gen.wav"
+
+            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_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)
+
+            ch_list_audio.change(
+                fn=get_audio_select,
+                inputs=[ch_list_audio],
+                outputs=[audio_ref_stream, audio_gen_stream],
+            )
+
             start_button.click(
                 fn=start_training,
                 inputs=[
@@ -1532,6 +1609,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],
             )
@@ -1583,6 +1661,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

src/f5_tts/train/train.py

@@ -83,6 +83,7 @@ def main():
         wandb_run_name=exp_name,
         wandb_resume_id=wandb_resume_id,
         last_per_steps=last_per_steps,
+        log_samples=True,
     )
 
     train_dataset = load_dataset(dataset_name, tokenizer, mel_spec_kwargs=mel_spec_kwargs)