Pseudo-color-3

app.py

@@ -1,16 +1,12 @@
 import torch
 import numpy as np
 import gradio as gr
-from PIL import Image, ImageOps
+from PIL import Image
 import os
 import json
-import glob
-import random
 import tifffile
 import re
-import imageio
 from torchvision import transforms, models
-import accelerate
 import shutil
 import time
 import torch.nn as nn
@@ -20,11 +16,12 @@ from collections import OrderedDict
 import tempfile
 import zipfile
 import matplotlib.cm as cm
+import matplotlib.pyplot as plt
+import io
 
 # --- Imports from both scripts ---
 from diffusers import DDPMScheduler, DDIMScheduler
 from transformers import CLIPTextModel, CLIPTokenizer
-from accelerate.state import AcceleratorState
 from transformers.utils import ContextManagers
 
 # --- Custom Model Imports ---
@@ -34,45 +31,37 @@ from models.controlnet import ControlNetModel
 from models.unet_2d_condition import UNet2DConditionModel
 from models.pipeline_controlnet import DDPMControlnetPipeline
 
-# --- New Import for Segmentation ---
+# --- Segmentation Imports ---
 from cellpose import models as cellpose_models
-from cellpose import plot as cellpose_plot
 from huggingface_hub import snapshot_download
 
 # --- 0. Configuration & Constants ---
 hf_token = os.environ.get("HF_TOKEN")
-# --- General ---
 MODEL_TITLE = "🔬 FluoGen: AI-Powered Fluorescence Microscopy Suite"
 MODEL_DESCRIPTION = """
 **Paper**: *Generative AI empowering fluorescence microscopy imaging and analysis*
 <br>
-Select a task from the tabs below: generate new images from text, enhance existing images using super-resolution, denoise them, generate training data from segmentation masks, perform cell segmentation, or classify cell images.
+Select a task below. 
+**Note**: The "Pseudocolor" option can be adjusted *after* generation for instant visualization. Use the "Download Raw Output" button to get the scientific 16-bit/Float data.
 """
 DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
 WEIGHT_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 LOGO_PATH = "utils/logo2_transparent.png"
-
-# --- Global switch to control example saving ---
 SAVE_EXAMPLES = False
 
 # --- Base directory for all models ---
 REPO_ID = "FluoGen-Group/FluoGen-demo-test-ckpts"
 MODELS_ROOT_DIR = snapshot_download(repo_id=REPO_ID, token=hf_token) 
 
-# --- Tab 1: Mask-to-Image Config ---
+# --- Paths Config ---
 M2I_CONTROLNET_PATH = f"{MODELS_ROOT_DIR}/ControlNet_M2I/checkpoint-30000"
 M2I_EXAMPLE_IMG_DIR = "example_images_m2i"
-
-# --- Tab 2: Text-to-Image Config ---
 T2I_EXAMPLE_IMG_DIR = "example_images"
 T2I_PRETRAINED_MODEL_PATH = f"{MODELS_ROOT_DIR}/stable-diffusion-v1-5"
 T2I_UNET_PATH = f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/checkpoint-285000"
-
-# --- Tab 3, 4: ControlNet-based Tasks Config ---
 CONTROLNET_CLIP_PATH = f"{MODELS_ROOT_DIR}/stable-diffusion-v1-5"
 CONTROLNET_UNET_PATH = f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/checkpoint-285000"
 
