Added options for collecting metadata that the conversion script doesn't care about as well as fixing error when importing PaperCut model

Author: CrudeDiatribe

backends/model_converter/no_pickle_fake_torch.py

@@ -11,7 +11,7 @@
 else:
  from astunparse import unparse
 
-NO_PICKLE_DEBUG = True
+NO_PICKLE_DEBUG = False
 
 def extract_weights_from_checkpoint(fb0):
  torch_weights = {}
@@ -22,32 +22,40 @@ def extract_weights_from_checkpoint(fb0):
  raise ValueError("Looks like the checkpoints file is in the wrong format")
  folder_name = folder_name[0].replace("/data.pkl" , "").replace("\\data.pkl" , "")
  with myzip.open(folder_name+'/data.pkl') as myfile:
- load_instructions, special_instructions = examine_pickle(myfile)
+ load_instructions = examine_pickle(myfile)
  for sd_key,load_instruction in load_instructions.items():
  with myzip.open(folder_name + f'/data/{load_instruction.obj_key}') as myfile:
  if (load_instruction.load_from_file_buffer(myfile)):
  torch_weights['state_dict'][sd_key] = load_instruction.get_data()
- if len(special_instructions) > 0:
- torch_weights['state_dict']['_metadata'] = {}
- for sd_key,special in special_instructions.items():
- torch_weights['state_dict']['_metadata'][sd_key] = special
+ #if len(special_instructions) > 0:
+ # torch_weights['state_dict']['_metadata'] = {}
+ # for sd_key,special in special_instructions.items():
+ # torch_weights['state_dict']['_metadata'][sd_key] = special
  return torch_weights
 
-def examine_pickle(fb0):
+def examine_pickle(fb0, return_special=False):
+ ## return_special:
+ ## A rabbit hole I chased trying to debug a model that wouldn't import that had 1300 useless metadata statements
+ ## If for some reason it's needed in the future turn it on. It is passed into the class AssignInstructions and 
+ ## if turned on collect_special will be True
+ ## 
 
+ #turn the pickle file into text we can parse
  decompiled = unparse(Pickled.load(fb0).ast).splitlines()
 
-## LINES WE CARE ABOUT:
-## 1: this defines a data file and what kind of data is in it
-## _var1 = _rebuild_tensor_v2(UNPICKLER.persistent_load(('storage', HalfStorage, '0', 'cpu', 11520)), 0, (320, 4, 3, 3), (36, 9, 3, 1), False, _var0)
-##
-## 2: this massive line assigns the previous data to dictionary entries 
-## _var2262 = {'model.diffusion_model.input_blocks.0.0.weight': _var1, [..... continue for ever]}
-##
-## 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 ]})
-##
-## that's it 
+ ## Parsing the decompiled pickle:
+ ## LINES WE CARE ABOUT:
+ ## 1: this defines a data file and what kind of data is in it
+ ## _var1 = _rebuild_tensor_v2(UNPICKLER.persistent_load(('storage', HalfStorage, '0', 'cpu', 11520)), 0, (320, 4, 3, 3), (36, 9, 3, 1), False, _var0)
+ ##
+ ## 2: this massive line assigns the previous data to dictionary entries 
+ ## _var2262 = {'model.diffusion_model.input_blocks.0.0.weight': _var1, [..... continue for ever]}
+ ##
+ ## 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 ]})
+ ##
+ ## 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+ = \{.*\}$')
@@ -62,15 +70,15 @@ def examine_pickle(fb0):
  line = line.strip()
  if re_rebuild.match(line):
  variable_name, load_instruction = line.split(' = ', 1)
- load_instructions[variable_name] = LoadInstruction(line)
+ load_instructions[variable_name] = LoadInstruction(line, variable_name)
  elif re_assign.match(line):
  assign_instructions.parse_assign_line(line)
  elif re_update.match(line):
  assign_instructions.parse_update_line(line)
  elif re_ordered_dict.match(line):
  #do nothing
  continue
- else:
+ elif NO_PICKLE_DEBUG:
  print(f'unmatched line: {line}')
 
 
@@ -79,14 +87,16 @@ def examine_pickle(fb0):
 
  assign_instructions.integrate(load_instructions)
 
- return assign_instructions.integrated_instructions, assign_instructions.special_instructions
- #return assign_instructions.integrated_instructions, {}
+ if return_special:
+ return assign_instructions.integrated_instructions, assign_instructions.special_instructions
+ return assign_instructions.integrated_instructions
 
 class AssignInstructions:
- def __init__(self):
+ def __init__(self, collect_special=False):
  self.instructions = {}
  self.special_instructions = {}
  self.integrated_instructions = {}
+ self.collect_special = collect_special;
 
  def parse_assign_line(self, line):
  # input looks like this:
@@ -115,7 +125,7 @@ def _add_assignment(self, assignment, re_var):
  if re_var.match(fickling_var):
  self.instructions[sd_key] = fickling_var
  return True
