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

Merge pull request #214 from craftmine1000/main

Browse Source
This commit is contained in:
Gong Junmin
2025-05-19 17:50:15 +08:00
committed by GitHub
parent f51063ee8b bfac11fecf
commit 516792aa82
3 changed files with 289 additions and 310 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acestep/pipeline_ace_step.py
+89 -310
View File
@@ -35,7 +35,7 @@ from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import (
from diffusers.utils.torch_utils import randn_tensor
from transformers import UMT5EncoderModel, AutoTokenizer
from acestep.language_segmentation import LangSegment
from acestep.language_segmentation import LangSegment, language_filters
from acestep.music_dcae.music_dcae_pipeline import MusicDCAE
from acestep.models.ace_step_transformer import ACEStepTransformer2DModel
from acestep.models.lyrics_utils.lyric_tokenizer import VoiceBpeTokenizer
@@ -88,6 +88,7 @@ def ensure_directory_exists(directory):
REPO_ID = "ACE-Step/ACE-Step-v1-3.5B"
REPO_ID_QUANT = REPO_ID + "-q4-K-M" # ??? update this i guess
# class ACEStepPipeline(DiffusionPipeline):
@@ -129,6 +130,8 @@ class ACEStepPipeline:
self.dtype = torch.float16
if device.type == "mps":
self.dtype = torch.float32
if 'ACE_PIPELINE_DTYPE' in os.environ and len(os.environ['ACE_PIPELINE_DTYPE']):
self.dtype = getattr(torch, os.environ['ACE_PIPELINE_DTYPE'])
self.device = device
self.loaded = False
self.torch_compile = torch_compile
@@ -151,8 +154,7 @@ class ACEStepPipeline:
import gc
gc.collect()
def load_checkpoint(self, checkpoint_dir=None, export_quantized_weights=False):
device = self.device
def get_checkpoint_path(self, checkpoint_dir, repo):
checkpoint_dir_models = None
if checkpoint_dir is not None:
@@ -170,160 +172,65 @@ class ACEStepPipeline:
if checkpoint_dir_models is None:
if checkpoint_dir is None:
logger.info(f"Download models from Hugging Face: {REPO_ID}")
checkpoint_dir_models = snapshot_download(REPO_ID)
logger.info(f"Download models from Hugging Face: {repo}")
checkpoint_dir_models = snapshot_download(repo)
else:
logger.info(f"Download models from Hugging Face: {REPO_ID}, cache to: {checkpoint_dir}")
checkpoint_dir_models = snapshot_download(REPO_ID, cache_dir=checkpoint_dir)
logger.info(f"Download models from Hugging Face: {repo}, cache to: {checkpoint_dir}")
checkpoint_dir_models = snapshot_download(repo, cache_dir=checkpoint_dir)
return checkpoint_dir_models
dcae_model_path = os.path.join(checkpoint_dir_models, "music_dcae_f8c8")
vocoder_model_path = os.path.join(checkpoint_dir_models, "music_vocoder")
ace_step_model_path = os.path.join(checkpoint_dir_models, "ace_step_transformer")
text_encoder_model_path = os.path.join(checkpoint_dir_models, "umt5-base")
dcae_checkpoint_path = dcae_model_path
vocoder_checkpoint_path = vocoder_model_path
ace_step_checkpoint_path = ace_step_model_path
text_encoder_checkpoint_path = text_encoder_model_path
self.music_dcae = MusicDCAE(
dcae_checkpoint_path=dcae_checkpoint_path,
vocoder_checkpoint_path=vocoder_checkpoint_path,
)
# self.music_dcae.to(device).eval().to(self.dtype)
if self.cpu_offload: # might be redundant
self.music_dcae = self.music_dcae.to("cpu").eval().to(self.dtype)
else:
self.music_dcae = self.music_dcae.to(device).eval().to(self.dtype)
def load_checkpoint(self, checkpoint_dir=None, export_quantized_weights=False):
checkpoint_dir = self.get_checkpoint_path(checkpoint_dir, REPO_ID)
dcae_checkpoint_path = os.path.join(checkpoint_dir, "music_dcae_f8c8")
vocoder_checkpoint_path = os.path.join(checkpoint_dir, "music_vocoder")
ace_step_checkpoint_path = os.path.join(checkpoint_dir, "ace_step_transformer")
text_encoder_checkpoint_path = os.path.join(checkpoint_dir, "umt5-base")
self.ace_step_transformer = ACEStepTransformer2DModel.from_pretrained(
ace_step_checkpoint_path, torch_dtype=self.dtype
)
# self.ace_step_transformer.to(device).eval().to(self.dtype)
# self.ace_step_transformer.to(self.device).eval().to(self.dtype)
if self.cpu_offload:
self.ace_step_transformer = (
self.ace_step_transformer.to("cpu").eval().to(self.dtype)
)
else:
self.ace_step_transformer = (
self.ace_step_transformer.to(device).eval().to(self.dtype)
self.ace_step_transformer.to(self.device).eval().to(self.dtype)
)
if self.torch_compile:
self.ace_step_transformer = torch.compile(self.ace_step_transformer)
self.music_dcae = MusicDCAE(
dcae_checkpoint_path=dcae_checkpoint_path,
vocoder_checkpoint_path=vocoder_checkpoint_path,
)
# self.music_dcae.to(self.device).eval().to(self.dtype)
if self.cpu_offload: # might be redundant
self.music_dcae = self.music_dcae.to("cpu").eval().to(self.dtype)
else:
self.music_dcae = self.music_dcae.to(self.device).eval().to(self.dtype)
if self.torch_compile:
self.music_dcae = torch.compile(self.music_dcae)
lang_segment = LangSegment()
lang_segment.setfilters(
[
"af",
"am",
"an",
"ar",
"as",
"az",
"be",
"bg",
"bn",
"br",
"bs",
"ca",
"cs",
"cy",
"da",
"de",
"dz",
"el",
"en",
"eo",
"es",
"et",
"eu",
"fa",
"fi",
"fo",
"fr",
"ga",
"gl",
"gu",
"he",
"hi",
"hr",
"ht",
"hu",
"hy",
"id",
"is",
"it",
"ja",
"jv",
"ka",
"kk",
"km",
"kn",
"ko",
"ku",
"ky",
"la",
"lb",
"lo",
"lt",
"lv",
"mg",
"mk",
"ml",
"mn",
"mr",
"ms",
"mt",
"nb",
"ne",
"nl",
"nn",
"no",
"oc",
"or",
"pa",
"pl",
"ps",
"pt",
"qu",
"ro",
"ru",
"rw",
"se",
"si",
"sk",
"sl",
"sq",
"sr",
"sv",
"sw",
"ta",
"te",
"th",
"tl",
"tr",
"ug",
"uk",
"ur",
"vi",
"vo",
"wa",
"xh",
"zh",
"zu",
]
)
lang_segment.setfilters(language_filters.default)
self.lang_segment = lang_segment
self.lyric_tokenizer = VoiceBpeTokenizer()
text_encoder_model = UMT5EncoderModel.from_pretrained(
text_encoder_checkpoint_path, torch_dtype=self.dtype
).eval()
# text_encoder_model = text_encoder_model.to(device).to(self.dtype)
# text_encoder_model = text_encoder_model.to(self.device).to(self.dtype)
if self.cpu_offload:
text_encoder_model = text_encoder_model.to("cpu").eval().to(self.dtype)
else:
text_encoder_model = text_encoder_model.to(device).eval().to(self.dtype)
text_encoder_model = text_encoder_model.to(self.device).eval().to(self.dtype)
text_encoder_model.requires_grad_(False)
self.text_encoder_model = text_encoder_model
if self.torch_compile:
self.text_encoder_model = torch.compile(self.text_encoder_model)
self.text_tokenizer = AutoTokenizer.from_pretrained(
text_encoder_checkpoint_path
)
@@ -331,10 +238,6 @@ class ACEStepPipeline:
# compile
if self.torch_compile:
self.music_dcae = torch.compile(self.music_dcae)
self.ace_step_transformer = torch.compile(self.ace_step_transformer)
self.text_encoder_model = torch.compile(self.text_encoder_model)
if export_quantized_weights:
from torchao.quantization import (
quantize_,
@@ -356,38 +259,31 @@ class ACEStepPipeline:
torch.save(
self.ace_step_transformer.state_dict(),
os.path.join(
ace_step_model_path, "diffusion_pytorch_model_int4wo.bin"
ace_step_checkpoint_path, "diffusion_pytorch_model_int4wo.bin"
),
)
print(
"Quantized Weights Saved to: ",
os.path.join(
ace_step_model_path, "diffusion_pytorch_model_int4wo.bin"
ace_step_checkpoint_path, "diffusion_pytorch_model_int4wo.bin"
),
)
torch.save(
self.text_encoder_model.state_dict(),
os.path.join(text_encoder_model_path, "pytorch_model_int4wo.bin"),
os.path.join(text_encoder_checkpoint_path, "pytorch_model_int4wo.bin"),
)
print(
"Quantized Weights Saved to: ",
os.path.join(text_encoder_model_path, "pytorch_model_int4wo.bin"),
os.path.join(text_encoder_checkpoint_path, "pytorch_model_int4wo.bin"),
)
def load_quantized_checkpoint(self, checkpoint_dir=None):
device = self.device
dcae_model_path = os.path.join(checkpoint_dir, "music_dcae_f8c8")
vocoder_model_path = os.path.join(checkpoint_dir, "music_vocoder")
ace_step_model_path = os.path.join(checkpoint_dir, "ace_step_transformer")
text_encoder_model_path = os.path.join(checkpoint_dir, "umt5-base")
dcae_checkpoint_path = dcae_model_path
vocoder_checkpoint_path = vocoder_model_path
ace_step_checkpoint_path = ace_step_model_path
text_encoder_checkpoint_path = text_encoder_model_path
checkpoint_dir = self.get_checkpoint_path(checkpoint_dir, REPO_ID_QUANT)
dcae_checkpoint_path = os.path.join(checkpoint_dir, "music_dcae_f8c8")
vocoder_checkpoint_path = os.path.join(checkpoint_dir, "music_vocoder")
ace_step_checkpoint_path = os.path.join(checkpoint_dir, "ace_step_transformer")
text_encoder_checkpoint_path = os.path.join(checkpoint_dir, "umt5-base")
self.music_dcae = MusicDCAE(
dcae_checkpoint_path=dcae_checkpoint_path,
@@ -416,7 +312,7 @@ class ACEStepPipeline:
self.text_encoder_model = torch.compile(self.text_encoder_model)
self.text_encoder_model.load_state_dict(
torch.load(
os.path.join(text_encoder_model_path, "pytorch_model_int4wo.bin"),
os.path.join(text_encoder_checkpoint_path, "pytorch_model_int4wo.bin"),
map_location=self.device,
),assign=True
)
@@ -427,114 +323,14 @@ class ACEStepPipeline:
)
lang_segment = LangSegment()
lang_segment.setfilters(
[
"af",
"am",
"an",
"ar",
"as",
"az",
"be",
"bg",
"bn",
"br",
"bs",
"ca",
"cs",
"cy",
"da",
"de",
"dz",
"el",
"en",
"eo",
"es",
"et",
"eu",
"fa",
"fi",
"fo",
"fr",
"ga",
"gl",
"gu",
"he",
"hi",
"hr",
"ht",
"hu",
"hy",
"id",
"is",
"it",
"ja",
"jv",
"ka",
"kk",
"km",
"kn",
"ko",
"ku",
"ky",
"la",
"lb",
"lo",
"lt",
"lv",
"mg",
"mk",
"ml",
"mn",
"mr",
"ms",
"mt",
"nb",
"ne",
"nl",
"nn",
"no",
"oc",
"or",
"pa",
"pl",
"ps",
"pt",
"qu",
"ro",
"ru",
"rw",
"se",
"si",
"sk",
"sl",
"sq",
"sr",
"sv",
"sw",
"ta",
"te",
"th",
"tl",
"tr",
"ug",
"uk",
"ur",
"vi",
"vo",
"wa",
"xh",
"zh",
"zu",
]
)
lang_segment.setfilters(language_filters.default)
self.lang_segment = lang_segment
self.lyric_tokenizer = VoiceBpeTokenizer()
self.loaded = True
@cpu_offload("text_encoder_model")
def get_text_embeddings(self, texts, device, text_max_length=256):
def get_text_embeddings(self, texts, text_max_length=256):
inputs = self.text_tokenizer(
texts,
return_tensors="pt",
@@ -542,9 +338,9 @@ class ACEStepPipeline:
truncation=True,
max_length=text_max_length,
)
inputs = {key: value.to(device) for key, value in inputs.items()}
if self.text_encoder_model.device != device:
self.text_encoder_model.to(device)
inputs = {key: value.to(self.device) for key, value in inputs.items()}
if self.text_encoder_model.device != self.device:
self.text_encoder_model.to(self.device)
with torch.no_grad():
outputs = self.text_encoder_model(**inputs)
last_hidden_states = outputs.last_hidden_state
@@ -553,7 +349,7 @@ class ACEStepPipeline:
@cpu_offload("text_encoder_model")
def get_text_embeddings_null(
self, texts, device, text_max_length=256, tau=0.01, l_min=8, l_max=10
self, texts, text_max_length=256, tau=0.01, l_min=8, l_max=10
):
inputs = self.text_tokenizer(
texts,
@@ -562,9 +358,9 @@ class ACEStepPipeline:
truncation=True,
max_length=text_max_length,
)
inputs = {key: value.to(device) for key, value in inputs.items()}
if self.text_encoder_model.device != device:
self.text_encoder_model.to(device)
inputs = {key: value.to(self.device) for key, value in inputs.items()}
if self.text_encoder_model.device != self.device:
self.text_encoder_model.to(self.device)
def forward_with_temperature(inputs, tau=0.01, l_min=8, l_max=10):
handlers = []
@@ -796,8 +592,6 @@ class ACEStepPipeline:
do_classifier_free_guidance = False
target_guidance_scale = guidance_scale
device = encoder_text_hidden_states.device
dtype = encoder_text_hidden_states.dtype
bsz = encoder_text_hidden_states.shape[0]
scheduler = FlowMatchEulerDiscreteScheduler(
@@ -807,10 +601,10 @@ class ACEStepPipeline:
T_steps = infer_steps
frame_length = src_latents.shape[-1]
attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)
attention_mask = torch.ones(bsz, frame_length, device=self.device, dtype=self.dtype)
timesteps, T_steps = retrieve_timesteps(
scheduler, T_steps, device, timesteps=None
scheduler, T_steps, self.device, timesteps=None
)
if do_classifier_free_guidance:
@@ -875,7 +669,7 @@ class ACEStepPipeline:
if i + 1 < len(timesteps):
t_im1 = (timesteps[i + 1]) / 1000
else:
t_im1 = torch.zeros_like(t_i).to(t_i.device)
t_im1 = torch.zeros_like(t_i).to(self.device)
if i < n_max:
# Calculate the average of the V predictions
@@ -884,8 +678,8 @@ class ACEStepPipeline:
fwd_noise = randn_tensor(
shape=x_src.shape,
generator=random_generators,
device=device,
dtype=dtype,
device=self.device,
dtype=self.dtype,
)
zt_src = (1 - t_i) * x_src + (t_i) * fwd_noise
@@ -912,16 +706,13 @@ class ACEStepPipeline:
attention_mask=attention_mask,
momentum_buffer=momentum_buffer,
)
V_delta_avg += (1 / n_avg) * (
Vt_tar - Vt_src
) # - (hfg-1)*( x_src))
V_delta_avg += (1 / n_avg) * (Vt_tar - Vt_src) # - (hfg - 1) * (x_src)
zt_edit = zt_edit.to(torch.float32)
zt_edit = zt_edit.to(torch.float32) # arbitrary, should be settable for compatibility
if scheduler_type != "pingpong":
# propagate direct ODE
zt_edit = zt_edit.to(torch.float32)
zt_edit = zt_edit + (t_im1 - t_i) * V_delta_avg
zt_edit = zt_edit.to(V_delta_avg.dtype)
zt_edit = zt_edit.to(self.dtype)
else:
# propagate pingpong SDE
zt_edit_denoised = zt_edit - t_i * V_delta_avg
@@ -933,8 +724,8 @@ class ACEStepPipeline:
fwd_noise = randn_tensor(
shape=x_src.shape,
generator=random_generators,
device=device,
dtype=dtype,
device=self.device,
dtype=self.dtype,
)
scheduler._init_step_index(t)
sigma = scheduler.sigmas[scheduler.step_index]
@@ -963,11 +754,10 @@ class ACEStepPipeline:
return_src_pred=False,
)
dtype = Vt_tar.dtype
xt_tar = xt_tar.to(torch.float32)
if scheduler_type != "pingpong":
prev_sample = xt_tar + (t_im1 - t_i) * Vt_tar
prev_sample = prev_sample.to(dtype)
prev_sample = prev_sample.to(self.dtype)
xt_tar = prev_sample
else:
prev_sample = xt_tar - t_i * Vt_tar
@@ -988,7 +778,6 @@ class ACEStepPipeline:
):
bsz = gt_latents.shape[0]
device = gt_latents.device
if scheduler_type == "euler":
scheduler = FlowMatchEulerDiscreteScheduler(
num_train_timesteps=1000,
@@ -1012,7 +801,7 @@ class ACEStepPipeline:
timesteps, num_inference_steps = retrieve_timesteps(
scheduler,
num_inference_steps=infer_steps,
device=device,
device=self.device,
timesteps=None,
)
noisy_image = gt_latents * (1 - scheduler.sigma_max) + noise * scheduler.sigma_max
@@ -1082,8 +871,6 @@ class ACEStepPipeline:
)
)
device = encoder_text_hidden_states.device
dtype = encoder_text_hidden_states.dtype
bsz = encoder_text_hidden_states.shape[0]
if scheduler_type == "euler":
@@ -1115,10 +902,10 @@ class ACEStepPipeline:
timesteps, num_inference_steps = retrieve_timesteps(
scheduler,
num_inference_steps=infer_steps,
device=device,
device=self.device,
timesteps=None,
)
new_timesteps = torch.zeros(len(oss_steps), dtype=dtype, device=device)
new_timesteps = torch.zeros(len(oss_steps), dtype=self.dtype, device=self.device)
for idx in range(len(oss_steps)):
new_timesteps[idx] = timesteps[oss_steps[idx] - 1]
num_inference_steps = len(oss_steps)
@@ -1126,7 +913,7 @@ class ACEStepPipeline:
timesteps, num_inference_steps = retrieve_timesteps(
scheduler,
num_inference_steps=num_inference_steps,
device=device,
device=self.device,
sigmas=sigmas,
)
logger.info(
@@ -1136,15 +923,15 @@ class ACEStepPipeline:
timesteps, num_inference_steps = retrieve_timesteps(
scheduler,
num_inference_steps=infer_steps,
device=device,
device=self.device,
timesteps=None,
)
target_latents = randn_tensor(
shape=(bsz, 8, 16, frame_length),
generator=random_generators,
device=device,
dtype=dtype,
device=self.device,
dtype=self.dtype,
)
is_repaint = False
@@ -1153,13 +940,13 @@ class ACEStepPipeline:
if add_retake_noise:
n_min = int(infer_steps * (1 - retake_variance))
retake_variance = (
torch.tensor(retake_variance * math.pi / 2).to(device).to(dtype)
torch.tensor(retake_variance * math.pi / 2).to(self.device).to(self.dtype)
)
retake_latents = randn_tensor(
shape=(bsz, 8, 16, frame_length),
generator=retake_random_generators,
device=device,
dtype=dtype,
device=self.device,
dtype=self.dtype,
)
repaint_start_frame = int(repaint_start * 44100 / 512 / 8)
repaint_end_frame = int(repaint_end * 44100 / 512 / 8)
@@ -1181,7 +968,7 @@ class ACEStepPipeline:
elif not is_extend:
# if repaint_end_frame
repaint_mask = torch.zeros(
(bsz, 8, 16, frame_length), device=device, dtype=dtype
(bsz, 8, 16, frame_length), device=self.device, dtype=self.dtype
)
repaint_mask[:, :, :, repaint_start_frame:repaint_end_frame] = 1.0
repaint_noise = (
@@ -1240,7 +1027,7 @@ class ACEStepPipeline:
gt_latents = extend_gt_latents
repaint_mask = torch.zeros(
(bsz, 8, 16, frame_length), device=device, dtype=dtype
(bsz, 8, 16, frame_length), device=self.device, dtype=self.dtype
)
if left_pad_frame_length > 0:
repaint_mask[:, :, :, :left_pad_frame_length] = 1.0
@@ -1279,7 +1066,7 @@ class ACEStepPipeline:
infer_steps=infer_steps,
)
attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)
attention_mask = torch.ones(bsz, frame_length, device=self.device, dtype=self.dtype)
# guidance interval
start_idx = int(num_inference_steps * ((1 - guidance_interval) / 2))
@@ -1530,11 +1317,10 @@ class ACEStepPipeline:
if i + 1 < len(timesteps):
t_im1 = (timesteps[i + 1]) / 1000
else:
t_im1 = torch.zeros_like(t_i).to(t_i.device)
dtype = noise_pred.dtype
t_im1 = torch.zeros_like(t_i).to(self.device)
target_latents = target_latents.to(torch.float32)
prev_sample = target_latents + (t_im1 - t_i) * noise_pred
prev_sample = prev_sample.to(dtype)
prev_sample = prev_sample.to(self.dtype)
target_latents = prev_sample
zt_src = (1 - t_im1) * x0 + (t_im1) * z0
target_latents = torch.where(
@@ -1621,8 +1407,7 @@ class ACEStepPipeline:
return None
input_audio, sr = self.music_dcae.load_audio(input_audio_path)
input_audio = input_audio.unsqueeze(0)
device, dtype = self.device, self.dtype
input_audio = input_audio.to(device=device, dtype=dtype)
input_audio = input_audio.to(device=self.device, dtype=self.dtype)
latents, _ = self.music_dcae.encode(input_audio, sr=sr)
return latents
@@ -1711,20 +1496,14 @@ class ACEStepPipeline:
oss_steps = []
texts = [prompt]
encoder_text_hidden_states, text_attention_mask = self.get_text_embeddings(
texts, self.device
)
encoder_text_hidden_states, text_attention_mask = self.get_text_embeddings(texts)
encoder_text_hidden_states = encoder_text_hidden_states.repeat(batch_size, 1, 1)
text_attention_mask = text_attention_mask.repeat(batch_size, 1)
encoder_text_hidden_states_null = None
if use_erg_tag:
encoder_text_hidden_states_null = self.get_text_embeddings_null(
texts, self.device
)
encoder_text_hidden_states_null = encoder_text_hidden_states_null.repeat(
batch_size, 1, 1
)
encoder_text_hidden_states_null = self.get_text_embeddings_null(texts)
encoder_text_hidden_states_null = encoder_text_hidden_states_null.repeat(batch_size, 1, 1)
# not support for released checkpoint
speaker_embeds = torch.zeros(batch_size, 512).to(self.device).to(self.dtype)
@@ -1785,7 +1564,7 @@ class ACEStepPipeline:
if task == "edit":
texts = [edit_target_prompt]
target_encoder_text_hidden_states, target_text_attention_mask = (
self.get_text_embeddings(texts, self.device)
self.get_text_embeddings(texts)
)
target_encoder_text_hidden_states = (
target_encoder_text_hidden_states.repeat(batch_size, 1, 1)