Upgrade the ways of importing timm modules (>=1.0.23).

birefnet.py

@@ -164,8 +164,8 @@ import torch
 import torch.nn as nn
 from functools import partial
 
-from timm.models.layers import DropPath, to_2tuple, trunc_normal_
-from timm.models.registry import register_model
+from timm.layers import DropPath, to_2tuple, trunc_normal_
+
 
 import math
 
@@ -545,7 +545,6 @@ def _conv_filter(state_dict, patch_size=16):
     return out_dict
 
 
-## @register_model
 class pvt_v2_b0(PyramidVisionTransformerImpr):
     def __init__(self, **kwargs):
         super(pvt_v2_b0, self).__init__(
@@ -555,7 +554,6 @@ class pvt_v2_b0(PyramidVisionTransformerImpr):
 
 
 
-## @register_model
 class pvt_v2_b1(PyramidVisionTransformerImpr):
     def __init__(self, **kwargs):
         super(pvt_v2_b1, self).__init__(
@@ -563,7 +561,6 @@ class pvt_v2_b1(PyramidVisionTransformerImpr):
             qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1],
             drop_rate=0.0, drop_path_rate=0.1)
 
-## @register_model
 class pvt_v2_b2(PyramidVisionTransformerImpr):
     def __init__(self, in_channels=3, **kwargs):
         super(pvt_v2_b2, self).__init__(
@@ -571,7 +568,6 @@ class pvt_v2_b2(PyramidVisionTransformerImpr):
             qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1],
             drop_rate=0.0, drop_path_rate=0.1, in_channels=in_channels)
 
-## @register_model
 class pvt_v2_b3(PyramidVisionTransformerImpr):
     def __init__(self, **kwargs):
         super(pvt_v2_b3, self).__init__(
@@ -579,7 +575,6 @@ class pvt_v2_b3(PyramidVisionTransformerImpr):
             qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1],
             drop_rate=0.0, drop_path_rate=0.1)
 
-## @register_model
 class pvt_v2_b4(PyramidVisionTransformerImpr):
     def __init__(self, **kwargs):
         super(pvt_v2_b4, self).__init__(
@@ -588,7 +583,6 @@ class pvt_v2_b4(PyramidVisionTransformerImpr):
             drop_rate=0.0, drop_path_rate=0.1)
 
 
-## @register_model
 class pvt_v2_b5(PyramidVisionTransformerImpr):
     def __init__(self, **kwargs):
         super(pvt_v2_b5, self).__init__(
@@ -612,7 +606,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torch.utils.checkpoint as checkpoint
 import numpy as np
-from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.layers import DropPath, to_2tuple, trunc_normal_
 
 # from config import Config
 
@@ -1193,7 +1187,7 @@ class SwinTransformer(nn.Module):
             # interpolate the position embedding to the corresponding size
             absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic')
             x = (x + absolute_pos_embed) # B Wh*Ww C
-            
+
         outs = []#x.contiguous()]
         x = x.flatten(2).transpose(1, 2)
         x = self.pos_drop(x)
@@ -1250,13 +1244,13 @@ class DeformableConv2d(nn.Module):
                  bias=False):
 
         super(DeformableConv2d, self).__init__()
-        
+
         assert type(kernel_size) == tuple or type(kernel_size) == int
 
         kernel_size = kernel_size if type(kernel_size) == tuple else (kernel_size, kernel_size)
         self.stride = stride if type(stride) == tuple else (stride, stride)
         self.padding = padding
-        
+
         self.offset_conv = nn.Conv2d(in_channels,
                                      2 * kernel_size[0] * kernel_size[1],
                                      kernel_size=kernel_size,
@@ -1266,7 +1260,7 @@ class DeformableConv2d(nn.Module):
 
         nn.init.constant_(self.offset_conv.weight, 0.)
         nn.init.constant_(self.offset_conv.bias, 0.)
-        
+
         self.modulator_conv = nn.Conv2d(in_channels,
                                      1 * kernel_size[0] * kernel_size[1],
                                      kernel_size=kernel_size,
@@ -1290,7 +1284,7 @@ class DeformableConv2d(nn.Module):
 
         offset = self.offset_conv(x)#.clamp(-max_offset, max_offset)
         modulator = 2. * torch.sigmoid(self.modulator_conv(x))
-        
+
         x = deform_conv2d(
             input=x,
             offset=offset,
@@ -1488,7 +1482,7 @@ class ResBlk(nn.Module):
 
         self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1)
         self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
-        
+
         self.conv_resi = nn.Conv2d(in_channels, out_channels, 1, 1, 0)
 
     def forward(self, x):
@@ -2139,7 +2133,7 @@ class Decoder(nn.Module):
                 self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
                 self.gdt_convs_pred_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
                 self.gdt_convs_pred_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
-                
+
                 self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
                 self.gdt_convs_attn_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
                 self.gdt_convs_attn_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))