add pipeline video

pipeline_video/data_utils/__init__.py

@@ -0,0 +1,29 @@
+# from .processors.builder import build_processors
+from .xgpt3_dataset import MultiModalDataset
+from mplug_owl_video.processing_mplug_owl import MplugOwlImageProcessor, MplugOwlProcessor
+
+def train_valid_test_datasets_provider(data_path, config, tokenizer, seq_length=1024, loss_objective = 'sequential'):
+    """Build train and valid datasets."""
+    print('> building train and validation datasets for mPLUG-Owl ...')
+    train_ds, valid_ds = build_train_valid_test_datasets(
+        input_file=data_path,  
+        tokenizer=tokenizer,
+        max_length=seq_length, 
+        config=config, loss_objective = loss_objective)  
+    print("> finished creating mPLUG-Owl datasets ...")
+
+    return train_ds, valid_ds
+
+
+def build_train_valid_test_datasets(input_file, tokenizer, max_length=80, config=None):
+
+    # train_processors = build_processors(config['train_processors'])
+    # valid_processors = build_processors(config['valid_processors'])
+    
+    image_processor = MplugOwlImageProcessor.from_pretrained(config['pretrained_ckpt'])
+    processor = MplugOwlProcessor(image_processor, tokenizer)
+
+    assert len(input_file) == 2 # If you have files more than 2, modify code at here or merger them into train and dev
+    train_ds = MultiModalDataset(input_file[0], tokenizer, processor, max_length, loss_objective = loss_objective)
+    valid_ds = MultiModalDataset(input_file[1], tokenizer, processor, max_length, loss_objective = loss_objective)
+    return (train_ds, valid_ds)

pipeline_video/data_utils/processors/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) Alibaba. All rights reserved.
+from .builder import PROCESSORS, build_processors
+from .default_processor import DefaultProcessor
+from .caption_processor import CaptionProcessor
+
+__all__ = [
+    'PROCESSORS', 'build_processors',
+    'DefaultProcessor', 'CaptionProcessor'
+]

pipeline_video/data_utils/processors/builder.py

@@ -0,0 +1,12 @@
+import os
+import numpy as np
+
+from data_utils.registry import Registry, build_from_cfg
+
+PROCESSORS = Registry('processors')
+
+def build_processors(processors_cfg):
+    processors = dict()
+    for task, processor in processors_cfg.items():
+        processors[task] = build_from_cfg(processor, PROCESSORS)
+    return processors

pipeline_video/data_utils/processors/caption_processor.py

@@ -0,0 +1,53 @@
+import torch
+from torchvision import transforms
+from PIL import Image
+import random
+
+from data_utils.randaugment import RandomAugment
+from .builder import PROCESSORS
+
+
+@PROCESSORS.register_module()
+class CaptionProcessor:
+    def __init__(self, image_size=224, min_scale = 0.5, randaug=False):
+        self.image_size = image_size
+        self.min_scale = min_scale
+
+        if randaug:
+            self.image_transform = transforms.Compose([                        
+                transforms.RandomResizedCrop(image_size,scale=(min_scale, 1.0), interpolation=Image.BICUBIC),
+                transforms.RandomHorizontalFlip(),
+                RandomAugment(2,7,isPIL=True,augs=['Identity','AutoContrast','Equalize','Brightness','Sharpness',
+                                                'ShearX', 'ShearY', 'TranslateX', 'TranslateY', 'Rotate']),     
+                transforms.ToTensor(),
+                transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+            ])
+        else:
+            self.image_transform = transforms.Compose([                        
+                transforms.RandomResizedCrop(image_size,scale=(min_scale, 1.0), interpolation=Image.BICUBIC),
+                transforms.RandomHorizontalFlip(),  
+                transforms.ToTensor(),
+                transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+            ])
+        self.text_transform = None
+
+    def __call__(self, image, text):
+        assert image or text
+        
+        if image:
+            image_input = self.image_transform(image)
+        else:
+            image_input = None
+
+        if text:
+            if isinstance(text["prompt"], list):
+                prompt = random.choice(text["prompt"])
+            else:
+                prompt = text["prompt"]
+            text_input = dict(
+                prompt=prompt,
+                completion=text["text"],
+            )
+        else:
+            text_input = None
+        return image_input, text_input

pipeline_video/data_utils/processors/default_processor.py

@@ -0,0 +1,42 @@
+import torch
+from torchvision import transforms
+from PIL import Image
+import random
+
+from data_utils.randaugment import RandomAugment
+from .builder import PROCESSORS
+
+
+@PROCESSORS.register_module()
+class DefaultProcessor:
+    def __init__(self, image_size=224):
+        self.image_size = image_size
+
+        self.image_transform = transforms.Compose([
+            transforms.Resize((image_size, image_size),interpolation=Image.BICUBIC),
+            transforms.ToTensor(),
+            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+        ])
+
+        self.text_transform = None
+
+    def __call__(self, image, text):
+        assert image or text
+
+        if image:
+            image_input = self.image_transform(image)
+        else:
+            image_input = None
+
+        if text:
+            if isinstance(text["prompt"], list):
+                prompt = random.choice(text["prompt"])
+            else:
+                prompt = text["prompt"]
+            text_input = dict(
+                prompt=prompt,
+                completion=text["text"],
+            )
+        else:
+            text_input = None
+        return image_input, text_input

pipeline_video/data_utils/randaugment.py

