| 994 | image = image.flatten(0, 1) |
| 995 | |
| 996 | class Phi3VisionWrapper(torch.nn.Module): |
| 997 | |
| 998 | def __init__(self, vision_model): |
| 999 | super().__init__() |
| 1000 | self.vision_model = vision_model |
| 1001 | |
| 1002 | def forward(self, pixel_values): |
| 1003 | return self.vision_model.get_img_features(pixel_values).reshape( |
| 1004 | 1, pixel_values.shape[0], -1, self.vision_model.image_dim_out) |
| 1005 | |
| 1006 | model = AutoModelForCausalLM.from_pretrained(args.model_path, |
| 1007 | dtype=torch.float16, |