Nodes/ComfyUI-HSWQ-Quantizer/ZIT HSWQ Calibration (DualMonitor V2)
ComfyUI Node

ZIT HSWQ Calibration (DualMonitor V2)

The same smart-quantization trick, now for Z-Image

By Shiba-2-shiba·Created 7 months ago·Updated 5 months ago· 3
ZIT HSWQ Calibration (DualMonitor V2)
  • model
  • model
save_folder_namezit_hswq_stats
file_prefixzit_calib
session_idsession_01
target_layerall_linear_conv
save_every_steps50
reset_sessionfalse

Z-Image (its architecture goes by "ZIT" - Z-Image Transformer - and the pack's aliases also say "NextDiT") is the 6B Alibaba model that quietly became the community's default successor to SDXL: near-Flux quality that runs on a 12GB card, Apache 2.0. If you're on that bandwagon and your card is groaning, the HSWQ route is the same one as for SDXL - calibrate, then convert to fp8 keeping the sensitive layers in fp16 - and this is the calibration half of it.

If you've read the SDXL calibration node article, most of this is familiar: same DualMonitor V2 machinery, same session accumulation, same atomic .pt saves. The ZIT version exists because a 6B transformer is not an SDXL UNet - its layer structure is different, and the stats collector needs to know what to hook. That's exactly what the one extra input here is for.

How it works

The collector clones your model, registers forward hooks, and per-layer tracks output variance (sensitivity) and input-channel importance in FP32-accumulated precision - the same two signals the pack's quantizer uses to rank layers and pick per-layer amax clipping points via weighted histogram MSE. Stats accumulate across runs into a session (keyed by folder/prefix/session name), restore from disk, and checkpoint atomically every save_every_steps so a crash can't corrupt the file. Nothing about your generation changes; you just generate and it watches.

The difference is target_layer, which filters which module families get hooked at all:

  • all_linear_conv (default) - hook every Linear and Conv2d, the "just collect everything" option.
  • attention_only - just attention/QKV layers, if you believe attention is where the sensitivity lives.
  • feed_forward_only - just the FFN blocks.
  • context_refiner - only the context-refiner component, the parts that handle conditioning.

Start with all_linear_conv; the narrower options are for when you're confident about the architecture or want to keep the stats file small.

Inputs and outputs

model is your Z-Image model, passed straight through to your KSampler. The naming trio is save_folder_name (default zit_hswq_stats), file_prefix (default zit_calib), and session_id (default session_01) - they build the output path output/<save_folder_name>/<file_prefix>_<session_id>.pt. save_every_steps (default 50) controls checkpoint frequency. And the footgun is the same one as its SDXL sibling: reset_session wipes accumulated stats every run when true, so keep it false unless you're deliberately recalibrating.

The single model output goes to your KSampler; that's the whole wiring.

Workflow and gotchas

Installation is the same pack, once: ComfyUI Manager (search "HSWQ") or

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

then restart. The pack needs a recent PyTorch with torch.float8_e4m3fn and current-mainline ComfyUI (it uses the newer node API).

Then run your normal generations with the node in the graph - the README's guidance of ~256 samples applies here too - and confirm the console shows the hooks armed ([ZITCollector] Armed N hooks) or you're collecting nothing. When you're done, the natural next step is the pack's ZIT HSWQ FP8 quantizer node, which reads these stats and does the actual conversion; make sure the folder/prefix/session names match what you calibrated under, because a missing stats file makes the quantizer fail silently and hand your model back unconverted. Calibrate on a spread of prompts, not one style - otherwise the sensitivity ranking is only "correct" for the images you fed it.

CategoryZIT/Quantization

Inputs (7)

NameTypeDefaultDescription
modelMODEL
save_folder_nameSTRINGzit_hswq_stats
file_prefixSTRINGzit_calib
session_idSTRINGsession_01
target_layerCOMBOall_linear_conv4 options: all_linear_conv, attention_only, feed_forward_only, context_refiner
save_every_stepsINT501–10000
reset_sessionBOOLEANfalse

Outputs (1)

NameTypeDescription
modelMODEL