add vram optimization cmdline args

acestep/music_dcae/music_dcae_pipeline.py

@@ -144,6 +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.")
         latents = latents / self.scale_factor + self.shift_factor
 
         pred_wavs = []

acestep/pipeline_ace_step.py

@@ -88,7 +88,6 @@ REPO_ID = "ACE-Step/ACE-Step-v1-3.5B"
 
 # class ACEStepPipeline(DiffusionPipeline):
 class ACEStepPipeline:
-
     def __init__(
         self,
         checkpoint_dir=None,
@@ -97,9 +96,9 @@ class ACEStepPipeline:
         text_encoder_checkpoint_path=None,
         persistent_storage_path=None,
         torch_compile=False,
-        cpu_offload=True,
+        cpu_offload=False,
         quantized=False,
-        overlapped_decode=True,
+        overlapped_decode=False,
         **kwargs,
     ):
         if not checkpoint_dir:
@@ -462,7 +461,6 @@ class ACEStepPipeline:
                 from torchao.quantization import (
                     quantize_,
                     Int4WeightOnlyConfig,
-                    Int8WeightOnlyConfig,
                 )
 
                 group_size = 128
@@ -498,6 +496,7 @@ class ACEStepPipeline:
                     os.path.join(text_encoder_model_path, "pytorch_model_int4wo.bin"),
                 )
 
+
     def load_quantized_checkpoint(self, checkpoint_dir=None):
         device = self.device
 
@@ -507,7 +506,6 @@ class ACEStepPipeline:
         text_encoder_model_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

infer.py

@@ -40,15 +40,23 @@ def sample_data(json_data):
 @click.option(
     "--torch_compile", type=bool, default=False, help="Whether to use 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)"
+)
 @click.option("--device_id", type=int, default=0, help="Device ID to use")
 @click.option("--output_path", type=str, default=None, help="Path to save the output")
-def main(checkpoint_path, bf16, torch_compile, device_id, output_path):
+def main(checkpoint_path, bf16, torch_compile, cpu_offload, overlapped_decode, device_id, output_path):
     os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
 
     model_demo = ACEStepPipeline(
         checkpoint_dir=checkpoint_path,
         dtype="bfloat16" if bf16 else "float32",
         torch_compile=torch_compile,
+        cpu_offload=cpu_offload,
+        overlapped_decode=overlapped_decode
     )
     print(model_demo)