-# Super-Resolution Models
 SR_CONTROLNET_MODELS = {
     "Checkpoint ER": f"{MODELS_ROOT_DIR}/ControlNet_SR/ER/checkpoint-30000",
     "Checkpoint Microtubules": f"{MODELS_ROOT_DIR}/ControlNet_SR/Microtubules/checkpoint-72500",
@@ -80,13 +69,10 @@ SR_CONTROLNET_MODELS = {
     "Checkpoint F-actin": f"{MODELS_ROOT_DIR}/ControlNet_SR/F-actin/checkpoint-35000",
 }
 SR_EXAMPLE_IMG_DIR = "example_images_sr"
-
-# Denoising Model
 DN_CONTROLNET_PATH = f"{MODELS_ROOT_DIR}/ControlNet_DN/checkpoint-10000"
 DN_PROMPT_RULES = {'MICE': 'mouse brain tissues', 'FISH': 'zebrafish embryos', 'BPAE_B': 'nucleus of BPAE', 'BPAE_R': 'mitochondria of BPAE', 'BPAE_G': 'F-actin of BPAE'}
 DN_EXAMPLE_IMG_DIR = "example_images_dn"
 
-# --- Tab 5: Cell Segmentation Config ---
 SEG_MODELS = {
     "DynamicNet Model": f"{MODELS_ROOT_DIR}/Cellpose/DynamicNet_baseline/CP_dynamic_ten_epoch_0100",
     "DynamicNet Model + FluoGen": f"{MODELS_ROOT_DIR}/Cellpose/DynamicNet_FluoGen/CP_dynamic_epoch_0300",
@@ -95,7 +81,6 @@ SEG_MODELS = {
 }
 SEG_EXAMPLE_IMG_DIR = "example_images_seg"
 
-# --- Tab 6: Classification Config ---
 CLS_MODEL_PATHS = OrderedDict({
     "5shot":        f"{MODELS_ROOT_DIR}/Classification/resnet50_hela_5_shot_re",
     "5shot+FluoGen":    f"{MODELS_ROOT_DIR}/Classification/resnet50_hela_5_shot_aug_re",
@@ -104,7 +89,7 @@ CLS_CLASS_NAMES = ['Nucleus', 'Endoplasmic Reticulum', 'Giantin', 'GPP130', 'Lys
 CLS_EXAMPLE_IMG_DIR = "example_images_cls"
 
 # --- Constants for Visualization ---
-COLOR_MAPS = ["Grayscale", "Green (GFP)", "Red (RFP)", "Blue (DAPI)", "Magenta", "Cyan", "Yellow", "Fire", "Viridis", "Inferno"]
+COLOR_MAPS = ["Grayscale", "Green (GFP)", "Red (RFP)", "Blue (DAPI)", "Magenta", "Cyan", "Yellow", "Fire", "Viridis", "Inferno", "Magma", "Plasma"]
 
 # --- Helper Functions ---
 def sanitize_prompt_for_filename(prompt):
@@ -116,57 +101,73 @@ def min_max_norm(x):
     if max_val - min_val < 1e-8: return np.zeros_like(x)
     return (x - min_val) / (max_val - min_val)
 
+def generate_colorbar_preview(color_name):
+    """Generates a small PIL image representing the colormap."""
+    if color_name == "Grayscale":
+        gradient = np.linspace(0, 1, 256).reshape(1, 256)
+        return Image.fromarray((gradient * 255).astype(np.uint8)).convert("RGB").resize((256, 30))
+    
+    gradient = np.linspace(0, 1, 256).reshape(1, 256)
+    rgb = np.zeros((1, 256, 3))
+    
+    if color_name == "Green (GFP)": rgb[..., 1] = gradient
+    elif color_name == "Red (RFP)": rgb[..., 0] = gradient
+    elif color_name == "Blue (DAPI)": rgb[..., 2] = gradient
+    elif color_name == "Magenta": rgb[..., 0] = gradient; rgb[..., 2] = gradient
+    elif color_name == "Cyan": rgb[..., 1] = gradient; rgb[..., 2] = gradient
+    elif color_name == "Yellow": rgb[..., 0] = gradient; rgb[..., 1] = gradient
+    else:
+        # Matplotlib maps
+        mpl_map_name = color_name.lower()
+        if color_name == "Fire": mpl_map_name = "gnuplot2"
+        try:
+            cmap = cm.get_cmap(mpl_map_name)
+            rgb = cmap(gradient)[..., :3]
+        except:
+            return generate_colorbar_preview("Grayscale") # Fallback
+            
+    img_np = (rgb * 255).astype(np.uint8)
+    return Image.fromarray(img_np).resize((256, 30))
+
 def apply_pseudocolor(image_np, color_name="Grayscale"):
     """
     Applies a pseudocolor to a single channel numpy image.
-    image_np: Single channel numpy array (any bit depth).
     Returns: PIL Image in RGB.
     """
+    if image_np is None: return None
+    
     # Normalize to 0-1 for processing
     norm_img = min_max_norm(np.squeeze(image_np))
     
     if color_name == "Grayscale":
-        # Just convert to uint8 L
         return Image.fromarray((norm_img * 255).astype(np.uint8)).convert("RGB")
     
-    # Create RGB canvas
     h, w = norm_img.shape
     rgb = np.zeros((h, w, 3), dtype=np.float32)
     
-    if color_name == "Green (GFP)":
-        rgb[..., 1] = norm_img
-    elif color_name == "Red (RFP)":
-        rgb[..., 0] = norm_img
-    elif color_name == "Blue (DAPI)":
-        rgb[..., 2] = norm_img
-    elif color_name == "Magenta":
-        rgb[..., 0] = norm_img
-        rgb[..., 2] = norm_img
-    elif color_name == "Cyan":
-        rgb[..., 1] = norm_img
-        rgb[..., 2] = norm_img
-    elif color_name == "Yellow":
-        rgb[..., 0] = norm_img
-        rgb[..., 1] = norm_img
-    elif color_name in ["Fire", "Viridis", "Inferno"]:
-        # Use matplotlib colormaps
-        cmap_map = {"Fire": "inferno", "Viridis": "viridis", "Inferno": "inferno"} # Fire looks like inferno usually
-        if color_name == "Fire": cmap = cm.get_cmap("gnuplot2") # Better fire approximation
-        else: cmap = cm.get_cmap(cmap_map[color_name])
-        
-        colored = cmap(norm_img) # Returns RGBA 0-1
-        rgb = colored[..., :3] # Drop Alpha
+    if color_name == "Green (GFP)": rgb[..., 1] = norm_img
+    elif color_name == "Red (RFP)": rgb[..., 0] = norm_img
+    elif color_name == "Blue (DAPI)": rgb[..., 2] = norm_img
+    elif color_name == "Magenta": rgb[..., 0] = norm_img; rgb[..., 2] = norm_img
+    elif color_name == "Cyan": rgb[..., 1] = norm_img; rgb[..., 2] = norm_img
+    elif color_name == "Yellow": rgb[..., 0] = norm_img; rgb[..., 1] = norm_img
+    else:
+        # Matplotlib maps
+        mpl_map_name = color_name.lower()
+        if color_name == "Fire": mpl_map_name = "gnuplot2"
+        try:
+            cmap = cm.get_cmap(mpl_map_name)
+            colored = cmap(norm_img)
+            rgb = colored[..., :3]
+        except:
+            return apply_pseudocolor(image_np, "Grayscale")
     
     return Image.fromarray((rgb * 255).astype(np.uint8))
 
 def save_temp_tiff(image_np, prefix="output"):
-    """Saves numpy array to a temp TIFF file and returns path."""
     tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".tif", prefix=f"{prefix}_")
-    # Ensure compatible type for Tiff (float32 or uint16 preferred for science)
-    if image_np.dtype == np.float16:
-        save_data = image_np.astype(np.float32)
-    else:
-        save_data = image_np
+    if image_np.dtype == np.float16: save_data = image_np.astype(np.float32)
+    else: save_data = image_np
     tifffile.imwrite(tfile.name, save_data)
     return tfile.name
 
@@ -176,8 +177,7 @@ def numpy_to_pil(image_np, target_mode="RGB"):
         if target_mode == "L" and image_np.mode != "L": return image_np.convert("L")
         return image_np
     squeezed_np = np.squeeze(image_np); 
-    if squeezed_np.dtype == np.uint8:
-        image_8bit = squeezed_np
+    if squeezed_np.dtype == np.uint8: image_8bit = squeezed_np
     else:
         normalized_np = min_max_norm(squeezed_np)
         image_8bit = (normalized_np * 255).astype(np.uint8)
@@ -211,8 +211,7 @@ def load_all_prompts():
                     if isinstance(data, list):
                         combined_prompts.extend(data)
                         for p in data: PROMPT_TO_MODEL_MAP[p] = cat["model"]
-                        print(f"✓ Loaded {len(data)} prompts from {cat['file']}")
-        except Exception as e: print(f"✗ Error loading {cat['file']}: {e}")
+        except Exception: pass
     if not combined_prompts: return ["F-actin of COS-7", "ER of COS-7"]
     return combined_prompts
 T2I_PROMPTS = load_all_prompts()
@@ -231,23 +230,22 @@ try:
     current_t2i_unet_path = T2I_UNET_PATH
     print("✓ Text-to-Image model loaded successfully!")
 except Exception as e:
-    print(f"!!!!!! FATAL: Text-to-Image Model Loading Failed !!!!!!\nError: {e}")
+    print(f"FATAL: Text-to-Image Model Loading Failed: {e}")
 
 try:
-    print("Loading shared ControlNet pipeline components...")
+    print("Loading shared ControlNet pipeline...")
     controlnet_unet = UNet2DConditionModel.from_pretrained(CONTROLNET_UNET_PATH, subfolder="unet").to(dtype=WEIGHT_DTYPE, device=DEVICE)
-    default_controlnet_path = M2I_CONTROLNET_PATH
-    controlnet_controlnet = ControlNetModel.from_pretrained(default_controlnet_path, subfolder="controlnet").to(dtype=WEIGHT_DTYPE, device=DEVICE)
+    controlnet_controlnet = ControlNetModel.from_pretrained(M2I_CONTROLNET_PATH, subfolder="controlnet").to(dtype=WEIGHT_DTYPE, device=DEVICE)
     controlnet_scheduler = DDIMScheduler(num_train_timesteps=1000, beta_schedule="linear", prediction_type="v_prediction", rescale_betas_zero_snr=False, timestep_spacing="trailing")
     controlnet_tokenizer = CLIPTokenizer.from_pretrained(CONTROLNET_CLIP_PATH, subfolder="tokenizer")
     with ContextManagers([]):
         controlnet_text_encoder = CLIPTextModel.from_pretrained(CONTROLNET_CLIP_PATH, subfolder="text_encoder").to(dtype=WEIGHT_DTYPE, device=DEVICE)
     controlnet_pipe = DDPMControlnetPipeline(unet=controlnet_unet, controlnet=controlnet_controlnet, scheduler=controlnet_scheduler, text_encoder=controlnet_text_encoder, tokenizer=controlnet_tokenizer)
     controlnet_pipe.to(dtype=WEIGHT_DTYPE, device=DEVICE)
-    controlnet_pipe.current_controlnet_path = default_controlnet_path
+    controlnet_pipe.current_controlnet_path = M2I_CONTROLNET_PATH
     print("✓ Shared ControlNet pipeline loaded successfully!")
 except Exception as e:
-    print(f"!!!!!! FATAL: ControlNet Pipeline Loading Failed !!!!!!\nError: {e}")
+    print(f"FATAL: ControlNet Pipeline Loading Failed: {e}")
 
 # --- 2. Core Logic Functions ---
 def swap_controlnet(pipe, target_path):
@@ -257,7 +255,7 @@ def swap_controlnet(pipe, target_path):
             pipe.controlnet = ControlNetModel.from_pretrained(target_path, subfolder="controlnet").to(dtype=WEIGHT_DTYPE, device=DEVICE)
             pipe.current_controlnet_path = target_path
         except Exception as e:
-            raise gr.Error(f"Failed to load ControlNet model '{target_path}'. Error: {e}")
+            raise gr.Error(f"Failed to load ControlNet model. Error: {e}")
     return pipe
 
 def swap_t2i_unet(pipe, target_unet_path):
@@ -269,245 +267,210 @@ def swap_t2i_unet(pipe, target_unet_path):
             new_unet = UNet2DModel.from_pretrained(target_unet_path, subfolder="unet").to(DEVICE)
             pipe.unet = new_unet
             current_t2i_unet_path = target_unet_path
-            print("✅ UNet swapped successfully.")
         except Exception as e:
-            raise gr.Error(f"Failed to load UNet from {target_unet_path}. Error: {e}")
+            raise gr.Error(f"Failed to load UNet. Error: {e}")
     return pipe
 
+# --- Dynamic Color Update Functions ---
+def update_single_image_color(raw_np_state, color_name):
+    if raw_np_state is None: return None, None
+    display_img = apply_pseudocolor(raw_np_state, color_name)
+    bar_img = generate_colorbar_preview(color_name)
+    return display_img, bar_img
+
+def update_gallery_color(raw_list_state, color_name):
+    if raw_list_state is None: return None, None
+    new_gallery = []
+    for img_np in raw_list_state:
+        new_gallery.append(apply_pseudocolor(img_np, color_name))
+    bar_img = generate_colorbar_preview(color_name)
+    return new_gallery, bar_img
+
+# --- Generation Functions ---
 @spaces.GPU(duration=120)
-def generate_t2i(prompt, num_inference_steps, colormap_choice):
+def generate_t2i(prompt, num_inference_steps, current_color):
     global t2i_pipe
     if t2i_pipe is None: raise gr.Error("Text-to-Image model is not loaded.")
-    target_model_path = PROMPT_TO_MODEL_MAP.get(prompt)
-    if not target_model_path:
-        target_model_path = f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/FULL-checkpoint-275000"
-        print(f"ℹ️ Prompt '{prompt}' not found in predefined list. Using Foundation (Full) model.")
+    target_model_path = PROMPT_TO_MODEL_MAP.get(prompt, f"{MODELS_ROOT_DIR}/UNET_T2I_CONTROLNET/FULL-checkpoint-275000")
     t2i_pipe = swap_t2i_unet(t2i_pipe, target_model_path)
     
-    print(f"\n🚀 Task started... | Prompt: '{prompt}' | Model: {current_t2i_unet_path}")
+    print(f"\n🚀 T2I Task started... | Prompt: '{prompt}'")
     image_np = t2i_pipe(prompt.lower(), generator=None, num_inference_steps=int(num_inference_steps), output_type="np").images
     
-    # 1. Save Raw Data
     raw_file_path = save_temp_tiff(image_np, prefix="t2i_raw")
+    display_image = apply_pseudocolor(image_np, current_color)
+    colorbar_img = generate_colorbar_preview(current_color)
     
-    # 2. Apply Pseudocolor for Display
-    display_image = apply_pseudocolor(image_np, colormap_choice)
-    
-    print("✓ Image generated")
     if SAVE_EXAMPLES:
         example_filepath = os.path.join(T2I_EXAMPLE_IMG_DIR, sanitize_prompt_for_filename(prompt))
-        if not os.path.exists(example_filepath):
-            display_image.save(example_filepath)
+        if not os.path.exists(example_filepath): display_image.save(example_filepath)
     
-    # Return Display Image AND Path to Raw File
-    return display_image, raw_file_path
+    # Return: Display, Raw Path, Raw State, Colorbar
+    return display_image, raw_file_path, image_np, colorbar_img
 
 @spaces.GPU(duration=120)
-def run_mask_to_image_generation(mask_file_obj, cell_type, num_images, steps, seed, colormap_choice):
+def run_mask_to_image_generation(mask_file_obj, cell_type, num_images, steps, seed, current_color):
     if controlnet_pipe is None: raise gr.Error("ControlNet pipeline is not loaded.")
-    if mask_file_obj is None: raise gr.Error("Please upload a segmentation mask TIF file.")
-    if not cell_type or not cell_type.strip(): raise gr.Error("Please enter a cell type.")
+    if mask_file_obj is None: raise gr.Error("Please upload a segmentation mask.")
     
     pipe = swap_controlnet(controlnet_pipe, M2I_CONTROLNET_PATH)
-    try:
-        mask_np = tifffile.imread(mask_file_obj.name)
-    except Exception as e:
-        raise gr.Error(f"Failed to read the TIF file. Error: {e}")
+    try: mask_np = tifffile.imread(mask_file_obj.name)
+    except Exception as e: raise gr.Error(f"Failed to read TIF. Error: {e}")
         
-    input_display_image = numpy_to_pil(mask_np, "L")
+    input_display = numpy_to_pil(mask_np, "L")
     mask_normalized = min_max_norm(mask_np)
-    image_tensor = torch.from_numpy(mask_normalized.astype(np.float32))
-    image_tensor = image_tensor.unsqueeze(0).unsqueeze(0).to(dtype=WEIGHT_DTYPE, device=DEVICE)
+    image_tensor = torch.from_numpy(mask_normalized.astype(np.float32)).unsqueeze(0).unsqueeze(0).to(dtype=WEIGHT_DTYPE, device=DEVICE)
     image_tensor = transforms.Resize((512, 512), antialias=True)(image_tensor)
     
     prompt = f"nuclei of {cell_type.strip()}"
-    print(f"\nTask started... | Task: Mask-to-Image | Prompt: '{prompt}' | Steps: {steps} | Images: {num_images}")
+    print(f"\nM2I Task started... | Prompt: '{prompt}'")
     
+    generated_raw_list = []
     generated_display_images = []
     generated_raw_files = []
-    
-    # Create a temp dir for the zip file
     temp_dir = tempfile.mkdtemp()
     
     for i in range(int(num_images)):
-        current_seed = int(seed) + i
-        generator = torch.Generator(device=DEVICE).manual_seed(current_seed)
+        generator = torch.Generator(device=DEVICE).manual_seed(int(seed) + i)
         with torch.autocast("cuda"):
             output_np = pipe(prompt=prompt, image_cond=image_tensor, num_inference_steps=int(steps), generator=generator, output_type="np").images
         
-        # Save individual raw file
+        # Save Raw State
+        generated_raw_list.append(output_np)
+        
+        # Save Temp File
         raw_name = f"m2i_sample_{i+1}.tif"
         raw_path = os.path.join(temp_dir, raw_name)
-        
-        # Ensure correct type saving
         save_data = output_np.astype(np.float32) if output_np.dtype == np.float16 else output_np
         tifffile.imwrite(raw_path, save_data)
         generated_raw_files.append(raw_path)
         
-        # Create display image
-        pil_image = apply_pseudocolor(output_np, colormap_choice)
-        generated_display_images.append(pil_image)
-        print(f"✓ Generated image {i+1}/{int(num_images)}")
+        # Display
+        generated_display_images.append(apply_pseudocolor(output_np, current_color))
     
-    # Create ZIP file
+    # ZIP
     zip_filename = os.path.join(temp_dir, "raw_output_images.zip")
     with zipfile.ZipFile(zip_filename, 'w') as zipf:
-        for file in generated_raw_files:
-            zipf.write(file, os.path.basename(file))
+        for file in generated_raw_files: zipf.write(file, os.path.basename(file))
             
-    return input_display_image, generated_display_images, zip_filename
+    colorbar_img = generate_colorbar_preview(current_color)
+    return input_display, generated_display_images, zip_filename, generated_raw_list, colorbar_img
 
 @spaces.GPU(duration=120)
-def run_super_resolution(low_res_file_obj, controlnet_model_name, prompt, steps, seed, colormap_choice):
+def run_super_resolution(low_res_file_obj, controlnet_model_name, prompt, steps, seed, current_color):
     if controlnet_pipe is None: raise gr.Error("ControlNet pipeline is not loaded.")
-    if low_res_file_obj is None: raise gr.Error("Please upload a low-resolution TIF file.")
+    if low_res_file_obj is None: raise gr.Error("Please upload a file.")
             
     target_path = SR_CONTROLNET_MODELS.get(controlnet_model_name)
-    if not target_path: raise gr.Error(f"ControlNet model '{controlnet_model_name}' not found.")
-    
     pipe = swap_controlnet(controlnet_pipe, target_path)
-    try:
-        image_stack_np = tifffile.imread(low_res_file_obj.name)
-    except Exception as e:
-        raise gr.Error(f"Failed to read the TIF file. Error: {e}")
+    
+    try: image_stack_np = tifffile.imread(low_res_file_obj.name)
+    except Exception as e: raise gr.Error(f"Failed to read TIF. Error: {e}")
         
     if image_stack_np.ndim != 3 or image_stack_np.shape[-3] != 9:
-        raise gr.Error(f"Invalid TIF shape. Expected 9 channels (shape 9, H, W), but got {image_stack_np.shape}.")
+        raise gr.Error(f"Invalid TIF shape. Expected 9 channels, got {image_stack_np.shape}.")
         
-    average_projection_np = np.mean(image_stack_np, axis=0)
-    input_display_image = numpy_to_pil(average_projection_np, "L")
-    
+    input_display = numpy_to_pil(np.mean(image_stack_np, axis=0), "L")
     image_tensor = torch.from_numpy(image_stack_np.astype(np.float32) / 65535.0).unsqueeze(0).to(dtype=WEIGHT_DTYPE, device=DEVICE)
     image_tensor = transforms.Resize((512, 512), antialias=True)(image_tensor)
     
+    print(f"\nSR Task started...")
     generator = torch.Generator(device=DEVICE).manual_seed(int(seed))
     with torch.autocast("cuda"):
         output_np = pipe(prompt=prompt, image_cond=image_tensor, num_inference_steps=int(steps), generator=generator, output_type="np").images
     
-    # Save Raw
     raw_file_path = save_temp_tiff(output_np, prefix="sr_raw")
-    # Display Color
-    output_display = apply_pseudocolor(output_np, colormap_choice)
+    output_display = apply_pseudocolor(output_np, current_color)
+    colorbar_img = generate_colorbar_preview(current_color)
         
-    return input_display_image, output_display, raw_file_path
+    return input_display, output_display, raw_file_path, output_np, colorbar_img
 
 @spaces.GPU(duration=120)
-def run_denoising(noisy_image_np, image_type, steps, seed, colormap_choice):
+def run_denoising(noisy_image_np, image_type, steps, seed, current_color):
     if controlnet_pipe is None: raise gr.Error("ControlNet pipeline is not loaded.")
-    if noisy_image_np is None: raise gr.Error("Please upload a noisy image.")
+    if noisy_image_np is None: raise gr.Error("Please upload an image.")
     
     pipe = swap_controlnet(controlnet_pipe, DN_CONTROLNET_PATH)
     prompt = DN_PROMPT_RULES.get(image_type, 'microscopy image')
-    print(f"\nTask started... | Task: Denoising | Prompt: '{prompt}' | Steps: {steps}")
     
-    image_tensor = torch.from_numpy(noisy_image_np.astype(np.float32) / 255.0)
-    image_tensor = image_tensor.unsqueeze(0).unsqueeze(0).to(dtype=WEIGHT_DTYPE, device=DEVICE)
+    print(f"\nDN Task started...")
+    image_tensor = torch.from_numpy(noisy_image_np.astype(np.float32) / 255.0).unsqueeze(0).unsqueeze(0).to(dtype=WEIGHT_DTYPE, device=DEVICE)
     image_tensor = transforms.Resize((512, 512), antialias=True)(image_tensor)
     
     generator = torch.Generator(device=DEVICE).manual_seed(int(seed))
     with torch.autocast("cuda"):
         output_np = pipe(prompt=prompt, image_cond=image_tensor, num_inference_steps=int(steps), generator=generator, output_type="np").images
     
-    # Save Raw
     raw_file_path = save_temp_tiff(output_np, prefix="dn_raw")
-    # Display Color
-    output_display = apply_pseudocolor(output_np, colormap_choice)
+    output_display = apply_pseudocolor(output_np, current_color)
+    colorbar_img = generate_colorbar_preview(current_color)
         
-    return numpy_to_pil(noisy_image_np, "L"), output_display, raw_file_path
+    return numpy_to_pil(noisy_image_np, "L"), output_display, raw_file_path, output_np, colorbar_img
 
+# --- Segmentation & Classification (Unchanged Logic, just wrapper) ---
 @spaces.GPU(duration=120)
 def run_segmentation(input_image_np, model_name, diameter, flow_threshold, cellprob_threshold):
-    # Segmentation remains mostly same, usually overlay is RGB anyway
-    if input_image_np is None: raise gr.Error("Please upload an image to segment.")
+    if input_image_np is None: raise gr.Error("Please upload an image.")
     model_path = SEG_MODELS.get(model_name)
-    if not model_path: raise gr.Error(f"Segmentation model '{model_name}' not found.")
     
-    print(f"\nTask started... | Task: Cell Segmentation | Model: '{model_name}'")
+    print(f"\nSeg Task started...")
     try:
-        use_gpu = torch.cuda.is_available()
-        model = cellpose_models.CellposeModel(gpu=use_gpu, pretrained_model=model_path)
-    except Exception as e:
-        raise gr.Error(f"Failed to load Cellpose model. Error: {e}")
-
-    diameter_to_use = model.diam_labels if diameter == 0 else float(diameter)
-    try:
-        masks, _, _ = model.eval([input_image_np], channels=[0, 0], diameter=diameter_to_use, flow_threshold=flow_threshold, cellprob_threshold=cellprob_threshold)
-        mask_output = masks[0]
-    except Exception as e:
-        raise gr.Error(f"Cellpose model evaluation failed. Error: {e}")
+        model = cellpose_models.CellposeModel(gpu=torch.cuda.is_available(), pretrained_model=model_path)
+        masks, _, _ = model.eval([input_image_np], channels=[0, 0], diameter=(model.diam_labels if diameter==0 else diameter), flow_threshold=flow_threshold, cellprob_threshold=cellprob_threshold)
+    except Exception as e: raise gr.Error(f"Segmentation failed: {e}")
 
     original_rgb = numpy_to_pil(input_image_np, "RGB")
-    original_rgb_np = np.array(original_rgb)
-    red_mask_layer = np.zeros_like(original_rgb_np)
-    red_mask_layer[mask_output > 0] = [139, 0, 0]
-    blended_image_np = ((0.6 * original_rgb_np + 0.4 * red_mask_layer).astype(np.uint8))
-    
-    return numpy_to_pil(input_image_np, "L"), numpy_to_pil(blended_image_np, "RGB")
+    red_mask = np.zeros_like(np.array(original_rgb)); red_mask[masks[0] > 0] = [139, 0, 0]
+    blended = ((0.6 * np.array(original_rgb) + 0.4 * red_mask).astype(np.uint8))
+    return numpy_to_pil(input_image_np, "L"), numpy_to_pil(blended, "RGB")
 
 @spaces.GPU(duration=120)
 def run_classification(input_image_np, model_name):
-    if input_image_np is None: raise gr.Error("Please upload an image to classify.")
-    model_dir = CLS_MODEL_PATHS.get(model_name)
-    if not model_dir: raise gr.Error(f"Classification model '{model_name}' not found.")
-    model_path = os.path.join(model_dir, "best_resnet50.pth")
+    if input_image_np is None: raise gr.Error("Please upload an image.")
+    model_path = os.path.join(CLS_MODEL_PATHS.get(model_name), "best_resnet50.pth")
     
-    print(f"\nTask started... | Task: Classification | Model: '{model_name}'")
+    print(f"\nCls Task started...")
     try:
-        model = models.resnet50(weights=None)
-        model.fc = nn.Linear(model.fc.in_features, len(CLS_CLASS_NAMES))
+        model = models.resnet50(weights=None); model.fc = nn.Linear(model.fc.in_features, len(CLS_CLASS_NAMES))
         model.load_state_dict(torch.load(model_path, map_location=DEVICE))
         model.to(DEVICE).eval()
-    except Exception as e:
-        raise gr.Error(f"Failed to load classification model. Error: {e}")
-
-    input_pil = numpy_to_pil(input_image_np, "RGB")
-    transform_test = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
-    input_tensor = transform_test(input_pil).unsqueeze(0).to(DEVICE)
+    except Exception as e: raise gr.Error(f"Classification failed: {e}")
 
+    input_tensor = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)])(numpy_to_pil(input_image_np, "RGB")).unsqueeze(0).to(DEVICE)
     with torch.no_grad():
-        outputs = model(input_tensor)
-        probabilities = F.softmax(outputs, dim=1).squeeze().cpu().numpy()
-
-    confidences = {name: float(prob) for name, prob in zip(CLS_CLASS_NAMES, probabilities)}
-    return numpy_to_pil(input_image_np, "L"), confidences
-
+        probs = F.softmax(model(input_tensor), dim=1).squeeze().cpu().numpy()
+    return numpy_to_pil(input_image_np, "L"), {name: float(p) for name, p in zip(CLS_CLASS_NAMES, probs)}
 
 # --- 3. Gradio UI Layout ---
 print("Building Gradio interface...")
-for d in [M2I_EXAMPLE_IMG_DIR, T2I_EXAMPLE_IMG_DIR, SR_EXAMPLE_IMG_DIR, DN_EXAMPLE_IMG_DIR, SEG_EXAMPLE_IMG_DIR, CLS_EXAMPLE_IMG_DIR]:
-    os.makedirs(d, exist_ok=True)
-
-# --- Load examples ---
-filename_to_prompt_map = { sanitize_prompt_for_filename(prompt): prompt for prompt in T2I_PROMPTS }
-t2i_gallery_examples = []
-for filename in os.listdir(T2I_EXAMPLE_IMG_DIR):
-    if filename in filename_to_prompt_map:
-        t2i_gallery_examples.append((os.path.join(T2I_EXAMPLE_IMG_DIR, filename), filename_to_prompt_map[filename]))
-
-def load_image_examples(example_dir, is_stack=False):
-    examples = []
-    if not os.path.exists(example_dir): return examples
-    for f in sorted(os.listdir(example_dir)):
+for d in [M2I_EXAMPLE_IMG_DIR, T2I_EXAMPLE_IMG_DIR, SR_EXAMPLE_IMG_DIR, DN_EXAMPLE_IMG_DIR, SEG_EXAMPLE_IMG_DIR, CLS_EXAMPLE_IMG_DIR]: os.makedirs(d, exist_ok=True)
+
+# Load Examples
+filename_to_prompt = { sanitize_prompt_for_filename(p): p for p in T2I_PROMPTS }
+t2i_examples = []
+for f in os.listdir(T2I_EXAMPLE_IMG_DIR):
+    if f in filename_to_prompt: t2i_examples.append((os.path.join(T2I_EXAMPLE_IMG_DIR, f), filename_to_prompt[f]))
+
+def load_examples(d, stack=False):
+    ex = []
+    if not os.path.exists(d): return ex
+    for f in sorted(os.listdir(d)):
         if f.lower().endswith(('.tif', '.tiff', '.png', '.jpg')):
-            filepath = os.path.join(example_dir, f)
+            fp = os.path.join(d, f)
             try:
-                if f.lower().endswith(('.tif', '.tiff')): img_np = tifffile.imread(filepath)
-                else: img_np = np.array(Image.open(filepath).convert("L"))
-                if is_stack and img_np.ndim == 3: img_np = np.mean(img_np, axis=0)
-                examples.append((numpy_to_pil(img_np, "L"), filepath))
+                img = tifffile.imread(fp) if f.endswith(('.tif','.tiff')) else np.array(Image.open(fp).convert("L"))
+                if stack and img.ndim == 3: img = np.mean(img, axis=0)
+                ex.append((numpy_to_pil(img, "L"), fp))
             except: pass
-    return examples
+    return ex
 
-m2i_gallery_examples = load_image_examples(M2I_EXAMPLE_IMG_DIR)
-sr_gallery_examples = load_image_examples(SR_EXAMPLE_IMG_DIR, is_stack=True)
-dn_gallery_examples = load_image_examples(DN_EXAMPLE_IMG_DIR)
-seg_gallery_examples = load_image_examples(SEG_EXAMPLE_IMG_DIR)
-cls_gallery_examples = load_image_examples(CLS_EXAMPLE_IMG_DIR)
-
-# --- Universal event handlers ---
-def select_example_prompt(evt: gr.SelectData): return evt.value['caption']
-def select_example_input_file(evt: gr.SelectData): return evt.value['caption']
+m2i_examples = load_examples(M2I_EXAMPLE_IMG_DIR)
+sr_examples = load_examples(SR_EXAMPLE_IMG_DIR, stack=True)
+dn_examples = load_examples(DN_EXAMPLE_IMG_DIR)
+seg_examples = load_examples(SEG_EXAMPLE_IMG_DIR)
+cls_examples = load_examples(CLS_EXAMPLE_IMG_DIR)
 
+# --- UI Builders ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with gr.Row():
         gr.Image(value=LOGO_PATH, width=300, height=200, container=False, interactive=False, show_download_button=False, show_fullscreen_button=False)
@@ -516,122 +479,160 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with gr.Tabs():
         # --- TAB 1: Text-to-Image ---
         with gr.Tab("Text-to-Image Generation", id="txt2img"):
+            # State to hold raw numpy data
+            t2i_raw_state = gr.State(None)
+            
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1, min_width=350):
-                    t2i_prompt_input = gr.Dropdown(choices=T2I_PROMPTS, value=T2I_PROMPTS[0], label="Search or Type a Prompt", filterable=True, allow_custom_value=True)
-                    t2i_steps_slider = gr.Slider(minimum=10, maximum=200, step=1, value=50, label="Inference Steps")
-                    # New: Color Selector
-                    t2i_color_input = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor (for Display)")
-                    t2i_generate_button = gr.Button("Generate", variant="primary")
+                    t2i_prompt = gr.Dropdown(choices=T2I_PROMPTS, value=T2I_PROMPTS[0], label="Search or Type a Prompt", filterable=True, allow_custom_value=True)
+                    t2i_steps = gr.Slider(10, 200, 50, step=1, label="Inference Steps")
+                    t2i_btn = gr.Button("Generate", variant="primary")
                 with gr.Column(scale=2):
