|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
""" |
|
Processor class for Llava. |
|
""" |
|
|
|
|
|
from typing import List, Optional, Union |
|
|
|
from transformers.feature_extraction_utils import BatchFeature |
|
from transformers.image_utils import ImageInput |
|
from transformers.tokenization_utils_base import ( |
|
PaddingStrategy, |
|
PreTokenizedInput, |
|
TextInput, |
|
TruncationStrategy, |
|
) |
|
from transformers.utils import TensorType |
|
import torch |
|
from open_clip.transform import PreprocessCfg, image_transform_v2 |
|
from modeling_llava import LlavaForConditionalGeneration |
|
|
|
|
|
class OpenCLIPImageProcessor: |
|
def __init__(self, config): |
|
cfg = PreprocessCfg(**config) |
|
transform = image_transform_v2(cfg=cfg, is_train=False) |
|
self.transform = transform |
|
|
|
def __call__(self, image, return_tensors): |
|
if isinstance(image, list): |
|
outputs = [] |
|
for item in image: |
|
outputs.append(self.transform(item)) |
|
return { |
|
"pixel_values": torch.tensor(outputs), |
|
} |
|
output = self.transform(image) |
|
return { |
|
"pixel_values": output.unsqueeze(0), |
|
} |
|
|
|
@property |
|
def model_input_names(self): |
|
return ["pixel_values"] |
|
|
|
|
|
class LlavaProcessor: |
|
def __init__(self, image_processor: OpenCLIPImageProcessor, tokenizer): |
|
self.image_processor = image_processor |
|
self.tokenizer = tokenizer |
|
|
|
def __call__( |
|
self, |
|
text: Union[ |
|
TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput] |
|
] = None, |
|
images: ImageInput = None, |
|
model: LlavaForConditionalGeneration = None, |
|
padding: Union[bool, str, PaddingStrategy] = False, |
|
truncation: Union[bool, str, TruncationStrategy] = None, |
|
max_length=None, |
|
return_tensors: Optional[Union[str, TensorType]] = TensorType.PYTORCH, |
|
) -> BatchFeature: |
|
if images is not None: |
|
pixel_values = self.image_processor(images, return_tensors=return_tensors)[ |
|
"pixel_values" |
|
] |
|
pixel_values = pixel_values.to(model.device).to(model.dtype) |
|
image_outputs = model.vision_model(pixel_values) |
|
image_features = model.multi_modal_projector(image_outputs) |
|
else: |
|
image_features = None |
|
text_inputs = self.tokenizer( |
|
text, |
|
return_tensors=return_tensors, |
|
padding=padding, |
|
truncation=truncation, |
|
max_length=max_length, |
|
) |
|
|
|
return BatchFeature(data={**text_inputs, "image_features": image_features}) |
|
|
|
def batch_decode(self, *args, **kwargs): |
|
return self.tokenizer.batch_decode(*args, **kwargs) |
|
|
|
def decode(self, *args, **kwargs): |
|
return self.tokenizer.decode(*args, **kwargs) |
|
|
|
@property |
|
def model_input_names(self): |
|
tokenizer_input_names = self.tokenizer.model_input_names |
|
image_processor_input_names = self.image_processor.model_input_names |
|
return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) |
|
|
