"F64": torch.float64,

"F32": torch.float32,

"F16": torch.float16,

"BF16": torch.bfloat16,

"I64": torch.int64,

"I32": torch.int32,

"I16": torch.int16,

"I8": torch.int8,

"U8": torch.uint8,

"BOOL": torch.bool,

"F8_E4M3": torch.float8_e4m3fn,

"F8_E5M2": torch.float8_e5m2,

"C64": torch.complex64,

"U64": torch.uint64,

"U32": torch.uint32,

"U16": torch.uint16,

import comfy_aimdo.model_mmap

file_size = os.path.getsize(ckpt)

if file_size < 8:

raise _incomplete_safetensors_error("The safetensors header is incomplete.", ckpt)

file_lock = threading.Lock()

model_mmap = comfy_aimdo.model_mmap.ModelMMAP(ckpt)

f = model_mmap.get_file_handle()

mv = memoryview((ctypes.c_uint8 * file_size).from_address(model_mmap.get()))

header_size = struct.unpack("<Q", mv[:8])[0]

if header_size > _SAFETENSORS_MAX_HEADER_SIZE:

raise _invalid_safetensors_error("The safetensors header is too large.", ckpt)

data_base_offset = 8 + header_size

if data_base_offset > file_size:

raise _incomplete_safetensors_error("The safetensors header is incomplete.", ckpt)

try:

header = json.loads(mv[8:data_base_offset].tobytes().decode("utf-8"))

except (UnicodeDecodeError, json.JSONDecodeError) as e:

raise _invalid_safetensors_error(str(e), ckpt) from e

if not isinstance(header, dict):

raise _invalid_safetensors_error("The safetensors header is invalid.", ckpt)

mv = mv[data_base_offset:]

data_size = len(mv)

sd = {}

for name, info in header.items():

if name == "__metadata__":

continue

start, end = info["data_offsets"]

dtype = _TYPES[info["dtype"]]

if start < 0 or end < start:

raise _invalid_safetensors_error("Tensor '{}' has invalid data offsets.".format(name), ckpt)

if end > data_size:

raise _incomplete_safetensors_error("Tensor '{}' extends past the end of the file.".format(name), ckpt)

if math.prod(info["shape"]) * dtype.itemsize != end - start:

raise _invalid_safetensors_error("Tensor '{}' does not match its declared shape and dtype.".format(name), ckpt)

if start == end:

sd[name] = torch.empty(info["shape"], dtype=dtype)

else:

with warnings.catch_warnings():

#We are working with read-only RAM by design

warnings.filterwarnings("ignore", message="The given buffer is not writable")

tensor = torch.frombuffer(mv[start:end], dtype=dtype).view(info["shape"])

storage = tensor.untyped_storage()

setattr(storage,

"_comfy_tensor_file_slice",

comfy.memory_management.TensorFileSlice(f, file_lock, data_base_offset + start, end - start))

setattr(storage, "_comfy_tensor_mmap_refs", (model_mmap, mv))

sd[name] = tensor

if device is None:

device = torch.device("cpu")

metadata = None

if ckpt.lower().endswith(".safetensors") or ckpt.lower().endswith(".sft"):

try:

if comfy.memory_management.aimdo_enabled:

sd, metadata = load_safetensors(ckpt)

if not return_metadata:

metadata = None

else:

with safetensors.safe_open(ckpt, framework="pt", device=device.type) as f:

sd = {}

for k in f.keys():

tensor = f.get_tensor(k)

if DISABLE_MMAP: # TODO: Not sure if this is the best way to bypass the mmap issues

tensor = tensor.to(device=device, copy=True)

sd[k] = tensor

if return_metadata:

metadata = f.metadata()

except Exception as e:

if len(e.args) > 0:

message = e.args[0]

if "HeaderTooLarge" in message:

raise _invalid_safetensors_error(message, ckpt)

if "MetadataIncompleteBuffer" in message:

raise _incomplete_safetensors_error(message, ckpt)

raise e

else:

torch_args = {}

if MMAP_TORCH_FILES:

torch_args["mmap"] = True

pl_sd = torch.load(ckpt, map_location=device, weights_only=True, **torch_args)

if "state_dict" in pl_sd:

sd = pl_sd["state_dict"]

else:

if len(pl_sd) == 1:

key = list(pl_sd.keys())[0]

sd = pl_sd[key]

if not isinstance(sd, dict):

sd = pl_sd

else:

sd = pl_sd

comfy.storage.annotate_state_dict(sd, ckpt)

if metadata is not None:

safetensors.torch.save_file(sd, ckpt, metadata=metadata)

else:

params = 0

for k in sd.keys():

if k.startswith(prefix):

w = sd[k]

params += w.nelement()

dtypes = {}

for k in sd.keys():

if k.startswith(prefix):

w = sd[k]

dtypes[w.dtype] = dtypes.get(w.dtype, 0) + w.numel()

if len(dtypes) == 0:

return None

for x in keys_to_replace:

if x in state_dict:

state_dict[keys_to_replace[x]] = state_dict.pop(x)

if filter_keys:

out = {}

else:

out = state_dict

for rp in replace_prefix:

