ZIT HSWQ FP8 Quantizer (Spec-aligned)

Z-Image is the model that made small cards feel like cheating - 6.15B parameters, happy at 12GB in BF16, still functional down around 4GB once you quantize. The catch: the community FP8 packages everyone links are a blind cast. Every layer gets clipped the same way, whether it's a fragile attention projection or a forgiving patch embedding. ZIT HSWQ FP8 Quantizer (Spec-aligned) is the calibrated alternative: it looks at how your model actually behaves, keeps the sensitive layers in FP16, and only clips the rest down to FP8.

"ZIT" is Z-Image Transformer - the diffusion transformer at the heart of NVIDIA's Z-Image (a NextDiT-style single-stream transformer). This node is one half of the ComfyUI-HSWQ-Quantizer pack, an unofficial ComfyUI port of Hybrid Sensitivity Weighted Quantization. It does for Z-Image what its sibling node does for SDXL.

How it works

The node reads a stats file produced by the pack's ZIT HSWQ Calibration (Dual Monitor V2) node, then walks the model's Linear and Conv2d layers:

1. Ranks layers by output variance (the sensitivity signal) and keeps the top keep_ratio of them in FP16 - normalizing BF16 → FP16 along the way, which some ComfyUI stacks handle better.
2. For the rest, optimizes each layer's FP8 clip value (amax) with a weighted histogram MSE search, weighting by per-channel input importance gathered during calibration. Layers with no stats, or that are already FP8, get skipped - safe to re-run.
3. Clips and casts to torch.float8_e4m3fn, optionally injecting comfy_quant / weight_scale buffers so ComfyUI loaders interpret the weights correctly.

Output is a MODEL that feeds straight into your sampler. Fair warning on scope: the conversion is in-memory - you don't get a reusable .safetensors out of it, just a lighter model for the current session.

The inputs that matter

• model - your Z-Image model (Base or Turbo, loaded via ComfyUI's Z-Image loader).
• stats_path - the .pt stats file, default output/zit_hswq_stats/zit_calib_session_01.pt. Resolved against ComfyUI's output directory if it isn't a valid path already.
• keep_ratio (0.25 default) - fraction of top-sensitivity layers kept in FP16. Lower saves more VRAM, higher protects quality; the original HSWQ guidance is ~0.25.
• bins / num_candidates / refinement_iterations - amax-search tuning. The defaults (8192 / 1000 / 10) are a fine starting point; the original repo suggests pushing refinement toward 20–25 for best quality.
• scaled (default off) - False is the compatible clip-then-cast path. True uses FP8's full dynamic range with a weight_scale buffer that only pays off if your loader honors it. Leave it off unless you're chasing spec-aligned V1 behavior.
• inject_comfy_metadata (default on) - writes the FP8 marker buffers downstream loaders expect. Keep it on.

Install

This pack is genuinely light - no model downloads, no requirements.txt. Just recent PyTorch (2.1+ for torch.float8_e4m3fn), plus lpips and open_clip_torch only if you want the pack's benchmark node.

cd ComfyUI/custom_nodes
git clone https://github.com/Shiba-2-shiba/ComfyUI-HSWQ-Quantizer

Restart ComfyUI. ComfyUI Manager works too - search for ComfyUI-HSWQ-Quantizer.

The workflow: calibrate first

Same trap as the SDXL node, and it's the pack's biggest gotcha: feed a default stats path and no stats file, and the model comes back unchanged (a printed [HSWQ] Error: Stats file not found is the only clue). HSWQ is a two-stage process and this is stage two.

1. Insert ZIT HSWQ Calibration (Dual Monitor V2) into your workflow ahead of the sampler and generate - a few sessions, roughly 256 samples total per the original guidance. Stats accumulate across runs with atomic saves.
2. Add this quantizer with the resulting stats path. One thing to know: stats are keyed to layer names, so a stats file only applies to the exact architecture you calibrated. Z-Image Base and Turbo aren't interchangeable - calibrate each separately.

Common issues

• Model passes through untouched - no stats file, or the path is wrong. Calibrate first, and double-check where the collector saved (the collector's save_folder_name / file_prefix should line up with this node's path).
• torch.float8_e4m3fn is not available - PyTorch older than 2.1; upgrade, or the pack refuses to run.
• No speed gain on 30-series - FP8 weights still shrink VRAM on older GPUs, but native fp8 compute is an RTX 40-series-and-newer thing. On a 3060/1050 Ti you're here for memory, not speed.

Import problems get logged to hswq_import_error.log inside the pack folder - check that file before assuming the worst.