zerogpu

video_to_video/video_to_video_model.py

@@ -14,6 +14,20 @@ from video_to_video.diffusion.schedules_sdedit import noise_schedule
 from video_to_video.utils.logger import get_logger
 
 from diffusers import AutoencoderKLTemporalDecoder
+import requests
+
+def download_model(url, model_path):
+    if not os.path.exists(model_path):
+        print(f"Model not found at {model_path}, downloading...")
+        response = requests.get(url, stream=True)
+        with open(model_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=1024):
+                if chunk:
+                    f.write(chunk)
+        print(f"Model downloaded to {model_path}")
+    else:
+        print(f"Model found at {model_path}, skipping download.")
+
 
 logger = get_logger()
 
@@ -34,6 +48,10 @@ class VideoToVideo_sr():
         generator.eval()
 
         cfg.model_path = opt.model_path
+        # download weight
+        model_url = 'https://huggingface.co/SherryX/STAR/resolve/main/I2VGen-XL-based/heavy_deg.pt'
+        download_model(model_url, cfg.model_path)
+
         load_dict = torch.load(cfg.model_path, map_location='cpu')
         if 'state_dict' in load_dict:
             load_dict = load_dict['state_dict']