diff --git a/infer/lib/infer_pack/modules.py b/infer/lib/infer_pack/modules.py index 593c301..a94c510 100644 --- a/infer/lib/infer_pack/modules.py +++ b/infer/lib/infer_pack/modules.py @@ -12,7 +12,7 @@ from torch.nn.utils import remove_weight_norm, weight_norm from rvc import utils from rvc.utils import get_padding, call_weight_data_normal_if_Conv -from infer.lib.infer_pack.transforms import piecewise_rational_quadratic_transform +from rvc.transforms import piecewise_rational_quadratic_transform LRELU_SLOPE = 0.1 @@ -583,7 +583,7 @@ class ConvFlow(nn.Module): reverse=False, ): x0, x1 = torch.split(x, [self.half_channels] * 2, 1) - h = self.pre(x0) + h: torch.Tensor = self.pre(x0) h = self.convs(h, x_mask, g=g) h = self.proj(h) * x_mask diff --git a/infer/lib/infer_pack/transforms.py b/rvc/transforms.py similarity index 88% rename from infer/lib/infer_pack/transforms.py rename to rvc/transforms.py index 6d07b3b..679882f 100644 --- a/infer/lib/infer_pack/transforms.py +++ b/rvc/transforms.py @@ -8,13 +8,13 @@ DEFAULT_MIN_DERIVATIVE = 1e-3 def piecewise_rational_quadratic_transform( - inputs, - unnormalized_widths, - unnormalized_heights, - unnormalized_derivatives, - inverse=False, - tails=None, - tail_bound=1.0, + inputs: torch.Tensor, + unnormalized_widths: torch.Tensor, + unnormalized_heights: torch.Tensor, + unnormalized_derivatives: torch.Tensor, + inverse: bool = False, + tails: str | None = None, + tail_bound: float = 1.0, min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_derivative=DEFAULT_MIN_DERIVATIVE, @@ -46,13 +46,13 @@ def searchsorted(bin_locations, inputs, eps=1e-6): def unconstrained_rational_quadratic_spline( - inputs, - unnormalized_widths, - unnormalized_heights, - unnormalized_derivatives, - inverse=False, - tails="linear", - tail_bound=1.0, + inputs: torch.Tensor, + unnormalized_widths: torch.Tensor, + unnormalized_heights: torch.Tensor, + unnormalized_derivatives: torch.Tensor, + inverse: bool = False, + tails: str = "linear", + tail_bound: float = 1.0, min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_derivative=DEFAULT_MIN_DERIVATIVE, @@ -96,11 +96,11 @@ def unconstrained_rational_quadratic_spline( def rational_quadratic_spline( - inputs, - unnormalized_widths, - unnormalized_heights, - unnormalized_derivatives, - inverse=False, + inputs: torch.Tensor, + unnormalized_widths: torch.Tensor, + unnormalized_heights: torch.Tensor, + unnormalized_derivatives: torch.Tensor, + inverse: bool = False, left=0.0, right=1.0, bottom=0.0,