update nothing important

app.py

@@ -2294,10 +2294,11 @@ with demo:
                 
                 run_hierarchical_button = gr.Button("🔴 RUN tree", elem_id="run_hierarchical", variant='primary')
                 with gr.Accordion("Hierarchical Structure Parameters:", open=True):
-                    num_sample_fps_slider = gr.Slider(1, 5000, step=1, label="FPS: num_sample", value=1000, elem_id="num_sample_fps")
-                    tsne_perplexity_slider = gr.Slider(1, 1000, step=1, label="t-SNE: perplexity", value=500, elem_id="perplexity_tsne")
+                    num_sample_fps_slider = gr.Slider(1, 10000, step=1, label="FPS: num_sample", value=1000, elem_id="num_sample_fps")
+                    tsne_perplexity_slider = gr.Slider(1, 500, step=1, label="t-SNE: perplexity", value=500, elem_id="perplexity_tsne")
+                    tsne_knn_slider = gr.Slider(1, 10, step=1, label="t-SNE: KNN", value=1, elem_id="knn_tsne")
                     fps_hc_seed_slider = gr.Slider(0, 1000, step=1, label="Seed", value=0, elem_id="fps_hc_seed")
-                    tree_method_radio = gr.Radio(["eigvecs", "tsne"], label="Tree Method (input type)", value="tsne", elem_id="tree_method", 
+                    tree_method_radio = gr.Radio(["eigvecs", "tsne"], label="Tree Method (input type)", value="tsne" if USE_HUGGINGFACE_ZEROGPU else "eigvecs", elem_id="tree_method", 
                                                 info="What's the input to build tree? `eigvecs` is loss-less, `tsne` is lossy; tsne make tree looks 'better', eigvecs is more accurate.")
                 tsne_plot = gr.Image(label="spectral-tSNE tree", elem_id="tsne_plot", interactive=False, format='png')
                 
@@ -2316,6 +2317,7 @@ with demo:
                     lengthes = np.linalg.norm(tsne_embed[edges[:, 0]] - tsne_embed[edges[:, 1]], axis=1)
                     max_length = lengthes[k:].max()
                     diag_length = np.linalg.norm(tsne_embed.max(axis=0) - tsne_embed.min(axis=0))
+                    
                     # draw the edges
                     # for i_edge in range(k, len(edges)):
                     #     edge = edges[i_edge]
@@ -2326,10 +2328,11 @@ with demo:
                     #         _do = 0.0
                     #     alpha = 0.7 * (1 - _do) + 0.0
                     #     ax.plot(tsne_embed[edge, 0], tsne_embed[edge, 1], 'k-', lw=1, alpha=alpha)
+                    
                     max_level = levels.max()
-                    line_widths = levels / max_level * 3
+                    line_widths = levels / max_level * 2
                     # flip line width
-                    line_widths = 3 - line_widths
+                    line_widths = 2 - line_widths
                     for i_edge in range(k, len(edges)):
                         edge = edges[i_edge]
                         level = levels[i_edge]
@@ -2450,7 +2453,7 @@ with demo:
                     return pil_image
                     
                 
-                def run_fps_tsne_hierarchical(image_gallery, eigvecs, num_sample_fps, perplexity_tsne, tsne3d_rgb, seed=0, tree_method='eigvecs', max_display_dots=300):
+                def run_fps_tsne_hierarchical(image_gallery, eigvecs, num_sample_fps, perplexity_tsne, knn_tsne, tsne3d_rgb, seed=0, tree_method='eigvecs', max_display_dots=300):
                     if len(eigvecs) == 0:
                         gr.Warning("Please run NCUT first.")
                         return
@@ -2476,12 +2479,19 @@ with demo:
                     fps_tsne3d_rgb = tsne3d_rgb[fps_idx]
                     
                     np.random.seed(seed)
-                    tsne_embed = TSNE(
-                        n_components=2,
-                        perplexity=perplexity_tsne,
-                        metric='cosine',
-                        random_state=seed,
-                    ).fit_transform(fps_eigvecs)
+                    # tsne_embed = TSNE(
+                    #     n_components=2,
+                    #     perplexity=perplexity_tsne,
+                    #     metric='cosine',
+                    #     random_state=seed,
+                    # ).fit_transform(fps_eigvecs)
+                    from ncut_pytorch import rgb_from_tsne_2d
+                    device = None
+                    # device = 'cpu'
+                    # if not USE_HUGGINGFACE_ZEROGPU and torch.cuda.is_available():
+                    #     device = 'cuda'
+                    tsne_embed, _ = rgb_from_tsne_2d(torch.tensor(fps_eigvecs), num_sample=1000, perplexity=perplexity_tsne,
+                                                        seed=seed, metric='cosine', knn=knn_tsne, device=device)
                     # normalize = [-1, 1]
                     tsne_embed[:, 0] = (tsne_embed[:, 0] - tsne_embed[:, 0].min()) / (tsne_embed[:, 0].max() - tsne_embed[:, 0].min()) * 2 - 1
                     tsne_embed[:, 1] = (tsne_embed[:, 1] - tsne_embed[:, 1].min()) / (tsne_embed[:, 1].max() - tsne_embed[:, 1].min()) * 2 - 1
@@ -2507,7 +2517,7 @@ with demo:
         
         run_hierarchical_button.click(
             run_fps_tsne_hierarchical,
-            inputs=[input_gallery, eigvecs, num_sample_fps_slider, tsne_perplexity_slider, tsne3d_rgb, fps_hc_seed_slider, tree_method_radio],
+            inputs=[input_gallery, eigvecs, num_sample_fps_slider, tsne_perplexity_slider, tsne_knn_slider, tsne3d_rgb, fps_hc_seed_slider, tree_method_radio],
             outputs=[tsne_2d_points, edges, levels, fps_eigvecs, fps_tsne_rgb, fps_indices, tsne_plot, tsne_image_plot],
         )
         with gr.Row():
@@ -2571,7 +2581,7 @@ with demo:
             with gr.Column(scale=5, min_width=200): 
                 prompt_radio = gr.Radio(["Tree [+Image]", "Image (NCUT)", "Image (Orig)"], label="Where to click on?", value="Tree [+Image]", elem_id="prompt_radio", show_label=True)
                 granularity_slider = gr.Slider(1, 1000, step=1, label="Cluster Granularity (k)", value=100, elem_id="granularity")
-                num_sample_fps_slider.change(fn=lambda x: gr.update(maximum=x, interactive=True), inputs=[num_sample_fps_slider], outputs=[granularity_slider])
+                num_sample_fps_slider.change(fn=lambda x: gr.update(maximum=x-1, interactive=True), inputs=[num_sample_fps_slider], outputs=[granularity_slider])
                 def updaste_tsne_plot_change_granularity(granularity, tsne_embed, edges, levels, fps_tsne_rgb, tsne_prompt_image):
                     # Plot the t-SNE points
                     pil_image = plot_tsne_tree(tsne_embed, edges, levels, fps_tsne_rgb, granularity)