text2video

.gitattributes

@@ -27,3 +27,5 @@
 *.jpeg filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
 *.ipynb filter=lfs diff=lfs merge=lfs -text
+Source[[:space:]]Code/annotator/ckpts/*.pth filter=lfs diff=lfs merge=lfs -text
+Source[[:space:]]Code/annotator/ckpts/*.pt filter=lfs diff=lfs merge=lfs -text

Source Code/annotator/uniformer/__init__.py

@@ -11,10 +11,17 @@ checkpoint_file = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/ann
 class UniformerDetector:
     def __init__(self):
         modelpath = os.path.join(annotator_ckpts_path, "upernet_global_small.pth")
+        if not os.path.exists(annotator_ckpts_path):
+            os.makedirs(annotator_ckpts_path)
+            
         if not os.path.exists(modelpath):
-            raise FileNotFoundError(f"Uniformer model not found in {annotator_ckpts_path}. Please ensure upernet_global_small.pth is present.")
+            from torch.hub import download_url_to_file
+            print(f"Downloading upernet_global_small from {checkpoint_file}...")
+            download_url_to_file(checkpoint_file, modelpath)
+            
         config_file = os.path.join(os.path.dirname(annotator_ckpts_path), "uniformer", "exp", "upernet_global_small", "config.py")
-        self.model = init_segmentor(config_file, modelpath).cuda()
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = init_segmentor(config_file, modelpath, device=device)
 
     def __call__(self, img):
         result = inference_segmentor(self.model, img)