@@ -26,12 +26,12 @@ def extract_weights_from_checkpoint(fb0):
26
26
with myzip .open (folder_name + f'/data/{ load_instruction .obj_key } ' ) as myfile :
27
27
if (load_instruction .load_from_file_buffer (myfile )):
28
28
torch_weights ['state_dict' ][sd_key ] = load_instruction .get_data ()
29
-
30
29
return torch_weights
31
30
32
31
def examine_pickle (fb0 ):
33
32
34
33
decompiled = unparse (Pickled .load (fb0 ).ast ).splitlines ()
34
+
35
35
## LINES WE CARE ABOUT:
36
36
## 1: this defines a data file and what kind of data is in it
37
37
## _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):
52
52
assign_instructions = AssignInstructions ()
53
53
54
54
for line in decompiled :
55
- ## see if line matches pattern of var =
55
+ ## see if line matches patterns of lines we care about:
56
56
line = line .strip ()
57
57
if re_rebuild .match (line ):
58
58
variable_name , load_instruction = line .split (' = ' , 1 )
@@ -69,10 +69,6 @@ def examine_pickle(fb0):
69
69
assign_instructions .integrate (load_instructions )
70
70
71
71
return assign_instructions .integrated_instructions
72
-
73
-
74
- #output = {}
75
- #output['state_dict'] = {}
76
72
77
73
class AssignInstructions :
78
74
def __init__ (self ):
@@ -122,13 +118,15 @@ def __init__(self, instruction_string):
122
118
self .ident = False
123
119
self .storage_type = False
124
120
self .obj_key = False
125
- self .location = False
121
+ self .location = False #unused
126
122
self .obj_size = False
127
- self .stride = False #args[3] -- unused, I think
123
+ self .stride = False #unused
128
124
self .data = False ;
129
125
self .parse_instruction (instruction_string )
130
126
131
127
def parse_instruction (self , instruction_string ):
128
+ ## this function could probably be cleaned up/shortened.
129
+
132
130
## this is the API def for _rebuild_tensor_v2:
133
131
## _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks):
134
132
#
@@ -145,7 +143,6 @@ def parse_instruction(self, instruction_string):
145
143
# etc = 0, (320, 4, 3, 3), (36, 9, 3, 1), False, _var0)
146
144
147
145
## call below maps to: ('storage', HalfStorage, '0', 'cpu', 11520)
148
- # ident, storage_type, obj_key, location, obj_size = args[0][0:5]
149
146
self .ident , self .storage_type , self .obj_key , self .location , self .obj_size = storage .split (', ' , 4 )
150
147
151
148
self .ident = self .ident .strip ("'" )
@@ -165,8 +162,8 @@ def parse_instruction(self, instruction_string):
165
162
stride = stride .strip ('(,' )
166
163
size = size .strip (',' )
167
164
168
-
169
165
if (size == '' ):
166
+ # rare case where there is an empty tuple. SDv1.4 has two of these.
170
167
self .size_tuple = ()
171
168
else :
172
169
self .size_tuple = tuple (map (int , size .split (', ' )))
@@ -194,24 +191,14 @@ def _torch_to_numpy(storage_type):
194
191
return np .int32
195
192
raise Exception ("Storage type not defined!" )
196
193
197
-
198
194
def load_from_file_buffer (self , fb ):
199
195
if self .data .dtype == "object" :
200
196
print (f"issue assigning object on { self .obj_key } " )
201
197
return False
202
198
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}))")
207
199
np .copyto (self .data , np .frombuffer (fb .read (), self .data .dtype ).reshape (self .size_tuple ))
208
200
return True
209
201
210
202
def get_data (self ):
211
203
return self .data
212
204
213
-
214
-
215
- #examine_pickle(open('classicanimation.archive/data.pkl', "rb"))
216
-
217
-
0 commit comments