app.py

import gradio as gr
import torch
import numpy as np
from PIL import Image, ImageDraw
try:
    from spaces import GPU
except ImportError:
    # Define a no-op decorator if running locally
    def GPU(func):
        return func
    
import os
import re
import json
import argparse
from datetime import datetime
from inference import GenerativeInferenceModel, get_inference_configs, get_imagenet_labels

# Parse command line arguments
parser = argparse.ArgumentParser(description='Predict Human Hallucinations Demo')
parser.add_argument('--port', type=int, default=7860, help='Port to run the server on')
args = parser.parse_args()

# Create model directories if they don't exist
os.makedirs("models", exist_ok=True)
os.makedirs("stimuli", exist_ok=True)
SAVED_RUNS_DIR = "saved_runs"
os.makedirs(SAVED_RUNS_DIR, exist_ok=True)

# Load ImageNet labels for biased-inference dropdown (1000 classes)
IMAGENET_LABELS = get_imagenet_labels()

# Check if running on Hugging Face Spaces
if "SPACE_ID" in os.environ:
    default_port = int(os.environ.get("PORT", 7860))
else:
    default_port = 8861  # Local default port

# Initialize model
model = GenerativeInferenceModel()

# Define example images and their parameters with updated values from the research
examples = [
    {
        "image": os.path.join("stimuli", "farm1.jpg"),
        "name": "farm1",
        "wiki": "https://en.wikipedia.org/wiki/Visual_perception",
        "papers": [
            "[Adversarially Robust Vision](https://github.com/MadryLab/robustness)",
            "[Generative Inference](https://doi.org/10.1016/j.tics.2003.08.003)"
        ],
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "all",
            "initial_noise": 0.0,
            "diffusion_noise": 0.02,
            "step_size": 1.0,
            "iterations": 501,
            "epsilon": 40.0
        },
        "inference_normalization": "off",
        "use_adaptive_eps": False,
        "use_adaptive_step": False,
        "mask_center_x": 0.0,
        "mask_center_y": 0.0,
        "mask_radius": 0.2,
        "mask_sigma": 0.3,
        "eps_max_mult": 300.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 10.0,
        "step_min_mult": 1.0,
    },
    {
        "image": os.path.join("stimuli", "ArtGallery1.jpg"),
        "name": "ArtGallery1",
        "wiki": "https://en.wikipedia.org/wiki/Visual_perception",
        "papers": [
            "[Adversarially Robust Vision](https://github.com/MadryLab/robustness)",
            "[Generative Inference](https://doi.org/10.1016/j.tics.2003.08.003)"
        ],
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "layer4",
            "initial_noise": 0.5,
            "diffusion_noise": 0.002,
            "step_size": 0.1,
            "iterations": 501,
            "epsilon": 40.0
        },
        "inference_normalization": "off",
        "use_adaptive_eps": False,
        "use_adaptive_step": True,
        "mask_center_x": 0.0,
        "mask_center_y": -1.0,
        "mask_radius": 0.1,
        "mask_sigma": 0.2,
        "eps_max_mult": 30.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 100.0,
        "step_min_mult": 1.0,
    },
    {
        "image": os.path.join("stimuli", "urbanoffice1.jpg"),
        "name": "UrbanOffice1",
        "wiki": "https://en.wikipedia.org/wiki/Visual_perception",
        "papers": [
            "[Adversarially Robust Vision](https://github.com/MadryLab/robustness)",
            "[Generative Inference](https://doi.org/10.1016/j.tics.2003.08.003)"
        ],
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "all",
            "initial_noise": 1.0,
            "diffusion_noise": 0.002,
            "step_size": 1.0,
            "iterations": 500,
            "epsilon": 40.0
        },
        "inference_normalization": "off",
        "use_adaptive_eps": False,
        "use_adaptive_step": True,
        "mask_center_x": 0.5,
        "mask_center_y": 0.0,
        "mask_radius": 0.2,
        "mask_sigma": 0.2,
        "eps_max_mult": 20.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 50.0,
        "step_min_mult": 0.2,
    },
    {
        "image": os.path.join("stimuli", "Neon_Color_Circle.jpg"),
        "name": "Neon Color Spreading",
        "wiki": "https://en.wikipedia.org/wiki/Neon_color_spreading",
        "papers": [
            "[Color Assimilation](https://doi.org/10.1016/j.visres.2000.200.1)",
            "[Perceptual Filling-in](https://doi.org/10.1016/j.tics.2003.08.003)"
        ],
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "layer3",
            "initial_noise": 0.8,
            "diffusion_noise": 0.003,
            "step_size": 1.0,
            "iterations": 101,
            "epsilon": 20.0
        },
        "use_adaptive_eps": False,
        "use_adaptive_step": False,
        "mask_center_x": 0.0,
        "mask_center_y": 0.0,
        "mask_radius": 0.2,
        "mask_sigma": 1.0,
        "eps_max_mult": 1.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 1.0,
        "step_min_mult": 1.0,
    },
    {
        "image": os.path.join("stimuli", "Kanizsa_square.jpg"),
        "name": "Kanizsa Square",
        "wiki": "https://en.wikipedia.org/wiki/Kanizsa_triangle",
        "papers": [
            "[Gestalt Psychology](https://en.wikipedia.org/wiki/Gestalt_psychology)",
            "[Neural Mechanisms](https://doi.org/10.1016/j.tics.2003.08.003)"
        ],
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "all",
            "initial_noise": 0.0,
            "diffusion_noise": 0.005,
            "step_size": 0.64,
            "iterations": 100,
            "epsilon": 5.0
        },
        "use_adaptive_eps": False,
        "use_adaptive_step": False,
        "mask_center_x": 0.0,
        "mask_center_y": 0.0,
        "mask_radius": 0.2,
        "mask_sigma": 1.0,
        "eps_max_mult": 1.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 1.0,
        "step_min_mult": 1.0,
    },
    {
        "image": os.path.join("stimuli", "CornsweetBlock.png"),
        "name": "Cornsweet Illusion",
        "wiki": "https://en.wikipedia.org/wiki/Cornsweet_illusion",
        "papers": [
            "[Brightness Perception](https://doi.org/10.1016/j.visres.2000.200.1)",
            "[Edge Effects](https://doi.org/10.1016/j.tics.2003.08.003)"
        ],
        "instructions": "Both blocks are gray in color (the same), use your finger to cover the middle line. Hit 'Load Parameters' and then hit 'Run Generative Inference' to see how the model sees the blocks.",
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "layer3",
            "initial_noise": 0.5,
            "diffusion_noise": 0.005,
            "step_size": 0.8,
            "iterations": 51,
            "epsilon": 20.0
        },
        "use_adaptive_eps": False,
        "use_adaptive_step": False,
        "mask_center_x": 0.0,
        "mask_center_y": 0.0,
        "mask_radius": 0.2,
        "mask_sigma": 1.0,
        "eps_max_mult": 1.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 1.0,
        "step_min_mult": 1.0,
    },
    {
        "image": os.path.join("stimuli", "face_vase.png"),
        "name": "Rubin's Face-Vase (Object Prior)",
        "wiki": "https://en.wikipedia.org/wiki/Rubin_vase",
        "papers": [
            "[Figure-Ground Perception](https://en.wikipedia.org/wiki/Figure-ground_(perception))",
            "[Bistable Perception](https://doi.org/10.1016/j.tics.2003.08.003)"
        ],
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "avgpool",
            "initial_noise": 0.9,
            "diffusion_noise": 0.003,
            "step_size": 0.58,
            "iterations": 100,
            "epsilon": 0.81
        },
        "use_adaptive_eps": False,
        "use_adaptive_step": False,
        "mask_center_x": 0.0,
        "mask_center_y": 0.0,
        "mask_radius": 0.2,
        "mask_sigma": 1.0,
        "eps_max_mult": 1.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 1.0,
        "step_min_mult": 1.0,
    },
    {
        "image": os.path.join("stimuli", "Confetti_illusion.png"),
        "name": "Confetti Illusion",
        "wiki": "https://www.youtube.com/watch?v=SvEiEi8O7QE",
        "papers": [
            "[Color Perception](https://doi.org/10.1016/j.visres.2000.200.1)",
            "[Context Effects](https://doi.org/10.1016/j.tics.2003.08.003)"
        ],
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "layer3",
            "initial_noise": 0.1,
            "diffusion_noise": 0.003,
            "step_size": 0.5,
            "iterations": 101,
            "epsilon": 20.0
        },
        "use_adaptive_eps": False,
        "use_adaptive_step": False,
        "mask_center_x": 0.0,
        "mask_center_y": 0.0,
        "mask_radius": 0.2,
        "mask_sigma": 1.0,
        "eps_max_mult": 1.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 1.0,
        "step_min_mult": 1.0,
    },
    {
        "image": os.path.join("stimuli", "EhresteinSingleColor.png"),
        "name": "Ehrenstein Illusion",
        "wiki": "https://en.wikipedia.org/wiki/Ehrenstein_illusion",
        "papers": [
            "[Subjective Contours](https://doi.org/10.1016/j.visres.2000.200.1)",
            "[Neural Processing](https://doi.org/10.1016/j.tics.2003.08.003)"
        ],
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "layer3",
            "initial_noise": 0.5,
            "diffusion_noise": 0.005,
            "step_size": 0.8,
            "iterations": 101,
            "epsilon": 20.0
        },
        "use_adaptive_eps": False,
        "use_adaptive_step": False,
        "mask_center_x": 0.0,
        "mask_center_y": 0.0,
        "mask_radius": 0.2,
        "mask_sigma": 1.0,
        "eps_max_mult": 1.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 1.0,
        "step_min_mult": 1.0,
    },
    {
        "image": os.path.join("stimuli", "GroupingByContinuity.png"),
        "name": "Grouping by Continuity",
        "wiki": "https://en.wikipedia.org/wiki/Principles_of_grouping",
        "papers": [
            "[Gestalt Principles](https://en.wikipedia.org/wiki/Gestalt_psychology)",
            "[Visual Organization](https://doi.org/10.1016/j.tics.2003.08.003)"
        ],
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "layer3",
            "initial_noise": 0.0,
            "diffusion_noise": 0.005,
            "step_size": 0.4,
            "iterations": 101,
            "epsilon": 4.0
        },
        "use_adaptive_eps": False,
        "use_adaptive_step": False,
        "mask_center_x": 0.0,
        "mask_center_y": 0.0,
        "mask_radius": 0.2,
        "mask_sigma": 1.0,
        "eps_max_mult": 1.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 1.0,
        "step_min_mult": 1.0,
    },
    {
        "image": os.path.join("stimuli", "figure_ground.png"),
        "name": "Figure-Ground Illusion",
        "wiki": "https://en.wikipedia.org/wiki/Figure-ground_(perception)",
        "papers": [
            "[Gestalt Principles](https://en.wikipedia.org/wiki/Gestalt_psychology)",
            "[Perceptual Organization](https://doi.org/10.1016/j.tics.2003.08.003)"
        ],
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust",
            "layer": "layer3",
            "initial_noise": 0.1,
            "diffusion_noise": 0.003,
            "step_size": 0.5,
            "iterations": 101,
            "epsilon": 3.0
        },
        "use_adaptive_eps": False,
        "use_adaptive_step": False,
        "mask_center_x": 0.0,
        "mask_center_y": 0.0,
        "mask_radius": 0.2,
        "mask_sigma": 1.0,
        "eps_max_mult": 1.0,
        "eps_min_mult": 1.0,
        "step_max_mult": 1.0,
        "step_min_mult": 1.0,
    }
]

