app.py

import gradio as gr
import torch
import numpy as np
from PIL import Image
try:
    from spaces import GPU
    print("Running on Hugging Face Spaces - GPU decorator available")
except ImportError:
    # Define a no-op decorator if running locally
    def GPU(func):
        """No-op decorator for local execution (GPU handling is automatic)"""
        return func
    print("Running locally - GPU decorator not available (using automatic GPU detection)")
    
import os
import argparse
from inference import GenerativeInferenceModel, get_inference_configs

# Parse command line arguments
parser = argparse.ArgumentParser(description='Run Generative Inference Demo')
parser.add_argument('--port', type=int, default=None, 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)

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

# Use command line port if provided, otherwise use default
server_port = args.port if args.port is not None else default_port

# Initialize model (lazy initialization to avoid startup failures)
try:
    model = GenerativeInferenceModel()
    print("Model manager initialized successfully")
except Exception as e:
    print(f"Warning: Error initializing model manager: {e}")
    print("Will attempt to initialize on first use")
    model = None

# Define example images and their parameters with updated values from the research
examples = [
    {
        "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_face",
            "layer": "layer4",
            "initial_noise": 0.0,
            "diffusion_noise": 0.006,
            "step_size": 0.18,
            "iterations": 100,
            "epsilon": 9.53
        }
    },
    {
        "image": os.path.join("stimuli", "RandomizedPhaseOvalGray.png"),
        "name": "Noise (Randomized Phase Oval)",
        "wiki": "https://en.wikipedia.org/wiki/Visual_noise",
        "papers": [
            "[Perceptual Organization](https://doi.org/10.1016/j.tics.2003.08.003)",
            "[Pattern Recognition](https://en.wikipedia.org/wiki/Pattern_recognition)"
        ],
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust_face",
            "layer": "all",
            "initial_noise": 0.0,
            "diffusion_noise": 0.05,
            "step_size": 1.12,
            "iterations": 428,
            "epsilon": 198.62
        }
    },
    {
        "image": os.path.join("stimuli", "Mooney_face1.png"),
        "name": "Mooney Face",
        "wiki": "https://en.wikipedia.org/wiki/Mooney_faces",
        "papers": [
            "[Face Recognition](https://en.wikipedia.org/wiki/Face_perception)",
            "[Perceptual Organization](https://doi.org/10.1016/j.tics.2003.08.003)"
        ],
        "method": "Prior-Guided Drift Diffusion",
        "reverse_diff": {
            "model": "resnet50_robust_face",
            "layer": "layer4",
            "initial_noise": 0.0,
            "diffusion_noise": 0.006,
            "step_size": 0.18,
            "iterations": 100,
            "epsilon": 9.53
        }
    },
    {
        "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
        }
    },
    {
        "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
        }
    },
    {
        "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
        }
    },
    {
        "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
        }
    },
    {
        "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
        }
    },
    {
        "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
        }
    },
    {
        "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
        }
    }
]

