Support dual-timestep reflow

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

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

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

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

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,19 @@ 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)
+            if step.dim() == 2:
+                step = step.unsqueeze(1)
 
         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)

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,18 @@ 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:
+            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)

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,23 @@ 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)
+            if step.dim() == 2:
+                step = step.unsqueeze(1)
+
         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))

modules/commons/common_layers.py

@@ -422,6 +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)
-        emb = x.unsqueeze(-1) * emb.unsqueeze(0)
+        emb = x.unsqueeze(-1) * emb
         emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
         return emb

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]

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)