replace = list(map(lambda a: (a, "{}{}".format(replace_prefix[rp], a[len(rp):])), filter(lambda a: a.startswith(rp), state_dict.keys())))

for x in replace:

w = state_dict.pop(x[0])

out[x[1]] = w

keys_to_replace = {

"{}positional_embedding": "{}embeddings.position_embedding.weight",

"{}token_embedding.weight": "{}embeddings.token_embedding.weight",

"{}ln_final.weight": "{}final_layer_norm.weight",

"{}ln_final.bias": "{}final_layer_norm.bias",

}

for k in keys_to_replace:

x = k.format(prefix_from)

if x in sd:

sd[keys_to_replace[k].format(prefix_to)] = sd.pop(x)

resblock_to_replace = {

"ln_1": "layer_norm1",

"ln_2": "layer_norm2",

"mlp.c_fc": "mlp.fc1",

"mlp.c_proj": "mlp.fc2",

"attn.out_proj": "self_attn.out_proj",

}

for resblock in range(number):

for x in resblock_to_replace:

for y in ["weight", "bias"]:

k = "{}transformer.resblocks.{}.{}.{}".format(prefix_from, resblock, x, y)

k_to = "{}encoder.layers.{}.{}.{}".format(prefix_to, resblock, resblock_to_replace[x], y)

if k in sd:

sd[k_to] = sd.pop(k)

for y in ["weight", "bias"]:

k_from = "{}transformer.resblocks.{}.attn.in_proj_{}".format(prefix_from, resblock, y)

if k_from in sd:

weights = sd.pop(k_from)

shape_from = weights.shape[0] // 3

for x in range(3):

p = ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"]

k_to = "{}encoder.layers.{}.{}.{}".format(prefix_to, resblock, p[x], y)

sd[k_to] = weights[shape_from*x:shape_from*(x + 1)]

sd = transformers_convert(sd, prefix_from, "{}text_model.".format(prefix_to), 32)

tp = "{}text_projection.weight".format(prefix_from)

if tp in sd:

sd["{}text_projection.weight".format(prefix_to)] = sd.pop(tp)

tp = "{}text_projection".format(prefix_from)

if tp in sd:

sd["{}text_projection.weight".format(prefix_to)] = sd.pop(tp).transpose(0, 1).contiguous()

"proj_in.weight",

"proj_in.bias",

"proj_out.weight",

"proj_out.bias",

"norm.weight",

"norm.bias",

"norm1.weight",

"norm1.bias",

"norm2.weight",

"norm2.bias",

"norm3.weight",

"norm3.bias",

"attn1.to_q.weight",

"attn1.to_k.weight",

"attn1.to_v.weight",

"attn1.to_out.0.weight",

"attn1.to_out.0.bias",

"attn2.to_q.weight",

"attn2.to_k.weight",

"attn2.to_v.weight",

"attn2.to_out.0.weight",

"attn2.to_out.0.bias",

"ff.net.0.proj.weight",

"ff.net.0.proj.bias",

"ff.net.2.weight",

"ff.net.2.bias",

"in_layers.2.weight": "conv1.weight",

"in_layers.2.bias": "conv1.bias",

"emb_layers.1.weight": "time_emb_proj.weight",

"emb_layers.1.bias": "time_emb_proj.bias",

"out_layers.3.weight": "conv2.weight",

"out_layers.3.bias": "conv2.bias",

"skip_connection.weight": "conv_shortcut.weight",

"skip_connection.bias": "conv_shortcut.bias",

"in_layers.0.weight": "norm1.weight",

"in_layers.0.bias": "norm1.bias",

"out_layers.0.weight": "norm2.weight",

"out_layers.0.bias": "norm2.bias",

("label_emb.0.0.weight", "class_embedding.linear_1.weight"),

("label_emb.0.0.bias", "class_embedding.linear_1.bias"),

("label_emb.0.2.weight", "class_embedding.linear_2.weight"),

("label_emb.0.2.bias", "class_embedding.linear_2.bias"),

("label_emb.0.0.weight", "add_embedding.linear_1.weight"),

("label_emb.0.0.bias", "add_embedding.linear_1.bias"),

("label_emb.0.2.weight", "add_embedding.linear_2.weight"),

("label_emb.0.2.bias", "add_embedding.linear_2.bias"),

("input_blocks.0.0.weight", "conv_in.weight"),

("input_blocks.0.0.bias", "conv_in.bias"),

("out.0.weight", "conv_norm_out.weight"),

("out.0.bias", "conv_norm_out.bias"),

("out.2.weight", "conv_out.weight"),

("out.2.bias", "conv_out.bias"),

("time_embed.0.weight", "time_embedding.linear_1.weight"),

("time_embed.0.bias", "time_embedding.linear_1.bias"),

("time_embed.2.weight", "time_embedding.linear_2.weight"),

("time_embed.2.bias", "time_embedding.linear_2.bias")

if "num_res_blocks" not in unet_config:

return {}

num_res_blocks = unet_config["num_res_blocks"]

channel_mult = unet_config["channel_mult"]

transformer_depth = unet_config["transformer_depth"][:]

transformer_depth_output = unet_config["transformer_depth_output"][:]

num_blocks = len(channel_mult)

