rixty
/
KoboldAI-Client
mirror of https://github.com/KoboldAI/KoboldAI-Client.git
1
1
Fork 0
Code Issues Projects Releases Wiki Activity

New pytorch zipfile support

This commit is contained in:
Henk 2022-12-10 19:11:07 +01:00
parent dd7363548c
commit 9a8e8a0005
2 changed files with 10 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
aiserver.py
View File

@ -2380,6 +2380,7 @@ def load_model(use_gpu=True, gpu_layers=None, disk_layers=None, initial_load=Fal
with zipfile.ZipFile(f, "r") as z:
try:
last_storage_key = None
zipfolder = os.path.basename(os.path.normpath(f)).split('.')[0]
f = None
current_offset = 0
able_to_pin_layers = True
@ -2391,7 +2392,10 @@ def load_model(use_gpu=True, gpu_layers=None, disk_layers=None, initial_load=Fal
last_storage_key = storage_key
if isinstance(f, zipfile.ZipExtFile):
f.close()
f = z.open(f"archive/data/{storage_key}")
try:
f = z.open(f"archive/data/{storage_key}")
except:
f = z.open(f"%s/data/{storage_key}" % zipfolder)
current_offset = 0
if current_offset != model_dict[key].seek_offset:
f.read(model_dict[key].seek_offset - current_offset)

6
tpu_mtj_backend.py
View File

@ -1251,6 +1251,7 @@ def load_model(path: str, driver_version="tpu_driver0.1_dev20210607", hf_checkpo
with zipfile.ZipFile(f, "r") as z:
try:
last_storage_key = None
zipfolder = os.path.basename(os.path.normpath(f)).split('.')[0]
f = None
current_offset = 0
if utils.current_shard == 0:
@ -1283,7 +1284,10 @@ def load_model(path: str, driver_version="tpu_driver0.1_dev20210607", hf_checkpo
last_storage_key = storage_key
if isinstance(f, zipfile.ZipExtFile):
f.close()
f = z.open(f"archive/data/{storage_key}")
try:
f = z.open(f"archive/data/{storage_key}")
except:
f = z.open(f"%s/data/{storage_key}" % zipfolder)
current_offset = 0
if current_offset != model_dict[key].seek_offset:
f.read(model_dict[key].seek_offset - current_offset)