Run kway ncut on CPU

dino_correspondence.py

@@ -24,6 +24,14 @@ def _get_compute_device(tensor: torch.Tensor) -> str:
     device_type = tensor.device.type
     return device_type if device_type in {"cpu", "cuda"} else "cpu"
 
+
+def _run_kway_ncut_on_cpu(eigenvectors: torch.Tensor, n_clusters: int) -> torch.Tensor:
+    """Keep NCut discretization on CPU to avoid unsupported CUDA kernels on HF ZeroGPU."""
+    return kway_ncut(
+        eigenvectors[:, :n_clusters].cpu(),
+        device="cpu",
+    )
+
 def ncut_tsne_multiple_images(image_embeds, n_eig=50, gamma=None, degree=0.5):
     """
     Apply NCut and t-SNE coloring to multiple image embeddings.
@@ -68,7 +76,7 @@ def _kway_cluster_single_image(image_embeds, n_clusters, gamma=None, degree=0.5)
         flattened_input, n_eig=n_eig, gamma=gamma, device=compute_device
     )
     
-    continuous_clusters = kway_ncut(eigenvectors[:, :n_clusters])
+    continuous_clusters = _run_kway_ncut_on_cpu(eigenvectors, n_clusters)
     return continuous_clusters
 
 
@@ -111,7 +119,7 @@ def kway_cluster_multiple_images(image_embeds, n_clusters, gamma=None, degree=0.
         flattened_input, n_eig=n_eig, gamma=gamma, device=compute_device
     )
     
-    continuous_clusters = kway_ncut(eigenvectors[:, :n_clusters])
+    continuous_clusters = _run_kway_ncut_on_cpu(eigenvectors, n_clusters)
     continuous_clusters = rearrange(
         continuous_clusters, '(b l) c -> b l c', b=batch_size
     )