Models v2 - phase 0

extensions_built_in/diffusion_models/mageflow/mageflow.py

@@ -629,18 +629,7 @@ class MageFlowModel(BaseModel):
     def get_transformer_block_names(self) -> Optional[List[str]]:
         return ["transformer_blocks"]
 
-    def convert_lora_weights_before_save(self, state_dict):
-        return {
-            k.replace("transformer.", "diffusion_model."): v
-            for k, v in state_dict.items()
-        }
-
-    def convert_lora_weights_before_load(self, state_dict):
-        return {
-            k.replace("diffusion_model.", "transformer."): v
-            for k, v in state_dict.items()
-        }
-
+    lora_keys_use_comfy_prefix = True
 
 class MageFlowEditModel(MageFlowModel):
     arch = "mageflow_edit"