From 481732a0ed609ccd86cb3d2df00e3d14c9133280 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Tue, 13 May 2025 04:32:16 -0700
Subject: [PATCH] Support official ACE Step loras. (#8094)

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

diff --git a/comfy/lora.py b/comfy/lora.py
index fff524be..ef110c16 100644
--- a/comfy/lora.py
+++ b/comfy/lora.py
@@ -286,6 +286,12 @@ def model_lora_keys_unet(model, key_map={}):
                     key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_")
                     key_map["lycoris_{}".format(key_lora)] = k #SimpleTuner lycoris format
 
+    if isinstance(model, comfy.model_base.ACEStep):
+        for k in sdk:
+            if k.startswith("diffusion_model.") and k.endswith(".weight"): #Official ACE step lora format
+                key_lora = k[len("diffusion_model."):-len(".weight")]
+                key_map["{}".format(key_lora)] = k
+
     return key_map