remove tv loss, using BCE loss now

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+from torch.nn import BCEWithLogitsLoss
 import numpy as np
 from transformers import ViTImageProcessor, ViTForImageClassification
 from PIL import Image
@@ -16,12 +17,7 @@ def get_encoder_activations(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, tv_weight, iterations, n_targets, seed):
+def process_image(input_image, learning_rate, iterations, n_targets, seed):
     if input_image is None:
         return None
 
@@ -32,20 +28,22 @@ def process_image(input_image, learning_rate, tv_weight, iterations, n_targets,
 
     
     torch.manual_seed(int(seed))
-    random_indices = torch.randperm(1000)[:int(n_targets)].to(pixel_values.device)
+    random_one_logits = torch.zeros(1000)
+    random_one_logits[torch.randperm(1000)[:n_targets]] = 1
+    random_one_logits = random_one_logits.to(pixel_values.device)
 
     for iteration in range(int(iterations)):
         model.zero_grad()
         if pixel_values.grad is not None:
             pixel_values.grad.data.zero_()
 
-        final_activations = get_encoder_activations(pixel_values)
-        logits = model.classifier(final_activations[0])
+        final_activations = get_encoder_activations(pixel_values.to('cuda'))
+    
+        logits = model.classifier(final_activations[0]).to(pixel_values.device)
+    
+        original_loss = BCEWithLogitsLoss(reduction='sum')(logits,random_one_logits)
     
-        original_loss = logits[random_indices].sum()
-        tv_loss = total_variation_loss(pixel_values)
-        total_loss = original_loss - tv_weight * tv_loss
-        total_loss.backward()
+        original_loss.backward()
 
         with torch.no_grad():
             pixel_values.data += learning_rate * pixel_values.grad.data
@@ -60,11 +58,10 @@ iface = gr.Interface(
     fn=process_image,
     inputs=[
         gr.Image(type="pil"), 
-        gr.Number(value=16.0, minimum=0, label="Learning Rate"),
-        gr.Number(value=0.0001, label="Total Variation Loss"), 
-        gr.Number(value=4, minimum=1, label="Iterations"),
+        gr.Number(value=1.0, minimum=0, label="Learning Rate"),
+        gr.Number(value=2, minimum=1, label="Iterations"),
         gr.Number(value=420, minimum=0, label="Seed"),
-        gr.Number(value=500, minimum=1, maximum=1000, label="Number of Random Target Class Activations to Maximise"),
+        gr.Number(value=250, minimum=1, maximum=1000, label="Number of Random Target Class Activations to Maximise"),
     ],
     outputs=[gr.Image(type="numpy", label="ViT-Dreamed Image")]
 )