Update model.py

model.py

@@ -1,8 +1,68 @@
+from torch import nn
+from transformers import AutoProcessor, CLIPVisionModel, AutoModelForCausalLM, AutoTokenizer
+from PIL import Image
+import torch
+import numpy as np
+import torch.nn.functional as F
+
 
 class ToyModel():
-    def __init__(self):
-        ...
+    """
+    CLIP + GPT2
+    """
+    def __init__(self, vision_model_path, language_model_path):
+        # load vision encoder
+        self.vision_encoder = CLIPVisionModel.from_pretrained(vision_model_path)
+        self.processor = AutoProcessor.from_pretrained(vision_model_path)
+
+        # load language encoder
+        self.language_model = AutoModelForCausalLM.from_pretrained(language_model_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(language_model_path)
+
+        # MLP connector
+        self.mlp = nn.Sequential(
+            nn.Linear(768, 768),
+            nn.ReLU(),
+            nn.Linear(768, 768),
+            nn.ReLU()
+        )
+
+    def encode_image(self, image):
+        image = self.processor(images=image, return_tensors="pt")
+        return self.vision_encoder(**image)
+
+    def encode_text(self, prompt):
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
+        prompt_embeddings = self.language_model.get_input_embeddings()(input_ids)
+        return prompt_embeddings
+
+
+    def chat(self, image, text):
+        # encode image
+        outputs = self.encode_image(image)
+        image_embeddings = outputs.last_hidden_state
+
+        # encode text
+        text_embeddings = self.encode_text(text)
+
+        # chat with image and text
+        # embedding fusion
+        image_embeddings = self.mlp(image_embeddings)
+        embedding = torch.cat((image_embeddings, text_embeddings), dim=1)
+        outputs = self.language_model(inputs_embeds=embedding)
+
+        # decode logits to text
+        logits = outputs.logits
+        preds = F.softmax(logits, dim=-1).argmax(dim=-1)
+        text_output = self.tokenizer.batch_decode(sequences=preds, skip_special_tokens=True)
+        return text_output
+
+
+if __name__ == '__main__':
+    model = ToyModel('/home/yuan/huggingface/model/clip-vit-base-patch32', '/home/yuan/huggingface/model/gpt2')
 
-    def chat(self, image_input, text_input):
-        return image_input, text_input
+    image = Image.open('/home/yuan/RS-VL-Perception/examples_v2/thief.png')
+    text = 'I am Iron Man'
 
+    print(model.chat(image, text))
+    # [",....\n.\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n 37\n\n\n 40 40 40 40\n'm a Man,"]