Blame for 14B/README.md

# Elastic Forcing — 14B

[Overview](../README.md) · [1.3B guide](../1.3B/README.md) · [Checkpoints](https://huggingface.co/Elastic-forcing/Elastic-forcing-checkpoint/tree/main/checkpoints/14B)

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/`:

```bash
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:

```python
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:

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

## Inference

```bash
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](configs/train.yaml), then run:

```bash
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:

```bash
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](https://github.com/Vchitect/VBench) implementation.

Run the CPU numerical tests from `14B/`:

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