mstar.model.lingbot_world.components.schedule

mstar.model.lingbot_world.components.schedule#

The causal-fast sampling schedule, ported exactly from the reference.

Provenance: the step picks and re-noising loop follow LingBot-World v1’s WanI2VFast.generate (Robbyant/lingbot-world, Apache-2.0); FlowUniPCMultistepScheduler’s shifted sigmas are Wan-Video/Wan2.1’s (Apache-2.0). The step indices and shift are the V2 1.3B checkpoint’s settings (config.py).

The reference builds a FlowUniPCMultistepScheduler (shift 1), calls set_timesteps(1000, shift=shift) and picks four entries by index. It never calls the solver’s step: each step converts the flow prediction to x0 and, except after the last step, re-noises x0 to the next timestep with add_noise. The tables below are built with the same numpy/torch ops, so the sigmas and int64 timesteps are bit-identical.

The int64 truncation makes neighbouring timesteps collide (about 0.3 timestep per index near t=1000 at shift 5), and the two lookups resolve a collision differently: the x0 conversion takes the FIRST match (argmin), add_noise the SECOND (index_for_timestep’s pos = 1). Both are kept as is.

Functions

add_noise(x0, noise, timestep, schedule)

(1 - sigma) * x0 + sigma * noise at the sample's dtype.

flow_to_x0(flow_pred, xt, timestep, schedule)

x0 = x_t - sigma_t * flow, in float64, cast back to the prediction's dtype.

make_fast_schedule(shift, timesteps_index[, ...])

Classes

FastSchedule(sigmas, timesteps, step_timesteps)

class mstar.model.lingbot_world.components.schedule.FastSchedule(sigmas: torch.Tensor, timesteps: torch.Tensor, step_timesteps: tuple[torch.Tensor, ...])[source]#

Bases: object

Parameters:
noise_sigma(timestep)[source]#

add_noise’s sigma: the second match when the timestep repeats.

Parameters:

timestep (Tensor)

Return type:

Tensor

sigmas: Tensor#
step_timesteps: tuple[Tensor, ...]#
timesteps: Tensor#
x0_sigma(timestep)[source]#

_convert_flow_pred_to_x0’s sigma: first-match argmin, float64.

Parameters:

timestep (Tensor)

Return type:

Tensor

mstar.model.lingbot_world.components.schedule.add_noise(x0, noise, timestep, schedule)[source]#

(1 - sigma) * x0 + sigma * noise at the sample’s dtype.

Parameters:
Return type:

Tensor

mstar.model.lingbot_world.components.schedule.flow_to_x0(flow_pred, xt, timestep, schedule)[source]#

x0 = x_t - sigma_t * flow, in float64, cast back to the prediction’s dtype.

Parameters:
Return type:

Tensor

mstar.model.lingbot_world.components.schedule.make_fast_schedule(shift, timesteps_index, num_train_timesteps=1000)[source]#
Parameters:
Return type:

FastSchedule