refactor quantization method selection

acestep/handler.py

@@ -247,15 +247,22 @@ class AceStepHandler:
                 if compile_model:
                     self.model = torch.compile(self.model)
                     
-                    if self.quantization == "int8_weight_only":
-                        from torchao.quantization import quantize_, Int8WeightOnlyConfig
-                        quantize_(self.model, Int8WeightOnlyConfig())
-                        logger.info("DiT quantized with Int8WeightOnlyConfig")
-                    elif self.quantization == "fp8_weight_only":
-                        from torchao.quantization import quantize_, Float8WeightOnlyConfig
-                        quantize_(self.model, Float8WeightOnlyConfig())                        
-                    elif self.quantization is not None:
-                        raise ValueError(f"Unsupported quantization type: {self.quantization}")
+                    if self.quantization is not None:
+                        from torchao.quantization import quantize_
+                        if self.quantization == "int8_weight_only":
+                            from torchao.quantization import Int8WeightOnlyConfig
+                            quant_config = Int8WeightOnlyConfig()
+                        elif self.quantization == "fp8_weight_only":
+                            from torchao.quantization import Float8WeightOnlyConfig
+                            quant_config = Float8WeightOnlyConfig()
+                        elif self.quantization == "w8a8_dynamic":
+                            from torchao.quantization import Int8DynamicActivationInt8WeightConfig, MappingType
+                            quant_config = Int8DynamicActivationInt8WeightConfig(act_mapping_type=MappingType.ASYMMETRIC)
+                        else:
+                            raise ValueError(f"Unsupported quantization type: {self.quantization}")
+                        
+                        quantize_(self.model, quant_config)
+                        logger.info("DiT quantized with:",self.quantization)
                     
                     
                 silence_latent_path = os.path.join(acestep_v15_checkpoint_path, "silence_latent.pt")

test.py

@@ -46,7 +46,7 @@ def main():
         compile_model=True,
         offload_to_cpu=True,
         offload_dit_to_cpu=False, # Keep DiT on GPU
-        quantization="fp8_weight_only", # Enable FP8 weight-only quantization
+        quantization="int8_weight_only", # Enable FP8 weight-only quantization
     )
     
     if not enabled:
@@ -108,7 +108,12 @@ def main():
         print(f"Generated Audio Codes (first 50 chars): {audio_codes[:50]}...")
     else:
         print("Skipping 5Hz LLM generation...")
-        metadata = {}
+        metadata = {
+            'bpm': 90,
+            'keyscale': 'A major',
+            'timesignature': '4',
+            'duration': 240,
+        }
         audio_codes = None
         lm_status = "Skipped"