Upload 2 files

some small fixes

app.py

@@ -79,6 +79,7 @@ examples = [
             "[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",
@@ -101,12 +102,12 @@ examples = [
         "method": "Prior-Guided Drift Diffusion",
         "reverse_diff": {
             "model": "resnet50_robust",
-            "layer": "layer4",
-            "initial_noise": 0.5,
-            "diffusion_noise": 0.01,
-            "step_size": 0.2,
-            "iterations": 301,
-            "epsilon": 40.0
+            "layer": "avgpool",
+            "initial_noise": 0.9,
+            "diffusion_noise": 0.003,
+            "step_size": 0.58,
+            "iterations": 100,
+            "epsilon": 0.81
         }
     },
     {
@@ -243,7 +244,8 @@ def apply_example(example):
         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
+        example["reverse_diff"]["layer"],  # Model layer
+        gr.Group(visible=True)  # Show parameters section
     ]
 
 # Define the interface
@@ -273,39 +275,45 @@ with gr.Blocks(title="Generative Inference Demo", css="""
             # Inputs
             image_input = gr.Image(label="Input Image", type="pil", value=os.path.join("stimuli", "Neon_Color_Circle.jpg"))
             
-            with gr.Row():
-                model_choice = gr.Dropdown(
-                    choices=["resnet50_robust", "standard_resnet50"], 
-                    value="resnet50_robust", 
-                    label="Model"
-                )
-                
-                inference_type = gr.Dropdown(
-                    choices=["Prior-Guided Drift Diffusion", "IncreaseConfidence"], 
-                    value="Prior-Guided Drift Diffusion", 
-                    label="Inference Method"
-                )
+            # Run Inference button right below the image
+            run_button = gr.Button("🪄 Run Generative Inference", variant="primary", elem_classes="purple-button")
             
-            with gr.Row():
-                eps_slider = gr.Slider(minimum=0.01, maximum=3.0, value=20.0, step=0.01, label="Epsilon (Perturbation Size)")
-                iterations_slider = gr.Slider(minimum=1, maximum=600, value=101, step=1, label="Number of Iterations")  # Updated max to 600
+            # Parameters toggle button
+            params_button = gr.Button("⚙️ Play with the parameters", variant="secondary")
             
-            with gr.Row():
-                initial_noise_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.8, step=0.01, 
-                                               label="Initial Noise Ratio")
-                diffusion_noise_slider = gr.Slider(minimum=0.0, maximum=0.05, value=0.003, step=0.001, 
-                                                label="Diffusion Noise Ratio")  # Corrected name
+            # 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():
-                step_size_slider = gr.Slider(minimum=0.01, maximum=2.0, value=1.0, step=0.01, 
-                                           label="Step Size")  # Added step size slider
-                layer_choice = gr.Dropdown(
-                    choices=["all", "conv1", "bn1", "relu", "maxpool", "layer1", "layer2", "layer3", "layer4", "avgpool"],
-                    value="layer3",
-                    label="Model Layer"
-                )
-            
-            run_button = gr.Button("🪄 Run Inference", variant="primary", elem_classes="purple-button")
+                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
@@ -332,7 +340,7 @@ with gr.Blocks(title="Generative Inference Demo", css="""
                         image_input, model_choice, inference_type, 
                         eps_slider, iterations_slider,
                         initial_noise_slider, diffusion_noise_slider,
-                        step_size_slider, layer_choice
+                        step_size_slider, layer_choice, params_section
                     ]
                 )
             
@@ -341,17 +349,10 @@ with gr.Blocks(title="Generative Inference Demo", css="""
                 gr.Markdown(f"### {ex['name']}")
                 gr.Markdown(f"[Read more on Wikipedia]({ex['wiki']})")
                 
-                gr.Markdown("**Generative Inference Parameters:**")
-                params_md = f"""
-                - **Method**: {ex['method']}
-                - **Model Layer**: {ex['reverse_diff']['layer']}
-                - **Initial Noise**: {ex['reverse_diff']['initial_noise']}
-                - **Diffusion Noise**: {ex['reverse_diff']['diffusion_noise']}
-                - **Step Size**: {ex['reverse_diff']['step_size']}
-                - **Iterations**: {ex['reverse_diff']['iterations']}
-                - **Epsilon**: {ex['reverse_diff']['epsilon']}
-                """
-                gr.Markdown(params_md)
+                # 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("---")
@@ -368,6 +369,15 @@ with gr.Blocks(title="Generative Inference Demo", css="""
         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
@@ -381,12 +391,12 @@ with gr.Blocks(title="Generative Inference Demo", css="""
     Moving away from the least likely class identified at iteration 0 (fast perception)
     
     ### Parameters:
-    - **Initial Noise Ratio**: Controls the amount of noise added to the image at the beginning
-    - **Diffusion Noise Ratio**: Controls the amount of noise added at each optimization step
-    - **Step Size**: Learning rate for the optimization process
+    - **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**: Controls the size of perturbation during optimization
+    - **Epsilon (Stimulus Fidelity)**: Controls the size of perturbation during optimization
     
     **Generative Inference was developed by [Tahereh Toosi](https://toosi.github.io).**
     """)

inference.py

@@ -25,7 +25,8 @@ 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_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]
@@ -162,7 +163,12 @@ 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
     
-    model_path = Path(f"models/{model_type}.pt")
+    # 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]