@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"):
    # Initialize model if not already initialized
    global model
    if model is None:
        try:
            model = GenerativeInferenceModel()
            print("Model manager initialized on first use")
        except Exception as e:
            return None, f"Error initializing model: {str(e)}. Please try again."
    
    # Check if image is provided
    if image is None:
        return None, "Please upload an image before running inference."
    
    # Convert eps to float
    eps = float(eps_value)
    
    # 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'] = float(step_size)  # Added step size parameter
        config['top_layer'] = model_layer
    
    # 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 with proper error handling
        try:
            # Ensure tensor is on CPU and detached
            if isinstance(step_image, torch.Tensor):
                step_image = step_image.detach().cpu()
                # Handle different tensor shapes
                if len(step_image.shape) == 4:  # [B, C, H, W]
                    step_image = step_image[0]  # Take first batch item
                elif len(step_image.shape) == 3:  # [C, H, W]
                    pass  # Already correct shape
                else:
                    raise ValueError(f"Unexpected tensor shape: {step_image.shape}")
                
                # Clamp values to [0, 1] range before converting
                step_image = torch.clamp(step_image, 0, 1)
                # Convert to numpy and ensure contiguous array
                step_np = step_image.permute(1, 2, 0).numpy()
                # Ensure it's a contiguous array with correct dtype
                step_np = np.ascontiguousarray(step_np, dtype=np.float32)
                # Convert to uint8
                step_np = (step_np * 255).astype(np.uint8)
                # Create PIL image
                step_pil = Image.fromarray(step_np, mode='RGB')
                frames.append(step_pil)
            else:
                print(f"Warning: step_image at index {i} is not a tensor: {type(step_image)}")
        except Exception as e:
            print(f"Error converting step {i} to PIL image: {e}, shape: {step_image.shape if hasattr(step_image, 'shape') else 'N/A'}")
            # Skip this frame if conversion fails
            continue
    
    # Convert the final output image to PIL
    try:
        if isinstance(output_image, torch.Tensor):
            output_image = output_image.detach().cpu()
            # Handle different tensor shapes
            if len(output_image.shape) == 4:  # [B, C, H, W]
                output_image = output_image[0]  # Take first batch item
            elif len(output_image.shape) == 3:  # [C, H, W]
                pass  # Already correct shape
            else:
                raise ValueError(f"Unexpected tensor shape: {output_image.shape}")
            
            # Clamp values to [0, 1] range before converting
            output_image = torch.clamp(output_image, 0, 1)
            # Convert to numpy and ensure contiguous array
            output_np = output_image.permute(1, 2, 0).numpy()
            # Ensure it's a contiguous array with correct dtype
            output_np = np.ascontiguousarray(output_np, dtype=np.float32)
            # Convert to uint8
            output_np = (output_np * 255).astype(np.uint8)
            # Create PIL image
            final_image = Image.fromarray(output_np, mode='RGB')
        else:
            raise ValueError(f"output_image is not a tensor: {type(output_image)}")
    except Exception as e:
        print(f"Error converting final image to PIL: {e}, shape: {output_image.shape if hasattr(output_image, 'shape') else 'N/A'}")
        # Return a black image as fallback
        final_image = Image.new('RGB', (224, 224), color='black')
    
    # Return the final inferred image and the animation frames directly
    return final_image, frames

# Helper function to apply example parameters
def apply_example(example):
    # Get the full path to the image file
    image_path = os.path.abspath(example["image"]) if os.path.exists(example["image"]) else example["image"]
    return [
        image_path,
        example["reverse_diff"]["model"],  # Model type from example
        example["method"],  # Inference type
        example["reverse_diff"]["epsilon"],  # Epsilon value
        example["reverse_diff"]["iterations"],  # Number of iterations
        example["reverse_diff"]["initial_noise"],  # Initial noise
        example["reverse_diff"]["diffusion_noise"],  # Diffusion noise value (corrected)
        example["reverse_diff"]["step_size"],  # Step size (added)
        example["reverse_diff"]["layer"],  # Model layer
        gr.Group(visible=True)  # Show parameters section
    ]

