Skip to main content

vidax

A Unified JAX Framework for Video Generative Models on Accelerator Meshes

Showcase

Wan2.1 14B1280x720
A cross-attention DiT conditioned by a UMT5-XXL text encoder and sampled with a Rectified-Flow Euler scheduler. The I2V variant adds a separate CLIP-image cross-attention branch per block.
Wan2.2 A14B832x480
Wan's DiT rebuilt as a two-expert Mixture-of-Experts routed by timestep, with per-token (not per-sample) modulation. TI2V-5B is instead one dense model, sharing a single checkpoint across T2V and I2V.
Cosmos-Predict2.5 14B1280x704
A DiT with AdaLN-LoRA modulation, conditioned by a 7B Qwen2.5-VL ("Reason1") text tower with every decoder layer concatenated, and sampled with a flow-matching UniPC multistep solver.
Cosmos3 Nano (16B)1280x704
An omnimodal Mixture-of-Transformers — a full weight set per modality — with dual causal / full-attention pathways, interleaved 3D mRoPE, no AdaLN, and a Karras-sigma UniPC sampler.
LTX-Video 13B dev1216x704
A from-scratch LTX port: a pixel-unshuffle-patchify VAE with a noise-conditioned decoder, a plain T5-XXL text encoder, and a Rectified-Flow Euler sampler on a linear-quadratic sigma schedule.
LTX-2.5 22B dev1216x704
A 22B LTX DiT with cross-attention AdaLN and per-head gated attention, an 8-layer embeddings connector, a Gemma-4 12B text encoder, and an ancestral (SDE) Euler sampler.
HunyuanVideo-1.5 8.3B (720p)1280x720
8.3B MMDiT that conditions by concatenating Qwen2.5-VL and byT5 text tokens directly onto the video sequence — no cross-attention. SigLIP adds the reference frame for I2V. Rectified-Flow Euler.
HunyuanVideo 13B1280x720
13B dual-stream + single-stream MMDiT with joint full-sequence self-attention, conditioned by a Llama-3-8B tower plus a pooled CLIP-L vector folded into AdaLN. Flow-matching Euler; token_replace for I2V.
CogVideoX1.5 5B1360x768
CogVideoX with temporal patchification (patch_size_t=2), a "slice"-grid 3D RoPE, and a native 1360×768 / 81-frame recipe run under DeepSpeed-Ulysses sequence parallelism.
Browse all samples in the Gallery →

Benchmark Explorer

TPUWeight
dtype
I/O
dtype
Cosmos3 Edge (4B)T2V832x480TPU v4bf16bf162.909274.80.44017.215FlyingGiraffe
CogVideoX 2BT2V720x480TPU v4bf16bf164.230519.20.09417.185FlyingGiraffe
LTX-2.5 22B (distilled)T2V1216x704TPU v4bf16bf164.663215.00.56315.355FlyingGiraffe
LTX-Video 13B (dev)T2V1216x704TPU v4bf16bf165.200347.60.34815.275FlyingGiraffe
LTX-Video 2B (distilled)T2V1216x704TPU v4bf16bf165.916162.60.7448.845FlyingGiraffe
Wan2.1 1.3BT2V832x480TPU v4bf16bf166.965456.40.17710.195FlyingGiraffe
Cosmos3 Nano (16B)T2V1280x704TPU v4bf16bf167.141437.50.21329.495FlyingGiraffe
LTX-2.5 22B (dev)T2V1216x704TPU v4bf16bf167.262409.70.29516.695FlyingGiraffe
CogVideoX 5BI2V720x480TPU v4bf16bf169.413851.10.05823.335FlyingGiraffe
CogVideoX 5BT2V720x480TPU v4bf16bf169.413835.30.05923.185FlyingGiraffe
Wan2.2 TI2V-5BT2V1280x704TPU v4fp32bf1610.518649.70.18618.295FlyingGiraffe
Wan2.2 TI2V-5BI2V800x1088TPU v4fp32bf1612.068665.40.18218.265FlyingGiraffe
LTX-Video 13B (distilled)T2V1216x704TPU v4bf16bf1613.028300.60.40315.275FlyingGiraffe
Wan2.1 14BT2V832x480TPU v4bf16bf1626.1361552.90.05217.175FlyingGiraffe
Wan2.1 14BI2V544x720TPU v4bf16bf1628.1321372.10.05922.085FlyingGiraffe
Cosmos-Predict2.5 2BT2V1280x704TPU v4bf16bf1638.7791497.30.06216.035FlyingGiraffe
Wan2.2 A14BT2V832x480TPU v4fp32bf1643.1822408.20.03428.385FlyingGiraffe
Wan2.2 A14BI2V544x720TPU v4fp32bf1644.4992063.20.03928.355FlyingGiraffe
Wan2.2 A14BI2V832x1104TPU v4fp32bf1649.0622163.70.01520.525FlyingGiraffe
CogVideoX1.5 5BI2V1360x768TPU v4bf16bf1652.7963400.80.02431.455FlyingGiraffe
CogVideoX1.5 5BT2V1360x768TPU v4bf16bf1652.7963380.20.02431.465FlyingGiraffe
LTX-2.5 22B (dev, diffusion VAE)T2V1216x704TPU v4bf16bf1695.3203444.30.03516.115FlyingGiraffe
HunyuanVideo-1.5 8.3B (480p)I2V544x720TPU v4bf16bf16112.8803838.10.03231.445FlyingGiraffe
HunyuanVideo-1.5 8.3B (480p)T2V832x480TPU v4bf16bf16119.4594015.20.03029.145FlyingGiraffe
Wan2.1 14BT2V1280x720TPU v4fp32bf16123.0116562.90.01223.005FlyingGiraffe
Wan2.1 14BI2V832x1104TPU v4fp32bf16127.2495651.60.01432.705FlyingGiraffe
Cosmos-Predict2.5 14BT2V1280x704TPU v4bf16bf16127.9884577.10.02014.705FlyingGiraffe
HunyuanVideo-1.5 8.3B (720p)I2V832x1104TPU v4bf16bf16219.1867125.00.01732.005FlyingGiraffe
HunyuanVideo-1.5 8.3B (720p)T2V1280x720TPU v4bf16bf16220.9937133.50.01730.285FlyingGiraffe
HunyuanVideo 13BT2V1280x720TPU v4bf16bf16299.61415620.70.00818.405FlyingGiraffe
HunyuanVideo 13BI2V832x1088TPU v4bf16bf16306.36916007.40.00818.415FlyingGiraffe
LTX-2.5 22B (distilled, diffusion VAE)T2V1216x704TPU v4bf16bf16335.0953247.40.03714.765FlyingGiraffe

Acknowledgments

Development supported by the Google TPU Research Cloud (TRC) program.