transformers_mid = unet_config.get("transformer_depth_middle", None)

diffusers_unet_map = {}

for x in range(num_blocks):

n = 1 + (num_res_blocks[x] + 1) * x

for i in range(num_res_blocks[x]):

for b in UNET_MAP_RESNET:

diffusers_unet_map["down_blocks.{}.resnets.{}.{}".format(x, i, UNET_MAP_RESNET[b])] = "input_blocks.{}.0.{}".format(n, b)

num_transformers = transformer_depth.pop(0)

if num_transformers > 0:

for b in UNET_MAP_ATTENTIONS:

diffusers_unet_map["down_blocks.{}.attentions.{}.{}".format(x, i, b)] = "input_blocks.{}.1.{}".format(n, b)

for t in range(num_transformers):

for b in TRANSFORMER_BLOCKS:

diffusers_unet_map["down_blocks.{}.attentions.{}.transformer_blocks.{}.{}".format(x, i, t, b)] = "input_blocks.{}.1.transformer_blocks.{}.{}".format(n, t, b)

n += 1

for k in ["weight", "bias"]:

diffusers_unet_map["down_blocks.{}.downsamplers.0.conv.{}".format(x, k)] = "input_blocks.{}.0.op.{}".format(n, k)

i = 0

for b in UNET_MAP_ATTENTIONS:

diffusers_unet_map["mid_block.attentions.{}.{}".format(i, b)] = "middle_block.1.{}".format(b)

for t in range(transformers_mid):

for b in TRANSFORMER_BLOCKS:

diffusers_unet_map["mid_block.attentions.{}.transformer_blocks.{}.{}".format(i, t, b)] = "middle_block.1.transformer_blocks.{}.{}".format(t, b)

for i, n in enumerate([0, 2]):

for b in UNET_MAP_RESNET:

diffusers_unet_map["mid_block.resnets.{}.{}".format(i, UNET_MAP_RESNET[b])] = "middle_block.{}.{}".format(n, b)

num_res_blocks = list(reversed(num_res_blocks))

for x in range(num_blocks):

n = (num_res_blocks[x] + 1) * x

l = num_res_blocks[x] + 1

for i in range(l):

c = 0

for b in UNET_MAP_RESNET:

diffusers_unet_map["up_blocks.{}.resnets.{}.{}".format(x, i, UNET_MAP_RESNET[b])] = "output_blocks.{}.0.{}".format(n, b)

c += 1

num_transformers = transformer_depth_output.pop()

if num_transformers > 0:

c += 1

for b in UNET_MAP_ATTENTIONS:

diffusers_unet_map["up_blocks.{}.attentions.{}.{}".format(x, i, b)] = "output_blocks.{}.1.{}".format(n, b)

for t in range(num_transformers):

for b in TRANSFORMER_BLOCKS:

diffusers_unet_map["up_blocks.{}.attentions.{}.transformer_blocks.{}.{}".format(x, i, t, b)] = "output_blocks.{}.1.transformer_blocks.{}.{}".format(n, t, b)

if i == l - 1:

for k in ["weight", "bias"]:

diffusers_unet_map["up_blocks.{}.upsamplers.0.conv.{}".format(x, k)] = "output_blocks.{}.{}.conv.{}".format(n, c, k)

n += 1

for k in UNET_MAP_BASIC:

diffusers_unet_map[k[1]] = k[0]

("context_embedder.bias", "context_embedder.bias"),

("context_embedder.weight", "context_embedder.weight"),

("t_embedder.mlp.0.bias", "time_text_embed.timestep_embedder.linear_1.bias"),

("t_embedder.mlp.0.weight", "time_text_embed.timestep_embedder.linear_1.weight"),

("t_embedder.mlp.2.bias", "time_text_embed.timestep_embedder.linear_2.bias"),

("t_embedder.mlp.2.weight", "time_text_embed.timestep_embedder.linear_2.weight"),

("x_embedder.proj.bias", "pos_embed.proj.bias"),

("x_embedder.proj.weight", "pos_embed.proj.weight"),

("y_embedder.mlp.0.bias", "time_text_embed.text_embedder.linear_1.bias"),

("y_embedder.mlp.0.weight", "time_text_embed.text_embedder.linear_1.weight"),

("y_embedder.mlp.2.bias", "time_text_embed.text_embedder.linear_2.bias"),

("y_embedder.mlp.2.weight", "time_text_embed.text_embedder.linear_2.weight"),

("pos_embed", "pos_embed.pos_embed"),

("final_layer.adaLN_modulation.1.bias", "norm_out.linear.bias", swap_scale_shift),

("final_layer.adaLN_modulation.1.weight", "norm_out.linear.weight", swap_scale_shift),

("final_layer.linear.bias", "proj_out.bias"),

("final_layer.linear.weight", "proj_out.weight"),

("context_block.adaLN_modulation.1.bias", "norm1_context.linear.bias"),

("context_block.adaLN_modulation.1.weight", "norm1_context.linear.weight"),

("context_block.attn.proj.bias", "attn.to_add_out.bias"),

("context_block.attn.proj.weight", "attn.to_add_out.weight"),

("context_block.mlp.fc1.bias", "ff_context.net.0.proj.bias"),

