Models v2 - phase 0

extensions_built_in/diffusion_models/chroma/chroma_model.py

@@ -443,21 +443,8 @@ class ChromaModel(BaseModel):
         batch = kwargs.get('batch')
         return (noise - batch.latents).detach()
     
-    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
+    lora_keys_use_comfy_prefix = True
 
-    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
     
     def get_base_model_version(self):
         return "chroma"