Upload 23 files

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+stimuli/Confetti_illusion.png filter=lfs diff=lfs merge=lfs -text
+stimuli/CornsweetBlock.png filter=lfs diff=lfs merge=lfs -text
+stimuli/EhresteinSingleColor.png filter=lfs diff=lfs merge=lfs -text
+stimuli/figure_ground.png filter=lfs diff=lfs merge=lfs -text
+stimuli/Neon_Color_Circle.jpg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 GenerativeInferenceDemo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE. 

README.md

@@ -1,13 +1,83 @@
 ---
-title: Generative Inference Faces
-emoji: 🚀
-colorFrom: yellow
-colorTo: indigo
+title: Generative Inference Demo
+emoji: 🧠
+colorFrom: indigo
+colorTo: purple
 sdk: gradio
-sdk_version: 5.49.1
+sdk_version: 5.23.1
 app_file: app.py
 pinned: false
-short_description: Generative inference on models trained for faces
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Generative Inference Demo
+
+This Gradio demo showcases how neural networks perceive visual illusions through generative inference. The demo uses both standard and robust ResNet50 models to reveal emergent perception of contours, figure-ground separation, and other visual phenomena.
+
+## Models
+
+- **Robust ResNet50**: A model trained with adversarial examples (ε=3.0), exhibiting more human-like visual perception
+- **Standard ResNet50**: A model trained without adversarial examples (ε=0.0)
+
+## Features
+
+- Upload your own images or use example illusions
+- Choose between robust and standard models
+- Adjust perturbation size (epsilon) and iteration count
+- Visualize how perception emerges over time
+- Includes classic illusions:
+  - Kanizsa shapes
+  - Face-Vase illusions
+  - Figure-Ground segmentation
+  - Neon color spreading
+
+## Usage
+
+1. Select an example image or upload your own
+2. Choose the model type (robust or standard)
+3. Adjust epsilon and iteration parameters
+4. Click "Run Inference" to see how the model perceives the image
+
+## About
+
+This demo is based on research showing how adversarially robust models develop more human-like visual representations. The generative inference process reveals these perceptual biases by optimizing the input to maximize the model's confidence.
+
+## Installation
+
+To run this demo locally:
+
+```bash
+# Clone the repository
+git clone [repo-url]
+cd GenerativeInferenceDemo
+
+# Install dependencies
+pip install -r requirements.txt
+
+# Run the app
+python app.py
+```
+
+The web app will be available at http://localhost:7860 (or another port if 7860 is busy).
+
+## About the Models
+
+- **Robust ResNet50**: A model trained with adversarial examples, making it more robust to small perturbations. These models often exhibit more human-like visual perception.
+- **Standard ResNet50**: A standard ImageNet-trained ResNet50 model.
+
+## How It Works
+
+1. The algorithm starts with an input image
+2. It iteratively updates the image to increase the model's confidence in its predictions
+3. These updates are constrained to a small neighborhood (controlled by epsilon) around the original image
+4. The resulting changes reveal how the network "sees" the image
+
+## Citation
+
+If you use this work in your research, please cite the original paper:
+
+[Citation information will be added here]
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.

app.py

@@ -1,7 +1,412 @@
 import gradio as gr
+import torch
+import numpy as np
+from PIL import Image
+try:
+    from spaces import GPU
+except ImportError:
+    # Define a no-op decorator if running locally
+    def GPU(func):
+        return func
+    
+import os
+import argparse
+from inference import GenerativeInferenceModel, get_inference_configs
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Parse command line arguments
+parser = argparse.ArgumentParser(description='Run Generative Inference Demo')
+parser.add_argument('--port', type=int, default=7860, help='Port to run the server on')
+args = parser.parse_args()
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# 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 = 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", "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", "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
+        }
+    },
+    {
+        "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"):
+    # 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
+        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))
+    
+    # Return the final inferred image and the animation frames directly
+    return final_image, frames
+
+# Helper function to apply example parameters
+def apply_example(example):
+    return [
+        example["image"],
+        "resnet50_robust",  # Model type
+        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 Demo", 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 Demo")
+    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
+            image_input = gr.Image(label="Input Image", type="pil", value=os.path.join("stimuli", "Neon_Color_Circle.jpg"))
+            
+            # 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" - 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=20.0, step=0.01, label="Epsilon (Stimulus Fidelity)")
+                    iterations_slider = gr.Slider(minimum=1, maximum=600, value=101, step=1, label="Number of Iterations")  # Updated max to 600
+                
+                with gr.Row():
+                    initial_noise_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.8, step=0.01, 
+                                                   label="Drift Noise")
+                    diffusion_noise_slider = gr.Slider(minimum=0.0, maximum=0.05, value=0.003, step=0.001, 
+                                                    label="Diffusion Noise")  # Corrected name
+                    
+                with gr.Row():
+                    step_size_slider = gr.Slider(minimum=0.01, maximum=2.0, value=1.0, 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="layer3",
+                        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 {args.port}")
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=args.port,
+        share=False,
+        debug=True
+    )

huggingface-metadata.json

@@ -0,0 +1,11 @@
+{
+  "title": "Generative Inference Demo",
+  "emoji": "🧠",
+  "colorFrom": "indigo",
+  "colorTo": "purple",
+  "sdk": "gradio",
+  "sdk_version": "3.32.0",
+  "app_file": "app.py",
+  "pinned": false,
+  "license": "mit"
+} 

inference.py

