ProDiff

repository·main·Indexed 19 days ago

https://github.com/rongjiehuang/prodiff

A PyTorch implementation of a progressive fast diffusion model for high-fidelity, efficient text-to-speech synthesis. ProDiff features a speed-quality trade-off mechanism for industrial deployment and supports a high-speed synthesis pipeline when combined with the FastDiff neural vocoder. The repository includes tools for data preprocessing, MFA alignment via train_mfa_align.py, and data binarization using binarize.py.

Tokens
1.7K
Snippets
5
Records
8
Agent score
16%

What's inside ProDiff

  1. Prepare data for training your own model

    main

    To train ProDiff on a custom dataset, follow these steps:

    1. Configure Paths: In your config file, set raw_data_dir, processed_data_dir, and binary_data_dir.
    2. Organize Raw Data: Download your dataset to raw_data_dir. The structure must follow the pattern expected by egs/datasets/audio/*/pre_align.py (or you must adapt pre_align.py for your specific format).
    3. Preprocessing Pipeline:
      • Unify structure: Run pre_align.py.
      • Alignment: Run MFA (Montreal Forced Aligner) via train_mfa_align.py.
      • Binarization: Convert data to binary format for fast I/O using binarize.py.
    # 1. Preprocess step: unify the file structure
    python data_gen/tts/bin/pre_align.py --config $path/to/config
    
    # 2. Align step: MFA alignment
    python data_gen/tts/runs/train_mfa_align.py --config $CONFIG_NAME
    
    # 3. Binarization step: Binarize data for fast IO
    CUDA_VISIBLE_DEVICES=$GPU python data_gen/tts/bin/binarize.py --config $path/to/config
  2. Download ProDiff pretrained models

    main

    To use ProDiff with pretrained weights, download the checkpoints from Hugging Face using snapshot_download and move them to the local checkpoints/ directory.

    Ensure the directory structure follows checkpoints/$Model/model_ckpt_steps_*.ckpt.

    Available models:

    • ProDiff Teacher: Uses modules/ProDiff/config/prodiff_teacher.yaml
    • ProDiff: Uses modules/ProDiff/config/prodiff.yaml
    from huggingface_hub import snapshot_download 
    downloaded_path = snapshot_download(repo_id="Rongjiehuang/ProDiff")
    # Move downloaded checkpoints to the local checkpoints folder
    # Replace ${downloaded_path} with the path returned by the python script
    # Note: The destination folder should be 'checkpoints/'
    mv ${downloaded_path}/checkpoints/  checkpoints/
  3. Train and run inference for ProDiff models

    main

    Use the tasks/run.py script to manage training and inference tasks. You must specify the --config, --exp_name, and use the --infer flag for inference.

    ProDiff Training & Inference

    • Train: python tasks/run.py --config modules/ProDiff/config/prodiff.yaml --exp_name ProDiff --reset
    • Infer: python tasks/run.py --config modules/ProDiff/config/prodiff.yaml --exp_name ProDiff --infer

    ProDiff Teacher Training & Inference

    • Train: python tasks/run.py --config modules/ProDiff/config/prodiff_teacher.yaml --exp_name ProDiff_Teacher --reset
    • Infer: python tasks/run.py --config modules/ProDiff/config/prodiff_teacher.yaml --exp_name ProDiff_Teacher --infer
  4. Perform fast text-to-speech inference

    main

    ProDiff provides a high-speed synthesis pipeline combining an acoustic model (ProDiff) and a neural vocoder (FastDiff).

    Setup

    1. Acoustic Model: Place LJSpeech checkpoints in checkpoints/ProDiff or checkpoints/ProDiff_Teacher.
    2. Vocoder: Place FastDiff LJSpeech checkpoints in checkpoints/FastDiff.

    Inference Options

    You can trade off speed and quality by adjusting the number of reverse sampling iterations N via the --hparams flag.

    Option 1: Extreme Speed (2-iter ProDiff + 4-iter FastDiff)

    CUDA_VISIBLE_DEVICES=$GPU python inference/ProDiff.py --config modules/ProDiff/config/prodiff.yaml --exp_name ProDiff --hparams="N=4,text='$txt'" --reset

    Option 2: Better Quality (4-iter ProDiff Teacher + 6-iter FastDiff)

    CUDA_VISIBLE_DEVICES=$GPU python inference/ProDiff_teacher.py --config modules/ProDiff/config/prodiff_teacher.yaml --exp_name ProDiff_Teacher --hparams="N=6,text='$txt'" --reset

    Generated .wav files are saved in the infer_out directory by default.

    # Example for extreme speed inference
    CUDA_VISIBLE_DEVICES=0 python inference/ProDiff.py --config modules/ProDiff/config/prodiff.yaml --exp_name ProDiff --hparams="N=4,text='hello world'" --reset
  5. Configure Multi-GPU training

    main
    By default, the implementation uses all available GPUs returned by torch.cuda.device_count(). To restrict training to specific GPUs, set the CUDA_DEVICES_AVAILABLE environment variable before running the training module.
  6. Binarize TTS data using the binarize.py CLI

    main

    The binarize.py script is used to convert raw TTS datasets into a binarized format suitable for training.

    It dynamically selects a binarizer class based on the binarizer_cls hyperparameter defined in your configuration. By default, it uses data_gen.tts.base_binarizer.BaseBinarizer. To use a custom binarizer, ensure your configuration file sets the binarizer_cls key to the full import path of your custom class.

    Note: The script sets OMP_NUM_THREADS=1 internally to manage CPU thread usage during processing.

    # To run the binarization process, execute the script via python.
    # Ensure your hyperparameters/config are correctly set up so the script can load the desired binarizer.
    python data_gen/tts/bin/binarize.py
  7. Run MFA alignment training via train_mfa_align.py

    main

    The train_mfa_align.py script is a CLI entrypoint used to trigger MFA (Montreal Forced Aligner) alignment training. It automates the execution of the underlying shell script usr/run_mfa_train_align.sh using the corpus directory name derived from your processed data path.

    Environment Variables

    • N_PROC: (Optional) Specifies the number of jobs/processes to use for the alignment task. If not set, it defaults to the total number of CPU cores available on the system.

    Configuration Requirements

    This script relies on the processed_data_dir hyperparameter to identify the corpus name. The script extracts the last component of the path provided in processed_data_dir to set the CORPUS environment variable for the alignment process.

    # Example: Running with a specific number of processes
    N_PROC=8 python data_gen/tts/bin/train_mfa_align.py