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

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

* Fix ruff.

Author: comfyanonymous

comfy/model_base.py

@@ -1024,6 +1024,8 @@ def extra_conds(self, **kwargs):
     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])

main.py

@@ -185,7 +185,7 @@ def prompt_worker(q, server_instance):
 
             current_time = time.perf_counter()
             execution_time = current_time - execution_start_time
-            
+
             # Log Time in a more readable way after 10 minutes
             if execution_time > 600:
                 execution_time = time.strftime("%H:%M:%S", time.gmtime(execution_time))