Modeling Specifics

Warning

The components documented on this page are internal model implementations and are subject to change between versions. These are advanced, under-the-hood components of World Engine. You should only interact with these directly if you are savvy and curious about the internal workings of the model architecture.

For normal usage, please refer to the main API Reference.

Internal Model Components

world_engine.model.base_model

class world_engine.model.base_model.BaseModel(*args, **kwargs)

Bases: Module

save_pretrained(path)

Save weights (.safetensors) and OmegaConf YAML.

Return type:

None

classmethod from_pretrained(path, cfg=None, device=None)

Load weights and OmegaConf YAML.

push_to_hub(uri, **kwargs)

static load_config(path)

world_engine.model.world_model

class world_engine.model.world_model.PromptEncoder(*args: Any, **kwargs: Any)

Bases: Module

os = <module 'os' (frozen)>

__init__(model_id='google/umt5-xl', dtype=torch.bfloat16)

encode(inputs)

forward(texts)

world_engine.model.world_model.rms_norm(x)

Return type:

Tensor

class world_engine.model.world_model.CFG(d_model, dropout)

Bases: Module

__init__(d_model, dropout)

forward(x, is_conditioned=None)

Return type:

Tensor

x: [B, L, D] is_conditioned:

• None: training-style random dropout

• bool: whole batch conditioned / unconditioned at sampling

class world_engine.model.world_model.MLP(*args: Any, **kwargs: Any)

Bases: Module

__init__(dim_in, dim_middle, dim_out)

forward(x)

class world_engine.model.world_model.ControllerInputEmbedding(*args: Any, **kwargs: Any)

Bases: Module

__init__(config)

forward(mouse, button, scroll)

class world_engine.model.world_model.MLPFusion(*args: Any, **kwargs: Any)

Bases: Module

Fuses per-group conditioning into tokens by applying an MLP to cat([x, cond])

__init__(config)

forward(x, cond)

Return type:

Tensor

class world_engine.model.world_model.CondHead(*args: Any, **kwargs: Any)

Bases: Module

Per-layer conditioning head: bias_in → SiLU → Linear → chunk(n_cond).

n_cond = 6

__init__(config)

forward(cond)

class world_engine.model.world_model.WorldDiTBlock(*args: Any, **kwargs: Any)

Bases: Module

__init__(config, layer_idx)

forward(x, pos_ids, cond, ctx, v, kv_cache=None)

0. Causal Frame Attention

1. Frame->CTX Cross Attention

2. MLP

class world_engine.model.world_model.WorldDiT(*args: Any, **kwargs: Any)

Bases: Module

__init__(config)

forward(x, pos_ids, cond, ctx, kv_cache=None)

class world_engine.model.world_model.WorldModel(*args: Any, **kwargs: Any)

Bases: BaseModel

WORLD: Wayfarer Operator-driven Rectified-flow Long-context Diffuser

Denoise a frame given - All previous frames - The prompt embedding - The controller input embedding - The current noise level

__init__(config)

forward(x, sigma, frame_timestamp, prompt_emb=None, prompt_pad_mask=None, mouse=None, button=None, scroll=None, kv_cache=None, ctrl_cond=None, prompt_cond=None)

x: [B, N, C, H, W], sigma: [B, N] frame_timestamp: [B, N] prompt_emb: [B, P, D] controller_inputs: [B, N, I] ctrl_cond: Inference only, whether to apply controller conditioning

world_engine.model.nn

class world_engine.model.nn.NoCastModule(*args, **kwargs)

Bases: Module

to(*args, **kwargs)

world_engine.model.nn.rms_norm(x)

Return type:

Tensor

class world_engine.model.nn.MLP(*args: Any, **kwargs: Any)

Bases: Module

__init__(dim_in, dim_middle, dim_out)

forward(x)

class world_engine.model.nn.AdaLN(*args: Any, **kwargs: Any)

Bases: Module

__init__(dim)

forward(x, cond)

world_engine.model.nn.ada_rmsnorm(x, scale, bias)

world_engine.model.nn.ada_gate(x, gate)

class world_engine.model.nn.NoiseConditioner(*args: Any, **kwargs: Any)

Bases: NoCastModule

Sigma -> logSNR -> Fourier Features -> Dense

__init__(dim, fourier_dim=512, base=10000.0)

forward(s, eps=torch.finfo.eps)

world_engine.model.attn

class world_engine.model.attn.RoPE(*args: Any, **kwargs: Any)

Bases: NoCastModule

__init__(config)

get_angles(pos_ids)

forward(x, pos_ids)

get_freqs(config)

class world_engine.model.attn.OrthoRoPE(*args: Any, **kwargs: Any)

