Changed running test code inline to running it in separate script

app.py

@@ -1,14 +1,8 @@
-from __future__ import absolute_import
-from __future__ import print_function
-from __future__ import division
-from torch.autograd import gradcheck
 import spaces
 import gradio as gr
 import copy
 import random
 import torch
-import torch.nn as nn
-import time
 import PIL
 from PIL import Image, ImageDraw, ImageFont
 import torchvision.transforms.functional as F
@@ -46,7 +40,12 @@ subprocess.run(
         "pip install MultiScaleDeformableAttention-1.0-cp310-cp310-linux_x86_64.whl"
     )
 )
-from functions.ms_deform_attn_func import MSDeformAttnFunction, ms_deform_attn_core_pytorch
+
+subprocess.run(
+    shlex.split(
+        "python test.py"
+    )
+)
 
 class AppSteps(Enum):
     JUST_TEXT = 1
@@ -56,74 +55,6 @@ class AppSteps(Enum):
 
 CONF_THRESH = 0.23
 
-@spaces.GPU
-def check_ms_deform_install():
-    
-    N, M, D = 1, 2, 2
-    Lq, L, P = 2, 2, 2
-    shapes = torch.as_tensor([(6, 4), (3, 2)], dtype=torch.long).cuda()
-    level_start_index = torch.cat((shapes.new_zeros((1, )), shapes.prod(1).cumsum(0)[:-1]))
-    S = sum([(H*W).item() for H, W in shapes])
-    
-    
-    torch.manual_seed(3)
-
-
-    @torch.no_grad()
-    def check_forward_equal_with_pytorch_double():
-        value = torch.rand(N, S, M, D).cuda() * 0.01
-        sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
-        attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
-        attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True)
-        im2col_step = 2
-        output_pytorch = ms_deform_attn_core_pytorch(value.double(), shapes, sampling_locations.double(), attention_weights.double()).detach().cpu()
-        output_cuda = MSDeformAttnFunction.apply(value.double(), shapes, level_start_index, sampling_locations.double(), attention_weights.double(), im2col_step).detach().cpu()
-        fwdok = torch.allclose(output_cuda, output_pytorch)
-        max_abs_err = (output_cuda - output_pytorch).abs().max()
-        max_rel_err = ((output_cuda - output_pytorch).abs() / output_pytorch.abs()).max()
-    
-        print(f'* {fwdok} check_forward_equal_with_pytorch_double: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
-    
-    
-    @torch.no_grad()
-    def check_forward_equal_with_pytorch_float():
-        value = torch.rand(N, S, M, D).cuda() * 0.01
-        sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
-        attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
-        attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True)
-        im2col_step = 2
-        output_pytorch = ms_deform_attn_core_pytorch(value, shapes, sampling_locations, attention_weights).detach().cpu()
-        output_cuda = MSDeformAttnFunction.apply(value, shapes, level_start_index, sampling_locations, attention_weights, im2col_step).detach().cpu()
-        fwdok = torch.allclose(output_cuda, output_pytorch, rtol=1e-2, atol=1e-3)
-        max_abs_err = (output_cuda - output_pytorch).abs().max()
-        max_rel_err = ((output_cuda - output_pytorch).abs() / output_pytorch.abs()).max()
-    
-        print(f'* {fwdok} check_forward_equal_with_pytorch_float: max_abs_err {max_abs_err:.2e} max_rel_err {max_rel_err:.2e}')
-
-
-    def check_gradient_numerical(channels=4, grad_value=True, grad_sampling_loc=True, grad_attn_weight=True):
-    
-        value = torch.rand(N, S, M, channels).cuda() * 0.01
-        sampling_locations = torch.rand(N, Lq, M, L, P, 2).cuda()
-        attention_weights = torch.rand(N, Lq, M, L, P).cuda() + 1e-5
-        attention_weights /= attention_weights.sum(-1, keepdim=True).sum(-2, keepdim=True)
-        im2col_step = 2
-        func = MSDeformAttnFunction.apply
-    
-        value.requires_grad = grad_value
-        sampling_locations.requires_grad = grad_sampling_loc
-        attention_weights.requires_grad = grad_attn_weight
-    
-        gradok = gradcheck(func, (value.double(), shapes, level_start_index, sampling_locations.double(), attention_weights.double(), im2col_step))
-    
-        print(f'* {gradok} check_gradient_numerical(D={channels})')
-    
-    check_forward_equal_with_pytorch_double()
-    check_forward_equal_with_pytorch_float()
-
-    for channels in [30, 32, 64, 71]:
-        check_gradient_numerical(channels, True, True, True)
-
 # MODEL:
 def get_args_parser():
     """