add snapshot_download

acestep/pipeline_ace_step.py

@@ -1546,6 +1546,19 @@ class ACEStepPipeline:
         latents, _ = self.music_dcae.encode(input_audio, sr=sr)
         return latents
     
+    def load_lora(self, lora_name_or_path):
+        if lora_name_or_path != "none":
+            if not os.path.exists(lora_name_or_path):
+                lora_download_path = snapshot_download(lora_name_or_path, cache_dir=self.checkpoint_dir)
+            else:
+                lora_download_path = lora_name_or_path
+            self.ace_step_transformer.load_lora_adapter(os.path.join(lora_download_path, "pytorch_lora_weights.safetensors"), adapter_name="zh_rap_lora", with_alpha=True)
+            logger.info(f"Loading lora weights from: {lora_name_or_path} download path is: {lora_download_path}")
+            self.lora_path = lora_name_or_path
+        elif self.lora_path != "none" and lora_name_or_path == "none":
+            logger.info("No lora weights to load.")
+            self.ace_step_transformer.unload_lora_weights()
+
     def __call__(
         self,
         audio_duration: float = 60.0,
@@ -1598,15 +1611,7 @@ class ACEStepPipeline:
                 self.load_quantized_checkpoint(self.checkpoint_dir)
             else:
                 self.load_checkpoint(self.checkpoint_dir)
-            # lora_path=lora_name_or_path
-            if lora_name_or_path != "none":
-                self.ace_step_transformer.load_lora_adapter(os.path.join(lora_name_or_path, "pytorch_lora_weights.safetensors"), adapter_name="zh_rap_lora", with_alpha=True)
-                logger.info(f"Loading lora weights from: {lora_name_or_path}")
-                self.lora_path = lora_name_or_path
-            elif self.lora_path != "none" and lora_name_or_path == "none":
-                logger.info("No lora weights to load.")
-                self.ace_step_transformer.unload_lora_weights()
-
+            self.load_lora(lora_name_or_path)
             load_model_cost = time.time() - start_time
             logger.info(f"Model loaded in {load_model_cost:.2f} seconds.")
 

acestep/ui/components.py

@@ -93,6 +93,7 @@ def create_text2music_ui(
                     label="Lora Name or Path",
                     choices=["ACE-Step/ACE-Step-v1-chinese-rap-LoRA", "none"],
                     value="none",
+                    allow_custom_value=True,
                 )
 
             ref_audio_input = gr.Audio(type="filepath", label="Reference Audio (for Audio2Audio)", visible=False, elem_id="ref_audio_input", show_download_button=True)
@@ -773,8 +774,8 @@ def create_text2music_ui(
                     outputs=extend_outputs + [extend_input_params_json],
                 )
 
-        def sample_data():
-            json_data = sample_data_func()
+        def sample_data(lora_name_or_path_):
+            json_data = sample_data_func(lora_name_or_path_)
             return (
                 json_data["audio_duration"],
                 json_data["prompt"],
@@ -817,11 +818,6 @@ def create_text2music_ui(
                     if "ref_audio_input" in json_data
                     else None
                 ),
-                (
-                    json_data["lora_name_or_path"]
-                    if "lora_name_or_path" in json_data
-                    else "none"
-                )
             )
 
         sample_bnt.click(