|
|
|
|
|
|
|
|
|
|
|
|
|
import torch |
|
import torch.nn as nn |
|
import torchvision.transforms as T |
|
from torchvision.transforms import InterpolationMode |
|
from transformers import LlamaTokenizer |
|
|
|
from .configuration_intern_vit import InternVisionConfig |
|
from .configuration_internvl import InternVLConfig |
|
from .modeling_intern_vit import InternVisionModel |
|
from .modeling_internvl import InternVL_C, InternVL_G, InternVLModel |
|
|
|
__all__ = ['InternVisionConfig', 'InternVisionModel', 'InternVLConfig', |
|
'InternVLModel', 'InternVL_C', 'InternVL_G'] |
|
|
|
|
|
|
|
class InternVLTokenizer(nn.Module): |
|
def __init__(self, model_path): |
|
super(InternVLTokenizer, self).__init__() |
|
self.tokenizer = LlamaTokenizer.from_pretrained(model_path) |
|
self.tokenizer.pad_token = ' ' |
|
self.tokenizer.add_eos_token = True |
|
|
|
def forward(self, text, prefix='summarize:'): |
|
if type(text) == str: |
|
text = prefix + text |
|
elif type(text) == list: |
|
text = [prefix + item for item in text] |
|
text = self.tokenizer(text, return_tensors='pt', max_length=80, truncation=True, padding='max_length').input_ids |
|
return text |
|
|
|
|
|
def build_transform(task, image_size=224, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]): |
|
if task == 'retrieval': |
|
transform = T.Compose([ |
|
T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img), |
|
T.Resize((image_size, image_size), interpolation=InterpolationMode.BICUBIC), |
|
T.ToTensor(), |
|
T.Normalize(mean=mean, std=std)]) |
|
else: |
|
transform = T.Compose([ |
|
T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img), |
|
T.Resize(image_size, interpolation=InterpolationMode.BICUBIC), |
|
T.CenterCrop(image_size), |
|
T.ToTensor(), |
|
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) |
|
return transform |
|
|
|
|
|
def load_internvl_c_huggingface(ckpt_path, device, task): |
|
model = InternVL_C.from_pretrained(ckpt_path, torch_dtype=torch.float16).to(device) |
|
if model.config.use_backbone_lora: |
|
model.vision_model.merge_and_unload() |
|
model.vision_model = model.vision_model.model |
|
if model.config.use_qllama_lora: |
|
model.qllama.merge_and_unload() |
|
model.qllama = model.qllama.model |
|
if model.config.force_image_size is not None: |
|
image_size = model.config.force_image_size |
|
else: |
|
image_size = model.config.vision_config.image_size |
|
transform = build_transform(task, image_size) |
|
tokenizer = InternVLTokenizer(ckpt_path) |
|
return model, transform, tokenizer |
|
|
|
|
|
def load_internvl_g_huggingface(ckpt_path, device, task): |
|
model = InternVL_G.from_pretrained(ckpt_path, torch_dtype=torch.float16).to(device) |
|
if model.config.use_backbone_lora: |
|
model.vision_model.merge_and_unload() |
|
model.vision_model = model.vision_model.model |
|
if model.config.use_qllama_lora: |
|
model.qllama.merge_and_unload() |
|
model.qllama = model.qllama.model |
|
if model.config.force_image_size is not None: |
|
image_size = model.config.force_image_size |
|
else: |
|
image_size = model.config.vision_config.image_size |
|
transform = build_transform(task, image_size) |
|
tokenizer = InternVLTokenizer(ckpt_path) |
|
return model, transform, tokenizer |
|
|
