72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
import torch
|
|
|
|
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:
|
|
model = MatchaFlow(
|
|
latent_channels=8,
|
|
speaker_channels=4,
|
|
channels=16,
|
|
num_heads=2,
|
|
dropout=0.0,
|
|
use_diff_attention=True,
|
|
)
|
|
target = torch.randn(2, 8, 7)
|
|
mu = torch.randn_like(target)
|
|
mask = torch.tensor(
|
|
[[[1, 1, 1, 1, 1, 1, 1]], [[1, 1, 1, 1, 1, 0, 0]]],
|
|
dtype=target.dtype,
|
|
)
|
|
speaker = torch.randn(2, 4, 1)
|
|
|
|
loss = model.compute_loss(target, mask, mu, speaker)
|
|
loss.backward()
|
|
|
|
assert loss.ndim == 0
|
|
assert torch.isfinite(loss)
|
|
assert any(parameter.grad is not None for parameter in model.parameters())
|
|
|
|
sample = model(mu, mask, n_timesteps=2, speaker=speaker)
|
|
assert sample.shape == target.shape
|
|
assert torch.count_nonzero(sample[1, :, 5:]) == 0
|