--- license: mit base_model: - hanshanxue/WorldStereo - Wan-AI/Wan2.1-I2V-14B-480P-Diffusers tags: - safetensors - fp8 - quantized - video - image-to-video - wan - worldstereo pipeline_tag: image-to-video --- # WorldStereo-memory-dmd-fp8_e4m3fn_scaled A scaled FP8 (e4m3fn) quantization of `hanshanxue/WorldStereo`'s `worldstereo-memory-dmd` variant. Drops VRAM from **35 GB → 20 GB** so the 14B WorldStereo DiT fits resident on 24 GB consumer GPUs (3090, 4090) instead of needing partial-load streaming. | | bf16 source | fp8_e4m3fn_scaled (this repo) | |---|---|---| | Size on disk | 34.86 GB | 20.35 GB (**-41.6%**) | | Tensors | 1799 (all BF16) | 551 FP8 + 1800 BF16 | | Fits in 24 GB VRAM | No (needs partial_load) | **Yes (resident)** | | Quality | reference | typical fp8_scaled drop: <1% perceptible on short clips, faint color drift / micro-detail loss possible on long sequences ([refs](#quality)) | | Speed on Ampere (3090) | reference | same (no native fp8 matmul; weight upcast to bf16 per matmul) | | Speed on Ada/Hopper (4090, H100) | reference | **>2x** via `torch._scaled_mm` | ## Quantization recipe (verbatim from [kijai/ComfyUI-WanVideoWrapper](https://github.com/kijai/ComfyUI-WanVideoWrapper)) For each tensor `T` in the source state-dict: 1. If `T`'s name contains any of these keywords -> **stored unchanged in BF16**: ``` norm, bias, time_in, time_, patch_embedding, img_emb, modulation, text_embedding, adapter, add, ref_conv, audio_proj ``` This keeps numerically-sensitive small tensors (norms, biases, modulation, I2V cross-attn projections) at full precision. 2. If `T` is `.weight` AND has rank >= 2 AND survived the exclusion check -> **cast to fp8_e4m3fn with per-tensor scale**: ```python FP8_MAX = 448.0 scale = T.float().abs().amax() / FP8_MAX T_fp8 = (T.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) ``` Stored as two tensors: - `` (the fp8 weight) - `.scale_weight` (bfloat16 scalar) 3. A marker tensor `scaled_fp8` (bf16 zero) is added so loaders that check `"scaled_fp8" in state_dict` auto-detect the scaled-fp8 format. ## Inference At runtime the scale must be re-applied. Two paths: ### Ampere (3090) - upcast on matmul ```python def fp8_scaled_linear(x, w_fp8, scale, bias): return torch.nn.functional.linear(x, w_fp8.to(x.dtype) * scale, bias) ``` ### Ada/Hopper (4090, H100) - native fp8 matmul ```python def fp8_scaled_linear_fast(x, w_fp8, scale_w, bias, base_dtype=torch.bfloat16): x = x.clamp(-448, 448).to(torch.float8_e4m3fn).contiguous() scale_in = torch.ones((), device=x.device, dtype=torch.float32) return torch._scaled_mm(x.view(-1, x.shape[-1]), w_fp8.t(), out_dtype=base_dtype, bias=bias, scale_a=scale_in, scale_b=scale_w) ``` Reference implementations: - `kijai/ComfyUI-WanVideoWrapper/fp8_optimization.py` (fast path) - `kijai/ComfyUI-WanVideoWrapper/nodes_model_loading.py` `_replace_linear` (slow path) ## Naming convention This file uses **diffusers-style keys** (`blocks.0.attn1.to_q.weight`, `blocks.0.ffn.net.0.proj.weight`), matching the original `hanshanxue/WorldStereo` source layout. This is different from Kijai's published Wan2.1 fp8 files which use **native Wan keys** (`blocks.0.self_attn.q.weight`). If your loader expects native names, run a key remap (see `diffusers_to_native_wan` in [ComfyUI-WorldStereo](https://github.com/pozzettiandrea/ComfyUI-WorldStereo)). ## What's in the file - **549 FP8 tensors** (14.5 GB): rank-2 Linear weights only (self-attn Q/K/V/O, FFN projections, controlnet linears). Conv weights stay BF16: no scaled-fp8 Conv kernel exists in PyTorch / ComfyUI - **1800 BF16 tensors** (5.84 GB): - 362 `norm` tensors (RMSNorm gains): 1.26 GB - 80 `add_*` tensors (I2V cross-attention to image embeddings): 4.19 GB - 3 `time_*` tensors (timestep MLP): 0.37 GB - 759 `bias` + others: ~0.02 GB - 551 `*.scale_weight` scalars: 1.1 MB - 1 marker tensor `scaled_fp8` ## Quality This recipe (Kijai's fp8_e4m3fn_scaled) is the community standard for Wan2.1 14B inference on 24 GB GPUs. Documented behavior: - **Subjective**: indistinguishable from bf16 on most short clips (Kijai's own statement: *"never any difference remotely [from fp16]"* in his per-tensor-scaled tests). - **Known artifacts on long clips**: faint color drift, occasional minor detail loss; usually not visible side-by-side without close inspection. References: - [Comfy-Org repackaged: fp8_e4m3fn vs fp8_scaled discussion](https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/discussions/22) - [Kijai WanVideo_comfy_fp8_scaled artifacts discussion](https://huggingface.co/Kijai/WanVideo_comfy_fp8_scaled/discussions/31) - **Round-trip on a sample 5120x5120 attention weight**: max_abs_err = 0.005 (3.7% of the tensor's max value, ~standard for e4m3fn-scaled). If you observe quality regressions specific to WorldStereo's camera/scene conditioning, fall back to bf16 with partial-load streaming. ## Provenance - **Source**: `hanshanxue/WorldStereo` (`worldstereo-memory-dmd` variant, commit `2adb716`) - **Source base model**: `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` (Apache-2.0) - Verified: WorldStereo's Wan base is **frozen** (sampled tensors differ from upstream only at fp16->bf16 ULPs). The novel content is the 1.76 GB controlnet branch under `controlnet.controlnet_blocks.*` -- WorldStereo's camera/scene-render conditioning. - **Quantization tool**: `cook_worldstereo_fp8.py` (in this repo) ## License - **WorldStereo overlay**: MIT (inherited from `hanshanxue/WorldStereo`) - **Wan2.1 base weights**: Apache-2.0 (inherited from `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers`) A NOTICE file in this repo provides Apache-2.0 attribution and a statement of changes per section 4 of the license. ## Files - `WorldStereo-memory-dmd-fp8_e4m3fn_scaled.safetensors` - the quantized model - `cook_worldstereo_fp8.py` - the quantization script (reproducible) - `config.json` - the source `worldstereo-memory-dmd` config (carries the scale_map, base_model path, sampling settings) - `README.md` - this file - `NOTICE` - Apache-2.0 attribution - `LICENSE` - MIT (for the WorldStereo overlay portion)