add mlp align

app.py

@@ -7,9 +7,9 @@ from functools import partial
 from io import BytesIO
 import json
 import os
-from pprint import pprint
 import uuid
 import zipfile
+import multiprocessing as mp
 
 from einops import rearrange
 from matplotlib import pyplot as plt
@@ -42,7 +42,6 @@ from ncut_pytorch.backbone import MODEL_DICT, LAYER_DICT, RES_DICT
 from ncut_pytorch import NCUT
 from ncut_pytorch import eigenvector_to_rgb, rotate_rgb_cube
 
-RUN_COUNT = 0
 
 DATASETS = {
     'Common': [
@@ -314,8 +313,7 @@ def blend_image_with_heatmap(image, heatmap, opacity1=0.5, opacity2=0.5):
     blended = (1 - opacity1) * image + opacity2 * heatmap
     return blended.astype(np.uint8)
 
-# preload the model
-load_model("CLIP(ViT-B-16/openai)")
+
 def segment_fg_bg(images):
 
     images = F.interpolate(images, (224, 224), mode="bilinear")
@@ -388,6 +386,8 @@ def segment_fg_bg(images):
 
 
 def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=False, clusters=50, eig_idx=None, title='cluster'):
+    if clusters == 0:
+        return [], []
     progress = gr.Progress()
     progress(progess_start, desc="Finding Clusters by FPS")
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -462,7 +462,7 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
         # fps_heatmaps[idx.item()] = heatmap.cpu()
         fps_heatmaps[idx.item()] = heatmap[mask_sort_idx[:6]].cpu()
         top3_image_idx[idx.item()] = mask_sort_idx[:3]
-        top10_image_idx[idx.item()] = mask_sort_idx[:5]
+        top10_image_idx[idx.item()] = mask_sort_idx[:6]
     # do the sorting
     _sort_idx = torch.tensor(sort_values).argsort(descending=True)
     fps_idx = fps_idx[_sort_idx]
@@ -480,51 +480,46 @@ def make_cluster_plot(eigvecs, images, h=64, w=64, progess_start=0.6, advanced=F
         fps_idx = fps_idx[torch.randperm(fps_idx.shape[0])]
     
         
-    fig_images = []
-    i_cluster = 0
-    num_plots = clusters // 5
-    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=f"Plotting {title}")
-        if not advanced:
-            fig, axs = plt.subplots(3, 5, figsize=(15, 9))
-        if advanced:
-            fig, axs = plt.subplots(5, 5, figsize=(15, 15))
+    def plot_cluster_images(fps_idx_chunk, chunk_idx):
+        fig, axs = plt.subplots(3, 5, figsize=(15, 9)) if not advanced else plt.subplots(6, 5, figsize=(15, 18))
         for ax in axs.flatten():
             ax.axis("off")
-        for j, idx in enumerate(fps_idx[i_fig*5:i_fig*5+5]):
+        for j, idx in enumerate(fps_idx_chunk):
             heatmap = fps_heatmaps[idx.item()]
-            # mask = (heatmap > 0.1).float()
-            # sorted_image_idxs = torch.argsort(mask.mean((1, 2)), descending=True)
             size = (images.shape[1], images.shape[2])
             heatmap = apply_reds_colormap(heatmap, size)
-            # for i, image_idx in enumerate(sorted_image_idxs[:3]):
             image_idxs = top3_image_idx[idx.item()] if not advanced else top10_image_idx[idx.item()]
             for i, image_idx in enumerate(image_idxs):
-                # _heatmap = blend_image_with_heatmap(images[image_idx], heatmap[image_idx])
                 _heatmap = blend_image_with_heatmap(images[image_idx], heatmap[i])
                 axs[i, j].imshow(_heatmap)
                 if i == 0:
-                    axs[i, j].set_title(f"{title} {i_cluster+1}", fontsize=24)
-                    i_cluster += 1
+                    axs[i, j].set_title(f"{title} {chunk_idx * 5 + j + 1}", fontsize=24)
         plt.tight_layout(h_pad=0.5, w_pad=0.3)
-
-        filename = uuid.uuid4()
+        filename = f"{datetime.now():%Y%m%d%H%M%S%f}_{uuid.uuid4().hex}"
         tmp_path = f"/tmp/{filename}.png"
         plt.savefig(tmp_path, bbox_inches='tight', dpi=72)
-        
-        img = Image.open(tmp_path)
-        img = img.convert("RGB")
-        img = copy.deepcopy(img)
-        
+        img = Image.open(tmp_path).convert("RGB")
         os.remove(tmp_path)
-
-        fig_images.append(img)
         plt.close()
+        return img
+
+    fig_images = []
+    num_plots = clusters // 5
+    plot_step_float = (1.0 - progess_start) / num_plots
+    fps_idx_chunks = [fps_idx[i*5:(i+1)*5] for i in range(num_plots)]
+
+    # with mp.Pool(processes=mp.cpu_count()) as pool:
+    #     results = [pool.apply_async(plot_cluster_images, args=(chunk, i)) for i, chunk in enumerate(fps_idx_chunks)]
+    #     for i, result in enumerate(results):
+    #         progress(progess_start + i * plot_step_float, desc=f"Plotted {title}")
+    #         fig_images.append(result.get())
+    for i, chunk in enumerate(fps_idx_chunks):
+        progress(progess_start + i * plot_step_float, desc=f"Plotted {title}")
+        fig_images.append(plot_cluster_images(chunk, i))
 
     return fig_images, ret_magnitude
 
-def make_cluster_plot_advanced(eigvecs, images, h=64, w=64, num_fg=100, num_bg=10, num_other=0, small_title=True):
+def make_cluster_plot_advanced(eigvecs, images, h=64, w=64):
     heatmap_fg, heatmap_bg = segment_fg_bg(images.clone())
     heatmap_bg = rearrange(heatmap_bg, 'b h w c -> b c h w')
     heatmap_fg = rearrange(heatmap_fg, 'b h w c -> b c h w')
@@ -545,15 +540,9 @@ def make_cluster_plot_advanced(eigvecs, images, h=64, w=64, num_fg=100, num_bg=1
     bg_idx = torch.arange(heatmap_bg.shape[0])[bg_mask]
     other_idx = torch.arange(heatmap_fg.shape[0])[other_mask]
     
-    fg_images = []
-    if num_fg > 0:
-        fg_images, _ = make_cluster_plot(eigvecs, images, h=h, w=w, advanced=True, clusters=num_fg, eig_idx=fg_idx, title="fg" if small_title else "cluster")
-    bg_images = []
-    if num_bg > 0:
-        bg_images, _ = make_cluster_plot(eigvecs, images, h=h, w=w, advanced=True, clusters=num_bg, eig_idx=bg_idx, title="bg" if small_title else "cluster")
-    other_images = []
-    if num_other > 0:
-        other_images, _ = make_cluster_plot(eigvecs, images, h=h, w=w, advanced=True, clusters=num_other, eig_idx=other_idx, title="other" if small_title else "cluster")
+    fg_images, _ = make_cluster_plot(eigvecs, images, h=h, w=w, advanced=True, clusters=100, eig_idx=fg_idx, title="fg")
+    bg_images, _ = make_cluster_plot(eigvecs, images, h=h, w=w, advanced=True, clusters=20, eig_idx=bg_idx, title="bg")
+    other_images, _ = make_cluster_plot(eigvecs, images, h=h, w=w, advanced=True, clusters=0, eig_idx=other_idx, title="other")
     
     cluster_images = fg_images + bg_images + other_images
     
@@ -842,7 +831,7 @@ def ncut_run(
         if advanced:
             cluster_images, eig_magnitude = make_cluster_plot_advanced(eigvecs, _images, h=h, w=w)
         else:
-            cluster_images, eig_magnitude = make_cluster_plot_advanced(eigvecs, _images, h=h, w=w, num_fg=20, num_bg=0, num_other=0, small_title=False)
+            cluster_images, eig_magnitude = make_cluster_plot(eigvecs, _images, h=h, w=w, progess_start=progress_start, advanced=False)
         logging_str += f"plot time: {time.time() - start:.2f}s\n"
     
     norm_images = None
@@ -903,10 +892,6 @@ if USE_HUGGINGFACE_ZEROGPU:
     def longer_run(*args, **kwargs):
         return _ncut_run(*args, **kwargs)
 
-    @spaces.GPU(duration=90)
-    def quite_long_run(*args, **kwargs):
-        return _ncut_run(*args, **kwargs)
-    
     @spaces.GPU(duration=120)
     def super_duper_long_run(*args, **kwargs):
         return _ncut_run(*args, **kwargs)
@@ -924,9 +909,6 @@ if not USE_HUGGINGFACE_ZEROGPU:
     def longer_run(*args, **kwargs):
         return _ncut_run(*args, **kwargs)
 
-    def quite_long_run(*args, **kwargs):
-        return _ncut_run(*args, **kwargs)
-
     def super_duper_long_run(*args, **kwargs):
         return _ncut_run(*args, **kwargs)
     
@@ -1241,10 +1223,7 @@ def run_fn(
         "alignedcut_eig_norm_plot": alignedcut_eig_norm_plot,
         "advanced": advanced,
     }
-    global RUN_COUNT
-    RUN_COUNT += 1
-    print(f"Run Count: {RUN_COUNT}")
-    pprint(kwargs)
+    # print(kwargs)
     
     try:
         # try to aquiare GPU, can fail if the user is out of GPU quota
@@ -1256,15 +1235,13 @@ def run_fn(
             return super_duper_long_run(model, images, **kwargs)
         
         num_images = len(images)
-        if num_images > 100:
+        if num_images >= 100:
             return super_duper_long_run(model, images, **kwargs)
         if 'diffusion' in model_name.lower():
             return super_duper_long_run(model, images, **kwargs)
         if recursion:
             return longer_run(model, images, **kwargs)
-        if num_images > 50:
-            return quite_long_run(model, images, **kwargs)
-        if num_images > 30:
+        if num_images >= 50:
             return longer_run(model, images, **kwargs)
         if old_school_ncut:
             return longer_run(model, images, **kwargs)
@@ -1284,7 +1261,7 @@ def run_fn(
     except gr.Error as e:
         # I assume this is a GPU quota error
         
-        info1 = 'Running out of HuggingFace GPU Quota?</br> Please try <a style="white-space: nowrap;text-underline-offset: 2px;color: var(--body-text-color)" href="https://ncut-pytorch.readthedocs.io/en/latest/demo/">Demo hosted at UPenn</a>, passcode is: `158.130.50.41`</br>'
+        info1 = 'Running out of HuggingFace GPU Quota?</br> Please try <a style="white-space: nowrap;text-underline-offset: 2px;color: var(--body-text-color)" href="https://ncut-pytorch.readthedocs.io/en/latest/demo/">Demo hosted at UPenn</a></br>'
         info2 = 'Or try use the Python package that powers this app: <a style="white-space: nowrap;text-underline-offset: 2px;color: var(--body-text-color)" href="https://ncut-pytorch.readthedocs.io/en/latest/">ncut-pytorch</a>'
         info = info1 + info2
         
@@ -1292,6 +1269,165 @@ def run_fn(
         raise gr.Error(message, duration=0)
 
 
+import torch
+from torch import nn
+from torch.utils.data import Dataset, DataLoader
+import pytorch_lightning as pl
+
+# Custom Dataset
+class TwoTensorDataset(Dataset):
+    def __init__(self, A, B):
+        self.A = A
+        self.B = B
+
+    def __len__(self):
+        return len(self.A)
+
+    def __getitem__(self, idx):
+        return self.A[idx], self.B[idx]
+
+# MLP model
+class MLP(pl.LightningModule):
+    def __init__(self, num_layer=3, width=512, lr=3e-4, fitting_steps=10000, seg_loss_lambda=1.0):
+        super().__init__()
+        layers = [nn.Linear(3, width), nn.GELU()]
+        for _ in range(num_layer - 1):
+            layers.append(nn.Linear(width, width))
+            layers.append(nn.GELU())
+        layers.append(nn.Linear(width, 3))
+        self.layers = nn.Sequential(*layers)
+        self.mse_loss = nn.MSELoss()
+        self.lr = lr
+        self.fitting_steps = fitting_steps
+        self.seg_loss_lambda = seg_loss_lambda
+        self.progress = gr.Progress()
+
+    def forward(self, x):
+        return self.layers(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self.forward(x)
+        loss = self.mse_loss(y_hat, y)
+        # loss = torch.nn.functional.mse_loss(torch.log(y_hat), torch.log(y))
+        self.log("train_loss", loss)
+        
+        # add segmentation constraint
+        bsz = x.shape[0]
+        sample_size = 1000
+        if bsz > sample_size:
+            idx = torch.randperm(bsz)[:sample_size]
+            x = x[idx]
+            y_hat = y_hat[idx]
+        
+        old_dist = torch.pdist(x, p=2)
+        new_dist = torch.pdist(y_hat, p=2)
+        # seg_loss = torch.log((old_dist - new_dist)).pow(2).mean()
+        seg_loss = self.mse_loss(old_dist, new_dist)
+        self.log("seg_loss", seg_loss)
+        loss += seg_loss * self.seg_loss_lambda
+        
+        step = self.global_step
+        if step % 100 == 0:
+            self.progress(step / self.fitting_steps, desc="Fitting MLP")
+            
+        return loss
+    
+    def predict_step(self, batch, batch_idx, dataloader_idx=None):
+        x = batch[0]
+        y_hat = self.forward(x)
+        return y_hat
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
+        return optimizer
+
+
+def fit_trans(rgb1, rgb2, num_layer=3, width=512, batch_size=256, lr=3e-4, fitting_steps=10000, fps_sample=4096, seg_loss_lambda=1.0):
+    A = rgb1.clone()
+    B = rgb2.clone()
+    
+    # FPS sample on the data
+    from ncut_pytorch.ncut_pytorch import farthest_point_sampling
+    A_idx = farthest_point_sampling(A, fps_sample)
+    B_idx = farthest_point_sampling(B, fps_sample)
+    A_B_idx = np.concatenate([A_idx, B_idx])
+    A = A[A_B_idx]
+    B = B[A_B_idx]
+
+    from torch.utils.data import DataLoader, TensorDataset
+    # Dataset and DataLoader
+    dataset = TwoTensorDataset(A, B)
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+    # Initialize model and trainer
+    mlp = MLP(num_layer=num_layer, width=width, lr=lr, fitting_steps=fitting_steps, seg_loss_lambda=seg_loss_lambda)
+    trainer = pl.Trainer(
+        auto_scale_batch_size='power', 
+        max_epochs=100000, 
+        gpus=1, 
+        max_steps=fitting_steps, 
+        enable_checkpointing=False,
+        enable_progress_bar=False,
+        gradient_clip_val=1.0
+    )
+
+    # Create a DataLoader for tensor A
+    batch_size = 256  # Define your batch size
+    data_loader = DataLoader(TensorDataset(rgb1), batch_size=batch_size, shuffle=False)
+    
+
+    # Train the model
+    trainer.fit(mlp, dataloader)
+
+    
+    results = trainer.predict(mlp, data_loader)
+    A_transformed = torch.cat(results, dim=0)
+    
+    return A_transformed
+
+if USE_HUGGINGFACE_ZEROGPU:
+    @spaces.GPU(duration=60)
+    def _run_mlp_fit(*args, **kwargs):
+        return fit_trans(*args, **kwargs)
+else:
+    def _run_mlp_fit(*args, **kwargs):
+        return fit_trans(*args, **kwargs)
+
+
+def run_mlp_fit(input_gallery, target_gallery, num_layer=3, width=512, batch_size=256, lr=3e-4, fitting_steps=10000, fps_sample=4096, seg_loss_lambda=1.0):
+    # print("Fitting MLP")
+    # print("Target Gallery Length:", len(target_gallery))
+    # print("Input Gallery Length:", len(input_gallery))
+    if target_gallery is None or len(target_gallery) == 0:
+        raise gr.Error("No target images selected. Please use the Mark button to select the target images.")
+    if input_gallery is None or len(input_gallery) == 0:
+        raise gr.Error("No input images selected.")
+    def gallery_to_rgb(gallery):
+        images = [tup[0] for tup in gallery]
+        rgb = []
+        for image in images:
+            if isinstance(image, str):
+                image = Image.open(image)
+            image = image.convert('RGB')
+            image = torch.tensor(np.array(image)).float()
+            image = image / 255
+            rgb.append(image)
+        rgb = torch.stack(rgb)
+        shape = rgb.shape
+        rgb = rgb.reshape(-1, 3)
+        return rgb, shape
+    
+    target_rgb, target_shape = gallery_to_rgb(target_gallery)
+    input_rgb, input_shape = gallery_to_rgb(input_gallery)
+    
+    input_transformed = _run_mlp_fit(input_rgb, target_rgb, num_layer=num_layer, width=width, batch_size=batch_size, lr=lr, 
+                                    fitting_steps=fitting_steps, fps_sample=fps_sample, seg_loss_lambda=seg_loss_lambda)
+    input_transformed = input_transformed.reshape(*input_shape)
+    pil_images = to_pil_images(input_transformed, resize=False)
+    return pil_images
+
+    
 def make_input_video_section():
     # gr.Markdown('### Input Video')
     input_gallery = gr.Video(value=None, label="Select video", elem_id="video-input", height="auto", show_share_button=False, interactive=True)
@@ -1426,7 +1562,7 @@ def make_input_images_section(rows=1, cols=3, height="auto", advanced=False, is_
         def make_example(name, images, dataset_name):
             with gr.Row():
                 button = gr.Button("Load\n"+name, elem_id=f"example-{name}", elem_classes="small-button", variant='secondary', size="sm", scale=1, min_width=60)
-                gallery = gr.Gallery(value=images, label=name, show_label=True, columns=[3], rows=[1], interactive=False, height=80, scale=8, object_fit="cover", min_width=140)
+                gallery = gr.Gallery(value=images, label=name, show_label=True, columns=[3], rows=[1], interactive=False, height=80, scale=8, object_fit="cover", min_width=140, allow_preview=False)
             button.click(fn=lambda: gr.update(value=load_dataset_images(True, dataset_name, 100, is_random=True, seed=42)), outputs=[input_gallery])
             return gallery, button
         example_items = [
@@ -1641,7 +1777,7 @@ def flip_rgb_gallery(images, axis=0):
     images = to_pil_images(images, resize=False)
     return images
 
-def add_output_images_buttons(output_gallery):
+def add_rotate_flip_buttons(output_gallery):
     with gr.Row():
         rotate_button = gr.Button("🔄 Rotate", elem_id="rotate_button", variant='secondary')
         rotate_button.click(sequence_rotate_rgb_gallery, inputs=[output_gallery], outputs=[output_gallery])
@@ -1760,7 +1896,7 @@ def add_download_button(gallery, filename_prefix="output"):
 def make_output_images_section():
     gr.Markdown('### Output Images')
     output_gallery = gr.Gallery(format='png', value=[], label="NCUT Embedding", show_label=True, elem_id="ncut", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, interactive=False)
-    add_output_images_buttons(output_gallery)
+    add_rotate_flip_buttons(output_gallery)
     return output_gallery
 
 def make_parameters_section(is_lisa=False, model_ratio=True):
@@ -1880,7 +2016,7 @@ with demo:
                 
             with gr.Column(scale=5, min_width=200):
                 output_gallery = make_output_images_section()
-                cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[2], object_fit="contain", height="auto", show_share_button=True, preview=True, interactive=True)
+                cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[2], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
                 [
                     model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
                     affinity_focal_gamma_slider, num_sample_ncut_slider, ncut_knn_slider, ncut_indirect_connection, ncut_make_orthogonal, 
@@ -2024,15 +2160,15 @@ with demo:
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #1)')
                 l1_gallery = gr.Gallery(format='png', value=[], label="Recursion #1", show_label=True, elem_id="ncut_l1", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
-                add_output_images_buttons(l1_gallery)
+                add_rotate_flip_buttons(l1_gallery)
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #2)')
                 l2_gallery = gr.Gallery(format='png', value=[], label="Recursion #2", show_label=True, elem_id="ncut_l2", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
-                add_output_images_buttons(l2_gallery)
+                add_rotate_flip_buttons(l2_gallery)
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #3)')
                 l3_gallery = gr.Gallery(format='png', value=[], label="Recursion #3", show_label=True, elem_id="ncut_l3", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
-                add_output_images_buttons(l3_gallery)
+                add_rotate_flip_buttons(l3_gallery)
         with gr.Row():
                 
             with gr.Column(scale=5, min_width=200):
@@ -2089,7 +2225,7 @@ with demo:
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #1)')
                 l1_gallery = gr.Gallery(format='png', value=[], label="Recursion #1", show_label=True, elem_id="ncut_l1", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
-                add_output_images_buttons(l1_gallery)
+                add_rotate_flip_buttons(l1_gallery)
                 add_download_button(l1_gallery, "ncut_embed_recur1")
                 l1_norm_gallery = gr.Gallery(value=[], label="Recursion #1 Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
                 add_download_button(l1_norm_gallery, "eig_norm_recur1")
@@ -2098,7 +2234,7 @@ with demo:
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #2)')
                 l2_gallery = gr.Gallery(format='png', value=[], label="Recursion #2", show_label=True, elem_id="ncut_l2", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
-                add_output_images_buttons(l2_gallery)
+                add_rotate_flip_buttons(l2_gallery)
                 add_download_button(l2_gallery, "ncut_embed_recur2")
                 l2_norm_gallery = gr.Gallery(value=[], label="Recursion #2 Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
                 add_download_button(l2_norm_gallery, "eig_norm_recur2")
@@ -2107,7 +2243,7 @@ with demo:
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #3)')
                 l3_gallery = gr.Gallery(format='png', value=[], label="Recursion #3", show_label=True, elem_id="ncut_l3", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
-                add_output_images_buttons(l3_gallery)
+                add_rotate_flip_buttons(l3_gallery)
                 add_download_button(l3_gallery, "ncut_embed_recur3")
                 l3_norm_gallery = gr.Gallery(value=[], label="Recursion #3 Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
                 add_download_button(l3_norm_gallery, "eig_norm_recur3")
@@ -2335,15 +2471,15 @@ with demo:
         #         add_output_images_buttons(l3_gallery)    
         gr.Markdown('### Output (Recursion #1)')
         l1_gallery = gr.Gallery(format='png', value=[], label="Recursion #1", show_label=True, elem_id="ncut_l1", columns=[100], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False, preview=True)
-        add_output_images_buttons(l1_gallery)
+        add_rotate_flip_buttons(l1_gallery)
         add_download_button(l1_gallery, "modelaligned_recur1")
         gr.Markdown('### Output (Recursion #2)')
         l2_gallery = gr.Gallery(format='png', value=[], label="Recursion #2", show_label=True, elem_id="ncut_l2", columns=[100], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False, preview=True)
-        add_output_images_buttons(l2_gallery)
+        add_rotate_flip_buttons(l2_gallery)
         add_download_button(l2_gallery, "modelaligned_recur2")
         gr.Markdown('### Output (Recursion #3)')
         l3_gallery = gr.Gallery(format='png', value=[], label="Recursion #3", show_label=True, elem_id="ncut_l3", columns=[100], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False, preview=True)
-        add_output_images_buttons(l3_gallery)  
+        add_rotate_flip_buttons(l3_gallery)  
         add_download_button(l3_gallery, "modelaligned_recur3")
     
         with gr.Row():
@@ -2413,7 +2549,7 @@ with demo:
                 gr.Markdown(f'### Output Images')
                 output_gallery = gr.Gallery(format='png', value=[], label="NCUT Embedding", show_label=False, elem_id=f"ncut{i_model}", columns=[3], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
                 submit_button = gr.Button("🔴 RUN", elem_id=f"submit_button{i_model}", variant='primary')
-                add_output_images_buttons(output_gallery)
+                add_rotate_flip_buttons(output_gallery)
                 [
                     model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
                     affinity_focal_gamma_slider, num_sample_ncut_slider, ncut_knn_slider, ncut_indirect_connection, ncut_make_orthogonal, 
@@ -2479,13 +2615,52 @@ with demo:
         buttons[-1].click(fn=lambda x: gr.update(visible=False), outputs=buttons[-1])
     
     with gr.Tab('Compare Models (Advanced)', visible=False) as tab_compare_models_advanced: 
+        
+        target_images = gr.State([])
+        input_images = gr.State([])
+        def add_mlp_fitting_buttons(output_gallery, mlp_gallery, target_images=target_images, input_images=input_images):
+            with gr.Row():
+                # mark_as_target_button = gr.Button("mark target", elem_id=f"mark_as_target_button_{output_gallery.elem_id}", variant='secondary')
+                # mark_as_input_button = gr.Button("mark input", elem_id=f"mark_as_input_button_{output_gallery.elem_id}", variant='secondary')
+                mark_as_target_button = gr.Button("🎯 Mark Target", elem_id=f"mark_as_target_button_{output_gallery.elem_id}", variant='secondary')
+            fit_to_target_button = gr.Button("🔴 [MLP] Fit", elem_id=f"fit_to_target_button_{output_gallery.elem_id}", variant='primary')
+            def mark_fn(images, text="target"):
+                if images is None:
+                    raise gr.Error("No images selected")
+                if len(images) == 0:
+                    raise gr.Error("No images selected")
+                num_images = len(images)
+                gr.Info(f"Marked {num_images} images as {text}")
+                images = [(Image.open(tup[0]), []) for tup in images]
+                return images
+            mark_as_target_button.click(partial(mark_fn, text="target"), inputs=[output_gallery], outputs=[target_images])
+            # mark_as_input_button.click(partial(mark_fn, text="input"), inputs=[output_gallery], outputs=[input_images])
+            
+            with gr.Accordion("➡️ MLP Parameters", open=False):
+                num_layers_slider = gr.Slider(2, 10, step=1, label="Number of Layers", value=3, elem_id=f"num_layers_slider_{output_gallery.elem_id}")
+                width_slider = gr.Slider(128, 4096, step=128, label="Width", value=512, elem_id=f"width_slider_{output_gallery.elem_id}")
+                batch_size_slider = gr.Slider(32, 4096, step=32, label="Batch Size", value=128, elem_id=f"batch_size_slider_{output_gallery.elem_id}")
+                lr_slider = gr.Slider(1e-6, 1, step=1e-6, label="Learning Rate", value=3e-4, elem_id=f"lr_slider_{output_gallery.elem_id}")
+                fitting_steps_slider = gr.Slider(1000, 100000, step=1000, label="Fitting Steps", value=30000, elem_id=f"fitting_steps_slider_{output_gallery.elem_id}")
+                fps_sample_slider = gr.Slider(128, 50000, step=128, label="FPS Sample", value=10240, elem_id=f"fps_sample_slider_{output_gallery.elem_id}")
+                segmentation_loss_lambda_slider = gr.Slider(0, 100, step=0.01, label="Segmentation Preserving Loss Lambda", value=1, elem_id=f"segmentation_loss_lambda_slider_{output_gallery.elem_id}")
+                
+            fit_to_target_button.click(
+                run_mlp_fit,
+                inputs=[output_gallery, target_images, num_layers_slider, width_slider, batch_size_slider, lr_slider, fitting_steps_slider, fps_sample_slider, segmentation_loss_lambda_slider],
+                outputs=[mlp_gallery],
+            )
+                
         def add_one_model(i_model=1):
             with gr.Column(scale=5, min_width=200) as col:
                 gr.Markdown(f'### Output Images')
                 output_gallery = gr.Gallery(format='png', value=[], label="NCUT Embedding", show_label=True, elem_id=f"ncut{i_model}", columns=[3], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
                 submit_button = gr.Button("🔴 RUN", elem_id=f"submit_button{i_model}", variant='primary')
-                add_output_images_buttons(output_gallery)
+                add_rotate_flip_buttons(output_gallery)
                 add_download_button(output_gallery, f"ncut_embed")
+                mlp_gallery = gr.Gallery(format='png', value=[], label="MLP color align", show_label=True, elem_id=f"mlp{i_model}", columns=[3], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
+                add_mlp_fitting_buttons(output_gallery, mlp_gallery)
+                add_download_button(mlp_gallery, f"mlp_color_align")
                 norm_gallery = gr.Gallery(value=[], label="Eigenvector Magnitude", show_label=True, elem_id=f"eig_norm{i_model}", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
                 add_download_button(norm_gallery, f"eig_norm")
                 cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id=f"clusters{i_model}", columns=[2], rows=[4], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
@@ -2515,8 +2690,12 @@ with demo:
                     outputs=[output_gallery, cluster_gallery, norm_gallery, logging_text]
                 )
                 
-                return col
+                output_gallery.change(lambda x: gr.update(value=x), inputs=[output_gallery], outputs=[mlp_gallery])
+                
+                return output_gallery
 
+        galleries = []
+        
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
                 input_gallery, submit_button, clear_images_button, dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_input_images_section(allow_download=True)
@@ -2524,7 +2703,8 @@ with demo:
 
                 
             for i in range(3):
-                add_one_model()
+                g = add_one_model()
+                galleries.append(g)
                 
         # Create rows and buttons in a loop
         rows = []
@@ -2537,7 +2717,8 @@ with demo:
             with row:
                 for j in range(4):
                     with gr.Column(scale=5, min_width=200):
-                        add_one_model()
+                        g = add_one_model()
+                        galleries.append(g)
 
             button = gr.Button("➕ Add Compare", elem_id=f"add_button_{i}", visible=False if i > 0 else True, scale=3)
             buttons.append(button)
@@ -2553,7 +2734,11 @@ with demo:
         # Last button only reveals the last row and hides itself
         buttons[-1].click(fn=lambda x: gr.update(visible=True), outputs=rows[-1])
         buttons[-1].click(fn=lambda x: gr.update(visible=False), outputs=buttons[-1])
-    
+
+        
+        # add MLP fitting buttons
+        
+        
     
     with gr.Tab('📄About'):
         with gr.Column():