diff --git a/acestep/gui.py b/acestep/gui.py index 0b241cf..a081498 100644 --- a/acestep/gui.py +++ b/acestep/gui.py @@ -46,7 +46,13 @@ from acestep.data_sampler import DataSampler @click.option( "--torch_compile", type=click.BOOL, default=False, help="Whether to use torch.compile." ) -def main(checkpoint_path, server_name, port, device_id, share, bf16, torch_compile): +@click.option( + "--cpu_offload", type=bool, default=False, help="Whether to use CPU offloading (only load current stage's model to GPU)" +) +@click.option( + "--overlapped_decode", type=bool, default=False, help="Whether to use overlapped decoding (run dcae and vocoder using sliding windows)" +) +def main(checkpoint_path, server_name, port, device_id, share, bf16, torch_compile, cpu_offload, overlapped_decode): """ Main function to launch the ACE Step pipeline demo. """ @@ -57,6 +63,8 @@ def main(checkpoint_path, server_name, port, device_id, share, bf16, torch_compi checkpoint_dir=checkpoint_path, dtype="bfloat16" if bf16 else "float32", torch_compile=torch_compile, + cpu_offload=cpu_offload, + overlapped_decode=overlapped_decode ) data_sampler = DataSampler() diff --git a/acestep/music_dcae/music_dcae_pipeline.py b/acestep/music_dcae/music_dcae_pipeline.py index c535ee9..c12c02b 100644 --- a/acestep/music_dcae/music_dcae_pipeline.py +++ b/acestep/music_dcae/music_dcae_pipeline.py @@ -144,7 +144,7 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin): @torch.no_grad() def decode_overlap(self, latents, audio_lengths=None, sr=None): - print("Using Overlapped DCAE and Vocoder Decoding.") + print("Using Overlapped DCAE and Vocoder ") latents = latents / self.scale_factor + self.shift_factor pred_wavs = [] diff --git a/acestep/pipeline_ace_step.py b/acestep/pipeline_ace_step.py index 1a4a258..e737906 100644 --- a/acestep/pipeline_ace_step.py +++ b/acestep/pipeline_ace_step.py @@ -30,7 +30,7 @@ from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3 import ( retrieve_timesteps, ) from diffusers.utils.torch_utils import randn_tensor -from transformers import UMT5EncoderModel, AutoTokenizer, AutoConfig +from transformers import UMT5EncoderModel, AutoTokenizer from acestep.language_segmentation import LangSegment from acestep.music_dcae.music_dcae_pipeline import MusicDCAE