("context_block.mlp.fc1.weight", "ff_context.net.0.proj.weight"),

("context_block.mlp.fc2.bias", "ff_context.net.2.bias"),

("context_block.mlp.fc2.weight", "ff_context.net.2.weight"),

("context_block.attn.ln_q.weight", "attn.norm_added_q.weight"),

("context_block.attn.ln_k.weight", "attn.norm_added_k.weight"),

("x_block.adaLN_modulation.1.bias", "norm1.linear.bias"),

("x_block.adaLN_modulation.1.weight", "norm1.linear.weight"),

("x_block.attn.proj.bias", "attn.to_out.0.bias"),

("x_block.attn.proj.weight", "attn.to_out.0.weight"),

("x_block.attn.ln_q.weight", "attn.norm_q.weight"),

("x_block.attn.ln_k.weight", "attn.norm_k.weight"),

("x_block.attn2.proj.bias", "attn2.to_out.0.bias"),

("x_block.attn2.proj.weight", "attn2.to_out.0.weight"),

("x_block.attn2.ln_q.weight", "attn2.norm_q.weight"),

("x_block.attn2.ln_k.weight", "attn2.norm_k.weight"),

("x_block.mlp.fc1.bias", "ff.net.0.proj.bias"),

("x_block.mlp.fc1.weight", "ff.net.0.proj.weight"),

("x_block.mlp.fc2.bias", "ff.net.2.bias"),

("x_block.mlp.fc2.weight", "ff.net.2.weight"),

key_map = {}

depth = mmdit_config.get("depth", 0)

num_blocks = mmdit_config.get("num_blocks", depth)

for i in range(num_blocks):

block_from = "transformer_blocks.{}".format(i)

block_to = "{}joint_blocks.{}".format(output_prefix, i)

offset = depth * 64

for end in ("weight", "bias"):

k = "{}.attn.".format(block_from)

qkv = "{}.x_block.attn.qkv.{}".format(block_to, end)

key_map["{}to_q.{}".format(k, end)] = (qkv, (0, 0, offset))

key_map["{}to_k.{}".format(k, end)] = (qkv, (0, offset, offset))

key_map["{}to_v.{}".format(k, end)] = (qkv, (0, offset * 2, offset))

qkv = "{}.context_block.attn.qkv.{}".format(block_to, end)

key_map["{}add_q_proj.{}".format(k, end)] = (qkv, (0, 0, offset))

key_map["{}add_k_proj.{}".format(k, end)] = (qkv, (0, offset, offset))

key_map["{}add_v_proj.{}".format(k, end)] = (qkv, (0, offset * 2, offset))

k = "{}.attn2.".format(block_from)

qkv = "{}.x_block.attn2.qkv.{}".format(block_to, end)

key_map["{}to_q.{}".format(k, end)] = (qkv, (0, 0, offset))

key_map["{}to_k.{}".format(k, end)] = (qkv, (0, offset, offset))

key_map["{}to_v.{}".format(k, end)] = (qkv, (0, offset * 2, offset))

for k in MMDIT_MAP_BLOCK:

key_map["{}.{}".format(block_from, k[1])] = "{}.{}".format(block_to, k[0])

map_basic = MMDIT_MAP_BASIC.copy()

map_basic.add(("joint_blocks.{}.context_block.adaLN_modulation.1.bias".format(depth - 1), "transformer_blocks.{}.norm1_context.linear.bias".format(depth - 1), swap_scale_shift))

map_basic.add(("joint_blocks.{}.context_block.adaLN_modulation.1.weight".format(depth - 1), "transformer_blocks.{}.norm1_context.linear.weight".format(depth - 1), swap_scale_shift))

for k in map_basic:

if len(k) > 2:

key_map[k[1]] = ("{}{}".format(output_prefix, k[0]), None, k[2])

else:

key_map[k[1]] = "{}{}".format(output_prefix, k[0])

("csize_embedder.mlp.0.weight", "adaln_single.emb.resolution_embedder.linear_1.weight"),

("csize_embedder.mlp.0.bias", "adaln_single.emb.resolution_embedder.linear_1.bias"),

("csize_embedder.mlp.2.weight", "adaln_single.emb.resolution_embedder.linear_2.weight"),

("csize_embedder.mlp.2.bias", "adaln_single.emb.resolution_embedder.linear_2.bias"),

("ar_embedder.mlp.0.weight", "adaln_single.emb.aspect_ratio_embedder.linear_1.weight"),

("ar_embedder.mlp.0.bias", "adaln_single.emb.aspect_ratio_embedder.linear_1.bias"),

("ar_embedder.mlp.2.weight", "adaln_single.emb.aspect_ratio_embedder.linear_2.weight"),

("ar_embedder.mlp.2.bias", "adaln_single.emb.aspect_ratio_embedder.linear_2.bias"),

("x_embedder.proj.weight", "pos_embed.proj.weight"),

("x_embedder.proj.bias", "pos_embed.proj.bias"),

("y_embedder.y_embedding", "caption_projection.y_embedding"),

("y_embedder.y_proj.fc1.weight", "caption_projection.linear_1.weight"),

("y_embedder.y_proj.fc1.bias", "caption_projection.linear_1.bias"),

