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

Update pipeline_ace_step.py

Browse Source
This commit is contained in:
mrfakename
2025-05-09 09:49:59 -07:00
committed by GitHub
parent 24f6e73013
commit 25b4c4bbcd
1 changed files with 78 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acestep/pipeline_ace_step.py
+76 -54
View File
@@ -44,6 +44,7 @@ from acestep.apg_guidance import (
cfg_double_condition_forward, cfg_double_condition_forward,
) )
import torchaudio import torchaudio
from .cpu_offload import cpu_offload
torch.backends.cudnn.benchmark = False torch.backends.cudnn.benchmark = False
@@ -96,6 +97,7 @@ class ACEStepPipeline:
text_encoder_checkpoint_path=None, text_encoder_checkpoint_path=None,
persistent_storage_path=None, persistent_storage_path=None,
torch_compile=False, torch_compile=False,
cpu_offload=False,
**kwargs, **kwargs,
): ):
if not checkpoint_dir: if not checkpoint_dir:
@@ -122,6 +124,7 @@ class ACEStepPipeline:
self.device = device self.device = device
self.loaded = False self.loaded = False
self.torch_compile = torch_compile self.torch_compile = torch_compile
self.cpu_offload = cpu_offload
def load_checkpoint(self, checkpoint_dir=None): def load_checkpoint(self, checkpoint_dir=None):
device = self.device device = self.device
@@ -143,12 +146,20 @@ class ACEStepPipeline:
dcae_checkpoint_path=dcae_checkpoint_path, dcae_checkpoint_path=dcae_checkpoint_path,
vocoder_checkpoint_path=vocoder_checkpoint_path, vocoder_checkpoint_path=vocoder_checkpoint_path,
) )
self.music_dcae.to(device).eval().to(self.dtype) # 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)
self.ace_step_transformer = ACEStepTransformer2DModel.from_pretrained( self.ace_step_transformer = ACEStepTransformer2DModel.from_pretrained(
ace_step_checkpoint_path, torch_dtype=self.dtype ace_step_checkpoint_path, torch_dtype=self.dtype
) )
self.ace_step_transformer.to(device).eval().to(self.dtype) # self.ace_step_transformer.to(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)
lang_segment = LangSegment() lang_segment = LangSegment()
@@ -258,7 +269,11 @@ class ACEStepPipeline:
text_encoder_model = UMT5EncoderModel.from_pretrained( text_encoder_model = UMT5EncoderModel.from_pretrained(
text_encoder_checkpoint_path, torch_dtype=self.dtype text_encoder_checkpoint_path, torch_dtype=self.dtype
).eval() ).eval()
text_encoder_model = text_encoder_model.to(device).to(self.dtype) # text_encoder_model = text_encoder_model.to(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.requires_grad_(False) text_encoder_model.requires_grad_(False)
self.text_encoder_model = text_encoder_model self.text_encoder_model = text_encoder_model
self.text_tokenizer = AutoTokenizer.from_pretrained( self.text_tokenizer = AutoTokenizer.from_pretrained(
@@ -272,6 +287,7 @@ class ACEStepPipeline:
self.ace_step_transformer = torch.compile(self.ace_step_transformer) self.ace_step_transformer = torch.compile(self.ace_step_transformer)
self.text_encoder_model = torch.compile(self.text_encoder_model) self.text_encoder_model = torch.compile(self.text_encoder_model)
@cpu_offload("text_encoder_model")
def get_text_embeddings(self, texts, device, text_max_length=256): def get_text_embeddings(self, texts, device, text_max_length=256):
inputs = self.text_tokenizer( inputs = self.text_tokenizer(
texts, texts,
@@ -289,6 +305,7 @@ class ACEStepPipeline:
attention_mask = inputs["attention_mask"] attention_mask = inputs["attention_mask"]
return last_hidden_states, attention_mask return last_hidden_states, attention_mask
@cpu_offload("text_encoder_model")
def get_text_embeddings_null( def get_text_embeddings_null(
self, texts, device, text_max_length=256, tau=0.01, l_min=8, l_max=10 self, texts, device, text_max_length=256, tau=0.01, l_min=8, l_max=10
): ):
@@ -331,28 +348,37 @@ class ACEStepPipeline:
return last_hidden_states return last_hidden_states
def set_seeds(self, batch_size, manual_seeds=None): def set_seeds(self, batch_size, manual_seeds=None):
seeds = None processed_input_seeds = None
if manual_seeds is not None: if manual_seeds is not None:
if isinstance(manual_seeds, str): if isinstance(manual_seeds, str):
if "," in manual_seeds: if "," in manual_seeds:
seeds = list(map(int, manual_seeds.split(","))) processed_input_seeds = list(map(int, manual_seeds.split(",")))
elif manual_seeds.isdigit(): elif manual_seeds.isdigit():
seeds = int(manual_seeds) processed_input_seeds = int(manual_seeds)
elif isinstance(manual_seeds, list) and all(isinstance(s, int) for s in manual_seeds):
if len(manual_seeds) > 0:
processed_input_seeds = list(manual_seeds)
elif isinstance(manual_seeds, int):
processed_input_seeds = manual_seeds
random_generators = [ random_generators = [
torch.Generator(device=self.device) for _ in range(batch_size) torch.Generator(device=self.device) for _ in range(batch_size)
] ]
actual_seeds = [] actual_seeds = []
for i in range(batch_size): for i in range(batch_size):
seed = None current_seed_for_generator = None
if seeds is None: if processed_input_seeds is None:
seed = torch.randint(0, 2**32, (1,)).item() current_seed_for_generator = torch.randint(0, 2**32, (1,)).item()
if isinstance(seeds, int): elif isinstance(processed_input_seeds, int):
seed = seeds current_seed_for_generator = processed_input_seeds
if isinstance(seeds, list): elif isinstance(processed_input_seeds, list):
seed = seeds[i] if i < len(processed_input_seeds):
random_generators[i].manual_seed(seed) current_seed_for_generator = processed_input_seeds[i]
actual_seeds.append(seed) else:
current_seed_for_generator = processed_input_seeds[-1]
if current_seed_for_generator is None:
current_seed_for_generator = torch.randint(0, 2**32, (1,)).item()
random_generators[i].manual_seed(current_seed_for_generator)
actual_seeds.append(current_seed_for_generator)
return random_generators, actual_seeds return random_generators, actual_seeds
def get_lang(self, text): def get_lang(self, text):
@@ -400,6 +426,7 @@ class ACEStepPipeline:
print("tokenize error", e, "for line", line, "major_language", lang) print("tokenize error", e, "for line", line, "major_language", lang)
return lyric_token_idx return lyric_token_idx
@cpu_offload("ace_step_transformer")
def calc_v( def calc_v(
self, self,
zt_src, zt_src,
@@ -688,6 +715,7 @@ class ACEStepPipeline:
target_latents = zt_edit if xt_tar is None else xt_tar target_latents = zt_edit if xt_tar is None else xt_tar
return target_latents return target_latents
@cpu_offload("ace_step_transformer")
@torch.no_grad() @torch.no_grad()
def text2music_diffusion_process( def text2music_diffusion_process(
self, self,
@@ -1204,6 +1232,7 @@ class ACEStepPipeline:
) )
return target_latents return target_latents
@cpu_offload("music_dcae")
def latents2audio( def latents2audio(
self, self,
latents, latents,
@@ -1211,29 +1240,20 @@ class ACEStepPipeline:
sample_rate=48000, sample_rate=48000,
save_path=None, save_path=None,
format="wav", format="wav",
do_save=True,
): ):
if do_save: output_audio_paths = []
output_audio_paths = [] bs = latents.shape[0]
bs = latents.shape[0] audio_lengths = [target_wav_duration_second * sample_rate] * bs
audio_lengths = [target_wav_duration_second * sample_rate] * bs pred_latents = latents
pred_latents = latents with torch.no_grad():
with torch.no_grad(): _, pred_wavs = self.music_dcae.decode(pred_latents, sr=sample_rate)
_, pred_wavs = self.music_dcae.decode(pred_latents, sr=sample_rate) pred_wavs = [pred_wav.cpu().float() for pred_wav in pred_wavs]
pred_wavs = [pred_wav.cpu().float() for pred_wav in pred_wavs] for i in tqdm(range(bs)):
for i in tqdm(range(bs)): output_audio_path = self.save_wav_file(
output_audio_path = self.save_wav_file( pred_wavs[i], i, save_path=save_path, sample_rate=sample_rate, format=format
pred_wavs[i], i, sample_rate=sample_rate )
) output_audio_paths.append(output_audio_path)
output_audio_paths.append(output_audio_path) return output_audio_paths
return output_audio_paths
else:
bs = latents.shape[0]
pred_latents = latents
with torch.no_grad():
_, pred_wavs = self.music_dcae.decode(pred_latents, sr=sample_rate)
pred_wavs = [pred_wav.cpu().float() for pred_wav in pred_wavs]
return pred_wavs
def save_wav_file( def save_wav_file(
self, target_wav, idx, save_path=None, sample_rate=48000, format="wav" self, target_wav, idx, save_path=None, sample_rate=48000, format="wav"
@@ -1242,20 +1262,25 @@ class ACEStepPipeline:
logger.warning("save_path is None, using default path ./outputs/") logger.warning("save_path is None, using default path ./outputs/")
base_path = f"./outputs" base_path = f"./outputs"
ensure_directory_exists(base_path) ensure_directory_exists(base_path)
output_path_wav = (
f"{base_path}/output_{time.strftime('%Y%m%d%H%M%S')}_{idx}.wav"
)
else: else:
base_path = save_path ensure_directory_exists(os.path.dirname(save_path))
ensure_directory_exists(base_path) if os.path.isdir(save_path):
logger.info(f"Provided save_path '{save_path}' is a directory. Appending timestamped filename.")
output_path_wav = os.path.join(save_path, f"output_{time.strftime('%Y%m%d%H%M%S')}_{idx}.wav")
else:
output_path_wav = save_path
output_path_wav = (
f"{base_path}/output_{time.strftime('%Y%m%d%H%M%S')}_{idx}.wav"
)
target_wav = target_wav.float() target_wav = target_wav.float()
print(target_wav) logger.info(f"Saving audio to {output_path_wav}")
torchaudio.save( torchaudio.save(
output_path_wav, target_wav, sample_rate=sample_rate, format=format output_path_wav, target_wav, sample_rate=sample_rate, format=format
) )
return output_path_wav return output_path_wav
@cpu_offload("music_dcae")
def infer_latents(self, input_audio_path): def infer_latents(self, input_audio_path):
if input_audio_path is None: if input_audio_path is None:
return None return None
@@ -1301,7 +1326,6 @@ class ACEStepPipeline:
format: str = "wav", format: str = "wav",
batch_size: int = 1, batch_size: int = 1,
debug: bool = False, debug: bool = False,
do_save: bool = True,
): ):
start_time = time.time() start_time = time.time()
@@ -1485,7 +1509,6 @@ class ACEStepPipeline:
target_wav_duration_second=audio_duration, target_wav_duration_second=audio_duration,
save_path=save_path, save_path=save_path,
format=format, format=format,
do_save=do_save,
) )
end_time = time.time() end_time = time.time()
@@ -1529,13 +1552,12 @@ class ACEStepPipeline:
"edit_target_lyrics": edit_target_lyrics, "edit_target_lyrics": edit_target_lyrics,
} }
# save input_params_json # save input_params_json
if do_save: for output_audio_path in output_paths:
for output_audio_path in output_paths: input_params_json_save_path = output_audio_path.replace(
input_params_json_save_path = output_audio_path.replace( f".{format}", "_input_params.json"
f".{format}", "_input_params.json" )
) input_params_json["audio_path"] = output_audio_path
input_params_json["audio_path"] = output_audio_path with open(input_params_json_save_path, "w", encoding="utf-8") as f:
with open(input_params_json_save_path, "w", encoding="utf-8") as f: json.dump(input_params_json, f, indent=4, ensure_ascii=False)
json.dump(input_params_json, f, indent=4, ensure_ascii=False)
return output_paths + [input_params_json] return output_paths + [input_params_json]