GTA world model: weights and inference code

Play it in your browser

An interactive world model. You give it one start image, then hold keys; it generates the video that follows, one frame at a time, and streams it to a browser while you play.

The architecture is Summaryformer. It is my own design: I arrived at it after testing a lot of ideas on consumer hardware, and I trained it from scratch. This checkpoint was trained on GG800, my new 800-hour GTA5 dataset.

This repository holds the weights and the complete inference runtime: the model, its memory, the sampler and the streaming server. There is no training code here. What I took from other people's work is listed under Credits.

The memory

The part of Summaryformer I most want people to read about is how a frame remembers the past.

A frame generated a second ago matters in detail: the next frame has to continue its motion pixel for pixel. A frame generated many seconds ago matters only in outline: what was there, roughly where. Keeping every past frame at full resolution spends almost all of the attention budget on detail nobody needs. Keeping only a short window forgets everything else.

So every frame in this model is stored twice, at two sizes, and the model itself writes both.

A frame is a record

Each frame enters the transformer as one record:

record(t) = [ raw(t) : H*W tokens ] ++ [ sum(t) : X tokens ]

• raw(t) is the frame's latent grid, one token per cell. At the served 384x672 that is 24x42 = 1008 tokens.
• sum(t) is X = 96 summary tokens. They start from a learned init (the same for every frame), have no position on the grid, and go through the same 28 blocks in the same forward pass as the raw tokens. By the last block they hold a compact description of the frame.

There is no separate summarizer network and no extra pass. The summary is a by-product of generating the frame.

What a frame attends to

While frame t is being generated, every one of its tokens attends to:

1. itself: the raw tokens and the summary tokens of frame t;
2. the last Y = 2 frames, as whole records (full resolution);
3. the older frames inside the window W = 24, as their 96 summary tokens only;
4. frame 0, the start image, as a whole record. It is never evicted.

Recent frames give local detail. Summaries give the long horizon. The start image anchors the scene.

One buffer, no mask

All of this is one preallocated key/value buffer per block:

| Z0 : record(0) | RAW : (Y+1) records, rotating | SUM : (W-Y) x X summary rows, rotating |
                    read = buffer[0 : end]

• Z0 holds the start image's record.
• RAW holds the frame being generated and the Y frames before it.
• SUM holds one X-row summary for each older frame in the window.

When a frame falls out of the raw window, the keys and values of its 96 summary tokens are copied from RAW into SUM, and its 1008 raw rows are overwritten by the next frame. That copy of 96 rows per block is the entire cost of moving a frame into long-term memory.

Attention is one dense read of the filled prefix of that buffer per block. There is no attention mask and no concatenation: what a frame may see is decided by what is in the buffer.

At the served resolution with the trained window, the buffer holds 6,528 keys. The same span with every frame kept at full resolution would be 26,208.

summary tokens per frame (X)	96
raw tokens per frame at 384x672	1008
frames kept whole (Y)	2
memory window in frames (W)	24 trained; the server default serves 16

One body frame of latents decodes to 4 RGB frames, so a window of 24 covers 96 RGB frames.

Telling a summary from a pixel: level planes

Raw tokens and summary tokens of the same frame share a temporal position, so rotary position embeddings alone cannot tell them apart. I did not add a token-type embedding. Instead the two lowest-frequency planes of the temporal RoPE band are used as a level axis.

Each attention head has 12 temporal planes. The two slowest rotate by less than half a radian over 1024 frames, and by a small fraction of that over the distances that occur inside the memory window, so for raw tokens they carry no usable position. Summary tokens get an exact quarter turn on those two planes. A query can then weight "summary of frame t" and "pixel of frame t" independently, using dimensions that were otherwise idle. The model checks at construction that those planes really are idle and refuses to build if they are not.

Running for a long time: the start image follows the window

Frame ids keep growing as the session goes on, but the start image was written at id 0. Left alone, the temporal distance between the start image and the live window would grow without bound.

RoPE keys are rotations, and rotation(id + T) = rotation(id) * rotation(T). So a stored key can be moved forward in time by multiplying it, in place, by one more rotation. Whenever the gap between the start image and the oldest frame in memory reaches 5, the start image's keys are rotated forward to bring the gap back to 1. Values are untouched. Nothing is recomputed.

What I believe is new

• Summary tokens produced by the generator's own forward pass, kept as the only trace of older frames, with full records kept for the newest few.
• One key/value buffer with a raw region and a summary region, fed by eviction, read with a single unmasked attention call.
• Level planes: marking token kind with a fixed rotation on RoPE planes that carry no position.
• Rotating the cached keys of the anchor frame so it stays next to a moving window.

The rest of the model

• Size. 630,205,008 parameters. 28 blocks, width 1280, 20 heads of 64, SwiGLU feed-forward.
• Latents. 48 channels on a /16 grid. The start image is one latent frame; every later latent frame is 4 RGB frames.
• Two entrances, two exits. The start image and the generated frames have separate token embeddings and separate output heads on the same trunk. One is a still image and the other is video.
• Conditioning. There is no text. The denoising timestep modulates every block. The player's action (5 numbers), the frame rate and the grid shape are added to that modulation, each through its own learned per-block gate.
• Sampling. Each frame takes three denoising passes (t = 999, 878, 669) and one more pass on the clean result, which rewrites that frame's entry in memory. The memory always holds clean frames.
• Positions. One rotary table with temporal, height and width bands, indexed in time by frame id. This part is adapted from Wan2.1 (see Credits); the level planes and the changes listed in rope.py are mine.
• Diffusion. v-prediction on a cosine noise schedule with a log-SNR shift. Inference needs only the signal and noise scale at the few timesteps a frame is denoised at; that is all of noise_schedule.py.
• Attention kernels. Two: PyTorch FlexAttention and FlashAttention's flash_attn_with_kvcache. The server uses FlexAttention.

