| 1 | # A unified script for inference process |
| 2 | # Make adjustments inside functions, and consider both gradio and cli scripts if need to change func output format |
| 3 | import os |
| 4 | import sys |
| 5 | from concurrent.futures import ThreadPoolExecutor |
| 6 | |
| 7 | |
| 8 | os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" # for MPS device compatibility |
| 9 | sys.path.append(f"{os.path.dirname(os.path.abspath(__file__))}/../../third_party/BigVGAN/") |
| 10 | |
| 11 | import hashlib |
| 12 | import re |
| 13 | import tempfile |
| 14 | from importlib.resources import files |
| 15 | |
| 16 | import matplotlib |
| 17 | |
| 18 | |
| 19 | matplotlib.use("Agg") |
| 20 | |
| 21 | import matplotlib.pylab as plt |
| 22 | import numpy as np |
| 23 | import torch |
| 24 | import torchaudio |
| 25 | import tqdm |
| 26 | from huggingface_hub import hf_hub_download |
| 27 | from pydub import AudioSegment, silence |
| 28 | from transformers import pipeline |
| 29 | from vocos import Vocos |
| 30 | |
| 31 | from f5_tts.model import CFM |
| 32 | from f5_tts.model.utils import convert_char_to_pinyin, get_tokenizer |
| 33 | |
| 34 | |
| 35 | _ref_audio_cache = {} |
| 36 | _ref_text_cache = {} |
| 37 | |
| 38 | device = ( |
| 39 | "cuda" |
| 40 | if torch.cuda.is_available() |
| 41 | else "xpu" |
| 42 | if torch.xpu.is_available() |
| 43 | else "mps" |
| 44 | if torch.backends.mps.is_available() |
| 45 | else "cpu" |
| 46 | ) |
| 47 | |
| 48 | tempfile_kwargs = {"delete_on_close": False} if sys.version_info >= (3, 12) else {"delete": False} |
| 49 | |
| 50 | # ----------------------------------------- |
| 51 | |
| 52 | target_sample_rate = 24000 |
| 53 | n_mel_channels = 100 |
| 54 | hop_length = 256 |
| 55 | win_length = 1024 |
| 56 | n_fft = 1024 |
| 57 | mel_spec_type = "vocos" |
| 58 | target_rms = 0.1 |
| 59 | cross_fade_duration = 0.15 |
| 60 | ode_method = "euler" |
| 61 | nfe_step = 32 # 16, 32 |
| 62 | cfg_strength = 2.0 |
| 63 | sway_sampling_coef = -1.0 |
| 64 | speed = 1.0 |
| 65 | fix_duration = None |
| 66 | |
| 67 | # ----------------------------------------- |
| 68 | |
| 69 | |
| 70 | # chunk text into smaller pieces |
| 71 | |
| 72 | |
| 73 | def chunk_text(text, max_chars=135): |
| 74 | """ |
| 75 | Splits the input text into chunks, each with a maximum number of characters. |
| 76 | |
| 77 | Args: |
| 78 | text (str): The text to be split. |
| 79 | max_chars (int): The maximum number of characters per chunk. |
| 80 | |
| 81 | Returns: |
| 82 | List[str]: A list of text chunks. |
| 83 | """ |
| 84 | chunks = [] |
| 85 | current_chunk = "" |
| 86 | # Split the text into sentences based on punctuation followed by whitespace |
| 87 | sentences = re.split(r"(?<=[;:,.!?])\s+|(?<=[；：，。！？])", text) |
| 88 | |
| 89 | for sentence in sentences: |
| 90 | if not sentence: |
| 91 | continue |
| 92 | if len(current_chunk.encode("utf-8")) + len(sentence.encode("utf-8")) <= max_chars: |
| 93 | current_chunk += sentence + " " if sentence and len(sentence[-1].encode("utf-8")) == 1 else sentence |
| 94 | else: |
| 95 | if current_chunk: |
| 96 | chunks.append(current_chunk.strip()) |
| 97 | current_chunk = sentence + " " if sentence and len(sentence[-1].encode("utf-8")) == 1 else sentence |
| 98 | |
| 99 | if current_chunk: |
| 100 | chunks.append(current_chunk.strip()) |
| 101 | |
| 102 | return chunks |
| 103 | |
| 104 | |
| 105 | # load vocoder |
| 106 | def load_vocoder(vocoder_name="vocos", is_local=False, local_path="", device=device, hf_cache_dir=None): |
| 107 | if vocoder_name == "vocos": |
| 108 | # vocoder = Vocos.from_pretrained("charactr/vocos-mel-24khz").to(device) |
| 109 | if is_local: |
| 110 | print(f"Load vocos from local path {local_path}") |
| 111 | config_path = f"{local_path}/config.yaml" |
| 112 | model_path = f"{local_path}/pytorch_model.bin" |
| 113 | else: |
| 114 | print("Download Vocos from huggingface charactr/vocos-mel-24khz") |
| 115 | repo_id = "charactr/vocos-mel-24khz" |
| 116 | config_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename="config.yaml") |
| 117 | model_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename="pytorch_model.bin") |
| 118 | vocoder = Vocos.from_hparams(config_path) |
| 119 | state_dict = torch.load(model_path, map_location="cpu", weights_only=True) |
| 120 | from vocos.feature_extractors import EncodecFeatures |
| 121 | |
| 122 | if isinstance(vocoder.feature_extractor, EncodecFeatures): |
| 123 | encodec_parameters = { |
| 124 | "feature_extractor.encodec." + key: value |
| 125 | for key, value in vocoder.feature_extractor.encodec.state_dict().items() |
| 126 | } |
| 127 | state_dict.update(encodec_parameters) |
| 128 | vocoder.load_state_dict(state_dict) |
| 129 | vocoder = vocoder.eval().to(device) |
| 130 | elif vocoder_name == "bigvgan": |
| 131 | try: |
| 132 | from third_party.BigVGAN import bigvgan |
| 133 | except ImportError: |
| 134 | print("You need to follow the README to init submodule and change the BigVGAN source code.") |
| 135 | if is_local: |
| 136 | # download generator from https://huggingface.co/nvidia/bigvgan_v2_24khz_100band_256x/tree/main |
| 137 | vocoder = bigvgan.BigVGAN.from_pretrained(local_path, use_cuda_kernel=False) |
| 138 | else: |
| 139 | vocoder = bigvgan.BigVGAN.from_pretrained( |
| 140 | "nvidia/bigvgan_v2_24khz_100band_256x", use_cuda_kernel=False, cache_dir=hf_cache_dir |
| 141 | ) |
| 142 | |
| 143 | vocoder.remove_weight_norm() |
| 144 | vocoder = vocoder.eval().to(device) |
| 145 | return vocoder |
| 146 | |
| 147 | |
| 148 | # load asr pipeline |
| 149 | |
| 150 | asr_pipe = None |
| 151 | |
| 152 | |
| 153 | def initialize_asr_pipeline(device: str = device, dtype=None): |
| 154 | if dtype is None: |
| 155 | dtype = ( |
| 156 | torch.float16 |
| 157 | if "cuda" in device |
| 158 | and torch.cuda.get_device_properties(device).major >= 7 |
| 159 | and not torch.cuda.get_device_name().endswith("[ZLUDA]") |
| 160 | else torch.float32 |
| 161 | ) |
| 162 | global asr_pipe |
| 163 | asr_pipe = pipeline( |
| 164 | "automatic-speech-recognition", |
| 165 | model="openai/whisper-large-v3-turbo", |
| 166 | torch_dtype=dtype, |
| 167 | device=device, |
| 168 | ) |
| 169 | |
| 170 | |
| 171 | # transcribe |
| 172 | |
| 173 | |
| 174 | def transcribe(ref_audio, language=None): |
| 175 | global asr_pipe |
| 176 | if asr_pipe is None: |
| 177 | initialize_asr_pipeline(device=device) |
| 178 | return asr_pipe( |
| 179 | ref_audio, |
| 180 | chunk_length_s=30, |
| 181 | batch_size=128, |
| 182 | generate_kwargs={"task": "transcribe", "language": language} if language else {"task": "transcribe"}, |
| 183 | return_timestamps=False, |
| 184 | )["text"].strip() |
| 185 | |
| 186 | |
| 187 | # load model checkpoint for inference |
| 188 | |
| 189 | |
| 190 | def load_checkpoint(model, ckpt_path, device: str, dtype=None, use_ema=True): |
| 191 | if dtype is None: |
| 192 | dtype = ( |
| 193 | torch.float16 |
| 194 | if "cuda" in device |
| 195 | and torch.cuda.get_device_properties(device).major >= 7 |
| 196 | and not torch.cuda.get_device_name().endswith("[ZLUDA]") |
| 197 | else torch.float32 |
| 198 | ) |
| 199 | model = model.to(dtype) |
| 200 | |
| 201 | ckpt_type = ckpt_path.split(".")[-1] |
| 202 | if ckpt_type == "safetensors": |
| 203 | from safetensors.torch import load_file |
| 204 | |
| 205 | checkpoint = load_file(ckpt_path, device=device) |
| 206 | else: |
| 207 | checkpoint = torch.load(ckpt_path, map_location=device, weights_only=True) |
| 208 | |
| 209 | if use_ema: |
| 210 | if ckpt_type == "safetensors": |
| 211 | checkpoint = {"ema_model_state_dict": checkpoint} |
| 212 | checkpoint["model_state_dict"] = { |
| 213 | k.replace("ema_model.", ""): v |
| 214 | for k, v in checkpoint["ema_model_state_dict"].items() |
| 215 | if k not in ["initted", "step"] |
| 216 | } |
| 217 | |
| 218 | # patch for backward compatibility, 305e3ea |
| 219 | for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]: |
| 220 | if key in checkpoint["model_state_dict"]: |
| 221 | del checkpoint["model_state_dict"][key] |
| 222 | |
| 223 | model.load_state_dict(checkpoint["model_state_dict"]) |
| 224 | else: |
| 225 | if ckpt_type == "safetensors": |
| 226 | checkpoint = {"model_state_dict": checkpoint} |
| 227 | model.load_state_dict(checkpoint["model_state_dict"]) |
| 228 | |
| 229 | del checkpoint |
| 230 | torch.cuda.empty_cache() |
| 231 | |
| 232 | return model.to(device) |
| 233 | |
| 234 | |
| 235 | # load model for inference |
| 236 | |
| 237 | |
| 238 | def load_model( |
| 239 | model_cls, |
| 240 | model_cfg, |
| 241 | ckpt_path, |
| 242 | mel_spec_type=mel_spec_type, |
| 243 | vocab_file="", |
| 244 | ode_method=ode_method, |
| 245 | use_ema=True, |
| 246 | device=device, |
| 247 | ): |
| 248 | if vocab_file == "": |
| 249 | vocab_file = str(files("f5_tts").joinpath("infer/examples/vocab.txt")) |
| 250 | tokenizer = "custom" |
| 251 | |
| 252 | print("\nvocab : ", vocab_file) |
| 253 | print("token : ", tokenizer) |
| 254 | print("model : ", ckpt_path, "\n") |
| 255 | |
| 256 | vocab_char_map, vocab_size = get_tokenizer(vocab_file, tokenizer) |
| 257 | model = CFM( |
| 258 | transformer=model_cls(**model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels), |
| 259 | mel_spec_kwargs=dict( |
| 260 | n_fft=n_fft, |
| 261 | hop_length=hop_length, |
| 262 | win_length=win_length, |
| 263 | n_mel_channels=n_mel_channels, |
| 264 | target_sample_rate=target_sample_rate, |
| 265 | mel_spec_type=mel_spec_type, |
| 266 | ), |
| 267 | odeint_kwargs=dict( |
| 268 | method=ode_method, |
| 269 | ), |
| 270 | vocab_char_map=vocab_char_map, |
| 271 | ).to(device) |
| 272 | |
| 273 | dtype = torch.float32 if mel_spec_type == "bigvgan" else None |
| 274 | model = load_checkpoint(model, ckpt_path, device, dtype=dtype, use_ema=use_ema) |
| 275 | |
| 276 | return model |
| 277 | |
| 278 | |
| 279 | def remove_silence_edges(audio, silence_threshold=-42): |
| 280 | # Remove silence from the start |
| 281 | non_silent_start_idx = silence.detect_leading_silence(audio, silence_threshold=silence_threshold) |
| 282 | audio = audio[non_silent_start_idx:] |
| 283 | |
| 284 | # Remove silence from the end |
| 285 | reversed_audio = audio.reverse() |
| 286 | non_silent_end_idx = silence.detect_leading_silence(reversed_audio, silence_threshold=silence_threshold) |
| 287 | if non_silent_end_idx > 0: |
| 288 | trimmed_audio = audio[: len(audio) - non_silent_end_idx] |
| 289 | else: |
| 290 | trimmed_audio = audio |
| 291 | |
| 292 | return trimmed_audio |
| 293 | |
| 294 | |
| 295 | # preprocess reference audio and text |
| 296 | |
| 297 | |
| 298 | def preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=print): |
| 299 | show_info("Converting audio...") |
| 300 | |
| 301 | # Compute a hash of the reference audio file |
| 302 | with open(ref_audio_orig, "rb") as audio_file: |
| 303 | audio_data = audio_file.read() |
| 304 | audio_hash = hashlib.md5(audio_data).hexdigest() |
| 305 | |
| 306 | global _ref_audio_cache |
| 307 | |
| 308 | if audio_hash in _ref_audio_cache: |
| 309 | show_info("Using cached preprocessed reference audio...") |
| 310 | ref_audio = _ref_audio_cache[audio_hash] |
| 311 | |
| 312 | else: # first pass, do preprocess |
| 313 | with tempfile.NamedTemporaryFile(suffix=".wav", **tempfile_kwargs) as f: |
| 314 | temp_path = f.name |
| 315 | |
| 316 | aseg = AudioSegment.from_file(ref_audio_orig) |
| 317 | |
| 318 | # 1. try to find long silence for clipping |
| 319 | non_silent_segs = silence.split_on_silence( |
| 320 | aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=1000, seek_step=10 |
| 321 | ) |
| 322 | non_silent_wave = AudioSegment.silent(duration=0) |
| 323 | for non_silent_seg in non_silent_segs: |
| 324 | if len(non_silent_wave) > 6000 and len(non_silent_wave + non_silent_seg) > 12000: |
| 325 | show_info("Audio is over 12s, clipping short. (1)") |
| 326 | break |
| 327 | non_silent_wave += non_silent_seg |
| 328 | |
| 329 | # 2. try to find short silence for clipping if 1. failed |
| 330 | if len(non_silent_wave) > 12000: |
| 331 | non_silent_segs = silence.split_on_silence( |
| 332 | aseg, min_silence_len=100, silence_thresh=-40, keep_silence=1000, seek_step=10 |
| 333 | ) |
| 334 | non_silent_wave = AudioSegment.silent(duration=0) |
| 335 | for non_silent_seg in non_silent_segs: |
| 336 | if len(non_silent_wave) > 6000 and len(non_silent_wave + non_silent_seg) > 12000: |
| 337 | show_info("Audio is over 12s, clipping short. (2)") |
| 338 | break |
| 339 | non_silent_wave += non_silent_seg |
| 340 | |
| 341 | aseg = non_silent_wave |
| 342 | |
| 343 | # 3. if no proper silence found for clipping |
| 344 | if len(aseg) > 12000: |
| 345 | aseg = aseg[:12000] |
| 346 | show_info("Audio is over 12s, clipping short. (3)") |
| 347 | |
| 348 | aseg = remove_silence_edges(aseg) + AudioSegment.silent(duration=50) |
| 349 | aseg.export(temp_path, format="wav") |
| 350 | ref_audio = temp_path |
| 351 | |
| 352 | # Cache the processed reference audio |
| 353 | _ref_audio_cache[audio_hash] = ref_audio |
| 354 | |
| 355 | if not ref_text.strip(): |
| 356 | global _ref_text_cache |
| 357 | if audio_hash in _ref_text_cache: |
| 358 | # Use cached asr transcription |
| 359 | show_info("Using cached reference text...") |
| 360 | ref_text = _ref_text_cache[audio_hash] |
| 361 | else: |
| 362 | show_info("No reference text provided, transcribing reference audio...") |
| 363 | ref_text = transcribe(ref_audio) |
| 364 | # Cache the transcribed text (not caching custom ref_text, enabling users to do manual tweak) |
| 365 | _ref_text_cache[audio_hash] = ref_text |
| 366 | else: |
| 367 | show_info("Using custom reference text...") |
| 368 | |
| 369 | # Ensure ref_text ends with a proper sentence-ending punctuation |
| 370 | if not ref_text.endswith(". ") and not ref_text.endswith("。"): |
| 371 | if ref_text.endswith("."): |
| 372 | ref_text += " " |
| 373 | else: |
| 374 | ref_text += ". " |
| 375 | |
| 376 | print("\nref_text ", ref_text) |
| 377 | |
| 378 | return ref_audio, ref_text |
| 379 | |
| 380 | |
| 381 | # infer process: chunk text -> infer batches [i.e. infer_batch_process()] |
| 382 | |
| 383 | |
| 384 | def infer_process( |
| 385 | ref_audio, |
| 386 | ref_text, |
| 387 | gen_text, |
| 388 | model_obj, |
| 389 | vocoder, |
| 390 | mel_spec_type=mel_spec_type, |
| 391 | show_info=print, |
| 392 | progress=tqdm, |
| 393 | target_rms=target_rms, |
| 394 | cross_fade_duration=cross_fade_duration, |
| 395 | nfe_step=nfe_step, |
| 396 | cfg_strength=cfg_strength, |
| 397 | sway_sampling_coef=sway_sampling_coef, |
| 398 | speed=speed, |
| 399 | fix_duration=fix_duration, |
| 400 | device=device, |
| 401 | ): |
| 402 | # Split the input text into batches |
| 403 | audio, sr = torchaudio.load(ref_audio) |
| 404 | max_chars = int(len(ref_text.encode("utf-8")) / (audio.shape[-1] / sr) * (22 - audio.shape[-1] / sr) * speed) |
| 405 | gen_text_batches = chunk_text(gen_text, max_chars=max_chars) |
| 406 | for i, gen_text_i in enumerate(gen_text_batches): |
| 407 | print(f"gen_text {i}", gen_text_i) |
| 408 | print("\n") |
| 409 | |
| 410 | show_info(f"Generating audio in {len(gen_text_batches)} batches...") |
| 411 | |
| 412 | if not gen_text_batches: |
| 413 | show_info("No text batches to generate.") |
| 414 | return None, target_sample_rate, None |
| 415 | |
| 416 | return next( |
| 417 | infer_batch_process( |
| 418 | (audio, sr), |
| 419 | ref_text, |
| 420 | gen_text_batches, |
| 421 | model_obj, |
| 422 | vocoder, |
| 423 | mel_spec_type=mel_spec_type, |
| 424 | progress=progress, |
| 425 | target_rms=target_rms, |
| 426 | cross_fade_duration=cross_fade_duration, |
| 427 | nfe_step=nfe_step, |
| 428 | cfg_strength=cfg_strength, |
| 429 | sway_sampling_coef=sway_sampling_coef, |
| 430 | speed=speed, |
| 431 | fix_duration=fix_duration, |
| 432 | device=device, |
| 433 | ) |
| 434 | ) |
| 435 | |
| 436 | |
| 437 | # infer batches |
| 438 | |
| 439 | |
| 440 | def infer_batch_process( |
| 441 | ref_audio, |
| 442 | ref_text, |
| 443 | gen_text_batches, |
| 444 | model_obj, |
| 445 | vocoder, |
| 446 | mel_spec_type="vocos", |
| 447 | progress=tqdm, |
| 448 | target_rms=0.1, |
| 449 | cross_fade_duration=0.15, |
| 450 | nfe_step=32, |
| 451 | cfg_strength=2.0, |
| 452 | sway_sampling_coef=-1, |
| 453 | speed=1, |
| 454 | fix_duration=None, |
| 455 | device=None, |
| 456 | streaming=False, |
| 457 | chunk_size=2048, |
| 458 | ): |
| 459 | audio, sr = ref_audio |
| 460 | if audio.shape[0] > 1: |
| 461 | audio = torch.mean(audio, dim=0, keepdim=True) |
| 462 | |
| 463 | rms = torch.sqrt(torch.mean(torch.square(audio))) |
| 464 | if rms < target_rms: |
| 465 | audio = audio * target_rms / rms |
| 466 | if sr != target_sample_rate: |
| 467 | resampler = torchaudio.transforms.Resample(sr, target_sample_rate) |
| 468 | audio = resampler(audio) |
| 469 | audio = audio.to(device) |
| 470 | |
| 471 | generated_waves = [] |
| 472 | spectrograms = [] |
| 473 | |
| 474 | if len(ref_text[-1].encode("utf-8")) == 1: |
| 475 | ref_text = ref_text + " " |
| 476 | |
| 477 | # Allocate fix_duration across chunks to avoid each chunk using the full-sentence duration (N× blow-up) |
| 478 | if fix_duration is not None and len(gen_text_batches) > 1: |
| 479 | ref_audio_len_frames = audio.shape[-1] // hop_length |
| 480 | ref_sec = ref_audio_len_frames * hop_length / target_sample_rate |
| 481 | target_total = fix_duration - ref_sec |
| 482 | weights = [len(c.encode("utf-8")) for c in gen_text_batches] |
| 483 | total_w = sum(weights) |
| 484 | allocatable = target_total + cross_fade_duration * (len(gen_text_batches) - 1) |
| 485 | fix_durations = [ref_sec + allocatable * w / total_w for w in weights] |
| 486 | else: |
| 487 | fix_durations = [fix_duration] * len(gen_text_batches) |
| 488 | |
| 489 | def _infer_basic(gen_text, fix_dur): |
| 490 | local_speed = speed |
| 491 | if len(gen_text.encode("utf-8")) < 10: |
| 492 | local_speed = 0.3 |
| 493 | |
| 494 | # Prepare the text |
| 495 | text_list = [ref_text + gen_text] |
| 496 | final_text_list = convert_char_to_pinyin(text_list) |
| 497 | |
| 498 | ref_audio_len = audio.shape[-1] // hop_length |
| 499 | if fix_dur is not None: |
| 500 | duration = int(fix_dur * target_sample_rate / hop_length) |
| 501 | else: |
| 502 | # Calculate duration |
| 503 | ref_text_len = len(ref_text.encode("utf-8")) |
| 504 | gen_text_len = len(gen_text.encode("utf-8")) |
| 505 | duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / local_speed) |
| 506 | |
| 507 | # inference |
| 508 | with torch.inference_mode(): |
| 509 | generated, _ = model_obj.sample( |
| 510 | cond=audio, |
| 511 | text=final_text_list, |
| 512 | duration=duration, |
| 513 | steps=nfe_step, |
| 514 | cfg_strength=cfg_strength, |
| 515 | sway_sampling_coef=sway_sampling_coef, |
| 516 | ) |
| 517 | del _ |
| 518 | |
| 519 | generated = generated.to(torch.float32) # generated mel spectrogram |
| 520 | generated = generated[:, ref_audio_len:, :] |
| 521 | generated = generated.permute(0, 2, 1) |
| 522 | if mel_spec_type == "vocos": |
| 523 | generated_wave = vocoder.decode(generated) |
| 524 | elif mel_spec_type == "bigvgan": |
| 525 | generated_wave = vocoder(generated) |
| 526 | if rms < target_rms: |
| 527 | generated_wave = generated_wave * rms / target_rms |
| 528 | |
| 529 | # wav -> numpy |
| 530 | generated_wave = generated_wave.squeeze().cpu().numpy() |
| 531 | |
| 532 | return generated_wave, generated |
| 533 | |
| 534 | def infer_single_process(gen_text, fix_dur): |
| 535 | generated_wave, generated = _infer_basic(gen_text, fix_dur) |
| 536 | generated_cpu = generated[0].cpu().numpy() |
| 537 | del generated |
| 538 | return generated_wave, generated_cpu |
| 539 | |
| 540 | def infer_single_process_streaming(gen_text, fix_dur): |
| 541 | # for src/f5_tts/socket_server.py |
| 542 | generated_wave, generated = _infer_basic(gen_text, fix_dur) |
| 543 | del generated |
| 544 | for j in range(0, len(generated_wave), chunk_size): |
| 545 | yield generated_wave[j : j + chunk_size], target_sample_rate |
| 546 | |
| 547 | if streaming: |
| 548 | batches_iter = progress.tqdm(gen_text_batches) if progress is not None else gen_text_batches |
| 549 | for gen_text, fix_dur in zip(batches_iter, fix_durations): |
| 550 | for chunk in infer_single_process_streaming(gen_text, fix_dur): |
| 551 | yield chunk |
| 552 | else: |
| 553 | with ThreadPoolExecutor() as executor: |
| 554 | futures = [executor.submit(infer_single_process, gt, fd) for gt, fd in zip(gen_text_batches, fix_durations)] |
| 555 | for future in progress.tqdm(futures) if progress is not None else futures: |
| 556 | result = future.result() |
| 557 | if result: |
| 558 | generated_wave, generated_mel_spec = result |
| 559 | generated_waves.append(generated_wave) |
| 560 | spectrograms.append(generated_mel_spec) |
| 561 | |
| 562 | if generated_waves: |
| 563 | if cross_fade_duration <= 0: |
| 564 | # Simply concatenate |
| 565 | final_wave = np.concatenate(generated_waves) |
| 566 | else: |
| 567 | # Combine all generated waves with cross-fading |
| 568 | final_wave = generated_waves[0] |
| 569 | for i in range(1, len(generated_waves)): |
| 570 | prev_wave = final_wave |
| 571 | next_wave = generated_waves[i] |
| 572 | |
| 573 | # Calculate cross-fade samples, ensuring it does not exceed wave lengths |
| 574 | cross_fade_samples = int(cross_fade_duration * target_sample_rate) |
| 575 | cross_fade_samples = min(cross_fade_samples, len(prev_wave), len(next_wave)) |
| 576 | |
| 577 | if cross_fade_samples <= 0: |
| 578 | # No overlap possible, concatenate |
| 579 | final_wave = np.concatenate([prev_wave, next_wave]) |
| 580 | continue |
| 581 | |
| 582 | # Overlapping parts |
| 583 | prev_overlap = prev_wave[-cross_fade_samples:] |
| 584 | next_overlap = next_wave[:cross_fade_samples] |
| 585 | |
| 586 | # Fade out and fade in |
| 587 | fade_out = np.linspace(1, 0, cross_fade_samples) |
| 588 | fade_in = np.linspace(0, 1, cross_fade_samples) |
| 589 | |
| 590 | # Cross-faded overlap |
| 591 | cross_faded_overlap = prev_overlap * fade_out + next_overlap * fade_in |
| 592 | |
| 593 | # Combine |
| 594 | new_wave = np.concatenate( |
| 595 | [prev_wave[:-cross_fade_samples], cross_faded_overlap, next_wave[cross_fade_samples:]] |
| 596 | ) |
| 597 | |
| 598 | final_wave = new_wave |
| 599 | |
| 600 | # Create a combined spectrogram |
| 601 | combined_spectrogram = np.concatenate(spectrograms, axis=1) |
| 602 | |
| 603 | yield final_wave, target_sample_rate, combined_spectrogram |
| 604 | |
| 605 | else: |
| 606 | yield None, target_sample_rate, None |
| 607 | |
| 608 | |
| 609 | # remove silence from generated wav |
| 610 | |
| 611 | |
| 612 | def remove_silence_for_generated_wav(filename): |
| 613 | aseg = AudioSegment.from_file(filename) |
| 614 | non_silent_segs = silence.split_on_silence( |
| 615 | aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500, seek_step=10 |
| 616 | ) |
| 617 | non_silent_wave = AudioSegment.silent(duration=0) |
| 618 | for non_silent_seg in non_silent_segs: |
| 619 | non_silent_wave += non_silent_seg |
| 620 | aseg = non_silent_wave |
| 621 | aseg.export(filename, format="wav") |
| 622 | |
| 623 | |
| 624 | # save spectrogram |
| 625 | |
| 626 | |
| 627 | def save_spectrogram(spectrogram, path): |
| 628 | plt.figure(figsize=(12, 4)) |
| 629 | plt.imshow(spectrogram, origin="lower", aspect="auto") |
| 630 | plt.colorbar() |
| 631 | plt.savefig(path) |
| 632 | plt.close() |
| 633 |