Upload modeling_tara.py with huggingface_hub

modeling_tara.py

@@ -0,0 +1,383 @@
+import os
+from abc import ABCMeta, abstractmethod
+from typing import Optional, Union, Dict, List
+from termcolor import colored
+import random
+
+
+import numpy as np
+import torch
+from transformers import (
+    AutoProcessor,
+    AutoTokenizer,
+    LlavaConfig, 
+    LlamaForCausalLM,
+)
+from torchvision.transforms.v2 import (
+    ToPILImage,
+)
+import decord
+from decord import VideoReader
+
+# TODO: need to use these directly
+from tarsier.modeling_tarsier import TarsierForConditionalGeneration
+from tarsier.processor import Processor
+# from utils.model import transform_pixel_values
+
+decord.bridge.set_bridge("torch")
+
+
+EOL_PROMPTS = {
+    'text': '<sent>\nSummary above sentence in one word:',
+    'image': '<image>\nSummary above image in one word:',
+    'video': '<video>\nSummary above video in one word:',
+}
+
+
+def transform_pixel_values(pixel_values: torch.Tensor | List[torch.Tensor]) -> torch.Tensor:
+    # NOTE: this function doesn't accept unbatched inputs
+    # pixel_values should be uint8 of (B, T, C, H, W)
+    if isinstance(pixel_values, list):
+        pixel_values = torch.stack(pixel_values)
+
+    if pixel_values.ndim == 4:
+        # pixel_values is (B, C, H, W)
+        # (B, C, H, W) -> (B, 1, C, H, W)
+        pixel_values = pixel_values.unsqueeze(1)
+    elif pixel_values.ndim == 5:
+        # pixel_values is (B, T, C, H, W)
+        pass
+    else:
+        raise ValueError(f"pixel_values should be 4D or 5D, got {pixel_values.ndim}D")
+    return pixel_values
+
+
+base_registry = {}
+class BaseModel(metaclass=ABCMeta):
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        # register model architecture
+        if hasattr(cls, 'ARCHITECTURE'):
+            base_registry[cls.ARCHITECTURE] = cls
+    
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        load_llm: bool = False,
+        device_map: Optional[Union[str, Dict[str, int]]] = None,
+        **kwargs):
+        print(colored(f'[ MODEL ] Loading {cls.__name__} from {model_name_or_path} [..............]', 'yellow'))
+
+        return cls(model_name_or_path, load_llm=load_llm, device_map=device_map, **kwargs)
+
+
+class BaseModelForTARA(BaseModel):
+    
+    ARCHITECTURE = "TarsierForConditionalGeneration"
+    LLM_CLASS = LlamaForCausalLM
+    MLLM_CLASS = TarsierForConditionalGeneration
+
+    @property
+    def describe_prompt(self):
+        return "Describe the video in detail."
+
+    @property
+    def text_eol_prompt(self):
+        prompt = f'USER: {EOL_PROMPTS["text"]} ASSISTANT: '
+        return prompt
+    
+    @property
+    def image_eol_prompt(self):
+        prompt = f'USER: {EOL_PROMPTS["image"]} ASSISTANT: '
+        return prompt
+    
+    @property
+    def video_eol_prompt(self):
+        prompt = f'USER: {EOL_PROMPTS["video"]} ASSISTANT: '
+        return prompt
+
+    def __init__(
+            self, 
+            model_name_or_path: str,
+            load_llm: Optional[bool] = None,
+            device_map: Optional[Union[str, Dict[str, int]]] = None,
+            **kwargs,
+        ):
+
+        MODEL_CLASS = self.LLM_CLASS if load_llm else self.MLLM_CLASS
+
+        if load_llm:
+            self.split_weights(model_name_or_path, model_name_or_path + '-llm')
+            model_name_or_path += '-llm'
+            model_config = None
+            self.processor = AutoProcessor.from_pretrained(model_name_or_path, use_fast=False)
+        else:
+            model_config = LlavaConfig.from_pretrained(
+                model_name_or_path,
+                # trust_remote_code=True,
+            )
+            self.processor = Processor(
+                model_name_or_path,
+                max_n_frames=32,
+            )
+        
+        self.tokenizer = self.processor.tokenizer
+
+        self.model = MODEL_CLASS.from_pretrained(
+            model_name_or_path,
+            config=model_config,
+            torch_dtype=kwargs.get("torch_dtype", torch.bfloat16),
+            device_map=device_map,
+            # trust_remote_code=True
+        )
+        
+        self.model.eval()
+
+    def split_weights(self, mllm_path, llm_path):
+        if os.path.exists(llm_path):
+            print(f'{llm_path} already exists. Skip splitting weights.')
+            return
+        print('Splitting LLM weights from MLLM.')
+        model = self.MLLM_CLASS.from_pretrained(mllm_path)
+        llm = model.language_model
+        processor = AutoProcessor.from_pretrained(mllm_path)
+        tokenizer = AutoTokenizer.from_pretrained(mllm_path)
+        llm.save_pretrained(llm_path)
+        processor.save_pretrained(llm_path)
+        tokenizer.save_pretrained(llm_path)
+
+
+encoder_registry = {}
+class EncodeMixin(metaclass=ABCMeta):
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        # register model architecture
+        if hasattr(cls, 'ARCHITECTURE'):
+            encoder_registry[cls.ARCHITECTURE] = cls
+
+    @abstractmethod
+    def encode_vision(self, pixel_values: torch.Tensor | List[torch.Tensor]) -> torch.Tensor:
+        """
+        Encodes vision data (images or videos) into a tensor representation.
+
+        Args:
+            pixel_values (torch.Tensor | List[torch.Tensor]): The input pixel values. 
+                - If a tensor, it should be of shape (B, C, H, W) for images or (B, T, C, H, W) for videos.
+                - If a list, it will be stacked into a tensor.
+
+        Returns:
+            torch.Tensor: The encoded tensor representation of the input vision data.
+
+        Raises:
+            ValueError: If `pixel_values` is not 4D or 5D.
+
+        ## Notes:
+            - This function does not accept unbatched inputs.
+            - `pixel_values` should be of type uint8.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def encode_text(self, text: str | List[str]) -> torch.Tensor:
+        """
+        Encodes the given text(s) into a tensor representation using the model.
+
+        Args:
+            text (str | List[str]): A single string or a list of strings to be encoded.
+
+        Returns:
+            torch.Tensor: The tensor representation of the encoded text(s).
+
+        ## Notes:
+            - The method uses a prompt to encode the text.
+            - If a single string is provided, it is converted into a list containing that string.
+            - The method processes the prompts and generates the tensor representation using the model.
+            - The output tensor contains the hidden states of the last token for each input text.
+        """
+        raise NotImplementedError
+
+
+class TARA(BaseModelForTARA, EncodeMixin):
+
+    def encode_vision(self, pixel_values: torch.Tensor | List[torch.Tensor]) -> torch.Tensor:
+
+        pixel_values = transform_pixel_values(pixel_values) # [B, T, C, H, W]
+        nframes = pixel_values.shape[1]
+        prompt = self.image_eol_prompt if nframes == 1 else self.video_eol_prompt
+        
+        to_image = ToPILImage()
+        batched_frames = []
+        for batch in pixel_values:
+            frames = [to_image(v) for v in batch]
+            batched_frames.append(frames)
+
+        generate_kwargs = {
+            "max_new_tokens": 1,
+            "output_hidden_states": True,
+            "return_dict_in_generate": True,
+        }
+
+        vision_embs = []
+
+        for frames in batched_frames:
+            input_prompt = prompt.replace("<video>", "<image>"*len(frames))
+            input_ids = self.processor.get_text_inputs(input_prompt)
+            frames = self.processor.get_pixel_values(frames)
+            inputs = {
+                "input_ids": input_ids,
+                "pixel_values": frames
+            }
+            inputs = {k:v.to(self.model.device) for k,v in inputs.items() if v is not None}
+            outputs = self.model.generate(
+                **inputs,
+                **generate_kwargs,
+            )
+            vision_embs.append(outputs.hidden_states[0][-1][:, -1, :])
+        
+        vision_embs = torch.cat(vision_embs)
+        return vision_embs
+    
+    def encode_text(self, text: str | List[str]) -> torch.Tensor:
+
+        prompt = self.text_eol_prompt
+
+        if isinstance(text, str):
+            text = [text]
+        
+        prompts = [prompt.replace('<sent>', t) for t in text]
+
+        generate_kwargs = {
+            "max_new_tokens": 1,
+            "output_hidden_states": True,
+            "return_dict_in_generate": True,
+        }
+
+        text_embs = []
+
+        for p in prompts:
+            text_inputs = self.processor.get_text_inputs(p)
+            inputs = {
+                "input_ids": text_inputs,
+            }
+            inputs = {k:v.to(self.model.device) for k,v in inputs.items() if v is not None}
+            outputs = self.model.generate(
+                **inputs,
+                **generate_kwargs,
+            )
+            text_embs.append(outputs.hidden_states[0][-1][:, -1, :])
+        
+        text_embs = torch.cat(text_embs)
+        return text_embs
+
+
+
+
+
+def get_frame_indices(num_frames, vlen, sample='rand', fix_start=None, input_fps=1, max_num_frames=-1):
+    if sample in ["rand", "middle"]: # uniform sampling
+        acc_samples = min(num_frames, vlen)
+        # split the video into `acc_samples` intervals, and sample from each interval.
+        intervals = np.linspace(start=0, stop=vlen, num=acc_samples + 1).astype(int)
+        ranges = []
+        for idx, interv in enumerate(intervals[:-1]):
+            ranges.append((interv, intervals[idx + 1] - 1))
+        if sample == 'rand':
+            try:
+                frame_indices = [random.choice(range(x[0], x[1])) for x in ranges]
+            except:
+                frame_indices = np.random.permutation(vlen)[:acc_samples]
+                frame_indices.sort()
+                frame_indices = list(frame_indices)
+        elif fix_start is not None:
+            frame_indices = [x[0] + fix_start for x in ranges]
+        elif sample == 'middle':
+            frame_indices = [(x[0] + x[1]) // 2 for x in ranges]
+        else:
+            raise NotImplementedError
+
+        if len(frame_indices) < num_frames:  # padded with last frame
+            padded_frame_indices = [frame_indices[-1]] * num_frames
+            padded_frame_indices[:len(frame_indices)] = frame_indices
+            frame_indices = padded_frame_indices
+    elif "fps" in sample:  # fps0.5, sequentially sample frames at 0.5 fps
+        output_fps = float(sample[3:])
+        duration = float(vlen) / input_fps
+        delta = 1 / output_fps  # gap between frames, this is also the clip length each frame represents
+        frame_seconds = np.arange(0 + delta / 2, duration + delta / 2, delta)
+        frame_indices = np.around(frame_seconds * input_fps).astype(int)
+        frame_indices = [e for e in frame_indices if e < vlen]
+        if max_num_frames > 0 and len(frame_indices) > max_num_frames:
+            frame_indices = frame_indices[:max_num_frames]
+            # frame_indices = np.linspace(0 + delta / 2, duration + delta / 2, endpoint=False, num=max_num_frames)
+    else:
+        raise ValueError
+    return frame_indices
+
+
+def read_frames_decord(
+        video_path, num_frames, sample='middle', fix_start=None, 
+        max_num_frames=-1, trimmed30=False, height=-1, width=-1
+    ):
+    decord.bridge.set_bridge('torch')
+
+    # num_threads = 1 if video_path.endswith('.webm') else 0 # make ssv2 happy
+    num_threads = 1
+    video_reader = VideoReader(video_path, num_threads=num_threads, height=height, width=width)
+    try:
+        vlen = len(video_reader)
+
+        fps = video_reader.get_avg_fps()
+        duration = vlen / float(fps)
+
+        # only use top 30 seconds
+        if trimmed30 and duration > 30:
+            duration = 30
+            vlen = int(30 * float(fps))
+
+        frame_indices = get_frame_indices(
+            num_frames, vlen, sample=sample, fix_start=fix_start,
+            input_fps=fps, max_num_frames=max_num_frames
+        )
+
+        frames = video_reader.get_batch(frame_indices)  # (T, H, W, C), torch.uint8
+        if not isinstance(frames, torch.Tensor):
+            frames = torch.from_numpy(frames.asnumpy())
+        frames = frames.permute(0, 3, 1, 2)  # (T, C, H, W), torch.uint8
+        return frames
+    finally:
+        # Explicitly release underlying resources to avoid file descriptor leaks
+        del video_reader
+
+
+if __name__ == "__main__":
+
+    # Load model
+    model = TARA.from_pretrained(
+        "/work/piyush/experiments/CaRe/Tarsier-7b/final-10112025/nli_9000+ego_1000+subj_replaced-seed_42/merged_checkpoint",
+        device_map='auto',
+        dtype=torch.bfloat16,
+    )
+    n_params = sum(p.numel() for p in model.model.parameters())
+    print(f"Number of parameters: {round(n_params/1e9, 3)}B")
+    
+    # Let's encode a sample video
+    print(colored("Testing video encoding...", 'cyan'))
+    video_path = "./assets/folding_paper.mp4"
+    video_tensor = read_frames_decord(video_path, num_frames=16)
+    video_tensor = video_tensor.unsqueeze(0)
+    video_tensor = video_tensor.to(model.model.device)
+    with torch.no_grad():
+        video_emb = model.encode_vision(video_tensor).cpu().squeeze(0).float()
+    print("Video shape:", video_tensor.shape) # torch.Size([1, 16, 3, 240, 426])
+    print("Video embedding shape:", video_emb.shape) # torch.Size([4096])
+    
+    # Let's encode a sample text
+    print(colored("Testing text encoding...", 'cyan'))
+    text = ['someone is folding a paper', 'cutting a paper', 'someone is folding a paper']
+    # NOTE: It can also take a single string
+    with torch.no_grad():
+        text_emb = model.encode_text(text).cpu().float()
+    print("Text:", text)
+    print("Text embedding shape:", text_emb.shape) # torch.Size([3, 4096])
+