(args)
| 729 | |
| 730 | |
| 731 | def build_fuyu_engine(args): |
| 732 | processor = FuyuProcessor.from_pretrained(args.model_path) |
| 733 | raw_image = Image.new('RGB', [10, 10]) |
| 734 | image = processor(text="dummy", images=raw_image, |
| 735 | return_tensors="pt")['image_patches'][0].to( |
| 736 | args.device, torch.float16).unsqueeze(0) |
| 737 | |
| 738 | class FuyuEncoderWrapper(torch.nn.Module): |
| 739 | |
| 740 | def __init__(self, linear): |
| 741 | super().__init__() |
| 742 | self.linear = linear.to(torch.float16) |
| 743 | |
| 744 | def forward(self, patches): |
| 745 | return self.linear(patches).flatten(0, 1) |
| 746 | |
| 747 | model = FuyuForCausalLM.from_pretrained(args.model_path, |
| 748 | dtype=torch.float16) |
| 749 | |
| 750 | vision_encoder = model.vision_embed_tokens |
| 751 | wrapper = FuyuEncoderWrapper(vision_encoder).to(args.device) |
| 752 | |
| 753 | export_onnx(wrapper, |
| 754 | image, |
| 755 | f'{args.output_dir}/onnx', |
| 756 | dynamic_axes={'input': { |
| 757 | 0: 'batch', |
| 758 | 2: 'patch' |
| 759 | }}) |
| 760 | build_trt_engine( |
| 761 | args.model_type, |
| 762 | # [nImgs, nImgPatches, nDims] |
| 763 | # nImgs is always one since each query has exactly one image |
| 764 | # nImgPatches depends on image size (patch size: 30x30) |
| 765 | # nDims is 30x30x3=2700 (patch size x color channels) |
| 766 | [[1, 1, 2700], [1, 500, 2700], [1, 4096, 2700]], |
| 767 | f'{args.output_dir}/onnx', |
| 768 | args.output_dir, |
| 769 | args.max_batch_size) |
| 770 | |
| 771 | |
| 772 | def build_neva_engine(args): |
no test coverage detected