- else:
+ elif self.collect_special:
  # now convert the string "{'version': 1}" into a dictionary {'version': 1}
  entries = fickling_var.split(',')
  special_dict = {}
@@ -133,14 +143,9 @@ def integrate(self, load_instructions):
  for sd_key, fickling_var in self.instructions.items():
  if fickling_var in load_instructions:
  self.integrated_instructions[sd_key] = load_instructions[fickling_var]
- if sd_key in self.special_instructions:
- if NO_PICKLE_DEBUG:
- print(f"Key found in both load and special instructions: {sd_key}")
  else:
- unfound_keys[sd_key] = True;
- #for sd_key, special in self.special_instructions.items():
- # if sd_key in unfound_keys:
- # #todo
+ if NO_PICKLE_DEBUG:
+ print(f"no load instruction found for {sd_key}")
 
  if NO_PICKLE_DEBUG:
  print(f"Have {len(self.integrated_instructions)} integrated load/assignment instructions")
@@ -164,14 +169,16 @@ def parse_update_line(self, line):
  print(f"Added/merged {update_count} updates. Total of {len(self.instructions)} assignment instructions")
 
 class LoadInstruction:
- def __init__(self, instruction_string):
+ def __init__(self, instruction_string, variable_name, extra_debugging = False):
  self.ident = False
  self.storage_type = False
  self.obj_key = False
  self.location = False #unused
  self.obj_size = False
  self.stride = False #unused
- self.data = False;
+ self.data = False
+ self.variable_name = variable_name
+ self.extra_debugging = extra_debugging
  self.parse_instruction(instruction_string)
 
  def parse_instruction(self, instruction_string):
@@ -185,12 +192,24 @@ def parse_instruction(self, instruction_string):
  #
  # the following comments will show the output of each string manipulation as if it started with the above.
 
+ if self.extra_debugging:
+ print(f"input: '{instruction_string}'")
+
  garbage, storage_etc = instruction_string.split('((', 1)
  # storage_etc = 'storage', HalfStorage, '0', 'cpu', 11520)), 0, (320, 4, 3, 3), (36, 9, 3, 1), False, _var0)
+
+ if self.extra_debugging:
+ print("storage_etc, reference: ''storage', HalfStorage, '0', 'cpu', 11520)), 0, (320, 4, 3, 3), (36, 9, 3, 1), False, _var0)'")
+ print(f"storage_etc, actual: '{storage_etc}'\n")
 
  storage, etc = storage_etc.split('))', 1)
  # storage = 'storage', HalfStorage, '0', 'cpu', 11520
- # etc = 0, (320, 4, 3, 3), (36, 9, 3, 1), False, _var0)
+ # etc = , 0, (320, 4, 3, 3), (36, 9, 3, 1), False, _var0)
+ if self.extra_debugging:
+ print("storage, reference: ''storage', HalfStorage, '0', 'cpu', 11520'")
+ print(f"storage, actual: '{storage}'\n")
+ print("etc, reference: ', 0, (320, 4, 3, 3), (36, 9, 3, 1), False, _var0)'")
+ print(f"etc, actual: '{etc}'\n")
 
  ## call below maps to: ('storage', HalfStorage, '0', 'cpu', 11520)
  self.ident, self.storage_type, self.obj_key, self.location, self.obj_size = storage.split(', ', 4)
@@ -201,10 +220,16 @@ def parse_instruction(self, instruction_string):
  self.obj_size = int(self.obj_size)
  self.storage_type = self._torch_to_numpy(self.storage_type)
 
+ if self.extra_debugging:
+ print(f"{self.ident}, {self.obj_key}, {self.location}, {self.obj_size}, {self.storage_type}")
+
  assert (self.ident == 'storage')
 
  garbage, etc = etc.split(', (', 1)
  # etc = 320, 4, 3, 3), (36, 9, 3, 1), False, _var0)
+ if self.extra_debugging:
+ print("etc, reference: '320, 4, 3, 3), (36, 9, 3, 1), False, _var0)'")
+ print(f"etc, actual: '{etc}'\n")
 
  size, stride, garbage = etc.split('), ', 2)
  # size = 320, 4, 3, 3
@@ -223,12 +248,19 @@ def parse_instruction(self, instruction_string):
  else:
  self.stride = tuple(map(int, stride.split(', ')))
 
+
+ if self.extra_debugging:
+ print(f"size: {self.size_tuple}, stride: {self.stride}")
+
  prod_size = prod(self.size_tuple)
  assert prod(self.size_tuple) == self.obj_size # does the size in the storage call match the size tuple
 
  # zero out the data
  self.data = np.zeros(self.size_tuple, dtype=self.storage_type)
 
+ def sayHi(self):
+ print(f"Hi, I'm an instance of LoadInstruction that will be used to load datafile {self.obj_key}")
+
  @staticmethod
  def _torch_to_numpy(storage_type):
  if storage_type == 'FloatStorage':