def _input_image_stem(image):
    """Return a safe filename stem from the input image: known name or 'user_img'."""
    if image is None:
        return "user_img"
    path = None
    if isinstance(image, str) and (os.path.isfile(image) or os.path.exists(image)):
        path = image
    if isinstance(image, dict) and image.get("path") and os.path.exists(image.get("path", "")):
        path = image["path"]
    if path:
        name = os.path.splitext(os.path.basename(path))[0]
        # Safe for filenames: alphanumeric, underscore, hyphen only; max length
        safe = re.sub(r"[^\w\-]", "_", name).strip("_") or "user_img"
        return safe[:80] if len(safe) > 80 else safe
    return "user_img"


def _get_image_path_for_stem(img):
    """Extract file path from Gradio image value (path string, dict with path, or PIL) for stem tracking."""
    if img is None:
        return ""
    if isinstance(img, str) and (os.path.isfile(img) or os.path.exists(img)):
        return img
    if isinstance(img, dict) and img.get("path"):
        p = img["path"]
        if isinstance(p, str) and os.path.exists(p):
            return p
    return ""


def _update_tracked_image_path(img):
    """Keep path only when it's a known stimulus (e.g. from stimuli/); else '' so stem is 'user_img'."""
    path = _get_image_path_for_stem(img)
    if path and "stimuli" in path:
        return path
    return ""


