mstar.model.waypoint.components.dit

mstar.model.waypoint.components.dit#

The Waypoint-1.5 DiT: 24 blocks, the 4+1 pass driver, and the world clock.

Five forwards per generated frame: four frozen Euler denoise passes over config.scheduler_sigmas, then one committing pass at sigma=0 that writes the ring. Two clocks are threaded separately: f_pos (buckets, slots, visibility) and t_pos = f_pos * config.ts_mult (RoPE time); conflating them is silent drift at any fps where ts_mult != 1.

Classes

WaypointDiT(config)

Conv2d patchify -> 24 blocks -> adaLN head -> unpatchify, plus the 4+1 per-frame driver.

WaypointDiTBlock(config, layer_idx)

One DiT block: adaLN-modulated causal frame attention, optional controller fusion, adaLN-modulated MLP.

WaypointPosIds(f_pos, t_pos, y_pos, x_pos)

The four position streams for one frame, per row of the step batch.
class mstar.model.waypoint.components.dit.WaypointDiT(config)[source]#

Bases: Module

Conv2d patchify -> 24 blocks -> adaLN head -> unpatchify, plus the 4+1 per-frame driver.

Built on the meta device and materialized by weight_loader in a fixed order:

with torch.device("meta"):
    dit = WaypointDiT(config)
dit.cast_serving_dtypes()   # bf16 + fp32 islands, on meta
dit.to_empty(device=device) # un-ties cond_proj
dit.retie_cond_proj()       # MUST follow to_empty
load_weights_into(dit, ...)

The KV ring is derived state owned by an engine resource, bound after materialization – not part of this module.

Parameters:

config (WaypointConfig)

append_frame(latent, frame_pos, *, mouse, button, scroll)[source]#

Prime the world state from a real (VAE-encoded) frame: the committing pass alone, one forward.

latent is already the settled x0, so there is nothing to solve and nothing to clone. Returned unchanged for the caller to decode.

Parameters:
Return type:

Tensor

cast_serving_dtypes()[source]#

bf16 everywhere, then the fp32 islands back to fp32.

Called on the meta module before to_empty(device) so storage is allocated directly in the serving dtype. .to(dtype) on meta preserves the cond_proj aliasing (to_empty does not), so no retie is needed here.

Return type:

WaypointDiT

compile_regions(**compile_kwargs)[source]#

Compile the denoise and cache passes, one region each. Idempotent.

materialize_runtime_tables must run first: the derived token grid, RoPE tables and conditioner LUT live outside the module tree, and a first touch inside a fullgraph=True region is at best a specialization and at worst a graph break.

Return type:

WaypointDiT

property dtype: dtype#

Bulk compute dtype (the non-island weights); callers cast latents to this.

forward(x, sigma, frame_pos, *, mouse, button, scroll, commit, cond_idx)[source]#

One pass over one latent frame; returns the rectified-flow velocity.

x [B, N, C, H, W] (N == 1), sigma [B, N], frame_pos [B] int64 ring clock. Returns [B, N, C, H, W]. commit says whether this pass keeps its K/V (False for the four denoise passes, True for the fifth); cond_idx selects the cached modulation row and must match sigma.

Parameters:
Return type:

Tensor

generate_frame(noise, frame_pos, *, mouse, button, scroll)[source]#

Denoise one frame from noise [B, N, C, H, W] and commit it.

Five forwards: 4 frozen + 1 committing, all sharing one frame_pos. The caller owns the ring clock and advances it once per committed frame.

Parameters:
Return type:

Tensor

materialize_runtime_tables(device)[source]#

Create every derived table after weight loading and before compile.

Parameters:

device (device | str)

Return type:

WaypointDiT

retie_cond_proj()[source]#

Alias blocks 1..23’s six cond_proj matrices onto block 0’s.

Public, and separate from __init__, because to_empty(device) destroys the tying (Module._apply has no cross-module memo, so each block would get 24 independent copies). bias_in is deliberately NOT tied – it is genuinely per-layer.

Return type:

WaypointDiT

class mstar.model.waypoint.components.dit.WaypointDiTBlock(config, layer_idx)[source]#

Bases: Module

One DiT block: adaLN-modulated causal frame attention, optional controller fusion, adaLN-modulated MLP.

The six modulation tensors come from this block’s cond_head, whose bias_in is per-layer and whose cond_proj matrices alias block 0’s.

Parameters:
forward(x, frame_pos, rope_angles, cond, ctrl_emb, v1, *, commit, cond_idx)[source]#

x [B, N*T, D], cond/ctrl_emb [B, N, D] (per frame) -> (x, v1). v1 is layer 0’s pre-lerp V, threaded down the stack; cond_idx is the scheduler_sigmas slot for the cond_head cache. Only f_pos reaches this far – RoPE angles are built once at the root.

Parameters:
Return type:

tuple[Tensor, Tensor]

class mstar.model.waypoint.components.dit.WaypointPosIds(f_pos, t_pos, y_pos, x_pos)[source]#

Bases: NamedTuple

The four position streams for one frame, per row of the step batch.

f_pos is the [B] int64 ring clock (the KV resource’s upsert takes it); t_pos is the [B, T] RoPE time coordinate f_pos * ts_mult; y_pos/x_pos are [B, T] token-grid coordinates, row = i // width and col = i % width.

Parameters:
f_pos: Tensor#

Alias for field number 0

t_pos: Tensor#

Alias for field number 1

x_pos: Tensor#

Alias for field number 3

y_pos: Tensor#

Alias for field number 2