fix

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)