Files

serve_ckpts/
  gta_model.pt                    the model weights
sae_compact.pt                    the autoencoder weights
latent_stats_*.pt                 per-channel mean / std of the autoencoder's latent
summary_server.py                 FastAPI server: start image in, H.264 out, keys in
summary_streaming_inference.py    one session: the per-frame loop
summaryformer.py                  the model
summaryformer_wrapper.py          the model's public surface
gta_streaming_kv_cache.py         the memory buffer
summaryformer_utils/
  transformer_block.py            the block
  rope.py                         rotary positions and level planes
  streaming_sink.py               moving the start image's keys
  latent_head.py                  output head
  visual_mlp.py                   SwiGLU feed-forward
  z0_focus_mask.py                optional focus region on the start image
  rollout_flex_attention.py       FlexAttention kernel
  fast_flash_kv.py                FlashAttention kernel
  recency_bias.py                 optional recency bias on summary keys
streaming_server_utils/
  ckpt_name_schema.py             checkpoint file names
  controls.py                     keys to the model's 5-number action
sae.py                            the autoencoder (parallel and autoregressive encode / decode)
noise_schedule.py                 the noise schedule and the timesteps a frame is denoised at
inference_common.py               small shared helpers

Running it

Requirements: a CUDA GPU; Python with torch (FlexAttention and torch.compile), numpy, fastapi, uvicorn, pillow, av, imageio. flash_attn is optional.

Everything the server needs is in this repository, at the paths its settings point to. The settings are the CONFIG dict at the top of summary_server.py:

• the model checkpoint, in checkpoints_dir (default serve_ckpts/): serve_ckpts/gta_model.pt;
• the autoencoder checkpoint (sae_ckpt, default sae_compact.pt): the encoder that turns the start image into a latent and the decoder that turns latents into frames;
• two latent-statistics files: latent_stats_z0.pt (start image) and latent_stats_z_rest.pt (generated frames). They hold the per-channel mean and standard deviation of the autoencoder's latent.

Download the repository and start the server from its root:

hf download lucidml/gta --local-dir gta
cd gta
python summary_server.py        # serves on 0.0.0.0:3002

At startup the server loads the weights, compiles the model and plays a short warm-up session, so the first start takes a while. One GPU process serves one player at a time.

The browser client is not in this repository. The protocol is small:

POST /z0/{user_id}	upload the start image (form field z0_image)
WS /ws/{user_id}	the session
GET /status, GET /models	occupancy and the available checkpoints

On the WebSocket the server sends binary messages, one byte (0x01 key frame, 0x00 delta frame) followed by an H.264 Annex-B access unit at 384x672. The client sends JSON with the keys currently held:

{"type": "keys", "pressed": {"w": true, "a": false, "s": false, "d": false, "shift": false, "space": false, "left": false, "right": true}}

Limits

• Sessions are finite. The position table is extended to 4096 frame ids at startup and a session stops when it reaches the last one. The server also has a wall-clock limit per session (ws_session_limit_s, default 180 s; 0 disables it).
• CUDA only.

Credits

Summaryformer is my own architecture and this model was trained from scratch. These are the pieces that came from other people:

• KV-cached inference, learned from Self-Forcing. The Self-Forcing repo (Apache-2.0; arXiv 2506.08009) is where I learned how a KV cache works for autoregressive video and how to run inference with one. Self-Forcing is itself built on Wan2.1. The memory here is a different design, but that is where I learned the technique, and the conditional_dict name at the model's boundary is from there.
• Rotary positions, from Wan2.1. I found the Wan2.1 code hard to follow at first. The part of it that mattered most for this model was its rotary positions, so that is the part I adapted: the 3-D rotary position code, and the way it handles time by giving every frame an integer id that indexes the temporal band of the table. Wan2.1 is Apache-2.0, Copyright 2024-2025 The Alibaba Wan Team Authors. summaryformer_utils/rope.py lists what I changed.
• Noise schedule and parameterization, from the literature. The cosine schedule is from Nichol & Dhariwal, Improved Denoising Diffusion Probabilistic Models (arXiv 2102.09672). The log-SNR shift is from Hoogeboom et al., simple diffusion (arXiv 2301.11093). v-prediction is from Salimans & Ho, Progressive Distillation for Fast Sampling of Diffusion Models (arXiv 2202.00512). noise_schedule.py is my own implementation of those formulas.
• Streaming convolutions in the autoencoder, from Wan2.1. The autoencoder is my own network and was trained from scratch. Its causal 3-D convolution, which carries a two-frame tail of the past from one call to the next, and the clip layout of one first frame followed by groups of 4, follow the Wan2.1 VAE (Apache-2.0).
• FlashAttention and PyTorch FlexAttention for the attention kernels.
• The optional recency bias follows Recency Forcing (arXiv 2609.19729).

License

Apache-2.0. See LICENSE.