AsadIsmail/CogVideoX-2b-ternary

CogVideoX-2B — Ternary Quantized (tritplane3)

Storage format note. This checkpoint is stored dequantized to FP16. Loaded with stock diffusers it runs like the base model — the memory/compute benefits of ternary are not realized in this format. It demonstrates that ternary PTQ preserves generation quality. The packed 2-bit ternary runtime (`ternary-quant`) currently targets transformers LLMs/VLMs, not diffusers video pipelines, so there is no accelerated ternary inference path for this model yet. Packed weights and a diffusers runtime are on the roadmap.

Ternary-quantized version of zai-org/CogVideoX-2b.

Produced with ternary-quant using component-aware tritplane3 quantization applied to the Diffusion Transformer (DiT) backbone.

Why this matters

This is a proof-of-concept for ternary post-training quantization on a diffusers-based DiT pipeline. It should be treated as an experimental artifact, not a benchmarked replacement for FP8, int8, or other production video quantization paths.

The same component-aware workflow can be tested on other text-to-video DiT models, but each architecture needs its own validation.

Model Specifications

Verified Working

Generated videos with the quantized pipeline:

• —Prompt: "a cat walking on green grass"
• —Resolution: 480×720, 9 frames
• —Steps: 25 (recommended; 5 is too few)
• —Seed: 42
• —Device: MPS (Apple Silicon), bfloat16

Output is coherent — shows a cat on green grass, natural anatomy, temporal consistency. See test_ternary_25steps.mp4 in the repo.

Quality vs FP16 original: Both produce valid outputs at 25 steps. Per-pixel PSNR is ~13 dB (expected — diffusion models produce stochastic outputs, pixel-level comparison is not meaningful between independent runs even with the same seed).

Size & Compression

Honest note: This repo ships the transformer with ternary-quantized weights dequantized back to FP16 for drop-in compatibility with the standard diffusers pipeline. On-disk size matches the original. The weights have ternary precision (~8 effective bits) but are stored in FP16 format. For actual 2× disk compression, weights would need to be saved in ternary-quant's packed tritplane format (requires custom inference wrapper).

This proves the quantization works end-to-end without requiring users to install anything beyond standard diffusers.

Memory Requirements (Inference)

Quickstart

pip install diffusers transformers accelerate tiktoken sentencepiece protobuf imageio imageio-ffmpeg

import torch
from diffusers import DiffusionPipeline
from diffusers.utils import export_to_video

pipe = DiffusionPipeline.from_pretrained(
    "AsadIsmail/CogVideoX-2b-ternary",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
)

# MPS workaround: cast float64 scheduler buffers to float32
device = "mps"  # or "cuda"
if device == "mps":
    for attr in ("alphas_cumprod", "betas", "alphas", "sigmas"):
        val = getattr(pipe.scheduler, attr, None)
        if torch.is_tensor(val) and val.dtype == torch.float64:
            setattr(pipe.scheduler, attr, val.float())

pipe.to(device)
pipe.enable_attention_slicing()

result = pipe(
    prompt="a cat walking on green grass",
    num_frames=9,
    num_inference_steps=25,   # use 25+, 5 is too few
    guidance_scale=6.0,
    height=480, width=720,
    generator=torch.Generator(device=device).manual_seed(42),
)

export_to_video(result.frames[0], "output.mp4", fps=8)

Limitations

• —Storage not reduced (dequantized FP16 format) — see honest note above
• —Compute not accelerated — standard FP16 GEMM, no specialized ternary kernels
• —5-step inference is too aggressive — anatomy artifacts (e.g., wrong eye counts). Use 25+ steps.
• —MPS workaround required (float64 scheduler buffers)
• —CPU inference is impractical (3+ hours for 5 frames on M4 Pro)

Reproduce

The quantization pipeline is at github.com/Asad-Ismail/ternary-models/tree/main/video.

# 1. Load pipeline
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained("zai-org/CogVideoX-2b", torch_dtype=torch.float16)

# 2. Extract DiT transformer
transformer = pipe.transformer

# 3. Capture activations on calibration forward pass
# 4. Quantize each nn.Linear using ProgressiveTritPlaneQuantizer
# 5. Replace weights with dequantized versions
# 6. Save pipeline

See the full script: `scripts/quantize_dit.py`.

Collection

Part of ternary-models — ternary-quantized VLMs, multimodal, audio, and video models.

GitHub: github.com/Asad-Ismail/ternary-models | Library: github.com/Asad-Ismail/ternary-quant

Citation

@software{ternary_quant,
  author = {Ismail, Asad},
  title = {ternary-quant: Post-training ternary quantization for HuggingFace generative models},
  url = {https://github.com/Asad-Ismail/ternary-quant},
  year = {2026}
}