Yueqian LinandSy03
c111f3a2d6
Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Sy03 <1370724210@qq.com>
506 lines
21 KiB
Python
506 lines
21 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
#
|
|
# Adapted from the AuK reference release (MIT), which in turn adapts NVIDIA
|
|
# BigVGAN (MIT), jik876/hifi-gan (MIT), junjun3518/alias-free-torch
|
|
# (Apache-2.0), EdwardDixon/snake (MIT) and adefossez/julius (MIT).
|
|
"""BigVGAN-flow VAE for AuK: 24 kHz mono waveform <-> 50 Hz, 64-channel latents.
|
|
|
|
Inference only. The reference module also carries a residual coupling flow, a KL
|
|
term and a discriminator; all three are training-only, so the released
|
|
``flow.*`` tensors are skipped (see :meth:`AuKVAE.load_weights`).
|
|
|
|
Parameter and buffer names follow the reference release, so the published
|
|
``vae.safetensors`` loads without a key map. That is also why the convolutions
|
|
are built with the legacy ``weight_norm`` parametrization (``weight_g`` /
|
|
``weight_v``): :meth:`remove_weight_norm` folds them into plain weights once
|
|
the checkpoint is in.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
import math
|
|
import warnings
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from safetensors.torch import load_file
|
|
from torch import nn
|
|
from torch.nn.utils import remove_weight_norm as _fold_weight_norm
|
|
from torch.nn.utils import weight_norm as _apply_weight_norm
|
|
from vllm.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
__all__ = ["AuKVAE"]
|
|
|
|
# Checkpoint prefix owned by the training-only coupling flow.
|
|
TRAINING_ONLY_PREFIX = "flow."
|
|
|
|
|
|
def _weight_norm(module: nn.Module) -> nn.Module:
|
|
"""Legacy ``weight_g``/``weight_v`` split, which is how the release stores these tensors."""
|
|
with warnings.catch_warnings():
|
|
warnings.filterwarnings("ignore", message=".*weight_norm.*deprecated.*", category=FutureWarning)
|
|
return _apply_weight_norm(module)
|
|
|
|
|
|
def _same_padding(kernel_size: int, dilation: int = 1) -> int:
|
|
return int((kernel_size * dilation - dilation) / 2)
|
|
|
|
|
|
def _kaiser_sinc_filter(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
|
|
"""Kaiser-windowed sinc low-pass, returned as a ``[1, 1, kernel_size]`` conv kernel."""
|
|
half_size = kernel_size // 2
|
|
attenuation = 2.285 * (half_size - 1) * math.pi * (4 * half_width) + 7.95
|
|
if attenuation > 50.0:
|
|
beta = 0.1102 * (attenuation - 8.7)
|
|
elif attenuation >= 21.0:
|
|
beta = 0.5842 * (attenuation - 21.0) ** 0.4 + 0.07886 * (attenuation - 21.0)
|
|
else:
|
|
beta = 0.0
|
|
window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
|
|
if kernel_size % 2 == 0:
|
|
time = torch.arange(-half_size, half_size) + 0.5
|
|
else:
|
|
time = torch.arange(kernel_size) - half_size
|
|
if cutoff == 0:
|
|
return torch.zeros_like(time).view(1, 1, kernel_size)
|
|
taps = 2 * cutoff * window * torch.sinc(2 * cutoff * time)
|
|
# Normalize to unit sum so a constant input passes through unscaled.
|
|
return (taps / taps.sum()).view(1, 1, kernel_size)
|
|
|
|
|
|
class Conv(nn.Conv1d):
|
|
"""1-D convolution over ``[B, C, T]``; ``causal`` swaps centred padding for a left pad."""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
kernel_size: int = 1,
|
|
stride: int = 1,
|
|
dilation: int = 1,
|
|
bias: bool = True,
|
|
causal: bool = False,
|
|
) -> None:
|
|
super().__init__(
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size,
|
|
stride=stride,
|
|
padding=0 if causal else _same_padding(kernel_size, dilation),
|
|
dilation=dilation,
|
|
bias=bias,
|
|
)
|
|
self.causal = causal
|
|
self.left_padding = dilation * (kernel_size - 1) if causal else 0
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
if self.causal:
|
|
x = F.pad(x, (self.left_padding, 0))
|
|
return super().forward(x)
|
|
|
|
|
|
class ConvTranspose(nn.ConvTranspose1d):
|
|
"""Transposed 1-D convolution; ``causal`` drops the trailing ``stride`` samples instead of padding."""
|
|
|
|
def __init__(
|
|
self, in_channels: int, out_channels: int, kernel_size: int, stride: int, causal: bool = False
|
|
) -> None:
|
|
if causal and kernel_size != 2 * stride:
|
|
raise ValueError(f"causal ConvTranspose needs kernel_size == 2 * stride, got {kernel_size} and {stride}")
|
|
super().__init__(
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size,
|
|
stride=stride,
|
|
padding=0 if causal else (kernel_size - stride) // 2,
|
|
)
|
|
self.causal = causal
|
|
self.trim = stride
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
x = super().forward(x)
|
|
return x[:, :, : -self.trim] if self.causal else x
|
|
|
|
|
|
class NormConv(nn.Module):
|
|
"""Weight-normalized conv held under ``layer``; the attribute name is load-bearing for the checkpoint."""
|
|
|
|
def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 1, stride: int = 1) -> None:
|
|
super().__init__()
|
|
self.layer = _weight_norm(
|
|
nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=(kernel_size - 1) // 2)
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.layer(x)
|
|
|
|
|
|
class ResStack(nn.Module):
|
|
"""Stack of dilated residual conv pairs. The leaky slope is torch's 0.01 default, as upstream."""
|
|
|
|
def __init__(self, channels: int, kernel_size: int = 3, base: int = 3, nums: int = 4) -> None:
|
|
super().__init__()
|
|
self.layers = nn.ModuleList(
|
|
[
|
|
nn.Sequential(
|
|
nn.LeakyReLU(),
|
|
_weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=base**i, padding=base**i)),
|
|
nn.LeakyReLU(),
|
|
_weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=1, padding=1)),
|
|
)
|
|
for i in range(nums)
|
|
]
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
for layer in self.layers:
|
|
x = x + layer(x)
|
|
return x
|
|
|
|
|
|
class Encoder(nn.Module):
|
|
"""Strided conv encoder over ``[B, 1, T]``, emitting ``[B, 2 * latent_dim, T / hop]`` mean/log-std pairs."""
|
|
|
|
def __init__(
|
|
self,
|
|
latent_dim: int,
|
|
channels: tuple[int, ...],
|
|
down_sample_factors: tuple[int, ...],
|
|
in_channels: int = 1,
|
|
proj_kernel_size: int = 3,
|
|
stack_kernel_size: int = 3,
|
|
stack_dilation_base: int = 2,
|
|
stacks: int = 6,
|
|
) -> None:
|
|
super().__init__()
|
|
layers: list[nn.Module] = [
|
|
NormConv(in_channels, channels[0], kernel_size=proj_kernel_size),
|
|
nn.LeakyReLU(0.2, True),
|
|
]
|
|
for in_c, out_c, factor in zip(channels[:-1], channels[1:], down_sample_factors, strict=True):
|
|
layers += [
|
|
NormConv(in_c, out_c, kernel_size=factor * 2, stride=factor),
|
|
ResStack(out_c, stack_kernel_size, stack_dilation_base, stacks),
|
|
nn.LeakyReLU(0.2, True),
|
|
]
|
|
layers.append(NormConv(channels[-1], latent_dim * 2, kernel_size=proj_kernel_size))
|
|
self.generator = nn.Sequential(*layers)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.generator(x)
|
|
|
|
|
|
class SnakeBeta(nn.Module):
|
|
"""``x + sin^2(alpha * x) / beta`` with per-channel learned alpha and beta (log-scale in AuK)."""
|
|
|
|
def __init__(self, channels: int, alpha_logscale: bool = False) -> None:
|
|
super().__init__()
|
|
self.alpha_logscale = alpha_logscale
|
|
init = torch.zeros(channels) if alpha_logscale else torch.ones(channels)
|
|
self.alpha = nn.Parameter(init.clone())
|
|
self.beta = nn.Parameter(init.clone())
|
|
self.eps = 1e-9
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
|
|
beta = self.beta.unsqueeze(0).unsqueeze(-1)
|
|
if self.alpha_logscale:
|
|
alpha = torch.exp(alpha)
|
|
beta = torch.exp(beta)
|
|
return x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2)
|
|
|
|
|
|
class LowPass(nn.Module):
|
|
"""Fixed FIR low-pass with replicate padding, applied per channel."""
|
|
|
|
def __init__(
|
|
self,
|
|
cutoff: float = 0.5,
|
|
half_width: float = 0.6,
|
|
stride: int = 1,
|
|
kernel_size: int = 12,
|
|
causal: bool = False,
|
|
) -> None:
|
|
super().__init__()
|
|
if not 0.0 <= cutoff <= 0.5:
|
|
raise ValueError(f"cutoff must lie in [0, 0.5], got {cutoff}")
|
|
if causal:
|
|
self.pad_left, self.pad_right = kernel_size - 1, 0
|
|
else:
|
|
even = kernel_size % 2 == 0
|
|
self.pad_left = kernel_size // 2 - int(even)
|
|
self.pad_right = kernel_size // 2
|
|
self.stride = stride
|
|
self.register_buffer("filter", _kaiser_sinc_filter(cutoff, half_width, kernel_size))
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
channels = x.size(1)
|
|
x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate")
|
|
return F.conv1d(x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels)
|
|
|
|
|
|
class Upsample(nn.Module):
|
|
"""Band-limited interpolation by ``ratio``. Always non-causal, matching how AuK builds it."""
|
|
|
|
def __init__(self, ratio: int = 2, kernel_size: int = 12) -> None:
|
|
super().__init__()
|
|
self.ratio = ratio
|
|
self.pad = kernel_size // ratio - 1
|
|
self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2
|
|
self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2
|
|
self.register_buffer("filter", _kaiser_sinc_filter(0.5 / ratio, 0.6 / ratio, kernel_size))
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
channels = x.size(1)
|
|
x = F.pad(x, (self.pad, self.pad), mode="replicate")
|
|
x = self.ratio * F.conv_transpose1d(x, self.filter.expand(channels, -1, -1), stride=self.ratio, groups=channels)
|
|
return x[..., self.pad_left : -self.pad_right]
|
|
|
|
|
|
class Downsample(nn.Module):
|
|
"""Decimation by ``ratio`` behind a low-pass; ``lowpass`` holds the filter buffer."""
|
|
|
|
def __init__(self, ratio: int = 2, kernel_size: int = 12, causal: bool = False) -> None:
|
|
super().__init__()
|
|
self.lowpass = LowPass(
|
|
cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=kernel_size, causal=causal
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.lowpass(x)
|
|
|
|
|
|
class AliasFreeActivation(nn.Module):
|
|
"""Oversample, apply the pointwise activation, decimate, so snake harmonics do not fold back."""
|
|
|
|
def __init__(
|
|
self,
|
|
activation: nn.Module,
|
|
up_ratio: int = 2,
|
|
down_ratio: int = 2,
|
|
up_kernel_size: int = 12,
|
|
down_kernel_size: int = 12,
|
|
causal: bool = False,
|
|
) -> None:
|
|
super().__init__()
|
|
self.act = activation
|
|
self.upsample = Upsample(up_ratio, up_kernel_size)
|
|
self.downsample = Downsample(down_ratio, down_kernel_size, causal=causal)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.downsample(self.act(self.upsample(x)))
|
|
|
|
|
|
class AmpBlock(nn.Module):
|
|
"""Anti-aliased multi-periodicity residual block: alternating alias-free snake and dilated convs."""
|
|
|
|
def __init__(
|
|
self,
|
|
channels: int,
|
|
kernel_size: int,
|
|
dilations: tuple[int, ...],
|
|
snake_logscale: bool = True,
|
|
causal: bool = True,
|
|
act_causal: bool = False,
|
|
) -> None:
|
|
super().__init__()
|
|
self.convs1 = nn.ModuleList(
|
|
[_weight_norm(Conv(channels, channels, kernel_size, 1, dilation=d, causal=causal)) for d in dilations]
|
|
)
|
|
self.convs2 = nn.ModuleList(
|
|
[_weight_norm(Conv(channels, channels, kernel_size, 1, dilation=1, causal=causal)) for _ in dilations]
|
|
)
|
|
self.activations = nn.ModuleList(
|
|
[
|
|
AliasFreeActivation(SnakeBeta(channels, alpha_logscale=snake_logscale), causal=act_causal)
|
|
for _ in range(len(self.convs1) + len(self.convs2))
|
|
]
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
acts1, acts2 = self.activations[::2], self.activations[1::2]
|
|
for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2, strict=True):
|
|
x = conv2(act2(conv1(act1(x)))) + x
|
|
return x
|
|
|
|
|
|
class AuKVAE(nn.Module):
|
|
"""AuK's audio codec: :meth:`encode` waveform to normalized latents, :meth:`decode` back.
|
|
|
|
Build from the released ``model.vae.model_init_kwargs`` with :meth:`from_config`, then
|
|
:meth:`load_weights`. Latents are normalized by the checkpoint's global statistics, which is
|
|
the space the rectified-flow transformer works in.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
upsample_rates: tuple[int, ...] = (5, 4, 3, 2, 2, 2),
|
|
upsample_kernel_sizes: tuple[int, ...] = (10, 8, 6, 4, 4, 4),
|
|
upsample_initial_channel: int = 1536,
|
|
resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11),
|
|
resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)),
|
|
downsample_rates: tuple[int, ...] = (2, 2, 2, 3, 4, 5),
|
|
downsample_channels: tuple[int, ...] = (12, 24, 48, 96, 192, 384, 768),
|
|
snake_logscale: bool = True,
|
|
latent_dim: int = 64,
|
|
use_vae: bool = True,
|
|
causal: bool = True,
|
|
act_causal: bool = True,
|
|
sample_rate: int = 24000,
|
|
) -> None:
|
|
super().__init__()
|
|
if not use_vae:
|
|
raise ValueError("AuKVAE requires use_vae=True: the encoder head emits mean/log-std pairs")
|
|
self.latent_dim = latent_dim
|
|
self.sample_rate = sample_rate
|
|
self.hop_size = math.prod(downsample_rates)
|
|
self.num_kernels = len(resblock_kernel_sizes)
|
|
self.num_upsamples = len(upsample_rates)
|
|
self._norm_folded = False
|
|
|
|
self.register_buffer("global_mean", torch.zeros(latent_dim, dtype=torch.float32))
|
|
self.register_buffer("global_log_std", torch.ones(latent_dim, dtype=torch.float32))
|
|
|
|
self.audio_encoder = Encoder(
|
|
latent_dim=latent_dim,
|
|
channels=tuple(downsample_channels),
|
|
down_sample_factors=tuple(downsample_rates),
|
|
)
|
|
|
|
self.conv_pre = _weight_norm(Conv(latent_dim, upsample_initial_channel, 7, 1, causal=False))
|
|
self.ups = nn.ModuleList()
|
|
for i, (rate, kernel) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)):
|
|
# One-element inner list: the checkpoint keys are ``ups.<i>.0.*``.
|
|
self.ups.append(
|
|
nn.ModuleList(
|
|
[
|
|
_weight_norm(
|
|
ConvTranspose(
|
|
upsample_initial_channel // (2**i),
|
|
upsample_initial_channel // (2 ** (i + 1)),
|
|
kernel,
|
|
rate,
|
|
causal=causal,
|
|
)
|
|
)
|
|
]
|
|
)
|
|
)
|
|
|
|
self.resblocks = nn.ModuleList()
|
|
channels = upsample_initial_channel
|
|
for i in range(self.num_upsamples):
|
|
channels = upsample_initial_channel // (2 ** (i + 1))
|
|
for kernel, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True):
|
|
self.resblocks.append(
|
|
AmpBlock(
|
|
channels,
|
|
kernel,
|
|
tuple(dilations),
|
|
snake_logscale=snake_logscale,
|
|
causal=causal,
|
|
act_causal=act_causal,
|
|
)
|
|
)
|
|
|
|
self.activation_post = AliasFreeActivation(
|
|
SnakeBeta(channels, alpha_logscale=snake_logscale), causal=act_causal
|
|
)
|
|
self.conv_post = _weight_norm(Conv(channels, 1, 7, 1, causal=causal, bias=False))
|
|
|
|
@classmethod
|
|
def from_config(cls, cfg: dict) -> AuKVAE:
|
|
"""Build from the released ``model_init_kwargs`` mapping, ignoring training-only entries."""
|
|
accepted = set(inspect.signature(cls.__init__).parameters) - {"self"}
|
|
return cls(**{key: value for key, value in cfg.items() if key in accepted})
|
|
|
|
def load_weights(
|
|
self, path: str | Path, *, device: str | torch.device | None = None, fold_weight_norm: bool = True
|
|
) -> tuple[list[str], list[str]]:
|
|
"""Load a released ``vae.safetensors``, then fold weight norm away. Returns (missing, unexpected).
|
|
|
|
``device`` defaults to where the module already lives, which is the placement that matters:
|
|
the fold reduces ``weight_v`` to a norm, and a CPU reduction lands up to one ULP away from the
|
|
accelerator's. The encoder amplifies its activations roughly fiftyfold, so that one ULP grows
|
|
into a ~4e-4 relative drift by the latent head. Move the module first, then load.
|
|
"""
|
|
target = self.global_mean.device if device is None else torch.device(device)
|
|
missing, unexpected = self.load_state_dict(load_file(str(path), device=str(target)), strict=False)
|
|
missing, unexpected = list(missing), list(unexpected)
|
|
stray = [key for key in missing + unexpected if not key.startswith(TRAINING_ONLY_PREFIX)]
|
|
if stray:
|
|
raise ValueError(f"AuK VAE weights at {path} do not match this module: {sorted(stray)[:8]}")
|
|
logger.info(
|
|
"Loaded AuK VAE from %s (skipped %d training-only %s* tensors)", path, len(unexpected), TRAINING_ONLY_PREFIX
|
|
)
|
|
if fold_weight_norm:
|
|
self.remove_weight_norm()
|
|
return missing, unexpected
|
|
|
|
def remove_weight_norm(self) -> None:
|
|
"""Fold ``weight_g``/``weight_v`` into plain weights, on whichever device the module is on.
|
|
|
|
Idempotent, and meant to run once the checkpoint is in. See :meth:`load_weights` for why the
|
|
device this runs on is a numeric decision rather than a convenience.
|
|
"""
|
|
if self._norm_folded:
|
|
return
|
|
folded = 0
|
|
for module in self.modules():
|
|
if hasattr(module, "weight_g"):
|
|
_fold_weight_norm(module)
|
|
folded += 1
|
|
self._norm_folded = True
|
|
logger.debug("Folded weight norm on %d AuK VAE convolutions", folded)
|
|
|
|
def encode(
|
|
self, wav: torch.Tensor, *, sample: bool = False, generator: torch.Generator | None = None
|
|
) -> torch.Tensor:
|
|
"""Encode a mono ``[1, T]`` waveform at :attr:`sample_rate` into ``[1, T / hop_size, latent_dim]``.
|
|
|
|
``sample=False`` takes the posterior mean, which is what a reproducible pipeline wants;
|
|
``sample=True`` draws ``mean + randn * exp(log_std)`` from ``generator``. Latents come back
|
|
normalized by the checkpoint's global statistics.
|
|
"""
|
|
if wav.dim() != 2 or wav.size(0) != 1:
|
|
raise ValueError(f"encode expects a mono [1, T] waveform, got {tuple(wav.shape)}")
|
|
# Statistics stay in fp32 even under an ambient autocast, as in the reference.
|
|
with torch.autocast(device_type=wav.device.type, enabled=False):
|
|
mean, log_std = self.audio_encoder(wav.unsqueeze(1)).chunk(2, dim=1)
|
|
if sample:
|
|
latents = mean + self._noise_like(mean, generator) * torch.exp(log_std)
|
|
else:
|
|
latents = mean
|
|
latents = latents.transpose(1, 2).float()
|
|
return (latents - self.global_mean.float()) / torch.sqrt(self.global_log_std.float())
|
|
|
|
def decode(self, latents: torch.Tensor) -> torch.Tensor:
|
|
"""Decode normalized ``[B, Np, latent_dim]`` latents into a ``[B, Np * hop_size]`` waveform in [-1, 1]."""
|
|
if latents.dim() != 3 or latents.size(-1) != self.latent_dim:
|
|
raise ValueError(f"decode expects [B, Np, {self.latent_dim}] latents, got {tuple(latents.shape)}")
|
|
x = latents.float() * torch.sqrt(self.global_log_std.float()) + self.global_mean.float()
|
|
x = self.conv_pre(x.permute(0, 2, 1))
|
|
for i in range(self.num_upsamples):
|
|
for up in self.ups[i]:
|
|
x = up(x)
|
|
first = i * self.num_kernels
|
|
residual = self.resblocks[first](x)
|
|
for j in range(1, self.num_kernels):
|
|
residual = residual + self.resblocks[first + j](x)
|
|
x = residual / self.num_kernels
|
|
x = self.conv_post(self.activation_post(x))
|
|
return torch.clamp(x, min=-1.0, max=1.0).squeeze(1)
|
|
|
|
@staticmethod
|
|
def _noise_like(reference: torch.Tensor, generator: torch.Generator | None) -> torch.Tensor:
|
|
"""Draw standard normal noise shaped like ``reference``; ``generator`` decides which RNG stream."""
|
|
device = reference.device if generator is None else generator.device
|
|
noise = torch.randn(reference.shape, generator=generator, dtype=reference.dtype, device=device)
|
|
return noise.to(reference.device)
|