(args)
| 356 | |
| 357 | |
| 358 | def build_blip2_engine(args): |
| 359 | processor = Blip2Processor.from_pretrained(args.model_path) |
| 360 | |
| 361 | raw_image = Image.new('RGB', [10, 10]) # dummy image |
| 362 | prompt = "Question: what is this? Answer:" |
| 363 | inputs = processor(raw_image, prompt, |
| 364 | return_tensors="pt").to(args.device, torch.float16) |
| 365 | image = inputs['pixel_values'] |
| 366 | |
| 367 | class Blip2VisionWrapper(torch.nn.Module): |
| 368 | |
| 369 | def __init__(self, vision_model, qformer, projector, query_tokens): |
| 370 | super().__init__() |
| 371 | self.vision_model = vision_model |
| 372 | self.qformer = qformer |
| 373 | self.projector = projector |
| 374 | self.query_tokens = query_tokens |
| 375 | |
| 376 | def forward(self, image): |
| 377 | features = self.vision_model(image)[0] |
| 378 | qformer_output = self.qformer(query_embeds=self.query_tokens, |
| 379 | encoder_hidden_states=features, |
| 380 | return_dict=True) |
| 381 | return self.projector(qformer_output.last_hidden_state) |
| 382 | |
| 383 | model = Blip2ForConditionalGeneration.from_pretrained(args.model_path, |
| 384 | dtype=torch.float16) |
| 385 | |
| 386 | blip2_llm = "" |
| 387 | if model.language_model.config.architectures[ |
| 388 | 0] == 'T5ForConditionalGeneration': |
| 389 | blip2_llm = "t5" |
| 390 | elif model.language_model.config.architectures[0] == 'OPTForCausalLM': |
| 391 | blip2_llm = "opt" |
| 392 | |
| 393 | wrapper = Blip2VisionWrapper(model.vision_model, model.qformer, |
| 394 | model.language_projection, model.query_tokens) |
| 395 | wrapper.to(args.device) |
| 396 | |
| 397 | export_onnx(wrapper, image, f'{args.output_dir}/onnx') |
| 398 | build_trt_engine( |
| 399 | args.model_type + "-" + blip2_llm, # blip2-t5 or blip2-opt |
| 400 | [image.shape[1], image.shape[2], image.shape[3]], # [3, H, W] |
| 401 | f'{args.output_dir}/onnx', |
| 402 | args.output_dir, |
| 403 | args.max_batch_size) |
| 404 | |
| 405 | |
| 406 | def build_interlm_xcomposer2_engine(args): |
no test coverage detected