Flow matchingを用いた高速・高品質なゼロショットTTS「ZipVoice」を日本語対応して学習・推論をする

AI

初めに

この前にZipVoiceを動かして英語の生成をしてみました。しかし日本語の対応がされていなかったので求めているユースケースとは異なります。

ayousanz.hatenadiary.jp

今回は日本語対応を行って追加学習を行い、日本語を話せるようにしてみます

以下がツクヨミちゃんデータセットで行った際にデモ動画になります

youtu.be

日本語対応のために行ったこと

以下が今回日本語対応および学習のために行ったことになります

日本語対応

  • pyopenjtalk-plusによる日本語G2P
  • 日英混合テキストの正規化
  • 日本語トークン追加時のembedding拡張

高速化

  • DataLoader改善
  • DDP最適化
  • torch.compileとTF32最適化

以下のブランチに対応した内容を公開しています

github.com

開発環境

学習環境構築

k2ライブラリがWindows非対応だったため、Docker環境(linux環境)が必要でした。このライブラリをいれないとNaN勾配、訓練が発散の問題が起きていました。

そのため以下のような Dockerファイルを作成しました

# Dockerfile for ZipVoice Japanese training with k2 support
FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV UV_SYSTEM_PYTHON=1

# Install system dependencies
RUN apt-get update && apt-get install -y \
    git \
    wget \
    curl \
    libsndfile1 \
    ffmpeg \
    && rm -rf /var/lib/apt/lists/*

# Install uv
RUN pip install uv

# Set working directory
WORKDIR /workspace

# Copy project files
COPY pyproject.toml /workspace/
COPY uv.lock /workspace/
COPY README.md /workspace/
COPY zipvoice /workspace/zipvoice

# Install project dependencies with uv sync (including cuda extras for k2)
RUN uv sync --extra cuda

# Add virtual environment to PATH
ENV PATH="/workspace/.venv/bin:$PATH"

# Default command
CMD ["bash"]

また 実行しやすいように docker-composeも作成しています

version: '3.8'

services:
  zipvoice-train:
    build:
      context: .
      dockerfile: Dockerfile
    image: zipvoice-japanese:latest
    container_name: zipvoice-japanese-train

    # GPU support
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

    # Mount volumes for data persistence
    volumes:
      - ./data:/workspace/data
      - ./download:/workspace/download
      - ./exp:/workspace/exp
      - ./egs:/workspace/egs

    # Environment variables
    environment:
      - WANDB_API_KEY=${WANDB_API_KEY}
      - CUDA_VISIBLE_DEVICES=0

    # Keep container running
    stdin_open: true
    tty: true

    # Working directory
    working_dir: /workspace

    # Default command for training
    # Using python directly since packages are installed via uv sync in Dockerfile
    command: >
      python -m zipvoice.bin.train_zipvoice
      --world-size 1
      --use-fp16 1
      --finetune 1
      --base-lr 0.0001
      --num-iters 10000
      --save-every-n 1000
      --max-duration 60
      --drop-last 0
      --model-config download/zipvoice/model.json
      --checkpoint download/zipvoice/model.pt
      --tokenizer japanese
      --token-file data/tokens_japanese_extended.txt
      --dataset custom
      --train-manifest data/fbank/tsukuyomi_cuts_train.jsonl.gz
      --dev-manifest data/fbank/tsukuyomi_cuts_dev.jsonl.gz
      --exp-dir exp/zipvoice_japanese
      --wandb-project zipvoice-japanese

  # Interactive shell for debugging
  zipvoice-shell:
    build:
      context: .
      dockerfile: Dockerfile
    image: zipvoice-japanese:latest
    container_name: zipvoice-japanese-shell

    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

    volumes:
      - ./data:/workspace/data
      - ./download:/workspace/download
      - ./exp:/workspace/exp
      - ./egs:/workspace/egs
      - ./zipvoice:/workspace/zipvoice

    environment:
      - WANDB_API_KEY=${WANDB_API_KEY}
      - CUDA_VISIBLE_DEVICES=0

    stdin_open: true
    tty: true
    working_dir: /workspace
    command: bash

このときの uv向けのpyproject.tomlは以下の通りです

[project]
name = "zipvoice"
version = "0.1.0"
description = "Fast and High-Quality Zero-Shot Text-to-Speech with Flow Matching"
readme = "README.md"
requires-python = ">=3.10,<3.12"
dependencies = [
    "cn2an>=0.5.23",
    "huggingface-hub>=1.2.3",
    "inflect>=7.5.0",
    "jaconv>=0.4.1",
    "jieba>=0.42.1",
    "lhotse>=1.32.1",
    "numpy>=1.26.0",
    "pydub>=0.25.1",
    "pyopenjtalk-plus>=0.4.1.post7",
    "pypinyin>=0.55.0",
    "safetensors>=0.7.0",
    "tensorboard>=2.20.0",
    "torch>=2.0.0",
    "torchaudio>=2.0.0",
    "vocos>=0.1.0",
    "wandb>=0.19.0",
    "piper-phonemize",
]

[project.optional-dependencies]
# k2 requires CUDA and is only available on Linux
cuda = [
    "k2==1.24.4.dev20241030+cuda12.4.torch2.5.1; sys_platform == 'linux'",
]

[[tool.uv.index]]
name = "pytorch-cu124"
url = "https://download.pytorch.org/whl/cu124"
explicit = true

[tool.uv.sources]
torch = { index = "pytorch-cu124" }
torchaudio = { index = "pytorch-cu124" }

[tool.uv]
find-links = [
    "https://k2-fsa.github.io/k2/cuda.html",
    "https://k2-fsa.github.io/icefall/piper_phonemize.html",
]

[tool.isort]
profile = "black"

[tool.black]
line-length = 88

以下のようなデータセットを作成します

utt_001   こんにちは、今日はいい天気ですね。 /path/to/audio/001.wav
utt_002 私の名前は田中です。  /path/to/audio/002.wav

データセットを作成した後に以下を実行します

uv run python -m zipvoice.bin.prepare_dataset \
    --tsv-path data/raw/custom_train.tsv \
    --prefix custom \
    --subset train \
    --num-jobs 4 \
    --output-dir data/manifests

トークン付加を行います

uv run python -m zipvoice.bin.prepare_tokens \
    --input-file data/manifests/custom_cuts_train.jsonl.gz \
    --output-file data/manifests/custom_cuts_train_tokens.jsonl.gz \
    --tokenizer japanese

Fbank特徴量計算を計算します

uv run python -m zipvoice.bin.compute_fbank \
    --source-dir data/manifests \
    --dest-dir data/fbank \
    --dataset custom \
    --subset train_tokens \
    --num-jobs 4

学習

以下を実行して学習を開始します

Dockerでファインチューニングを開始します

docker-compose build
docker-compose up zipvoice-train

推論

以下のコマンドで作成をしたモデルで推論を行います

# モデルディレクトリ準備
mkdir -p exp/custom/infer
cp exp/custom/checkpoint-10000.pt exp/custom/infer/model.pt
cp download/zipvoice/model.json exp/custom/infer/model.json
cp data/tokens_japanese_extended.txt exp/custom/infer/tokens.txt

# 推論実行
docker-compose run --rm zipvoice-shell python -m zipvoice.bin.infer_zipvoice \
    --model-dir exp/custom/infer \
    --tokenizer japanese \
    --prompt-wav data/prompt.wav \
    --prompt-text "プロンプト音声の書き起こし" \
    --text "合成したいテキスト" \
    --res-wav-path output.wav