tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
tuna2134
2026-07-20 21:25:49 +09:00
parent 0c2be00f0a
commit b8ce11605c
17 changed files with 496 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
colab.ipynb
View File

@@ -279,6 +279,13 @@
"# JP-Extra （日本語特化版）を使うかどうか。日本語の能力が向上する代わりに英語と中国語は使えなくなります。\n",
"use_jp_extra = True\n",
"\n",
"# Matcha Flow + Differential Attention V2だけを既存モデルへFTする場合はTrue\n",
"# 初回はTrueを推奨。収束後の全体FTではFalseに戻し、学習率を下げてください。\n",
"matcha_only = True\n",
"\n",
"# Matcha-only FTの学習率。全体FTへ移行するときは2e-5程度まで下げることを推奨\n",
"learning_rate = 1e-4\n",
"\n",
"# 学習のバッチサイズ。VRAMのはみ出具合に応じて調整してください。\n",
"batch_size = 4\n",
"\n",
@@ -344,6 +351,7 @@
" val_per_lang=0,\n",
" log_interval=200,\n",
" yomi_error=yomi_error,\n",
" matcha_only=matcha_only,\n",
")"
]
},
@@ -381,6 +389,16 @@
"dataset_path = str(paths.dataset_path)\n",
"config_path = str(paths.config_path)\n",
"\n",
"from pathlib import Path\n",
"from configure_matcha_ft import configure_matcha_ft\n",
"configure_matcha_ft(\n",
" Path(config_path),\n",
" matcha_only=matcha_only,\n",
" learning_rate=learning_rate,\n",
" epochs=epochs,\n",
" save_every_steps=save_every_steps,\n",
")\n",
"\n",
"with open(\"default_config.yml\", \"r\", encoding=\"utf-8\") as f:\n",
" yml_data = yaml.safe_load(f)\n",
"yml_data[\"model_name\"] = model_name\n",

2
configs/config.json
View File

@@ -17,6 +17,7 @@
"c_mel": 45,
"c_kl": 1.0,
"c_matcha": 1.0,
"matcha_only": false,
"skip_optimizer": false,
"freeze_ZH_bert": false,
"freeze_JP_bert": false,
@@ -55,6 +56,7 @@
"matcha_dropout": 0.05,
"matcha_sigma_min": 0.0001,
"matcha_n_timesteps": 8,
"matcha_use_diff_attention": true,
"inter_channels": 192,
"hidden_channels": 192,
"filter_channels": 768,

2
configs/config_jp_extra.json
View File

@@ -18,6 +18,7 @@
"c_mel": 45,
"c_kl": 1.0,
"c_matcha": 1.0,
"matcha_only": false,
"c_commit": 100,
"skip_optimizer": false,
"freeze_ZH_bert": false,
@@ -55,6 +56,7 @@
"matcha_dropout": 0.05,
"matcha_sigma_min": 0.0001,
"matcha_n_timesteps": 8,
"matcha_use_diff_attention": true,
"inter_channels": 192,
"hidden_channels": 192,
"filter_channels": 768,

79
configure_matcha_ft.py Normal file
View File

@@ -0,0 +1,79 @@
"""Configure a dataset for Matcha-only or joint fine-tuning."""
import argparse
import json
import shutil
from pathlib import Path
from typing import Optional
def configure_matcha_ft(
config_path: Path,
matcha_only: bool = True,
learning_rate: Optional[float] = None,
epochs: Optional[int] = None,
save_every_steps: Optional[int] = None,
) -> dict:
if not config_path.is_file():
raise FileNotFoundError(f"Config file not found: {config_path}")
with config_path.open(encoding="utf-8") as stream:
config = json.load(stream)
train = config.setdefault("train", {})
model = config.setdefault("model", {})
train["matcha_only"] = matcha_only
model["use_matcha"] = True
model["matcha_use_diff_attention"] = True
if learning_rate is not None:
if learning_rate <= 0:
raise ValueError("Learning rate must be positive")
train["learning_rate"] = learning_rate
if epochs is not None:
if epochs < 1:
raise ValueError("Epochs must be at least 1")
train["epochs"] = epochs
if save_every_steps is not None:
if save_every_steps < 1:
raise ValueError("Save interval must be at least 1")
train["eval_interval"] = save_every_steps
backup_path = config_path.with_suffix(config_path.suffix + ".bak")
if not backup_path.exists():
shutil.copy2(config_path, backup_path)
with config_path.open("w", encoding="utf-8") as stream:
json.dump(config, stream, indent=2, ensure_ascii=False)
stream.write("\n")
return config
def main() -> None:
parser = argparse.ArgumentParser(
description="Configure Matcha Flow fine-tuning for an existing dataset"
)
parser.add_argument("--config", type=Path, required=True)
parser.add_argument(
"--joint",
action="store_true",
help="Jointly fine-tune the full model instead of Matcha-only training",
)
parser.add_argument("--learning-rate", type=float)
parser.add_argument("--epochs", type=int)
parser.add_argument("--save-every-steps", type=int)
args = parser.parse_args()
config = configure_matcha_ft(
config_path=args.config,
matcha_only=not args.joint,
learning_rate=args.learning_rate,
epochs=args.epochs,
save_every_steps=args.save_every_steps,
)
mode = "joint fine-tuning" if args.joint else "Matcha-only fine-tuning"
print(f"Configured {args.config} for {mode}")
print(f"learning_rate={config['train']['learning_rate']}")
print(f"epochs={config['train']['epochs']}")
print(f"save_every_steps={config['train']['eval_interval']}")
if __name__ == "__main__":
main()