("y_embedder.y_proj.fc2.weight", "caption_projection.linear_2.weight"),

("y_embedder.y_proj.fc2.bias", "caption_projection.linear_2.bias"),

("t_embedder.mlp.0.weight", "adaln_single.emb.timestep_embedder.linear_1.weight"),

("t_embedder.mlp.0.bias", "adaln_single.emb.timestep_embedder.linear_1.bias"),

("t_embedder.mlp.2.weight", "adaln_single.emb.timestep_embedder.linear_2.weight"),

("t_embedder.mlp.2.bias", "adaln_single.emb.timestep_embedder.linear_2.bias"),

("t_block.1.weight", "adaln_single.linear.weight"),

("t_block.1.bias", "adaln_single.linear.bias"),

("final_layer.linear.weight", "proj_out.weight"),

("final_layer.linear.bias", "proj_out.bias"),

("final_layer.scale_shift_table", "scale_shift_table"),

("scale_shift_table", "scale_shift_table"),

("attn.proj.weight", "attn1.to_out.0.weight"),

("attn.proj.bias", "attn1.to_out.0.bias"),

("mlp.fc1.weight", "ff.net.0.proj.weight"),

("mlp.fc1.bias", "ff.net.0.proj.bias"),

("mlp.fc2.weight", "ff.net.2.weight"),

("mlp.fc2.bias", "ff.net.2.bias"),

("cross_attn.proj.weight" ,"attn2.to_out.0.weight"),

("cross_attn.proj.bias" ,"attn2.to_out.0.bias"),

key_map = {}

depth = mmdit_config.get("depth", 0)

offset = mmdit_config.get("hidden_size", 1152)

for i in range(depth):

block_from = "transformer_blocks.{}".format(i)

block_to = "{}blocks.{}".format(output_prefix, i)

for end in ("weight", "bias"):

s = "{}.attn1.".format(block_from)

qkv = "{}.attn.qkv.{}".format(block_to, end)

key_map["{}to_q.{}".format(s, end)] = (qkv, (0, 0, offset))

key_map["{}to_k.{}".format(s, end)] = (qkv, (0, offset, offset))

key_map["{}to_v.{}".format(s, end)] = (qkv, (0, offset * 2, offset))

s = "{}.attn2.".format(block_from)

q = "{}.cross_attn.q_linear.{}".format(block_to, end)

kv = "{}.cross_attn.kv_linear.{}".format(block_to, end)

key_map["{}to_q.{}".format(s, end)] = q

key_map["{}to_k.{}".format(s, end)] = (kv, (0, 0, offset))

key_map["{}to_v.{}".format(s, end)] = (kv, (0, offset, offset))

for k in PIXART_MAP_BLOCK:

key_map["{}.{}".format(block_from, k[1])] = "{}.{}".format(block_to, k[0])

for k in PIXART_MAP_BASIC:

key_map[k[1]] = "{}{}".format(output_prefix, k[0])

n_double_layers = mmdit_config.get("n_double_layers", 0)

n_layers = mmdit_config.get("n_layers", 0)

key_map = {}

for i in range(n_layers):

if i < n_double_layers:

index = i

prefix_from = "joint_transformer_blocks"

prefix_to = "{}double_layers".format(output_prefix)

block_map = {

"attn.to_q.weight": "attn.w2q.weight",

"attn.to_k.weight": "attn.w2k.weight",

"attn.to_v.weight": "attn.w2v.weight",

"attn.to_out.0.weight": "attn.w2o.weight",

"attn.add_q_proj.weight": "attn.w1q.weight",

"attn.add_k_proj.weight": "attn.w1k.weight",

"attn.add_v_proj.weight": "attn.w1v.weight",

"attn.to_add_out.weight": "attn.w1o.weight",

"ff.linear_1.weight": "mlpX.c_fc1.weight",

"ff.linear_2.weight": "mlpX.c_fc2.weight",

"ff.out_projection.weight": "mlpX.c_proj.weight",

"ff_context.linear_1.weight": "mlpC.c_fc1.weight",

"ff_context.linear_2.weight": "mlpC.c_fc2.weight",

"ff_context.out_projection.weight": "mlpC.c_proj.weight",

"norm1.linear.weight": "modX.1.weight",

"norm1_context.linear.weight": "modC.1.weight",

}

else:

index = i - n_double_layers

prefix_from = "single_transformer_blocks"

prefix_to = "{}single_layers".format(output_prefix)

block_map = {

"attn.to_q.weight": "attn.w1q.weight",

"attn.to_k.weight": "attn.w1k.weight",

"attn.to_v.weight": "attn.w1v.weight",

"attn.to_out.0.weight": "attn.w1o.weight",

"norm1.linear.weight": "modCX.1.weight",

"ff.linear_1.weight": "mlp.c_fc1.weight",

"ff.linear_2.weight": "mlp.c_fc2.weight",

"ff.out_projection.weight": "mlp.c_proj.weight"

}

for k in block_map:

key_map["{}.{}.{}".format(prefix_from, index, k)] = "{}.{}.{}".format(prefix_to, index, block_map[k])

