From 2f2797e50212bc00d25b7af6da9db47f424764ff Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Sun, 21 Apr 2024 21:12:15 +0200
Subject: [PATCH] PixArt Allow missing state dict keys

Fixes #21
---
 PixArt/diffusers_convert.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/PixArt/diffusers_convert.py b/PixArt/diffusers_convert.py
index 7151773..2bab19e 100644
--- a/PixArt/diffusers_convert.py
+++ b/PixArt/diffusers_convert.py
@@ -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