3
docs/CLI.md
View File

@@ -62,7 +62,7 @@ Optional
## 2. Preprocess
```bash
python preprocess_all.py -m <model_name> [--use_jp_extra] [-b <batch_size>] [-e <epochs>] [-s <save_every_steps>] [--num_processes <num_processes>] [--normalize] [--trim] [--val_per_lang <val_per_lang>] [--log_interval <log_interval>] [--freeze_EN_bert] [--freeze_JP_bert] [--freeze_ZH_bert] [--freeze_style] [--freeze_decoder] [--yomi_error <yomi_error>]
python preprocess_all.py -m <model_name> [--use_jp_extra] [--matcha_only] [-b <batch_size>] [-e <epochs>] [-s <save_every_steps>] [--num_processes <num_processes>] [--normalize] [--trim] [--val_per_lang <val_per_lang>] [--log_interval <log_interval>] [--freeze_EN_bert] [--freeze_JP_bert] [--freeze_ZH_bert] [--freeze_style] [--freeze_decoder] [--yomi_error <yomi_error>]
```
Required:
@@ -81,6 +81,7 @@ Optional:
- `--freeze_style`: Freeze style vector.
- `--freeze_decoder`: Freeze decoder.
- `--use_jp_extra`: Use JP-Extra model.
- `--matcha_only`: Enable Matcha-only fine-tuning and Differential Attention V2.
- `--val_per_lang`: Validation data per language (default: 0).
- `--log_interval`: Log interval (default: 200).
- `--yomi_error`: How to handle yomi errors (default: `raise`: raise an error after preprocessing all texts, `skip`: skip the texts with errors, `use`: use the texts with errors by ignoring unknown characters).

31
docs/MATCHA_FLOW.md
View File

@@ -18,7 +18,38 @@ Relevant settings:
- `model.matcha_sigma_min`: minimum flow path noise.
- `model.matcha_n_timesteps`: Euler steps at inference; more steps trade speed
for refinement quality.
- `model.matcha_use_diff_attention`: enables Differential Attention V2 in the
U-Net transformer blocks. It uses explicit PyTorch matrix multiplication and
softmax, with no FlashAttention, CUDA-only kernel, or extra dependency.
The Matcha branch has new parameters, so enabling it on an existing checkpoint
requires fine-tuning before inference. It is trained on the same random latent
segments used by the waveform generator to keep memory use bounded.
## Matcha-only fine-tuning
Set `train.matcha_only` to `true` to train only `net_g.matcha`. All legacy
generator parameters are frozen, and adversarial, duration-discriminator,
WavLM-discriminator, mel, duration, and KL optimization steps are skipped. The
existing VITS normalizing-flow latent `z_p` remains the detached CFM target.
This mode requires `model.use_matcha: true`. Optimizer state from a full-model
checkpoint is intentionally not restored because its parameter groups differ;
model weights are still loaded. After the Matcha loss has converged, set
`train.matcha_only` back to `false` and use a lower learning rate for joint
fine-tuning.
For an already preprocessed dataset (including Google Colab), the mode can be
configured without rerunning preprocessing:
```bash
python configure_matcha_ft.py --config Data/MyModel/config.json \
--learning-rate 0.0001 --epochs 20 --save-every-steps 100
```
To switch to the second-stage joint fine-tuning:
```bash
python configure_matcha_ft.py --config Data/MyModel/config.json \
--joint --learning-rate 0.00002
```

