xinyun/ComfyUI-CogVideoXWrapper
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-CogVideoXWrapper.git synced 2025-12-16 00:14:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

expose chunk_size for encode node

Browse Source
This commit is contained in:
kijai 2024-08-28 21:23:07 +03:00
parent fe1dded986
commit 097faeeff0
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
nodes.py
View File

@ -169,6 +169,9 @@ class CogVideoImageEncode:
"pipeline": ("COGVIDEOPIPE",), "pipeline": ("COGVIDEOPIPE",),
"image": ("IMAGE", ), "image": ("IMAGE", ),
}, },
"optional": {
"chunk_size": ("INT", {"default": 16, "min": 1}),
},
} }
RETURN_TYPES = ("LATENT",) RETURN_TYPES = ("LATENT",)
@ -176,7 +179,7 @@ class CogVideoImageEncode:
FUNCTION = "encode" FUNCTION = "encode"
CATEGORY = "CogVideoWrapper" CATEGORY = "CogVideoWrapper"
def encode(self, pipeline, image): def encode(self, pipeline, image, chunk_size=16):
device = mm.get_torch_device() device = mm.get_torch_device()
offload_device = mm.unet_offload_device() offload_device = mm.unet_offload_device()
generator = torch.Generator(device=device).manual_seed(0) 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.to(vae.dtype).to(device)
input_image = input_image.unsqueeze(0).permute(0, 4, 1, 2, 3) # B, C, T, H, W 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 B, C, T, H, W = input_image.shape
chunk_size = 16
latents_list = [] latents_list = []
# Loop through the temporal dimension in chunks of 16 # Loop through the temporal dimension in chunks of 16
for i in range(0, T, chunk_size): for i in range(0, T, chunk_size):