298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-07-26 01:17:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix exception when using a noise mask with cosmos predict2. (#8621)

Browse Source 
* Fix exception when using a noise mask with cosmos predict2.

* Fix ruff.
This commit is contained in:
comfyanonymous 2025-06-21 00:30:39 -07:00 committed by GitHub
parent 31ca603ccb
commit 1883e70b43
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
comfy/model_base.py
View File

@ -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])