CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

What this repo is

Echo-Memory is the release code for a paper that runs a controlled study of memory mechanisms in action-conditioned video world models. Every experiment shares the same Wan 2.1 1.3B DiT backbone and the same two-chunk training/eval setup; only the memory pathway changes between rows. The goal is to compare how four memory families (Context, Compression, Spatial, State-Space) preserve scene identity/layout/viewpoint when a generated trajectory leaves a view and later revisits it.

This is a public release of a larger internal codebase. Private benchmark launchers, cluster submit files, logs, generated outputs, and machine-local absolute paths have been intentionally removed. Do not reintroduce them (see Public-repo constraints below).

Environment & required variables

conda env create -f environment.yml && conda activate echo-memory
pip install -r requirements.txt

Almost every script is path-portable via environment variables:

export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B   # must contain diffusion_pytorch_model.safetensors, models_t5_umt5-xxl-enc-bf16.pth, Wan2.1_VAE.pth
export DATASET_BASE_PATH=data/Context-as-Memory-Dataset   # training/eval pool root (code default if unset)
export PYTHONPATH=$PWD:${PYTHONPATH:-}
export CKPT=./ckpts/<row_id>/epoch-0.safetensors   # for eval/inference
export OUTPUT_BASE_ROOT=$PWD/outputs               # training output root (optional)

Critical convention: the CKPT path must keep the row folder name (e.g. spatial_mem/, context_k20/). env/memory_baseline_runtime.py and inference/unified_inference.py infer the memory profile (which flags to set) by matching substrings in that path. Renaming the folder breaks memory selection.

Common commands

Run a single test (tests are standalone scripts, not pytest — there is no pytest config):

PYTHONPATH=. python3 tests/test_context_chunk_utils.py
PYTHONPATH=. python3 tests/test_framepack_memory_align.py
PYTHONPATH=. python3 tests/test_two_chunk_anchor_readout.py

These check memory/context plumbing (context-frame selection, FramePack alignment, two-chunk anchor read-out).

Train a memory row (run from repo root; uses accelerate launch src/model_training/train.py under the hood):

bash train/memory_baselines_basic/run_spatial_memory_baseline.sh   # spatial / SSM / compression rows
bash train/context_learning/run_pre_qkv_ctx20.sh                   # raw-context K=1/5/20 rows
bash train/dynamic_spatialvid/run_dyn_spatial_mem.sh               # SpatialVID dynamic pool rows

Inference (single-chunk generation, any memory family):

python inference/unified_inference.py --ckpt $CKPT --memory_type auto \
  --context_image assets/opendomain_revisit/1774363417.png \
  --action_path env/action_rotation_left_45.json \
  --prompt "A toy bear on a table" --output_path output.mp4

--memory_type auto detects the family from the ckpt path; see python inference/unified_inference.py --help for the full list.

Evaluate (in-domain replay/revisit; requires the static pool):

bash eval/v2/run_basic_replay_gt.sh                            # fast single-video GT fidelity check (~5 min)
bash eval/v2/run_static_consistency_loop_and_revisit.sh        # full paper eval bundle (loop closure + revisit)
PHASE=stage1 OOD_DIR=assets/opendomain_revisit bash eval/v2/revisit_suite/run_one_click_revisit_eval.sh

Outputs go to ${CKPT_DIR}/evals_v2/.

Architecture: the big picture

Layered structure (each layer reads from the one below):

