refactor: deduplicate dual-timestep embedding interpolation

modules/backbones/lynxnet.py

@@ -6,7 +6,13 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, Transpose, AdamWConv1d
+from modules.commons.common_layers import (
+    AdamWConv1d,
+    SinusoidalPosEmb,
+    SwiGLU,
+    Transpose,
+    interpolate_dual_timestep_embedding,
+)
 from modules.commons.common_layers import KaimingNormalConv1d as Conv1d
 from utils.hparams import hparams
 
@@ -127,14 +133,12 @@ def forward(self, spec, diffusion_step, cond, diffusion_step_2=None, mask=None):
         if not self.strong_cond:
             x = F.gelu(x)
 
-        if mask is not None:
-            step = torch.cat((diffusion_step, diffusion_step_2), dim=0)
-            step = self.diffusion_embedding(step)
-            step, step_2 = torch.split(step, x.shape[0], dim=0) #[B, 1, C]
-            mask = mask.to(x).unsqueeze(-1) # [B, T, 1]
-            step = step + (step_2 - step) * mask
-        else:
-            step = self.diffusion_embedding(diffusion_step)
+        step = interpolate_dual_timestep_embedding(
+            self.diffusion_embedding,
+            diffusion_step,
+            diffusion_step_2,
+            mask,
+        )
 
         for layer in self.residual_layers:
             x = layer(x, cond, step.transpose(1, 2), front_cond_inject=self.strong_cond)

modules/backbones/lynxnet2.py

@@ -2,7 +2,14 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from modules.commons.common_layers import SinusoidalPosEmb, SwiGLU, ATanGLU, Transpose, AdamWLinear
+from modules.commons.common_layers import (
+    ATanGLU,
+    AdamWLinear,
+    SinusoidalPosEmb,
+    SwiGLU,
+    Transpose,
+    interpolate_dual_timestep_embedding,
+)
 from utils.hparams import hparams
 
 
@@ -95,15 +102,14 @@ def forward(self, spec, diffusion_step, cond, diffusion_step_2=None, mask=None):
             x = x + self.conditioner_projection(cond).transpose(1, 2)
         else:
             x = x + self.conditioner_projection(cond.transpose(1, 2))
-
-        if mask is not None:
-            step = torch.cat((diffusion_step, diffusion_step_2), dim=0)
-            step = self.diffusion_embedding(step)
-            step, step_2 = torch.split(step, x.shape[0], dim=0) #[B, 1, C]
-            mask = mask.to(x).unsqueeze(-1) # [B, T, 1]
-            x = x + step + (step_2 - step) * mask
-        else:
-            x = x + self.diffusion_embedding(diffusion_step)
+
+        step = interpolate_dual_timestep_embedding(
+            self.diffusion_embedding,
+            diffusion_step,
+            diffusion_step_2,
+            mask,
+        )
+        x = x + step
 
         for layer in self.residual_layers:
             x = layer(x)

modules/backbones/wavenet.py

@@ -5,7 +5,11 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from modules.commons.common_layers import SinusoidalPosEmb, AdamWConv1d
+from modules.commons.common_layers import (
+    AdamWConv1d,
+    SinusoidalPosEmb,
+    interpolate_dual_timestep_embedding,
+)
 from modules.commons.common_layers import KaimingNormalConv1d as Conv1d
 from utils.hparams import hparams
 
@@ -84,18 +88,14 @@ def forward(self, spec, diffusion_step, cond, diffusion_step_2=None, mask=None):
         x = self.input_projection(x)  # [B, C, T]
 
         x = F.relu(x)
-
-        if mask is not None:
-            step = torch.cat((diffusion_step, diffusion_step_2), dim=0)
-            step = self.diffusion_embedding(step)
-            step = self.mlp(step)
-            step, step_2 = torch.split(step, x.shape[0], dim=0) #[B, 1, C]
-            mask = mask.to(x).unsqueeze(-1) # [B, T, 1]
-            step = step + (step_2 - step) * mask
-        else:
-            step = self.diffusion_embedding(diffusion_step)
-            step = self.mlp(step)
-
+
+        step = interpolate_dual_timestep_embedding(
+            lambda timestep: self.mlp(self.diffusion_embedding(timestep)),
+            diffusion_step,
+            diffusion_step_2,
+            mask,
+        )
+
         skip = []
         for layer in self.residual_layers:
             x, skip_connection = layer(x, cond, step)

modules/commons/common_layers.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import math
+from typing import Callable
+
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -197,6 +199,25 @@ def forward(self, x):
         return x.transpose(*self.dims)
 
 
+def interpolate_dual_timestep_embedding(
+        embedding: nn.Module | Callable[[torch.Tensor], torch.Tensor],
+        timestep: torch.Tensor,
+        timestep_2: torch.Tensor | None = None,
+        mask: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Embed one or two timesteps and interpolate the second on masked frames."""
+    if mask is None:
+        return embedding(timestep)
+    if timestep_2 is None:
+        raise ValueError('timestep_2 is required when mask is provided')
+
+    batch_size = timestep.shape[0]
+    embedded = embedding(torch.cat((timestep, timestep_2), dim=0))
+    step, step_2 = torch.split(embedded, batch_size, dim=0)
+    frame_mask = mask.to(step).unsqueeze(-1)
+    return step + (step_2 - step) * frame_mask
+
+
 class Mixed_LayerNorm(nn.Module):
     def __init__(
             self,