Hugging Face's logo

amonshano
/
echo

Text-to-Video
world-model
video-generation
action-conditioned
memory
ablation
research-artifact
Model card Files
xet
 Community
echo / code /train /memory_baselines_basic /README.md
amonshano's picture
amonshano
Add Echo-Memory codebase used for this run (CC BY 4.0, JD Echo Team) (part 4)
c335050 verified
preview code
|
Raw
History Blame Contribute Delete
6.32 kB

Memory Baselines Basic: FramePack, Spatial, and State-Space

This folder contains the public training recipes for the paper's controlled memory-design matrix. See ../../doc/memory_mechanisms.md for the concise paper-row to implementation map. The scripts vary only the memory/context profile while keeping the backbone, optimizer, action conditioning, and evaluation interface aligned.

Paper Rows and Code Mapping

Paper row Mechanism Main implementation
FramePack-Weight Per-frame temporal decay and global scaling over context tokens; context length is unchanged. diffsynth/models/memory/framepack_weight.py
FramePack-Length Temporal mean pooling over context latents with matched RT-action padding and pooling. diffsynth/models/memory/framepack_length.py
Hybrid FramePack Length compression plus token weighting. wan_video_new.py memory path plus FramePack helpers
Token-grid baseline (spatial_mem) Time-mean context summary to learned spatial grid tokens. This row has no depth/3D reconstruction. diffsynth/models/memory/spatial_grid_memory.py
Geometry-grounded Spatial Memory TSDF-fused static point-cloud renders are VAE-encoded and summarized into geometry conditioning tokens. diffsynth/models/memory/geometry_spatial_memory.py
Block-wise SSM Paper-aligned recurrent state inside selected DiT blocks. diffsynth/models/memory/block_wise_ssm.py + --use_block_wise_ssm
VideoSSM hybrid Legacy lightweight temporal-convolution state-space baseline; kept separate from Block-wise SSM. diffsynth/models/memory/videossm_hybrid.py + --use_videossm_hybrid

Two-Chunk Ablation Scripts

The two-chunk scripts source common_env.sh and common_sampling_two_chunk.sh. The monitor uses left_45 followed by right_45, writes sampling_videos/step_*_two_chunk_memory*.mp4, and stores metadata beside the videos.

Script Purpose
run_ablation_no_memory_baseline_two_chunk.sh Anchor/no-extra-memory reference.
run_ablation_framepack_weight_two_chunk.sh FramePack token weighting.
run_ablation_framepack_len_r2_two_chunk.sh Length compression with ratio 2.
run_ablation_framepack_len_r4_two_chunk.sh Length compression with ratio 4.
run_ablation_framepack_hybrid_r2_weight_two_chunk.sh Ratio-2 length compression plus token weighting.
run_ablation_framepack_hybrid_r4_weight_two_chunk.sh Ratio-4 length compression plus token weighting.
run_ablation_spatial_inject_none_two_chunk.sh Spatial tokens are stored but not injected.
run_ablation_spatial_concat_text_two_chunk.sh Spatial tokens are appended to text cross-attention keys/values.
run_ablation_spatial_cross_attn_readout_two_chunk.sh Spatial tokens are read through a dedicated cross-attention read-out.
run_ablation_videossm_hybrid_two_chunk.sh Legacy VideoSSM hybrid with the two-chunk monitor.
run_ablation_block_wise_ssm_two_chunk.sh Paper-aligned block-wise SSM.
run_all_ablations_two_chunk.sh Sequential launcher for the full ablation set.

Representative Baselines

The non-ablation baseline scripts are still useful for representative rows and quick training checks:

  • run_framepack_baseline.sh: FramePack weight-only baseline.
  • run_framepack_lencompress_r2.sh: FramePack length compression ratio 2.
  • run_framepack_lencompress_r4.sh: FramePack length compression ratio 4.
  • run_spatial_memory_baseline.sh: representative Spatial Memory baseline.
  • run_geometry_spatial_memory_baseline.sh: geometry-grounded adaptation of Video World Models with Long-term Spatial Memory. It requires a metadata geometry_memory column containing paths to rendered static point-cloud videos such as the reference implementation's Vid_masktarget.mp4.
  • run_videossm_hybrid_baseline.sh: legacy VideoSSM hybrid baseline.

These scripts expose common overrides through environment variables: CKPT_INTERVAL, TIMESTEP_SHIFT, SAMPLING_INTERVAL_STEPS, SAMPLING_NUM_INFERENCE_STEPS, SAMPLING_HEIGHT, SAMPLING_WIDTH, and SAMPLING_NUM_FRAMES.

Context Construction

  • --context_source fov: FOV-overlap retrieval over historical frames.
  • --context_source replay: replay-style context construction aligned with env/run_replay_loop_two_chunk.py.
  • --context_source prev_chunk_tail: continuous frames from [start_frame - N, start_frame) on disk.

The shared implementation lives in src/model_training/context_chunk_utils.py.

Evaluation Alignment

Evaluation scripts source env/eval_infer_alignment_env.sh and call env/memory_baseline_runtime.py to infer the correct runtime memory flags from checkpoint paths. Keep output suffixes stable if you add a new script, or update env/memory_baseline_runtime.py so evaluation can recover the matching memory profile.

Generated training monitor videos are useful for fast visual checks, but paper-quality revisit panels should be produced through eval/v2/revisit_suite, which stores first frames, revisit-tail frames, change maps, and generated MP4 files for each case.

Preparing Geometry Memory

The geometry extractor is not SpatialGridMemory. Follow the official spmem/spmem preprocessing path:

  1. Recover RGB, metric depth, camera intrinsics, and camera-to-world poses (Mega-SaM for offline training data; CUT3R is used by the reference work for recurrent online reconstruction).
  2. Run TSDF fusion and render the static point cloud along the target camera trajectory. The reference script writes Vid_masktarget.mp4.
SPMEM_ROOT=/path/to/spmem \
INPUT_NPZ=/path/to/reconstructed_clip.npz \
CLIP_NAME=sample_id \
GEOMETRY_OUTPUT_ROOT=/path/to/tsdf/outputs \
bash scripts/run_spmem_tsdf_preprocess.sh
  1. Add those video paths to metadata:
python scripts/add_geometry_memory_column.py \
  --metadata /path/to/metadata_full.csv \
  --geometry_root /path/to/tsdf/outputs \
  --output /path/to/metadata_geometry.csv
  1. Train with:
METADATA_NAME=metadata_geometry.csv \
GEOMETRY_MEMORY_ROOT=/path/to/tsdf/outputs \
bash train/memory_baselines_basic/run_geometry_spatial_memory_baseline.sh