-                    t2i_generated_output = gr.Image(label="Generated Image (Pseudocolor)", type="pil", interactive=False)
-                    # New: Raw Download Button
-                    t2i_raw_download = gr.DownloadButton(label="Download Raw Output (.tif)", visible=True)
-            t2i_gallery = gr.Gallery(value=t2i_gallery_examples, label="Examples", columns=6, object_fit="contain", height="auto")
+                    # Image with Download Button DISABLED
+                    t2i_out = gr.Image(label="Generated Image", type="pil", interactive=False, show_download_button=False)
+                    
+                    with gr.Row(equal_height=True):
+                        # Color Controls
+                        with gr.Column(scale=2):
+                            t2i_color = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor (Adjust after generation)")
+                        with gr.Column(scale=2):
+                            t2i_colorbar = gr.Image(label="Colorbar", show_label=False, container=False, height=40, show_download_button=False, interactive=False)
+                        with gr.Column(scale=1):
+                            t2i_dl = gr.DownloadButton(label="Download Raw (.tif)")
+                            
+            t2i_gal = gr.Gallery(value=t2i_examples, label="Examples", columns=6, height="auto")
+
+            # Events
+            t2i_btn.click(
+                fn=generate_t2i, 
+                inputs=[t2i_prompt, t2i_steps, t2i_color], 
+                outputs=[t2i_out, t2i_dl, t2i_raw_state, t2i_colorbar]
+            )
+            # Real-time color update using State
+            t2i_color.change(
+                fn=update_single_image_color,
+                inputs=[t2i_raw_state, t2i_color],
+                outputs=[t2i_out, t2i_colorbar]
+            )
+            t2i_gal.select(lambda e: e.value['caption'], None, t2i_prompt)
 
         # --- TAB 2: Super-Resolution ---
         with gr.Tab("Super-Resolution", id="super_res"):
