From 6e7dc43b12b4fa06de8da10a2cc1e9a1dcf9e13c Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Wed, 28 May 2025 18:46:31 -0400
Subject: [PATCH] Support HiDream SimpleTuner loras.

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

diff --git a/comfy/lora.py b/comfy/lora.py
index ef110c164..387d5c52a 100644
--- a/comfy/lora.py
+++ b/comfy/lora.py
@@ -283,8 +283,9 @@ def model_lora_keys_unet(model, key_map={}):
         for k in sdk:
             if k.startswith("diffusion_model."):
                 if k.endswith(".weight"):
-                    key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_")
-                    key_map["lycoris_{}".format(key_lora)] = k #SimpleTuner lycoris format
+                    key_lora = k[len("diffusion_model."):-len(".weight")]
+                    key_map["lycoris_{}".format(key_lora.replace(".", "_"))] = k #SimpleTuner lycoris format
+                    key_map["transformer.{}".format(key_lora)] = k #SimpleTuner regular format
 
     if isinstance(model, comfy.model_base.ACEStep):
         for k in sdk: