| 517 | args.device, torch.float16)[0] |
| 518 | |
| 519 | class LlavaNextVisionWrapper(torch.nn.Module): |
| 520 | |
| 521 | def __init__(self, vision_tower, projector): |
| 522 | super().__init__() |
| 523 | self.vision_tower = vision_tower |
| 524 | self.projector = projector |
| 525 | |
| 526 | def forward(self, pixel_values): |
| 527 | image_features = self.vision_tower(pixel_values, |
| 528 | output_hidden_states=True) |
| 529 | selected_image_feature = image_features.hidden_states[-2][:, 1:] |
| 530 | image_features = self.projector(selected_image_feature) |
| 531 | return image_features # (bs, 576, c) |
| 532 | |
| 533 | hf_config = AutoConfig.from_pretrained(args.model_path) |
| 534 | hf_config.vision_config._attn_implementation = "eager" |
no outgoing calls
no test coverage detected