From 3baf92d120a91842c84a9907883eda64882e90d6 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Wed, 15 Jan 2025 17:19:59 -0500
Subject: [PATCH] CosmosImageToVideoLatent batch_size now does something.

---
 comfy_extras/nodes_cosmos.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/comfy_extras/nodes_cosmos.py b/comfy_extras/nodes_cosmos.py
index b76ff950..bd35ddb0 100644
--- a/comfy_extras/nodes_cosmos.py
+++ b/comfy_extras/nodes_cosmos.py
@@ -71,8 +71,8 @@ class CosmosImageToVideoLatent:
             mask[:, :, -latent_temp.shape[-3]:] *= 0.0
 
         out_latent = {}
-        out_latent["samples"] = latent
-        out_latent["noise_mask"] = mask
+        out_latent["samples"] = latent.repeat((batch_size, ) + (1,) * (latent.ndim - 1))
+        out_latent["noise_mask"] = mask.repeat((batch_size, ) + (1,) * (mask.ndim - 1))
         return (out_latent,)