Generalized weight search path.

This commit is contained in:
Light 2023-04-13 21:43:32 +08:00
parent 97e67d136b
commit da74cd7c44

View File

@ -108,14 +108,10 @@ def load_quantized(model_name):
path_to_model = Path(f'{shared.args.model_dir}/{model_name}') path_to_model = Path(f'{shared.args.model_dir}/{model_name}')
pt_path = None pt_path = None
priority_name_list = [ priority_name_list = [
Path(f'{shared.args.model_dir}/{model_name}/{shared.args.wbits}bit-{shared.args.groupsize}g.safetensors'), Path(f'{shared.args.model_dir}/{model_name}{hyphen}{shared.args.wbits}bit{group}{ext}')
Path(f'{shared.args.model_dir}/{model_name}/{shared.args.wbits}bit-{shared.args.groupsize}g.pt'), for ext in ['.safetensors', '.pt']
Path(f'{shared.args.model_dir}/{model_name}/{shared.args.wbits}bit.safetensors'), for group in ([f'-{shared.args.groupsize}g', ''] if shared.args.groupsize > 0 else [''])
Path(f'{shared.args.model_dir}/{model_name}/{shared.args.wbits}bit.pt'), for hyphen in ['-', f'/{model_name}-', '/']
Path(f'{shared.args.model_dir}/{model_name}-{shared.args.wbits}bit-{shared.args.groupsize}g.safetensors'),
Path(f'{shared.args.model_dir}/{model_name}-{shared.args.wbits}bit-{shared.args.groupsize}g.pt'),
Path(f'{shared.args.model_dir}/{model_name}-{shared.args.wbits}bit.safetensors'),
Path(f'{shared.args.model_dir}/{model_name}-{shared.args.wbits}bit.pt'),
] ]
for path in priority_name_list: for path in priority_name_list:
if path.exists(): if path.exists():