MAP_BASIC = {

("positional_encoding", "pos_embed.pos_embed"),

("register_tokens", "register_tokens"),

("t_embedder.mlp.0.weight", "time_step_proj.linear_1.weight"),

("t_embedder.mlp.0.bias", "time_step_proj.linear_1.bias"),

("t_embedder.mlp.2.weight", "time_step_proj.linear_2.weight"),

("t_embedder.mlp.2.bias", "time_step_proj.linear_2.bias"),

("cond_seq_linear.weight", "context_embedder.weight"),

("init_x_linear.weight", "pos_embed.proj.weight"),

("init_x_linear.bias", "pos_embed.proj.bias"),

("final_linear.weight", "proj_out.weight"),

("modF.1.weight", "norm_out.linear.weight", swap_scale_shift),

}

for k in MAP_BASIC:

if len(k) > 2:

key_map[k[1]] = ("{}{}".format(output_prefix, k[0]), None, k[2])

else:

key_map[k[1]] = "{}{}".format(output_prefix, k[0])

n_double_layers = mmdit_config.get("depth", 0)

n_single_layers = mmdit_config.get("depth_single_blocks", 0)

hidden_size = mmdit_config.get("hidden_size", 0)

key_map = {}

for index in range(n_double_layers):

prefix_from = "transformer_blocks.{}".format(index)

prefix_to = "{}double_blocks.{}".format(output_prefix, index)

for end in ("weight", "bias"):

k = "{}.attn.".format(prefix_from)

qkv = "{}.img_attn.qkv.{}".format(prefix_to, end)

key_map["{}to_q.{}".format(k, end)] = (qkv, (0, 0, hidden_size))

