Correct the inference code.

.gitattributes

@@ -37,3 +37,4 @@ tokenizer.json filter=lfs diff=lfs merge=lfs -text
 inference_demo/examples/1408px-Lilium_philadelphicum_var._philadelphicum.jpg filter=lfs diff=lfs merge=lfs -text
 inference_demo/examples/stock-footage-pictorial-upper-view-sunset-beams-light-waterfall-stone-cascade-and-fresh-green-tropical-trees.mp4 filter=lfs diff=lfs merge=lfs -text
 examples/stock-footage-timelapse-of-stormy-clouds-over-open-sea-and-snowcapped-mountain.mp4 filter=lfs diff=lfs merge=lfs -text
+inference_demo/examples/stock-footage-timelapse-of-stormy-clouds-over-open-sea-and-snowcapped-mountain.mp4 filter=lfs diff=lfs merge=lfs -text

inference_demo/examples/stock-footage-timelapse-of-stormy-clouds-over-open-sea-and-snowcapped-mountain.mp4

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

inference_demo/inference.py

@@ -0,0 +1,141 @@
+import torch
+from transformers import AutoTokenizer, AutoProcessor
+from qwen_vl_utils import process_vision_info
+from modeling_unite import UniteQwen2VL
+
+
+model_path = 'friedrichor/Unite-Base-Qwen2-VL-2B'
+model = UniteQwen2VL.from_pretrained(
+    model_path, 
+    torch_dtype=torch.bfloat16,
+    device_map="cuda"
+)
+
+# We recommend enabling flash_attention_2 for better acceleration and memory saving.
+# model = UniteQwen2VL.from_pretrained(
+#     model_path,
+#     device_map="cuda",
+#     torch_dtype=torch.bfloat16,
+#     attn_implementation='flash_attention_2', 
+#     low_cpu_mem_usage=True,
+# )
+
+tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+processor = AutoProcessor.from_pretrained(model_path, min_pixels=256*28*28, max_pixels=1280*28*28)
+
+def process_messages(msg):
+    text = processor.apply_chat_template(msg, tokenize=False, add_generation_prompt=True) + "<|endoftext|>"
+    image_inputs, video_inputs = process_vision_info(msg)
+    inputs = processor(
+        text=[text], 
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    )
+    inputs = inputs.to("cuda")
+
+    return inputs
+
+## ============================== Text-Image ==============================
+messages_txt = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "text", "text": "The book titled 'Riding with Reindeer - A Bicycle Odyssey through Finland, Lapland, and the Arctic' provides a detailed account of a journey that explores the regions of Lapland and the Arctic, focusing on the experience of riding with reindeer."},
+            {"type": "text", "text": "\nSummary above sentence in one word:"},
+        ],
+    }
+]
+
+messages_img = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "image", "image": "./examples/518L0uDGe0L.jpg"},
+            {"type": "text", "text": "\nSummary above image in one word:"},
+        ],
+    }
+]
+
+inputs_txt = process_messages(messages_txt)
+inputs_img = process_messages(messages_img)
+
+with torch.no_grad():
+    embeddings_txt = model(**inputs_txt)  # [1, 1536]
+    embeddings_img = model(**inputs_img)  # [1, 1536]
+    print(f"embeddings_txt.shape: {embeddings_txt.shape}")
+    print(f"embeddings_img.shape: {embeddings_img.shape}")
+
+    print(torch.matmul(embeddings_txt, embeddings_img.T))
+    # tensor([[0.7500]], dtype=torch.bfloat16)
+
+## ============================== Text-Video ==============================
+messages_txt = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "text", "text": "Timelapse of stormy clouds over open sea and snowcapped mountain"},
+            {"type": "text", "text": "\nSummary above sentence in one word:"},
+        ],
+    }
+]
+
+messages_vid = [
+    {
+        "role": "user",
+        "content": [
+            {
+                "type": "video", 
+                "video": "./examples/stock-footage-timelapse-of-stormy-clouds-over-open-sea-and-snowcapped-mountain.mp4",
+                "max_pixels": 360 * 420, 
+                "fps": 1,
+                "max_frames": 32
+            },
+            {"type": "text", "text": "\nSummary above video in one word:"},
+        ],
+    }
+]
+
+inputs_txt = process_messages(messages_txt)
+inputs_vid = process_messages(messages_vid)
+
+with torch.no_grad():
+    embeddings_txt = model(**inputs_txt)  # [1, 1536]
+    embeddings_vid = model(**inputs_vid)  # [1, 1536]
+
+    print(torch.matmul(embeddings_txt, embeddings_vid.T))
+    # tensor([[0.5664]], dtype=torch.bfloat16)
+
+## ============================== Fused Modal ==============================
+messages_qry = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "image", "image": "./examples/oven_05011373.jpg"},
+            {"type": "text", "text": "What is the name of this place?"},
+            {"type": "text", "text": "\nSummary above sentence and image in one word:"},
+        ],
+    }
+]
+
+messages_tgt = [
+    {
+        "role": "user",
+        "content": [
+            {"type": "image", "image": "./examples/Q673659.jpg"},
+            {"type": "text", "text": "Marina Beach."},
+            {"type": "text", "text": "\nSummary above sentence and image in one word:"},
+        ],
+    }
+]
+
+inputs_qry = process_messages(messages_qry)
+inputs_tgt = process_messages(messages_tgt)
+
+with torch.no_grad():
+    embeddings_qry = model(**inputs_qry)  # [1, 1536]
+    embeddings_tgt = model(**inputs_tgt)  # [1, 1536]
+
+    print(torch.matmul(embeddings_qry, embeddings_tgt.T))
+    # tensor([[0.7695]], dtype=torch.bfloat16)