@@ -0,0 +1,923 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+import torchvision.transforms as transforms
+from torchvision.models.resnet import ResNet50_Weights
+from PIL import Image
+import numpy as np
+import os
+import requests
+import time
+import copy
+from collections import OrderedDict
+from pathlib import Path
+
+# Check for available hardware acceleration
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    device = torch.device("mps")  # Use Apple Metal Performance Shaders for M-series Macs
+else:
+    device = torch.device("cpu")
+print(f"Using device: {device}")
+
+# Constants
+MODEL_URLS = {
+    'resnet50_robust': 'https://huggingface.co/madrylab/robust-imagenet-models/resolve/main/resnet50_l2_eps3.ckpt',
+    'resnet50_standard': 'https://huggingface.co/madrylab/robust-imagenet-models/resolve/main/resnet50_l2_eps0.ckpt',
+    'resnet50_robust_face': 'https://huggingface.co/ttoosi/resnet50_robust_face/blob/main/100_checkpoint.pt'
+}
+
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
+
+# Define the transforms based on whether normalization is on or off
+def get_transform(input_size=224, normalize=False, norm_mean=IMAGENET_MEAN, norm_std=IMAGENET_STD):
+    if normalize:
+        return transforms.Compose([
+            transforms.Resize(input_size),
+            transforms.CenterCrop(input_size),
+            transforms.ToTensor(),
+            transforms.Normalize(norm_mean, norm_std),
+        ])
+    else:
+        return transforms.Compose([
+            transforms.Resize(input_size),
+            transforms.CenterCrop(input_size),
+            transforms.ToTensor(),
+        ])
+
+# Default transform without normalization
+transform = transforms.Compose([
+    transforms.Resize(224),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+])
+
+normalize_transform = transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD)
+
+def extract_middle_layers(model, layer_index):
+    """
+    Extract a subset of the model up to a specific layer.
+    
+    Args:
+        model: The neural network model
+        layer_index: String 'all' for the full model, or a layer identifier (string or int)
+                    For ResNet: integers 0-8 representing specific layers
+                    For ViT: strings like 'encoder.layers.encoder_layer_3'
+    
+    Returns:
+        A modified model that outputs features from the specified layer
+    """
+    if isinstance(layer_index, str) and layer_index == 'all':
+        return model
+    
+    # Special case for ViT's encoder layers with DataParallel wrapper
+    if isinstance(layer_index, str) and layer_index.startswith('encoder.layers.encoder_layer_'):
+        try:
+            target_layer_idx = int(layer_index.split('_')[-1])
+            
+            # Create a deep copy of the model to avoid modifying the original
+            new_model = copy.deepcopy(model)
+            
+            # For models wrapped in DataParallel
+            if hasattr(new_model, 'module'):
+                # Create a subset of encoder layers up to the specified index
+                encoder_layers = nn.Sequential()
+                for i in range(target_layer_idx + 1):
+                    layer_name = f"encoder_layer_{i}"
+                    if hasattr(new_model.module.encoder.layers, layer_name):
+                        encoder_layers.add_module(layer_name, 
+                                               getattr(new_model.module.encoder.layers, layer_name))
+                
+                # Replace the encoder layers with our truncated version
+                new_model.module.encoder.layers = encoder_layers
+                
+                # Remove the heads since we're stopping at the encoder layer
+                new_model.module.heads = nn.Identity()
+                
+                return new_model
+            else:
+                # Direct model access (not DataParallel)
+                encoder_layers = nn.Sequential()
+                for i in range(target_layer_idx + 1):
+                    layer_name = f"encoder_layer_{i}"
+                    if hasattr(new_model.encoder.layers, layer_name):
+                        encoder_layers.add_module(layer_name, 
+                                               getattr(new_model.encoder.layers, layer_name))
+                
+                # Replace the encoder layers with our truncated version
+                new_model.encoder.layers = encoder_layers
+                
+                # Remove the heads since we're stopping at the encoder layer
+                new_model.heads = nn.Identity()
+                
+                return new_model
+                
+        except (ValueError, IndexError) as e:
+            raise ValueError(f"Invalid ViT layer specification: {layer_index}. Error: {e}")
+    
+    # Handling for ViT whole blocks
+    elif hasattr(model, 'blocks') or (hasattr(model, 'module') and hasattr(model.module, 'blocks')):
+        # Check for DataParallel wrapper
+        base_model = model.module if hasattr(model, 'module') else model
+        
+        # Create a deep copy to avoid modifying the original
+        new_model = copy.deepcopy(model)
+        base_new_model = new_model.module if hasattr(new_model, 'module') else new_model
+        
+        # Add the desired number of transformer blocks
+        if isinstance(layer_index, int):
+            # Truncate the blocks
+            base_new_model.blocks = base_new_model.blocks[:layer_index+1]
+            
+        return new_model
+    
+    else:
+        # Original ResNet/VGG handling
+        modules = list(model.named_children())
+        print(f"DEBUG - extract_middle_layers - Looking for '{layer_index}' in {[name for name, _ in modules]}")
+        
+        cutoff_idx = next((i for i, (name, _) in enumerate(modules) 
+                          if name == str(layer_index)), None)
+        
+        if cutoff_idx is not None:
+            # Keep modules up to and including the target
+            new_model = nn.Sequential(OrderedDict(modules[:cutoff_idx+1]))
+            return new_model
+        else:
+            raise ValueError(f"Module {layer_index} not found in model")
+
+# Get ImageNet labels
+def get_imagenet_labels():
+    url = "https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json"
+    response = requests.get(url)
+    if response.status_code == 200:
+        return response.json()
+    else:
+        raise RuntimeError("Failed to fetch ImageNet labels")
+
+# Download model if needed
+def download_model(model_type):
+    if model_type not in MODEL_URLS or MODEL_URLS[model_type] is None:
+        return None  # Use PyTorch's pretrained model
+    
+    # Handle special case for face model
+    if model_type == 'resnet50_robust_face':
+        model_path = Path("models/resnet50_robust_face_100_checkpoint.pt")
+    else:
+        model_path = Path(f"models/{model_type}.pt")
+    
+    if not model_path.exists():
+        print(f"Downloading {model_type} model...")
+        url = MODEL_URLS[model_type]
+        response = requests.get(url, stream=True)
+        if response.status_code == 200:
+            with open(model_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
+            print(f"Model downloaded and saved to {model_path}")
+        else:
+            raise RuntimeError(f"Failed to download model: {response.status_code}")
+    return model_path
+
+class NormalizeByChannelMeanStd(nn.Module):
+    def __init__(self, mean, std):
+        super(NormalizeByChannelMeanStd, self).__init__()
+        if not isinstance(mean, torch.Tensor):
+            mean = torch.tensor(mean)
+        if not isinstance(std, torch.Tensor):
+            std = torch.tensor(std)
+        self.register_buffer("mean", mean)
+        self.register_buffer("std", std)
+        
+    def forward(self, tensor):
+        return self.normalize_fn(tensor, self.mean, self.std)
+    
+    def normalize_fn(self, tensor, mean, std):
+        """Differentiable version of torchvision.functional.normalize"""
+        # here we assume the color channel is at dim=1
+        mean = mean[None, :, None, None]
+        std = std[None, :, None, None]
+        return tensor.sub(mean).div(std)
+
+class InferStep:
+    def __init__(self, orig_image, eps, step_size):
+        self.orig_image = orig_image
+        self.eps = eps
+        self.step_size = step_size
+
+    def project(self, x):
+        diff = x - self.orig_image
+        diff = torch.clamp(diff, -self.eps, self.eps)
+        return torch.clamp(self.orig_image + diff, 0, 1)
+
+    def step(self, x, grad):
+        l = len(x.shape) - 1
+        grad_norm = torch.norm(grad.view(grad.shape[0], -1), dim=1).view(-1, *([1]*l))
+        scaled_grad = grad / (grad_norm + 1e-10)
+        return scaled_grad * self.step_size
+
+def get_iterations_to_show(n_itr):
+    """Generate a dynamic list of iterations to show based on total iterations."""
+    if n_itr <= 50:
+        return [1, 5, 10, 20, 30, 40, 50, n_itr]
+    elif n_itr <= 100:
+        return [1, 5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, n_itr]
+    elif n_itr <= 200:
+        return [1, 5, 10, 20, 30, 40, 50, 75, 100, 125, 150, 175, 200, n_itr]
+    elif n_itr <= 500:
+        return [1, 5, 10, 20, 30, 40, 50, 75, 100, 150, 200, 250, 300, 350, 400, 450, 500, n_itr]
+    else:
+        # For very large iterations, show more evenly distributed points
+        return [1, 5, 10, 20, 30, 40, 50, 75, 100, 150, 200, 250, 300, 350, 400, 450, 500, 
+                int(n_itr*0.6), int(n_itr*0.7), int(n_itr*0.8), int(n_itr*0.9), n_itr]
+
+def get_inference_configs(inference_type='IncreaseConfidence', eps=0.5, n_itr=50, step_size=1.0):
+    """Generate inference configuration with customizable parameters.
+    
+    Args:
+        inference_type (str): Type of inference ('IncreaseConfidence' or 'Prior-Guided Drift Diffusion')
+        eps (float): Maximum perturbation size
+        n_itr (int): Number of iterations
+        step_size (float): Step size for each iteration
+    """
+    
+    # Base configuration common to all inference types
+    config = {
+        'loss_infer': inference_type,  # How to guide the optimization
+        'n_itr': n_itr,  # Number of iterations
+        'eps': eps,  # Maximum perturbation size
+        'step_size': step_size,  # Step size for each iteration
+        'diffusion_noise_ratio': 0.0,  # No diffusion noise
+        'initial_inference_noise_ratio': 0.0,  # No initial noise
+        'top_layer': 'all',  # Use all layers of the model
+        'inference_normalization': False,  # Apply normalization during inference
+        'recognition_normalization': False,  # Apply normalization during recognition
+        'iterations_to_show': get_iterations_to_show(n_itr),  # Dynamic iterations to visualize
+        'misc_info': {'keep_grads': False}  # Additional configuration
+    }
+    
+    # Customize based on inference type
+    if inference_type == 'IncreaseConfidence':
+        config['loss_function'] = 'CE'  # Cross Entropy
+    
+    elif inference_type == 'Prior-Guided Drift Diffusion':
+        config['loss_function'] = 'MSE'  # Mean Square Error
+        config['initial_inference_noise_ratio'] = 0.05  # Initial noise for diffusion
+        config['diffusion_noise_ratio'] = 0.01  # Add noise during diffusion
+    
+    elif inference_type == 'GradModulation':
+        config['loss_function'] = 'CE'  # Cross Entropy
+        config['misc_info']['grad_modulation'] = 0.5  # Gradient modulation strength
+    
+    elif inference_type == 'CompositionalFusion':
+        config['loss_function'] = 'CE'  # Cross Entropy
+        config['misc_info']['positive_classes'] = []  # Classes to maximize
+        config['misc_info']['negative_classes'] = []  # Classes to minimize
+    
+    return config
+
+class GenerativeInferenceModel:
+    def __init__(self):
+        self.models = {}
+        self.normalizer = NormalizeByChannelMeanStd(IMAGENET_MEAN, IMAGENET_STD).to(device)
+        self.labels = get_imagenet_labels()
+        
+    def verify_model_integrity(self, model, model_type):
+        """
+        Verify model integrity by running a test input through it.
+        Returns whether the model passes basic integrity check.
+        """
+        try:
+            print(f"\n=== Running model integrity check for {model_type} ===")
+            # Create a deterministic test input directly on the correct device
+            test_input = torch.zeros(1, 3, 224, 224, device=device)
+            test_input[0, 0, 100:124, 100:124] = 0.5  # Red square
+            
+            # Run forward pass
+            with torch.no_grad():
+                output = model(test_input)
+            
+            # Check output shape
+            if output.shape != (1, 1000):
+                print(f"❌ Unexpected output shape: {output.shape}, expected (1, 1000)")
+                return False
+                
+            # Get top prediction
+            probs = torch.nn.functional.softmax(output, dim=1)
+            confidence, prediction = torch.max(probs, 1)
+            
+            # Calculate basic statistics on output
+            mean = output.mean().item()
+            std = output.std().item()
+            min_val = output.min().item()
+            max_val = output.max().item()
+            
+            print(f"Model integrity check results:")
+            print(f"- Output shape: {output.shape}")
+            print(f"- Top prediction: Class {prediction.item()} with {confidence.item()*100:.2f}% confidence")
+            print(f"- Output statistics: mean={mean:.3f}, std={std:.3f}, min={min_val:.3f}, max={max_val:.3f}")
+            
+            # Basic sanity checks
+            if torch.isnan(output).any():
+                print("❌ Model produced NaN outputs")
+                return False
+                
+            if output.std().item() < 0.1:
+                print("⚠️ Low output variance, model may not be discriminative")
+                
+            print("✅ Model passes basic integrity check")
+            return True
+            
+        except Exception as e:
+            print(f"❌ Model integrity check failed with error: {e}")
+            # Rather than failing completely, we'll continue
+            return True
+    
+    def load_model(self, model_type):
+        """Load model from checkpoint or use pretrained model."""
+        if model_type in self.models:
+            print(f"Using cached {model_type} model")
+            return self.models[model_type]
+        
+        # Record loading time for performance analysis
+        start_time = time.time()
+        model_path = download_model(model_type)
+        
+        # Create a sequential model with normalizer and ResNet50
+        resnet = models.resnet50()
+        model = nn.Sequential(
+            self.normalizer,  # Normalizer is part of the model sequence
+            resnet
+        )
+        
+        # Load the model checkpoint
+        if model_path:
+            print(f"Loading {model_type} model from {model_path}...")
+            try:
+                checkpoint = torch.load(model_path, map_location=device)
+                
+                # Print checkpoint structure for better understanding
+                print("\n=== Analyzing checkpoint structure ===")
+                if isinstance(checkpoint, dict):
+                    print(f"Checkpoint contains keys: {list(checkpoint.keys())}")
+                    
+                    # Examine 'model' structure if it exists
+                    if 'model' in checkpoint and isinstance(checkpoint['model'], dict):
+                        model_dict = checkpoint['model']
+                        # Get sample of keys to understand structure
+                        first_keys = list(model_dict.keys())[:5]
+                        print(f"'model' contains keys like: {first_keys}")
+                        
+                        # Check for common prefixes in the model dict
+                        prefixes = set()
+                        for key in list(model_dict.keys())[:100]:  # Check first 100 keys
+                            parts = key.split('.')
+                            if len(parts) > 1:
+                                prefixes.add(parts[0])
+                        if prefixes:
+                            print(f"Common prefixes in model dict: {prefixes}")
+                else:
+                    print(f"Checkpoint is not a dictionary, but a {type(checkpoint)}")
+                
+                # Handle different checkpoint formats
+                if 'model' in checkpoint:
+                    # Format from madrylab robust models
+                    state_dict = checkpoint['model']
+                    print("Using 'model' key from checkpoint")
+                elif 'state_dict' in checkpoint:
+                    state_dict = checkpoint['state_dict']
+                    print("Using 'state_dict' key from checkpoint")
+                else:
+                    # Direct state dict
+                    state_dict = checkpoint
+                    print("Using checkpoint directly as state_dict")
+                
+                # Handle prefix in state dict keys for ResNet part
+                resnet_state_dict = {}
+                prefixes_to_try = ['', 'module.', 'model.', 'attacker.model.']
+                resnet_keys = set(resnet.state_dict().keys())
+                
+                # First check if we can find keys directly in the attacker.model path
+                print("\n=== Phase 1: Checking for specific model structures ===")
+                
+                # Check for 'module.model' structure (seen in actual checkpoint)
+                module_model_keys = [key for key in state_dict.keys() if key.startswith('module.model.')]
+                if module_model_keys:
+                    print(f"Found 'module.model' structure with {len(module_model_keys)} parameters")
+                    # Extract all parameters from module.model
+                    for source_key, value in state_dict.items():
+                        if source_key.startswith('module.model.'):
+                            target_key = source_key[len('module.model.'):]
+                            resnet_state_dict[target_key] = value
+                            
+                    print(f"Extracted {len(resnet_state_dict)} parameters from module.model")
+                
+                # Check for 'attacker.model' structure
+                attacker_model_keys = [key for key in state_dict.keys() if key.startswith('attacker.model.')]
+                if attacker_model_keys:
+                    print(f"Found 'attacker.model' structure with {len(attacker_model_keys)} parameters")
+                    # Extract all parameters from attacker.model
+                    for source_key, value in state_dict.items():
+                        if source_key.startswith('attacker.model.'):
+                            target_key = source_key[len('attacker.model.'):]
+                            resnet_state_dict[target_key] = value
+                            
+                    print(f"Extracted {len(resnet_state_dict)} parameters from attacker.model")
+                    
+                    # Check if 'model' (not attacker.model) exists as a fallback
+                    model_keys = [key for key in state_dict.keys() if key.startswith('model.') and not key.startswith('attacker.model.')]
+                    if model_keys and len(resnet_state_dict) < len(resnet_keys):
+                        print(f"Found additional 'model.' structure with {len(model_keys)} parameters")
+                        # Try to complete missing parameters
+                        for source_key, value in state_dict.items():
+                            if source_key.startswith('model.'):
+                                target_key = source_key[len('model.'):]
+                                if target_key in resnet_keys and target_key not in resnet_state_dict:
+                                    resnet_state_dict[target_key] = value
+                                    
+                else:
+                    # Check for other known structures
+                    structure_found = False
+                    
+                    # Check for 'model.' prefix
+                    model_keys = [key for key in state_dict.keys() if key.startswith('model.')]
+                    if model_keys:
+                        print(f"Found 'model.' structure with {len(model_keys)} parameters")
+                        for source_key, value in state_dict.items():
+                            if source_key.startswith('model.'):
+                                target_key = source_key[len('model.'):]
+                                resnet_state_dict[target_key] = value
+                        structure_found = True
+                    
+                    # Check for ResNet parameters at the top level
+                    top_level_resnet_keys = 0
+                    for key in resnet_keys:
+                        if key in state_dict:
+                            top_level_resnet_keys += 1
+                            
+                    if top_level_resnet_keys > 0:
+                        print(f"Found {top_level_resnet_keys} ResNet parameters at top level")
+                        for target_key in resnet_keys:
+                            if target_key in state_dict:
+                                resnet_state_dict[target_key] = state_dict[target_key]
+                        structure_found = True
+                    
+                    # If no structure was recognized, try the prefix mapping approach
+                    if not structure_found:
+                        print("No standard model structure found, trying prefix mappings...")
+                        for target_key in resnet_keys:
+                            for prefix in prefixes_to_try:
+                                source_key = prefix + target_key
+                                if source_key in state_dict:
+                                    resnet_state_dict[target_key] = state_dict[source_key]
+                                    break
+                
+                # If we still can't find enough keys, try a final approach of removing prefixes
+                if len(resnet_state_dict) < len(resnet_keys):
+                    print(f"Found only {len(resnet_state_dict)}/{len(resnet_keys)} parameters, trying prefix removal...")
+                    
+                    # Track matches found through prefix removal
+                    prefix_matches = {prefix: 0 for prefix in ['module.', 'model.', 'attacker.model.', 'attacker.']}
+                    layer_matches = {}  # Track matches by layer type
+                    
+                    # Count parameter keys by layer type for analysis
+                    for key in resnet_keys:
+                        layer_name = key.split('.')[0] if '.' in key else key
+                        if layer_name not in layer_matches:
+                            layer_matches[layer_name] = {'total': 0, 'matched': 0}
+                        layer_matches[layer_name]['total'] += 1
+                    
+                    # Try keys with common prefixes
+                    for source_key, value in state_dict.items():
+                        # Skip if already found
+                        target_key = source_key
+                        matched_prefix = None
+                        
+                        # Try removing various prefixes
+                        for prefix in ['module.', 'model.', 'attacker.model.', 'attacker.']:
+                            if source_key.startswith(prefix):
+                                target_key = source_key[len(prefix):]
+                                matched_prefix = prefix
+                                break
+                        
+                        # If the target key is in the ResNet keys, add it to the state dict
+                        if target_key in resnet_keys and target_key not in resnet_state_dict:
+                            resnet_state_dict[target_key] = value
+                            
+                            # Update match statistics
+                            if matched_prefix:
+                                prefix_matches[matched_prefix] += 1
+                            
+                            # Update layer matches
+                            layer_name = target_key.split('.')[0] if '.' in target_key else target_key
+                            if layer_name in layer_matches:
+                                layer_matches[layer_name]['matched'] += 1
+                    
+                    # Print detailed prefix removal statistics
+                    print("\n=== Prefix Removal Statistics ===")
+                    total_matches = sum(prefix_matches.values())
+                    print(f"Total parameters matched through prefix removal: {total_matches}/{len(resnet_keys)} ({(total_matches/len(resnet_keys))*100:.1f}%)")
+                    
+                    # Show matches by prefix
+                    print("\nMatches by prefix:")
+                    for prefix, count in sorted(prefix_matches.items(), key=lambda x: x[1], reverse=True):
+                        if count > 0:
+                            print(f"  {prefix}: {count} parameters")
+                    
+                    # Show matches by layer type
+                    print("\nMatches by layer type:")
+                    for layer, stats in sorted(layer_matches.items(), key=lambda x: x[1]['total'], reverse=True):
+                        match_percent = (stats['matched'] / stats['total']) * 100 if stats['total'] > 0 else 0
+                        print(f"  {layer}: {stats['matched']}/{stats['total']} ({match_percent:.1f}%)")
+                    
+                    # Check for specific important layers (conv1, layer1, etc.)
+                    critical_layers = ['conv1', 'bn1', 'layer1', 'layer2', 'layer3', 'layer4', 'fc']
+                    print("\nStatus of critical layers:")
+                    for layer in critical_layers:
+                        if layer in layer_matches:
+                            match_percent = (layer_matches[layer]['matched'] / layer_matches[layer]['total']) * 100
+                            status = "✅ COMPLETE" if layer_matches[layer]['matched'] == layer_matches[layer]['total'] else "⚠️ INCOMPLETE"
+                            print(f"  {layer}: {layer_matches[layer]['matched']}/{layer_matches[layer]['total']} ({match_percent:.1f}%) - {status}")
+                        else:
+                            print(f"  {layer}: Not found in model")
+                
+                # Load the ResNet state dict
+                if resnet_state_dict:
+                    try:
+                        # Use strict=False to allow missing keys
+                        result = resnet.load_state_dict(resnet_state_dict, strict=False)
+                        missing_keys, unexpected_keys = result
+                        
+                        # Generate detailed information with better formatting
+                        loading_report = []
+                        loading_report.append(f"\n===== MODEL LOADING REPORT: {model_type} =====")
+                        loading_report.append(f"Total parameters in checkpoint: {len(resnet_state_dict):,}")
+                        loading_report.append(f"Total parameters in model: {len(resnet.state_dict()):,}")
+                        loading_report.append(f"Missing keys: {len(missing_keys):,} parameters")
+                        loading_report.append(f"Unexpected keys: {len(unexpected_keys):,} parameters")
+
+                        # Calculate percentage of parameters loaded
+                        loaded_keys = set(resnet_state_dict.keys()) - set(unexpected_keys)
+                        loaded_percent = (len(loaded_keys) / len(resnet.state_dict())) * 100
+                        
+                        # Determine loading success status
+                        if loaded_percent >= 99.5:
+                            status = "✅ COMPLETE - All important parameters loaded"
+                        elif loaded_percent >= 90:
+                            status = "🟡 PARTIAL - Most parameters loaded, should still function"
+                        elif loaded_percent >= 50:
+                            status = "⚠️ INCOMPLETE - Many parameters missing, may not function properly"
+                        else:
+                            status = "❌ FAILED - Critical parameters missing, will not function properly"
+                            
+                        loading_report.append(f"Successfully loaded: {len(loaded_keys):,} parameters ({loaded_percent:.1f}%)")
+                        loading_report.append(f"Loading status: {status}")
+                        
+                        # If loading is severely incomplete, fall back to PyTorch's pretrained model
+                        if loaded_percent < 50:
+                            loading_report.append("\n⚠️ WARNING: Loading from checkpoint is too incomplete.")
+                            loading_report.append("⚠️ Falling back to PyTorch's pretrained model to avoid broken inference.")
+                            
+                            # Create a new ResNet model with pretrained weights
+                            resnet = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
+                            model = nn.Sequential(self.normalizer, resnet)
+                            loading_report.append("✅ Successfully loaded PyTorch's pretrained ResNet50 model")
+                        
+                        # Show missing keys by layer type
+                        if missing_keys:
+                            loading_report.append("\nMissing keys by layer type:")
+                            layer_types = {}
+                            for key in missing_keys:
+                                # Extract layer type (e.g., 'conv', 'bn', 'layer1', etc.)
+                                parts = key.split('.')
+                                if len(parts) > 0:
+                                    layer_type = parts[0]
+                                    if layer_type not in layer_types:
+                                        layer_types[layer_type] = 0
+                                    layer_types[layer_type] += 1
+                            
+                            # Add counts by layer type
+                            for layer_type, count in sorted(layer_types.items(), key=lambda x: x[1], reverse=True):
+                                loading_report.append(f"  {layer_type}: {count:,} parameters")
+                                
+                            loading_report.append("\nFirst 10 missing keys:")
+                            for i, key in enumerate(sorted(missing_keys)[:10]):
+                                loading_report.append(f"  {i+1}. {key}")
+                        
+                        # Show unexpected keys if any
+                        if unexpected_keys:
+                            loading_report.append("\nFirst 10 unexpected keys:")
+                            for i, key in enumerate(sorted(unexpected_keys)[:10]):
+                                loading_report.append(f"  {i+1}. {key}")
+                                
+                        loading_report.append("========================================")
+                        
+                        # Convert report to string and print it
+                        report_text = "\n".join(loading_report)
+                        print(report_text)
+                        
+                        # Also save to a file for reference
+                        os.makedirs("logs", exist_ok=True)
+                        with open(f"logs/model_loading_{model_type}.log", "w") as f:
+                            f.write(report_text)
+                        
+                        # Look for normalizer parameters as well
+                        if any(key.startswith('attacker.normalize.') for key in state_dict.keys()):
+                            norm_state_dict = {}
+                            for key, value in state_dict.items():
+                                if key.startswith('attacker.normalize.'):
+                                    norm_key = key[len('attacker.normalize.'):]
+                                    norm_state_dict[norm_key] = value
+                            
+                            if norm_state_dict:
+                                try:
+                                    self.normalizer.load_state_dict(norm_state_dict, strict=False)
+                                    print("Successfully loaded normalizer parameters")
+                                except Exception as e:
+                                    print(f"Warning: Could not load normalizer parameters: {e}")
+                    except Exception as e:
+                        print(f"Warning: Error loading ResNet parameters: {e}")
+                        # Fall back to loading without normalizer
+                        model = resnet  # Use just the ResNet model without normalizer
+            except Exception as e:
+                print(f"Error loading model checkpoint: {e}")
+                # Fallback to PyTorch's pretrained model
+                print("Falling back to PyTorch's pretrained model")
+                resnet = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
+                model = nn.Sequential(self.normalizer, resnet)
+        else:
+            # Fallback to PyTorch's pretrained model
+            print("No checkpoint available, using PyTorch's pretrained model")
+            resnet = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
+            model = nn.Sequential(self.normalizer, resnet)
+        
+        model = model.to(device)
+        model.eval()  # Set to evaluation mode
+        
+        # Verify model integrity
+        self.verify_model_integrity(model, model_type)
+        
+        # Store the model for future use
+        self.models[model_type] = model
+        end_time = time.time()
+        load_time = end_time - start_time
+        print(f"Model {model_type} loaded in {load_time:.2f} seconds")
+        return model
+    
+    def inference(self, image, model_type, config):
+        """Run generative inference on the image."""
+        # Time the entire inference process
+        inference_start = time.time()
+        
+        # Load model if not already loaded
+        model = self.load_model(model_type)
+        
+        # Check if image is a file path
+        if isinstance(image, str):
+            if os.path.exists(image):
+                image = Image.open(image).convert('RGB')
+            else:
+                raise ValueError(f"Image path does not exist: {image}")
+        elif isinstance(image, torch.Tensor):
+            raise ValueError(f"Image type {type(image)}, looks like already a transformed tensor")
+        
+        # Prepare image tensor - match original code's conditional transform
+        load_start = time.time()
+        use_norm = config['inference_normalization'] == 'on'
+        custom_transform = get_transform(
+            input_size=224, 
+            normalize=use_norm, 
+            norm_mean=IMAGENET_MEAN, 
+            norm_std=IMAGENET_STD
+        )
+        
+        # Special handling for GradModulation as in original
+        if config['loss_infer'] == 'GradModulation' and 'misc_info' in config and 'grad_modulation' in config['misc_info']:
+            grad_modulation = config['misc_info']['grad_modulation']
+            image_tensor = custom_transform(image).unsqueeze(0).to(device)
+            image_tensor = image_tensor * (1-grad_modulation) + grad_modulation * torch.randn_like(image_tensor).to(device)
+        else:
+            image_tensor = custom_transform(image).unsqueeze(0).to(device)
+            
+        image_tensor.requires_grad = True
+        print(f"Image loaded and processed in {time.time() - load_start:.2f} seconds")
+        
+        # Check model structure
+        is_sequential = isinstance(model, nn.Sequential)
+        
+        # Get original predictions
+        with torch.no_grad():
+            # If the model is sequential with a normalizer, skip the normalization step
+            if is_sequential and isinstance(model[0], NormalizeByChannelMeanStd):
+                print("Model is sequential with normalization")
+                # Get the core model part (typically at index 1 in Sequential)
+                core_model = model[1]
+                if config['inference_normalization']:
+                    output_original = model(image_tensor)  # Model includes normalization
+                else:
+                    output_original = core_model(image_tensor)  # Model includes normalization
+                
+            else:
+                print("Model is not sequential with normalization")
+                # Use manual normalization for non-sequential models
+                if config['inference_normalization']:
+                    normalized_tensor = normalize_transform(image_tensor)
+                    output_original = model(normalized_tensor)
+                else:
+                    output_original = model(image_tensor)
+                core_model = model
+            
+            probs_orig = F.softmax(output_original, dim=1)
+            conf_orig, classes_orig = torch.max(probs_orig, 1)
+            
+            # Get least confident classes
+            _, least_confident_classes = torch.topk(probs_orig, k=100, largest=False)
+        
+        # Initialize inference step
+        infer_step = InferStep(image_tensor, config['eps'], config['step_size'])
+        
+        # Storage for inference steps
+        # Create a new tensor that requires gradients
+        x = image_tensor.clone().detach().requires_grad_(True)
+        all_steps = [image_tensor[0].detach().cpu()]
+        
+        # For Prior-Guided Drift Diffusion, extract selected layer and initialize with noisy features
+        noisy_features = None
+        layer_model = None
+        if config['loss_infer'] == 'Prior-Guided Drift Diffusion':
+            print(f"Setting up Prior-Guided Drift Diffusion with layer {config['top_layer']} and noise {config['initial_inference_noise_ratio']}...")
+            
+            # Extract model up to the specified layer
+            try:
+                # Start by finding the actual model to use
+                base_model = model
+                
+                # Handle DataParallel wrapper if present
+                if hasattr(base_model, 'module'):
+                    base_model = base_model.module
+                
+                # Log the initial model structure
+                print(f"DEBUG - Initial model structure: {type(base_model)}")
+                
+                # If we have a Sequential model (which is likely our normalizer + model structure)
+                if isinstance(base_model, nn.Sequential):
+                    print(f"DEBUG - Sequential model with {len(list(base_model.children()))} children")
+                    
+                    # If this is our NormalizeByChannelMeanStd + ResNet pattern
+                    if len(list(base_model.children())) >= 2:
+                        # The actual ResNet model is the second component (index 1)
+                        actual_model = list(base_model.children())[1]
+                        print(f"DEBUG - Using ResNet component: {type(actual_model)}")
+                        print(f"DEBUG - Available layers: {[name for name, _ in actual_model.named_children()]}")
+                        
+                        # Extract from the actual ResNet
+                        layer_model = extract_middle_layers(actual_model, config['top_layer'])
+                    else:
+                        # Just a single component Sequential
+                        layer_model = extract_middle_layers(base_model, config['top_layer'])
+                else:
+                    # Not Sequential, might be direct model
+                    print(f"DEBUG - Available layers: {[name for name, _ in base_model.named_children()]}")
+                    layer_model = extract_middle_layers(base_model, config['top_layer'])
+                
+                print(f"Successfully extracted model up to layer: {config['top_layer']}")
+            except ValueError as e:
+                print(f"Layer extraction failed: {e}. Using full model.")
+                layer_model = model
+            
+            # Add noise to the image - exactly match original code
+            added_noise = config['initial_inference_noise_ratio'] * torch.randn_like(image_tensor).to(device)
+            noisy_image_tensor = image_tensor + added_noise
+            
+            # Compute noisy features - simplified to match original code
+            noisy_features = layer_model(noisy_image_tensor)
+            
+            print(f"Noisy features computed for Prior-Guided Drift Diffusion target with shape: {noisy_features.shape if hasattr(noisy_features, 'shape') else 'unknown'}")
+        
+        # Main inference loop
+        print(f"Starting inference loop with {config['n_itr']} iterations for {config['loss_infer']}...")
+        loop_start = time.time()
+        for i in range(config['n_itr']):
+            # Reset gradients
+            x.grad = None
+            
+            # Forward pass - use layer_model for Prior-Guided Drift Diffusion, full model otherwise
+            if config['loss_infer'] == 'Prior-Guided Drift Diffusion' and layer_model is not None:
+                # Use the extracted layer model for Prior-Guided Drift Diffusion
+                # In original code, normalization is handled at transform time, not during forward pass
+                output = layer_model(x)
+            else:
+                # Standard forward pass with full model
+                # Simplified to match original code's approach
+                output = model(x)
+            
+            # Calculate loss and gradients based on inference type
+            try:
+                if config['loss_infer'] == 'Prior-Guided Drift Diffusion':
+                    # Use MSE loss to match the noisy features
+                    assert config['loss_function'] == 'MSE', "Reverse Diffusion loss function must be MSE"
+                    if noisy_features is not None:
+                        loss = F.mse_loss(output, noisy_features)
+                        grad = torch.autograd.grad(loss, x)[0]  # Removed retain_graph=True to match original
+                    else:
+                        raise ValueError("Noisy features not computed for Prior-Guided Drift Diffusion")
+                
+                else:  # Default 'IncreaseConfidence' approach
+                    # Get the least confident classes
+                    num_classes = min(10, least_confident_classes.size(1))
+                    target_classes = least_confident_classes[0, :num_classes]
+                    
+                    # Create targets for least confident classes
+                    targets = torch.tensor([idx.item() for idx in target_classes], device=device)
+                    
+                    # Use a combined loss to increase confidence
+                    loss = 0
+                    for target in targets:
+                        # Create one-hot target
+                        one_hot = torch.zeros_like(output)
+                        one_hot[0, target] = 1
+                        # Use loss to maximize confidence
+                        loss = loss + F.mse_loss(F.softmax(output, dim=1), one_hot)
+                    
+                    grad = torch.autograd.grad(loss, x, retain_graph=True)[0]
+                
+                if grad is None:
+                    print("Warning: Direct gradient calculation failed")
+                    # Fall back to random perturbation
+                    random_noise = (torch.rand_like(x) - 0.5) * 2 * config['step_size']
+                    x = infer_step.project(x + random_noise)
+                else:
+                    # Update image with gradient - do this exactly as in original code
+                    adjusted_grad = infer_step.step(x, grad)
+                    
+                    # Add diffusion noise if specified
+                    diffusion_noise = config['diffusion_noise_ratio'] * torch.randn_like(x).to(device)
+                    
+                    # Apply gradient and noise in one operation before projecting, exactly as in original
+                    x = infer_step.project(x.clone() + adjusted_grad + diffusion_noise)
+                    
+            except Exception as e:
+                print(f"Error in gradient calculation: {e}")
+                # Fall back to random perturbation - match original code
+                random_noise = (torch.rand_like(x) - 0.5) * 2 * config['step_size']
+                x = infer_step.project(x.clone() + random_noise)
+            
+            # Store step if in iterations_to_show
+            if i+1 in config['iterations_to_show'] or i+1 == config['n_itr']:
+                all_steps.append(x[0].detach().cpu())
+        
+        # Print some info about the inference
+        with torch.no_grad():
+            if is_sequential and isinstance(model[0], NormalizeByChannelMeanStd):
+                if config['inference_normalization']:
+                    final_output = model(x)
+                else:
+                    final_output = core_model(x)
+            else:
+                if config['inference_normalization']:
+                    normalized_x = normalize_transform(x)
+                    final_output = model(normalized_x)
+                else:
+                    final_output = model(x)
+                
+            final_probs = F.softmax(final_output, dim=1)
+            final_conf, final_classes = torch.max(final_probs, 1)
+            
+            # Calculate timing information
+            loop_time = time.time() - loop_start
+            total_time = time.time() - inference_start
+            avg_iter_time = loop_time / config['n_itr'] if config['n_itr'] > 0 else 0
+            
+            print(f"Original top class: {classes_orig.item()} ({conf_orig.item():.4f})")
+            print(f"Final top class: {final_classes.item()} ({final_conf.item():.4f})")
+            print(f"Inference loop completed in {loop_time:.2f} seconds ({avg_iter_time:.4f} sec/iteration)")
+            print(f"Total inference time: {total_time:.2f} seconds")
+            
+        # Return results in format compatible with both old and new code
+        return {
+            'final_image': x[0].detach().cpu(),
+            'steps': all_steps,
+            'original_class': classes_orig.item(),
+            'original_confidence': conf_orig.item(),
+            'final_class': final_classes.item(),
+            'final_confidence': final_conf.item()
+        }
+
+# Utility function to show inference steps
+def show_inference_steps(steps, figsize=(15, 10)):
+    import matplotlib.pyplot as plt
+    
+    n_steps = len(steps)
+    fig, axes = plt.subplots(1, n_steps, figsize=figsize)
+    
+    for i, step_img in enumerate(steps):
+        img = step_img.permute(1, 2, 0).numpy()
+        axes[i].imshow(img)
+        axes[i].set_title(f"Step {i}")
+        axes[i].axis('off')
+    
+    plt.tight_layout()
+    return fig 

logs/model_loading_resnet50_robust.log

@@ -0,0 +1,9 @@
+
+===== MODEL LOADING REPORT: resnet50_robust =====
+Total parameters in checkpoint: 320
+Total parameters in model: 320
+Missing keys: 0 parameters
+Unexpected keys: 0 parameters
+Successfully loaded: 320 parameters (100.0%)
+Loading status: ✅ COMPLETE - All important parameters loaded
+========================================

logs/model_loading_resnet50_robust_face.log

@@ -0,0 +1,46 @@
+
+===== MODEL LOADING REPORT: resnet50_robust_face =====
+Total parameters in checkpoint: 320
+Total parameters in model: 320
+Missing keys: 267 parameters
+Unexpected keys: 320 parameters
+Successfully loaded: 0 parameters (0.0%)
+Loading status: ❌ FAILED - Critical parameters missing, will not function properly
+
+⚠️ WARNING: Loading from checkpoint is too incomplete.
+⚠️ Falling back to PyTorch's pretrained model to avoid broken inference.
+✅ Successfully loaded PyTorch's pretrained ResNet50 model
+
+Missing keys by layer type:
+  layer3: 95 parameters
+  layer2: 65 parameters
+  layer1: 50 parameters
+  layer4: 50 parameters
+  bn1: 4 parameters
+  fc: 2 parameters
+  conv1: 1 parameters
+
+First 10 missing keys:
+  1. bn1.bias
+  2. bn1.running_mean
+  3. bn1.running_var
+  4. bn1.weight
+  5. conv1.weight
+  6. fc.bias
+  7. fc.weight
+  8. layer1.0.bn1.bias
+  9. layer1.0.bn1.running_mean
+  10. layer1.0.bn1.running_var
+
+First 10 unexpected keys:
+  1. model.bn1.bias
+  2. model.bn1.num_batches_tracked
+  3. model.bn1.running_mean
+  4. model.bn1.running_var
+  5. model.bn1.weight
+  6. model.conv1.weight
+  7. model.fc.bias
+  8. model.fc.weight
+  9. model.layer1.0.bn1.bias
+  10. model.layer1.0.bn1.num_batches_tracked
+========================================

logs/model_loading_robust_resnet50.log

@@ -0,0 +1,9 @@
+
+===== MODEL LOADING REPORT: robust_resnet50 =====
+Total parameters in checkpoint: 320
+Total parameters in model: 320
+Missing keys: 0 parameters
+Unexpected keys: 0 parameters
+Successfully loaded: 320 parameters (100.0%)
+Loading status: ✅ COMPLETE - All important parameters loaded
+========================================

logs/model_loading_standard_resnet50.log

@@ -0,0 +1,9 @@
+
+===== MODEL LOADING REPORT: standard_resnet50 =====
+Total parameters in checkpoint: 320
+Total parameters in model: 320
+Missing keys: 0 parameters
+Unexpected keys: 0 parameters
+Successfully loaded: 320 parameters (100.0%)
+Loading status: ✅ COMPLETE - All important parameters loaded
+========================================

models/resnet50_robust.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:380b14e6f9750bffa1447cf7017f65da4dc5ce71a3dd112f107515dcf7b14d9d
+size 204818947

models/resnet50_robust_face_100_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c48a5c16ca0d5ac4cb20f1b98e2128838746f18b658728ac661f1ffd589c37bf
+size 196695413

models/robust_resnet50.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:380b14e6f9750bffa1447cf7017f65da4dc5ce71a3dd112f107515dcf7b14d9d
+size 204818947

models/standard_resnet50.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72d4a99582db5d7fa86c3fd2a089f0bfd6a10f69d635bca51f6ad72ac6b458f0
+size 204818947

requirements.txt

@@ -0,0 +1,9 @@
+torch
+torchvision
+numpy
+pillow
+gradio
+matplotlib
+requests
+tqdm
+huggingface_hub