slider

app.py

@@ -63,7 +63,7 @@ def download_bin(model):
         os.rename(model_bin, bin_path)
 
 
-def run_attn(input_img, model="vc1-base"):
+def run_attn(input_img, model="vc1-base",discard_ratio=0.89):
     download_bin(model)
     model, embedding_dim, transform, metadata = get_model(model)
     if input_img.shape[0] != 3:
@@ -75,7 +75,7 @@ def run_attn(input_img, model="vc1-base"):
     input_img = resize_transform(input_img)
     x = transform(input_img)
 
-    attention_rollout = VITAttentionGradRollout(model,head_fusion="max",discard_ratio=0.89)
+    attention_rollout = VITAttentionGradRollout(model,head_fusion="max",discard_ratio=discard_ratio)
     
     y = model(x)
     mask = attention_rollout.get_attn_mask()
@@ -85,10 +85,12 @@ def run_attn(input_img, model="vc1-base"):
 model_type = gr.Dropdown(
             ["vc1-base", "vc1-large"], label="Model Size", value="vc1-base")
 input_img = gr.Image(shape=(250,250))
+discard_ratio = gr.Slider(0,1,value=0.89)
+
 output_img = gr.Image(shape=(250,250))
 css = "#component-2, .input-image, .image-preview {height: 240px !important}"
 markdown ="This is a demo for the Visual Cortex models. When passed an image input, it displays the attention(green) of the last layer of the transformer."
 demo = gr.Interface(fn=run_attn, title="Visual Cortex Model", description=markdown,
-                    examples=[[os.path.join('./imgs',x),None]for x in os.listdir(os.path.join(os.getcwd(),'imgs')) if 'jpg' in x],
-                    inputs=[input_img,model_type],outputs=output_img,css=css)
+                    examples=[[os.path.join('./imgs',x),None,None]for x in os.listdir(os.path.join(os.getcwd(),'imgs')) if 'jpg' in x],
+                    inputs=[input_img,model_type,discard_ratio],outputs=output_img,css=css)
 demo.launch()