+            sr_raw_state = gr.State(None)
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1, min_width=350):
-                    sr_input_file = gr.File(label="Upload 9-Channel TIF Stack", file_types=['.tif', '.tiff'])
-                    sr_model_selector = gr.Dropdown(choices=list(SR_CONTROLNET_MODELS.keys()), value=list(SR_CONTROLNET_MODELS.keys())[-1], label="Select Super-Resolution Model")
-                    sr_prompt_input = gr.Textbox(label="Prompt", value="F-actin of COS-7", interactive=False)
-                    sr_steps_slider = gr.Slider(minimum=5, maximum=50, step=1, value=10, label="Inference Steps")
-                    sr_seed_input = gr.Number(label="Seed", value=42)
-                    sr_color_input = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor (for Display)")
-                    sr_generate_button = gr.Button("Generate Super-Resolution", variant="primary")
+                    sr_file = gr.File(label="Upload 9-Channel TIF Stack", file_types=['.tif', '.tiff'])
+                    sr_model = gr.Dropdown(choices=list(SR_CONTROLNET_MODELS.keys()), value=list(SR_CONTROLNET_MODELS.keys())[-1], label="Model")
+                    sr_prompt = gr.Textbox(label="Prompt", value="F-actin of COS-7", interactive=False)
+                    sr_steps = gr.Slider(5, 50, 10, step=1, label="Steps")
+                    sr_seed = gr.Number(label="Seed", value=42)
+                    sr_btn = gr.Button("Generate", variant="primary")
                 with gr.Column(scale=2):
                     with gr.Row():
