Models v2 - phase 0

extensions_built_in/audio_models/base_audio_model.py

@@ -69,21 +69,7 @@ class BaseAudioModel(BaseModel):
             "save_model is not implemented for this model. Use the pipeline directly instead."
         )
 
-    def convert_lora_weights_before_save(self, state_dict):
-        # currently starte with transformer. but needs to start with diffusion_model. for comfyui
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("transformer.", "diffusion_model.")
-            new_sd[new_key] = value
-        return new_sd
-
-    def convert_lora_weights_before_load(self, state_dict):
-        # saved as diffusion_model. but needs to be transformer. for ai-toolkit
-        new_sd = {}
-        for key, value in state_dict.items():
-            new_key = key.replace("diffusion_model.", "transformer.")
-            new_sd[new_key] = value
-        return new_sd
+    lora_keys_use_comfy_prefix = True
 
     def encode_images(self, image_list: torch.Tensor, device=None, dtype=None):
         # make it more obvious for audio models