From fc57b98501263a7e3c13c9dc8d0e188d0da4fe14 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sat, 14 Dec 2024 17:23:18 +0100 Subject: [PATCH] PixArt Diffusers convert logic --- comfy/model_detection.py | 10 +++++ comfy/supported_models.py | 2 +- comfy/utils.py | 81 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 92 insertions(+), 1 deletion(-) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index 18315d7d6..11d3bd5ba 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -183,6 +183,10 @@ def detect_unet_config(state_dict, key_prefix): 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 '{}adaln_single.emb.timestep_embedder.linear_1.bias'.format(key_prefix) in state_dict_keys: #Lightricks ltxv dit_config = {} dit_config["image_model"] = "ltxv" @@ -205,6 +209,9 @@ def detect_unet_config(state_dict, 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 + else: + dit_config["input_size"] = 128 # 1024 + dit_config["pe_interpolation"] = 2 ar_key = "{}ar_embedder.mlp.0.weight".format(key_prefix) if ar_key in state_dict_keys: @@ -580,6 +587,9 @@ def convert_diffusers_mmdit(state_dict, output_prefix=""): num_joint = count_blocks(state_dict, 'joint_transformer_blocks.{}.') num_single = count_blocks(state_dict, 'single_transformer_blocks.{}.') sd_map = comfy.utils.auraflow_to_diffusers({"n_double_layers": num_joint, "n_layers": num_joint + num_single}, output_prefix=output_prefix) + elif 'adaln_single.emb.timestep_embedder.linear_2.weight' in state_dict: # PixArt + num_blocks = count_blocks(state_dict, 'transformer_blocks.{}.') + sd_map = comfy.utils.pixart_to_diffusers({"depth": num_blocks}, output_prefix=output_prefix) elif 'x_embedder.weight' in state_dict: #Flux depth = count_blocks(state_dict, 'transformer_blocks.{}.') depth_single_blocks = count_blocks(state_dict, 'single_transformer_blocks.{}.') diff --git a/comfy/supported_models.py b/comfy/supported_models.py index acaa266f0..be8196e04 100644 --- a/comfy/supported_models.py +++ b/comfy/supported_models.py @@ -619,7 +619,7 @@ class PixArtAlpha(supported_models_base.BASE): class PixArtSigma(PixArtAlpha): unet_config = { - "image_model": "pixart_alpha", + "image_model": "pixart_sigma", } latent_format = latent_formats.SDXL diff --git a/comfy/utils.py b/comfy/utils.py index 3add621ff..1e021ac7a 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -384,6 +384,87 @@ def mmdit_to_diffusers(mmdit_config, output_prefix=""): return key_map +PIXART_MAP_BASIC = { + # Resolution + ("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"), + # Aspect ratio + ("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"), + # Patch embeddings + ("x_embedder.proj.weight", "pos_embed.proj.weight"), + ("x_embedder.proj.bias", "pos_embed.proj.bias"), + # Caption projection + ("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"), + # AdaLN-single LN + ("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"), + # Shared norm + ("t_block.1.weight", "adaln_single.linear.weight"), + ("t_block.1.bias", "adaln_single.linear.bias"), + # Final block + ("final_layer.linear.weight", "proj_out.weight"), + ("final_layer.linear.bias", "proj_out.bias"), + ("final_layer.scale_shift_table", "scale_shift_table"), +} + +PIXART_MAP_BLOCK = { + (f"scale_shift_table", f"scale_shift_table"), + # Projection + (f"attn.proj.weight", f"attn1.to_out.0.weight"), + (f"attn.proj.bias", f"attn1.to_out.0.bias"), + # Feed-forward + (f"mlp.fc1.weight", f"ff.net.0.proj.weight"), + (f"mlp.fc1.bias", f"ff.net.0.proj.bias"), + (f"mlp.fc2.weight", f"ff.net.2.weight"), + (f"mlp.fc2.bias", f"ff.net.2.bias"), + # Cross-attention (proj) + (f"cross_attn.proj.weight" ,f"attn2.to_out.0.weight"), + (f"cross_attn.proj.bias" ,f"attn2.to_out.0.bias"), +} + +def pixart_to_diffusers(mmdit_config, output_prefix=""): + 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]) + + return key_map def auraflow_to_diffusers(mmdit_config, output_prefix=""): n_double_layers = mmdit_config.get("n_double_layers", 0)