VL-JEPA → Cosmos context predictor (jepa-guided-diffusion)

VL-JEPA is a lightweight predictor bridge that maps a WTS traffic video (plus a text query) into the frozen Cosmos-Reason1 (512, 1024) context space, so the predicted embedding can condition Cosmos-Predict2.5 video generation without running the 7B Reason1 text encoder at inference time.

• Input: video / frame directory + a WTS-format query.
• Output: a (512, 1024) context tensor (the Cosmos cross-attention conditioning set).
• Objective: permutation-invariant set-match Hungarian MSE — Cosmos consumes the context as an unordered set, so the predicted 512 tokens are matched to the target 512 tokens (Hungarian assignment) before MSE.

Architecture

Part	Role	Source
X-Encoder (frozen)	VJEPA 2.1 ViT-L/384 (vjepa2_1_vit_large_384), video mode → visual tokens	torch.hub (facebookresearch/vjepa2)
Compressor (trained)	BLIP-2-style QFormer: 512 query tokens, 4 layers, cross-attn every layer	in model.pt
Backbone (trained)	last 4 layers of Llama-3.2-1B, made bidirectional	arch from hub, weights in model.pt
Head (trained)	output_projection → (512, 1024)	in model.pt
Y-Encoder (frozen, target only)	Cosmos-Reason1-7B full_concat + crossattn_proj	not needed at inference

Only the compressor, the last 4 Llama layers, the vision projection, and the output projection are trained; VJEPA2 and Reason1 are frozen.

Repo layout

config.yaml                         # full model/eval config the loader reads
best_model/model.pt                 # trained weights (predictor only)
best_model/model_config.yaml        # model config sidecar
best_model/query_tokenizer/         # Llama query tokenizer

Usage

Install the VL-JEPA package (uv-managed), then pass this repo id as --checkpoint — it is downloaded and cached automatically:

# whole test set: writes <scenario>/embedding.pt for the Cosmos handoff
uv run vl-jepa seq-infer --checkpoint AlterraLaniakea/jepa-guided-diffusion --test-root wts/test

# single scenario -> print the (512,1024) embedding
uv run vl-jepa infer \
  --checkpoint AlterraLaniakea/jepa-guided-diffusion \
  --visual-path wts/test/<scenario>/input \
  --output embedding \
  --query "The current preset is from overhead view with 2 target subjects. Pedestrian: A man in his 30s stands on the road facing the oncoming vehicle. Vehicle: The vehicle goes straight at a constant speed."

Query format. For in-distribution embeddings, use the same query format as training (seq-infer builds this automatically from caption.json):

The current preset is from <overhead|vehicle> view with <N> target subject(s). Pedestrian: <ped caption> Vehicle: <veh caption>

• <overhead|vehicle> — overhead for CCTV, vehicle for dashcam (video*) clips.
• <N> — subject types present: 1 (pedestrian or vehicle) or 2 (both).

Requirements

• Gated Llama-3.2-1B: run hf auth login (or export HF_TOKEN) before first use — the loader rebuilds the Llama backbone architecture from the hub and overlays the trained weights. (Set hf_token: true in config.yaml, or provide the token via env.)
• VJEPA2 is fetched from facebookresearch/vjepa2 (torch.hub / GitHub source).
• A GPU is recommended; VJEPA2 runs at crop size 384.

Downstream

Feed the predicted embedding.pt into Cosmos-Predict2.5 via scripts/generate_cosmos.py (see the repo README) to generate the WTS video for each scenario.