Address review: move per-frame timestep into forward(), robust t handling

Author: Siyi Chen

fastgen/networks/cosmos_predict2/network.py

@@ -1232,20 +1232,29 @@ 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
                 cond_with_mask = condition
                 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,