fred-dev/comfy_ui_ali
0
1import json2import comfy.supported_models3import comfy.supported_models_base4import comfy.utils5import math6import logging7import torch8 9def count_blocks(state_dict_keys, prefix_string):10 count = 011 while True:12 c = False13 for k in state_dict_keys:14 if k.startswith(prefix_string.format(count)):15 c = True16 break17 if c == False:18 break19 count += 120 return count21 22def calculate_transformer_depth(prefix, state_dict_keys, state_dict):23 context_dim = None24 use_linear_in_transformer = False25 26 transformer_prefix = prefix + "1.transformer_blocks."27 transformer_keys = sorted(list(filter(lambda a: a.startswith(transformer_prefix), state_dict_keys)))28 if len(transformer_keys) > 0:29 last_transformer_depth = count_blocks(state_dict_keys, transformer_prefix + '{}')30 context_dim = state_dict['{}0.attn2.to_k.weight'.format(transformer_prefix)].shape[1]31 use_linear_in_transformer = len(state_dict['{}1.proj_in.weight'.format(prefix)].shape) == 232 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_dict33 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_dict34 return last_transformer_depth, context_dim, use_linear_in_transformer, time_stack, time_stack_cross35 return None36 37def detect_unet_config(state_dict, key_prefix, metadata=None):38 state_dict_keys = list(state_dict.keys())39 40 if '{}joint_blocks.0.context_block.attn.qkv.weight'.format(key_prefix) in state_dict_keys: #mmdit model41 unet_config = {}42 unet_config["in_channels"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[1]43 patch_size = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[2]44 unet_config["patch_size"] = patch_size45 final_layer = '{}final_layer.linear.weight'.format(key_prefix)46 if final_layer in state_dict:47 unet_config["out_channels"] = state_dict[final_layer].shape[0] // (patch_size * patch_size)48 49 unet_config["depth"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[0] // 6450 unet_config["input_size"] = None51 y_key = '{}y_embedder.mlp.0.weight'.format(key_prefix)52 if y_key in state_dict_keys:53 unet_config["adm_in_channels"] = state_dict[y_key].shape[1]54 55 context_key = '{}context_embedder.weight'.format(key_prefix)56 if context_key in state_dict_keys:57 in_features = state_dict[context_key].shape[1]58 out_features = state_dict[context_key].shape[0]59 unet_config["context_embedder_config"] = {"target": "torch.nn.Linear", "params": {"in_features": in_features, "out_features": out_features}}60 num_patches_key = '{}pos_embed'.format(key_prefix)61 if num_patches_key in state_dict_keys:62 num_patches = state_dict[num_patches_key].shape[1]63 unet_config["num_patches"] = num_patches64 unet_config["pos_embed_max_size"] = round(math.sqrt(num_patches))65 66 rms_qk = '{}joint_blocks.0.context_block.attn.ln_q.weight'.format(key_prefix)67 if rms_qk in state_dict_keys:68 unet_config["qk_norm"] = "rms"69 70 unet_config["pos_embed_scaling_factor"] = None #unused for inference71 context_processor = '{}context_processor.layers.0.attn.qkv.weight'.format(key_prefix)72 if context_processor in state_dict_keys:73 unet_config["context_processor_layers"] = count_blocks(state_dict_keys, '{}context_processor.layers.'.format(key_prefix) + '{}.')74 unet_config["x_block_self_attn_layers"] = []75 for key in state_dict_keys:76 if key.startswith('{}joint_blocks.'.format(key_prefix)) and key.endswith('.x_block.attn2.qkv.weight'):77 layer = key[len('{}joint_blocks.'.format(key_prefix)):-len('.x_block.attn2.qkv.weight')]78 unet_config["x_block_self_attn_layers"].append(int(layer))79 return unet_config80 81 if '{}clf.1.weight'.format(key_prefix) in state_dict_keys: #stable cascade82 unet_config = {}83 text_mapper_name = '{}clip_txt_mapper.weight'.format(key_prefix)84 if text_mapper_name in state_dict_keys:85 unet_config['stable_cascade_stage'] = 'c'86 w = state_dict[text_mapper_name]87 if w.shape[0] == 1536: #stage c lite88 unet_config['c_cond'] = 153689 unet_config['c_hidden'] = [1536, 1536]90 unet_config['nhead'] = [24, 24]91 unet_config['blocks'] = [[4, 12], [12, 4]]92 elif w.shape[0] == 2048: #stage c full93 unet_config['c_cond'] = 204894 elif '{}clip_mapper.weight'.format(key_prefix) in state_dict_keys:95 unet_config['stable_cascade_stage'] = 'b'96 w = state_dict['{}down_blocks.1.0.channelwise.0.weight'.format(key_prefix)]97 if w.shape[-1] == 640:98 unet_config['c_hidden'] = [320, 640, 1280, 1280]99 unet_config['nhead'] = [-1, -1, 20, 20]100 unet_config['blocks'] = [[2, 6, 28, 6], [6, 28, 6, 2]]101 unet_config['block_repeat'] = [[1, 1, 1, 1], [3, 3, 2, 2]]102 elif w.shape[-1] == 576: #stage b lite103 unet_config['c_hidden'] = [320, 576, 1152, 1152]104 unet_config['nhead'] = [-1, 9, 18, 18]105 unet_config['blocks'] = [[2, 4, 14, 4], [4, 14, 4, 2]]106 unet_config['block_repeat'] = [[1, 1, 1, 1], [2, 2, 2, 2]]107 return unet_config108 109 if '{}transformer.rotary_pos_emb.inv_freq'.format(key_prefix) in state_dict_keys: #stable audio dit110 unet_config = {}111 unet_config["audio_model"] = "dit1.0"112 return unet_config113 114 if '{}double_layers.0.attn.w1q.weight'.format(key_prefix) in state_dict_keys: #aura flow dit115 unet_config = {}116 unet_config["max_seq"] = state_dict['{}positional_encoding'.format(key_prefix)].shape[1]117 unet_config["cond_seq_dim"] = state_dict['{}cond_seq_linear.weight'.format(key_prefix)].shape[1]118 double_layers = count_blocks(state_dict_keys, '{}double_layers.'.format(key_prefix) + '{}.')119 single_layers = count_blocks(state_dict_keys, '{}single_layers.'.format(key_prefix) + '{}.')120 unet_config["n_double_layers"] = double_layers121 unet_config["n_layers"] = double_layers + single_layers122 return unet_config123 124 if '{}mlp_t5.0.weight'.format(key_prefix) in state_dict_keys: #Hunyuan DiT125 unet_config = {}126 unet_config["image_model"] = "hydit"127 unet_config["depth"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')128 unet_config["hidden_size"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[0]129 if unet_config["hidden_size"] == 1408 and unet_config["depth"] == 40: #DiT-g/2130 unet_config["mlp_ratio"] = 4.3637131 if state_dict['{}extra_embedder.0.weight'.format(key_prefix)].shape[1] == 3968:132 unet_config["size_cond"] = True133 unet_config["use_style_cond"] = True134 unet_config["image_model"] = "hydit1"135 return unet_config136 137 if '{}txt_in.individual_token_refiner.blocks.0.norm1.weight'.format(key_prefix) in state_dict_keys: #Hunyuan Video138 dit_config = {}139 dit_config["image_model"] = "hunyuan_video"140 dit_config["in_channels"] = state_dict['{}img_in.proj.weight'.format(key_prefix)].shape[1] #SkyReels img2video has 32 input channels141 dit_config["patch_size"] = [1, 2, 2]142 dit_config["out_channels"] = 16143 dit_config["vec_in_dim"] = 768144 dit_config["context_in_dim"] = 4096145 dit_config["hidden_size"] = 3072146 dit_config["mlp_ratio"] = 4.0147 dit_config["num_heads"] = 24148 dit_config["depth"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.')149 dit_config["depth_single_blocks"] = count_blocks(state_dict_keys, '{}single_blocks.'.format(key_prefix) + '{}.')150 dit_config["axes_dim"] = [16, 56, 56]151 dit_config["theta"] = 256152 dit_config["qkv_bias"] = True153 guidance_keys = list(filter(lambda a: a.startswith("{}guidance_in.".format(key_prefix)), state_dict_keys))154 dit_config["guidance_embed"] = len(guidance_keys) > 0155 return dit_config156 157 if '{}double_blocks.0.img_attn.norm.key_norm.scale'.format(key_prefix) in state_dict_keys: #Flux158 dit_config = {}159 dit_config["image_model"] = "flux"160 dit_config["in_channels"] = 16161 patch_size = 2162 dit_config["patch_size"] = patch_size163 in_key = "{}img_in.weight".format(key_prefix)164 if in_key in state_dict_keys:165 dit_config["in_channels"] = state_dict[in_key].shape[1] // (patch_size * patch_size)166 dit_config["out_channels"] = 16167 dit_config["vec_in_dim"] = 768168 dit_config["context_in_dim"] = 4096169 dit_config["hidden_size"] = 3072170 dit_config["mlp_ratio"] = 4.0171 dit_config["num_heads"] = 24172 dit_config["depth"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.')173 dit_config["depth_single_blocks"] = count_blocks(state_dict_keys, '{}single_blocks.'.format(key_prefix) + '{}.')174 dit_config["axes_dim"] = [16, 56, 56]175 dit_config["theta"] = 10000176 dit_config["qkv_bias"] = True177 dit_config["guidance_embed"] = "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys178 return dit_config179 180 if '{}t5_yproj.weight'.format(key_prefix) in state_dict_keys: #Genmo mochi preview181 dit_config = {}182 dit_config["image_model"] = "mochi_preview"183 dit_config["depth"] = 48184 dit_config["patch_size"] = 2185 dit_config["num_heads"] = 24186 dit_config["hidden_size_x"] = 3072187 dit_config["hidden_size_y"] = 1536188 dit_config["mlp_ratio_x"] = 4.0189 dit_config["mlp_ratio_y"] = 4.0190 dit_config["learn_sigma"] = False191 dit_config["in_channels"] = 12192 dit_config["qk_norm"] = True193 dit_config["qkv_bias"] = False194 dit_config["out_bias"] = True195 dit_config["attn_drop"] = 0.0196 dit_config["patch_embed_bias"] = True197 dit_config["posenc_preserve_area"] = True198 dit_config["timestep_mlp_bias"] = True199 dit_config["attend_to_padding"] = False200 dit_config["timestep_scale"] = 1000.0201 dit_config["use_t5"] = True202 dit_config["t5_feat_dim"] = 4096203 dit_config["t5_token_length"] = 256204 dit_config["rope_theta"] = 10000.0205 return dit_config206 207 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:208 # PixArt diffusers209 return None210 211 if '{}adaln_single.emb.timestep_embedder.linear_1.bias'.format(key_prefix) in state_dict_keys: #Lightricks ltxv212 dit_config = {}213 dit_config["image_model"] = "ltxv"214 if metadata is not None and "config" in metadata:215 dit_config.update(json.loads(metadata["config"]).get("transformer", {}))216 return dit_config217 218 if '{}t_block.1.weight'.format(key_prefix) in state_dict_keys: # PixArt219 patch_size = 2220 dit_config = {}221 dit_config["num_heads"] = 16222 dit_config["patch_size"] = patch_size223 dit_config["hidden_size"] = 1152224 dit_config["in_channels"] = 4225 dit_config["depth"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')226 227 y_key = "{}y_embedder.y_embedding".format(key_prefix)228 if y_key in state_dict_keys:229 dit_config["model_max_length"] = state_dict[y_key].shape[0]230 231 pe_key = "{}pos_embed".format(key_prefix)232 if pe_key in state_dict_keys:233 dit_config["input_size"] = int(math.sqrt(state_dict[pe_key].shape[1])) * patch_size234 dit_config["pe_interpolation"] = dit_config["input_size"] // (512//8) # guess235 236 ar_key = "{}ar_embedder.mlp.0.weight".format(key_prefix)237 if ar_key in state_dict_keys:238 dit_config["image_model"] = "pixart_alpha"239 dit_config["micro_condition"] = True240 else:241 dit_config["image_model"] = "pixart_sigma"242 dit_config["micro_condition"] = False243 return dit_config244 245 if '{}blocks.block0.blocks.0.block.attn.to_q.0.weight'.format(key_prefix) in state_dict_keys: # Cosmos246 dit_config = {}247 dit_config["image_model"] = "cosmos"248 dit_config["max_img_h"] = 240249 dit_config["max_img_w"] = 240250 dit_config["max_frames"] = 128251 concat_padding_mask = True252 dit_config["in_channels"] = (state_dict['{}x_embedder.proj.1.weight'.format(key_prefix)].shape[1] // 4) - int(concat_padding_mask)253 dit_config["out_channels"] = 16254 dit_config["patch_spatial"] = 2255 dit_config["patch_temporal"] = 1256 dit_config["model_channels"] = state_dict['{}blocks.block0.blocks.0.block.attn.to_q.0.weight'.format(key_prefix)].shape[0]257 dit_config["block_config"] = "FA-CA-MLP"258 dit_config["concat_padding_mask"] = concat_padding_mask259 dit_config["pos_emb_cls"] = "rope3d"260 dit_config["pos_emb_learnable"] = False261 dit_config["pos_emb_interpolation"] = "crop"262 dit_config["block_x_format"] = "THWBD"263 dit_config["affline_emb_norm"] = True264 dit_config["use_adaln_lora"] = True265 dit_config["adaln_lora_dim"] = 256266 267 if dit_config["model_channels"] == 4096:268 # 7B269 dit_config["num_blocks"] = 28270 dit_config["num_heads"] = 32271 dit_config["extra_per_block_abs_pos_emb"] = True272 dit_config["rope_h_extrapolation_ratio"] = 1.0273 dit_config["rope_w_extrapolation_ratio"] = 1.0274 dit_config["rope_t_extrapolation_ratio"] = 2.0275 dit_config["extra_per_block_abs_pos_emb_type"] = "learnable"276 else: # 5120277 # 14B278 dit_config["num_blocks"] = 36279 dit_config["num_heads"] = 40280 dit_config["extra_per_block_abs_pos_emb"] = True281 dit_config["rope_h_extrapolation_ratio"] = 2.0282 dit_config["rope_w_extrapolation_ratio"] = 2.0283 dit_config["rope_t_extrapolation_ratio"] = 2.0284 dit_config["extra_h_extrapolation_ratio"] = 2.0285 dit_config["extra_w_extrapolation_ratio"] = 2.0286 dit_config["extra_t_extrapolation_ratio"] = 2.0287 dit_config["extra_per_block_abs_pos_emb_type"] = "learnable"288 return dit_config289 290 if '{}cap_embedder.1.weight'.format(key_prefix) in state_dict_keys: # Lumina 2291 dit_config = {}292 dit_config["image_model"] = "lumina2"293 dit_config["patch_size"] = 2294 dit_config["in_channels"] = 16295 dit_config["dim"] = 2304296 dit_config["cap_feat_dim"] = 2304297 dit_config["n_layers"] = 26298 dit_config["n_heads"] = 24299 dit_config["n_kv_heads"] = 8300 dit_config["qk_norm"] = True301 dit_config["axes_dims"] = [32, 32, 32]302 dit_config["axes_lens"] = [300, 512, 512]303 return dit_config304 305 if '{}head.modulation'.format(key_prefix) in state_dict_keys: # Wan 2.1306 dit_config = {}307 dit_config["image_model"] = "wan2.1"308 dim = state_dict['{}head.modulation'.format(key_prefix)].shape[-1]309 dit_config["dim"] = dim310 dit_config["num_heads"] = dim // 128311 dit_config["ffn_dim"] = state_dict['{}blocks.0.ffn.0.weight'.format(key_prefix)].shape[0]312 dit_config["num_layers"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')313 dit_config["patch_size"] = (1, 2, 2)314 dit_config["freq_dim"] = 256315 dit_config["window_size"] = (-1, -1)316 dit_config["qk_norm"] = True317 dit_config["cross_attn_norm"] = True318 dit_config["eps"] = 1e-6319 dit_config["in_dim"] = state_dict['{}patch_embedding.weight'.format(key_prefix)].shape[1]320 if '{}img_emb.proj.0.bias'.format(key_prefix) in state_dict_keys:321 dit_config["model_type"] = "i2v"322 else:323 dit_config["model_type"] = "t2v"324 return dit_config325 326 if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:327 return None328 329 unet_config = {330 "use_checkpoint": False,331 "image_size": 32,332 "use_spatial_transformer": True,333 "legacy": False334 }335 336 y_input = '{}label_emb.0.0.weight'.format(key_prefix)337 if y_input in state_dict_keys:338 unet_config["num_classes"] = "sequential"339 unet_config["adm_in_channels"] = state_dict[y_input].shape[1]340 else:341 unet_config["adm_in_channels"] = None342 343 model_channels = state_dict['{}input_blocks.0.0.weight'.format(key_prefix)].shape[0]344 in_channels = state_dict['{}input_blocks.0.0.weight'.format(key_prefix)].shape[1]345 346 out_key = '{}out.2.weight'.format(key_prefix)347 if out_key in state_dict:348 out_channels = state_dict[out_key].shape[0]349 else:350 out_channels = 4351 352 num_res_blocks = []353 channel_mult = []354 transformer_depth = []355 transformer_depth_output = []356 context_dim = None357 use_linear_in_transformer = False358 359 video_model = False360 video_model_cross = False361 362 current_res = 1363 count = 0364 365 last_res_blocks = 0366 last_channel_mult = 0367 368 input_block_count = count_blocks(state_dict_keys, '{}input_blocks'.format(key_prefix) + '.{}.')369 for count in range(input_block_count):370 prefix = '{}input_blocks.{}.'.format(key_prefix, count)371 prefix_output = '{}output_blocks.{}.'.format(key_prefix, input_block_count - count - 1)372 373 block_keys = sorted(list(filter(lambda a: a.startswith(prefix), state_dict_keys)))374 if len(block_keys) == 0:375 break376 377 block_keys_output = sorted(list(filter(lambda a: a.startswith(prefix_output), state_dict_keys)))378 379 if "{}0.op.weight".format(prefix) in block_keys: #new layer380 num_res_blocks.append(last_res_blocks)381 channel_mult.append(last_channel_mult)382 383 current_res *= 2384 last_res_blocks = 0385 last_channel_mult = 0386 out = calculate_transformer_depth(prefix_output, state_dict_keys, state_dict)387 if out is not None:388 transformer_depth_output.append(out[0])389 else:390 transformer_depth_output.append(0)391 else:392 res_block_prefix = "{}0.in_layers.0.weight".format(prefix)393 if res_block_prefix in block_keys:394 last_res_blocks += 1395 last_channel_mult = state_dict["{}0.out_layers.3.weight".format(prefix)].shape[0] // model_channels396 397 out = calculate_transformer_depth(prefix, state_dict_keys, state_dict)398 if out is not None:399 transformer_depth.append(out[0])400 if context_dim is None:401 context_dim = out[1]402 use_linear_in_transformer = out[2]403 video_model = out[3]404 video_model_cross = out[4]405 else:406 transformer_depth.append(0)407 408 res_block_prefix = "{}0.in_layers.0.weight".format(prefix_output)409 if res_block_prefix in block_keys_output:410 out = calculate_transformer_depth(prefix_output, state_dict_keys, state_dict)411 if out is not None:412 transformer_depth_output.append(out[0])413 else:414 transformer_depth_output.append(0)415 416 417 num_res_blocks.append(last_res_blocks)418 channel_mult.append(last_channel_mult)419 if "{}middle_block.1.proj_in.weight".format(key_prefix) in state_dict_keys:420 transformer_depth_middle = count_blocks(state_dict_keys, '{}middle_block.1.transformer_blocks.'.format(key_prefix) + '{}')421 elif "{}middle_block.0.in_layers.0.weight".format(key_prefix) in state_dict_keys:422 transformer_depth_middle = -1423 else:424 transformer_depth_middle = -2425 426 unet_config["in_channels"] = in_channels427 unet_config["out_channels"] = out_channels428 unet_config["model_channels"] = model_channels429 unet_config["num_res_blocks"] = num_res_blocks430 unet_config["transformer_depth"] = transformer_depth431 unet_config["transformer_depth_output"] = transformer_depth_output432 unet_config["channel_mult"] = channel_mult433 unet_config["transformer_depth_middle"] = transformer_depth_middle434 unet_config['use_linear_in_transformer'] = use_linear_in_transformer435 unet_config["context_dim"] = context_dim436 437 if video_model:438 unet_config["extra_ff_mix_layer"] = True439 unet_config["use_spatial_context"] = True440 unet_config["merge_strategy"] = "learned_with_images"441 unet_config["merge_factor"] = 0.0442 unet_config["video_kernel_size"] = [3, 1, 1]443 unet_config["use_temporal_resblock"] = True444 unet_config["use_temporal_attention"] = True445 unet_config["disable_temporal_crossattention"] = not video_model_cross446 else:447 unet_config["use_temporal_resblock"] = False448 unet_config["use_temporal_attention"] = False449 450 return unet_config451 452def model_config_from_unet_config(unet_config, state_dict=None):453 for model_config in comfy.supported_models.models:454 if model_config.matches(unet_config, state_dict):455 return model_config(unet_config)456 457 logging.error("no match {}".format(unet_config))458 return None459 460def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=False, metadata=None):461 unet_config = detect_unet_config(state_dict, unet_key_prefix, metadata=metadata)462 if unet_config is None:463 return None464 model_config = model_config_from_unet_config(unet_config, state_dict)465 if model_config is None and use_base_if_no_match:466 model_config = comfy.supported_models_base.BASE(unet_config)467 468 scaled_fp8_key = "{}scaled_fp8".format(unet_key_prefix)469 if scaled_fp8_key in state_dict:470 scaled_fp8_weight = state_dict.pop(scaled_fp8_key)471 model_config.scaled_fp8 = scaled_fp8_weight.dtype472 if model_config.scaled_fp8 == torch.float32:473 model_config.scaled_fp8 = torch.float8_e4m3fn474 if scaled_fp8_weight.nelement() == 2:475 model_config.optimizations["fp8"] = False476 else:477 model_config.optimizations["fp8"] = True478 479 return model_config480 481def unet_prefix_from_state_dict(state_dict):482 candidates = ["model.diffusion_model.", #ldm/sgm models483 "model.model.", #audio models484 "net.", #cosmos485 ]486 counts = {k: 0 for k in candidates}487 for k in state_dict:488 for c in candidates:489 if k.startswith(c):490 counts[c] += 1491 break492 493 top = max(counts, key=counts.get)494 if counts[top] > 5:495 return top496 else:497 return "model." #aura flow and others498 499 500def convert_config(unet_config):501 new_config = unet_config.copy()502 num_res_blocks = new_config.get("num_res_blocks", None)503 channel_mult = new_config.get("channel_mult", None)504 505 if isinstance(num_res_blocks, int):506 num_res_blocks = len(channel_mult) * [num_res_blocks]507 508 if "attention_resolutions" in new_config:509 attention_resolutions = new_config.pop("attention_resolutions")510 transformer_depth = new_config.get("transformer_depth", None)511 transformer_depth_middle = new_config.get("transformer_depth_middle", None)512 513 if isinstance(transformer_depth, int):514 transformer_depth = len(channel_mult) * [transformer_depth]515 if transformer_depth_middle is None:516 transformer_depth_middle = transformer_depth[-1]517 t_in = []518 t_out = []519 s = 1520 for i in range(len(num_res_blocks)):521 res = num_res_blocks[i]522 d = 0523 if s in attention_resolutions:524 d = transformer_depth[i]525 526 t_in += [d] * res527 t_out += [d] * (res + 1)528 s *= 2529 transformer_depth = t_in530 new_config["transformer_depth"] = t_in531 new_config["transformer_depth_output"] = t_out532 new_config["transformer_depth_middle"] = transformer_depth_middle533 534 new_config["num_res_blocks"] = num_res_blocks535 return new_config536 537 538def unet_config_from_diffusers_unet(state_dict, dtype=None):539 match = {}540 transformer_depth = []541 542 attn_res = 1543 down_blocks = count_blocks(state_dict, "down_blocks.{}")544 for i in range(down_blocks):545 attn_blocks = count_blocks(state_dict, "down_blocks.{}.attentions.".format(i) + '{}')546 res_blocks = count_blocks(state_dict, "down_blocks.{}.resnets.".format(i) + '{}')547 for ab in range(attn_blocks):548 transformer_count = count_blocks(state_dict, "down_blocks.{}.attentions.{}.transformer_blocks.".format(i, ab) + '{}')549 transformer_depth.append(transformer_count)550 if transformer_count > 0:551 match["context_dim"] = state_dict["down_blocks.{}.attentions.{}.transformer_blocks.0.attn2.to_k.weight".format(i, ab)].shape[1]552 553 attn_res *= 2554 if attn_blocks == 0:555 for i in range(res_blocks):556 transformer_depth.append(0)557 558 match["transformer_depth"] = transformer_depth559 560 match["model_channels"] = state_dict["conv_in.weight"].shape[0]561 match["in_channels"] = state_dict["conv_in.weight"].shape[1]562 match["adm_in_channels"] = None563 if "class_embedding.linear_1.weight" in state_dict:564 match["adm_in_channels"] = state_dict["class_embedding.linear_1.weight"].shape[1]565 elif "add_embedding.linear_1.weight" in state_dict:566 match["adm_in_channels"] = state_dict["add_embedding.linear_1.weight"].shape[1]567 568 SDXL = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,569 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,570 'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 10,571 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],572 'use_temporal_attention': False, 'use_temporal_resblock': False}573 574 SDXL_refiner = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,575 'num_classes': 'sequential', 'adm_in_channels': 2560, 'dtype': dtype, 'in_channels': 4, 'model_channels': 384,576 'num_res_blocks': [2, 2, 2, 2], 'transformer_depth': [0, 0, 4, 4, 4, 4, 0, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 4,577 'use_linear_in_transformer': True, 'context_dim': 1280, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 4, 4, 4, 4, 4, 4, 0, 0, 0],578 'use_temporal_attention': False, 'use_temporal_resblock': False}579 580 SD21 = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,581 'adm_in_channels': None, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': [2, 2, 2, 2],582 'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1, 'use_linear_in_transformer': True,583 'context_dim': 1024, 'num_head_channels': 64, 'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],584 'use_temporal_attention': False, 'use_temporal_resblock': False}585 586 SD21_uncliph = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,587 'num_classes': 'sequential', 'adm_in_channels': 2048, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,588 'num_res_blocks': [2, 2, 2, 2], 'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1,589 'use_linear_in_transformer': True, 'context_dim': 1024, 'num_head_channels': 64, 'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],590 'use_temporal_attention': False, 'use_temporal_resblock': False}591 592 SD21_unclipl = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,593 'num_classes': 'sequential', 'adm_in_channels': 1536, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,594 'num_res_blocks': [2, 2, 2, 2], 'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1,595 'use_linear_in_transformer': True, 'context_dim': 1024, 'num_head_channels': 64, 'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],596 'use_temporal_attention': False, 'use_temporal_resblock': False}597 598 SD15 = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, 'adm_in_channels': None,599 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': [2, 2, 2, 2], 'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0],600 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1, 'use_linear_in_transformer': False, 'context_dim': 768, 'num_heads': 8,601 'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],602 'use_temporal_attention': False, 'use_temporal_resblock': False}603 604 SDXL_mid_cnet = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,605 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,606 'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 0, 0, 1, 1], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 1,607 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 0, 0, 0, 1, 1, 1],608 'use_temporal_attention': False, 'use_temporal_resblock': False}609 610 SDXL_small_cnet = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,611 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,612 'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 0, 0, 0, 0], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 0,613 'use_linear_in_transformer': True, 'num_head_channels': 64, 'context_dim': 1, 'transformer_depth_output': [0, 0, 0, 0, 0, 0, 0, 0, 0],614 'use_temporal_attention': False, 'use_temporal_resblock': False}615 616 SDXL_diffusers_inpaint = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,617 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 9, 'model_channels': 320,618 'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 10,619 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],620 'use_temporal_attention': False, 'use_temporal_resblock': False}621 622 SDXL_diffusers_ip2p = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,623 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 8, 'model_channels': 320,624 'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 10,625 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],626 'use_temporal_attention': False, 'use_temporal_resblock': False}627 628 SSD_1B = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,629 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,630 'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 4, 4], 'transformer_depth_output': [0, 0, 0, 1, 1, 2, 10, 4, 4],631 'channel_mult': [1, 2, 4], 'transformer_depth_middle': -1, 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,632 'use_temporal_attention': False, 'use_temporal_resblock': False}633 634 Segmind_Vega = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,635 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,636 'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 1, 1, 2, 2], 'transformer_depth_output': [0, 0, 0, 1, 1, 1, 2, 2, 2],637 'channel_mult': [1, 2, 4], 'transformer_depth_middle': -1, 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,638 'use_temporal_attention': False, 'use_temporal_resblock': False}639 640 KOALA_700M = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,641 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,642 'num_res_blocks': [1, 1, 1], 'transformer_depth': [0, 2, 5], 'transformer_depth_output': [0, 0, 2, 2, 5, 5],643 'channel_mult': [1, 2, 4], 'transformer_depth_middle': -2, 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,644 'use_temporal_attention': False, 'use_temporal_resblock': False}645 646 KOALA_1B = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,647 'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,648 'num_res_blocks': [1, 1, 1], 'transformer_depth': [0, 2, 6], 'transformer_depth_output': [0, 0, 2, 2, 6, 6],649 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 6, 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,650 'use_temporal_attention': False, 'use_temporal_resblock': False}651 652 SD09_XS = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,653 'adm_in_channels': None, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': [1, 1, 1],654 'transformer_depth': [1, 1, 1], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': -2, 'use_linear_in_transformer': True,655 'context_dim': 1024, 'num_head_channels': 64, 'transformer_depth_output': [1, 1, 1, 1, 1, 1],656 'use_temporal_attention': False, 'use_temporal_resblock': False, 'disable_self_attentions': [True, False, False]}657 658 SD_XS = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,659 'adm_in_channels': None, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': [1, 1, 1],660 'transformer_depth': [0, 1, 1], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': -2, 'use_linear_in_transformer': False,661 'context_dim': 768, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 1, 1, 1, 1],662 'use_temporal_attention': False, 'use_temporal_resblock': False}663 664 SD15_diffusers_inpaint = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, 'adm_in_channels': None,665 'dtype': dtype, 'in_channels': 9, 'model_channels': 320, 'num_res_blocks': [2, 2, 2, 2], 'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0],666 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1, 'use_linear_in_transformer': False, 'context_dim': 768, 'num_heads': 8,667 'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],668 'use_temporal_attention': False, 'use_temporal_resblock': False}669 670 671 supported_models = [SDXL, SDXL_refiner, SD21, SD15, SD21_uncliph, SD21_unclipl, SDXL_mid_cnet, SDXL_small_cnet, SDXL_diffusers_inpaint, SSD_1B, Segmind_Vega, KOALA_700M, KOALA_1B, SD09_XS, SD_XS, SDXL_diffusers_ip2p, SD15_diffusers_inpaint]672 673 for unet_config in supported_models:674 matches = True675 for k in match:676 if match[k] != unet_config[k]:677 matches = False678 break679 if matches:680 return convert_config(unet_config)681 return None682 683def model_config_from_diffusers_unet(state_dict):684 unet_config = unet_config_from_diffusers_unet(state_dict)685 if unet_config is not None:686 return model_config_from_unet_config(unet_config)687 return None688 689def convert_diffusers_mmdit(state_dict, output_prefix=""):690 out_sd = {}691 692 if 'joint_transformer_blocks.0.attn.add_k_proj.weight' in state_dict: #AuraFlow693 num_joint = count_blocks(state_dict, 'joint_transformer_blocks.{}.')694 num_single = count_blocks(state_dict, 'single_transformer_blocks.{}.')695 sd_map = comfy.utils.auraflow_to_diffusers({"n_double_layers": num_joint, "n_layers": num_joint + num_single}, output_prefix=output_prefix)696 elif 'adaln_single.emb.timestep_embedder.linear_1.bias' in state_dict and 'pos_embed.proj.bias' in state_dict: # PixArt697 num_blocks = count_blocks(state_dict, 'transformer_blocks.{}.')698 sd_map = comfy.utils.pixart_to_diffusers({"depth": num_blocks}, output_prefix=output_prefix)699 elif 'x_embedder.weight' in state_dict: #Flux700 depth = count_blocks(state_dict, 'transformer_blocks.{}.')701 depth_single_blocks = count_blocks(state_dict, 'single_transformer_blocks.{}.')702 hidden_size = state_dict["x_embedder.bias"].shape[0]703 sd_map = comfy.utils.flux_to_diffusers({"depth": depth, "depth_single_blocks": depth_single_blocks, "hidden_size": hidden_size}, output_prefix=output_prefix)704 elif 'transformer_blocks.0.attn.add_q_proj.weight' in state_dict: #SD3705 num_blocks = count_blocks(state_dict, 'transformer_blocks.{}.')706 depth = state_dict["pos_embed.proj.weight"].shape[0] // 64707 sd_map = comfy.utils.mmdit_to_diffusers({"depth": depth, "num_blocks": num_blocks}, output_prefix=output_prefix)708 else:709 return None710 711 for k in sd_map:712 weight = state_dict.get(k, None)713 if weight is not None:714 t = sd_map[k]715 716 if not isinstance(t, str):717 if len(t) > 2:718 fun = t[2]719 else:720 fun = lambda a: a721 offset = t[1]722 if offset is not None:723 old_weight = out_sd.get(t[0], None)724 if old_weight is None:725 old_weight = torch.empty_like(weight)726 if old_weight.shape[offset[0]] < offset[1] + offset[2]:727 exp = list(weight.shape)728 exp[offset[0]] = offset[1] + offset[2]729 new = torch.empty(exp, device=weight.device, dtype=weight.dtype)730 new[:old_weight.shape[0]] = old_weight731 old_weight = new732 733 w = old_weight.narrow(offset[0], offset[1], offset[2])734 else:735 old_weight = weight736 w = weight737 w[:] = fun(weight)738 t = t[0]739 out_sd[t] = old_weight740 else:741 out_sd[t] = weight742 state_dict.pop(k)743 744 return out_sd745