Merge branch 'restricted-unpickler' into no-pickle-importer

Author: CrudeDiatribe

backends/model_converter/fake_torch.py

@@ -15,8 +15,16 @@ def __new__(cls, *args):
       assert ident == 'storage'
 
       assert prod(args[2]) == obj_size
+
+
       ret = np.zeros(args[2], dtype=storage_type)
-      key_prelookup[obj_key] = (storage_type, obj_size, ret, args[2], args[3])
+      if obj_key not in key_prelookup:
+        key_prelookup[obj_key] = []
+
+      key_prelookup[obj_key].append((storage_type, obj_size, ret, args[2], args[3]))
+
+      #print(f"File: {obj_key}, references: {len(key_prelookup[obj_key])}, size: {args[2]}, storage_type: {storage_type}")
+
       return ret
 
   class HackParameter:
@@ -62,12 +70,17 @@ def fake_torch_load_zipped(fb0, load_weights=True):
     with myzip.open(folder_name+'/data.pkl') as myfile:
       ret = my_unpickle(myfile)
     if load_weights:
-      for k,v in ret[1].items():
+      for k, v_arr in ret[1].items():
         with myzip.open(folder_name + f'/data/{k}') as myfile:
-          if v[2].dtype == "object":
-            print(f"issue assigning object on {k}")
-            continue
-          np.copyto(v[2], np.frombuffer(myfile.read(), v[2].dtype).reshape(v[3]))
+          #print(f"Eating data file {k} now")
+          file_data = myfile.read()
+          for v in v_arr:
+            if v[2].dtype == "object":
+              print(f"issue assigning object on {k}")
+              continue
+            #weight = np.frombuffer(file_data, v[2].dtype).reshape(v[3])
+            #np.copyto(v[2], weight)
+            np.copyto(v[2], np.frombuffer(file_data, v[2].dtype).reshape(v[3]))
   return ret[0]
 
 def fake_torch_load(b0):