cretz/FastWan2.2-TI2V-5B-ONNX-sharded

FastWan 2.2 TI2V-5B - ONNX (browser / onnxruntime-web)

ONNX export of FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers, sharded and quantized for in-browser inference via onnxruntime-web on WebGPU. The base model is a DMD-distilled Wan 2.2 TI2V-5B that runs at 4 UniPC steps. Pipeline-side we match KingNish/wan2-2-fast: UniPCMultistepScheduler with flow sigmas, predict_x0, bh2, solver_order=2, corrector, flow_shift=8.0 (the Space overrides the config's 5.0). Output is 81 frames @ 16 fps - 5 second clips.

Status: known quality issues

This export runs end-to-end in a browser tab on desktop WebGPU, but output quality does not match the reference HF Space at the same prompt. The result is coherent (recognizable scenes) but blocky and topically off. Two known contributors:

• Text encoder q4f16 drift. CPU-ORT vs PyTorch shows the 24-layer UMT5-XXL composes to ~47% relative error at q4f16 (MatMulNBits is too lossy for UMT5's 50k-magnitude intermediates, even at --accuracy-level 1). The fp16 text encoder passes (max|diff|=0.002) and is included in this repo for callers that can afford the bandwidth and memory.
• WebGPU-side numeric divergence somewhere in transformer or VAE. CPU ORT inference of the transformer blocks is faithful to PyTorch (29 of 30 blocks pass cleanly, one block has a single-element outlier with negligible mean error). So the blockiness lives on the GPU kernel side or in the VAE. Not yet root-caused.

If you need quality parity with the HF Space today, run the original diffusers checkpoint. This export is for browser deployment, where those are not options.

Repo layout

onnx/
  vae_decoder-480.onnx              (29 MB, LightTAE fp16, latent 30x30)
  vae_decoder-576.onnx              (33 MB, LightTAE fp16, latent 36x36)

  text-encoder/                     UMT5-XXL fp16, per-layer
    embedding.bin                   raw fp16 embedding table [256384, 4096]
    layer_00.onnx ... layer_23.onnx (+ .data sidecars)
    shell_post.onnx                 final LayerNorm

  text-encoder-q4f16/               UMT5-XXL q4f16, per-layer
    embedding_q8.bin                int8 embedding rows [256384, 4096]
    embedding_scales.bin            fp16 per-row scale [256384]
    layer_00.onnx ... layer_23.onnx (+ .data sidecars)
    shell_post.onnx                 final LayerNorm

  transformer-480/                  fp16 transformer at 480x480
    shell_pre.onnx (+ .data)        patch embed + RoPE + cond projection
    block_00.onnx ... block_29.onnx (+ .data sidecars)
    shell_post.onnx                 final norm + unpatchify
  transformer-576/                  fp16 transformer at 576x576
    (same layout as transformer-480)

  transformer-q4f16-480/            q4f16 transformer at 480x480
    shell_pre.onnx                  fp16 (intentionally NOT quantized,
                                    see "shell_pre is fp16" below)
    block_00.onnx ... block_29.onnx (+ .data sidecars, q4f16)
    shell_post.onnx                 fp16
  transformer-q4f16-576/            q4f16 transformer at 576x576
    (same layout as transformer-q4f16-480)

tokenizer/
  tokenizer.json
  special_tokens_map.json

Resolution-bound dirs use a -${res} suffix; text-encoder and tokenizer are resolution-agnostic.

How each artifact was made

All export and quantization scripts are plain Python on top of torch.onnx.export + onnxruntime.tools.MatMulNBitsQuantizer. CPU-only torch is sufficient.

LightTAE VAE decoder (vae_decoder-${res}.onnx)

lightx2v/lighttaew2_2 tiny decoder, which decodes Wan 2.2 latents at a fraction of the cost of the full AutoencoderKLWan. Exported with a decoder-only key filter, dynamo=False (Windows-Unicode workaround), pixel_shuffle reshaped 4D->5D around the call site. fp16 throughout. Per-resolution because the spatial shape is baked in at export time.

The full AutoencoderKLWan was also exported (Conv3D decomposed into 3xConv2D + temporal sum to dodge ORT-web's slow / TDR-prone Conv3DNaive kernel) and a 32-tensor cache I/O contract between decoder_init + decoder_step graphs let it stream frames. An A/B at 480x832 found it did not improve quality over LightTAE, so this repo ships LightTAE only.

UMT5-XXL text encoder

24 layers of UMT5-XXL exported one at a time. Each layer takes hidden_states [B, seq, 4096] + attention_mask [B, 1, 1, seq] and returns hidden_states_out. Two patches were required:

• aten::rms_norm is unsupported in opset 23, so RMSNorm was monkey-patched to its decomposed form before tracing.
• UMT5's LayerNorm overflows fp16 variance on long sequences during the trace. Patched to a fp32 decomposition for the variance reduction, casting back to fp16.

The 2.1 GB embedding table is extracted to a raw binary (embedding.bin) and looked up JS-side rather than living in the ONNX graph - WebGPU's maxBufferSize on mobile cannot hold a single 2.1 GB buffer.

For text-encoder-q4f16/, the same per-layer ONNX files are then quantized with MatMulNBitsQuantizer(block_size=32, symmetric=True, accuracy_level=1). The embedding is separately quantized per-row to int8 symmetric (2.1 GB fp16 -> 1.05 GB int8 + 0.5 MB fp16 scales). JS-side lookup: row[j] = int8[id*4096 + j] * scale_fp16[id].

Transformer

Per-block instantiation: each of the 30 DiT blocks is built and exported in isolation, with peak export RAM around 2 GB. accelerate's disk-offload hooks were avoided because they trigger an aten::view(Tensor, int) tracer assert during ONNX export. fp16 throughout.

shell_pre.onnx produces the per-token tensors the blocks consume (tokens, enc_proj, timestep_proj, temb, freqs_cos, freqs_sin). shell_post.onnx does the final norm + unpatchify back to [B, 48, T, H, W].

attn1 in every block is rewritten at export time to chunk Q along the sequence axis: each chunk runs MatMul + Softmax + MatMul independently and the results are Concat'd. Required because at larger resolutions the Q.K^T intermediate is bigger than the 2 GiB WebGPU maxBufferSize (e.g. 3.07 GiB at seq=8190); the original op silently short-allocates and produces garbage. The chunker is auto-skipped at shapes where Q.K^T fits (e.g. 480x480 / seq=4725).

For transformer-q4f16-${res}/, blocks are then quantized with MatMulNBitsQuantizer(block_size=32, symmetric=True, accuracy_level=4).

shell_pre is fp16 (intentional)

In transformer-q4f16-${res}/, shell_pre.onnx is shipped as unquantized fp16 (180 MB) instead of q4 (52 MB). ORT-web's WebGPU MatMulNBits kernel at accuracy_level=4 produces NaN in shell_pre's small q4 MatMuls - one column is NaN at every token position, on WebGPU only. CPU ORT at both fp16 and q4 was clean. Until that kernel is fixed, fp16 shell_pre is the workaround.

shell_post.onnx is also fp16: it's only 1.2 MB, too small to bother quantizing.

Tokenizer

tokenizer.json and special_tokens_map.json are copied straight from the upstream FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers repo, loaded via transformers.js T5Tokenizer.

Inference contract

Tensor names and dtypes in this repo match the diffusers source. The key shapes (default seq_len = T * (H/p_h) * (W/p_w) = 21 * 15 * 26 = 8190 for 480x832; at 480x480 it's 21 * 15 * 15 = 4725):

• transformer/shell_pre.onnx: [B, 48, T, H, W] latent + [B, seq_len] int64 timestep + [B, 512, 4096] text embeds in; tokens [B, seq_len, 3072], enc_proj [B, 512, 3072], timestep_proj [B, seq_len, 6, 3072], temb [B, seq_len, 3072], freqs_cos [1, seq_len, 1, 128] fp32, freqs_sin [1, seq_len, 1, 128] fp32 out.
• transformer/block_NN.onnx: same per-block inputs, same shape hidden_states_out [B, seq_len, 3072] out.
• transformer/shell_post.onnx: hidden_states [B, seq_len, 3072] + temb + scalar int64 ppf, pph, ppw (= T, H/p_h, W/p_w) -> noise_pred [B, 48, T, H, W].
• vae_decoder-${res}.onnx: latent [1, 16, 1, h, w] -> frames [1, 3, 81, res, res] in [-1, 1].

Note freqs_cos / freqs_sin are fp32 (not fp16) by design - rotary position embeddings need the precision.

Pipeline notes for callers

A few details that are easy to get wrong porting the diffusers pipeline:

• Timestep is [B, seq_len], not scalar. Wan 2.2 TI2V gives every token its own timestep value.
• Latents need per-channel denormalization before VAE decode: x = x * latents_std[c] + latents_mean[c]. Constants live in the upstream model's vae/config.json.
• Padded text-embed positions must be zero-filled (the transformer blocks have no encoder_attention_mask input; without zeroing, the prompt drowns in padding, symptom: prompt-agnostic fabric texture).
• fp16 readback in JS: modern Chrome's Float16Array is not bit-compatible with new Uint16Array(tensor.data) - the latter numerically converts instead of reinterpreting bits. Use a bit-copy helper.

License

Inherits the base model's Apache 2.0 license. See the upstream model card for usage terms.