echo / code /doc /checkpoints.md
amonshano's picture
Add Echo-Memory codebase used for this run (CC BY 4.0, JD Echo Team) (part 2)
eafbe80 verified
|
Raw
History Blame Contribute Delete
3.43 kB

Checkpoints (Hugging Face)

Repo: Echo-Team/Echo-Memory

Fine-tuned DiT weights on top of Wan-AI/Wan2.1-T2V-1.3B. Released rows are saved as {row_id}/epoch-0.safetensors after 1 epoch / 30,000 steps on the static in-domain pool (640×352, 81-frame chunks). Mechanism names follow memory_mechanisms.md.

Checkpoint index

Family Paper row HF path Steps Echo-Memory recipe
Raw context Context K=1 context_k1/epoch-0.safetensors 30,000 train/context_learning/run_pre_qkv_ctx1.sh
Raw context Context K=20 TODO TODO train/context_learning/run_pre_qkv_ctx20.sh
Spatial Spatial Memory TODO TODO train/memory_baselines_basic/run_spatial_memory_baseline.sh
State-space Block-wise SSM TODO TODO train/memory_baselines_basic/run_ablation_block_wise_ssm_two_chunk.sh
State-space Legacy Hybrid (VideoSSM) TODO TODO train/memory_baselines_basic/run_videossm_hybrid_baseline.sh
Spatial concat text (ablation) TODO TODO train/memory_baselines_basic/run_ablation_spatial_concat_text_two_chunk.sh
Spatial inject none (ablation) TODO TODO train/memory_baselines_basic/run_ablation_spatial_inject_none_two_chunk.sh
Spatial cross-attn t32 (ablation) TODO TODO train/memory_baselines_basic/run_ablation_spatial_cross_attn_readout_two_chunk.sh
State-space SSM ctx1 / every4 / hint21 TODO TODO SSM ablation
State-space SSM ctx5 / every1 / hint21 TODO TODO SSM ablation
State-space SSM ctx5 / every4 / hint81 TODO TODO SSM ablation

Context K=5, Context K=20, Spatial memory, FramePack compression, and State-space / SSM rows are TODO and not yet released as epoch-0 weights.

Download

pip install -U "huggingface_hub[cli]"

# one row (keeps HF folder layout under ./ckpts/)
huggingface-cli download Echo-Team/Echo-Memory context_k1/epoch-0.safetensors --local-dir ./ckpts

# all currently released rows
huggingface-cli download Echo-Team/Echo-Memory --local-dir ./ckpts

Keep the subdirectory name in the local path (e.g. ./ckpts/context_k1/epoch-0.safetensors). Eval scripts use env/memory_baseline_runtime.py to infer memory flags from path substrings; Spatial and SSM checkpoint rows remain TODO.

Use with Echo-Memory

Set the Wan backbone, static in-domain data pool, and checkpoint path:

export WAN_BASE_MODEL=/path/to/Wan2.1-T2V-1.3B
export DATASET_BASE_PATH=data/Context-as-Memory-Dataset
export PYTHONPATH=$PWD:${PYTHONPATH:-}
export CKPT=./ckpts/context_k1/epoch-0.safetensors

In-domain replay + revisit (paper bundle):

bash eval/v2/run_static_consistency_loop_and_revisit.sh
bash eval/v2/run_basic_replay_gt.sh

Open-domain revisit (first frames already in assets/opendomain_revisit/):

PHASE=stage1 OOD_DIR=assets/opendomain_revisit \
  bash eval/v2/revisit_suite/run_one_click_revisit_eval.sh

Visual comparison (fixed prompt + first frame):

python eval/metrics/run_visual_eval.py \
  --ckpt "$CKPT" \
  --output_root ./evals_visual

See eval/v2/README.md and eval/metrics/README.md for full options.