@@ -0,0 +1,345 @@
+import cv2
+import numpy as np
+from PIL import Image
+
+
+## aug functions
+def identity_func(img):
+    return img
+
+
+def autocontrast_func(img, cutoff=0):
+    '''
+        same output as PIL.ImageOps.autocontrast
+    '''
+    n_bins = 256
+
+    def tune_channel(ch):
+        n = ch.size
+        cut = cutoff * n // 100
+        if cut == 0:
+            high, low = ch.max(), ch.min()
+        else:
+            hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
+            low = np.argwhere(np.cumsum(hist) > cut)
+            low = 0 if low.shape[0] == 0 else low[0]
+            high = np.argwhere(np.cumsum(hist[::-1]) > cut)
+            high = n_bins - 1 if high.shape[0] == 0 else n_bins - 1 - high[0]
+        if high <= low:
+            table = np.arange(n_bins)
+        else:
+            scale = (n_bins - 1) / (high - low)
+            offset = -low * scale
+            table = np.arange(n_bins) * scale + offset
+            table[table < 0] = 0
+            table[table > n_bins - 1] = n_bins - 1
+        table = table.clip(0, 255).astype(np.uint8)
+        return table[ch]
+
+    channels = [tune_channel(ch) for ch in cv2.split(img)]
+    out = cv2.merge(channels)
+    return out
+
+
+def equalize_func(img):
+    '''
+        same output as PIL.ImageOps.equalize
+        PIL's implementation is different from cv2.equalize
+    '''
+    n_bins = 256
+
+    def tune_channel(ch):
+        hist = cv2.calcHist([ch], [0], None, [n_bins], [0, n_bins])
+        non_zero_hist = hist[hist != 0].reshape(-1)
+        step = np.sum(non_zero_hist[:-1]) // (n_bins - 1)
+        if step == 0: return ch
+        n = np.empty_like(hist)
+        n[0] = step // 2
+        n[1:] = hist[:-1]
+        table = (np.cumsum(n) // step).clip(0, 255).astype(np.uint8)
+        return table[ch]
+
+    channels = [tune_channel(ch) for ch in cv2.split(img)]
+    out = cv2.merge(channels)
+    return out
+
+
+def rotate_func(img, degree, fill=(0, 0, 0)):
+    '''
+    like PIL, rotate by degree, not radians
+    '''
+    H, W = img.shape[0], img.shape[1]
+    center = W / 2, H / 2
+    M = cv2.getRotationMatrix2D(center, degree, 1)
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill)
+    return out
+
+
+def solarize_func(img, thresh=128):
+    '''
+        same output as PIL.ImageOps.posterize
+    '''
+    table = np.array([el if el < thresh else 255 - el for el in range(256)])
+    table = table.clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def color_func(img, factor):
+    '''
+        same output as PIL.ImageEnhance.Color
+    '''
+    ## implementation according to PIL definition, quite slow
+    #  degenerate = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[:, :, np.newaxis]
+    #  out = blend(degenerate, img, factor)
+    #  M = (
+    #      np.eye(3) * factor
+    #      + np.float32([0.114, 0.587, 0.299]).reshape(3, 1) * (1. - factor)
+    #  )[np.newaxis, np.newaxis, :]
+    M = (
+            np.float32([
+                [0.886, -0.114, -0.114],
+                [-0.587, 0.413, -0.587],
+                [-0.299, -0.299, 0.701]]) * factor
+            + np.float32([[0.114], [0.587], [0.299]])
+    )
+    out = np.matmul(img, M).clip(0, 255).astype(np.uint8)
+    return out
+
+
+def contrast_func(img, factor):
+    """
+        same output as PIL.ImageEnhance.Contrast
+    """
+    mean = np.sum(np.mean(img, axis=(0, 1)) * np.array([0.114, 0.587, 0.299]))
+    table = np.array([(
+        el - mean) * factor + mean
+        for el in range(256)
+    ]).clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def brightness_func(img, factor):
+    '''
+        same output as PIL.ImageEnhance.Contrast
+    '''
+    table = (np.arange(256, dtype=np.float32) * factor).clip(0, 255).astype(np.uint8)
+    out = table[img]
+    return out
+
+
+def sharpness_func(img, factor):
+    '''
+    The differences the this result and PIL are all on the 4 boundaries, the center
+    areas are same
+    '''
+    kernel = np.ones((3, 3), dtype=np.float32)
+    kernel[1][1] = 5
+    kernel /= 13
+    degenerate = cv2.filter2D(img, -1, kernel)
+    if factor == 0.0:
+        out = degenerate
+    elif factor == 1.0:
+        out = img
+    else:
+        out = img.astype(np.float32)
+        degenerate = degenerate.astype(np.float32)[1:-1, 1:-1, :]
+        out[1:-1, 1:-1, :] = degenerate + factor * (out[1:-1, 1:-1, :] - degenerate)
+        out = out.astype(np.uint8)
+    return out
+
+
+def shear_x_func(img, factor, fill=(0, 0, 0)):
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, factor, 0], [0, 1, 0]])
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR).astype(np.uint8)
+    return out
+
+
+def translate_x_func(img, offset, fill=(0, 0, 0)):
+    '''
+        same output as PIL.Image.transform
+    '''
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, -offset], [0, 1, 0]])
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR).astype(np.uint8)
+    return out
+
+
+def translate_y_func(img, offset, fill=(0, 0, 0)):
+    '''
+        same output as PIL.Image.transform
+    '''
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, 0], [0, 1, -offset]])
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR).astype(np.uint8)
+    return out
+
+
+def posterize_func(img, bits):
+    '''
+        same output as PIL.ImageOps.posterize
+    '''
+    out = np.bitwise_and(img, np.uint8(255 << (8 - bits)))
+    return out
+
+
+def shear_y_func(img, factor, fill=(0, 0, 0)):
+    H, W = img.shape[0], img.shape[1]
+    M = np.float32([[1, 0, 0], [factor, 1, 0]])
+    out = cv2.warpAffine(img, M, (W, H), borderValue=fill, flags=cv2.INTER_LINEAR).astype(np.uint8)
+    return out
+
+
+def cutout_func(img, pad_size, replace=(0, 0, 0)):
+    replace = np.array(replace, dtype=np.uint8)
+    H, W = img.shape[0], img.shape[1]
+    rh, rw = np.random.random(2)
+    pad_size = pad_size // 2
+    ch, cw = int(rh * H), int(rw * W)
+    x1, x2 = max(ch - pad_size, 0), min(ch + pad_size, H)
+    y1, y2 = max(cw - pad_size, 0), min(cw + pad_size, W)
+    out = img.copy()
+    out[x1:x2, y1:y2, :] = replace
+    return out
+
+
+### level to args
+def enhance_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        return ((level / MAX_LEVEL) * 1.8 + 0.1,)
+    return level_to_args
+
+
+def shear_level_to_args(MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * 0.3
+        if np.random.random() > 0.5: level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def translate_level_to_args(translate_const, MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * float(translate_const)
+        if np.random.random() > 0.5: level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def cutout_level_to_args(cutout_const, MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * cutout_const)
+        return (level, replace_value)
+
+    return level_to_args
+
+
+def solarize_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * 256)
+        return (level, )
+    return level_to_args
+
+
+def none_level_to_args(level):
+    return ()
+
+
+def posterize_level_to_args(MAX_LEVEL):
+    def level_to_args(level):
+        level = int((level / MAX_LEVEL) * 4)
+        return (level, )
+    return level_to_args
+
+
+def rotate_level_to_args(MAX_LEVEL, replace_value):
+    def level_to_args(level):
+        level = (level / MAX_LEVEL) * 30
+        if np.random.random() < 0.5:
+            level = -level
+        return (level, replace_value)
+
+    return level_to_args
+
+
+func_dict = {
+    'Identity': identity_func,
+    'AutoContrast': autocontrast_func,
+    'Equalize': equalize_func,
+    'Rotate': rotate_func,
+    'Solarize': solarize_func,
+    'Color': color_func,
+    'Contrast': contrast_func,
+    'Brightness': brightness_func,
+    'Sharpness': sharpness_func,
+    'ShearX': shear_x_func,
+    'TranslateX': translate_x_func,
+    'TranslateY': translate_y_func,
+    'Posterize': posterize_func,
+    'ShearY': shear_y_func,
+}
+
+translate_const = 10
+MAX_LEVEL = 10
+replace_value = (128, 128, 128)
+arg_dict = {
+    'Identity': none_level_to_args,
+    'AutoContrast': none_level_to_args,
+    'Equalize': none_level_to_args,
+    'Rotate': rotate_level_to_args(MAX_LEVEL, replace_value),
+    'Solarize': solarize_level_to_args(MAX_LEVEL),
+    'Color': enhance_level_to_args(MAX_LEVEL),
+    'Contrast': enhance_level_to_args(MAX_LEVEL),
+    'Brightness': enhance_level_to_args(MAX_LEVEL),
+    'Sharpness': enhance_level_to_args(MAX_LEVEL),
+    'ShearX': shear_level_to_args(MAX_LEVEL, replace_value),
+    'TranslateX': translate_level_to_args(
+        translate_const, MAX_LEVEL, replace_value
+    ),
+    'TranslateY': translate_level_to_args(
+        translate_const, MAX_LEVEL, replace_value
+    ),
+    'Posterize': posterize_level_to_args(MAX_LEVEL),
+    'ShearY': shear_level_to_args(MAX_LEVEL, replace_value),
+}
+
+
+class RandomAugment(object):
+
+    def __init__(self, N=2, M=10, isPIL=False, returnPIL=False, augs=[]):
+        self.N = N
+        self.M = M
+        self.isPIL = isPIL
+        self.returnPIL = returnPIL
+        if augs:
+            self.augs = augs       
+        else:
+            self.augs = list(arg_dict.keys())
+
+    def get_random_ops(self):
+        sampled_ops = np.random.choice(self.augs, self.N)
+        return [(op, 0.5, self.M) for op in sampled_ops]
+
+    def __call__(self, img):
+        if self.isPIL:
+            img = np.array(img)            
+        ops = self.get_random_ops()
+        for name, prob, level in ops:
+            if np.random.random() > prob:
+                continue
+            args = arg_dict[name](level)
+            img = func_dict[name](img, *args) 
+        if self.returnPIL:
+            img = img.astype('uint8')
+            img = Image.fromarray(img)
+        return img
+
+
+if __name__ == '__main__':
+    a = RandomAugment()
+    img = np.random.randn(32, 32, 3)
+    a(img)

pipeline_video/data_utils/registry.py

@@ -0,0 +1,422 @@
+# Copyright (c) Alibaba. All rights reserved.
+import inspect
+import warnings
+import functools
+from functools import partial
+from typing import Any, Dict, Optional
+from collections import abc
+from inspect import getfullargspec
+
+
+def is_seq_of(seq, expected_type, seq_type=None):
+    """Check whether it is a sequence of some type.
+    Args:
+        seq (Sequence): The sequence to be checked.
+        expected_type (type): Expected type of sequence items.
+        seq_type (type, optional): Expected sequence type.
+    Returns:
+        bool: Whether the sequence is valid.
+    """
+    if seq_type is None:
+        exp_seq_type = abc.Sequence
+    else:
+        assert isinstance(seq_type, type)
+        exp_seq_type = seq_type
+    if not isinstance(seq, exp_seq_type):
+        return False
+    for item in seq:
+        if not isinstance(item, expected_type):
+            return False
+    return True
+
+
+def deprecated_api_warning(name_dict, cls_name=None):
+    """A decorator to check if some arguments are deprecate and try to replace
+    deprecate src_arg_name to dst_arg_name.
+    Args:
+        name_dict(dict):
+            key (str): Deprecate argument names.
+            val (str): Expected argument names.
+    Returns:
+        func: New function.
+    """
+
+    def api_warning_wrapper(old_func):
+
+        @functools.wraps(old_func)
+        def new_func(*args, **kwargs):
+            # get the arg spec of the decorated method
+            args_info = getfullargspec(old_func)
+            # get name of the function
+            func_name = old_func.__name__
+            if cls_name is not None:
+                func_name = f'{cls_name}.{func_name}'
+            if args:
+                arg_names = args_info.args[:len(args)]
+                for src_arg_name, dst_arg_name in name_dict.items():
+                    if src_arg_name in arg_names:
+                        warnings.warn(
+                            f'"{src_arg_name}" is deprecated in '
+                            f'`{func_name}`, please use "{dst_arg_name}" '
+                            'instead', DeprecationWarning)
+                        arg_names[arg_names.index(src_arg_name)] = dst_arg_name
+            if kwargs:
+                for src_arg_name, dst_arg_name in name_dict.items():
+                    if src_arg_name in kwargs:
+
+                        assert dst_arg_name not in kwargs, (
+                            f'The expected behavior is to replace '
+                            f'the deprecated key `{src_arg_name}` to '
+                            f'new key `{dst_arg_name}`, but got them '
+                            f'in the arguments at the same time, which '
+                            f'is confusing. `{src_arg_name} will be '
+                            f'deprecated in the future, please '
+                            f'use `{dst_arg_name}` instead.')
+
+                        warnings.warn(
+                            f'"{src_arg_name}" is deprecated in '
+                            f'`{func_name}`, please use "{dst_arg_name}" '
+                            'instead', DeprecationWarning)
+                        kwargs[dst_arg_name] = kwargs.pop(src_arg_name)
+
+            # apply converted arguments to the decorated method
+            output = old_func(*args, **kwargs)
+            return output
+
+        return new_func
+
+    return api_warning_wrapper
+
+
+
+def build_from_cfg(cfg: Dict,
+                   registry: 'Registry',
+                   default_args: Optional[Dict] = None) -> Any:
+    """Build a module from config dict when it is a class configuration, or
+    call a function from config dict when it is a function configuration.
+
+    Example:
+        >>> MODELS = Registry('models')
+        >>> @MODELS.register_module()
+        >>> class ResNet:
+        >>>     pass
+        >>> resnet = build_from_cfg(dict(type='Resnet'), MODELS)
+        >>> # Returns an instantiated object
+        >>> @MODELS.register_module()
+        >>> def resnet50():
+        >>>     pass
+        >>> resnet = build_from_cfg(dict(type='resnet50'), MODELS)
+        >>> # Return a result of the calling function
+
+    Args:
+        cfg (dict): Config dict. It should at least contain the key "type".
+        registry (:obj:`Registry`): The registry to search the type from.
+        default_args (dict, optional): Default initialization arguments.
+
+    Returns:
+        object: The constructed object.
+    """
+    if not isinstance(cfg, dict):
+        raise TypeError(f'cfg must be a dict, but got {type(cfg)}')
+    if 'type' not in cfg:
+        if default_args is None or 'type' not in default_args:
+            raise KeyError(
+                '`cfg` or `default_args` must contain the key "type", '
+                f'but got {cfg}\n{default_args}')
+    if not isinstance(registry, Registry):
+        raise TypeError('registry must be an mmcv.Registry object, '
+                        f'but got {type(registry)}')
+    if not (isinstance(default_args, dict) or default_args is None):
+        raise TypeError('default_args must be a dict or None, '
+                        f'but got {type(default_args)}')
+
+    args = cfg.copy()
+
+    if default_args is not None:
+        for name, value in default_args.items():
+            args.setdefault(name, value)
+
+    obj_type = args.pop('type')
+    if isinstance(obj_type, str):
+        obj_cls = registry.get(obj_type)
+        if obj_cls is None:
+            raise KeyError(
+                f'{obj_type} is not in the {registry.name} registry')
+    elif inspect.isclass(obj_type) or inspect.isfunction(obj_type):
+        obj_cls = obj_type
+    else:
+        raise TypeError(
+            f'type must be a str or valid type, but got {type(obj_type)}')
+    try:
+        return obj_cls(**args)
+    except Exception as e:
+        # Normal TypeError does not print class name.
+        raise type(e)(f'{obj_cls.__name__}: {e}')
+
+
+class Registry:
+    """A registry to map strings to classes or functions.
+
+    Registered object could be built from registry. Meanwhile, registered
+    functions could be called from registry.
+
+    Example:
+        >>> MODELS = Registry('models')
+        >>> @MODELS.register_module()
+        >>> class ResNet:
+        >>>     pass
+        >>> resnet = MODELS.build(dict(type='ResNet'))
+        >>> @MODELS.register_module()
+        >>> def resnet50():
+        >>>     pass
+        >>> resnet = MODELS.build(dict(type='resnet50'))
+
+    Please refer to
+    https://mmcv.readthedocs.io/en/latest/understand_mmcv/registry.html for
+    advanced usage.
+
+    Args:
+        name (str): Registry name.
+        build_func(func, optional): Build function to construct instance from
+            Registry, func:`build_from_cfg` is used if neither ``parent`` or
+            ``build_func`` is specified. If ``parent`` is specified and
+            ``build_func`` is not given,  ``build_func`` will be inherited
+            from ``parent``. Default: None.
+        parent (Registry, optional): Parent registry. The class registered in
+            children registry could be built from parent. Default: None.
+        scope (str, optional): The scope of registry. It is the key to search
+            for children registry. If not specified, scope will be the name of
+            the package where class is defined, e.g. mmdet, mmcls, mmseg.
+            Default: None.
+    """
+
+    def __init__(self, name, build_func=None, parent=None, scope=None):
+        self._name = name
+        self._module_dict = dict()
+        self._children = dict()
+        self._scope = self.infer_scope() if scope is None else scope
+
+        # self.build_func will be set with the following priority:
+        # 1. build_func
+        # 2. parent.build_func
+        # 3. build_from_cfg
+        if build_func is None:
+            if parent is not None:
+                self.build_func = parent.build_func
+            else:
+                self.build_func = build_from_cfg
+        else:
+            self.build_func = build_func
+        if parent is not None:
+            assert isinstance(parent, Registry)
+            parent._add_children(self)
+            self.parent = parent
+        else:
+            self.parent = None
+
+    def __len__(self):
+        return len(self._module_dict)
+
+    def __contains__(self, key):
+        return self.get(key) is not None
+
+    def __repr__(self):
+        format_str = self.__class__.__name__ + \
+                     f'(name={self._name}, ' \
+                     f'items={self._module_dict})'
+        return format_str
+
+    @staticmethod
+    def infer_scope():
+        """Infer the scope of registry.
+
+        The name of the package where registry is defined will be returned.
+
+        Example:
+            >>> # in mmdet/models/backbone/resnet.py
+            >>> MODELS = Registry('models')
+            >>> @MODELS.register_module()
+            >>> class ResNet:
+            >>>     pass
+            The scope of ``ResNet`` will be ``mmdet``.
+
+        Returns:
+            str: The inferred scope name.
+        """
+        # We access the caller using inspect.currentframe() instead of
+        # inspect.stack() for performance reasons. See details in PR #1844
+        frame = inspect.currentframe()
+        # get the frame where `infer_scope()` is called
+        infer_scope_caller = frame.f_back.f_back
+        filename = inspect.getmodule(infer_scope_caller).__name__
+        split_filename = filename.split('.')
+        return split_filename[0]
+
+    @staticmethod
+    def split_scope_key(key):
+        """Split scope and key.
+
+        The first scope will be split from key.
+
+        Examples:
+            >>> Registry.split_scope_key('mmdet.ResNet')
+            'mmdet', 'ResNet'
+            >>> Registry.split_scope_key('ResNet')
+            None, 'ResNet'
+
+        Return:
+            tuple[str | None, str]: The former element is the first scope of
+            the key, which can be ``None``. The latter is the remaining key.
+        """
+        split_index = key.find('.')
+        if split_index != -1:
+            return key[:split_index], key[split_index + 1:]
+        else:
+            return None, key
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def scope(self):
+        return self._scope
+
+    @property
+    def module_dict(self):
+        return self._module_dict
+
+    @property
+    def children(self):
+        return self._children
+
+    def get(self, key):
+        """Get the registry record.
+
+        Args:
+            key (str): The class name in string format.
+
+        Returns:
+            class: The corresponding class.
+        """
+        scope, real_key = self.split_scope_key(key)
+        if scope is None or scope == self._scope:
+            # get from self
+            if real_key in self._module_dict:
+                return self._module_dict[real_key]
+        else:
+            # get from self._children
+            if scope in self._children:
+                return self._children[scope].get(real_key)
+            else:
+                # goto root
+                parent = self.parent
+                while parent.parent is not None:
+                    parent = parent.parent
+                return parent.get(key)
+
+    def build(self, *args, **kwargs):
+        return self.build_func(*args, **kwargs, registry=self)
+
+    def _add_children(self, registry):
+        """Add children for a registry.
+
+        The ``registry`` will be added as children based on its scope.
+        The parent registry could build objects from children registry.
+
+        Example:
+            >>> models = Registry('models')
+            >>> mmdet_models = Registry('models', parent=models)
+            >>> @mmdet_models.register_module()
+            >>> class ResNet:
+            >>>     pass
+            >>> resnet = models.build(dict(type='mmdet.ResNet'))
+        """
+
+        assert isinstance(registry, Registry)
+        assert registry.scope is not None
+        assert registry.scope not in self.children, \
+            f'scope {registry.scope} exists in {self.name} registry'
+        self.children[registry.scope] = registry
+
+    @deprecated_api_warning(name_dict=dict(module_class='module'))
+    def _register_module(self, module, module_name=None, force=False):
+        if not inspect.isclass(module) and not inspect.isfunction(module):
+            raise TypeError('module must be a class or a function, '
+                            f'but got {type(module)}')
+
+        if module_name is None:
+            module_name = module.__name__
+        if isinstance(module_name, str):
+            module_name = [module_name]
+        for name in module_name:
+            if not force and name in self._module_dict:
+                raise KeyError(f'{name} is already registered '
+                               f'in {self.name}')
+            self._module_dict[name] = module
+
+    def deprecated_register_module(self, cls=None, force=False):
+        warnings.warn(
+            'The old API of register_module(module, force=False) '
+            'is deprecated and will be removed, please use the new API '
+            'register_module(name=None, force=False, module=None) instead.',
+            DeprecationWarning)
+        if cls is None:
+            return partial(self.deprecated_register_module, force=force)
+        self._register_module(cls, force=force)
+        return cls
+
+    def register_module(self, name=None, force=False, module=None):
+        """Register a module.
+
+        A record will be added to `self._module_dict`, whose key is the class
+        name or the specified name, and value is the class itself.
+        It can be used as a decorator or a normal function.
+
+        Example:
+            >>> backbones = Registry('backbone')
+            >>> @backbones.register_module()
+            >>> class ResNet:
+            >>>     pass
+
+            >>> backbones = Registry('backbone')
+            >>> @backbones.register_module(name='mnet')
+            >>> class MobileNet:
+            >>>     pass
+
+            >>> backbones = Registry('backbone')
+            >>> class ResNet:
+            >>>     pass
+            >>> backbones.register_module(ResNet)
+
+        Args:
+            name (str | None): The module name to be registered. If not
+                specified, the class name will be used.
+            force (bool, optional): Whether to override an existing class with
+                the same name. Default: False.
+            module (type): Module class or function to be registered.
+        """
+        if not isinstance(force, bool):
+            raise TypeError(f'force must be a boolean, but got {type(force)}')
+        # NOTE: This is a walkaround to be compatible with the old api,
+        # while it may introduce unexpected bugs.
+        if isinstance(name, type):
+            return self.deprecated_register_module(name, force=force)
+
+        # raise the error ahead of time
+        if not (name is None or isinstance(name, str) or is_seq_of(name, str)):
+            raise TypeError(
+                'name must be either of None, an instance of str or a sequence'
+                f'  of str, but got {type(name)}')
+
+        # use it as a normal method: x.register_module(module=SomeClass)
+        if module is not None:
+            self._register_module(module=module, module_name=name, force=force)
+            return module
+
+        # use it as a decorator: @x.register_module()
+        def _register(module):
+            self._register_module(module=module, module_name=name, force=force)
+            return module
+
+        return _register

pipeline_video/data_utils/xgpt3_dataset.py

@@ -0,0 +1,204 @@
+import json
+import logging
+import os
+import random
+import re
+import time
+import traceback
+import warnings
+from io import BytesIO
+import pandas as pd
+import h5py
+import numpy as np
+import torch
+from icecream import ic
+from PIL import Image, ImageFile
+from torch.utils.data import Dataset, Subset
+
+from utils import get_args
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+ImageFile.MAX_IMAGE_PIXELS = None
+Image.MAX_IMAGE_PIXELS = None
+
+logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
+                    datefmt='%m/%d/%Y %H:%M:%S',
+                    level=logging.INFO)
+warnings.filterwarnings("ignore")
+logger = logging.getLogger(__name__)
+
+
+def load_jsonl(filename):
+    with open(filename, "r", encoding="utf-8") as f:
+        return [json.loads(l.strip("\n")) for l in f.readlines()]
+
+
+class MultiModalDataset(Dataset):
+    """MultiModal dataset"""
+
+    def __init__(self, input_file, tokenizer, processor,
+                 max_length=2048,
+                 media_tokens=['<image>', '<|video|>'], loss_objective = 'sequential'):
+        
+        args = get_args()
+
+        self.loss_objective = loss_objective
+        if 'sequential' in self.loss_objective:
+            self.dataset = pd.read_csv(input_file) 
+            self.dataset = self.dataset.dropna()
+        else:
+            raise NotImplementedError('dataset loader not implemented for other loss objectives')
+        
+        self.dataset = pd.read_csv(input_file)
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.processor = processor
+        self.media_tokens = {k: -int(i+1) for i, k in enumerate(media_tokens)}
+        self.media_lengths = {'<image>': 1+64,'<|video|>': 1+64}
+        print("num_media_token: ", self.media_lengths)
+        print(len(self.dataset))
+        self.bucket = {}
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, index):
+        
+        data = self.dataset.iloc[index]
+        videopath = data['videopath']
+        caption   = data['caption']
+        video_input = self.processor(videos=[videopath], num_frames=32, return_tensors='pt') # video_pixel_values
+        text_input = self._extract_text_token_from_conversation(caption, self.max_length, index)
+        item  = {'video': video_input, 'text': text_input, 'videopath': videopath, 'caption': caption}
+        return item
+
+    def _extract_text_token_from_conversation(self, data, max_length, index):
+        # output enc_chunk
+        enc_chunk = []
+
+        if self.tokenizer.bos_token_id > 0:
+            prompt_chunk = [self.tokenizer.bos_token_id]
+        else:
+            prompt_chunk = []
+
+        # conversation = data["completion"]
+        conversation = data
+
+        # For Text only data
+        if all([media_token not in conversation for media_token in self.media_tokens.keys()]):
+            pattern = '|'.join(map(re.escape, ['AI: ', '\nHuman: ']))
+            chunk_strs = re.split(f'({pattern})', conversation)
+            prompt_length = -1
+            stop_flag = False
+            for idx, chunk_str in enumerate(chunk_strs):
+                if idx == 0:
+                    enc_chunk = prompt_chunk + \
+                        self.tokenizer(chunk_str, add_special_tokens=False)[
+                            'input_ids']
+                    enc_length = len(enc_chunk)
+                    label_chunk = [0] * enc_length
+                else:
+                    if chunk_strs[idx-1] == 'AI: ':
+                        curr_chunk = self.tokenizer(
+                            chunk_str, add_special_tokens=False)['input_ids']
+                        if enc_length + len(curr_chunk) >= max_length:
+                            curr_chunk = curr_chunk[:max_length-enc_length]
+                            stop_flag = True
+                        curr_chunk += [self.tokenizer.eos_token_id]
+                        enc_length += len(curr_chunk)
+                        enc_chunk += curr_chunk
+                        label_chunk += [1] * len(curr_chunk)
+                    else:
+                        curr_chunk = self.tokenizer(
+                            chunk_str, add_special_tokens=False)['input_ids']
+                        if enc_length + len(curr_chunk) >= max_length + 1:
+                            curr_chunk = curr_chunk[:max_length+1-enc_length]
+                            stop_flag = True
+                        enc_length += len(curr_chunk)
+                        enc_chunk += curr_chunk
+                        label_chunk += [0] * len(curr_chunk)
+                    if stop_flag:
+                        break
+
+        # For Image-Text Data
+        else:
+            enc_length = 0
+            prompt_length = -2
+            pattern = '|'.join(
+                map(re.escape, list(self.media_tokens.keys()) + ['AI: ', '\nHuman: ']))
+            chunk_strs = re.split(f'({pattern})', conversation)
+            chunk_strs = [x for x in chunk_strs if len(x) > 0]
+            for idx, chunk_str in enumerate(chunk_strs):
+                if enc_length >= max_length + 1:
+                    break
+
+                if idx == 0:
+                    enc_chunk = prompt_chunk + \
+                        self.tokenizer(chunk_str, add_special_tokens=False)[
+                            'input_ids']
+                    enc_length = len(enc_chunk)
+                    label_chunk = [0] * enc_length
+                else:
+                    if chunk_str in self.media_tokens:
+                        # [CLS] + 256 + [EOS]
+                        if enc_length + self.media_lengths[chunk_str] > max_length + 1:
+                            break
+                        else:
+                            enc_chunk += [self.media_tokens[chunk_str]
+                                          ] * self.media_lengths[chunk_str]
+                            enc_length += self.media_lengths[chunk_str]
+                            label_chunk += [0] * self.media_lengths[chunk_str]
+                    else:
+
+                        if chunk_strs[idx-1] == 'AI: ':
+                            curr_chunk = self.tokenizer(
+                                chunk_str, add_special_tokens=False)['input_ids']
+                            if enc_length + len(curr_chunk) >= max_length:
+                                curr_chunk = curr_chunk[:max_length-enc_length]
+                            curr_chunk += [self.tokenizer.eos_token_id]
+                            enc_length += len(curr_chunk)
+                            enc_chunk += curr_chunk
+                            label_chunk += [1] * len(curr_chunk)
+                        else:
+                            curr_chunk = self.tokenizer(
+                                chunk_str, add_special_tokens=False)['input_ids']
+                            if enc_length + len(curr_chunk) >= max_length + 1:
+                                curr_chunk = curr_chunk[:max_length +
+                                                        1-enc_length]
+                            enc_length += len(curr_chunk)
+                            enc_chunk += curr_chunk
+                            label_chunk += [0] * len(curr_chunk)
+
+        if enc_length < max_length + 1:
+            padding_chunk = [self.tokenizer.pad_token_id] * \
+                (max_length + 1 - enc_length)
+            padding_length = len(padding_chunk)
+            label_chunk += [0] * (max_length + 1 - enc_length)
+            enc_chunk = enc_chunk + padding_chunk
+        else:
+            padding_length = 0
+
+        assert enc_length + padding_length == max_length + \
+            1, (index, prompt_length, enc_length,
+                padding_length, max_length + 1)
+        assert len(label_chunk) == max_length + \
+            1, (len(label_chunk), max_length + 1)
+        non_padding_mask = [1 if i < enc_length -
+                            1 else 0 for i in range(max_length)]
+
+        enc_chunk = torch.tensor(enc_chunk).long()
+        non_padding_mask = torch.tensor(non_padding_mask).long()
+        prompt_mask = torch.tensor(label_chunk)[1:].long()
+        prompt_length = torch.tensor([prompt_length]).long()
+
+        # Create loss mask
+        if all([media_token not in conversation for media_token in self.media_tokens.keys()]):
+            non_media_mask = torch.ones_like(non_padding_mask).long()
+        else:
+            tmp_enc_chunk = enc_chunk.clone()
+            tmp_enc_chunk[tmp_enc_chunk >= 0] = 1
+            tmp_enc_chunk[tmp_enc_chunk < 0] = 0
+            non_media_mask = torch.tensor(tmp_enc_chunk).long()
+            non_media_mask = non_media_mask[1:].long()
+        return {'input_ids': enc_chunk, "prompt_length": prompt_length, 'seq_length': enc_length,
+                "non_padding_mask": non_padding_mask, 'non_media_mask': non_media_mask, 'prompt_mask': prompt_mask}

