use n random target classes to maximise activation for

app.py

@@ -11,7 +11,7 @@ model = ViTForImageClassification.from_pretrained('google/vit-large-patch32-384'
 model.to(device)
 model.eval()
 
-def process_image(input_image, learning_rate, iterations):
+def process_image(input_image, learning_rate, iterations, n_targets, seed):
     if input_image is None:
         return None
     
@@ -25,13 +25,18 @@ def process_image(input_image, learning_rate, iterations):
     pixel_values = pixel_values.to(device)
     pixel_values.requires_grad_(True)
 
+    
+    torch.manual_seed(int(seed))
+    random_indices = torch.randperm(1000)[:n_targets].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)
-        target_sum = final_activations.sum()
+        logits = model.classifier(final_activations[0])
+        target_sum = logits[random_indices].sum()
         target_sum.backward()
 
         with torch.no_grad():
@@ -48,7 +53,9 @@ iface = gr.Interface(
     inputs=[
         gr.Image(type="pil"), 
         gr.Number(value=4.0, label="Learning Rate"), 
-        gr.Number(value=4, label="Iterations")
+        gr.Number(value=4, label="Iterations"),
+        gr.Number(value=420, label="Seed"),
+        gr.Number(value=50, minimum=1, maximum=1000, label="n target classes"),
     ],
     outputs=[gr.Image(type="numpy", label="ViT-Dreamed Image")]
 )