From 411a2c5bfade92682ea829933d6569cbe5a535b4 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sat, 21 Jun 2025 03:24:52 -0400 Subject: [PATCH] Fix exception when using a noise mask with cosmos predict2. --- comfy/model_base.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/comfy/model_base.py b/comfy/model_base.py index cb7689e84..75ec42699 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -1024,6 +1024,8 @@ class CosmosPredict2(BaseModel): def process_timestep(self, timestep, x, denoise_mask=None, **kwargs): if denoise_mask is None: return timestep + if denoise_mask.ndim <= 4: + return timestep condition_video_mask_B_1_T_1_1 = denoise_mask.mean(dim=[1, 3, 4], keepdim=True) c_noise_B_1_T_1_1 = 0.0 * (1.0 - condition_video_mask_B_1_T_1_1) + timestep.reshape(timestep.shape[0], 1, 1, 1, 1) * condition_video_mask_B_1_T_1_1 out = c_noise_B_1_T_1_1.squeeze(dim=[1, 3, 4])