update

trainer.py

@@ -11,6 +11,7 @@ import sys
 import gradio as gr
 import slugify
 import torch
+import huggingface_hub
 from huggingface_hub import HfApi
 from omegaconf import OmegaConf
 
@@ -33,16 +34,20 @@ class Trainer:
         self.checkpoint_dir = pathlib.Path('checkpoints')
         self.checkpoint_dir.mkdir(exist_ok=True)
 
-    def download_base_model(self, base_model_id: str) -> str:
+    def download_base_model(self, base_model_id: str, token=None) -> str:
         model_dir = self.checkpoint_dir / base_model_id
         if not model_dir.exists():
             org_name = base_model_id.split('/')[0]
             org_dir = self.checkpoint_dir / org_name
             org_dir.mkdir(exist_ok=True)
         print(f'https://huggingface.co/{base_model_id}')
-        subprocess.run(shlex.split(
-            f'git clone https://huggingface.co/{base_model_id}'),
-                        cwd=org_dir)
+        try:
+            subprocess.run(shlex.split(
+                f'git clone https://huggingface.co/{base_model_id}'),
+                            cwd=org_dir)
+        except:
+            temp_path = huggingface_hub.snapshot_download(base_model_id, use_auth_token=token)
+            subprocess.run(shlex.split(f'mv {temp_path} {org_dir}'))
         return model_dir.as_posix()
 
     def join_model_library_org(self, token: str) -> None:
@@ -241,7 +246,7 @@ class Trainer:
                 self.hf_token if self.hf_token else input_token)
 
         config = OmegaConf.load('Video-P2P/configs/man-skiing.yaml')
-        config.pretrained_model_path = self.download_base_model(tuned_model)
+        config.pretrained_model_path = self.download_base_model(tuned_model, token=input_token)
         config.output_dir = output_dir.as_posix()
         config.train_data.video_path = training_video.name  # type: ignore
         config.train_data.prompt = training_prompt