MeanCache for Z-Image (ComfyUI)

February 4, 2026 · View on GitHub

Training-free inference acceleration for Z-Image Flow Matching models based on UnicomAI MeanCache. z_00128_

Features

  • Training-free: No model fine-tuning required
  • JVP-based velocity correction: Uses average velocity instead of instantaneous velocity for accurate ODE trajectory
  • PSSP scheduling: Peak-Suppressed Shortest Path algorithm for optimal compute budget allocation
  • Preset profiles: Quality / Balanced / Speed / Turbo presets for easy configuration
  • ~1.4x-2.0x speedup: Inference acceleration while maintaining image quality

Installation

Copy the comfyui-meancache-z folder to your ComfyUI custom_nodes directory.

Usage

  1. Load your Z-Image model
  2. Connect it to the MeanCache (Z-Image) node
  3. Select a preset or use Custom mode
  4. Connect the patched model output to your sampler image

Presets

PresetSpeedupDescription
Quality~1.25xConservative, minimal skipping, best quality
Balanced~1.7xGood speed/quality tradeoff (default)
Speed~1.75xAggressive skipping
Turbo~2.0xMaximum speed, may reduce quality
Custom-Manual parameter control

Preset Parameters

Presetrel_l1_threshskip_budgetstart_steppeak_thresholdgamma
Quality0.150.2030.083.0
Balanced0.300.4020.152.0
Speed0.500.5510.351.5
Turbo0.550.6010.451.0

Custom Mode Parameters

ParameterDefaultRangeDescription
rel_l1_thresh0.30.05-0.70Skip threshold (lower=quality, higher=speed)
skip_budget0.30.0-0.75Max fraction of steps to skip
start_step20-20Step to begin caching (protect early structure)
end_step-1-1 or 0+Step to end caching (-1=until end)
enable_psspTrue-Enable PSSP trajectory scheduling
peak_threshold0.150.05-0.60Max single-step velocity deviation
gamma2.00.5-3.0PSSP peak suppression exponent
cache_devicecpucpu/cudaDevice for velocity cache
debugFalse-Enable debug logging

Sampling Summary

When sampling completes, a summary is printed to console:

[MeanCache] Sampling complete (Balanced): 35 steps, 14 skipped, 21 computed (40.0% skip rate, ~1.67x speedup)

Algorithm

MeanCache improves Flow Matching inference by:

  1. Computing JVP (Jacobian-Vector Product) approximation via finite differences:

    JVP_{r→t} ≈ (v_t - v_r) / (t - r)
    
  2. Using average velocity instead of instantaneous velocity:

    û(z_t, t, s) = v(z_t, t) + (s - t) · JVP_{r→t}
    
  3. Stability deviation metric (L_K) for adaptive skip decision:

    L_K = ||v_current - (v_prev + dt · JVP)|| / ||v_current||
    
  4. PSSP scheduling with dynamic programming:

    π* = argmin Σ C(e)^γ   s.t. |π| ≤ B
    
  5. Intelligently skipping steps when velocity is stable, using cached JVP-corrected velocity

File Structure

comfyui-meancache-z/
├── __init__.py              # Plugin entry point (V2/V3 compatible)
├── nodes/
│   └── meancache_node.py    # MeanCache_ZImage node definition
├── patch/
│   └── model_patch.py       # Model wrapper with MeanCache logic
├── core/
│   ├── meancache_state.py   # State management per prediction
│   ├── velocity_cache.py    # JVP computation utilities
│   └── trajectory_scheduler.py  # PSSP scheduling algorithm
├── web/js/
│   └── meancache_preset.js  # Frontend preset widget sync
└── README.md

References

  • MeanCache Paper (UnicomAI)
  • "From Instantaneous to Average Velocity for Accelerating Flow Matching Inference"

License

MIT License