pipeline_video/entailment_inference.py

@@ -0,0 +1,122 @@
+import os
+import csv
+import json
+import torch
+import argparse
+import pandas as pd
+import torch.nn as nn
+from tqdm import tqdm
+from collections import defaultdict
+from transformers.models.llama.tokenization_llama import LlamaTokenizer
+from torch.utils.data import DataLoader
+from mplug_owl_video.modeling_mplug_owl import MplugOwlForConditionalGeneration
+from mplug_owl_video.processing_mplug_owl import MplugOwlImageProcessor, MplugOwlProcessor
+from peft import LoraConfig, get_peft_model
+from data_utils.xgpt3_dataset import MultiModalDataset
+from utils import batchify
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument('--input_csv', type = str, required = True, help = 'input json file')
+parser.add_argument('--output_csv', type = str, help = 'output csv with scores')
+parser.add_argument('--pretrained_ckpt', type = str, required = True, help = 'pretrained ckpt')
+parser.add_argument('--trained_ckpt', type = str, help = 'trained ckpt')
+parser.add_argument('--lora_r', type = int, default = 32)
+parser.add_argument('--use_lora', action = 'store_true', help = 'lora model')
+parser.add_argument('--all-params', action = 'store_true', help = 'use all params of the model')
+parser.add_argument('--batch_size', type = int, default = 32)
+
+args = parser.parse_args()
+softmax = nn.Softmax(dim=2)
+
+def get_entail(logits, input_ids, tokenizer):
+    logits = softmax(logits)
+    token_id_yes = tokenizer.encode('Yes', add_special_tokens = False)[0]
+    token_id_no  = tokenizer.encode('No', add_special_tokens = False)[0]
+    entailment = []
+    for j in range(len(logits)):
+        for i in range(len(input_ids[j])):
+            if input_ids[j][i] == tokenizer.pad_token_id: # pad token if the answer is not present
+                i = i - 1
+                break
+            elif i == len(input_ids[j]) - 1:
+                break
+        score = logits[j][i][token_id_yes] / (logits[j][i][token_id_yes] + logits[j][i][token_id_no])
+        entailment.append(score)
+    entailment = torch.stack(entailment)
+    return entailment
+
+def get_scores(model, tokenizer, dataloader):
+
+    with torch.no_grad():
+        for index, inputs in tqdm(enumerate(dataloader)):
+            for k, v in inputs.items():
+                if torch.is_tensor(v):
+                    if v.dtype == torch.float:
+                        inputs[k] = v.bfloat16()
+                    inputs[k] = inputs[k].to(model.device)
+            outputs = model(pixel_values = inputs['pixel_values'], video_pixel_values = inputs['video_pixel_values'], labels = None, \
+                                num_images = inputs['num_images'], num_videos = inputs['num_videos'], input_ids = inputs['input_ids'], non_padding_mask = inputs['non_padding_mask'], \
+                                non_media_mask = inputs['non_media_mask'], prompt_mask = inputs['prompt_mask'])
+            logits = outputs['logits']
+            entail_scores = get_entail(logits, inputs['input_ids'], tokenizer)
+            for m in range(len(entail_scores)):
+                with open(args.output_csv, 'a') as f:
+                    writer = csv.writer(f)
+                    writer.writerow([inputs['videopaths'][m], inputs['captions'][m], entail_scores[m].item()])
+            print(f"Batch {index} Done")
+
+def main():
+    
+    pretrained_ckpt = args.pretrained_ckpt
+
+    # Processors
+    tokenizer = LlamaTokenizer.from_pretrained(pretrained_ckpt)
+    image_processor = MplugOwlImageProcessor.from_pretrained(pretrained_ckpt)
+    processor = MplugOwlProcessor(image_processor, tokenizer)
+
+    valid_data = MultiModalDataset(args.input_csv, tokenizer, processor, max_length = 256, loss_objective = 'sequential')
+    dataloader = DataLoader(valid_data, batch_size=args.batch_size, pin_memory=True, collate_fn=batchify)
+    
+    # Instantiate model
+    model = MplugOwlForConditionalGeneration.from_pretrained(
+        pretrained_ckpt,
+        torch_dtype=torch.bfloat16,
+        device_map={'':0}
+    )
+
+    if args.use_lora:
+        for name, param in model.named_parameters():
+            param.requires_grad = False
+        if args.all_params:
+            peft_config = LoraConfig(
+                target_modules=r'.*language_model.*\.(q_proj|v_proj|k_proj|o_proj|gate_proj|down_proj|up_proj)', 
+                inference_mode=True, 
+                r=args.lora_r, 
+                lora_alpha=16, 
+                lora_dropout=0.05
+            )
+        else:
+            peft_config = LoraConfig(
+                target_modules=r'.*language_model.*\.(q_proj|v_proj|k_proj|o_proj)', 
+                inference_mode=True, 
+                r=args.lora_r, 
+                lora_alpha=16, 
+                lora_dropout=0.05
+            )
+            
+        model = get_peft_model(model, peft_config)
+        model.print_trainable_parameters()
+
+        with open(args.trained_ckpt, 'rb') as f:
+            ckpt = torch.load(f, map_location = torch.device(f"cuda:0"))
+        model.load_state_dict(ckpt)
+        model = model.to(torch.bfloat16)
+        print('Model Loaded')
+
+    model.eval()
+
+    get_scores(model, tokenizer, dataloader)
+
+if __name__  == "__main__":
+    main()

pipeline_video/mplug_owl_video/__init__.py

@@ -0,0 +1,77 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from transformers.utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available
+
+
+_import_structure = {
+    "configuration_mplug_owl": ["MPLUG_OWL_PRETRAINED_CONFIG_ARCHIVE_MAP", "MplugOwlConfig"],
+    "processing_mplug_owl": ["MplugOwlImageProcessor", "MplugOwlProcessor"],
+    "tokenization_mplug_owl": ["MplugOwlTokenizer"],
+}
+
+try:
+    if not is_tokenizers_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+
+
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    pass
+else:
+    _import_structure["modeling_mplug_owl"] = [
+        "MPLUG_OWL_PRETRAINED_MODEL_ARCHIVE_LIST",
+        "MplugOwlForConditionalGeneration",
+        "MplugOwlModel",
+    ]
+
+
+if TYPE_CHECKING:
+    from .configuration_mplug_owl import MPLUG_OWL_PRETRAINED_CONFIG_ARCHIVE_MAP, MplugOwlConfig
+    from .tokenization_mplug_owl import MplugOwlTokenizer
+
+    try:
+        if not is_tokenizers_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        pass
+    else:
+        from .modeling_mplug_owl import (
+            MPLUG_OWL_PRETRAINED_MODEL_ARCHIVE_LIST,
+            MplugOwlForConditionalGeneration,
+            MplugOwlModel,
+            MplugOwlPreTrainedModel,
+        )
+
+
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
+
+from .configuration_mplug_owl import *
+from .modeling_mplug_owl import *
+from .processing_mplug_owl import *
+from .tokenization_mplug_owl import *

pipeline_video/mplug_owl_video/configuration_mplug_owl.py

