Upload 4 files

configuration_cogagent.py

@@ -1,4 +1,5 @@
 from typing import Literal
+
 from transformers import PretrainedConfig
 
 
@@ -6,27 +7,26 @@ class CogAgentConfig(PretrainedConfig):
     _auto_class = "AutoConfig"
 
     def __init__(
-            self,
-            vocab_size=32000,
-            hidden_size=4096,
-            cross_hidden_size=1024,
-            cross_compute_hidden_size=1024,
-            cross_image_size=1120,
-            intermediate_size=11008,
-            num_hidden_layers=32,
-            num_attention_heads=32,
-            hidden_act='silu',
-            max_position_embeddings=2048,
-            initializer_range=0.02,
-            rms_norm_eps=1e-06,
-            template_version: Literal["base", "chat"] = "chat",
-
-            pad_token_id=0,
-            bos_token_id=1,
-            eos_token_id=2,
-            tie_word_embeddings=False,
-            use_cache=True,
-            **kwargs,
+        self,
+        vocab_size=32000,
+        hidden_size=4096,
+        cross_hidden_size=1024,
+        cross_compute_hidden_size=1024,
+        cross_image_size=1120,
+        intermediate_size=11008,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        hidden_act="silu",
+        max_position_embeddings=2048,
+        initializer_range=0.02,
+        rms_norm_eps=1e-06,
+        template_version: Literal["base", "chat"] = "chat",
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        use_cache=True,
+        **kwargs,
     ):
         self.hidden_size = hidden_size
         self.cross_hidden_size = cross_hidden_size

cross_visual.py

@@ -1,30 +1,36 @@
+import logging
 from math import pi
+
 import torch
-from torch import nn
 from einops import rearrange, repeat
-import logging
+from torch import nn
+
 
-def broadcat(tensors, dim = -1):
+def broadcat(tensors, dim=-1):
     num_tensors = len(tensors)
     shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
-    assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions'
+    assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
     shape_len = list(shape_lens)[0]
     dim = (dim + shape_len) if dim < 0 else dim
     dims = list(zip(*map(lambda t: list(t.shape), tensors)))
     expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
-    assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]), 'invalid dimensions for broadcastable concatentation'
+    assert all(
+        [*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]
+    ), "invalid dimensions for broadcastable concatentation"
     max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
     expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
     expanded_dims.insert(dim, (dim, dims[dim]))
     expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
     tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
-    return torch.cat(tensors, dim = dim)
+    return torch.cat(tensors, dim=dim)
+
 
 def rotate_half(x):
-    x = rearrange(x, '... (d r) -> ... d r', r = 2)
-    x1, x2 = x.unbind(dim = -1)
-    x = torch.stack((-x2, x1), dim = -1)
-    return rearrange(x, '... d r -> ... (d r)')
+    x = rearrange(x, "... (d r) -> ... d r", r=2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, "... d r -> ... (d r)")
+
 
 class VisionRotaryEmbeddingFast(nn.Module):
     def __init__(
@@ -32,31 +38,34 @@ class VisionRotaryEmbeddingFast(nn.Module):
         dim,
         pt_seq_len,
         ft_seq_len=None,
-        custom_freqs = None,
-        freqs_for = 'lang',
-        theta = 10000,
-        max_freq = 10,
-        num_freqs = 1,
-        patch_dropout = 0.
+        custom_freqs=None,
+        freqs_for="lang",
+        theta=10000,
+        max_freq=10,
+        num_freqs=1,
+        patch_dropout=0.0,
     ):
         super().__init__()
         if custom_freqs:
             freqs = custom_freqs
-        elif freqs_for == 'lang':
-            freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
-        elif freqs_for == 'pixel':
-            freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
-        elif freqs_for == 'constant':
+        elif freqs_for == "lang":
+            freqs = 1.0 / (
+                theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)
+            )
+        elif freqs_for == "pixel":
+            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
+        elif freqs_for == "constant":
             freqs = torch.ones(num_freqs).float()
         else:
-            raise ValueError(f'unknown modality {freqs_for}')
+            raise ValueError(f"unknown modality {freqs_for}")
 
-        if ft_seq_len is None: ft_seq_len = pt_seq_len
+        if ft_seq_len is None:
+            ft_seq_len = pt_seq_len
         t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
 
-        freqs = torch.einsum('..., f -> ... f', t, freqs)
-        freqs = repeat(freqs, '... n -> ... (n r)', r = 2)
-        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim = -1)
+        freqs = torch.einsum("..., f -> ... f", t, freqs)
+        freqs = repeat(freqs, "... n -> ... (n r)", r=2)
+        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
 
         freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
         freqs_sin = freqs.sin().view(-1, freqs.shape[-1])
@@ -66,7 +75,7 @@ class VisionRotaryEmbeddingFast(nn.Module):
         self.register_buffer("freqs_cos", freqs_cos)
         self.register_buffer("freqs_sin", freqs_sin)
 
-        logging.info(f'Shape of rope freq: {self.freqs_cos.shape}')
+        logging.info(f"Shape of rope freq: {self.freqs_cos.shape}")
 
     def forward(self, t, patch_indices_keep=None):
         if patch_indices_keep is not None:
@@ -74,44 +83,46 @@ class VisionRotaryEmbeddingFast(nn.Module):
             batch_indices = torch.arange(batch)
             batch_indices = batch_indices[..., None]
 
-            freqs_cos = repeat(self.freqs_cos, 'i j -> n i m j', n=t.shape[0], m=t.shape[1])
-            freqs_sin = repeat(self.freqs_sin, 'i j -> n i m j', n=t.shape[0], m=t.shape[1])
+            freqs_cos = repeat(
+                self.freqs_cos, "i j -> n i m j", n=t.shape[0], m=t.shape[1]
+            )
+            freqs_sin = repeat(
+                self.freqs_sin, "i j -> n i m j", n=t.shape[0], m=t.shape[1]
+            )
 
             freqs_cos = freqs_cos[batch_indices, patch_indices_keep]
-            freqs_cos = rearrange(freqs_cos, 'n i m j -> n m i j')
+            freqs_cos = rearrange(freqs_cos, "n i m j -> n m i j")
             freqs_sin = freqs_sin[batch_indices, patch_indices_keep]
-            freqs_sin = rearrange(freqs_sin, 'n i m j -> n m i j')
+            freqs_sin = rearrange(freqs_sin, "n i m j -> n m i j")
 
-            return  t * freqs_cos + rotate_half(t) * freqs_sin
+            return t * freqs_cos + rotate_half(t) * freqs_sin
 
-        return  t * self.freqs_cos + rotate_half(t) * self.freqs_sin
+        return t * self.freqs_cos + rotate_half(t) * self.freqs_sin
 
-import torch.nn as nn
-import os
-from dataclasses import dataclass
-from typing import Optional, Tuple, Union
-from functools import partial
 
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torch import nn
+import logging
 
 # --------------------------------------------------------
 # Adapted from  https://github.com/microsoft/unilm/tree/master/beit
 # --------------------------------------------------------
 import math
 import os
+from dataclasses import dataclass
 from functools import partial
+from typing import Optional, Tuple, Union
+
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import logging
+from torch import nn
+
 try:
     from timm.models.layers import drop_path, to_2tuple, trunc_normal_
 except:
     from timm.layers import drop_path, to_2tuple, trunc_normal_
-    
+
+
 class PatchDropout(nn.Module):
     """
     https://arxiv.org/abs/2212.00794
@@ -119,13 +130,13 @@ class PatchDropout(nn.Module):
 
     def __init__(self, prob, exclude_first_token=True):
         super().__init__()
-        assert 0 <= prob < 1.
+        assert 0 <= prob < 1.0
         self.prob = prob
         self.exclude_first_token = exclude_first_token  # exclude CLS token
         logging.info(f"os.getenv('RoPE')={os.getenv('RoPE')}")
 
     def forward(self, x):
-        if not self.training or self.prob == 0.:
+        if not self.training or self.prob == 0.0:
             return x
 
         if self.exclude_first_token:
@@ -150,12 +161,13 @@ class PatchDropout(nn.Module):
         if self.exclude_first_token:
             x = torch.cat((cls_tokens, x), dim=1)
 
-        if self.training and os.getenv('RoPE') == '1':
+        if self.training and os.getenv("RoPE") == "1":
             return x, patch_indices_keep
 
         return x
 
-if os.getenv('ENV_TYPE') == 'deepspeed':
+
+if os.getenv("ENV_TYPE") == "deepspeed":
     try:
         from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
     except:
@@ -165,32 +177,32 @@ else:
 
 import xformers.ops as xops
 
+
 class DropPath(nn.Module):
-    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
-    """
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
+
     def __init__(self, drop_prob=None):
         super(DropPath, self).__init__()
         self.drop_prob = drop_prob
 
     def forward(self, x):
         return drop_path(x, self.drop_prob, self.training)
-    
+
     def extra_repr(self) -> str:
-        return 'p={}'.format(self.drop_prob)
+        return "p={}".format(self.drop_prob)
 
 
 class Mlp(nn.Module):
     def __init__(
-        self, 
-        in_features, 
-        hidden_features=None, 
-        out_features=None, 
-        act_layer=nn.GELU, 
-        norm_layer=nn.LayerNorm, 
-        drop=0.,
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        drop=0.0,
         subln=False,
-
-        ):
+    ):
         super().__init__()
         out_features = out_features or in_features
         hidden_features = hidden_features or in_features
@@ -206,16 +218,25 @@ class Mlp(nn.Module):
         x = self.fc1(x)
         x = self.act(x)
         # x = self.drop(x)
-        # commit this for the orignal BERT implement 
+        # commit this for the orignal BERT implement
         x = self.ffn_ln(x)
 
         x = self.fc2(x)
         x = self.drop(x)
         return x
 
+
 class SwiGLU(nn.Module):
-    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0., 
-                norm_layer=nn.LayerNorm, subln=False):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.SiLU,
+        drop=0.0,
+        norm_layer=nn.LayerNorm,
+        subln=False,
+    ):
         super().__init__()
         out_features = out_features or in_features
         hidden_features = hidden_features or in_features
@@ -226,7 +247,7 @@ class SwiGLU(nn.Module):
         self.act = act_layer()
         self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
         self.w3 = nn.Linear(hidden_features, out_features)
-        
+
         self.drop = nn.Dropout(drop)
 
     def forward(self, x):
@@ -238,17 +259,30 @@ class SwiGLU(nn.Module):
         x = self.drop(x)
         return x
 
+
 class Attention(nn.Module):
     def __init__(
-            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
-            proj_drop=0., window_size=None, attn_head_dim=None, xattn=False, rope=None, subln=False, norm_layer=nn.LayerNorm):
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        window_size=None,
+        attn_head_dim=None,
+        xattn=False,
+        rope=None,
+        subln=False,
+        norm_layer=nn.LayerNorm,
+    ):
         super().__init__()
         self.num_heads = num_heads
         head_dim = dim // num_heads
         if attn_head_dim is not None:
             head_dim = attn_head_dim
         all_head_dim = head_dim * self.num_heads
-        self.scale = qk_scale or head_dim ** -0.5
+        self.scale = qk_scale or head_dim**-0.5
 
         self.subln = subln
         if self.subln:
@@ -267,9 +301,12 @@ class Attention(nn.Module):
 
         if window_size:
             self.window_size = window_size
-            self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+            self.num_relative_distance = (2 * window_size[0] - 1) * (
+                2 * window_size[1] - 1
+            ) + 3
             self.relative_position_bias_table = nn.Parameter(
-                torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+                torch.zeros(self.num_relative_distance, num_heads)
+            )  # 2*Wh-1 * 2*Ww-1, nH
             # cls to token & token 2 cls & cls to cls
 
             # get pair-wise relative position index for each token inside the window
@@ -277,13 +314,19 @@ class Attention(nn.Module):
             coords_w = torch.arange(window_size[1])
             coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
             coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
-            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
-            relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+            relative_coords = (
+                coords_flatten[:, :, None] - coords_flatten[:, None, :]
+            )  # 2, Wh*Ww, Wh*Ww
+            relative_coords = relative_coords.permute(
+                1, 2, 0
+            ).contiguous()  # Wh*Ww, Wh*Ww, 2
             relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
             relative_coords[:, :, 1] += window_size[1] - 1
             relative_coords[:, :, 0] *= 2 * window_size[1] - 1
-            relative_position_index = \
-                torch.zeros(size=(window_size[0] * window_size[1] + 1, ) * 2, dtype=relative_coords.dtype)
+            relative_position_index = torch.zeros(
+                size=(window_size[0] * window_size[1] + 1,) * 2,
+                dtype=relative_coords.dtype,
+            )
             relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
             relative_position_index[0, 0:] = self.num_relative_distance - 3
             relative_position_index[0:, 0] = self.num_relative_distance - 2
@@ -307,28 +350,88 @@ class Attention(nn.Module):
 
     def forward(self, x, rel_pos_bias=None, attn_mask=None):
         B, N, C = x.shape
-        if self.subln: 
+        if self.subln:
+            print("self.q_proj.weight.dtype=", self.q_proj.weight.dtype)
             if self.q_proj.weight.dtype == torch.uint8:
                 import bitsandbytes as bnb
-                q = bnb.matmul_4bit(x, self.q_proj.weight.t(), bias=self.q_bias, quant_state=self.q_proj.weight.quant_state)
-                k = bnb.matmul_4bit(x, self.k_proj.weight.t(), bias=None, quant_state=self.k_proj.weight.quant_state)
-                v = bnb.matmul_4bit(x, self.v_proj.weight.t(), bias=self.v_bias, quant_state=self.v_proj.weight.quant_state)
+
+                # print("self.q_proj.weight.quant_state=", self.q_proj.weight.quant_state)
+                # print("self.k_proj.weight.quant_state=", self.q_proj.weight.quant_state)
+                # print("self.v_proj.weight.quant_state=", self.q_proj.weight.quant_state)
+                q = bnb.matmul_4bit(
+                    x,
+                    self.q_proj.weight.t(),
+                    bias=self.q_bias,
+                    quant_state=self.q_proj.weight.quant_state,
+                )
+                k = bnb.matmul_4bit(
+                    x,
+                    self.k_proj.weight.t(),
+                    bias=None,
+                    quant_state=self.k_proj.weight.quant_state,
+                )
+                v = bnb.matmul_4bit(
+                    x,
+                    self.v_proj.weight.t(),
+                    bias=self.v_bias,
+                    quant_state=self.v_proj.weight.quant_state,
+                )
+            elif self.q_proj.weight.dtype == torch.int8:
+                import bitsandbytes as bnb
+
+                def make_state(weight_v):
+                    state = bnb.MatmulLtState()
+                    state.threshold = 0
+                    state.has_fp16_weights = weight_v.has_fp16_weights
+                    state.memory_efficient_backward = False
+                    state.CB = weight_v.CB
+                    state.SCB = weight_v.SCB
+                    return state
+
+                q = bnb.matmul(
+                    x,
+                    self.q_proj.weight,
+                    bias=self.q_bias,
+                    state=make_state(self.q_proj.weight),
+                )
+                k = bnb.matmul(
+                    x,
+                    self.k_proj.weight,
+                    bias=None,
+                    state=make_state(self.k_proj.weight),
+                )
+                v = bnb.matmul(
+                    x,
+                    self.v_proj.weight,
+                    bias=self.v_bias,
+                    state=make_state(self.v_proj.weight),
+                )
             else:
                 q = F.linear(input=x, weight=self.q_proj.weight, bias=self.q_bias)
                 k = F.linear(input=x, weight=self.k_proj.weight, bias=None)
                 v = F.linear(input=x, weight=self.v_proj.weight, bias=self.v_bias)
 
-            q = q.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)     # B, num_heads, N, C
-            k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)  
-            v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3) 
-        else: 
+            q = q.reshape(B, N, self.num_heads, -1).permute(
+                0, 2, 1, 3
+            )  # B, num_heads, N, C
+            k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
+            v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
+        else:
 
             qkv_bias = None
             if self.q_bias is not None:
-                qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
-            
+                qkv_bias = torch.cat(
+                    (
+                        self.q_bias,
+                        torch.zeros_like(self.v_bias, requires_grad=False),
+                        self.v_bias,
+                    )
+                )
+
             qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
-            qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)   # 3, B, num_heads, N, C
+            qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(
+                2, 0, 3, 1, 4
+            )  # 3, B, num_heads, N, C
             q, k, v = qkv[0], qkv[1], qkv[2]
 
         if self.rope:
@@ -342,29 +445,36 @@ class Attention(nn.Module):
             k = torch.cat((k[:, :, :1, :], ro_k_t), -2).type_as(v)
 
         if self.xattn:
-            q = q.permute(0, 2, 1, 3)   # B, num_heads, N, C -> B, N, num_heads, C
+            q = q.permute(0, 2, 1, 3)  # B, num_heads, N, C -> B, N, num_heads, C
             k = k.permute(0, 2, 1, 3)
             v = v.permute(0, 2, 1, 3)
 
             x = xops.memory_efficient_attention(
-                q, k, v,
+                q,
+                k,
+                v,
                 p=self.xattn_drop,
                 scale=self.scale,
-                )
+            )
             x = x.reshape(B, N, -1)
             x = self.inner_attn_ln(x)
             x = self.proj(x)
             x = self.proj_drop(x)
         else:
             q = q * self.scale
-            attn = (q @ k.transpose(-2, -1))
+            attn = q @ k.transpose(-2, -1)
 
             if self.relative_position_bias_table is not None:
-                relative_position_bias = \
-                    self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
-                        self.window_size[0] * self.window_size[1] + 1,
-                        self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
-                relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+                relative_position_bias = self.relative_position_bias_table[
+                    self.relative_position_index.view(-1)
+                ].view(
+                    self.window_size[0] * self.window_size[1] + 1,
+                    self.window_size[0] * self.window_size[1] + 1,
+                    -1,
+                )  # Wh*Ww,Wh*Ww,nH
+                relative_position_bias = relative_position_bias.permute(
+                    2, 0, 1
+                ).contiguous()  # nH, Wh*Ww, Wh*Ww
                 attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)
 
             if rel_pos_bias is not None:
@@ -373,7 +483,7 @@ class Attention(nn.Module):
             if attn_mask is not None:
                 attn_mask = attn_mask.bool()
                 attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))
-            
+
             attn = attn.softmax(dim=-1)
             attn = self.attn_drop(attn)
 
@@ -386,40 +496,71 @@ class Attention(nn.Module):
 
 class Block(nn.Module):
 
-    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
-                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
-                 window_size=None, attn_head_dim=None, xattn=False, rope=None, postnorm=False,
-                 subln=False, naiveswiglu=False):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        init_values=None,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        window_size=None,
+        attn_head_dim=None,
+        xattn=False,
+        rope=None,
+        postnorm=False,
+        subln=False,
+        naiveswiglu=False,
+    ):
         super().__init__()
         self.norm1 = norm_layer(dim)
         self.attn = Attention(
-            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
-            attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim,
-            xattn=xattn, rope=rope, subln=subln, norm_layer=norm_layer)
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+            window_size=window_size,
+            attn_head_dim=attn_head_dim,
+            xattn=xattn,
+            rope=rope,
+            subln=subln,
+            norm_layer=norm_layer,
+        )
         # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
         self.norm2 = norm_layer(dim)
         mlp_hidden_dim = int(dim * mlp_ratio)
 
         if naiveswiglu:
             self.mlp = SwiGLU(
-                in_features=dim, 
-                hidden_features=mlp_hidden_dim, 
+                in_features=dim,
+                hidden_features=mlp_hidden_dim,
                 subln=subln,
                 norm_layer=norm_layer,
             )
         else:
             self.mlp = Mlp(
-                in_features=dim, 
-                hidden_features=mlp_hidden_dim, 
+                in_features=dim,
+                hidden_features=mlp_hidden_dim,
                 act_layer=act_layer,
                 subln=subln,
-                drop=drop
+                drop=drop,
             )
 
         if init_values is not None and init_values > 0:
-            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
-            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+            self.gamma_1 = nn.Parameter(
+                init_values * torch.ones((dim)), requires_grad=True
+            )
+            self.gamma_2 = nn.Parameter(
+                init_values * torch.ones((dim)), requires_grad=True
+            )
         else:
             self.gamma_1, self.gamma_2 = None, None
 
@@ -428,24 +569,42 @@ class Block(nn.Module):
     def forward(self, x, rel_pos_bias=None, attn_mask=None):
         if self.gamma_1 is None:
             if self.postnorm:
-                x = x + self.drop_path(self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+                x = x + self.drop_path(
+                    self.norm1(
+                        self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)
+                    )
+                )
                 x = x + self.drop_path(self.norm2(self.mlp(x)))
             else:
-                x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+                x = x + self.drop_path(
+                    self.attn(
+                        self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask
+                    )
+                )
                 x = x + self.drop_path(self.mlp(self.norm2(x)))
         else:
             if self.postnorm:
-                x = x + self.drop_path(self.gamma_1 * self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+                x = x + self.drop_path(
+                    self.gamma_1
+                    * self.norm1(
+                        self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)
+                    )
+                )
                 x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
             else:
-                x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+                x = x + self.drop_path(
+                    self.gamma_1
+                    * self.attn(
+                        self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask
+                    )
+                )
                 x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
         return x
 
 
 class PatchEmbed(nn.Module):
-    """ Image to Patch Embedding
-    """
+    """Image to Patch Embedding"""
+
     def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
         super().__init__()
         img_size = to_2tuple(img_size)
@@ -456,13 +615,16 @@ class PatchEmbed(nn.Module):
         self.patch_size = patch_size
         self.num_patches = num_patches
 
-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        self.proj = nn.Conv2d(
+            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
 
     def forward(self, x, **kwargs):
         B, C, H, W = x.shape
         # FIXME look at relaxing size constraints
-        assert H == self.img_size[0] and W == self.img_size[1], \
-            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        assert (
+            H == self.img_size[0] and W == self.img_size[1]
+        ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
         x = self.proj(x).flatten(2).transpose(1, 2)
         return x
 
@@ -472,9 +634,12 @@ class RelativePositionBias(nn.Module):
     def __init__(self, window_size, num_heads):
         super().__init__()
         self.window_size = window_size
-        self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+        self.num_relative_distance = (2 * window_size[0] - 1) * (
+            2 * window_size[1] - 1
+        ) + 3
         self.relative_position_bias_table = nn.Parameter(
-            torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+            torch.zeros(self.num_relative_distance, num_heads)
+        )  # 2*Wh-1 * 2*Ww-1, nH
         # cls to token & token 2 cls & cls to cls
 
         # get pair-wise relative position index for each token inside the window
@@ -482,13 +647,18 @@ class RelativePositionBias(nn.Module):
         coords_w = torch.arange(window_size[1])
         coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
         coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
-        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
-        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords = (
+            coords_flatten[:, :, None] - coords_flatten[:, None, :]
+        )  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(
+            1, 2, 0
+        ).contiguous()  # Wh*Ww, Wh*Ww, 2
         relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
         relative_coords[:, :, 1] += window_size[1] - 1
         relative_coords[:, :, 0] *= 2 * window_size[1] - 1
-        relative_position_index = \
-            torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
+        relative_position_index = torch.zeros(
+            size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype
+        )
         relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
         relative_position_index[0, 0:] = self.num_relative_distance - 3
         relative_position_index[0:, 0] = self.num_relative_distance - 2
@@ -497,29 +667,64 @@ class RelativePositionBias(nn.Module):
         self.register_buffer("relative_position_index", relative_position_index)
 
     def forward(self):
-        relative_position_bias = \
-            self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
-                self.window_size[0] * self.window_size[1] + 1,
-                self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = self.relative_position_bias_table[
+            self.relative_position_index.view(-1)
+        ].view(
+            self.window_size[0] * self.window_size[1] + 1,
+            self.window_size[0] * self.window_size[1] + 1,
+            -1,
+        )  # Wh*Ww,Wh*Ww,nH
         return relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
 
 
 class EVAVisionTransformer(nn.Module):
-    """ Vision Transformer with support for patch or hybrid CNN input stage
-    """
-    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
-                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
-                 drop_path_rate=0., norm_layer=nn.LayerNorm, init_values=None, patch_dropout=0.,
-                 use_abs_pos_emb=True, use_rel_pos_bias=False, use_shared_rel_pos_bias=False, rope=False,
-                 use_mean_pooling=True, init_scale=0.001, grad_checkpointing=False, xattn=False, postnorm=False,
-                 pt_hw_seq_len=16, intp_freq=False, naiveswiglu=False, subln=False):
+    """Vision Transformer with support for patch or hybrid CNN input stage"""
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        num_classes=1000,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        norm_layer=nn.LayerNorm,
+        init_values=None,
+        patch_dropout=0.0,
+        use_abs_pos_emb=True,
+        use_rel_pos_bias=False,
+        use_shared_rel_pos_bias=False,
+        rope=False,
+        use_mean_pooling=True,
+        init_scale=0.001,
+        grad_checkpointing=False,
+        xattn=False,
+        postnorm=False,
+        pt_hw_seq_len=16,
+        intp_freq=False,
+        naiveswiglu=False,
+        subln=False,
+    ):
         super().__init__()
         self.image_size = img_size
         self.num_classes = num_classes
-        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.num_features = self.embed_dim = (
+            embed_dim  # num_features for consistency with other models
+        )
 
         self.patch_embed = PatchEmbed(
-            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
         num_patches = self.patch_embed.num_patches
 
         self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
@@ -531,7 +736,9 @@ class EVAVisionTransformer(nn.Module):
         self.pos_drop = nn.Dropout(p=drop_rate)
 
         if use_shared_rel_pos_bias:
-            self.rel_pos_bias = RelativePositionBias(window_size=self.patch_embed.patch_shape, num_heads=num_heads)
+            self.rel_pos_bias = RelativePositionBias(
+                window_size=self.patch_embed.patch_shape, num_heads=num_heads
+            )
         else:
             self.rel_pos_bias = None
 
@@ -544,40 +751,64 @@ class EVAVisionTransformer(nn.Module):
                 ft_seq_len=hw_seq_len if intp_freq else None,
                 # patch_dropout=patch_dropout
             )
-        else: 
+        else:
             self.rope = None
 
         self.naiveswiglu = naiveswiglu
 
-        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, depth)
+        ]  # stochastic depth decay rule
         self.use_rel_pos_bias = use_rel_pos_bias
-        self.blocks = nn.ModuleList([
-            Block(
-                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
-                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
-                init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None,
-                xattn=xattn, rope=self.rope, postnorm=postnorm, subln=subln, naiveswiglu=naiveswiglu)
-            for i in range(depth)])
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    init_values=init_values,
+                    window_size=(
+                        self.patch_embed.patch_shape if use_rel_pos_bias else None
+                    ),
+                    xattn=xattn,
+                    rope=self.rope,
+                    postnorm=postnorm,
+                    subln=subln,
+                    naiveswiglu=naiveswiglu,
+                )
+                for i in range(depth)
+            ]
+        )
         self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
         self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
-        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+        self.head = (
+            nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+        )
 
         if self.pos_embed is not None:
-            trunc_normal_(self.pos_embed, std=.02)
+            trunc_normal_(self.pos_embed, std=0.02)
 
-        trunc_normal_(self.cls_token, std=.02)
+        trunc_normal_(self.cls_token, std=0.02)
         # trunc_normal_(self.mask_token, std=.02)
 
         self.apply(self._init_weights)
         self.fix_init_weight()
 
         if isinstance(self.head, nn.Linear):
-            trunc_normal_(self.head.weight, std=.02)
+            trunc_normal_(self.head.weight, std=0.02)
             self.head.weight.data.mul_(init_scale)
             self.head.bias.data.mul_(init_scale)
 
         # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
-        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
+        self.patch_dropout = (
+            PatchDropout(patch_dropout) if patch_dropout > 0.0 else nn.Identity()
+        )
 
         self.grad_checkpointing = grad_checkpointing
 
@@ -597,7 +828,7 @@ class EVAVisionTransformer(nn.Module):
 
     def _init_weights(self, m):
         if isinstance(m, nn.Linear):
-            trunc_normal_(m.weight, std=.02)
+            trunc_normal_(m.weight, std=0.02)
             if m.bias is not None:
                 nn.init.constant_(m.bias, 0)
         elif isinstance(m, nn.LayerNorm):
@@ -606,9 +837,11 @@ class EVAVisionTransformer(nn.Module):
 
     def get_num_layers(self):
         return len(self.blocks)
-    
+
     def lock(self, unlocked_groups=0, freeze_bn_stats=False):
-        assert unlocked_groups == 0, 'partial locking not currently supported for this model'
+        assert (
+            unlocked_groups == 0
+        ), "partial locking not currently supported for this model"
         for param in self.parameters():
             param.requires_grad = False
 
@@ -618,31 +851,37 @@ class EVAVisionTransformer(nn.Module):
 
     @torch.jit.ignore
     def no_weight_decay(self):
-        return {'pos_embed', 'cls_token'}
+        return {"pos_embed", "cls_token"}
 
     def get_classifier(self):
         return self.head
 
-    def reset_classifier(self, num_classes, global_pool=''):
+    def reset_classifier(self, num_classes, global_pool=""):
         self.num_classes = num_classes
-        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+        self.head = (
+            nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+        )
 
     def forward_features(self, x, return_all_features=False):
-        
+
         x = self.patch_embed(x)
         batch_size, seq_len, _ = x.size()
 
-        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        cls_tokens = self.cls_token.expand(
+            batch_size, -1, -1
+        )  # stole cls_tokens impl from Phil Wang, thanks
         x = torch.cat((cls_tokens, x), dim=1)
         if self.pos_embed is not None:
             x = x + self.pos_embed
         x = self.pos_drop(x)
 
         # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
-        if os.getenv('RoPE') == '1':
+        if os.getenv("RoPE") == "1":
             if self.training and not isinstance(self.patch_dropout, nn.Identity):
                 x, patch_indices_keep = self.patch_dropout(x)
-                self.rope.forward = partial(self.rope.forward, patch_indices_keep=patch_indices_keep)
+                self.rope.forward = partial(
+                    self.rope.forward, patch_indices_keep=patch_indices_keep
+                )
             else:
                 self.rope.forward = partial(self.rope.forward, patch_indices_keep=None)
                 x = self.patch_dropout(x)
@@ -651,7 +890,7 @@ class EVAVisionTransformer(nn.Module):
 
         rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None
         for i, blk in enumerate(self.blocks):
-            if i == len(self.blocks)-1:
+            if i == len(self.blocks) - 1:
                 continue
             if self.grad_checkpointing:
                 x = checkpoint(blk, x, (rel_pos_bias,))
@@ -673,6 +912,7 @@ class EVAVisionTransformer(nn.Module):
         x = self.head(x)
         return x
 
+
 class LayerNorm(nn.LayerNorm):
     """Subclass torch's LayerNorm (with cast back to input dtype)."""
 
@@ -681,6 +921,7 @@ class LayerNorm(nn.LayerNorm):
         x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
         return x.to(orig_type)
 
+
 try:
     from apex.normalization import FusedLayerNorm
 except:
@@ -697,30 +938,41 @@ class CLIPVisionCfg:
     patch_size: int = 16
     image_size: Union[Tuple[int, int], int] = 224
     ls_init_value: Optional[float] = None  # layer scale initial value
-    patch_dropout: float = 0. # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
-    global_average_pool: bool = False # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
+    patch_dropout: float = (
+        0.0  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+    )
+    global_average_pool: bool = (
+        False  # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
+    )
     drop_path_rate: Optional[float] = None  # drop path rate
-    timm_model_name: str = None  # a valid model name overrides layers, width, patch_size
-    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
-    timm_pool: str = 'avg'  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
-    timm_proj: str = 'linear'  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_model_name: str = (
+        None  # a valid model name overrides layers, width, patch_size
+    )
+    timm_model_pretrained: bool = (
+        False  # use (imagenet) pretrained weights for named model
+    )
+    timm_pool: str = (
+        "avg"  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    )
+    timm_proj: str = (
+        "linear"  # linear projection for timm model output ('linear', 'mlp', '')
+    )
     timm_proj_bias: bool = False  # enable bias final projection
-    eva_model_name: str = None # a valid eva model name overrides layers, width, patch_size
+    eva_model_name: str = (
+        None  # a valid eva model name overrides layers, width, patch_size
+    )
     qkv_bias: bool = True
     fusedLN: bool = False
     xattn: bool = False
     postnorm: bool = False
     rope: bool = False
-    pt_hw_seq_len: int = 16   # 224/14
+    pt_hw_seq_len: int = 16  # 224/14
     intp_freq: bool = False
     naiveswiglu: bool = False
     subln: bool = False
 
 
-def _build_vision_tower(
-        embed_dim: int,
-        vision_cfg: CLIPVisionCfg
-):
+def _build_vision_tower(embed_dim: int, vision_cfg: CLIPVisionCfg):
     if isinstance(vision_cfg, dict):
         vision_cfg = CLIPVisionCfg(**vision_cfg)
 
@@ -731,7 +983,7 @@ def _build_vision_tower(
             img_size=vision_cfg.image_size,
             patch_size=vision_cfg.patch_size,
             num_classes=embed_dim,
-            use_mean_pooling=vision_cfg.global_average_pool, #False
+            use_mean_pooling=vision_cfg.global_average_pool,  # False
             init_values=vision_cfg.ls_init_value,
             patch_dropout=vision_cfg.patch_dropout,
             embed_dim=vision_cfg.width,
@@ -740,18 +992,23 @@ def _build_vision_tower(
             mlp_ratio=vision_cfg.mlp_ratio,
             qkv_bias=vision_cfg.qkv_bias,
             drop_path_rate=vision_cfg.drop_path_rate,
-            norm_layer= partial(FusedLayerNorm, eps=1e-6) if vision_cfg.fusedLN else partial(norm_layer, eps=1e-6),
+            norm_layer=(
+                partial(FusedLayerNorm, eps=1e-6)
+                if vision_cfg.fusedLN
+                else partial(norm_layer, eps=1e-6)
+            ),
             xattn=vision_cfg.xattn,
             rope=vision_cfg.rope,
             postnorm=vision_cfg.postnorm,
-            pt_hw_seq_len= vision_cfg.pt_hw_seq_len,   # 224/14
-            intp_freq= vision_cfg.intp_freq,
-            naiveswiglu= vision_cfg.naiveswiglu,
-            subln= vision_cfg.subln
+            pt_hw_seq_len=vision_cfg.pt_hw_seq_len,  # 224/14
+            intp_freq=vision_cfg.intp_freq,
+            naiveswiglu=vision_cfg.naiveswiglu,
+            subln=vision_cfg.subln,
         )
 
     return visual
 
+
 class Eva2LargeEncoder(nn.Module):
     def __init__(self, image_size=224):
         super(Eva2LargeEncoder, self).__init__()
@@ -772,26 +1029,38 @@ class Eva2LargeEncoder(nn.Module):
                 "pt_hw_seq_len": 16,
                 "intp_freq": True,
                 "naiveswiglu": True,
-                "subln": True
-            }
+                "subln": True,
+            },
         }
-        self.config['vision_cfg']['image_size'] = image_size
-        
+        self.config["vision_cfg"]["image_size"] = image_size
+
         import os
-        os.environ['delRoPE'] = '1' # to avoid error in rope params when changing image size
-        self.model = _build_vision_tower(**self.config)
 
+        os.environ["delRoPE"] = (
+            "1"  # to avoid error in rope params when changing image size
+        )
+        self.model = _build_vision_tower(**self.config)
 
     def forward(self, images):
         encode = self.model(images, return_all_features=True)[:, 1:, :]
         return encode
 
+
 class CrossVisionModel(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.vit = Eva2LargeEncoder(image_size=config.cross_image_size)
-        self.pos_embed = nn.Parameter(torch.zeros((self.vit.config['vision_cfg']['image_size'] // self.vit.config['vision_cfg']['patch_size']) ** 2, self.vit.config['vision_cfg']['width']))
+        self.pos_embed = nn.Parameter(
+            torch.zeros(
+                (
+                    self.vit.config["vision_cfg"]["image_size"]
+                    // self.vit.config["vision_cfg"]["patch_size"]
+                )
+                ** 2,
+                self.vit.config["vision_cfg"]["width"],
+            )
+        )
 
     def forward(self, images):
         enc = self.vit(images)
-        return enc + self.pos_embed.to(enc.device).unsqueeze(0)
+        return enc + self.pos_embed.to(enc.device).unsqueeze(0)

modeling_cogagent.py

@@ -1,24 +1,28 @@
 """largely copy from llama and adapt for CogAgent"""
-import warnings
-from typing import TYPE_CHECKING, Optional, Tuple, List, Union, Literal, Dict, Any
 
 import math
+import warnings
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union
+
 import torch
+from einops import rearrange
 from torch import nn
+
+# from .util import FastRotaryEmbedding
 from torch.nn import CrossEntropyLoss
+from torch.nn import functional as F
 from torchvision import transforms
-from einops import rearrange
-
 from transformers import PreTrainedModel, PreTrainedTokenizer
-from transformers.utils.logging import get_logger
 from transformers.activations import ACT2FN
-from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.utils.logging import get_logger
 
 from .configuration_cogagent import CogAgentConfig
-# from .util import FastRotaryEmbedding
-from torch.nn import functional as F
-from .visual import EVA2CLIPModel
 from .cross_visual import CrossVisionModel
+from .visual import EVA2CLIPModel
 
 if TYPE_CHECKING:
     from transformers.utils import ModelOutput
@@ -31,7 +35,10 @@ VISION_TOKEN_TYPE = 1
 
 # Copied from transformers.models.bart.modeling_bart._make_causal_mask
 def _make_causal_mask(
-        input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
+    input_ids_shape: torch.Size,
+    dtype: torch.dtype,
+    device: torch.device,
+    past_key_values_length: int = 0,
 ):
     """
     Make causal mask used for bi-directional self-attention.
@@ -43,8 +50,18 @@ def _make_causal_mask(
     mask = mask.to(dtype)
 
     if past_key_values_length > 0:
-        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
-    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
+        mask = torch.cat(
+            [
+                torch.zeros(
+                    tgt_len, past_key_values_length, dtype=dtype, device=device
+                ),
+                mask,
+            ],
+            dim=-1,
+        )
+    return mask[None, None, :, :].expand(
+        bsz, 1, tgt_len, tgt_len + past_key_values_length
+    )
 
 
 # Copied from transformers.models.bart.modeling_bart._expand_mask
@@ -59,7 +76,9 @@ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int]
 
     inverted_mask = 1.0 - expanded_mask
 
-    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
+    return inverted_mask.masked_fill(
+        inverted_mask.to(torch.bool), torch.finfo(dtype).min
+    )
 
 
 class RMSNorm(nn.Module):
@@ -91,9 +110,13 @@ class MLP(nn.Module):
         return down_proj
 
 
-def get_expert_mask(token_type_ids: "torch.LongTensor(B, L)") -> "[torch.BoolTensor(B, L), torch.BoolTensor(B, L)]":
+def get_expert_mask(
+    token_type_ids: "torch.LongTensor(B, L)",
+) -> "[torch.BoolTensor(B, L), torch.BoolTensor(B, L)]":
     vision_token_mask = torch.zeros_like(token_type_ids, dtype=torch.bool)
-    vision_token_mask[:, :-1] = (token_type_ids[:, :-1] == VISION_TOKEN_TYPE) & (token_type_ids[:, 1:] == VISION_TOKEN_TYPE)
+    vision_token_mask[:, :-1] = (token_type_ids[:, :-1] == VISION_TOKEN_TYPE) & (
+        token_type_ids[:, 1:] == VISION_TOKEN_TYPE
+    )
     language_token_mask = ~vision_token_mask
     return vision_token_mask, language_token_mask
 
@@ -104,47 +127,71 @@ class VisionExpertMLP(nn.Module):
         self.language_mlp = MLP(config)
         self.vision_mlp = MLP(config)
 
-    def forward(self, hidden_states: "torch.Tensor(B, L, D)", token_type_ids: "torch.LongTensor(B, L)"):
-        output = torch.empty(hidden_states.shape, dtype=hidden_states.dtype, device=hidden_states.device)
+    def forward(
+        self,
+        hidden_states: "torch.Tensor(B, L, D)",
+        token_type_ids: "torch.LongTensor(B, L)",
+    ):
+        output = torch.empty(
+            hidden_states.shape, dtype=hidden_states.dtype, device=hidden_states.device
+        )
         vision_token_mask, language_token_mask = get_expert_mask(token_type_ids)
         output[vision_token_mask] = self.vision_mlp(hidden_states[vision_token_mask])
-        output[language_token_mask] = self.language_mlp(hidden_states[language_token_mask])
+        output[language_token_mask] = self.language_mlp(
+            hidden_states[language_token_mask]
+        )
         return output
 
 
 def attention_fn(
-        query_layer: "torch.tensor(B, H, L, HD)",
-        key_layer: "torch.tensor(B, H, L, HD)",
-        value_layer: "torch.tensor(B, H, L, HD)",
-        attention_mask: "torch.tensor(B, H, L, HD)",
-        *,
-        scaling_attention_score: bool = True,
-        attention_dropout: nn.Module = None
+    query_layer: "torch.tensor(B, H, L, HD)",
+    key_layer: "torch.tensor(B, H, L, HD)",
+    value_layer: "torch.tensor(B, H, L, HD)",
+    attention_mask: "torch.tensor(B, H, L, HD)",
+    *,
+    scaling_attention_score: bool = True,
+    attention_dropout: nn.Module = None,
 ):
-    attention_mask_bool = (attention_mask == 0)
-    is_low_triangle = (attention_mask_bool == torch.ones_like(attention_mask_bool, dtype=torch.float).tril()).all()
+    attention_mask_bool = attention_mask == 0
+    is_low_triangle = (
+        attention_mask_bool
+        == torch.ones_like(attention_mask_bool, dtype=torch.float).tril()
+    ).all()
     is_full = (attention_mask_bool > 0).all()
-    if not (int(torch.__version__.split('.')[0]) >= 2):
+    if not (int(torch.__version__.split(".")[0]) >= 2):
         warnings.warn("It's recommended to use torch2.0 or higher.")
-    if int(torch.__version__.split('.')[0]) >= 2 and scaling_attention_score and (is_full or is_low_triangle):
-        dropout_p = 0. if attention_dropout is None or not attention_dropout.training else attention_dropout.p
+    if (
+        int(torch.__version__.split(".")[0]) >= 2
+        and scaling_attention_score
+        and (is_full or is_low_triangle)
+    ):
+        dropout_p = (
+            0.0
+            if attention_dropout is None or not attention_dropout.training
+            else attention_dropout.p
+        )
         return torch.nn.functional.scaled_dot_product_attention(
-            query_layer, key_layer, value_layer,
+            query_layer,
+            key_layer,
+            value_layer,
             attn_mask=None,
             dropout_p=dropout_p,
-            is_causal=not is_full
+            is_causal=not is_full,
         )
     else:
         if scaling_attention_score:
             query_layer = query_layer / math.sqrt(query_layer.shape[-1])
         attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
         attention_scores = attention_scores + attention_mask
-        attention_scores = nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32).to(query_layer.dtype)
+        attention_scores = nn.functional.softmax(
+            attention_scores, dim=-1, dtype=torch.float32
+        ).to(query_layer.dtype)
         if attention_dropout is not None:
             attention_scores = attention_dropout(attention_scores)
         context_layer = torch.matmul(attention_scores, value_layer)
         return context_layer
 
+
 class RotaryEmbedding(torch.nn.Module):
     def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
         super().__init__()
@@ -158,19 +205,24 @@ class RotaryEmbedding(torch.nn.Module):
 
     def _compute_inv_freq(self, device=None):
         return 1.0 / (
-                self.base
-                ** (torch.arange(0, self.dim, 2, device=device) / self.dim)
+            self.base ** (torch.arange(0, self.dim, 2, device=device) / self.dim)
         )
 
     def _set_cos_sin_cache(self, seq_len, device, dtype):
         self.max_seq_len_cached = seq_len
-        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+        t = torch.arange(
+            self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
+        )
 
         freqs = torch.einsum("i,j->ij", t, self.inv_freq)
         # Different from paper, but it uses a different permutation in order to obtain the same calculation
         emb = torch.cat((freqs, freqs), dim=-1)
-        self.register_buffer("cos_cached", emb.cos()[:, None, :].to(dtype), persistent=False)
-        self.register_buffer("sin_cached", emb.sin()[:, None, :].to(dtype), persistent=False)
+        self.register_buffer(
+            "cos_cached", emb.cos()[:, None, :].to(dtype), persistent=False
+        )
+        self.register_buffer(
+            "sin_cached", emb.sin()[:, None, :].to(dtype), persistent=False
+        )
 
     def forward(self, x, seq_len):
         # x: [bs, num_attention_heads, seq_len, head_size]
@@ -184,17 +236,19 @@ class RotaryEmbedding(torch.nn.Module):
 
 
 def rotate_half(x):
-    x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
     return torch.cat((-x2, x1), dim=x1.ndim - 1)
 
 
 def apply_rotary_pos_emb_index_bhs(q, k, cos, sin, position_id):
     # batch_size, num_head, seq_len, hidden_size
-    cos, sin = F.embedding(position_id, cos.squeeze(1)).unsqueeze(1), \
-        F.embedding(position_id, sin.squeeze(1)).unsqueeze(1)
+    cos, sin = F.embedding(position_id, cos.squeeze(1)).unsqueeze(1), F.embedding(
+        position_id, sin.squeeze(1)
+    ).unsqueeze(1)
     q, k = (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
     return q, k
 
+
 class VisionExpertAttention(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -205,10 +259,18 @@ class VisionExpertAttention(nn.Module):
         self.max_position_embeddings = config.max_position_embeddings
 
         self.rotary_emb = RotaryEmbedding(self.head_dim)
-        self.vision_expert_query_key_value = nn.Linear(self.hidden_size, self.hidden_size * 3, bias=False)
-        self.vision_expert_dense = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
-        self.language_expert_query_key_value = nn.Linear(self.hidden_size, self.hidden_size * 3, bias=False)
-        self.language_expert_dense = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.vision_expert_query_key_value = nn.Linear(
+            self.hidden_size, self.hidden_size * 3, bias=False
+        )
+        self.vision_expert_dense = nn.Linear(
+            self.hidden_size, self.hidden_size, bias=False
+        )
+        self.language_expert_query_key_value = nn.Linear(
+            self.hidden_size, self.hidden_size * 3, bias=False
+        )
+        self.language_expert_dense = nn.Linear(
+            self.hidden_size, self.hidden_size, bias=False
+        )
 
     def _transpose_for_scores(self, tensor):
         """Transpose a 3D tensor [B, L, H*HD] into a 4D tensor with size [B H L HD]."""
@@ -217,25 +279,33 @@ class VisionExpertAttention(nn.Module):
         return tensor.permute(0, 2, 1, 3)
 
     def forward(
-            self,
-            hidden_states: torch.Tensor,
-            token_type_ids: torch.LongTensor,
-            position_ids: torch.LongTensor,
-            attention_mask: Optional[torch.Tensor] = None,
-            past_key_value: Optional[Tuple[torch.Tensor]] = None,
-            output_attentions: bool = False,
-            use_cache: bool = False,
+        self,
+        hidden_states: torch.Tensor,
+        token_type_ids: torch.LongTensor,
+        position_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         bsz, q_len, _ = hidden_states.size()
         vision_token_mask, language_token_mask = get_expert_mask(token_type_ids)
 
         shape = list(hidden_states.shape)
         shape[-1] = shape[-1] * 3
-        mixed_raw_layer = torch.empty(shape, dtype=hidden_states.dtype, device=hidden_states.device)
-        mixed_raw_layer[vision_token_mask] = self.vision_expert_query_key_value(hidden_states[vision_token_mask])
-        mixed_raw_layer[language_token_mask] = self.language_expert_query_key_value(hidden_states[language_token_mask])
+        mixed_raw_layer = torch.empty(
+            shape, dtype=hidden_states.dtype, device=hidden_states.device
+        )
+        mixed_raw_layer[vision_token_mask] = self.vision_expert_query_key_value(
+            hidden_states[vision_token_mask]
+        )
+        mixed_raw_layer[language_token_mask] = self.language_expert_query_key_value(
+            hidden_states[language_token_mask]
+        )
 
-        query_states, key_states, value_states = torch.split(mixed_raw_layer, self.hidden_size, dim=-1)
+        query_states, key_states, value_states = torch.split(
+            mixed_raw_layer, self.hidden_size, dim=-1
+        )
         query_states = self._transpose_for_scores(query_states)  # B, H, L, HD
         key_states = self._transpose_for_scores(key_states)  # B, H, L, HD
         value_states = self._transpose_for_scores(value_states)  # B, H, L, HD
@@ -245,7 +315,9 @@ class VisionExpertAttention(nn.Module):
             kv_seq_len += past_key_value[0].shape[-2]
 
         cos, sin = self.rotary_emb(value_states, seq_len=position_ids.max() + 1)
-        query_states, key_states = apply_rotary_pos_emb_index_bhs(query_states, key_states, cos, sin, position_ids)
+        query_states, key_states = apply_rotary_pos_emb_index_bhs(
+            query_states, key_states, cos, sin, position_ids
+        )
 
         if past_key_value is not None:
             key_states = torch.cat([past_key_value[0], key_states], dim=2)
@@ -254,24 +326,40 @@ class VisionExpertAttention(nn.Module):
         past_key_value = (key_states, value_states) if use_cache else None
 
         context_layer = attention_fn(
-            query_layer=query_states, key_layer=key_states, value_layer=value_states, attention_mask=attention_mask,
-            scaling_attention_score=True, attention_dropout=None)
+            query_layer=query_states,
+            key_layer=key_states,
+            value_layer=value_states,
+            attention_mask=attention_mask,
+            scaling_attention_score=True,
+            attention_dropout=None,
+        )
         if context_layer.size() != (bsz, self.num_heads, q_len, self.head_dim):
             raise ValueError(
                 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
                 f" {context_layer.size()}"
             )
-        context_layer = context_layer.transpose(1, 2).contiguous().reshape(bsz, q_len, self.hidden_size)
+        context_layer = (
+            context_layer.transpose(1, 2)
+            .contiguous()
+            .reshape(bsz, q_len, self.hidden_size)
+        )
 
-        attn_output = torch.empty(context_layer.shape, dtype=hidden_states.dtype, device=hidden_states.device)
-        attn_output[vision_token_mask] = self.vision_expert_dense(context_layer[vision_token_mask])
-        attn_output[language_token_mask] = self.language_expert_dense(context_layer[language_token_mask])
+        attn_output = torch.empty(
+            context_layer.shape, dtype=hidden_states.dtype, device=hidden_states.device
+        )
+        attn_output[vision_token_mask] = self.vision_expert_dense(
+            context_layer[vision_token_mask]
+        )
+        attn_output[language_token_mask] = self.language_expert_dense(
+            context_layer[language_token_mask]
+        )
 
         if output_attentions:
             warnings.warn("output_attentions is not implemented.")
 
         return attn_output, None, past_key_value
 
+
 class CrossAttention(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -284,9 +372,15 @@ class CrossAttention(nn.Module):
         self.cross_head_dim = self.cross_compute_hidden_size // self.num_heads
         self.max_position_embeddings = config.max_position_embeddings
 
-        self.query = nn.Linear(self.hidden_size, self.cross_compute_hidden_size, bias=False)
-        self.key_value = nn.Linear(self.cross_hidden_size, self.cross_compute_hidden_size * 2, bias=False)
-        self.dense = nn.Linear(self.cross_compute_hidden_size, self.hidden_size, bias=False)
+        self.query = nn.Linear(
+            self.hidden_size, self.cross_compute_hidden_size, bias=False
+        )
+        self.key_value = nn.Linear(
+            self.cross_hidden_size, self.cross_compute_hidden_size * 2, bias=False
+        )
+        self.dense = nn.Linear(
+            self.cross_compute_hidden_size, self.hidden_size, bias=False
+        )
 
     def _transpose_for_scores(self, tensor):
         """Transpose a 3D tensor [B, L, H*HD] into a 4D tensor with size [B H L HD]."""
@@ -295,13 +389,13 @@ class CrossAttention(nn.Module):
         return tensor.permute(0, 2, 1, 3)
 
     def forward(
-            self,
-            hidden_states: torch.Tensor,
-            encoder_outputs: torch.LongTensor,
-            attention_mask: Optional[torch.Tensor] = None,
-            past_key_value: Optional[Tuple[torch.Tensor]] = None,
-            output_attentions: bool = False,
-            use_cache: bool = False,
+        self,
+        hidden_states: torch.Tensor,
+        encoder_outputs: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         bsz, q_len, _ = hidden_states.size()
 
@@ -311,7 +405,9 @@ class CrossAttention(nn.Module):
         mixed_query_layer = self.query(hidden_states)
         if past_key_value is None:
             mixed_x_layer = self.key_value(encoder_outputs)
-            mixed_key_layer, mixed_value_layer = torch.split(mixed_x_layer, self.cross_compute_hidden_size, dim=-1)
+            mixed_key_layer, mixed_value_layer = torch.split(
+                mixed_x_layer, self.cross_compute_hidden_size, dim=-1
+            )
             key_states = self._transpose_for_scores(mixed_key_layer)  # B, H, L, HD
             value_states = self._transpose_for_scores(mixed_value_layer)  # B, H, L, HD
         else:
@@ -322,14 +418,23 @@ class CrossAttention(nn.Module):
         past_key_value = (key_states, value_states) if use_cache else None
 
         context_layer = attention_fn(
-            query_layer=query_states, key_layer=key_states, value_layer=value_states, attention_mask=attention_mask,
-            scaling_attention_score=True, attention_dropout=None)
+            query_layer=query_states,
+            key_layer=key_states,
+            value_layer=value_states,
+            attention_mask=attention_mask,
+            scaling_attention_score=True,
+            attention_dropout=None,
+        )
         if context_layer.size() != (bsz, self.num_heads, q_len, self.cross_head_dim):
             raise ValueError(
                 f"`cross_attn_output` should be of size {(bsz, self.num_heads, q_len, self.cross_head_dim)}, but is"
                 f" {context_layer.size()}"
             )
-        context_layer = context_layer.transpose(1, 2).contiguous().reshape(bsz, q_len, self.cross_hidden_size)
+        context_layer = (
+            context_layer.transpose(1, 2)
+            .contiguous()
+            .reshape(bsz, q_len, self.cross_hidden_size)
+        )
 
         attn_output = self.dense(context_layer)
 
@@ -338,6 +443,7 @@ class CrossAttention(nn.Module):
 
         return attn_output, None, past_key_value
 
+
 class CogAgentDecoderLayer(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -346,21 +452,27 @@ class CogAgentDecoderLayer(nn.Module):
         self.cross_attn = CrossAttention(config=config)
         self.mlp = VisionExpertMLP(config)
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_cross_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_cross_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
 
     def forward(
-            self,
-            hidden_states: torch.Tensor,
-            encoder_outputs: torch.Tensor,
-            token_type_ids: torch.LongTensor,
-            position_ids: torch.LongTensor,
-            attention_mask: Optional[torch.Tensor] = None,
-            cross_attention_mask: Optional[torch.Tensor] = None,
-            past_key_value: Optional[Tuple[torch.Tensor]] = None,
-            output_attentions: Optional[bool] = False,
-            use_cache: Optional[bool] = False,
-    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        self,
+        hidden_states: torch.Tensor,
+        encoder_outputs: torch.Tensor,
+        token_type_ids: torch.LongTensor,
+        position_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+    ) -> Tuple[
+        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+    ]:
         residual = hidden_states
 
         hidden_states = self.input_layernorm(hidden_states)
@@ -379,14 +491,18 @@ class CogAgentDecoderLayer(nn.Module):
 
         cross_input = self.post_cross_attention_layernorm(hidden_states)
         # Fully Connected
-        attention_output, self_cross_attn_weights, present_cross_key_value = self.cross_attn(
-            hidden_states=cross_input,
-            encoder_outputs=encoder_outputs,
-            attention_mask=cross_attention_mask,
-            past_key_value=past_key_value[-2:] if past_key_value is not None else None,
-            output_attentions=output_attentions,
-            use_cache=use_cache,
+        attention_output, self_cross_attn_weights, present_cross_key_value = (
+            self.cross_attn(
+                hidden_states=cross_input,
+                encoder_outputs=encoder_outputs,
+                attention_mask=cross_attention_mask,
+                past_key_value=(
+                    past_key_value[-2:] if past_key_value is not None else None
+                ),
+                output_attentions=output_attentions,
+                use_cache=use_cache,
             )
+        )
         hidden_states = hidden_states + attention_output
         mlp_input = self.post_attention_layernorm(hidden_states)
         mlp_output = self.mlp(mlp_input, token_type_ids=token_type_ids)
@@ -398,7 +514,7 @@ class CogAgentDecoderLayer(nn.Module):
             outputs += (self_attn_weights,)
 
         if use_cache:
-            outputs += (present_key_value+present_cross_key_value,)
+            outputs += (present_key_value + present_cross_key_value,)
 
         return outputs  # type: ignore
 
@@ -407,7 +523,7 @@ class CogAgentPreTrainedModel(PreTrainedModel):
     config_class = CogAgentConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = False
-    _no_split_modules = ["CogAgentDecoderLayer", 'TransformerLayer', 'Block']
+    _no_split_modules = ["CogAgentDecoderLayer", "TransformerLayer", "Block"]
     _skip_keys_device_placement = "past_key_values"
 
     def _init_weights(self, module):
@@ -431,7 +547,10 @@ def is_empty(images_list: Optional[List[List[torch.Tensor]]]):
     return True
 
 
-def build_position_ids(x: "torch.BoolTensor(B, L)", attention_mask: Optional["torch.BoolTensor(B, L)"] = None) -> "torch.LongTensor(B, L)":
+def build_position_ids(
+    x: "torch.BoolTensor(B, L)",
+    attention_mask: Optional["torch.BoolTensor(B, L)"] = None,
+) -> "torch.LongTensor(B, L)":
     if attention_mask is not None:
         tmp = x.clone()
         tmp[~(attention_mask.bool())] = -1
@@ -439,14 +558,20 @@ def build_position_ids(x: "torch.BoolTensor(B, L)", attention_mask: Optional["to
         tmp = x.clone()
     # image boi eoi token as LANGUAGE_TOKEN_TYPE
     is_boi_eoi = torch.zeros_like(x, dtype=torch.bool)
-    is_boi_eoi[:, 1:] |= (tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE)
-    is_boi_eoi[:, 0] |= (tmp[:, 0] == VISION_TOKEN_TYPE)
-    is_boi_eoi[:, :-1] |= (tmp[:, :-1] == VISION_TOKEN_TYPE) & (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE)
-    is_boi_eoi[:, -1] |= (tmp[:, -1] == VISION_TOKEN_TYPE)
+    is_boi_eoi[:, 1:] |= (tmp[:, 1:] == VISION_TOKEN_TYPE) & (
+        tmp[:, :-1] == LANGUAGE_TOKEN_TYPE
+    )
+    is_boi_eoi[:, 0] |= tmp[:, 0] == VISION_TOKEN_TYPE
+    is_boi_eoi[:, :-1] |= (tmp[:, :-1] == VISION_TOKEN_TYPE) & (
+        tmp[:, 1:] == LANGUAGE_TOKEN_TYPE
+    )
+    is_boi_eoi[:, -1] |= tmp[:, -1] == VISION_TOKEN_TYPE
     tmp[is_boi_eoi] = LANGUAGE_TOKEN_TYPE
     # final position ids
     y = torch.zeros_like(x, dtype=torch.long)
-    y[:, 1:] = (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE) | ((tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE))
+    y[:, 1:] = (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE) | (
+        (tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE)
+    )
     y = y.cumsum(dim=-1)
     return y
 
@@ -457,8 +582,12 @@ class CogAgentModel(CogAgentPreTrainedModel):
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
 
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
-        self.layers = nn.ModuleList([CogAgentDecoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [CogAgentDecoderLayer(config) for _ in range(config.num_hidden_layers)]
+        )
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
         self.vision = EVA2CLIPModel(config)
@@ -491,22 +620,22 @@ class CogAgentModel(CogAgentPreTrainedModel):
         images = torch.stack(images)
         encoder_outputs = self.cross_vision(images)
         return encoder_outputs
-    
+
     def forward(
-            self,
-            input_ids: torch.LongTensor = None,
-            images: List[List[torch.Tensor]] = None,
-            cross_images: List[List[torch.Tensor]] = None,
-            token_type_ids: Optional[torch.LongTensor] = None,
-            attention_mask: Optional[torch.Tensor] = None,
-            cross_attention_mask: Optional[torch.Tensor] = None,
-            position_ids: Optional[torch.LongTensor] = None,
-            past_key_values: Optional[List[torch.FloatTensor]] = None,
-            inputs_embeds: Optional[torch.FloatTensor] = None,
-            use_cache: Optional[bool] = None,
-            output_attentions: Optional[bool] = None,
-            output_hidden_states: Optional[bool] = None,
-            return_dict: Optional[bool] = None,
+        self,
+        input_ids: torch.LongTensor = None,
+        images: List[List[torch.Tensor]] = None,
+        cross_images: List[List[torch.Tensor]] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         """take care of image_encode, token_type_ids, position_ids and (attention_mask = None is fine)"""
 
@@ -515,20 +644,35 @@ class CogAgentModel(CogAgentPreTrainedModel):
             # generate mode with past_key_values. the image features are already mapped
         else:
             # not allow for inputs_embeds, because we want to process image feature
-            assert input_ids is not None and inputs_embeds is None, f"{input_ids} {inputs_embeds}"
+            assert (
+                input_ids is not None and inputs_embeds is None
+            ), f"{input_ids} {inputs_embeds}"
             if not is_empty(images):  # multi-modality
-                assert token_type_ids is not None, f"multi-modality requires `token_type_ids`!"
+                assert (
+                    token_type_ids is not None
+                ), f"multi-modality requires `token_type_ids`!"
                 assert len(input_ids) == len(images), f"{len(input_ids)} {len(images)}"
                 inputs_embeds = self.embed_tokens(input_ids)
                 images_features = self.encode_images(images)
                 encoder_outputs = self.encode_cross_images(cross_images)
-                images_features = rearrange(images_features, 'b n d -> (b n) d')
-                images_features = images_features.to(dtype=inputs_embeds.dtype, device=inputs_embeds.device)
-                inputs_embeds = inputs_embeds.index_put([token_type_ids == VISION_TOKEN_TYPE], images_features)
+                images_features = rearrange(images_features, "b n d -> (b n) d")
+                images_features = images_features.to(
+                    dtype=inputs_embeds.dtype, device=inputs_embeds.device
+                )
+                inputs_embeds = inputs_embeds.index_put(
+                    [token_type_ids == VISION_TOKEN_TYPE], images_features
+                )
             else:  # single-modality
                 if token_type_ids is None:
-                    token_type_ids = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device) * LANGUAGE_TOKEN_TYPE
-                assert not (token_type_ids == VISION_TOKEN_TYPE).any(), f"{(token_type_ids == VISION_TOKEN_TYPE).sum()}"
+                    token_type_ids = (
+                        torch.ones_like(
+                            input_ids, dtype=torch.long, device=input_ids.device
+                        )
+                        * LANGUAGE_TOKEN_TYPE
+                    )
+                assert not (
+                    token_type_ids == VISION_TOKEN_TYPE
+                ).any(), f"{(token_type_ids == VISION_TOKEN_TYPE).sum()}"
                 inputs_embeds = self.embed_tokens(input_ids)
                 encoder_outputs = None
 
@@ -552,38 +696,50 @@ class CogAgentModel(CogAgentPreTrainedModel):
         )
 
     def llm_forward(
-            self,
-            input_ids: torch.LongTensor = None,
-            encoder_outputs: torch.LongTensor = None,
-            token_type_ids: torch.LongTensor = None,
-            attention_mask: Optional[torch.Tensor] = None,
-            cross_attention_mask: Optional[torch.Tensor] = None,
-            position_ids: Optional[torch.LongTensor] = None,
-            past_key_values: Optional[List[torch.FloatTensor]] = None,
-            inputs_embeds: Optional[torch.FloatTensor] = None,
-            use_cache: Optional[bool] = None,
-            output_attentions: Optional[bool] = None,
-            output_hidden_states: Optional[bool] = None,
-            return_dict: Optional[bool] = None,
+        self,
+        input_ids: torch.LongTensor = None,
+        encoder_outputs: torch.LongTensor = None,
+        token_type_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         """largely copy from llama forward and adapt for CogAgent with `token_type_ids`"""
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        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
+            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
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
 
         # retrieve input_ids and inputs_embeds
         if input_ids is not None and inputs_embeds is not None:
-            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+            raise ValueError(
+                "You cannot specify both decoder_input_ids and decoder_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 decoder_input_ids or decoder_inputs_embeds")
+            raise ValueError(
+                "You have to specify either decoder_input_ids or decoder_inputs_embeds"
+            )
 
         seq_length_with_past = seq_length
         past_key_values_length = 0
@@ -595,7 +751,10 @@ class CogAgentModel(CogAgentPreTrainedModel):
         if position_ids is None:
             device = input_ids.device if input_ids is not None else inputs_embeds.device
             position_ids = torch.arange(
-                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+                past_key_values_length,
+                seq_length + past_key_values_length,
+                dtype=torch.long,
+                device=device,
             )
             position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
         else:
@@ -606,14 +765,19 @@ class CogAgentModel(CogAgentPreTrainedModel):
         # embed positions
         if attention_mask is None:
             attention_mask = torch.ones(
-                (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
+                (batch_size, seq_length_with_past),
+                dtype=torch.bool,
+                device=inputs_embeds.device,
             )
         if cross_attention_mask is None:
             cross_attention_mask = torch.ones(
                 (batch_size, 1), dtype=torch.bool, device=inputs_embeds.device
             )
         attention_mask = self._prepare_decoder_attention_mask(
-            attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+            attention_mask,
+            (batch_size, seq_length),
+            inputs_embeds,
+            past_key_values_length,
         )
 
         hidden_states = inputs_embeds
@@ -627,7 +791,9 @@ class CogAgentModel(CogAgentPreTrainedModel):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
-            past_key_value = past_key_values[idx] if past_key_values is not None else None
+            past_key_value = (
+                past_key_values[idx] if past_key_values is not None else None
+            )
             layer_outputs = decoder_layer(
                 hidden_states,
                 encoder_outputs=encoder_outputs,
@@ -655,7 +821,11 @@ class CogAgentModel(CogAgentPreTrainedModel):
 
         next_cache = next_decoder_cache if use_cache else None
         if not return_dict:
-            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+            return tuple(
+                v
+                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
+                if v is not None
+            )
         return BaseModelOutputWithPast(
             last_hidden_state=hidden_states,
             past_key_values=next_cache,
@@ -671,7 +841,9 @@ class CogAgentModel(CogAgentPreTrainedModel):
 
     # noinspection PyMethodMayBeStatic
     # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
-    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
+    def _prepare_decoder_attention_mask(
+        self, attention_mask, input_shape, inputs_embeds, past_key_values_length
+    ):
         # create causal mask
         # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
         combined_attention_mask = None
@@ -685,15 +857,18 @@ class CogAgentModel(CogAgentPreTrainedModel):
 
         if attention_mask is not None:
             # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
-            expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
-                inputs_embeds.device
-            )
+            expanded_attn_mask = _expand_mask(
+                attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+            ).to(inputs_embeds.device)
             combined_attention_mask = (
-                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+                expanded_attn_mask
+                if combined_attention_mask is None
+                else expanded_attn_mask + combined_attention_mask
             )
 
         return combined_attention_mask
 
+
 def vqa_history_to_prompt(history, query):
     # Only support single round chat in vqa mode
     prompt = "<EOI>Question: "
@@ -702,6 +877,7 @@ def vqa_history_to_prompt(history, query):
     prompt += query + " Short answer:"
     return prompt
 
+
 def chat_old_history_to_prompt(history, query):
     prompt = "<EOI>Question: "
     for i, (old_query, response) in enumerate(history):
@@ -709,6 +885,7 @@ def chat_old_history_to_prompt(history, query):
     prompt += query + " Answer:"
     return prompt
 
+
 def chat_history_to_prompt(history, query):
     prompt = " [INST] "
     for i, (old_query, response) in enumerate(history):
@@ -726,7 +903,7 @@ _history_to_prompt = {
     "base": base_history_to_prompt,
     "chat": chat_history_to_prompt,
     "chat_old": chat_old_history_to_prompt,
-    "vqa": vqa_history_to_prompt
+    "vqa": vqa_history_to_prompt,
 }
 
 
@@ -761,26 +938,34 @@ class CogAgentForCausalLM(CogAgentPreTrainedModel):
         return self.model
 
     def forward(
-            self,
-            input_ids: torch.LongTensor = None,
-            images: List[List[torch.Tensor]] = None,
-            cross_images: List[List[torch.Tensor]] = None,
-            token_type_ids: Optional[torch.LongTensor] = None,
-            attention_mask: Optional[torch.Tensor] = None,
-            position_ids: Optional[torch.LongTensor] = None,
-            past_key_values: Optional[List[torch.FloatTensor]] = None,
-            inputs_embeds: Optional[torch.FloatTensor] = None,
-            use_cache: Optional[bool] = None,
-            output_attentions: Optional[bool] = None,
-            output_hidden_states: Optional[bool] = None,
-            return_dict: Optional[bool] = None,
-            labels: Optional[torch.LongTensor] = None,
+        self,
+        input_ids: torch.LongTensor = None,
+        images: List[List[torch.Tensor]] = None,
+        cross_images: List[List[torch.Tensor]] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        labels: Optional[torch.LongTensor] = None,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        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
+            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
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         outputs = self.model(
@@ -828,15 +1013,23 @@ class CogAgentForCausalLM(CogAgentPreTrainedModel):
         )
 
     def _prepare_attention_mask_for_generation(
-            self,
-            inputs: torch.Tensor,
-            pad_token_id: Optional[int],
-            eos_token_id: Optional[Union[int, List[int]]],
+        self,
+        inputs: torch.Tensor,
+        pad_token_id: Optional[int],
+        eos_token_id: Optional[Union[int, List[int]]],
     ) -> torch.LongTensor:
         return torch.ones(inputs.shape[:2], dtype=torch.long, device=inputs.device)  # type: ignore
 
     def prepare_inputs_for_generation(
-            self, input_ids, token_type_ids, images=None, cross_images=None, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+        self,
+        input_ids,
+        token_type_ids,
+        images=None,
+        cross_images=None,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        **kwargs,
     ):
         # build position_ids if needed
         position_ids = kwargs.get("position_ids", None)
@@ -868,11 +1061,11 @@ class CogAgentForCausalLM(CogAgentPreTrainedModel):
         return model_inputs
 
     def _update_model_kwargs_for_generation(
-            self,
-            outputs: "ModelOutput",
-            model_kwargs: Dict[str, Any],
-            is_encoder_decoder: bool = False,
-            standardize_cache_format: bool = False,
+        self,
+        outputs: "ModelOutput",
+        model_kwargs: Dict[str, Any],
+        is_encoder_decoder: bool = False,
+        standardize_cache_format: bool = False,
     ) -> Dict[str, Any]:
         # update past_key_values
         model_kwargs["past_key_values"] = self._extract_past_from_model_output(
@@ -884,22 +1077,40 @@ class CogAgentForCausalLM(CogAgentPreTrainedModel):
         # update token_type_ids with last value
         if "token_type_ids" in model_kwargs:
             token_type_ids = model_kwargs["token_type_ids"]
-            new_token_type_ids = torch.ones(size=(token_type_ids.shape[0], 1), dtype=token_type_ids.dtype, device=token_type_ids.device) * LANGUAGE_TOKEN_TYPE
-            model_kwargs["token_type_ids"] = torch.cat([token_type_ids, new_token_type_ids], dim=-1)
+            new_token_type_ids = (
+                torch.ones(
+                    size=(token_type_ids.shape[0], 1),
+                    dtype=token_type_ids.dtype,
+                    device=token_type_ids.device,
+                )
+                * LANGUAGE_TOKEN_TYPE
+            )
+            model_kwargs["token_type_ids"] = torch.cat(
+                [token_type_ids, new_token_type_ids], dim=-1
+            )
 
         if not is_encoder_decoder:
             # update attention mask
             if "attention_mask" in model_kwargs:
                 attention_mask = model_kwargs["attention_mask"]
                 model_kwargs["attention_mask"] = torch.cat(
-                    [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
+                    [
+                        attention_mask,
+                        attention_mask.new_ones((attention_mask.shape[0], 1)),
+                    ],
+                    dim=-1,
                 )
         else:
             # update decoder attention mask
             if "decoder_attention_mask" in model_kwargs:
                 decoder_attention_mask = model_kwargs["decoder_attention_mask"]
                 model_kwargs["decoder_attention_mask"] = torch.cat(
-                    [decoder_attention_mask, decoder_attention_mask.new_ones((decoder_attention_mask.shape[0], 1))],
+                    [
+                        decoder_attention_mask,
+                        decoder_attention_mask.new_ones(
+                            (decoder_attention_mask.shape[0], 1)
+                        ),
+                    ],
                     dim=-1,
                 )
 
@@ -909,22 +1120,25 @@ class CogAgentForCausalLM(CogAgentPreTrainedModel):
         reordered_past = ()
         for layer_past in past_key_values:
             reordered_past += (
-                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+                tuple(
+                    past_state.index_select(0, beam_idx.to(past_state.device))
+                    for past_state in layer_past
+                ),
             )
         return reordered_past
 
     def build_conversation_input_ids(
-            self,
-            tokenizer: "PreTrainedTokenizer",
-            *,
-            query: str,
-            history: Optional[List[Tuple[str, str]]] = None,
-            images: Optional[List["PIL.Image"]] = None,
-            template_version: Optional[Literal["base", "chat", "vqa"]] = None,
+        self,
+        tokenizer: "PreTrainedTokenizer",
+        *,
+        query: str,
+        history: Optional[List[Tuple[str, str]]] = None,
+        images: Optional[List["PIL.Image"]] = None,
+        template_version: Optional[Literal["base", "chat", "vqa"]] = None,
     ):
-        image_size: int = self.config.vision_config['image_size']
+        image_size: int = self.config.vision_config["image_size"]
         cross_image_size: int = self.config.cross_image_size
-        patch_size: int = self.config.vision_config['patch_size']
+        patch_size: int = self.config.vision_config["patch_size"]
         template_version = template_version or self.config.template_version
         assert images is None or len(images) <= 1, f"not support multi images by now."
         history = history or []
@@ -938,25 +1152,35 @@ class CogAgentForCausalLM(CogAgentPreTrainedModel):
             transform = transforms.Compose(
                 [
                     transforms.Resize(
-                        (image_size, image_size), interpolation=transforms.InterpolationMode.BICUBIC
+                        (image_size, image_size),
+                        interpolation=transforms.InterpolationMode.BICUBIC,
                     ),
                     transforms.ToTensor(),
-                    transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+                    transforms.Normalize(
+                        (0.48145466, 0.4578275, 0.40821073),
+                        (0.26862954, 0.26130258, 0.27577711),
+                    ),
                 ]
             )
             images = [transform(ori[0])]
             cross_transform = transforms.Compose(
                 [
                     transforms.Resize(
-                        (cross_image_size, cross_image_size), interpolation=transforms.InterpolationMode.BICUBIC
+                        (cross_image_size, cross_image_size),
+                        interpolation=transforms.InterpolationMode.BICUBIC,
                     ),
                     transforms.ToTensor(),
-                    transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+                    transforms.Normalize(
+                        (0.48145466, 0.4578275, 0.40821073),
+                        (0.26862954, 0.26130258, 0.27577711),
+                    ),
                 ]
             )
             cross_images = [cross_transform(ori[0])]
             # language
-            vision_token_num = (image_size // patch_size) * (image_size // patch_size) + 2
+            vision_token_num = (image_size // patch_size) * (
+                image_size // patch_size
+            ) + 2
             input_ids += [tokenizer.pad_token_id] * vision_token_num
             token_type_ids += [VISION_TOKEN_TYPE] * vision_token_num
         text_ids = tokenizer.encode(text, add_special_tokens=False)
@@ -966,9 +1190,9 @@ class CogAgentForCausalLM(CogAgentPreTrainedModel):
         attention_mask = [1] * len(input_ids)
 
         return {
-            'input_ids': torch.tensor(input_ids, dtype=torch.long),
-            'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long),
-            'attention_mask': torch.tensor(attention_mask, dtype=torch.long),
-            'images': images,
-            'cross_images': cross_images
+            "input_ids": torch.tensor(input_ids, dtype=torch.long),
+            "token_type_ids": torch.tensor(token_type_ids, dtype=torch.long),
+            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
+            "images": images,
+            "cross_images": cross_images,
         }

visual.py

@@ -1,14 +1,20 @@
-import torch
-from torch import nn
 from argparse import Namespace
+
+import torch
 import xformers.ops as xops
+from torch import nn
 from transformers.activations import ACT2FN
 
 
 class PatchEmbedding(nn.Module):
     def __init__(self, config):
         super().__init__()
-        self.proj = nn.Conv2d(config.in_channels, config.hidden_size, kernel_size=config.patch_size, stride=config.patch_size)
+        self.proj = nn.Conv2d(
+            config.in_channels,
+            config.hidden_size,
+            kernel_size=config.patch_size,
+            stride=config.patch_size,
+        )
         self.cls_embedding = nn.Parameter(torch.zeros(1, config.hidden_size))
         self.position_embedding = nn.Embedding(config.num_positions, config.hidden_size)
 
@@ -26,7 +32,7 @@ class Attention(nn.Module):
         super().__init__()
         self.num_heads = config.num_heads
         head_dim = config.hidden_size // config.num_heads
-        self.scale = head_dim ** -0.5
+        self.scale = head_dim**-0.5
         self.query_key_value = nn.Linear(config.hidden_size, config.hidden_size * 3)
         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
         self.output_dropout = torch.nn.Dropout(config.dropout_prob)
@@ -34,11 +40,16 @@ class Attention(nn.Module):
     def forward(self, x: "tensor(B, L, D)") -> "tensor(B, L, D)":
         B, L, _ = x.shape
         qkv = self.query_key_value(x)
-        qkv = qkv.reshape(B, L, 3, self.num_heads, -1).permute(2, 0, 1, 3, 4)  # 3, B, L, H, D
+        qkv = qkv.reshape(B, L, 3, self.num_heads, -1).permute(
+            2, 0, 1, 3, 4
+        )  # 3, B, L, H, D
         q, k, v = qkv[0], qkv[1], qkv[2]
 
         out = xops.memory_efficient_attention(
-            q, k, v, scale=self.scale,
+            q,
+            k,
+            v,
+            scale=self.scale,
         )
         output = self.dense(out.view(B, L, -1))
         output = self.output_dropout(output)
@@ -69,10 +80,14 @@ class MLP(nn.Module):
 class TransformerLayer(nn.Module):
     def __init__(self, config):
         super().__init__()
-        self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.input_layernorm = nn.LayerNorm(
+            config.hidden_size, eps=config.layer_norm_eps
+        )
         self.attention = Attention(config)
         self.mlp = MLP(config)
-        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.post_attention_layernorm = nn.LayerNorm(
+            config.hidden_size, eps=config.layer_norm_eps
+        )
 
     def forward(self, hidden_states):
         attention_input = hidden_states
@@ -87,7 +102,9 @@ class TransformerLayer(nn.Module):
 class Transformer(nn.Module):
     def __init__(self, config):
         super().__init__()
-        self.layers = nn.ModuleList([TransformerLayer(config) for _ in range(config.num_hidden_layers)])
+        self.layers = nn.ModuleList(
+            [TransformerLayer(config) for _ in range(config.num_hidden_layers)]
+        )
 
     def forward(self, hidden_states):
         for layer_module in self.layers:
@@ -102,9 +119,15 @@ class GLU(nn.Module):
         self.norm1 = nn.LayerNorm(config.hidden_size)
         self.act1 = nn.GELU()
         self.act2 = nn.functional.silu
-        self.dense_h_to_4h = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.dense_4h_to_h = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.dense_h_to_4h = nn.Linear(
+            config.hidden_size, config.intermediate_size, bias=False
+        )
+        self.gate_proj = nn.Linear(
+            config.hidden_size, config.intermediate_size, bias=False
+        )
+        self.dense_4h_to_h = nn.Linear(
+            config.intermediate_size, config.hidden_size, bias=False
+        )
 
     def forward(self, x):
         x = self.linear_proj(x)
@@ -123,7 +146,12 @@ class EVA2CLIPModel(nn.Module):
         self.linear_proj = GLU(config, in_features=vision_config.hidden_size)
         self.boi = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
         self.eoi = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
-        self.pos_embed = nn.Parameter(torch.zeros((vision_config.image_size // vision_config.patch_size) ** 2, vision_config.hidden_size))
+        self.pos_embed = nn.Parameter(
+            torch.zeros(
+                (vision_config.image_size // vision_config.patch_size) ** 2,
+                vision_config.hidden_size,
+            )
+        )
 
     def forward(self, images: "tensor(B, C, H, W)") -> "tensor(B, L, D)":
         x = self.patch_embedding(images)