Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
csy2077 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

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

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
81c4cb3
Fix per-frame timestep allocation for video2world mode in CosmosPredict2
csy2077 May 16, 2026
794b3e2
Address review: move per-frame timestep into forward(), robust t hand…
csy2077 Jul 4, 2026
972300b
Apply suggestions from code review
csy2077 Jul 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 32 additions & 5 deletions fastgen/networks/cosmos_predict2/network.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]]:
"""
Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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:
Comment thread
csy2077 marked this conversation as resolved.
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,
Expand Down