Update briarmbg.py

briarmbg.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from huggingface_hub import PyTorchModelHubMixin
 
 class REBNCONV(nn.Module):
     def __init__(self,in_ch=3,out_ch=3,dirate=1,stride=1):
@@ -344,11 +345,12 @@ class myrebnconv(nn.Module):
         return self.rl(self.bn(self.conv(x)))
 
 
-class BriaRMBG(nn.Module):
+class BriaRMBG(nn.Module, PyTorchModelHubMixin):
 
-    def __init__(self,in_ch=3,out_ch=1):
+    def __init__(self,config:dict={"in_ch":3,"out_ch":1}):
         super(BriaRMBG,self).__init__()
-
+        in_ch=config["in_ch"]
+        out_ch=config["out_ch"]
         self.conv_in = nn.Conv2d(in_ch,64,3,stride=2,padding=1)
         self.pool_in = nn.MaxPool2d(2,stride=2,ceil_mode=True)
 
@@ -451,5 +453,4 @@ class BriaRMBG(nn.Module):
         d6 = self.side6(hx6)
         d6 = _upsample_like(d6,x)
 
-        return [F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)],[hx1d,hx2d,hx3d,hx4d,hx5d,hx6]
-
+        return [F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)],[hx1d,hx2d,hx3d,hx4d,hx5d,hx6]