forestcalled/text-generation-webui

1from pathlib import Path2 3from accelerate.utils import is_xpu_available4from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig5 6import modules.shared as shared7from modules.logging_colors import logger8from modules.models import get_max_memory_dict9 10 11def load_quantized(model_name):12    path_to_model = Path(f'{shared.args.model_dir}/{model_name}')13    pt_path = None14 15    # Find the model checkpoint16    if shared.args.checkpoint:17        pt_path = Path(shared.args.checkpoint)18    else:19        for ext in ['.safetensors', '.pt', '.bin']:20            found = list(path_to_model.glob(f"*{ext}"))21            if len(found) > 0:22                if len(found) > 1:23                    logger.warning(f'More than one {ext} model has been found. The last one will be selected. It could be wrong.')24 25                pt_path = found[-1]26                break27 28    if pt_path is None:29        logger.error("The model could not be loaded because its checkpoint file in .bin/.pt/.safetensors format could not be located.")30        return31 32    use_safetensors = pt_path.suffix == '.safetensors'33    if not (path_to_model / "quantize_config.json").exists():34        quantize_config = BaseQuantizeConfig(35            bits=bits if (bits := shared.args.wbits) > 0 else 4,36            group_size=gs if (gs := shared.args.groupsize) > 0 else -1,37            desc_act=shared.args.desc_act38        )39    else:40        quantize_config = None41 42    # Define the params for AutoGPTQForCausalLM.from_quantized43    params = {44        'model_basename': pt_path.stem,45        'device': "xpu:0" if is_xpu_available() else "cuda:0" if not shared.args.cpu else "cpu",46        'use_triton': shared.args.triton,47        'inject_fused_attention': not shared.args.no_inject_fused_attention,48        'inject_fused_mlp': not shared.args.no_inject_fused_mlp,49        'use_safetensors': use_safetensors,50        'trust_remote_code': shared.args.trust_remote_code,51        'max_memory': get_max_memory_dict(),52        'quantize_config': quantize_config,53        'use_cuda_fp16': not shared.args.no_use_cuda_fp16,54        'disable_exllama': shared.args.disable_exllama,55        'disable_exllamav2': shared.args.disable_exllamav2,56    }57 58    logger.info(f"The AutoGPTQ params are: {params}")59    model = AutoGPTQForCausalLM.from_quantized(path_to_model, **params)60 61    # These lines fix the multimodal extension when used with AutoGPTQ62    if hasattr(model, 'model'):63        if not hasattr(model, 'dtype'):64            if hasattr(model.model, 'dtype'):65                model.dtype = model.model.dtype66 67        if hasattr(model.model, 'model') and hasattr(model.model.model, 'embed_tokens'):68            if not hasattr(model, 'embed_tokens'):69                model.embed_tokens = model.model.model.embed_tokens70 71            if not hasattr(model.model, 'embed_tokens'):72                model.model.embed_tokens = model.model.model.embed_tokens73 74    return model75