test

app.py

@@ -89,30 +89,30 @@ model_option = gr.Radio(options, value="dino16",
 
 @spaces.GPU
 def upsample_features(image, model_option):
-    
-    subprocess.check_call(
-        ["pip", "install", "git+https://github.com/mhamilton723/FeatUp"])
-
-    from featup.util import norm, unnorm
-    models = {o: torch.hub.load("mhamilton723/FeatUp", o) for o in options}
-
-    # Image preprocessing
-    input_size = 224
-    transform = T.Compose([
-        T.Resize(input_size),
-        T.CenterCrop((input_size, input_size)),
-        T.ToTensor(),
-        norm
-    ])
-    image_tensor = transform(image).unsqueeze(0).cuda()
-
-    # Load the selected model
-    upsampler = models[model_option].cuda()
-    hr_feats = upsampler(image_tensor)
-    lr_feats = upsampler.model(image_tensor)
-    upsampler.cpu()
-
-    return plot_feats(unnorm(image_tensor)[0], lr_feats[0], hr_feats[0])
+    with torch.no_grad():
+        subprocess.check_call(
+            ["pip", "install", "git+https://github.com/mhamilton723/FeatUp"])
+
+        from featup.util import norm, unnorm
+        models = {o: torch.hub.load("mhamilton723/FeatUp", o) for o in options}
+
+        # Image preprocessing
+        input_size = 224
+        transform = T.Compose([
+            T.Resize(input_size),
+            T.CenterCrop((input_size, input_size)),
+            T.ToTensor(),
+            norm
+        ])
+        image_tensor = transform(image).unsqueeze(0).cuda()
+
+        # Load the selected model
+        upsampler = models[model_option].cuda()
+        hr_feats = upsampler(image_tensor)
+        lr_feats = upsampler.model(image_tensor)
+        upsampler.cpu()
+
+        return plot_feats(unnorm(image_tensor)[0], lr_feats[0], hr_feats[0])
 
 
 demo = gr.Interface(fn=upsample_features,