logger.info(f"Loading {model_name}...")

t0 = time.time()

shared.is_seq2seq = False

load_func_map = {

'Transformers': huggingface_loader,

'AutoGPTQ': AutoGPTQ_loader,

'GPTQ-for-LLaMa': GPTQ_loader,

'llama.cpp': llamacpp_loader,

'llamacpp_HF': llamacpp_HF_loader,

'RWKV': RWKV_loader,

'ExLlama': ExLlama_loader,

'ExLlama_HF': ExLlama_HF_loader,

'ExLlamav2': ExLlamav2_loader,

'ExLlamav2_HF': ExLlamav2_HF_loader,

'ctransformers': ctransformers_loader,

'AutoAWQ': AutoAWQ_loader,

}

if loader is None:

if shared.args.loader is not None:

loader = shared.args.loader

else:

loader = get_model_metadata(model_name)['loader']

if loader is None:

logger.error('The path to the model does not exist. Exiting.')

return None, None

shared.args.loader = loader

output = load_func_map[loader](model_name)

if type(output) is tuple:

model, tokenizer = output

else:

model = output

if model is None:

return None, None

else:

tokenizer = load_tokenizer(model_name, model)

# Hijack attention with xformers

if any((shared.args.xformers, shared.args.sdp_attention)):

llama_attn_hijack.hijack_llama_attention()

logger.info(f"Loaded the model in {(time.time()-t0):.2f} seconds.")

tokenizer = None

path_to_model = Path(f"{shared.args.model_dir}/{model_name}/")

if any(s in model_name.lower() for s in ['gpt-4chan', 'gpt4chan']) and Path(f"{shared.args.model_dir}/gpt-j-6B/").exists():

tokenizer = AutoTokenizer.from_pretrained(Path(f"{shared.args.model_dir}/gpt-j-6B/"))

elif path_to_model.exists():

if shared.args.use_fast:

logger.info('Loading the tokenizer with use_fast=True.')

tokenizer = AutoTokenizer.from_pretrained(

path_to_model,

trust_remote_code=shared.args.trust_remote_code,

use_fast=shared.args.use_fast

)

path_to_model = Path(f'{shared.args.model_dir}/{model_name}')

params = {

'low_cpu_mem_usage': True,

'trust_remote_code': shared.args.trust_remote_code,

'torch_dtype': torch.bfloat16 if shared.args.bf16 else torch.float16

}

config = AutoConfig.from_pretrained(path_to_model, trust_remote_code=params['trust_remote_code'])

if 'chatglm' in model_name.lower():

LoaderClass = AutoModel

else:

if config.to_dict().get('is_encoder_decoder', False):

LoaderClass = AutoModelForSeq2SeqLM

shared.is_seq2seq = True

else:

LoaderClass = AutoModelForCausalLM

# Load the model in simple 16-bit mode by default

if not any([shared.args.cpu, shared.args.load_in_8bit, shared.args.load_in_4bit, shared.args.auto_devices, shared.args.disk, shared.args.deepspeed, shared.args.gpu_memory is not None, shared.args.cpu_memory is not None, shared.args.compress_pos_emb > 1, shared.args.alpha_value > 1, shared.args.disable_exllama]):

model = LoaderClass.from_pretrained(path_to_model, **params)

if torch.backends.mps.is_available():

device = torch.device('mps')

model = model.to(device)

elif is_xpu_available():

device = torch.device("xpu")

model = model.to(device)

else:

model = model.cuda()

# DeepSpeed ZeRO-3

elif shared.args.deepspeed:

model = LoaderClass.from_pretrained(path_to_model, torch_dtype=params['torch_dtype'])

model = deepspeed.initialize(model=model, config_params=ds_config, model_parameters=None, optimizer=None, lr_scheduler=None)[0]

model.module.eval() # Inference

logger.info(f'DeepSpeed ZeRO-3 is enabled: {is_deepspeed_zero3_enabled()}')

# Load with quantization and/or offloading

