int4 weight only quantized model using torchao

acestep/cpu_offload.py

@@ -10,11 +10,13 @@ class CpuOffloader:
         self.original_dtype = model.dtype
     
     def __enter__(self):
-        self.model.to(self.original_device, dtype=self.original_dtype)
+        if not hasattr(self.model,"torchao_quantized"):
+            self.model.to(self.original_device, dtype=self.original_dtype)
         return self.model
     
     def __exit__(self, *args):
-        self.model.to("cpu")
+        if not hasattr(self.model,"torchao_quantized"):
+            self.model.to("cpu")
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
             torch.cuda.synchronize()

acestep/music_dcae/music_dcae_pipeline.py

@@ -124,10 +124,8 @@ class MusicDCAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
             if sr is not None:
                 resampler = (
                     torchaudio.transforms.Resample(44100, sr)
-                    .to(latents.device)
-                    .to(latents.dtype)
                 )
-                wav = resampler(wav)
+                wav = resampler(wav.cpu().float()).to(latent.device).to(latent.dtype)
             else:
                 sr = 44100
             pred_wavs.append(wav)

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
+from transformers import UMT5EncoderModel, AutoTokenizer, AutoConfig
 
 from acestep.language_segmentation import LangSegment
 from acestep.music_dcae.music_dcae_pipeline import MusicDCAE
@@ -97,12 +97,15 @@ class ACEStepPipeline:
         text_encoder_checkpoint_path=None,
         persistent_storage_path=None,
         torch_compile=False,
-        cpu_offload=False,
+        cpu_offload=True,
+        quantized=True,
         **kwargs,
     ):
         if not checkpoint_dir:
             if persistent_storage_path is None:
-                checkpoint_dir = os.path.join(os.path.expanduser("~"), ".cache/ace-step/checkpoints")
+                checkpoint_dir = os.path.join(
+                    os.path.expanduser("~"), ".cache/ace-step/checkpoints"
+                )
                 os.makedirs(checkpoint_dir, exist_ok=True)
             else:
                 checkpoint_dir = os.path.join(persistent_storage_path, "checkpoints")
@@ -125,8 +128,9 @@ class ACEStepPipeline:
         self.loaded = False
         self.torch_compile = torch_compile
         self.cpu_offload = cpu_offload
+        self.quantized = quantized
 
-    def load_checkpoint(self, checkpoint_dir=None):
+    def load_checkpoint(self, checkpoint_dir=None, export_quantized_weights=False):
         device = self.device
 
         dcae_model_path = os.path.join(checkpoint_dir, "music_dcae_f8c8")
@@ -248,22 +252,52 @@ class ACEStepPipeline:
             )
 
             # Verify files were downloaded correctly
