Update video_to_video/video_to_video_model.py

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(os.path.join(model_path, 'model.pt')):
+    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(os.path.join(model_path, 'model.pt'), '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)
@@ -54,7 +54,7 @@ class VideoToVideo_sr():
         download_model(model_url, cfg.model_path)
 
         # 拼接完整路径
-        model_file_path = os.path.join(cfg.model_path, 'model.pt')
+        model_file_path = os.path.join('pretrained_weight', 'I2VGen-XL-based', 'heavy_deg.pt')
         print('model_file_path:', model_file_path)
 
         # 加载模型