split

tests/test_mel_synthesizer.py

@@ -0,0 +1,136 @@
+import sys
+import types
+
+import torch
+
+
+def maximum_path(score: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    path = torch.zeros_like(score)
+    for batch in range(score.shape[0]):
+        mel_length = int(mask[batch].any(-1).sum())
+        text_length = int(mask[batch].any(0).sum())
+        for frame in range(mel_length):
+            token = min(frame * text_length // mel_length, text_length - 1)
+            path[batch, frame, token] = 1
+    return path
+
+
+alignment_stub = types.ModuleType(
+    "style_bert_vits2.models.monotonic_alignment"
+)
+alignment_stub.maximum_path = maximum_path
+sys.modules["style_bert_vits2.models.monotonic_alignment"] = alignment_stub
+
+from style_bert_vits2.models.mel_synthesizer import (
+    JointMelSynthesizer,
+    MelAcousticModel,
+    MelVocoder,
+)
+
+
+class DummyEncoder(torch.nn.Module):
+    def __init__(self, hidden: int, n_mels: int) -> None:
+        super().__init__()
+        self.embedding = torch.nn.Embedding(16, hidden)
+        self.projection = torch.nn.Conv1d(hidden, n_mels * 2, 1)
+
+    def forward(self, x, x_lengths, *_args, g=None):
+        hidden = self.embedding(x).transpose(1, 2)
+        mask = (
+            torch.arange(x.shape[1], device=x.device)[None, :] < x_lengths[:, None]
+        ).unsqueeze(1).to(hidden)
+        mu, logs = self.projection(hidden * mask).chunk(2, dim=1)
+        return hidden, mu * mask, logs.tanh() * mask, mask
+
+
+class DummyDurationPredictor(torch.nn.Module):
+    def __init__(self, hidden: int) -> None:
+        super().__init__()
+        self.projection = torch.nn.Conv1d(hidden, 1, 1)
+
+    def forward(self, x, mask, g=None):
+        return self.projection(x) * mask
+
+
+class DummyStochasticDurationPredictor(DummyDurationPredictor):
+    def forward(self, x, mask, w=None, g=None, reverse=False, noise_scale=1.0):
+        prediction = super().forward(x, mask, g)
+        if reverse:
+            return prediction
+        return ((prediction - torch.log(w + 1e-6)) ** 2 * mask).sum()
+
+
+class DummyVocoder(torch.nn.Module):
+    def __init__(self, n_mels: int) -> None:
+        super().__init__()
+        self.projection = torch.nn.Conv1d(n_mels, 1, 1)
+
+    def forward(self, mel, g=None):
+        return self.projection(mel)
+
+
+def make_acoustic() -> MelAcousticModel:
+    hidden, n_mels, speaker_channels = 8, 4, 2
+    return MelAcousticModel(
+        DummyEncoder(hidden, n_mels),
+        DummyStochasticDurationPredictor(hidden),
+        DummyDurationPredictor(hidden),
+        n_mel_channels=n_mels,
+        n_speakers=2,
+        gin_channels=speaker_channels,
+        hidden_channels=hidden,
+        matcha_channels=8,
+        matcha_num_heads=2,
+        matcha_dropout=0.0,
+        matcha_sigma_min=1e-4,
+        matcha_n_timesteps=2,
+        matcha_use_diff_attention=False,
+    )
+
+
+def test_acoustic_model_outputs_mel_and_losses() -> None:
+    model = make_acoustic()
+    x = torch.tensor([[1, 2, 3], [1, 2, 0]])
+    x_lengths = torch.tensor([3, 2])
+    mel = torch.randn(2, 4, 7)
+    mel_lengths = torch.tensor([7, 5])
+    sid = torch.tensor([0, 1])
+
+    output = model(
+        x, x_lengths, mel, mel_lengths, sid, out_size=4, generate=True
+    )
+    loss = (
+        output["duration_loss"] + output["prior_loss"] + output["flow_loss"]
+        + output["generated_mel"].abs().mean()
+    )
+    loss.backward()
+
+    assert output["generated_mel"].shape == (2, 4, 4)
+    assert output["target_mel"].shape == (2, 4, 4)
+    assert all(torch.isfinite(output[name]) for name in (
+        "duration_loss", "prior_loss", "flow_loss"
+    ))
+    assert model.enc_p.embedding.weight.grad is not None
+    assert any(p.grad is not None for p in model.matcha.parameters())
+
+
+def test_joint_model_backpropagates_through_vocoder_to_acoustic() -> None:
+    acoustic = make_acoustic()
+    vocoder = MelVocoder(DummyVocoder(4), n_speakers=2, gin_channels=2)
+    model = JointMelSynthesizer(acoustic, vocoder)
+    waveform, output = model(
+        torch.tensor([[1, 2, 3]]),
+        torch.tensor([3]),
+        torch.randn(1, 4, 6),
+        torch.tensor([6]),
+        torch.tensor([0]),
+        out_size=4,
+        n_timesteps=2,
+    )
+
+    waveform.square().mean().backward()
+
+    assert waveform.shape == (1, 1, 4)
+    assert model.vocoder.generator.projection.weight.grad is not None
+    assert any(p.grad is not None for p in acoustic.matcha.parameters())
+    assert output["generated_mel"].grad_fn is not None