ComfyUI Extension: Asymmetric Quantization Sampler

Authored by lee09lee26

Created

Updated

1 stars

Run ComfyUI workflows without the setup

No installs, no CUDA version roulette, no GPU sitting idle on your bill. Bring a workflow and run it in the browser.

Stepwise asymmetric quantization for flow matching models. Use low-bit models for early denoising steps and high-bit models for later steps — with no quality loss.

README

ComfyUI-AsymQuantSampler

Stepwise Asymmetric Quantization for Flow Matching Models

Paper ComfyUI License: MIT

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.

Overview

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)

Workflow Screenshot

  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

Run ComfyUI workflows without the setup

No installs, no CUDA version roulette, no GPU sitting idle on your bill. Bring a workflow and run it in the browser.

Learn more