| 184 | return answer |
| 185 | |
| 186 | class MiniCPMV2_6: |
| 187 | def __init__(self, model_path, multi_gpus=False) -> None: |
| 188 | |
| 189 | print('torch_version:', torch.__version__) |
| 190 | if multi_gpus: # inference on multi-gpus |
| 191 | from accelerate import load_checkpoint_and_dispatch, init_empty_weights, infer_auto_device_map |
| 192 | with init_empty_weights(): |
| 193 | model = AutoModel.from_pretrained(model_path, trust_remote_code=True, |
| 194 | attn_implementation='sdpa', torch_dtype=torch.bfloat16) |
| 195 | |
| 196 | device_map = infer_auto_device_map(model, max_memory={0: "10GB", 1: "10GB"}, |
| 197 | no_split_module_classes=['SiglipVisionTransformer', 'Qwen2DecoderLayer']) |
| 198 | device_id = device_map["llm.model.embed_tokens"] |
| 199 | device_map["llm.lm_head"] = device_id # first and last layer of llm should be in the same device |
| 200 | device_map["vpm"] = device_id |
| 201 | device_map["resampler"] = device_id |
| 202 | device_id2 = device_map["llm.model.layers.26"] |
| 203 | device_map["llm.model.layers.8"] = device_id2 |
| 204 | device_map["llm.model.layers.9"] = device_id2 |
| 205 | device_map["llm.model.layers.10"] = device_id2 |
| 206 | device_map["llm.model.layers.11"] = device_id2 |
| 207 | device_map["llm.model.layers.12"] = device_id2 |
| 208 | device_map["llm.model.layers.13"] = device_id2 |
| 209 | device_map["llm.model.layers.14"] = device_id2 |
| 210 | device_map["llm.model.layers.15"] = device_id2 |
| 211 | device_map["llm.model.layers.16"] = device_id2 |
| 212 | print(device_map) |
| 213 | |
| 214 | self.model = load_checkpoint_and_dispatch(model, model_path, dtype=torch.bfloat16, device_map=device_map) |
| 215 | self.model.eval() |
| 216 | else: |
| 217 | self.model = AutoModel.from_pretrained(model_path, trust_remote_code=True, |
| 218 | attn_implementation='sdpa', torch_dtype=torch.bfloat16) |
| 219 | self.model.eval().cuda() |
| 220 | |
| 221 | self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) |
| 222 | |
| 223 | def chat(self, input): |
| 224 | image = None |
| 225 | if "image" in input and len(input["image"]) > 10: # legacy API |
| 226 | try: |
| 227 | image = Image.open(io.BytesIO(base64.b64decode(input['image']))).convert('RGB') |
| 228 | except Exception as e: |
| 229 | return "Image decode error" |
| 230 | |
| 231 | msgs = json.loads(input["question"]) |
| 232 | |
| 233 | for msg in msgs: |
| 234 | contents = msg.pop('content') # support str or List[Dict] |
| 235 | if isinstance(contents, str): |
| 236 | contents = [contents] |
| 237 | |
| 238 | new_cnts = [] |
| 239 | for c in contents: |
| 240 | if isinstance(c, dict): |
| 241 | if c['type'] == 'text': |
| 242 | c = c['pairs'] |
| 243 | elif c['type'] == 'image': |