@@ -0,0 +1,296 @@
+# coding=utf-8
+# Copyright 2022 x-plug and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" MplugOwl model configuration """
+import copy
+import os
+from typing import Union
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+from transformers.utils import logging
+from transformers.models.auto import CONFIG_MAPPING
+
+
+logger = logging.get_logger(__name__)
+
+MPLUG_OWL_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "MAGAer13/mplug-owl-llama-7b": "https://huggingface.co/MAGAer13/mplug-owl-llama-7b/resolve/main/config.json",
+    # See all MplugOwl models at https://huggingface.co/models?filter=mplug_owl
+}
+
+
+class MplugOwlVisionConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MplugOwlVisionModel`]. It is used to instantiate a
+     mPLUG-Owl vision encoder according to the specified arguments, defining the model architecture. Instantiating a
+     configuration defaults will yield a similar configuration to that of the mPLUG-Owl
+     [x-plug/x_plug-llama-7b](https://huggingface.co/x-plug/x_plug-llama-7b) architecture.
+
+     Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+     documentation from [`PretrainedConfig`] for more information.
+
+     Args:
+         hidden_size (`int`, *optional*, defaults to 768):
+             Dimensionality of the encoder layers and the pooler layer.
+         intermediate_size (`int`, *optional*, defaults to 3072):
+             Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+         num_hidden_layers (`int`, *optional*, defaults to 12):
+             Number of hidden layers in the Transformer encoder.
+         num_attention_heads (`int`, *optional*, defaults to 12):
+             Number of attention heads for each attention layer in the Transformer encoder.
+         image_size (`int`, *optional*, defaults to 224):
+             The size (resolution) of each image.
+         patch_size (`int`, *optional*, defaults to 32):
+             The size (resolution) of each patch.
+         hidden_act (`str` or `function`, *optional*, defaults to `"quick_gelu"`):
+             The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+             `"relu"`, `"selu"` and `"gelu_new"` ``"quick_gelu"` are supported.
+         layer_norm_eps (`float`, *optional*, defaults to 1e-5):
+             The epsilon used by the layer normalization layers.
+         attention_dropout (`float`, *optional*, defaults to 0.0):
+             The dropout ratio for the attention probabilities.
+         initializer_range (`float`, *optional*, defaults to 0.02):
+             The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+         initializer_factor (`float`, *optional*, defaults to 1):
+             A factor for initializing all weight matrices (should be kept to 1, used internally for initialization
+             testing).
+
+
+     ```"""
+
+    model_type = "mplug_owl_vision_model"
+
+    def __init__(
+        self,
+        hidden_size=1024,
+        intermediate_size=4096,
+        projection_dim=768,
+        num_hidden_layers=24,
+        num_attention_heads=16,
+        num_channels=3,
+        image_size=224,
+        patch_size=14,
+        hidden_act="quick_gelu",
+        layer_norm_eps=1e-6,
+        attention_dropout=0.0,
+        initializer_range=0.02,
+        initializer_factor=1.0,
+        use_flash_attn=False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.projection_dim = projection_dim
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.initializer_range = initializer_range
+        self.initializer_factor = initializer_factor
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.use_flash_attn = use_flash_attn
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        # get the vision config dict if we are loading from MplugOwlConfig
+        if config_dict.get("model_type") == "mplug-owl":
+            config_dict = config_dict["vision_config"]
+
+        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class MplugOwlVisualAbstractorConfig(PretrainedConfig):
+    model_type = "mplug_owl_visual_abstract"
+
+    def __init__(
+        self,
+        hidden_size=1024,  #
+        num_hidden_layers=6,  #
+        num_attention_heads=16,  #
+        intermediate_size=4096,  #
+        attention_probs_dropout_prob=0.1,  #
+        initializer_range=0.02,
+        layer_norm_eps=1e-6,  #
+        encoder_hidden_size=1024,  #
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.encoder_hidden_size = encoder_hidden_size
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        # get the visual_abstractor config dict if we are loading from MplugOwlConfig
+        if config_dict.get("model_type") == "mplug-owl":
+            config_dict = config_dict["abstractor_config"]
+
+        if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class MplugOwlConfig(PretrainedConfig):
+    r"""
+    [`MplugOwlConfig`] is the configuration class to store the configuration of a [`MplugOwlForConditionalGeneration`]. It is
+    used to instantiate a mPLUG-Owl model according to the specified arguments, defining the vision model, Q-Former model
+    and language model configs. Instantiating a configuration with the defaults will yield a similar configuration to
+    that of the mPLUG-Owl [x-plug/x_plug-llama-7b](https://huggingface.co/x-plug/x_plug-llama-7b) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        vision_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize [`MplugOwlVisionConfig`].
+        visual_abstractor_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize [`MplugOwlVisualAbstractorConfig`].
+        text_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize any [`PretrainedConfig`].
+        num_query_tokens (`int`, *optional*, defaults to 32):
+            The number of query tokens passed through the Transformer.
+
+        kwargs (*optional*):
+            Dictionary of keyword arguments.
+
+    Example:
+
+    ```python
+    >>> from transformers import (
+    ...     MplugOwlVisionConfig,
+    ...     MplugOwlVisualAbstractorConfig,
+    ...     OPTConfig,
+    ...     MplugOwlConfig,
+    ...     MplugOwlForConditionalGeneration,
+    ... )
+
+    >>> # Initializing a MplugOwlConfig with x-plug/x_plug-llama-7b style configuration
+    >>> configuration = MplugOwlConfig()
+
+    >>> # Initializing a MplugOwlForConditionalGeneration (with random weights) from the x-plug/x_plug-llama-7b style configuration
+    >>> model = MplugOwlForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+
+    >>> # We can also initialize a MplugOwlConfig from a MplugOwlVisionConfig, MplugOwlVisualAbstractorConfig and any PretrainedConfig
+
+    >>> # Initializing mPLUG-Owl vision, mPLUG-Owl Q-Former and language model configurations
+    >>> vision_config = MplugOwlVisionConfig()
+    >>> visual_abstractor_config = MplugOwlVisualAbstractorConfig()
+    >>> text_config = OPTConfig()
+
+    >>> config = MplugOwlConfig.from_text_vision_configs(vision_config, visual_abstractor_config, text_config)
+    ```"""
+    model_type = "mplug-owl"
+    is_composition = True
+
+    def __init__(
+        self, vision_config=None, visual_abstractor_config=None, text_config=None, num_query_tokens=64, **kwargs
+    ):
+        super().__init__(**kwargs)
+        if vision_config is None:
+            vision_config = MplugOwlVisionConfig().to_dict()
+            logger.info("vision_config is None.")
+
+        if visual_abstractor_config is None:
+            visual_abstractor_config = {}
+            logger.info("abstractor_config is None. ")
+
+        if text_config is None:
+            # we use LLAMA 7b by default
+            from ..llama.configuration_llama import LlamaConfig
+
+            text_config = LlamaConfig(pad_token_id=2).to_dict()
+            logger.info("text_config is None.")
+
+        self.vision_config = MplugOwlVisionConfig(**vision_config)
+        self.visual_abstractor_config = MplugOwlVisualAbstractorConfig(**visual_abstractor_config)
+        # self.visual_abstractor_config.layer_norm_eps = 1e-6
+        text_model_type = text_config["model_type"] if "model_type" in text_config else "llama"
+        self.text_config = CONFIG_MAPPING[text_model_type](**text_config)
+
+        self.tie_word_embeddings = self.text_config.tie_word_embeddings
+        self.is_encoder_decoder = self.text_config.is_encoder_decoder
+
+        self.num_query_tokens = num_query_tokens
+        # self.visual_abstractor_config.encoder_hidden_size = self.vision_config.hidden_size
+        self.use_decoder_only_language_model = self.text_config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+        self.initializer_factor = 1.0
+        self.initializer_range = 0.02
+
+        for attr in dir(self.text_config):
+            if not hasattr(self, attr):
+                setattr(self, attr, getattr(self.text_config, attr))
+
+    @classmethod
+    def from_vision_visual_abstractor_text_configs(
+        cls,
+        vision_config: MplugOwlVisionConfig,
+        visual_abstractor_config: MplugOwlVisualAbstractorConfig,
+        text_config: PretrainedConfig,
+        **kwargs,
+    ):
+        r"""
+        Instantiate a [`MplugOwlConfig`] (or a derived class) from a mPLUG-Owl vision model, Q-Former and language model
+        configurations.
+
+        Returns:
+            [`MplugOwlConfig`]: An instance of a configuration object
+        """
+
+        return cls(
+            vision_config=vision_config.to_dict(),
+            visual_abstractor_config=visual_abstractor_config.to_dict(),
+            text_config=text_config.to_dict(),
+            **kwargs,
+        )
+
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
+
+        Returns:
+            `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+        """
+        output = copy.deepcopy(self.__dict__)
+        output["vision_config"] = self.vision_config.to_dict()
+        output["visual_abstractor_config"] = self.visual_abstractor_config.to_dict()
+        output["text_config"] = self.text_config.to_dict()
+        output["model_type"] = self.__class__.model_type
+        return output

pipeline_video/mplug_owl_video/modeling_mplug_owl.py

@@ -0,0 +1,1938 @@
+# coding=utf-8
+# Copyright 2022 x-plug The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch MplugOwl model. """
+
+import logging
+import math
+from typing import Any, Optional, Tuple, Union
+
+try:
+    from flash_attn.flash_attn_interface import flash_attn_unpadded_func
+
+    flash_attn_func = flash_attn_unpadded_func
+except:
+    flash_attn_func = None
+    print("install flash-attn first.")
+import math
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+import einops
+
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    BaseModelOutputWithPooling,
+    BaseModelOutputWithPastAndCrossAttentions
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
+from transformers.utils import (
+    ModelOutput,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.models.auto import AutoModelForCausalLM
+from .configuration_mplug_owl import MplugOwlConfig, MplugOwlVisionConfig, MplugOwlVisualAbstractorConfig
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "MAGAer13/mplug-owl-llama-7b"
+_CONFIG_FOR_DOC = "MplugOwlConfig"
+
+
+MPLUG_OWL_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "MAGAer13/mplug-owl-llama-7b",
+    # See all MplugOwl models at https://huggingface.co/models?filter=mplug_owl
+]
+
+
+@dataclass
+class MplugOwlForConditionalGenerationModelOutput(ModelOutput):
+    """
+    Class defining the outputs of [`MPlugOwlForConditionalGeneration`].
+
+    Args:
+        loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
+            Language modeling loss from the language model.
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head of the language model.
+        vision_outputs (`BaseModelOutputWithPooling`):
+            Outputs of the vision encoder.
+
+        language_model_outputs (`CausalLMOutputWithPast` or `Seq2SeqLMOutput`):
+            Outputs of the language model.
+    """
+
+    loss: Optional[Tuple[torch.FloatTensor]] = None
+    logits: Optional[Tuple[torch.FloatTensor]] = None
+    vision_outputs: Optional[torch.FloatTensor] = None
+    language_model_outputs: Optional[Tuple[torch.FloatTensor]] = None
+
+    def to_tuple(self) -> Tuple[Any]:
+        return tuple(
+            self[k] if k not in ["vision_outputs", "language_model_outputs"] else getattr(self, k).to_tuple()
+            for k in self.keys()
+        )
+
+
+def get_ltor_masks_and_position_ids_from_embeddings(data):
+    """Build masks and position id for left to right model."""
+
+    # Extract batch size and sequence length.
+    micro_batch_size, seq_length = data.size()[:2]
+
+    # Attention mask (lower triangular).
+    att_mask_batch = 1
+    attention_mask = torch.tril(torch.ones((att_mask_batch, seq_length, seq_length), device=data.device)).view(
+        att_mask_batch, 1, seq_length, seq_length
+    )
+
+    # Loss mask.
+    loss_mask = torch.ones(data.size()[:2], dtype=torch.float, device=data.device)
+
+    # Position ids.
+    position_ids = torch.arange(seq_length, dtype=torch.long, device=data.device)
+    position_ids = position_ids.unsqueeze(0).expand_as(data[..., 0])
+
+    # Convert attention mask to binary:
+    attention_mask = attention_mask < 0.5
+
+    return attention_mask, loss_mask, position_ids
+
+
+class MplugOwlVisionEmbeddings(nn.Module):
+    def __init__(self, config: MplugOwlVisionConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+
+        self.cls_token = nn.Parameter(torch.randn(1, 1, self.hidden_size))
+
+        self.patch_embed = nn.Conv2d(
+            in_channels=3,
+            out_channels=self.hidden_size,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            bias=False,
+        )
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+
+        self.position_embedding = nn.Parameter(torch.randn(1, self.num_patches + 1, self.hidden_size))
+
+        self.pre_layernorm = LayerNormFp32(self.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+        # [B, C, T, H, W] or [B, C, H, W]
+        batch_size = pixel_values.size(0)
+        T = pixel_values.size(2) if pixel_values.dim() > 4 else 1
+        if T > 1:
+            pixel_values = einops.rearrange(pixel_values, 'b c t h w -> (b t) c h w')
+        image_embeds = self.patch_embed(pixel_values)
+        image_embeds = image_embeds.flatten(2).transpose(1, 2)
+
+        class_embeds = self.cls_token.expand(batch_size * T, 1, -1).to(image_embeds.dtype)
+        embeddings = torch.cat([class_embeds, image_embeds], dim=1)
+        embeddings = embeddings + self.position_embedding[:, : embeddings.size(1)].to(image_embeds.dtype)
+        embeddings = self.pre_layernorm(embeddings)
+        embeddings = einops.rearrange(embeddings, '(b t) n d -> b t n d', b=batch_size)
+        return embeddings
+
+
+class LayerNormFp32(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16 (by casting to float32 and back)."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x: torch.Tensor):
+        output = torch.nn.functional.layer_norm(
+            x.float(),
+            self.normalized_shape,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        )
+        return output.type_as(x)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class MplugOwlVisionLocalTemporal(nn.Module):
+    def __init__(self, config):
+        super(MplugOwlVisionLocalTemporal, self).__init__()
+
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+        self.num_patches = 1 + (self.image_size // self.patch_size) ** 2
+        self.hidden_size = config.hidden_size
+        d_bottleneck = self.hidden_size // 2
+        
+        self.ln = LayerNormFp32(self.hidden_size)
+        self.down_proj = nn.Conv3d(self.hidden_size, d_bottleneck, kernel_size=1, stride=1, padding=0)
+        self.conv = nn.Conv3d(d_bottleneck, d_bottleneck, kernel_size=(3, 1, 1), stride=1, padding=(1, 0, 0), groups=d_bottleneck)
+        self.up_proj = nn.Conv3d(d_bottleneck, self.hidden_size, kernel_size=1, stride=1, padding=0)
+
+        nn.init.constant_(self.up_proj.weight, 0)
+        nn.init.constant_(self.up_proj.bias, 0)
+
+        self.activation_func = QuickGELU()
+
+    def forward(self, x):
+        # [b, t, s, c]
+        T = x.size(1)
+        H = int((self.num_patches - 1)**0.5) 
+        cls_token, x = x[:, :, 0:1], x[:, :, 1:]
+        x = self.ln(x)
+        x = einops.rearrange(x, 'b t (h w) c -> b c t h w', h=H)
+        x = self.down_proj(x)
+        _device = x.device
+        self = self.to('cpu') # hack: cpu offloading since bfloat16 on gpu gives error with conv_depthwise3d 
+        x = x.to('cpu')
+        x = self.conv(x)
+        self = self.to(_device)
+        x = x.to(_device)
+        x = self.activation_func(x)
+        x = self.up_proj(x)
+        x = einops.rearrange(x, 'b c t h w -> b t (h w) c')
+        x = torch.cat([cls_token, x], dim = 2)
+        return x
+
+
+class MplugOwlVisionAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        if self.head_dim * self.num_heads != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+        self.scale = self.head_dim**-0.5
+        self.dropout = nn.Dropout(config.attention_dropout)
+
+        self.query_key_value = nn.Linear(self.hidden_size, 3 * self.hidden_size)
+        self.dense = nn.Linear(self.hidden_size, self.hidden_size)
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        head_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """Input shape: Batch x Time x Channel"""
+
+        bsz, seq_len, embed_dim = hidden_states.size()
+
+        mixed_qkv = self.query_key_value(hidden_states)
+
+        mixed_qkv = mixed_qkv.reshape(bsz, seq_len, self.num_heads, 3, embed_dim // self.num_heads).permute(
+            3, 0, 2, 1, 4
+        )  # [3, b, np, sq, hn]
+        query_states, key_states, value_states = (
+            mixed_qkv[0],
+            mixed_qkv[1],
+            mixed_qkv[2],
+        )
+        # if self.config.use_flash_attn and flash_attn_func is not None:
+        if False:
+            # [b*sq, np, hn]
+            query_states = query_states.permute(0, 2, 1, 3).contiguous()
+            query_states = query_states.view(query_states.size(0) * query_states.size(1), query_states.size(2), -1)
+
+            key_states = key_states.permute(0, 2, 1, 3).contiguous()
+            key_states = key_states.view(key_states.size(0) * key_states.size(1), key_states.size(2), -1)
+
+            value_states = value_states.permute(0, 2, 1, 3).contiguous()
+            value_states = value_states.view(value_states.size(0) * value_states.size(1), value_states.size(2), -1)
+
+            cu_seqlens = torch.arange(
+                0, (bsz + 1) * seq_len, step=seq_len, dtype=torch.int32, device=query_states.device
+            )
+
+            context_layer = flash_attn_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens,
+                cu_seqlens,
+                seq_len,
+                seq_len,
+                self.dropout if self.training else 0.0,
+                softmax_scale=self.scale,
+                causal=False,
+                return_attn_probs=False,
+            )
+            # [b*sq, np, hn] => [b, sq, np, hn]
+            context_layer = context_layer.view(bsz, seq_len, context_layer.size(1), context_layer.size(2))
+        else:
+            # Take the dot product between "query" and "key" to get the raw attention scores.
+            attention_scores = torch.matmul(query_states, key_states.transpose(-1, -2))
+
+            attention_scores = attention_scores * self.scale
+
+            # Normalize the attention scores to probabilities.
+            attention_probs = torch.softmax(attention_scores, dim=-1)
+
+            # This is actually dropping out entire tokens to attend to, which might
+            # seem a bit unusual, but is taken from the original Transformer paper.
+            attention_probs = self.dropout(attention_probs)
+
+            # Mask heads if we want to
+            if head_mask is not None:
+                attention_probs = attention_probs * head_mask
+
+            context_layer = torch.matmul(attention_probs, value_states).permute(0, 2, 1, 3)
+
+        new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size,)
+        context_layer = context_layer.reshape(new_context_layer_shape)
+
+        output = self.dense(context_layer)
+
+        outputs = (output, attention_probs) if output_attentions else (output, None)
+
+        return outputs
+
+
+class MplugOwlMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.activation_fn = QuickGELU()
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+class MplugOwlVisionEncoderLayer(nn.Module):
+    def __init__(self, config: MplugOwlVisionConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.temporal = MplugOwlVisionLocalTemporal(config)
+        self.self_attn = MplugOwlVisionAttention(config)
+        self.input_layernorm = LayerNormFp32(self.hidden_size, eps=config.layer_norm_eps)
+        self.mlp = MplugOwlMLP(config)
+        self.post_attention_layernorm = LayerNormFp32(self.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.FloatTensor]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, time, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+                `(config.encoder_attention_heads,)`.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+        """
+        B, T = hidden_states.size(0), hidden_states.size(1)
+        if T > 1:
+            hidden_states = hidden_states + self.temporal(hidden_states)
+        hidden_states = einops.rearrange(hidden_states, 'b t n d -> (b t) n d')
+        
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            head_mask=attention_mask,
+            output_attentions=output_attentions,
+        )
+        hidden_states = hidden_states + residual
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+
+        hidden_states = hidden_states + residual
+        hidden_states = einops.rearrange(hidden_states, '(b t) n d -> b t n d', b=B)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
+
+
+class MplugOwlPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = MplugOwlConfig
+    base_model_prefix = "mplug_owl"
+    supports_gradient_checkpointing = True
+    _keys_to_ignore_on_load_missing = [
+        r"position_ids",
+        r"language_model.encoder.embed_tokens.weight",
+        r"language_model.decoder.embed_tokens.weight",
+        r"language_model.lm_head.weight",
+    ]
+    _no_split_modules = [
+        "MplugOwlVisionEncoderLayer",
+        "LlamaDecoderLayer",
+        "MplugOwlVisualAbstractorLayer",
+        "LlamaForCausalLM",
+        "Parameter",
+    ]
+    _keep_in_fp32_modules = ["wo"]
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        factor = self.config.initializer_range
+        if isinstance(module, nn.Conv2d) or isinstance(module, nn.Embedding) or isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=factor)
+            if hasattr(module, "bias") and module.bias is not None:
+                module.bias.data.zero_()
+
+        if isinstance(module, MplugOwlVisionEmbeddings):
+            if hasattr(self.config, "vision_config"):
+                factor = self.config.vision_config.initializer_range
+            nn.init.trunc_normal_(module.position_embedding, mean=0.0, std=factor)
+            nn.init.trunc_normal_(module.cls_token, mean=0.0, std=factor)
+
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        elif isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+        elif isinstance(module, nn.Parameter):
+            raise ValueError
+            nn.init.trunc_normal_(module.data, mean=0.0, std=factor)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, MplugOwlVisionEncoder):
+            module.gradient_checkpointing = value
+
+
+MPLUG_OWL_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`MplugOwlConfig`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+MPLUG_OWL_VISION_INPUTS_DOCSTRING = r"""
+    Args:
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Pixel values. Pixel values can be obtained using [`MplugOwlProcessor`]. See [`MplugOwlProcessor.__call__`] for
+            details.
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+MPLUG_OWL_TEXT_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+            [What are attention masks?](../glossary#attention-mask)
+        decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+            Indices of decoder input sequence tokens in the vocabulary.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+            T5 uses the `pad_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values`
+            is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`).
+
+            To know more on how to prepare `decoder_input_ids` for pretraining take a look at [T5
+            Training](./t5#training).
+        decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+            Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+            be used by default.
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+MPLUG_OWL_INPUTS_DOCSTRING = r"""
+    Args:
+        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            Pixel values. Pixel values can be obtained using [`MplugOwlProcessor`]. See [`MplugOwlProcessor.__call__`] for
+            details.
+
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of input sequence tokens in the vocabulary of the language model. Input tokens can optionally be
+            provided to serve as text prompt, which the language model can continue.
+
+            Indices can be obtained using [`MplugOwlProcessor`]. See [`MplugOwlProcessor.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+        decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+            Indices of decoder input sequence tokens in the vocabulary of the language model. Only relevant in case an
+            encoder-decoder language model (like T5) is used.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details. [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+        decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+            Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+            be used by default.
+
+            Only relevant in case an encoder-decoder language model (like T5) is used.
+
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+class MplugOwlVisionEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`MplugOwlVisionEncoderLayer`].
+
+    Args:
+        config (`MplugOwlVisionConfig`):
+            The corresponding vision configuration for the `MplugOwlEncoder`.
+    """
+
+    def __init__(self, config: MplugOwlVisionConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([MplugOwlVisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        inputs_embeds,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Embedded representation of the inputs. Should be float, not int tokens.
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        hidden_states = inputs_embeds
+        for idx, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(encoder_layer),
+                    hidden_states,
+                    attention_mask,
+                )
+            else:
+                layer_outputs = encoder_layer(
+                    hidden_states,
+                    attention_mask,
+                    output_attentions=output_attentions,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
+        )
+
+
+class MplugOwlVisionModel(MplugOwlPreTrainedModel):
+    main_input_name = "pixel_values"
+    config_class = MplugOwlVisionConfig
+
+    def __init__(self, config: MplugOwlVisionConfig):
+        super().__init__(config)
+        self.config = config
+        self.hidden_size = config.hidden_size
+
+        self.embeddings = MplugOwlVisionEmbeddings(config)
+        self.encoder = MplugOwlVisionEncoder(config)
+        self.post_layernorm = LayerNormFp32(self.hidden_size, eps=config.layer_norm_eps)
+
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(MPLUG_OWL_VISION_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=MplugOwlVisionConfig)
+    def forward(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        r"""
+        Returns:
+
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if pixel_values is None:
+            raise ValueError("You have to specify pixel_values")
+
+        hidden_states = self.embeddings(pixel_values) # [B, T, N, D]
+
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        last_hidden_state = encoder_outputs[0]
+        last_hidden_state = self.post_layernorm(last_hidden_state)
+
+        pooled_output = last_hidden_state[:, :, 0, :].mean(1)
+        pooled_output = self.post_layernorm(pooled_output)
+
+        if not return_dict:
+            return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+
+class MplugOwlVisualAbstractorMLP(nn.Module):
+    def __init__(self, config: MplugOwlVisualAbstractorConfig):
+        super().__init__()
+        self.config = config
+        in_features = config.hidden_size
+        hidden_features = config.intermediate_size
+        if hidden_features != 2816:
+            hidden_features = int(2 * hidden_features / 3)
+            multiple_of = 256
+            hidden_features = multiple_of * ((hidden_features + multiple_of - 1) // multiple_of)
+        self.act = nn.SiLU()
+
+        self.w1 = nn.Linear(in_features, hidden_features)
+        self.w2 = nn.Linear(hidden_features, in_features)
+        self.w3 = nn.Linear(in_features, hidden_features)
+        self.ffn_ln = LayerNormFp32(hidden_features, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.act(self.w1(hidden_states)) * self.w3(hidden_states)
+        hidden_states = self.ffn_ln(hidden_states)
+        hidden_states = self.w2(hidden_states)
+        return hidden_states
+
+
+class MplugOwlVisualAbstractorMultiHeadAttention(nn.Module):
+    def __init__(self, config: MplugOwlVisualAbstractorConfig):
+        super().__init__()
+        self.config = config
+        if config.hidden_size % config.num_attention_heads != 0:
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention heads (%d)"
+                % (config.hidden_size, config.num_attention_heads)
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        self.key = nn.Linear(config.encoder_hidden_size, self.all_head_size)
+        self.value = nn.Linear(config.encoder_hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+        self.save_attention = False
+
+    def save_attn_gradients(self, attn_gradients):
+        self.attn_gradients = attn_gradients
+
+    def get_attn_gradients(self):
+        return self.attn_gradients
+
+    def save_attention_map(self, attention_map):
+        self.attention_map = attention_map
+
+    def get_attention_map(self):
+        return self.attention_map
+
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+    ):
+        # If this is instantiated as a cross-attention module, the keys
+        # and values come from an encoder; the attention mask needs to be
+        # such that the encoder's padding tokens are not attended to.
+        key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
+        value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
+        attention_mask = encoder_attention_mask
+
+        mixed_query_layer = self.query(hidden_states)
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+
+        past_key_value = (key_layer, value_layer)
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
+            attention_scores = attention_scores + attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+
+        if self.save_attention:
+            self.save_attention_map(attention_probs)
+            attention_probs.register_hook(self.save_attn_gradients)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs_dropped = self.dropout(attention_probs)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attention_probs_dropped = attention_probs_dropped * head_mask
+
+        context_layer = torch.matmul(attention_probs_dropped, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+
+        outputs = outputs + (past_key_value,)
+        return outputs
+
+
+class MplugOwlVisualAbstractorCrossOutput(nn.Module):
+    def __init__(self, config: MplugOwlVisualAbstractorConfig):
+        super().__init__()
+        dim = config.hidden_size
+        self.out_proj = nn.Linear(dim, dim, bias=True)
+        self.norm2 = LayerNormFp32(dim)
+        self.mlp = MplugOwlVisualAbstractorMLP(config)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        input_tensor = input_tensor + self.out_proj(hidden_states)
+        input_tensor = input_tensor + self.mlp(self.norm2(input_tensor))
+        return input_tensor
+
+
+class MplugOwlVisualAbstractorAttention(nn.Module):
+    def __init__(self, config: MplugOwlVisualAbstractorConfig):
+        super().__init__()
+        self.attention = MplugOwlVisualAbstractorMultiHeadAttention(config)
+        self.output = MplugOwlVisualAbstractorCrossOutput(config)
+        self.pruned_heads = set()
+        self.norm1 = LayerNormFp32(config.hidden_size)
+        self.normk = LayerNormFp32(config.hidden_size)
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads
+        )
+
+        # Prune linear layers
+        self.attention.query = prune_linear_layer(self.attention.query, index)
+        self.attention.key = prune_linear_layer(self.attention.key, index)
+        self.attention.value = prune_linear_layer(self.attention.value, index)
+        self.output.dense = prune_linear_layer(self.output.out_proj, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)
+        self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        output_attentions: Optional[bool] = False,
+    ) -> Tuple[torch.Tensor]:
+        # HACK we apply norm on q and k
+        hidden_states = self.norm1(hidden_states)
+        encoder_hidden_states = self.normk(encoder_hidden_states)
+        encoder_hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
+        encoder_attention_mask = torch.cat([attention_mask, encoder_attention_mask], dim=-1)
+        self_outputs = self.attention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            past_key_value,
+            output_attentions,
+        )
+        attention_output = self.output(self_outputs[0], hidden_states)
+        # add attentions if we output them
+        outputs = (attention_output,) + self_outputs[1:]
+        return outputs
+
+
+class MplugOwlVisualAbstractorLayer(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+
+        self.layer_idx = layer_idx
+
+        self.crossattention = MplugOwlVisualAbstractorAttention(config)
+        self.has_cross_attention = True
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=False,
+    ):
+        if encoder_hidden_states is None:
+            raise ValueError("encoder_hidden_states must be given for cross-attention layers")
+        cross_attention_outputs = self.crossattention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            output_attentions=output_attentions,
+        )
+        query_attention_output = cross_attention_outputs[0]
+
+        outputs = (query_attention_output,)
+        return outputs
+
+
+class MplugOwlVisualAbstractorEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList(
+            [MplugOwlVisualAbstractorLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        output_attentions=False,
+        output_hidden_states=False,
+        return_dict=True,
+    ):
+        all_hidden_states = () if output_hidden_states else None
+
+        for i in range(self.config.num_hidden_layers):
+            layer_module = self.layers[i]
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            past_key_value = past_key_values[i] if past_key_values is not None else None
+
+            if getattr(self.config, "gradient_checkpointing", False) and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, past_key_value, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    output_attentions,
+                )
+
+            hidden_states = layer_outputs[0]
+
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+        )
+
+
+class MplugOwlVisualAbstractorModel(MplugOwlPreTrainedModel):
+    def __init__(self, config: MplugOwlVisualAbstractorConfig, language_hidden_size):
+        super().__init__(config)
+        self.config = config
+
+        self.encoder = MplugOwlVisualAbstractorEncoder(config)
+        self.visual_fc = torch.nn.Linear(config.hidden_size, language_hidden_size)
+        self.temporal_visual_fc = torch.nn.Linear(config.hidden_size, language_hidden_size)
+        self.vit_eos = torch.nn.Parameter(torch.randn(1, 1, language_hidden_size))
+        nn.init.trunc_normal_(self.vit_eos, mean=0.0, std=self.config.initializer_range)
+        self.post_init()
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    def get_extended_attention_mask(
+        self,
+        attention_mask: torch.Tensor,
+        input_shape: Tuple[int],
+        device: torch.device,
+    ) -> torch.Tensor:
+        """
+        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
+
+        Arguments:
+            attention_mask (`torch.Tensor`):
+                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
+            input_shape (`Tuple[int]`):
+                The shape of the input to the model.
+            device: (`torch.device`):
+                The device of the input to the model.
+
+        Returns:
+            `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
+        """
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        if attention_mask.dim() == 3:
+            extended_attention_mask = attention_mask[:, None, :, :]
+        elif attention_mask.dim() == 2:
+            # Provided a padding mask of dimensions [batch_size, seq_length]
+            # - the model is an encoder, so make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
+            extended_attention_mask = attention_mask[:, None, None, :]
+        else:
+            raise ValueError(
+                "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
+                    input_shape, attention_mask.shape
+                )
+            )
+
+        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+        # masked positions, this operation will create a tensor which is 0.0 for
+        # positions we want to attend and -10000.0 for masked positions.
+        # Since we are adding it to the raw scores before the softmax, this is
+        # effectively the same as removing these entirely.
+        extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+        return extended_attention_mask
+
+    def forward(
+        self,
+        query_embeds,
+        temporal_query_embeds=None,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        encoder_hidden_states  (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of:
+            shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): Contains precomputed key and
+            value hidden states of the attention blocks. Can be used to speed up decoding. If `past_key_values` are
+            used, the user can optionally input only the last `decoder_input_ids` (those that don't have their past key
+            value states given to this model) of shape `(batch_size, 1)` instead of all `decoder_input_ids` of shape
+            `(batch_size, sequence_length)`.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        T = encoder_hidden_states.size(1)
+        if T == 1 or temporal_query_embeds is None:
+            embedding_output = query_embeds
+        else:
+            embedding_output = torch.cat([query_embeds, temporal_query_embeds], dim=1)
+        input_shape = embedding_output.size()[:-1]
+        batch_size, seq_length = input_shape
+        device = embedding_output.device
+
+        encoder_hidden_states = einops.rearrange(
+            encoder_hidden_states, 'b t n d -> b (t n) d'
+        )
+
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                (embedding_output.shape[0], embedding_output.shape[1]), dtype=torch.long, device=embedding_output.device
+            )
+        extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape, device)
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if encoder_hidden_states is not None:
+            if type(encoder_hidden_states) == list:
+                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].size()
+            else:
+                (
+                    encoder_batch_size,
+                    encoder_sequence_length,
+                    _,
+                ) = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+
+            if type(encoder_attention_mask) == list:
+                encoder_extended_attention_mask = [self.invert_attention_mask(mask) for mask in encoder_attention_mask]
+            elif encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+            else:
+                encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_extended_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_extended_attention_mask,
+            past_key_values=past_key_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = sequence_output[:, 0, :]
+
+        if T == 1 or temporal_query_embeds is None:
+            temporal_sequence_output = None
+        else:
+            temporal_sequence_output = sequence_output[:, query_embeds.size(1):]
+            sequence_output = sequence_output[:, :query_embeds.size(1)]
+
+        sequence_output = self.visual_fc(sequence_output)
+        if temporal_sequence_output is not None:
+            sequence_output += self.temporal_visual_fc(temporal_sequence_output)
+        sequence_output = torch.cat([sequence_output, self.vit_eos.repeat(sequence_output.shape[0], 1, 1)], dim=1)
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+        )
+
+
+@add_start_docstrings(
+    """
+    mPLUG-Owl Model for generating text and image features. The model consists of a vision encoder, Querying Transformer
+    (Q-Former) and a language model.
+    """,
+    MPLUG_OWL_START_DOCSTRING,
+)
+class MplugOwlModel(MplugOwlPreTrainedModel):
+    config_class = MplugOwlConfig
+    main_input_name = "pixel_values"
+
+    def __init__(self, config: MplugOwlConfig, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+
+        self.vision_model = MplugOwlVisionModel(config.vision_config)
+
+        self.query_tokens = nn.Parameter(
+            torch.zeros(1, config.num_query_tokens, config.visual_abstractor_config.hidden_size)
+        )
+        self.temporal_query_tokens = nn.Parameter(
+            torch.zeros(1, config.num_query_tokens, config.visual_abstractor_config.hidden_size)
+        )
+        self.abstractor = MplugOwlVisualAbstractorModel(
+            config.visual_abstractor_config, config.text_config.hidden_size
+        )
+
+        # if config.use_decoder_only_language_model:
+        # from llama.modeling_llama import LlamaForCausalLM
+        language_model = AutoModelForCausalLM.from_config(config.text_config)
+        # else:
+        #     language_model = AutoModelForSeq2SeqLM.from_config(config.text_config)
+        self.language_model = language_model
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def set_output_embeddings(self, new_embeddings):
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.language_model.get_output_embeddings()
+
+    def get_encoder(self):
+        return self.language_model.get_encoder()
+
+    def get_decoder(self):
+        return self.language_model.get_decoder()
+
+    def _tie_weights(self):
+        if not self.config.use_decoder_only_language_model:
+            self.language_model.encoder.embed_tokens = self.language_model.shared
+            self.language_model.decoder.embed_tokens = self.language_model.shared
+
+    def get_text_features(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        decoder_input_ids: Optional[torch.Tensor] = None,
+        decoder_attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if self.config.use_decoder_only_language_model:
+            text_outputs = self.language_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+        else:
+            inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+            text_outputs = self.language_model(
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                decoder_input_ids=decoder_input_ids,
+                decoder_attention_mask=decoder_attention_mask,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                labels=labels,
+            )
+
+        return text_outputs
+
+    def get_image_features(
+        self,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        return vision_outputs
+
+
+def get_media_indices(my_list):
+    if isinstance(my_list, torch.Tensor):
+        my_list = my_list.cpu().tolist()
+    result = []
+    for i in range(len(my_list)):
+        if i == 0 and my_list[i] < 0:
+            result.append(i)
+        elif my_list[i] != my_list[i - 1] and my_list[i] < 0:
+            result.append(i)
+    return result
+
+def get_media_types(tensors, positions):
+    if isinstance(tensors, torch.Tensor):
+        tensors = tensors.cpu().tolist()
+    result = []
+    for pos in positions:
+        result.append(tensors[pos])
+    return result
+
+
+@add_start_docstrings(
+    """
+    mPLUG-Owl Model for generating text given an image and an optional text prompt.
+    """,
+    MPLUG_OWL_START_DOCSTRING,
+)
+class MplugOwlForConditionalGeneration(MplugOwlPreTrainedModel):
+    config_class = MplugOwlConfig
+    main_input_name = "pixel_values"
+
+    def __init__(self, config: MplugOwlConfig):
+        super().__init__(config)
+
+        self.vision_model = MplugOwlVisionModel(config.vision_config)
+
+        self.query_tokens = nn.Parameter(
+            torch.zeros(1, config.num_query_tokens, config.visual_abstractor_config.hidden_size)
+        )
+        self.temporal_query_tokens = nn.Parameter(
+            torch.zeros(1, config.num_query_tokens, config.visual_abstractor_config.hidden_size)
+        )
+        self.abstractor = MplugOwlVisualAbstractorModel(
+            config.visual_abstractor_config, config.text_config.hidden_size
+        )
+
+        # if config.use_decoder_only_language_model:
+        # from llama.modeling_llama import LlamaForCausalLM
+        language_model = AutoModelForCausalLM.from_config(config.text_config)
+        # else:
+        #     language_model = AutoModelForSeq2SeqLM.from_config(config.text_config)
+        self.language_model = language_model
+
+        # Initialize weights and apply final processing
+        self.post_init()
+        self.main_input_name = "input_ids"
+        from transformers import GenerationConfig
+
+        self.generation_config = GenerationConfig(
+            max_length=512, do_sample=True, top_k=3, pad_token_id=0, unk_token_id=0, bos_token_id=1, eos_token_id=2
+        )
+
+        # Hack Bloom
+        if config.text_config.model_type == 'bloom':
+            bound_method = bloom_forward.__get__(self.language_model.transformer, self.language_model.transformer.__class__)
+            setattr(self.language_model.transformer, 'forward', bound_method)
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    def set_output_embeddings(self, new_embeddings):
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.language_model.get_output_embeddings()
+
+    def get_encoder(self):
+        return self.language_model.get_encoder()
+
+    def get_decoder(self):
+        return self.language_model.get_decoder()
+
+    def _tie_weights(self):
+        if not self.config.use_decoder_only_language_model:
+            self.language_model.encoder.embed_tokens = self.language_model.shared
+            self.language_model.decoder.embed_tokens = self.language_model.shared
+
+    def _preprocess_accelerate(self):
+        r"""
+        Some pre-processing hacks to make the model `accelerate` compatible. Check
+        https://github.com/huggingface/transformers/pull/21707 for more details.
+        """
+        hf_device_map = self.hf_device_map
+
+        if len(hf_device_map) > 1 and "language_model" not in hf_device_map and torch.cuda.device_count() > 1:
+            # warn users about unexpected behavior when using multi-GPU + mPLUG-Owl + `accelerate`.
+            logger.warning(
+                "The `language_model` is not in the `hf_device_map` dictionary and you are running your script"
+                " in a multi-GPU environment. this may lead to unexpected behavior when using `accelerate`."
+                " Please pass a `device_map` that contains `language_model` to remove this warning."
+                " Please refer to https://github.com/huggingface/blog/blob/main/accelerate-large-models.md for"
+                " more details on creating a `device_map` for large models.",
+            )
+
+        if hasattr(self.language_model, "_hf_hook"):
+            self.language_model._hf_hook.io_same_device = True  # For `generate` compatibility
+
+    @add_start_docstrings_to_model_forward(MPLUG_OWL_INPUTS_DOCSTRING)
+    @replace_return_docstrings(
+        output_type=MplugOwlForConditionalGenerationModelOutput, config_class=MplugOwlVisionConfig
+    )
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        video_pixel_values: torch.FloatTensor,
+        input_ids: torch.FloatTensor,
+        num_images,
+        num_videos,
+        non_padding_mask: Optional[torch.LongTensor] = None,
+        non_media_mask: Optional[torch.LongTensor] = None,
+        prompt_mask: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        decoder_attention_mask: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        labels: Optional[torch.LongTensor] = None,
+        return_dict: Optional[bool] = None,
+        **forward_kwargs,
+    ) -> Union[Tuple, MplugOwlForConditionalGenerationModelOutput]:
+        r"""
+        Returns:
+
+        Examples:
+
+        Image captioning (without providing a text prompt):
+
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import MplugOwlProcessor, MplugOwlForConditionalGeneration
+        >>> import torch
+
+        >>> device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        >>> processor = MplugOwlProcessor.from_pretrained("x-plug/x_plug-llama-7b")
+        >>> model = MplugOwlForConditionalGeneration.from_pretrained(
+        ...     "x-plug/x_plug-llama-7b", torch_dtype=torch.float16
+        ... )
+        >>> model.to(device)  # doctest: +IGNORE_RESULT
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+
+        >>> inputs = processor(images=image, return_tensors="pt").to(device, torch.float16)
+
+        >>> generated_ids = model.generate(**inputs)
+        >>> generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+        >>> print(generated_text)
+        two cats laying on a couch
+        ```
+
+        Visual question answering (prompt = question):
+
+        ```python
+        >>> from PIL import Image
+        >>> import requests
+        >>> from transformers import MplugOwlProcessor, MplugOwlForConditionalGeneration
+        >>> import torch
+
+        >>> device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        >>> processor = MplugOwlProcessor.from_pretrained("x-plug/x_plug-llama-7b")
+        >>> model = MplugOwlForConditionalGeneration.from_pretrained(
+        ...     "x-plug/x_plug-llama-7b", torch_dtype=torch.float16
+        ... )
+        >>> model.to(device)  # doctest: +IGNORE_RESULT
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> image = Image.open(requests.get(url, stream=True).raw)
+
+        >>> prompt = "Question: how many cats are there? Answer:"
+        >>> inputs = processor(images=image, text=prompt, return_tensors="pt").to(device, torch.float16)
+
+        >>> generated_ids = model.generate(**inputs)
+        >>> generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+        >>> print(generated_text)
+        two
+        ```"""
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if attention_mask is None:
+            attention_mask = input_ids.new_ones(*input_ids.shape)
+
+        # get text embedding
+        text_tokens_ = input_ids.clone()
+        batch_size = input_ids.shape[0]
+
+        media_token_indices = [
+            # [:-1] since we would not use the last token for embedding
+            get_media_indices(text_tokens_[i][:-1])
+            for i in range(batch_size)
+        ]
+
+        media_token_types = [
+            get_media_types(text_tokens_[i][:-1], media_token_indices[i])
+            for i in range(batch_size)
+        ]
+
+        text_tokens_[text_tokens_ < 0] = 1  # Not used
+        inputs_embeds = self.get_input_embeddings()(text_tokens_)  # Temporally Embedding
+
+        if hasattr(self.language_model, 'transformer') and hasattr(self.language_model.transformer, 'word_embeddings_layernorm'):
+            inputs_embeds = self.language_model.transformer.word_embeddings_layernorm(inputs_embeds)
+
+        if pixel_values is not None:
+            image_embeds = self.vision_model(pixel_values, return_dict=True).last_hidden_state
+
+            image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device)
+            query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
+            temporal_query_tokens = self.temporal_query_tokens.expand(image_embeds.shape[0], -1, -1)
+
+            query_features = self.abstractor(
+                query_embeds=query_tokens,
+                encoder_hidden_states=image_embeds,
+                encoder_attention_mask=image_attention_mask,
+            )["last_hidden_state"]
+            img_seq_length = query_features.shape[1]
+
+        if video_pixel_values is not None:
+            video_embeds = self.vision_model(video_pixel_values, return_dict=True).last_hidden_state
+
+            video_attention_mask = torch.ones(video_embeds.size()[:-1], dtype=torch.long, device=video_embeds.device)
+            video_attention_mask = einops.rearrange(
+                video_attention_mask, 'b t n -> b (t n)'
+            )
+            query_tokens = self.query_tokens.expand(video_embeds.shape[0], -1, -1)
+            temporal_query_tokens = self.temporal_query_tokens.expand(video_embeds.shape[0], -1, -1)
+
+            video_query_features = self.abstractor(
+                query_embeds=query_tokens,
+                temporal_query_embeds=temporal_query_tokens,
+                encoder_hidden_states=video_embeds,
+                encoder_attention_mask=video_attention_mask,
+            )["last_hidden_state"]
+            video_embeds   = video_query_features
+            vid_seq_length = video_query_features.shape[1]
+
+        num_images_per_sample = num_images.long().cpu().tolist()
+        num_videos_per_sample = num_videos.long().cpu().tolist()
+
+        text_chunk_embeds = []
+        text_chunk_attns = []
+        img_idx = 0
+        vid_idx = 0
+        for b in range(batch_size):
+            start = 0
+            result = []
+            result_attn = []
+            for i, pos in enumerate(media_token_indices[b]):
+                curr_image_idx, curr_video_idx = 0, 0
+                if pos > start:
+                    result.append(inputs_embeds[b, start:pos])
+                    result_attn.append(attention_mask[b, start:pos])
+                if media_token_types[b][i] == -1:
+                    result.append(image_embeds[img_idx + curr_image_idx])
+                    result_attn.append(torch.ones(image_embeds[img_idx + curr_image_idx].shape[0], device=inputs_embeds.device))
+                    start = pos + img_seq_length
+                    curr_image_idx += 1
+                else:
+                    result.append(video_embeds[vid_idx + curr_video_idx])
+                    result_attn.append(torch.ones(video_embeds[vid_idx + curr_video_idx].shape[0], device=inputs_embeds.device))
+                    start = pos + vid_seq_length
+                    curr_video_idx += 1
+            if start < inputs_embeds.shape[1]:
+                result.append(inputs_embeds[b, start:])
+                result_attn.append(attention_mask[b, start:])
+
+            img_idx += num_images_per_sample[b]
+            vid_idx += num_videos_per_sample[b]
+            text_chunk_embeds.append(torch.cat(result, dim=0))
+            text_chunk_attns.append(torch.cat(result_attn, dim=0))
+
+        inputs_embeds  = torch.stack(text_chunk_embeds, dim=0)
+        attention_mask = torch.stack(text_chunk_attns, dim=0)
+        
+        if labels is not None:
+            # Create causal mask and position ids
+            _, loss_mask, position_ids = get_ltor_masks_and_position_ids_from_embeddings(inputs_embeds)
+
+            # Calculate the loss_mask
+            non_padding_mask = non_padding_mask.long()
+            non_media_mask = non_media_mask.long()
+            prompt_mask = prompt_mask.long()  # TODO How to deal with prompt mask
+            loss_mask = loss_mask[:, :-1]
+
+            loss_mask = loss_mask * non_padding_mask * non_media_mask * prompt_mask
+            labels[:, 1:][loss_mask != 1] = -100
+                
+        # Forward into GPT
+        outputs = self.language_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            labels=labels,
+            return_dict=return_dict,
+            output_attentions=self.config.output_attentions,
+        )
+
+        return outputs
+
+    @torch.no_grad()
+    def generate(
+        self,
+        pixel_values: torch.FloatTensor = None,
+        video_pixel_values: torch.FloatTensor = None,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        isdecoder=True,
+        get_logits_only=False,
+        **generate_kwargs,
+    ) -> torch.LongTensor:
+        """
+        Overrides `generate` function to be able to use the model as a conditional generator.
+
+        Args:
+            pixel_values (`torch.FloatTensor` of shape (batch_size, num_channels, height, width)):
+                Input images to be processed.
+            input_ids (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*):
+                The sequence used as a prompt for the generation.
+            attention_mask (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*):
+                Mask to avoid performing attention on padding token indices
+
+        Returns:
+            captions (list): A list of strings of length batch_size * num_captions.
+        """
+        if input_ids is None:
+            return self.language_model.generate(attention_mask=attention_mask, **generate_kwargs)
+
+        if attention_mask is None:
+            attention_mask = input_ids.new_ones(*input_ids.shape)
+
+        batch_size = input_ids.size(0)
+        media_token_indices = [get_media_indices(input_ids[i]) for i in range(batch_size)]
+        media_token_types = [
+            get_media_types(input_ids[i], media_token_indices[i])
+            for i in range(batch_size)
+        ]
+        num_images_per_sample = [len([y for y in x if y==-1]) for x in media_token_types]
+        num_videos_per_sample = [len([y for y in x if y<-1]) for x in media_token_types]
+        input_ids = input_ids.clone()  # prevent inplace modify
+        input_ids[input_ids < 0] = 0  # Not used
+
+        if hasattr(self, "hf_device_map"):
+            # preprocess for `accelerate`
+            self._preprocess_accelerate()
+
+        batch_size = input_ids.shape[0]
+        # get text embedding
+        inputs_embeds = self.get_input_embeddings()(input_ids)
+        if hasattr(self.language_model, 'transformer') and hasattr(self.language_model.transformer, 'word_embeddings_layernorm'):
+            inputs_embeds = self.language_model.transformer.word_embeddings_layernorm(inputs_embeds)
+        # get visual embedding
+        if pixel_values is not None:
+            pixel_values = pixel_values.to(input_ids.device)
+            with torch.no_grad():
+                image_embeds = self.vision_model(pixel_values, return_dict=True).last_hidden_state
+                image_attention_mask = torch.ones(
+                    image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device
+                )
+                query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
+                query_outputs = self.abstractor(
+                    query_embeds=query_tokens,
+                    encoder_hidden_states=image_embeds,
+                    encoder_attention_mask=image_attention_mask,
+                    return_dict=True,
+                )
+                query_output = query_outputs["last_hidden_state"]
+                image_embeds = query_output
+            img_seq_length = image_embeds.shape[1]
+
+        if video_pixel_values is not None:
+            video_pixel_values = video_pixel_values.to(input_ids.device)
+            with torch.no_grad():
+                video_embeds = self.vision_model(video_pixel_values, return_dict=True).last_hidden_state
+                video_attention_mask = torch.ones(
+                    video_embeds.size()[:-1], dtype=torch.long, device=video_embeds.device
+                )
+                video_attention_mask = einops.rearrange(
+                    video_attention_mask, 'b t n -> b (t n)'
+                )
+                query_tokens = self.query_tokens.expand(video_embeds.shape[0], -1, -1)
+                temporal_query_tokens = self.temporal_query_tokens.expand(video_embeds.shape[0], -1, -1)
+                query_outputs = self.abstractor(
+                    query_embeds=query_tokens,
+                    temporal_query_embeds=temporal_query_tokens,
+                    encoder_hidden_states=video_embeds,
+                    encoder_attention_mask=video_attention_mask,
+                    return_dict=True,
+                )
+                query_output = query_outputs["last_hidden_state"]
+                video_embeds = query_output
+            vid_seq_length = video_embeds.shape[1]
+
+        # ===================
+        # Get actual input embeddings
+        # ===================
+        text_chunk_embeds = []
+        text_chunk_attns = []
+        img_idx = 0
+        vid_idx = 0
+
+        for b in range(batch_size):
+            start = 0
+            result = []
+            result_attn = []
+            for i, pos in enumerate(media_token_indices[b]):
+                curr_image_idx, curr_video_idx = 0, 0
+                if pos > start:
+                    result.append(inputs_embeds[b, start:pos])
+                    result_attn.append(attention_mask[b, start:pos])
+                if media_token_types[b][i] == -1:
+                    result.append(image_embeds[img_idx + curr_image_idx])
+                    result_attn.append(torch.ones(image_embeds[img_idx + curr_image_idx].shape[0], device=inputs_embeds.device))
+                    start = pos + img_seq_length
+                    curr_image_idx += 1
+                else:
+                    result.append(video_embeds[vid_idx + curr_video_idx])
+                    result_attn.append(torch.ones(video_embeds[vid_idx + curr_video_idx].shape[0], device=inputs_embeds.device))
+                    start = pos + vid_seq_length
+                    curr_video_idx += 1
+            if start < inputs_embeds.shape[1]:
+                result.append(inputs_embeds[b, start:])
+                result_attn.append(attention_mask[b, start:])
+
+            img_idx += num_images_per_sample[b]
+            vid_idx += num_videos_per_sample[b]
+            text_chunk_embeds.append(torch.cat(result, dim=0))
+            text_chunk_attns.append(torch.cat(result_attn, dim=0))
+        inputs_embeds = torch.stack(text_chunk_embeds, dim=0)
+        attention_mask = torch.stack(text_chunk_attns, dim=0)
+
+        if get_logits_only:
+            outputs = self.language_model(
+                        inputs_embeds=inputs_embeds,
+                        attention_mask=attention_mask,
+                        return_dict=True,
+                        output_attentions=self.config.output_attentions,
+                    )
+        else:
+            outputs = self.language_model.generate(
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                **generate_kwargs,
+            )
+            
+        return outputs
+
+    def prepare_inputs_for_generation(
+        self, input_ids, pixel_values=None, video_pixel_values=None,
+        past_key_values=None, attention_mask=None, **model_kwargs
+    ):
+        input_shape = input_ids.shape
+        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
+        if attention_mask is None:
+            attention_mask = input_ids.new_ones(input_shape)
+
+        # # cut decoder_input_ids if past_key_values is used
+        # if past_key_values is not None:
+        #     input_ids = input_ids[:, -1:]
+
+        return {
+            "input_ids": input_ids,
+            "pixel_values": pixel_values,
+            "video_pixel_values": video_pixel_values,
+            "attention_mask": attention_mask,
+            # "past_key_values": past_key_values,
+            # "encoder_hidden_states": model_kwargs.get("encoder_hidden_states", None),
+            # "encoder_attention_mask": model_kwargs.get("encoder_attention_mask", None),
+            "is_decoder": True,
+        }
+
+
+def bloom_forward(
+    self,
+    input_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    head_mask: Optional[torch.LongTensor] = None,
+    inputs_embeds: Optional[torch.LongTensor] = None,
+    use_cache: Optional[bool] = None,
+    output_attentions: Optional[bool] = None,
+    output_hidden_states: Optional[bool] = None,
+    return_dict: Optional[bool] = None,
+    **deprecated_arguments,
+) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
+    if deprecated_arguments.pop("position_ids", False) is not False:
+        # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`
+        warnings.warn(
+            "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"
+            " passing `position_ids`.",
+            FutureWarning,
+        )
+    if len(deprecated_arguments) > 0:
+        raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")
+
+    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+    output_hidden_states = (
+        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+    )
+    use_cache = use_cache if use_cache is not None else self.config.use_cache
+    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+    if input_ids is not None and inputs_embeds is not None:
+        raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+    elif input_ids is not None:
+        batch_size, seq_length = input_ids.shape
+    elif inputs_embeds is not None:
+        batch_size, seq_length, _ = inputs_embeds.shape
+    else:
+        raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+    if past_key_values is None:
+        past_key_values = tuple([None] * len(self.h))
+
+    # Prepare head mask if needed
+    # 1.0 in head_mask indicate we keep the head
+    # attention_probs has shape batch_size x num_heads x N x N
+    # head_mask has shape n_layer x batch x num_heads x N x N
+    head_mask = self.get_head_mask(head_mask, self.config.n_layer)
+
+    if inputs_embeds is None:
+        inputs_embeds = self.word_embeddings(input_ids)
+        inputs_embeds = self.word_embeddings_layernorm(inputs_embeds)
+    
+    hidden_states = inputs_embeds
+
+    presents = () if use_cache else None
+    all_self_attentions = () if output_attentions else None
+    all_hidden_states = () if output_hidden_states else None
+
+    if self.gradient_checkpointing and self.training:
+        if use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+            )
+            use_cache = False
+
+    # Compute alibi tensor: check build_alibi_tensor documentation
+    seq_length_with_past = seq_length
+    past_key_values_length = 0
+    if past_key_values[0] is not None:
+        past_key_values_length = past_key_values[0][0].shape[2]
+        seq_length_with_past = seq_length_with_past + past_key_values_length
+    if attention_mask is None:
+        attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)
+    else:
+        attention_mask = attention_mask.to(hidden_states.device)
+
+    alibi = self.build_alibi_tensor(attention_mask, self.num_heads, dtype=hidden_states.dtype)
+
+    causal_mask = self._prepare_attn_mask(
+        attention_mask,
+        input_shape=(batch_size, seq_length),
+        past_key_values_length=past_key_values_length,
+    )
+
+    for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if self.gradient_checkpointing and self.training:
+
+            def create_custom_forward(module):
+                def custom_forward(*inputs):
+                    # None for past_key_value
+                    return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)
+
+                return custom_forward
+
+            outputs = torch.utils.checkpoint.checkpoint(
+                create_custom_forward(block),
+                hidden_states,
+                alibi,
+                causal_mask,
+                layer_past,
+                head_mask[i],
+            )
+        else:
+            outputs = block(
+                hidden_states,
+                layer_past=layer_past,
+                attention_mask=causal_mask,
+                head_mask=head_mask[i],
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                alibi=alibi,
+            )
+
+        hidden_states = outputs[0]
+        if use_cache is True:
+            presents = presents + (outputs[1],)
+
+        if output_attentions:
+            all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
+
+    # Add last hidden state
+    hidden_states = self.ln_f(hidden_states)
+
+    if output_hidden_states:
+        all_hidden_states = all_hidden_states + (hidden_states,)
+
+    if not return_dict:
+        return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
+
+    return BaseModelOutputWithPastAndCrossAttentions(
+        last_hidden_state=hidden_states,
+        past_key_values=presents,
+        hidden_states=all_hidden_states,
+        attentions=all_self_attentions,
+    )

