ggmworld

# GIM: Geometry-Aware Implicit Memory for Video World Models Training and inference code for **GIM** (Geometry-Aware Implicit Memory), a memory mechanism for autoregressive video world models, together with three baselines under a matched history-token budget: | Method | History handling | Entry point | |---|---|---| | **GIM** (ours) | Information-guided pruning + implicit memory encoder + camera-queryable geometry supervision | `train/train_gim.py` | | **SSM** | Recurrent state-space memory baseline (uniform history, no geometry supervision) | `train/train_gim.py` + `ssm.yaml` | | **CaM** | FoV-overlap retrieval of explicit context frames | `train/train_cam.py` | | **FramePack** | Recent contiguous window, older frames compressed with progressively coarser patchify kernels | `train/train_framepack.py` | All methods build on the Wan2.1-T2V-1.3B DiT (vendored in `wan/`), with per-block camera injection and per-frame action conditioning, trained with flow matching on the MIND dataset. ## Repository layout ``` train/ train_gim.py # GIM training (attention or SSM memory encoder) train_cam.py # Context-as-Memory baseline training train_framepack.py # FramePack baseline training rollout.py # autoregressive long-video rollout / mem_test export evaluate.py # multi-GPU quantitative evaluation validate.py # single-sample validation & sampling utilities configs/ gim.yaml # GIM (attention memory + geometry supervision) ssm.yaml # SSM memory baseline context_as_memory.yaml framepack.yaml data/ precompute_mind.py # offline latent / camera / text preprocessing mind_dataset.py # dataset + history selection (uniform/mi/fov/recent) models/ gim_dit.py # DiT forward with implicit memory tokens implicit_memory.py # attention & SSM memory compressors geometry_supervision.py # camera-queryable readout + VGGT distill loss wan_with_context.py # DiT forward with explicit context frames + camera framepack.py # FramePack multi-kernel history encoder action_condition.py # per-frame action conditioning utils/ memory_pruning.py # information-guided (MI) history pruning fov_retrieval.py # FoV-overlap context retrieval (CaM) precompute_vggt_encoder_features.py # VGGT encoder feature precompute scripts/ # single-node launch scripts (all local GPUs) wan/ # vendored Wan2.1 model code vggt/ # vendored VGGT (geometry supervision only) ``` ## Installation Python >= 3.10, CUDA GPU required. ```bash pip install -r requirements.txt ``` Downloads should use `./hfd.sh` by default. If it is missing, install it in the project root: ```bash wget -O hfd.sh https://hf-mirror.com/hfd/hfd.sh chmod a+x hfd.sh export HF_ENDPOINT=${HF_ENDPOINT:-https://hf-mirror.com} ``` Download the Wan2.1-T2V-1.3B base model next to this repository: ```bash ./hfd.sh Wan-AI/Wan2.1-T2V-1.3B --local-dir ../Wan2.1-T2V-1.3B ``` Place the MIND dataset at `../MIND` (or pass `--mind_root`). Use `./hfd.sh CSU-JPG/MIND --dataset --local-dir ../MIND` if downloading it from HuggingFace. ## Local assets, archive size, and deployment On the current machine, the project root is: ```bash /mnt/workspacev2/zhanggy/gim-world_pre ``` The dataset and model weights are available from this project through symlinks: | Project path | Real path | Current size | |---|---|---:| | `MIND` | `/mnt/workspacev2/zhanggy/Dataset/MIND` | 410G | | `Wan2.1-T2V-1.3B` | `/mnt/workspacev2/zhanggy/checkpoint/wan_models/Wan2.1-T2V-1.3B` | 17G | | `VGGT-1B` | `/mnt/workspacev2/zhanggy/checkpoint/VGGT-1B` | 9.4G | | `checkpoints/` | local project directory | 23G | The MIND directory includes both raw data and precomputed training caches: | Path | Current size | |---|---:| | `MIND/1st_data` | 17G | | `MIND/3rd_data` | 18G | | `MIND/latent_continuous` | 343G | Inside `latent_continuous`, the largest files are: | File type | Count | Total size | |---|---:|---:| | `vggt_encoder_tokens.pt` | 100 | 257.49G | | `latents.pt` | 250 | 82.82G | | `text_emb.pt` | 250 | 1.95G | Archive-size estimates: - Archive this repo without following symlinks: about 23G, mostly from `checkpoints/gim_smoke`. - Archive this repo without `checkpoints/` and without following symlinks: about 2.7M. - Archive this repo while following symlinks to include MIND and both model directories: about 458G before compression. Compression is expected to save little because `.pt`, `.safetensors`, and `.mp4` files are already dense or compressed. Recommended archive commands: ```bash # Run from the parent directory, for example: cd /mnt/workspacev2/zhanggy # Code/config/docs only. tar -czf gim-world_pre.code.tgz --exclude='checkpoints' gim-world_pre # Code plus local smoke checkpoint, but symlink targets are not included. tar -czf gim-world_pre.with_smoke_ckpt.tgz gim-world_pre # Full self-contained archive. Very large: about 430-460G compressed. tar --dereference -czf gim-world_pre.full_with_data_and_models.tgz gim-world_pre ``` ### Deploy on another machine Optional: install the Codex CLI on the target server if you want to use Codex from the terminal while setting up or maintaining the project. If you need the npm-based install, install Node.js/npm first, then install `@openai/codex`. Recommended npm install path for Linux/macOS servers: ```bash # 1. Install nvm, which installs Node.js and npm per user. curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.5/install.sh | bash # 2. Load nvm in the current shell. If this fails, close and reopen the shell. export NVM_DIR="$([ -z "${XDG_CONFIG_HOME-}" ] && printf %s "${HOME}/.nvm" || printf %s "${XDG_CONFIG_HOME}/nvm")" [ -s "$NVM_DIR/nvm.sh" ] && . "$NVM_DIR/nvm.sh" # 3. Install and use the latest LTS Node.js with npm. nvm install --lts --latest-npm nvm use --lts nvm alias default 'lts/*' # 4. Verify Node.js and npm. node -v npm -v # 5. Install Codex CLI with npm. npm install -g @openai/codex # 6. Verify the install, then start Codex. codex --version codex ``` If `codex` is not found after the npm install, add npm's global bin directory to `PATH`: ```bash export PATH="$(npm config get prefix)/bin:$PATH" echo 'export PATH="$(npm config get prefix)/bin:$PATH"' >> ~/.bashrc ``` On a headless server, use device-code login: ```bash codex login --device-auth ``` The official standalone installer is still available if npm is not required: ```bash # macOS/Linux standalone installer. curl -fsSL https://chatgpt.com/codex/install.sh | sh # Unattended install for remote setup scripts. curl -fsSL https://chatgpt.com/codex/install.sh | CODEX_NON_INTERACTIVE=1 sh ``` If the standalone installer works but the shell cannot find `codex`: ```bash export PATH="$HOME/.local/bin:$PATH" echo 'export PATH="$HOME/.local/bin:$PATH"' >> ~/.bashrc ``` Official references: `https://developers.openai.com/codex/cli`, `https://developers.openai.com/codex/quickstart`, `https://developers.openai.com/codex/auth`, `https://docs.npmjs.com/downloading-and-installing-node-js-and-npm`, and `https://github.com/nvm-sh/nvm`. Pick paths on the new machine first. The commands below use the same layout as the current machine: ```bash export PROJECT_ROOT=/mnt/workspacev2/zhanggy/gim-world_pre export MIND_ROOT=/mnt/workspacev2/zhanggy/Dataset/MIND export WAN_MODEL_DIR=/mnt/workspacev2/zhanggy/checkpoint/wan_models/Wan2.1-T2V-1.3B export VGGT_MODEL_DIR=/mnt/workspacev2/zhanggy/checkpoint/VGGT-1B ``` Clone the repository: ```bash mkdir -p "$(dirname "$PROJECT_ROOT")" git clone https://github.com/nagara214/gim-world_pre "$PROJECT_ROOT" cd "$PROJECT_ROOT" ``` Create or activate the Python environment, then install requirements: ```bash # Example environment; adjust to your cluster's conda location. conda create -n worldplay python=3.10 -y conda activate worldplay pip install -r requirements.txt ``` Download the raw MIND dataset and required model weights. Use `./hfd.sh` by default: ```bash mkdir -p "$(dirname "$MIND_ROOT")" "$(dirname "$WAN_MODEL_DIR")" "$(dirname "$VGGT_MODEL_DIR")" if [ ! -x ./hfd.sh ]; then wget -O hfd.sh https://hf-mirror.com/hfd/hfd.sh chmod a+x hfd.sh fi export HF_ENDPOINT=${HF_ENDPOINT:-https://hf-mirror.com} ./hfd.sh CSU-JPG/MIND \ --dataset \ --local-dir "$MIND_ROOT" ./hfd.sh Wan-AI/Wan2.1-T2V-1.3B \ --local-dir "$WAN_MODEL_DIR" ./hfd.sh facebook/VGGT-1B \ --local-dir "$VGGT_MODEL_DIR" ``` If a repository requires authentication, pass `--hf_username` and `--hf_token` to `./hfd.sh`. For example: ```bash ./hfd.sh ORG/REPO --hf_username "$HF_USERNAME" --hf_token "$HF_TOKEN" --local-dir /path/to/dir ``` Create project symlinks: ```bash cd "$PROJECT_ROOT" ln -sfn "$MIND_ROOT" MIND ln -sfn "$WAN_MODEL_DIR" Wan2.1-T2V-1.3B ln -sfn "$VGGT_MODEL_DIR" VGGT-1B ``` The default training configs currently refer to `../MIND` and `../Wan2.1-T2V-1.3B`, so either edit those paths in `train/configs/*.yaml`, or also create parent-directory symlinks: ```bash ln -sfn "$MIND_ROOT" "$(dirname "$PROJECT_ROOT")/MIND" ln -sfn "$WAN_MODEL_DIR" "$(dirname "$PROJECT_ROOT")/Wan2.1-T2V-1.3B" ``` The raw HuggingFace MIND download does not include GIM's precomputed `latent_continuous` cache. Either copy it from an existing machine: ```bash rsync -a --info=progress2 \ old_machine:/mnt/workspacev2/zhanggy/Dataset/MIND/latent_continuous \ "$MIND_ROOT"/ ``` or regenerate it: ```bash cd "$PROJECT_ROOT" MIND_ROOT="$MIND_ROOT" \ WAN_MODEL_DIR="$WAN_MODEL_DIR" \ OUTPUT_DIR="$MIND_ROOT/latent_continuous" \ bash scripts/precompute_mind.sh MIND_ROOT="$MIND_ROOT" \ PRECOMPUTED_DIR="$MIND_ROOT/latent_continuous" \ VGGT_MODEL_NAME="$VGGT_MODEL_DIR" \ bash scripts/precompute_vggt_encoder.sh ``` Expected precomputed-file counts after a complete setup: ```bash find -L MIND/latent_continuous -name latents.pt | wc -l # 250 find -L MIND/latent_continuous -name camera_rt.pt | wc -l # 250 find -L MIND/latent_continuous -name actions.pt | wc -l # 250 find -L MIND/latent_continuous -name text_emb.pt | wc -l # 250 find -L MIND/latent_continuous -name meta.pt | wc -l # 250 find -L MIND/latent_continuous -name vggt_encoder_tokens.pt | wc -l # 100 ``` ## 1. Data preprocessing Encode videos to VAE latents and precompute camera trajectories and T5 text embeddings (one-time, sharded across all local GPUs): ```bash bash scripts/precompute_mind.sh ``` Outputs go to `../MIND/latent_continuous/{1st_data,3rd_data}/...` with `latents.pt`, `camera_rt.pt`, `actions.pt`, `text_emb.pt`, `meta.pt` per video. ### Local preprocessing warning check For the local MIND preprocessing run under `/mnt/workspacev2/zhanggy/Dataset/MIND/latent_continuous`, all 250 videos produced the required `latents.pt`, `camera_rt.pt`, `actions.pt`, `text_emb.pt`, and `meta.pt` files. Two warning types from `train/data/precompute_mind.py` were checked: - `WARNING: action.json has ... entries but video has ... frames`: found in one sample, `1st_data/train/data-99`. Its `action.json` has 6891 entries, while the decoded video has 6962 frames, so preprocessing used the shared aligned length `T_effective=6891` and ignored the final 71 video frames. - `WARNING: VAE latents cover ... raw frames`: found in 187 of 250 samples. This is expected with continuous VAE temporal alignment because `encoded_raw = 1 + 4 * (T_latent - 1)` may leave the last 1-3 raw frames uncovered. The observed distribution was 65 samples dropping 1 frame, 65 samples dropping 2 frames, and 57 samples dropping 3 frames. No fatal preprocessing error was found from the generated metadata. The training smoke test log at `checkpoints/gim_smoke/train.log` also completed successfully. For GIM's geometry supervision, additionally precompute VGGT encoder features for the training split: ```bash bash scripts/precompute_vggt_encoder.sh ``` This writes `vggt_encoder_tokens.pt` per training video. VGGT-1B weights are pulled from `facebook/VGGT-1B` by default (`--vggt_model_name` to override). ## 2. Training All training scripts are single-node and use every local GPU via `accelerate`. Append `--resume_from_checkpoint <dir>` to resume. ```bash # GIM (attention memory encoder + MI pruning + geometry supervision) bash scripts/train_gim.sh # SSM memory baseline (shares the GIM trainer, compressor_type: ssm) bash scripts/train_ssm.sh # Context-as-Memory baseline (FoV-overlap frame retrieval) bash scripts/train_cam.sh # FramePack baseline (progressive history compression) bash scripts/train_framepack.sh ``` Checkpoints are written to `./checkpoints/{gim,ssm,cam,framepack}/`, each containing `model_weights/` (DiT), `cam_proj.pt`, `action_condition.pt` and the method-specific module (`compressor.pt` or `framepack.pt`). ## 3. Autoregressive rollout (inference) Generates long videos chunk-by-chunk, feeding generated frames back as history, and exports them in the layout expected by the MIND evaluator: ```bash # GIM / SSM baseline bash scripts/rollout.sh \ --ckpt_dir checkpoints/gim/step_XXXX \ --output_root ./exports --model_name gim \ --mode memory --perspective 1st # CaM (explicit context, FoV retrieval at inference) bash scripts/rollout.sh \ --ckpt_dir checkpoints/cam/step_XXXX \ --output_root ./exports --model_name cam \ --mode context --context_selection fov --perspective 1st # FramePack bash scripts/rollout.sh \ --ckpt_dir checkpoints/framepack/step_XXXX \ --output_root ./exports --model_name framepack \ --mode framepack --perspective 1st ``` For the SSM variant add `--compressor_type ssm`. Videos are sharded across local GPUs; outputs land at `{output_root}/{model_name}/{perspective}_data/mem_test/{video}/video.mp4`. ## 4. Evaluation Quantitative evaluation (PSNR / LPIPS over memory-probing chunks) on the validation split: ```bash bash scripts/evaluate.sh \ --ckpt_dir checkpoints/gim/step_XXXX \ --output_dir eval_results --perspective 1st ``` The exported `mem_test` videos from step 3 can also be scored directly with the official MIND evaluator. ## License Code built on Wan2.1 is subject to the Apache 2.0 License (`LICENSE.txt`).