guysss/ACE-Step
0
1# Copyright 2024 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14from dataclasses import dataclass15from typing import Any, Dict, Optional, Tuple, List, Union16 17import torch18import torch.nn.functional as F19from torch import nn20 21from diffusers.configuration_utils import ConfigMixin, register_to_config22from diffusers.utils import BaseOutput, is_torch_version23from diffusers.models.modeling_utils import ModelMixin24from diffusers.models.embeddings import TimestepEmbedding, Timesteps25from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin26 27 28from .attention import LinearTransformerBlock, t2i_modulate29from .lyrics_utils.lyric_encoder import ConformerEncoder as LyricEncoder30 31 32def cross_norm(hidden_states, controlnet_input):33 # input N x T x c34 mean_hidden_states, std_hidden_states = hidden_states.mean(dim=(1,2), keepdim=True), hidden_states.std(dim=(1,2), keepdim=True)35 mean_controlnet_input, std_controlnet_input = controlnet_input.mean(dim=(1,2), keepdim=True), controlnet_input.std(dim=(1,2), keepdim=True)36 controlnet_input = (controlnet_input - mean_controlnet_input) * (std_hidden_states / (std_controlnet_input + 1e-12)) + mean_hidden_states37 return controlnet_input38 39 40# Copied from transformers.models.mixtral.modeling_mixtral.MixtralRotaryEmbedding with Mixtral->Qwen241class Qwen2RotaryEmbedding(nn.Module):42 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):43 super().__init__()44 45 self.dim = dim46 self.max_position_embeddings = max_position_embeddings47 self.base = base48 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))49 self.register_buffer("inv_freq", inv_freq, persistent=False)50 51 # Build here to make `torch.jit.trace` work.52 self._set_cos_sin_cache(53 seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()54 )55 56 def _set_cos_sin_cache(self, seq_len, device, dtype):57 self.max_seq_len_cached = seq_len58 t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)59 60 freqs = torch.outer(t, self.inv_freq)61 # Different from paper, but it uses a different permutation in order to obtain the same calculation62 emb = torch.cat((freqs, freqs), dim=-1)63 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)64 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)65 66 def forward(self, x, seq_len=None):67 # x: [bs, num_attention_heads, seq_len, head_size]68 if seq_len > self.max_seq_len_cached:69 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)70 71 return (72 self.cos_cached[:seq_len].to(dtype=x.dtype),73 self.sin_cached[:seq_len].to(dtype=x.dtype),74 )75 76 77class T2IFinalLayer(nn.Module):78 """79 The final layer of Sana.80 """81 82 def __init__(self, hidden_size, patch_size=[16, 1], out_channels=256):83 super().__init__()84 self.norm_final = nn.RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6)85 self.linear = nn.Linear(hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True)86 self.scale_shift_table = nn.Parameter(torch.randn(2, hidden_size) / hidden_size**0.5)87 self.out_channels = out_channels88 self.patch_size = patch_size89 90 def unpatchfy(91 self,92 hidden_states: torch.Tensor,93 width: int,94 ):95 # 4 unpatchify96 new_height, new_width = 1, hidden_states.size(1)97 hidden_states = hidden_states.reshape(98 shape=(hidden_states.shape[0], new_height, new_width, self.patch_size[0], self.patch_size[1], self.out_channels)99 ).contiguous()100 hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)101 output = hidden_states.reshape(102 shape=(hidden_states.shape[0], self.out_channels, new_height * self.patch_size[0], new_width * self.patch_size[1])103 ).contiguous()104 if width > new_width:105 output = torch.nn.functional.pad(output, (0, width - new_width, 0, 0), 'constant', 0)106 elif width < new_width:107 output = output[:, :, :, :width]108 return output109 110 def forward(self, x, t, output_length):111 shift, scale = (self.scale_shift_table[None] + t[:, None]).chunk(2, dim=1)112 x = t2i_modulate(self.norm_final(x), shift, scale)113 x = self.linear(x)114 # unpatchify115 output = self.unpatchfy(x, output_length)116 return output117 118 119class PatchEmbed(nn.Module):120 """2D Image to Patch Embedding"""121 122 def __init__(123 self,124 height=16,125 width=4096,126 patch_size=(16, 1),127 in_channels=8,128 embed_dim=1152,129 bias=True,130 ):131 super().__init__()132 patch_size_h, patch_size_w = patch_size133 self.early_conv_layers = nn.Sequential(134 nn.Conv2d(in_channels, in_channels*256, kernel_size=patch_size, stride=patch_size, padding=0, bias=bias),135 torch.nn.GroupNorm(num_groups=32, num_channels=in_channels*256, eps=1e-6, affine=True),136 nn.Conv2d(in_channels*256, embed_dim, kernel_size=1, stride=1, padding=0, bias=bias)137 )138 self.patch_size = patch_size139 self.height, self.width = height // patch_size_h, width // patch_size_w140 self.base_size = self.width141 142 def forward(self, latent):143 # early convolutions, N x C x H x W -> N x 256 * sqrt(patch_size) x H/patch_size x W/patch_size144 latent = self.early_conv_layers(latent)145 latent = latent.flatten(2).transpose(1, 2) # BCHW -> BNC146 return latent147 148 149@dataclass150class Transformer2DModelOutput(BaseOutput):151 152 sample: torch.FloatTensor153 proj_losses: Optional[Tuple[Tuple[str, torch.Tensor]]] = None154 155 156class ACEStepTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):157 _supports_gradient_checkpointing = True158 159 @register_to_config160 def __init__(161 self,162 in_channels: Optional[int] = 8,163 num_layers: int = 28,164 inner_dim: int = 1536,165 attention_head_dim: int = 64,166 num_attention_heads: int = 24,167 mlp_ratio: float = 4.0,168 out_channels: int = 8,169 max_position: int = 32768,170 rope_theta: float = 1000000.0,171 speaker_embedding_dim: int = 512,172 text_embedding_dim: int = 768,173 ssl_encoder_depths: List[int] = [9, 9],174 ssl_names: List[str] = ["mert", "m-hubert"],175 ssl_latent_dims: List[int] = [1024, 768],176 lyric_encoder_vocab_size: int = 6681,177 lyric_hidden_size: int = 1024,178 patch_size: List[int] = [16, 1],179 max_height: int = 16,180 max_width: int = 4096,181 **kwargs,182 ):183 super().__init__()184 185 self.num_attention_heads = num_attention_heads186 self.attention_head_dim = attention_head_dim187 inner_dim = num_attention_heads * attention_head_dim188 self.inner_dim = inner_dim189 self.out_channels = out_channels190 self.max_position = max_position191 self.patch_size = patch_size192 193 self.rope_theta = rope_theta194 195 self.rotary_emb = Qwen2RotaryEmbedding(196 dim=self.attention_head_dim,197 max_position_embeddings=self.max_position,198 base=self.rope_theta,199 )200 201 # 2. Define input layers202 self.in_channels = in_channels203 204 # 3. Define transformers blocks205 self.transformer_blocks = nn.ModuleList(206 [207 LinearTransformerBlock(208 dim=self.inner_dim,209 num_attention_heads=self.num_attention_heads,210 attention_head_dim=attention_head_dim,211 mlp_ratio=mlp_ratio,212 add_cross_attention=True,213 add_cross_attention_dim=self.inner_dim,214 )215 for i in range(self.config.num_layers)216 ]217 )218 self.num_layers = num_layers219 220 self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)221 self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=self.inner_dim)222 self.t_block = nn.Sequential(nn.SiLU(), nn.Linear(self.inner_dim, 6 * self.inner_dim, bias=True))223 224 # speaker225 self.speaker_embedder = nn.Linear(speaker_embedding_dim, self.inner_dim)226 227 # genre228 self.genre_embedder = nn.Linear(text_embedding_dim, self.inner_dim)229 230 # lyric231 self.lyric_embs = nn.Embedding(lyric_encoder_vocab_size, lyric_hidden_size)232 self.lyric_encoder = LyricEncoder(input_size=lyric_hidden_size, static_chunk_size=0)233 self.lyric_proj = nn.Linear(lyric_hidden_size, self.inner_dim)234 235 projector_dim = 2 * self.inner_dim236 237 self.projectors = nn.ModuleList([238 nn.Sequential(239 nn.Linear(self.inner_dim, projector_dim),240 nn.SiLU(),241 nn.Linear(projector_dim, projector_dim),242 nn.SiLU(),243 nn.Linear(projector_dim, ssl_dim),244 ) for ssl_dim in ssl_latent_dims245 ])246 247 self.ssl_latent_dims = ssl_latent_dims248 self.ssl_encoder_depths = ssl_encoder_depths249 250 self.cosine_loss = torch.nn.CosineEmbeddingLoss(margin=0.0, reduction='mean')251 self.ssl_names = ssl_names252 253 self.proj_in = PatchEmbed(254 height=max_height,255 width=max_width,256 patch_size=patch_size,257 embed_dim=self.inner_dim,258 bias=True,259 )260 261 self.final_layer = T2IFinalLayer(self.inner_dim, patch_size=patch_size, out_channels=out_channels)262 self.gradient_checkpointing = False263 264 # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking265 def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:266 """267 Sets the attention processor to use [feed forward268 chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).269 270 Parameters:271 chunk_size (`int`, *optional*):272 The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually273 over each tensor of dim=`dim`.274 dim (`int`, *optional*, defaults to `0`):275 The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)276 or dim=1 (sequence length).277 """278 if dim not in [0, 1]:279 raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")280 281 # By default chunk size is 1282 chunk_size = chunk_size or 1283 284 def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):285 if hasattr(module, "set_chunk_feed_forward"):286 module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)287 288 for child in module.children():289 fn_recursive_feed_forward(child, chunk_size, dim)290 291 for module in self.children():292 fn_recursive_feed_forward(module, chunk_size, dim)293 294 def _set_gradient_checkpointing(self, module, value=False):295 if hasattr(module, "gradient_checkpointing"):296 module.gradient_checkpointing = value297 298 def forward_lyric_encoder(299 self,300 lyric_token_idx: Optional[torch.LongTensor] = None,301 lyric_mask: Optional[torch.LongTensor] = None,302 ):303 # N x T x D304 lyric_embs = self.lyric_embs(lyric_token_idx)305 prompt_prenet_out, _mask = self.lyric_encoder(lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1)306 prompt_prenet_out = self.lyric_proj(prompt_prenet_out)307 return prompt_prenet_out308 309 def encode(310 self,311 encoder_text_hidden_states: Optional[torch.Tensor] = None,312 text_attention_mask: Optional[torch.LongTensor] = None,313 speaker_embeds: Optional[torch.FloatTensor] = None,314 lyric_token_idx: Optional[torch.LongTensor] = None,315 lyric_mask: Optional[torch.LongTensor] = None,316 ):317 318 bs = encoder_text_hidden_states.shape[0]319 device = encoder_text_hidden_states.device320 321 # speaker embedding322 encoder_spk_hidden_states = self.speaker_embedder(speaker_embeds).unsqueeze(1)323 speaker_mask = torch.ones(bs, 1, device=device)324 325 # genre embedding326 encoder_text_hidden_states = self.genre_embedder(encoder_text_hidden_states)327 328 # lyric329 encoder_lyric_hidden_states = self.forward_lyric_encoder(330 lyric_token_idx=lyric_token_idx,331 lyric_mask=lyric_mask,332 )333 334 encoder_hidden_states = torch.cat([encoder_spk_hidden_states, encoder_text_hidden_states, encoder_lyric_hidden_states], dim=1)335 encoder_hidden_mask = torch.cat([speaker_mask, text_attention_mask, lyric_mask], dim=1)336 return encoder_hidden_states, encoder_hidden_mask337 338 def decode(339 self,340 hidden_states: torch.Tensor,341 attention_mask: torch.Tensor,342 encoder_hidden_states: torch.Tensor,343 encoder_hidden_mask: torch.Tensor,344 timestep: Optional[torch.Tensor],345 ssl_hidden_states: Optional[List[torch.Tensor]] = None,346 output_length: int = 0,347 block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,348 controlnet_scale: Union[float, torch.Tensor] = 1.0,349 return_dict: bool = True,350 ):351 352 embedded_timestep = self.timestep_embedder(self.time_proj(timestep).to(dtype=hidden_states.dtype))353 temb = self.t_block(embedded_timestep)354 355 hidden_states = self.proj_in(hidden_states)356 357 # controlnet logic358 if block_controlnet_hidden_states is not None:359 control_condi = cross_norm(hidden_states, block_controlnet_hidden_states)360 hidden_states = hidden_states + control_condi * controlnet_scale361 362 inner_hidden_states = []363 364 rotary_freqs_cis = self.rotary_emb(hidden_states, seq_len=hidden_states.shape[1])365 encoder_rotary_freqs_cis = self.rotary_emb(encoder_hidden_states, seq_len=encoder_hidden_states.shape[1])366 367 for index_block, block in enumerate(self.transformer_blocks):368 369 if self.training and self.gradient_checkpointing:370 371 def create_custom_forward(module, return_dict=None):372 def custom_forward(*inputs):373 if return_dict is not None:374 return module(*inputs, return_dict=return_dict)375 else:376 return module(*inputs)377 378 return custom_forward379 380 ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}381 hidden_states = torch.utils.checkpoint.checkpoint(382 create_custom_forward(block),383 hidden_states=hidden_states,384 attention_mask=attention_mask,385 encoder_hidden_states=encoder_hidden_states,386 encoder_attention_mask=encoder_hidden_mask,387 rotary_freqs_cis=rotary_freqs_cis,388 rotary_freqs_cis_cross=encoder_rotary_freqs_cis,389 temb=temb,390 **ckpt_kwargs,391 )392 393 else:394 hidden_states = block(395 hidden_states=hidden_states,396 attention_mask=attention_mask,397 encoder_hidden_states=encoder_hidden_states,398 encoder_attention_mask=encoder_hidden_mask,399 rotary_freqs_cis=rotary_freqs_cis,400 rotary_freqs_cis_cross=encoder_rotary_freqs_cis,401 temb=temb,402 )403 404 for ssl_encoder_depth in self.ssl_encoder_depths:405 if index_block == ssl_encoder_depth:406 inner_hidden_states.append(hidden_states)407 408 proj_losses = []409 if len(inner_hidden_states) > 0 and ssl_hidden_states is not None and len(ssl_hidden_states) > 0:410 411 for inner_hidden_state, projector, ssl_hidden_state, ssl_name in zip(inner_hidden_states, self.projectors, ssl_hidden_states, self.ssl_names):412 if ssl_hidden_state is None:413 continue414 # 1. N x T x D1 -> N x D x D2415 est_ssl_hidden_state = projector(inner_hidden_state)416 # 3. projection loss417 bs = inner_hidden_state.shape[0]418 proj_loss = 0.0419 for i, (z, z_tilde) in enumerate(zip(ssl_hidden_state, est_ssl_hidden_state)):420 # 2. interpolate421 z_tilde = F.interpolate(z_tilde.unsqueeze(0).transpose(1, 2), size=len(z), mode='linear', align_corners=False).transpose(1, 2).squeeze(0)422 423 z_tilde = torch.nn.functional.normalize(z_tilde, dim=-1)424 z = torch.nn.functional.normalize(z, dim=-1)425 # T x d -> T x 1 -> 1426 target = torch.ones(z.shape[0], device=z.device)427 proj_loss += self.cosine_loss(z, z_tilde, target)428 proj_losses.append((ssl_name, proj_loss / bs))429 430 output = self.final_layer(hidden_states, embedded_timestep, output_length)431 if not return_dict:432 return (output, proj_losses)433 434 return Transformer2DModelOutput(sample=output, proj_losses=proj_losses)435 436 # @torch.compile437 def forward(438 self,439 hidden_states: torch.Tensor,440 attention_mask: torch.Tensor,441 encoder_text_hidden_states: Optional[torch.Tensor] = None,442 text_attention_mask: Optional[torch.LongTensor] = None,443 speaker_embeds: Optional[torch.FloatTensor] = None,444 lyric_token_idx: Optional[torch.LongTensor] = None,445 lyric_mask: Optional[torch.LongTensor] = None,446 timestep: Optional[torch.Tensor] = None,447 ssl_hidden_states: Optional[List[torch.Tensor]] = None,448 block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,449 controlnet_scale: Union[float, torch.Tensor] = 1.0,450 return_dict: bool = True,451 ):452 encoder_hidden_states, encoder_hidden_mask = self.encode(453 encoder_text_hidden_states=encoder_text_hidden_states,454 text_attention_mask=text_attention_mask,455 speaker_embeds=speaker_embeds,456 lyric_token_idx=lyric_token_idx,457 lyric_mask=lyric_mask,458 )459 460 output_length = hidden_states.shape[-1]461 462 output = self.decode(463 hidden_states=hidden_states,464 attention_mask=attention_mask,465 encoder_hidden_states=encoder_hidden_states,466 encoder_hidden_mask=encoder_hidden_mask,467 timestep=timestep,468 ssl_hidden_states=ssl_hidden_states,469 output_length=output_length,470 block_controlnet_hidden_states=block_controlnet_hidden_states,471 controlnet_scale=controlnet_scale,472 return_dict=return_dict,473 )474 475 return output476