all inference code
This commit is contained in:
sean
@@ -0,0 +1,2 @@
|
||||
from .music_dcae import MusicDCAE
|
||||
from .music_log_mel import LogMelSpectrogram, get_mel_transform
|
||||
@@ -0,0 +1,92 @@
|
||||
import torch
|
||||
from torch.autograd import grad
|
||||
|
||||
|
||||
class Balancer:
|
||||
"""
|
||||
Balancer for dynamically re-weighting multiple losses based on gradient norms.
|
||||
|
||||
Args:
|
||||
weights (dict): Predefined weights for each loss.
|
||||
Example: {"mse_loss": 1.0, "adv_loss": 1.0}
|
||||
ema_decay (float): Decay factor for exponential moving average (default: 0.99).
|
||||
epsilon (float): Small value to avoid division by zero (default: 1e-8).
|
||||
"""
|
||||
def __init__(self, weights, ema_decay=0.99, epsilon=1e-8):
|
||||
self.weights = weights
|
||||
self.ema_decay = ema_decay
|
||||
self.epsilon = epsilon
|
||||
self.ema_values = {key: 0.0 for key in weights} # Initialize EMA for each loss
|
||||
|
||||
def forward(self, losses, grad_inputs):
|
||||
"""
|
||||
Re-weight the input losses based on gradient norms and return a combined loss.
|
||||
|
||||
Args:
|
||||
losses (dict): Dictionary of losses with names as keys and loss tensors as values.
|
||||
Example: {"mse_loss": mse_loss, "adv_loss": adv_loss}
|
||||
grad_inputs (dict): Dictionary of inputs for autograd.grad corresponding to each loss.
|
||||
Example: {"mse_loss": recon_mels, "adv_loss": recon_mels}
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Combined weighted loss.
|
||||
"""
|
||||
# Validate inputs
|
||||
if set(losses.keys()) != set(grad_inputs.keys()):
|
||||
raise ValueError("Keys of losses and grad_inputs must match.")
|
||||
|
||||
norm_values = {}
|
||||
|
||||
# Compute gradient norms for each loss
|
||||
for name, loss in losses.items():
|
||||
loss_grad, = grad(loss.mean(), [grad_inputs[name]], create_graph=True)
|
||||
dims = tuple(range(1, loss_grad.ndim)) # Exclude batch dimension
|
||||
grad_norm = torch.linalg.vector_norm(loss_grad, ord=2, dim=dims).mean()
|
||||
|
||||
# Update EMA for the gradient norm
|
||||
if self.ema_values[name] == 0.0:
|
||||
self.ema_values[name] = grad_norm.item()
|
||||
else:
|
||||
self.ema_values[name] = (
|
||||
self.ema_values[name] * self.ema_decay + grad_norm.item() * (1 - self.ema_decay)
|
||||
)
|
||||
|
||||
# Normalize gradient norm
|
||||
norm_values[name] = grad_norm / (self.ema_values[name] + self.epsilon)
|
||||
|
||||
# Compute dynamic weights
|
||||
total_norm = sum(norm_values.values())
|
||||
dynamic_weights = {name: norm / total_norm for name, norm in norm_values.items()}
|
||||
|
||||
# Combine losses with dynamic weights
|
||||
loss = 0.0
|
||||
log_weights = {}
|
||||
for name in losses:
|
||||
loss = loss + self.weights[name] * dynamic_weights[name] * losses[name]
|
||||
log_weights[f"{name}_weight"] = dynamic_weights[name]
|
||||
return loss, log_weights
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example usage
|
||||
mel_real = torch.randn(1, 80, 10)
|
||||
generator = torch.nn.Linear(10, 10)
|
||||
recon_mels = generator(mel_real)
|
||||
discriminator = torch.nn.Linear(10, 1)
|
||||
disc_out = discriminator(recon_mels)
|
||||
|
||||
mse_loss = torch.nn.functional.mse_loss(recon_mels, mel_real).mean()
|
||||
adv_loss = torch.nn.functional.softplus(-disc_out).mean()
|
||||
losses = {"mse_loss": mse_loss, "adv_loss": adv_loss}
|
||||
grad_inputs = {"mse_loss": recon_mels, "adv_loss": recon_mels}
|
||||
print("losses", losses)
|
||||
# Define predefined weights for each loss
|
||||
weights = {"mse_loss": 1.0, "adv_loss": 1.0}
|
||||
|
||||
# Initialize balancer
|
||||
balancer = Balancer(weights)
|
||||
|
||||
# Forward pass
|
||||
loss, log_weights = balancer.forward(losses, grad_inputs)
|
||||
print("Combined Loss:", loss)
|
||||
print("Dynamic Weights:", log_weights)
|
||||
@@ -0,0 +1,69 @@
|
||||
{
|
||||
"_class_name": "AutoencoderDC",
|
||||
"_diffusers_version": "0.32.1",
|
||||
"_name_or_path": "mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers",
|
||||
"attention_head_dim": 32,
|
||||
"decoder_act_fns": "silu",
|
||||
"decoder_block_out_channels": [
|
||||
128,
|
||||
256,
|
||||
512,
|
||||
1024
|
||||
],
|
||||
"decoder_block_types": [
|
||||
"ResBlock",
|
||||
"ResBlock",
|
||||
"ResBlock",
|
||||
"EfficientViTBlock"
|
||||
],
|
||||
"decoder_layers_per_block": [
|
||||
3,
|
||||
3,
|
||||
3,
|
||||
3
|
||||
],
|
||||
"decoder_norm_types": "rms_norm",
|
||||
"decoder_qkv_multiscales": [
|
||||
[],
|
||||
[],
|
||||
[
|
||||
5
|
||||
],
|
||||
[
|
||||
5
|
||||
]
|
||||
],
|
||||
"downsample_block_type": "Conv",
|
||||
"encoder_block_out_channels": [
|
||||
128,
|
||||
256,
|
||||
512,
|
||||
1024
|
||||
],
|
||||
"encoder_block_types": [
|
||||
"ResBlock",
|
||||
"ResBlock",
|
||||
"ResBlock",
|
||||
"EfficientViTBlock"
|
||||
],
|
||||
"encoder_layers_per_block": [
|
||||
2,
|
||||
2,
|
||||
3,
|
||||
3
|
||||
],
|
||||
"encoder_qkv_multiscales": [
|
||||
[],
|
||||
[],
|
||||
[
|
||||
5
|
||||
],
|
||||
[
|
||||
5
|
||||
]
|
||||
],
|
||||
"in_channels": 2,
|
||||
"latent_channels": 8,
|
||||
"scaling_factor": 0.41407,
|
||||
"upsample_block_type": "interpolate"
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
# All rights reserved.
|
||||
#
|
||||
# This source code is licensed under the license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
"""Torch distributed utilities."""
|
||||
|
||||
import typing as tp
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def rank():
|
||||
if torch.distributed.is_initialized():
|
||||
return torch.distributed.get_rank()
|
||||
else:
|
||||
return 0
|
||||
|
||||
|
||||
def world_size():
|
||||
if torch.distributed.is_initialized():
|
||||
return torch.distributed.get_world_size()
|
||||
else:
|
||||
return 1
|
||||
|
||||
|
||||
def is_distributed():
|
||||
return world_size() > 1
|
||||
|
||||
|
||||
def all_reduce(tensor: torch.Tensor, op=torch.distributed.ReduceOp.SUM):
|
||||
if is_distributed():
|
||||
return torch.distributed.all_reduce(tensor, op)
|
||||
|
||||
|
||||
def _is_complex_or_float(tensor):
|
||||
return torch.is_floating_point(tensor) or torch.is_complex(tensor)
|
||||
|
||||
|
||||
def _check_number_of_params(params: tp.List[torch.Tensor]):
|
||||
# utility function to check that the number of params in all workers is the same,
|
||||
# and thus avoid a deadlock with distributed all reduce.
|
||||
if not is_distributed() or not params:
|
||||
return
|
||||
tensor = torch.tensor([len(params)], device=params[0].device, dtype=torch.long)
|
||||
all_reduce(tensor)
|
||||
if tensor.item() != len(params) * world_size():
|
||||
# If not all the workers have the same number, for at least one of them,
|
||||
# this inequality will be verified.
|
||||
raise RuntimeError(f"Mismatch in number of params: ours is {len(params)}, "
|
||||
"at least one worker has a different one.")
|
||||
|
||||
|
||||
def broadcast_tensors(tensors: tp.Iterable[torch.Tensor], src: int = 0):
|
||||
"""Broadcast the tensors from the given parameters to all workers.
|
||||
This can be used to ensure that all workers have the same model to start with.
|
||||
"""
|
||||
if not is_distributed():
|
||||
return
|
||||
tensors = [tensor for tensor in tensors if _is_complex_or_float(tensor)]
|
||||
_check_number_of_params(tensors)
|
||||
handles = []
|
||||
for tensor in tensors:
|
||||
handle = torch.distributed.broadcast(tensor.data, src=src, async_op=True)
|
||||
handles.append(handle)
|
||||
for handle in handles:
|
||||
handle.wait()
|
||||
|
||||
|
||||
def sync_buffer(buffers, average=True):
|
||||
"""
|
||||
Sync grad for buffers. If average is False, broadcast instead of averaging.
|
||||
"""
|
||||
if not is_distributed():
|
||||
return
|
||||
handles = []
|
||||
for buffer in buffers:
|
||||
if torch.is_floating_point(buffer.data):
|
||||
if average:
|
||||
handle = torch.distributed.all_reduce(
|
||||
buffer.data, op=torch.distributed.ReduceOp.SUM, async_op=True)
|
||||
else:
|
||||
handle = torch.distributed.broadcast(
|
||||
buffer.data, src=0, async_op=True)
|
||||
handles.append((buffer, handle))
|
||||
for buffer, handle in handles:
|
||||
handle.wait()
|
||||
if average:
|
||||
buffer.data /= world_size
|
||||
|
||||
|
||||
def sync_grad(params):
|
||||
"""
|
||||
Simpler alternative to DistributedDataParallel, that doesn't rely
|
||||
on any black magic. For simple models it can also be as fast.
|
||||
Just call this on your model parameters after the call to backward!
|
||||
"""
|
||||
if not is_distributed():
|
||||
return
|
||||
handles = []
|
||||
for p in params:
|
||||
if p.grad is not None:
|
||||
handle = torch.distributed.all_reduce(
|
||||
p.grad.data, op=torch.distributed.ReduceOp.SUM, async_op=True)
|
||||
handles.append((p, handle))
|
||||
for p, handle in handles:
|
||||
handle.wait()
|
||||
p.grad.data /= world_size()
|
||||
|
||||
|
||||
def average_metrics(metrics: tp.Dict[str, float], count=1.):
|
||||
"""Average a dictionary of metrics across all workers, using the optional
|
||||
`count` as unnormalized weight.
|
||||
"""
|
||||
if not is_distributed():
|
||||
return metrics
|
||||
keys, values = zip(*metrics.items())
|
||||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
tensor = torch.tensor(list(values) + [1], device=device, dtype=torch.float32)
|
||||
tensor *= count
|
||||
all_reduce(tensor)
|
||||
averaged = (tensor[:-1] / tensor[-1]).cpu().tolist()
|
||||
return dict(zip(keys, averaged))
|
||||
@@ -0,0 +1,78 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from diffusers import AutoencoderDC
|
||||
import json
|
||||
|
||||
|
||||
DEFAULT_CONFIG_PATH = "/root/sag_train/music_dcae/config_f32c32_large.json"
|
||||
|
||||
class MusicDCAE(nn.Module):
|
||||
def __init__(self, config_path=DEFAULT_CONFIG_PATH):
|
||||
super(MusicDCAE, self).__init__()
|
||||
with open(config_path) as f:
|
||||
config = json.load(f)
|
||||
self.dcae = AutoencoderDC(**config)
|
||||
|
||||
def encode(self, x):
|
||||
return self.dcae.encode(x).latent
|
||||
|
||||
def decode(self, latent):
|
||||
sample = self.dcae.decode(latent).sample
|
||||
return sample
|
||||
|
||||
def forward(self, x):
|
||||
sample = self.dcae(x).sample
|
||||
return sample
|
||||
|
||||
def return_middle_layers(self):
|
||||
last_down_block = self.dcae.encoder.down_blocks[-1]
|
||||
encoder_conv_out = self.dcae.encoder.conv_out
|
||||
decoder_conv_in = self.dcae.decoder.conv_in
|
||||
decoder_up_blocks = self.dcae.decoder.up_blocks[0]
|
||||
middle_layers = [last_down_block, encoder_conv_out, decoder_conv_in, decoder_up_blocks]
|
||||
return middle_layers
|
||||
|
||||
def return_head_layers(self):
|
||||
decoder_up_blocks = self.dcae.decoder.up_blocks[-1]
|
||||
conv_out = self.dcae.decoder.conv_out
|
||||
head_layers = [decoder_up_blocks, conv_out]
|
||||
return head_layers
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
model = MusicDCAE("/root/sag_train/music_dcae/config_f8c8_large.json")
|
||||
|
||||
x = torch.randn(1, 2, 128, 1024)
|
||||
# mask = None
|
||||
# if mask is None:
|
||||
# mask = torch.ones(x.shape[0], 1, x.shape[2], x.shape[3]).to(x.device)
|
||||
# # N x 1024
|
||||
# elif len(mask.shape) == 2:
|
||||
# mask = mask.unsqueeze(1).unsqueeze(1).float()
|
||||
# mask = mask.repeat(1, 1, x.shape[2], 1)
|
||||
latent = model.encode(x)
|
||||
print("latent shape: ", latent.shape)
|
||||
y = model(x)
|
||||
print("y", y.shape)
|
||||
total_params = sum(p.numel() for p in model.parameters())
|
||||
print(f"模型参数总数: {total_params / 1e6:.2f}M")
|
||||
|
||||
# middle_layers = model.return_middle_layers()
|
||||
# middle_params_count = 0
|
||||
# for layer in middle_layers:
|
||||
# for name, param in layer.named_parameters():
|
||||
# layer_param_count = param.numel()
|
||||
# middle_params_count += layer_param_count
|
||||
# print(f"{name}: {param.shape}, 参数量: {layer_param_count/1e6:.2f}M")
|
||||
|
||||
# print(f"中间层总参数量: {middle_params_count/1e6:.2f}M")
|
||||
|
||||
# head_layers = model.return_head_layers()
|
||||
# head_params_count = 0
|
||||
# for layer in head_layers:
|
||||
# for name, param in layer.named_parameters():
|
||||
# layer_param_count = param.numel()
|
||||
# head_params_count += layer_param_count
|
||||
# print(f"{name}: {param.shape}, 参数量: {layer_param_count/1e6:.2f}M")
|
||||
|
||||
# print(f"头部层总参数量: {head_params_count/1e6:.2f}M")
|
||||
@@ -0,0 +1,155 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from diffusers import AutoencoderDC
|
||||
import torchaudio
|
||||
import torchvision.transforms as transforms
|
||||
import torchaudio
|
||||
|
||||
try:
|
||||
from .music_log_mel import get_mel_transform
|
||||
from .music_vocoder import ADaMoSHiFiGANV1
|
||||
except ImportError:
|
||||
from music_log_mel import get_mel_transform
|
||||
from music_vocoder import ADaMoSHiFiGANV1
|
||||
|
||||
import os
|
||||
|
||||
root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
DEFAULT_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_dcae_f8c8")
|
||||
VOCODER_PRETRAINED_PATH = os.path.join(root_dir, "checkpoints", "music_vocoder.pt")
|
||||
|
||||
|
||||
class MusicDCAE(nn.Module):
|
||||
def __init__(self, pretrained_path=DEFAULT_PRETRAINED_PATH, encoder_only=False, source_sample_rate=None):
|
||||
super(MusicDCAE, self).__init__()
|
||||
dcae = AutoencoderDC.from_pretrained(pretrained_path)
|
||||
self.encoder_only = encoder_only
|
||||
|
||||
self.mel_transform = get_mel_transform()
|
||||
if encoder_only:
|
||||
self.encoder = dcae.encoder
|
||||
else:
|
||||
self.encoder = dcae.encoder
|
||||
self.decoder = dcae.decoder
|
||||
self.vocoder = ADaMoSHiFiGANV1(VOCODER_PRETRAINED_PATH).eval()
|
||||
|
||||
if source_sample_rate is None:
|
||||
source_sample_rate = 48000
|
||||
|
||||
self.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)
|
||||
|
||||
self.transform = transforms.Compose([
|
||||
transforms.Normalize(0.5, 0.5),
|
||||
])
|
||||
self.min_mel_value = -11.0
|
||||
self.max_mel_value = 3.0
|
||||
self.audio_chunk_size = int(round((1024 * 512 / 44100 * 48000)))
|
||||
self.mel_chunk_size = 1024
|
||||
self.time_dimention_multiple = 8
|
||||
self.latent_chunk_size = self.mel_chunk_size // self.time_dimention_multiple
|
||||
self.scale_factor = 0.1786
|
||||
self.shift_factor = -1.9091
|
||||
|
||||
def load_audio(self, audio_path):
|
||||
audio, sr = torchaudio.load(audio_path)
|
||||
return audio, sr
|
||||
|
||||
def forward_mel(self, audios):
|
||||
mels = []
|
||||
for i in range(len(audios)):
|
||||
image = self.mel_transform(audios[i])
|
||||
mels.append(image)
|
||||
mels = torch.stack(mels)
|
||||
return mels
|
||||
|
||||
@torch.no_grad()
|
||||
def encode(self, audios, audio_lengths=None, sr=None):
|
||||
if audio_lengths is None:
|
||||
audio_lengths = torch.tensor([audios.shape[2]] * audios.shape[0])
|
||||
audio_lengths = audio_lengths.to(audios.device)
|
||||
|
||||
# audios: N x 2 x T, 48kHz
|
||||
device = audios.device
|
||||
dtype = audios.dtype
|
||||
|
||||
if sr is None:
|
||||
sr = 48000
|
||||
resampler = self.resampler
|
||||
else:
|
||||
resampler = torchaudio.transforms.Resample(sr, 44100).to(device).to(dtype)
|
||||
|
||||
audio = resampler(audios)
|
||||
|
||||
max_audio_len = audio.shape[-1]
|
||||
if max_audio_len % (8 * 512) != 0:
|
||||
audio = torch.nn.functional.pad(audio, (0, 8 * 512 - max_audio_len % (8 * 512)))
|
||||
|
||||
mels = self.forward_mel(audio)
|
||||
mels = (mels - self.min_mel_value) / (self.max_mel_value - self.min_mel_value)
|
||||
mels = self.transform(mels)
|
||||
latents = []
|
||||
for mel in mels:
|
||||
latent = self.encoder(mel.unsqueeze(0))
|
||||
latents.append(latent)
|
||||
latents = torch.cat(latents, dim=0)
|
||||
latent_lengths = (audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple).long()
|
||||
latents = (latents - self.shift_factor) * self.scale_factor
|
||||
return latents, latent_lengths
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(self, latents, audio_lengths=None, sr=None):
|
||||
latents = latents / self.scale_factor + self.shift_factor
|
||||
|
||||
mels = []
|
||||
|
||||
for latent in latents:
|
||||
mel = self.decoder(latent.unsqueeze(0))
|
||||
mels.append(mel)
|
||||
mels = torch.cat(mels, dim=0)
|
||||
|
||||
mels = mels * 0.5 + 0.5
|
||||
mels = mels * (self.max_mel_value - self.min_mel_value) + self.min_mel_value
|
||||
bsz, channels, num_mel, mel_width = mels.shape
|
||||
pred_wavs = []
|
||||
for i in range(bsz):
|
||||
mel = mels[i]
|
||||
wav = self.vocoder.decode(mel).squeeze(1)
|
||||
pred_wavs.append(wav)
|
||||
|
||||
pred_wavs = torch.stack(pred_wavs)
|
||||
|
||||
if sr is not None:
|
||||
resampler = torchaudio.transforms.Resample(44100, sr).to(latents.device).to(latents.dtype)
|
||||
pred_wavs = [resampler(wav) for wav in pred_wavs]
|
||||
else:
|
||||
sr = 44100
|
||||
if audio_lengths is not None:
|
||||
pred_wavs = [wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)]
|
||||
return sr, pred_wavs
|
||||
|
||||
def forward(self, audios, audio_lengths=None, sr=None):
|
||||
latents, latent_lengths = self.encode(audios=audios, audio_lengths=audio_lengths, sr=sr)
|
||||
sr, pred_wavs = self.decode(latents=latents, audio_lengths=audio_lengths, sr=sr)
|
||||
return sr, pred_wavs, latents, latent_lengths
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
audio, sr = torchaudio.load("/root/data/repo/gongjunmin/sag_train/orig2.wav")
|
||||
audio_lengths = torch.tensor([audio.shape[1]])
|
||||
audios = audio.unsqueeze(0)
|
||||
|
||||
# test encode only
|
||||
model = MusicDCAE()
|
||||
# latents, latent_lengths = model.encode(audios, audio_lengths)
|
||||
# print("latents shape: ", latents.shape)
|
||||
# print("latent_lengths: ", latent_lengths)
|
||||
|
||||
# test encode and decode
|
||||
sr, pred_wavs, latents, latent_lengths = model(audios, audio_lengths, sr)
|
||||
print("reconstructed wavs: ", pred_wavs[0].shape)
|
||||
print("latents shape: ", latents.shape)
|
||||
print("latent_lengths: ", latent_lengths)
|
||||
print("sr: ", sr)
|
||||
torchaudio.save("/root/data/repo/gongjunmin/sag_train/reconstructed.wav", pred_wavs[0], sr)
|
||||
print("reconstructed wav saved to /root/data/repo/gongjunmin/sag_train/reconstructed.wav")
|
||||
@@ -0,0 +1,551 @@
|
||||
from typing import Tuple, Union, Optional, Dict, Any
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from diffusers.models.autoencoders.autoencoder_dc import DCUpBlock2d, get_block, RMSNorm, Decoder
|
||||
from diffusers.models.transformers.sana_transformer import SanaTransformerBlock
|
||||
from diffusers.models.embeddings import get_2d_sincos_pos_embed
|
||||
from diffusers.models.normalization import AdaLayerNormSingle, RMSNorm
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.attention_processor import AttentionProcessor
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.utils import is_torch_version
|
||||
from diffusers.models.unets import UNet2DModel
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int = 32,
|
||||
out_channels: int = 8,
|
||||
attention_head_dim: int = 32,
|
||||
block_out_channels: Tuple[int] = (512, 1024, 2048),
|
||||
layers_per_block: Tuple[int] = (3, 3, 3),
|
||||
block_type: str = "EfficientViTBlock",
|
||||
norm_type: str = "rms_norm",
|
||||
act_fn: str = "silu",
|
||||
qkv_multiscales: tuple = (5,),
|
||||
):
|
||||
super(Encoder, self).__init__()
|
||||
|
||||
num_blocks = len(block_out_channels)
|
||||
|
||||
self.dump_encoder = False
|
||||
if num_blocks == 0:
|
||||
self.dump_encoder = True
|
||||
return
|
||||
|
||||
self.conv_in = nn.Conv2d(in_channels, block_out_channels[-1], kernel_size=3, stride=1, padding=1)
|
||||
|
||||
up_blocks = []
|
||||
for i, (out_channel, num_layers) in reversed(list(enumerate(zip(block_out_channels, layers_per_block)))):
|
||||
up_block_list = []
|
||||
|
||||
if i < num_blocks - 1 and num_layers > 0:
|
||||
upsample_block = DCUpBlock2d(
|
||||
block_out_channels[i + 1],
|
||||
out_channel,
|
||||
interpolate=True,
|
||||
shortcut=True,
|
||||
)
|
||||
up_block_list.append(upsample_block)
|
||||
|
||||
for _ in range(num_layers):
|
||||
block = get_block(
|
||||
block_type,
|
||||
out_channel,
|
||||
out_channel,
|
||||
attention_head_dim=attention_head_dim,
|
||||
norm_type=norm_type,
|
||||
act_fn=act_fn,
|
||||
qkv_mutliscales=qkv_multiscales,
|
||||
)
|
||||
up_block_list.append(block)
|
||||
|
||||
up_blocks.insert(0, nn.Sequential(*up_block_list))
|
||||
|
||||
self.up_blocks = nn.ModuleList(up_blocks)
|
||||
|
||||
self.norm_out = RMSNorm(block_out_channels[0], 1e-5, elementwise_affine=True, bias=True)
|
||||
self.conv_act = nn.ReLU()
|
||||
self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
if self.dump_encoder:
|
||||
return hidden_states
|
||||
|
||||
hidden_states = self.conv_in(hidden_states)
|
||||
i = 0
|
||||
for up_block in reversed(self.up_blocks):
|
||||
hidden_states = up_block(hidden_states)
|
||||
i += 1
|
||||
|
||||
hidden_states = self.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
|
||||
hidden_states = self.conv_act(hidden_states)
|
||||
hidden_states = self.conv_out(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class PatchEmbed(nn.Module):
|
||||
"""
|
||||
2D Image to Patch Embedding with support for SD3 cropping.
|
||||
|
||||
Args:
|
||||
height (`int`, defaults to `224`): The height of the image.
|
||||
width (`int`, defaults to `224`): The width of the image.
|
||||
patch_size (`int`, defaults to `16`): The size of the patches.
|
||||
in_channels (`int`, defaults to `3`): The number of input channels.
|
||||
embed_dim (`int`, defaults to `768`): The output dimension of the embedding.
|
||||
layer_norm (`bool`, defaults to `False`): Whether or not to use layer normalization.
|
||||
flatten (`bool`, defaults to `True`): Whether or not to flatten the output.
|
||||
bias (`bool`, defaults to `True`): Whether or not to use bias.
|
||||
interpolation_scale (`float`, defaults to `1`): The scale of the interpolation.
|
||||
pos_embed_type (`str`, defaults to `"sincos"`): The type of positional embedding.
|
||||
pos_embed_max_size (`int`, defaults to `None`): The maximum size of the positional embedding.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
height=16,
|
||||
width=128,
|
||||
patch_size=(16,1),
|
||||
in_channels=16,
|
||||
embed_dim=768,
|
||||
layer_norm=False,
|
||||
flatten=True,
|
||||
bias=True,
|
||||
interpolation_scale=1,
|
||||
pos_embed_type="sincos",
|
||||
pos_embed_max_size=None, # For SD3 cropping
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
num_patches = (height // patch_size[0]) * (width // patch_size[1])
|
||||
self.flatten = flatten
|
||||
self.layer_norm = layer_norm
|
||||
self.pos_embed_max_size = pos_embed_max_size
|
||||
|
||||
self.proj = nn.Conv2d(
|
||||
in_channels, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias
|
||||
)
|
||||
if layer_norm:
|
||||
self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6)
|
||||
else:
|
||||
self.norm = None
|
||||
|
||||
self.patch_size = patch_size
|
||||
self.height, self.width = height // patch_size[0], width // patch_size[1]
|
||||
self.base_size = height // patch_size[1]
|
||||
self.interpolation_scale = interpolation_scale
|
||||
|
||||
# Calculate positional embeddings based on max size or default
|
||||
if pos_embed_max_size:
|
||||
grid_size = pos_embed_max_size
|
||||
else:
|
||||
grid_size = int(num_patches**0.5)
|
||||
|
||||
if pos_embed_type is None:
|
||||
self.pos_embed = None
|
||||
elif pos_embed_type == "sincos":
|
||||
pos_embed = get_2d_sincos_pos_embed(
|
||||
embed_dim,
|
||||
grid_size,
|
||||
base_size=self.base_size,
|
||||
interpolation_scale=self.interpolation_scale,
|
||||
output_type="pt",
|
||||
)
|
||||
persistent = True if pos_embed_max_size else False
|
||||
self.register_buffer("pos_embed", pos_embed.float().unsqueeze(0), persistent=persistent)
|
||||
else:
|
||||
raise ValueError(f"Unsupported pos_embed_type: {pos_embed_type}")
|
||||
|
||||
def cropped_pos_embed(self, height, width):
|
||||
"""Crops positional embeddings for SD3 compatibility."""
|
||||
if self.pos_embed_max_size is None:
|
||||
raise ValueError("`pos_embed_max_size` must be set for cropping.")
|
||||
|
||||
height = height // self.patch_size
|
||||
width = width // self.patch_size
|
||||
if height > self.pos_embed_max_size:
|
||||
raise ValueError(
|
||||
f"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."
|
||||
)
|
||||
if width > self.pos_embed_max_size:
|
||||
raise ValueError(
|
||||
f"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."
|
||||
)
|
||||
|
||||
top = (self.pos_embed_max_size - height) // 2
|
||||
left = (self.pos_embed_max_size - width) // 2
|
||||
spatial_pos_embed = self.pos_embed.reshape(1, self.pos_embed_max_size, self.pos_embed_max_size, -1)
|
||||
spatial_pos_embed = spatial_pos_embed[:, top : top + height, left : left + width, :]
|
||||
spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1])
|
||||
return spatial_pos_embed
|
||||
|
||||
def forward(self, latent):
|
||||
if self.pos_embed_max_size is not None:
|
||||
height, width = latent.shape[-2:]
|
||||
else:
|
||||
height, width = latent.shape[-2] // self.patch_size[0], latent.shape[-1] // self.patch_size[1]
|
||||
latent = self.proj(latent)
|
||||
if self.flatten:
|
||||
latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC
|
||||
if self.layer_norm:
|
||||
latent = self.norm(latent)
|
||||
if self.pos_embed is None:
|
||||
return latent.to(latent.dtype)
|
||||
# Interpolate or crop positional embeddings as needed
|
||||
if self.pos_embed_max_size:
|
||||
pos_embed = self.cropped_pos_embed(height, width)
|
||||
else:
|
||||
if self.height != height or self.width != width:
|
||||
pos_embed = get_2d_sincos_pos_embed(
|
||||
embed_dim=self.pos_embed.shape[-1],
|
||||
grid_size=(height, width),
|
||||
base_size=self.base_size,
|
||||
interpolation_scale=self.interpolation_scale,
|
||||
device=latent.device,
|
||||
output_type="pt",
|
||||
)
|
||||
pos_embed = pos_embed.float().unsqueeze(0)
|
||||
else:
|
||||
pos_embed = self.pos_embed
|
||||
|
||||
return (latent + pos_embed).to(latent.dtype)
|
||||
|
||||
|
||||
class DiTDecoder(ModelMixin, ConfigMixin):
|
||||
|
||||
_supports_gradient_checkpointing = True
|
||||
@register_to_config
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sample_size: Tuple[int, int] = (16, 128),
|
||||
in_channels: int = 16,
|
||||
out_channels: int = 8,
|
||||
patch_size: Tuple[int, int] = (16, 1),
|
||||
inner_dim: int = 1152,
|
||||
num_attention_heads: int = 36,
|
||||
attention_head_dim: int = 32,
|
||||
dropout: float = 0.0,
|
||||
cross_attention_dim: Optional[int] = None,
|
||||
num_cross_attention_heads: Optional[int] = None,
|
||||
cross_attention_head_dim: Optional[int] = None,
|
||||
attention_bias: bool = False,
|
||||
norm_elementwise_affine: bool = False,
|
||||
norm_eps: float = 1e-6,
|
||||
interpolation_scale: int = 1,
|
||||
mlp_ratio: float = 2.5,
|
||||
num_layers: int = 12,
|
||||
):
|
||||
super(DiTDecoder, self).__init__()
|
||||
interpolation_scale = interpolation_scale if interpolation_scale is not None else max(sample_size // 64, 1)
|
||||
self.interpolation_scale = interpolation_scale
|
||||
|
||||
self.patch_embed = PatchEmbed(
|
||||
height=sample_size[0],
|
||||
width=sample_size[1],
|
||||
patch_size=patch_size,
|
||||
in_channels=in_channels,
|
||||
embed_dim=inner_dim,
|
||||
interpolation_scale=interpolation_scale,
|
||||
)
|
||||
|
||||
self.time_embed = AdaLayerNormSingle(inner_dim)
|
||||
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
SanaTransformerBlock(
|
||||
inner_dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
dropout=dropout,
|
||||
num_cross_attention_heads=num_cross_attention_heads,
|
||||
cross_attention_head_dim=cross_attention_head_dim,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attention_bias=attention_bias,
|
||||
norm_elementwise_affine=norm_elementwise_affine,
|
||||
norm_eps=norm_eps,
|
||||
mlp_ratio=mlp_ratio,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim ** 0.5)
|
||||
self.norm_out = nn.LayerNorm(inner_dim, eps=1e-6, elementwise_affine=False)
|
||||
self.proj_out = nn.Linear(inner_dim, patch_size[0] * patch_size[1] * out_channels)
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if hasattr(module, "gradient_checkpointing"):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
@property
|
||||
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
|
||||
def attn_processors(self) -> Dict[str, AttentionProcessor]:
|
||||
r"""
|
||||
Returns:
|
||||
`dict` of attention processors: A dictionary containing all attention processors used in the model with
|
||||
indexed by its weight name.
|
||||
"""
|
||||
# set recursively
|
||||
processors = {}
|
||||
|
||||
def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]):
|
||||
if hasattr(module, "get_processor"):
|
||||
processors[f"{name}.processor"] = module.get_processor()
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
|
||||
|
||||
return processors
|
||||
|
||||
for name, module in self.named_children():
|
||||
fn_recursive_add_processors(name, module, processors)
|
||||
|
||||
return processors
|
||||
|
||||
# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
|
||||
def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
|
||||
r"""
|
||||
Sets the attention processor to use to compute attention.
|
||||
|
||||
Parameters:
|
||||
processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
|
||||
The instantiated processor class or a dictionary of processor classes that will be set as the processor
|
||||
for **all** `Attention` layers.
|
||||
|
||||
If `processor` is a dict, the key needs to define the path to the corresponding cross attention
|
||||
processor. This is strongly recommended when setting trainable attention processors.
|
||||
|
||||
"""
|
||||
count = len(self.attn_processors.keys())
|
||||
|
||||
if isinstance(processor, dict) and len(processor) != count:
|
||||
raise ValueError(
|
||||
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
|
||||
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
|
||||
)
|
||||
|
||||
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
|
||||
if hasattr(module, "set_processor"):
|
||||
if not isinstance(processor, dict):
|
||||
module.set_processor(processor)
|
||||
else:
|
||||
module.set_processor(processor.pop(f"{name}.processor"))
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
|
||||
|
||||
for name, module in self.named_children():
|
||||
fn_recursive_attn_processor(name, module, processor)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
timestep: Optional[int] = None,
|
||||
return_dict: bool = True,
|
||||
):
|
||||
|
||||
# 1. Input
|
||||
batch_size, num_channels, height, width = hidden_states.shape
|
||||
patch_size = self.config.patch_size
|
||||
|
||||
post_patch_height, post_patch_width = height // patch_size[0], width // patch_size[1]
|
||||
|
||||
hidden_states = self.patch_embed(hidden_states)
|
||||
|
||||
timestep, embedded_timestep = self.time_embed(
|
||||
timestep, batch_size=batch_size, hidden_dtype=hidden_states.dtype
|
||||
)
|
||||
|
||||
# 2. Transformer blocks
|
||||
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
def custom_forward(*inputs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
|
||||
|
||||
for block in self.transformer_blocks:
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(block),
|
||||
hidden_states,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
timestep,
|
||||
post_patch_height,
|
||||
post_patch_width,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
|
||||
else:
|
||||
for block in self.transformer_blocks:
|
||||
hidden_states = block(
|
||||
hidden_states,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
timestep,
|
||||
post_patch_height,
|
||||
post_patch_width,
|
||||
)
|
||||
|
||||
# 3. Normalization
|
||||
shift, scale = (
|
||||
self.scale_shift_table[None] + embedded_timestep[:, None].to(self.scale_shift_table.device)
|
||||
).chunk(2, dim=1)
|
||||
|
||||
# 4. Modulation
|
||||
hidden_states = hidden_states * (1 + scale) + shift
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
|
||||
# 5. Unpatchify
|
||||
hidden_states = hidden_states.reshape(
|
||||
batch_size, post_patch_height, post_patch_width, self.config.patch_size[0], self.config.patch_size[1], -1
|
||||
)
|
||||
hidden_states = hidden_states.permute(0, 5, 1, 3, 2, 4)
|
||||
output = hidden_states.reshape(batch_size, -1, post_patch_height * patch_size[0], post_patch_width * patch_size[1])
|
||||
|
||||
if not return_dict:
|
||||
return (output,)
|
||||
|
||||
return Transformer2DModelOutput(sample=output)
|
||||
|
||||
|
||||
class MusicDcaeRefiner(ModelMixin, ConfigMixin):
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int = 32,
|
||||
attention_head_dim: int = 32,
|
||||
block_out_channels: Tuple[int] = (512, 1024, 2048),
|
||||
layers_per_block: Tuple[int] = (3, 3, 3),
|
||||
conv_block_out_channels: Tuple[int] = (224, 448, 672, 896),
|
||||
out_channels: int = 8,
|
||||
block_type: str = "EfficientViTBlock",
|
||||
norm_type: str = "rms_norm",
|
||||
act_fn: str = "silu",
|
||||
qkv_multiscales: tuple = (5,),
|
||||
sample_size: Tuple[int, int] = (16, 128),
|
||||
patch_size: Tuple[int, int] = (16, 1),
|
||||
inner_dim: int = 1152,
|
||||
num_attention_heads: int = 36,
|
||||
dropout: float = 0.0,
|
||||
cross_attention_dim: Optional[int] = None,
|
||||
num_cross_attention_heads: Optional[int] = None,
|
||||
cross_attention_head_dim: Optional[int] = None,
|
||||
attention_bias: bool = False,
|
||||
norm_elementwise_affine: bool = False,
|
||||
norm_eps: float = 1e-6,
|
||||
interpolation_scale: int = 1,
|
||||
mlp_ratio: float = 2.5,
|
||||
num_layers: int = 12,
|
||||
decoder_type: str = "ConvDecoder",
|
||||
|
||||
):
|
||||
super(MusicDcaeRefiner, self).__init__()
|
||||
|
||||
self.encoder = Encoder(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
attention_head_dim=attention_head_dim,
|
||||
block_out_channels=block_out_channels,
|
||||
layers_per_block=layers_per_block,
|
||||
block_type=block_type,
|
||||
norm_type=norm_type,
|
||||
act_fn=act_fn,
|
||||
qkv_multiscales=qkv_multiscales,
|
||||
)
|
||||
if decoder_type == "DiTDecoder":
|
||||
self.decoder = DiTDecoder(
|
||||
sample_size=sample_size,
|
||||
in_channels=out_channels * 2,
|
||||
out_channels=out_channels,
|
||||
patch_size=patch_size,
|
||||
inner_dim=inner_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=attention_head_dim,
|
||||
dropout=dropout,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
num_cross_attention_heads=num_cross_attention_heads,
|
||||
cross_attention_head_dim=cross_attention_head_dim,
|
||||
attention_bias=attention_bias,
|
||||
norm_elementwise_affine=norm_elementwise_affine,
|
||||
norm_eps=norm_eps,
|
||||
interpolation_scale=interpolation_scale,
|
||||
mlp_ratio=mlp_ratio,
|
||||
num_layers=num_layers,
|
||||
)
|
||||
else:
|
||||
self.decoder = UNet2DModel(
|
||||
sample_size=sample_size,
|
||||
in_channels=out_channels * 2,
|
||||
out_channels=out_channels,
|
||||
block_out_channels=conv_block_out_channels,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
timestep: Optional[int] = None,
|
||||
return_dict: bool = True
|
||||
):
|
||||
encoder_hidden_states = self.encoder(encoder_hidden_states)
|
||||
hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
|
||||
output = self.decoder(hidden_states, timestep=timestep, return_dict=return_dict)
|
||||
return output
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# f32c32 -> f8c8
|
||||
# model = MusicDcaeRefiner()
|
||||
|
||||
# x = torch.randn(1, 8, 16, 128)
|
||||
# encoder_x = torch.randn(1, 32, 4, 32)
|
||||
# timestep = 0
|
||||
# y = model(x, encoder_x, timestep=timestep)
|
||||
# print("y", y.sample.shape)
|
||||
# total_params = sum(p.numel() for p in model.parameters())
|
||||
# print(f"模型参数总数: {total_params / 1e6:.2f}M")
|
||||
|
||||
# # 分别计算encoder和decoder的参数量
|
||||
# encoder_params_count = sum(p.numel() for p in model.encoder.parameters())
|
||||
# decoder_params_count = sum(p.numel() for p in model.decoder.parameters())
|
||||
# print(f"encoder参数量: {encoder_params_count/1e6:.2f}M")
|
||||
# print(f"decoder参数量: {decoder_params_count/1e6:.2f}M")
|
||||
|
||||
|
||||
# f8c8 -> mel
|
||||
import json
|
||||
with open("music_dcae/config_f8c8_to_mel_refiner.json", "r") as f:
|
||||
config = json.load(f)
|
||||
model = MusicDcaeRefiner(**config)
|
||||
|
||||
x = torch.randn(1, 2, 128, 1024)
|
||||
encoder_x = torch.randn(1, 2, 128, 1024)
|
||||
timestep = 0
|
||||
y = model(x, encoder_x, timestep=timestep)
|
||||
print("y", y.sample.shape)
|
||||
total_params = sum(p.numel() for p in model.parameters())
|
||||
print(f"模型参数总数: {total_params / 1e6:.2f}M")
|
||||
|
||||
# 分别计算encoder和decoder的参数量
|
||||
encoder_params_count = sum(p.numel() for p in model.encoder.parameters())
|
||||
decoder_params_count = sum(p.numel() for p in model.decoder.parameters())
|
||||
print(f"encoder参数量: {encoder_params_count/1e6:.2f}M")
|
||||
print(f"decoder参数量: {decoder_params_count/1e6:.2f}M")
|
||||
@@ -0,0 +1,157 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from diffusers import AutoencoderDC
|
||||
import json
|
||||
import torchvision.transforms as transforms
|
||||
import torchaudio
|
||||
|
||||
try:
|
||||
from .music_vocoder import ADaMoSHiFiGANV1
|
||||
except ImportError:
|
||||
from music_vocoder import ADaMoSHiFiGANV1
|
||||
|
||||
|
||||
DEFAULT_CONFIG_PATH = "/root/sag_train/music_dcae/config_f32c32_large.json"
|
||||
DCAE_PRETRAINED_PATH = "/root/sag_train/checkpoints/music_dcae_f32c32"
|
||||
VOCODER_PRETRAINED_PATH = "/root/sag_train/checkpoints/music_vocoder.pt"
|
||||
|
||||
|
||||
class MusicDCAEVocoder(nn.Module):
|
||||
def __init__(self, config_path=DEFAULT_CONFIG_PATH, pretrained_path=DCAE_PRETRAINED_PATH):
|
||||
super(MusicDCAEVocoder, self).__init__()
|
||||
if pretrained_path is None:
|
||||
with open(config_path) as f:
|
||||
config = json.load(f)
|
||||
self.dcae = AutoencoderDC(**config)
|
||||
else:
|
||||
self.dcae = AutoencoderDC.from_pretrained(pretrained_path)
|
||||
self.vocoder = ADaMoSHiFiGANV1(VOCODER_PRETRAINED_PATH)
|
||||
self.freeze_vocoder()
|
||||
self.transform = transforms.Compose([
|
||||
transforms.Normalize(0.5, 0.5),
|
||||
])
|
||||
self.min_mel_value = -11.0
|
||||
self.max_mel_value = 3.0
|
||||
self.target_sr = 44100
|
||||
|
||||
def load_audio(self, audio_path):
|
||||
audio, sr = torchaudio.load(audio_path)
|
||||
if audio.shape[0] == 1:
|
||||
audio = torch.cat([audio, audio], dim=0)
|
||||
return audio, sr
|
||||
|
||||
def resample_audio(self, audio, sr=48000):
|
||||
resampler = torchaudio.transforms.Resample(sr, self.target_sr)
|
||||
resampler = resampler.to(audio.device)
|
||||
audio = resampler(audio)
|
||||
return audio
|
||||
|
||||
def forward_mel(self, audios):
|
||||
mels = []
|
||||
for i in range(len(audios)):
|
||||
image = self.vocoder.mel_transform(audios[i])
|
||||
mels.append(image)
|
||||
mels = torch.stack(mels)
|
||||
return mels
|
||||
|
||||
def norm_mel(self, mels):
|
||||
normed_mels = (mels - self.min_mel_value) / (self.max_mel_value - self.min_mel_value)
|
||||
normed_mels = self.transform(normed_mels)
|
||||
return normed_mels
|
||||
|
||||
def denorm_mel(self, normed_mels):
|
||||
mels = normed_mels * 0.5 + 0.5
|
||||
mels = mels * (self.max_mel_value - self.min_mel_value) + self.min_mel_value
|
||||
return mels
|
||||
|
||||
def encode_latent(self, normed_mels):
|
||||
# N x 2 x 128 x W -> N x C x 128//F x W//F
|
||||
latent = self.dcae.encode(normed_mels).latent
|
||||
return latent
|
||||
|
||||
def decode_mel(self, latent):
|
||||
# N x C x 128//F x W//F -> N x 2 x 128 x W
|
||||
normed_mels = self.dcae.decode(latent).sample
|
||||
return normed_mels
|
||||
|
||||
def decode_audio(self, mels):
|
||||
# mels: N x 2 x 128 x W -> 2N x 128 x W
|
||||
bs = mels.shape[0]
|
||||
mono_mels = mels.reshape(-1, 128, mels.shape[-1])
|
||||
mono_audios = self.vocoder(mono_mels)
|
||||
audios = mono_audios.reshape(bs, 2, -1)
|
||||
return audios
|
||||
|
||||
def encode(self, audios):
|
||||
mels = self.forward_mel(audios)
|
||||
normed_mels = self.norm_mel(mels)
|
||||
latent = self.encode_latent(normed_mels)
|
||||
return latent, mels
|
||||
|
||||
def decode(self, latent):
|
||||
recon_normed_mels = self.decode_mel(latent)
|
||||
recon_mels = self.denorm_mel(recon_normed_mels)
|
||||
recon_audios = self.decode_audio(recon_mels)
|
||||
return recon_audios, recon_mels
|
||||
|
||||
def forward(self, audios):
|
||||
audios_len = audios.shape[-1]
|
||||
latent, mels = self.encode(audios)
|
||||
recon_audios, recon_mels = self.decode(latent)
|
||||
if recon_audios.shape[-1] > audios_len:
|
||||
recon_audios = recon_audios[:, :, :audios_len]
|
||||
elif recon_audios.shape[-1] < audios_len:
|
||||
recon_audios = F.pad(recon_audios, (0, audios_len - recon_audios.shape[-1]))
|
||||
return recon_audios, mels, recon_mels, latent
|
||||
|
||||
def freeze_vocoder(self):
|
||||
self.vocoder.eval()
|
||||
self.vocoder.requires_grad_(False)
|
||||
|
||||
def unfreeze_vocoder(self):
|
||||
self.vocoder.train()
|
||||
self.vocoder.requires_grad_(True)
|
||||
|
||||
def return_middle_layers(self):
|
||||
last_down_block = self.dcae.encoder.down_blocks[-1]
|
||||
encoder_conv_out = self.dcae.encoder.conv_out
|
||||
decoder_conv_in = self.dcae.decoder.conv_in
|
||||
decoder_up_blocks = self.dcae.decoder.up_blocks[0]
|
||||
middle_layers = [last_down_block, encoder_conv_out, decoder_conv_in, decoder_up_blocks]
|
||||
return middle_layers
|
||||
|
||||
def return_head_layers(self):
|
||||
decoder_up_blocks = self.dcae.decoder.up_blocks[-1]
|
||||
conv_out = self.dcae.decoder.conv_out
|
||||
head_layers = [decoder_up_blocks, conv_out]
|
||||
return head_layers
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
model = MusicDCAEVocoder()
|
||||
|
||||
audio_path = "/root/sag_train/orig2.wav"
|
||||
audio, sr = model.load_audio(audio_path)
|
||||
audio = model.resample_audio(audio, sr)
|
||||
|
||||
model.eval()
|
||||
model = model.to("cuda:0")
|
||||
audio = audio.to("cuda:0")
|
||||
with torch.no_grad():
|
||||
audios_len = audio.shape[-1]
|
||||
min_frame = 512 * 32
|
||||
if audios_len % min_frame != 0:
|
||||
padding = torch.zeros(audio.shape[0], 2, min_frame - audios_len % min_frame).to(audios.device)
|
||||
audios = torch.cat([audio, padding], dim=-1)
|
||||
recon_audios, mels, recon_mels, latent = model(audio.unsqueeze(0))
|
||||
recon_audios = recon_audios[:, :, :audios_len]
|
||||
|
||||
print("latent shape: ", latent.shape)
|
||||
print("recon_audios", recon_audios.shape)
|
||||
print("mels", mels.shape, "min:", mels.min(), "max:", mels.max(), "mean:", mels.mean(), "std:", mels.std())
|
||||
print("recon_mels", recon_mels.shape, "min:", recon_mels.min(), "max:", recon_mels.max(), "mean:", recon_mels.mean(), "std:", recon_mels.std())
|
||||
total_params = sum(p.numel() for p in model.parameters())
|
||||
print(f"模型参数总数: {total_params / 1e6:.2f}M")
|
||||
|
||||
torchaudio.save("/root/sag_train/recon2.wav", recon_audios[0].cpu(), 44100)
|
||||
Executable
+119
@@ -0,0 +1,119 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch import Tensor
|
||||
from torchaudio.transforms import MelScale
|
||||
|
||||
|
||||
class LinearSpectrogram(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_fft=2048,
|
||||
win_length=2048,
|
||||
hop_length=512,
|
||||
center=False,
|
||||
mode="pow2_sqrt",
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.n_fft = n_fft
|
||||
self.win_length = win_length
|
||||
self.hop_length = hop_length
|
||||
self.center = center
|
||||
self.mode = mode
|
||||
|
||||
self.register_buffer("window", torch.hann_window(win_length))
|
||||
|
||||
def forward(self, y: Tensor) -> Tensor:
|
||||
if y.ndim == 3:
|
||||
y = y.squeeze(1)
|
||||
|
||||
y = torch.nn.functional.pad(
|
||||
y.unsqueeze(1),
|
||||
(
|
||||
(self.win_length - self.hop_length) // 2,
|
||||
(self.win_length - self.hop_length + 1) // 2,
|
||||
),
|
||||
mode="reflect",
|
||||
).squeeze(1)
|
||||
dtype = y.dtype
|
||||
spec = torch.stft(
|
||||
y.float(),
|
||||
self.n_fft,
|
||||
hop_length=self.hop_length,
|
||||
win_length=self.win_length,
|
||||
window=self.window,
|
||||
center=self.center,
|
||||
pad_mode="reflect",
|
||||
normalized=False,
|
||||
onesided=True,
|
||||
return_complex=True,
|
||||
)
|
||||
spec = torch.view_as_real(spec)
|
||||
|
||||
if self.mode == "pow2_sqrt":
|
||||
spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
|
||||
spec = spec.to(dtype)
|
||||
return spec
|
||||
|
||||
|
||||
class LogMelSpectrogram(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
sample_rate=44100,
|
||||
n_fft=2048,
|
||||
win_length=2048,
|
||||
hop_length=512,
|
||||
n_mels=128,
|
||||
center=False,
|
||||
f_min=0.0,
|
||||
f_max=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.sample_rate = sample_rate
|
||||
self.n_fft = n_fft
|
||||
self.win_length = win_length
|
||||
self.hop_length = hop_length
|
||||
self.center = center
|
||||
self.n_mels = n_mels
|
||||
self.f_min = f_min
|
||||
self.f_max = f_max or sample_rate // 2
|
||||
|
||||
self.spectrogram = LinearSpectrogram(n_fft, win_length, hop_length, center)
|
||||
self.mel_scale = MelScale(
|
||||
self.n_mels,
|
||||
self.sample_rate,
|
||||
self.f_min,
|
||||
self.f_max,
|
||||
self.n_fft // 2 + 1,
|
||||
"slaney",
|
||||
"slaney",
|
||||
)
|
||||
|
||||
def compress(self, x: Tensor) -> Tensor:
|
||||
return torch.log(torch.clamp(x, min=1e-5))
|
||||
|
||||
def decompress(self, x: Tensor) -> Tensor:
|
||||
return torch.exp(x)
|
||||
|
||||
def forward(self, x: Tensor, return_linear: bool = False) -> Tensor:
|
||||
linear = self.spectrogram(x)
|
||||
x = self.mel_scale(linear)
|
||||
x = self.compress(x)
|
||||
# print(x.shape)
|
||||
if return_linear:
|
||||
return x, self.compress(linear)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def get_mel_transform():
|
||||
return LogMelSpectrogram(
|
||||
sample_rate=44100,
|
||||
n_fft=2048,
|
||||
win_length=2048,
|
||||
hop_length=512,
|
||||
f_min=40,
|
||||
f_max=16000,
|
||||
n_mels=128,
|
||||
)
|
||||
Executable
+565
@@ -0,0 +1,565 @@
|
||||
import librosa
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from functools import partial
|
||||
from math import prod
|
||||
from typing import Callable, Tuple, List
|
||||
|
||||
import numpy as np
|
||||
import torch.nn.functional as F
|
||||
from torch.nn import Conv1d
|
||||
from torch.nn.utils import weight_norm
|
||||
from torch.nn.utils.parametrize import remove_parametrizations as remove_weight_norm
|
||||
|
||||
try:
|
||||
from music_log_mel import LogMelSpectrogram
|
||||
except ImportError:
|
||||
from .music_log_mel import LogMelSpectrogram
|
||||
|
||||
|
||||
def drop_path(
|
||||
x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
|
||||
):
|
||||
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
|
||||
|
||||
This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
|
||||
the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
|
||||
See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
|
||||
changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
|
||||
'survival rate' as the argument.
|
||||
|
||||
""" # noqa: E501
|
||||
|
||||
if drop_prob == 0.0 or not training:
|
||||
return x
|
||||
keep_prob = 1 - drop_prob
|
||||
shape = (x.shape[0],) + (1,) * (
|
||||
x.ndim - 1
|
||||
) # work with diff dim tensors, not just 2D ConvNets
|
||||
random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
|
||||
if keep_prob > 0.0 and scale_by_keep:
|
||||
random_tensor.div_(keep_prob)
|
||||
return x * random_tensor
|
||||
|
||||
|
||||
class DropPath(nn.Module):
|
||||
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" # noqa: E501
|
||||
|
||||
def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
|
||||
super(DropPath, self).__init__()
|
||||
self.drop_prob = drop_prob
|
||||
self.scale_by_keep = scale_by_keep
|
||||
|
||||
def forward(self, x):
|
||||
return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
|
||||
|
||||
def extra_repr(self):
|
||||
return f"drop_prob={round(self.drop_prob,3):0.3f}"
|
||||
|
||||
|
||||
class LayerNorm(nn.Module):
|
||||
r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
|
||||
The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
|
||||
shape (batch_size, height, width, channels) while channels_first corresponds to inputs
|
||||
with shape (batch_size, channels, height, width).
|
||||
""" # noqa: E501
|
||||
|
||||
def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(normalized_shape))
|
||||
self.bias = nn.Parameter(torch.zeros(normalized_shape))
|
||||
self.eps = eps
|
||||
self.data_format = data_format
|
||||
if self.data_format not in ["channels_last", "channels_first"]:
|
||||
raise NotImplementedError
|
||||
self.normalized_shape = (normalized_shape,)
|
||||
|
||||
def forward(self, x):
|
||||
if self.data_format == "channels_last":
|
||||
return F.layer_norm(
|
||||
x, self.normalized_shape, self.weight, self.bias, self.eps
|
||||
)
|
||||
elif self.data_format == "channels_first":
|
||||
u = x.mean(1, keepdim=True)
|
||||
s = (x - u).pow(2).mean(1, keepdim=True)
|
||||
x = (x - u) / torch.sqrt(s + self.eps)
|
||||
x = self.weight[:, None] * x + self.bias[:, None]
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXtBlock(nn.Module):
|
||||
r"""ConvNeXt Block. There are two equivalent implementations:
|
||||
(1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
|
||||
(2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
|
||||
We use (2) as we find it slightly faster in PyTorch
|
||||
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
drop_path (float): Stochastic depth rate. Default: 0.0
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0.
|
||||
kernel_size (int): Kernel size for depthwise conv. Default: 7.
|
||||
dilation (int): Dilation for depthwise conv. Default: 1.
|
||||
""" # noqa: E501
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
drop_path: float = 0.0,
|
||||
layer_scale_init_value: float = 1e-6,
|
||||
mlp_ratio: float = 4.0,
|
||||
kernel_size: int = 7,
|
||||
dilation: int = 1,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.dwconv = nn.Conv1d(
|
||||
dim,
|
||||
dim,
|
||||
kernel_size=kernel_size,
|
||||
padding=int(dilation * (kernel_size - 1) / 2),
|
||||
groups=dim,
|
||||
) # depthwise conv
|
||||
self.norm = LayerNorm(dim, eps=1e-6)
|
||||
self.pwconv1 = nn.Linear(
|
||||
dim, int(mlp_ratio * dim)
|
||||
) # pointwise/1x1 convs, implemented with linear layers
|
||||
self.act = nn.GELU()
|
||||
self.pwconv2 = nn.Linear(int(mlp_ratio * dim), dim)
|
||||
self.gamma = (
|
||||
nn.Parameter(layer_scale_init_value *
|
||||
torch.ones((dim)), requires_grad=True)
|
||||
if layer_scale_init_value > 0
|
||||
else None
|
||||
)
|
||||
self.drop_path = DropPath(
|
||||
drop_path) if drop_path > 0.0 else nn.Identity()
|
||||
|
||||
def forward(self, x, apply_residual: bool = True):
|
||||
input = x
|
||||
|
||||
x = self.dwconv(x)
|
||||
x = x.permute(0, 2, 1) # (N, C, L) -> (N, L, C)
|
||||
x = self.norm(x)
|
||||
x = self.pwconv1(x)
|
||||
x = self.act(x)
|
||||
x = self.pwconv2(x)
|
||||
|
||||
if self.gamma is not None:
|
||||
x = self.gamma * x
|
||||
|
||||
x = x.permute(0, 2, 1) # (N, L, C) -> (N, C, L)
|
||||
x = self.drop_path(x)
|
||||
|
||||
if apply_residual:
|
||||
x = input + x
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class ParallelConvNeXtBlock(nn.Module):
|
||||
def __init__(self, kernel_sizes: List[int], *args, **kwargs):
|
||||
super().__init__()
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
ConvNeXtBlock(kernel_size=kernel_size, *args, **kwargs)
|
||||
for kernel_size in kernel_sizes
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.stack(
|
||||
[block(x, apply_residual=False) for block in self.blocks] + [x],
|
||||
dim=1,
|
||||
).sum(dim=1)
|
||||
|
||||
|
||||
class ConvNeXtEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_channels=3,
|
||||
depths=[3, 3, 9, 3],
|
||||
dims=[96, 192, 384, 768],
|
||||
drop_path_rate=0.0,
|
||||
layer_scale_init_value=1e-6,
|
||||
kernel_sizes: Tuple[int] = (7,),
|
||||
):
|
||||
super().__init__()
|
||||
assert len(depths) == len(dims)
|
||||
|
||||
self.channel_layers = nn.ModuleList()
|
||||
stem = nn.Sequential(
|
||||
nn.Conv1d(
|
||||
input_channels,
|
||||
dims[0],
|
||||
kernel_size=7,
|
||||
padding=3,
|
||||
padding_mode="replicate",
|
||||
),
|
||||
LayerNorm(dims[0], eps=1e-6, data_format="channels_first"),
|
||||
)
|
||||
self.channel_layers.append(stem)
|
||||
|
||||
for i in range(len(depths) - 1):
|
||||
mid_layer = nn.Sequential(
|
||||
LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
|
||||
nn.Conv1d(dims[i], dims[i + 1], kernel_size=1),
|
||||
)
|
||||
self.channel_layers.append(mid_layer)
|
||||
|
||||
block_fn = (
|
||||
partial(ConvNeXtBlock, kernel_size=kernel_sizes[0])
|
||||
if len(kernel_sizes) == 1
|
||||
else partial(ParallelConvNeXtBlock, kernel_sizes=kernel_sizes)
|
||||
)
|
||||
|
||||
self.stages = nn.ModuleList()
|
||||
drop_path_rates = [
|
||||
x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
|
||||
]
|
||||
|
||||
cur = 0
|
||||
for i in range(len(depths)):
|
||||
stage = nn.Sequential(
|
||||
*[
|
||||
block_fn(
|
||||
dim=dims[i],
|
||||
drop_path=drop_path_rates[cur + j],
|
||||
layer_scale_init_value=layer_scale_init_value,
|
||||
)
|
||||
for j in range(depths[i])
|
||||
]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
self.norm = LayerNorm(dims[-1], eps=1e-6, data_format="channels_first")
|
||||
self.apply(self._init_weights)
|
||||
|
||||
def _init_weights(self, m):
|
||||
if isinstance(m, (nn.Conv1d, nn.Linear)):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
for channel_layer, stage in zip(self.channel_layers, self.stages):
|
||||
x = channel_layer(x)
|
||||
x = stage(x)
|
||||
|
||||
return self.norm(x)
|
||||
|
||||
|
||||
def init_weights(m, mean=0.0, std=0.01):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find("Conv") != -1:
|
||||
m.weight.data.normal_(mean, std)
|
||||
|
||||
|
||||
def get_padding(kernel_size, dilation=1):
|
||||
return (kernel_size * dilation - dilation) // 2
|
||||
|
||||
|
||||
class ResBlock1(torch.nn.Module):
|
||||
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
|
||||
super().__init__()
|
||||
|
||||
self.convs1 = nn.ModuleList(
|
||||
[
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[0],
|
||||
padding=get_padding(kernel_size, dilation[0]),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[1],
|
||||
padding=get_padding(kernel_size, dilation[1]),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=dilation[2],
|
||||
padding=get_padding(kernel_size, dilation[2]),
|
||||
)
|
||||
),
|
||||
]
|
||||
)
|
||||
self.convs1.apply(init_weights)
|
||||
|
||||
self.convs2 = nn.ModuleList(
|
||||
[
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=1,
|
||||
padding=get_padding(kernel_size, 1),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=1,
|
||||
padding=get_padding(kernel_size, 1),
|
||||
)
|
||||
),
|
||||
weight_norm(
|
||||
Conv1d(
|
||||
channels,
|
||||
channels,
|
||||
kernel_size,
|
||||
1,
|
||||
dilation=1,
|
||||
padding=get_padding(kernel_size, 1),
|
||||
)
|
||||
),
|
||||
]
|
||||
)
|
||||
self.convs2.apply(init_weights)
|
||||
|
||||
def forward(self, x):
|
||||
for c1, c2 in zip(self.convs1, self.convs2):
|
||||
xt = F.silu(x)
|
||||
xt = c1(xt)
|
||||
xt = F.silu(xt)
|
||||
xt = c2(xt)
|
||||
x = xt + x
|
||||
return x
|
||||
|
||||
def remove_weight_norm(self):
|
||||
for conv in self.convs1:
|
||||
remove_weight_norm(conv)
|
||||
for conv in self.convs2:
|
||||
remove_weight_norm(conv)
|
||||
|
||||
|
||||
class HiFiGANGenerator(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
hop_length: int = 512,
|
||||
upsample_rates: Tuple[int] = (8, 8, 2, 2, 2),
|
||||
upsample_kernel_sizes: Tuple[int] = (16, 16, 8, 2, 2),
|
||||
resblock_kernel_sizes: Tuple[int] = (3, 7, 11),
|
||||
resblock_dilation_sizes: Tuple[Tuple[int]] = (
|
||||
(1, 3, 5), (1, 3, 5), (1, 3, 5)),
|
||||
num_mels: int = 128,
|
||||
upsample_initial_channel: int = 512,
|
||||
use_template: bool = True,
|
||||
pre_conv_kernel_size: int = 7,
|
||||
post_conv_kernel_size: int = 7,
|
||||
post_activation: Callable = partial(nn.SiLU, inplace=True),
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
assert (
|
||||
prod(upsample_rates) == hop_length
|
||||
), f"hop_length must be {prod(upsample_rates)}"
|
||||
|
||||
self.conv_pre = weight_norm(
|
||||
nn.Conv1d(
|
||||
num_mels,
|
||||
upsample_initial_channel,
|
||||
pre_conv_kernel_size,
|
||||
1,
|
||||
padding=get_padding(pre_conv_kernel_size),
|
||||
)
|
||||
)
|
||||
|
||||
self.num_upsamples = len(upsample_rates)
|
||||
self.num_kernels = len(resblock_kernel_sizes)
|
||||
|
||||
self.noise_convs = nn.ModuleList()
|
||||
self.use_template = use_template
|
||||
self.ups = nn.ModuleList()
|
||||
|
||||
for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
|
||||
c_cur = upsample_initial_channel // (2 ** (i + 1))
|
||||
self.ups.append(
|
||||
weight_norm(
|
||||
nn.ConvTranspose1d(
|
||||
upsample_initial_channel // (2**i),
|
||||
upsample_initial_channel // (2 ** (i + 1)),
|
||||
k,
|
||||
u,
|
||||
padding=(k - u) // 2,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
if not use_template:
|
||||
continue
|
||||
|
||||
if i + 1 < len(upsample_rates):
|
||||
stride_f0 = np.prod(upsample_rates[i + 1:])
|
||||
self.noise_convs.append(
|
||||
Conv1d(
|
||||
1,
|
||||
c_cur,
|
||||
kernel_size=stride_f0 * 2,
|
||||
stride=stride_f0,
|
||||
padding=stride_f0 // 2,
|
||||
)
|
||||
)
|
||||
else:
|
||||
self.noise_convs.append(Conv1d(1, c_cur, kernel_size=1))
|
||||
|
||||
self.resblocks = nn.ModuleList()
|
||||
for i in range(len(self.ups)):
|
||||
ch = upsample_initial_channel // (2 ** (i + 1))
|
||||
for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes):
|
||||
self.resblocks.append(ResBlock1(ch, k, d))
|
||||
|
||||
self.activation_post = post_activation()
|
||||
self.conv_post = weight_norm(
|
||||
nn.Conv1d(
|
||||
ch,
|
||||
1,
|
||||
post_conv_kernel_size,
|
||||
1,
|
||||
padding=get_padding(post_conv_kernel_size),
|
||||
)
|
||||
)
|
||||
self.ups.apply(init_weights)
|
||||
self.conv_post.apply(init_weights)
|
||||
|
||||
def forward(self, x, template=None):
|
||||
x = self.conv_pre(x)
|
||||
|
||||
for i in range(self.num_upsamples):
|
||||
x = F.silu(x, inplace=True)
|
||||
x = self.ups[i](x)
|
||||
|
||||
if self.use_template:
|
||||
x = x + self.noise_convs[i](template)
|
||||
|
||||
xs = None
|
||||
|
||||
for j in range(self.num_kernels):
|
||||
if xs is None:
|
||||
xs = self.resblocks[i * self.num_kernels + j](x)
|
||||
else:
|
||||
xs += self.resblocks[i * self.num_kernels + j](x)
|
||||
|
||||
x = xs / self.num_kernels
|
||||
|
||||
x = self.activation_post(x)
|
||||
x = self.conv_post(x)
|
||||
x = torch.tanh(x)
|
||||
|
||||
return x
|
||||
|
||||
def remove_weight_norm(self):
|
||||
for up in self.ups:
|
||||
remove_weight_norm(up)
|
||||
for block in self.resblocks:
|
||||
block.remove_weight_norm()
|
||||
remove_weight_norm(self.conv_pre)
|
||||
remove_weight_norm(self.conv_post)
|
||||
|
||||
|
||||
class ADaMoSHiFiGANV1(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
checkpoint_path: str = "checkpoints/adamos-generator-1640000.pth",
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.backbone = ConvNeXtEncoder(
|
||||
input_channels=128,
|
||||
depths=[3, 3, 9, 3],
|
||||
dims=[128, 256, 384, 512],
|
||||
drop_path_rate=0,
|
||||
kernel_sizes=(7,),
|
||||
)
|
||||
|
||||
self.head = HiFiGANGenerator(
|
||||
hop_length=512,
|
||||
upsample_rates=(4, 4, 2, 2, 2, 2, 2),
|
||||
upsample_kernel_sizes=(8, 8, 4, 4, 4, 4, 4),
|
||||
resblock_kernel_sizes=(3, 7, 11, 13),
|
||||
resblock_dilation_sizes=(
|
||||
(1, 3, 5), (1, 3, 5), (1, 3, 5), (1, 3, 5)),
|
||||
num_mels=512,
|
||||
upsample_initial_channel=1024,
|
||||
use_template=False,
|
||||
pre_conv_kernel_size=13,
|
||||
post_conv_kernel_size=13,
|
||||
)
|
||||
self.sampling_rate = 44100
|
||||
|
||||
ckpt_state = torch.load(checkpoint_path, map_location="cpu")
|
||||
|
||||
if "state_dict" in ckpt_state:
|
||||
ckpt_state = ckpt_state["state_dict"]
|
||||
|
||||
if any(k.startswith("generator.") for k in ckpt_state):
|
||||
ckpt_state = {
|
||||
k.replace("generator.", ""): v
|
||||
for k, v in ckpt_state.items()
|
||||
if k.startswith("generator.")
|
||||
}
|
||||
|
||||
self.load_state_dict(ckpt_state)
|
||||
self.eval()
|
||||
|
||||
self.mel_transform = LogMelSpectrogram(
|
||||
sample_rate=44100,
|
||||
n_fft=2048,
|
||||
win_length=2048,
|
||||
hop_length=512,
|
||||
f_min=40,
|
||||
f_max=16000,
|
||||
n_mels=128,
|
||||
)
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(self, mel):
|
||||
y = self.backbone(mel)
|
||||
y = self.head(y)
|
||||
return y
|
||||
|
||||
@torch.no_grad()
|
||||
def encode(self, x):
|
||||
return self.mel_transform(x)
|
||||
|
||||
def forward(self, mel):
|
||||
y = self.backbone(mel)
|
||||
y = self.head(y)
|
||||
return y
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import soundfile as sf
|
||||
|
||||
x = "./test.wav"
|
||||
model = ADaMoSHiFiGANV1(checkpoint_path='./step_001640000.pth')
|
||||
|
||||
wav, sr = librosa.load(x, sr=44100, mono=True)
|
||||
wav = torch.from_numpy(wav).float()[None]
|
||||
mel = model.encode(wav)
|
||||
|
||||
wav = model.decode(mel)[0].mT
|
||||
sf.write("test_out.wav", wav.cpu().numpy(), 44100)
|
||||
Reference in New Issue
Block a user