xsl
24c526af68
- 模型 17/17 文件经镜像全量缓存（gated+token，直连本机不通；sfx.md/setup.md/asset-pipeline.md 镜像结论同步更新） - 修复迁移遗留：stable_audio_3/models/ 源码子包被 .gitignore models/ 规则误伤丢失，从上游恢复并加白名单 - radio：arg_muffled/brake_screech 等叙事件落地（接线原有 L111/L154 零改动）+ phone_ring SYN→SA3 升级；dungeon 13 源 A/B 对齐（RMS 偏差 ≤0.3dB）全量同名替换，SYN 原版备份 tools/asset_fix_backup/sa3_ab_20260902/ - 关键玩法语义达标：echo 距离档 0.83<1.43<2.18<2.98s 严格递增；one-shot 真峰值全 ≤-3dBFS（解码回读法） - check_all_assets 8 游戏全绿（radio 254/254、dungeon 303/303）；playtest radio 221/dungeon 1610 全绿
639 lines
31 KiB
Python
639 lines
31 KiB
Python
import torch
|
|
|
|
from torch import nn
|
|
from torch.nn.utils import weight_norm
|
|
from torchaudio import transforms as T
|
|
from einops import rearrange
|
|
|
|
from ..inference.audio_utils import prepare_audio
|
|
from .transformer import TransformerBlock
|
|
|
|
def WNConv1d(*args, **kwargs):
|
|
return weight_norm(nn.Conv1d(*args, **kwargs))
|
|
|
|
def checkpoint(function, *args, **kwargs):
|
|
kwargs.setdefault("use_reentrant", False)
|
|
return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)
|
|
|
|
|
|
class Transpose(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
def forward(self, x, **kwargs):
|
|
return rearrange(x, '... a b -> ... b a')
|
|
|
|
def _zero_pad_modulo_sequence(x, size, dim=-2):
|
|
input_len = x.shape[dim]
|
|
pad_len = (size - input_len % size) % size
|
|
if pad_len > 0:
|
|
pad_shape = list(x.shape)
|
|
pad_shape[dim] = pad_len
|
|
x = torch.cat([x, torch.zeros(pad_shape, device=x.device, dtype=x.dtype)], dim=dim)
|
|
return x
|
|
|
|
class TransformerResamplingBlock(nn.Module):
|
|
def __init__(self, in_channels, out_channels, stride, sliding_window = None, chunk_size = 128, chunk_midpoint_shift = False, type = 'encoder', transformer_depth = 3, checkpointing = False,
|
|
conformer = False, layer_scale = False, dim_heads = 128, differential = True, variable_stride = False, feat_scale = False,
|
|
sinusoidal_blocks = 0, mask_noise = 0, ff_mult = 3, mapping_bias = True, cross_attn = False, dyt = True, conv_mapping = False, freeze_backbone = False, **kwargs):
|
|
super().__init__()
|
|
if type not in ['encoder', 'decoder']:
|
|
raise ValueError(f"Unknown type {type}. Must be 'encoder' or 'decoder'")
|
|
|
|
self.checkpointing = checkpointing
|
|
|
|
transformer_dim = out_channels if type == 'encoder' else in_channels
|
|
transformers = []
|
|
self.in_channels = in_channels
|
|
self.out_channels = out_channels
|
|
self.variable_stride = variable_stride
|
|
self.stride = stride
|
|
self.mapping = WNConv1d(in_channels, out_channels, 3 if conv_mapping else 1, padding = 'same', bias = mapping_bias) if in_channels != out_channels else nn.Identity()
|
|
self.chunk_size = chunk_size
|
|
self.chunk_midpoint_shift = chunk_midpoint_shift
|
|
self.type = type
|
|
self.mask_noise = mask_noise
|
|
self.sliding_window_latents = sliding_window
|
|
|
|
self.sliding_window_seq = self._get_sliding_window_size(sliding_window, stride)
|
|
self.input_seg_size, self.output_seg_size, self.sub_chunk_size = self._get_seg_sizes(stride)
|
|
self.transformer_depth = transformer_depth
|
|
for i in range(transformer_depth):
|
|
sinusoidal = True if ((transformer_depth - i) < sinusoidal_blocks) else False
|
|
transformers.append(TransformerBlock(transformer_dim,
|
|
dim_heads = dim_heads,
|
|
causal = False,
|
|
zero_init_branch_outputs = True if not layer_scale else False,
|
|
norm_type = 'dyt' if dyt else 'rms_norm',
|
|
conformer = conformer,
|
|
layer_scale = layer_scale,
|
|
add_rope = True,
|
|
attn_kwargs={'qk_norm': "dyt" if dyt else "rms", "qk_norm_eps": 1e-3, "differential": differential, "feat_scale": feat_scale},
|
|
ff_kwargs={'mult': ff_mult, 'no_bias': False, "sinusoidal": sinusoidal},
|
|
norm_kwargs = {'eps': 1e-3},
|
|
cross_attend = cross_attn))
|
|
|
|
self.new_tokens = nn.Parameter(1e-5 * torch.randn(1, self.output_seg_size if not self.variable_stride else 1, out_channels if type == 'encoder' else in_channels))
|
|
|
|
self.transformers = nn.ModuleList(transformers)
|
|
|
|
if freeze_backbone:
|
|
for param in self.transformers.parameters():
|
|
param.requires_grad = False
|
|
self.new_tokens.requires_grad = False
|
|
|
|
def _get_sliding_window_size(self, window, stride, prepend_cond_length = 0):
|
|
if window is None:
|
|
return None
|
|
else:
|
|
return [(win * (stride + 1 + prepend_cond_length)) for win in window]
|
|
|
|
def _get_seg_sizes(self, stride, prepend_cond_length = 0):
|
|
sub_chunk_size = stride + 1 + prepend_cond_length
|
|
if self.sliding_window_latents is None:
|
|
assert (self.chunk_size % stride) == 0, f"Stride must fit evenly into chunk size:{self.chunk_size}"
|
|
input_seg_size = stride if self.type == 'encoder' else 1
|
|
output_seg_size = 1 if self.type == 'encoder' else stride
|
|
return input_seg_size, output_seg_size, sub_chunk_size
|
|
|
|
#@torch.compile
|
|
def forward(self, x, stride = None, return_features = False, override_new_tokens = None, prepend_cond = None, cross_attn_cond = None):
|
|
batch_size = x.shape[0]
|
|
input_length = x.shape[-1]
|
|
if return_features:
|
|
features = []
|
|
|
|
if stride == None:
|
|
input_seg_size = self.input_seg_size
|
|
output_seg_size = self.output_seg_size
|
|
sub_chunk_size = self.sub_chunk_size
|
|
sliding_window = self.sliding_window_seq
|
|
else:
|
|
if not self.variable_stride:
|
|
print("cannot override stride if variable_stride is not set")
|
|
prepend_cond_length = prepend_cond.shape[-2] if prepend_cond is not None else 0
|
|
input_seg_size, output_seg_size, sub_chunk_size = self._get_seg_sizes(stride, prepend_cond_length)
|
|
sliding_window = self._get_sliding_window_size(self.sliding_window_latents, stride, prepend_cond_length)
|
|
|
|
if self.type == 'encoder':
|
|
# Pad before mapping so silence zeros get projected through the mapping,
|
|
# rather than inserting raw zeros in the mapped space
|
|
if self.transformer_depth > 0:
|
|
if sliding_window is None:
|
|
pad_modulo = self.chunk_size
|
|
else:
|
|
pad_modulo = input_seg_size
|
|
x = _zero_pad_modulo_sequence(x, pad_modulo, dim=-1)
|
|
x = self.mapping(x)
|
|
|
|
if self.transformer_depth > 0:
|
|
x = rearrange(x, '... a b -> ... b a')
|
|
if return_features:
|
|
features.append(x)
|
|
if self.type != 'encoder':
|
|
if sliding_window is None:
|
|
active_stride = stride if stride is not None else self.stride
|
|
pad_modulo = self.chunk_size // active_stride
|
|
x = _zero_pad_modulo_sequence(x, pad_modulo)
|
|
else:
|
|
x = _zero_pad_modulo_sequence(x, input_seg_size)
|
|
x = rearrange(x, 'b (n c) d -> (b n) c d', c = input_seg_size)
|
|
new_token_seq_dim = -1 if not self.variable_stride else output_seg_size
|
|
new_tokens = self.new_tokens.expand([x.shape[0],new_token_seq_dim,-1])
|
|
if override_new_tokens is not None:
|
|
#print(f"Using override new tokens with shape {override_new_tokens.shape}, new tokens shape {new_tokens.shape}, x shape {x.shape}")
|
|
override_new_tokens = rearrange(override_new_tokens, 'b (n c) d -> (b n) c d', c = output_seg_size)
|
|
new_tokens = new_tokens + override_new_tokens
|
|
elif self.mask_noise > 0:
|
|
new_tokens = new_tokens + torch.randn_like(new_tokens) * self.mask_noise
|
|
x = torch.cat([x,new_tokens], dim = -2)
|
|
if prepend_cond is not None:
|
|
n = x.shape[0] // batch_size
|
|
cond_folded = prepend_cond.unsqueeze(1).expand(batch_size, n, prepend_cond.shape[-2], x.shape[-1]).reshape(n * batch_size, prepend_cond.shape[-2], x.shape[-1])
|
|
x = torch.cat([cond_folded, x], dim=-2)
|
|
x = rearrange(x, '(b n) c d -> b (n c) d', b=batch_size)#.contiguous()
|
|
|
|
# Fold into contiguous chunks if no sliding window
|
|
if sliding_window is None:
|
|
prepend_cond_length = prepend_cond.shape[-2] if prepend_cond is not None else 0
|
|
effective_chunk_size = self.chunk_size + self.chunk_size * (1 + prepend_cond_length) // (stride if stride is not None else self.stride)
|
|
|
|
if sliding_window is None and self.chunk_midpoint_shift:
|
|
split = self.transformer_depth // 2
|
|
shift = effective_chunk_size // 2
|
|
|
|
# First half: standard chunks
|
|
nc = x.shape[1] // effective_chunk_size
|
|
x = rearrange(x, 'b (nc cc) d -> (b nc) cc d', cc=effective_chunk_size)
|
|
cross_attn_first = None
|
|
if cross_attn_cond is not None:
|
|
cross_attn_first = cross_attn_cond.repeat_interleave(nc, dim=0)
|
|
for layer in self.transformers[:split]:
|
|
if self.checkpointing:
|
|
x = checkpoint(layer, x, context = cross_attn_first, self_attention_flash_sliding_window = None)
|
|
else:
|
|
x = layer(x, context = cross_attn_first)
|
|
if return_features:
|
|
features.append(rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size))
|
|
x = rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size)
|
|
|
|
# Second half: shifted chunks with sub-chunk repeat padding
|
|
# shift is always a multiple of sub_chunk_size, so slicing
|
|
# by shift gives whole sub-chunks in their original order
|
|
x = torch.cat([x[:, :shift, :], x, x[:, -shift:, :]], dim=1)
|
|
nc_shifted = x.shape[1] // effective_chunk_size
|
|
x = rearrange(x, 'b (nc cc) d -> (b nc) cc d', cc=effective_chunk_size)
|
|
cross_attn_second = None
|
|
if cross_attn_cond is not None:
|
|
cross_attn_second = cross_attn_cond.repeat_interleave(nc_shifted, dim=0)
|
|
for layer in self.transformers[split:]:
|
|
if self.checkpointing:
|
|
x = checkpoint(layer, x, context = cross_attn_second, self_attention_flash_sliding_window = None)
|
|
else:
|
|
x = layer(x, context = cross_attn_second)
|
|
if return_features:
|
|
feat = rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size)
|
|
features.append(feat[:, shift:-shift, :])
|
|
x = rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size)
|
|
x = x[:, shift:-shift, :]
|
|
else:
|
|
if sliding_window is None:
|
|
x = rearrange(x, 'b (nc cc) d -> (b nc) cc d', cc=effective_chunk_size)
|
|
|
|
for layer in self.transformers:
|
|
if self.checkpointing:
|
|
x = checkpoint(layer, x, context = cross_attn_cond, self_attention_flash_sliding_window = sliding_window)
|
|
else:
|
|
x = layer(x, context = cross_attn_cond, self_attention_flash_sliding_window = sliding_window)
|
|
if return_features:
|
|
features.append(x)
|
|
|
|
# Unfold chunks back to original batch
|
|
if sliding_window is None:
|
|
x = rearrange(x, '(b nc) cc d -> b (nc cc) d', b=batch_size)
|
|
|
|
x = rearrange(x, 'b (n c) d -> (b n) c d', c=sub_chunk_size)
|
|
x = x[:,-output_seg_size:,:]
|
|
x = rearrange(x, '(b n) c d -> b d (n c)', b=batch_size)
|
|
if self.type == 'decoder':
|
|
x = self.mapping(x)
|
|
if return_features:
|
|
return x, features
|
|
else:
|
|
return x
|
|
|
|
|
|
class SAMEEncoder(nn.Module):
|
|
def __init__(self,
|
|
in_channels=2,
|
|
channels=128,
|
|
latent_dim=32,
|
|
c_mults = [1, 2, 4, 8],
|
|
strides = [2, 4, 8, 8],
|
|
transformer_depths = [3,3,3,3],
|
|
sliding_window = None,
|
|
checkpointing = False,
|
|
conformer = False,
|
|
layer_scale = False,
|
|
causal = False,
|
|
differential = True,
|
|
variable_stride = False,
|
|
mask_noise = 0.0,
|
|
conv_mapping = False,
|
|
freeze_backbone = False,
|
|
**kwargs
|
|
):
|
|
super().__init__()
|
|
self.in_channels = in_channels
|
|
self.strides = strides
|
|
|
|
channel_dims = [c * channels for c in c_mults]
|
|
channel_dims = [in_channels] + channel_dims
|
|
|
|
self.depth = len(c_mults)
|
|
|
|
layers = []
|
|
|
|
for i in range(self.depth):
|
|
layers += [TransformerResamplingBlock(in_channels=channel_dims[i], out_channels=channel_dims[i+1], stride=strides[i], transformer_depth = transformer_depths[i],
|
|
sliding_window = sliding_window, checkpointing = checkpointing, conformer = conformer, layer_scale = layer_scale, causal = causal,
|
|
differential = differential, variable_stride = variable_stride, mask_noise = mask_noise, conv_mapping = conv_mapping,
|
|
freeze_backbone = freeze_backbone, **kwargs)]
|
|
|
|
layers += [Transpose(), nn.Linear(channel_dims[-1], latent_dim), Transpose()]
|
|
self.layers = nn.ModuleList(layers)
|
|
|
|
if freeze_backbone:
|
|
for param in self.layers[-2].parameters():
|
|
param.requires_grad = False
|
|
|
|
def forward(self, x, override_stride = None, return_features = False, **kwargs):
|
|
if override_stride != None:
|
|
assert isinstance(override_stride, list), "override_stride must be a list"
|
|
assert len(override_stride) == self.depth, "override_stride must be a list containing strides for every layer"
|
|
for i, layer in enumerate(self.layers):
|
|
if isinstance(layer, TransformerResamplingBlock):
|
|
if override_stride != None:
|
|
stride = override_stride[i]
|
|
else:
|
|
stride = None
|
|
if return_features:
|
|
x, features = layer(x, stride = stride, return_features = True)
|
|
else:
|
|
x = layer(x, stride = stride)
|
|
else:
|
|
x = layer(x)
|
|
if return_features:
|
|
return x, features
|
|
else:
|
|
return x
|
|
|
|
class SAMEDecoder(nn.Module):
|
|
def __init__(self,
|
|
out_channels=2,
|
|
channels=128,
|
|
latent_dim=32,
|
|
c_mults = [1, 2, 4, 8],
|
|
strides = [2, 4, 8, 8],
|
|
transformer_depths = [3,3,3,3],
|
|
sliding_window = None,
|
|
checkpointing = False,
|
|
conformer = False,
|
|
layer_scale = False,
|
|
causal = False,
|
|
differential = True,
|
|
variable_stride = False,
|
|
sinusoidal_blocks = [0,0,0,0],
|
|
mask_noise = 0.0,
|
|
conv_mapping = False,
|
|
freeze_backbone = False,
|
|
**kwargs
|
|
):
|
|
super().__init__()
|
|
|
|
channel_dims = [c * channels for c in c_mults]
|
|
channel_dims = [out_channels] + channel_dims
|
|
|
|
self.depth = len(c_mults)
|
|
|
|
layers = [Transpose(), nn.Linear(latent_dim, channel_dims[-1]), Transpose()]
|
|
|
|
for i in range(self.depth, 0, -1):
|
|
layers += [TransformerResamplingBlock(in_channels=channel_dims[i], out_channels=channel_dims[i-1], stride=strides[i-1], type = 'decoder', transformer_depth = transformer_depths[i-1],
|
|
sliding_window = sliding_window, checkpointing = checkpointing, conformer = conformer, layer_scale = layer_scale, causal = causal, differential = differential,
|
|
variable_stride = variable_stride, sinusoidal_blocks = sinusoidal_blocks[i-1], mask_noise = mask_noise, conv_mapping = conv_mapping,
|
|
freeze_backbone = freeze_backbone, **kwargs)]
|
|
|
|
self.layers = nn.ModuleList(layers)
|
|
|
|
if freeze_backbone:
|
|
for param in self.layers[1].parameters():
|
|
param.requires_grad = False
|
|
|
|
def forward(self, x, override_stride = None, **kwargs):
|
|
if override_stride != None:
|
|
assert isinstance(override_stride, list), "override_stride must be a list"
|
|
assert len(override_stride) == self.depth, "override_stride must be a list containing strides for every layer"
|
|
|
|
transformer_layer_index = 0
|
|
for i, layer in enumerate(self.layers):
|
|
if isinstance(layer, TransformerResamplingBlock):
|
|
if override_stride != None:
|
|
stride = override_stride[transformer_layer_index]
|
|
else:
|
|
stride = None
|
|
x = layer(x, stride = stride)
|
|
transformer_layer_index += 1
|
|
else:
|
|
x = layer(x)
|
|
return x
|
|
|
|
|
|
class AudioAutoencoder(nn.Module):
|
|
def __init__(
|
|
self,
|
|
encoder,
|
|
decoder,
|
|
latent_dim,
|
|
downsampling_ratio,
|
|
sample_rate,
|
|
io_channels=2,
|
|
bottleneck: nn.Module = None,
|
|
pretransform: nn.Module = None,
|
|
in_channels = None,
|
|
out_channels = None,
|
|
soft_clip = False,
|
|
freeze_pretransform = False
|
|
):
|
|
super().__init__()
|
|
self.downsampling_ratio = downsampling_ratio
|
|
self.sample_rate = sample_rate
|
|
|
|
self.latent_dim = latent_dim
|
|
self.io_channels = io_channels
|
|
self.in_channels = io_channels
|
|
self.out_channels = io_channels
|
|
|
|
self.min_length = self.downsampling_ratio
|
|
|
|
if in_channels is not None:
|
|
self.in_channels = in_channels
|
|
|
|
if out_channels is not None:
|
|
self.out_channels = out_channels
|
|
|
|
self.bottleneck = bottleneck
|
|
|
|
self.encoder = encoder
|
|
|
|
self.decoder = decoder
|
|
|
|
self.pretransform = pretransform
|
|
|
|
self.freeze_pretransform = freeze_pretransform
|
|
if self.pretransform is not None:
|
|
if self.freeze_pretransform:
|
|
for p in self.pretransform.parameters():
|
|
p.requires_grad = False
|
|
else:
|
|
for p in self.pretransform.parameters():
|
|
p.requires_grad = True
|
|
|
|
self.soft_clip = soft_clip
|
|
|
|
self.is_discrete = False
|
|
|
|
def encode(self, audio, return_info=False, skip_pretransform=False, iterate_batch=False, return_pretransform = False, **kwargs):
|
|
|
|
info = {}
|
|
if self.pretransform is not None and not skip_pretransform:
|
|
if self.pretransform.enable_grad:
|
|
if iterate_batch:
|
|
audios = []
|
|
for i in range(audio.shape[0]):
|
|
audios.append(self.pretransform.encode(audio[i:i+1]))
|
|
audio = torch.cat(audios, dim=0)
|
|
else:
|
|
audio = self.pretransform.encode(audio)
|
|
else:
|
|
with torch.no_grad():
|
|
if iterate_batch:
|
|
audios = []
|
|
for i in range(audio.shape[0]):
|
|
audios.append(self.pretransform.encode(audio[i:i+1]))
|
|
audio = torch.cat(audios, dim=0)
|
|
else:
|
|
audio = self.pretransform.encode(audio)
|
|
if self.encoder is not None:
|
|
if iterate_batch:
|
|
latents = []
|
|
for i in range(audio.shape[0]):
|
|
latents.append(self.encoder(audio[i:i+1], **kwargs))
|
|
latents = torch.cat(latents, dim=0)
|
|
else:
|
|
latents = self.encoder(audio, **kwargs)
|
|
else:
|
|
latents = audio
|
|
if self.bottleneck is not None:
|
|
# TODO: Add iterate batch logic, needs to merge the info dicts
|
|
latents, bottleneck_info = self.bottleneck.encode(latents, return_info=True, **kwargs)
|
|
|
|
info.update(bottleneck_info)
|
|
|
|
if return_info and return_pretransform:
|
|
return latents, info, audio
|
|
elif return_info:
|
|
return latents, info
|
|
elif return_pretransform:
|
|
return latents, audio
|
|
else:
|
|
return latents
|
|
|
|
def decode(self, latents, iterate_batch=False, return_loss = False, **kwargs):
|
|
if self.bottleneck is not None:
|
|
if iterate_batch:
|
|
decoded = []
|
|
for i in range(latents.shape[0]):
|
|
decoded.append(self.bottleneck.decode(latents[i:i+1]))
|
|
latents = torch.cat(decoded, dim=0)
|
|
else:
|
|
latents = self.bottleneck.decode(latents)
|
|
if iterate_batch:
|
|
decoded = []
|
|
for i in range(latents.shape[0]):
|
|
decoded.append(self.decoder(latents[i:i+1], **kwargs))
|
|
decoded = torch.cat(decoded, dim=0)
|
|
else:
|
|
if return_loss:
|
|
decoded, loss = self.decoder(latents, **kwargs)
|
|
else:
|
|
decoded = self.decoder(latents, **kwargs)
|
|
if self.pretransform is not None:
|
|
if self.pretransform.enable_grad:
|
|
if iterate_batch:
|
|
decodeds = []
|
|
for i in range(decoded.shape[0]):
|
|
decodeds.append(self.pretransform.decode(decoded[i:i+1]))
|
|
decoded = torch.cat(decodeds, dim=0)
|
|
else:
|
|
decoded = self.pretransform.decode(decoded)
|
|
else:
|
|
with torch.no_grad():
|
|
if iterate_batch:
|
|
decodeds = []
|
|
for i in range(latents.shape[0]):
|
|
decodeds.append(self.pretransform.decode(decoded[i:i+1]))
|
|
decoded = torch.cat(decodeds, dim=0)
|
|
else:
|
|
decoded = self.pretransform.decode(decoded)
|
|
|
|
if self.soft_clip:
|
|
decoded = torch.tanh(decoded)
|
|
if return_loss:
|
|
return decoded, loss
|
|
else:
|
|
return decoded
|
|
|
|
|
|
def preprocess_audio_for_encoder(self, audio, in_sr):
|
|
'''
|
|
Preprocess single audio tensor (Channels x Length) to be compatible with the encoder.
|
|
If the model is mono, stereo audio will be converted to mono.
|
|
Audio will be silence-padded to be a multiple of the model's downsampling ratio.
|
|
Audio will be resampled to the model's sample rate.
|
|
The output will have batch size 1 and be shape (1 x Channels x Length)
|
|
'''
|
|
return self.preprocess_audio_list_for_encoder([audio], [in_sr])
|
|
|
|
def preprocess_audio_list_for_encoder(self, audio_list, in_sr_list):
|
|
'''
|
|
Preprocess a [list] of audio (Channels x Length) into a batch tensor to be compatable with the encoder.
|
|
The audio in that list can be of different lengths and channels.
|
|
in_sr can be an integer or list. If it's an integer it will be assumed it is the input sample_rate for every audio.
|
|
All audio will be resampled to the model's sample rate.
|
|
Audio will be silence-padded to the longest length, and further padded to be a multiple of the model's downsampling ratio.
|
|
If the model is mono, all audio will be converted to mono.
|
|
The output will be a tensor of shape (Batch x Channels x Length)
|
|
'''
|
|
batch_size = len(audio_list)
|
|
if isinstance(in_sr_list, int):
|
|
in_sr_list = [in_sr_list]*batch_size
|
|
assert len(in_sr_list) == batch_size, "list of sample rates must be the same length of audio_list"
|
|
new_audio = []
|
|
max_length = 0
|
|
# resample & find the max length
|
|
for i in range(batch_size):
|
|
audio = audio_list[i]
|
|
in_sr = in_sr_list[i]
|
|
if len(audio.shape) == 3 and audio.shape[0] == 1:
|
|
# batchsize 1 was given by accident. Just squeeze it.
|
|
audio = audio.squeeze(0)
|
|
elif len(audio.shape) == 1:
|
|
# Mono signal, channel dimension is missing, unsqueeze it in
|
|
audio = audio.unsqueeze(0)
|
|
assert len(audio.shape)==2, "Audio should be shape (Channels x Length) with no batch dimension"
|
|
# Resample audio
|
|
if in_sr != self.sample_rate:
|
|
resample_tf = T.Resample(in_sr, self.sample_rate).to(audio.device)
|
|
audio = resample_tf(audio)
|
|
new_audio.append(audio)
|
|
if audio.shape[-1] > max_length:
|
|
max_length = audio.shape[-1]
|
|
# Pad every audio to the same length, multiple of model's downsampling ratio
|
|
padded_audio_length = max_length + (self.min_length - (max_length % self.min_length)) % self.min_length
|
|
for i in range(batch_size):
|
|
# Pad it & if necessary, mixdown/duplicate stereo/mono channels to support model
|
|
new_audio[i] = prepare_audio(new_audio[i], in_sr=in_sr, target_sr=in_sr, target_length=padded_audio_length,
|
|
target_channels=self.in_channels, device=new_audio[i].device).squeeze(0)
|
|
# convert to tensor
|
|
return torch.stack(new_audio)
|
|
|
|
def encode_audio(self, audio, chunked=False, overlap=32, chunk_size=128, **kwargs):
|
|
'''
|
|
Encode audios into latents. Audios should already be preprocesed by preprocess_audio_for_encoder.
|
|
If chunked is True, split the audio into chunks of a given maximum size chunk_size, with given overlap.
|
|
Overlap and chunk_size params are both measured in number of latents (not audio samples)
|
|
# and therefore you likely could use the same values with decode_audio.
|
|
A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size.
|
|
Every autoencoder will have a different receptive field size, and thus ideal overlap.
|
|
You can determine it empirically by diffing unchunked vs chunked output and looking at maximum diff.
|
|
The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks.
|
|
Smaller chunk_size uses less memory, but more compute.
|
|
The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version
|
|
For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks
|
|
'''
|
|
samples_per_latent = int(self.downsampling_ratio)
|
|
if not chunked or audio.shape[-1] < chunk_size * samples_per_latent:
|
|
return self.encode(audio, **kwargs)
|
|
|
|
# chunk_size/overlap are in latent units; scale to samples for slicing.
|
|
chunk_size_samples = chunk_size * samples_per_latent
|
|
hop_samples = (chunk_size - overlap) * samples_per_latent
|
|
total_samples = audio.shape[-1]
|
|
|
|
# Anchor the final chunk to the signal end (may overlap more than `overlap`).
|
|
chunk_starts = list(range(0, total_samples - chunk_size_samples + 1, hop_samples))
|
|
if chunk_starts[-1] != total_samples - chunk_size_samples:
|
|
chunk_starts.append(total_samples - chunk_size_samples)
|
|
|
|
encoded_chunks = [self.encode(audio[..., s:s + chunk_size_samples]) for s in chunk_starts]
|
|
|
|
total_latents = total_samples // samples_per_latent
|
|
half_overlap_latents = overlap // 2
|
|
output = audio.new_zeros(*encoded_chunks[0].shape[:-1], total_latents)
|
|
num_chunks = len(chunk_starts)
|
|
|
|
# Trim half the overlap off inner edges (edges facing a neighbour chunk).
|
|
for i, (start_sample, chunk) in enumerate(zip(chunk_starts, encoded_chunks)):
|
|
is_first = i == 0
|
|
is_last = i == num_chunks - 1
|
|
out_start = (total_latents - chunk_size) if is_last else (start_sample // samples_per_latent)
|
|
left = 0 if is_first else half_overlap_latents
|
|
right = chunk_size if is_last else chunk_size - half_overlap_latents
|
|
output[..., out_start + left : out_start + right] = chunk[..., left:right]
|
|
|
|
return output
|
|
|
|
def decode_audio(self, latents, chunked=False, overlap=32, chunk_size=128, **kwargs):
|
|
'''
|
|
Decode latents to audio.
|
|
If chunked is True, split the latents into chunks of a given maximum size chunk_size, with given overlap, both of which are measured in number of latents.
|
|
A overlap of zero will cause discontinuity artefacts. Overlap should be => receptive field size.
|
|
Every autoencoder will have a different receptive field size, and thus ideal overlap.
|
|
You can determine it empirically by diffing unchunked vs chunked audio and looking at maximum diff.
|
|
The final chunk may have a longer overlap in order to keep chunk_size consistent for all chunks.
|
|
Smaller chunk_size uses less memory, but more compute.
|
|
The chunk_size vs memory tradeoff isn't linear, and possibly depends on the GPU and CUDA version
|
|
For example, on a A6000 chunk_size 128 is overall faster than 256 and 512 even though it has more chunks
|
|
'''
|
|
if not chunked or latents.shape[-1] < chunk_size:
|
|
return self.decode(latents, **kwargs)
|
|
|
|
# chunk_size/overlap are in latent units — same as `latents`, no scaling.
|
|
samples_per_latent = int(self.downsampling_ratio)
|
|
hop_latents = chunk_size - overlap
|
|
total_latents = latents.shape[-1]
|
|
|
|
# Anchor the final chunk to the signal end (may overlap more than `overlap`).
|
|
chunk_starts = list(range(0, total_latents - chunk_size + 1, hop_latents))
|
|
if chunk_starts[-1] != total_latents - chunk_size:
|
|
chunk_starts.append(total_latents - chunk_size)
|
|
|
|
decoded_chunks = [self.decode(latents[..., s:s + chunk_size]) for s in chunk_starts]
|
|
|
|
total_samples = total_latents * samples_per_latent
|
|
chunk_size_samples = chunk_size * samples_per_latent
|
|
half_overlap_samples = (overlap // 2) * samples_per_latent
|
|
output = latents.new_zeros(*decoded_chunks[0].shape[:-1], total_samples)
|
|
num_chunks = len(chunk_starts)
|
|
|
|
# Trim half the overlap off inner edges (edges facing a neighbour chunk).
|
|
for i, (start_latent, chunk) in enumerate(zip(chunk_starts, decoded_chunks)):
|
|
is_first = i == 0
|
|
is_last = i == num_chunks - 1
|
|
out_start = (total_samples - chunk_size_samples) if is_last else (start_latent * samples_per_latent)
|
|
left = 0 if is_first else half_overlap_samples
|
|
right = chunk_size_samples if is_last else chunk_size_samples - half_overlap_samples
|
|
output[..., out_start + left : out_start + right] = chunk[..., left:right]
|
|
|
|
return output
|
|
|