Elastic-forcing's picture
Update documentation
066102e verified
|
Raw History Blame Contribute Delete
4.42 kB

Elastic Forcing — 14B

Overview · 1.3B guide · Checkpoints

The 14B model uses Wan2.1-T2V-14B with DINOv3, V-JEPA 2, and VideoMAEv2 for distribution matching. It generates 81-frame videos at 832 × 480 resolution and 16 FPS. Inference uses four causal denoising steps followed by terminal cleanup.

Environment

Use Linux, Python 3.10, and NVIDIA GPUs with BF16 support. Run these commands from 14B/:

conda create -n elastic-forcing-14b python=3.10 -y
conda activate elastic-forcing-14b
python -m pip install torch==2.5.1 torchvision==0.20.1 \
  --index-url https://download.pytorch.org/whl/cu121
python -m pip install -r requirements.txt
python -m pip install flash-attn==2.7.4.post1 --no-build-isolation
export PYTHON_BIN="$(command -v python)"

Models and assets

Inference requires the Wan2.1-T2V-14B base model, T5 text encoder, tokenizer, VAE, and the Elastic Forcing generator checkpoint. Place the base-model bundle under $ASSET_ROOT/models/Wan2.1-T2V-14B/.

The generator checkpoint is stored as eight FP32 safetensors shards. Download them and create the .pt file expected by the inference launcher:

import json
from pathlib import Path
import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file

folder = Path(snapshot_download(
    "Elastic-forcing/Elastic-forcing-checkpoint",
    allow_patterns=["checkpoints/14B/*"],
)) / "checkpoints/14B"
index = json.loads((folder / "model.safetensors.index.json").read_text())
state = {}
for shard in sorted(set(index["weight_map"].values())):
    state.update(load_file(str(folder / shard), device="cpu"))
torch.save({"generator": state}, "/path/to/local/model.pt")

Conversion preserves tensor values and precision. Allow enough host memory and disk space for the approximately 57 GB FP32 generator and the output file.

For inference, prepare the base-model assets and generate the local configuration:

export ASSET_ROOT=/path/to/assets
python scripts/prepare_config.py --asset-root "$ASSET_ROOT" --inference-only

Inference

CHECKPOINT=/path/to/local/model.pt \
GPUS=0 \
PROMPTS=code/prompts/test_8.txt \
OUTPUT_DIR=outputs/demo \
bash scripts/infer.sh

The defaults are seed 0 and five samples per prompt with non-EMA generator weights. Set NUM_SAMPLES=1 for a single sample. The generator executes one video at a time while retaining the sampling order across prompts. Model loading uses local base-model assets.

Training

Training requires the causal ODE initialization, frozen encoders, matched prompt/video manifest, reference feature bank, and Nyström summaries. Prepare the model and reference assets at the locations specified in configs/train.yaml, then run:

python scripts/prepare_config.py --asset-root "$ASSET_ROOT"
GPUS=0,1,2,3,4,5,6,7 LOGDIR=runs/train bash scripts/train.sh
Setting Value
Training updates 80
Distribution / differentiation / reference batch 256 / 128 / 128
Per-GPU generated / differentiated videos 32 / 16
Reference videos 8,000
Encoder weights DINOv3 : V-JEPA 2 : VideoMAEv2 = 10 : 1 : 1
Loss divisor 3
Nyström / Monte Carlo weights 2 : 1
Optimizer AdamW, betas (0, 0.999), weight decay 0
Learning rate Cosine from 2.5e-6 to 5e-7 over 200 updates
Gradient clipping Maximum norm 10
Checkpoint interval Every ten updates
Memory management FSDP, activation checkpointing, optimizer and encoder CPU offload

The learning-rate schedule spans 200 updates while training stops at update 80. Use a separate LOGDIR for each run. Keep checkpoint companion files for resuming training with the same distributed topology and asset configuration.

To resume:

python scripts/prepare_config.py --asset-root "$ASSET_ROOT" \
  --resume-checkpoint /path/to/checkpoint/model.pt
GPUS=0,1,2,3,4,5,6,7 LOGDIR=runs/resume bash scripts/train.sh

Evaluation and tests

Generate benchmark videos with consistent prompt ordering and sampling settings, then evaluate them with the official VBench implementation.

Run the CPU numerical tests from 14B/:

python -m pip install pytest
python scripts/test_cpu.py