From c43093611418c887b3dca04d032e1d0370a93335 Mon Sep 17 00:00:00 2001 From: KakaruHayate Date: Fri, 31 Jul 2026 16:42:58 +0800 Subject: [PATCH] refactor: share dual-timestep embedding interpolation --- modules/backbones/lynxnet.py | 22 +++++++++++++--------- modules/backbones/lynxnet2.py | 26 ++++++++++++++++---------- modules/backbones/wavenet.py | 26 +++++++++++++------------- modules/commons/common_layers.py | 21 +++++++++++++++++++++ 4 files changed, 63 insertions(+), 32 deletions(-) diff --git a/modules/backbones/lynxnet.py b/modules/backbones/lynxnet.py index 94d2ad5f..2e839997 100644 --- a/modules/backbones/lynxnet.py +++ b/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) diff --git a/modules/backbones/lynxnet2.py b/modules/backbones/lynxnet2.py index 151fb42f..aae47c1a 100644 --- a/modules/backbones/lynxnet2.py +++ b/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) diff --git a/modules/backbones/wavenet.py b/modules/backbones/wavenet.py index 3e1a769a..634cca74 100644 --- a/modules/backbones/wavenet.py +++ b/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) diff --git a/modules/commons/common_layers.py b/modules/commons/common_layers.py index e37e6de6..10bc6af2 100644 --- a/modules/commons/common_layers.py +++ b/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,