update advanced plots

app.py

@@ -257,15 +257,19 @@ def blend_image_with_heatmap(image, heatmap, opacity1=0.5, opacity2=0.5):
     blended = (1 - opacity1) * image + opacity2 * heatmap
     return blended.astype(np.uint8)
 
-def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6):
+def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=False):
     progress = gr.Progress()
     progress(progess_start, desc="Finding Clusters by FPS")
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     eigvecs = eigvecs.to(device)
     from ncut_pytorch.ncut_pytorch import farthest_point_sampling
     magnitude = torch.norm(eigvecs, dim=-1)
-    p = 0.8
-    top_p_idx = magnitude.argsort(descending=True)[:int(p * magnitude.shape[0])]
+    
+    # gr.Info("Finding Clusters by FPS, no magnitude filtering")
+    top_p_idx = torch.arange(eigvecs.shape[0])
+    # gr.Info("Finding Clusters by FPS, with magnitude filtering")
+    # p = 0.8
+    # top_p_idx = magnitude.argsort(descending=True)[:int(p * magnitude.shape[0])]
     
     ret_magnitude = magnitude.reshape(-1, h, w)
     
@@ -283,7 +287,7 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6):
     right = F.normalize(right, dim=-1)
     heatmap = left @ right.T
     heatmap = F.normalize(heatmap, dim=-1)
-    num_samples = 80
+    num_samples = 80 if not advanced else 130
     if num_samples > fps_idx.shape[0]:
         num_samples = fps_idx.shape[0]
     r2_fps_idx = farthest_point_sampling(heatmap, num_samples)
@@ -328,13 +332,19 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6):
     # reverse the fps_idx
     # fps_idx = fps_idx.flip(0)
     # discard the big clusters
-    fps_idx = fps_idx[10:]
-    # shuffle the fps_idx
-    fps_idx = fps_idx[torch.randperm(fps_idx.shape[0])]
+    
+    # gr.Info("Discarding the biggest 10 clusters")
+    # fps_idx = fps_idx[10:]
+    # gr.Info("Not discarding the biggest 10 clusters")
+    # gr.Info("Discarding the smallest 30 out of 80 sampled clusters")
+    
+    if not advanced:
+        # shuffle the fps_idx
+        fps_idx = fps_idx[torch.randperm(fps_idx.shape[0])]
     
     fig_images = []
     i_cluster = 0
-    num_plots = 10
+    num_plots = 10 if not advanced else 20
     plot_step_float = (1.0 - progess_start) / num_plots
     for i_fig in range(num_plots):    
         progress(progess_start + i_fig * plot_step_float, desc="Plotting Clusters")
@@ -609,7 +619,8 @@ def ncut_run(
         if torch.cuda.is_available():
             images = images.cuda()
         _images = reverse_transform_image(images, stablediffusion="stable" in model_name.lower())
-        cluster_images, eig_magnitude = make_cluster_plot(eigvecs, _images, h=h, w=w, progess_start=progress_start)
+        advanced = kwargs.get("advanced", False)
+        cluster_images, eig_magnitude = make_cluster_plot(eigvecs, _images, h=h, w=w, progess_start=progress_start, advanced=advanced)
         logging_str += f"plot time: {time.time() - start:.2f}s\n"
     
     norm_images = None
@@ -622,8 +633,10 @@ def ncut_run(
         colormap = matplotlib.colormaps['Reds']
         for i_image in range(eig_magnitude.shape[0]):
             norm_image = colormap(eig_magnitude[i_image])
-            norm_image = (norm_image[..., :3] * 255).astype(np.uint8)
-            norm_images.append(Image.fromarray(norm_image))
+            # norm_image = (norm_image[..., :3] * 255).astype(np.uint8)
+            # norm_images.append(Image.fromarray(norm_image))
+            norm_images.append(torch.tensor(norm_image[..., :3]))
+        norm_images = to_pil_images(norm_images)
         logging_str += "Eigenvector Magnitude\n"
         logging_str += f"Min: {vmin:.2f}, Max: {vmax:.2f}\n"
         gr.Info(f"Eigenvector Magnitude:</br> Min: {vmin:.2f}, Max: {vmax:.2f}", duration=0)
@@ -855,6 +868,7 @@ def run_fn(
     n_ret=1,
     plot_clusters=False,
     alignedcut_eig_norm_plot=False,
+    advanced=False,
 ):
     
     progress=gr.Progress()
@@ -987,6 +1001,7 @@ def run_fn(
         "n_ret": n_ret,
         "plot_clusters": plot_clusters,
         "alignedcut_eig_norm_plot": alignedcut_eig_norm_plot,
+        "advanced": advanced,
     }
     # print(kwargs)
     
@@ -1416,7 +1431,7 @@ with demo:
         no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
         
         submit_button.click(
-            partial(run_fn, n_ret=3, plot_clusters=True, alignedcut_eig_norm_plot=True),
+            partial(run_fn, n_ret=3, plot_clusters=True, alignedcut_eig_norm_plot=True, advanced=True),
             inputs=[
                 input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
                 positive_prompt, negative_prompt,
@@ -1802,7 +1817,7 @@ with demo:
         no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
         
         submit_button.click(
-            partial(run_fn, n_ret=3),
+            partial(run_fn, n_ret=3, advanced=True),
             inputs=[
                 input_gallery, model_dropdown, layer_slider, l1_num_eig_slider, node_type_dropdown, 
                 positive_prompt, negative_prompt,