count = 0

while True:

c = False

for k in state_dict_keys:

if k.startswith(prefix_string.format(count)):

c = True

break

if c == False:

break

count += 1

for x in suffix_list:

if "{}{}{}".format(prefix, main, x) in keys:

return True

context_dim = None

use_linear_in_transformer = False

transformer_prefix = prefix + "1.transformer_blocks."

transformer_keys = sorted(list(filter(lambda a: a.startswith(transformer_prefix), state_dict_keys)))

if len(transformer_keys) > 0:

last_transformer_depth = count_blocks(state_dict_keys, transformer_prefix + '{}')

context_dim = state_dict['{}0.attn2.to_k.weight'.format(transformer_prefix)].shape[1]

use_linear_in_transformer = len(state_dict['{}1.proj_in.weight'.format(prefix)].shape) == 2

time_stack = '{}1.time_stack.0.attn1.to_q.weight'.format(prefix) in state_dict or '{}1.time_mix_blocks.0.attn1.to_q.weight'.format(prefix) in state_dict

time_stack_cross = '{}1.time_stack.0.attn2.to_q.weight'.format(prefix) in state_dict or '{}1.time_mix_blocks.0.attn2.to_q.weight'.format(prefix) in state_dict

return last_transformer_depth, context_dim, use_linear_in_transformer, time_stack, time_stack_cross

state_dict_keys = list(state_dict.keys())

if (

'{}cond_layer_logits'.format(key_prefix) in state_dict_keys

and '{}latent_conditioners.0.weight'.format(key_prefix) in state_dict_keys

and '{}diffusion_transformer.transformer.layers.0.self_attn.to_qkv.weight'.format(key_prefix) in state_dict_keys

):

return {"audio_model": "minimax_music3"}

if '{}joint_blocks.0.context_block.attn.qkv.weight'.format(key_prefix) in state_dict_keys: #mmdit model

unet_config = {}