key_map["{}to_k.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))

key_map["{}to_v.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))

k = "{}.attn.".format(prefix_from)

qkv = "{}.txt_attn.qkv.{}".format(prefix_to, end)

key_map["{}add_q_proj.{}".format(k, end)] = (qkv, (0, 0, hidden_size))

key_map["{}add_k_proj.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))

key_map["{}add_v_proj.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))

block_map = {

"attn.to_out.0.weight": "img_attn.proj.weight",

"attn.to_out.0.bias": "img_attn.proj.bias",

"norm1.linear.weight": "img_mod.lin.weight",

"norm1.linear.bias": "img_mod.lin.bias",

"norm1_context.linear.weight": "txt_mod.lin.weight",

"norm1_context.linear.bias": "txt_mod.lin.bias",

"attn.to_add_out.weight": "txt_attn.proj.weight",

"attn.to_add_out.bias": "txt_attn.proj.bias",

"ff.net.0.proj.weight": "img_mlp.0.weight",

"ff.net.0.proj.bias": "img_mlp.0.bias",

"ff.net.2.weight": "img_mlp.2.weight",

"ff.net.2.bias": "img_mlp.2.bias",

"ff_context.net.0.proj.weight": "txt_mlp.0.weight",

"ff_context.net.0.proj.bias": "txt_mlp.0.bias",

"ff_context.net.2.weight": "txt_mlp.2.weight",

"ff_context.net.2.bias": "txt_mlp.2.bias",

"ff.linear_in.weight": "img_mlp.0.weight", # LyCoris LoKr

"ff.linear_in.bias": "img_mlp.0.bias",

"ff.linear_out.weight": "img_mlp.2.weight",

"ff.linear_out.bias": "img_mlp.2.bias",

"ff_context.linear_in.weight": "txt_mlp.0.weight",

"ff_context.linear_in.bias": "txt_mlp.0.bias",

"ff_context.linear_out.weight": "txt_mlp.2.weight",

"ff_context.linear_out.bias": "txt_mlp.2.bias",

"attn.norm_q.weight": "img_attn.norm.query_norm.weight",

"attn.norm_k.weight": "img_attn.norm.key_norm.weight",

"attn.norm_added_q.weight": "txt_attn.norm.query_norm.weight",

"attn.norm_added_k.weight": "txt_attn.norm.key_norm.weight",

}

for k in block_map:

key_map["{}.{}".format(prefix_from, k)] = "{}.{}".format(prefix_to, block_map[k])

for index in range(n_single_layers):

prefix_from = "single_transformer_blocks.{}".format(index)

prefix_to = "{}single_blocks.{}".format(output_prefix, index)

for end in ("weight", "bias"):

k = "{}.attn.".format(prefix_from)

qkv = "{}.linear1.{}".format(prefix_to, end)

key_map["{}to_q.{}".format(k, end)] = (qkv, (0, 0, hidden_size))

key_map["{}to_k.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))

key_map["{}to_v.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))

key_map["{}.proj_mlp.{}".format(prefix_from, end)] = (qkv, (0, hidden_size * 3, hidden_size * 4))

block_map = {

"norm.linear.weight": "modulation.lin.weight",

"norm.linear.bias": "modulation.lin.bias",

"proj_out.weight": "linear2.weight",

"proj_out.bias": "linear2.bias",

"attn.norm_q.weight": "norm.query_norm.weight",

"attn.norm_k.weight": "norm.key_norm.weight",

"attn.to_qkv_mlp_proj.weight": "linear1.weight", # Flux 2

"attn.to_out.weight": "linear2.weight", # Flux 2

}

for k in block_map:

key_map["{}.{}".format(prefix_from, k)] = "{}.{}".format(prefix_to, block_map[k])

MAP_BASIC = {

("final_layer.linear.bias", "proj_out.bias"),

("final_layer.linear.weight", "proj_out.weight"),

("img_in.bias", "x_embedder.bias"),

("img_in.weight", "x_embedder.weight"),

("time_in.in_layer.bias", "time_text_embed.timestep_embedder.linear_1.bias"),

("time_in.in_layer.weight", "time_text_embed.timestep_embedder.linear_1.weight"),

("time_in.out_layer.bias", "time_text_embed.timestep_embedder.linear_2.bias"),

("time_in.out_layer.weight", "time_text_embed.timestep_embedder.linear_2.weight"),

("txt_in.bias", "context_embedder.bias"),

("txt_in.weight", "context_embedder.weight"),

("vector_in.in_layer.bias", "time_text_embed.text_embedder.linear_1.bias"),

("vector_in.in_layer.weight", "time_text_embed.text_embedder.linear_1.weight"),

("vector_in.out_layer.bias", "time_text_embed.text_embedder.linear_2.bias"),

("vector_in.out_layer.weight", "time_text_embed.text_embedder.linear_2.weight"),

("guidance_in.in_layer.bias", "time_text_embed.guidance_embedder.linear_1.bias"),

("guidance_in.in_layer.weight", "time_text_embed.guidance_embedder.linear_1.weight"),

("guidance_in.out_layer.bias", "time_text_embed.guidance_embedder.linear_2.bias"),

("guidance_in.out_layer.weight", "time_text_embed.guidance_embedder.linear_2.weight"),

("final_layer.adaLN_modulation.1.bias", "norm_out.linear.bias", swap_scale_shift),

("final_layer.adaLN_modulation.1.weight", "norm_out.linear.weight", swap_scale_shift),

("pos_embed_input.bias", "controlnet_x_embedder.bias"),

("pos_embed_input.weight", "controlnet_x_embedder.weight"),

}

for k in MAP_BASIC:

if len(k) > 2:

key_map[k[1]] = ("{}{}".format(output_prefix, k[0]), None, k[2])

else:

key_map[k[1]] = "{}{}".format(output_prefix, k[0])

n_layers = mmdit_config.get("n_layers", 0)

hidden_size = mmdit_config.get("dim", 0)

n_context_refiner = mmdit_config.get("n_refiner_layers", 2)

n_noise_refiner = mmdit_config.get("n_refiner_layers", 2)

key_map = {}

def add_block_keys(prefix_from, prefix_to, has_adaln=True):

for end in ("weight", "bias"):

k = "{}.attention.".format(prefix_from)

qkv = "{}.attention.qkv.{}".format(prefix_to, end)

key_map["{}to_q.{}".format(k, end)] = (qkv, (0, 0, hidden_size))

key_map["{}to_k.{}".format(k, end)] = (qkv, (0, hidden_size, hidden_size))

key_map["{}to_v.{}".format(k, end)] = (qkv, (0, hidden_size * 2, hidden_size))

block_map = {

"attention.norm_q.weight": "attention.q_norm.weight",

"attention.norm_k.weight": "attention.k_norm.weight",

"attention.to_out.0.weight": "attention.out.weight",

"attention.to_out.0.bias": "attention.out.bias",

"attention_norm1.weight": "attention_norm1.weight",

"attention_norm2.weight": "attention_norm2.weight",

"feed_forward.w1.weight": "feed_forward.w1.weight",

"feed_forward.w2.weight": "feed_forward.w2.weight",

"feed_forward.w3.weight": "feed_forward.w3.weight",

"ffn_norm1.weight": "ffn_norm1.weight",

"ffn_norm2.weight": "ffn_norm2.weight",

}

if has_adaln:

block_map["adaLN_modulation.0.weight"] = "adaLN_modulation.0.weight"

block_map["adaLN_modulation.0.bias"] = "adaLN_modulation.0.bias"

for k, v in block_map.items():

key_map["{}.{}".format(prefix_from, k)] = "{}.{}".format(prefix_to, v)

for i in range(n_layers):

add_block_keys("layers.{}".format(i), "{}layers.{}".format(output_prefix, i))

for i in range(n_context_refiner):

add_block_keys("context_refiner.{}".format(i), "{}context_refiner.{}".format(output_prefix, i))

for i in range(n_noise_refiner):

add_block_keys("noise_refiner.{}".format(i), "{}noise_refiner.{}".format(output_prefix, i))

MAP_BASIC = [

("final_layer.linear.weight", "all_final_layer.2-1.linear.weight"),

("final_layer.linear.bias", "all_final_layer.2-1.linear.bias"),

("final_layer.adaLN_modulation.1.weight", "all_final_layer.2-1.adaLN_modulation.1.weight"),

("final_layer.adaLN_modulation.1.bias", "all_final_layer.2-1.adaLN_modulation.1.bias"),

("x_embedder.weight", "all_x_embedder.2-1.weight"),

("x_embedder.bias", "all_x_embedder.2-1.bias"),

("x_pad_token", "x_pad_token"),

("cap_embedder.0.weight", "cap_embedder.0.weight"),

("cap_embedder.1.weight", "cap_embedder.1.weight"),

("cap_embedder.1.bias", "cap_embedder.1.bias"),

("cap_pad_token", "cap_pad_token"),

("t_embedder.mlp.0.weight", "t_embedder.mlp.0.weight"),

("t_embedder.mlp.0.bias", "t_embedder.mlp.0.bias"),

("t_embedder.mlp.2.weight", "t_embedder.mlp.2.weight"),

("t_embedder.mlp.2.bias", "t_embedder.mlp.2.bias"),

]

for c, diffusers in MAP_BASIC:

key_map[diffusers] = "{}{}".format(output_prefix, c)

n_layers = mmdit_config.get("layers", 0)

n_txt_layerwise = 2 # TextFusionTransformer hardcodes 2 layerwise + 2 refiner blocks

n_txt_refiner = 2

key_map = {}

def add_block(prefix_to, prefix_from):

block_map = {

"attn.to_q": "attn.wq", "attn.to_k": "attn.wk", "attn.to_v": "attn.wv",

"attn.to_gate": "attn.gate", "attn.to_out.0": "attn.wo",

"attn.to_out": "attn.wo", # some tools drop the ".0" on to_out

"ff.gate": "mlp.gate", "ff.up": "mlp.up", "ff.down": "mlp.down",

}

for d, c in block_map.items():

key_map["{}.{}.weight".format(prefix_to, d)] = "{}{}.{}.weight".format(output_prefix, prefix_from, c)

for i in range(n_layers):

add_block("transformer_blocks.{}".format(i), "blocks.{}".format(i))

for i in range(n_txt_layerwise):

add_block("text_fusion.layerwise_blocks.{}".format(i), "txtfusion.layerwise_blocks.{}".format(i))

for i in range(n_txt_refiner):

add_block("text_fusion.refiner_blocks.{}".format(i), "txtfusion.refiner_blocks.{}".format(i))

MAP_BASIC = [

("img_in", "first"),

("time_embed.linear_1", "tmlp.0"),

("time_embed.linear_2", "tmlp.2"),

("time_mod_proj", "tproj.1"),

("txt_in.linear_1", "txtmlp.1"),

("txt_in.linear_2", "txtmlp.3"),

("text_fusion.projector", "txtfusion.projector"),

("final_layer.linear", "last.linear"),

]

for d, c in MAP_BASIC:

key_map["{}.weight".format(d)] = "{}{}.weight".format(output_prefix, c)

if tensor.shape[dim] > batch_size:

return tensor.narrow(dim, 0, batch_size)

elif tensor.shape[dim] < batch_size:

return tensor.repeat(dim * [1] + [math.ceil(batch_size / tensor.shape[dim])] + [1] * (len(tensor.shape) - 1 - dim)).narrow(dim, 0, batch_size)

in_batch_size = tensor.shape[0]

if in_batch_size == batch_size:

return tensor

if batch_size <= 1:

return tensor[:batch_size]

output = torch.empty([batch_size] + list(tensor.shape)[1:], dtype=tensor.dtype, device=tensor.device)

if batch_size < in_batch_size:

scale = (in_batch_size - 1) / (batch_size - 1)

for i in range(batch_size):

output[i] = tensor[min(round(i * scale), in_batch_size - 1)]

else:

scale = in_batch_size / batch_size

for i in range(batch_size):

output[i] = tensor[min(math.floor((i + 0.5) * scale), in_batch_size - 1)]

in_batch_size = len(l)

if in_batch_size == batch_size or in_batch_size == 0:

return l

if batch_size <= 1:

return l[:batch_size]

output = []

if batch_size < in_batch_size:

scale = (in_batch_size - 1) / (batch_size - 1)

for i in range(batch_size):

output.append(l[min(round(i * scale), in_batch_size - 1)])

else:

scale = in_batch_size / batch_size

for i in range(batch_size):

output.append(l[min(math.floor((i + 0.5) * scale), in_batch_size - 1)])

keys = list(state_dict.keys())

for k in keys:

state_dict[k] = state_dict[k].to(dtype)

with open(safetensors_path, "rb") as f:

header = f.read(8)

length_of_header = struct.unpack('<Q', header)[0]

if length_of_header > max_size:

return None

attrs = attr.split(".")

for name in attrs[:-1]:

obj = getattr(obj, name)

obj, name = resolve_attr(obj, attr)

prev = getattr(obj, name, ATTR_UNSET)

if value is ATTR_UNSET:

delattr(obj, name)

else:

setattr(obj, name, value)

# Clone inference tensors (created under torch.inference_mode) since

# their version counter is frozen and nn.Parameter() cannot wrap them.

if (not torch.is_inference_mode_enabled()) and value.is_inference():

value = value.clone()

obj, name = resolve_attr(obj, attr)

prev = getattr(obj, name, ATTR_UNSET)

persistent = name not in getattr(obj, "_non_persistent_buffers_set", set())

obj.register_buffer(name, value, persistent=persistent)

# inplace update tensor instead of replacing it

attrs = attr.split(".")

for name in attrs[:-1]:

obj = getattr(obj, name)

prev = getattr(obj, attrs[-1])

"""Retrieves a nested attribute from an object using dot notation.