fix repo_name, add css

app.py

@@ -57,9 +57,10 @@ def download_bin(model):
     else:
         raise NameError("model not found: " + model)
     
+    repo_name = 'facebook/' + model
     bin_path = os.path.join(MODEL_DIR,bin_file)
     if not os.path.isfile(bin_path):
-        model_bin = hf_hub_download(repo_id=REPO_ID, filename='pytorch_model.bin',local_dir=MODEL_DIR,local_dir_use_symlinks=True,token=HF_TOKEN)
+        model_bin = hf_hub_download(repo_id=repo_name, filename='pytorch_model.bin',local_dir=MODEL_DIR,local_dir_use_symlinks=True,token=HF_TOKEN)
         os.rename(model_bin, bin_path)
 
 
@@ -95,11 +96,12 @@ input_img = gr.Image(shape=(250,250))
 input_button = gr.Radio(["min", "max", "mean"], value="min",label="Attention Head Fusion", info="How to combine the last layer attention across all 12 heads of the transformer.")
 output_img = gr.Image(shape=(250,250))
 output_plot = gr.Plot()
+css = ".output-image, .input-image, .image-preview {height: 600px !important}"
 
 markdown ="This is a demo for the Visual Cortex models. When passed an image input, it displays the attention of the last layer of the transformer.\n \
          The user can decide how the attention heads will be combined. \
          Along with the attention heatmap, it also displays the embedding values reshaped to a 16x48 or 16x64 grid."
 demo = gr.Interface(fn=run_attn, title="Visual Cortex Large Model", description=markdown,
                     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,input_button],outputs=[output_img,output_plot])
+                    inputs=[input_img,model_type,input_button],outputs=[output_img,output_plot],css=css)
 demo.launch()