Bases: RoPE

RoPE for rotation across orthogonal axes: time, height, and width Time: Geometric Spectrum – rotates 1/2 of head dim Height / Width: Linear Spectrum – rotates 1/4th of head dim each (1/2 combined)

get_freqs(config)

class world_engine.model.attn.Attn(*args: Any, **kwargs: Any)

Bases: Module

__init__(config, layer_idx)

forward(x, pos_ids, v1, kv_cache)

class world_engine.model.attn.CrossAttention(*args: Any, **kwargs: Any)

Bases: Module

__init__(config, context_dim=None)

forward(x, context, context_pad_mask=None)

world_engine.model.kv_cache

world_engine.model.kv_cache.make_block_mask(T, L, written)

T: Q length for this frame L: KV capacity == written.numel() written: [L] bool, True where there is valid KV data

Return type:

BlockMask

class world_engine.model.kv_cache.LayerKVCache(B, H, L, Dh, dtype, tokens_per_frame, pinned_dilation=1)

Bases: Module

Ring-buffer KV cache with fixed capacity L (tokens) for history plus one extra frame (tokens_per_frame) at the tail holding the current frame.

__init__(B, H, L, Dh, dtype, tokens_per_frame, pinned_dilation=1)

reset()

upsert(kv, pos_ids, is_frozen)

kv: [2, B, H, T, Dh] for a single frame (T = tokens_per_frame) t_pos: [B, T], all equal per frame (ignoring -1)

class world_engine.model.kv_cache.StaticKVCache(*args: Any, **kwargs: Any)

Bases: Module

__init__(config, batch_size, dtype)

reset()

set_frozen(is_frozen)

upsert(k, v, pos_ids, layer)

Patch Model

world_engine.patch_model

class world_engine.patch_model.CachedDenoiseStepEmb(base, sigmas)

Bases: Module

bf16 sigma -> bf16 embedding via 64k LUT; invalid sigma => OOB index error (no silent wrong).

__init__(base, sigmas)

forward(sigma)

Return type:

Tensor

class world_engine.patch_model.CachedCondHead(base, cached_denoise_step_emb, max_key_dims=8)

Bases: Module

bf16 cond -> cached (s0,b0,g0,s1,b1,g1); invalid cond => OOB index error (no silent wrong).

__init__(base, cached_denoise_step_emb, max_key_dims=8)

forward(cond)

world_engine.patch_model.patch_cached_noise_conditioning(model)

Return type:

None

class world_engine.patch_model.MergedQKVAttn(src, config)

Bases: Attn

__init__(src, config)

forward(x, pos_ids, v1, kv_cache)

world_engine.patch_model.patch_Attn_merge_qkv(model)

Return type:

None

class world_engine.patch_model.SplitMLPFusion(src)

Bases: Module

Packed MLPFusion -> split linears (no cat, quant-friendly).

__init__(src)

forward(x, cond)

Return type:

Tensor

world_engine.patch_model.patch_MLPFusion_split(model)

Return type:

None

world_engine.patch_model.apply_inference_patches(model)

Return type:

None

Autoencoder Neural Network Components

world_engine.ae_nn

world_engine.ae_nn.WeightNormConv2d(*args, **kwargs)

class world_engine.ae_nn.ResBlock(ch)

Bases: Module

__init__(ch)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.LandscapeToSquare(ch_in, ch_out)

Bases: Module

__init__(ch_in, ch_out)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.Downsample(ch_in, ch_out)

Bases: Module

__init__(ch_in, ch_out)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.DownBlock(ch_in, ch_out, num_res=1)

Bases: Module

__init__(ch_in, ch_out, num_res=1)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.SpaceToChannel(ch_in, ch_out)

Bases: Module

__init__(ch_in, ch_out)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.ChannelAverage(ch_in, ch_out)

Bases: Module

__init__(ch_in, ch_out)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.SquareToLandscape(ch_in, ch_out)

Bases: Module

__init__(ch_in, ch_out)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.Upsample(ch_in, ch_out)

Bases: Module

__init__(ch_in, ch_out)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.UpBlock(ch_in, ch_out, num_res=1)

Bases: Module

__init__(ch_in, ch_out, num_res=1)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.ChannelToSpace(ch_in, ch_out)

Bases: Module

__init__(ch_in, ch_out)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.ChannelDuplication(ch_in, ch_out)

Bases: Module

__init__(ch_in, ch_out)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.Encoder(config)

Bases: Module

__init__(config)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.Decoder(config)

Bases: Module

__init__(config)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class world_engine.ae_nn.AutoEncoder(encoder_config, decoder_config=None)

Bases: Module

__init__(encoder_config, decoder_config=None)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.