From 097faeeff0a663cc71dd7dfbf7cb8e3389dffe7a Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 28 Aug 2024 21:23:07 +0300
Subject: [PATCH] expose chunk_size for encode node

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

diff --git a/nodes.py b/nodes.py
index 8c21e49..078192c 100644
--- a/nodes.py
+++ b/nodes.py
@@ -169,6 +169,9 @@ class CogVideoImageEncode:
             "pipeline": ("COGVIDEOPIPE",),
             "image": ("IMAGE", ),
             },
+            "optional": {
+                "chunk_size": ("INT", {"default": 16, "min": 1}),
+            },
         }
 
     RETURN_TYPES = ("LATENT",)
@@ -176,7 +179,7 @@ class CogVideoImageEncode:
     FUNCTION = "encode"
     CATEGORY = "CogVideoWrapper"
 
-    def encode(self, pipeline, image):
+    def encode(self, pipeline, image, chunk_size=16):
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
         generator = torch.Generator(device=device).manual_seed(0)
@@ -187,7 +190,7 @@ class CogVideoImageEncode:
         input_image = input_image.to(vae.dtype).to(device)
         input_image = input_image.unsqueeze(0).permute(0, 4, 1, 2, 3) # B, C, T, H, W
         B, C, T, H, W = input_image.shape
-        chunk_size = 16
+
         latents_list = []
         # Loop through the temporal dimension in chunks of 16
         for i in range(0, T, chunk_size):