diff --git a/colab.ipynb b/colab.ipynb index 2dfc841..2159ca8 100644 --- a/colab.ipynb +++ b/colab.ipynb @@ -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", diff --git a/configs/config.json b/configs/config.json index 755f857..e06c9a9 100644 --- a/configs/config.json +++ b/configs/config.json @@ -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, diff --git a/configs/config_jp_extra.json b/configs/config_jp_extra.json index fd2d4fe..5f01d44 100644 --- a/configs/config_jp_extra.json +++ b/configs/config_jp_extra.json @@ -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, diff --git a/configure_matcha_ft.py b/configure_matcha_ft.py new file mode 100644 index 0000000..bb49f43 --- /dev/null +++ b/configure_matcha_ft.py @@ -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() diff --git a/docs/CLI.md b/docs/CLI.md index 726c42d..9d0c6e1 100644 --- a/docs/CLI.md +++ b/docs/CLI.md @@ -62,7 +62,7 @@ Optional ## 2. Preprocess ```bash -python preprocess_all.py -m [--use_jp_extra] [-b ] [-e ] [-s ] [--num_processes ] [--normalize] [--trim] [--val_per_lang ] [--log_interval ] [--freeze_EN_bert] [--freeze_JP_bert] [--freeze_ZH_bert] [--freeze_style] [--freeze_decoder] [--yomi_error ] +python preprocess_all.py -m [--use_jp_extra] [--matcha_only] [-b ] [-e ] [-s ] [--num_processes ] [--normalize] [--trim] [--val_per_lang ] [--log_interval ] [--freeze_EN_bert] [--freeze_JP_bert] [--freeze_ZH_bert] [--freeze_style] [--freeze_decoder] [--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). diff --git a/docs/MATCHA_FLOW.md b/docs/MATCHA_FLOW.md index e917986..dbd332c 100644 --- a/docs/MATCHA_FLOW.md +++ b/docs/MATCHA_FLOW.md @@ -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 +``` diff --git a/gradio_tabs/train.py b/gradio_tabs/train.py index fbf8522..61dfda8 100644 --- a/gradio_tabs/train.py +++ b/gradio_tabs/train.py @@ -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 diff --git a/preprocess_all.py b/preprocess_all.py index 1ec7d09..89cfa05 100644 --- a/preprocess_all.py +++ b/preprocess_all.py @@ -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, ) diff --git a/style_bert_vits2/models/hyper_parameters.py b/style_bert_vits2/models/hyper_parameters.py index b2081f8..afbb619 100644 --- a/style_bert_vits2/models/hyper_parameters.py +++ b/style_bert_vits2/models/hyper_parameters.py @@ -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 diff --git a/style_bert_vits2/models/infer.py b/style_bert_vits2/models/infer.py index c718fff..835a042 100644 --- a/style_bert_vits2/models/infer.py +++ b/style_bert_vits2/models/infer.py @@ -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, diff --git a/style_bert_vits2/models/matcha_flow.py b/style_bert_vits2/models/matcha_flow.py index 7729ac0..0a4a50e 100644 --- a/style_bert_vits2/models/matcha_flow.py +++ b/style_bert_vits2/models/matcha_flow.py @@ -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( diff --git a/style_bert_vits2/models/models.py b/style_bert_vits2/models/models.py index 22ccd23..d270ff3 100644 --- a/style_bert_vits2/models/models.py +++ b/style_bert_vits2/models/models.py @@ -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( diff --git a/style_bert_vits2/models/models_jp_extra.py b/style_bert_vits2/models/models_jp_extra.py index ed90d07..221e11e 100644 --- a/style_bert_vits2/models/models_jp_extra.py +++ b/style_bert_vits2/models/models_jp_extra.py @@ -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( diff --git a/tests/test_configure_matcha_ft.py b/tests/test_configure_matcha_ft.py new file mode 100644 index 0000000..a997e1f --- /dev/null +++ b/tests/test_configure_matcha_ft.py @@ -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 diff --git a/tests/test_matcha_flow.py b/tests/test_matcha_flow.py index d4e735f..f0e4903 100644 --- a/tests/test_matcha_flow.py +++ b/tests/test_matcha_flow.py @@ -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) diff --git a/train_ms.py b/train_ms.py index 8fc1dee..74dd9ab 100644 --- a/train_ms.py +++ b/train_ms.py @@ -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, diff --git a/train_ms_jp_extra.py b/train_ms_jp_extra.py index c73083d..170c7ae 100644 --- a/train_ms_jp_extra.py +++ b/train_ms_jp_extra.py @@ -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,