refactor: align PixArt weights with upstream layout (#2061)

src/name_conversion.cpp

@@ -932,6 +932,68 @@ static bool is_diffusers_controlnet_name(const std::string& name) {
     return false;
 }
 
+static std::string convert_diffusers_dit_to_original_pixart(std::string name) {
+    static const std::vector<std::pair<std::string, std::string>> prefix_map = {
+        {"pos_embed.proj.", "x_embedder.proj."},
+        {"adaln_single.emb.timestep_embedder.linear_1.", "t_embedder.mlp.0."},
+        {"adaln_single.emb.timestep_embedder.linear_2.", "t_embedder.mlp.2."},
+        {"adaln_single.emb.resolution_embedder.linear_1.", "csize_embedder.mlp.0."},
+        {"adaln_single.emb.resolution_embedder.linear_2.", "csize_embedder.mlp.2."},
+        {"adaln_single.emb.aspect_ratio_embedder.linear_1.", "ar_embedder.mlp.0."},
+        {"adaln_single.emb.aspect_ratio_embedder.linear_2.", "ar_embedder.mlp.2."},
+        {"adaln_single.linear.", "t_block.1."},
+        {"caption_projection.linear_1.", "y_embedder.y_proj.fc1."},
+        {"caption_projection.linear_2.", "y_embedder.y_proj.fc2."},
+        {"proj_out.", "final_layer.linear."},
+    };
+    for (const auto& entry : prefix_map) {
+        if (starts_with(name, entry.first)) {
+            return entry.second + name.substr(entry.first.size());
+        }
+    }
+    if (name == "scale_shift_table") {
+        return "final_layer.scale_shift_table";
+    }
+    const std::string block_prefix = "transformer_blocks.";
+    if (!starts_with(name, block_prefix)) {
+        return name;
+    }
+    const size_t block_end = name.find('.', block_prefix.size());
+    if (block_end == std::string::npos) {
+        return name;
+    }
+    const std::string prefix                                                = "blocks." + name.substr(block_prefix.size(), block_end - block_prefix.size()) + ".";
+    name                                                                    = name.substr(block_end + 1);
+    static const std::vector<std::pair<std::string, std::string>> block_map = {
+        {"attn1.to_q.", "attn.qkv."},
+        {"attn1.to_out.0.", "attn.proj."},
+        {"attn2.to_q.", "cross_attn.q_linear."},
+        {"attn2.to_k.", "cross_attn.kv_linear."},
+        {"attn2.to_out.0.", "cross_attn.proj."},
+        {"ff.net.0.proj.", "mlp.fc1."},
+        {"ff.net.2.", "mlp.fc2."},
+    };
+    for (const auto& entry : block_map) {
+        if (starts_with(name, entry.first)) {
+            return prefix + entry.second + name.substr(entry.first.size());
+        }
+    }
+    static const std::vector<std::pair<std::string, std::string>> part_map = {
+        {"attn1.to_k.weight", "attn.qkv.weight.1"},
+        {"attn1.to_k.bias", "attn.qkv.bias.1"},
+        {"attn1.to_v.weight", "attn.qkv.weight.2"},
+        {"attn1.to_v.bias", "attn.qkv.bias.2"},
+        {"attn2.to_v.weight", "cross_attn.kv_linear.weight.1"},
+        {"attn2.to_v.bias", "cross_attn.kv_linear.bias.1"},
+    };
+    for (const auto& entry : part_map) {
+        if (name == entry.first || starts_with(name, entry.first + ".")) {
+            return prefix + entry.second + name.substr(entry.first.size());
+        }
+    }
+    return prefix + name;
+}
+
 std::string convert_diffusion_model_name(std::string name, std::string prefix, SDVersion version) {
     if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
         name = convert_diffusers_unet_to_original_sd1(name);
@@ -951,6 +1013,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
         name = convert_other_dit_to_original_anima(name);
     } else if (sd_version_is_krea2(version)) {
         name = convert_diffusers_dit_to_original_krea2(name);
+    } else if (sd_version_is_pixart(version)) {
+        name = convert_diffusers_dit_to_original_pixart(name);
     }
     return name;
 }