Views
No views yet
1import torch
2from PIL import Image
3from transformers import AutoModel, CLIPImageProcessor
4
5class CustomModel(nn.Module):
6 def __init__(self, base_model, num_classes=2):
7 super(CustomModel, self).__init__()
8 self.base_model = base_model
9 self.classifier = nn.Linear(base_model.config.hidden_size, num_classes).to(torch.bfloat16)
10
11 def forward(self, x):
12 outputs = self.base_model(x)
13 pooled_output = outputs.pooler_output
14 logits = self.classifier(pooled_output)
15 return logits
16
17base_model = AutoModel.from_pretrained(
18 'OpenGVLab/InternViT-6B-448px-V1-5',
19 torch_dtype=torch.bfloat16,
20 low_cpu_mem_usage=True,
21 trust_remote_code=True).cuda().eval()
22
23model = CustomModel(base_model, num_classes=2).to(device).eval()
24model.classifier.load_state_dict(torch.load("checkpoints/classifier_weights.pth"))
25
26image = Image.open('./examples/image1.jpg').convert('RGB')
27
28image_processor = CLIPImageProcessor.from_pretrained('OpenGVLab/InternViT-6B-448px-V1-5')
29
30pixel_values = image_processor(images=image, return_tensors='pt').pixel_values.to(torch.bfloat16).cuda()
31
32with torch.no_grad():
33 outputs = model(pixel_values)
34