Upload model

eradio_model.py

@@ -19,19 +19,25 @@ from timm.models.layers import trunc_normal_, DropPath, LayerNorm2d
 import numpy as np
 import torch.nn.functional as F
 from .block import C2f
-TRT = False # should help for TRT
+
+TRT = False  # should help for TRT
 
 import pickle
+
 global bias_indx
 bias_indx = -1
 DEBUG = False
 
 
-
 def pixel_unshuffle(data, factor=2):
     # performs nn.PixelShuffle(factor) in reverse, torch has some bug for ONNX and TRT, so doing it manually
     B, C, H, W = data.shape
-    return data.view(B, C, factor, H//factor, factor, W//factor).permute(0,1,2,4,3,5).reshape(B, -1, H//factor, W//factor)
+    return (
+        data.view(B, C, factor, H // factor, factor, W // factor)
+        .permute(0, 1, 2, 4, 3, 5)
+        .reshape(B, -1, H // factor, W // factor)
+    )
+
 
 class SwiGLU(nn.Module):
     # should be more advanced, but doesnt improve results so far
@@ -51,7 +57,7 @@ def window_partition(x, window_size):
     """
     B, C, H, W = x.shape
 
-    if window_size == 0 or (window_size==H and window_size==W):
+    if window_size == 0 or (window_size == H and window_size == W):
         windows = x.flatten(2).transpose(1, 2)
         Hp, Wp = H, W
     else:
@@ -62,23 +68,38 @@ def window_partition(x, window_size):
         Hp, Wp = H + pad_h, W + pad_w
 
         x = x.view(B, C, Hp // window_size, window_size, Wp // window_size, window_size)
-        windows = x.permute(0, 2, 4, 3, 5, 1).reshape(-1, window_size*window_size, C)
+        windows = x.permute(0, 2, 4, 3, 5, 1).reshape(-1, window_size * window_size, C)
 
     return windows, (Hp, Wp)
 
+
 class Conv2d_BN(nn.Module):
-    '''
+    """
     Conv2d + BN layer with folding capability to speed up inference
-    '''
-    def __init__(self, a, b, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bn_weight_init=1, bias=False):
+    """
+
+    def __init__(
+        self,
+        a,
+        b,
+        kernel_size=1,
+        stride=1,
+        padding=0,
+        dilation=1,
+        groups=1,
+        bn_weight_init=1,
+        bias=False,
+    ):
         super().__init__()
-        self.conv = torch.nn.Conv2d(a, b, kernel_size, stride, padding, dilation, groups, bias=False)
+        self.conv = torch.nn.Conv2d(
+            a, b, kernel_size, stride, padding, dilation, groups, bias=False
+        )
         if 1:
             self.bn = torch.nn.BatchNorm2d(b)
             torch.nn.init.constant_(self.bn.weight, bn_weight_init)
             torch.nn.init.constant_(self.bn.bias, 0)
 
-    def forward(self,x):
+    def forward(self, x):
         x = self.conv(x)
         x = self.bn(x)
         return x
@@ -91,14 +112,12 @@ class Conv2d_BN(nn.Module):
             c, bn = self.conv, self.bn
             w = bn.weight / (bn.running_var + bn.eps) ** 0.5
             w = c.weight * w[:, None, None, None]
-            b = bn.bias - bn.running_mean * bn.weight / \
-                (bn.running_var + bn.eps)**0.5
+            b = bn.bias - bn.running_mean * bn.weight / (bn.running_var + bn.eps) ** 0.5
             self.conv.weight.data.copy_(w)
             self.conv.bias = nn.Parameter(b)
             self.bn = nn.Identity()
 
 
-
 def window_reverse(windows, window_size, H, W, pad_hw):
     """
     Args:
@@ -113,48 +132,63 @@ def window_reverse(windows, window_size, H, W, pad_hw):
     """
     # print(f"window_reverse, windows.shape {windows.shape}")
     Hp, Wp = pad_hw
-    if window_size == 0 or (window_size==H and window_size==W):
+    if window_size == 0 or (window_size == H and window_size == W):
         B = int(windows.shape[0] / (Hp * Wp / window_size / window_size))
         x = windows.transpose(1, 2).view(B, -1, H, W)
     else:
         B = int(windows.shape[0] / (Hp * Wp / window_size / window_size))
-        x = windows.view(B, Hp // window_size, Wp // window_size, window_size, window_size, -1)
-        x = x.permute(0, 5, 1, 3, 2, 4).reshape(B,windows.shape[2], Hp, Wp)
+        x = windows.view(
+            B, Hp // window_size, Wp // window_size, window_size, window_size, -1
+        )
+        x = x.permute(0, 5, 1, 3, 2, 4).reshape(B, windows.shape[2], Hp, Wp)
 
         if Hp > H or Wp > W:
-            x = x[:, :, :H, :W, ].contiguous()
+            x = x[:, :, :H, :W,].contiguous()
 
     return x
 
 
-
 class PosEmbMLPSwinv2D(nn.Module):
-    def __init__(self, window_size, pretrained_window_size, num_heads, seq_length, no_log=False):
+    def __init__(
+        self, window_size, pretrained_window_size, num_heads, seq_length, no_log=False
+    ):
         super().__init__()
         self.window_size = window_size
         self.num_heads = num_heads
         # mlp to generate continuous relative position bias
-        self.cpb_mlp = nn.Sequential(nn.Linear(2, 512, bias=True),
-                                     nn.ReLU(inplace=True),
-                                     nn.Linear(512, num_heads, bias=False))
+        self.cpb_mlp = nn.Sequential(
+            nn.Linear(2, 512, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Linear(512, num_heads, bias=False),
+        )
 
         # get relative_coords_table
-        relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32)
-        relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32)
-        relative_coords_table = torch.stack(
-            torch.meshgrid([relative_coords_h,
-                            relative_coords_w])).permute(1, 2, 0).contiguous().unsqueeze(0)  # 1, 2*Wh-1, 2*Ww-1, 2
+        relative_coords_h = torch.arange(
+            -(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32
+        )
+        relative_coords_w = torch.arange(
+            -(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32
+        )
+        relative_coords_table = (
+            torch.stack(torch.meshgrid([relative_coords_h, relative_coords_w]))
+            .permute(1, 2, 0)
+            .contiguous()
+            .unsqueeze(0)
+        )  # 1, 2*Wh-1, 2*Ww-1, 2
         if pretrained_window_size[0] > 0:
-            relative_coords_table[:, :, :, 0] /= (pretrained_window_size[0] - 1)
-            relative_coords_table[:, :, :, 1] /= (pretrained_window_size[1] - 1)
+            relative_coords_table[:, :, :, 0] /= pretrained_window_size[0] - 1
+            relative_coords_table[:, :, :, 1] /= pretrained_window_size[1] - 1
         else:
-            relative_coords_table[:, :, :, 0] /= (self.window_size[0] - 1)
-            relative_coords_table[:, :, :, 1] /= (self.window_size[1] - 1)
+            relative_coords_table[:, :, :, 0] /= self.window_size[0] - 1
+            relative_coords_table[:, :, :, 1] /= self.window_size[1] - 1
 
         if not no_log:
             relative_coords_table *= 8  # normalize to -8, 8
-            relative_coords_table = torch.sign(relative_coords_table) * torch.log2(
-                torch.abs(relative_coords_table) + 1.0) / np.log2(8)
+            relative_coords_table = (
+                torch.sign(relative_coords_table)
+                * torch.log2(torch.abs(relative_coords_table) + 1.0)
+                / np.log2(8)
+            )
 
         self.register_buffer("relative_coords_table", relative_coords_table)
 
@@ -163,8 +197,12 @@ class PosEmbMLPSwinv2D(nn.Module):
         coords_w = torch.arange(self.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] += self.window_size[0] - 1  # shift to start from 0
         relative_coords[:, :, 1] += self.window_size[1] - 1
         relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
@@ -177,7 +215,7 @@ class PosEmbMLPSwinv2D(nn.Module):
 
         relative_bias = torch.zeros(1, num_heads, seq_length, seq_length)
         self.seq_length = seq_length
-        self.register_buffer("relative_bias", relative_bias) #for EMA
+        self.register_buffer("relative_bias", relative_bias)  # for EMA
 
     def switch_to_deploy(self):
         self.deploy = True
@@ -191,7 +229,8 @@ class PosEmbMLPSwinv2D(nn.Module):
         # if not (input_tensor.shape[1:] == self.relative_bias.shape[1:]):
         #     self.grid_exists = False
 
-        if self.training: self.grid_exists = False
+        if self.training:
+            self.grid_exists = False
 
         if self.deploy and self.grid_exists:
             input_tensor += self.relative_bias
@@ -200,12 +239,20 @@ class PosEmbMLPSwinv2D(nn.Module):
         if not self.grid_exists:
             self.grid_exists = True
 
-            relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads)
-            relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view(
-                self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[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_table = self.cpb_mlp(
+                self.relative_coords_table
+            ).view(-1, self.num_heads)
+            relative_position_bias = relative_position_bias_table[
+                self.relative_position_index.view(-1)
+            ].view(
+                self.window_size[0] * self.window_size[1],
+                self.window_size[0] * self.window_size[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 = 16 * torch.sigmoid(relative_position_bias)
 
             self.relative_bias = relative_position_bias.unsqueeze(0)
@@ -214,14 +261,25 @@ class PosEmbMLPSwinv2D(nn.Module):
         return input_tensor
 
 
-
 class GRAAttentionBlock(nn.Module):
-    def __init__(self, window_size, dim_in, dim_out,
-                 num_heads, drop_path=0., qk_scale=None, qkv_bias=False,
-                 norm_layer=nn.LayerNorm, layer_scale=None,
-                  use_swiglu=True,
-                  subsample_ratio=1, dim_ratio=1, conv_base=False,
-                  do_windowing=True, multi_query=False) -> None:
+    def __init__(
+        self,
+        window_size,
+        dim_in,
+        dim_out,
+        num_heads,
+        drop_path=0.0,
+        qk_scale=None,
+        qkv_bias=False,
+        norm_layer=nn.LayerNorm,
+        layer_scale=None,
+        use_swiglu=True,
+        subsample_ratio=1,
+        dim_ratio=1,
+        conv_base=False,
+        do_windowing=True,
+        multi_query=False,
+    ) -> None:
         super().__init__()
 
         dim = dim_in
@@ -232,49 +290,131 @@ class GRAAttentionBlock(nn.Module):
 
         if do_windowing:
             if SHUFFLE:
-                self.downsample_op = torch.nn.PixelUnshuffle(subsample_ratio) if subsample_ratio>1 else torch.nn.Identity()
-                self.downsample_mixer = nn.Conv2d(dim_in * (subsample_ratio * subsample_ratio), dim_in * (dim_ratio), kernel_size=1, stride=1, padding=0, bias=False) if dim*dim_ratio != dim * subsample_ratio * subsample_ratio else torch.nn.Identity()
+                self.downsample_op = (
+                    torch.nn.PixelUnshuffle(subsample_ratio)
+                    if subsample_ratio > 1
+                    else torch.nn.Identity()
+                )
+                self.downsample_mixer = (
+                    nn.Conv2d(
+                        dim_in * (subsample_ratio * subsample_ratio),
+                        dim_in * (dim_ratio),
+                        kernel_size=1,
+                        stride=1,
+                        padding=0,
+                        bias=False,
+                    )
+                    if dim * dim_ratio != dim * subsample_ratio * subsample_ratio
+                    else torch.nn.Identity()
+                )
             else:
                 if conv_base:
-                    self.downsample_op = nn.Conv2d(dim_in, dim_out, kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()
+                    self.downsample_op = (
+                        nn.Conv2d(
+                            dim_in,
+                            dim_out,
+                            kernel_size=subsample_ratio,
+                            stride=subsample_ratio,
+                        )
+                        if subsample_ratio > 1
+                        else nn.Identity()
+                    )
                     self.downsample_mixer = nn.Identity()
                 else:
-                    self.downsample_op = nn.AvgPool2d(kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()
-                    self.downsample_mixer = Conv2d_BN(dim_in, dim_out, kernel_size=1, stride=1) if subsample_ratio > 1 else nn.Identity()
-
+                    self.downsample_op = (
+                        nn.AvgPool2d(
+                            kernel_size=subsample_ratio, stride=subsample_ratio
+                        )
+                        if subsample_ratio > 1
+                        else nn.Identity()
+                    )
+                    self.downsample_mixer = (
+                        Conv2d_BN(dim_in, dim_out, kernel_size=1, stride=1)
+                        if subsample_ratio > 1
+                        else nn.Identity()
+                    )
 
         if do_windowing:
             if SHUFFLE:
-                self.upsample_mixer =nn.Conv2d(dim_in * dim_ratio, dim_in * (subsample_ratio * subsample_ratio), kernel_size=1, stride=1, padding=0, bias=False) if dim*dim_ratio != dim * subsample_ratio * subsample_ratio else torch.nn.Identity()
-                self.upsample_op = torch.nn.PixelShuffle(subsample_ratio) if subsample_ratio>1 else torch.nn.Identity()
+                self.upsample_mixer = (
+                    nn.Conv2d(
+                        dim_in * dim_ratio,
+                        dim_in * (subsample_ratio * subsample_ratio),
+                        kernel_size=1,
+                        stride=1,
+                        padding=0,
+                        bias=False,
+                    )
+                    if dim * dim_ratio != dim * subsample_ratio * subsample_ratio
+                    else torch.nn.Identity()
+                )
+                self.upsample_op = (
+                    torch.nn.PixelShuffle(subsample_ratio)
+                    if subsample_ratio > 1
+                    else torch.nn.Identity()
+                )
             else:
                 if conv_base:
                     self.upsample_mixer = nn.Identity()
-                    self.upsample_op = nn.ConvTranspose2d(dim_in, dim_out, kernel_size=subsample_ratio, stride=subsample_ratio) if subsample_ratio > 1 else nn.Identity()
+                    self.upsample_op = (
+                        nn.ConvTranspose2d(
+                            dim_in,
+                            dim_out,
+                            kernel_size=subsample_ratio,
+                            stride=subsample_ratio,
+                        )
+                        if subsample_ratio > 1
+                        else nn.Identity()
+                    )
                 else:
-                    self.upsample_mixer = nn.Upsample(scale_factor=subsample_ratio, mode='nearest') if subsample_ratio > 1 else nn.Identity()
-                    self.upsample_op = Conv2d_BN(dim_in, dim_out, kernel_size=1, stride=1, padding=0, bias=False) if subsample_ratio > 1 else nn.Identity()
+                    self.upsample_mixer = (
+                        nn.Upsample(scale_factor=subsample_ratio, mode="nearest")
+                        if subsample_ratio > 1
+                        else nn.Identity()
+                    )
+                    self.upsample_op = (
+                        Conv2d_BN(
+                            dim_in,
+                            dim_out,
+                            kernel_size=1,
+                            stride=1,
+                            padding=0,
+                            bias=False,
+                        )
+                        if subsample_ratio > 1
+                        else nn.Identity()
+                    )
 
         self.window_size = window_size
 
         self.norm1 = norm_layer(dim_in)
         if DEBUG:
-            print(f"GRAAttentionBlock: input_resolution: , window_size: {window_size}, dim_in: {dim_in}, dim_out: {dim_out}, num_heads: {num_heads}, drop_path: {drop_path}, qk_scale: {qk_scale}, qkv_bias: {qkv_bias}, layer_scale: {layer_scale}")
-
+            print(
+                f"GRAAttentionBlock: input_resolution: , window_size: {window_size}, dim_in: {dim_in}, dim_out: {dim_out}, num_heads: {num_heads}, drop_path: {drop_path}, qk_scale: {qk_scale}, qkv_bias: {qkv_bias}, layer_scale: {layer_scale}"
+            )
 
         self.attn = WindowAttention(
             dim_in,
-            num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
             resolution=window_size,
-            seq_length=window_size**2, dim_out=dim_in, multi_query=multi_query)
+            seq_length=window_size ** 2,
+            dim_out=dim_in,
+            multi_query=multi_query,
+        )
         if DEBUG:
-            print(f"Attention: dim_in: {dim_in}, num_heads: {num_heads}, qkv_bias: {qkv_bias}, qk_scale: {qk_scale}, resolution: {window_size}, seq_length: {window_size**2}, dim_out: {dim_in}")
+            print(
+                f"Attention: dim_in: {dim_in}, num_heads: {num_heads}, qkv_bias: {qkv_bias}, qk_scale: {qk_scale}, resolution: {window_size}, seq_length: {window_size**2}, dim_out: {dim_in}"
+            )
             print(f"drop_path: {drop_path}, layer_scale: {layer_scale}")
 
-        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
 
         use_layer_scale = layer_scale is not None and type(layer_scale) in [int, float]
-        self.gamma1 = nn.Parameter(layer_scale * torch.ones(dim_in))  if use_layer_scale else 1
+        self.gamma1 = (
+            nn.Parameter(layer_scale * torch.ones(dim_in)) if use_layer_scale else 1
+        )
 
         ### mlp layer
         mlp_ratio = 4
@@ -282,17 +422,27 @@ class GRAAttentionBlock(nn.Module):
         mlp_hidden_dim = int(dim_in * mlp_ratio)
 
         activation = nn.GELU if not use_swiglu else SwiGLU
-        mlp_hidden_dim = int((4 * dim_in * 1 / 2) / 64) * 64 if use_swiglu else mlp_hidden_dim
-
-        self.mlp = Mlp(in_features=dim_in, hidden_features=mlp_hidden_dim, act_layer=activation, use_swiglu=use_swiglu)
-
-        self.gamma2 = nn.Parameter(layer_scale * torch.ones(dim_in)) if layer_scale else 1
-        self.drop_path2=DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        mlp_hidden_dim = (
+            int((4 * dim_in * 1 / 2) / 64) * 64 if use_swiglu else mlp_hidden_dim
+        )
+
+        self.mlp = Mlp(
+            in_features=dim_in,
+            hidden_features=mlp_hidden_dim,
+            act_layer=activation,
+            use_swiglu=use_swiglu,
+        )
+
+        self.gamma2 = (
+            nn.Parameter(layer_scale * torch.ones(dim_in)) if layer_scale else 1
+        )
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
         if DEBUG:
-            print(f"MLP layer: dim_in: {dim_in}, dim_out: {dim_in}, mlp_hidden_dim: {mlp_hidden_dim}")
+            print(
+                f"MLP layer: dim_in: {dim_in}, dim_out: {dim_in}, mlp_hidden_dim: {mlp_hidden_dim}"
+            )
             print(f"drop_path: {drop_path}, layer_scale: {layer_scale}")
 
-
     def forward(self, x):
         skip_connection = x
 
@@ -301,15 +451,15 @@ class GRAAttentionBlock(nn.Module):
             x = self.downsample_op(x)
             x = self.downsample_mixer(x)
 
-            if self.window_size>0:
+            if self.window_size > 0:
                 H, W = x.shape[2], x.shape[3]
 
             x, pad_hw = window_partition(x, self.window_size)
 
         # window attention
-        x = x + self.drop_path1(self.gamma1*self.attn(self.norm1(x)))
+        x = x + self.drop_path1(self.gamma1 * self.attn(self.norm1(x)))
         # mlp layer
-        x = x + self.drop_path2(self.gamma2*self.mlp(self.norm2(x)))
+        x = x + self.drop_path2(self.gamma2 * self.mlp(self.norm2(x)))
 
         if self.do_windowing:
             if self.window_size > 0:
@@ -318,9 +468,19 @@ class GRAAttentionBlock(nn.Module):
             x = self.upsample_mixer(x)
             x = self.upsample_op(x)
 
-
-            if x.shape[2] != skip_connection.shape[2] or x.shape[3] != skip_connection.shape[3]:
-                x = torch.nn.functional.pad(x, ( 0, -x.shape[3] + skip_connection.shape[3], 0, -x.shape[2] + skip_connection.shape[2]))
+            if (
+                x.shape[2] != skip_connection.shape[2]
+                or x.shape[3] != skip_connection.shape[3]
+            ):
+                x = torch.nn.functional.pad(
+                    x,
+                    (
+                        0,
+                        -x.shape[3] + skip_connection.shape[3],
+                        0,
+                        -x.shape[2] + skip_connection.shape[2],
+                    ),
+                )
         # need to add skip connection because downsampling and upsampling will break residual connection
         # 0.5 is needed to make sure that the skip connection is not too strong
         # in case of no downsample / upsample we can show that 0.5 compensates for the residual connection
@@ -329,8 +489,6 @@ class GRAAttentionBlock(nn.Module):
         return x
 
 
-
-
 class MultiResolutionAttention(nn.Module):
     """
     MultiResolutionAttention (MRA) module
@@ -340,12 +498,23 @@ class MultiResolutionAttention(nn.Module):
 
     """
 
-    def __init__(self, window_size, sr_ratio,
-                 dim, dim_ratio, num_heads,
-                 do_windowing=True,
-                 layer_scale=1e-5, norm_layer=nn.LayerNorm,
-                 drop_path = 0, qkv_bias=False, qk_scale=1.0,
-                 use_swiglu=True, multi_query=False, conv_base=False) -> None:
+    def __init__(
+        self,
+        window_size,
+        sr_ratio,
+        dim,
+        dim_ratio,
+        num_heads,
+        do_windowing=True,
+        layer_scale=1e-5,
+        norm_layer=nn.LayerNorm,
+        drop_path=0,
+        qkv_bias=False,
+        qk_scale=1.0,
+        use_swiglu=True,
+        multi_query=False,
+        conv_base=False,
+    ) -> None:
         """
         Args:
             input_resolution: input image resolution
@@ -357,10 +526,8 @@ class MultiResolutionAttention(nn.Module):
 
         depth = len(sr_ratio)
 
-
         self.attention_blocks = nn.ModuleList()
 
-
         for i in range(depth):
             subsample_ratio = sr_ratio[i]
             if len(window_size) > i:
@@ -368,15 +535,25 @@ class MultiResolutionAttention(nn.Module):
             else:
                 window_size_local = window_size[0]
 
-            self.attention_blocks.append(GRAAttentionBlock(window_size=window_size_local,
-                                            dim_in=dim, dim_out=dim, num_heads=num_heads,
-                                            qkv_bias=qkv_bias, qk_scale=qk_scale, norm_layer=norm_layer,
-                                            layer_scale=layer_scale, drop_path=drop_path,
-                                            use_swiglu=use_swiglu, subsample_ratio=subsample_ratio, dim_ratio=dim_ratio,
-                                            do_windowing=do_windowing, multi_query=multi_query, conv_base=conv_base),
-                                        )
-
-
+            self.attention_blocks.append(
+                GRAAttentionBlock(
+                    window_size=window_size_local,
+                    dim_in=dim,
+                    dim_out=dim,
+                    num_heads=num_heads,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    norm_layer=norm_layer,
+                    layer_scale=layer_scale,
+                    drop_path=drop_path,
+                    use_swiglu=use_swiglu,
+                    subsample_ratio=subsample_ratio,
+                    dim_ratio=dim_ratio,
+                    do_windowing=do_windowing,
+                    multi_query=multi_query,
+                    conv_base=conv_base,
+                ),
+            )
 
     def forward(self, x):
 
@@ -386,19 +563,20 @@ class MultiResolutionAttention(nn.Module):
         return x
 
 
-
 class Mlp(nn.Module):
     """
     Multi-Layer Perceptron (MLP) block
     """
 
-    def __init__(self,
-                 in_features,
-                 hidden_features=None,
-                 out_features=None,
-                 act_layer=nn.GELU,
-                 use_swiglu=True,
-                 drop=0.):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        use_swiglu=True,
+        drop=0.0,
+    ):
         """
         Args:
             in_features: input features dimension.
@@ -411,7 +589,9 @@ class Mlp(nn.Module):
         super().__init__()
         out_features = out_features or in_features
         hidden_features = hidden_features or in_features
-        self.fc1 = nn.Linear(in_features, hidden_features * (2 if use_swiglu else 1), bias=False)
+        self.fc1 = nn.Linear(
+            in_features, hidden_features * (2 if use_swiglu else 1), bias=False
+        )
         self.act = act_layer()
         self.fc2 = nn.Linear(hidden_features, out_features, bias=False)
         # self.drop = GaussianDropout(drop)
@@ -427,6 +607,7 @@ class Mlp(nn.Module):
         x = x.view(x_size)
         return x
 
+
 class Downsample(nn.Module):
     """
     Down-sampling block
@@ -434,10 +615,9 @@ class Downsample(nn.Module):
     Pixel Unshuffle is used for down-sampling, works great accuracy - wise but takes 10% more TRT time
     """
 
-    def __init__(self,
-                 dim,
-                 shuffle = False,
-                 ):
+    def __init__(
+        self, dim, shuffle=False,
+    ):
         """
         Args:
             dim: feature size dimension.
@@ -450,14 +630,13 @@ class Downsample(nn.Module):
 
         if shuffle:
             self.norm = lambda x: pixel_unshuffle(x, factor=2)
-            self.reduction = Conv2d_BN(dim*4, dim_out, 1, 1, 0, bias=False)
+            self.reduction = Conv2d_BN(dim * 4, dim_out, 1, 1, 0, bias=False)
         else:
-            #removed layer norm for better, in this formulation we are getting 10% better speed
+            # removed layer norm for better, in this formulation we are getting 10% better speed
             # LayerNorm for high resolution inputs will be a pain as it pools over the entire spatial dimension
             self.norm = nn.Identity()
             self.reduction = Conv2d_BN(dim, dim_out, 3, 2, 1, bias=False)
 
-
     def forward(self, x):
         x = self.norm(x)
         x = self.reduction(x)
@@ -486,8 +665,8 @@ class PatchEmbed(nn.Module):
                 Conv2d_BN(in_chans, in_dim, 3, 2, 1, bias=False),
                 nn.ReLU(),
                 Conv2d_BN(in_dim, dim, 3, 2, 1, bias=False),
-                nn.ReLU()
-                )
+                nn.ReLU(),
+            )
         else:
             self.proj = lambda x: pixel_unshuffle(x, factor=4)
 
@@ -496,9 +675,9 @@ class PatchEmbed(nn.Module):
             #                                Conv2d_BN(in_dim, dim, 3, 1, 1),
             #                                nn.SiLU(),
             #                                )
-            self.conv_down = nn.Sequential(Conv2d_BN(in_chans*16, dim, 3, 1, 1),
-                                           nn.ReLU(),
-                                           )
+            self.conv_down = nn.Sequential(
+                Conv2d_BN(in_chans * 16, dim, 3, 1, 1), nn.ReLU(),
+            )
 
     def forward(self, x):
         x = self.proj(x)
@@ -506,7 +685,6 @@ class PatchEmbed(nn.Module):
         return x
 
 
-
 class ConvBlock(nn.Module):
     """
     Convolutional block, used in first couple of stages
@@ -514,24 +692,30 @@ class ConvBlock(nn.Module):
     Experimented with RepVGG, dont see significant improvement in accuracy
     Finally, YOLOv8 idea seem to work fine (resnet-18 like block with squeezed feature dimension, and feature concatendation at the end)
     """
-    def __init__(self, dim,
-                 drop_path=0.,
-                 layer_scale=None,
-                 kernel_size=3,
-                 rep_vgg=False):
+
+    def __init__(
+        self, dim, drop_path=0.0, layer_scale=None, kernel_size=3, rep_vgg=False
+    ):
         super().__init__()
         self.rep_vgg = rep_vgg
         if not rep_vgg:
-            self.conv1 = Conv2d_BN(dim, dim, kernel_size=kernel_size, stride=1, padding=1)
+            self.conv1 = Conv2d_BN(
+                dim, dim, kernel_size=kernel_size, stride=1, padding=1
+            )
             self.act1 = nn.GELU()
         else:
-            self.conv1 = RepVGGBlock(dim, dim, kernel_size=kernel_size, stride=1, padding=1, groups=1)
-
+            self.conv1 = RepVGGBlock(
+                dim, dim, kernel_size=kernel_size, stride=1, padding=1, groups=1
+            )
 
         if not rep_vgg:
-            self.conv2 = Conv2d_BN(dim, dim, kernel_size=kernel_size, stride=1, padding=1)
+            self.conv2 = Conv2d_BN(
+                dim, dim, kernel_size=kernel_size, stride=1, padding=1
+            )
         else:
-            self.conv2 = RepVGGBlock(dim, dim, kernel_size=kernel_size, stride=1, padding=1, groups=1)
+            self.conv2 = RepVGGBlock(
+                dim, dim, kernel_size=kernel_size, stride=1, padding=1, groups=1
+            )
 
         self.layer_scale = layer_scale
         if layer_scale is not None and type(layer_scale) in [int, float]:
@@ -539,7 +723,7 @@ class ConvBlock(nn.Module):
             self.layer_scale = True
         else:
             self.layer_scale = False
-        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()
 
     def forward(self, x):
         input = x
@@ -563,11 +747,21 @@ class WindowAttention(nn.Module):
     # look into palm: https://github.com/lucidrains/PaLM-pytorch/blob/main/palm_pytorch/palm_pytorch.py
     # single kv attention, mlp in parallel (didnt improve speed)
 
-    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, resolution=0,
-                 seq_length=0, dim_out=None, multi_query=False):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        resolution=0,
+        seq_length=0,
+        dim_out=None,
+        multi_query=False,
+    ):
         # taken from EdgeViT and tweaked with attention bias.
         super().__init__()
-        if not dim_out: dim_out = dim
+        if not dim_out:
+            dim_out = dim
         self.multi_query = multi_query
         self.num_heads = num_heads
         head_dim = dim // num_heads
@@ -584,14 +778,16 @@ class WindowAttention(nn.Module):
             else:
                 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
         else:
-            self.qkv = nn.Linear(dim, dim + 2*self.head_dim, bias=qkv_bias)
+            self.qkv = nn.Linear(dim, dim + 2 * self.head_dim, bias=qkv_bias)
 
         self.proj = nn.Linear(dim, dim_out, bias=False)
         # attention positional bias
-        self.pos_emb_funct = PosEmbMLPSwinv2D(window_size=[resolution, resolution],
-                                              pretrained_window_size=[resolution, resolution],
-                                              num_heads=num_heads,
-                                              seq_length=seq_length)
+        self.pos_emb_funct = PosEmbMLPSwinv2D(
+            window_size=[resolution, resolution],
+            pretrained_window_size=[resolution, resolution],
+            num_heads=num_heads,
+            seq_length=seq_length,
+        )
 
         self.resolution = resolution
 
@@ -600,17 +796,37 @@ class WindowAttention(nn.Module):
 
         if not self.multi_query:
             if TRT:
-                q = self.q(x).reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
-                k = self.k(x).reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
-                v = self.v(x).reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+                q = (
+                    self.q(x)
+                    .reshape(B, -1, self.num_heads, C // self.num_heads)
+                    .permute(0, 2, 1, 3)
+                )
+                k = (
+                    self.k(x)
+                    .reshape(B, -1, self.num_heads, C // self.num_heads)
+                    .permute(0, 2, 1, 3)
+                )
+                v = (
+                    self.v(x)
+                    .reshape(B, -1, self.num_heads, C // self.num_heads)
+                    .permute(0, 2, 1, 3)
+                )
             else:
-                qkv = self.qkv(x).reshape(B, -1, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+                qkv = (
+                    self.qkv(x)
+                    .reshape(B, -1, 3, self.num_heads, C // self.num_heads)
+                    .permute(2, 0, 3, 1, 4)
+                )
                 q, k, v = qkv[0], qkv[1], qkv[2]
         else:
             qkv = self.qkv(x)
-            (q, k, v) = qkv.split([self.dim_internal, self.head_dim, self.head_dim], dim=2)
+            (q, k, v) = qkv.split(
+                [self.dim_internal, self.head_dim, self.head_dim], dim=2
+            )
 
-            q = q.reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+            q = q.reshape(B, -1, self.num_heads, C // self.num_heads).permute(
+                0, 2, 1, 3
+            )
             k = k.reshape(B, -1, 1, C // self.num_heads).permute(0, 2, 1, 3)
             v = v.reshape(B, -1, 1, C // self.num_heads).permute(0, 2, 1, 3)
 
@@ -624,35 +840,34 @@ class WindowAttention(nn.Module):
         return x
 
 
-
 class FasterViTLayer(nn.Module):
     """
     fastervitlayer
     """
 
-    def __init__(self,
-                 dim,
-                 depth,
-                 num_heads,
-                 window_size,
-                 conv=False,
-                 downsample=True,
-                 mlp_ratio=4.,
-                 qkv_bias=False,
-                 qk_scale=None,
-                 norm_layer=nn.LayerNorm,
-                 drop_path=0.,
-                 layer_scale=None,
-                 layer_scale_conv=None,
-                 sr_dim_ratio=1,
-                 sr_ratio=1,
-                 multi_query=False,
-                 use_swiglu=True,
-                 rep_vgg=False,
-                 yolo_arch=False,
-                 downsample_shuffle=False,
-                 conv_base=False,
-
+    def __init__(
+        self,
+        dim,
+        depth,
+        num_heads,
+        window_size,
+        conv=False,
+        downsample=True,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        norm_layer=nn.LayerNorm,
+        drop_path=0.0,
+        layer_scale=None,
+        layer_scale_conv=None,
+        sr_dim_ratio=1,
+        sr_ratio=1,
+        multi_query=False,
+        use_swiglu=True,
+        rep_vgg=False,
+        yolo_arch=False,
+        downsample_shuffle=False,
+        conv_base=False,
     ):
         """
         Args:
@@ -674,23 +889,33 @@ class FasterViTLayer(nn.Module):
 
         super().__init__()
         self.conv = conv
-        self.yolo_arch=False
+        self.yolo_arch = False
         if conv:
             if not yolo_arch:
-                self.blocks = nn.ModuleList([
-                    ConvBlock(dim=dim,
-                            drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
-                            layer_scale=layer_scale_conv, rep_vgg=rep_vgg)
-                    for i in range(depth)])
+                self.blocks = nn.ModuleList(
+                    [
+                        ConvBlock(
+                            dim=dim,
+                            drop_path=drop_path[i]
+                            if isinstance(drop_path, list)
+                            else drop_path,
+                            layer_scale=layer_scale_conv,
+                            rep_vgg=rep_vgg,
+                        )
+                        for i in range(depth)
+                    ]
+                )
             else:
-                self.blocks = C2f(dim,dim,n=depth,shortcut=True,e=0.5)
-                self.yolo_arch=True
+                self.blocks = C2f(dim, dim, n=depth, shortcut=True, e=0.5)
+                self.yolo_arch = True
         else:
-            if not isinstance(window_size, list): window_size = [window_size]
+            if not isinstance(window_size, list):
+                window_size = [window_size]
             self.window_size = window_size[0]
             self.do_single_windowing = True
-            if not isinstance(sr_ratio, list): sr_ratio = [sr_ratio]
-            if any([sr!=1 for sr in sr_ratio]) or len(set(window_size))>1:
+            if not isinstance(sr_ratio, list):
+                sr_ratio = [sr_ratio]
+            if any([sr != 1 for sr in sr_ratio]) or len(set(window_size)) > 1:
                 self.do_single_windowing = False
                 do_windowing = True
             else:
@@ -701,29 +926,31 @@ class FasterViTLayer(nn.Module):
             for i in range(depth):
 
                 self.blocks.append(
-                    MultiResolutionAttention(window_size=window_size,
-                                             sr_ratio=sr_ratio,
-                                             dim=dim,
-                                             dim_ratio = sr_dim_ratio,
-                                             num_heads=num_heads,
-                                             norm_layer=norm_layer,
-                                             drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
-                                             layer_scale=layer_scale,
-                                             qkv_bias=qkv_bias,
-                                             qk_scale=qk_scale,
-                                             use_swiglu=use_swiglu,
-                                             do_windowing=do_windowing,
-                                             multi_query=multi_query,
-                                             conv_base=conv_base,
-                    ))
+                    MultiResolutionAttention(
+                        window_size=window_size,
+                        sr_ratio=sr_ratio,
+                        dim=dim,
+                        dim_ratio=sr_dim_ratio,
+                        num_heads=num_heads,
+                        norm_layer=norm_layer,
+                        drop_path=drop_path[i]
+                        if isinstance(drop_path, list)
+                        else drop_path,
+                        layer_scale=layer_scale,
+                        qkv_bias=qkv_bias,
+                        qk_scale=qk_scale,
+                        use_swiglu=use_swiglu,
+                        do_windowing=do_windowing,
+                        multi_query=multi_query,
+                        conv_base=conv_base,
+                    )
+                )
 
         self.transformer = not conv
 
-
-        self.downsample = None if not downsample else Downsample(dim=dim, shuffle=downsample_shuffle)
-
-
-
+        self.downsample = (
+            None if not downsample else Downsample(dim=dim, shuffle=downsample_shuffle)
+        )
 
     def forward(self, x):
         B, C, H, W = x.shape
@@ -741,11 +968,10 @@ class FasterViTLayer(nn.Module):
         if self.transformer and self.do_single_windowing:
             x = window_reverse(x, self.window_size, H, W, pad_hw)
 
-
         if self.downsample is None:
             return x, x
 
-        return self.downsample(x), x #changing to output pre downsampled features
+        return self.downsample(x), x  # changing to output pre downsampled features
 
 
 class FasterViT(nn.Module):
@@ -753,37 +979,39 @@ class FasterViT(nn.Module):
     FasterViT
     """
 
-    def __init__(self,
-                 dim,
-                 in_dim,
-                 depths,
-                 window_size,
-                 mlp_ratio,
-                 num_heads,
-                 drop_path_rate=0.2,
-                 in_chans=3,
-                 num_classes=1000,
-                 qkv_bias=False,
-                 qk_scale=None,
-                 layer_scale=None,
-                 layer_scale_conv=None,
-                 layer_norm_last=False,
-                 sr_ratio = [1, 1, 1, 1],
-                 max_depth = -1,
-                 conv_base=False,
-                 use_swiglu=False,
-                 multi_query=False,
-                 norm_layer=nn.LayerNorm,
-                 rep_vgg=False,
-                 drop_uniform=False,
-                 yolo_arch=False,
-                 shuffle_down=False,
-                 downsample_shuffle=False,
-                 return_full_features=False,
-                 full_features_head_dim=128,
-                 neck_start_stage=1,
-                 use_neck=False,
-                 **kwargs):
+    def __init__(
+        self,
+        dim,
+        in_dim,
+        depths,
+        window_size,
+        mlp_ratio,
+        num_heads,
+        drop_path_rate=0.2,
+        in_chans=3,
+        num_classes=1000,
+        qkv_bias=False,
+        qk_scale=None,
+        layer_scale=None,
+        layer_scale_conv=None,
+        layer_norm_last=False,
+        sr_ratio=[1, 1, 1, 1],
+        max_depth=-1,
+        conv_base=False,
+        use_swiglu=False,
+        multi_query=False,
+        norm_layer=nn.LayerNorm,
+        rep_vgg=False,
+        drop_uniform=False,
+        yolo_arch=False,
+        shuffle_down=False,
+        downsample_shuffle=False,
+        return_full_features=False,
+        full_features_head_dim=128,
+        neck_start_stage=1,
+        use_neck=False,
+        **kwargs,
+    ):
         """
         Args:
             dim: feature size dimension.
@@ -811,7 +1039,9 @@ class FasterViT(nn.Module):
 
         num_features = int(dim * 2 ** (len(depths) - 1))
         self.num_classes = num_classes
-        self.patch_embed = PatchEmbed(in_chans=in_chans, in_dim=in_dim, dim=dim, shuffle_down=shuffle_down)
+        self.patch_embed = PatchEmbed(
+            in_chans=in_chans, in_dim=in_dim, dim=dim, shuffle_down=shuffle_down
+        )
         # set return_full_features true if we want to return full features from all stages
         self.return_full_features = return_full_features
         self.use_neck = use_neck
@@ -820,32 +1050,35 @@ class FasterViT(nn.Module):
         if drop_uniform:
             dpr = [drop_path_rate for x in range(sum(depths))]
 
-        if not isinstance(max_depth, list): max_depth = [max_depth] * len(depths)
+        if not isinstance(max_depth, list):
+            max_depth = [max_depth] * len(depths)
 
         self.levels = nn.ModuleList()
         for i in range(len(depths)):
             conv = True if (i == 0 or i == 1) else False
 
-            level = FasterViTLayer(dim=int(dim * 2 ** i),
-                                   depth=depths[i],
-                                   num_heads=num_heads[i],
-                                   window_size=window_size[i],
-                                   mlp_ratio=mlp_ratio,
-                                   qkv_bias=qkv_bias,
-                                   qk_scale=qk_scale,
-                                   conv=conv,
-                                   drop_path=dpr[sum(depths[:i]):sum(depths[:i + 1])],
-                                   downsample=(i < 3),
-                                   layer_scale=layer_scale,
-                                   layer_scale_conv=layer_scale_conv,
-                                   sr_ratio=sr_ratio[i],
-                                   use_swiglu=use_swiglu,
-                                   multi_query=multi_query,
-                                   norm_layer=norm_layer,
-                                   rep_vgg=rep_vgg,
-                                   yolo_arch=yolo_arch,
-                                   downsample_shuffle=downsample_shuffle,
-                                   conv_base=conv_base)
+            level = FasterViTLayer(
+                dim=int(dim * 2 ** i),
+                depth=depths[i],
+                num_heads=num_heads[i],
+                window_size=window_size[i],
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                conv=conv,
+                drop_path=dpr[sum(depths[:i]) : sum(depths[: i + 1])],
+                downsample=(i < 3),
+                layer_scale=layer_scale,
+                layer_scale_conv=layer_scale_conv,
+                sr_ratio=sr_ratio[i],
+                use_swiglu=use_swiglu,
+                multi_query=multi_query,
+                norm_layer=norm_layer,
+                rep_vgg=rep_vgg,
+                yolo_arch=yolo_arch,
+                downsample_shuffle=downsample_shuffle,
+                conv_base=conv_base,
+            )
 
             self.levels.append(level)
 
@@ -857,50 +1090,84 @@ class FasterViT(nn.Module):
             for i in range(len(depths)):
                 level_n_features_output = int(dim * 2 ** i)
 
-                if self.neck_start_stage > i: continue
+                if self.neck_start_stage > i:
+                    continue
 
-                if (upsample_ratio > 1) or full_features_head_dim!=level_n_features_output:
+                if (
+                    upsample_ratio > 1
+                ) or full_features_head_dim != level_n_features_output:
                     feature_projection = nn.Sequential()
                     # feature_projection.add_module("norm",LayerNorm2d(level_n_features_output)) #slow, but better
 
-
-                    if 0    :
+                    if 0:
                         # Train: 0 [1900/10009 ( 19%)]  Loss: 6.113 (6.57)  Time: 0.548s,  233.40/s  (0.549s,  233.04/s)  LR: 1.000e-05  Data: 0.015 (0.013)
-                        feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output)) #fast, but worse
-                        feature_projection.add_module("dconv", nn.ConvTranspose2d(level_n_features_output,
-                                                                                full_features_head_dim, kernel_size=upsample_ratio, stride=upsample_ratio))
+                        feature_projection.add_module(
+                            "norm", nn.BatchNorm2d(level_n_features_output)
+                        )  # fast, but worse
+                        feature_projection.add_module(
+                            "dconv",
+                            nn.ConvTranspose2d(
+                                level_n_features_output,
+                                full_features_head_dim,
+                                kernel_size=upsample_ratio,
+                                stride=upsample_ratio,
+                            ),
+                        )
                     else:
                         # pixel shuffle based upsampling
                         # Train: 0 [1950/10009 ( 19%)]  Loss: 6.190 (6.55)  Time: 0.540s,  236.85/s  (0.548s,  233.38/s)  LR: 1.000e-05  Data: 0.015 (0.013)
-                        feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output)) #fast, but worse
-                        feature_projection.add_module("conv", nn.Conv2d(level_n_features_output,
-                                                                                full_features_head_dim*upsample_ratio*upsample_ratio, kernel_size=1, stride=1))
-                        feature_projection.add_module("upsample_pixelshuffle", nn.PixelShuffle(upsample_ratio))
+                        feature_projection.add_module(
+                            "norm", nn.BatchNorm2d(level_n_features_output)
+                        )  # fast, but worse
+                        feature_projection.add_module(
+                            "conv",
+                            nn.Conv2d(
+                                level_n_features_output,
+                                full_features_head_dim
+                                * upsample_ratio
+                                * upsample_ratio,
+                                kernel_size=1,
+                                stride=1,
+                            ),
+                        )
+                        feature_projection.add_module(
+                            "upsample_pixelshuffle", nn.PixelShuffle(upsample_ratio)
+                        )
 
                 else:
                     feature_projection = nn.Sequential()
-                    feature_projection.add_module("norm",nn.BatchNorm2d(level_n_features_output))
-
+                    feature_projection.add_module(
+                        "norm", nn.BatchNorm2d(level_n_features_output)
+                    )
 
                 self.neck_features_proj.append(feature_projection)
 
-                if i>0 and self.levels[i-1].downsample is not None:
+                if i > 0 and self.levels[i - 1].downsample is not None:
                     upsample_ratio *= 2
 
-
-        num_features = full_features_head_dim if (self.return_full_features or self.use_neck) else num_features
+        num_features = (
+            full_features_head_dim
+            if (self.return_full_features or self.use_neck)
+            else num_features
+        )
 
         self.num_features = num_features
 
-        self.norm = LayerNorm2d(num_features) if layer_norm_last else nn.BatchNorm2d(num_features)
+        self.norm = (
+            LayerNorm2d(num_features)
+            if layer_norm_last
+            else nn.BatchNorm2d(num_features)
+        )
         self.avgpool = nn.AdaptiveAvgPool2d(1)
-        self.head = nn.Linear(num_features, num_classes) if num_classes > 0 else nn.Identity()
+        self.head = (
+            nn.Linear(num_features, num_classes) if num_classes > 0 else nn.Identity()
+        )
         self.apply(self._init_weights)
         # pass
 
     def _init_weights(self, m):
         if isinstance(m, nn.Linear):
-            trunc_normal_(m.weight, std=.02)
+            trunc_normal_(m.weight, std=0.02)
             if isinstance(m, nn.Linear) and m.bias is not None:
                 nn.init.constant_(m.bias, 0)
         elif isinstance(m, nn.LayerNorm):
@@ -915,7 +1182,7 @@ class FasterViT(nn.Module):
 
     @torch.jit.ignore
     def no_weight_decay_keywords(self):
-        return {'rpb'}
+        return {"rpb"}
 
     def forward_features(self, x):
         x = self.patch_embed(x)
@@ -924,18 +1191,34 @@ class FasterViT(nn.Module):
             x, pre_downsample_x = level(x)
 
             if self.return_full_features or self.use_neck:
-                if self.neck_start_stage > il: continue
+                if self.neck_start_stage > il:
+                    continue
                 if full_features is None:
-                    full_features = self.neck_features_proj[il - self.neck_start_stage](pre_downsample_x)
+                    full_features = self.neck_features_proj[il - self.neck_start_stage](
+                        pre_downsample_x
+                    )
                 else:
-                    #upsample torch tensor x to match full_features size, and add to full_features
-                    feature_projection = self.neck_features_proj[il - self.neck_start_stage](pre_downsample_x)
-                    if feature_projection.shape[2] != full_features.shape[2] or feature_projection.shape[3] != full_features.shape[3]:
-                        feature_projection = torch.nn.functional.pad(feature_projection, ( 0, -feature_projection.shape[3] + full_features.shape[3], 0, -feature_projection.shape[2] + full_features.shape[2]))
+                    # upsample torch tensor x to match full_features size, and add to full_features
+                    feature_projection = self.neck_features_proj[
+                        il - self.neck_start_stage
+                    ](pre_downsample_x)
+                    if (
+                        feature_projection.shape[2] != full_features.shape[2]
+                        or feature_projection.shape[3] != full_features.shape[3]
+                    ):
+                        feature_projection = torch.nn.functional.pad(
+                            feature_projection,
+                            (
+                                0,
+                                -feature_projection.shape[3] + full_features.shape[3],
+                                0,
+                                -feature_projection.shape[2] + full_features.shape[2],
+                            ),
+                        )
                     full_features += feature_projection
 
         # x = self.norm(full_features if (self.return_full_features or self.use_neck) else x)
-        x = self.norm(x) # new version for
+        x = self.norm(x)  # new version for
         x = self.avgpool(x)
         x = torch.flatten(x, 1)
 
@@ -952,384 +1235,435 @@ class FasterViT(nn.Module):
         return x
 
     def switch_to_deploy(self):
-        '''
+        """
         A method to perform model self-compression
         merges BN into conv layers
         converts MLP relative positional bias into precomputed buffers
-        '''
+        """
         for level in [self.patch_embed, self.levels, self.head]:
             for module in level.modules():
-                if hasattr(module, 'switch_to_deploy'):
+                if hasattr(module, "switch_to_deploy"):
                     module.switch_to_deploy()
 
+
 @register_model
-def fastervit2_small(pretrained=False, **kwargs): #,
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=96,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.2,
-                     sr_ratio=[1, 1, [1, 2], 1],
-                     use_swiglu=False,
-                     downsample_shuffle=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     **kwargs)
+def fastervit2_small(pretrained=False, **kwargs):  # ,
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=96,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.2,
+        sr_ratio=[1, 1, [1, 2], 1],
+        use_swiglu=False,
+        downsample_shuffle=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
+
 @register_model
-def fastervit2_tiny(pretrained=False, **kwargs): #,
-    model = FasterViT(depths=[1, 3, 4, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=80,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.2,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     downsample_shuffle=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     **kwargs)
+def fastervit2_tiny(pretrained=False, **kwargs):  # ,
+    model = FasterViT(
+        depths=[1, 3, 4, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=80,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.2,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        downsample_shuffle=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
+
 @register_model
 def fastervit2_base(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=128,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.2,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     conv_base=True,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=128,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.2,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
+
 @register_model
 def fastervit2_base_fullres1(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=128,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.2,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     conv_base=True,
-                     use_neck=True,
-                     full_features_head_dim=1024,
-                     neck_start_stage=2,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=128,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.2,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=1024,
+        neck_start_stage=2,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
+
 @register_model
 def fastervit2_base_fullres2(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=128,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.2,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     conv_base=True,
-                     use_neck=True,
-                     full_features_head_dim=512,
-                     neck_start_stage=1,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=128,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.2,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=512,
+        neck_start_stage=1,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
+
 @register_model
 def fastervit2_base_fullres3(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=128,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.2,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     conv_base=True,
-                     use_neck=True,
-                     full_features_head_dim=256,
-                     neck_start_stage=1,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=128,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.2,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=256,
+        neck_start_stage=1,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
+
 @register_model
 def fastervit2_base_fullres4(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=128,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.2,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     conv_base=True,
-                     use_neck=True,
-                     full_features_head_dim=256,
-                     neck_start_stage=2,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=128,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.2,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=256,
+        neck_start_stage=2,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
+
 @register_model
 def fastervit2_base_fullres5(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=128,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.2,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     conv_base=True,
-                     use_neck=True,
-                     full_features_head_dim=512,
-                     neck_start_stage=2,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=128,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.2,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=512,
+        neck_start_stage=2,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
-#pyt: 1934, 4202 TRT
+
+# pyt: 1934, 4202 TRT
 @register_model
 def fastervit2_large(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=128+64,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.2,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=128 + 64,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.2,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
+
 @register_model
 def fastervit2_large_fullres(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[None, None, [7, 7], 7],
-                     dim=192,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     conv_base=True,
-                     use_neck=True,
-                     full_features_head_dim=1536,
-                     neck_start_stage=2,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[None, None, [7, 7], 7],
+        dim=192,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=1536,
+        neck_start_stage=2,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
+
 @register_model
 def fastervit2_large_fullres_ws8(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[None, None, [8, 8], 8],
-                     dim=192,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     conv_base=True,
-                     use_neck=True,
-                     full_features_head_dim=1536,
-                     neck_start_stage=2,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[None, None, [8, 8], 8],
+        dim=192,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=1536,
+        neck_start_stage=2,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
+
 @register_model
 def fastervit2_large_fullres_ws16(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[None, None, [16, 16], 16],
-                     dim=192,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     conv_base=True,
-                     use_neck=True,
-                     full_features_head_dim=1536,
-                     neck_start_stage=2,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[None, None, [16, 16], 16],
+        dim=192,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=1536,
+        neck_start_stage=2,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
+
 @register_model
 def fastervit2_large_fullres_ws32(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[None, None, [32, 32], 32],
-                     dim=192,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     conv_base=True,
-                     use_neck=True,
-                     full_features_head_dim=1536,
-                     neck_start_stage=2,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[None, None, [32, 32], 32],
+        dim=192,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        conv_base=True,
+        use_neck=True,
+        full_features_head_dim=1536,
+        neck_start_stage=2,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
-#pyt: 897
+
+# pyt: 897
 @register_model
 def fastervit2_xlarge(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=128+128+64,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.2,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=128 + 128 + 64,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.2,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
 
-#pyt:
+# pyt:
 @register_model
 def fastervit2_huge(pretrained=False, **kwargs):
-    model = FasterViT(depths=[3, 3, 5, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=128+128+128+64,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.2,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     **kwargs)
+    model = FasterViT(
+        depths=[3, 3, 5, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=128 + 128 + 128 + 64,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.2,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
 
 @register_model
-def fastervit2_xtiny(pretrained=False, **kwargs): #,
-    model = FasterViT(depths=[1, 3, 4, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=64,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.1,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     downsample_shuffle=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     **kwargs)
+def fastervit2_xtiny(pretrained=False, **kwargs):  # ,
+    model = FasterViT(
+        depths=[1, 3, 4, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=64,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.1,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        downsample_shuffle=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
 
 @register_model
-def fastervit2_xxtiny_5(pretrained=False, **kwargs): #,
-    model = FasterViT(depths=[1, 3, 4, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=48,
-                     in_dim=64,
-                     mlp_ratio=4,
-                     drop_path_rate=0.05,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     downsample_shuffle=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     **kwargs)
+def fastervit2_xxtiny_5(pretrained=False, **kwargs):  # ,
+    model = FasterViT(
+        depths=[1, 3, 4, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=48,
+        in_dim=64,
+        mlp_ratio=4,
+        drop_path_rate=0.05,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        downsample_shuffle=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
 
+
 @register_model
-def fastervit2_xxxtiny(pretrained=False, **kwargs): #,
-    model = FasterViT(depths=[1, 3, 4, 5],
-                     num_heads=[2, 4, 8, 16],
-                     window_size=[8, 8, [7, 7], 7],
-                     dim=32,
-                     in_dim=32,
-                     mlp_ratio=4,
-                     drop_path_rate=0.0,
-                     sr_ratio=[1, 1, [2, 1], 1],
-                     use_swiglu=False,
-                     downsample_shuffle=False,
-                     yolo_arch=True,
-                     shuffle_down=False,
-                     **kwargs)
+def fastervit2_xxxtiny(pretrained=False, **kwargs):  # ,
+    model = FasterViT(
+        depths=[1, 3, 4, 5],
+        num_heads=[2, 4, 8, 16],
+        window_size=[8, 8, [7, 7], 7],
+        dim=32,
+        in_dim=32,
+        mlp_ratio=4,
+        drop_path_rate=0.0,
+        sr_ratio=[1, 1, [2, 1], 1],
+        use_swiglu=False,
+        downsample_shuffle=False,
+        yolo_arch=True,
+        shuffle_down=False,
+        **kwargs,
+    )
     if pretrained:
         model.load_state_dict(torch.load(pretrained))
     return model
@@ -1337,4 +1671,4 @@ def fastervit2_xxxtiny(pretrained=False, **kwargs): #,
 
 @register_model
 def eradio(pretrained=False, **kwargs):
-    return fastervit2_large_fullres(pretrained=pretrained, **kwargs)
+    return fastervit2_large_fullres_ws16(pretrained=pretrained, **kwargs)

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:115b8f54d0d4999c180718ce138f8078127af9815b6cb507b253e5db10a5723c
-size 1057766065
+oid sha256:3840092575224b5ff90adf3b7970a5a5e379f8988241ee8145969e27c32c17e7
+size 1105844337