commit demo to space

config.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+
+PREFIX_MAP = {
+    "openai/clip-vit-base-patch32": 50,
+    "openai/clip-vit-large-patch14": 257
+}
+
+
+@dataclass
+class LinearMappingConfig:
+    image_model: str = "openai/clip-vit-base-patch32"
+    freeze_image_model: bool = True
+    text_model: str = "gpt2-large"
+    freeze_text_model: bool = True
+    image_hidden_size: int = 768
+    text_hidden_size: int = 1280
+    linear_mapping_type: int = "linear"
+    max_seq_length: int = 2048
+    image_resize: int = 224
+    add_image_token: bool = True
+    freeze_ln: bool = False
+
+    def __post_init__(self):
+        self.prefix_length = PREFIX_MAP[self.image_model]

demo.py

@@ -0,0 +1,56 @@
+import gradio as gr
+from linear_mapping import LinearMapping, LinearMappingConfig, LinearMappingProcessor
+import os
+import torch
+
+os.environ['CURL_CA_BUNDLE'] = ''
+
+config = LinearMappingConfig()
+model = LinearMapping(config)
+model.load_state_dict(torch.load("pytorch_model.bin"))
+processor = LinearMappingProcessor(config)
+processor.tokenizer.padding_side = 'left'
+processor.tokenizer.pad_token_id = processor.tokenizer.eos_token_id
+
+title = "Generate Image Captions With CLIP And GPT2"
+
+
+def generate_image_captions(image, text):
+    inputs = processor(images=image, texts=text, return_tensors="pt")
+    input_ids = inputs.get("input_ids", None)
+    pixel_values = inputs.get("pixel_values", None)
+    attention_mask = inputs.get("attention_mask", None)
+    prediction = model.generate(
+        pixel_values=pixel_values,
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        max_new_tokens=50
+    )
+
+    prediction_text = processor.decode(prediction[0], num_beams=5, skip_special_tokens=True)
+    return prediction_text
+
+
+article = "This demo is originated from this paper: [original paper](https://arxiv.org/abs/2209.15162)"
+description = """
+### Expand GPT2's language capabilities to vision with CLIP!
+"""
+demo = gr.Interface(
+    fn=generate_image_captions,
+    inputs=[
+        gr.Image(),
+        gr.Textbox(placeholder="A picture of", lines=3)
+    ],
+    outputs="text",
+    examples=[
+        [os.path.join(os.getcwd(), 'two_bear.png'), ""],
+        [os.path.join(os.getcwd(), 'cat_with_food.png'), "Describe the picture:"],
+        [os.path.join(os.getcwd(), 'dog_with_frisbee.png'), "What is the color of the frisbee in the photo? Answer:"],
+        [os.path.join(os.getcwd(), 'stop_sign.png'), "What does the sign in the picture say? Answer:"]
+    ],
+    article=article,
+    title=title,
+    description=description
+)
+
+demo.launch(share=True)

linear_mapping.py

@@ -0,0 +1,278 @@
+from config import LinearMappingConfig
+from transformers import (
+    GPT2TokenizerFast, GPT2LMHeadModel, AutoModel,
+    CLIPVisionModel, AutoProcessor, BatchEncoding,
+)
+from transformers.models.gpt2.modeling_gpt2 import GPT2DoubleHeadsModelOutput
+import torch
+import torch.nn as nn
+from typing import List, Optional, Union, Tuple, Dict
+from torchvision.transforms import CenterCrop, ConvertImageDtype, Normalize, Resize
+from torchvision.transforms.functional import InterpolationMode
+
+
+class Transform(torch.nn.Module):
+    def __init__(self, image_size, mean, std):
+        super().__init__()
+        self.transforms = torch.nn.Sequential(
+            Resize([image_size], interpolation=InterpolationMode.BICUBIC, antialias=True),
+            CenterCrop(image_size),
+            ConvertImageDtype(torch.float32),
+            Normalize(mean, std),
+        )
+
+    def forward(self, x) -> torch.Tensor:
+        """`x` should be an instance of `PIL.Image.Image`"""
+        with torch.no_grad():
+            x = self.transforms(x)
+        return x
+
+
+class LinearMappingProcessor:
+    """
+    A combination of ImageProcessor and GPT2TokenizerFast
+    """
+
+    def __init__(self, config: LinearMappingConfig):
+        self.image_processor = AutoProcessor.from_pretrained(config.image_model)
+        self.tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
+        self.add_image_token = config.add_image_token
+        if config.add_image_token:
+            self.tokenizer.add_special_tokens({"cls_token": "|<image>|"})
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        self.tokenizer.padding_side = "right"
+        self.prefix_length = config.prefix_length
+
+    def __call__(self, texts=None, images=None, return_tensors="pt", **kwargs):
+        """
+        The processor assumes that images and texts are of the same number
+        """
+
+        if len(texts) == 0:     # empty strings should be None
+            texts = None
+
+        if images is not None:
+            image_features = self.image_processor(images=images, return_tensors=return_tensors, **kwargs)
+            image_features["attention_mask"] = torch.ones(image_features.pixel_values.size(0),
+                                                          self.prefix_length).to(dtype=torch.int64)
+            if texts is None and self.add_image_token:
+                texts = [self.tokenizer.cls_token for _ in range(image_features.pixel_values.size(0))]
+            elif texts is not None and self.add_image_token:
+                if isinstance(texts, str):
+                    texts = [texts]
+                texts = [self.tokenizer.cls_token + text for text in texts]
+
+        elif texts is None:
+            texts = self.tokenizer.bos_token
+
+        if texts is not None:
+            encoding = self.tokenizer(texts, return_tensors=return_tensors, **kwargs)
+
+        if texts is not None and images is not None:
+            encoding["pixel_values"] = image_features.pixel_values
+
+            encoding["attention_mask"] = torch.cat([
+                image_features["attention_mask"],
+                encoding["attention_mask"]
+            ], dim=1).to(dtype=torch.long)  # create attention mask for images
+            return encoding
+
+        elif texts is not None:
+            return encoding
+        else:
+            return BatchEncoding(data=dict(**image_features), tensor_type=return_tensors)
+
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to GPT2TokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
+        refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to GPT2TokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
+        the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+
+class ImagePrefix(nn.Module):
+    """
+    Converts pixel values to prefix image prompts that are later fed to a LLM
+    """
+
+    def __init__(self, config: LinearMappingConfig):
+        super().__init__()
+        self.encoder = AutoModel.from_pretrained(config.image_model)
+        if "clip" in config.image_model:
+            self.encoder = CLIPVisionModel.from_pretrained(config.image_model)
+
+        if config.freeze_image_model:
+            for param in self.encoder.parameters():
+                param.requires_grad = False
+
+        self.linear = nn.Linear(config.image_hidden_size, config.text_hidden_size)
+        self.ln = nn.LayerNorm(config.text_hidden_size)
+
+    def forward(
+            self, pixel_values: torch.Tensor  # B x C x H x W
+    ) -> torch.Tensor:
+        prefixes = self.encoder(pixel_values).last_hidden_state  # B x N x D
+        prefix_prompts = self.linear(prefixes)
+        return self.ln(prefix_prompts)
+
+
+class LinearMapping(nn.Module):
+
+    def __init__(self, config: LinearMappingConfig):
+        super().__init__()
+        self.image_prefix = ImagePrefix(config)
+        self.language_model = GPT2LMHeadModel.from_pretrained(config.text_model)
+        self.processor = LinearMappingProcessor(config)
+        self.tokenizer = self.processor.tokenizer
+        self.image_processor = self.processor.image_processor
+        self.add_image_token = config.add_image_token
+        if config.add_image_token:
+            self.language_model.resize_token_embeddings(len(self.tokenizer))
+
+        if config.freeze_text_model:
+            for module in self.language_model.modules():
+                if not isinstance(module, nn.LayerNorm) or config.freeze_ln:
+                    for param in module.parameters():
+                        param.requires_grad = False
+            if config.add_image_token:
+                # create a gradient mask for the lm_head weight and bias and hook it
+                self.language_model.lm_head.weight.requires_grad = True
+                self.weight_gradient_mask = nn.Parameter(torch.zeros_like(self.language_model.lm_head.weight),
+                                                         requires_grad=False)
+                self.weight_gradient_mask[-1, :] = 1.0
+                self.language_model.lm_head.weight.register_hook(lambda grad: grad.mul_(self.weight_gradient_mask))
+
+    def prepare_text_inputs(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.language_model.transformer.wte(input_ids.to(dtype=torch.int64))
+
+    def prepare_inputs(
+            self,
+            input_ids: Optional[torch.Tensor],
+            pixel_values: Optional[torch.Tensor]
+    ) -> Dict:
+        """
+        Prepare captions and pixel values for training.
+        It takes the captions' input ids and turn them into input embeddings
+        and turns pixel values into prefix prompts.
+        Then it concatenates them into one whole prompt batch.
+        """
+        if input_ids is not None and pixel_values is not None:
+
+            text_embeddings = self.prepare_text_inputs(input_ids)  # B x T x D
+            prefix_prompts = self.image_prefix(pixel_values)  # B x V x D
+            inputs_embeddings = torch.cat([prefix_prompts, text_embeddings], dim=1)
+
+            prefix_labels = torch.zeros(prefix_prompts.shape[:2], device=prefix_prompts.device) - 100
+            labels = torch.cat([prefix_labels, input_ids], dim=1)  # B x (V + T)
+
+            for label in labels:
+                for k, token in enumerate(label):
+                    if token == self.tokenizer.eos_token_id:
+                        label[k + 1:] = -100
+                        break
+            return {"hidden_states": inputs_embeddings, "labels": labels.to(dtype=torch.int64)}
+
+        elif pixel_values is not None:
+            prefix_prompts = self.image_prefix(pixel_values)  # B x V x D
+            prefix_labels = torch.zeros(prefix_prompts.shape[:2], device=prefix_prompts.device) - 100
+            return {"hidden_states": prefix_prompts, "labels": prefix_labels.to(dtype=torch.int64)}
+
+        elif input_ids is not None:
+            text_embeddings = self.prepare_text_inputs(input_ids)
+            labels = input_ids.clone()
+            for label in labels:
+                for k, token in enumerate(label):
+                    if token == self.tokenizer.eos_token_id:
+                        label[k + 1:] = -100
+                        break
+            return {"hidden_states": text_embeddings, "labels": labels.to(dtype=torch.int64)}
+        else:
+            return {"hidden_states": None, "labels": None}
+
+    @torch.no_grad()
+    def generate(
+            self,
+            input_ids: Optional[torch.Tensor] = None,
+            pixel_values: Optional[torch.Tensor] = None,
+            **kwargs
+    ):
+        if pixel_values is None:
+            return self.language_model.generate(
+                input_ids=input_ids,
+                **kwargs
+            )
+        batch_size = pixel_values.size(0)
+        past_input_ids = None
+        if input_ids is None:
+            if self.add_image_token:
+                input_ids = torch.tensor([self.tokenizer.cls_token_id for _ in range(batch_size)]).view(batch_size, -1)
+            else:
+                input_ids = torch.tensor([self.tokenizer.bos_token_id for _ in range(batch_size)]).view(batch_size, -1)
+        if input_ids.size(-1) <= 1:
+            first_forward_outputs = self.forward(
+                pixel_values=pixel_values
+            )
+        else:
+            first_forward_outputs = self.forward(
+                pixel_values=pixel_values,
+                input_ids=input_ids[:, :-1]
+            )
+            past_input_ids = input_ids[:, :-1]
+            input_ids = input_ids[:, -1].view(batch_size, -1)
+
+        past_key_values = first_forward_outputs.past_key_values
+
+        if kwargs.get("attention_mask", None) is None:
+            attention_mask_size = (past_key_values[0][0].size(0), past_key_values[0][0].size(-2))
+
+            attention_mask = torch.ones(attention_mask_size, dtype=torch.int64)
+        else:
+            attention_mask = kwargs.pop("attention_mask")
+
+        generated_token_ids = self.language_model.generate(
+            past_key_values=past_key_values,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            **kwargs
+        )
+        if past_input_ids is not None:
+            generated_token_ids = torch.cat([past_input_ids, generated_token_ids], dim=-1)
+        return generated_token_ids
+
+    def forward(
+            self,
+            input_ids: Optional[torch.Tensor] = None,
+            pixel_values: Optional[torch.Tensor] = None,
+            labels: Optional[torch.Tensor] = None,
+            inputs_embeds: Optional[torch.Tensor] = None,
+            output_hidden_states: bool = True,
+            output_attentions: bool = True,
+            attention_mask: Optional[torch.Tensor] = None,
+            return_dict: Optional[bool] = True,
+            **kwargs
+    ) -> Union[GPT2DoubleHeadsModelOutput, Tuple]:
+        if (pixel_values is None and input_ids is None) and inputs_embeds is None:
+            raise ValueError("You have to specify inputs")
+        if inputs_embeds is not None and (pixel_values is not None or input_ids is not None):
+            raise ValueError("Either inputs_embeds or (pixel_values and input_ids) should be specified, not both")
+
+        inputs = self.prepare_inputs(input_ids, pixel_values)
+        hidden_states = inputs.get('hidden_states', None) if inputs_embeds is None else inputs_embeds
+        labels = inputs.get('labels', None) if labels is None else labels
+
+        return self.language_model(
+            inputs_embeds=hidden_states,
+            labels=labels,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            attention_mask=attention_mask,
+            return_dict=return_dict,
+            **kwargs
+        )

main.py

@@ -0,0 +1,116 @@
+from datasets import load_dataset
+from linear_mapping import LinearMapping, LinearMappingProcessor, LinearMappingConfig, Transform
+import torch
+from torchvision.io import ImageReadMode, read_image
+from transformers import Trainer, TrainingArguments
+import os
+from PIL import Image
+os.environ["WANDB_DISABLED"] = "true"
+
+DATA_DIR = os.path.join(os.getcwd(), "coco")
+CAPTION_COLUMN = "caption"
+IMAGE_COLUMN = "image_path"
+
+
+def main():
+    ds = load_dataset("ydshieh/coco_dataset_script", "2017", DATA_DIR)
+    config = LinearMappingConfig()
+    processor = LinearMappingProcessor(config)
+
+    def collate_fn(batch):
+        return {
+            'pixel_values': torch.stack([x['pixel_values'] for x in batch]),
+            'input_ids': torch.tensor([x['input_ids'] for x in batch], dtype=torch.long),
+            'attention_mask': torch.stack([x["attention_mask"] for x in batch]),
+        }
+
+    def tokenize_fn(examples):
+        texts = list(examples[CAPTION_COLUMN])
+        if config.add_image_token:
+            texts = list(processor.tokenizer.cls_token + text for text in texts)
+        inputs = processor.tokenizer(
+            texts, padding="max_length", max_length=77,
+            return_tensors="pt", truncation=True
+        )
+        examples["input_ids"] = inputs.input_ids
+        examples["attention_mask"] = inputs.attention_mask
+        return examples
+
+    image_transformations = Transform(
+        config.image_resize,
+        [0.48145466, 0.4578275, 0.40821073],
+        [0.26862954, 0.26130258, 0.27577711]
+    )
+    image_transformations = torch.jit.script(image_transformations)
+
+    def transform_images(examples):
+        images = [read_image(image_file, mode=ImageReadMode.RGB) for image_file in examples[IMAGE_COLUMN]]
+        examples["pixel_values"] = [image_transformations(image) for image in images]
+
+        examples["attention_mask"] = torch.cat([
+            torch.ones(len(images), config.prefix_length),
+            torch.tensor(examples["attention_mask"])
+        ], dim=1).to(dtype=torch.long)
+        return examples
+
+    def preprocess_fn(examples):
+
+        texts = list(examples[CAPTION_COLUMN])
+
+        images = [read_image(image_file, mode=ImageReadMode.RGB) for image_file in examples[IMAGE_COLUMN]]
+        inputs = processor(
+            texts=texts, images=images, padding="max_length", truncation=True, max_length=77, return_tensors="pt"
+        )
+        return inputs
+
+    def filter_corrupt_images(examples):
+        """remove problematic images"""
+        valid_images = []
+        for image_file in examples[IMAGE_COLUMN]:
+            try:
+                Image.open(image_file)
+                valid_images.append(True)
+            except Exception:
+                valid_images.append(False)
+        return valid_images
+
+    train_dataset = ds["train"]
+
+    train_dataset = train_dataset.filter(
+        function=filter_corrupt_images,
+        batched=True
+    )
+    train_dataset = train_dataset.map(
+        function=tokenize_fn,
+        batched=True,
+        remove_columns=[col for col in train_dataset.column_names if col != IMAGE_COLUMN and col != CAPTION_COLUMN],
+        load_from_cache_file=True
+    )
+    train_dataset.set_transform(transform_images)
+
+    training_args = TrainingArguments(
+        learning_rate=5e-4,
+        lr_scheduler_type='cosine',
+        output_dir='clip-gpt2-image-captioner',
+        do_train=True,
+        logging_steps=50,
+        num_train_epochs=5,
+        logging_dir='runs',
+        remove_unused_columns=False,
+        max_grad_norm=1.0,
+        per_device_train_batch_size=16,
+        save_total_limit=3,
+        warmup_steps=500
+    )
+    model = LinearMapping(config)
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        data_collator=collate_fn
+    )
+    trainer.train()
+
+
+if __name__ == '__main__':
+    main()

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f817ef4696fa1ccb00cf19e71ed36660d9c52212fd1e953dbf52f923a7553ca0
+size 3707484877