delete irrelevant code
This commit is contained in:
chuxij
@@ -1,92 +0,0 @@
|
||||
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)
|
||||
@@ -1,69 +0,0 @@
|
||||
{
|
||||
"_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"
|
||||
}
|
||||
@@ -1,124 +0,0 @@
|
||||
# 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))
|
||||
@@ -1,78 +0,0 @@
|
||||
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")
|
||||
@@ -135,7 +135,7 @@ class MusicDCAE(nn.Module):
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
audio, sr = torchaudio.load("/root/data/repo/gongjunmin/sag_train/orig2.wav")
|
||||
audio, sr = torchaudio.load("test.wav")
|
||||
audio_lengths = torch.tensor([audio.shape[1]])
|
||||
audios = audio.unsqueeze(0)
|
||||
|
||||
@@ -151,5 +151,5 @@ if __name__ == "__main__":
|
||||
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")
|
||||
torchaudio.save("test_reconstructed.flac", pred_wavs[0], sr)
|
||||
print("/test_reconstructed.flac")
|
||||
|
||||
@@ -1,551 +0,0 @@
|
||||
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")
|
||||
@@ -1,157 +0,0 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user