pipeline_video/mplug_owl_video/processing_mplug_owl.py

@@ -0,0 +1,246 @@
+import re
+import random
+import torch
+import torch.utils.checkpoint
+
+from transformers.processing_utils import ProcessorMixin
+from transformers.tokenization_utils_base import BatchEncoding
+from transformers.models.clip.image_processing_clip import CLIPImageProcessor
+from .tokenization_mplug_owl import MplugOwlTokenizer
+
+from decord import VideoReader
+import numpy as np
+from PIL import Image
+
+def get_index(num_frames, num_segments):
+    seg_size = float(num_frames - 1) / num_segments
+    start = int(seg_size / 2)
+    offsets = np.array([
+        start + int(np.round(seg_size * idx)) for idx in range(num_segments)
+    ])
+    return offsets
+
+def load_video(path, num_frames=4):
+    vr = VideoReader(path, height=224, width=224)
+    total_frames = len(vr)
+    frame_indices = get_index(total_frames, num_frames)
+    images_group = list()
+    for frame_index in frame_indices:
+        img = Image.fromarray(vr[frame_index].asnumpy()).convert('RGB')
+        images_group.append(img)
+    return images_group
+
+class MplugOwlProcessor(ProcessorMixin):
+    attributes = []
+    tokenizer_class = ("MplugOwlTokenizer")
+    
+    def __init__(self, image_processor=None, tokenizer=None, **kwargs):
+    
+        super().__init__(**kwargs)
+        self.tokens_to_generate = 0
+        self.image_processor = image_processor
+        self.tokenizer = tokenizer
+        self.add_BOS = True
+
+    def __call__(self, videos=None, text=None, num_frames=4, return_tensors=None, **kwargs):
+
+        if text is not None:
+            encoding = tokenize_prompts(
+                prompts=text,
+                tokens_to_generate=self.tokens_to_generate,
+                add_BOS=self.add_BOS,
+                tokenizer=self.tokenizer,
+                ignore_dist=True,
+                **kwargs,
+            )
+
+        if videos is not None:
+            video_features = []
+            for video in videos:
+                video_frames = load_video(video, num_frames)
+                video_feature = self.image_processor(video_frames, return_tensors=return_tensors, **kwargs)['pixel_values']
+                video_features.append(video_feature)
+            video_features = torch.stack(video_features, dim=0)
+            video_features = video_features.permute(0, 2, 1, 3, 4)
+    
+        if text is not None and videos is not None:
+            encoding["video_pixel_values"] = video_features
+            return encoding
+        if text is not None and videos is None:
+            return encoding
+            
+        return video_features
+
+    def batch_decode(self, skip_special_tokens=True, *args, **kwargs):
+        """
+        This method forwards all its arguments to CLIPTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
+        refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, skip_special_tokens=skip_special_tokens, **kwargs)
+
+    def decode(self, skip_special_tokens=True, *args, **kwargs):
+        """
+        This method forwards all its arguments to CLIPTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
+        the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, skip_special_tokens=skip_special_tokens, **kwargs)
+
+
+class MplugOwlImageProcessor(CLIPImageProcessor):
+    pass
+
+
+def detokenize_generations(tokens_gpu_tensor, lengths_gpu_tensor, return_segments, tokenizer):
+    """Detokenize the generated tokens."""
+
+    prompts_plus_generations = []
+    if return_segments:
+        prompts_plus_generations_segments = []
+
+    tokens = tokens_gpu_tensor.cpu().numpy().tolist()
+    lengths = lengths_gpu_tensor.cpu().numpy().tolist()
+    for sequence_tokens, length in zip(tokens, lengths):
+        sequence_tokens = sequence_tokens[:length]
+        prompts_plus_generations.append(tokenizer.detokenize(sequence_tokens))
+        if return_segments:
+            from tokenizers.decoders import Metaspace
+
+            if hasattr(tokenizer, "tokenizer"):
+                if isinstance(tokenizer.tokenizer.decoder, Metaspace):
+                    words = tokenizer.tokenizer.decode(sequence_tokens)
+                else:
+                    words = []
+                    for token in sequence_tokens:
+                        word = tokenizer.tokenizer.decoder[token]
+                        word = bytearray([tokenizer.tokenizer.byte_decoder[c] for c in word]).decode(
+                            "utf-8", errors="replace"
+                        )
+                        words.append(word)
+                prompts_plus_generations_segments.append(words)
+            else:
+                words = tokenizer.detokenize(sequence_tokens)
+                # else:
+                #     words = []
+                #     for token in sequence_tokens:
+                #         word = tokenizer.tokenizer.decoder[token]
+                #         word = bytearray(
+                #             [tokenizer.tokenizer.byte_decoder[c] for c in word]).decode(
+                #                 'utf-8', errors='replace')
+                #         words.append(word)
+                prompts_plus_generations_segments.append(words)
+
+    if return_segments:
+        return tokens, prompts_plus_generations, prompts_plus_generations_segments
+
+    return tokens, prompts_plus_generations
+
+
+def tokenize_prompts(
+    prompts=None, tokens_to_generate=None, add_BOS=None, rank=0, tokenizer=None, ignore_dist=False, **kwargs
+):
+    """Tokenize prompts and make them avaiable on all ranks."""
+
+    # On all ranks set to None so we can pass them to functions
+    prompts_tokens_cuda_long_tensor = None
+    prompts_length_cuda_long_tensor = None
+
+    # On the specified rank, build the above.
+    attention_mask = None
+    if ignore_dist or torch.distributed.get_rank() == rank:
+        assert prompts is not None
+        assert tokens_to_generate is not None
+        # Tensor of tokens padded and their unpadded length.
+        prompts_tokens_cuda_long_tensor, prompts_length_cuda_long_tensor, attention_mask = _tokenize_prompts_and_batch(
+            prompts, tokens_to_generate, add_BOS, tokenizer, **kwargs
+        )
+        # We need the sizes of these tensors for the boradcast
+        [
+            prompts_tokens_cuda_long_tensor.size(0),  # Batch size
+            prompts_tokens_cuda_long_tensor.size(1),
+        ]  # Sequence lenght
+
+    return {
+        "input_ids": prompts_tokens_cuda_long_tensor,
+        "attention_mask": attention_mask,
+        # "prompt_length": prompts_length_cuda_long_tensor,
+    }
+
+
+def _tokenize_prompts_and_batch(prompts, tokens_to_generate, add_BOS, tokenizer, **kwargs):
+    """Given a set of prompts and number of tokens to generate:
+    - tokenize prompts
+    - set the sequence length to be the max of length of prompts
+      plus the number of tokens we would like to generate
+    - pad all the sequences to this length so we can convert them
+      into a 2D tensor.
+    """
+
+    # Tokenize all the prompts.
+    # if add_BOS:
+    #     prompts_tokens = [[tokenizer.bos] + tokenizer.tokenize(prompt)
+    #                       for prompt in prompts]
+    # else:
+    #     prompts_tokens = [tokenizer.tokenize(prompt) for prompt in prompts]
+
+    prompts_tokens = [_tokenize_prompt(prompt, tokenizer, add_BOS, **kwargs) for prompt in prompts]
+
+    # Now we have a list of list of tokens which each list has a different
+    # size. We want to extend this list to:
+    #   - incorporate the tokens that need to be generated
+    #   - make all the sequences equal length.
+    # Get the prompts length.
+    prompts_length = [len(prompt_tokens) for prompt_tokens in prompts_tokens]
+    # Get the max prompts length.
+    max_prompt_len = max(prompts_length)
+    # Number of tokens in the each sample of the batch.
+    samples_length = max_prompt_len + tokens_to_generate
+    # Now update the list of list to be of the same size: samples_length.
+    for prompt_tokens, prompt_length in zip(prompts_tokens, prompts_length):
+        padding_size = samples_length - prompt_length
+        prompt_tokens.extend([tokenizer.eos_token_id] * padding_size)
+
+    # Now we are in a structured format, we can convert to tensors.
+    prompts_tokens_tensor = torch.LongTensor(prompts_tokens)
+    prompts_length_tensor = torch.LongTensor(prompts_length)
+    attention_mask = torch.zeros(prompts_tokens_tensor.shape[:2])
+    for i, l in enumerate(prompts_length_tensor):
+        attention_mask[i, :l] = 1
+    return prompts_tokens_tensor, prompts_length_tensor, attention_mask
+
+
+def _tokenize_prompt(
+        prompt, tokenizer, add_BOS=False, 
+        media_info={"<image>": 65, "<|video|>": 65},
+        **kwargs
+    ):
+    media_tokens = {k: -int(i + 1) for i, k in enumerate(media_info.keys())}
+    media_lengths = media_info.copy()
+
+    if add_BOS:
+        prompt_chunk = [tokenizer.bos_token_id]
+    else:
+        prompt_chunk = []
+
+    # Pure Text
+    if all([media_token not in prompt for media_token in media_tokens.keys()]):
+        enc_chunk = prompt_chunk + tokenizer(prompt, add_special_tokens=False, **kwargs)["input_ids"]
+
+    # Multi-Modal Text
+    else:
+        enc_chunk = prompt_chunk
+        pattern = "|".join(map(re.escape, list(media_tokens.keys())))
+        chunk_strs = re.split(f"({pattern})", prompt)
+        chunk_strs = [x for x in chunk_strs if len(x) > 0]
+        for idx, chunk_str in enumerate(chunk_strs):
+            if chunk_str in media_tokens:
+                enc_chunk += [media_tokens[chunk_str]] * media_lengths[chunk_str]
+            else:
+                tmp_chunk = tokenizer(chunk_str, add_special_tokens=False)["input_ids"]
+                # if idx < len(chunk_strs) - 1: # Last chunk should not have eos
+                #     tmp_chunk += [tokenizer.eod_id]
+                enc_chunk += tmp_chunk
+    return enc_chunk
+
+
+if __name__ == "__main__":
+    pass