-                        sr_input_display = gr.Image(label="Input (Avg)", type="pil", interactive=False)
-                        sr_output_image = gr.Image(label="Super-Resolved Image", type="pil", interactive=False)
-                    sr_raw_download = gr.DownloadButton(label="Download Raw Output (.tif)", visible=True)
-            sr_gallery = gr.Gallery(value=sr_gallery_examples, label="Examples", columns=6, object_fit="contain", height="auto")
+                        sr_in_disp = gr.Image(label="Input (Avg)", type="pil", interactive=False, show_download_button=False)
+                        sr_out_disp = gr.Image(label="Output", type="pil", interactive=False, show_download_button=False)
+                    with gr.Row(equal_height=True):
+                        with gr.Column(scale=2):
+                            sr_color = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor")
+                        with gr.Column(scale=2):
+                            sr_colorbar = gr.Image(label="Colorbar", show_label=False, container=False, height=40, show_download_button=False, interactive=False)
+                        with gr.Column(scale=1):
+                            sr_dl = gr.DownloadButton(label="Download Raw (.tif)")
+            
+            sr_gal = gr.Gallery(value=sr_examples, label="Examples", columns=6, height="auto")
+            
+            sr_model.change(update_sr_prompt, sr_model, sr_prompt)
+            sr_btn.click(run_super_resolution, [sr_file, sr_model, sr_prompt, sr_steps, sr_seed, sr_color], [sr_in_disp, sr_out_disp, sr_dl, sr_raw_state, sr_colorbar])
+            sr_color.change(update_single_image_color, [sr_raw_state, sr_color], [sr_out_disp, sr_colorbar])
+            sr_gal.select(lambda e: e.value['caption'], None, sr_file)
 
         # --- TAB 3: Denoising ---
         with gr.Tab("Denoising", id="denoising"):
