Fix: per-frame timestep allocation for video2world (image2world) mode in CosmosPredict2

[csy2077]

Problem

In CosmosPredict2.forward(), when running in video2world (image2world) mode, the conditioning (clean) frames were not assigned a special timestep of 0.0. Instead, the model received the same noisy timestep for all frames — including the clean conditioning frame(s). This caused the model to treat the clean first frame as a fully-noised frame, leading to:

• Dramatic quality degradation in generated videos
• No temporal coherence to the conditioning frame
• Effectively broken video2world / image2world distillation

Fix

After replacing the conditioning frames in model_input, expand t to per-frame shape (B, T) and zero out the timestep for conditioning frames (indicated by condition_mask). This tells the transformer that those frames are already clean and require no denoising:

t_expanded = t.unsqueeze(1).expand(B, T)
mask_B_T = condition_mask[:, 0, :, 0, 0]  # (B, T)
t = t_expanded * (1 - mask_B_T)

The transformer already accepts timesteps_B_T of shape (B, T), so no other changes are needed.

Impact

Without this fix, CosmosPredict2 video2world distillation produces incoherent videos that ignore the conditioning frame. With this fix, the model correctly preserves the conditioning frame and generates temporally consistent video.