fix eig norm

app.py

@@ -870,6 +870,7 @@ def ncut_run(
     if not directed:
         only_eigvecs = kwargs.get("only_eigvecs", False)
         return_eigvec_and_rgb = kwargs.get("return_eigvec_and_rgb", False)
+        normalize_eigvec_return = kwargs.get("normalize_eigvec_return", False)
         
         rgb, _logging_str, eigvecs = compute_ncut(
             features,
@@ -893,12 +894,16 @@ def ncut_run(
         
         
         if only_eigvecs:
+            if normalize_eigvec_return:
+                eigvecs = F.normalize(eigvecs, dim=-1)
             eigvecs = eigvecs.to("cpu").reshape(features.shape[:-1] + (num_eig,))
             eigvecs = eigvecs.detach().numpy()
             logging_str += _logging_str
             return eigvecs, logging_str
         
         if return_eigvec_and_rgb:
+            if normalize_eigvec_return:
+                eigvecs = F.normalize(eigvecs, dim=-1)
             eigvecs = eigvecs.to("cpu").reshape(features.shape[:-1] + (num_eig,))
             eigvecs = eigvecs.detach().numpy()
             rgb = rgb.cpu().numpy()
@@ -1249,6 +1254,7 @@ def run_fn(
     directed=False,
     only_eigvecs=False,
     return_eigvec_and_rgb=False,
+    normalize_eigvec_return=False,
 ):
     # print(node_type2, head_index_text, make_symmetric)
     progress=gr.Progress()
@@ -1391,6 +1397,7 @@ def run_fn(
         "make_symmetric": make_symmetric,
         "only_eigvecs": only_eigvecs,
         "return_eigvec_and_rgb": return_eigvec_and_rgb,
+        "normalize_eigvec_return": normalize_eigvec_return,
     }
     # print(kwargs)
     
@@ -2232,16 +2239,16 @@ with demo:
                 def __run_fn(*args, **kwargs):
                     eigvecs, rgb, logging_str = run_fn(*args, **kwargs)
                     rgb_gallery = to_pil_images(rgb)
-                    # normalize the eigvecs
-                    eigvecs = torch.tensor(eigvecs)
-                    if torch.cuda.is_available():
-                        eigvecs = eigvecs.cuda()
-                    eigvecs = F.normalize(eigvecs, p=2, dim=-1)
-                    eigvecs = eigvecs.cpu().numpy()
+                    # # normalize the eigvecs
+                    # eigvecs = torch.tensor(eigvecs)
+                    # if torch.cuda.is_available():
+                    #     eigvecs = eigvecs.cuda()
+                    # eigvecs = F.normalize(eigvecs, p=2, dim=-1)
+                    # eigvecs = eigvecs.cpu().numpy()
                     return eigvecs, rgb, rgb_gallery, logging_str
                 
                 submit_button.click(
-                    partial(__run_fn, n_ret=2, return_eigvec_and_rgb=True),
+                    partial(__run_fn, n_ret=2, return_eigvec_and_rgb=True, normalize_eigvec_return=True),
                     inputs=[
                         input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
                         positive_prompt, negative_prompt,