wehub-resource-sync
124 lines
4.1 KiB
Python
124 lines
4.1 KiB
Python
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
|