rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

possible fix for cache dl thing

Browse Source
This commit is contained in:
somebody
2023-04-14 20:25:03 -05:00
parent 334c09606b
commit 38c53191d3
1 changed files with 12 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
modeling/inference_models/generic_hf_torch.py
View File

@@ -171,16 +171,19 @@ class GenericHFTorchInferenceModel(HFTorchInferenceModel):
if utils.num_shards is None:
# Save the pytorch_model.bin or model.safetensors of an unsharded model
for possible_weight_name in [
any_success = False
possible_checkpoint_names = [
transformers.modeling_utils.WEIGHTS_NAME,
"model.safetensors",
]:
]
for possible_checkpoint_name in possible_checkpoint_names:
try:
shutil.move(
os.path.realpath(
huggingface_hub.hf_hub_download(
self.model_name,
possible_weight_name,
possible_checkpoint_name,
revision=utils.koboldai_vars.revision,
cache_dir="cache",
local_files_only=True,
@@ -191,12 +194,15 @@ class GenericHFTorchInferenceModel(HFTorchInferenceModel):
self.get_local_model_path(
ignore_existance=True
),
possible_weight_name,
possible_checkpoint_name,
),
)
any_success = True
except Exception:
if possible_weight_name == "model.safetensors":
raise
pass
if not any_success:
raise RuntimeError(f"Couldn't find any of {possible_checkpoint_names} in cache for {self.model_name} @ '{utils.koboldai_vars.revisison}'")
else:
# Handle saving sharded models