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

reduce device and dtype vars

Browse Source
This commit is contained in:
Michael Hedman
2025-05-19 07:34:13 +02:00
parent d470cd903f
commit 0b379b6163
1 changed files with 42 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acestep/pipeline_ace_step.py
+42 -58
View File
@@ -154,7 +154,6 @@ class ACEStepPipeline:
gc.collect() gc.collect()
def load_checkpoint(self, checkpoint_dir=None, export_quantized_weights=False): def load_checkpoint(self, checkpoint_dir=None, export_quantized_weights=False):
device = self.device
checkpoint_dir_models = None checkpoint_dir_models = None
if checkpoint_dir is not None: if checkpoint_dir is not None:
@@ -186,14 +185,14 @@ class ACEStepPipeline:
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(self.device).eval().to(self.dtype)
if self.cpu_offload: if self.cpu_offload:
self.ace_step_transformer = ( self.ace_step_transformer = (
self.ace_step_transformer.to("cpu").eval().to(self.dtype) self.ace_step_transformer.to("cpu").eval().to(self.dtype)
) )
else: else:
self.ace_step_transformer = ( 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: if self.torch_compile:
self.ace_step_transformer = torch.compile(self.ace_step_transformer) self.ace_step_transformer = torch.compile(self.ace_step_transformer)
@@ -202,11 +201,11 @@ 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(self.device).eval().to(self.dtype)
if self.cpu_offload: # might be redundant if self.cpu_offload: # might be redundant
self.music_dcae = self.music_dcae.to("cpu").eval().to(self.dtype) self.music_dcae = self.music_dcae.to("cpu").eval().to(self.dtype)
else: else:
self.music_dcae = self.music_dcae.to(device).eval().to(self.dtype) self.music_dcae = self.music_dcae.to(self.device).eval().to(self.dtype)
if self.torch_compile: if self.torch_compile:
self.music_dcae = torch.compile(self.music_dcae) self.music_dcae = torch.compile(self.music_dcae)
@@ -218,11 +217,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(self.device).to(self.dtype)
if self.cpu_offload: if self.cpu_offload:
text_encoder_model = text_encoder_model.to("cpu").eval().to(self.dtype) text_encoder_model = text_encoder_model.to("cpu").eval().to(self.dtype)
else: 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) text_encoder_model.requires_grad_(False)
self.text_encoder_model = text_encoder_model self.text_encoder_model = text_encoder_model
if self.torch_compile: if self.torch_compile:
@@ -276,7 +275,6 @@ class ACEStepPipeline:
def load_quantized_checkpoint(self, checkpoint_dir=None): def load_quantized_checkpoint(self, checkpoint_dir=None):
device = self.device
dcae_checkpoint_path = os.path.join(checkpoint_dir, "music_dcae_f8c8") dcae_checkpoint_path = os.path.join(checkpoint_dir, "music_dcae_f8c8")
vocoder_checkpoint_path = os.path.join(checkpoint_dir, "music_vocoder") vocoder_checkpoint_path = os.path.join(checkpoint_dir, "music_vocoder")
@@ -328,7 +326,7 @@ class ACEStepPipeline:
self.loaded = True self.loaded = True
@cpu_offload("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, text_max_length=256):
inputs = self.text_tokenizer( inputs = self.text_tokenizer(
texts, texts,
return_tensors="pt", return_tensors="pt",
@@ -336,9 +334,9 @@ class ACEStepPipeline:
truncation=True, truncation=True,
max_length=text_max_length, max_length=text_max_length,
) )
inputs = {key: value.to(device) for key, value in inputs.items()} inputs = {key: value.to(self.device) for key, value in inputs.items()}
if self.text_encoder_model.device != device: if self.text_encoder_model.device != self.device:
self.text_encoder_model.to(device) self.text_encoder_model.to(self.device)
with torch.no_grad(): with torch.no_grad():
outputs = self.text_encoder_model(**inputs) outputs = self.text_encoder_model(**inputs)
last_hidden_states = outputs.last_hidden_state last_hidden_states = outputs.last_hidden_state
@@ -347,7 +345,7 @@ class ACEStepPipeline:
@cpu_offload("text_encoder_model") @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, text_max_length=256, tau=0.01, l_min=8, l_max=10
): ):
inputs = self.text_tokenizer( inputs = self.text_tokenizer(
texts, texts,
@@ -356,9 +354,9 @@ class ACEStepPipeline:
truncation=True, truncation=True,
max_length=text_max_length, max_length=text_max_length,
) )
inputs = {key: value.to(device) for key, value in inputs.items()} inputs = {key: value.to(self.device) for key, value in inputs.items()}
if self.text_encoder_model.device != device: if self.text_encoder_model.device != self.device:
self.text_encoder_model.to(device) self.text_encoder_model.to(self.device)
def forward_with_temperature(inputs, tau=0.01, l_min=8, l_max=10): def forward_with_temperature(inputs, tau=0.01, l_min=8, l_max=10):
handlers = [] handlers = []
@@ -590,8 +588,6 @@ class ACEStepPipeline:
do_classifier_free_guidance = False do_classifier_free_guidance = False
target_guidance_scale = guidance_scale target_guidance_scale = guidance_scale
device = encoder_text_hidden_states.device
dtype = encoder_text_hidden_states.dtype
bsz = encoder_text_hidden_states.shape[0] bsz = encoder_text_hidden_states.shape[0]
scheduler = FlowMatchEulerDiscreteScheduler( scheduler = FlowMatchEulerDiscreteScheduler(
@@ -601,10 +597,10 @@ class ACEStepPipeline:
T_steps = infer_steps T_steps = infer_steps
frame_length = src_latents.shape[-1] 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( timesteps, T_steps = retrieve_timesteps(
scheduler, T_steps, device, timesteps=None scheduler, T_steps, self.device, timesteps=None
) )
if do_classifier_free_guidance: if do_classifier_free_guidance:
@@ -669,7 +665,7 @@ class ACEStepPipeline:
if i + 1 < len(timesteps): if i + 1 < len(timesteps):
t_im1 = (timesteps[i + 1]) / 1000 t_im1 = (timesteps[i + 1]) / 1000
else: 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: if i < n_max:
# Calculate the average of the V predictions # Calculate the average of the V predictions
@@ -678,8 +674,8 @@ class ACEStepPipeline:
fwd_noise = randn_tensor( fwd_noise = randn_tensor(
shape=x_src.shape, shape=x_src.shape,
generator=random_generators, generator=random_generators,
device=device, device=self.device,
dtype=dtype, dtype=self.dtype,
) )
zt_src = (1 - t_i) * x_src + (t_i) * fwd_noise zt_src = (1 - t_i) * x_src + (t_i) * fwd_noise
@@ -727,8 +723,8 @@ class ACEStepPipeline:
fwd_noise = randn_tensor( fwd_noise = randn_tensor(
shape=x_src.shape, shape=x_src.shape,
generator=random_generators, generator=random_generators,
device=device, device=self.device,
dtype=dtype, dtype=self.dtype,
) )
scheduler._init_step_index(t) scheduler._init_step_index(t)
sigma = scheduler.sigmas[scheduler.step_index] sigma = scheduler.sigmas[scheduler.step_index]
@@ -757,11 +753,10 @@ class ACEStepPipeline:
return_src_pred=False, return_src_pred=False,
) )
dtype = Vt_tar.dtype
xt_tar = xt_tar.to(torch.float32) xt_tar = xt_tar.to(torch.float32)
if scheduler_type != "pingpong": if scheduler_type != "pingpong":
prev_sample = xt_tar + (t_im1 - t_i) * Vt_tar 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 xt_tar = prev_sample
else: else:
prev_sample = xt_tar - t_i * Vt_tar prev_sample = xt_tar - t_i * Vt_tar
@@ -782,7 +777,6 @@ class ACEStepPipeline:
): ):
bsz = gt_latents.shape[0] bsz = gt_latents.shape[0]
device = gt_latents.device
if scheduler_type == "euler": if scheduler_type == "euler":
scheduler = FlowMatchEulerDiscreteScheduler( scheduler = FlowMatchEulerDiscreteScheduler(
num_train_timesteps=1000, num_train_timesteps=1000,
@@ -806,7 +800,7 @@ class ACEStepPipeline:
timesteps, num_inference_steps = retrieve_timesteps( timesteps, num_inference_steps = retrieve_timesteps(
scheduler, scheduler,
num_inference_steps=infer_steps, num_inference_steps=infer_steps,
device=device, device=self.device,
timesteps=None, timesteps=None,
) )
noisy_image = gt_latents * (1 - scheduler.sigma_max) + noise * scheduler.sigma_max noisy_image = gt_latents * (1 - scheduler.sigma_max) + noise * scheduler.sigma_max
@@ -876,8 +870,6 @@ class ACEStepPipeline:
) )
) )
device = encoder_text_hidden_states.device
dtype = encoder_text_hidden_states.dtype
bsz = encoder_text_hidden_states.shape[0] bsz = encoder_text_hidden_states.shape[0]
if scheduler_type == "euler": if scheduler_type == "euler":
@@ -909,10 +901,10 @@ class ACEStepPipeline:
timesteps, num_inference_steps = retrieve_timesteps( timesteps, num_inference_steps = retrieve_timesteps(
scheduler, scheduler,
num_inference_steps=infer_steps, num_inference_steps=infer_steps,
device=device, device=self.device,
timesteps=None, 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)): for idx in range(len(oss_steps)):
new_timesteps[idx] = timesteps[oss_steps[idx] - 1] new_timesteps[idx] = timesteps[oss_steps[idx] - 1]
num_inference_steps = len(oss_steps) num_inference_steps = len(oss_steps)
@@ -920,7 +912,7 @@ class ACEStepPipeline:
timesteps, num_inference_steps = retrieve_timesteps( timesteps, num_inference_steps = retrieve_timesteps(
scheduler, scheduler,
num_inference_steps=num_inference_steps, num_inference_steps=num_inference_steps,
device=device, device=self.device,
sigmas=sigmas, sigmas=sigmas,
) )
logger.info( logger.info(
@@ -930,15 +922,15 @@ class ACEStepPipeline:
timesteps, num_inference_steps = retrieve_timesteps( timesteps, num_inference_steps = retrieve_timesteps(
scheduler, scheduler,
num_inference_steps=infer_steps, num_inference_steps=infer_steps,
device=device, device=self.device,
timesteps=None, timesteps=None,
) )
target_latents = randn_tensor( target_latents = randn_tensor(
shape=(bsz, 8, 16, frame_length), shape=(bsz, 8, 16, frame_length),
generator=random_generators, generator=random_generators,
device=device, device=self.device,
dtype=dtype, dtype=self.dtype,
) )
is_repaint = False is_repaint = False
@@ -947,13 +939,13 @@ class ACEStepPipeline:
if add_retake_noise: if add_retake_noise:
n_min = int(infer_steps * (1 - retake_variance)) n_min = int(infer_steps * (1 - retake_variance))
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( retake_latents = randn_tensor(
shape=(bsz, 8, 16, frame_length), shape=(bsz, 8, 16, frame_length),
generator=retake_random_generators, generator=retake_random_generators,
device=device, device=self.device,
dtype=dtype, dtype=self.dtype,
) )
repaint_start_frame = int(repaint_start * 44100 / 512 / 8) repaint_start_frame = int(repaint_start * 44100 / 512 / 8)
repaint_end_frame = int(repaint_end * 44100 / 512 / 8) repaint_end_frame = int(repaint_end * 44100 / 512 / 8)
@@ -975,7 +967,7 @@ class ACEStepPipeline:
elif not is_extend: elif not is_extend:
# if repaint_end_frame # if repaint_end_frame
repaint_mask = torch.zeros( 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_mask[:, :, :, repaint_start_frame:repaint_end_frame] = 1.0
repaint_noise = ( repaint_noise = (
@@ -1034,7 +1026,7 @@ class ACEStepPipeline:
gt_latents = extend_gt_latents gt_latents = extend_gt_latents
repaint_mask = torch.zeros( 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: if left_pad_frame_length > 0:
repaint_mask[:, :, :, :left_pad_frame_length] = 1.0 repaint_mask[:, :, :, :left_pad_frame_length] = 1.0
@@ -1073,7 +1065,7 @@ class ACEStepPipeline:
infer_steps=infer_steps, 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 # guidance interval
start_idx = int(num_inference_steps * ((1 - guidance_interval) / 2)) start_idx = int(num_inference_steps * ((1 - guidance_interval) / 2))
@@ -1324,11 +1316,10 @@ class ACEStepPipeline:
if i + 1 < len(timesteps): if i + 1 < len(timesteps):
t_im1 = (timesteps[i + 1]) / 1000 t_im1 = (timesteps[i + 1]) / 1000
else: else:
t_im1 = torch.zeros_like(t_i).to(t_i.device) t_im1 = torch.zeros_like(t_i).to(self.device)
dtype = noise_pred.dtype
target_latents = target_latents.to(torch.float32) target_latents = target_latents.to(torch.float32)
prev_sample = target_latents + (t_im1 - t_i) * noise_pred 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 target_latents = prev_sample
zt_src = (1 - t_im1) * x0 + (t_im1) * z0 zt_src = (1 - t_im1) * x0 + (t_im1) * z0
target_latents = torch.where( target_latents = torch.where(
@@ -1415,8 +1406,7 @@ class ACEStepPipeline:
return None return None
input_audio, sr = self.music_dcae.load_audio(input_audio_path) input_audio, sr = self.music_dcae.load_audio(input_audio_path)
input_audio = input_audio.unsqueeze(0) input_audio = input_audio.unsqueeze(0)
device, dtype = self.device, self.dtype input_audio = input_audio.to(device=self.device, dtype=self.dtype)
input_audio = input_audio.to(device=device, dtype=dtype)
latents, _ = self.music_dcae.encode(input_audio, sr=sr) latents, _ = self.music_dcae.encode(input_audio, sr=sr)
return latents return latents
@@ -1505,20 +1495,14 @@ class ACEStepPipeline:
oss_steps = [] oss_steps = []
texts = [prompt] texts = [prompt]
encoder_text_hidden_states, text_attention_mask = self.get_text_embeddings( encoder_text_hidden_states, text_attention_mask = self.get_text_embeddings(texts)
texts, self.device
)
encoder_text_hidden_states = encoder_text_hidden_states.repeat(batch_size, 1, 1) encoder_text_hidden_states = encoder_text_hidden_states.repeat(batch_size, 1, 1)
text_attention_mask = text_attention_mask.repeat(batch_size, 1) text_attention_mask = text_attention_mask.repeat(batch_size, 1)
encoder_text_hidden_states_null = None encoder_text_hidden_states_null = None
if use_erg_tag: if use_erg_tag:
encoder_text_hidden_states_null = self.get_text_embeddings_null( encoder_text_hidden_states_null = self.get_text_embeddings_null(texts)
texts, self.device encoder_text_hidden_states_null = encoder_text_hidden_states_null.repeat(batch_size, 1, 1)
)
encoder_text_hidden_states_null = encoder_text_hidden_states_null.repeat(
batch_size, 1, 1
)
# not support for released checkpoint # not support for released checkpoint
speaker_embeds = torch.zeros(batch_size, 512).to(self.device).to(self.dtype) speaker_embeds = torch.zeros(batch_size, 512).to(self.device).to(self.dtype)
@@ -1579,7 +1563,7 @@ class ACEStepPipeline:
if task == "edit": if task == "edit":
texts = [edit_target_prompt] texts = [edit_target_prompt]
target_encoder_text_hidden_states, target_text_attention_mask = ( 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 = (
target_encoder_text_hidden_states.repeat(batch_size, 1, 1) target_encoder_text_hidden_states.repeat(batch_size, 1, 1)