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

add edit

Browse Source
This commit is contained in:
chuxij
2025-04-29 11:12:02 +00:00
parent 3eb0c178a9
commit 175797250c
2 changed files with 466 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pipeline_ace_step.py
+329 -35
View File
@@ -56,7 +56,7 @@ REPO_ID = "ACE-Step/ACE-Step-v1-3.5B"
# class ACEStepPipeline(DiffusionPipeline):
class ACEStepPipeline:
def __init__(self, checkpoint_dir=None, device_id=0, dtype="bfloat16", text_encoder_checkpoint_path=None, **kwargs):
def __init__(self, checkpoint_dir=None, device_id=0, dtype="bfloat16", text_encoder_checkpoint_path=None, torch_compile=False, **kwargs):
# check checkpoint dir exist
if not checkpoint_dir:
checkpoint_dir = os.path.join(os.path.dirname(__file__), "checkpoints")
@@ -72,6 +72,7 @@ class ACEStepPipeline:
self.dtype = torch.float16 if dtype == "bfloat16" else torch.float32
self.device = device
self.loaded = False
self.torch_compile = torch_compile
def load_checkpoint(self, checkpoint_dir=None):
device = self.device
@@ -105,9 +106,10 @@ class ACEStepPipeline:
self.loaded = True
# 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 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)
def get_text_embeddings(self, texts, device, text_max_length=256):
inputs = self.text_tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=text_max_length)
@@ -215,6 +217,249 @@ class ACEStepPipeline:
print("tokenize error", e, "for line", line, "major_language", lang)
return lyric_token_idx
def calc_v(
self,
zt_src,
zt_tar,
t,
encoder_text_hidden_states,
text_attention_mask,
target_encoder_text_hidden_states,
target_text_attention_mask,
speaker_embds,
target_speaker_embeds,
lyric_token_ids,
lyric_mask,
target_lyric_token_ids,
target_lyric_mask,
do_classifier_free_guidance=False,
guidance_scale=1.0,
target_guidance_scale=1.0,
cfg_type="apg",
attention_mask=None,
momentum_buffer=None,
momentum_buffer_tar=None,
return_src_pred=True
):
noise_pred_src = None
if return_src_pred:
src_latent_model_input = torch.cat([zt_src, zt_src]) if do_classifier_free_guidance else zt_src
timestep = t.expand(src_latent_model_input.shape[0])
# source
noise_pred_src = self.ace_step_transformer(
hidden_states=src_latent_model_input,
attention_mask=attention_mask,
encoder_text_hidden_states=encoder_text_hidden_states,
text_attention_mask=text_attention_mask,
speaker_embeds=speaker_embds,
lyric_token_idx=lyric_token_ids,
lyric_mask=lyric_mask,
timestep=timestep,
).sample
if do_classifier_free_guidance:
noise_pred_with_cond_src, noise_pred_uncond_src = noise_pred_src.chunk(2)
if cfg_type == "apg":
noise_pred_src = apg_forward(
pred_cond=noise_pred_with_cond_src,
pred_uncond=noise_pred_uncond_src,
guidance_scale=guidance_scale,
momentum_buffer=momentum_buffer,
)
elif cfg_type == "cfg":
noise_pred_src = cfg_forward(
cond_output=noise_pred_with_cond_src,
uncond_output=noise_pred_uncond_src,
cfg_strength=guidance_scale,
)
tar_latent_model_input = torch.cat([zt_tar, zt_tar]) if do_classifier_free_guidance else zt_tar
timestep = t.expand(tar_latent_model_input.shape[0])
# target
noise_pred_tar = self.ace_step_transformer(
hidden_states=tar_latent_model_input,
attention_mask=attention_mask,
encoder_text_hidden_states=target_encoder_text_hidden_states,
text_attention_mask=target_text_attention_mask,
speaker_embeds=target_speaker_embeds,
lyric_token_idx=target_lyric_token_ids,
lyric_mask=target_lyric_mask,
timestep=timestep,
).sample
if do_classifier_free_guidance:
noise_pred_with_cond_tar, noise_pred_uncond_tar = noise_pred_tar.chunk(2)
if cfg_type == "apg":
noise_pred_tar = apg_forward(
pred_cond=noise_pred_with_cond_tar,
pred_uncond=noise_pred_uncond_tar,
guidance_scale=target_guidance_scale,
momentum_buffer=momentum_buffer_tar,
)
elif cfg_type == "cfg":
noise_pred_tar = cfg_forward(
cond_output=noise_pred_with_cond_tar,
uncond_output=noise_pred_uncond_tar,
cfg_strength=target_guidance_scale,
)
return noise_pred_src, noise_pred_tar
@torch.no_grad()
def flowedit_diffusion_process(
self,
encoder_text_hidden_states,
text_attention_mask,
speaker_embds,
lyric_token_ids,
lyric_mask,
target_encoder_text_hidden_states,
target_text_attention_mask,
target_speaker_embeds,
target_lyric_token_ids,
target_lyric_mask,
src_latents,
random_generators=None,
infer_steps=60,
guidance_scale=15.0,
n_min=0,
n_max=1.0,
n_avg=1,
):
do_classifier_free_guidance = True
if guidance_scale == 0.0 or guidance_scale == 1.0:
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(
num_train_timesteps=1000,
shift=3.0,
)
T_steps = infer_steps
frame_length = src_latents.shape[-1]
attention_mask = torch.ones(bsz, frame_length, device=device, dtype=dtype)
timesteps, T_steps = retrieve_timesteps(scheduler, T_steps, device, timesteps=None)
if do_classifier_free_guidance:
attention_mask = torch.cat([attention_mask] * 2, dim=0)
encoder_text_hidden_states = torch.cat([encoder_text_hidden_states, torch.zeros_like(encoder_text_hidden_states)], 0)
text_attention_mask = torch.cat([text_attention_mask] * 2, dim=0)
target_encoder_text_hidden_states = torch.cat([target_encoder_text_hidden_states, torch.zeros_like(target_encoder_text_hidden_states)], 0)
target_text_attention_mask = torch.cat([target_text_attention_mask] * 2, dim=0)
speaker_embds = torch.cat([speaker_embds, torch.zeros_like(speaker_embds)], 0)
target_speaker_embeds = torch.cat([target_speaker_embeds, torch.zeros_like(target_speaker_embeds)], 0)
lyric_token_ids = torch.cat([lyric_token_ids, torch.zeros_like(lyric_token_ids)], 0)
lyric_mask = torch.cat([lyric_mask, torch.zeros_like(lyric_mask)], 0)
target_lyric_token_ids = torch.cat([target_lyric_token_ids, torch.zeros_like(target_lyric_token_ids)], 0)
target_lyric_mask = torch.cat([target_lyric_mask, torch.zeros_like(target_lyric_mask)], 0)
momentum_buffer = MomentumBuffer()
momentum_buffer_tar = MomentumBuffer()
x_src = src_latents
zt_edit = x_src.clone()
xt_tar = None
n_min = int(infer_steps * n_min)
n_max = int(infer_steps * n_max)
logger.info("flowedit start from {} to {}".format(n_min, n_max))
for i, t in tqdm(enumerate(timesteps), total=T_steps):
if i < n_min:
continue
t_i = t/1000
if i+1 < len(timesteps):
t_im1 = (timesteps[i+1])/1000
else:
t_im1 = torch.zeros_like(t_i).to(t_i.device)
if i < n_max:
# Calculate the average of the V predictions
V_delta_avg = torch.zeros_like(x_src)
for k in range(n_avg):
fwd_noise = randn_tensor(shape=x_src.shape, generator=random_generators, device=device, dtype=dtype)
zt_src = (1 - t_i) * x_src + (t_i) * fwd_noise
zt_tar = zt_edit + zt_src - x_src
Vt_src, Vt_tar = self.calc_v(
zt_src=zt_src,
zt_tar=zt_tar,
t=t,
encoder_text_hidden_states=encoder_text_hidden_states,
text_attention_mask=text_attention_mask,
target_encoder_text_hidden_states=target_encoder_text_hidden_states,
target_text_attention_mask=target_text_attention_mask,
speaker_embds=speaker_embds,
target_speaker_embeds=target_speaker_embeds,
lyric_token_ids=lyric_token_ids,
lyric_mask=lyric_mask,
target_lyric_token_ids=target_lyric_token_ids,
target_lyric_mask=target_lyric_mask,
do_classifier_free_guidance=do_classifier_free_guidance,
guidance_scale=guidance_scale,
target_guidance_scale=target_guidance_scale,
attention_mask=attention_mask,
momentum_buffer=momentum_buffer
)
V_delta_avg += (1 / n_avg) * (Vt_tar - Vt_src) # - (hfg-1)*( x_src))
# 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)
else: # i >= T_steps-n_min # regular sampling for last n_min steps
if i == n_max:
fwd_noise = randn_tensor(shape=x_src.shape, generator=random_generators, device=device, dtype=dtype)
scheduler._init_step_index(t)
sigma = scheduler.sigmas[scheduler.step_index]
xt_src = sigma * fwd_noise + (1.0 - sigma) * x_src
xt_tar = zt_edit + xt_src - x_src
_, Vt_tar = self.calc_v(
zt_src=None,
zt_tar=xt_tar,
t=t,
encoder_text_hidden_states=encoder_text_hidden_states,
text_attention_mask=text_attention_mask,
target_encoder_text_hidden_states=target_encoder_text_hidden_states,
target_text_attention_mask=target_text_attention_mask,
speaker_embds=speaker_embds,
target_speaker_embeds=target_speaker_embeds,
lyric_token_ids=lyric_token_ids,
lyric_mask=lyric_mask,
target_lyric_token_ids=target_lyric_token_ids,
target_lyric_mask=target_lyric_mask,
do_classifier_free_guidance=do_classifier_free_guidance,
guidance_scale=guidance_scale,
target_guidance_scale=target_guidance_scale,
attention_mask=attention_mask,
momentum_buffer_tar=momentum_buffer_tar,
return_src_pred=False,
)
dtype = Vt_tar.dtype
xt_tar = xt_tar.to(torch.float32)
prev_sample = xt_tar + (t_im1 - t_i) * Vt_tar
prev_sample = prev_sample.to(dtype)
xt_tar = prev_sample
target_latents = zt_edit if xt_tar is None else xt_tar
return target_latents
@torch.no_grad()
def text2music_diffusion_process(
@@ -593,6 +838,11 @@ class ACEStepPipeline:
repaint_start: int = 0,
repaint_end: int = 0,
src_audio_path: str = None,
edit_target_prompt: str = None,
edit_target_lyrics: str = None,
edit_n_min: float = 0.0,
edit_n_max: float = 1.0,
edit_n_avg: int = 1,
save_path: str = None,
format: str = "flac",
batch_size: int = 1,
@@ -653,40 +903,76 @@ class ACEStepPipeline:
repaint_end = audio_duration
src_latents = None
if task == "repaint":
assert src_audio_path is not None, "src_audio_path is required for repaint task"
if src_audio_path is not None:
assert src_audio_path is not None and task in ("repaint", "edit"), "src_audio_path is required for repaint task"
assert os.path.exists(src_audio_path), f"src_audio_path {src_audio_path} does not exist"
src_latents = self.infer_latents(src_audio_path)
target_latents = self.text2music_diffusion_process(
duration=audio_duration,
encoder_text_hidden_states=encoder_text_hidden_states,
text_attention_mask=text_attention_mask,
speaker_embds=speaker_embeds,
lyric_token_ids=lyric_token_idx,
lyric_mask=lyric_mask,
guidance_scale=guidance_scale,
omega_scale=omega_scale,
infer_steps=infer_step,
random_generators=random_generators,
scheduler_type=scheduler_type,
cfg_type=cfg_type,
guidance_interval=guidance_interval,
guidance_interval_decay=guidance_interval_decay,
min_guidance_scale=min_guidance_scale,
oss_steps=oss_steps,
encoder_text_hidden_states_null=encoder_text_hidden_states_null,
use_erg_lyric=use_erg_lyric,
use_erg_diffusion=use_erg_diffusion,
retake_random_generators=retake_random_generators,
retake_variance=retake_variance,
add_retake_noise=add_retake_noise,
guidance_scale_text=guidance_scale_text,
guidance_scale_lyric=guidance_scale_lyric,
repaint_start=repaint_start,
repaint_end=repaint_end,
src_latents=src_latents,
)
if task == "edit":
texts = [edit_target_prompt]
target_encoder_text_hidden_states, target_text_attention_mask = self.get_text_embeddings(texts, self.device)
target_encoder_text_hidden_states = target_encoder_text_hidden_states.repeat(batch_size, 1, 1)
target_text_attention_mask = target_text_attention_mask.repeat(batch_size, 1)
target_lyric_token_idx = torch.tensor([0]).repeat(batch_size, 1).to(self.device).long()
target_lyric_mask = torch.tensor([0]).repeat(batch_size, 1).to(self.device).long()
if len(edit_target_lyrics) > 0:
target_lyric_token_idx = self.tokenize_lyrics(edit_target_lyrics, debug=True)
target_lyric_mask = [1] * len(target_lyric_token_idx)
target_lyric_token_idx = torch.tensor(target_lyric_token_idx).unsqueeze(0).to(self.device).repeat(batch_size, 1)
target_lyric_mask = torch.tensor(target_lyric_mask).unsqueeze(0).to(self.device).repeat(batch_size, 1)
target_speaker_embeds = speaker_embeds.clone()
target_latents = self.flowedit_diffusion_process(
encoder_text_hidden_states=encoder_text_hidden_states,
text_attention_mask=text_attention_mask,
speaker_embds=speaker_embeds,
lyric_token_ids=lyric_token_idx,
lyric_mask=lyric_mask,
target_encoder_text_hidden_states=target_encoder_text_hidden_states,
target_text_attention_mask=target_text_attention_mask,
target_speaker_embeds=target_speaker_embeds,
target_lyric_token_ids=target_lyric_token_idx,
target_lyric_mask=target_lyric_mask,
src_latents=src_latents,
random_generators=random_generators,
infer_steps=infer_step,
guidance_scale=guidance_scale,
n_min=edit_n_min,
n_max=edit_n_max,
n_avg=edit_n_avg,
)
else:
target_latents = self.text2music_diffusion_process(
duration=audio_duration,
encoder_text_hidden_states=encoder_text_hidden_states,
text_attention_mask=text_attention_mask,
speaker_embds=speaker_embeds,
lyric_token_ids=lyric_token_idx,
lyric_mask=lyric_mask,
guidance_scale=guidance_scale,
omega_scale=omega_scale,
infer_steps=infer_step,
random_generators=random_generators,
scheduler_type=scheduler_type,
cfg_type=cfg_type,
guidance_interval=guidance_interval,
guidance_interval_decay=guidance_interval_decay,
min_guidance_scale=min_guidance_scale,
oss_steps=oss_steps,
encoder_text_hidden_states_null=encoder_text_hidden_states_null,
use_erg_lyric=use_erg_lyric,
use_erg_diffusion=use_erg_diffusion,
retake_random_generators=retake_random_generators,
retake_variance=retake_variance,
add_retake_noise=add_retake_noise,
guidance_scale_text=guidance_scale_text,
guidance_scale_lyric=guidance_scale_lyric,
repaint_start=repaint_start,
repaint_end=repaint_end,
src_latents=src_latents,
)
end_time = time.time()
diffusion_time_cost = end_time - start_time
@@ -730,6 +1016,14 @@ class ACEStepPipeline:
"retake_variance": retake_variance,
"guidance_scale_text": guidance_scale_text,
"guidance_scale_lyric": guidance_scale_lyric,
"repaint_start": repaint_start,
"repaint_end": repaint_end,
"edit_n_min": edit_n_min,
"edit_n_max": edit_n_max,
"edit_n_avg": edit_n_avg,
"src_audio_path": src_audio_path,
"edit_target_prompt": edit_target_prompt,
"edit_target_lyrics": edit_target_lyrics,
}
# save input_params_json
for output_audio_path in output_paths: