fix: detect SDXL when the UNet is loaded standalone from diffusers tensor names (#2247)

This commit is contained in:
dubs-dev authored and GitHub committed 2026-10-29 23:33:44 +08:00
1 parent 997478a4cb
commit e7b7ec6d6f
1 file changed
+15
+15
View File
@@ -469,6 +469,7 @@ SDVersion ModelLoader::get_sd_version() const {
bool has_output_block_71 = false; bool has_output_block_71 = false;
bool has_attn_1024 = false; bool has_attn_1024 = false;
bool has_distilled_guidance = false; bool has_distilled_guidance = false;
bool has_sdxl_add_embedding = false;
for (auto& [name, tensor_storage] : tensor_storage_map) { for (auto& [name, tensor_storage] : tensor_storage_map) {
if (tensor_storage.name.find("model.diffusion_model.double_blocks.") != std::string::npos || if (tensor_storage.name.find("model.diffusion_model.double_blocks.") != std::string::npos ||
@@ -597,6 +598,7 @@ SDVersion ModelLoader::get_sd_version() const {
has_img_emb = true; has_img_emb = true;
} }
if (tensor_storage.name.find("model.diffusion_model.input_blocks.") != std::string::npos || if (tensor_storage.name.find("model.diffusion_model.input_blocks.") != std::string::npos ||
tensor_storage.name.find("model.diffusion_model.down_blocks.") != std::string::npos ||
tensor_storage.name.find("unet.down_blocks.") != std::string::npos) { tensor_storage.name.find("unet.down_blocks.") != std::string::npos) {
is_unet = true; is_unet = true;
if (has_multiple_encoders) { if (has_multiple_encoders) {
@@ -611,18 +613,27 @@ SDVersion ModelLoader::get_sd_version() const {
is_xl = true; is_xl = true;
} }
} }
// A standalone SDXL UNet (--diffusion-model) identifies the family by its added
// text/time embedding; its clip_g never carries a checkpoint prefix.
if (tensor_storage.name.find("model.diffusion_model.add_embedding.") != std::string::npos ||
tensor_storage.name.find("model.diffusion_model.label_emb.0.0.") != std::string::npos) {
has_sdxl_add_embedding = true;
}
if (tensor_storage.name.find("model.diffusion_model.input_blocks.8.0.time_mixer.mix_factor") != std::string::npos) { if (tensor_storage.name.find("model.diffusion_model.input_blocks.8.0.time_mixer.mix_factor") != std::string::npos) {
return VERSION_SVD; return VERSION_SVD;
} }
if (tensor_storage.name.find("model.diffusion_model.middle_block.1.") != std::string::npos || if (tensor_storage.name.find("model.diffusion_model.middle_block.1.") != std::string::npos ||
tensor_storage.name.find("model.diffusion_model.mid_block.resnets.1.") != std::string::npos ||
tensor_storage.name.find("unet.mid_block.resnets.1.") != std::string::npos) { tensor_storage.name.find("unet.mid_block.resnets.1.") != std::string::npos) {
has_middle_block_1 = true; has_middle_block_1 = true;
} }
if (tensor_storage.name.find("model.diffusion_model.output_blocks.3.1.transformer_blocks.1") != std::string::npos || if (tensor_storage.name.find("model.diffusion_model.output_blocks.3.1.transformer_blocks.1") != std::string::npos ||
tensor_storage.name.find("model.diffusion_model.up_blocks.1.attentions.0.transformer_blocks.1") != std::string::npos ||
tensor_storage.name.find("unet.up_blocks.1.attentions.0.transformer_blocks.1") != std::string::npos) { tensor_storage.name.find("unet.up_blocks.1.attentions.0.transformer_blocks.1") != std::string::npos) {
has_output_block_311 = true; has_output_block_311 = true;
} }
if (tensor_storage.name.find("model.diffusion_model.output_blocks.7.1") != std::string::npos || if (tensor_storage.name.find("model.diffusion_model.output_blocks.7.1") != std::string::npos ||
tensor_storage.name.find("model.diffusion_model.up_blocks.2.attentions.1") != std::string::npos ||
tensor_storage.name.find("unet.up_blocks.2.attentions.1") != std::string::npos) { tensor_storage.name.find("unet.up_blocks.2.attentions.1") != std::string::npos) {
has_output_block_71 = true; has_output_block_71 = true;
if (tensor_storage.name.find("model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_k.weight") != std::string::npos) { if (tensor_storage.name.find("model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_k.weight") != std::string::npos) {
@@ -640,6 +651,7 @@ SDVersion ModelLoader::get_sd_version() const {
// break; // break;
} }
if (tensor_storage.name == "model.diffusion_model.input_blocks.0.0.weight" || if (tensor_storage.name == "model.diffusion_model.input_blocks.0.0.weight" ||
tensor_storage.name == "model.diffusion_model.conv_in.weight" ||
tensor_storage.name == "model.diffusion_model.img_in.weight" || tensor_storage.name == "model.diffusion_model.img_in.weight" ||
tensor_storage.name == "unet.conv_in.weight") { tensor_storage.name == "unet.conv_in.weight") {
input_block_weight = tensor_storage; input_block_weight = tensor_storage;
@@ -663,6 +675,9 @@ SDVersion ModelLoader::get_sd_version() const {
} }
bool is_inpaint = input_block_weight.ne[2] == 9; bool is_inpaint = input_block_weight.ne[2] == 9;
bool is_ip2p = input_block_weight.ne[2] == 8; bool is_ip2p = input_block_weight.ne[2] == 8;
if (is_unet && has_sdxl_add_embedding) {
is_xl = true;
}
if (is_xl) { if (is_xl) {
if (is_inpaint) { if (is_inpaint) {
return VERSION_SDXL_INPAINT; return VERSION_SDXL_INPAINT;