From 81c4cb3e17a1e833f9fa625f02c81110a249e357 Mon Sep 17 00:00:00 2001 From: csy2077 Date: Sat, 16 May 2026 14:29:36 +0000 Subject: [PATCH 1/3] Fix per-frame timestep allocation for video2world mode in CosmosPredict2 In video2world (image2world) mode, conditioning frames were receiving the same noisy timestep as all other frames. This caused the transformer to treat the clean conditioning frame as a fully-noised input, breaking temporal coherence and causing severe quality degradation. Fix: after replacing conditioning frames in model_input, expand t to shape (B, T) and zero out timesteps for frames where condition_mask=1, signaling to the model that those frames are already clean (t=0). Co-Authored-By: Claude Sonnet 4.6 --- fastgen/networks/cosmos_predict2/network.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/fastgen/networks/cosmos_predict2/network.py b/fastgen/networks/cosmos_predict2/network.py index ff85360..65b1536 100644 --- a/fastgen/networks/cosmos_predict2/network.py +++ b/fastgen/networks/cosmos_predict2/network.py @@ -1352,6 +1352,13 @@ def forward( conditioning_latents_full = conditioning_latents # Replace conditioning frames model_input = conditioning_latents_full * condition_mask_C + x_t * (1 - condition_mask_C) + # Per-frame timesteps: assign timestep 0 to conditioning (clean) frames so the + # model knows they are noise-free and should not be denoised. Without this, the + # model treats the clean first frame as a fully-noised frame, causing incoherent + # video2world (image2world) generation. + 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) model_outputs = self.transformer( x_B_C_T_H_W=model_input, From 794b3e277974d3cebb2f4e823d99470049ed3874 Mon Sep 17 00:00:00 2001 From: Siyi Chen Date: Sat, 4 Jul 2026 03:48:12 -0500 Subject: [PATCH 2/3] Address review: move per-frame timestep into forward(), robust t handling --- fastgen/networks/cosmos_predict2/network.py | 34 ++++++++++++++++----- 1 file changed, 27 insertions(+), 7 deletions(-) diff --git a/fastgen/networks/cosmos_predict2/network.py b/fastgen/networks/cosmos_predict2/network.py index 65b1536..fa16cbe 100644 --- a/fastgen/networks/cosmos_predict2/network.py +++ b/fastgen/networks/cosmos_predict2/network.py @@ -1232,8 +1232,17 @@ def sample( t_per_frame = t # Wrap condition with mask for forward() to use - cond_with_mask = {"text_embeds": condition, "condition_mask": condition_mask} - neg_cond_with_mask = {"text_embeds": neg_condition, "condition_mask": condition_mask} + # cond_with_mask = {"text_embeds": condition, "condition_mask": condition_mask} + cond_with_mask = { + "text_embeds": condition, + "conditioning_latents": conditioning_latents, + "condition_mask": condition_mask, + } + neg_cond_with_mask = { + "text_embeds": neg_condition, + "conditioning_latents": conditioning_latents, + "condition_mask": condition_mask, + } else: model_input = latents t_per_frame = t @@ -1241,11 +1250,11 @@ def sample( neg_cond_with_mask = neg_condition # Forward pass - velocity_pred = self(model_input, t_per_frame, cond_with_mask, fps=fps) + velocity_pred = self(model_input, t_per_frame, cond_with_mask, fps=fps, conditional_frame_timestep=conditional_frame_timestep) # Classifier-free guidance if guidance_scale > 1.0: - velocity_uncond = self(model_input, t_per_frame, neg_cond_with_mask, fps=fps) + velocity_uncond = self(model_input, t_per_frame, neg_cond_with_mask, fps=fps, conditional_frame_timestep=conditional_frame_timestep) velocity_pred = velocity_uncond + guidance_scale * (velocity_pred - velocity_uncond) # Replace velocity for conditioning frames with analytical velocity: v = noise - x0 @@ -1278,6 +1287,7 @@ def forward( fps: Optional[torch.Tensor] = None, padding_mask: Optional[torch.Tensor] = None, skip_layers: Optional[List[int]] = None, + conditional_frame_timestep: float = 0.0, **fwd_kwargs, ) -> Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: """ @@ -1304,6 +1314,8 @@ def forward( fps: Frames per second tensor padding_mask: Padding mask tensor skip_layers: List of block indices to skip during forward pass + conditional_frame_timestep: Timestep value for conditioning frames (default 0.0). + Use 0.0 to indicate clean frames with no noise. Returns: Depending on the arguments: @@ -1340,6 +1352,7 @@ def forward( # Video2world training: replace input frames with conditioning latents model_input = x_t + transformer_t = t if conditioning_latents is not None and condition_mask is not None: B, C, T, H, W = x_t.shape # Expand condition_mask to channel dimension @@ -1356,13 +1369,20 @@ def forward( # model knows they are noise-free and should not be denoised. Without this, the # model treats the clean first frame as a fully-noised frame, causing incoherent # video2world (image2world) generation. - t_expanded = t.unsqueeze(1).expand(B, T) + # t_expanded = t.unsqueeze(1).expand(B, T) + if t.ndim == 1: + t_expanded = t.unsqueeze(1).expand(B, T) + else: + t_expanded = t.expand(B, T) mask_B_T = condition_mask[:, 0, :, 0, 0] # (B, T) - t = t_expanded * (1 - mask_B_T) + transformer_t = conditional_frame_timestep * mask_B_T + t_expanded * (1 - mask_B_T) + + # Reshape for convert_model_output broadcasting against (B, C, T, H, W) + t = transformer_t.reshape(B, 1, T, 1, 1) model_outputs = self.transformer( x_B_C_T_H_W=model_input, - timesteps_B_T=t, + timesteps_B_T=transformer_t, crossattn_emb=text_embeds, fps=fps, padding_mask=padding_mask, From 972300bef1a9760be9d7b286cce9db230bdb6df9 Mon Sep 17 00:00:00 2001 From: Siyi Chen <133584269+csy2077@users.noreply.github.com> Date: Sun, 5 Jul 2026 09:21:16 +0800 Subject: [PATCH 3/3] Apply suggestions from code review Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- fastgen/networks/cosmos_predict2/network.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/fastgen/networks/cosmos_predict2/network.py b/fastgen/networks/cosmos_predict2/network.py index fa16cbe..adbfba6 100644 --- a/fastgen/networks/cosmos_predict2/network.py +++ b/fastgen/networks/cosmos_predict2/network.py @@ -1234,10 +1234,10 @@ def sample( # Wrap condition with mask for forward() to use # cond_with_mask = {"text_embeds": condition, "condition_mask": condition_mask} cond_with_mask = { - "text_embeds": condition, - "conditioning_latents": conditioning_latents, - "condition_mask": condition_mask, - } + "text_embeds": condition, + "conditioning_latents": conditioning_latents, + "condition_mask": condition_mask, + } neg_cond_with_mask = { "text_embeds": neg_condition, "conditioning_latents": conditioning_latents,