# Define the interface
with gr.Blocks(title="Generative Inference for Psychiatry Demo (in development, not ready for public use)", css="""
.purple-button {
    background-color: #8B5CF6 !important;
    color: white !important;
    border: none !important;
}
.purple-button:hover {
    background-color: #7C3AED !important;
}
""") as demo:
    gr.Markdown("# Generative Inference for Psychiatry Demo (in development, not ready for public use)")
    gr.Markdown("This demo showcases how neural networks can perceive visual illusions and develop Gestalt principles of perceptual organization through generative inference.")
    
    gr.Markdown("""
    **How to use this demo:**
    - **Load pre-configured examples**: Click on any visual illusion below and hit "Load Parameters" to automatically set up the optimal parameters for that illusion
    - **Run the inference**: After loading parameters or setting your own, hit "Run Inference" to start the generative inference process
    - **You can also upload your own images** and experiment with different parameters to see how they affect the generative inference process
    """)
    
    # Main processing interface
    with gr.Row():
        with gr.Column(scale=1):
            # Inputs
            # Use absolute path for default image to avoid directory errors
            default_image_path = os.path.abspath(os.path.join("stimuli", "face_vase.png")) if os.path.exists(os.path.join("stimuli", "face_vase.png")) else None
            image_input = gr.Image(label="Input Image", type="pil", value=default_image_path)
            
            # Run Inference button right below the image
            run_button = gr.Button("🪄 Run Generative Inference", 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"], # "resnet50_robust_face" - hidden for deployment
                        value="resnet50_robust_face", 
                        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=200.0, value=9.53, step=0.01, label="Epsilon (Stimulus Fidelity)")
                    iterations_slider = gr.Slider(minimum=1, maximum=600, value=100, step=1, label="Number of Iterations")  # Updated max to 600
                
                with gr.Row():
                    initial_noise_slider = gr.Slider(minimum=0.0, maximum=5.0, value=0.0, step=0.01, 
                                                   label="Drift Noise")
                    diffusion_noise_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.006, step=0.001, 
                                                    label="Diffusion Noise")  # Corrected name
                    
                with gr.Row():
                    step_size_slider = gr.Slider(minimum=0.0, maximum=10.0, value=0.18, step=0.01, 
                                               label="Update Rate")  # Added step size slider
                    layer_choice = gr.Dropdown(
                        choices=["all", "conv1", "bn1", "relu", "maxpool", "layer1", "layer2", "layer3", "layer4", "avgpool"],
                        value="layer4",
                        label="Model Layer"
                    )
            
        with gr.Column(scale=2):
            # Outputs
            output_image = gr.Image(label="Final Inferred Image")
            output_frames = gr.Gallery(label="Inference Steps", columns=5, rows=2)
    
    # Examples section with integrated explanations
    gr.Markdown("## Visual Illusion Examples")
    gr.Markdown("Select an illusion to load its parameters and see how generative inference reveals perceptual effects")
    
    # 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, params_section
                    ]
                )
            
            # Right column for the explanation
            with gr.Column(scale=2):
                gr.Markdown(f"### {ex['name']}")
                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
        ],
        outputs=[output_image, output_frames]
    )
    
    # Toggle parameters visibility
    def toggle_params():
        return gr.Group(visible=True)
    
    params_button.click(
        fn=toggle_params,
        outputs=[params_section]
    )
    
    # About section
    gr.Markdown("""
    ## About Generative Inference
    
    Generative inference is a technique that reveals how neural networks perceive visual stimuli. This demo primarily uses the Prior-Guided Drift Diffusion method.
    
    ### Prior-Guided Drift Diffusion
    Moving away from a noisy representation of the input images
    
    ### IncreaseConfidence
    Moving away from the least likely class identified at iteration 0 (fast perception)
    
    ### Parameters:
    - **Drift Noise**: Controls the amount of noise added to the image at the beginning
    - **Diffusion Noise**: Controls the amount of noise added at each optimization step
    - **Update Rate**: Learning rate for the optimization process
    - **Number of Iterations**: How many optimization steps to perform
    - **Model Layer**: Select a specific layer of the ResNet50 model to extract features from
    - **Epsilon (Stimulus Fidelity)**: Controls the size of perturbation during optimization
    
    **Generative Inference was developed by [Tahereh Toosi](https://toosi.github.io).**
    """)

# Launch the demo
if __name__ == "__main__":
    print(f"Starting server on port {server_port}")
    # On Hugging Face Spaces, don't specify server_name/server_port
    if "SPACE_ID" in os.environ:
        demo.launch(share=False, debug=False)
    else:
        demo.launch(
            server_name="0.0.0.0",
            server_port=server_port,
            share=False,
            debug=True
        )