unet_config["in_channels"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[1]

patch_size = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[2]

unet_config["patch_size"] = patch_size

final_layer = '{}final_layer.linear.weight'.format(key_prefix)

if final_layer in state_dict:

unet_config["out_channels"] = state_dict[final_layer].shape[0] // (patch_size * patch_size)

unet_config["depth"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[0] // 64

unet_config["input_size"] = None

y_key = '{}y_embedder.mlp.0.weight'.format(key_prefix)

if y_key in state_dict_keys:

unet_config["adm_in_channels"] = state_dict[y_key].shape[1]

context_key = '{}context_embedder.weight'.format(key_prefix)

if context_key in state_dict_keys:

in_features = state_dict[context_key].shape[1]

out_features = state_dict[context_key].shape[0]

unet_config["context_embedder_config"] = {"target": "torch.nn.Linear", "params": {"in_features": in_features, "out_features": out_features}}

num_patches_key = '{}pos_embed'.format(key_prefix)

if num_patches_key in state_dict_keys:

num_patches = state_dict[num_patches_key].shape[1]

unet_config["num_patches"] = num_patches

unet_config["pos_embed_max_size"] = round(math.sqrt(num_patches))

rms_qk = '{}joint_blocks.0.context_block.attn.ln_q.weight'.format(key_prefix)

if rms_qk in state_dict_keys:

unet_config["qk_norm"] = "rms"

unet_config["pos_embed_scaling_factor"] = None #unused for inference

context_processor = '{}context_processor.layers.0.attn.qkv.weight'.format(key_prefix)

if context_processor in state_dict_keys:

unet_config["context_processor_layers"] = count_blocks(state_dict_keys, '{}context_processor.layers.'.format(key_prefix) + '{}.')

unet_config["x_block_self_attn_layers"] = []

for key in state_dict_keys:

if key.startswith('{}joint_blocks.'.format(key_prefix)) and key.endswith('.x_block.attn2.qkv.weight'):

layer = key[len('{}joint_blocks.'.format(key_prefix)):-len('.x_block.attn2.qkv.weight')]

unet_config["x_block_self_attn_layers"].append(int(layer))

return unet_config

if '{}clf.1.weight'.format(key_prefix) in state_dict_keys: #stable cascade

unet_config = {}

text_mapper_name = '{}clip_txt_mapper.weight'.format(key_prefix)

if text_mapper_name in state_dict_keys:

unet_config['stable_cascade_stage'] = 'c'

w = state_dict[text_mapper_name]

if w.shape[0] == 1536: #stage c lite

unet_config['c_cond'] = 1536

unet_config['c_hidden'] = [1536, 1536]

unet_config['nhead'] = [24, 24]

unet_config['blocks'] = [[4, 12], [12, 4]]

elif w.shape[0] == 2048: #stage c full

unet_config['c_cond'] = 2048

elif '{}clip_mapper.weight'.format(key_prefix) in state_dict_keys:

unet_config['stable_cascade_stage'] = 'b'

w = state_dict['{}down_blocks.1.0.channelwise.0.weight'.format(key_prefix)]

if w.shape[-1] == 640:

unet_config['c_hidden'] = [320, 640, 1280, 1280]

unet_config['nhead'] = [-1, -1, 20, 20]

unet_config['blocks'] = [[2, 6, 28, 6], [6, 28, 6, 2]]

unet_config['block_repeat'] = [[1, 1, 1, 1], [3, 3, 2, 2]]

elif w.shape[-1] == 576: #stage b lite

unet_config['c_hidden'] = [320, 576, 1152, 1152]

unet_config['nhead'] = [-1, 9, 18, 18]

unet_config['blocks'] = [[2, 4, 14, 4], [4, 14, 4, 2]]

unet_config['block_repeat'] = [[1, 1, 1, 1], [2, 2, 2, 2]]

return unet_config

shape_key = '{}img2shape.t_embedder.mlp.0.weight'.format(key_prefix)

tex_key = '{}shape2txt.t_embedder.mlp.0.weight'.format(key_prefix)

if shape_key in state_dict_keys or tex_key in state_dict_keys: # trellis2 / pixal3d

has_shape = shape_key in state_dict_keys

has_tex = tex_key in state_dict_keys

unet_config = {

"image_model": "trellis2",

"resolution": 32 if (metadata or {}).get("is_512") else 64,

"init_txt_model": has_tex,

"txt_only": has_tex and not has_shape,

}

# Per-submodel projection head (Pixal3D adds `proj_linear`; Trellis2 doesn't).

for sub, name, proj_in_channels in (("img2shape", "shape", 2048),

("shape2txt", "texture", 2048),

("structure_model", "structure", 1024)):

key = '{}{}.blocks.0.cross_attn.proj_linear.weight'.format(key_prefix, sub)

if key in state_dict_keys:

unet_config["image_attn_mode_{}".format(name)] = "proj"

unet_config["proj_in_channels_{}".format(name)] = proj_in_channels

return unet_config

if '{}transformer.rotary_pos_emb.inv_freq'.format(key_prefix) in state_dict_keys: #stable audio dit

unet_config = {}

unet_config["audio_model"] = "dit1.0"

unet_config["global_cond_dim"] = state_dict['{}to_global_embed.0.weight'.format(key_prefix)].shape[1]

cond_embed = state_dict['{}to_cond_embed.0.weight'.format(key_prefix)]

unet_config["project_cond_tokens"] = cond_embed.shape[0] != cond_embed.shape[1]

unet_config["embed_dim"] = state_dict['{}to_timestep_embed.0.weight'.format(key_prefix)].shape[0]

mem_tokens = state_dict.get('{}transformer.memory_tokens'.format(key_prefix), None)

to_qkv = state_dict.get('{}transformer.layers.0.self_attn.to_qkv.weight'.format(key_prefix), None)

differential = False

if to_qkv is not None:

if to_qkv.shape[0] == to_qkv.shape[1] * 5:

differential = True

if mem_tokens is not None:

unet_config["num_memory_tokens"] = mem_tokens.shape[0]

if '{}transformer.layers.0.self_attn.q_norm.weight'.format(key_prefix) in state_dict:

unet_config["attn_kwargs"] = {"qk_norm": "ln", "feat_scale": True}

rms_norm = state_dict.get('{}transformer.layers.0.self_attn.q_norm.gamma'.format(key_prefix), None)

if rms_norm is not None:

unet_config["attn_kwargs"] = {"qk_norm": "rms", "differential": differential}

unet_config["norm_type"] = "rms_norm"

unet_config["num_heads"] = unet_config["embed_dim"] // rms_norm.shape[0]

if '{}timestep_features.weight'.format(key_prefix) in state_dict:

unet_config["timestep_features_type"] = "learned"

else:

unet_config["timestep_features_type"] = "expo"

io_channels = state_dict['{}postprocess_conv.weight'.format(key_prefix)].shape[0]

unet_config["io_channels"] = io_channels

unet_config["input_concat_dim"] = state_dict['{}transformer.project_in.weight'.format(key_prefix)].shape[1] - io_channels

local_add_cond = state_dict.get('{}transformer.layers.0.to_local_embed.0.weight'.format(key_prefix), None)

if local_add_cond is not None:

unet_config["local_add_cond_dim"] = local_add_cond.shape[1]

global_cond_embed = state_dict.get('{}transformer.global_cond_embedder.0.weight'.format(key_prefix), None)

if global_cond_embed is not None:

unet_config["global_cond_shared_embed"] = True

unet_config["global_cond_type"] = "adaLN"

unet_config["depth"] = count_blocks(state_dict_keys, '{}transformer.layers.'.format(key_prefix) + '{}.')

return unet_config

if '{}double_layers.0.attn.w1q.weight'.format(key_prefix) in state_dict_keys: #aura flow dit

unet_config = {}

unet_config["max_seq"] = state_dict['{}positional_encoding'.format(key_prefix)].shape[1]

unet_config["cond_seq_dim"] = state_dict['{}cond_seq_linear.weight'.format(key_prefix)].shape[1]

double_layers = count_blocks(state_dict_keys, '{}double_layers.'.format(key_prefix) + '{}.')

single_layers = count_blocks(state_dict_keys, '{}single_layers.'.format(key_prefix) + '{}.')

unet_config["n_double_layers"] = double_layers

unet_config["n_layers"] = double_layers + single_layers

return unet_config

if '{}mlp_t5.0.weight'.format(key_prefix) in state_dict_keys: #Hunyuan DiT

unet_config = {}

unet_config["image_model"] = "hydit"

unet_config["depth"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')

unet_config["hidden_size"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[0]

if unet_config["hidden_size"] == 1408 and unet_config["depth"] == 40: #DiT-g/2

unet_config["mlp_ratio"] = 4.3637

if state_dict['{}extra_embedder.0.weight'.format(key_prefix)].shape[1] == 3968:

unet_config["size_cond"] = True

unet_config["use_style_cond"] = True

unet_config["image_model"] = "hydit1"

return unet_config

if '{}txt_in.individual_token_refiner.blocks.0.norm1.weight'.format(key_prefix) in state_dict_keys: #Hunyuan Video

dit_config = {}

in_w = state_dict['{}img_in.proj.weight'.format(key_prefix)]

out_w = state_dict['{}final_layer.linear.weight'.format(key_prefix)]

dit_config["image_model"] = "hunyuan_video"

dit_config["in_channels"] = in_w.shape[1] #SkyReels img2video has 32 input channels

dit_config["patch_size"] = list(in_w.shape[2:])

dit_config["out_channels"] = out_w.shape[0] // math.prod(dit_config["patch_size"])

if any(s.startswith('{}vector_in.'.format(key_prefix)) for s in state_dict_keys):

dit_config["vec_in_dim"] = 768

else:

dit_config["vec_in_dim"] = None

if len(dit_config["patch_size"]) == 2:

dit_config["axes_dim"] = [64, 64]

else:

dit_config["axes_dim"] = [16, 56, 56]

if any(s.startswith('{}time_r_in.'.format(key_prefix)) for s in state_dict_keys):

dit_config["meanflow"] = True

else:

dit_config["meanflow"] = False

dit_config["context_in_dim"] = state_dict['{}txt_in.input_embedder.weight'.format(key_prefix)].shape[1]

dit_config["hidden_size"] = in_w.shape[0]

dit_config["mlp_ratio"] = 4.0

dit_config["num_heads"] = in_w.shape[0] // 128

dit_config["depth"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.')

dit_config["depth_single_blocks"] = count_blocks(state_dict_keys, '{}single_blocks.'.format(key_prefix) + '{}.')

dit_config["theta"] = 256

dit_config["qkv_bias"] = True

if '{}byt5_in.fc1.weight'.format(key_prefix) in state_dict:

dit_config["byt5"] = True

else:

dit_config["byt5"] = False

guidance_keys = list(filter(lambda a: a.startswith("{}guidance_in.".format(key_prefix)), state_dict_keys))

dit_config["guidance_embed"] = len(guidance_keys) > 0

# HunyuanVideo 1.5

if '{}cond_type_embedding.weight'.format(key_prefix) in state_dict_keys:

dit_config["use_cond_type_embedding"] = True

else:

dit_config["use_cond_type_embedding"] = False

if '{}vision_in.proj.0.weight'.format(key_prefix) in state_dict_keys:

dit_config["vision_in_dim"] = state_dict['{}vision_in.proj.0.weight'.format(key_prefix)].shape[0]

dit_config["meanflow_sum"] = True

else:

dit_config["vision_in_dim"] = None

dit_config["meanflow_sum"] = False

return dit_config

if any_suffix_in(state_dict_keys, key_prefix, 'double_blocks.0.img_attn.norm.key_norm.', ["weight", "scale"]) and ('{}img_in.weight'.format(key_prefix) in state_dict_keys or any_suffix_in(state_dict_keys, key_prefix, 'distilled_guidance_layer.norms.0.', ["weight", "scale"])): #Flux, Chroma or Chroma Radiance (has no img_in.weight)

dit_config = {}

if '{}double_stream_modulation_img.lin.weight'.format(key_prefix) in state_dict_keys:

dit_config["image_model"] = "flux2"

dit_config["axes_dim"] = [32, 32, 32, 32]

dit_config["num_heads"] = 48

dit_config["mlp_ratio"] = 3.0

dit_config["theta"] = 2000

dit_config["out_channels"] = 128

dit_config["global_modulation"] = True

dit_config["mlp_silu_act"] = True

dit_config["qkv_bias"] = False

dit_config["ops_bias"] = False

dit_config["default_ref_method"] = "index"

dit_config["ref_index_scale"] = 10.0

dit_config["txt_ids_dims"] = [3]

patch_size = 1

else:

dit_config["image_model"] = "flux"

dit_config["axes_dim"] = [16, 56, 56]

dit_config["num_heads"] = 24

dit_config["mlp_ratio"] = 4.0

dit_config["theta"] = 10000

dit_config["out_channels"] = 16

dit_config["qkv_bias"] = True

dit_config["txt_ids_dims"] = []

patch_size = 2

dit_config["in_channels"] = 16

dit_config["hidden_size"] = 3072

dit_config["context_in_dim"] = 4096

dit_config["patch_size"] = patch_size

in_key = "{}img_in.weight".format(key_prefix)

if in_key in state_dict_keys:

w = state_dict[in_key]

dit_config["in_channels"] = w.shape[1] // (patch_size * patch_size)

dit_config["hidden_size"] = w.shape[0]

txt_in_key = "{}txt_in.weight".format(key_prefix)

if txt_in_key in state_dict_keys:

w = state_dict[txt_in_key]

dit_config["context_in_dim"] = w.shape[1]

dit_config["hidden_size"] = w.shape[0]

vec_in_key = '{}vector_in.in_layer.weight'.format(key_prefix)

if vec_in_key in state_dict_keys:

dit_config["vec_in_dim"] = state_dict[vec_in_key].shape[1]

else:

dit_config["vec_in_dim"] = None

dit_config["num_heads"] = dit_config["hidden_size"] // sum(dit_config["axes_dim"])

dit_config["depth"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.')

dit_config["depth_single_blocks"] = count_blocks(state_dict_keys, '{}single_blocks.'.format(key_prefix) + '{}.')

if any_suffix_in(state_dict_keys, key_prefix, 'distilled_guidance_layer.0.norms.0.', ["weight", "scale"]) or any_suffix_in(state_dict_keys, key_prefix, 'distilled_guidance_layer.norms.0.', ["weight", "scale"]): #Chroma

dit_config["image_model"] = "chroma"

dit_config["in_channels"] = 64

dit_config["out_channels"] = 64

dit_config["in_dim"] = 64

dit_config["out_dim"] = 3072

dit_config["hidden_dim"] = 5120

dit_config["n_layers"] = 5

if any_suffix_in(state_dict_keys, key_prefix, 'nerf_blocks.0.norm.', ["weight", "scale"]): #Chroma Radiance

dit_config["image_model"] = "chroma_radiance"

dit_config["in_channels"] = 3

dit_config["out_channels"] = 3

dit_config["patch_size"] = state_dict.get('{}img_in_patch.weight'.format(key_prefix)).size(dim=-1)

dit_config["nerf_hidden_size"] = 64

dit_config["nerf_mlp_ratio"] = 4

dit_config["nerf_depth"] = 4

dit_config["nerf_max_freqs"] = 8

dit_config["nerf_tile_size"] = 512

dit_config["nerf_final_head_type"] = "conv" if any_suffix_in(state_dict_keys, key_prefix, 'nerf_final_layer_conv.norm.', ["weight", "scale"]) else "linear"

dit_config["nerf_embedder_dtype"] = torch.float32

if "{}__x0__".format(key_prefix) in state_dict_keys: # x0 pred

dit_config["use_x0"] = True

else:

dit_config["use_x0"] = False

if "{}__sequential__".format(key_prefix) in state_dict_keys: # sequential txt_ids

dit_config["use_sequential_txt_ids"] = True

else:

dit_config["use_sequential_txt_ids"] = False

else:

dit_config["guidance_embed"] = "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys

dit_config["yak_mlp"] = '{}double_blocks.0.img_mlp.gate_proj.weight'.format(key_prefix) in state_dict_keys

dit_config["txt_norm"] = any_suffix_in(state_dict_keys, key_prefix, 'txt_norm.', ["weight", "scale"])

if dit_config["yak_mlp"] and dit_config["txt_norm"]: # Ovis model

dit_config["txt_ids_dims"] = [1, 2]

if dit_config.get("context_in_dim") == 3584 and dit_config["vec_in_dim"] is None: # LongCat-Image

dit_config["txt_ids_dims"] = [1, 2]

return dit_config

if '{}t5_yproj.weight'.format(key_prefix) in state_dict_keys: #Genmo mochi preview

dit_config = {}

dit_config["image_model"] = "mochi_preview"

dit_config["depth"] = 48

dit_config["patch_size"] = 2

dit_config["num_heads"] = 24

dit_config["hidden_size_x"] = 3072

dit_config["hidden_size_y"] = 1536

dit_config["mlp_ratio_x"] = 4.0

dit_config["mlp_ratio_y"] = 4.0

dit_config["learn_sigma"] = False

dit_config["in_channels"] = 12

dit_config["qk_norm"] = True

dit_config["qkv_bias"] = False

dit_config["out_bias"] = True

dit_config["attn_drop"] = 0.0

dit_config["patch_embed_bias"] = True

dit_config["posenc_preserve_area"] = True

dit_config["timestep_mlp_bias"] = True

dit_config["attend_to_padding"] = False

dit_config["timestep_scale"] = 1000.0

dit_config["use_t5"] = True

dit_config["t5_feat_dim"] = 4096

dit_config["t5_token_length"] = 256

dit_config["rope_theta"] = 10000.0

return dit_config

if '{}adaln_single.emb.timestep_embedder.linear_1.bias'.format(key_prefix) in state_dict_keys and '{}pos_embed.proj.bias'.format(key_prefix) in state_dict_keys:

# PixArt diffusers

return None

if '{}video_patch_proj.weight'.format(key_prefix) in state_dict_keys and '{}audio_patch_proj.weight'.format(key_prefix) in state_dict_keys: # MiniMax H3

dit_config = {}

dit_config["image_model"] = "minimax_h3"

dit_config["num_layers"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')

dit_config["token_refiner_num_layers"] = count_blocks(state_dict_keys, '{}token_refiner.blocks.'.format(key_prefix) + '{}.')

dit_config["hidden_size"] = state_dict['{}video_patch_proj.weight'.format(key_prefix)].shape[0]

dit_config["latents_dim"] = state_dict['{}final_layer.video_out.weight'.format(key_prefix)].shape[0] // 4 # patch 1x2x2

dit_config["audio_latents_dim"] = state_dict['{}final_layer.audio_out.weight'.format(key_prefix)].shape[0]

dit_config["attention_head_dim"] = state_dict['{}blocks.0.attn.q_norm.weight'.format(key_prefix)].shape[0]

qkv = state_dict['{}blocks.0.attn.qkv_proj.weight'.format(key_prefix)]

dit_config["num_attention_heads"] = qkv.shape[0] // (3 * dit_config["attention_head_dim"])

dit_config["ffn_hidden_size"] = state_dict['{}blocks.0.mlp.fc1.weight'.format(key_prefix)].shape[0] // 2

dit_config["text_dim"] = state_dict['{}condition_proj.weight'.format(key_prefix)].shape[1]

table_key = '{}adaln_t_table'.format(key_prefix)

if table_key in state_dict_keys:

# adaln shipped over a precomputed curve basis: the adaln linears span a small shared basis of the time-embedding curve (no time embedder)

table = state_dict[table_key].shape # [grid, k]

dit_config["adaln_curve_grid"] = table[0]

dit_config["time_embed_dim"] = table[1]

else:

te = state_dict['{}time_embedder.proj_in.weight'.format(key_prefix)]

dit_config["timestep_input_dim"] = te.shape[1]

dit_config["time_embed_hidden_size"] = te.shape[0]

dit_config["time_embed_dim"] = state_dict['{}time_embedder.proj_out.weight'.format(key_prefix)].shape[0]

dit_config["rope_inv_freq_len"] = state_dict['{}rope.inv_freq'.format(key_prefix)].shape[0]

dit_config["gate_compress"] = '{}blocks.0.attn.to_gate_compress.weight'.format(key_prefix) in state_dict_keys # VSA-trained

if metadata is not None and "config" in metadata:

dit_config.update(json.loads(metadata["config"]).get("transformer", {}))

return dit_config

if '{}adaln_single.emb.timestep_embedder.linear_1.bias'.format(key_prefix) in state_dict_keys: #Lightricks ltxv

dit_config = {}

dit_config["image_model"] = "ltxav" if f'{key_prefix}audio_adaln_single.linear.weight' in state_dict_keys else "ltxv"

dit_config["num_layers"] = count_blocks(state_dict_keys, '{}transformer_blocks.'.format(key_prefix) + '{}.')

shape = state_dict['{}transformer_blocks.0.attn2.to_k.weight'.format(key_prefix)].shape

dit_config["attention_head_dim"] = shape[0] // 32

dit_config["cross_attention_dim"] = shape[1]

if metadata is not None and "config" in metadata:

dit_config.update(json.loads(metadata["config"]).get("transformer", {}))

dit_config["use_keyframes_abs_pos_embedding"] = '{}keyframes_abs_pos_embedding'.format(key_prefix) in state_dict_keys

return dit_config

if '{}genre_embedder.weight'.format(key_prefix) in state_dict_keys: #ACE-Step model

dit_config = {}

dit_config["audio_model"] = "ace"

dit_config["attention_head_dim"] = 128

dit_config["in_channels"] = 8

dit_config["inner_dim"] = 2560

dit_config["max_height"] = 16

dit_config["max_position"] = 32768

dit_config["max_width"] = 32768

dit_config["mlp_ratio"] = 2.5

dit_config["num_attention_heads"] = 20

dit_config["num_layers"] = 24

dit_config["out_channels"] = 8

dit_config["patch_size"] = [16, 1]

dit_config["rope_theta"] = 1000000.0

dit_config["speaker_embedding_dim"] = 512

dit_config["text_embedding_dim"] = 768

dit_config["ssl_encoder_depths"] = [8, 8]

dit_config["ssl_latent_dims"] = [1024, 768]

dit_config["ssl_names"] = ["mert", "m-hubert"]

dit_config["lyric_encoder_vocab_size"] = 6693

dit_config["lyric_hidden_size"] = 1024

return dit_config

if '{}t_block.1.weight'.format(key_prefix) in state_dict_keys: # PixArt

patch_size = 2

dit_config = {}

dit_config["num_heads"] = 16

dit_config["patch_size"] = patch_size

dit_config["hidden_size"] = 1152

dit_config["in_channels"] = 4

dit_config["depth"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')

y_key = "{}y_embedder.y_embedding".format(key_prefix)

if y_key in state_dict_keys:

dit_config["model_max_length"] = state_dict[y_key].shape[0]

pe_key = "{}pos_embed".format(key_prefix)

if pe_key in state_dict_keys:

dit_config["input_size"] = int(math.sqrt(state_dict[pe_key].shape[1])) * patch_size

dit_config["pe_interpolation"] = dit_config["input_size"] // (512//8) # guess

ar_key = "{}ar_embedder.mlp.0.weight".format(key_prefix)

if ar_key in state_dict_keys:

dit_config["image_model"] = "pixart_alpha"

dit_config["micro_condition"] = True

else:

dit_config["image_model"] = "pixart_sigma"

dit_config["micro_condition"] = False

return dit_config

if '{}blocks.block0.blocks.0.block.attn.to_q.0.weight'.format(key_prefix) in state_dict_keys: # Cosmos

dit_config = {}

dit_config["image_model"] = "cosmos"

dit_config["max_img_h"] = 240

dit_config["max_img_w"] = 240

dit_config["max_frames"] = 128

concat_padding_mask = True

dit_config["in_channels"] = (state_dict['{}x_embedder.proj.1.weight'.format(key_prefix)].shape[1] // 4) - int(concat_padding_mask)

dit_config["out_channels"] = 16

dit_config["patch_spatial"] = 2

dit_config["patch_temporal"] = 1

dit_config["model_channels"] = state_dict['{}blocks.block0.blocks.0.block.attn.to_q.0.weight'.format(key_prefix)].shape[0]

dit_config["block_config"] = "FA-CA-MLP"

dit_config["concat_padding_mask"] = concat_padding_mask

dit_config["pos_emb_cls"] = "rope3d"

dit_config["pos_emb_learnable"] = False

dit_config["pos_emb_interpolation"] = "crop"

dit_config["block_x_format"] = "THWBD"

dit_config["affline_emb_norm"] = True

dit_config["use_adaln_lora"] = True

dit_config["adaln_lora_dim"] = 256

if dit_config["model_channels"] == 4096:

# 7B

dit_config["num_blocks"] = 28

dit_config["num_heads"] = 32

dit_config["extra_per_block_abs_pos_emb"] = True

dit_config["rope_h_extrapolation_ratio"] = 1.0

dit_config["rope_w_extrapolation_ratio"] = 1.0

dit_config["rope_t_extrapolation_ratio"] = 2.0

dit_config["extra_per_block_abs_pos_emb_type"] = "learnable"

else: # 5120

# 14B

dit_config["num_blocks"] = 36

dit_config["num_heads"] = 40

dit_config["extra_per_block_abs_pos_emb"] = True

dit_config["rope_h_extrapolation_ratio"] = 2.0

dit_config["rope_w_extrapolation_ratio"] = 2.0

dit_config["rope_t_extrapolation_ratio"] = 2.0

dit_config["extra_h_extrapolation_ratio"] = 2.0

dit_config["extra_w_extrapolation_ratio"] = 2.0

dit_config["extra_t_extrapolation_ratio"] = 2.0

dit_config["extra_per_block_abs_pos_emb_type"] = "learnable"

return dit_config

# PiD (Pixel Diffusion Decoder). Must check BEFORE plain PixelDiT_T2I.

_lq_w_key = '{}lq_proj.latent_proj.0.weight'.format(key_prefix)

if _lq_w_key in state_dict_keys:

latent_proj_in_channels = int(state_dict[_lq_w_key].shape[1])

hidden_dim = int(state_dict[_lq_w_key].shape[0])

_gate_prefix = '{}lq_proj.gate_modules.'.format(key_prefix)

num_gates = len({k[len(_gate_prefix):].split('.')[0]

for k in state_dict_keys if k.startswith(_gate_prefix)})

pid_v1_5 = '{}lq_proj.pit_head.weight'.format(key_prefix) in state_dict_keys

dit_config = {"image_model": "pid",

"lq_hidden_dim": hidden_dim}

if num_gates > 0:

dit_config["lq_interval"] = (14 + num_gates - 1) // num_gates

if pid_v1_5:

pid_v1_5_variants = {

16: { # Flux and QwenImage

"lq_latent_channels": 16,

"latent_spatial_down_factor": 8,

"lq_latent_unpatchify_factor": 1,

},

32: { # Flux2 after 2x latent unpatchify

"lq_latent_channels": 128,

"latent_spatial_down_factor": 16,

"lq_latent_unpatchify_factor": 2,

},

}

variant = pid_v1_5_variants.get(latent_proj_in_channels)

if variant is None:

raise ValueError(f"Unsupported PiD v1.5 latent projection with {latent_proj_in_channels} input channels")

gate_weight = state_dict['{}lq_proj.gate_modules.0.content_proj.weight'.format(key_prefix)]

dit_config.update(variant)

dit_config.update({

"lq_conv_padding_mode": "replicate",

"lq_gate_per_token": gate_weight.shape[0] == 1,

"pit_lq_inject": True,

"rope_ref_h": 2048,

"rope_ref_w": 2048,

})

else:

dit_config.update({

"lq_latent_channels": latent_proj_in_channels,

"latent_spatial_down_factor": 16 if latent_proj_in_channels >= 64 else 8,

})

return dit_config

if '{}core.pixel_embedder.proj.weight'.format(key_prefix) in state_dict_keys: # PixelDiT T2I

return {"image_model": "pixeldit_t2i"}

if '{}cap_embedder.1.weight'.format(key_prefix) in state_dict_keys and '{}noise_refiner.0.attention.k_norm.weight'.format(key_prefix) in state_dict_keys: # Lumina 2

dit_config = {}

dit_config["image_model"] = "lumina2"

dit_config["patch_size"] = 2

dit_config["in_channels"] = 16

w = state_dict['{}cap_embedder.1.weight'.format(key_prefix)]

dit_config["dim"] = w.shape[0]

dit_config["cap_feat_dim"] = w.shape[1]

dit_config["n_layers"] = count_blocks(state_dict_keys, '{}layers.'.format(key_prefix) + '{}.')

dit_config["qk_norm"] = True

if dit_config["dim"] == 2304: # Original Lumina 2

dit_config["n_heads"] = 24

dit_config["n_kv_heads"] = 8

dit_config["axes_dims"] = [32, 32, 32]

dit_config["axes_lens"] = [300, 512, 512]

dit_config["rope_theta"] = 10000.0

dit_config["ffn_dim_multiplier"] = 4.0

ctd_weight = state_dict.get('{}clip_text_pooled_proj.0.weight'.format(key_prefix), None)

if ctd_weight is not None: # NewBie

dit_config["clip_text_dim"] = ctd_weight.shape[0]

# NewBie also sets axes_lens = [1024, 512, 512] but it's not used in ComfyUI

elif dit_config["dim"] == 3840: # Z image

dit_config["n_heads"] = 30

dit_config["n_kv_heads"] = 30

dit_config["axes_dims"] = [32, 48, 48]

dit_config["axes_lens"] = [1536, 512, 512]

dit_config["rope_theta"] = 256.0

dit_config["ffn_dim_multiplier"] = (8.0 / 3.0)

dit_config["z_image_modulation"] = True

dit_config["time_scale"] = 1000.0

try:

dit_config["allow_fp16"] = torch.std(state_dict['{}layers.{}.ffn_norm1.weight'.format(key_prefix, dit_config["n_layers"] - 2)], unbiased=False).item() < 0.42

except Exception:

pass

dec_cond_key = '{}dec_net.cond_embed.weight'.format(key_prefix)

if '{}cap_pad_token'.format(key_prefix) in state_dict_keys:

dit_config["pad_tokens_multiple"] = 32

# Ming-Image Design has no learned pad tokens; Layer needs a saved marker or repack metadata.

ming_metadata = metadata is not None and "config" in metadata and json.loads(metadata["config"]).get("transformer", {}).get("image_model") == "ming_image"

if '{}__ming_image__'.format(key_prefix) in state_dict_keys or ming_metadata or ("pad_tokens_multiple" not in dit_config and dec_cond_key not in state_dict_keys):

dit_config["image_model"] = "ming_image"

if "pad_tokens_multiple" not in dit_config:

dit_config["masked_pad_multiple"] = 32

sig_weight = state_dict.get('{}siglip_embedder.0.weight'.format(key_prefix), None)

if sig_weight is not None:

dit_config["siglip_feat_dim"] = sig_weight.shape[0]

if dec_cond_key in state_dict_keys: # pixel-space variant

dit_config["image_model"] = "zimage_pixel"

# patch_size and in_channels are derived from x_embedder:

# x_embedder: Linear(patch_size * patch_size * in_channels, dim)

# The decoder also receives the full flat patch, so decoder_in_channels = x_embedder input dim.

x_emb_in = state_dict['{}x_embedder.weight'.format(key_prefix)].shape[1]

dec_out = state_dict['{}dec_net.final_layer.linear.weight'.format(key_prefix)].shape[0]

# patch_size: infer from decoder final layer output matching x_embedder input

# in_channels: infer from dec_net input_embedder (in_features = dec_in_ch + max_freqs^2)

embedder_w = state_dict['{}dec_net.input_embedder.embedder.0.weight'.format(key_prefix)]

dec_in_ch = dec_out # decoder in == decoder out (same pixel space)

dit_config["patch_size"] = round((x_emb_in / 3) ** 0.5) # assume RGB (in_channels=3)

dit_config["in_channels"] = 3

dit_config["decoder_in_channels"] = dec_in_ch

dit_config["decoder_hidden_size"] = state_dict[dec_cond_key].shape[0]

dit_config["decoder_num_res_blocks"] = count_blocks(

state_dict_keys, '{}dec_net.res_blocks.'.format(key_prefix) + '{}.'

)

dit_config["decoder_max_freqs"] = int((embedder_w.shape[1] - dec_in_ch) ** 0.5)

if '{}__x0__'.format(key_prefix) in state_dict_keys:

dit_config["use_x0"] = True

return dit_config

if '{}blocks.0.norm1.linear.weight'.format(key_prefix) in state_dict_keys: # CogVideoX

dit_config = {}

dit_config["image_model"] = "cogvideox"

# Extract config from weight shapes

norm1_weight = state_dict['{}blocks.0.norm1.linear.weight'.format(key_prefix)]

time_embed_dim = norm1_weight.shape[1]

dim = norm1_weight.shape[0] // 6

dit_config["num_attention_heads"] = dim // 64

dit_config["attention_head_dim"] = 64

dit_config["time_embed_dim"] = time_embed_dim

dit_config["num_layers"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')

# Detect in_channels from patch_embed

patch_proj_key = '{}patch_embed.proj.weight'.format(key_prefix)

if patch_proj_key in state_dict_keys:

w = state_dict[patch_proj_key]

if w.ndim == 4:

# Conv2d: [out, in, kh, kw] — CogVideoX 1.0

dit_config["in_channels"] = w.shape[1]

dit_config["patch_size"] = w.shape[2]

elif w.ndim == 2:

# Linear: [out, in_channels * patch_size * patch_size * patch_size_t] — CogVideoX 1.5

dit_config["patch_size"] = 2

dit_config["patch_size_t"] = 2

dit_config["in_channels"] = w.shape[1] // (2 * 2 * 2) # 256 // 8 = 32

text_proj_key = '{}patch_embed.text_proj.weight'.format(key_prefix)

if text_proj_key in state_dict_keys:

dit_config["text_embed_dim"] = state_dict[text_proj_key].shape[1]

# Detect OFS embedding

ofs_key = '{}ofs_embedding_linear_1.weight'.format(key_prefix)

if ofs_key in state_dict_keys:

dit_config["ofs_embed_dim"] = state_dict[ofs_key].shape[1]

# Detect positional embedding type

pos_key = '{}patch_embed.pos_embedding'.format(key_prefix)

if pos_key in state_dict_keys:

dit_config["use_learned_positional_embeddings"] = True

dit_config["use_rotary_positional_embeddings"] = False

else:

dit_config["use_learned_positional_embeddings"] = False

dit_config["use_rotary_positional_embeddings"] = True

return dit_config

seedvr2_7b_separate_key = "{}blocks.35.mlp.vid.proj_out.weight".format(key_prefix)

if seedvr2_7b_separate_key in state_dict_keys and state_dict[seedvr2_7b_separate_key].shape[0] == 3072: # seedvr2 7b

dit_config = {}

dit_config["image_model"] = "seedvr2"

dit_config["vid_dim"] = 3072

dit_config["heads"] = 24

dit_config["num_layers"] = 36

# This checkpoint uses separate vid/txt MMModule keys in every block.

dit_config["mm_layers"] = 36

dit_config["norm_eps"] = 1e-5

dit_config["rope_type"] = "rope3d"

dit_config["rope_dim"] = 64

dit_config["mlp_type"] = "normal"

return dit_config

if "{}blocks.35.mlp.all.proj_in_gate.weight".format(key_prefix) in state_dict_keys: # seedvr2 7b

dit_config = {}

dit_config["image_model"] = "seedvr2"

dit_config["vid_dim"] = 3072

dit_config["heads"] = 24

dit_config["num_layers"] = 36

# This checkpoint uses shared all.* MMModule keys after the initial blocks.

dit_config["mm_layers"] = 10

dit_config["norm_eps"] = 1e-5

dit_config["rope_type"] = "rope3d"

dit_config["rope_dim"] = 64

dit_config["mlp_type"] = "swiglu"

return dit_config

if "{}blocks.31.mlp.all.proj_in_gate.weight".format(key_prefix) in state_dict_keys: # seedvr2 3b

dit_config = {}

dit_config["image_model"] = "seedvr2"

dit_config["vid_dim"] = 2560

dit_config["heads"] = 20

dit_config["num_layers"] = 32

dit_config["norm_eps"] = 1.0e-05

dit_config["mlp_type"] = "swiglu"

dit_config["vid_out_norm"] = True

return dit_config

if '{}head.modulation'.format(key_prefix) in state_dict_keys: # Wan 2.1

dit_config = {}

dit_config["image_model"] = "wan2.1"

dim = state_dict['{}head.modulation'.format(key_prefix)].shape[-1]

out_dim = state_dict['{}head.head.weight'.format(key_prefix)].shape[0] // 4

dit_config["dim"] = dim

dit_config["out_dim"] = out_dim

dit_config["num_heads"] = dim // 128

dit_config["ffn_dim"] = state_dict['{}blocks.0.ffn.0.weight'.format(key_prefix)].shape[0]

dit_config["num_layers"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')

dit_config["patch_size"] = (1, 2, 2)

dit_config["freq_dim"] = 256

dit_config["window_size"] = (-1, -1)

dit_config["qk_norm"] = True

dit_config["cross_attn_norm"] = True

dit_config["eps"] = 1e-6

dit_config["in_dim"] = state_dict['{}patch_embedding.weight'.format(key_prefix)].shape[1]

if '{}vace_patch_embedding.weight'.format(key_prefix) in state_dict_keys:

dit_config["model_type"] = "vace"

dit_config["vace_in_dim"] = state_dict['{}vace_patch_embedding.weight'.format(key_prefix)].shape[1]

dit_config["vace_layers"] = count_blocks(state_dict_keys, '{}vace_blocks.'.format(key_prefix) + '{}.')

if '{}img_emb.proj.0.bias'.format(key_prefix) in state_dict_keys: # ID-V2V, vace on an i2v model

dit_config["vace_image_input"] = True

elif '{}control_adapter.conv.weight'.format(key_prefix) in state_dict_keys:

if '{}img_emb.proj.0.bias'.format(key_prefix) in state_dict_keys:

dit_config["model_type"] = "camera"

else:

dit_config["model_type"] = "camera_2.2"

elif '{}casual_audio_encoder.encoder.final_linear.weight'.format(key_prefix) in state_dict_keys:

dit_config["model_type"] = "s2v"

elif '{}audio_proj.audio_proj_glob_1.layer.bias'.format(key_prefix) in state_dict_keys:

dit_config["model_type"] = "humo"

elif '{}face_adapter.fuser_blocks.0.k_norm.weight'.format(key_prefix) in state_dict_keys:

dit_config["model_type"] = "animate"

elif '{}patch_embedding_mask.weight'.format(key_prefix) in state_dict_keys:

dit_config["model_type"] = "scail2"

elif '{}patch_embedding_pose.weight'.format(key_prefix) in state_dict_keys:

dit_config["model_type"] = "scail"

elif '{}patch_embedding_global.weight'.format(key_prefix) in state_dict_keys:

dit_config["model_type"] = "wandancer"

else:

if '{}img_emb.proj.0.bias'.format(key_prefix) in state_dict_keys:

dit_config["model_type"] = "i2v"

else:

dit_config["model_type"] = "t2v"

flf_weight = state_dict.get('{}img_emb.emb_pos'.format(key_prefix))

if flf_weight is not None:

dit_config["flf_pos_embed_token_number"] = flf_weight.shape[1]

ref_conv_weight = state_dict.get('{}ref_conv.weight'.format(key_prefix))

if ref_conv_weight is not None:

dit_config["in_dim_ref_conv"] = ref_conv_weight.shape[1]

if metadata is not None and "config" in metadata:

dit_config.update(json.loads(metadata["config"]).get("transformer", {}))

return dit_config

if '{}latent_in.weight'.format(key_prefix) in state_dict_keys: # Hunyuan 3D

in_shape = state_dict['{}latent_in.weight'.format(key_prefix)].shape

dit_config = {}

dit_config["image_model"] = "hunyuan3d2"

dit_config["in_channels"] = in_shape[1]

dit_config["context_in_dim"] = state_dict['{}cond_in.weight'.format(key_prefix)].shape[1]

dit_config["hidden_size"] = in_shape[0]

dit_config["mlp_ratio"] = 4.0

dit_config["num_heads"] = 16

dit_config["depth"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.')

dit_config["depth_single_blocks"] = count_blocks(state_dict_keys, '{}single_blocks.'.format(key_prefix) + '{}.')

dit_config["qkv_bias"] = True

dit_config["guidance_embed"] = "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys

return dit_config

if f"{key_prefix}t_embedder.mlp.2.weight" in state_dict_keys and f"{key_prefix}blocks.0.attn1.k_norm.weight" in state_dict_keys: # Hunyuan 3D 2.1

dit_config = {}

dit_config["image_model"] = "hunyuan3d2_1"

dit_config["in_channels"] = state_dict[f"{key_prefix}x_embedder.weight"].shape[1]

dit_config["context_dim"] = 1024

dit_config["hidden_size"] = state_dict[f"{key_prefix}x_embedder.weight"].shape[0]

dit_config["mlp_ratio"] = 4.0

dit_config["num_heads"] = 16

dit_config["depth"] = count_blocks(state_dict_keys, f"{key_prefix}blocks.{{}}")

dit_config["qkv_bias"] = False

dit_config["guidance_cond_proj_dim"] = None#f"{key_prefix}t_embedder.cond_proj.weight" in state_dict_keys

return dit_config

if '{}cam_out_layer.weight'.format(key_prefix) in state_dict_keys and '{}repo_layers.0.final_map.weight'.format(key_prefix) in state_dict_keys: # TripoSplat

return {"image_model": "triposplat"}

if '{}t_embedder1.mlp.0.weight'.format(key_prefix) in state_dict_keys and '{}x_embedder.proj1.weight'.format(key_prefix) in state_dict_keys: # HiDream-O1

return {"image_model": "hidream_o1"}

vision_key = f"{key_prefix}fm_modules.vision_model_mot_gen.embeddings.patch_embedding.weight"

query_key = f"{key_prefix}language_model.model.layers.0.self_attn.q_proj_mot_gen.weight"

if (

vision_key in state_dict

and query_key in state_dict

and state_dict[vision_key].shape[0] == 1024

and state_dict[query_key].shape[0] == 4096

): # SenseNova U1.5

return {"image_model": "sensenova_u15"}

if '{}caption_projection.0.linear.weight'.format(key_prefix) in state_dict_keys: # HiDream

dit_config = {}

dit_config["image_model"] = "hidream"

dit_config["attention_head_dim"] = 128

dit_config["axes_dims_rope"] = [64, 32, 32]

dit_config["caption_channels"] = [4096, 4096]

dit_config["max_resolution"] = [128, 128]

dit_config["in_channels"] = 16

dit_config["llama_layers"] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31, 31]

dit_config["num_attention_heads"] = 20

dit_config["num_routed_experts"] = 4

dit_config["num_activated_experts"] = 2

dit_config["num_layers"] = 16

dit_config["num_single_layers"] = 32

dit_config["out_channels"] = 16

dit_config["patch_size"] = 2

dit_config["text_emb_dim"] = 2048

return dit_config

if '{}blocks.0.mlp.layer1.weight'.format(key_prefix) in state_dict_keys: # Cosmos predict2

dit_config = {}

dit_config["image_model"] = "cosmos_predict2"

if "{}llm_adapter.blocks.0.cross_attn.q_proj.weight".format(key_prefix) in state_dict_keys:

dit_config["image_model"] = "anima"

dit_config["max_img_h"] = 240

dit_config["max_img_w"] = 240

dit_config["max_frames"] = 128

concat_padding_mask = True

dit_config["in_channels"] = (state_dict['{}x_embedder.proj.1.weight'.format(key_prefix)].shape[1] // 4) - int(concat_padding_mask)

dit_config["out_channels"] = 16

dit_config["patch_spatial"] = 2

dit_config["patch_temporal"] = 1

dit_config["model_channels"] = state_dict['{}x_embedder.proj.1.weight'.format(key_prefix)].shape[0]

dit_config["concat_padding_mask"] = concat_padding_mask

dit_config["crossattn_emb_channels"] = 1024

dit_config["pos_emb_cls"] = "rope3d"

dit_config["pos_emb_learnable"] = True