diff --git a/fastgen/networks/cosmos_predict2/network.py b/fastgen/networks/cosmos_predict2/network.py index ff85360..adbfba6 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 @@ -1352,10 +1365,24 @@ 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) + 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) + 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,