def _config_to_json_serializable(c):
    """Return a copy of config with only JSON-serializable values."""
    if isinstance(c, dict):
        return {k: _config_to_json_serializable(v) for k, v in c.items()}
    if isinstance(c, (list, tuple)):
        return [_config_to_json_serializable(x) for x in c]
    if isinstance(c, (bool, int, float, str, type(None))):
        return c
    if hasattr(c, "item"):  # e.g. numpy scalar
        return c.item()
    return str(c)


@GPU 
def run_inference(image, model_type, inference_type, eps_value, num_iterations, 
                 initial_noise=0.05, diffusion_noise=0.3, step_size=0.8, model_layer="layer3",
                 use_adaptive_eps=False, use_adaptive_step=False,
                 mask_center_x=0.0, mask_center_y=0.0, mask_radius=0.3, mask_sigma=0.2,
                 eps_max_mult=4.0, eps_min_mult=1.0, step_max_mult=4.0, step_min_mult=1.0,
                 use_biased_inference=False, biased_class_name="",
                 current_image_path=""):
    # Check if image is provided
    if image is None:
        return None, [], "Please upload an image before running inference.", None
    
    # Convert eps to float
    eps = float(eps_value)
    step_size_f = float(step_size)
    # Coerce checkbox values (Gradio can send bool, None, or other)
    use_adaptive_eps = bool(use_adaptive_eps) if use_adaptive_eps is not None else False
    use_adaptive_step = bool(use_adaptive_step) if use_adaptive_step is not None else False
    
    # Load inference configuration based on the selected type
    config = get_inference_configs(inference_type=inference_type, eps=eps, n_itr=int(num_iterations))
    
    # Handle Prior-Guided Drift Diffusion specific parameters
    if inference_type == "Prior-Guided Drift Diffusion":
        config['initial_inference_noise_ratio'] = float(initial_noise)
        config['diffusion_noise_ratio'] = float(diffusion_noise)
        config['step_size'] = step_size_f
        config['top_layer'] = model_layer
    
    # Inference normalization off (option removed from UI)
    config['inference_normalization'] = 'off'
    config['recognition_normalization'] = 'off'
    
    # Adaptive epsilon (Gaussian mask)
    if use_adaptive_eps:
        config['adaptive_epsilon'] = {
            'enabled': True,
            'base_epsilon': eps,
            'center_x': float(mask_center_x),
            'center_y': float(mask_center_y),
            'flat_radius': float(mask_radius),
            'sigma': float(mask_sigma),
            'max_multiplier': float(eps_max_mult),
            'min_multiplier': float(eps_min_mult),
        }
    else:
        config['adaptive_epsilon'] = None
    
    # Adaptive step size (same Gaussian mask location/size, different multipliers)
    if use_adaptive_step:
        config['adaptive_step_size'] = {
            'enabled': True,
            'base_step_size': step_size_f,
            'center_x': float(mask_center_x),
            'center_y': float(mask_center_y),
            'flat_radius': float(mask_radius),
            'sigma': float(mask_sigma),
            'max_multiplier': float(step_max_mult),
            'min_multiplier': float(step_min_mult),
        }
    else:
        config['adaptive_step_size'] = None
    
    # Biased inference: bias perception toward a target ImageNet class
    use_biased_inference = bool(use_biased_inference) if use_biased_inference is not None else False
    biased_class_name = (biased_class_name or "").strip() if biased_class_name else ""
    if use_biased_inference and biased_class_name:
        config['biased_inference'] = {'enable': True, 'class': biased_class_name}
    else:
        config['biased_inference'] = config.get('biased_inference') or {'enable': False, 'class': None}
    
    # Run generative inference
    result = model.inference(image, model_type, config)
    
    # Extract results based on return type
    if isinstance(result, tuple):
        # Old format returning (output_image, all_steps)
        output_image, all_steps = result
    else:
        # New format returning dictionary
        output_image = result['final_image']
        all_steps = result['steps']
    
    # Create animation frames
    frames = []
    for i, step_image in enumerate(all_steps):
        # Convert tensor to PIL image
        step_pil = Image.fromarray((step_image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
        frames.append(step_pil)
    
    # Convert the final output image to PIL
    final_image = Image.fromarray((output_image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))
    
    # Always save GIF and config and offer as downloads (browser will ask where to save)
    save_status = ""
    files_for_download = None
    if frames:
        # Use tracked path when available (e.g. from Load Parameters); else derive from image (PIL loses path)
        stem = _input_image_stem(current_image_path if (current_image_path and current_image_path.strip()) else image)
        unique_id = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{stem}"
        gif_path = os.path.join(SAVED_RUNS_DIR, f"{unique_id}.gif")
        config_path = os.path.join(SAVED_RUNS_DIR, f"{unique_id}_config.json")
        try:
            frames[0].save(
                gif_path,
                save_all=True,
                append_images=frames[1:],
                loop=0,
                duration=200,
            )
            save_config = {
                "model_type": model_type,
                "input_image_name": stem,
                **_config_to_json_serializable(config),
            }
            with open(config_path, "w") as f:
                json.dump(save_config, f, indent=2)
            files_for_download = [gif_path, config_path]
            save_status = "**Download results** — Use the links below to save the GIF and config to your device (your browser may ask where to save)."
        except Exception as e:
            save_status = f"Save failed: {e}"
    
    return final_image, frames, save_status, files_for_download

def _image_to_pil(img):
    """Convert Gradio image value (PIL, numpy, path, or dict) to PIL Image; return None if invalid."""
    if img is None:
        return None
    if isinstance(img, Image.Image):
        return img
    if isinstance(img, np.ndarray):
        return Image.fromarray(img.astype(np.uint8))
    if isinstance(img, dict) and "path" in img:
        return Image.open(img["path"]).convert("RGB")
    if isinstance(img, str) and os.path.exists(img):
        return Image.open(img).convert("RGB")
    return None


def _image_size(img):
    """Return (width, height) from Gradio image value, or (1, 1) if unknown."""
    pil = _image_to_pil(img)
    if pil is not None:
        return pil.size
    return (1, 1)


def image_click_to_center(evt: gr.SelectData):
    """Convert click (x, y) on image to normalized mask center (-1 to 1). Returns (center_x, center_y)."""
    if evt.index is None or (isinstance(evt.index, (list, tuple)) and len(evt.index) < 2):
        return 0.0, 0.0
    x, y = float(evt.index[0]), float(evt.index[1])
    w, h = _image_size(evt.value)
    if w <= 0 or h <= 0:
        return 0.0, 0.0
    center_x = (x / w) * 2.0 - 1.0
    center_y = (y / h) * 2.0 - 1.0
    center_x = max(-1.0, min(1.0, center_x))
    center_y = max(-1.0, min(1.0, center_y))
    return center_x, center_y


def draw_mask_overlay(image, center_x, center_y, radius):
    """Draw the Gaussian mask center and radius on a copy of the image. Returns PIL or None."""
    pil = _image_to_pil(image)
    if pil is None:
        return None
    img = pil.convert("RGB").copy()
    w, h = img.size
    cx_px = (float(center_x) + 1.0) / 2.0 * w
    cy_px = (float(center_y) + 1.0) / 2.0 * h
    radius_px = float(radius) * min(w, h) / 2.0
    draw = ImageDraw.Draw(img)
    # Circle for radius
    draw.ellipse(
        [cx_px - radius_px, cy_px - radius_px, cx_px + radius_px, cy_px + radius_px],
        outline="#E11D48",
        width=2 * max(2, min(w, h) // 150),
    )
    # Center dot
    r = max(2, min(w, h) // 80)
    draw.ellipse([cx_px - r, cy_px - r, cx_px + r, cy_px + r], fill="#E11D48", outline="#FFF")
    return img


# Helper function to apply example parameters (adaptive mask off by default unless example defines it)
def apply_example(example):
    rd = example["reverse_diff"]
    mcx = example.get("mask_center_x", 0.0)
    mcy = example.get("mask_center_y", 0.0)
    mrad = example.get("mask_radius", 0.3)
    mask_img = draw_mask_overlay(example["image"], mcx, mcy, mrad)
    return [
        example["image"],
        rd.get("model", "resnet50_robust"),
        example["method"],
        rd["epsilon"],
        rd["iterations"],
        rd["initial_noise"],
        rd["diffusion_noise"],
        rd["step_size"],
        rd["layer"],
        example.get("use_adaptive_eps", False),
        example.get("use_adaptive_step", False),
        mcx,
        mcy,
        example.get("mask_radius", 0.3),
        example.get("mask_sigma", 0.2),
        example.get("eps_max_mult", 4.0),
        example.get("eps_min_mult", 1.0),
        example.get("step_max_mult", 4.0),
        example.get("step_min_mult", 1.0),
        example.get("use_biased_inference", False),
        example.get("biased_class_name", ""),
        example["image"],  # keep path for save filename (e.g. UrbanOffice1 -> urbanoffice1)
        mask_img,
        gr.Group(visible=True),
    ]

# Define the interface
with gr.Blocks(title="Human Hallucination Prediction", css="""
.purple-button {
    background-color: #8B5CF6 !important;
    color: white !important;
    border: none !important;
}
.purple-button:hover {
    background-color: #7C3AED !important;
}
""") as demo:
    gr.Markdown("# Human Hallucination Prediction")
    gr.Markdown("**Predict what visual hallucinations humans may experience** using neural networks.")
    
    gr.Markdown("""
    **How to predict hallucinations:**
    1. **Select an example image** below and click "Load Parameters" to set the prediction settings
    2. **Click "Predict Hallucinations"** to predict what hallucination humans may perceive
    3. **View the prediction**: Watch as the model reveals the perceptual structures it expects—matching what humans typically hallucinate
    4. **You can upload your own images**
    5. **You can download the results** as a .gif file together with the configs.json
    """)
    with gr.Row():
        with gr.Column(scale=1):
            # Inputs (track path so save filenames use stimulus name when from example)
            default_image_path = os.path.join("stimuli", "urbanoffice1.jpg")
            image_input = gr.Image(label="Input Image (click to set mask center)", type="pil", value=default_image_path)
            current_image_path_state = gr.State(value=default_image_path)
            mask_preview = gr.Image(
                label="Mask center preview (click to set center — circle shows mask)",
                type="pil",
                interactive=False,
            )
            # Run Inference button right below the image
            run_button = gr.Button("🔮 Predict Hallucination", variant="primary", elem_classes="purple-button")
            
            # Parameters toggle button
            params_button = gr.Button("⚙️ Play with the parameters", variant="secondary")
            
            # Parameters section (initially hidden)
            with gr.Group(visible=False) as params_section:
                with gr.Row():
                    model_choice = gr.Dropdown(
                        choices=["resnet50_robust", "standard_resnet50"], # "resnet50_robust_face" - hidden for deployment
                        value="resnet50_robust", 
                        label="Model"
                    )
                    
                    inference_type = gr.Dropdown(
                        choices=["Prior-Guided Drift Diffusion", "IncreaseConfidence"], 
                        value="Prior-Guided Drift Diffusion", 
                        label="Inference Method"
                    )
                
                with gr.Row():
                    eps_slider = gr.Slider(minimum=0.0, maximum=40.0, value=40.0, step=0.01, label="Epsilon (Stimulus Fidelity)")
                    iterations_slider = gr.Slider(minimum=1, maximum=600, value=500, step=1, label="Number of Iterations")
                
                with gr.Row():
                    initial_noise_slider = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.01, 
                                                   label="Drift Noise")
                    diffusion_noise_slider = gr.Slider(minimum=0.0, maximum=0.05, value=0.002, step=0.001, 
                                                    label="Diffusion Noise")
                    
                with gr.Row():
                    step_size_slider = gr.Slider(minimum=0.01, maximum=2.0, value=1.0, step=0.01, 
                                               label="Update Rate")
                    layer_choice = gr.Dropdown(
                        choices=["all", "conv1", "bn1", "relu", "maxpool", "layer1", "layer2", "layer3", "layer4", "avgpool"],
                        value="all",
                        label="Model Layer"
                    )
                
                gr.Markdown("### 🎯 Adaptive Gaussian mask (spatially varying constraint)")
                gr.Markdown("Define where on the image the mask is centered and how large its radius is. Coordinates: **-1** = left/top, **1** = right/bottom, **0** = center.")
                with gr.Row():
                    use_adaptive_eps_check = gr.Checkbox(value=False, label="Use adaptive epsilon (stronger/weaker constraint by region)")
                    use_adaptive_step_check = gr.Checkbox(value=True, label="Use adaptive step size (stronger/weaker updates by region)")
                with gr.Row():
                    mask_center_x_slider = gr.Slider(minimum=-1.0, maximum=1.0, value=0.5, step=0.05, label="Mask center X")
                    mask_center_y_slider = gr.Slider(minimum=-1.0, maximum=1.0, value=0.0, step=0.05, label="Mask center Y")
                with gr.Row():
                    mask_radius_slider = gr.Slider(minimum=0.01, maximum=1.0, value=0.2, step=0.01, label="Mask radius (flat region size)")
                    mask_sigma_slider = gr.Slider(minimum=0.05, maximum=1.0, value=0.2, step=0.01, label="Mask sigma (fall-off outside radius)")
                with gr.Row():
                    eps_max_mult_slider = gr.Slider(minimum=0.1, maximum=350.0, value=20.0, step=0.1, label="Epsilon: multiplier at center")
                    eps_min_mult_slider = gr.Slider(minimum=0.1, maximum=10.0, value=1.0, step=0.1, label="Epsilon: multiplier at periphery")
                with gr.Row():
                    step_max_mult_slider = gr.Slider(minimum=0.1, maximum=150.0, value=50.0, step=0.1, label="Step size: multiplier at center")
                    step_min_mult_slider = gr.Slider(minimum=0.1, maximum=10.0, value=0.2, step=0.1, label="Step size: multiplier at periphery")
                gr.Markdown("### 🎯 Biased inference")
                gr.Markdown("Bias the prediction toward a specific ImageNet category (1000 classes).")
                with gr.Row():
                    use_biased_inference_check = gr.Checkbox(value=False, label="Use biased inference (bias toward a target class)")
                    biased_class_dropdown = gr.Dropdown(
                        choices=[("— No bias —", "")] + [(label, label) for label in sorted(IMAGENET_LABELS)],
                        value="",
                        label="Biased toward category",
                        allow_custom_value=False,
                        filterable=True,
                    )
        with gr.Column(scale=2):
            # Outputs
            output_image = gr.Image(label="Predicted Hallucination")
            output_frames = gr.Gallery(label="Hallucination Prediction Process", columns=5, rows=2)
            save_status_md = gr.Markdown(value="")
            download_files = gr.File(label="Download results (GIF + config)", file_count="multiple")
    
    # Examples section with integrated explanations
    gr.Markdown("## Examples")
    gr.Markdown("Select an example and click Load Parameters to apply its settings")
    
    # For each example, create a row with the image and explanation side by side
    for i, ex in enumerate(examples):
        with gr.Row():
            # Left column for the image
            with gr.Column(scale=1):
                # Display the example image
                example_img = gr.Image(value=ex["image"], type="filepath", label=f"{ex['name']}")
                load_btn = gr.Button(f"Load Parameters", variant="primary")
                
                # Set up the load button to apply this example's parameters
                load_btn.click(
                    fn=lambda ex=ex: apply_example(ex),
                    outputs=[
                        image_input, model_choice, inference_type,
                        eps_slider, iterations_slider,
                        initial_noise_slider, diffusion_noise_slider,
                        step_size_slider, layer_choice,
                        use_adaptive_eps_check, use_adaptive_step_check,
                        mask_center_x_slider, mask_center_y_slider,
                        mask_radius_slider, mask_sigma_slider,
                        eps_max_mult_slider, eps_min_mult_slider,
                        step_max_mult_slider, step_min_mult_slider,
                        use_biased_inference_check, biased_class_dropdown,
                        current_image_path_state,
                        mask_preview,
                        params_section,
                    ],
                )
            
            # Right column for the explanation
            with gr.Column(scale=2):
                gr.Markdown(f"### {ex['name']}")
                if ex["name"] not in ("farm1", "ArtGallery1", "UrbanOffice1"):
                    gr.Markdown(f"[Read more on Wikipedia]({ex['wiki']})")
                
                # Show instructions if they exist
                if "instructions" in ex:
                    gr.Markdown(f"**Instructions:** {ex['instructions']}")
                
        
        if i < len(examples) - 1:  # Don't add separator after the last example
            gr.Markdown("---")
    
    # Set up event handler for the main inference
    run_button.click(
        fn=run_inference,
        inputs=[
            image_input, model_choice, inference_type,
            eps_slider, iterations_slider,
            initial_noise_slider, diffusion_noise_slider,
            step_size_slider, layer_choice,
            use_adaptive_eps_check, use_adaptive_step_check,
            mask_center_x_slider, mask_center_y_slider,
            mask_radius_slider, mask_sigma_slider,
            eps_max_mult_slider, eps_min_mult_slider,
            step_max_mult_slider, step_min_mult_slider,
            use_biased_inference_check, biased_class_dropdown,
            current_image_path_state,
        ],
        outputs=[output_image, output_frames, save_status_md, download_files]
    )
    
    # Toggle parameters visibility
    def toggle_params():
        return gr.Group(visible=True)

    params_button.click(
        fn=toggle_params,
        outputs=[params_section],
    )

    # Click on input image or mask preview to set Gaussian mask center
    def _mask_preview_inputs():
        return [image_input, mask_center_x_slider, mask_center_y_slider, mask_radius_slider]

    image_input.select(
        fn=image_click_to_center,
        outputs=[mask_center_x_slider, mask_center_y_slider],
    )
    mask_preview.select(
        fn=image_click_to_center,
        outputs=[mask_center_x_slider, mask_center_y_slider],
    )
    # Update mask preview when image or center/radius change
    image_input.change(
        fn=draw_mask_overlay,
        inputs=_mask_preview_inputs(),
        outputs=[mask_preview],
    )
    # Keep tracked path for save filename: known stimulus name or clear so stem becomes 'user_img'
    image_input.change(
        fn=_update_tracked_image_path,
        inputs=[image_input],
        outputs=[current_image_path_state],
    )
    mask_center_x_slider.change(
        fn=draw_mask_overlay,
        inputs=_mask_preview_inputs(),
        outputs=[mask_preview],
    )
    mask_center_y_slider.change(
        fn=draw_mask_overlay,
        inputs=_mask_preview_inputs(),
        outputs=[mask_preview],
    )
    mask_radius_slider.change(
        fn=draw_mask_overlay,
        inputs=_mask_preview_inputs(),
        outputs=[mask_preview],
    )
    # Populate mask preview on load
    demo.load(
        fn=draw_mask_overlay,
        inputs=_mask_preview_inputs(),
        outputs=[mask_preview],
    )
    
    # About section
    gr.Markdown("""
    ## 🧠 About Hallucination Prediction
    
    This tool predicts human visual hallucinations using **generative inference** with adversarially robust neural networks. Robust models develop human-like perceptual biases, allowing them to forecast what perceptual structures humans will experience.
    
    ### Prediction Methods:
    
    **Prior-Guided Drift Diffusion** (Primary Method)  
    Starting from a noisy representation, the model converges toward what it expects to perceive—revealing predicted hallucinations
    
    **IncreaseConfidence**  
    Moving away from unlikely interpretations to reveal the most probable perceptual experience
    
    ### Parameters:
    - **Drift Noise**: Initial uncertainty in the prediction process
    - **Diffusion Noise**: Stochastic exploration during prediction
    - **Update Rate**: Speed of convergence to the predicted hallucination
    - **Number of Iterations**: How many prediction steps to perform
    - **Model Layer**: Which perceptual level to predict from (early edges vs. high-level objects)
    - **Epsilon (Stimulus Fidelity)**: How closely the prediction must match the input stimulus
    
    ### Why Does This Work?
    
    Adversarially robust neural networks develop perceptual representations similar to human vision. When we use generative inference to reveal what these networks "expect" to see, it matches what humans hallucinate in ambiguous images—allowing us to predict human perception.
    
    **Developed by [Tahereh Toosi](https://toosi.github.io)**
    """)

# Launch the demo
if __name__ == "__main__":
    print(f"Starting server on port {args.port}")
    demo.launch(
        server_name="0.0.0.0",
        server_port=args.port,
        share=False,
        debug=True
    )