add total variation loss + tuning changes

app.py

@@ -11,14 +11,19 @@ model = ViTForImageClassification.from_pretrained('google/vit-large-patch32-384'
 model.to(device)
 model.eval()
 
+def get_encoder_activations(x):
+    encoder_output = model.vit(x)
+    final_activations = encoder_output.last_hidden_state[:,0,:]
+    return final_activations
+
+def total_variation_loss(img):
+    pixel_dif1 = img[:, :, 1:, :] - img[:, :, :-1, :]
+    pixel_dif2 = img[:, :, :, 1:] - img[:, :, :, :-1]
+    return (torch.sum(torch.abs(pixel_dif1)) + torch.sum(torch.abs(pixel_dif2)))
+
 def process_image(input_image, learning_rate, iterations, n_targets, seed):
     if input_image is None:
         return None
-    
-    def get_encoder_activations(x):
-        encoder_output = model.vit(x)
-        final_activations = encoder_output.last_hidden_state[:,0,:]
-        return final_activations
 
     image = input_image.convert('RGB')
     pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
@@ -36,8 +41,11 @@ def process_image(input_image, learning_rate, iterations, n_targets, seed):
 
         final_activations = get_encoder_activations(pixel_values)
         logits = model.classifier(final_activations[0])
-        target_sum = logits[random_indices].sum()
-        target_sum.backward()
+    
+        original_loss = -logits[random_indices].sum()
+        tv_loss = total_variation_loss(pixel_values)
+        total_loss = original_loss + 0.00625 * tv_loss
+        total_loss.backward()
 
         with torch.no_grad():
             pixel_values.data += learning_rate * pixel_values.grad.data
@@ -52,9 +60,10 @@ iface = gr.Interface(
     fn=process_image,
     inputs=[
         gr.Image(type="pil"), 
-        gr.Number(value=4.0, label="Learning Rate"), 
-        gr.Number(value=4, label="Iterations"),
-        gr.Number(value=420, label="Seed"),
+        gr.Number(value=10.0, minimum=0, label="Learning Rate"),
+        gr.Number(value=0.00625, label="Total Variation Loss"), 
+        gr.Number(value=1, minimum=1, label="Iterations"),
+        gr.Number(value=420, minimum=0, label="Seed"),
         gr.Number(value=50, minimum=1, maximum=1000, label="Number of Random Target Class Activations to Maximise"),
     ],
     outputs=[gr.Image(type="numpy", label="ViT-Dreamed Image")]