guysss/ACE-Step
0
1import random2import time3import os4import re5import spaces6import torch7import torch.nn as nn8from loguru import logger9from tqdm import tqdm10import json11import math12from huggingface_hub import hf_hub_download, snapshot_download13 14# from diffusers.pipelines.pipeline_utils import DiffusionPipeline15from schedulers.scheduling_flow_match_euler_discrete import (16 FlowMatchEulerDiscreteScheduler,17)18from schedulers.scheduling_flow_match_heun_discrete import (19 FlowMatchHeunDiscreteScheduler,20)21from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import (22 retrieve_timesteps,23)24from diffusers.utils.torch_utils import randn_tensor25from transformers import UMT5EncoderModel, AutoTokenizer26 27from language_segmentation import LangSegment28from music_dcae.music_dcae_pipeline import MusicDCAE29from models.ace_step_transformer import ACEStepTransformer2DModel30from models.lyrics_utils.lyric_tokenizer import VoiceBpeTokenizer31from apg_guidance import (32 apg_forward,33 MomentumBuffer,34 cfg_forward,35 cfg_zero_star,36 cfg_double_condition_forward,37)38import torchaudio39import torio40 41 42torch.backends.cudnn.benchmark = False43torch.set_float32_matmul_precision("high")44torch.backends.cudnn.deterministic = True45torch.backends.cuda.matmul.allow_tf32 = True46os.environ["TOKENIZERS_PARALLELISM"] = "false"47 48 49SUPPORT_LANGUAGES = {50 "en": 259,51 "de": 260,52 "fr": 262,53 "es": 284,54 "it": 285,55 "pt": 286,56 "pl": 294,57 "tr": 295,58 "ru": 267,59 "cs": 293,60 "nl": 297,61 "ar": 5022,62 "zh": 5023,63 "ja": 5412,64 "hu": 5753,65 "ko": 6152,66 "hi": 6680,67}68 69structure_pattern = re.compile(r"\[.*?\]")70 71 72def ensure_directory_exists(directory):73 directory = str(directory)74 if not os.path.exists(directory):75 os.makedirs(directory)76 77 78REPO_ID = "ACE-Step/ACE-Step-v1-3.5B"79 80 81# class ACEStepPipeline(DiffusionPipeline):82class ACEStepPipeline:83 84 def __init__(85 self,86 checkpoint_dir=None,87 device_id=0,88 dtype="bfloat16",89 text_encoder_checkpoint_path=None,90 persistent_storage_path=None,91 torch_compile=False,92 **kwargs,93 ):94 if not checkpoint_dir:95 if persistent_storage_path is None:96 checkpoint_dir = os.path.join(os.path.dirname(__file__), "checkpoints")97 else:98 checkpoint_dir = os.path.join(persistent_storage_path, "checkpoints")99 ensure_directory_exists(checkpoint_dir)100 self.checkpoint_dir = checkpoint_dir101 device = (102 torch.device(f"cuda:{device_id}")103 if torch.cuda.is_available()104 else torch.device("cpu")105 )106 if device.type == "cpu" and torch.backends.mps.is_available():107 device = torch.device("mps")108 self.dtype = torch.bfloat16 if dtype == "bfloat16" else torch.float32109 if device.type == "mps":110 self.dtype = torch.float32111 self.device = device112 self.loaded = False113 self.torch_compile = torch_compile114 self.lora_path = "none"115 116 def load_lora(self, lora_name_or_path):117 if lora_name_or_path != self.lora_path and lora_name_or_path != "none":118 if not os.path.exists(lora_name_or_path):119 lora_download_path = snapshot_download(120 lora_name_or_path, cache_dir=self.checkpoint_dir121 )122 else:123 lora_download_path = lora_name_or_path124 if self.lora_path != "none":125 self.ace_step_transformer.unload_lora()126 self.ace_step_transformer.load_lora_adapter(127 os.path.join(lora_download_path, "pytorch_lora_weights.safetensors"),128 adapter_name="zh_rap_lora",129 with_alpha=True,130 )131 logger.info(132 f"Loading lora weights from: {lora_name_or_path} download path is: {lora_download_path}"133 )134 self.lora_path = lora_name_or_path135 elif self.lora_path != "none" and lora_name_or_path == "none":136 logger.info("No lora weights to load.")137 self.ace_step_transformer.unload_lora()138 139 def load_checkpoint(self, checkpoint_dir=None):140 device = self.device141 142 dcae_model_path = os.path.join(checkpoint_dir, "music_dcae_f8c8")143 vocoder_model_path = os.path.join(checkpoint_dir, "music_vocoder")144 ace_step_model_path = os.path.join(checkpoint_dir, "ace_step_transformer")145 text_encoder_model_path = os.path.join(checkpoint_dir, "umt5-base")146 147 files_exist = (148 os.path.exists(os.path.join(dcae_model_path, "config.json"))149 and os.path.exists(150 os.path.join(dcae_model_path, "diffusion_pytorch_model.safetensors")151 )152 and os.path.exists(os.path.join(vocoder_model_path, "config.json"))153 and os.path.exists(154 os.path.join(vocoder_model_path, "diffusion_pytorch_model.safetensors")155 )156 and os.path.exists(os.path.join(ace_step_model_path, "config.json"))157 and os.path.exists(158 os.path.join(ace_step_model_path, "diffusion_pytorch_model.safetensors")159 )160 and os.path.exists(os.path.join(text_encoder_model_path, "config.json"))161 and os.path.exists(162 os.path.join(text_encoder_model_path, "model.safetensors")163 )164 and os.path.exists(165 os.path.join(text_encoder_model_path, "special_tokens_map.json")166 )167 and os.path.exists(168 os.path.join(text_encoder_model_path, "tokenizer_config.json")169 )170 and os.path.exists(os.path.join(text_encoder_model_path, "tokenizer.json"))171 )172 173 if not files_exist:174 logger.info(175 f"Checkpoint directory {checkpoint_dir} is not complete, downloading from Hugging Face Hub"176 )177 178 # download music dcae model179 os.makedirs(dcae_model_path, exist_ok=True)180 hf_hub_download(181 repo_id=REPO_ID,182 subfolder="music_dcae_f8c8",183 filename="config.json",184 local_dir=checkpoint_dir,185 local_dir_use_symlinks=False,186 )187 hf_hub_download(188 repo_id=REPO_ID,189 subfolder="music_dcae_f8c8",190 filename="diffusion_pytorch_model.safetensors",191 local_dir=checkpoint_dir,192 local_dir_use_symlinks=False,193 )194 195 # download vocoder model196 os.makedirs(vocoder_model_path, exist_ok=True)197 hf_hub_download(198 repo_id=REPO_ID,199 subfolder="music_vocoder",200 filename="config.json",201 local_dir=checkpoint_dir,202 local_dir_use_symlinks=False,203 )204 hf_hub_download(205 repo_id=REPO_ID,206 subfolder="music_vocoder",207 filename="diffusion_pytorch_model.safetensors",208 local_dir=checkpoint_dir,209 local_dir_use_symlinks=False,210 )211 212 # download ace_step transformer model213 os.makedirs(ace_step_model_path, exist_ok=True)214 hf_hub_download(215 repo_id=REPO_ID,216 subfolder="ace_step_transformer",217 filename="config.json",218 local_dir=checkpoint_dir,219 local_dir_use_symlinks=False,220 )221 hf_hub_download(222 repo_id=REPO_ID,223 subfolder="ace_step_transformer",224 filename="diffusion_pytorch_model.safetensors",225 local_dir=checkpoint_dir,226 local_dir_use_symlinks=False,227 )228 229 # download text encoder model230 os.makedirs(text_encoder_model_path, exist_ok=True)231 hf_hub_download(232 repo_id=REPO_ID,233 subfolder="umt5-base",234 filename="config.json",235 local_dir=checkpoint_dir,236 local_dir_use_symlinks=False,237 )238 hf_hub_download(239 repo_id=REPO_ID,240 subfolder="umt5-base",241 filename="model.safetensors",242 local_dir=checkpoint_dir,243 local_dir_use_symlinks=False,244 )245 hf_hub_download(246 repo_id=REPO_ID,247 subfolder="umt5-base",248 filename="special_tokens_map.json",249 local_dir=checkpoint_dir,250 local_dir_use_symlinks=False,251 )252 hf_hub_download(253 repo_id=REPO_ID,254 subfolder="umt5-base",255 filename="tokenizer_config.json",256 local_dir=checkpoint_dir,257 local_dir_use_symlinks=False,258 )259 hf_hub_download(260 repo_id=REPO_ID,261 subfolder="umt5-base",262 filename="tokenizer.json",263 local_dir=checkpoint_dir,264 local_dir_use_symlinks=False,265 )266 267 logger.info("Models downloaded")268 269 dcae_checkpoint_path = dcae_model_path270 vocoder_checkpoint_path = vocoder_model_path271 ace_step_checkpoint_path = ace_step_model_path272 text_encoder_checkpoint_path = text_encoder_model_path273 274 self.music_dcae = MusicDCAE(275 dcae_checkpoint_path=dcae_checkpoint_path,276 vocoder_checkpoint_path=vocoder_checkpoint_path,277 )278 self.music_dcae.to(device).eval().to(self.dtype)279 280 self.ace_step_transformer = ACEStepTransformer2DModel.from_pretrained(281 ace_step_checkpoint_path, torch_dtype=self.dtype282 )283 self.ace_step_transformer.to(device).eval().to(self.dtype)284 285 lang_segment = LangSegment()286 287 lang_segment.setfilters(288 [289 "af",290 "am",291 "an",292 "ar",293 "as",294 "az",295 "be",296 "bg",297 "bn",298 "br",299 "bs",300 "ca",301 "cs",302 "cy",303 "da",304 "de",305 "dz",306 "el",307 "en",308 "eo",309 "es",310 "et",311 "eu",312 "fa",313 "fi",314 "fo",315 "fr",316 "ga",317 "gl",318 "gu",319 "he",320 "hi",321 "hr",322 "ht",323 "hu",324 "hy",325 "id",326 "is",327 "it",328 "ja",329 "jv",330 "ka",331 "kk",332 "km",333 "kn",334 "ko",335 "ku",336 "ky",337 "la",338 "lb",339 "lo",340 "lt",341 "lv",342 "mg",343 "mk",344 "ml",345 "mn",346 "mr",347 "ms",348 "mt",349 "nb",350 "ne",351 "nl",352 "nn",353 "no",354 "oc",355 "or",356 "pa",357 "pl",358 "ps",359 "pt",360 "qu",361 "ro",362 "ru",363 "rw",364 "se",365 "si",366 "sk",367 "sl",368 "sq",369 "sr",370 "sv",371 "sw",372 "ta",373 "te",374 "th",375 "tl",376 "tr",377 "ug",378 "uk",379 "ur",380 "vi",381 "vo",382 "wa",383 "xh",384 "zh",385 "zu",386 ]387 )388 self.lang_segment = lang_segment389 self.lyric_tokenizer = VoiceBpeTokenizer()390 text_encoder_model = UMT5EncoderModel.from_pretrained(391 text_encoder_checkpoint_path, torch_dtype=self.dtype392 ).eval()393 text_encoder_model = text_encoder_model.to(device).to(self.dtype)394 text_encoder_model.requires_grad_(False)395 self.text_encoder_model = text_encoder_model396 self.text_tokenizer = AutoTokenizer.from_pretrained(397 text_encoder_checkpoint_path398 )399 self.loaded = True400 401 # compile402 if self.torch_compile:403 self.music_dcae = torch.compile(self.music_dcae)404 self.ace_step_transformer = torch.compile(self.ace_step_transformer)405 self.text_encoder_model = torch.compile(self.text_encoder_model)406 407 def get_text_embeddings(self, texts, device, text_max_length=256):408 inputs = self.text_tokenizer(409 texts,410 return_tensors="pt",411 padding=True,412 truncation=True,413 max_length=text_max_length,414 )415 inputs = {key: value.to(device) for key, value in inputs.items()}416 if self.text_encoder_model.device != device:417 self.text_encoder_model.to(device)418 with torch.no_grad():419 outputs = self.text_encoder_model(**inputs)420 last_hidden_states = outputs.last_hidden_state421 attention_mask = inputs["attention_mask"]422 return last_hidden_states, attention_mask423 424 def get_text_embeddings_null(425 self, texts, device, text_max_length=256, tau=0.01, l_min=8, l_max=10426 ):427 inputs = self.text_tokenizer(428 texts,429 return_tensors="pt",430 padding=True,431 truncation=True,432 max_length=text_max_length,433 )434 inputs = {key: value.to(device) for key, value in inputs.items()}435 if self.text_encoder_model.device != device:436 self.text_encoder_model.to(device)437 438 def forward_with_temperature(inputs, tau=0.01, l_min=8, l_max=10):439 handlers = []440 441 def hook(module, input, output):442 output[:] *= tau443 return output444 445 for i in range(l_min, l_max):446 handler = (447 self.text_encoder_model.encoder.block[i]448 .layer[0]449 .SelfAttention.q.register_forward_hook(hook)450 )451 handlers.append(handler)452 453 with torch.no_grad():454 outputs = self.text_encoder_model(**inputs)455 last_hidden_states = outputs.last_hidden_state456 457 for hook in handlers:458 hook.remove()459 460 return last_hidden_states461 462 last_hidden_states = forward_with_temperature(inputs, tau, l_min, l_max)463 return last_hidden_states464 465 def set_seeds(self, batch_size, manual_seeds=None):466 processed_input_seeds = None467 if manual_seeds is not None:468 if isinstance(manual_seeds, str):469 if "," in manual_seeds:470 processed_input_seeds = list(map(int, manual_seeds.split(",")))471 elif manual_seeds.isdigit():472 processed_input_seeds = int(manual_seeds)473 elif isinstance(manual_seeds, list) and all(474 isinstance(s, int) for s in manual_seeds475 ):476 if len(manual_seeds) > 0:477 processed_input_seeds = list(manual_seeds)478 elif isinstance(manual_seeds, int):479 processed_input_seeds = manual_seeds480 random_generators = [481 torch.Generator(device=self.device) for _ in range(batch_size)482 ]483 actual_seeds = []484 for i in range(batch_size):485 current_seed_for_generator = None486 if processed_input_seeds is None:487 current_seed_for_generator = torch.randint(0, 2**32, (1,)).item()488 elif isinstance(processed_input_seeds, int):489 current_seed_for_generator = processed_input_seeds490 elif isinstance(processed_input_seeds, list):491 if i < len(processed_input_seeds):492 current_seed_for_generator = processed_input_seeds[i]493 else:494 current_seed_for_generator = processed_input_seeds[-1]495 if current_seed_for_generator is None:496 current_seed_for_generator = torch.randint(0, 2**32, (1,)).item()497 random_generators[i].manual_seed(current_seed_for_generator)498 actual_seeds.append(current_seed_for_generator)499 return random_generators, actual_seeds500 501 def get_lang(self, text):502 language = "en"503 try:504 _ = self.lang_segment.getTexts(text)505 langCounts = self.lang_segment.getCounts()506 language = langCounts[0][0]507 if len(langCounts) > 1 and language == "en":508 language = langCounts[1][0]509 except Exception as err:510 language = "en"511 return language512 513 def tokenize_lyrics(self, lyrics, debug=False):514 lines = lyrics.split("\n")515 lyric_token_idx = [261]516 for line in lines:517 line = line.strip()518 if not line:519 lyric_token_idx += [2]520 continue521 522 lang = self.get_lang(line)523 524 if lang not in SUPPORT_LANGUAGES:525 lang = "en"526 if "zh" in lang:527 lang = "zh"528 if "spa" in lang:529 lang = "es"530 531 try:532 if structure_pattern.match(line):533 token_idx = self.lyric_tokenizer.encode(line, "en")534 else:535 token_idx = self.lyric_tokenizer.encode(line, lang)536 if debug:537 toks = self.lyric_tokenizer.batch_decode(538 [[tok_id] for tok_id in token_idx]539 )540 logger.info(f"debbug {line} --> {lang} --> {toks}")541 lyric_token_idx = lyric_token_idx + token_idx + [2]542 except Exception as e:543 print("tokenize error", e, "for line", line, "major_language", lang)544 return lyric_token_idx545 546 def calc_v(547 self,548 zt_src,549 zt_tar,550 t,551 encoder_text_hidden_states,552 text_attention_mask,553 target_encoder_text_hidden_states,554 target_text_attention_mask,555 speaker_embds,556 target_speaker_embeds,557 lyric_token_ids,558 lyric_mask,559 target_lyric_token_ids,560 target_lyric_mask,561 do_classifier_free_guidance=False,562 guidance_scale=1.0,563 target_guidance_scale=1.0,564 cfg_type="apg",565 attention_mask=None,566 momentum_buffer=None,567 momentum_buffer_tar=None,568 return_src_pred=True,569 ):570 noise_pred_src = None571 if return_src_pred:572 src_latent_model_input = (573 torch.cat([zt_src, zt_src]) if do_classifier_free_guidance else zt_src574 )575 timestep = t.expand(src_latent_model_input.shape[0])576 # source577 noise_pred_src = self.ace_step_transformer(578 hidden_states=src_latent_model_input,579 attention_mask=attention_mask,580 encoder_text_hidden_states=encoder_text_hidden_states,581 text_attention_mask=text_attention_mask,582 speaker_embeds=speaker_embds,583 lyric_token_idx=lyric_token_ids,584 lyric_mask=lyric_mask,585 timestep=timestep,586 ).sample587 588 if do_classifier_free_guidance:589 noise_pred_with_cond_src, noise_pred_uncond_src = noise_pred_src.chunk(590 2591 )592 if cfg_type == "apg":593 noise_pred_src = apg_forward(594 pred_cond=noise_pred_with_cond_src,595 pred_uncond=noise_pred_uncond_src,596 guidance_scale=guidance_scale,597 momentum_buffer=momentum_buffer,598 )599 elif cfg_type == "cfg":600 noise_pred_src = cfg_forward(601 cond_output=noise_pred_with_cond_src,602 uncond_output=noise_pred_uncond_src,603 cfg_strength=guidance_scale,604 )605 606 tar_latent_model_input = (607 torch.cat([zt_tar, zt_tar]) if do_classifier_free_guidance else zt_tar608 )609 timestep = t.expand(tar_latent_model_input.shape[0])610 # target611 noise_pred_tar = self.ace_step_transformer(612 hidden_states=tar_latent_model_input,613 attention_mask=attention_mask,614 encoder_text_hidden_states=target_encoder_text_hidden_states,615 text_attention_mask=target_text_attention_mask,616 speaker_embeds=target_speaker_embeds,617 lyric_token_idx=target_lyric_token_ids,618 lyric_mask=target_lyric_mask,619 timestep=timestep,620 ).sample621 622 if do_classifier_free_guidance:623 noise_pred_with_cond_tar, noise_pred_uncond_tar = noise_pred_tar.chunk(2)624 if cfg_type == "apg":625 noise_pred_tar = apg_forward(626 pred_cond=noise_pred_with_cond_tar,627 pred_uncond=noise_pred_uncond_tar,628 guidance_scale=target_guidance_scale,629 momentum_buffer=momentum_buffer_tar,630 )631 elif cfg_type == "cfg":632 noise_pred_tar = cfg_forward(633 cond_output=noise_pred_with_cond_tar,634 uncond_output=noise_pred_uncond_tar,635 cfg_strength=target_guidance_scale,636 )637 return noise_pred_src, noise_pred_tar638 639 @torch.no_grad()640 def flowedit_diffusion_process(641 self,642 encoder_text_hidden_states,643 text_attention_mask,644 speaker_embds,645 lyric_token_ids,646 lyric_mask,647 target_encoder_text_hidden_states,648 target_text_attention_mask,649 target_speaker_embeds,650 target_lyric_token_ids,651 target_lyric_mask,652 src_latents,653 random_generators=None,654 infer_steps=60,655 guidance_scale=15.0,656 n_min=0,657 n_max=1.0,658 n_avg=1,659 ):660 661 do_classifier_free_guidance = True662 if guidance_scale == 0.0 or guidance_scale == 1.0:663 do_classifier_free_guidance = False664 665 target_guidance_scale = guidance_scale666 device = encoder_text_hidden_states.device667 dtype = encoder_text_hidden_states.dtype668 bsz = encoder_text_hidden_states.shape[0]669 670 scheduler = FlowMatchEulerDiscreteScheduler(671 num_train_timesteps=1000,672 shift=3.0,673 )674 675 T_steps = infer_steps676 frame_length = src_latents.shape[-1]677 attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)678 679 timesteps, T_steps = retrieve_timesteps(680 scheduler, T_steps, device, timesteps=None681 )682 683 if do_classifier_free_guidance:684 attention_mask = torch.cat([attention_mask] * 2, dim=0)685 686 encoder_text_hidden_states = torch.cat(687 [688 encoder_text_hidden_states,689 torch.zeros_like(encoder_text_hidden_states),690 ],691 0,692 )693 text_attention_mask = torch.cat([text_attention_mask] * 2, dim=0)694 695 target_encoder_text_hidden_states = torch.cat(696 [697 target_encoder_text_hidden_states,698 torch.zeros_like(target_encoder_text_hidden_states),699 ],700 0,701 )702 target_text_attention_mask = torch.cat(703 [target_text_attention_mask] * 2, dim=0704 )705 706 speaker_embds = torch.cat(707 [speaker_embds, torch.zeros_like(speaker_embds)], 0708 )709 target_speaker_embeds = torch.cat(710 [target_speaker_embeds, torch.zeros_like(target_speaker_embeds)], 0711 )712 713 lyric_token_ids = torch.cat(714 [lyric_token_ids, torch.zeros_like(lyric_token_ids)], 0715 )716 lyric_mask = torch.cat([lyric_mask, torch.zeros_like(lyric_mask)], 0)717 718 target_lyric_token_ids = torch.cat(719 [target_lyric_token_ids, torch.zeros_like(target_lyric_token_ids)], 0720 )721 target_lyric_mask = torch.cat(722 [target_lyric_mask, torch.zeros_like(target_lyric_mask)], 0723 )724 725 momentum_buffer = MomentumBuffer()726 momentum_buffer_tar = MomentumBuffer()727 x_src = src_latents728 zt_edit = x_src.clone()729 xt_tar = None730 n_min = int(infer_steps * n_min)731 n_max = int(infer_steps * n_max)732 733 logger.info("flowedit start from {} to {}".format(n_min, n_max))734 735 for i, t in tqdm(enumerate(timesteps), total=T_steps):736 737 if i < n_min:738 continue739 740 t_i = t / 1000741 742 if i + 1 < len(timesteps):743 t_im1 = (timesteps[i + 1]) / 1000744 else:745 t_im1 = torch.zeros_like(t_i).to(t_i.device)746 747 if i < n_max:748 # Calculate the average of the V predictions749 V_delta_avg = torch.zeros_like(x_src)750 for k in range(n_avg):751 fwd_noise = randn_tensor(752 shape=x_src.shape,753 generator=random_generators,754 device=device,755 dtype=dtype,756 )757 758 zt_src = (1 - t_i) * x_src + (t_i) * fwd_noise759 760 zt_tar = zt_edit + zt_src - x_src761 762 Vt_src, Vt_tar = self.calc_v(763 zt_src=zt_src,764 zt_tar=zt_tar,765 t=t,766 encoder_text_hidden_states=encoder_text_hidden_states,767 text_attention_mask=text_attention_mask,768 target_encoder_text_hidden_states=target_encoder_text_hidden_states,769 target_text_attention_mask=target_text_attention_mask,770 speaker_embds=speaker_embds,771 target_speaker_embeds=target_speaker_embeds,772 lyric_token_ids=lyric_token_ids,773 lyric_mask=lyric_mask,774 target_lyric_token_ids=target_lyric_token_ids,775 target_lyric_mask=target_lyric_mask,776 do_classifier_free_guidance=do_classifier_free_guidance,777 guidance_scale=guidance_scale,778 target_guidance_scale=target_guidance_scale,779 attention_mask=attention_mask,780 momentum_buffer=momentum_buffer,781 )782 V_delta_avg += (1 / n_avg) * (783 Vt_tar - Vt_src784 ) # - (hfg-1)*( x_src))785 786 # propagate direct ODE787 zt_edit = zt_edit.to(torch.float32)788 zt_edit = zt_edit + (t_im1 - t_i) * V_delta_avg789 zt_edit = zt_edit.to(V_delta_avg.dtype)790 else: # i >= T_steps-n_min # regular sampling for last n_min steps791 if i == n_max:792 fwd_noise = randn_tensor(793 shape=x_src.shape,794 generator=random_generators,795 device=device,796 dtype=dtype,797 )798 scheduler._init_step_index(t)799 sigma = scheduler.sigmas[scheduler.step_index]800 xt_src = sigma * fwd_noise + (1.0 - sigma) * x_src801 xt_tar = zt_edit + xt_src - x_src802 803 _, Vt_tar = self.calc_v(804 zt_src=None,805 zt_tar=xt_tar,806 t=t,807 encoder_text_hidden_states=encoder_text_hidden_states,808 text_attention_mask=text_attention_mask,809 target_encoder_text_hidden_states=target_encoder_text_hidden_states,810 target_text_attention_mask=target_text_attention_mask,811 speaker_embds=speaker_embds,812 target_speaker_embeds=target_speaker_embeds,813 lyric_token_ids=lyric_token_ids,814 lyric_mask=lyric_mask,815 target_lyric_token_ids=target_lyric_token_ids,816 target_lyric_mask=target_lyric_mask,817 do_classifier_free_guidance=do_classifier_free_guidance,818 guidance_scale=guidance_scale,819 target_guidance_scale=target_guidance_scale,820 attention_mask=attention_mask,821 momentum_buffer_tar=momentum_buffer_tar,822 return_src_pred=False,823 )824 825 dtype = Vt_tar.dtype826 xt_tar = xt_tar.to(torch.float32)827 prev_sample = xt_tar + (t_im1 - t_i) * Vt_tar828 prev_sample = prev_sample.to(dtype)829 xt_tar = prev_sample830 831 target_latents = zt_edit if xt_tar is None else xt_tar832 return target_latents833 834 def add_latents_noise(835 self,836 gt_latents,837 variance,838 noise,839 scheduler,840 ):841 842 bsz = gt_latents.shape[0]843 u = torch.tensor([variance] * bsz, dtype=gt_latents.dtype)844 indices = (u * scheduler.config.num_train_timesteps).long()845 timesteps = scheduler.timesteps.unsqueeze(1).to(gt_latents.dtype)846 indices = indices.to(timesteps.device).to(gt_latents.dtype).unsqueeze(1)847 nearest_idx = torch.argmin(torch.cdist(indices, timesteps), dim=1)848 sigma = (849 scheduler.sigmas[nearest_idx]850 .flatten()851 .to(gt_latents.device)852 .to(gt_latents.dtype)853 )854 while len(sigma.shape) < gt_latents.ndim:855 sigma = sigma.unsqueeze(-1)856 noisy_image = sigma * noise + (1.0 - sigma) * gt_latents857 init_timestep = indices[0]858 return noisy_image, init_timestep859 860 @torch.no_grad()861 def text2music_diffusion_process(862 self,863 duration,864 encoder_text_hidden_states,865 text_attention_mask,866 speaker_embds,867 lyric_token_ids,868 lyric_mask,869 random_generators=None,870 infer_steps=60,871 guidance_scale=15.0,872 omega_scale=10.0,873 scheduler_type="euler",874 cfg_type="apg",875 zero_steps=1,876 use_zero_init=True,877 guidance_interval=0.5,878 guidance_interval_decay=1.0,879 min_guidance_scale=3.0,880 oss_steps=[],881 encoder_text_hidden_states_null=None,882 use_erg_lyric=False,883 use_erg_diffusion=False,884 retake_random_generators=None,885 retake_variance=0.5,886 add_retake_noise=False,887 guidance_scale_text=0.0,888 guidance_scale_lyric=0.0,889 repaint_start=0,890 repaint_end=0,891 src_latents=None,892 audio2audio_enable=False,893 ref_audio_strength=0.5,894 ref_latents=None,895 ):896 897 logger.info(898 "cfg_type: {}, guidance_scale: {}, omega_scale: {}".format(899 cfg_type, guidance_scale, omega_scale900 )901 )902 do_classifier_free_guidance = True903 if guidance_scale == 0.0 or guidance_scale == 1.0:904 do_classifier_free_guidance = False905 906 do_double_condition_guidance = False907 if (908 guidance_scale_text is not None909 and guidance_scale_text > 1.0910 and guidance_scale_lyric is not None911 and guidance_scale_lyric > 1.0912 ):913 do_double_condition_guidance = True914 logger.info(915 "do_double_condition_guidance: {}, guidance_scale_text: {}, guidance_scale_lyric: {}".format(916 do_double_condition_guidance,917 guidance_scale_text,918 guidance_scale_lyric,919 )920 )921 922 device = encoder_text_hidden_states.device923 dtype = encoder_text_hidden_states.dtype924 bsz = encoder_text_hidden_states.shape[0]925 926 if scheduler_type == "euler":927 scheduler = FlowMatchEulerDiscreteScheduler(928 num_train_timesteps=1000,929 shift=3.0,930 )931 elif scheduler_type == "heun":932 scheduler = FlowMatchHeunDiscreteScheduler(933 num_train_timesteps=1000,934 shift=3.0,935 )936 937 frame_length = int(duration * 44100 / 512 / 8)938 if src_latents is not None:939 frame_length = src_latents.shape[-1]940 941 if ref_latents is not None:942 frame_length = ref_latents.shape[-1]943 944 if len(oss_steps) > 0:945 infer_steps = max(oss_steps)946 scheduler.set_timesteps947 timesteps, num_inference_steps = retrieve_timesteps(948 scheduler,949 num_inference_steps=infer_steps,950 device=device,951 timesteps=None,952 )953 new_timesteps = torch.zeros(len(oss_steps), dtype=dtype, device=device)954 for idx in range(len(oss_steps)):955 new_timesteps[idx] = timesteps[oss_steps[idx] - 1]956 num_inference_steps = len(oss_steps)957 sigmas = (new_timesteps / 1000).float().cpu().numpy()958 timesteps, num_inference_steps = retrieve_timesteps(959 scheduler,960 num_inference_steps=num_inference_steps,961 device=device,962 sigmas=sigmas,963 )964 logger.info(965 f"oss_steps: {oss_steps}, num_inference_steps: {num_inference_steps} after remapping to timesteps {timesteps}"966 )967 else:968 timesteps, num_inference_steps = retrieve_timesteps(969 scheduler,970 num_inference_steps=infer_steps,971 device=device,972 timesteps=None,973 )974 975 target_latents = randn_tensor(976 shape=(bsz, 8, 16, frame_length),977 generator=random_generators,978 device=device,979 dtype=dtype,980 )981 982 is_repaint = False983 is_extend = False984 if add_retake_noise:985 n_min = int(infer_steps * (1 - retake_variance))986 retake_variance = (987 torch.tensor(retake_variance * math.pi / 2).to(device).to(dtype)988 )989 retake_latents = randn_tensor(990 shape=(bsz, 8, 16, frame_length),991 generator=retake_random_generators,992 device=device,993 dtype=dtype,994 )995 repaint_start_frame = int(repaint_start * 44100 / 512 / 8)996 repaint_end_frame = int(repaint_end * 44100 / 512 / 8)997 x0 = src_latents998 # retake999 is_repaint = repaint_end_frame - repaint_start_frame != frame_length1000 1001 is_extend = (repaint_start_frame < 0) or (repaint_end_frame > frame_length)1002 if is_extend:1003 is_repaint = True1004 1005 # TODO: train a mask aware repainting controlnet1006 # to make sure mean = 0, std = 11007 if not is_repaint:1008 target_latents = (1009 torch.cos(retake_variance) * target_latents1010 + torch.sin(retake_variance) * retake_latents1011 )1012 elif not is_extend:1013 # if repaint_end_frame1014 repaint_mask = torch.zeros(1015 (bsz, 8, 16, frame_length), device=device, dtype=dtype1016 )1017 repaint_mask[:, :, :, repaint_start_frame:repaint_end_frame] = 1.01018 repaint_noise = (1019 torch.cos(retake_variance) * target_latents1020 + torch.sin(retake_variance) * retake_latents1021 )1022 repaint_noise = torch.where(1023 repaint_mask == 1.0, repaint_noise, target_latents1024 )1025 zt_edit = x0.clone()1026 z0 = repaint_noise1027 elif is_extend:1028 to_right_pad_gt_latents = None1029 to_left_pad_gt_latents = None1030 gt_latents = src_latents1031 src_latents_length = gt_latents.shape[-1]1032 max_infer_fame_length = int(240 * 44100 / 512 / 8)1033 left_pad_frame_length = 01034 right_pad_frame_length = 01035 right_trim_length = 01036 left_trim_length = 01037 if repaint_start_frame < 0:1038 left_pad_frame_length = abs(repaint_start_frame)1039 frame_length = left_pad_frame_length + gt_latents.shape[-1]1040 extend_gt_latents = torch.nn.functional.pad(1041 gt_latents, (left_pad_frame_length, 0), "constant", 01042 )1043 if frame_length > max_infer_fame_length:1044 right_trim_length = frame_length - max_infer_fame_length1045 extend_gt_latents = extend_gt_latents[1046 :, :, :, :max_infer_fame_length1047 ]1048 to_right_pad_gt_latents = extend_gt_latents[1049 :, :, :, -right_trim_length:1050 ]1051 frame_length = max_infer_fame_length1052 repaint_start_frame = 01053 gt_latents = extend_gt_latents1054 1055 if repaint_end_frame > src_latents_length:1056 right_pad_frame_length = repaint_end_frame - gt_latents.shape[-1]1057 frame_length = gt_latents.shape[-1] + right_pad_frame_length1058 extend_gt_latents = torch.nn.functional.pad(1059 gt_latents, (0, right_pad_frame_length), "constant", 01060 )1061 if frame_length > max_infer_fame_length:1062 left_trim_length = frame_length - max_infer_fame_length1063 extend_gt_latents = extend_gt_latents[1064 :, :, :, -max_infer_fame_length:1065 ]1066 to_left_pad_gt_latents = extend_gt_latents[1067 :, :, :, :left_trim_length1068 ]1069 frame_length = max_infer_fame_length1070 repaint_end_frame = frame_length1071 gt_latents = extend_gt_latents1072 1073 repaint_mask = torch.zeros(1074 (bsz, 8, 16, frame_length), device=device, dtype=dtype1075 )1076 if left_pad_frame_length > 0:1077 repaint_mask[:, :, :, :left_pad_frame_length] = 1.01078 if right_pad_frame_length > 0:1079 repaint_mask[:, :, :, -right_pad_frame_length:] = 1.01080 x0 = gt_latents1081 padd_list = []1082 if left_pad_frame_length > 0:1083 padd_list.append(retake_latents[:, :, :, :left_pad_frame_length])1084 padd_list.append(1085 target_latents[1086 :,1087 :,1088 :,1089 left_trim_length : target_latents.shape[-1] - right_trim_length,1090 ]1091 )1092 if right_pad_frame_length > 0:1093 padd_list.append(retake_latents[:, :, :, -right_pad_frame_length:])1094 target_latents = torch.cat(padd_list, dim=-1)1095 assert (1096 target_latents.shape[-1] == x0.shape[-1]1097 ), f"{target_latents.shape=} {x0.shape=}"1098 zt_edit = x0.clone()1099 z0 = target_latents1100 1101 init_timestep = 10001102 if audio2audio_enable and ref_latents is not None:1103 target_latents, init_timestep = self.add_latents_noise(1104 gt_latents=ref_latents,1105 variance=(1 - ref_audio_strength),1106 noise=target_latents,1107 scheduler=scheduler,1108 )1109 1110 attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)1111 1112 # guidance interval1113 start_idx = int(num_inference_steps * ((1 - guidance_interval) / 2))1114 end_idx = int(num_inference_steps * (guidance_interval / 2 + 0.5))1115 logger.info(1116 f"start_idx: {start_idx}, end_idx: {end_idx}, num_inference_steps: {num_inference_steps}"1117 )1118 1119 momentum_buffer = MomentumBuffer()1120 1121 def forward_encoder_with_temperature(self, inputs, tau=0.01, l_min=4, l_max=6):1122 handlers = []1123 1124 def hook(module, input, output):1125 output[:] *= tau1126 return output1127 1128 for i in range(l_min, l_max):1129 handler = self.ace_step_transformer.lyric_encoder.encoders[1130 i1131 ].self_attn.linear_q.register_forward_hook(hook)1132 handlers.append(handler)1133 1134 encoder_hidden_states, encoder_hidden_mask = (1135 self.ace_step_transformer.encode(**inputs)1136 )1137 1138 for hook in handlers:1139 hook.remove()1140 1141 return encoder_hidden_states1142 1143 # P(speaker, text, lyric)1144 encoder_hidden_states, encoder_hidden_mask = self.ace_step_transformer.encode(1145 encoder_text_hidden_states,1146 text_attention_mask,1147 speaker_embds,1148 lyric_token_ids,1149 lyric_mask,1150 )1151 1152 if use_erg_lyric:1153 # P(null_speaker, text_weaker, lyric_weaker)1154 encoder_hidden_states_null = forward_encoder_with_temperature(1155 self,1156 inputs={1157 "encoder_text_hidden_states": (1158 encoder_text_hidden_states_null1159 if encoder_text_hidden_states_null is not None1160 else torch.zeros_like(encoder_text_hidden_states)1161 ),1162 "text_attention_mask": text_attention_mask,1163 "speaker_embeds": torch.zeros_like(speaker_embds),1164 "lyric_token_idx": lyric_token_ids,1165 "lyric_mask": lyric_mask,1166 },1167 )1168 else:1169 # P(null_speaker, null_text, null_lyric)1170 encoder_hidden_states_null, _ = self.ace_step_transformer.encode(1171 torch.zeros_like(encoder_text_hidden_states),1172 text_attention_mask,1173 torch.zeros_like(speaker_embds),1174 torch.zeros_like(lyric_token_ids),1175 lyric_mask,1176 )1177 1178 encoder_hidden_states_no_lyric = None1179 if do_double_condition_guidance:1180 # P(null_speaker, text, lyric_weaker)1181 if use_erg_lyric:1182 encoder_hidden_states_no_lyric = forward_encoder_with_temperature(1183 self,1184 inputs={1185 "encoder_text_hidden_states": encoder_text_hidden_states,1186 "text_attention_mask": text_attention_mask,1187 "speaker_embeds": torch.zeros_like(speaker_embds),1188 "lyric_token_idx": lyric_token_ids,1189 "lyric_mask": lyric_mask,1190 },1191 )1192 # P(null_speaker, text, no_lyric)1193 else:1194 encoder_hidden_states_no_lyric, _ = self.ace_step_transformer.encode(1195 encoder_text_hidden_states,1196 text_attention_mask,1197 torch.zeros_like(speaker_embds),1198 torch.zeros_like(lyric_token_ids),1199 lyric_mask,1200 )