nodes/city96-ComfyUI_ExtraModels
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

PixArt Allow missing state dict keys

Browse Source 
Fixes #21
This commit is contained in:
City
2024-04-21 21:12:15 +02:00
parent f9e1178de7
commit 2f2797e502
1 changed files with 12 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
PixArt/diffusers_convert.py
+12 -2
View File
@@ -72,7 +72,8 @@ def convert_state_dict(state_dict):
else:
cmap = conversion_map
new_state_dict = {k: state_dict[v] for k,v in cmap}
missing = [k for k,v in cmap if v not in state_dict]
new_state_dict = {k: state_dict[v] for k,v in cmap if k not in missing}
matched = list(v for k,v in cmap if v in state_dict.keys())
for depth in range(28):
@@ -96,6 +97,10 @@ def convert_state_dict(state_dict):
print(f"PixArt: UNET conversion has leftover keys! ({len(matched)} vs {len(state_dict)})")
print(list( set(state_dict.keys()) - set(matched) ))
if len(missing) > 0:
print(f"PixArt: UNET conversion has missing keys!")
print(missing)
return new_state_dict
# Same as above but for LoRA weights:
@@ -118,7 +123,8 @@ def convert_lora_state_dict(state_dict):
cmap.append((rep_ak(k), rep_ap(v)))
cmap.append((rep_bk(k), rep_bp(v)))
new_state_dict = {k: state_dict[v] for k,v in cmap}
missing = [k for k,v in cmap if v not in state_dict]
new_state_dict = {k: state_dict[v] for k,v in cmap if k not in missing}
matched = list(v for k,v in cmap if v in state_dict.keys())
for fp, fk in ((rep_ap, rep_ak),(rep_bp, rep_bk)):
@@ -142,4 +148,8 @@ def convert_lora_state_dict(state_dict):
print(f"PixArt: LoRA conversion has leftover keys! ({len(matched)} vs {len(state_dict)})")
print(list( set(state_dict.keys()) - set(matched) ))
if len(missing) > 0:
print(f"PixArt: LoRA conversion has missing keys!")
print(missing)
return new_state_dict