piper-plusのVITS DecoderをHiFi-GANからMB-iSTFTに移行してCPU推論を2.21倍高速化

初めに

今回は、piper-plusのVITSモデルでボコーダ部分(潜在表現から実際の波形を作る部分)を、HiFi-GANからMB-iSTFT (Multi-Band inverse STFT) に置き換えた話を書きます。

PR #320 (Issue #268) でこの置き換えを行い、CPU ONNX推論で 2.21倍の高速化(100音素 p50で 168.2ms → 76.2ms)、Latency P50で 43.3ms → 26.9ms (-38%) を達成しました。さらにONNXの入出力形状を変えていないので、C++/Rust/C#/Go/WASMの推論ランタイムは無修正でそのまま動きます。

「ボコーダを差し替えるだけで2倍以上速くなった」というのが今回の話です。なぜそんなことが可能なのか、信号処理側の発想を中心に解説していきます。

github.com

前提知識のミニ整理

VITSの内部構造に詳しくない方向けに、関係する部分だけ簡単に整理します。

  • VITS — テキストから直接波形を生成するEnd-to-EndのTTSモデルです[2]。テキストから音素潜在表現を作る前段(テキストエンコーダ + Flow + Posterior Encoder + Duration Predictor)と、潜在表現から波形を作る後段(Decoder)に大きく分かれています。今回いじるのは後段のDecoderだけです。
  • HiFi-GAN — メルスペクトログラムや潜在表現から波形を作るボコーダの定番[3]ConvTranspose1d でアップサンプルしながら最終的にサンプル数を稼ぐ構造で、VITSのオリジナル実装でもDecoderとして採用されています。
  • iSTFT (inverse STFT) — 短時間フーリエ変換の逆操作。マグニチュードと位相の組 から 波形 を復元します。「周波数領域 → 時間領域」の変換です。
  • PQMF (Pseudo Quadrature Mirror Filterbank) — フルバンド信号を等幅のサブバンドに分解 (analysis) したり、サブバンドからフルバンドに合成 (synthesis) したりするフィルタバンク。今回は4分割で使います。

「Decoderの最後の方の重い畳み込みを、iSTFTとPQMFという信号処理の道具に置き換える」 — これが今回の置き換えの核です。

なぜHiFi-GANは遅いのか

VITSのDecoderは、潜在表現 [B, 192, T_frames](フレームレベル)を受け取り、サンプルレベルの波形 [B, 1, T] に変換します。22,050Hzの音声・hop_length=256の場合、T = T_frames * 256 です。つまり時間軸方向に 256倍にアップサンプル する必要があります。

HiFi-GANではこの256倍を、4段の ConvTranspose1d で稼いでいました。

HiFi-GAN Generator のアップサンプル構造
[B, 192, T_frames]                           # 潜在表現
  ↓ conv_pre (Conv1d)
[B, 512, T_frames]
  ↓ ConvTranspose1d ×8 + ResBlock
[B, 256, T_frames * 8]                       # 8倍
  ↓ ConvTranspose1d ×8 + ResBlock
[B, 128, T_frames * 64]                      # 64倍
  ↓ ConvTranspose1d ×2 + ResBlock
[B, 64, T_frames * 128]                      # 128倍
  ↓ ConvTranspose1d ×2 + ResBlock
[B, 32, T_frames * 256]                      # 256倍 ← 時間軸ここまで伸びる
  ↓ Conv1d → tanh
[B, 1, T_frames * 256]                       # 波形

CPU推論で重いのは、後半の段(時間軸が既に長くなっているところ)の畳み込みです。サンプル数が多いほど畳み込みのコストが線形に増えるので、最終的な T = T_frames * 256 の長さで畳み込みをかけるのは避けたいわけです。

scripts/benchmark.py で計測した旧 HiFi-GAN ベースのCPU推論時間は、100音素 p50で約168.2ms。Decoder単体のプロファイリングでも、後段のアップサンプル畳み込みが支配的でした。

MB-iSTFT-VITS2のアイデア

MB-iSTFT-VITS2は、Kawamura らによって ICASSP 2023 で提案された軽量化手法です[1]。位置づけとしては、iSTFTNet[4] の「iSTFTで時間軸アップサンプルの一部を肩代わりする」発想と、Multi-Band MelGAN[5] の「サブバンド分解で生成サンプル数を1/4に減らす」発想を、VITSのDecoderに同時適用したものです。