else:

if not any((shared.args.cpu, torch.cuda.is_available(), is_xpu_available(), torch.backends.mps.is_available())):

logger.warning('torch.cuda.is_available() and is_xpu_available() returned False. This means that no GPU has been detected. Falling back to CPU mode.')

shared.args.cpu = True

if shared.args.cpu:

params['torch_dtype'] = torch.float32

else:

params['device_map'] = 'auto'

params['max_memory'] = get_max_memory_dict()

if shared.args.load_in_4bit:

# See https://github.com/huggingface/transformers/pull/23479/files

# and https://huggingface.co/blog/4bit-transformers-bitsandbytes

quantization_config_params = {

'load_in_4bit': True,

'bnb_4bit_compute_dtype': eval("torch.{}".format(shared.args.compute_dtype)) if shared.args.compute_dtype in ["bfloat16", "float16", "float32"] else None,

'bnb_4bit_quant_type': shared.args.quant_type,

'bnb_4bit_use_double_quant': shared.args.use_double_quant,

}

logger.info('Using the following 4-bit params: ' + str(quantization_config_params))

params['quantization_config'] = BitsAndBytesConfig(**quantization_config_params)

elif shared.args.load_in_8bit:

if any((shared.args.auto_devices, shared.args.gpu_memory)):

params['quantization_config'] = BitsAndBytesConfig(load_in_8bit=True, llm_int8_enable_fp32_cpu_offload=True)

else:

params['quantization_config'] = BitsAndBytesConfig(load_in_8bit=True)

if params['max_memory'] is not None:

with init_empty_weights():

model = LoaderClass.from_config(config, trust_remote_code=params['trust_remote_code'])

model.tie_weights()

params['device_map'] = infer_auto_device_map(

model,

dtype=torch.int8,

max_memory=params['max_memory'],

no_split_module_classes=model._no_split_modules

)

if shared.args.disk:

params['offload_folder'] = shared.args.disk_cache_dir

if shared.args.disable_exllama:

try:

gptq_config = GPTQConfig(bits=config.quantization_config.get('bits', 4), disable_exllama=True)

params['quantization_config'] = gptq_config

logger.info('Loading with ExLlama kernel disabled.')

except:

exc = traceback.format_exc()

logger.error('Failed to disable exllama. Does the config.json for this model contain the necessary quantization info?')

print(exc)

if shared.args.compress_pos_emb > 1:

params['rope_scaling'] = {'type': 'linear', 'factor': shared.args.compress_pos_emb}

elif shared.args.alpha_value > 1:

params['rope_scaling'] = {'type': 'dynamic', 'factor': RoPE.get_alpha_value(shared.args.alpha_value, shared.args.rope_freq_base)}

model = LoaderClass.from_pretrained(path_to_model, **params)

from modules.llamacpp_model import LlamaCppModel

path = Path(f'{shared.args.model_dir}/{model_name}')

if path.is_file():

model_file = path

else:

model_file = list(Path(f'{shared.args.model_dir}/{model_name}').glob('*.gguf'))[0]

logger.info(f"llama.cpp weights detected: {model_file}")

model, tokenizer = LlamaCppModel.from_pretrained(model_file)

from modules.llamacpp_hf import LlamacppHF

for fname in [model_name, "oobabooga_llama-tokenizer", "llama-tokenizer"]:

path = Path(f'{shared.args.model_dir}/{fname}')

if all((path / file).exists() for file in ['tokenizer_config.json', 'special_tokens_map.json', 'tokenizer.model']):

logger.info(f'Using tokenizer from: {path}')

break

else:

logger.error("Could not load the model because a tokenizer in transformers format was not found. Please download oobabooga/llama-tokenizer.")

return None, None

if shared.args.use_fast:

logger.info('Loading the tokenizer with use_fast=True.')

tokenizer = AutoTokenizer.from_pretrained(

path,

trust_remote_code=shared.args.trust_remote_code,

use_fast=shared.args.use_fast

)

