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

reduce path vars

Browse Source
This commit is contained in:
Michael Hedman
2025-05-19 06:47:04 +02:00
parent 21efcd3905
commit d470cd903f
1 changed files with 13 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
acestep/pipeline_ace_step.py
+13 -24
View File
@@ -178,15 +178,10 @@ class ACEStepPipeline:
logger.info(f"Download models from Hugging Face: {REPO_ID}, cache to: {checkpoint_dir}") 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) checkpoint_dir_models = snapshot_download(REPO_ID, cache_dir=checkpoint_dir)
dcae_model_path = os.path.join(checkpoint_dir_models, "music_dcae_f8c8") dcae_checkpoint_path = os.path.join(checkpoint_dir_models, "music_dcae_f8c8")
vocoder_model_path = os.path.join(checkpoint_dir_models, "music_vocoder") vocoder_checkpoint_path = os.path.join(checkpoint_dir_models, "music_vocoder")
ace_step_model_path = os.path.join(checkpoint_dir_models, "ace_step_transformer") ace_step_checkpoint_path = os.path.join(checkpoint_dir_models, "ace_step_transformer")
text_encoder_model_path = os.path.join(checkpoint_dir_models, "umt5-base") text_encoder_checkpoint_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.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
@@ -261,38 +256,32 @@ class ACEStepPipeline:
torch.save( torch.save(
self.ace_step_transformer.state_dict(), self.ace_step_transformer.state_dict(),
os.path.join( os.path.join(
ace_step_model_path, "diffusion_pytorch_model_int4wo.bin" ace_step_checkpoint_path, "diffusion_pytorch_model_int4wo.bin"
), ),
) )
print( print(
"Quantized Weights Saved to: ", "Quantized Weights Saved to: ",
os.path.join( os.path.join(
ace_step_model_path, "diffusion_pytorch_model_int4wo.bin" ace_step_checkpoint_path, "diffusion_pytorch_model_int4wo.bin"
), ),
) )
torch.save( torch.save(
self.text_encoder_model.state_dict(), 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( print(
"Quantized Weights Saved to: ", "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): def load_quantized_checkpoint(self, checkpoint_dir=None):
device = self.device device = self.device
dcae_model_path = os.path.join(checkpoint_dir, "music_dcae_f8c8") dcae_checkpoint_path = os.path.join(checkpoint_dir, "music_dcae_f8c8")
vocoder_model_path = os.path.join(checkpoint_dir, "music_vocoder") vocoder_checkpoint_path = os.path.join(checkpoint_dir, "music_vocoder")
ace_step_model_path = os.path.join(checkpoint_dir, "ace_step_transformer") ace_step_checkpoint_path = os.path.join(checkpoint_dir, "ace_step_transformer")
text_encoder_model_path = os.path.join(checkpoint_dir, "umt5-base") text_encoder_checkpoint_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
self.music_dcae = MusicDCAE( self.music_dcae = MusicDCAE(
dcae_checkpoint_path=dcae_checkpoint_path, dcae_checkpoint_path=dcae_checkpoint_path,
@@ -321,7 +310,7 @@ class ACEStepPipeline:
self.text_encoder_model = torch.compile(self.text_encoder_model) self.text_encoder_model = torch.compile(self.text_encoder_model)
self.text_encoder_model.load_state_dict( self.text_encoder_model.load_state_dict(
torch.load( 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, map_location=self.device,
),assign=True ),assign=True
) )