updates

video_to_video/video_to_video_model.py

@@ -17,10 +17,10 @@ from diffusers import AutoencoderKLTemporalDecoder
 import requests
 
 def download_model(url, model_path):
-    if not os.path.exists(model_path):
+    if not os.path.exists(os.path.join(model_path, 'heavy_deg.pt')):
         print(f"Model not found at {model_path}, downloading...")
         response = requests.get(url, stream=True)
-        with open(model_path, 'wb') as f:
+        with open(os.path.join(model_path, 'heavy_deg.pt'), 'wb') as f:
             for chunk in response.iter_content(chunk_size=1024):
                 if chunk:
                     f.write(chunk)
@@ -37,7 +37,7 @@ class VideoToVideo_sr():
         self.device = device # torch.device(f'cuda:0')
 
         # text_encoder
-        text_encoder = FrozenOpenCLIPEmbedder(device=self.device, pretrained="laion2b_s32b_b79k")
+        text_encoder = FrozenOpenCLIPEmbedder(device=self.device, pretrained="/home/test/Workspace/yhliu/VSR/ours/checkpoints/open_clip_pytorch_model.bin")
         text_encoder.model.to(self.device)
         self.text_encoder = text_encoder
         logger.info(f'Build encoder with FrozenOpenCLIPEmbedder')
@@ -80,7 +80,7 @@ class VideoToVideo_sr():
 
         # Temporal VAE
         vae = AutoencoderKLTemporalDecoder.from_pretrained(
-            "stabilityai/stable-video-diffusion-img2vid", subfolder="vae", variant="fp16"
+            "/home/test/Workspace/yhliu/VSR/ours/checkpoints/stable-video-diffusion-img2vid", subfolder="vae", variant="fp16"
         )
         vae.eval()
         vae.requires_grad_(False)