gui add cmdline args

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()
 

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 = []

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