Update app.py

app.py

@@ -21,6 +21,8 @@ from detectron2.utils.visualizer import ColorMode, Visualizer
 from detectron2.data import MetadataCatalog
 from detectron2.engine import DefaultPredictor
 
+from huggingface_hub import hf_hub_download
+
 import gradio as gr
 
 
@@ -30,7 +32,8 @@ add_vit_config(cfg)
 cfg.merge_from_file("cascade_dit_base.yml")
 
 # Step 2: add model weights URL to config
-cfg.MODEL.WEIGHTS = "publaynet_dit-b_cascade.pth"
+filepath = hf_hub_download(repo_id="Sebas6k/DiT_weights", filename="publaynet_dit-b_cascade.pth", repo_type="model")
+cfg.MODEL.WEIGHTS = filepath
 
 # Step 3: set device
 cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"