(state_dict)
| 29 | |
| 30 | |
| 31 | def convert_transformer(state_dict): |
| 32 | converted_state_dict = {} |
| 33 | state_dict_keys = list(state_dict.keys()) |
| 34 | |
| 35 | converted_state_dict["register_tokens"] = state_dict.pop("model.register_tokens") |
| 36 | converted_state_dict["pos_embed.pos_embed"] = state_dict.pop("model.positional_encoding") |
| 37 | converted_state_dict["pos_embed.proj.weight"] = state_dict.pop("model.init_x_linear.weight") |
| 38 | converted_state_dict["pos_embed.proj.bias"] = state_dict.pop("model.init_x_linear.bias") |
| 39 | |
| 40 | converted_state_dict["time_step_proj.linear_1.weight"] = state_dict.pop("model.t_embedder.mlp.0.weight") |
| 41 | converted_state_dict["time_step_proj.linear_1.bias"] = state_dict.pop("model.t_embedder.mlp.0.bias") |
| 42 | converted_state_dict["time_step_proj.linear_2.weight"] = state_dict.pop("model.t_embedder.mlp.2.weight") |
| 43 | converted_state_dict["time_step_proj.linear_2.bias"] = state_dict.pop("model.t_embedder.mlp.2.bias") |
| 44 | |
| 45 | converted_state_dict["context_embedder.weight"] = state_dict.pop("model.cond_seq_linear.weight") |
| 46 | |
| 47 | mmdit_layers = calculate_layers(state_dict_keys, key_prefix="double_layers") |
| 48 | single_dit_layers = calculate_layers(state_dict_keys, key_prefix="single_layers") |
| 49 | |
| 50 | # MMDiT blocks 🎸. |
| 51 | for i in range(mmdit_layers): |
| 52 | # feed-forward |
| 53 | path_mapping = {"mlpX": "ff", "mlpC": "ff_context"} |
| 54 | weight_mapping = {"c_fc1": "linear_1", "c_fc2": "linear_2", "c_proj": "out_projection"} |
| 55 | for orig_k, diffuser_k in path_mapping.items(): |
| 56 | for k, v in weight_mapping.items(): |
| 57 | converted_state_dict[f"joint_transformer_blocks.{i}.{diffuser_k}.{v}.weight"] = state_dict.pop( |
| 58 | f"model.double_layers.{i}.{orig_k}.{k}.weight" |
| 59 | ) |
| 60 | |
| 61 | # norms |
| 62 | path_mapping = {"modX": "norm1", "modC": "norm1_context"} |
| 63 | for orig_k, diffuser_k in path_mapping.items(): |
| 64 | converted_state_dict[f"joint_transformer_blocks.{i}.{diffuser_k}.linear.weight"] = state_dict.pop( |
| 65 | f"model.double_layers.{i}.{orig_k}.1.weight" |
| 66 | ) |
| 67 | |
| 68 | # attns |
| 69 | x_attn_mapping = {"w2q": "to_q", "w2k": "to_k", "w2v": "to_v", "w2o": "to_out.0"} |
| 70 | context_attn_mapping = {"w1q": "add_q_proj", "w1k": "add_k_proj", "w1v": "add_v_proj", "w1o": "to_add_out"} |
| 71 | for attn_mapping in [x_attn_mapping, context_attn_mapping]: |
| 72 | for k, v in attn_mapping.items(): |
| 73 | converted_state_dict[f"joint_transformer_blocks.{i}.attn.{v}.weight"] = state_dict.pop( |
| 74 | f"model.double_layers.{i}.attn.{k}.weight" |
| 75 | ) |
| 76 | |
| 77 | # Single-DiT blocks. |
| 78 | for i in range(single_dit_layers): |
| 79 | # feed-forward |
| 80 | mapping = {"c_fc1": "linear_1", "c_fc2": "linear_2", "c_proj": "out_projection"} |
| 81 | for k, v in mapping.items(): |
| 82 | converted_state_dict[f"single_transformer_blocks.{i}.ff.{v}.weight"] = state_dict.pop( |
| 83 | f"model.single_layers.{i}.mlp.{k}.weight" |
| 84 | ) |
| 85 | |
| 86 | # norms |
| 87 | converted_state_dict[f"single_transformer_blocks.{i}.norm1.linear.weight"] = state_dict.pop( |
| 88 | f"model.single_layers.{i}.modCX.1.weight" |
no test coverage detected
searching dependent graphs…