Update inference.py

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