biased inferneceimplemented, next up: saving ourput as a gif and the config json

app.py

@@ -11,7 +11,7 @@ except ImportError:
     
 import os
 import argparse
-from inference import GenerativeInferenceModel, get_inference_configs
+from inference import GenerativeInferenceModel, get_inference_configs, get_imagenet_labels
 
 # Parse command line arguments
 parser = argparse.ArgumentParser(description='Run Generative Inference Demo')
@@ -22,6 +22,9 @@ args = parser.parse_args()
 os.makedirs("models", exist_ok=True)
 os.makedirs("stimuli", exist_ok=True)
 
+# Load ImageNet labels for biased-inference dropdown (1000 classes)
+IMAGENET_LABELS = get_imagenet_labels()
+
 # Check if running on Hugging Face Spaces
 if "SPACE_ID" in os.environ:
     default_port = int(os.environ.get("PORT", 7860))
@@ -363,7 +366,8 @@ def run_inference(image, model_type, inference_type, eps_value, num_iterations,
                  initial_noise=0.05, diffusion_noise=0.3, step_size=0.8, model_layer="layer3",
                  use_adaptive_eps=False, use_adaptive_step=False,
                  mask_center_x=0.0, mask_center_y=0.0, mask_radius=0.3, mask_sigma=0.2,
-                 eps_max_mult=4.0, eps_min_mult=1.0, step_max_mult=4.0, step_min_mult=1.0):
+                 eps_max_mult=4.0, eps_min_mult=1.0, step_max_mult=4.0, step_min_mult=1.0,
+                 use_biased_inference=False, biased_class_name=""):
     # Check if image is provided
     if image is None:
         return None, "Please upload an image before running inference."
@@ -419,6 +423,14 @@ def run_inference(image, model_type, inference_type, eps_value, num_iterations,
     else:
         config['adaptive_step_size'] = None
     
+    # Biased inference: bias perception toward a target ImageNet class
+    use_biased_inference = bool(use_biased_inference) if use_biased_inference is not None else False
+    biased_class_name = (biased_class_name or "").strip() if biased_class_name else ""
+    if use_biased_inference and biased_class_name:
+        config['biased_inference'] = {'enable': True, 'class': biased_class_name}
+    else:
+        config['biased_inference'] = config.get('biased_inference') or {'enable': False, 'class': None}
+    
     # Run generative inference
     result = model.inference(image, model_type, config)
     
@@ -532,6 +544,8 @@ def apply_example(example):
         example.get("eps_min_mult", 1.0),
         example.get("step_max_mult", 4.0),
         example.get("step_min_mult", 1.0),
+        example.get("use_biased_inference", False),
+        example.get("biased_class_name", ""),
         mask_img,
         gr.Group(visible=True),
     ]
@@ -623,6 +637,17 @@ with gr.Blocks(title="Human Hallucination Prediction", css="""
                 with gr.Row():
                     step_max_mult_slider = gr.Slider(minimum=0.1, maximum=150.0, value=50.0, step=0.1, label="Step size: multiplier at center")
                     step_min_mult_slider = gr.Slider(minimum=0.1, maximum=10.0, value=0.2, step=0.1, label="Step size: multiplier at periphery")
+                gr.Markdown("### 🎯 Biased inference")
+                gr.Markdown("Bias the prediction toward a specific ImageNet category (1000 classes).")
+                with gr.Row():
+                    use_biased_inference_check = gr.Checkbox(value=False, label="Use biased inference (bias toward a target class)")
+                    biased_class_dropdown = gr.Dropdown(
+                        choices=[("— No bias —", "")] + [(label, label) for label in sorted(IMAGENET_LABELS)],
+                        value="",
+                        label="Target class",
+                        allow_custom_value=False,
+                        filterable=True,
+                    )
             
         with gr.Column(scale=2):
             # Outputs
@@ -655,6 +680,7 @@ with gr.Blocks(title="Human Hallucination Prediction", css="""
                         mask_radius_slider, mask_sigma_slider,
                         eps_max_mult_slider, eps_min_mult_slider,
                         step_max_mult_slider, step_min_mult_slider,
+                        use_biased_inference_check, biased_class_dropdown,
                         mask_preview,
                         params_section,
                     ],