pipeline_video/mplug_owl_video/tokenization_mplug_owl.py

@@ -0,0 +1,62 @@
+# coding=utf-8
+# Copyright 2022 x-plug and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tokenization classes for MplugOwl."""
+
+from transformers.utils import logging
+from transformers.models.llama.tokenization_llama import LlamaTokenizer
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "MAGAer13/mplug-owl-llama-7b": "https://huggingface.co/MAGAer13/mplug-owl-llama-7b/resolve/main/vocab.txt",
+    },
+}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "MAGAer13/mplug-owl-llama-7b": 2048,
+}
+
+
+class MplugOwlTokenizer(LlamaTokenizer):
+    def __init__(
+        self,
+        vocab_file,
+        unk_token="<unk>",
+        bos_token="<s>",
+        eos_token="</s>",
+        pad_token="<unk>",
+        sp_model_kwargs=None,
+        add_bos_token=False,
+        add_eos_token=False,
+        clean_up_tokenization_spaces=False,
+        **kwargs,
+    ):
+        super().__init__(
+            vocab_file,
+            unk_token,
+            bos_token,
+            eos_token,
+            pad_token,
+            sp_model_kwargs,
+            add_bos_token,
+            add_eos_token,
+            clean_up_tokenization_spaces,
+            **kwargs,
+        )
+        self.eod_id = self.eos_token_id

pipeline_video/nle_inference.py

@@ -0,0 +1,126 @@
+import os
+import csv
+import json
+import torch
+import argparse
+import pandas as pd
+from tqdm import tqdm
+from peft import LoraConfig, get_peft_model
+from torch.utils.data import Dataset, DataLoader
+from transformers.models.llama.tokenization_llama import LlamaTokenizer
+from mplug_owl_video.modeling_mplug_owl import MplugOwlForConditionalGeneration
+from mplug_owl_video.processing_mplug_owl import MplugOwlImageProcessor, MplugOwlProcessor
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument('--input_file', type = str, required = True, help = 'input csv file')
+parser.add_argument('--output_file', type = str, help = 'output csv file')
+parser.add_argument('--pretrained_ckpt', type = str, required = True, help = 'pretrained ckpt')
+parser.add_argument('--trained_ckpt', type = str, help = 'trained ckpt')
+parser.add_argument('--lora_r', type = int, default = 32)
+parser.add_argument('--use_lora', action = 'store_true', help = 'lora model')
+parser.add_argument('--all_params', action = 'store_true', help = 'all params')
+parser.add_argument('--batch_size', type = int, default = 1)
+parser.add_argument('--num_frames', type = int, default = 32)
+
+args = parser.parse_args()
+
+PROMPT_FEEDBACK = '''The following is a conversation between a curious human and AI assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.
+Human: <|video|>
+Human: What is the misalignment between this video and the description: "{caption}"?
+AI: '''
+
+generate_kwargs = {
+    'do_sample': True,
+    'top_k': 5,
+    'max_length': 512
+}
+
+class VideoCaptionDataset(Dataset):
+
+    def __init__(self, input_file):
+        self.data = pd.read_csv(input_file)
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, index):
+        item = {}
+        item['videopath'] = self.data.iloc[index]['videopath']
+        item['neg_caption'] = self.data.iloc[index]['neg_caption']
+        return item
+
+def get_nle(args, model, processor, tokenizer, dataloader):
+
+    with torch.no_grad():
+        for _, batch in tqdm(enumerate(dataloader)):
+            videopaths = batch['videopath']
+            neg_caption = batch['neg_caption'][0]
+            prompts = [PROMPT_FEEDBACK.format(caption = neg_caption)] 
+            inputs = processor(text=prompts, videos=videopaths, num_frames=args.num_frames, return_tensors='pt')
+            inputs = {k: v.bfloat16() if v.dtype == torch.float else v for k, v in inputs.items()}
+            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+            res = model.generate(**inputs, **generate_kwargs)
+            generated_nle = tokenizer.decode(res.tolist()[0], skip_special_tokens=True)
+
+            with open(args.output_file, 'a') as f:
+                writer = csv.writer(f)
+                writer.writerow([videopaths[0], neg_caption, generated_nle])
+
+def main():
+
+    # Create dataloader
+    dataset = VideoCaptionDataset(args.input_file)
+    dataloader = DataLoader(dataset, batch_size = args.batch_size)
+
+    pretrained_ckpt = args.pretrained_ckpt
+
+    # Processors
+    tokenizer = LlamaTokenizer.from_pretrained(pretrained_ckpt)
+    image_processor = MplugOwlImageProcessor.from_pretrained(pretrained_ckpt)
+    processor = MplugOwlProcessor(image_processor, tokenizer)
+
+    # Instantiate model
+    model = MplugOwlForConditionalGeneration.from_pretrained(
+        pretrained_ckpt,
+        torch_dtype=torch.bfloat16,
+        device_map={'':0}
+    )
+
+    if args.use_lora:
+        for name, param in model.named_parameters():
+            param.requires_grad = False
+        if args.all_params:
+            peft_config = LoraConfig(
+                target_modules=r'.*language_model.*\.(q_proj|v_proj|k_proj|o_proj|gate_proj|down_proj|up_proj)', 
+                inference_mode=True, 
+                r=args.lora_r, 
+                lora_alpha=16, 
+                lora_dropout=0.05
+            )
+        else:
+            peft_config = LoraConfig(
+                target_modules=r'.*language_model.*\.(q_proj|v_proj|k_proj|o_proj)', 
+                inference_mode=True, 
+                r=args.lora_r, 
+                lora_alpha=16, 
+                lora_dropout=0.05
+            )
+
+        model = get_peft_model(model, peft_config)
+        model.print_trainable_parameters()
+        with open(args.trained_ckpt, 'rb') as f:
+            ckpt = torch.load(f, map_location = torch.device(f"cuda:0"))
+        model.load_state_dict(ckpt)
+        model = model.to(torch.bfloat16)
+        print('Model Loaded')
+        
+    model.eval()
+
+    # get nle
+    get_nle(args, model, processor, tokenizer, dataloader)
+
+
+
+if __name__  == "__main__":
+    main()

pipeline_video/train.py

@@ -0,0 +1,263 @@
+import argparse
+from functools import partial
+import os
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from sconf import Config
+from icecream import ic
+from peft import LoraConfig, get_peft_model
+from transformers import Trainer
+from transformers.training_args import TrainingArguments
+
+from mplug_owl_video.modeling_mplug_owl import MplugOwlForConditionalGeneration
+from transformers.models.llama.tokenization_llama import LlamaTokenizer
+from data_utils import train_valid_test_datasets_provider
+from utils import batchify, set_args
+
+
+parser = argparse.ArgumentParser()
+# Model
+parser.add_argument('--pretrained-ckpt', type=str, default='MAGAer13/mplug-owl-llama-7b-pt',
+                    help='Path to the pretrained checkpoint.')
+parser.add_argument('--finetuned-ckpt', type=str, default=None,
+                    help='Path to the finetuned checkpoint.')
+parser.add_argument('--inference_mode', type=bool, default=False,
+                    help='The inference mode.')
+parser.add_argument('--seq-length', type=int, default=1024,
+                    help='Maximum sequence length to process.')
+
+parser.add_argument('--use-lora', action='store_true', help='LORA.')
+parser.add_argument('--all-params', action='store_true', help='All params in LORA')
+parser.add_argument('--lora-r', type=int, default=8,
+                    help='curvature.')
+parser.add_argument('--lora-alpha', type=int, default=32,
+                    help='The initialization coefficient of lora-alpha.')  
+parser.add_argument('--lora-dropout', type=int, default=0.05,
+                    help='The initialization coefficient of lora_dropout.')
+parser.add_argument('--bf16', action='store_true', default=False,
+                    help='Run model in bfloat16 mode.')
+
+parser.add_argument('--wandb_run_name', type=str, default="test", help='wandb run name.')
+
+# Data
+parser.add_argument('--mm-config', type=str, default=None, help='Multimodal Config.')
+parser.add_argument('--num-workers', type=int, default=8,
+                    help="Dataloader number of workers.")  
+
+# Training HyperParameters
+parser.add_argument('--train-epochs', type=int, default=3,
+                    help='Total number of epochs to train over all '
+                    'training runs.')
+parser.add_argument('--micro-batch-size', type=int, default=None,
+                    help='Batch size per model instance (local batch size). '
+                    'Global batch size is local batch size times data '
+                    'parallel size times number of micro batches.')
+parser.add_argument('--lr', type=float, default=None,
+                    help='Initial learning rate. Depending on decay style '
+                    'and initial warmup, the learing rate at each '
+                    'iteration would be different.')
+parser.add_argument('--min-lr', type=float, default=1e-6,
+                    help='Minumum value for learning rate. The scheduler'
+                    'clip values below this threshold.')
+parser.add_argument('--weight-decay', type=float, default=0.01,
+                    help='Weight decay coefficient for L2 regularization.')
+parser.add_argument('--gradient-accumulation-steps', type=int, default=8,
+                    help='The gradient accumulation steps.')
+parser.add_argument('--clip-grad', type=float, default=1.0,
+                    help='Gradient clipping based on global L2 norm.')
+parser.add_argument('--adam-beta1', type=float, default=0.9,
+                    help='First coefficient for computing running averages '
+                    'of gradient and its square')
+parser.add_argument('--adam-beta2', type=float, default=0.999,
+                    help='Second coefficient for computing running averages '
+                    'of gradient and its square')
+parser.add_argument('--adam-eps', type=float, default=1e-08,
+                    help='Term added to the denominator to improve'
+                    'numerical stability')
+
+parser.add_argument('--num-warmup-steps', type=int, default=50,
+                    help='The number of warmup steps.')
+parser.add_argument('--num-training-steps', type=int, default=4236,
+                    help='The number of total training steps for lr scheduler.')
+parser.add_argument('--loss_objective', default = 'sequential', choices = ['sequential'], help = 'toggle loss objectives')
+
+# Evaluation & Save
+parser.add_argument('--save-path', type=str, default=None,
+                    help='Output directory to save checkpoints to.')
+parser.add_argument('--save-interval', type=int, default=None,
+                    help='Number of iterations between checkpoint saves.')
+parser.add_argument('--eval-iters', type=int, default=100,
+                    help='Number of iterations to run for evaluation'
+                    'validation/test for.')
+
+# Other
+parser.add_argument('--gradient-checkpointing', action='store_true',
+                    help='The gradient checkpointing.')
+parser.add_argument('--logging-nan-inf-filter', action='store_true',
+                    help='The logging nan inf filter.')
+parser.add_argument('--ddp-find-unused-parameters', action='store_true',
+                    help='unused parameters finding.')
+parser.add_argument('--do-train', action='store_true', default=True,
+                    help='Whether to do training.')  
+parser.add_argument('--local_rank', type=int, default=-1,
+                    help='Local rank')
+
+softmax = nn.Softmax(dim=2)
+sigm = torch.nn.Sigmoid()
+
+
+class CustomTrainer(Trainer):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def get_train_dataloader(self) -> DataLoader:
+        dataset = self.train_dataset
+        sampler = DistributedSampler(dataset)
+        return torch.utils.data.DataLoader(
+            dataset, batch_size=self._train_batch_size,
+            sampler=sampler,
+            num_workers=self.args.dataloader_num_workers,
+            drop_last=True,
+            pin_memory=True,
+            collate_fn=batchify)
+
+    def get_eval_dataloader(self, eval_dataset) -> DataLoader:
+        dataset = self.eval_dataset
+        sampler = DistributedSampler(dataset, shuffle=False)
+        return torch.utils.data.DataLoader(
+            dataset, batch_size=self._train_batch_size,
+            sampler=sampler,
+            num_workers=self.args.dataloader_num_workers,
+            drop_last=True,
+            pin_memory=True,
+            collate_fn=batchify)
+
+    def compute_loss(self, model, inputs, return_outputs = False):
+        outputs = model(pixel_values = inputs['pixel_values'], video_pixel_values = inputs['video_pixel_values'], labels = inputs['labels'], 
+                            num_images = inputs['num_images'], num_videos = inputs['num_videos'], input_ids = inputs['input_ids'], non_padding_mask = inputs['non_padding_mask'], \
+                            non_media_mask = inputs['non_media_mask'], prompt_mask = inputs['prompt_mask'])
+        loss = outputs.loss
+        return loss
+        
+    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys = None):
+        for k, v in inputs.items():
+            if torch.is_tensor(v):
+                if v.dtype == torch.float:
+                    inputs[k] = v.bfloat16()
+                inputs[k] = inputs[k].to(model.device)
+        with torch.no_grad():
+            loss = self.compute_loss(model, inputs)
+            loss = loss.detach()
+        return loss, None, None
+
+def main():
+    args, left_argv = parser.parse_known_args()  
+    ic(left_argv)
+    config = Config(args.mm_config)
+
+    set_args(args)
+    print(args.pretrained_ckpt)
+    model = MplugOwlForConditionalGeneration.from_pretrained(
+        args.pretrained_ckpt,
+        torch_dtype=torch.bfloat16 if args.bf16 else torch.half,
+    )
+    tokenizer = LlamaTokenizer.from_pretrained(args.pretrained_ckpt)
+    if args.use_lora:
+        for name, param in model.named_parameters():
+            param.requires_grad = False
+        if args.all_params:
+            peft_config = LoraConfig(
+                target_modules=r'.*language_model.*\.(q_proj|v_proj|k_proj|o_proj|gate_proj|down_proj|up_proj)', 
+                inference_mode=args.inference_mode, 
+                r=args.lora_r, 
+                lora_alpha=args.lora_alpha, 
+                lora_dropout=args.lora_dropout
+            )
+        else:
+            peft_config = LoraConfig(
+                target_modules=r'.*language_model.*\.(q_proj|v_proj|k_proj|o_proj)', 
+                inference_mode=args.inference_mode, 
+                r=args.lora_r, 
+                lora_alpha=args.lora_alpha, 
+                lora_dropout=args.lora_dropout
+            )
+        model = get_peft_model(model, peft_config)
+        model.print_trainable_parameters()
+
+        if args.gradient_checkpointing:
+            def make_inputs_require_grad(module, input, output):
+                output.requires_grad_(True)
+            model.language_model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
+            model.language_model.apply(
+                partial(model.language_model._set_gradient_checkpointing, value=True))
+
+    else:
+        for name, param in model.named_parameters():
+            if 'language_model' in name:
+                param.requires_grad = True
+            else:
+                param.requires_grad = False
+        if args.gradient_checkpointing:
+            model.language_model.apply(
+                partial(model.language_model._set_gradient_checkpointing, value=True))
+
+    model.train()
+
+    train_data, valid_data = train_valid_test_datasets_provider(
+        config.data_files, config=config, 
+        tokenizer=tokenizer, seq_length=args.seq_length, loss_objective = args.loss_objective
+    )
+
+    if len(valid_data) > 500:
+        valid_data = torch.utils.data.Subset(valid_data, range(500))
+        
+    trainer = CustomTrainer(
+        model=model,
+        train_dataset=train_data,
+        eval_dataset=valid_data,
+        args=TrainingArguments(
+            learning_rate=args.lr,
+            warmup_steps=args.num_warmup_steps,
+            do_train=args.do_train,
+            do_eval=True,
+            num_train_epochs=args.train_epochs,
+            output_dir=args.save_path,
+            save_strategy='epoch',
+            evaluation_strategy='steps',
+            eval_steps=args.eval_iters,
+            per_device_train_batch_size=args.micro_batch_size,
+            max_grad_norm=args.clip_grad,
+            weight_decay=args.weight_decay,
+            bf16=args.bf16,
+            fp16=not args.bf16,
+            gradient_accumulation_steps=args.gradient_accumulation_steps,
+            gradient_checkpointing=args.gradient_checkpointing,
+            logging_steps=args.eval_iters//10,
+            logging_dir=args.save_path,            
+            logging_nan_inf_filter=args.logging_nan_inf_filter,
+            ddp_find_unused_parameters=args.ddp_find_unused_parameters,
+            run_name=args.wandb_run_name,
+            prediction_loss_only=True,
+        ),
+    )
+    trainer.loss_objective = args.loss_objective
+    trainer.tokenizer = tokenizer
+
+    if torch.__version__ >= "2" and sys.platform != "win32":
+        model = torch.compile(model)
+    
+    if args.local_rank == 0:
+        with open(os.path.join(args.save_path, "params.txt"), "w") as file:
+            for key in sorted(vars(args)):
+                value = getattr(args, key)
+                file.write(f"{key}: {value}\n")
+
+    trainer.train()
+
+    model.save_pretrained(args.save_path)
+
+if __name__ == '__main__':
+    main()

pipeline_video/utils.py

@@ -0,0 +1,160 @@
+import math
+import random
+import torch
+import numpy as np
+from icecream import ic
+
+def print_rank_0(message):
+    """If distributed is initialized, print only on rank 0."""
+    if torch.distributed.is_initialized():
+        if torch.distributed.get_rank() == 0:
+            print(message, flush=True)
+    else:
+        print(message, flush=True)
+
+ARGS = None
+def set_args(args):
+    global ARGS
+    ARGS = args
+
+def get_args():
+    return ARGS
+
+TOKENIZER = None
+def set_tokenizer(tokenizer):
+    global TOKENIZER
+    TOKENIZER = tokenizer
+
+def get_tokenizer():
+    return TOKENIZER
+from torch import distributed as dist
+
+class worker_init:
+    def __init__(self, epoch_id):
+        self.epoch_id = epoch_id
+    def _worker_init_fn(self, worker_id):
+        random.seed(worker_id + self.epoch_id*1e4 + dist.get_rank()*1e8)
+
+
+def batchify(batch):
+    # collate_fn
+    video = [data["video"] if data["video"] is not None else None for data in batch]
+    if all([img is None for img in video]):
+        video = None
+    else:
+        video = torch.cat([img for img in video if img is not None], dim=0)
+    num_videos_per_sample = torch.LongTensor([data["video"].size(0) if data['video'] is not None else 0 for data in batch])
+    num_images_per_sample = torch.LongTensor([0 for data in batch])
+
+    text = torch.stack([torch.LongTensor(data["text"]['input_ids']) for data in batch], dim=0)
+    non_padding_mask = torch.stack([torch.LongTensor(data["text"]['non_padding_mask']) for data in batch], dim=0)
+    non_media_mask = torch.stack([torch.LongTensor(data["text"]['non_media_mask']) for data in batch], dim=0)
+    prompt_mask = torch.stack([torch.LongTensor(data["text"]['prompt_mask']) for data in batch], dim=0)
+    videopaths = [data["videopath"] for data in batch]
+    captions = [data["caption"] for data in batch] 
+    output_batch = {
+        "pixel_values": None,
+        "video_pixel_values": video,
+        "input_ids": text.long(),
+        "labels": text.long().clone(),
+        "num_images": num_images_per_sample.long(),
+        "num_videos": num_videos_per_sample.long(),
+        "non_padding_mask": non_padding_mask.long(),
+        "non_media_mask": non_media_mask.long(),
+        "prompt_mask": prompt_mask.long(),
+        "videopaths": videopaths,
+        "captions": captions,
+    }
+    
+    return output_batch
+
+
+def get_param_groups(modules,
+                     no_weight_decay_cond,
+                     scale_lr_cond,
+                     lr_mult):
+    """creates param groups based on weight decay condition (regularized vs non regularized)
+       and learning rate scale condition (args.lr vs lr_mult * args.lr)
+       scale_lr_cond is used during finetuning where head of the network requires a scaled
+       version of the base learning rate. 
+    """
+    wd_no_scale_lr = []
+    wd_scale_lr = []
+    no_wd_no_scale_lr = []
+    no_wd_scale_lr = []
+    for module in modules:
+        for name, param in module.named_parameters():
+            if not param.requires_grad:
+                continue
+
+            if no_weight_decay_cond is not None:
+                no_wd = no_weight_decay_cond(name, param)
+            else:
+                # do not regularize biases nor Norm parameters
+                no_wd = name.endswith(".bias") or len(param.shape) == 1
+
+            if scale_lr_cond is not None:
+                scale_lr = scale_lr_cond(name, param)
+            else:
+                scale_lr = False
+
+            if not no_wd and not scale_lr:
+                wd_no_scale_lr.append(param)
+            elif not no_wd and scale_lr:
+                wd_scale_lr.append(param)
+            elif no_wd and not scale_lr:
+                no_wd_no_scale_lr.append(param)
+            else:
+                no_wd_scale_lr.append(param)
+
+    param_groups = []
+    if len(wd_no_scale_lr):
+        param_groups.append(
+            {'params': wd_no_scale_lr, 'wd_mult': 1.0, 'lr_mult': 1.0})
+    if len(wd_scale_lr):
+        param_groups.append(
+            {'params': wd_scale_lr, 'wd_mult': 1.0, 'lr_mult': lr_mult})
+    if len(no_wd_no_scale_lr):
+        param_groups.append({'params': no_wd_no_scale_lr,
+                            'wd_mult': 0.0, 'lr_mult': 1.0})
+    if len(no_wd_scale_lr):
+        param_groups.append(
+            {'params': no_wd_scale_lr, 'wd_mult': 0.0, 'lr_mult': lr_mult})
+
+    return param_groups
+
+def get_cosine_schedule_with_warmup(
+        optimizer, lr, min_lr, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1
+    ):
+        """
+        Create a schedule with a learning rate that decreases following the values of the cosine function between the
+        initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the
+        initial lr set in the optimizer.
+
+        Args:
+            optimizer ([`~torch.optim.Optimizer`]):
+                The optimizer for which to schedule the learning rate.
+            num_warmup_steps (`int`):
+                The number of steps for the warmup phase.
+            num_training_steps (`int`):
+                The total number of training steps.
+            num_cycles (`float`, *optional*, defaults to 0.5):
+                The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0
+                following a half-cosine).
+            last_epoch (`int`, *optional*, defaults to -1):
+                The index of the last epoch when resuming training.
+
+        Return:
+            `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+        """
+
+        delta_min_lr = (lr-min_lr)/lr  # 0.95
+
+        def lr_lambda(current_step):
+            if current_step < num_warmup_steps:
+                return (1-delta_min_lr) + delta_min_lr * float(current_step) / float(max(1, num_warmup_steps))
+            progress = float(current_step - num_warmup_steps) / \
+                float(max(1, num_training_steps - num_warmup_steps))
+            return delta_min_lr + (1-delta_min_lr) * max(0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))
+        from torch.optim.lr_scheduler import LambdaLR
+        return LambdaLR(optimizer, lr_lambda, last_epoch)