text2video

Source Code/annotator/openpose/__init__.py

@@ -18,8 +18,18 @@ class OpenposeDetector:
         body_modelpath = os.path.join(annotator_ckpts_path, "body_pose_model.pth")
         hand_modelpath = os.path.join(annotator_ckpts_path, "hand_pose_model.pth")
 
-        if not os.path.exists(hand_modelpath) or not os.path.exists(body_modelpath):
-             raise FileNotFoundError(f"Openpose models not found in {annotator_ckpts_path}. Please ensure body_pose_model.pth and hand_pose_model.pth are present.")
+        if not os.path.exists(annotator_ckpts_path):
+            os.makedirs(annotator_ckpts_path)
+
+        if not os.path.exists(body_modelpath):
+            from torch.hub import download_url_to_file
+            print(f"Downloading body_pose_model from {body_model_path}...")
+            download_url_to_file(body_model_path, body_modelpath)
+
+        if not os.path.exists(hand_modelpath):
+            from torch.hub import download_url_to_file
+            print(f"Downloading hand_pose_model from {hand_model_path}...")
+            download_url_to_file(hand_model_path, hand_modelpath)
 
         self.body_estimation = Body(body_modelpath)
         self.hand_estimation = Hand(hand_modelpath)