Spaces:
Build error
Build error
| import json | |
| import os | |
| import sys | |
| from pathlib import Path | |
| import shutil | |
| import torchaudio | |
| from datasets import load_dataset | |
| from datasets.arrow_writer import ArrowWriter | |
| from tqdm import tqdm | |
| import soundfile as sf | |
| import csv | |
| import subprocess | |
| import argparse | |
| def save_dataset_to_local_disk(output_dir, base_model, audio_header, text_header): | |
| """ | |
| Saves a dataset to a local directory. | |
| Args: | |
| output_dir (str): The directory to save the dataset to. | |
| base_model (str): The base model to load the dataset from. | |
| audio_header (str): The header for the audio data in the dataset. | |
| text_header (str): The header for the text data in the dataset. | |
| """ | |
| wavs_dir = os.path.join(output_dir, "wavs") | |
| metadata_path = os.path.join(output_dir, "metadata.csv") | |
| os.makedirs(wavs_dir, exist_ok=True) | |
| try: | |
| ds = load_dataset(base_model)['train'] | |
| except Exception as e: | |
| print(f"Error loading dataset: {e}", file=sys.stderr) | |
| return | |
| metadata = [] | |
| for idx, sample in tqdm(enumerate(ds), total=len(ds), desc="Saving samples to directory"): | |
| try: | |
| audio_array = sample[audio_header]['array'] | |
| sampling_rate = sample[audio_header]['sampling_rate'] | |
| filename = f"audio_{idx:06d}.wav" | |
| sf.write(os.path.join(wavs_dir, filename), audio_array, sampling_rate) | |
| metadata.append([f"wavs/{filename}", sample[text_header]]) | |
| except Exception as e: | |
| print(f"Error processing sample {idx}: {e}", file=sys.stderr) | |
| continue | |
| try: | |
| with open(metadata_path, 'w', newline='', encoding='utf-8') as f: | |
| csv.writer(f, delimiter='|').writerows(metadata) | |
| print(f"Dataset saved to {output_dir}") | |
| except Exception as e: | |
| print(f"Error writing metadata: {e}", file=sys.stderr) | |
| def run_preprocess(input_dir, output_dir, workers): | |
| """ | |
| Runs the preprocessing script with real-time output. | |
| Args: | |
| input_dir (str): Input directory for preprocessing. | |
| output_dir (str): Output directory for processed data. | |
| workers (int): Number of parallel processes. | |
| """ | |
| script_path = "./src/f5_tts/train/datasets/prepare_csv_wavs.py" | |
| if not os.path.exists(script_path): | |
| print(f"Preprocessing script not found at {script_path}", file=sys.stderr) | |
| return | |
| command = [ | |
| "python", script_path, | |
| input_dir, output_dir, | |
| "--workers", str(workers) | |
| ] | |
| try: | |
| process = subprocess.Popen( | |
| command, | |
| stdout=subprocess.PIPE, | |
| stderr=subprocess.PIPE, | |
| text=True, | |
| bufsize=1, # Line buffered | |
| universal_newlines=True | |
| ) | |
| # Real-time output for stdout and stderr | |
| while True: | |
| stdout_line = process.stdout.readline() | |
| stderr_line = process.stderr.readline() | |
| if stdout_line: | |
| print(stdout_line, end='', flush=True) | |
| if stderr_line: | |
| print(stderr_line, end='', flush=True, file=sys.stderr) | |
| if process.poll() is not None: | |
| break | |
| # Capture any remaining output | |
| stdout, stderr = process.communicate() | |
| if stdout: | |
| print(stdout, end='', flush=True) | |
| if stderr: | |
| print(stderr, end='', flush=True, file=sys.stderr) | |
| if process.returncode == 0: | |
| print("\nPreprocessing completed successfully.") | |
| else: | |
| print(f"\nPreprocessing failed with return code {process.returncode}.", file=sys.stderr) | |
| except Exception as e: | |
| print(f"Error during preprocessing: {e}", file=sys.stderr) | |
| if __name__ == "__main__": | |
| # Set up argument parsing | |
| parser = argparse.ArgumentParser(description="Prepare dataset for training.") | |
| # parser.add_argument("--command", type=str, choices=["save", "preprocess"], required=True, | |
| # help="Command to execute: 'save' or 'preprocess'") | |
| parser.add_argument("--output_dir", type=str, default="./data/vin100h-preprocessed-v2", | |
| help="Output directory for save command") | |
| parser.add_argument("--base_model", type=str, default="htdung167/vin100h-preprocessed-v2", | |
| help="Base model for save command") | |
| parser.add_argument("--audio_header", type=str, default="audio", | |
| help="Audio header for save command") | |
| parser.add_argument("--text_header", type=str, default="preprocessed_sentence_v2", | |
| help="Text header for save command") | |
| parser.add_argument("--prepare_csv_input_dir", type=str, | |
| default="./data/vin100h-preprocessed-v2", | |
| help="Input directory for preprocess command") | |
| parser.add_argument("--prepare_csv_output_dir", type=str, | |
| default="./data/vin100h-preprocessed-v2_pinyin", | |
| help="Output directory for preprocess command") | |
| parser.add_argument("--workers", type=int, default=4, | |
| help="Number of parallel processes for preprocess command") | |
| args = parser.parse_args() | |
| # if args.command == "save": | |
| save_dataset_to_local_disk(args.output_dir, args.base_model, args.audio_header, args.text_header) | |
| # elif args.command == "preprocess": | |
| run_preprocess(args.prepare_csv_input_dir, args.prepare_csv_output_dir, args.workers) |