-            if not all([
+            if not all(
-                os.path.exists(os.path.join(dcae_model_path, "config.json")),
+                [
-                os.path.exists(os.path.join(dcae_model_path, "diffusion_pytorch_model.safetensors")),
+                    os.path.exists(os.path.join(dcae_model_path, "config.json")),
-                os.path.exists(os.path.join(vocoder_model_path, "config.json")),
+                    os.path.exists(
-                os.path.exists(os.path.join(vocoder_model_path, "diffusion_pytorch_model.safetensors")),
+                        os.path.join(
-                os.path.exists(os.path.join(ace_step_model_path, "config.json")),
+                            dcae_model_path, "diffusion_pytorch_model.safetensors"
-                os.path.exists(os.path.join(ace_step_model_path, "diffusion_pytorch_model.safetensors")),
+                        )
-                os.path.exists(os.path.join(text_encoder_model_path, "config.json")),
+                    ),
-                os.path.exists(os.path.join(text_encoder_model_path, "model.safetensors")),
+                    os.path.exists(os.path.join(vocoder_model_path, "config.json")),
-                os.path.exists(os.path.join(text_encoder_model_path, "special_tokens_map.json")),
+                    os.path.exists(
-            ]):
+                        os.path.join(
-                logger.error("Failed to download all required model files. Please check your internet connection and try again.")
+                            vocoder_model_path, "diffusion_pytorch_model.safetensors"
-                logger.info(f"DCAE model path: {dcae_model_path}, files exist: {os.path.exists(os.path.join(dcae_model_path, 'config.json'))}")
+                        )
-                logger.info(f"Vocoder model path: {vocoder_model_path}, files exist: {os.path.exists(os.path.join(vocoder_model_path, 'config.json'))}")
+                    ),
-                logger.info(f"ACE-Step model path: {ace_step_model_path}, files exist: {os.path.exists(os.path.join(ace_step_model_path, 'config.json'))}")
+                    os.path.exists(os.path.join(ace_step_model_path, "config.json")),
-                logger.info(f"Text encoder model path: {text_encoder_model_path}, files exist: {os.path.exists(os.path.join(text_encoder_model_path, 'config.json'))}")
+                    os.path.exists(
+                        os.path.join(
+                            ace_step_model_path, "diffusion_pytorch_model.safetensors"
+                        )
+                    ),
+                    os.path.exists(
+                        os.path.join(text_encoder_model_path, "config.json")
+                    ),
+                    os.path.exists(
+                        os.path.join(text_encoder_model_path, "model.safetensors")
+                    ),
+                    os.path.exists(
+                        os.path.join(text_encoder_model_path, "special_tokens_map.json")
+                    ),
+                ]
+            ):
+                logger.error(
+                    "Failed to download all required model files. Please check your internet connection and try again."
+                )
+                logger.info(
+                    f"DCAE model path: {dcae_model_path}, files exist: {os.path.exists(os.path.join(dcae_model_path, 'config.json'))}"
+                )
+                logger.info(
+                    f"Vocoder model path: {vocoder_model_path}, files exist: {os.path.exists(os.path.join(vocoder_model_path, 'config.json'))}"
+                )
+                logger.info(
+                    f"ACE-Step model path: {ace_step_model_path}, files exist: {os.path.exists(os.path.join(ace_step_model_path, 'config.json'))}"
+                )
+                logger.info(
+                    f"Text encoder model path: {text_encoder_model_path}, files exist: {os.path.exists(os.path.join(text_encoder_model_path, 'config.json'))}"
+                )
                 raise RuntimeError("Model download failed. See logs for details.")
 
             logger.info("Models downloaded successfully")
@@ -278,7 +312,7 @@ class ACEStepPipeline:
             vocoder_checkpoint_path=vocoder_checkpoint_path,
         )
         # self.music_dcae.to(device).eval().to(self.dtype)
-        if self.cpu_offload: # might be redundant
+        if self.cpu_offload:  # might be redundant
             self.music_dcae = self.music_dcae.to("cpu").eval().to(self.dtype)
         else:
             self.music_dcae = self.music_dcae.to(device).eval().to(self.dtype)
@@ -288,9 +322,13 @@ class ACEStepPipeline:
         )
         # self.ace_step_transformer.to(device).eval().to(self.dtype)
         if self.cpu_offload:
-            self.ace_step_transformer = self.ace_step_transformer.to("cpu").eval().to(self.dtype)
+            self.ace_step_transformer = (
+                self.ace_step_transformer.to("cpu").eval().to(self.dtype)
+            )
         else:
-            self.ace_step_transformer = self.ace_step_transformer.to(device).eval().to(self.dtype)
+            self.ace_step_transformer = (
+                self.ace_step_transformer.to(device).eval().to(self.dtype)
+            )
 
         lang_segment = LangSegment()
 
@@ -418,6 +456,205 @@ class ACEStepPipeline:
             self.ace_step_transformer = torch.compile(self.ace_step_transformer)
             self.text_encoder_model = torch.compile(self.text_encoder_model)
 
