Avoid overwriting existing keys in state_dict

Prevent overwriting existing keys when mapping t_embedder_2 keys to t_embedder.
This commit is contained in:
azazeal04 2026-04-04 20:05:17 +02:00 committed by GitHub
parent 6965b83a2a
commit b2e08bb963
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -825,8 +825,9 @@ def twinflow_z_image_key_mapping(state_dict, key):
"""
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
if new_key not in state_dict:
state_dict[new_key] = state_dict.pop(key)
return state_dic
def repeat_to_batch_size(tensor, batch_size, dim=0):
if tensor.shape[dim] > batch_size: