From e978c61a799493da84da0cedc6062ef72675faf7 Mon Sep 17 00:00:00 2001
From: azazeal04 <132445160+azazeal04@users.noreply.github.com>
Date: Sat, 4 Apr 2026 16:06:41 +0200
Subject: [PATCH] Implement twinflow_z_image_key_mapping function

Add twinflow_z_image_key_mapping function for key mapping.
---
 comfy/utils.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/comfy/utils.py b/comfy/utils.py
index 78c491b98..0a5aaf3d0 100644
--- a/comfy/utils.py
+++ b/comfy/utils.py
@@ -818,6 +818,16 @@ def z_image_to_diffusers(mmdit_config, output_prefix=""):
 
     return key_map
 
+def twinflow_z_image_key_mapping(state_dict, key):
+    """
+    TwinFlow-Z-Image key mapping.
+    Maps t_embedder_2 keys to t_embedder for weight loading.
+    """
+    if key.startswith("t_embedder_2."):
+        new_key = key.replace("t_embedder_2.", "t_embedder.", 1)
+        state_dict[new_key] = state_dict.pop(key)
+    return state_dict
+
 def repeat_to_batch_size(tensor, batch_size, dim=0):
     if tensor.shape[dim] > batch_size:
         return tensor.narrow(dim, 0, batch_size)