From 25b51b1a8b6be9c3dadd1d755c78394009c4d1d4 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 25 Jul 2024 22:42:54 -0400
Subject: [PATCH] Hunyuan DiT lora support.

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

diff --git a/comfy/lora.py b/comfy/lora.py
index e36b354f..fdc128c0 100644
--- a/comfy/lora.py
+++ b/comfy/lora.py
@@ -282,4 +282,10 @@ def model_lora_keys_unet(model, key_map={}):
                 key_lora = "transformer.{}".format(k[:-len(".weight")]) #simpletrainer and probably regular diffusers lora format
                 key_map[key_lora] = to
 
+    if isinstance(model, comfy.model_base.HunyuanDiT):
+        for k in sdk:
+            if k.startswith("diffusion_model.") and k.endswith(".weight"):
+                key_lora = k[len("diffusion_model."):-len(".weight")]
+                key_map["base_model.model.{}".format(key_lora)] = k #official hunyuan lora format
+
     return key_map