inference_demo/modeling_unite.py

@@ -0,0 +1,88 @@
+from typing import Optional, List
+
+import torch
+import torch.nn as nn
+from transformers import Qwen2VLConfig, Qwen2VLForConditionalGeneration, Qwen2VLModel
+
+
+class UniteQwen2VLConfig(Qwen2VLConfig):
+    model_type = "unite_qwen2_vl"
+
+
+class UniteQwen2VL(Qwen2VLForConditionalGeneration):
+    def __init__(self, config):
+        Qwen2VLForConditionalGeneration.__init__(self, config)
+        config.model_type = "unite_qwen2_vl"
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+        self.normalize = True
+    
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        pooling_mask: Optional[torch.LongTensor] = None,
+        **kwargs
+    ) -> torch.Tensor:
+        if inputs_embeds is None:
+            inputs_embeds = self.model.embed_tokens(input_ids)
+            if pixel_values is not None:
+                pixel_values = pixel_values.type(self.visual.get_dtype())
+                image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
+                image_mask = (
+                    (input_ids == self.config.image_token_id)
+                    .unsqueeze(-1)
+                    .expand_as(inputs_embeds)
+                    .to(inputs_embeds.device)
+                )
+                image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+                inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+            if pixel_values_videos is not None:
+                pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype())
+                video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
+                video_mask = (
+                    (input_ids == self.config.video_token_id)
+                    .unsqueeze(-1)
+                    .expand_as(inputs_embeds)
+                    .to(inputs_embeds.device)
+                )
+                video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+                inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(inputs_embeds.device)
+
+        outputs = self.model(
+            input_ids=None,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+        )
+
+        pooling_mask = attention_mask if pooling_mask is None else pooling_mask
+        left_padding = (pooling_mask[:, -1].sum() == pooling_mask.shape[0])  # TODO
+        
+        if left_padding:
+            embeddings = outputs.last_hidden_state[:, -1]
+        else:
+            sequence_lengths = pooling_mask.sum(dim=1) - 1
+            batch_size = outputs.last_hidden_state.shape[0]
+            embeddings = outputs.last_hidden_state[torch.arange(
+                batch_size, device=outputs.last_hidden_state.device
+            ), sequence_lengths]
+        if self.normalize:
+            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+
+        return embeddings.contiguous()
+
+