@@ -26,12 +26,12 @@ def extract_weights_from_checkpoint(fb0):
2626 with myzip .open (folder_name + f'/data/{ load_instruction .obj_key } ' ) as myfile :
2727 if (load_instruction .load_from_file_buffer (myfile )):
2828 torch_weights ['state_dict' ][sd_key ] = load_instruction .get_data ()
29-
3029 return torch_weights
3130
3231def examine_pickle (fb0 ):
3332
3433 decompiled = unparse (Pickled .load (fb0 ).ast ).splitlines ()
34+
3535## LINES WE CARE ABOUT:
3636## 1: this defines a data file and what kind of data is in it
3737## _var1 = _rebuild_tensor_v2(UNPICKLER.persistent_load(('storage', HalfStorage, '0', 'cpu', 11520)), 0, (320, 4, 3, 3), (36, 9, 3, 1), False, _var0)
@@ -52,7 +52,7 @@ def examine_pickle(fb0):
5252 assign_instructions = AssignInstructions ()
5353
5454 for line in decompiled :
55- ## see if line matches pattern of var =
55+ ## see if line matches patterns of lines we care about:
5656 line = line .strip ()
5757 if re_rebuild .match (line ):
5858 variable_name , load_instruction = line .split (' = ' , 1 )
@@ -69,10 +69,6 @@ def examine_pickle(fb0):
6969 assign_instructions .integrate (load_instructions )
7070
7171 return assign_instructions .integrated_instructions
72-
73-
74- #output = {}
75- #output['state_dict'] = {}
7672
7773class AssignInstructions :
7874 def __init__ (self ):
@@ -122,13 +118,15 @@ def __init__(self, instruction_string):
122118 self .ident = False
123119 self .storage_type = False
124120 self .obj_key = False
125- self .location = False
121+ self .location = False #unused
126122 self .obj_size = False
127- self .stride = False #args[3] -- unused, I think
123+ self .stride = False #unused
128124 self .data = False ;
129125 self .parse_instruction (instruction_string )
130126
131127 def parse_instruction (self , instruction_string ):
128+ ## this function could probably be cleaned up/shortened.
129+
132130 ## this is the API def for _rebuild_tensor_v2:
133131 ## _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks):
134132 #
@@ -145,7 +143,6 @@ def parse_instruction(self, instruction_string):
145143 # etc = 0, (320, 4, 3, 3), (36, 9, 3, 1), False, _var0)
146144
147145 ## call below maps to: ('storage', HalfStorage, '0', 'cpu', 11520)
148- # ident, storage_type, obj_key, location, obj_size = args[0][0:5]
149146 self .ident , self .storage_type , self .obj_key , self .location , self .obj_size = storage .split (', ' , 4 )
150147
151148 self .ident = self .ident .strip ("'" )
@@ -165,8 +162,8 @@ def parse_instruction(self, instruction_string):
165162 stride = stride .strip ('(,' )
166163 size = size .strip (',' )
167164
168-
169165 if (size == '' ):
166+ # rare case where there is an empty tuple. SDv1.4 has two of these.
170167 self .size_tuple = ()
171168 else :
172169 self .size_tuple = tuple (map (int , size .split (', ' )))
@@ -194,24 +191,14 @@ def _torch_to_numpy(storage_type):
194191 return np .int32
195192 raise Exception ("Storage type not defined!" )
196193
197-
198194 def load_from_file_buffer (self , fb ):
199195 if self .data .dtype == "object" :
200196 print (f"issue assigning object on { self .obj_key } " )
201197 return False
202198 else :
203- #key_prelookup[obj_key] = (storage_type, obj_size, ret, args[2], args[3])
204- #maps to: where v is the right side of the above assignment
205- #np.copyto(v[2], np.frombuffer(myfile.read(), v[2].dtype).reshape(v[3]))
206- #print(f"np.copyto(self.data, np.frombuffer(fb.read(), {self.data.dtype}).reshape({self.size_tuple}))")
207199 np .copyto (self .data , np .frombuffer (fb .read (), self .data .dtype ).reshape (self .size_tuple ))
208200 return True
209201
210202 def get_data (self ):
211203 return self .data
212204
213-
214-
215- #examine_pickle(open('classicanimation.archive/data.pkl', "rb"))
216-
217-
0 commit comments