+            if export_quantized_weights:
+                from torchao.quantization import (
+                    quantize_,
+                    Int4WeightOnlyConfig,
+                    Int8WeightOnlyConfig,
+                )
+
+                group_size = 128
+                use_hqq = True
+                quantize_(
+                    self.ace_step_transformer,
+                    Int4WeightOnlyConfig(group_size=group_size, use_hqq=use_hqq),
+                )
+                quantize_(
+                    self.text_encoder_model,
+                    Int4WeightOnlyConfig(group_size=group_size, use_hqq=use_hqq),
+                )
+
+                # save quantized weights
+                torch.save(
+                    self.ace_step_transformer.state_dict(),
+                    os.path.join(
+                        ace_step_model_path, "diffusion_pytorch_model_int4wo.bin"
+                    ),
+                )
+                print(
+                    "Quantized Weights Saved to: ",
+                    os.path.join(
+                        ace_step_model_path, "diffusion_pytorch_model_int4wo.bin"
+                    ),
+                )
+                torch.save(
+                    self.text_encoder_model.state_dict(),
+                    os.path.join(text_encoder_model_path, "pytorch_model_int4wo.bin"),
+                )
+                print(
+                    "Quantized Weights Saved to: ",
+                    os.path.join(text_encoder_model_path, "pytorch_model_int4wo.bin"),
+                )
+
+    def load_quantized_checkpoint(self, checkpoint_dir=None):
+        device = self.device
+
+        dcae_model_path = os.path.join(checkpoint_dir, "music_dcae_f8c8")
+        vocoder_model_path = os.path.join(checkpoint_dir, "music_vocoder")
+        ace_step_model_path = os.path.join(checkpoint_dir, "ace_step_transformer")
+        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
+        text_encoder_checkpoint_path = text_encoder_model_path
+
+        self.music_dcae = MusicDCAE(
+            dcae_checkpoint_path=dcae_checkpoint_path,
+            vocoder_checkpoint_path=vocoder_checkpoint_path,
+        )
+        if self.cpu_offload:
+            self.music_dcae.eval().to(self.dtype).to(self.device)
+        else:
+            self.music_dcae.eval().to(self.dtype).to('cpu')
+        self.music_dcae = torch.compile(self.music_dcae)
+
+
+        self.ace_step_transformer = ACEStepTransformer2DModel.from_pretrained(ace_step_checkpoint_path)
+        self.ace_step_transformer.eval().to(self.dtype).to('cpu')
+        self.ace_step_transformer = torch.compile(self.ace_step_transformer)
+        self.ace_step_transformer.load_state_dict(
+            torch.load(
+                os.path.join(ace_step_checkpoint_path, "diffusion_pytorch_model_int4wo.bin"),
+                map_location=self.device,
+            ),assign=True
+        )
+        self.ace_step_transformer.torchao_quantized = True
+
+        self.text_encoder_model = UMT5EncoderModel.from_pretrained(text_encoder_checkpoint_path)
+        self.text_encoder_model.eval().to(self.dtype).to('cpu')
+        self.text_encoder_model = torch.compile(self.text_encoder_model)
+        self.text_encoder_model.load_state_dict(
+            torch.load(
+                os.path.join(text_encoder_model_path, "pytorch_model_int4wo.bin"),
+                map_location=self.device,
+            ),assign=True
+        )
+        self.text_encoder_model.torchao_quantized = True
+
+        self.text_tokenizer = AutoTokenizer.from_pretrained(
+            text_encoder_checkpoint_path
+        )
+
+        lang_segment = LangSegment()
+        lang_segment.setfilters(
+            [
+                "af",
+                "am",
+                "an",
+                "ar",
+                "as",
+                "az",
+                "be",
+                "bg",
+                "bn",
+                "br",
+                "bs",
+                "ca",
+                "cs",
+                "cy",
+                "da",
+                "de",
+                "dz",
+                "el",
+                "en",
+                "eo",
+                "es",
+                "et",
+                "eu",
+                "fa",
+                "fi",
+                "fo",
+                "fr",
+                "ga",
+                "gl",
+                "gu",
+                "he",
+                "hi",
+                "hr",
+                "ht",
+                "hu",
+                "hy",
+                "id",
+                "is",
+                "it",
+                "ja",
+                "jv",
+                "ka",
+                "kk",
+                "km",
+                "kn",
+                "ko",
+                "ku",
+                "ky",
+                "la",
+                "lb",
+                "lo",
+                "lt",
+                "lv",
+                "mg",
+                "mk",
+                "ml",
+                "mn",
+                "mr",
+                "ms",
+                "mt",
+                "nb",
+                "ne",
+                "nl",
+                "nn",
+                "no",
+                "oc",
+                "or",
+                "pa",
+                "pl",
+                "ps",
+                "pt",
+                "qu",
+                "ro",
+                "ru",
+                "rw",
+                "se",
+                "si",
+                "sk",
+                "sl",
+                "sq",
+                "sr",
+                "sv",
+                "sw",
+                "ta",
+                "te",
+                "th",
+                "tl",
+                "tr",
+                "ug",
+                "uk",
+                "ur",
+                "vi",
+                "vo",
+                "wa",
+                "xh",
+                "zh",
+                "zu",
+            ]
+        )
+        self.lang_segment = lang_segment
+        self.lyric_tokenizer = VoiceBpeTokenizer()
+
+        self.loaded = True
+
     @cpu_offload("text_encoder_model")
     def get_text_embeddings(self, texts, device, text_max_length=256):
         inputs = self.text_tokenizer(
@@ -486,7 +723,9 @@ class ACEStepPipeline:
                     processed_input_seeds = list(map(int, manual_seeds.split(",")))
                 elif manual_seeds.isdigit():
                     processed_input_seeds = int(manual_seeds)
-            elif isinstance(manual_seeds, list) and all(isinstance(s, int) for s in manual_seeds):
+            elif isinstance(manual_seeds, list) and all(
+                isinstance(s, int) for s in manual_seeds
+            ):
                 if len(manual_seeds) > 0:
                     processed_input_seeds = list(manual_seeds)
             elif isinstance(manual_seeds, int):
@@ -507,7 +746,7 @@ class ACEStepPipeline:
                 else:
                     current_seed_for_generator = processed_input_seeds[-1]
             if current_seed_for_generator is None:
-                 current_seed_for_generator = torch.randint(0, 2**32, (1,)).item()
+                current_seed_for_generator = torch.randint(0, 2**32, (1,)).item()
             random_generators[i].manual_seed(current_seed_for_generator)
             actual_seeds.append(current_seed_for_generator)
         return random_generators, actual_seeds
@@ -1381,7 +1620,11 @@ class ACEStepPipeline:
         pred_wavs = [pred_wav.cpu().float() for pred_wav in pred_wavs]
         for i in tqdm(range(bs)):
             output_audio_path = self.save_wav_file(
-                pred_wavs[i], i, save_path=save_path, sample_rate=sample_rate, format=format
+                pred_wavs[i],
+                i,
+                save_path=save_path,
+                sample_rate=sample_rate,
+                format=format,
             )
             output_audio_paths.append(output_audio_path)
         return output_audio_paths
@@ -1399,8 +1642,12 @@ class ACEStepPipeline:
         else:
             ensure_directory_exists(os.path.dirname(save_path))
             if os.path.isdir(save_path):
-                logger.info(f"Provided save_path '{save_path}' is a directory. Appending timestamped filename.")
+                logger.info(
-                output_path_wav = os.path.join(save_path, f"output_{time.strftime('%Y%m%d%H%M%S')}_{idx}.wav")
+                    f"Provided save_path '{save_path}' is a directory. Appending timestamped filename."
+                )
+                output_path_wav = os.path.join(
+                    save_path, f"output_{time.strftime('%Y%m%d%H%M%S')}_{idx}.wav"
+                )
             else:
                 output_path_wav = save_path
 
@@ -1463,7 +1710,10 @@ class ACEStepPipeline:
 
         if not self.loaded:
             logger.warning("Checkpoint not loaded, loading checkpoint...")
-            self.load_checkpoint(self.checkpoint_dir)
+            if self.quantized:
+                self.load_quantized_checkpoint(self.checkpoint_dir)
+            else:
+                self.load_checkpoint(self.checkpoint_dir)
             load_model_cost = time.time() - start_time
             logger.info(f"Model loaded in {load_model_cost:.2f} seconds.")
 

infer.py

@@ -83,7 +83,8 @@ def main(checkpoint_path, bf16, torch_compile, device_id, output_path):
         audio_duration,
         prompt,
         lyrics,
-        infer_step,
+        # infer_step,
+        10,
         guidance_scale,
         scheduler_type,
         cfg_type,