From bc84acb5747e39c35d719fc912423a782126edf2 Mon Sep 17 00:00:00 2001 From: chuxij Date: Mon, 5 May 2025 04:51:58 +0000 Subject: [PATCH] add torch_dtype when load to avoice oom for windows --- pipeline_ace_step.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pipeline_ace_step.py b/pipeline_ace_step.py index 4281721..9f3b671 100644 --- a/pipeline_ace_step.py +++ b/pipeline_ace_step.py @@ -143,7 +143,7 @@ class ACEStepPipeline: 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) - self.ace_step_transformer = ACEStepTransformer2DModel.from_pretrained(ace_step_checkpoint_path) + self.ace_step_transformer = ACEStepTransformer2DModel.from_pretrained(ace_step_checkpoint_path, torch_dtype=self.dtype) self.ace_step_transformer.to(device).eval().to(self.dtype) lang_segment = LangSegment()