| 572 | args.device, torch.float16)[0] |
| 573 | |
| 574 | class LlavaOnevisionVisionWrapper(torch.nn.Module): |
| 575 | |
| 576 | def __init__(self, vision_tower, projector, config): |
| 577 | super().__init__() |
| 578 | self.vision_tower = vision_tower |
| 579 | self.projector = projector |
| 580 | self.config = config |
| 581 | |
| 582 | def forward(self, pixel_values): |
| 583 | image_features = self.vision_tower(pixel_values, |
| 584 | output_hidden_states=True) |
| 585 | selected_image_feature = image_features.hidden_states[ |
| 586 | self.config.vision_feature_layer] |
| 587 | image_features = self.projector(selected_image_feature) |
| 588 | return image_features # (sigma(bs, patches_i), 729, c) |
| 589 | |
| 590 | model = LlavaOnevisionForConditionalGeneration.from_pretrained( |
| 591 | args.model_path, dtype=torch.float16) |
no outgoing calls
no test coverage detected