+            dn_raw_state = gr.State(None)
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1, min_width=350):
-                    dn_input_image = gr.Image(type="numpy", label="Upload Noisy Image", image_mode="L")
-                    dn_image_type_selector = gr.Dropdown(choices=list(DN_PROMPT_RULES.keys()), value='MICE', label="Select Image Type")
-                    dn_steps_slider = gr.Slider(minimum=5, maximum=50, step=1, value=10, label="Inference Steps")
-                    dn_seed_input = gr.Number(label="Seed", value=42)
-                    dn_color_input = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor (for Display)")
-                    dn_generate_button = gr.Button("Denoise Image", variant="primary")
+                    dn_img = gr.Image(type="numpy", label="Upload Noisy Image", image_mode="L")
+                    dn_type = gr.Dropdown(choices=list(DN_PROMPT_RULES.keys()), value='MICE', label="Image Type")
+                    dn_steps = gr.Slider(5, 50, 10, step=1, label="Steps")
+                    dn_seed = gr.Number(label="Seed", value=42)
+                    dn_btn = gr.Button("Denoise", variant="primary")
                 with gr.Column(scale=2):
                     with gr.Row():
-                        dn_original_display = gr.Image(label="Original", type="pil", interactive=False)
-                        dn_output_image = gr.Image(label="Denoised Image", type="pil", interactive=False)
-                    dn_raw_download = gr.DownloadButton(label="Download Raw Output (.tif)", visible=True)
-            dn_gallery = gr.Gallery(value=dn_gallery_examples, label="Examples", columns=6, object_fit="contain", height="auto")
+                        dn_orig = gr.Image(label="Original", type="pil", interactive=False, show_download_button=False)
+                        dn_out = gr.Image(label="Denoised", type="pil", interactive=False, show_download_button=False)
+                    with gr.Row(equal_height=True):
+                        with gr.Column(scale=2):
+                            dn_color = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor")
+                        with gr.Column(scale=2):
+                            dn_colorbar = gr.Image(label="Colorbar", show_label=False, container=False, height=40, show_download_button=False, interactive=False)
+                        with gr.Column(scale=1):
+                            dn_dl = gr.DownloadButton(label="Download Raw (.tif)")
+            
+            dn_gal = gr.Gallery(value=dn_examples, label="Examples", columns=6, height="auto")
+            
+            dn_btn.click(run_denoising, [dn_img, dn_type, dn_steps, dn_seed, dn_color], [dn_orig, dn_out, dn_dl, dn_raw_state, dn_colorbar])
+            dn_color.change(update_single_image_color, [dn_raw_state, dn_color], [dn_out, dn_colorbar])
+            dn_gal.select(lambda e: e.value['caption'], None, dn_img)
 
         # --- TAB 4: Mask-to-Image ---
         with gr.Tab("Mask-to-Image", id="mask2img"):
