split

style_bert_vits2/models/matcha_flow.py

@@ -270,8 +270,7 @@ class MatchaFlow(nn.Module):
         error = (prediction - velocity).square() * mask
         return error.sum() / (mask.sum().clamp_min(1) * target.shape[1])
 
-    @torch.inference_mode()
-    def forward(
+    def sample(
         self,
         mu: torch.Tensor,
         mask: torch.Tensor,
@@ -289,3 +288,15 @@ class MatchaFlow(nn.Module):
             )
             x = x + dt * self.estimator(x, mask, mu, time, speaker=speaker)
         return x * mask
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        mu: torch.Tensor,
+        mask: torch.Tensor,
+        n_timesteps: int,
+        temperature: float = 1.0,
+        speaker: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Sample without building a graph (normal inference path)."""
+        return self.sample(mu, mask, n_timesteps, temperature, speaker)