From 72981e4d6a290228380d1fe9b31d6c22fd37becc Mon Sep 17 00:00:00 2001 From: yxlllc Date: Thu, 30 Jul 2026 22:51:09 +0800 Subject: [PATCH 1/2] support dual-timestep reflow --- configs/acoustic.yaml | 1 + configs/templates/config_acoustic.yaml | 1 + configs/templates/config_variance.yaml | 1 + configs/variance.yaml | 1 + modules/backbones/lynxnet.py | 14 +++++++++++--- modules/backbones/lynxnet2.py | 12 ++++++++++-- modules/backbones/wavenet.py | 20 +++++++++++++++----- modules/commons/common_layers.py | 4 +++- modules/core/reflow.py | 25 ++++++++++++++++++------- modules/losses/reflow_loss.py | 4 ++-- 10 files changed, 63 insertions(+), 20 deletions(-) diff --git a/configs/acoustic.yaml b/configs/acoustic.yaml index fad75600e..293a0695f 100644 --- a/configs/acoustic.yaml +++ b/configs/acoustic.yaml @@ -61,6 +61,7 @@ use_key_shift_embed: false use_speed_embed: false diffusion_type: reflow +use_dual_timestep: true time_scale_factor: 1000 timesteps: 1000 max_beta: 0.02 diff --git a/configs/templates/config_acoustic.yaml b/configs/templates/config_acoustic.yaml index e344fb450..6425ecc32 100644 --- a/configs/templates/config_acoustic.yaml +++ b/configs/templates/config_acoustic.yaml @@ -72,6 +72,7 @@ augmentation_args: # diffusion and shallow diffusion diffusion_type: reflow +use_dual_timestep: true enc_ffn_kernel_size: 3 use_rope: true rope_interleaved: false diff --git a/configs/templates/config_variance.yaml b/configs/templates/config_variance.yaml index 116154ac7..d9676177d 100644 --- a/configs/templates/config_variance.yaml +++ b/configs/templates/config_variance.yaml @@ -90,6 +90,7 @@ glide_types: [up, down] glide_embed_scale: 11.313708498984760 # sqrt(128) diffusion_type: reflow +use_dual_timestep: true pitch_prediction_args: pitd_norm_min: -8.0 diff --git a/configs/variance.yaml b/configs/variance.yaml index d4e203670..17fb1c89c 100644 --- a/configs/variance.yaml +++ b/configs/variance.yaml @@ -106,6 +106,7 @@ lambda_pitch_loss: 1.0 lambda_var_loss: 1.0 diffusion_type: reflow # ddpm +use_dual_timestep: true time_scale_factor: 1000 schedule_type: 'linear' K_step: 1000 diff --git a/modules/backbones/lynxnet.py b/modules/backbones/lynxnet.py index 9529d1efe..94d2ad5f1 100644 --- a/modules/backbones/lynxnet.py +++ b/modules/backbones/lynxnet.py @@ -2,6 +2,7 @@ # https://github.com/CNChTu/Diffusion-SVC/blob/v2.0_dev/diffusion/naive_v2/model_conformer_naive.py # https://github.com/CNChTu/Diffusion-SVC/blob/v2.0_dev/diffusion/naive_v2/naive_v2_diff.py +import torch import torch.nn as nn import torch.nn.functional as F @@ -109,7 +110,7 @@ def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansio self.strong_cond = strong_cond nn.init.zeros_(self.output_projection.weight) - def forward(self, spec, diffusion_step, cond): + def forward(self, spec, diffusion_step, cond, diffusion_step_2=None, mask=None): """ :param spec: [B, F, M, T] :param diffusion_step: [B, 1] @@ -126,10 +127,17 @@ def forward(self, spec, diffusion_step, cond): if not self.strong_cond: x = F.gelu(x) - diffusion_step = self.diffusion_embedding(diffusion_step).unsqueeze(-1) + 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) for layer in self.residual_layers: - x = layer(x, cond, diffusion_step, front_cond_inject=self.strong_cond) + x = layer(x, cond, step.transpose(1, 2), front_cond_inject=self.strong_cond) # post-norm x = self.norm(x.transpose(1, 2)).transpose(1, 2) diff --git a/modules/backbones/lynxnet2.py b/modules/backbones/lynxnet2.py index 6e55d5f87..151fb42f9 100644 --- a/modules/backbones/lynxnet2.py +++ b/modules/backbones/lynxnet2.py @@ -77,7 +77,7 @@ def __init__(self, in_dims, n_feats, *, num_layers=6, num_channels=512, expansio nn.init.kaiming_normal_(self.conditioner_projection.weight) nn.init.zeros_(self.output_projection.weight) - def forward(self, spec, diffusion_step, cond): + def forward(self, spec, diffusion_step, cond, diffusion_step_2=None, mask=None): """ :param spec: [B, F, M, T] :param diffusion_step: [B, 1] @@ -95,7 +95,15 @@ def forward(self, spec, diffusion_step, cond): x = x + self.conditioner_projection(cond).transpose(1, 2) else: x = x + self.conditioner_projection(cond.transpose(1, 2)) - x = x + self.diffusion_embedding(diffusion_step).unsqueeze(1) + + 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) for layer in self.residual_layers: x = layer(x) diff --git a/modules/backbones/wavenet.py b/modules/backbones/wavenet.py index 77ccc6430..3e1a769a2 100644 --- a/modules/backbones/wavenet.py +++ b/modules/backbones/wavenet.py @@ -26,7 +26,7 @@ def __init__(self, encoder_hidden, residual_channels, dilation): self.output_projection = nn.Conv1d(residual_channels, 2 * residual_channels, 1) def forward(self, x, conditioner, diffusion_step): - diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1) + diffusion_step = self.diffusion_projection(diffusion_step).transpose(1, 2) conditioner = self.conditioner_projection(conditioner) y = x + diffusion_step @@ -67,7 +67,7 @@ def __init__(self, in_dims, n_feats, *, num_layers=20, num_channels=256, dilatio self.output_projection = AdamWConv1d(num_channels, in_dims * n_feats, 1) nn.init.zeros_(self.output_projection.weight) - def forward(self, spec, diffusion_step, cond): + def forward(self, spec, diffusion_step, cond, diffusion_step_2=None, mask=None): """ :param spec: [B, F, M, T] :param diffusion_step: [B, 1] @@ -84,11 +84,21 @@ def forward(self, spec, diffusion_step, cond): x = self.input_projection(x) # [B, C, T] x = F.relu(x) - diffusion_step = self.diffusion_embedding(diffusion_step) - diffusion_step = self.mlp(diffusion_step) + + 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) + skip = [] for layer in self.residual_layers: - x, skip_connection = layer(x, cond, diffusion_step) + x, skip_connection = layer(x, cond, step) skip.append(skip_connection) x = torch.sum(torch.stack(skip), dim=0) / sqrt(len(self.residual_layers)) diff --git a/modules/commons/common_layers.py b/modules/commons/common_layers.py index 4da65693a..e37e6de65 100644 --- a/modules/commons/common_layers.py +++ b/modules/commons/common_layers.py @@ -422,6 +422,8 @@ def forward(self, x): half_dim = self.dim // 2 emb = math.log(10000) / (half_dim - 1) emb = torch.exp(torch.arange(half_dim, device=device) * -emb) - emb = x.unsqueeze(-1) * emb.unsqueeze(0) + if x.dim() == 1: + x = x.unsqueeze(-1) + emb = x.unsqueeze(-1) * emb emb = torch.cat((emb.sin(), emb.cos()), dim=-1) return emb diff --git a/modules/core/reflow.py b/modules/core/reflow.py index c2f9f1ef7..2a7bd8441 100644 --- a/modules/core/reflow.py +++ b/modules/core/reflow.py @@ -18,6 +18,7 @@ def __init__(self, out_dims, num_feats=1, t_start=0., time_scale_factor=1000, self.velocity_fn: nn.Module = build_backbone(out_dims, num_feats, backbone_type, backbone_args) self.out_dims = out_dims self.num_feats = num_feats + self.use_dual_timestep = hparams.get('use_dual_timestep', False) self.use_shallow_diffusion = hparams.get('use_shallow_diffusion', False) if self.use_shallow_diffusion: assert 0. <= t_start <= 1., 'T_start should be in [0, 1].' @@ -33,24 +34,34 @@ def __init__(self, out_dims, num_feats=1, t_start=0., time_scale_factor=1000, self.register_buffer('spec_min', spec_min, persistent=False) self.register_buffer('spec_max', spec_max, persistent=False) - def p_losses(self, x_end, t, cond): + def p_losses(self, x_end, t1, cond, t2=None, mask=None): + t = t1 if mask is None else t1 + (t2 - t1) * mask x_start = torch.randn_like(x_end) - x_t = x_start + t[:, None, None, None] * (x_end - x_start) - v_pred = self.velocity_fn(x_t, t * self.time_scale_factor, cond) + x_t = x_start + t[:, None, None,:] * (x_end - x_start) + s1 = t1 * self.time_scale_factor + s2 = None if t2 is None else t2 * self.time_scale_factor + v_pred = self.velocity_fn(x_t, s1, cond, s2, mask) - return v_pred, x_end - x_start + return v_pred, x_end - x_start, t def forward(self, condition, gt_spec=None, src_spec=None, infer=True): cond = condition.transpose(1, 2) - b, device = condition.shape[0], condition.device + b, _, n_frames = cond.shape + device = condition.device if not infer: # gt_spec: [B, T, M] or [B, F, T, M] spec = self.norm_spec(gt_spec).transpose(-2, -1) # [B, M, T] or [B, F, M, T] if self.num_feats == 1: spec = spec[:, None, :, :] # [B, F=1, M, T] - t = self.t_start + (1.0 - self.t_start) * torch.rand((b,), device=device) - v_pred, v_gt = self.p_losses(spec, t, cond=cond) + t1 = self.t_start + (1.0 - self.t_start) * torch.rand((b, 1), device=device) + if self.use_dual_timestep: + t2 = self.t_start + (1.0 - self.t_start) * torch.rand((b, 1), device=device) + mask = (torch.rand(b, n_frames, device=device) < 0.25).float() + else: + t2 = None + mask = None + v_pred, v_gt, t = self.p_losses(spec, t1, cond=cond, t2=t2, mask=mask) return v_pred, v_gt, t else: # src_spec: [B, T, M] or [B, F, T, M] diff --git a/modules/losses/reflow_loss.py b/modules/losses/reflow_loss.py index 4917dce2e..640062552 100644 --- a/modules/losses/reflow_loss.py +++ b/modules/losses/reflow_loss.py @@ -31,7 +31,7 @@ def get_weights(t): weights = 0.398942 / t / (1 - t) * torch.exp( -0.5 * torch.log(t / (1 - t)) ** 2 ) + eps - return weights[:, None, None, None] + return weights[:, None, None, :] def _forward(self, v_pred, v_gt, t=None): if self.log_norm: @@ -43,7 +43,7 @@ def forward(self, v_pred: Tensor, v_gt: Tensor, t: Tensor, non_padding: Tensor = """ :param v_pred: [B, 1, M, T] :param v_gt: [B, 1, M, T] - :param t: [B,] + :param t: [B, 1] or [B, T] :param non_padding: [B, T, M] """ v_pred, v_gt = self._mask_non_padding(v_pred, v_gt, non_padding) From 4764cab403906e6876d464ad5a54d9c6013e4285 Mon Sep 17 00:00:00 2001 From: Kakaru <97896816+KakaruHayate@users.noreply.github.com> Date: Sun, 2 Aug 2026 21:13:27 +0800 Subject: [PATCH 2/2] fix: preserve timestep rank in sinusoidal embedding (#316) --- modules/backbones/lynxnet.py | 2 ++ modules/backbones/lynxnet2.py | 5 ++++- modules/backbones/wavenet.py | 2 ++ modules/commons/common_layers.py | 2 -- 4 files changed, 8 insertions(+), 3 deletions(-) diff --git a/modules/backbones/lynxnet.py b/modules/backbones/lynxnet.py index 94d2ad5f1..80201fe33 100644 --- a/modules/backbones/lynxnet.py +++ b/modules/backbones/lynxnet.py @@ -135,6 +135,8 @@ def forward(self, spec, diffusion_step, cond, diffusion_step_2=None, mask=None): step = step + (step_2 - step) * mask else: step = self.diffusion_embedding(diffusion_step) + if step.dim() == 2: + step = step.unsqueeze(1) 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 151fb42f9..0bab77f57 100644 --- a/modules/backbones/lynxnet2.py +++ b/modules/backbones/lynxnet2.py @@ -103,7 +103,10 @@ def forward(self, spec, diffusion_step, cond, diffusion_step_2=None, mask=None): 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 = self.diffusion_embedding(diffusion_step) + if step.dim() == 2: + step = step.unsqueeze(1) + 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 3e1a769a2..b31f341ad 100644 --- a/modules/backbones/wavenet.py +++ b/modules/backbones/wavenet.py @@ -95,6 +95,8 @@ def forward(self, spec, diffusion_step, cond, diffusion_step_2=None, mask=None): else: step = self.diffusion_embedding(diffusion_step) step = self.mlp(step) + if step.dim() == 2: + step = step.unsqueeze(1) skip = [] for layer in self.residual_layers: diff --git a/modules/commons/common_layers.py b/modules/commons/common_layers.py index e37e6de65..33429a291 100644 --- a/modules/commons/common_layers.py +++ b/modules/commons/common_layers.py @@ -422,8 +422,6 @@ def forward(self, x): half_dim = self.dim // 2 emb = math.log(10000) / (half_dim - 1) emb = torch.exp(torch.arange(half_dim, device=device) * -emb) - if x.dim() == 1: - x = x.unsqueeze(-1) emb = x.unsqueeze(-1) * emb emb = torch.cat((emb.sin(), emb.cos()), dim=-1) return emb