mrleo1nid/ACE-Step
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

work on pip package

Browse Source
This commit is contained in:
mrfakename
2025-05-06 18:59:32 -07:00
parent 54da683d36
commit a5746eaab6
36 changed files with 2928 additions and 1251 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acestep/models/ace_step_transformer.py
+155 -46
View File
@@ -31,9 +31,15 @@ from .lyrics_utils.lyric_encoder import ConformerEncoder as LyricEncoder
def cross_norm(hidden_states, controlnet_input):
# input N x T x c
mean_hidden_states, std_hidden_states = hidden_states.mean(dim=(1,2), keepdim=True), hidden_states.std(dim=(1,2), keepdim=True)
mean_controlnet_input, std_controlnet_input = controlnet_input.mean(dim=(1,2), keepdim=True), controlnet_input.std(dim=(1,2), keepdim=True)
controlnet_input = (controlnet_input - mean_controlnet_input) * (std_hidden_states / (std_controlnet_input + 1e-12)) + mean_hidden_states
mean_hidden_states, std_hidden_states = hidden_states.mean(
dim=(1, 2), keepdim=True
), hidden_states.std(dim=(1, 2), keepdim=True)
mean_controlnet_input, std_controlnet_input = controlnet_input.mean(
dim=(1, 2), keepdim=True
), controlnet_input.std(dim=(1, 2), keepdim=True)
controlnet_input = (controlnet_input - mean_controlnet_input) * (
std_hidden_states / (std_controlnet_input + 1e-12)
) + mean_hidden_states
return controlnet_input
@@ -45,17 +51,27 @@ class Qwen2RotaryEmbedding(nn.Module):
self.dim = dim
self.max_position_embeddings = max_position_embeddings
self.base = base
inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
inv_freq = 1.0 / (
self.base
** (
torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device)
/ self.dim
)
)
self.register_buffer("inv_freq", inv_freq, persistent=False)
# Build here to make `torch.jit.trace` work.
self._set_cos_sin_cache(
seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
seq_len=max_position_embeddings,
device=self.inv_freq.device,
dtype=torch.get_default_dtype(),
)
def _set_cos_sin_cache(self, seq_len, device, dtype):
self.max_seq_len_cached = seq_len
t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
t = torch.arange(
self.max_seq_len_cached, device=device, dtype=torch.int64
).type_as(self.inv_freq)
freqs = torch.outer(t, self.inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
@@ -82,8 +98,12 @@ class T2IFinalLayer(nn.Module):
def __init__(self, hidden_size, patch_size=[16, 1], out_channels=256):
super().__init__()
self.norm_final = nn.RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)
self.linear = nn.Linear(hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True)
self.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5)
self.linear = nn.Linear(
hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True
)
self.scale_shift_table = nn.Parameter(
torch.randn(2, hidden_size) / hidden_size**0.5
)
self.out_channels = out_channels
self.patch_size = patch_size
@@ -95,14 +115,28 @@ class T2IFinalLayer(nn.Module):
# 4 unpatchify
new_height, new_width = 1, hidden_states.size(1)
hidden_states = hidden_states.reshape(
shape=(hidden_states.shape[0], new_height, new_width, self.patch_size[0], self.patch_size[1], self.out_channels)
shape=(
hidden_states.shape[0],
new_height,
new_width,
self.patch_size[0],
self.patch_size[1],
self.out_channels,
)
).contiguous()
hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
output = hidden_states.reshape(
shape=(hidden_states.shape[0], self.out_channels, new_height * self.patch_size[0], new_width * self.patch_size[1])
shape=(
hidden_states.shape[0],
self.out_channels,
new_height * self.patch_size[0],
new_width * self.patch_size[1],
)
).contiguous()
if width > new_width:
output = torch.nn.functional.pad(output, (0, width - new_width, 0, 0), 'constant', 0)
output = torch.nn.functional.pad(
output, (0, width - new_width, 0, 0), "constant", 0
)
elif width < new_width:
output = output[:, :, :, :width]
return output
@@ -131,9 +165,25 @@ class PatchEmbed(nn.Module):
super().__init__()
patch_size_h, patch_size_w = patch_size
self.early_conv_layers = nn.Sequential(
nn.Conv2d(in_channels, in_channels*256, kernel_size=patch_size, stride=patch_size, padding=0, bias=bias),
torch.nn.GroupNorm(num_groups=32, num_channels=in_channels*256, eps=1e-6, affine=True),
nn.Conv2d(in_channels*256, embed_dim, kernel_size=1, stride=1, padding=0, bias=bias)
nn.Conv2d(
in_channels,
in_channels * 256,
kernel_size=patch_size,
stride=patch_size,
padding=0,
bias=bias,
),
torch.nn.GroupNorm(
num_groups=32, num_channels=in_channels * 256, eps=1e-6, affine=True
),
nn.Conv2d(
in_channels * 256,
embed_dim,
kernel_size=1,
stride=1,
padding=0,
bias=bias,
),
)
self.patch_size = patch_size
self.height, self.width = height // patch_size_h, width // patch_size_w
@@ -153,7 +203,9 @@ class Transformer2DModelOutput(BaseOutput):
proj_losses: Optional[Tuple[Tuple[str, torch.Tensor]]] = None
class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
class ACEStepTransformer2DModel(
ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin
):
_supports_gradient_checkpointing = True
@register_to_config
@@ -217,9 +269,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
)
self.num_layers = num_layers
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=self.inner_dim)
self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True))
self.time_proj = Timesteps(
num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0
)
self.timestep_embedder = TimestepEmbedding(
in_channels=256, time_embed_dim=self.inner_dim
)
self.t_block = nn.Sequential(
nn.SiLU(), nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True)
)
# speaker
self.speaker_embedder = nn.Linear(speaker_embedding_dim, self.inner_dim)
@@ -229,25 +287,30 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
# lyric
self.lyric_embs = nn.Embedding(lyric_encoder_vocab_size, lyric_hidden_size)
self.lyric_encoder = LyricEncoder(input_size=lyric_hidden_size, static_chunk_size=0)
self.lyric_encoder = LyricEncoder(
input_size=lyric_hidden_size, static_chunk_size=0
)
self.lyric_proj = nn.Linear(lyric_hidden_size, self.inner_dim)
projector_dim = 2 * self.inner_dim
self.projectors = nn.ModuleList([
nn.Sequential(
nn.Linear(self.inner_dim, projector_dim),
nn.SiLU(),
nn.Linear(projector_dim, projector_dim),
nn.SiLU(),
nn.Linear(projector_dim, ssl_dim),
) for ssl_dim in ssl_latent_dims
])
self.projectors = nn.ModuleList(
[
nn.Sequential(
nn.Linear(self.inner_dim, projector_dim),
nn.SiLU(),
nn.Linear(projector_dim, projector_dim),
nn.SiLU(),
nn.Linear(projector_dim, ssl_dim),
)
for ssl_dim in ssl_latent_dims
]
)
self.ssl_latent_dims = ssl_latent_dims
self.ssl_encoder_depths = ssl_encoder_depths
self.cosine_loss = torch.nn.CosineEmbeddingLoss(margin=0.0, reduction='mean')
self.cosine_loss = torch.nn.CosineEmbeddingLoss(margin=0.0, reduction="mean")
self.ssl_names = ssl_names
self.proj_in = PatchEmbed(
@@ -258,11 +321,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
bias=True,
)
self.final_layer = T2IFinalLayer(self.inner_dim, patch_size=patch_size, out_channels=out_channels)
self.final_layer = T2IFinalLayer(
self.inner_dim, patch_size=patch_size, out_channels=out_channels
)
self.gradient_checkpointing = False
# Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
def enable_forward_chunking(
self, chunk_size: Optional[int] = None, dim: int = 0
) -> None:
"""
Sets the attention processor to use [feed forward
chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
@@ -281,7 +348,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
# By default chunk size is 1
chunk_size = chunk_size or 1
def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
def fn_recursive_feed_forward(
module: torch.nn.Module, chunk_size: int, dim: int
):
if hasattr(module, "set_chunk_feed_forward"):
module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
@@ -302,7 +371,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
):
# N x T x D
lyric_embs = self.lyric_embs(lyric_token_idx)
prompt_prenet_out, _mask = self.lyric_encoder(lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1)
prompt_prenet_out, _mask = self.lyric_encoder(
lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1
)
prompt_prenet_out = self.lyric_proj(prompt_prenet_out)
return prompt_prenet_out
@@ -317,7 +388,7 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
bs = encoder_text_hidden_states.shape[0]
device = encoder_text_hidden_states.device
# speaker embedding
encoder_spk_hidden_states = self.speaker_embedder(speaker_embeds).unsqueeze(1)
speaker_mask = torch.ones(bs, 1, device=device)
@@ -331,8 +402,17 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
lyric_mask=lyric_mask,
)
encoder_hidden_states = torch.cat([encoder_spk_hidden_states, encoder_text_hidden_states, encoder_lyric_hidden_states], dim=1)
encoder_hidden_mask = torch.cat([speaker_mask, text_attention_mask, lyric_mask], dim=1)
encoder_hidden_states = torch.cat(
[
encoder_spk_hidden_states,
encoder_text_hidden_states,
encoder_lyric_hidden_states,
],
dim=1,
)
encoder_hidden_mask = torch.cat(
[speaker_mask, text_attention_mask, lyric_mask], dim=1
)
return encoder_hidden_states, encoder_hidden_mask
def decode(
@@ -344,12 +424,16 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
timestep: Optional[torch.Tensor],
ssl_hidden_states: Optional[List[torch.Tensor]] = None,
output_length: int = 0,
block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
block_controlnet_hidden_states: Optional[
Union[List[torch.Tensor], torch.Tensor]
] = None,
controlnet_scale: Union[float, torch.Tensor] = 1.0,
return_dict: bool = True,
):
embedded_timestep = self.timestep_embedder(self.time_proj(timestep).to(dtype=hidden_states.dtype))
embedded_timestep = self.timestep_embedder(
self.time_proj(timestep).to(dtype=hidden_states.dtype)
)
temb = self.t_block(embedded_timestep)
hidden_states = self.proj_in(hidden_states)
@@ -361,8 +445,12 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
inner_hidden_states = []
rotary_freqs_cis = self.rotary_emb(hidden_states, seq_len=hidden_states.shape[1])
encoder_rotary_freqs_cis = self.rotary_emb(encoder_hidden_states, seq_len=encoder_hidden_states.shape[1])
rotary_freqs_cis = self.rotary_emb(
hidden_states, seq_len=hidden_states.shape[1]
)
encoder_rotary_freqs_cis = self.rotary_emb(
encoder_hidden_states, seq_len=encoder_hidden_states.shape[1]
)
for index_block, block in enumerate(self.transformer_blocks):
@@ -377,7 +465,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
return custom_forward
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
ckpt_kwargs: Dict[str, Any] = (
{"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
)
hidden_states = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states=hidden_states,
@@ -406,9 +496,15 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
inner_hidden_states.append(hidden_states)
proj_losses = []
if len(inner_hidden_states) > 0 and ssl_hidden_states is not None and len(ssl_hidden_states) > 0:
if (
len(inner_hidden_states) > 0
and ssl_hidden_states is not None
and len(ssl_hidden_states) > 0
):
for inner_hidden_state, projector, ssl_hidden_state, ssl_name in zip(inner_hidden_states, self.projectors, ssl_hidden_states, self.ssl_names):
for inner_hidden_state, projector, ssl_hidden_state, ssl_name in zip(
inner_hidden_states, self.projectors, ssl_hidden_states, self.ssl_names
):
if ssl_hidden_state is None:
continue
# 1. N x T x D1 -> N x D x D2
@@ -416,9 +512,20 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
# 3. projection loss
bs = inner_hidden_state.shape[0]
proj_loss = 0.0
for i, (z, z_tilde) in enumerate(zip(ssl_hidden_state, est_ssl_hidden_state)):
for i, (z, z_tilde) in enumerate(
zip(ssl_hidden_state, est_ssl_hidden_state)
):
# 2. interpolate
z_tilde = F.interpolate(z_tilde.unsqueeze(0).transpose(1, 2), size=len(z), mode='linear', align_corners=False).transpose(1, 2).squeeze(0)
z_tilde = (
F.interpolate(
z_tilde.unsqueeze(0).transpose(1, 2),
size=len(z),
mode="linear",
align_corners=False,
)
.transpose(1, 2)
.squeeze(0)
)
z_tilde = torch.nn.functional.normalize(z_tilde, dim=-1)
z = torch.nn.functional.normalize(z, dim=-1)
@@ -445,7 +552,9 @@ class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromO
lyric_mask: Optional[torch.LongTensor] = None,
timestep: Optional[torch.Tensor] = None,
ssl_hidden_states: Optional[List[torch.Tensor]] = None,
block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
block_controlnet_hidden_states: Optional[
Union[List[torch.Tensor], torch.Tensor]
] = None,
controlnet_scale: Union[float, torch.Tensor] = 1.0,
return_dict: bool = True,
):