Update

app.py

@@ -1,5 +1,6 @@
-from collections import namedtuple
+import gc
 import io
+from collections import namedtuple
 from typing import Tuple
 
 import gradio as gr
@@ -48,7 +49,11 @@ def get_model(name: str):
 def get_transform(name: str):
     if any(x in name for x in ('dino', 'mocov3', 'convnext', )):
         normalize = transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
-        transform = transforms.Compose([transforms.ToTensor(), normalize])
+        transform = transforms.Compose([
+            transforms.Resize(size=256, interpolation=TF.InterpolationMode.BICUBIC, max_size=384),
+            transforms.ToTensor(), 
+            normalize
+        ])
     else:
         raise NotImplementedError()
     return transform
@@ -60,17 +65,28 @@ def get_diagonal(W: scipy.sparse.csr_matrix, threshold: float = 1e-12):
     D = scipy.sparse.diags(D)
     return D
 
+# Cache
+torch.cuda.empty_cache()
 
 # Parameters
 model_name = 'dino_vitb16'  # TODO: Figure out how to make this user-editable
 K = 5
 
-# Fixed parameters
-MAX_SIZE = 384
-
 # Load model
 model, val_transform, patch_size, num_heads = get_model(model_name)
 
+# Add hook
+which_block = -1
+if 'dino' in model_name or 'mocov3' in model_name:
+    feat_out = {}
+    def hook_fn_forward_qkv(module, input, output):
+        feat_out["qkv"] = output
+    handle: RemovableHandle = model._modules["blocks"][which_block]._modules["attn"]._modules["qkv"].register_forward_hook(
+        hook_fn_forward_qkv
+    )
+else:
+    raise ValueError(model_name)
+
 
 # GPU
 if torch.cuda.is_available():
@@ -90,18 +106,6 @@ def segment(inp: Image):
     images: torch.Tensor = val_transform(inp)
     images = images.unsqueeze(0).to(device)
 
-    # Add hook
-    which_block = -1
-    if 'dino' in model_name or 'mocov3' in model_name:
-        feat_out = {}
-        def hook_fn_forward_qkv(module, input, output):
-            feat_out["qkv"] = output
-        handle: RemovableHandle = model._modules["blocks"][which_block]._modules["attn"]._modules["qkv"].register_forward_hook(
-            hook_fn_forward_qkv
-        )
-    else:
-        raise ValueError(model_name)
-
     # Reshape image
     P = patch_size
     B, C, H, W = images.shape
@@ -119,9 +123,6 @@ def segment(inp: Image):
         feats = output_qkv[1].transpose(1, 2).reshape(B, T, -1)[:, 1:, :].squeeze(0)
     else:
         raise ValueError(model_name)
-
-    # Remove hook from the model
-    handle.remove()
     
     # Normalize features
     normalize = True
@@ -160,7 +161,7 @@ def segment(inp: Image):
             eigenvectors[k] = 0 - eigenvectors[k]
 
     # Arrange eigenvectors into grid
-    cmap = get_cmap('viridis')
+    # cmap = get_cmap('viridis')
     output_images = []
     # eigenvectors_upscaled = []
     for i in range(1, K + 1):
@@ -209,6 +210,11 @@ def segment(inp: Image):
     # # Postprocess for Gradio
     # output_images = np.array(TF.to_pil_image(output_images))
     print(f'{len(output_images)=}')
+
+    # Garbage collection and other memory-related things
+    gc.collect()
+    del eigenvector, eigenvector_vis, eigenvectors, W_comb, D_comb
+
     return output_images
 
 # Placeholders