--- license: apache-2.0 --- アニメイラストのセリフや擬音を検出するモデルです ベースモデルにInternViT-6B-448px-V1-5を使用しています。 https://huggingface.co/OpenGVLab/InternViT-6B-448px-V1-5を使用しています。 ベースモデルのpooler_output層にこんな感じに繋げば使えると思います。 ```python import torch from PIL import Image from transformers import AutoModel, CLIPImageProcessor class CustomModel(nn.Module): def __init__(self, base_model, num_classes=2): super(CustomModel, self).__init__() self.base_model = base_model self.classifier = nn.Linear(base_model.config.hidden_size, num_classes).to(torch.bfloat16) def forward(self, x): outputs = self.base_model(x) pooled_output = outputs.pooler_output logits = self.classifier(pooled_output) return logits base_model = AutoModel.from_pretrained( 'OpenGVLab/InternViT-6B-448px-V1-5', torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, trust_remote_code=True).cuda().eval() model = CustomModel(base_model, num_classes=2).to(device).eval() model.classifier.load_state_dict(torch.load("checkpoints/classifier_weights.pth")) image = Image.open('./examples/image1.jpg').convert('RGB') image_processor = CLIPImageProcessor.from_pretrained('OpenGVLab/InternViT-6B-448px-V1-5') pixel_values = image_processor(images=image, return_tensors='pt').pixel_values.to(torch.bfloat16).cuda() with torch.no_grad(): outputs = model(pixel_values) ```