(
model_path, device, num_gpus, max_gpu_memory=None, load_8bit=False, lora_weights=None
)
| 130 | return gpu_memory |
| 131 | |
| 132 | def load_model( |
| 133 | model_path, device, num_gpus, max_gpu_memory=None, load_8bit=False, lora_weights=None |
| 134 | ): |
| 135 | if device == "cpu": |
| 136 | kwargs = {} |
| 137 | elif device == "cuda": |
| 138 | kwargs = {"torch_dtype": torch.float16} |
| 139 | if num_gpus == "auto": |
| 140 | kwargs["device_map"] = "auto" |
| 141 | else: |
| 142 | num_gpus = int(num_gpus) |
| 143 | if num_gpus != 1: |
| 144 | kwargs["device_map"] = "auto" |
| 145 | if max_gpu_memory is None: |
| 146 | kwargs[ |
| 147 | "device_map" |
| 148 | ] = "sequential" # This is important for not the same VRAM sizes |
| 149 | available_gpu_memory = get_gpu_memory(num_gpus) |
| 150 | kwargs["max_memory"] = { |
| 151 | i: str(int(available_gpu_memory[i] * 0.85)) + "GiB" |
| 152 | for i in range(num_gpus) |
| 153 | } |
| 154 | else: |
| 155 | kwargs["max_memory"] = {i: max_gpu_memory for i in range(num_gpus)} |
| 156 | else: |
| 157 | raise ValueError(f"Invalid device: {device}") |
| 158 | |
| 159 | tokenizer = LlamaTokenizer.from_pretrained( |
| 160 | model_path, use_fast=False) |
| 161 | |
| 162 | if lora_weights is None: |
| 163 | model = HuskyForConditionalGeneration.from_pretrained( |
| 164 | model_path, low_cpu_mem_usage=True, **kwargs |
| 165 | ) |
| 166 | else: |
| 167 | kwargs["device_map"] = "auto" |
| 168 | model = HuskyForConditionalGeneration.from_pretrained( |
| 169 | model_path, low_cpu_mem_usage=True, **kwargs |
| 170 | ) |
| 171 | model.language_model = PeftModel.from_pretrained( |
| 172 | model.language_model, |
| 173 | lora_weights, |
| 174 | **kwargs |
| 175 | ) |
| 176 | |
| 177 | if load_8bit: |
| 178 | compress_module(model, device) |
| 179 | |
| 180 | if (device == "cuda" and num_gpus == 1) or device == "mps": |
| 181 | model.to(device) |
| 182 | |
| 183 | model = model.eval() |
| 184 | return model, tokenizer |
| 185 | |
| 186 | def read_image(image_file): |
| 187 | if image_file.startswith('http') or image_file.startswith('https'): |
no test coverage detected