added safetensors support

Author: divamgupta

backends/model_converter/convert_model.py

@@ -8,15 +8,21 @@
 # pyinstaller convert_model.py --onefile --noconfirm --clean # build using intel machine so that its cross platform lol
 
 
-
+from safetensor_wrapper import SafetensorWrapper
 from fake_torch import extract_weights_from_checkpoint
 from sd_shapes import get_model_type , possible_model_shapes , ctdict_ids
 from tdict import TDict
 
 
 
 def convert_model(checkpoint_filename, out_filename ):
- torch_weights = extract_weights_from_checkpoint(open(checkpoint_filename, "rb"))
+
+ if checkpoint_filename.lower().endswith(".ckpt"):
+ torch_weights = extract_weights_from_checkpoint(open(checkpoint_filename, "rb"))
+ elif checkpoint_filename.lower().endswith(".safetensors"):
+ torch_weights = SafetensorWrapper(checkpoint_filename)
+ else:
+ raise ValueError("Invalid import format")
 
  if 'state_dict' in torch_weights:
  state_dict = torch_weights['state_dict']
@@ -65,6 +71,13 @@ def convert_model(checkpoint_filename, out_filename ):
  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")
 
+ if "float16" in model_type:
+ cur_dtype = "float16"
+ elif "float32" in model_type:
+ cur_dtype = "float32"
+ else:
+ assert False
+
  print("model type " , model_type)
 
  model_shapes = possible_model_shapes[model_type]
@@ -76,6 +89,8 @@ def convert_model(checkpoint_filename, out_filename ):
 
  for k in model_shapes:
  np_arr = state_dict[k]
+ if "float" in str(np_arr.dtype):
+ np_arr = np_arr.astype(cur_dtype)
  shape = list(np_arr.shape)
  assert tuple(shape) == tuple(model_shapes[k]), ( "shape mismatch at" , k , shape , SD_SHAPES[k] )
  outfile.write_key(key=k , tensor=np_arr)
@@ -113,4 +128,6 @@ def usage():
  checkpoint_filename = args[0]
  out_filename = args[1]
 
+ convert_model(checkpoint_filename , out_filename )
+
 

backends/model_converter/fake_torch.py

@@ -140,13 +140,17 @@ def examine_pickle(fb0, return_special=False):
  ## 3: this massive line also assigns values to keys, but does so differently
  ## _var2262.update({ 'cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm2.bias': _var2001, [ .... and on and on ]})
  ##
+ ## 4: in some pruned models, the last line is instead a combination of 2/3 into the final variable:
+ ## result = {'model.diffusion_model.input_blocks.0.0.weight': _var1, 'model.diffusion_model.input_blocks.0.0.bias': _var3, }
+ ##
  ## that's it 
 
  # make some REs to match the above.
  re_rebuild = re.compile('^_var\d+ = _rebuild_tensor_v2\(UNPICKLER\.persistent_load\(\(.*\)$')
  re_assign = re.compile('^_var\d+ = \{.*\}$')
  re_update = re.compile('^_var\d+\.update\(\{.*\}\)$')
  re_ordered_dict = re.compile('^_var\d+ = OrderedDict\(\)$')
+ re_result = re.compile('^result = \{.*\}$')
 
  load_instructions = {}
  assign_instructions = AssignInstructions()
@@ -157,7 +161,7 @@ def examine_pickle(fb0, return_special=False):
  if re_rebuild.match(line):
  variable_name, load_instruction = line.split(' = ', 1)
  load_instructions[variable_name] = LoadInstruction(line, variable_name)
- elif re_assign.match(line):
+ elif re_assign.match(line) or re_result.match(line):
  assign_instructions.parse_assign_line(line)
  elif re_update.match(line):
  assign_instructions.parse_update_line(line)
@@ -184,11 +188,34 @@ def __init__(self, collect_special=False):
  self.integrated_instructions = {}
  self.collect_special = collect_special;
 
+ def parse_result_line(self, line):
+ garbage, huge_mess = line.split(' = {', 1)
+ assignments = huge_mess.split(', ')
+ del huge_mess
+ assignments[-1] = assignments[-1].strip('}')
+
+ #compile RE here to avoid doing it every loop iteration:
+ re_var = re.compile('^_var\d+$')
+
+ assignment_count = 0
+ for a in assignments:
+ if self._add_assignment(a, re_var):
+ assignment_count = assignment_count + 1
+ if NO_PICKLE_DEBUG:
+ print(f"Added/merged {assignment_count} assignments. Total of {len(self.instructions)} assignment instructions")
+
  def parse_assign_line(self, line):
  # input looks like this:
  # _var2262 = {'model.diffusion_model.input_blocks.0.0.weight': _var1, 'model.diffusion_model.input_blocks.0.0.bias': _var3,\
  # ...\
  # 'cond_stage_model.transformer.text_model.encoder.layers.3.layer_norm2.weight': _var1999}
+
+ # input looks like the above, but with 'result' in place of _var2262:
+ # result = {'model.diffusion_model.input_blocks.0.0.weight': _var1, ... }
+ #
+ # or also look like: 
+ # result = {'state_dict': _var2314}
+ # ... which will be ignored later
  garbage, huge_mess = line.split(' = {', 1)
  assignments = huge_mess.split(', ')
  del huge_mess
@@ -211,7 +238,7 @@ def _add_assignment(self, assignment, re_var):
  # 'embedding_manager.embedder.transformer.text_model.encoder.layers.6.mlp.fc1': {'version': 1}
  sd_key, fickling_var = assignment.split(': ', 1)
  sd_key = sd_key.strip("'")
- if re_var.match(fickling_var):
+ if sd_key != 'state_dict' and re_var.match(fickling_var):
  self.instructions[sd_key] = fickling_var
  return True
  elif self.collect_special:
