Update app.py

app.py

@@ -71,9 +71,9 @@ def simple_generative_inference(image, mode, model, n_iterations=10, step_size=0
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
     ])
-    image_tensor = transform(image).unsqueeze(0)
+    image_original = transform(image).unsqueeze(0)
 
-    image_tensor = image_tensor.detach().clone()
+    image_tensor = image_original.detach().clone()
     image_tensor.requires_grad = True
     # image_tensor.retain_grad()  # Ensure gradients are retained for non-leaf tensor
 
@@ -119,7 +119,9 @@ def simple_generative_inference(image, mode, model, n_iterations=10, step_size=0
 
         # Update image tensor
         image_tensor = image_tensor + step_size * grad
-        image_tensor = torch.clamp(image_tensor, -eps, eps)   # Keep within range
+        delta = image_tensor - image_original
+        delta = torch.clamp(delta, -eps, eps)   # Keep within range
+        image_tensor = torch.clamp(image_original + diff, 0, 1)
         image_tensor = image_tensor.clone().detach().requires_grad_(True)  # Ensure it's a new leaf tensor
         image_tensor.retain_grad()