Update app.py

fix params and grad

app.py

@@ -64,74 +64,56 @@ from torchvision import transforms
 from PIL import Image
 import numpy as np
 
-# Simple Generative Inference function
-def simple_generative_inference(
-    image, mode, model, n_iterations=10, step_size=0.01, noise_ratio=0.1, eps=0.1
-):
-    """
-    Perform Generative Perceptual Inference on the input image.
-    :param image: Input image as a PIL image.
-    :param mode: Either 'increase confidence' or 'ReverseDiffuse'.
-    :param model: Pretrained PyTorch model.
-    :param n_iterations: Number of inference iterations.
-    :param step_size: Step size for gradient-based updates.
-    :param noise_ratio: Ratio of noise to be added in ReverseDiffuse mode.
-    :param eps: Constraint on perturbation magnitude.
-    :return: Processed image and gradient visualization.
-    """
+def simple_generative_inference(image, mode, model, n_iterations=10, step_size=0.01, eps=0.1, noise_ratio=0.1):
     # Preprocess image
     transform = transforms.Compose([
-        transforms.Resize((224, 224)),
+        transforms.Resize((224, 224)),  # Enforce fixed size
         transforms.ToTensor(),
-        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # Adjust normalization as needed
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
     ])
-    image_tensor = transform(image).unsqueeze(0)
-    image_tensor.requires_grad_(True)  # Enable gradient computation for the image tensor
-
-    optimizer = torch.optim.SGD([image_tensor], lr=step_size)
-
-    # Define least likely classes for "increase confidence" mode
-    if mode == "increase confidence":
-        with torch.no_grad():
-            output = model(image_tensor)
-            probs = torch.nn.functional.softmax(output, dim=1)
-            _, least_likely_classes = torch.topk(probs, k=5, largest=False, dim=1)
-
-    # Create noisy image (only for ReverseDiffuse mode)
-    if mode == "ReverseDiffuse":
-        noisy_image = image_tensor + torch.randn_like(image_tensor) * noise_ratio
+    image_tensor = transform(image).unsqueeze(0).requires_grad_(True)
+    image_tensor.retain_grad()  # Ensure gradients are retained for non-leaf tensor
 
     for _ in range(n_iterations):
-        optimizer.zero_grad()
+        # Zero gradients
+        if image_tensor.grad is not None:
+            image_tensor.grad.zero_()
+
+        # Forward pass
         output = model(image_tensor)
 
         # Define inference loss based on mode
         if mode == "increase confidence":
+            probs = torch.nn.functional.softmax(output, dim=1)
+            _, least_likely_indices = torch.topk(probs, k=2, largest=False)
             losses = []
-            for idx in least_likely_classes[0]:  # Iterate over least likely classes
+            for idx in least_likely_indices[0]:
                 target = torch.full((1,), idx, dtype=torch.long, device=output.device)
-                loss = torch.nn.functional.cross_entropy(output, target)
+                loss = torch.nn.CrossEntropyLoss()(output, target)
                 losses.append(loss)
-            loss = torch.stack(losses).mean()  # Average loss over least likely classes
+            loss = torch.stack(losses).mean()
         elif mode == "ReverseDiffuse":
+            noisy_image = image_tensor + torch.randn_like(image_tensor) * noise_ratio
             loss = torch.nn.functional.mse_loss(image_tensor, noisy_image)
         else:
             raise ValueError("Invalid mode selected. Choose 'increase confidence' or 'ReverseDiffuse'.")
 
-        # Compute gradients and update the image
+        # Backward pass
         loss.backward()
-        grad = image_tensor.grad.data
+
+        # Access gradient
+        grad = image_tensor.grad  # Gradient is now retained
         grad_norm = grad.view(grad.shape[0], -1).norm(dim=1, keepdim=True).view(grad.shape[0], 1, 1, 1)
-        scaled_grad = grad / (grad_norm + 1e-10)
-        image_tensor = torch.clamp(
-            image_tensor + step_size * scaled_grad,
-            min=image_tensor - eps,
-            max=image_tensor + eps
-        )
+        grad = grad / (grad_norm + 1e-10)  # Avoid division by zero
+
+        # Update image tensor
+        with torch.no_grad():
+            image_tensor += step_size * grad
+            image_tensor.clamp_(-eps, eps)  # Keep within range
 
     # Generate gradient visualization
-    grad_visualization = image_tensor.grad.abs().mean(dim=1).squeeze().cpu().numpy()
-    grad_image = (grad_visualization - grad_visualization.min()) / (grad_visualization.max() - grad_visualization.min())
+    grad_image = grad.abs().mean(dim=1).squeeze().cpu().numpy()
+    grad_image = (grad_image - grad_image.min()) / (grad_image.max() - grad_image.min())
     grad_image = Image.fromarray((grad_image * 255).astype(np.uint8))
 
     # Convert final processed image back to PIL format
@@ -149,17 +131,17 @@ iface = gr.Interface(
     inputs=[
         gr.Image(type="pil", label="Input Image"),  # Input image
         gr.Radio(["increase confidence", "ReverseDiffuse"], label="Inference Mode"),  # Mode selection
-        gr.Slider(0.001, 1.0, value=0.01, step=0.001, label="Step Size"),  # Step size
-        gr.Slider(0.001, 0.5, value=0.1, step=0.001, label="Epsilon (eps)"),  # Epsilon constraint
-        gr.Slider(0.0, 0.5, value=0.1, step=0.01, label="Noise Ratio"),  # Noise ratio
-        gr.Slider(1, 100, value=10, step=1, label="Number of Iterations"),  # Number of iterations
+        gr.Slider(0.1, 20, value=1, step=0.1, label="Step Size"),  # Step size
+        gr.Slider(0.1, 40, value=0.5, step=0.1, label="Epsilon (eps)"),  # Epsilon constraint
+        gr.Slider(0.0, 1.0, value=0.5, step=0.1, label="Noise Ratio"),  # Noise ratio
+        gr.Slider(1, 1000, value=100, step=1, label="Number of Iterations"),  # Number of iterations
     ],
     outputs=[
         gr.Image(label="Processed Image"),  # Processed image
         gr.Image(label="Gradient Visualization")  # Gradient visualization
     ],
-    title="Generative Perceptual Inference (GPI)",
-    description="Perform GPI on input images using adjustable parameters such as step size, epsilon, noise ratio, and number of iterations."
+    title="Generative Inference",
+    description="Perform generative inference on input images using adjustable parameters such as step size, epsilon, noise ratio, and number of iterations."
 )