sb2

app.py

@@ -24,22 +24,21 @@ net_c=net_c.to('cpu')
 #net_c=net_c.to('cuda')
 net_i.eval().to('cuda')
 net_c.eval().to('cuda')
-image_tensor = torch.zeros((1, 3, 420, 420), device='cuda')
+
 def align(x1):
     h, w = x1.shape[:2]
     h, w = h // 32 * 32, w // 32 * 32
     x1 = np.resize(x1, (h, w, x1.shape[2]))
     return x1
 
-def pipe(img):
-    img=align(img)
-    img=TF.to_tensor(img)
-    image_tensor=img.unsqueeze(0)
-    return predict(image_tensor)
+
 
 @spaces.GPU(duration=120)
-def predict(image_tensor):
+def predict(img):
     with torch.no_grad():
+        img=align(img)
+        img=TF.to_tensor(img)
+        image_tensor=img.unsqueeze(0)
         ipt=net_c(image_tensor)
         image_tensor = image_tensor.half()
         ipt = ipt.half()