Use different keys for detection between original model and converted model

2024-06-19 09:45:33 +01:00
parent 91340dedc7
commit 712f57c915
2 changed files with 2 additions and 2 deletions
@@ -17,7 +17,7 @@ conversion_map_ms = [ # for multi_scale_train (MS)
 ]

 def get_depth(state_dict):
-	return sum(key.endswith('.scale_shift_table') for key in state_dict.keys())
+	return sum(key.endswith('.attn1.to_k.bias') for key in state_dict.keys())

 def get_conversion_map(state_dict):
 	conversion_map = [  # main SD conversion map (PixArt reference, HF Diffusers)