weihang44/Causal-Forcing-Memory-Checkpoints

Causal Forcing CR10 checkpoints: 2k, 3k, 4k, 5k

This public release contains four CR10 model checkpoints from the same Causal Forcing compressor training run: steps 2000, 3000, 4000, and 5000. Tensor files are preserved byte-for-byte with SHA256 checksums.

Available checkpoints

Step	Checkpoint directory	Model size
2000	checkpoints/cf-cr10-lr1e4-cos2e5-5k/iter_000002000	5.682 GB
3000	checkpoints/cf-cr10-lr1e4-cos2e5-5k/iter_000003000	5.682 GB
4000	checkpoints/cf-cr10-lr1e4-cos2e5-5k/iter_000004000	5.682 GB
5000	checkpoints/cf-cr10-lr1e4-cos2e5-5k/iter_000005000	5.682 GB

Total model size is approximately 22.73 GB. The memory ratio is 0.1 (CR10). Original training provenance and exact file sizes are recorded in inventory.json.

Each model/ directory is a PyTorch Distributed Checkpoint (DCP) containing both frozen net.* generator state and trained net_compressor.* state. Keep .metadata and all .distcp shards together, including empty rank shards. This release contains model weights and a training config; optimizer, scheduler, and trainer state are not included.

The saved config contains original cluster paths that must be adapted on another installation. Its credential fields are empty.

Download

from huggingface_hub import snapshot_download

step = 2000  # 2000, 3000, 4000, or 5000
root = snapshot_download(
    repo_id="weihang44/Causal-Forcing-Memory-Checkpoints",
    allow_patterns=[
        f"checkpoints/cf-cr10-lr1e4-cos2e5-5k/iter_{step:09d}/model/*",
        "checkpoints/cf-cr10-lr1e4-cos2e5-5k/config.yaml",
        "manifest.json",
    ],
    local_dir="cf-models",
)

To download all four checkpoints, omit allow_patterns.

Load the compressor in the matching inference repository

Use the supported inference environment and the source pinned at video-compress commit bdbf3f2e903392e855638f19b1d555c150904d96. Its loader reads DCP checkpoints without requiring the original training world size, extracts net_compressor.*, and checks the target state dictionary strictly. It also handles the optional unused AE task embedding in older checkpoints.

From that repository root, after downloading the checkpoint:

import sys
import torch

sys.path.insert(0, "scripts/context_compression")
from generate_compressed_context import build_compressor, load_compressor_weights

compressor = build_compressor("1.3B", 0.1, 1024)
load_compressor_weights(
    compressor,
    "cf-models/checkpoints/cf-cr10-lr1e4-cos2e5-5k/iter_000002000",
)
compressor = compressor.to(device="cuda", dtype=torch.bfloat16).eval()

This loads model weights; generation still requires the CF generator, text encoder, VAE, and the chosen cache policy. Selection-free fullraw21, sink count, and recent-frame count are inference policies, not different checkpoint identities. See the published CF method and review index for the exact evaluated setup and its limitations.

Integrity and upstream dependencies

manifest.json and SHA256SUMS cover every model/config file in this release. After downloading all four checkpoints, run sha256sum -c SHA256SUMS from the download directory.

The original standalone generator can be obtained from zhuhz22/Causal-Forcing. Its SHA256 is cf75ee5cc6f4e2e336c59c973f5544655d8f0aa481761efe6de1b9cb2eb0cd9d. The matching inference setup also requires its text encoder and VAE.

These are CF-trained CR10 checkpoints. The training run, iteration, and file hashes identify each checkpoint. See NOTICE.md for upstream attribution.