support for SD2 added, debug mode added

Author: divamgupta

backends/model_converter/convert_model.py

@@ -51,6 +51,13 @@ def convert_model(checkpoint_filename, out_filename ):
  state_dict[k + "._split_1"] = pp[:pp.shape[0]//2].copy() 
  state_dict[k + "._split_2"] = pp[pp.shape[0]//2:].copy() 
 
+ elif "attn.in_proj" in k and state_dict[k].shape[0] == 3072 :
+ pp = state_dict[k]
+ state_dict[k + "._split_1"] = pp[:pp.shape[0]//3].copy()
+ state_dict[k + "._split_2"] = pp[pp.shape[0]//3:2*pp.shape[0]//3].copy()
+ state_dict[k + "._split_3"] = pp[2*pp.shape[0]//3: ].copy()
+
+
  keys_list = list(state_dict.keys())
  mid_key = keys_list[len(keys_list)//2]
 
@@ -69,7 +76,7 @@ def convert_model(checkpoint_filename, out_filename ):
  model_type = get_model_type(state_dict)
 
  if model_type is None:
- raise ValueError("The model is not supported. Please make sure it is a valid SD 1.4/1.5 .ckpt file")
+ raise ValueError("The model is not supported. Please make sure it is a valid SD 1.4/1.5/2.1 .ckpt/safetensor file")
 
  if "float16" in model_type:
  cur_dtype = "float16"
@@ -88,7 +95,9 @@ def convert_model(checkpoint_filename, out_filename ):
  outfile.init_write(ctdict_version=ctdict_id )
 
  for k in model_shapes:
- np_arr = state_dict[k]
+ np_arr = np.copy(state_dict[k])
+ np_arr = np.reshape(np_arr , model_shapes[k] )
+
  if "float" in str(np_arr.dtype):
  np_arr = np_arr.astype(cur_dtype)
  shape = list(np_arr.shape)

backends/model_converter/sd_shapes.py

@@ -1,5 +1,5 @@
 
-from sd_shapes_consts import shapes_unet , shapes_encoder, shapes_decoder , shapes_text_encoder, shapes_params
+from sd_shapes_consts import shapes_unet , shapes_encoder, shapes_decoder , shapes_text_encoder, shapes_params , shapes_unet_v2 , text_encoder_open_clip
 import copy
 from collections import Counter
 
@@ -10,11 +10,19 @@ def add_aux_shapes(d):
 
  if ".ff." in k:
  sh = list(d[k])
- sh[0] /= 2
+ sh[0] = sh[0] // 2
  sh = tuple(sh)
  d[k + "._split_1"] = sh
  d[k + "._split_2"] = sh
 
+ elif "attn.in_proj" in k and d[k][0] == 3072 :
+ sh = list(d[k])
+ sh[0] = sh[0] // 3
+ sh = tuple(sh)
+ d[k + "._split_1"] = sh
+ d[k + "._split_2"] = sh
+ d[k + "._split_3"] = sh
+
  for i in range(1,21):
  nn = 320*i
  d["zeros_"+str(nn)] = (nn,)
@@ -41,16 +49,35 @@ def add_aux_shapes(d):
 add_aux_shapes(sd_1x_inpaint_shapes)
 
 
+
+
+sd_2x_shapes = {}
+sd_2x_shapes.update(shapes_unet_v2)
+sd_2x_shapes.update(shapes_encoder)
+sd_2x_shapes.update(shapes_decoder)
+sd_2x_shapes.update(text_encoder_open_clip)
+sd_2x_shapes.update(shapes_params)
+
+
+add_aux_shapes(sd_2x_shapes)
+
+
 possible_model_shapes = {"SD_1x_float32": sd_1x_shapes , 
  "SD_1x_inpaint_float32": sd_1x_inpaint_shapes, 
  "SD_1x_float16": sd_1x_shapes , 
+
+ "SD_2x_float16": sd_2x_shapes , 
+ "SD_2x_float32": sd_2x_shapes , 
+
  "SD_1x_inpaint_float16": sd_1x_inpaint_shapes}
 
 ctdict_ids = {"SD_1x_float32": 12 , 
  "SD_1x_inpaint_float32": 13, 
+ "SD_2x_float32": 15 ,
  "SD_1x_float16": 1012 , 
  "SD_1x_inpaint_float16": 1013 , 
- "SD_1x_just_controlnet_16" : 1014}
+ "SD_1x_just_controlnet_16" : 1014,
+ "SD_2x_float16": 1015 }
 
 
 extra_keys = ['temb_coefficients_fp32' , 'temb_coefficients_fp16' , 'causal_mask' , 'aux_output_conv.weight' , 'aux_output_conv.bias', 'alphas_cumprod']
@@ -63,8 +90,9 @@ def are_shapes_matching(state_dict , template_shapes , name=None):
  print("key", k , "not found in state_dict" , state_dict.keys())
  return False
  if tuple(template_shapes[k]) != tuple(state_dict[k].shape):
- print("shape mismatch", k , tuple(template_shapes[k]) ,tuple(state_dict[k].shape) )
- return False 
+ if tuple(template_shapes[k]) != tuple(state_dict[k].shape) + (1,1):
+ print("shape mismatch", k , tuple(template_shapes[k]) ,tuple(state_dict[k].shape) )
+ return False 
 
  return True
 
@@ -101,6 +129,9 @@ def get_model_type(state_dict):
  if are_shapes_matching(state_dict , sd_1x_shapes) :
  shapes = sd_1x_shapes
  mname = "SD_1x"
+ elif are_shapes_matching(state_dict , sd_2x_shapes) :
+ shapes = sd_2x_shapes
+ mname = "SD_2x"
  elif are_shapes_matching(state_dict , sd_1x_inpaint_shapes) :
  shapes = sd_1x_inpaint_shapes
  mname = "SD_1x_inpaint"