arxiv.org

Kawamura, M., Shirahata, Y., Yamamoto, R., & Tachibana, K. (2023). Lightweight and High-Fidelity End-to-End Text-to-Speech with Multi-Band Generation and Inverse Short-Time Fourier Transform. In ICASSP 2023. arXiv:2210.15975. https://arxiv.org/abs/2210.15975 公式実装: https://github.com/MasayaKawamura/MB-iSTFT-VITS

発想はとてもシンプルで、256倍のアップサンプルを「畳み込み」だけで稼ぐのをやめて、3段に分担する というものです。

担当 倍率 何をするか
ConvTranspose1d (2段) 16倍 従来通り。ただし4倍×2の浅い構成
iSTFT 4倍 「マグニチュードと位相」から波形を復元すると、hop_length=4分の長さが一気に出る
PQMF synthesis 4倍 4本のサブバンド信号から1本のフルバンド信号を合成
合計 256倍 HiFi-GANと同じ総倍率

ポイントは2つあります。

  1. iSTFTは畳み込みではない — 周波数領域 (mag/phase) から時間領域 (波形) への変換を、「フレームごとのIFFT + overlap-add」という決まったアルゴリズムで行います。深いネットワークではなく単発の信号処理なので非常に軽いです。
  2. PQMFで「最後の4倍」をサブバンドに分散 — 「フルバンドの波形を直接生成する」のではなく「4本のサブバンド波形を生成して合成する」形にします。各サブバンド波形は周波数帯域が狭いのでサンプルレートも1/4で済み、ニューラルネット側で生成すべきサンプル数も1/4になります。

論文の Table 2 では、オリジナルVITS (RTF 0.27) に対して MB-iSTFT-VITS が RTF 0.078 (3.46倍速)、MS-iSTFT-VITS が RTF 0.066 (4.1倍速) と、エンドツーエンドで 3.46〜4.1倍速 が報告されています[1, Table 2]。さらに同論文の Table 1 では、Decoder が VITS 全体の推論時間の 96%以上を占める(VITS全体 RTF 0.849 のうち Decoder が 0.819)ことが示されており[1, Table 1]、ここを iSTFT + PQMF に置き換えるだけで全体の速度がここまで伸びる、というのが論文の主張です。

なお論文では MS-iSTFT-VITS (4.1倍速) の方が高速ですが、サブバンドごとに異なるFFTサイズの iSTFT を並列適用する構成のため、ONNXグラフが複雑化します。piper-plus では学習の安定性とエクスポートのシンプルさを優先して、まずは等幅4サブバンドの MB 版を採用しました。

piper-plus 側でも、HiFi-GAN ベースの旧Decoder と MB-iSTFT 化した新Decoder で同条件比較したところ、エンドツーエンドのCPU推論で 2.21倍の高速化 を実機で確認しました(後述のベンチマーク参照)。

データフローで比較する

新旧のDecoder末端でどう処理が変わったかを見てみます。

旧 HiFi-GAN
[B, 192, T]
  ↓ ConvTranspose1d ×(8,8,2,2) + ResBlocks  ★後段の畳み込みが重い
[B, 32, T * 256]
  ↓ Conv1d → tanh
[B, 1, T * 256]                              # 直接フルバンド波形

新 MB-iSTFT
[B, 192, T]
  ↓ ConvTranspose1d ×(4,4) + ResBlocks       ★アップサンプルは16倍までで打ち切り
[B, 64, T * 16]
  ↓ subband_conv_post (Conv1d)
[B, 4 × (n_fft+2), T * 16]                   # subbands × (mag + phase) の係数を一気に生成
  ↓ reshape して mag = exp(...) / phase = sin(...) * π
