From e1a784d92bd2dd3fbb5e39a485af5ad5e6377e4b Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 7 Oct 2024 22:15:06 +0300
Subject: [PATCH] error out if trying to use I2V model with other frame counts
 than 49

---
 nodes.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index 1125ae8..7af169b 100644
--- a/nodes.py
+++ b/nodes.py
@@ -419,7 +419,8 @@ class DownloadAndLoadCogVideoModel:
             "base_path": base_path,
             "onediff": True if compile == "onediff" else False,
             "cpu_offloading": enable_sequential_cpu_offload,
-            "scheduler_config": scheduler_config
+            "scheduler_config": scheduler_config,
+            "model_name": model
         }
 
         return (pipeline,)
@@ -588,7 +589,8 @@ class DownloadAndLoadCogVideoGGUFModel:
             "base_path": model,
             "onediff": True if compile == "onediff" else False,
             "cpu_offloading": enable_sequential_cpu_offload,
-            "scheduler_config": scheduler_config
+            "scheduler_config": scheduler_config,
+            "model_name": model
         }
 
         return (pipeline,)
@@ -868,6 +870,7 @@ class CogVideoSampler:
         base_path = pipeline["base_path"]
 
         assert "fun" not in base_path.lower(), "'Fun' models not supported in 'CogVideoSampler', use the 'CogVideoXFunSampler'"
+        assert ("I2V" not in pipeline["model_name"] or num_frames == 49 or context_options is not None), "I2V model can only do 49 frames"
 
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()