ComfyUI-InferenceTimeScaling

ComfyUI-InferenceTimeScaling

Version

A ComfyUI extension implementing "Inference-time scaling for diffusion models beyond scaling denoising steps" (Ma et al., 2025). This extension provides inference-time optimization techniques to enhance diffusion-based image generation quality through random search and zero-order optimization algorithms, along with an ensemble verification system.

Athletic Cat

Prompt: "Photo of an athlete cat explaining it's latest scandal at a press conference to journalists."

Features

• Implementation of two search algorithms from the paper:
  • Random search optimization
  • Zero-order optimization
• Ensemble verification system using three verifiers:
  • CLIP Score verification
  • ImageReward verification
  • VLM grading using Qwen2.5-VL-7B
• Automated model downloading and management

Example Workflow

Simple example using FLUX-1 Dev with random search for 6 rounds

View example workflows

How It Works

This extension implements two different search algorithms to find the best possible image for your prompt:

1. Random Search: The simplest approach - generates multiple images with different random noises and evaluates them to explore the noise space.

2. Zero-Order Search: A more sophisticated approach that performs local optimization. It starts with a random noise, generates nearby variations by perturbing noise, and iteratively moves toward better results based on evaluation.

To explore the noise space, the quality of generated images is evaluated using an ensemble of three verifiers:

• CLIP Score: Measures how well the image matches the text prompt using OpenAI's CLIP model
• ImageReward: Evaluates image quality and prompt alignment using a specialized reward model
• Qwen VLM: Uses a large vision-language model to provide detailed scoring across multiple aspects (visual quality, creativity, prompt accuracy, etc.)

By exploring the noise space and using these verifiers to guide the search, it can produce images of higher quality and better prompt alignment than simply increasing denoising steps, with the tradeoff being increased time and compute during inference.

For more detailed information about the algorithms and methodology, please refer to the original paper from Google DeepMind: "Inference-time scaling for diffusion models beyond scaling denoising steps".

Installation

Prerequisites

1. Install ComfyUI
2. Install ComfyUI-Manager

Hardware Requirements

This was developed and tested on a system with:

• Single NVIDIA L40S GPU (48GB VRAM)
• 62GB System RAM

Installation Methods

Using ComfyUI-Manager (Recommended)

1. Open ComfyUI
2. Open the Manager panel
3. Search for "ComfyUI-InferenceTimeScaling"
4. Click "Install"
5. Restart ComfyUI

Manual Installation

cd ComfyUI/custom_nodes
git clone https://github.com/YRIKKA/ComfyUI-InferenceTimeScaling
cd ComfyUI-InferenceTimeScaling
pip install -e .

Usage

The extension adds the following nodes to ComfyUI:

InferenceTimeScaler Node

This is the main node implementing the random search and zero-order optimization algorithms from the paper.

InferenceTimeScaler Node

Inputs:

• model: (MODEL) Denoising model
• seed: (INT) Random seed for reproducibility
• steps: (INT) Number of denoising steps
• cfg: (FLOAT) Classifier-Free Guidance scale
• sampler_name: (SAMPLER) Sampling algorithm
• scheduler: (SCHEDULER) Noise scheduler
• positive: (CONDITIONING) Positive prompt conditioning
• negative: (CONDITIONING) Negative prompt conditioning
• latent_image: (LATENT) Input latent image
• denoise: (FLOAT) Amount of denoising to appl