mirrors/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-09-10 19:46:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Properly support SDXL diffusers unet with UNETLoader node.

Browse Source
This commit is contained in:
comfyanonymous
2023-07-21 14:38:56 -04:00
parent 0115018695
commit 58b2364f58
2 changed files with 23 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
comfy/sd.py
View File

@@ -1139,12 +1139,14 @@ def load_unet(unet_path): #load unet in diffusers format
fp16 = model_management.should_use_fp16(model_params=parameters)
match = {}
match["context_dim"] = sd["down_blocks.0.attentions.1.transformer_blocks.0.attn2.to_k.weight"].shape[1]
match["context_dim"] = sd["down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_k.weight"].shape[1]
match["model_channels"] = sd["conv_in.weight"].shape[0]
match["in_channels"] = sd["conv_in.weight"].shape[1]
match["adm_in_channels"] = None
if "class_embedding.linear_1.weight" in sd:
match["adm_in_channels"] = sd["class_embedding.linear_1.weight"].shape[1]
elif "add_embedding.linear_1.weight" in sd:
match["adm_in_channels"] = sd["add_embedding.linear_1.weight"].shape[1]
SDXL = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
'num_classes': 'sequential', 'adm_in_channels': 2816, 'use_fp16': fp16, 'in_channels': 4, 'model_channels': 320,
@@ -1198,6 +1200,7 @@ def load_unet(unet_path): #load unet in diffusers format
model = model.to(offload_device)
model.load_model_weights(new_sd, "")
return ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=offload_device)
print("ERROR UNSUPPORTED UNET", unet_path)
def save_checkpoint(output_path, model, clip, vae, metadata=None):
try: