(
self,
dim: int,
num_heads: int,
mlp_ratio: float = 4.0,
qkv_bias: bool = True,
proj_bias: bool = True,
ffn_bias: bool = True,
drop: float = 0.0,
attn_drop: float = 0.0,
init_values=None,
drop_path: float = 0.0,
act_layer: Callable[..., nn.Module] = nn.GELU,
norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
ffn_layer: Callable[..., nn.Module] = Mlp,
qk_norm: bool = False,
rope=None,
kv_cache_sliding_window: int = 64,
kv_cache_scale_frames: int = 8,
kv_cache_cross_frame_special: bool = True,
kv_cache_include_scale_frames: bool = True,
kv_cache_camera_only: bool = False,
)
| 156 | """ |
| 157 | |
| 158 | def __init__( |
| 159 | self, |
| 160 | dim: int, |
| 161 | num_heads: int, |
| 162 | mlp_ratio: float = 4.0, |
| 163 | qkv_bias: bool = True, |
| 164 | proj_bias: bool = True, |
| 165 | ffn_bias: bool = True, |
| 166 | drop: float = 0.0, |
| 167 | attn_drop: float = 0.0, |
| 168 | init_values=None, |
| 169 | drop_path: float = 0.0, |
| 170 | act_layer: Callable[..., nn.Module] = nn.GELU, |
| 171 | norm_layer: Callable[..., nn.Module] = nn.LayerNorm, |
| 172 | ffn_layer: Callable[..., nn.Module] = Mlp, |
| 173 | qk_norm: bool = False, |
| 174 | rope=None, |
| 175 | kv_cache_sliding_window: int = 64, |
| 176 | kv_cache_scale_frames: int = 8, |
| 177 | kv_cache_cross_frame_special: bool = True, |
| 178 | kv_cache_include_scale_frames: bool = True, |
| 179 | kv_cache_camera_only: bool = False, |
| 180 | ) -> None: |
| 181 | super().__init__() |
| 182 | |
| 183 | self.norm1 = norm_layer(dim) |
| 184 | self.attn = FlashInferAttention( |
| 185 | dim=dim, |
| 186 | num_heads=num_heads, |
| 187 | qk_norm=qk_norm, |
| 188 | qkv_bias=qkv_bias, |
| 189 | proj_bias=proj_bias, |
| 190 | attn_drop=attn_drop, |
| 191 | proj_drop=drop, |
| 192 | rope=rope, |
| 193 | kv_cache_sliding_window=kv_cache_sliding_window, |
| 194 | kv_cache_scale_frames=kv_cache_scale_frames, |
| 195 | kv_cache_cross_frame_special=kv_cache_cross_frame_special, |
| 196 | kv_cache_include_scale_frames=kv_cache_include_scale_frames, |
| 197 | kv_cache_camera_only=kv_cache_camera_only, |
| 198 | ) |
| 199 | |
| 200 | self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() |
| 201 | self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 202 | |
| 203 | self.norm2 = norm_layer(dim) |
| 204 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 205 | self.mlp = ffn_layer( |
| 206 | in_features=dim, |
| 207 | hidden_features=mlp_hidden_dim, |
| 208 | act_layer=act_layer, |
| 209 | drop=drop, |
| 210 | bias=ffn_bias |
| 211 | ) |
| 212 | self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity() |
| 213 | self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 214 | |
| 215 | self.sample_drop_ratio = drop_path |
nothing calls this directly
no test coverage detected