+            m2i_raw_state = gr.State(None) # Stores list of numpy arrays
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1, min_width=350):
-                    m2i_input_file = gr.File(label="Upload Segmentation Mask (.tif)", file_types=['.tif', '.tiff'])
-                    m2i_cell_type_input = gr.Textbox(label="Cell Type", placeholder="e.g., CoNSS, HeLa")
-                    m2i_num_images_slider = gr.Slider(minimum=1, maximum=10, step=1, value=5, label="Number of Images")
-                    m2i_steps_slider = gr.Slider(minimum=5, maximum=50, step=1, value=10, label="Inference Steps")
-                    m2i_seed_input = gr.Number(label="Seed", value=42)
-                    m2i_color_input = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor (for Display)")
-                    m2i_generate_button = gr.Button("Generate Samples", variant="primary")
+                    m2i_file = gr.File(label="Upload Mask (.tif)", file_types=['.tif', '.tiff'])
+                    m2i_type = gr.Textbox(label="Cell Type", placeholder="e.g., HeLa")
+                    m2i_num = gr.Slider(1, 10, 5, step=1, label="Count")
+                    m2i_steps = gr.Slider(5, 50, 10, step=1, label="Steps")
+                    m2i_seed = gr.Number(label="Seed", value=42)
+                    m2i_btn = gr.Button("Generate", variant="primary")
                 with gr.Column(scale=2):