9
gradio_tabs/train.py
View File

@@ -58,6 +58,7 @@ def initialize(
freeze_decoder: bool,
use_jp_extra: bool,
log_interval: int,
matcha_only: bool = False,
):
global logger_handler
paths = get_path(model_name)
@@ -70,7 +71,7 @@ def initialize(
logger_handler = logger.add(paths.dataset_path / file_name)
logger.info(
f"Step 1: start initialization...\nmodel_name: {model_name}, batch_size: {batch_size}, epochs: {epochs}, save_every_steps: {save_every_steps}, freeze_ZH_bert: {freeze_ZH_bert}, freeze_JP_bert: {freeze_JP_bert}, freeze_EN_bert: {freeze_EN_bert}, freeze_style: {freeze_style}, freeze_decoder: {freeze_decoder}, use_jp_extra: {use_jp_extra}"
f"Step 1: start initialization...\nmodel_name: {model_name}, batch_size: {batch_size}, epochs: {epochs}, save_every_steps: {save_every_steps}, freeze_ZH_bert: {freeze_ZH_bert}, freeze_JP_bert: {freeze_JP_bert}, freeze_EN_bert: {freeze_EN_bert}, freeze_style: {freeze_style}, freeze_decoder: {freeze_decoder}, use_jp_extra: {use_jp_extra}, matcha_only: {matcha_only}"
)
default_config_path = (
@@ -86,6 +87,10 @@ def initialize(
config["train"]["epochs"] = epochs
config["train"]["eval_interval"] = save_every_steps
config["train"]["log_interval"] = log_interval
config["train"]["matcha_only"] = matcha_only
if matcha_only:
config["model"]["use_matcha"] = True
config["model"]["matcha_use_diff_attention"] = True
config["train"]["freeze_EN_bert"] = freeze_EN_bert
config["train"]["freeze_JP_bert"] = freeze_JP_bert
@@ -275,6 +280,7 @@ def preprocess_all(
val_per_lang: int,
log_interval: int,
yomi_error: str,
matcha_only: bool = False,
):
if model_name == "":
return False, "Error: モデル名を入力してください"
@@ -290,6 +296,7 @@ def preprocess_all(
freeze_decoder=freeze_decoder,
use_jp_extra=use_jp_extra,
log_interval=log_interval,
matcha_only=matcha_only,
)
if not success:
return False, message

6
preprocess_all.py
View File

@@ -72,6 +72,11 @@ if __name__ == "__main__":
action="store_true",
help="Use JP-Extra model",
)
parser.add_argument(
"--matcha_only",
action="store_true",
help="Train only the Matcha Flow and Differential Attention modules",
)
parser.add_argument(
"--val_per_lang",
type=int,
@@ -110,4 +115,5 @@ if __name__ == "__main__":
val_per_lang=args.val_per_lang,
log_interval=args.log_interval,
yomi_error=args.yomi_error,
matcha_only=args.matcha_only,
)

2
style_bert_vits2/models/hyper_parameters.py
View File

@@ -28,6 +28,7 @@ class HyperParametersTrain(BaseModel):
c_mel: int = 45
c_kl: float = 1.0
c_matcha: float = 1.0
matcha_only: bool = False
c_commit: int = 100
skip_optimizer: bool = False
freeze_ZH_bert: bool = False
@@ -83,6 +84,7 @@ class HyperParametersModel(BaseModel):
matcha_dropout: float = 0.05
matcha_sigma_min: float = 0.0001
matcha_n_timesteps: int = 8
matcha_use_diff_attention: bool = False
inter_channels: int = 192
hidden_channels: int = 192
filter_channels: int = 768

2
style_bert_vits2/models/infer.py
View File

@@ -41,6 +41,7 @@ def get_net_g(
matcha_dropout=hps.model.matcha_dropout,
matcha_sigma_min=hps.model.matcha_sigma_min,
matcha_n_timesteps=hps.model.matcha_n_timesteps,
matcha_use_diff_attention=hps.model.matcha_use_diff_attention,
inter_channels=hps.model.inter_channels,
hidden_channels=hps.model.hidden_channels,
filter_channels=hps.model.filter_channels,
@@ -78,6 +79,7 @@ def get_net_g(
matcha_dropout=hps.model.matcha_dropout,
matcha_sigma_min=hps.model.matcha_sigma_min,
matcha_n_timesteps=hps.model.matcha_n_timesteps,
matcha_use_diff_attention=hps.model.matcha_use_diff_attention,
inter_channels=hps.model.inter_channels,
hidden_channels=hps.model.hidden_channels,
filter_channels=hps.model.filter_channels,

113
style_bert_vits2/models/matcha_flow.py
View File

@@ -13,6 +13,21 @@ from torch import nn
from torch.nn import functional as F
def configure_matcha_only_training(model: nn.Module) -> int:
"""Freeze a synthesizer except for its Matcha branch.
Returns the number of trainable parameters after configuration.
"""
matcha = getattr(model, "matcha", None)
if not getattr(model, "use_matcha", False) or matcha is None:
raise ValueError("Matcha-only training requires an enabled Matcha branch")
for parameter in model.parameters():
parameter.requires_grad = False
for parameter in matcha.parameters():
parameter.requires_grad = True
return sum(parameter.numel() for parameter in matcha.parameters())
class SinusoidalTimeEmbedding(nn.Module):
def __init__(self, channels: int) -> None:
super().__init__()
@@ -58,13 +73,24 @@ class MaskedResBlock(nn.Module):
class MaskedTransformerBlock(nn.Module):
def __init__(self, channels: int, num_heads: int, dropout: float) -> None:
def __init__(
self,
channels: int,
num_heads: int,
dropout: float,
use_diff_attention: bool,
) -> None:
super().__init__()
if channels % num_heads:
raise ValueError("Matcha channels must be divisible by attention heads")
self.norm1 = nn.LayerNorm(channels)
self.attention = nn.MultiheadAttention(
channels, num_heads, dropout=dropout, batch_first=True
self.use_diff_attention = use_diff_attention
self.attention = (
DifferentialAttentionV2(channels, num_heads, dropout)
if use_diff_attention
else nn.MultiheadAttention(
channels, num_heads, dropout=dropout, batch_first=True
)
)
self.norm2 = nn.LayerNorm(channels)
self.feed_forward = nn.Sequential(
@@ -78,14 +104,68 @@ class MaskedTransformerBlock(nn.Module):
x = x.transpose(1, 2)
valid = mask[:, 0].bool()
hidden = self.norm1(x)
hidden, _ = self.attention(
hidden, hidden, hidden, key_padding_mask=~valid, need_weights=False
)
if self.use_diff_attention:
hidden = self.attention(hidden, valid)
else:
hidden, _ = self.attention(
hidden, hidden, hidden, key_padding_mask=~valid, need_weights=False
)
x = x + hidden
x = x + self.feed_forward(self.norm2(x))
return x.transpose(1, 2) * mask
class DifferentialAttentionV2(nn.Module):
"""Differential Attention V2 without FlashAttention or custom kernels.
Each logical head has two query heads that share one key/value head. Their
contexts are combined as ``context_1 - sigmoid(lambda) * context_2``, where
lambda is projected per token and logical head.
"""
def __init__(self, channels: int, num_heads: int, dropout: float) -> None:
super().__init__()
if channels % num_heads:
raise ValueError("Channels must be divisible by attention heads")
self.num_heads = num_heads
self.head_channels = channels // num_heads
self.scale = self.head_channels**-0.5
self.query_projection = nn.Linear(channels, channels * 2)
self.key_projection = nn.Linear(channels, channels)
self.value_projection = nn.Linear(channels, channels)
self.lambda_projection = nn.Linear(channels, num_heads)
self.output_projection = nn.Linear(channels, channels)
self.attention_dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
batch, length, channels = x.shape
queries = self.query_projection(x).view(
batch, length, self.num_heads, 2, self.head_channels
)
keys = self.key_projection(x).view(
batch, length, self.num_heads, self.head_channels
)
values = self.value_projection(x).view(
batch, length, self.num_heads, self.head_channels
)
queries = queries.permute(0, 2, 3, 1, 4)
keys = keys.permute(0, 2, 1, 3)
values = values.permute(0, 2, 1, 3)
scores = torch.matmul(queries, keys.unsqueeze(2).transpose(-1, -2))
scores = scores * self.scale
scores = scores.masked_fill(~valid[:, None, None, None, :], float("-inf"))
attention = self.attention_dropout(torch.softmax(scores, dim=-1))
contexts = torch.matmul(attention, values.unsqueeze(2))
lambda_value = torch.sigmoid(self.lambda_projection(x))
lambda_value = lambda_value.permute(0, 2, 1).unsqueeze(-1)
contexts = contexts[:, :, 0] - lambda_value * contexts[:, :, 1]
contexts = contexts.permute(0, 2, 1, 3).reshape(batch, length, channels)
contexts = self.output_projection(contexts)
return contexts * valid.unsqueeze(-1).to(contexts.dtype)
class MatchaEstimator(nn.Module):
"""A compact masked 1-D U-Net velocity estimator."""
@@ -96,6 +176,7 @@ class MatchaEstimator(nn.Module):
channels: int,
num_heads: int,
dropout: float,
use_diff_attention: bool,
) -> None:
super().__init__()
time_channels = channels * 4
@@ -107,12 +188,18 @@ class MatchaEstimator(nn.Module):
nn.Linear(time_channels, time_channels),
)
self.down_block = MaskedResBlock(input_channels, channels, time_channels)
self.down_attention = MaskedTransformerBlock(channels, num_heads, dropout)
self.down_attention = MaskedTransformerBlock(
channels, num_heads, dropout, use_diff_attention
)
self.downsample = nn.Conv1d(channels, channels, 3, stride=2, padding=1)
self.mid_block = MaskedResBlock(channels, channels, time_channels)
self.mid_attention = MaskedTransformerBlock(channels, num_heads, dropout)
self.mid_attention = MaskedTransformerBlock(
channels, num_heads, dropout, use_diff_attention
)
self.up_block = MaskedResBlock(channels * 2, channels, time_channels)
self.up_attention = MaskedTransformerBlock(channels, num_heads, dropout)
self.up_attention = MaskedTransformerBlock(
channels, num_heads, dropout, use_diff_attention
)
self.output = nn.Conv1d(channels, latent_channels, 1)
def forward(
@@ -152,11 +239,17 @@ class MatchaFlow(nn.Module):
num_heads: int = 2,
dropout: float = 0.05,
sigma_min: float = 1e-4,
use_diff_attention: bool = False,
) -> None:
super().__init__()
self.sigma_min = sigma_min
self.estimator = MatchaEstimator(
latent_channels, speaker_channels, channels, num_heads, dropout
latent_channels,
speaker_channels,
channels,
num_heads,
dropout,
use_diff_attention,
)
def compute_loss(

1
style_bert_vits2/models/models.py
View File

@@ -961,6 +961,7 @@ class SynthesizerTrn(nn.Module):
num_heads=kwargs.get("matcha_num_heads", n_heads),
dropout=kwargs.get("matcha_dropout", 0.05),
sigma_min=kwargs.get("matcha_sigma_min", 1e-4),
use_diff_attention=kwargs.get("matcha_use_diff_attention", False),
)
def forward(

1
style_bert_vits2/models/models_jp_extra.py
View File

@@ -1018,6 +1018,7 @@ class SynthesizerTrn(nn.Module):
num_heads=kwargs.get("matcha_num_heads", n_heads),
dropout=kwargs.get("matcha_dropout", 0.05),
sigma_min=kwargs.get("matcha_sigma_min", 1e-4),
use_diff_attention=kwargs.get("matcha_use_diff_attention", False),
)
def forward(

33
tests/test_configure_matcha_ft.py Normal file
View File

@@ -0,0 +1,33 @@
import json
from pathlib import Path
from configure_matcha_ft import configure_matcha_ft
def test_configure_matcha_ft_and_joint_mode(tmp_path: Path) -> None:
config_path = tmp_path / "config.json"
config_path.write_text(
json.dumps(
{
"train": {"learning_rate": 0.0002, "epochs": 100},
"model": {"use_matcha": False},
}
),
encoding="utf-8",
)
configured = configure_matcha_ft(
config_path,
learning_rate=0.0001,
epochs=20,
save_every_steps=50,
)
assert configured["train"]["matcha_only"] is True
assert configured["train"]["learning_rate"] == 0.0001
assert configured["train"]["eval_interval"] == 50
assert configured["model"]["use_matcha"] is True
assert configured["model"]["matcha_use_diff_attention"] is True
assert config_path.with_suffix(".json.bak").is_file()
configured = configure_matcha_ft(config_path, matcha_only=False)
assert configured["train"]["matcha_only"] is False

42
tests/test_matcha_flow.py
View File

@@ -1,6 +1,45 @@
import torch
from style_bert_vits2.models.matcha_flow import MatchaFlow
from style_bert_vits2.models.matcha_flow import (
DifferentialAttentionV2,
MatchaFlow,
configure_matcha_only_training,
)
class DummySynthesizer(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.use_matcha = True
self.legacy = torch.nn.Linear(4, 4)
self.matcha = torch.nn.Linear(4, 4)
def test_matcha_only_freezes_legacy_parameters() -> None:
model = DummySynthesizer()
trainable = configure_matcha_only_training(model)
assert trainable == sum(parameter.numel() for parameter in model.matcha.parameters())
assert all(parameter.requires_grad for parameter in model.matcha.parameters())
assert not any(parameter.requires_grad for parameter in model.legacy.parameters())
def test_differential_attention_v2_on_cpu() -> None:
attention = DifferentialAttentionV2(channels=16, num_heads=4, dropout=0.0)
inputs = torch.randn(2, 7, 16, requires_grad=True)
valid = torch.tensor(
[[True, True, True, True, True, True, True],
[True, True, True, True, True, False, False]]
)
output = attention(inputs, valid)
output.square().mean().backward()
assert output.shape == inputs.shape
assert output.device.type == "cpu"
assert torch.count_nonzero(output[1, 5:]) == 0
assert attention.lambda_projection.weight.grad is not None
def test_matcha_flow_loss_and_sampling_with_odd_length() -> None:
@@ -10,6 +49,7 @@ def test_matcha_flow_loss_and_sampling_with_odd_length() -> None:
channels=16,
num_heads=2,
dropout=0.0,
use_diff_attention=True,
)
target = torch.randn(2, 8, 7)
mu = torch.randn_like(target)

91
train_ms.py
View File

@@ -28,6 +28,7 @@ from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
from style_bert_vits2.logging import logger
from style_bert_vits2.models import commons, utils
from style_bert_vits2.models.hyper_parameters import HyperParameters
from style_bert_vits2.models.matcha_flow import configure_matcha_only_training
from style_bert_vits2.models.models import (
DurationDiscriminator,
MultiPeriodDiscriminator,
@@ -42,8 +43,7 @@ torch.backends.cudnn.allow_tf32 = (
True # If encontered training problem,please try to disable TF32.
)
torch.set_float32_matmul_precision("medium")
torch.backends.cuda.sdp_kernel("flash")
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_flash_sdp(False)
torch.backends.cuda.enable_mem_efficient_sdp(
True
) # Not available if torch version is lower than 2.0
@@ -286,7 +286,7 @@ def run():
logger.info("Using normal MAS for VITS1")
mas_noise_scale_initial = 0.0
noise_scale_delta = 0.0
if hps.model.use_duration_discriminator is True:
if hps.model.use_duration_discriminator is True and not hps.train.matcha_only:
logger.info("Using duration discriminator for VITS2")
net_dur_disc = DurationDiscriminator(
hps.model.hidden_channels,
@@ -322,6 +322,7 @@ def run():
matcha_dropout=hps.model.matcha_dropout,
matcha_sigma_min=hps.model.matcha_sigma_min,
matcha_n_timesteps=hps.model.matcha_n_timesteps,
matcha_use_diff_attention=hps.model.matcha_use_diff_attention,
inter_channels=hps.model.inter_channels,
hidden_channels=hps.model.hidden_channels,
filter_channels=hps.model.filter_channels,
@@ -365,6 +366,10 @@ def run():
for param in net_g.dec.parameters():
param.requires_grad = False
if hps.train.matcha_only:
trainable = configure_matcha_only_training(net_g)
logger.info(f"Matcha-only training enabled ({trainable:,} trainable parameters)")
net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(local_rank)
optim_g = torch.optim.AdamW(
filter(lambda p: p.requires_grad, net_g.parameters()),
@@ -401,7 +406,7 @@ def run():
utils.checkpoints.get_latest_checkpoint_path(model_dir, "DUR_*.pth"),
net_dur_disc,
optim_dur_disc,
skip_optimizer=hps.train.skip_optimizer,
skip_optimizer=hps.train.skip_optimizer or hps.train.matcha_only,
)
if not optim_dur_disc.param_groups[0].get("initial_lr"):
optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
@@ -409,13 +414,13 @@ def run():
utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"),
net_g,
optim_g,
skip_optimizer=hps.train.skip_optimizer,
skip_optimizer=hps.train.skip_optimizer or hps.train.matcha_only,
)
_, optim_d, d_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
utils.checkpoints.get_latest_checkpoint_path(model_dir, "D_*.pth"),
net_d,
optim_d,
skip_optimizer=hps.train.skip_optimizer,
skip_optimizer=hps.train.skip_optimizer or hps.train.matcha_only,
)
if not optim_g.param_groups[0].get("initial_lr"):
optim_g.param_groups[0]["initial_lr"] = g_resume_lr
@@ -618,10 +623,17 @@ def train_and_evaluate(
train_loader.batch_sampler.set_epoch(epoch)
global global_step
net_g.train()
net_d.train()
if net_dur_disc is not None:
net_dur_disc.train()
if hps.train.matcha_only:
net_g.eval()
net_g.module.matcha.train()
net_d.eval()
if net_dur_disc is not None:
net_dur_disc.eval()
else:
net_g.train()
net_d.train()
if net_dur_disc is not None:
net_dur_disc.train()
for batch_idx, (
x,
x_lengths,
@@ -637,7 +649,7 @@ def train_and_evaluate(
en_bert,
style_vec,
) in enumerate(train_loader):
if net_g.module.use_noise_scaled_mas:
if net_g.module.use_noise_scaled_mas and not hps.train.matcha_only:
current_mas_noise_scale = (
net_g.module.mas_noise_scale_initial
- net_g.module.noise_scale_delta * global_step
@@ -684,6 +696,63 @@ def train_and_evaluate(
en_bert,
style_vec,
)
if hps.train.matcha_only:
loss_matcha = matcha_loss.float() * hps.train.c_matcha
if hps.train.matcha_only:
optim_g.zero_grad()
scaler.scale(loss_matcha).backward()
scaler.unscale_(optim_g)
grad_norm_g = commons.clip_grad_value_(
net_g.module.matcha.parameters(), None
)
scaler.step(optim_g)
scaler.update()
if rank == 0:
if global_step % hps.train.log_interval == 0 and not hps.speedup:
utils.summarize(
writer=writer,
global_step=global_step,
scalars={
"loss/g/total": loss_matcha,
"loss/g/matcha": loss_matcha,
"learning_rate": optim_g.param_groups[0]["lr"],
"grad_norm_g": grad_norm_g,
},
)
if (
global_step % hps.train.eval_interval == 0
and global_step != 0
and initial_step != global_step
):
assert hps.model_dir is not None
utils.checkpoints.save_checkpoint(
net_g,
optim_g,
hps.train.learning_rate,
epoch,
os.path.join(hps.model_dir, f"G_{global_step}.pth"),
)
utils.safetensors.save_safetensors(
net_g,
epoch,
os.path.join(
config.out_dir,
f"{config.model_name}_e{epoch}_s{global_step}.safetensors",
),
for_infer=True,
)
global_step += 1
if pbar is not None:
pbar.set_description(
f"Matcha FT {epoch}({100.0 * batch_idx / len(train_loader):.0f}%)/{hps.train.epochs}"
)
pbar.update()
continue
with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16):
mel = spec_to_mel_torch(
spec,
hps.data.filter_length,

99
train_ms_jp_extra.py
View File

@@ -28,6 +28,7 @@ from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
from style_bert_vits2.logging import logger
from style_bert_vits2.models import commons, utils
from style_bert_vits2.models.hyper_parameters import HyperParameters
from style_bert_vits2.models.matcha_flow import configure_matcha_only_training
from style_bert_vits2.models.models_jp_extra import (
DurationDiscriminator,
MultiPeriodDiscriminator,
@@ -44,8 +45,7 @@ torch.backends.cudnn.allow_tf32 = (
)
torch.set_num_threads(1)
torch.set_float32_matmul_precision("medium")
torch.backends.cuda.sdp_kernel("flash")
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_flash_sdp(False)
torch.backends.cuda.enable_mem_efficient_sdp(
True
) # Not available if torch version is lower than 2.0
@@ -291,7 +291,7 @@ def run():
logger.info("Using normal MAS for VITS1")
mas_noise_scale_initial = 0.0
noise_scale_delta = 0.0
if hps.model.use_duration_discriminator is True:
if hps.model.use_duration_discriminator is True and not hps.train.matcha_only:
logger.info("Using duration discriminator for VITS2")
net_dur_disc = DurationDiscriminator(
hps.model.hidden_channels,
@@ -302,7 +302,7 @@ def run():
).cuda(local_rank)
else:
net_dur_disc = None
if hps.model.use_wavlm_discriminator is True:
if hps.model.use_wavlm_discriminator is True and not hps.train.matcha_only:
net_wd = WavLMDiscriminator(
hps.model.slm.hidden, hps.model.slm.nlayers, hps.model.slm.initial_channel
).cuda(local_rank)
@@ -335,6 +335,7 @@ def run():
matcha_dropout=hps.model.matcha_dropout,
matcha_sigma_min=hps.model.matcha_sigma_min,
matcha_n_timesteps=hps.model.matcha_n_timesteps,
matcha_use_diff_attention=hps.model.matcha_use_diff_attention,
inter_channels=hps.model.inter_channels,
hidden_channels=hps.model.hidden_channels,
filter_channels=hps.model.filter_channels,
@@ -367,6 +368,10 @@ def run():
for param in net_g.dec.parameters():
param.requires_grad = False
if hps.train.matcha_only:
trainable = configure_matcha_only_training(net_g)
logger.info(f"Matcha-only training enabled ({trainable:,} trainable parameters)")
net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(local_rank)
optim_g = torch.optim.AdamW(
filter(lambda p: p.requires_grad, net_g.parameters()),
@@ -430,7 +435,7 @@ def run():
),
net_dur_disc,
optim_dur_disc,
skip_optimizer=hps.train.skip_optimizer,
skip_optimizer=hps.train.skip_optimizer or hps.train.matcha_only,
)
if not optim_dur_disc.param_groups[0].get("initial_lr"):
optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
@@ -462,13 +467,13 @@ def run():
utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"),
net_g,
optim_g,
skip_optimizer=hps.train.skip_optimizer,
skip_optimizer=hps.train.skip_optimizer or hps.train.matcha_only,
)
_, optim_d, d_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
utils.checkpoints.get_latest_checkpoint_path(model_dir, "D_*.pth"),
net_d,
optim_d,
skip_optimizer=hps.train.skip_optimizer,
skip_optimizer=hps.train.skip_optimizer or hps.train.matcha_only,
)
if not optim_g.param_groups[0].get("initial_lr"):
optim_g.param_groups[0]["initial_lr"] = g_resume_lr
@@ -706,12 +711,21 @@ def train_and_evaluate(
# train_loader.batch_sampler.set_epoch(epoch)
global global_step
net_g.train()
net_d.train()
if net_dur_disc is not None:
net_dur_disc.train()
if net_wd is not None:
net_wd.train()
if hps.train.matcha_only:
net_g.eval()
net_g.module.matcha.train()
net_d.eval()
if net_dur_disc is not None:
net_dur_disc.eval()
if net_wd is not None:
net_wd.eval()
else:
net_g.train()
net_d.train()
if net_dur_disc is not None:
net_dur_disc.train()
if net_wd is not None:
net_wd.train()
for batch_idx, (
x,
x_lengths,
@@ -725,7 +739,7 @@ def train_and_evaluate(
bert,
style_vec,
) in enumerate(train_loader):
if net_g.module.use_noise_scaled_mas:
if net_g.module.use_noise_scaled_mas and not hps.train.matcha_only:
current_mas_noise_scale = (
net_g.module.mas_noise_scale_initial
- net_g.module.noise_scale_delta * global_step
@@ -769,6 +783,63 @@ def train_and_evaluate(
bert,
style_vec,
)
if hps.train.matcha_only:
loss_matcha = matcha_loss.float() * hps.train.c_matcha
if hps.train.matcha_only:
optim_g.zero_grad()
scaler.scale(loss_matcha).backward()
scaler.unscale_(optim_g)
grad_norm_g = commons.clip_grad_value_(
net_g.module.matcha.parameters(), None
)
scaler.step(optim_g)
scaler.update()
if rank == 0:
if global_step % hps.train.log_interval == 0 and not hps.speedup:
utils.summarize(
writer=writer,
global_step=global_step,
scalars={
"loss/g/total": loss_matcha,
"loss/g/matcha": loss_matcha,
"learning_rate": optim_g.param_groups[0]["lr"],
"grad_norm_g": grad_norm_g,
},
)
if (
global_step % hps.train.eval_interval == 0
and global_step != 0
and initial_step != global_step
):
assert hps.model_dir is not None
utils.checkpoints.save_checkpoint(
net_g,
optim_g,
hps.train.learning_rate,
epoch,
os.path.join(hps.model_dir, f"G_{global_step}.pth"),
)
utils.safetensors.save_safetensors(
net_g,
epoch,
os.path.join(
config.out_dir,
f"{config.model_name}_e{epoch}_s{global_step}.safetensors",
),
for_infer=True,
)
global_step += 1
if pbar is not None:
pbar.set_description(
f"Matcha FT {epoch}({100.0 * batch_idx / len(train_loader):.0f}%)/{hps.train.epochs}"
)
pbar.update()
continue
with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16):
mel = spec_to_mel_torch(
spec,
hps.data.filter_length,