fix various total variation bugs

app.py

@@ -21,7 +21,7 @@ def total_variation_loss(img):
     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_loss, iterations, n_targets, seed):
+def process_image(input_image, learning_rate, tv_weight, iterations, n_targets, seed):
     if input_image is None:
         return None
 
@@ -42,9 +42,9 @@ def process_image(input_image, learning_rate, tv_loss, iterations, n_targets, se
         final_activations = get_encoder_activations(pixel_values)
         logits = model.classifier(final_activations[0])
     
-        original_loss = -logits[random_indices].sum()
+        original_loss = logits[random_indices].sum()
         tv_loss = total_variation_loss(pixel_values)
-        total_loss = original_loss + 0.00625 * tv_loss
+        total_loss = original_loss - tv_weight * tv_loss
         total_loss.backward()
 
         with torch.no_grad():
@@ -60,11 +60,11 @@ iface = gr.Interface(
     fn=process_image,
     inputs=[
         gr.Image(type="pil"), 
-        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=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=420, minimum=0, label="Seed"),
-        gr.Number(value=50, minimum=1, maximum=1000, label="Number of Random Target Class Activations to Maximise"),
+        gr.Number(value=500, minimum=1, maximum=1000, label="Number of Random Target Class Activations to Maximise"),
     ],
     outputs=[gr.Image(type="numpy", label="ViT-Dreamed Image")]
 )