# Copyright (c) 2025 Tsinghua Univ. (authors: Xingchen Song) # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ Example Usage: see README.md """ import argparse import json import os import random import sys import time from concurrent.futures import ThreadPoolExecutor from datetime import datetime import numpy as np import onnxruntime import s3tokenizer import torch import torch.distributed as dist import torchaudio import torchaudio.compliance.kaldi as kaldi from torch.utils.data import DataLoader, Dataset, DistributedSampler from tqdm import tqdm from flashcosyvoice.config import Config, CosyVoice2LLMConfig, SamplingParams from flashcosyvoice.cosyvoice2 import CosyVoice2 from flashcosyvoice.utils.audio import mel_spectrogram def set_all_random_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def save_file_async( wav, prompt_speech_tokens, generated_speech_tokens, info, timing_stats ): """Save audio asynchronously.""" try: os.makedirs(os.path.dirname(info['wav']), exist_ok=True) if wav is not None: wav = wav.cpu() torchaudio.save(info['wav'], wav, 24000) duration = wav.shape[-1] / 24000.0 rtf = ((timing_stats['dataloader_time'] + timing_stats['model_inference_time']) / timing_stats['batch_size']) / duration timing_stats['rtf'] = rtf else: duration = 0.0 info['timing_stats'] = timing_stats info['prompt_speech_tokens'] = prompt_speech_tokens info['generated_speech_tokens'] = generated_speech_tokens with open(f"{info['wav'].replace('.wav', '.json')}", "w") as f: json.dump(info, f, ensure_ascii=False, indent=4) return duration except Exception as e: timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f"[{timestamp}] - [ERROR] - Error saving audio {info.get('key', 'unknown')}: {e}") return 0.0 class AudioDataset(Dataset): def __init__(self, text_norm, text_tokenizer, data_list, model_config: Config): self.datas = [] self.text_norm = text_norm self.model_config = model_config """Example data_list: ``` {"key": "uttid_1", "prompt_text": "你好，我是小明。", "text": "你好，我是小红。", "prompt_wav": "/mnt/data/audio/00000000.wav", "wav": "/mnt/data/audio_synthetic/uttid_1.wav"} {"key": "uttid_2", "prompt_text": "你好，我是小红。", "text": "你好，我是小明。", "prompt_wav": "/mnt/data/audio/00000001.wav", "wav": "/mnt/data/audio_synthetic/uttid_2.wav"} ``` Note: - `key` is the key of this sample. - `prompt_text` is the text used for prompt. - `text` is the text used for generating real audio. - `prompt_wav` is the audio used for prompt. - `wav` is the path to the generated audio to be saved (we highly recommend to pre-define the save path before running the script). """ missing = 0 with open(data_list, 'r', encoding='utf-8') as f: lines = f.readlines() total_lines = len(lines) if torch.distributed.get_node_local_rank() == 0: iterator = tqdm(lines, desc='Loading data') else: iterator = lines for line in iterator: data = json.loads(line.strip()) valid = True for k in ['key', 'prompt_text', 'text', 'prompt_wav']: if k not in data: valid = False break if data[k] is None: valid = False break if not os.path.exists(data['prompt_wav']): valid = False if valid: self.datas.append(data) else: missing += 1 if torch.distributed.get_node_local_rank() == 0: timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f'[{timestamp}] - [INFO] - Loaded {total_lines} lines, found {missing} missing lines, total valid lines == {len(self.datas)}.') self.text_tokenizer = text_tokenizer option = onnxruntime.SessionOptions() option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL option.intra_op_num_threads = 1 self.spk_model = onnxruntime.InferenceSession(f"{self.model_config.model}/campplus.onnx", sess_options=option, providers=["CPUExecutionProvider"]) def __len__(self): return len(self.datas) def __getitem__(self, idx): data = self.datas[idx] try: # 1. feature for s3tokenizer audio = s3tokenizer.load_audio(data['prompt_wav'], sr=16000) # [T] log_mel = s3tokenizer.log_mel_spectrogram(audio) # [num_mels, T] # 2. feature for speaker embedding spk_feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000) spk_feat = spk_feat - spk_feat.mean(dim=0, keepdim=True) spk_emb = self.spk_model.run( None, {self.spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()} )[0].flatten().tolist() # 3. feature for flow audio, sample_rate = torchaudio.load(data['prompt_wav'], backend='soundfile') audio = audio.mean(dim=0, keepdim=True) # [1, T] if sample_rate != 24000: audio = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=24000)(audio) mel = mel_spectrogram(audio).transpose(1, 2).squeeze(0) # [T, num_mels] mel_len = mel.shape[0] # 4. feature for llm if self.text_norm is not None: prompt_texts = [i["text"] for i in json.loads(self.text_norm.do_voicegen_frd(data['prompt_text'].strip()))["sentences"]] prompt_text = ''.join(prompt_texts) texts = [i["text"] for i in json.loads(self.text_norm.do_voicegen_frd(data['text'].strip()))["sentences"]] text = ''.join(texts) else: prompt_text = data['prompt_text'] text = data['text'] prompt_text_ids = self.text_tokenizer.encode(prompt_text) prompt_text_ids = [i + self.model_config.hf_config.speech_vocab_size + 2 for i in prompt_text_ids] text_ids = self.text_tokenizer.encode(text) text_ids = [i + self.model_config.hf_config.speech_vocab_size + 2 for i in text_ids] item = { "prompt_text_tokens": prompt_text_ids, "text_tokens": text_ids, "spk_emb": spk_emb, "mel": mel, "mel_len": mel_len, "log_mel": log_mel, "info": data, "min_tokens": len(text_ids) * self.model_config.min_token_text_ratio, "max_tokens": len(text_ids) * self.model_config.max_token_text_ratio, } except Exception as e: timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f"[{timestamp}] - [WARNING] - Error processing data item {data.get('key', idx)}: {e}") return None return item def collate_fn(batch): prompt_mels_for_llm = [item["log_mel"] for item in batch if item is not None] prompt_mels_for_llm, prompt_mels_lens_for_llm = s3tokenizer.padding(prompt_mels_for_llm) # [B, num_mels=128, T] prompt_text_tokens_for_llm = [item["prompt_text_tokens"] for item in batch if item is not None] text_tokens_for_llm = [item["text_tokens"] for item in batch if item is not None] prompt_mels_for_flow = [item["mel"] for item in batch if item is not None] prompt_mels_for_flow = torch.nn.utils.rnn.pad_sequence(prompt_mels_for_flow, batch_first=True, padding_value=0) # [B, T', num_mels=80] prompt_mels_lens_for_flow = [item["mel_len"] for item in batch if item is not None] prompt_mels_lens_for_flow = torch.tensor(prompt_mels_lens_for_flow) spk_emb_for_flow = [item["spk_emb"] for item in batch if item is not None] spk_emb_for_flow = torch.tensor(spk_emb_for_flow) sampling_params = [SamplingParams(min_tokens=item["min_tokens"], max_tokens=item["max_tokens"], use_ras=True) for item in batch if item is not None] infos = [item["info"] for item in batch if item is not None] return { "prompt_mels_for_llm": prompt_mels_for_llm, "prompt_mels_lens_for_llm": prompt_mels_lens_for_llm, "prompt_text_tokens_for_llm": prompt_text_tokens_for_llm, "text_tokens_for_llm": text_tokens_for_llm, "prompt_mels_for_flow": prompt_mels_for_flow, "prompt_mels_lens_for_flow": prompt_mels_lens_for_flow, "spk_emb_for_flow": spk_emb_for_flow, "sampling_params": sampling_params, "infos": infos, } def init_distributed(): world_size = int(os.environ.get('WORLD_SIZE', 1)) local_rank = int(os.environ.get('LOCAL_RANK', 0)) rank = int(os.environ.get('RANK', 0)) timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f'[{timestamp}] - [INFO] - Inference on multiple gpus, this gpu {local_rank}, rank {rank}, world_size {world_size}') torch.cuda.set_device(local_rank) dist.init_process_group("nccl") return world_size, local_rank, rank def get_args(): parser = argparse.ArgumentParser(description='FlashCosyVoice') parser.add_argument('--model_path', required=True, type=str, help='model path') parser.add_argument('--data_list', required=True, type=str, help='data list') parser.add_argument('--batch_size_dataloader', required=True, type=int, help='batch size (per-device) for dataloading') parser.add_argument('--batch_size_flow', required=True, type=int, help='batch size (per-device) for flow-matching') parser.add_argument('--num_workers', type=int, default=4, help='workers for dataloader') parser.add_argument('--prefetch', type=int, default=5, help='prefetch for dataloader') parser.add_argument('--enable_tn', action='store_true', help='enable text normalization') parser.add_argument('--only_llm', action='store_true', help='only generate speech tokens from llm') parser.add_argument('--fp16_flow', action='store_true', help='enable fp16 flow') parser.add_argument('--seed', type=int, default=1986, help='random seed for generation') args = parser.parse_args() return args def main(): args = get_args() if args.enable_tn: # Check python version, if == 3.10, use ttsfrd if sys.version_info.major == 3 and sys.version_info.minor == 10: # Check if ttsfrd is installed try: import ttsfrd from cosyvoice_ttsfrd import get_resource_path except ImportError as e: raise ImportError("ttsfrd is not installed, please install it first, see `https://github.com/xingchensong/CosyVoice-ttsfrd` for installation guide.") from e text_norm = ttsfrd.TtsFrontendEngine() text_norm.initialize(get_resource_path()) text_norm.set_lang_type('pinyinvg') else: timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f"[{timestamp}] - [WARNING] - Only python 3.10 is supported for ttsfrd, see `https://github.com/xingchensong/CosyVoice-ttsfrd` for more info. Setting enable_tn to False...") # TODO: maybe we should use wetext if python version is not 3.10? args.enable_tn = False text_norm = None else: text_norm = None assert (torch.cuda.is_available()) world_size, local_rank, rank = init_distributed() config = Config(model=args.model_path, enforce_eager=True, tensor_parallel_size=1, max_num_seqs=args.batch_size_dataloader, hf_config=CosyVoice2LLMConfig(fp16_flow=args.fp16_flow), rank=local_rank) model = CosyVoice2(config) set_all_random_seed(args.seed) dataset = AudioDataset(text_norm, model.llm.tokenizer, args.data_list, config) sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank) dataloader = DataLoader(dataset, batch_size=args.batch_size_dataloader, num_workers=args.num_workers, pin_memory=True, sampler=sampler, shuffle=False, prefetch_factor=args.prefetch, collate_fn=collate_fn) total_steps = len(dataset) if local_rank == 0: timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f"[{timestamp}] - [INFO] - {args}") progress_bar = tqdm(total=total_steps, desc="Processing samples", unit="wav", position=0, leave=True, dynamic_ncols=True) cpu_counts = os.cpu_count() executor = ThreadPoolExecutor(max_workers=min(args.batch_size_dataloader, cpu_counts // 8)) pending_futures = [] dataloader_iter = iter(dataloader) succeed_duration = 0.01 # avoid division by zero start_time = time.time() estimated_total_wavs = 0 succeed_wavs = 0 failed_wavs = 0 last_print_time = start_time while True: try: dataloader_start = time.time() batch = next(dataloader_iter) dataloader_time = time.time() - dataloader_start if len(batch['infos']) == 0: timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f"[{timestamp}] - [WARNING] - rank {rank} of {world_size}: No valid batch found, skipping this batch...") continue model_start = time.time() results_dict, timing_stats = model(**batch, batch_size_flow=args.batch_size_flow, only_llm=args.only_llm) model_time = time.time() - model_start estimated_total_wavs += len(results_dict['generated_wavs']) timing_stats['dataloader_time'] = dataloader_time timing_stats['model_inference_time'] = model_time if args.only_llm: results_dict['generated_wavs'] = [None] * len(results_dict['prompt_speech_tokens']) for i in range(len(results_dict['generated_wavs'])): future = executor.submit( save_file_async, results_dict['generated_wavs'][i], results_dict['prompt_speech_tokens'][i], results_dict['generated_speech_tokens'][i], batch['infos'][i].copy(), timing_stats.copy() ) pending_futures.append(future) completed_futures = [] for future in pending_futures: if future.done(): try: duration = future.result() succeed_duration += duration succeed_wavs += 1 except Exception as e: failed_wavs += 1 timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f"[{timestamp}] - [ERROR] - rank {rank} of {world_size}: Error in async save task: {e}") completed_futures.append(future) for future in completed_futures: pending_futures.remove(future) if local_rank == 0: update_n = world_size * len(batch["prompt_text_tokens_for_llm"]) if progress_bar.n + update_n > progress_bar.total: progress_bar.update(progress_bar.total - progress_bar.n) else: progress_bar.update(update_n) current_time = time.time() if current_time - last_print_time >= 120 and not args.only_llm: elapsed_time = current_time - start_time avg_duration = succeed_duration / succeed_wavs if succeed_wavs > 0 else 0 estimated_total_duration = avg_duration * estimated_total_wavs current_rtf = elapsed_time / estimated_total_duration if estimated_total_duration > 0.01 else 0 timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f"[{timestamp}] - [INFO] - rank {rank} of {world_size}: Estimated total wavs: {estimated_total_wavs} ({estimated_total_wavs - succeed_wavs} pending to save), Succeed wavs: {succeed_wavs}, Failed wavs: {failed_wavs}, Estimated total duration: {estimated_total_duration:.2f}s ({estimated_total_duration / 3600:.2f} h), Estimated RTF: {current_rtf:.5f}, Elapsed time: {elapsed_time:.2f}s") # noqa last_print_time = current_time except StopIteration: break except Exception as e: failed_wavs += 1 timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f"[{timestamp}] - [ERROR] - rank {rank} of {world_size}: Error in main loop: {e}") continue total_time = time.time() - start_time if local_rank == 0: timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f"[{timestamp}] - [INFO] - Waiting for {len(pending_futures)} pending save tasks to complete...") for future in pending_futures: try: duration = future.result(timeout=60) succeed_duration += duration succeed_wavs += 1 except Exception as e: failed_wavs += 1 timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f"[{timestamp}] - [ERROR] - rank {rank} of {world_size}: Error in final async save task: {e}") executor.shutdown(wait=True) if local_rank == 0: timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f"[{timestamp}] - [INFO] - All async save tasks completed.") progress_bar.close() if not args.only_llm: timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S,%f')[:-3] tqdm.write(f"[{timestamp}] - [INFO] - rank {rank} of {world_size}: Final Report - Succeed wavs: {succeed_wavs}, Failed wavs: {failed_wavs}, Total duration: {succeed_duration:.2f}s ({succeed_duration / 3600:.2f} h), RTF: {total_time / succeed_duration:.5f}") # noqa dist.barrier() dist.destroy_process_group() if __name__ == "__main__": main()