add prompt diffusion

app.py

@@ -351,19 +351,19 @@ def ncut_run(
         return to_pil_images(rgb), logging_str
 
 def _ncut_run(*args, **kwargs):
-    # try:
-    #     ret = ncut_run(*args, **kwargs)
-    #     if torch.cuda.is_available():
-    #         torch.cuda.empty_cache()
-    #     return ret
-    # except Exception as e:
-    #     gr.Error(str(e))
-    #     if torch.cuda.is_available():
-    #         torch.cuda.empty_cache()
-    #     return [], "Error: " + str(e)
-
-    ret = ncut_run(*args, **kwargs)
-    return ret
+    try:
+        ret = ncut_run(*args, **kwargs)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return ret
+    except Exception as e:
+        gr.Error(str(e))
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        return [], "Error: " + str(e)
+
+    # ret = ncut_run(*args, **kwargs)
+    # return ret
 
 if USE_HUGGINGFACE_ZEROGPU:
     @spaces.GPU(duration=20)
@@ -445,12 +445,17 @@ def load_alignedthreemodel():
     # model = torch.load(save_path)
     return model
 
+promptable_models = ["Diffusion(stabilityai/stable-diffusion-2)", "Diffusion(CompVis/stable-diffusion-v1-4)"]
+
+
 def run_fn(
     images,
     model_name="SAM(sam_vit_b)",
     layer=-1,
     num_eig=100,
     node_type="block",
+    positive_prompt="",
+    negative_prompt="",
     affinity_focal_gamma=0.3,
     num_sample_ncut=10000,
     knn_ncut=10,
@@ -503,6 +508,10 @@ def run_fn(
     if "stable" in model_name.lower() and "diffusion" in model_name.lower():
         model.timestep = layer
         layer = 1
+        
+    if model_name in promptable_models:
+        model.positive_prompt = positive_prompt
+        model.negative_prompt = negative_prompt
     
     kwargs = {
         "model_name": model_name,
@@ -694,6 +703,10 @@ def make_parameters_section():
     model_names = sorted(model_names)
     model_dropdown = gr.Dropdown(model_names, label="Backbone", value="DiNO(dino_vitb8_448)", elem_id="model_name")
     layer_slider = gr.Slider(1, 12, step=1, label="Backbone: Layer index", value=10, elem_id="layer")
+    positive_prompt = gr.Textbox(label="Prompt (Positive)", elem_id="prompt", placeholder="e.g. 'a photo of Gibson Les Pual guitar'")
+    positive_prompt.visible = False
+    negative_prompt = gr.Textbox(label="Prompt (Negative)", elem_id="prompt", placeholder="e.g. 'a photo from egocentric view'")
+    negative_prompt.visible = False
     node_type_dropdown = gr.Dropdown(["attn: attention output", "mlp: mlp output", "block: sum of residual"], label="Backbone: Layer type", value="block: sum of residual", elem_id="node_type", info="which feature to take from each layer?")
     num_eig_slider = gr.Slider(1, 1000, step=1, label="NCUT: Number of eigenvectors", value=100, elem_id="num_eig", info='increase for more clusters')
 
@@ -716,6 +729,14 @@ def make_parameters_section():
                     gr.Slider(1, value, step=1, label="Backbone: Layer index", value=value, elem_id="layer", visible=True))
     model_dropdown.change(fn=change_layer_slider, inputs=model_dropdown, outputs=[layer_slider, node_type_dropdown])
     
+    def change_prompt_text(model_name):
+        if model_name in promptable_models:
+            return (gr.Textbox(label="Prompt (Positive)", elem_id="prompt", placeholder="e.g. 'a photo of Gibson Les Pual guitar'", visible=True),
+                    gr.Textbox(label="Prompt (Negative)", elem_id="prompt", placeholder="e.g. 'a photo from egocentric view'", visible=True))
+        return (gr.Textbox(label="Prompt (Positive)", elem_id="prompt", placeholder="e.g. 'a photo of Gibson Les Pual guitar'", visible=False),
+                gr.Textbox(label="Prompt (Negative)", elem_id="prompt", placeholder="e.g. 'a photo from egocentric view'", visible=False))
+    model_dropdown.change(fn=change_prompt_text, inputs=model_dropdown, outputs=[positive_prompt, negative_prompt])
+    
     with gr.Accordion("➡️ Click to expand: more parameters", open=False):
         gr.Markdown("<a href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Docs: How to Get Better Segmentation</a>")
         affinity_focal_gamma_slider = gr.Slider(0.01, 1, step=0.01, label="NCUT: Affinity focal gamma", value=0.5, elem_id="affinity_focal_gamma", info="decrease for shaper segmentation")
@@ -732,7 +753,7 @@ def make_parameters_section():
             affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
             embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
             perplexity_slider, n_neighbors_slider, min_dist_slider, 
-            sampling_method_dropdown]
+            sampling_method_dropdown, positive_prompt, negative_prompt]
 
 demo = gr.Blocks(
     theme=gr.themes.Base(spacing_size='md', text_size='lg', primary_hue='blue', neutral_hue='slate', secondary_hue='pink'),
@@ -754,7 +775,7 @@ with demo:
                     affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                     embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
-                    sampling_method_dropdown
+                    sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section()
                 # logging text box
                 logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
@@ -766,6 +787,7 @@ with demo:
             run_fn,
             inputs=[
                 input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                positive_prompt, negative_prompt,
                 affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                 embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                 perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown
@@ -809,7 +831,7 @@ with demo:
                     affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                     embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
-                    sampling_method_dropdown
+                    sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section()
                 old_school_ncut_checkbox = gr.Checkbox(label="Old school NCut", value=True, elem_id="old_school_ncut")
                 invisible_list = [old_school_ncut_checkbox, num_sample_ncut_slider, knn_ncut_slider,
@@ -824,6 +846,7 @@ with demo:
             run_fn,
             inputs=[
                 input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                positive_prompt, negative_prompt,
                 affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                 embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                 perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown, 
@@ -870,7 +893,7 @@ with demo:
                     affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                     embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
-                    sampling_method_dropdown
+                    sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section()
                 num_eig_slider.visible = False
                 affinity_focal_gamma_slider.visible = False
@@ -893,6 +916,7 @@ with demo:
             run_fn,
             inputs=[
                 input_gallery, model_dropdown, layer_slider, l1_num_eig_slider, node_type_dropdown, 
+                positive_prompt, negative_prompt,
                 affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                 embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                 perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown,
@@ -916,7 +940,7 @@ with demo:
                     affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                     embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
-                    sampling_method_dropdown
+                    sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section()
                 num_sample_tsne_slider.value = 1000
                 perplexity_slider.value = 500
@@ -931,6 +955,7 @@ with demo:
             run_fn,
             inputs=[
                 video_input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                positive_prompt, negative_prompt,
                 affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                 embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                 perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown,
@@ -971,7 +996,7 @@ with demo:
                     affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                     embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
-                    sampling_method_dropdown
+                    sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section()
                 model_dropdown.value = "AlignedThreeModelAttnNodes"
                 model_dropdown.visible = False
@@ -995,6 +1020,7 @@ with demo:
             run_fn,
             inputs=[
                 input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                positive_prompt, negative_prompt,
                 affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                 embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                 perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown
@@ -1013,7 +1039,7 @@ with demo:
                     affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                     embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                     perplexity_slider, n_neighbors_slider, min_dist_slider, 
-                    sampling_method_dropdown
+                    sampling_method_dropdown, positive_prompt, negative_prompt
                 ] = make_parameters_section()
                 # logging text box
                 logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
@@ -1021,6 +1047,7 @@ with demo:
                     run_fn,
                     inputs=[
                         input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                        positive_prompt, negative_prompt,
                         affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider, 
                         embedding_method_dropdown, num_sample_tsne_slider, knn_tsne_slider, 
                         perplexity_slider, n_neighbors_slider, min_dist_slider, sampling_method_dropdown

requirements.txt

@@ -14,4 +14,4 @@ segment-anything @ git+https://github.com/facebookresearch/segment-anything.git@
 mobile-sam @ git+https://github.com/ChaoningZhang/MobileSAM.git@c12dd83
 timm
 open-clip-torch==2.20.0
-ncut-pytorch>=1.3.8
+ncut-pytorch>=1.3.10