ComfyUI-AsymQuantSampler

Stepwise Asymmetric Quantization for Flow Matching Models

Use a lightweight low-bit model (e.g., Q2_K) for early denoising steps and switch to a high-quality model (e.g., Q5_1) for later steps — with no quality loss. Our research shows the first 3 steps of FLUX.1-dev are robust to extreme quantization (2.5-bit), producing images that match or exceed 8-bit quality.

TL;DR: Early denoising steps don't need precision. Save compute by using a tiny model for steps 1–3, then switch to a quality model for the rest.

Key Findings

Our paper evaluates 2,000 images across 8 GGUF quantization levels on FLUX.1-dev and finds:

1. Steps 1–3 are quantization-robust. Replacing them with a 2.5-bit model (Q2_K) preserves full image quality (Quality Score 0.0308 vs Q8_0's 0.0306; ImageReward 1.291 vs 1.252).

2. Reference-based metrics are misleading for model switching. LPIPS and FID measure trajectory divergence, not quality degradation. Configs with FID=62.3 achieve quality matching 8-bit baselines.

3. CLIP Score is invariant to quantization. Semantic composition is preserved even at 2-bit precision (0.358–0.362 across all levels).

Configuration	Quality Score	ImageReward	LPIPS vs Q8_0	CLIP Score
Q8_0 (baseline)	0.0306	1.252	—	0.361
Asymmetric k=3 (Q2_K→Q5_1)	0.0308	1.291	0.233	0.360
Q5_1 (uniform)	0.0300	1.241	0.066	0.361
Q2_K (uniform)	0.0247	1.022	0.304	0.358

Nodes Included

This pack provides 5 custom nodes in two categories:

Asymmetric Sampling

Node	Description
Dual Model KSampler	Two-model step switching. Use model_a for steps 1–k, model_b for steps k+1–N. Supports 2-phase and 3-phase (sandwich) modes.
Dual Model KSampler Batch	Batch experiment node. Test multiple switch points (e.g., "3,5,7,10,15") in one run and output a LATENT batch.

Analysis Tools

Node	Description
Stepwise KSampler	Standard sampling with per-step latent saving. Saves intermediate .pt files for analysis.
Stepwise Latent Loader	Load a saved intermediate latent from a specific step.
Stepwise Latent Preview	Decode and display a grid of intermediate latents across steps using a VAE.

Installation

Option 1: ComfyUI Manager (Recommended)

Search for "Asymmetric Quantization Sampler" in ComfyUI Manager and click Install.

Option 2: Git Clone

cd ComfyUI/custom_nodes/
git clone https://github.com/lee09lee26/ComfyUI-AsymQuantSampler.git

No additional pip dependencies required — uses only ComfyUI built-in modules.

Option 3: Comfy CLI

comfy node install asym-quant-sampler

Usage

Basic 2-Phase Workflow (Recommended Starting Point)

1. Load two GGUF checkpoints of the same model at different quantization levels (e.g., flux1-dev-Q2_K.gguf and flux1-dev-Q5_1.gguf).
2. Connect both models to the Dual Model KSampler node.
3. Set switch_step = 3, steps = 20, cfg = 1.0.
4. Connect your prompt conditioning and an empty latent image as usual.
5. The node runs steps 1–3 with model_a (Q2_K), then steps 4–20 with model_b (Q5_1).

Node inputs:

Input	Type	Description
model_a	MODEL	Lightweight model for early steps (e.g., Q2_K)
model_b	MODEL	Quality model for later steps (e.g., Q5_1)
switch_step	INT	Step at which to switch from model_a to model_b
switch_end	INT (optional)	If > switch_step, enables 3-phase sandwich mode: A→B→A
steps	INT	Total sampling steps
seed, cfg, sampler_name, scheduler, denoise	—	Standard KSampler parameters

Outputs:

Output	Type	Description
latent	LATENT	Final denoised latent
info	STRING	Phase configuration description (e.g., "2-Phase: A(1-3) → B(4-20)")

3-Phase Sandwich Mode

Set switch_end to a value greater than switch_step. For example, switch_step=3, switch_end=17, steps=20 runs:

• Steps 1–3: model_a (Q2_K)
• Steps 4–17: model_b (Q5_1)
• Steps 18–20: model_a (Q2_K)

Batch Experimentation

Use Dual Model KSampler Batch to test multiple switch points at once. Set switch_steps_csv to something like "3,5,7,10,15" and the node will output a batched LATENT with one image per switch point. Connect to a PreviewImage or SaveImage node to compare results side by side.

Stepwise Analysis

To reproduce the step-by-step analysis from our paper:

1. Connect your model to Stepwise KSampler.
2. Set save_steps = "all" (or specific steps like "1,3,5,10,20").
3. Run the workflow. Intermediate latents are saved as .pt files.
4. Use Stepwise Latent Loader to load any saved step.
5. Use Stepwise Latent Preview with a VAE to generate a comparison grid.

Example Workflows

Pre-built workflow JSON files are in example_workflows/:

File	Description
dual_model_basic.json	Basic Q2_K → Q5_1 two-phase setup
dual_model_sandwich.json	Three-phase sandwich mode
dual_model_batch.json	Batch switch-point comparison
stepwise_analysis.json	Full step-by-step latent analysis

Evaluation Toolkit

The eval/ directory contains scripts to reproduce the paper's results:

# Generate benchmark images across all quantization levels
python eval/benchmark_generate.py --output_dir ./results

# Compute metrics (LPIPS, FID, CLIP Score, Quality Score, ImageReward)
python eval/benchmark_measure.py --input_dir ./results --output metrics.json

See eval/README.md for detailed instructions and required dependencies.

Recommended Configurations

Based on our experiments with FLUX.1-dev, 20-step Euler sampling, CFG 1.0:

Goal	model_a	model_b	switch_step	Notes
Best quality/efficiency	Q2_K	Q5_1	3	Matches Q8_0 quality
Safe conservative	Q4_0	Q5_1	5	Minimal risk
Maximum savings	Q2_K	Q4_1	3	Slight quality trade-off
Analysis/debugging	Any	Any	varies	Use Batch node to sweep

How It Works

Flow matching models like FLUX use straight ODE trajectories from noise to data. The key insight is that early denoising steps (high noise) establish coarse structure that is robust to quantization error, while later steps (low noise) refine fine details that require higher precision.

Step 1  ──────  Step 3  ──────────────────  Step 20
 ┌─────────────┐ ┌──────────────────────────────┐
 │  model_a     │ │        model_b                │
 │  (Q2_K)      │ │        (Q5_1)                 │
 │  Structure   │ │        Refinement             │
 │  formation   │ │        & detail               │
 └─────────────┘ └──────────────────────────────┘
  ← Robust to  →  ← Sensitive to quantization →
    quantization

When Q2_K handles early steps, the trajectory diverges from Q8_0's path but lands on an equally valid point in the learned distribution. This is why reference-based metrics (LPIPS, FID) show large distances while reference-free quality metrics (Quality Score, ImageReward) show no degradation.

Paper

📄 Stepwise Asymmetric Quantization for Flow Matching Models: Exploiting Non-Uniform Error Sensitivity Across Denoising Steps

If you find this work useful, please consider citing:

@article{lee09lee262026asymquant,
  title={Stepwise Asymmetric Quantization for Flow Matching Models: Exploiting Non-Uniform Error Sensitivity Across Denoising Steps},
  author={lee09lee26},
  year={2026},
  url={https://github.com/lee09lee26/ComfyUI-AsymQuantSampler}
}

Known Limitations

• Model switching overhead: Loading/unloading models adds ~7s on RTX 4070. On systems with >16GB VRAM where both models fit simultaneously, this overhead is minimal.
• Peak VRAM: Determined by the larger model regardless of switch point, since it must be loaded at step k.
• Tested on FLUX.1-dev only: The quantization-robust zone (steps 1–3) may differ for other flow matching models or step counts.
• GGUF format only: Other quantization formats (GPTQ, AWQ) have not been tested.

Changelog

v1.0.0

• Initial release
• Dual Model KSampler (2-phase and 3-phase sandwich)
• Dual Model KSampler Batch
• Stepwise KSampler, Latent Loader, Latent Preview
• Evaluation toolkit

Acknowledgments

Built on ComfyUI by comfyanonymous. GGUF quantization support via city96/gguf.

License

MIT License

##DOI

10.5281/zenodo.19496644