model = LlamacppHF.from_pretrained(model_name)

from modules.ctransformers_model import CtransformersModel

path = Path(f'{shared.args.model_dir}/{model_name}')

ctrans = CtransformersModel()

if ctrans.model_type_is_auto():

model_file = path

else:

if path.is_file():

model_file = path

else:

entries = Path(f'{shared.args.model_dir}/{model_name}')

gguf = list(entries.glob('*.gguf'))

bin = list(entries.glob('*.bin'))

if len(gguf) > 0:

model_file = gguf[0]

elif len(bin) > 0:

model_file = bin[0]

else:

logger.error("Could not find a model for ctransformers.")

return None, None

logger.info(f'ctransformers weights detected: {model_file}')

model, tokenizer = ctrans.from_pretrained(model_file)

from awq import AutoAWQForCausalLM

model_dir = Path(f'{shared.args.model_dir}/{model_name}')

model = AutoAWQForCausalLM.from_quantized(

quant_path=model_dir,

max_new_tokens=shared.args.max_seq_len,

trust_remote_code=shared.args.trust_remote_code,

fuse_layers=not shared.args.no_inject_fused_attention,

max_memory=get_max_memory_dict(),

batch_size=1,

safetensors=any(model_dir.glob('*.safetensors')),

)

# Monkey patch

if shared.args.monkey_patch:

logger.warning("Applying the monkey patch for using LoRAs with GPTQ models. It may cause undefined behavior outside its intended scope.")

from modules.monkey_patch_gptq_lora import load_model_llama

model, _ = load_model_llama(model_name)

# No monkey patch

else:

import modules.GPTQ_loader

model = modules.GPTQ_loader.load_quantized(model_name)

from modules.exllama import ExllamaModel

model, tokenizer = ExllamaModel.from_pretrained(model_name)

from modules.exllamav2 import Exllamav2Model

model, tokenizer = Exllamav2Model.from_pretrained(model_name)

'''

from modules.RWKV import RWKVModel, RWKVTokenizer

model = RWKVModel.from_pretrained(

Path(f'{shared.args.model_dir}/{model_name}'),

dtype="fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16",

device="cpu" if shared.args.cpu else "xpu" if is_xpu_available() else "cuda"

)

tokenizer = RWKVTokenizer.from_pretrained(Path(shared.args.model_dir))

max_memory = {}

if shared.args.gpu_memory:

memory_map = list(map(lambda x: x.strip(), shared.args.gpu_memory))

for i in range(len(memory_map)):

max_memory[i] = f'{memory_map[i]}GiB' if not re.match('.*ib$', memory_map[i].lower()) else memory_map[i]

max_cpu_memory = shared.args.cpu_memory.strip() if shared.args.cpu_memory is not None else '99GiB'

max_memory['cpu'] = f'{max_cpu_memory}GiB' if not re.match('.*ib$', max_cpu_memory.lower()) else max_cpu_memory

# If --auto-devices is provided standalone, try to get a reasonable value

# for the maximum memory of device :0

elif shared.args.auto_devices:

if is_xpu_available():

total_mem = (torch.xpu.get_device_properties(0).total_memory / (1024 * 1024))

else:

total_mem = (torch.cuda.get_device_properties(0).total_memory / (1024 * 1024))

suggestion = round((total_mem - 1000) / 1000) * 1000

if total_mem - suggestion < 800:

suggestion -= 1000

suggestion = int(round(suggestion / 1000))

logger.warning(f"Auto-assiging --gpu-memory {suggestion} for your GPU to try to prevent out-of-memory errors. You can manually set other values.")

max_memory = {0: f'{suggestion}GiB', 'cpu': f'{shared.args.cpu_memory or 99}GiB'}

gc.collect()

if not shared.args.cpu:

if is_xpu_available():

torch.xpu.empty_cache()

else:

shared.model = shared.tokenizer = None

shared.lora_names = []

shared.model_dirty_from_training = False