| 1 | import os |
| 2 | |
| 3 | |
| 4 | os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" # for MPS device compatibility |
| 5 | |
| 6 | from importlib.resources import files |
| 7 | |
| 8 | import torch |
| 9 | import torch.nn.functional as F |
| 10 | import torchaudio |
| 11 | from cached_path import cached_path |
| 12 | from hydra.utils import get_class |
| 13 | from omegaconf import OmegaConf |
| 14 | |
| 15 | from f5_tts.infer.utils_infer import load_checkpoint, load_vocoder, save_spectrogram |
| 16 | from f5_tts.model import CFM |
| 17 | from f5_tts.model.utils import convert_char_to_pinyin, get_tokenizer |
| 18 | |
| 19 | |
| 20 | device = ( |
| 21 | "cuda" |
| 22 | if torch.cuda.is_available() |
| 23 | else "xpu" |
| 24 | if torch.xpu.is_available() |
| 25 | else "mps" |
| 26 | if torch.backends.mps.is_available() |
| 27 | else "cpu" |
| 28 | ) |
| 29 | |
| 30 | |
| 31 | # ---------------------- infer setting ---------------------- # |
| 32 | |
| 33 | seed = None # int | None |
| 34 | |
| 35 | exp_name = "F5TTS_v1_Base" # F5TTS_v1_Base | E2TTS_Base |
| 36 | ckpt_step = 1250000 |
| 37 | |
| 38 | nfe_step = 32 # 16, 32 |
| 39 | cfg_strength = 2.0 |
| 40 | ode_method = "euler" # euler | midpoint |
| 41 | sway_sampling_coef = -1.0 |
| 42 | speed = 1.0 |
| 43 | target_rms = 0.1 |
| 44 | |
| 45 | |
| 46 | model_cfg = OmegaConf.load(str(files("f5_tts").joinpath(f"configs/{exp_name}.yaml"))) |
| 47 | model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}") |
| 48 | model_arc = model_cfg.model.arch |
| 49 | |
| 50 | dataset_name = model_cfg.datasets.name |
| 51 | tokenizer = model_cfg.model.tokenizer |
| 52 | |
| 53 | mel_spec_type = model_cfg.model.mel_spec.mel_spec_type |
| 54 | target_sample_rate = model_cfg.model.mel_spec.target_sample_rate |
| 55 | n_mel_channels = model_cfg.model.mel_spec.n_mel_channels |
| 56 | hop_length = model_cfg.model.mel_spec.hop_length |
| 57 | win_length = model_cfg.model.mel_spec.win_length |
| 58 | n_fft = model_cfg.model.mel_spec.n_fft |
| 59 | |
| 60 | |
| 61 | # ckpt_path = str(files("f5_tts").joinpath("../../")) + f"/ckpts/{exp_name}/model_{ckpt_step}.safetensors" |
| 62 | ckpt_path = str(cached_path(f"hf://SWivid/F5-TTS/{exp_name}/model_{ckpt_step}.safetensors")) |
| 63 | output_dir = "tests" |
| 64 | |
| 65 | |
| 66 | # [leverage https://github.com/MahmoudAshraf97/ctc-forced-aligner to get char level alignment] |
| 67 | # pip install git+https://github.com/MahmoudAshraf97/ctc-forced-aligner.git |
| 68 | # [write the origin_text into a file, e.g. tests/test_edit.txt] |
| 69 | # ctc-forced-aligner --audio_path "src/f5_tts/infer/examples/basic/basic_ref_en.wav" --text_path "tests/test_edit.txt" --language "zho" --romanize --split_size "char" |
| 70 | # [result will be saved at same path of audio file] |
| 71 | # [--language "zho" for Chinese, "eng" for English] |
| 72 | # [if local ckpt, set --alignment_model "../checkpoints/mms-300m-1130-forced-aligner"] |
| 73 | |
| 74 | audio_to_edit = str(files("f5_tts").joinpath("infer/examples/basic/basic_ref_en.wav")) |
| 75 | origin_text = "Some call me nature, others call me mother nature." |
| 76 | target_text = "Some call me optimist, others call me realist." |
| 77 | parts_to_edit = [ |
| 78 | [1.42, 2.44], |
| 79 | [4.04, 4.9], |
| 80 | ] # stard_ends of "nature" & "mother nature", in seconds |
| 81 | fix_duration = [ |
| 82 | 1.2, |
| 83 | 1, |
| 84 | ] # fix duration for "optimist" & "realist", in seconds |
| 85 | |
| 86 | # audio_to_edit = "src/f5_tts/infer/examples/basic/basic_ref_zh.wav" |
| 87 | # origin_text = "对，这就是我，万人敬仰的太乙真人。" |
| 88 | # target_text = "对，那就是你，万人敬仰的太白金星。" |
| 89 | # parts_to_edit = [[0.84, 1.4], [1.92, 2.4], [4.26, 6.26], ] |
| 90 | # fix_duration = None # use origin text duration |
| 91 | |
| 92 | # audio_to_edit = "src/f5_tts/infer/examples/basic/basic_ref_zh.wav" |
| 93 | # origin_text = "对，这就是我，万人敬仰的太乙真人。" |
| 94 | # target_text = "对，这就是你，万人敬仰的李白金星。" |
| 95 | # parts_to_edit = [[1.500, 2.784], [4.083, 6.760]] |
| 96 | # fix_duration = [1.284, 2.677] |
| 97 | |
| 98 | |
| 99 | # -------------------------------------------------# |
| 100 | |
| 101 | use_ema = True |
| 102 | |
| 103 | if not os.path.exists(output_dir): |
| 104 | os.makedirs(output_dir) |
| 105 | |
| 106 | # Vocoder model |
| 107 | local = False |
| 108 | if mel_spec_type == "vocos": |
| 109 | vocoder_local_path = "../checkpoints/charactr/vocos-mel-24khz" |
| 110 | elif mel_spec_type == "bigvgan": |
| 111 | vocoder_local_path = "../checkpoints/bigvgan_v2_24khz_100band_256x" |
| 112 | vocoder = load_vocoder(vocoder_name=mel_spec_type, is_local=local, local_path=vocoder_local_path) |
| 113 | |
| 114 | # Tokenizer |
| 115 | vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer) |
| 116 | |
| 117 | # Model |
| 118 | model = CFM( |
| 119 | transformer=model_cls(**model_arc, text_num_embeds=vocab_size, mel_dim=n_mel_channels), |
| 120 | mel_spec_kwargs=dict( |
| 121 | n_fft=n_fft, |
| 122 | hop_length=hop_length, |
| 123 | win_length=win_length, |
| 124 | n_mel_channels=n_mel_channels, |
| 125 | target_sample_rate=target_sample_rate, |
| 126 | mel_spec_type=mel_spec_type, |
| 127 | ), |
| 128 | odeint_kwargs=dict( |
| 129 | method=ode_method, |
| 130 | ), |
| 131 | vocab_char_map=vocab_char_map, |
| 132 | ).to(device) |
| 133 | |
| 134 | dtype = torch.float32 if mel_spec_type == "bigvgan" else None |
| 135 | model = load_checkpoint(model, ckpt_path, device, dtype=dtype, use_ema=use_ema) |
| 136 | |
| 137 | # Audio |
| 138 | audio, sr = torchaudio.load(audio_to_edit) |
| 139 | if audio.shape[0] > 1: |
| 140 | audio = torch.mean(audio, dim=0, keepdim=True) |
| 141 | rms = torch.sqrt(torch.mean(torch.square(audio))) |
| 142 | if rms < target_rms: |
| 143 | audio = audio * target_rms / rms |
| 144 | if sr != target_sample_rate: |
| 145 | resampler = torchaudio.transforms.Resample(sr, target_sample_rate) |
| 146 | audio = resampler(audio) |
| 147 | |
| 148 | # Convert to mel spectrogram FIRST (on clean original audio) |
| 149 | # This avoids boundary artifacts from mel windows straddling zeros and real audio |
| 150 | audio = audio.to(device) |
| 151 | with torch.inference_mode(): |
| 152 | original_mel = model.mel_spec(audio) # (batch, n_mel, n_frames) |
| 153 | original_mel = original_mel.permute(0, 2, 1) # (batch, n_frames, n_mel) |
| 154 | |
| 155 | # Build mel_cond and edit_mask at FRAME level |
| 156 | # Insert zero frames in mel domain instead of zero samples in wav domain |
| 157 | offset_frame = 0 |
| 158 | mel_cond = torch.zeros(1, 0, n_mel_channels, device=device) |
| 159 | edit_mask = torch.zeros(1, 0, dtype=torch.bool, device=device) |
| 160 | fix_dur_list = fix_duration.copy() if fix_duration is not None else None |
| 161 | |
| 162 | for part in parts_to_edit: |
| 163 | start, end = part |
| 164 | part_dur_sec = end - start if fix_dur_list is None else fix_dur_list.pop(0) |
| 165 | |
| 166 | # Convert to frames (this is the authoritative unit) |
| 167 | start_frame = round(start * target_sample_rate / hop_length) |
| 168 | end_frame = round(end * target_sample_rate / hop_length) |
| 169 | part_dur_frames = round(part_dur_sec * target_sample_rate / hop_length) |
| 170 | |
| 171 | # Number of frames for the kept (non-edited) region |
| 172 | keep_frames = start_frame - offset_frame |
| 173 | |
| 174 | # Build mel_cond: original mel frames + zero frames for edit region |
| 175 | mel_cond = torch.cat( |
| 176 | ( |
| 177 | mel_cond, |
| 178 | original_mel[:, offset_frame:start_frame, :], |
| 179 | torch.zeros(1, part_dur_frames, n_mel_channels, device=device), |
| 180 | ), |
| 181 | dim=1, |
| 182 | ) |
| 183 | edit_mask = torch.cat( |
| 184 | ( |
| 185 | edit_mask, |
| 186 | torch.ones(1, keep_frames, dtype=torch.bool, device=device), |
| 187 | torch.zeros(1, part_dur_frames, dtype=torch.bool, device=device), |
| 188 | ), |
| 189 | dim=-1, |
| 190 | ) |
| 191 | offset_frame = end_frame |
| 192 | |
| 193 | # Append remaining mel frames after last edit |
| 194 | mel_cond = torch.cat((mel_cond, original_mel[:, offset_frame:, :]), dim=1) |
| 195 | edit_mask = F.pad(edit_mask, (0, mel_cond.shape[1] - edit_mask.shape[-1]), value=True) |
| 196 | |
| 197 | # Text |
| 198 | text_list = [target_text] |
| 199 | if tokenizer == "pinyin": |
| 200 | final_text_list = convert_char_to_pinyin(text_list) |
| 201 | else: |
| 202 | final_text_list = [text_list] |
| 203 | print(f"text : {text_list}") |
| 204 | print(f"pinyin: {final_text_list}") |
| 205 | |
| 206 | # Duration - use mel_cond length (not raw audio length) |
| 207 | duration = mel_cond.shape[1] |
| 208 | |
| 209 | # Inference - pass mel_cond directly (not wav) |
| 210 | with torch.inference_mode(): |
| 211 | generated, trajectory = model.sample( |
| 212 | cond=mel_cond, # Now passing mel directly, not wav |
| 213 | text=final_text_list, |
| 214 | duration=duration, |
| 215 | steps=nfe_step, |
| 216 | cfg_strength=cfg_strength, |
| 217 | sway_sampling_coef=sway_sampling_coef, |
| 218 | seed=seed, |
| 219 | edit_mask=edit_mask, |
| 220 | ) |
| 221 | print(f"Generated mel: {generated.shape}") |
| 222 | |
| 223 | # Final result |
| 224 | generated = generated.to(torch.float32) |
| 225 | gen_mel_spec = generated.permute(0, 2, 1) |
| 226 | if mel_spec_type == "vocos": |
| 227 | generated_wave = vocoder.decode(gen_mel_spec).cpu() |
| 228 | elif mel_spec_type == "bigvgan": |
| 229 | generated_wave = vocoder(gen_mel_spec).squeeze(0).cpu() |
| 230 | |
| 231 | if rms < target_rms: |
| 232 | generated_wave = generated_wave * rms / target_rms |
| 233 | |
| 234 | save_spectrogram(gen_mel_spec[0].cpu().numpy(), f"{output_dir}/speech_edit_out.png") |
| 235 | torchaudio.save(f"{output_dir}/speech_edit_out.wav", generated_wave, target_sample_rate) |
| 236 | print(f"Generated wav: {generated_wave.shape}") |
| 237 |