-                    m2i_output_gallery = gr.Gallery(label="Generated Samples", columns=5, object_fit="contain", height="auto")
-                    m2i_raw_download = gr.DownloadButton(label="Download All Raw Samples (.zip)", visible=True)
-                    m2i_input_display = gr.Image(label="Input Mask", type="pil", interactive=False)
-            m2i_gallery = gr.Gallery(value=m2i_gallery_examples, label="Examples", columns=6, object_fit="contain", height="auto")
+                    m2i_gal_out = gr.Gallery(label="Generated Samples", columns=5, height="auto")
+                    with gr.Row(equal_height=True):
+                        with gr.Column(scale=2):
+                            m2i_color = gr.Dropdown(choices=COLOR_MAPS, value="Grayscale", label="Pseudocolor")
+                        with gr.Column(scale=2):
+                            m2i_colorbar = gr.Image(label="Colorbar", show_label=False, container=False, height=40, show_download_button=False, interactive=False)
+                        with gr.Column(scale=1):
+                            m2i_dl = gr.DownloadButton(label="Download ZIP")
+                    m2i_in_disp = gr.Image(label="Input Mask", type="pil", interactive=False, show_download_button=False)
+            
+            m2i_gal = gr.Gallery(value=m2i_examples, label="Examples", columns=6, height="auto")
+            
+            m2i_btn.click(run_mask_to_image_generation, [m2i_file, m2i_type, m2i_num, m2i_steps, m2i_seed, m2i_color], [m2i_in_disp, m2i_gal_out, m2i_dl, m2i_raw_state, m2i_colorbar])
+            m2i_color.change(update_gallery_color, [m2i_raw_state, m2i_color], [m2i_gal_out, m2i_colorbar])
+            m2i_gal.select(lambda e: e.value['caption'], None, m2i_file)
 
         # --- TAB 5: Cell Segmentation ---
         with gr.Tab("Cell Segmentation", id="segmentation"):
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1, min_width=350):
-                    seg_input_image = gr.Image(type="numpy", label="Upload Image", image_mode="L")
-                    seg_model_selector = gr.Dropdown(choices=list(SEG_MODELS.keys()), value=list(SEG_MODELS.keys())[0], label="Model")
-                    seg_diameter_input = gr.Number(label="Cell Diameter (0=auto)", value=30)
-                    seg_flow_slider = gr.Slider(minimum=0.0, maximum=3.0, step=0.1, value=0.4, label="Flow Threshold")
-                    seg_cellprob_slider = gr.Slider(minimum=-6.0, maximum=6.0, step=0.5, value=0.0, label="Cell Probability Threshold")
-                    seg_generate_button = gr.Button("Segment Cells", variant="primary")
+                    seg_img = gr.Image(type="numpy", label="Upload Image", image_mode="L")
+                    seg_model = gr.Dropdown(choices=list(SEG_MODELS.keys()), value=list(SEG_MODELS.keys())[0], label="Model")
+                    seg_diam = gr.Number(label="Diameter (0=auto)", value=30)
+                    seg_flow = gr.Slider(0.0, 3.0, 0.4, step=0.1, label="Flow Thresh")
+                    seg_prob = gr.Slider(-6.0, 6.0, 0.0, step=0.5, label="Prob Thresh")
+                    seg_btn = gr.Button("Segment", variant="primary")
                 with gr.Column(scale=2):
                     with gr.Row():
-                        seg_original_display = gr.Image(label="Original", type="pil", interactive=False)
-                        seg_output_image = gr.Image(label="Segmented Overlay", type="pil", interactive=False)
-            seg_gallery = gr.Gallery(value=seg_gallery_examples, label="Examples", columns=6, object_fit="contain", height="auto")
-            
+                        seg_orig = gr.Image(label="Original", type="pil", interactive=False, show_download_button=False)
+                        seg_out = gr.Image(label="Overlay", type="pil", interactive=False, show_download_button=False)
+            seg_gal = gr.Gallery(value=seg_examples, label="Examples", columns=6, height="auto")
+            seg_btn.click(run_segmentation, [seg_img, seg_model, seg_diam, seg_flow, seg_prob], [seg_orig, seg_out])
+            seg_gal.select(lambda e: e.value['caption'], None, seg_img)
+
         # --- TAB 6: Classification ---
         with gr.Tab("Classification", id="classification"):
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1, min_width=350):
-                    cls_input_image = gr.Image(type="numpy", label="Upload Image", image_mode="L")
-                    cls_model_selector = gr.Dropdown(choices=list(CLS_MODEL_PATHS.keys()), value=list(CLS_MODEL_PATHS.keys())[0], label="Model")
-                    cls_generate_button = gr.Button("Classify Image", variant="primary")
+                    cls_img = gr.Image(type="numpy", label="Upload Image", image_mode="L")
+                    cls_model = gr.Dropdown(choices=list(CLS_MODEL_PATHS.keys()), value=list(CLS_MODEL_PATHS.keys())[0], label="Model")
+                    cls_btn = gr.Button("Classify", variant="primary")
                 with gr.Column(scale=2):
-                    cls_original_display = gr.Image(label="Input Image", type="pil", interactive=False)
-                    cls_output_label = gr.Label(label="Results", num_top_classes=len(CLS_CLASS_NAMES))
-            cls_gallery = gr.Gallery(value=cls_gallery_examples, label="Examples", columns=6, object_fit="contain", height="auto")
-
-
-    # --- Event Handlers ---
-    m2i_generate_button.click(fn=run_mask_to_image_generation, inputs=[m2i_input_file, m2i_cell_type_input, m2i_num_images_slider, m2i_steps_slider, m2i_seed_input, m2i_color_input], outputs=[m2i_input_display, m2i_output_gallery, m2i_raw_download])
-    m2i_gallery.select(fn=select_example_input_file, outputs=m2i_input_file)
-    
-    t2i_generate_button.click(fn=generate_t2i, inputs=[t2i_prompt_input, t2i_steps_slider, t2i_color_input], outputs=[t2i_generated_output, t2i_raw_download])
-    t2i_gallery.select(fn=select_example_prompt, outputs=t2i_prompt_input)
-    
-    sr_model_selector.change(fn=update_sr_prompt, inputs=sr_model_selector, outputs=sr_prompt_input)
-    sr_generate_button.click(fn=run_super_resolution, inputs=[sr_input_file, sr_model_selector, sr_prompt_input, sr_steps_slider, sr_seed_input, sr_color_input], outputs=[sr_input_display, sr_output_image, sr_raw_download])
-    sr_gallery.select(fn=select_example_input_file, outputs=sr_input_file)
-    
-    dn_generate_button.click(fn=run_denoising, inputs=[dn_input_image, dn_image_type_selector, dn_steps_slider, dn_seed_input, dn_color_input], outputs=[dn_original_display, dn_output_image, dn_raw_download])
-    dn_gallery.select(fn=select_example_input_file, outputs=dn_input_image)
-
-    seg_generate_button.click(fn=run_segmentation, inputs=[seg_input_image, seg_model_selector, seg_diameter_input, seg_flow_slider, seg_cellprob_slider], outputs=[seg_original_display, seg_output_image])
-    seg_gallery.select(fn=select_example_input_file, outputs=seg_input_image)
-
-    cls_generate_button.click(fn=run_classification, inputs=[cls_input_image, cls_model_selector], outputs=[cls_original_display, cls_output_label])
-    cls_gallery.select(fn=select_example_input_file, outputs=cls_input_image)
-
+                    cls_orig = gr.Image(label="Input", type="pil", interactive=False, show_download_button=False)
+                    cls_res = gr.Label(label="Results", num_top_classes=10)
+            cls_gal = gr.Gallery(value=cls_examples, label="Examples", columns=6, height="auto")
+            cls_btn.click(run_classification, [cls_img, cls_model], [cls_orig, cls_res])
+            cls_gal.select(lambda e: e.value['caption'], None, cls_img)
 
-# --- 4. Launch Application ---
 if __name__ == "__main__":
-    print("Interface built. Launching server...")
     demo.launch()