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)