1919

app.py

@@ -8,48 +8,45 @@ from models.arch.RDnet_ import FullNet_NLP
 from models.arch.classifier import PretrainedConvNext
 import torchvision.transforms.functional as TF
 
-
-channels = [64, 128, 256, 512]
-layers = [2, 2, 4, 2]
-num_subnet = 4
-net_i = FullNet_NLP(channels, layers, num_subnet, 4,num_classes=1000, drop_path=0,save_memory=True, inter_supv=True, head_init_scale=None,kernel_size=3)
-for param in net_i.parameters():
-    param.data = param.data.to(torch.float16)
-net_i.load_state_dict(torch.load('./fp16_check.pt')['icnn'])
-net_i=net_i.to('cpu')
-#net_i=net_i.to('cuda')
-net_c = PretrainedConvNext("convnext_small_in22k")
-net_c.load_state_dict(torch.load('./classifier_32.pt')['icnn'])
-net_c=net_c.to('cpu')
-#net_c=net_c.to('cuda')
-net_i.eval().to('cuda')
-net_c.eval().to('cuda')
-
-def align(x1):
-
-    return x1
-
+class Pipe:
+    def __init__(self):
+        channels = [64, 128, 256, 512]
+        layers = [2, 2, 4, 2]
+        num_subnet = 4
+        self.net_i = FullNet_NLP(channels, layers, num_subnet, 4,num_classes=1000, drop_path=0,save_memory=True, inter_supv=True, head_init_scale=None,kernel_size=3)
+        for param in self.net_i.parameters():
+            param.data = param.data.to(torch.float16)
+        self.net_i.load_state_dict(torch.load('./fp16_check.pt')['icnn'])
+        self.net_i = self.net_i.to('cpu')
+        self.net_c = PretrainedConvNext("convnext_small_in22k")
+        self.net_c.load_state_dict(torch.load('./classifier_32.pt')['icnn'])
+        self.net_c=self.net_c.to('cpu')
+        #net_c=net_c.to('cuda')
+        self.net_i.eval().to('cuda')
+        self.net_c.eval().to('cuda')
+        self.output = None
+
+    def __call__(self, img):
+        with torch.no_grad():
+            image_tensor = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0)
+            h, w = image_tensor.shape[-2], image_tensor.shape[-1]
+            h, w = h // 32 * 32, w // 32 * 32
+            image_tensor = torch.nn.functional.interpolate(image_tensor, size=(h, w), mode='bilinear')
+            ipt=net_c(image_tensor)
+            image_tensor = image_tensor.half()
+            ipt = ipt.half()
+            output_i, output_j=net_i(image_tensor,ipt,prompt=True)
+            output_j_out=[]
+            clean = output[-1][:, 3:, ...]
+            clean=torch.clamp(clean, 0, 1)
+            self.output = clean
+
+pipe = Pipe()
 
 
 @spaces.GPU(duration=120)
 def predict(img):
-    with torch.no_grad():
-        image_tensor = torch.from_numpy(img).permute(2, 0, 1).float().unsqueeze(0)
-        h, w = image_tensor.shape[-2], image_tensor.shape[-1]
-        h, w = h // 32 * 32, w // 32 * 32
-        image_tensor = torch.nn.functional.interpolate(image_tensor, size=(h, w), mode='bilinear')
-        ipt=net_c(image_tensor)
-        image_tensor = image_tensor.half()
-        ipt = ipt.half()
-        output_i, output_j=net_i(image_tensor,ipt,prompt=True)
-        output_j_out=[]
-        for i in range(4):
-            out_reflection, out_clean = output_j[i][:, :3, ...], output_j[i][:, 3:, ...]
-            output_j_out.append(out_clean)  
-            output_j_out.append(out_reflection)
-        clean = output_j_out[6]
-        clean=torch.clamp(clean, 0, 1)
-    return clean
+    return pipe(img).output
 
 demo=gr.Interface(predict, gr.Image(), "image")