[B*4, n_fft//2+1, T * 16]                    # iSTFT入力 (subband軸をbatch軸に畳む)
  ↓ OnnxISTFT (conv_transpose1d でiSTFT)
[B, 4, T * 64]                               # 4本のサブバンド波形
  ↓ PQMF synthesis (conv1d 系)
[B, 1, T * 256]                              # フルバンド波形

新Decoderの最終出力 [B, 1, T*256] は旧HiFi-GANと同じ形です。ここがONNX互換性のカギ で、推論側のランタイムから見るとモデルの入出力契約は何も変わっていません。中身が「全部ニューラルネット」から「ニューラルネット + iSTFT + PQMF」になっただけです。

3つの新規コンポーネント

実装は src/python/piper_train/vits/ 配下に3ファイル追加しました。

ファイル 中身 役割
mb_istft.py PQMF, MBiSTFTGenerator 新Decoder本体 + サブバンド分解/合成
stft_onnx.py OnnxISTFT ONNXに乗るiSTFT
stft_loss.py MultiResolutionSTFTLoss 学習用のサブバンドSTFT損失

順に見ていきます。

PQMF — 音声を4つの周波数帯域に分けるフィルタ

PQMF (Pseudo Quadrature Mirror Filterbank) は、ざっくり言うと 音声を周波数帯域ごとに4分割する道具 です。スピーカーで「低音はウーファー、中低音はミッドバス、中高音はミッド、高音はツイーター」と帯域を分けて鳴らすイメージに近く、音声波形を4本のサブバンド信号に分けたり、逆に4本のサブバンドを足し合わせて元の波形に戻したりできます。

なぜわざわざ分けるかというと、各サブバンドは元の1/4の周波数帯域しか持たないので、サンプリング定理から1/4のサンプル数で表現できる からです。フルバンドで T サンプル必要な情報が、4本のサブバンドそれぞれは T/4 サンプルで済む。ニューラルネットが生成すべきサンプル数も1/4で済むので、計算量も減ります。

PQMFには嬉しい性質が2つあります。

  • 学習しない固定フィルタ — フィルタ係数は信号処理の理論で決まっていて、学習中もずっと同じ値です。重みではなく register_buffer で持つので optimizer.step() で更新されません。
  • ほぼ完全再構成性analysis(分解)synthesis(合成) を順に通すと、ほぼ元の波形に戻ります。「分けたものを正しく戻せる」保証があるので、Decoder末尾で安心して使えます。

入出力サイズの整理

PQMFには analysis()synthesis() の2つのメソッドがあり、ニューラルネットの中での使い場面が違います。

メソッド 入力 出力 いつ使う
analysis(x) フルバンド [B, 1, T] サブバンド [B, 4, T/4] 学習時にGT波形を分解 (損失計算用)
synthesis(x) サブバンド [B, 4, T_sub] フルバンド [B, 1, T_sub * 4] Decoder末尾、推論時もここを通る

学習時は双方向で使い、推論時は synthesis() だけ通ります。

内部実装の概略

PQMFのフィルタ係数は、信号処理で標準的な手順で構築されます。式や係数の意味に踏み込まなくても、3ステップ覚えておけば十分です。

  1. プロトタイプローパス — Kaiser窓 × sinc関数で、まず1本の理想ローパスフィルタを作る
  2. コサイン変調 — このプロトタイプを「中心周波数をずらした4本のフィルタ」にコピーする (1/8、3/8、5/8、7/8 の位置)
  3. synthesisフィルタは時間反転 — analysis用フィルタを時間反転すれば synthesis用になる

実際のコードはこんな形です(信号処理に詳しくない方は「固定の係数を作って register_buffer に置いている」だけ把握すればOKです)。

class PQMF(nn.Module):
    def __init__(self, subbands=4, taps=62, cutoff_ratio=0.15, beta=9.0):
        super().__init__()
        self.subbands = subbands
        # ... Kaiser窓 × sinc のプロトタイプローパスを作る ...

        # コサイン変調で 4本のフィルタを生成
        analysis_filter = np.zeros((subbands, 1, filter_length), dtype=np.float64)
        for k in range(subbands):
            for n in range(filter_length):
                analysis_filter[k, 0, n] = (
                    2.0 * prototype[n]
                    * np.cos((2*k + 1) * np.pi / (2*subbands) * (n - subbands/2))
                )

        # synthesis フィルタは時間反転
        synthesis_filter = analysis_filter[:, :, ::-1].copy()

        # 学習対象ではなく固定バッファとして登録
        self.register_buffer("analysis_filter",
                             torch.from_numpy(analysis_filter).float())
        self.register_buffer("synthesis_filter",
                             torch.from_numpy(synthesis_filter).float())

フィルタを実際に当てる処理は、PyTorchの F.conv1d / F.conv_transpose1d 1発で書けます。

  • analysis(x) — フィルタで4チャンネルに分解 → ストライド4でダウンサンプル
  • synthesis(x) — ストライド4でアップサンプル → 合成フィルタを当てる

ニューラルネット側は「サブバンド [B, 4, T_sub] を作る」ところまで担当し、PQMFの synthesis() がそれをフルバンド [B, 1, T_sub * 4] に戻す、という分業になっています。ニューラルネットがサンプル数1/4で済む のはこの分業のおかげです。

MBiSTFTGenerator — 新しいDecoder本体

メインのDecoderです。前半部(conv_pre → アップサンプル + ResBlocks)はHiFi-GANと同じ構造で、違うのは末尾だけです。

class MBiSTFTGenerator(nn.Module):
    def __init__(
        self,
        initial_channel,
        resblock,
        resblock_kernel_sizes,
        resblock_dilation_sizes,
        upsample_rates=(4, 4),                    # ←(8,8,2,2)から(4,4)に
        upsample_initial_channel=256,
        upsample_kernel_sizes=(16, 16),
        gin_channels=0,
        n_fft=16,
        hop_length=4,
        subbands=4,
        pqmf=None,
    ):
        super().__init__()
        # ... conv_pre, ups, resblocks (HiFi-GAN同等) ...

        # 末尾: subband × (magnitude + phase) を一気に出すConv1d
        post_in_channels = upsample_initial_channel // (2 ** len(upsample_rates))
        self.subband_conv_post = Conv1d(
            post_in_channels, subbands * (n_fft + 2), 7, padding=3
        )

        self.istft = OnnxISTFT(n_fft=n_fft, hop_length=hop_length)
        self.pqmf = pqmf if pqmf is not None else PQMF(subbands=subbands)

forward() の末尾で何をしているか、データの形と一緒に追います。

def forward(self, x, g=None):
    # ... 前半のアップサンプルとResBlocks (省略) ...
    # この時点で x: [B, 64, T_frames * 16]

    x = F.leaky_relu(x, LRELU_SLOPE)
    x = self.subband_conv_post(x)
    # x: [B, 4 × (n_fft+2), T_frames * 16]
    #         └ subbands × (magnitude bins + phase bins) のチャンネル

    B = x.size(0)
    T_frames = x.size(-1)
    n_half = self.n_fft // 2 + 1   # n_fft=16 → 9

    # サブバンドとmag/phaseに分けてreshape
    x = x.reshape(B, self.subbands, self.n_fft + 2, T_frames)

    # magnitude側: 必ず正の値にしたいので exp
    mag = torch.exp(x[:, :, :n_half, :])
    # phase側: [-π, π] に押し込めたいので sin × π
    phase = torch.sin(x[:, :, n_half:, :]) * math.pi

    # サブバンド軸をバッチ軸に畳んでiSTFTに渡す
    mag = mag.reshape(B * self.subbands, n_half, T_frames)
    phase = phase.reshape(B * self.subbands, n_half, T_frames)
    sub_wav = self.istft(mag, phase)
    # [B*subbands, 1, T_sub_raw] → サブバンド軸を戻す
    subbands_signal = sub_wav.reshape(B, self.subbands, -1)

    # 期待長にトリム
    expected_sub_T = T_frames * self.hop_length
    subbands_signal = subbands_signal[..., :expected_sub_T]

    # サブバンド → フルバンド
    fullband = self.pqmf.synthesis(subbands_signal)

    if self.onnx_export_mode:
        return fullband
    return fullband, subbands_signal

ポイント:

  • subband_conv_post「サブバンド × (mag + phase)」のチャンネル を一気に出す層です。subbands=4n_fft=16 のとき、出力チャンネルは 4 * (16+2) = 72 になります。
  • magnitudeに exp をかけているのは、必ず正の値にしたいから(マグニチュードは負になりえない)。phaseに sin(...) * π をかけているのは [-π, π] の範囲に収めるためです。これらの非線形でネットワークの出力に物理的な意味を持たせています。
  • 末尾の if self.onnx_export_mode: 分岐がもう一つの肝です。学習時は (fullband, subbands) のタプル を返してサブバンド側を損失計算に使い、ONNXエクスポート時は fullband のみ を返して旧HiFi-GANと同じ出力形状 [B, 1, T] に揃えます。

OnnxISTFT — ONNXに乗るiSTFT

ここが今回の実装で一番頭を使ったところです。

iSTFT自体はPyTorchに torch.istft がありますが、ONNX opset 15には iSTFT op が存在しません。普通にexportすると失敗するか、Aten opとして埋め込まれて多くのランタイムで動かなくなります。

そこで、iSTFTを F.conv_transpose1d 1発で表現する ことにしました。考え方は次の通りです。

通常のiSTFTは「フレームごとのIFFT → 合成窓掛け → overlap-add → 窓正規化」と複数ステップに分かれています。これらを 行列演算1つに事前融合 し、その行列を conv_transpose1d の重みとして使います。stride=hop_length がそのまま overlap-add の働きをします。

class OnnxISTFT(nn.Module):
    def __init__(self, n_fft=16, hop_length=4):
        super().__init__()
        # 事前計算した「逆DFT行列 × Hann窓 ÷ 正規化定数」
        inverse_basis = self._build_inverse_basis(n_fft, hop_length)
        # shape: (n_fft + 2, 1, n_fft)  例: (18, 1, 16)
        self.register_buffer("inverse_basis", inverse_basis)
        self.hop_length = hop_length
        self.n_fft = n_fft

    def forward(self, magnitude, phase):
        # mag/phase → 実部/虚部
        real = magnitude * torch.cos(phase)
        imag = magnitude * torch.sin(phase)
        combined = torch.cat([real, imag], dim=1)  # (B, n_fft+2, T)

        # 事前計算した基底とのconv_transpose1dで波形を一発復元
        waveform = F.conv_transpose1d(
            combined, self.inverse_basis, stride=self.hop_length
        )
        return waveform

_build_inverse_basis() の中で、

  1. 逆DFTの合成行列 S を作る (one-sided なので hermitian対称性で内側のbinに係数2倍を入れる)
  2. 周期Hann窓 を掛ける
  3. WSS (Window-Sum-of-Squares) で割る ことでoverlap-add時の正規化を吸収
  4. conv_transpose1d の重み形状 (in_channels, out_channels/groups, kernel_size) に整形

という処理を行っています。

@staticmethod
def _build_inverse_basis(n_fft, hop_length):
    cutoff = n_fft // 2 + 1

    # one-sided iDFT 合成行列 S
    n_idx = np.arange(n_fft)
    k_idx = np.arange(cutoff)
    angle = 2.0 * np.pi * np.outer(n_idx, k_idx) / n_fft
    cos_table = np.cos(angle)
    sin_table = np.sin(angle)

    S_re = np.zeros((n_fft, cutoff))
    S_im = np.zeros((n_fft, cutoff))

    # DC bin: 係数 1/N
    S_re[:, 0] = 1.0 / n_fft
    # Nyquist bin: 係数 1/N
    S_re[:, cutoff - 1] = cos_table[:, cutoff - 1] / n_fft
    S_im[:, cutoff - 1] = -sin_table[:, cutoff - 1] / n_fft
    # 内側 bin: hermitian対称性で 2/N
    for ki in range(1, cutoff - 1):
        S_re[:, ki] = 2.0 * cos_table[:, ki] / n_fft
        S_im[:, ki] = -2.0 * sin_table[:, ki] / n_fft

    S = np.hstack([S_re, S_im])  # (n_fft, 2*cutoff)

    # Hann窓 と 正規化定数を吸収
    window = np.hanning(n_fft + 1)[:n_fft]
    wss = np.sum(window**2) * hop_length / n_fft
    inverse_basis = S * (window[:, np.newaxis] / wss)

    # conv_transpose1d の重み形状に整形
    inverse_basis = inverse_basis.T[:, np.newaxis, :]  # (2*cutoff, 1, n_fft)
    return torch.FloatTensor(inverse_basis)

ランタイム側では inverse_basis学習対象ではない固定バッファ です。Conv1d / ConvTranspose1d という、どのONNXランタイムでもネイティブにサポートされている演算子だけで iSTFT が表現できる、という点が大事なところです。

MultiResolutionSTFTLoss — 学習時にサブバンドの品質を保つロス

サブバンド波形の品質を上げるため、学習時にはフルバンドの再構成ロス(mel loss / adversarial loss など)に加えて、サブバンド側のSTFTロス を計算します。論文に沿って3つの異なる解像度(FFTサイズ違い)で計算し、平均をとります。

class MultiResolutionSTFTLoss(nn.Module):
    def __init__(
        self,
        fft_sizes=(171, 384, 683),
        hop_sizes=(10, 30, 60),
        win_sizes=(60, 150, 300),
    ):
        super().__init__()
        self.stft_losses = nn.ModuleList()
        for fs, hs, ws in zip(fft_sizes, hop_sizes, win_sizes, strict=False):
            self.stft_losses.append(STFTLoss(fs, hs, ws))

    def forward(self, x, y):
        # サブバンド軸をバッチ軸に畳んで一気に計算
        if x.dim() == 3:
            B, S, T = x.shape
            x = x.reshape(B * S, T)
            y = y.reshape(B * S, T)

        loss = 0.0
        for stft_loss in self.stft_losses:
            loss += stft_loss(x, y)
        return loss / len(self.stft_losses)

各解像度ごとに、

  • Spectral Convergence Loss — 予測マグニチュードと正解マグニチュードのフロベニウスノルム比。形が合っているかを見ます。
  • Log STFT Magnitude Loss — 対数マグニチュードのL1距離。低エネルギー帯の細部をしっかり合わせます。

の2つを計算します。SpectralConvergenceLoss は分母が0になり得るので clamp(min=1e-7) でゼロ除算を防いでいます。

class SpectralConvergenceLoss(nn.Module):
    def forward(self, x_mag, y_mag):
        return torch.norm(y_mag - x_mag, p="fro") / torch.norm(y_mag, p="fro").clamp(min=1e-7)

このロスは 学習時のみ 計算されます。推論用ONNXグラフには含まれないので、推論コストは増えません。

学習パイプラインへの統合

SynthesizerTrn (models.py) では、decMBiSTFTGenerator に置き換えました。HiFi-GAN Generator クラスは完全に削除しているので、if mb_istft: のような分岐は残っていません。

# models.py (抜粋)
self.dec = MBiSTFTGenerator(
    initial_channel=inter_channels,
    resblock=resblock,
    resblock_kernel_sizes=resblock_kernel_sizes,
    resblock_dilation_sizes=resblock_dilation_sizes,
    upsample_rates=upsample_rates,
    upsample_initial_channel=upsample_initial_channel,
    upsample_kernel_sizes=upsample_kernel_sizes,
    gin_channels=gin_channels,
)

VitsModel (lightning.py) ではPQMFと sub-band STFT loss を常時初期化します。PQMFインスタンスは MBiSTFTGenerator 側と共有させて、フィルタ係数バッファの重複を防ぎます。

# lightning.py (__init__ 抜粋)
self.pqmf = PQMF(subbands=4)
# Generator と同じインスタンスを使う (バッファ重複回避)
self.model_g.dec.pqmf = self.pqmf
self.sub_stft_loss = MultiResolutionSTFTLoss(
    fft_sizes=self.hparams.sub_stft_fft_sizes,
    hop_sizes=self.hparams.sub_stft_hop_sizes,
    win_sizes=self.hparams.sub_stft_win_sizes,
)

学習ステップでは、Generator出力のサブバンド o_mbGT波形をPQMF.analysisで分解したもの y_mb を比較して、sub-band STFT lossをトータルロスに足し込みます。

# lightning.py (training_step 抜粋)
loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl

if o_mb is not None:
    y_mb = self.pqmf.analysis(y)              # GTを4サブバンドに分解
    loss_sub_stft = self.sub_stft_loss(o_mb, y_mb) * self.hparams.c_sub_stft
    loss_gen_all = loss_gen_all + loss_sub_stft

ONNXエクスポートと互換性

ONNXエクスポート時には set_export_mode() ユーティリティで全モジュールの onnx_export_mode を一括Trueにします。これによって MBiSTFTGenerator.forward() がfullbandのみを返すようになり、出力形状 [B, 1, T] で書き出されます。

# export_onnx.py (抜粋)
def set_export_mode(model, mode=True):
    """onnx_export_mode を持つ全サブモジュールに一括設定"""
    for m in model.modules():
        if hasattr(m, "onnx_export_mode"):
            m.onnx_export_mode = mode

# ...
set_export_mode(model_g, True)

結果として、推論側ランタイムから見た契約はこうなります。

項目 旧 HiFi-GAN ONNX 新 MB-iSTFT ONNX
入力 input_ids, input_lengths, scales ... 同じ
出力 [B, 1, T] の波形 同じ [B, 1, T] の波形
opset 15 15
必要な op Conv1d / ConvTranspose1d / 標準演算 Conv1d / ConvTranspose1d / 標準演算

つまり、C++/Rust/C#/Go/WASM/Pythonの推論コードは 1行も変えずに 新モデルを読めます。既存のHiFi-GAN ONNXファイルもそのまま動き続けます。

開発環境

学習環境:

  • OS: Ubuntu 22.04
  • GPU: NVIDIA V100 × 4 (DDP)
  • Python: 3.12
  • パッケージマネージャー: uv
  • PyTorch Lightning: 2.x
  • piper-plus: dev (PR #320)

推論ベンチ環境:

  • CPU: Intel Xeon E5-2650 v4 @ 2.20GHz (48コア)
  • OS: Linux x86_64
  • ONNX Runtime: 1.24
  • 計測パラメータ: warmup 5回 / 計測30回 / intra-op threads = auto

環境構築

リポジトリをクローンして uv sync で依存をインストールします。

git clone https://github.com/ayutaz/piper-plus.git
cd piper-plus
uv sync

学習コマンドとCLIの変更点

--mb-istft フラグは廃止され、常時有効になりました。upsample_rates=(4, 4) / upsample_kernel_sizes=(16, 16) も自動設定されるため、ユーザー側で意識する必要はありません。--quality medium --quality high どちらでも MB-iSTFT が動作します。

新規追加されたCLIオプションは sub-band STFT loss のチューニング用です。

オプション デフォルト 説明
--c-sub-stft 1.0 sub-band STFT loss の重み
--sub-stft-fft-sizes 171,384,683 sub-band Multi-resolution STFT loss の FFT サイズ (3解像度)
--sub-stft-hop-sizes 10,30,60 hop サイズ
--sub-stft-win-sizes 60,150,300 窓サイズ

学習コマンドの例 (シングルGPU):

uv run python -m piper_train \
  --dataset-dir /path/to/dataset \
  --accelerator gpu --devices 1 --precision 16-mixed \
  --max_epochs 200 --batch-size 16 \
  --quality medium \
  --prosody-dim 16 \
  --ema-decay 0.9995

学習自体のコマンドは変わっていません。Decoderだけが裏で MB-iSTFT に差し替わっています。

ONNXエクスポート

学習済みチェックポイントからONNXにエクスポートします。手順は旧来と同じです。

CUDA_VISIBLE_DEVICES="" uv run python -m piper_train.export_onnx \
  ./training_dir/lightning_logs/version_0/checkpoints/last.ckpt \
  ./output/model.onnx

config.json を同じディレクトリに配置します。

cp ./training_dir/config.json ./output/config.json

推論の実行

エクスポートしたモデルで音声合成します。

uv run python -m piper \
  -m ./output/model.onnx \
  -f ./output/test.wav \
  "MB-iSTFTデコーダによる高速な音声合成のテストです。"

ベンチマークスクリプトでも計測できます。

uv run python scripts/benchmark.py \
  --model ./output/model.onnx \
  --config ./output/config.json \
  --language ja \
  --text "MB-iSTFTデコーダによる高速な音声合成のテストです。" \
  --n-warmup 5 --n-runs 30 --format markdown

実行結果

エンドツーエンドのCPU推論時間

scripts/benchmark.py で 100 音素 p50 を計測した結果です。旧 HiFi-GAN と新 MB-iSTFT、それぞれを同条件で scratch から学習・エクスポートして比較しました。

指標 HiFi-GAN MB-iSTFT
Inference time (ms, 100 phoneme p50) 168.2 76.2 2.21x
RTF 0.066 0.037 -44%
Latency P50 (ms) 43.3 26.9 -38%

つくよみちゃんファインチューンモデル(500 epoch)では、MB-iSTFTで 61.9 ms (RTF 0.046) を確認しています。

なお、論文 Table 2 ではオリジナルVITSとの比較で MB-iSTFT-VITS が RTF 0.27 → 0.078 (約3.46倍速)、MS-iSTFT-VITS が RTF 0.066 (4.1倍速) を報告しています[1, Table 2]。piper-plus の比較対象は VITS そのものではなく Piper の HiFi-GAN ベース(CPU向けに既に軽量化された構成)なので、論文と同じ比率にはなりませんが、「Decoder の置き換えで全体が大幅に速くなる」という傾向は同じ向きで再現できました。

他システムとの比較

同じ計測環境(Intel Xeon E5-2650 v4 / Linux / ORT 1.24)で、競合システムとも比較しました。

システム RTF ↓ Latency P50 (ms) サイズ (MB) RAM (MB) 初回起動 (ms)
piper-plus (MB-iSTFT) 0.078 27 38 208 1633
Piper本家 (archived) 0.066 35 60 185 2510
sherpa-onnx (VITS Piper-fmt) 0.075 53 60 202 2554

Latency P50 で Piper本家比 -23%、sherpa-onnx比 -49% です。モデルサイズも 38MB と最小クラスを維持しています。

学習済みモデル

PR #320のマージに合わせて、MB-iSTFT対応のベースモデル・追加モデルを再公開しました。

モデル エポック 完了日 公開先
6lang MB-iSTFT base (571話者, 6言語, scratch学習) 75 2026-04-16 piper-plus-base
つくよみちゃん MB-iSTFT FT (6lang base から) 500 2026-05-02 piper-plus-tsukuyomi-chan
CSS10 Japanese MB-iSTFT FT (6lang base から, 6,841発話) 50 2026-05-03 piper-plus-css10-ja-6lang

所感

「畳み込みだけで256倍を稼ぐ」のをやめ、信号処理の道具(iSTFT + PQMF)に 最後の16倍ぶん を任せたら、エンドツーエンドのCPU推論で2.21倍の高速化が得られました。しかも推論側ランタイムは無修正です。論文値 (3.46〜4.1倍速、対 オリジナルVITS) には及びませんが、比較対象が CPU向けに既に軽量化された Piper 系列の HiFi-GAN ベースなので、これくらいが相場という印象です。

特にONNX互換iSTFTを「逆DFT行列 × Hann窓 ÷ 正規化定数」を事前計算して conv_transpose1d の重みに焼き込む形にしたことで、限定的なopsetの中でもネイティブ演算子だけで成立しているのが面白いところです。同じ手法は他のVITS派生モデルや、ONNXに乗せたい音声処理モジュールにも応用できると思います。

参考文献

[1] Kawamura, M., Shirahata, Y., Yamamoto, R., & Tachibana, K. (2023). Lightweight and High-Fidelity End-to-End Text-to-Speech with Multi-Band Generation and Inverse Short-Time Fourier Transform. In ICASSP 2023. arXiv:2210.15975. https://arxiv.org/abs/2210.15975 — 公式実装: MasayaKawamura/MB-iSTFT-VITS

[2] Kim, J., Kong, J., & Son, J. (2021). Conditional Variational Autoencoder with Adversarial Learning for End-to-End Text-to-Speech (VITS). In ICML 2021. arXiv:2106.06103. https://arxiv.org/abs/2106.06103

[3] Kong, J., Kim, J., & Bae, J. (2020). HiFi-GAN: Generative Adversarial Networks for Efficient and High Fidelity Speech Synthesis. In NeurIPS 2020. arXiv:2010.05646. https://arxiv.org/abs/2010.05646

[4] Kaneko, T., Tanaka, K., Kameoka, H., & Seki, S. (2022). iSTFTNet: Fast and Lightweight Mel-Spectrogram Vocoder Incorporating Inverse Short-Time Fourier Transform. In ICASSP 2022. arXiv:2203.02395. https://arxiv.org/abs/2203.02395 — MB-iSTFT-VITS の iSTFT 適用パートのベース手法

[5] Yang, G., Yang, S., Liu, K., Fang, P., Chen, W., & Xie, L. (2020). Multi-Band MelGAN: Faster Waveform Generation for High-Quality Text-to-Speech. arXiv:2005.05106. https://arxiv.org/abs/2005.05106 — MB-iSTFT-VITS のサブバンド分解パートのベース手法