chore: import upstream snapshot with attribution

2026-07-13 12:35:17 +08:00
commit 344816a5d8
136 changed files with 25044 additions and 0 deletions
@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+from typing import List, Tuple
+
+import torch
+
+from modules.core import (
+    RectifiedFlow, PitchRectifiedFlow, MultiVarianceRectifiedFlow
+)
+
+
+class RectifiedFlowONNX(RectifiedFlow):
+    @property
+    def backbone(self):
+        return self.velocity_fn
+
+    # We give up the setter for the property `backbone` because this will cause TorchScript to fail
+    # @backbone.setter
+    @torch.jit.unused
+    def set_backbone(self, value):
+        self.velocity_fn = value
+
+    def sample_euler(self, x, t, dt: float, cond):
+        x += self.velocity_fn(x, t * self.time_scale_factor, cond) * dt
+        return x
+
+    def norm_spec(self, x):
+        k = (self.spec_max - self.spec_min) / 2.
+        b = (self.spec_max + self.spec_min) / 2.
+        return (x - b) / k
+
+    def denorm_spec(self, x):
+        k = (self.spec_max - self.spec_min) / 2.
+        b = (self.spec_max + self.spec_min) / 2.
+        return x * k + b
+
+    def forward(self, condition, x_end=None, depth=None, steps: int = 10):
+        condition = condition.transpose(1, 2)  # [1, T, H] => [1, H, T]
+        device = condition.device
+        n_frames = condition.shape[2]
+        noise = torch.randn((1, self.num_feats, self.out_dims, n_frames), device=device)
+        if x_end is None:
+            t_start = 0.
+            x = noise
+        else:
+            t_start = torch.max(1 - depth, torch.tensor(self.t_start, dtype=torch.float32, device=device))
+            x_end = self.norm_spec(x_end).transpose(-2, -1)
+            if self.num_feats == 1:
+                x_end = x_end[:, None, :, :]
+            if t_start <= 0.:
+                x = noise
+            elif t_start >= 1.:
+                x = x_end
+            else:
+                x = t_start * x_end + (1 - t_start) * noise
+
+        t_width = 1. - t_start
+        if t_width >= 0.:
+            dt = t_width / max(1, steps)
+            for t in torch.arange(steps, dtype=torch.long, device=device)[:, None].float() * dt + t_start:
+                x = self.sample_euler(x, t, dt, condition)
+
+        if self.num_feats == 1:
+            x = x.squeeze(1).permute(0, 2, 1)  # [B, 1, M, T] => [B, T, M]
+        else:
+            x = x.permute(0, 1, 3, 2)  # [B, F, M, T] => [B, F, T, M]
+        x = self.denorm_spec(x)
+        return x
+
+
+class PitchRectifiedFlowONNX(RectifiedFlowONNX, PitchRectifiedFlow):
+    def __init__(self, vmin: float, vmax: float,
+                 cmin: float, cmax: float, repeat_bins,
+                 time_scale_factor=1000,
+                 backbone_type=None, backbone_args=None):
+        self.vmin = vmin
+        self.vmax = vmax
+        self.cmin = cmin
+        self.cmax = cmax
+        super(PitchRectifiedFlow, self).__init__(
+            vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
+            time_scale_factor=time_scale_factor,
+            backbone_type=backbone_type, backbone_args=backbone_args
+        )
+
+    def clamp_spec(self, x):
+        return x.clamp(min=self.cmin, max=self.cmax)
+
+    def denorm_spec(self, x):
+        d = (self.spec_max - self.spec_min) / 2.
+        m = (self.spec_max + self.spec_min) / 2.
+        x = x * d + m
+        x = x.mean(dim=-1)
+        return x
+
+
+class MultiVarianceRectifiedFlowONNX(RectifiedFlowONNX, MultiVarianceRectifiedFlow):
+    def __init__(
+            self, ranges: List[Tuple[float, float]],
+            clamps: List[Tuple[float | None, float | None] | None],
+            repeat_bins, time_scale_factor=1000,
+            backbone_type=None, backbone_args=None
+    ):
+        assert len(ranges) == len(clamps)
+        self.clamps = clamps
+        vmin = [r[0] for r in ranges]
+        vmax = [r[1] for r in ranges]
+        if len(vmin) == 1:
+            vmin = vmin[0]
+        if len(vmax) == 1:
+            vmax = vmax[0]
+        super(MultiVarianceRectifiedFlow, self).__init__(
+            vmin=vmin, vmax=vmax, repeat_bins=repeat_bins,
+            time_scale_factor=time_scale_factor,
+            backbone_type=backbone_type, backbone_args=backbone_args
+        )
+
+    def denorm_spec(self, x):
+        d = (self.spec_max - self.spec_min) / 2.
+        m = (self.spec_max + self.spec_min) / 2.
+        x = x * d + m
+        x = x.mean(dim=-1)
+        return x