fix

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",

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,

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,

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()

docs/CLI.md

@@ -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).

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
+```

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

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,
     )

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

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,

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(
+        self.use_diff_attention = use_diff_attention
-            channels, num_heads, dropout=dropout, batch_first=True
+        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(
+        if self.use_diff_attention:
-            hidden, hidden, hidden, key_padding_mask=~valid, need_weights=False
+            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(

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(

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(

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

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)

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(False)
-torch.backends.cuda.enable_flash_sdp(True)
 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()
+    if hps.train.matcha_only:
-    net_d.train()
+        net_g.eval()
-    if net_dur_disc is not None:
+        net_g.module.matcha.train()
-        net_dur_disc.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,

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(False)
-torch.backends.cuda.enable_flash_sdp(True)
 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()
+    if hps.train.matcha_only:
-    net_d.train()
+        net_g.eval()
-    if net_dur_disc is not None:
+        net_g.module.matcha.train()
-        net_dur_disc.train()
+        net_d.eval()
-    if net_wd is not None:
+        if net_dur_disc is not None:
-        net_wd.train()
+            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,