add LISA

app.py

@@ -1,5 +1,6 @@
 # Author: Huzheng Yang
 # %%
+import copy
 import os
 USE_HUGGINGFACE_ZEROGPU = os.getenv("USE_HUGGINGFACE_ZEROGPU", "False").lower() in ["true", "1", "yes"]
 DOWNLOAD_ALL_MODELS_DATASETS = os.getenv("DOWNLOAD_ALL_MODELS_DATASETS", "False").lower() in ["true", "1", "yes"]
@@ -232,6 +233,10 @@ def ncut_run(
     recursion_l2_gamma=0.5,
     recursion_l3_gamma=0.5,
     video_output=False,
+    is_lisa=False,
+    lisa_prompt1="",
+    lisa_prompt2="",
+    lisa_prompt3="",
 ):
     logging_str = ""
     if "AlignedThreeModelAttnNodes" == model_name:
@@ -256,6 +261,24 @@ def ncut_run(
     if "AlignedThreeModelAttnNodes" == model_name:
         # dirty patch for the alignedcut paper
         features = run_alignedthreemodelattnnodes(images, model, batch_size=BATCH_SIZE)
+    elif is_lisa == True:
+        # dirty patch for the LISA model
+        features = []
+        with torch.no_grad():
+            model = model.cuda()
+            images = images.cuda()
+            lisa_prompts = [lisa_prompt1, lisa_prompt2, lisa_prompt3]
+            for prompt in lisa_prompts:
+                import bleach
+                prompt = bleach.clean(prompt)
+                prompt = prompt.strip()
+                # print(prompt)
+                # # copy the sting to a new string
+                # copy_s = copy.copy(prompt)
+                feature = model(images, input_str=prompt)[node_type][0]
+                feature = F.normalize(feature, dim=-1)
+                features.append(feature.cpu().float())
+            features = torch.stack(features)
     else:
         features = extract_features(
             images, model, node_type=node_type, layer=layer-1, batch_size=BATCH_SIZE
@@ -340,6 +363,14 @@ def ncut_run(
                 galleries.append(to_pil_images(_rgb, target_size=56))
             return *galleries, logging_str
         
+        if is_lisa == True:
+            # dirty patch for the LISA model
+            galleries = []
+            for i_prompt in range(len(lisa_prompts)):
+                _rgb = rgb[i_prompt]
+                galleries.append(to_pil_images(_rgb, target_size=256))
+            return *galleries, logging_str
+        
         rgb = dont_use_too_much_green(rgb)
         
     
@@ -451,7 +482,8 @@ def load_alignedthreemodel():
     # model = torch.load(save_path)
     return model
 
-promptable_models = ["Diffusion(stabilityai/stable-diffusion-2)", "Diffusion(CompVis/stable-diffusion-v1-4)"]
+promptable_diffusion_models = ["Diffusion(stabilityai/stable-diffusion-2)", "Diffusion(CompVis/stable-diffusion-v1-4)"]
+promptable_segmentation_models = ["LISA(xinlai/LISA-7B-v1)"]
 
 
 def run_fn(
@@ -462,6 +494,10 @@ def run_fn(
     node_type="block",
     positive_prompt="",
     negative_prompt="",
+    is_lisa=False,
+    lisa_prompt1="",
+    lisa_prompt2="",
+    lisa_prompt3="",
     affinity_focal_gamma=0.3,
     num_sample_ncut=10000,
     knn_ncut=10,
@@ -515,10 +551,10 @@ def run_fn(
         model.timestep = layer
         layer = 1
         
-    if model_name in promptable_models:
+    if model_name in promptable_diffusion_models:
         model.positive_prompt = positive_prompt
         model.negative_prompt = negative_prompt
-    
+        
     kwargs = {
         "model_name": model_name,
         "layer": layer,
@@ -543,11 +579,18 @@ def run_fn(
         "recursion_l2_gamma": recursion_l2_gamma,
         "recursion_l3_gamma": recursion_l3_gamma,
         "video_output": video_output,
+        "lisa_prompt1": lisa_prompt1,
+        "lisa_prompt2": lisa_prompt2,
+        "lisa_prompt3": lisa_prompt3,
+        "is_lisa": is_lisa,
     }
     # print(kwargs)
     
     if old_school_ncut:
-        super_duper_long_run(model, images, **kwargs)
+        return super_duper_long_run(model, images, **kwargs)
+    
+    if is_lisa:
+        return super_duper_long_run(model, images, **kwargs)
     
     num_images = len(images)
     if num_images >= 100:
@@ -702,18 +745,26 @@ def make_output_images_section():
     output_gallery = gr.Gallery(value=[], label="NCUT Embedding", show_label=False, elem_id="ncut", columns=[3], rows=[1], object_fit="contain", height="auto")
     return output_gallery
 
-def make_parameters_section():
+def make_parameters_section(is_lisa=False):
     gr.Markdown("### Parameters <a style='color: #0044CC;' href='https://ncut-pytorch.readthedocs.io/en/latest/how_to_get_better_segmentation/' target='_blank'>Help</a>")
     from ncut_pytorch.backbone import list_models, get_demo_model_names
     model_names = list_models()
     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?")
+    if is_lisa:
+        model_dropdown = gr.Dropdown(["LISA(xinlai/LISA-7B-v1)"], label="Backbone", value="LISA(xinlai/LISA-7B-v1)", elem_id="model_name")
+        layer_slider = gr.Slider(1, 6, step=1, label="LISA decoder: Layer index", value=6, elem_id="layer", visible=False)        
+        layer_names = ["dec_0_input", "dec_0_attn", "dec_0_block", "dec_1_input", "dec_1_attn", "dec_1_block"]
+        positive_prompt = gr.Textbox(label="Prompt (Positive)", elem_id="prompt", placeholder="e.g. 'a photo of Gibson Les Pual guitar'", visible=False)
+        negative_prompt = gr.Textbox(label="Prompt (Negative)", elem_id="prompt", placeholder="e.g. 'a photo from egocentric view'", visible=False)
+        node_type_dropdown = gr.Dropdown(layer_names, label="LISA (SAM) decoder: Layer and Node", value="dec_1_block", elem_id="node_type")
+    else:
+        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')
 
     def change_layer_slider(model_name):
@@ -724,6 +775,12 @@ def make_parameters_section():
             return (gr.Slider(1, 49, step=1, label="Diffusion: Timestep (Noise)", value=5, elem_id="layer", visible=True, info="Noise level, 50 is max noise"),
                     gr.Dropdown(SD_KEY_DICT[model_name], label="Diffusion: Layer and Node", value=default_layer, elem_id="node_type", info="U-Net (v1, v2) or DiT (v3)"))
         
+        if model_name == "LISSL(xinlai/LISSL-7B-v1)":
+            layer_names = ["dec_0_input", "dec_0_attn", "dec_0_block", "dec_1_input", "dec_1_attn", "dec_1_block"]
+            default_layer = "dec_1_block"
+            return (gr.Slider(1, 6, step=1, label="LISA decoder: Layer index", value=6, elem_id="layer", visible=False),
+                    gr.Dropdown(layer_names, label="LISA decoder: Layer and Node", value=default_layer, elem_id="node_type"))
+
         layer_dict = LAYER_DICT
         if model_name in layer_dict:
             value = layer_dict[model_name]
@@ -736,7 +793,7 @@ def make_parameters_section():
     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:
+        if model_name in promptable_diffusion_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),
@@ -788,12 +845,15 @@ with demo:
 
         clear_images_button.click(lambda x: ([], []), outputs=[input_gallery, output_gallery])
         
+        false_placeholder = gr.Checkbox(label="False", value=False, elem_id="false_placeholder", visible=False)
+        no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
         
         submit_button.click(
             run_fn,
             inputs=[
                 input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
                 positive_prompt, negative_prompt,
+                false_placeholder, no_prompt, no_prompt, no_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
@@ -848,11 +908,15 @@ with demo:
                 logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
             
         clear_images_button.click(lambda x: ([], []), outputs=[input_gallery, output_gallery])
+        false_placeholder = gr.Checkbox(label="False", value=False, elem_id="false_placeholder", visible=False)
+        no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
+        
         submit_button.click(
             run_fn,
             inputs=[
                 input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
                 positive_prompt, negative_prompt,
+                false_placeholder, no_prompt, no_prompt, no_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, 
@@ -918,11 +982,15 @@ with demo:
         number_placeholder = gr.Number(0, label="Number placeholder", elem_id="number_placeholder")
         number_placeholder.visible = False
         clear_images_button.click(lambda x: ([],), outputs=[input_gallery])
+        false_placeholder = gr.Checkbox(label="False", value=False, elem_id="false_placeholder", visible=False)
+        no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
+        
         submit_button.click(
             run_fn,
             inputs=[
-                input_gallery, model_dropdown, layer_slider, l1_num_eig_slider, node_type_dropdown, 
+                input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
                 positive_prompt, negative_prompt,
+                false_placeholder, no_prompt, no_prompt, no_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,
@@ -957,11 +1025,15 @@ with demo:
         clear_video_button.click(lambda x: (None, None), outputs=[video_input_gallery, video_output_gallery])
         place_holder_false = gr.Checkbox(label="Place holder", value=False, elem_id="place_holder_false")
         place_holder_false.visible = False
+        false_placeholder = gr.Checkbox(label="False", value=False, elem_id="false_placeholder", visible=False)
+        no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
+        
         submit_button.click(
             run_fn,
             inputs=[
-                video_input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
                 positive_prompt, negative_prompt,
+                false_placeholder, no_prompt, no_prompt, no_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,
@@ -979,6 +1051,57 @@ with demo:
             from draft_gradio_app_text import make_demo
         make_demo()
         
+    with gr.Tab('Vision-Language'):
+        gr.Markdown('[LISA]((https://arxiv.org/pdf/2308.00692)) is a vision-language model. Input a text prompt and image, LISA generate segmentation masks.')
+        gr.Markdown('In the mask decoder layers, LISA updates the image features w.r.t. the text prompt')
+        gr.Markdown('This page aims to see how the text prompt affects the image features')
+        gr.Markdown('---')
+        gr.Markdown('<p style="text-align: center;">Color is <b>aligned</b> across 3 prompts. NCUT is computed on the concatenated features from 3 prompts.</p>')
+        with gr.Row():
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown('### Output (Prompt #1)')
+                l1_gallery = gr.Gallery(value=[], label="Prompt #1", show_label=False, elem_id="ncut_p1", columns=[3], rows=[5], object_fit="contain", height="auto")
+                prompt1 = gr.Textbox(label="Input Prompt #1", elem_id="prompt1", value="where is the person, include the clothes, don't include the guitar and chair", lines=3)
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown('### Output (Prompt #2)')
+                l2_gallery = gr.Gallery(value=[], label="Prompt #2", show_label=False, elem_id="ncut_p2", columns=[3], rows=[5], object_fit="contain", height="auto")
+                prompt2 = gr.Textbox(label="Input Prompt #2", elem_id="prompt2", value="where is the Gibson Les Pual guitar", lines=3)
+            with gr.Column(scale=5, min_width=200):
+                gr.Markdown('### Output (Prompt #3)')
+                l3_gallery = gr.Gallery(value=[], label="Prompt #3", show_label=False, elem_id="ncut_p3", columns=[3], rows=[5], object_fit="contain", height="auto")
+                prompt3 = gr.Textbox(label="Input Prompt #3", elem_id="prompt3", value="where is the floor", lines=3)
+                
+        with gr.Row():
+            with gr.Column(scale=5, min_width=200):
+                input_gallery, submit_button, clear_images_button = make_input_images_section()
+                dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_dataset_images_section(advanced=False)
+                clear_images_button.click(lambda x: ([], [], [], []), outputs=[input_gallery, l1_gallery, l2_gallery, l3_gallery])
+                
+            with gr.Column(scale=5, min_width=200):
+                [
+                    model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
+                    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, positive_prompt, negative_prompt
+                ] = make_parameters_section(is_lisa=True)
+                logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
+        
+        galleries = [l1_gallery, l2_gallery, l3_gallery]
+        true_placeholder = gr.Checkbox(label="True placeholder", value=True, elem_id="true_placeholder", visible=False)
+        submit_button.click(
+            run_fn,
+            inputs=[
+                input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
+                positive_prompt, negative_prompt,
+                true_placeholder, prompt1, prompt2, prompt3,
+                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
+            ],
+            outputs=galleries + [logging_text],
+        )
+                
     with gr.Tab('Model Aligned'): 
         gr.Markdown('This page reproduce the results from the paper [AlignedCut](https://arxiv.org/abs/2406.18344)')
         gr.Markdown('---')
@@ -1022,11 +1145,16 @@ with demo:
                 
 
         clear_images_button.click(lambda x: [] * (len(galleries) + 1), outputs=[input_gallery] + galleries)
+        
+        false_placeholder = gr.Checkbox(label="False", value=False, elem_id="false_placeholder", visible=False)
+        no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
+        
         submit_button.click(
             run_fn,
             inputs=[
                 input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
                 positive_prompt, negative_prompt,
+                false_placeholder, no_prompt, no_prompt, no_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
@@ -1049,11 +1177,15 @@ with demo:
                 ] = make_parameters_section()
                 # logging text box
                 logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information")
+                false_placeholder = gr.Checkbox(label="False", value=False, elem_id="false_placeholder", visible=False)
+                no_prompt = gr.Textbox("", label="", elem_id="empty_placeholder", type="text", placeholder="", visible=False)
+                
                 submit_button.click(
                     run_fn,
                     inputs=[
                         input_gallery, model_dropdown, layer_slider, num_eig_slider, node_type_dropdown, 
                         positive_prompt, negative_prompt,
+                        false_placeholder, no_prompt, no_prompt, no_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

@@ -2,7 +2,7 @@ torch
 torchvision 
 opencv-python
 decord
-transformers
+transformers==4.31.0
 datasets
 diffusers
 accelerate
@@ -12,6 +12,7 @@ pillow==9.4.0
 SAM-2 @ git+https://github.com/huzeyann/segment-anything-2.git
 segment-anything @ git+https://github.com/facebookresearch/segment-anything.git@6fdee8f
 mobile-sam @ git+https://github.com/ChaoningZhang/MobileSAM.git@c12dd83
+lisa @ git+https://github.com/huzeyann/LISA.git
 timm
 open-clip-torch==2.20.0
-ncut-pytorch>=1.3.10
+ncut-pytorch>=1.3.13