2023-05-17 10:12:12 -04:00
|
|
|
import logging
|
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
|
|
|
|
from auto_gptq import AutoGPTQForCausalLM
|
|
|
|
|
|
|
|
|
|
import modules.shared as shared
|
|
|
|
|
from modules.models import get_max_memory_dict
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def load_quantized(model_name):
|
|
|
|
|
path_to_model = Path(f'{shared.args.model_dir}/{model_name}')
|
|
|
|
|
pt_path = None
|
|
|
|
|
use_safetensors = False
|
|
|
|
|
|
|
|
|
|
# Find the model checkpoint
|
|
|
|
|
found_pts = list(path_to_model.glob("*.pt"))
|
|
|
|
|
found_safetensors = list(path_to_model.glob("*.safetensors"))
|
|
|
|
|
if len(found_safetensors) > 0:
|
2023-05-17 10:16:37 -04:00
|
|
|
if len(found_safetensors) > 1:
|
2023-05-17 10:12:12 -04:00
|
|
|
logging.warning('More than one .safetensors model has been found. The last one will be selected. It could be wrong.')
|
|
|
|
|
|
|
|
|
|
use_safetensors = True
|
|
|
|
|
pt_path = found_safetensors[-1]
|
|
|
|
|
elif len(found_pts) > 0:
|
|
|
|
|
if len(found_pts) > 1:
|
|
|
|
|
logging.warning('More than one .pt model has been found. The last one will be selected. It could be wrong.')
|
|
|
|
|
|
|
|
|
|
pt_path = found_pts[-1]
|
|
|
|
|
|
|
|
|
|
# Define the params for AutoGPTQForCausalLM.from_quantized
|
|
|
|
|
params = {
|
|
|
|
|
'model_basename': pt_path.stem,
|
|
|
|
|
'device': "cuda:0" if not shared.args.cpu else "cpu",
|
|
|
|
|
'use_triton': shared.args.triton,
|
|
|
|
|
'use_safetensors': use_safetensors,
|
|
|
|
|
'max_memory': get_max_memory_dict()
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
logging.warning(f"The AutoGPTQ params are: {params}")
|
|
|
|
|
model = AutoGPTQForCausalLM.from_quantized(path_to_model, **params)
|
|
|
|
|
return model
|
