added stylematte, fix gpu issue

app.py

@@ -3,7 +3,7 @@ from test import inference_img
 from models import *
 
 device='cpu'
-model = MaskForm()
+model = StyleMatte()
 model = model.to(device)
 checkpoint = f"stylematte.pth"
 state_dict = torch.load(checkpoint, map_location=f'{device}')

models.py

@@ -8,7 +8,7 @@ from typing import List
 from itertools import chain
 
 from transformers import SegformerForSemanticSegmentation,Mask2FormerForUniversalSegmentation
-
+device='cpu'
 class EncoderDecoder(nn.Module):
     def __init__(
         self,
@@ -284,9 +284,9 @@ class SegForm(nn.Module):
         return upsampled_logits
  
 
-class MaskForm(nn.Module):
+class StyleMatte(nn.Module):
     def __init__(self):
-        super(MaskForm, self).__init__()
+        super(StyleMatte, self).__init__()
 #         configuration = SegformerConfig.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
 #         configuration.num_labels = 1 ## set output as 1 
         self.fpn = FPN_fuse(feature_channels=[256, 256, 256, 256],fpn_out=256)
@@ -476,6 +476,6 @@ class BoxFilter(nn.Module):
         return diff_y(diff_x(x.cumsum(dim=2), self.r).cumsum(dim=3), self.r)
     
 if __name__ == '__main__':
-    model = MaskForm().cuda()
-    out=model(torch.randn(1,3,640,480).cuda())
+    model = StyleMatte().to(device)
+    out=model(torch.randn(1,3,640,480).to(devuce))
     print(out.shape)

test.py

@@ -26,7 +26,7 @@ import logging
 import time
 from omegaconf import OmegaConf
 config = OmegaConf.load("base.yaml")
-device = "cuda"
+device = "cpu"
 
 def conv3x3(in_planes, out_planes, stride=1):
     "3x3 convolution with padding"
@@ -720,7 +720,7 @@ def get_masked_local_from_global_test(global_result, local_result):
     return fusion_result
 def inference_once( model, scale_img, scale_trimap=None):
     pred_list = []
-    tensor_img = torch.from_numpy(scale_img[:, :, :]).permute(2, 0, 1).cuda()
+    tensor_img = torch.from_numpy(scale_img[:, :, :]).permute(2, 0, 1).to(device)
     input_t = tensor_img
     input_t = input_t/255.0
     normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],