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

rearrange model loading order, interleave compilation with loading

Browse Source
This commit is contained in:
Michael Hedman
2025-05-16 17:32:18 +02:00
parent 6c14becf98
commit c50acb83a6
1 changed files with 17 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acestep/pipeline_ace_step.py
+17 -14
View File
@@ -171,16 +171,6 @@ class ACEStepPipeline:
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)
self.ace_step_transformer = ACEStepTransformer2DModel.from_pretrained(
ace_step_checkpoint_path, torch_dtype=self.dtype
)
@@ -193,6 +183,20 @@ class ACEStepPipeline:
self.ace_step_transformer = (
self.ace_step_transformer.to(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(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)
if self.torch_compile:
self.music_dcae = torch.compile(self.music_dcae)
lang_segment = LangSegment()
@@ -309,6 +313,9 @@ class ACEStepPipeline:
text_encoder_model = text_encoder_model.to(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
)
@@ -316,10 +323,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_,