@@ -686,6 +712,7 @@ with gr.Blocks(title="Human Hallucination Prediction", css="""
             mask_radius_slider, mask_sigma_slider,
             eps_max_mult_slider, eps_min_mult_slider,
             step_max_mult_slider, step_min_mult_slider,
+            use_biased_inference_check, biased_class_dropdown,
         ],
         outputs=[output_image, output_frames]
     )

inference.py

@@ -280,12 +280,16 @@ class InferStep:
             self.step_size = step_size
 
     def project(self, x):
-        # Per-pixel deviation from original; clamp by eps (or eps_map). No add/subtract of mask.
-        diff = x - self.orig_image
+        # Per-pixel L2 (color) constraint: scale diff so color_norm <= eps (or eps_map).
+        # Constraint is per-pixel L2 on RGB (spherical), not L∞ per channel.
+        diff = x - self.orig_image  # (B, C, H, W)
+        color_norm = torch.norm(diff, dim=1, keepdim=True)  # (B, 1, H, W)
         if self.use_adaptive_eps:
-            diff = torch.clamp(diff, -self.eps_map, self.eps_map)
+            # eps_map has shape (1, 1, H, W), broadcasts to (B, 1, H, W)
+            scale = torch.clamp(self.eps_map / (color_norm + 1e-10), max=1.0)
         else:
-            diff = torch.clamp(diff, -self.eps, self.eps)
+            scale = torch.clamp(self.eps / (color_norm + 1e-10), max=1.0)
+        diff = diff * scale
         return torch.clamp(self.orig_image + diff, 0, 1)
 
     def step(self, x, grad):
@@ -334,7 +338,8 @@ def get_inference_configs(inference_type='IncreaseConfidence', eps=0.5, n_itr=50
         'inference_normalization': 'off',  # 'on' or 'off' (match psychiatry implementation)
         'recognition_normalization': 'off',
         'iterations_to_show': get_iterations_to_show(n_itr),  # Dynamic iterations to visualize
-        'misc_info': {'keep_grads': False}  # Additional configuration
+        'misc_info': {'keep_grads': False},  # Additional configuration
+        'biased_inference': {'enable': False, 'class': None}  # Bias perception toward a target class (ImageNet label)
     }
     
     # Customize based on inference type
@@ -829,6 +834,30 @@ class GenerativeInferenceModel:
             adaptive_step_config=adaptive_step,
         )
         
+        # Biased inference: resolve class name to index (ImageNet simple labels, case-insensitive)
+        biased_inference_config = config.get('biased_inference', {'enable': False, 'class': None})
+        biased_class_index = None
+        biased_class_tensor = None
+        if biased_inference_config.get('enable', False):
+            class_name = biased_inference_config.get('class') or biased_inference_config.get('class_name')
+            if class_name:
+                try:
+                    biased_class_index = next(
+                        i for i, label in enumerate(self.labels)
+                        if label.lower() == class_name.lower()
+                    )
+                    biased_class_tensor = torch.tensor(
+                        [biased_class_index], device=device, dtype=torch.long
+                    )
+                    print(f"Biased inference: biasing toward class '{self.labels[biased_class_index]}' (index {biased_class_index})")
+                except StopIteration:
+                    raise ValueError(
+                        f"biased_inference class '{class_name}' not found in ImageNet simple labels. "
+                        "Use a label from imagenet-simple-labels (e.g. 'goldfish', 'tabby cat')."
+                    )
+            else:
+                print("Biased inference enabled but no class specified; ignoring.")
+        
         # Storage for inference steps
         # Create a new tensor that requires gradients
         x = image_tensor.clone().detach().requires_grad_(True)
@@ -934,6 +963,13 @@ class GenerativeInferenceModel:
                     
                     grad = torch.autograd.grad(loss, x, retain_graph=True)[0]
                 
+                # Biased inference: subtract gradient that pushes toward target class
+                if biased_inference_config.get('enable', False) and biased_class_tensor is not None:
+                    output_full = model(x)
+                    loss_biased = F.cross_entropy(output_full, biased_class_tensor)
+                    grad_biased = torch.autograd.grad(loss_biased, x)[0]
+                    grad = grad - grad_biased
+                
                 if grad is None:
                     print("Warning: Direct gradient calculation failed")
                     # Fall back to random perturbation