@@ -225,7 +252,8 @@ def _add_assignment(self, assignment, re_var):
  v = v.strip("'")
  special_dict[k] = v
  self.special_instructions[sd_key] = special_dict
- return False
+
+ return False
 
  def integrate(self, load_instructions):
  unfound_keys = {}

backends/model_converter/safetensor_wrapper.py

@@ -0,0 +1,30 @@
+from safetensors import safe_open
+
+
+class SafetensorWrapper:
+
+def __init__(self , fname ):
+self.file = safe_open(fname, framework="np", device="cpu")
+self.new_items = {}
+
+def keys(self):
+return list(self.file.keys()) + list(self.new_items.keys())
+
+def __contains__(self, k):
+if k in self.file.keys():
+return True 
+if k in self.new_items:
+return True 
+return False
+
+def __getitem__(self , k):
+if k in self.new_items:
+return self.new_items[k]
+else:
+return self.file.get_tensor(k)
+
+def __setitem__(self, key , item ):
+self.new_items[key] = item
+
+def __iter__(self):
+return iter(self.keys())

backends/model_converter/sd_shapes.py

@@ -1,7 +1,7 @@
 
 from sd_shapes_consts import shapes_unet , shapes_encoder, shapes_decoder , shapes_text_encoder, shapes_params
 import copy
-
+from collections import Counter
 
 def add_aux_shapes(d):
  for k in list(d.keys()):
@@ -57,7 +57,7 @@ def add_aux_shapes(d):
 
 
 
-def are_shapes_matching(state_dict , template_shapes):
+def are_shapes_matching(state_dict , template_shapes , name=None):
  for k in template_shapes:
  if k not in state_dict:
  print("key", k , "not found in state_dict" , state_dict.keys())
@@ -68,28 +68,50 @@ def are_shapes_matching(state_dict , template_shapes):
 
  return True
 
-def are_shapes_dtype(state_dict, template_shapes , dtype):
+
+def get_dtype(state_dict, template_shapes ):
+ c = Counter()
+
  for k in state_dict:
  if k in extra_keys:
  continue
  if k not in template_shapes:
  continue
- if state_dict[k].dtype != dtype:
- return False 
 
- return True 
+ if 'float' in str(state_dict[k].dtype):
+ c[ str(state_dict[k].dtype)] += 1 
+ print(c.most_common())
+ return c.most_common(1)[0][0]
+
+
+
+def check_shapes_float(state_dict, template_shapes ):
+ for k in state_dict:
+ if k in extra_keys:
+ continue
+ if k not in template_shapes:
+ continue
+ 
+ assert 'float' in str(state_dict[k].dtype )
+
 
 
 def get_model_type(state_dict):
- if are_shapes_matching(state_dict , sd_1x_shapes) and are_shapes_dtype(state_dict , sd_1x_shapes, "float32"):
- return "SD_1x_float32"
- elif are_shapes_matching(state_dict , sd_1x_inpaint_shapes) and are_shapes_dtype(state_dict , sd_1x_inpaint_shapes , "float32"):
- return "SD_1x_inpaint_float32"
- elif are_shapes_matching(state_dict , sd_1x_shapes) and are_shapes_dtype(state_dict , sd_1x_shapes , "float16"):
- return "SD_1x_float16"
- elif are_shapes_matching(state_dict , sd_1x_inpaint_shapes) and are_shapes_dtype(state_dict , sd_1x_inpaint_shapes, "float16"):
- return "SD_1x_inpaint_float16"
+
+ if are_shapes_matching(state_dict , sd_1x_shapes) :
+ shapes = sd_1x_shapes
+ mname = "SD_1x"
+ elif are_shapes_matching(state_dict , sd_1x_inpaint_shapes) :
+ shapes = sd_1x_inpaint_shapes
+ mname = "SD_1x_inpaint"
  else:
  return None
 
+ check_shapes_float(state_dict , shapes)
+ c_dtype = get_dtype(state_dict , shapes)
+ if c_dtype not in ["float32" , "float16"]:
+ raise ValueError("The weights should either be float32 or float16, but these are " + c_dtype)
+
+ return mname + "_" + c_dtype
+
 

backends/stable_diffusion/.gitingore renamed to backends/stable_diffusion/.gitignore

@@ -1,7 +1,7 @@
 {
  "name": "DiffusionBee",
- "version": "1.6.0",
- "build_number": "0017",
+ "version": "1.7.0",
+ "build_number": "0018",
  "website": "https://diffusionbee.com",
  "description": "Diffusion Bee - Stable Diffusion App.",
  "is_dev": false,

electron_app/package.json

@@ -81,7 +81,7 @@ export default {
  add_model(){
  let that = this;
  
- let pytorch_model_path = window.ipcRenderer.sendSync('file_dialog', "ckpt_file" );
+ let pytorch_model_path = window.ipcRenderer.sendSync('file_dialog', "weights_file" );
  if(!pytorch_model_path)
  return;
 

electron_app/src/components/Settings.vue

@@ -61,10 +61,10 @@ ipcMain.on('file_dialog', (event, arg) => {
  properties = ['openFile' ]
  options = { filters :[ {name: 'Images', extensions: ['jpg', 'jpeg', 'png', 'bmp']}] , properties: properties } ;
  }
- else if(arg == 'ckpt_file') // single image file 
+ else if(arg == 'weights_file') // single image file 
  {
  properties = ['openFile' ]
- options = { filters :[ {name: 'Checkpoints', extensions: ['ckpt']}] , properties: properties } ;
+ options = { filters :[ {name: 'Checkpoints', extensions: ['ckpt' , 'safetensors' ]}] , properties: properties } ;
  }
  else if(arg == 'img_files') // multi image files
  {