train/, inference/, eval/ — bash launchers that set flags and call into the Python entrypoints. The paper's experiment matrix lives here as a set of run_*.sh scripts, one per memory row. They are the canonical, reproducible interface — prefer editing/extending an existing launcher over writing new Python entrypoints.
src/model_training/train.py — the single training entrypoint. All memory mechanisms are toggled by command-line flags (--use_spatial_memory, --use_framepack_memory, --use_block_wise_ssm, --context_memory_frames, etc.), so the same script trains every row. It writes the active memory config into the pipeline.
inference/unified_inference.py + env/loop_utils.py — load a checkpoint, reconstruct the matching memory module (inferring structure from checkpoint key names / weight shapes), and run generation.
diffsynth/ — a vendored DiffSynth model stack (Wan + many other diffusion model families). The Echo-Memory-specific code is concentrated in:
- diffsynth/models/memory/ — the four memory families: spatial_grid_memory.py, framepack_length.py, framepack_weight.py, block_wise_ssm.py (paper-aligned recurrent SSM), videossm_hybrid.py (legacy temporal-conv baseline).
- diffsynth/pipelines/wan_video_new.py — the action-conditioned Wan pipeline; the raw-context and FramePack read-out paths live here.
- diffsynth/models/wan_video_*.py — DiT, VAE, text/image encoders, camera/motion controllers.
src/model_training/context_chunk_utils.py, context_retrieval.py, etc. — context-frame selection and multi-chunk sampling helpers (the logic the tests cover).
env/ — runtime glue: memory_baseline_runtime.py (CKPT path → memory flags), loop_utils.py (pipeline + checkpoint loading), action JSONs, replay-loop driver.

The two-chunk paradigm (central to the whole repo)

All memory baselines are trained and evaluated in a two-chunk setup that simulates the revisit scenario:

Chunk 1 (context): a clean reference segment, VAE-encoded, with matched camera RT actions. Concatenated at the suffix position (CONTEXT_POSITION=suffix).
Chunk 2 (target): the noisy segment the model learns to denoise. It can only access chunk-1 information through the memory mechanism — this is what forces the memory pathway to do work, and what evaluation later probes.

The ablations are designed to change only the memory read/write pathway while keeping backbone, action conditioning, resolution (640×352), chunk length (81 frames), and training schedule fixed. See doc/memory_mechanisms.md for the authoritative paper-row → code-path → training-script mapping.

Memory rows at a glance

Family	Rows	Code path
Raw context	`context_k1/k5/k20`	`wan_video_new.py` context latent path
Compression	`framepack_weight`, `framepack_len_r2/r4`, `framepack_hybrid_*`	`memory/framepack_{weight,length}.py`
Spatial	`spatial_mem`, `spatial_inject_none/concat_text/cross_attn_readout`	`memory/spatial_grid_memory.py` (same storage, different read-out)
State-space	`block_wise_ssm` (paper), `videossm_hybrid` (legacy)	`memory/block_wise_ssm.py`, `memory/videossm_hybrid.py`

Two dataset pools (same on-disk layout)

DATASET_BASE_PATH points at one of two interchangeable pools, each with frames/ jsons/ overlap_labels/ metadata_full.csv [latents/]:

Static in-domain pool — data/Context-as-Memory-Dataset (default). Used for all in-domain training and eval.
Dynamic training pool — data/dynamic-spatialvid-motion60/mixed (SpatialVID subset). Training + inference only; dynamic eval is TODO.

Metadata is regenerated with scripts/run_generate_metadata.sh; optional latent precompute with scripts/run_precompute_ctx_target_latents.sh.

Public-repo constraints (the codebase is split public/private)

No machine-local absolute paths (e.g. /pfs/…, cluster paths) in any committed file — use env vars with sensible defaults.
No checkpoint-upload scripts, internal benchmark names, or cluster submit files. scripts/upload_hf_checkpoints.sh and related are gitignored; do not add them.
In public-facing markdown use the released pool names ("static in-domain pool", "dynamic training pool"), not internal codenames.
data/, outputs/, logs/, wandb/, checkpoints/, and binary weights are gitignored. Versioned binaries are limited to README/site visual assets under assets/ and docs/.
Keep diffs minimal and match the existing bash/Python patterns in train/ and eval/v2/.

Cursor skills

Maintainer workflows are documented as Cursor skills in .cursor/skills/ (echo-memory-train, echo-memory-eval, echo-memory-release, echo-memory-agent) and in doc/DEVELOPER.md. Consult doc/memory_mechanisms.md, doc/checkpoints.md, and doc/dataset_preprocessing.md for row/checkpoint/dataset details.