From ceee378d86b6f40a4856fb3e7c294f9301bc0e09 Mon Sep 17 00:00:00 2001
From: rickard <rickard@mindemia.com>
Date: Wed, 25 Dec 2024 18:04:56 +0100
Subject: [PATCH] replace 'diffusion_model' with 'transformer'

---
 comfy/lora.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/comfy/lora.py b/comfy/lora.py
index d6ad5fb8..e3ddb904 100644
--- a/comfy/lora.py
+++ b/comfy/lora.py
@@ -401,9 +401,9 @@ def model_lora_keys_unet(model, key_map={}):
 
     if isinstance(model, comfy.model_base.LTXV):
         for k in sdk:
-            if k.startswith("transformer.") and k.endswith(".weight"): #Official Mochi lora format
-                key_lora = k[len("transformer."):-len(".weight")]
-                key_map["{}".format(key_lora)] = k
+            if k.startswith("diffusion_model.") and k.endswith(".weight"):
+                key_lora = k[len("diffusion_model."):-len(".weight")]
+                key_map["transformer.{}".format(key_lora)] = k
 
     return key_map