removed redundancies

app.py

@@ -90,16 +90,14 @@ def plot_attention(image, layer_num):
     attn = attn_call.locals['attn'][0]
     with torch.inference_mode():
         # loop over attention heads
-        attention_block_num = 0
         for h_weights in attn:
             h_weights = h_weights.mean(axis=-2)  # average over all attention keys
             h_weights = h_weights[1:]  # skip the [class] token
-            attention_block_num += 1
-            output_img = plot_weights(input_data, h_weights, attention_block_num)
+            output_img = plot_weights(input_data, h_weights)
             attention_map_outputs.append(output_img)
     return attention_map_outputs
 
-def plot_weights(input_data, patch_weights, num_attention_block):
+def plot_weights(input_data, patch_weights):
     """Display the image: Brighter the patch, higher is the attention."""
     # multiply each patch of the input image by the corresponding weight
     plot = inv_normalize(input_data.clone())
@@ -111,7 +109,7 @@ def plot_weights(input_data, patch_weights, num_attention_block):
     attn_map_img = attn_map_img.resize((224, 224), Image.ANTIALIAS)
     return attn_map_img
 
-attention_maps = plot_attention(random_image, 6)
+attention_maps = plot_attention(random_image, "6")
 
 title_classify = "Image Based Plant Disease Identification 🍃🤓"
 
@@ -152,5 +150,7 @@ attention_interface = gr.Interface(
 )
 
 demo = gr.TabbedInterface([classify_interface, attention_interface], 
-                          ["Identify Disease", "Visualize Attention Map"], 
-                          title="NatureAI Diagnostics🧑🩺").launch(debug=False, share=True)
+                          ["Identify Disease", "Visualize Attention Map"], title="NatureAI Diagnostics🧑🩺")
+
+if __name__ == "__main__":
+    demo.launch()