nodes/YRIKKA-ComfyUI-InferenceTimeScaling
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

readme

Browse Source
This commit is contained in:
MaximClouser
2025-02-26 16:50:56 +00:00
parent 2ac0f2c48b
commit 7068a87d4e
1 changed files with 121 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+121 -49
View File
@@ -1,24 +1,125 @@
# InferenceScale
# ComfyUI-InferenceScale
ComfyUI-InferenceTimeScaling
A ComfyUI extension implementing "Inference-time scaling for diffusion models beyond scaling denoising steps" (Ma et al., 2025). This extension provides advanced inference-time optimization techniques to improve image generation quality through random search and zero-order optimization algorithms, along with an ensemble verification system.
> [!NOTE]
> This projected was created with a [cookiecutter](https://github.com/Comfy-Org/cookiecutter-comfy-extension) template. It helps you start writing custom nodes without worrying about the Python setup.
## Features
## Quickstart
- 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
1. Install [ComfyUI](https://docs.comfy.org/get_started).
1. Install [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager)
1. Look up this extension in ComfyUI-Manager. If you are installing manually, clone this repository under `ComfyUI/custom_nodes`.
1. Restart ComfyUI.
## Installation
# Features
### Prerequisites
- A list of features
1. Install [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
2. Install [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager)
## Develop
### Installation Methods
To install the dev dependencies and pre-commit (will run the ruff hook), do:
#### Using ComfyUI-Manager (Recommended)
1. Open ComfyUI
2. Open the Manager panel
3. Search for "InferenceScale"
4. Click "Install"
5. Restart ComfyUI
#### Manual Installation
```bash
cd ComfyUI/custom_nodes
git clone https://github.com/your-username/ComfyUI-InferenceScale
cd ComfyUI-InferenceScale
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.
#### 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 apply
- `text_prompt_to_compare`: (STRING) Text prompt for verifiers
- `search_rounds`: (INT) Number of search rounds (random seeds for random search, or iterations for zero-order search)
- `vae`: (VAE) VAE model for decoding latents
- `view_top_k`: (INT) Number of top images to show in grid
- `search_algorithm`: Choice between "random" and "zero-order"
- `num_neighbors`: (INT) Number of neighbors per iteration in zero-order search (only used if search_algorithm is "zero-order")
- `lambda_threshold`: (FLOAT) Perturbation step size for zero-order search (only used if search_algorithm is "zero-order")
#### Optional Inputs:
- `loaded_clip_score_verifier`: (CS_VERIFIER) CLIP model for scoring
- `loaded_image_reward_verifier`: (IR_VERIFIER) ImageReward model
- `loaded_qwen_verifier`: (QWN_VERIFIER) Qwen VLM model
**Note:** At least one verifier must be included!
#### Outputs:
- `Best Image`: The highest-scoring generated image
- `Top-k Grid`: Grid view of the top-k scoring images
- `Top-k Score(s)`: JSON string containing detailed scores
### LoadQwenVLMVerifier Node
Loads the Qwen VLM verifier model for image evaluation.
#### Inputs:
- `qwen_verifier_id`: Model identifier (default: "Qwen/Qwen2.5-VL-7B-Instruct")
- `device`: Device to load model on ("cuda" or "cpu")
#### Outputs:
- `qwen_verifier_instance`: Loaded Qwen verifier instance
### LoadCLIPScoreVerifier Node
Loads the CLIP model for computing text-image similarity scores.
#### Inputs:
- `clip_verifier_id`: Model identifier (default: "openai/clip-vit-base-patch32")
- `device`: Device to load model on ("cuda" or "cpu")
#### Outputs:
- `clip_verifier_instance`: Loaded CLIP verifier instance
### LoadImageRewardVerifier Node
Loads the ImageReward model for image quality assessment.
#### Inputs:
- `ir_verifier_id`: Model identifier (default: "ImageReward-v1.0")
- `device`: Device to load model on ("cuda" or "cpu")
#### Outputs:
- `image_reward_verifier_instance`: Loaded ImageReward verifier instance
## Example Workflow
1. Load your model using the standard ComfyUI checkpoint loader
2. Load the verifier models using their respective loader nodes
3. Connect everything to the InferenceTimeScaler node
4. Set your desired search algorithm and parameters
5. Generate optimized images with improved quality
## Development
To install development dependencies:
```bash
cd inferencescale
@@ -26,42 +127,13 @@ pip install -e .[dev]
pre-commit install
```
The `-e` flag above will result in a "live" install, in the sense that any changes you make to your node extension will automatically be picked up the next time you run ComfyUI.
## License
## Publish to Github
[MIT License](LICENSE)
Install Github Desktop or follow these [instructions](https://docs.github.com/en/authentication/connecting-to-github-with-ssh/generating-a-new-ssh-key-and-adding-it-to-the-ssh-agent) for ssh.
1. Create a Github repository that matches the directory name.
2. Push the files to Git
```
git add .
git commit -m "project scaffolding"
git push
```
## Writing custom nodes
An example custom node is located in [node.py](src/inferencescale/nodes.py). To learn more, read the [docs](https://docs.comfy.org/essentials/custom_node_overview).
## Tests
This repo contains unit tests written in Pytest in the `tests/` directory. It is recommended to unit test your custom node.
- [build-pipeline.yml](.github/workflows/build-pipeline.yml) will run pytest and linter on any open PRs
- [validate.yml](.github/workflows/validate.yml) will run [node-diff](https://github.com/Comfy-Org/node-diff) to check for breaking changes
## Publishing to Registry
If you wish to share this custom node with others in the community, you can publish it to the registry. We've already auto-populated some fields in `pyproject.toml` under `tool.comfy`, but please double-check that they are correct.
You need to make an account on https://registry.comfy.org and create an API key token.
- [ ] Go to the [registry](https://registry.comfy.org). Login and create a publisher id (everything after the `@` sign on your registry profile).
- [ ] Add the publisher id into the pyproject.toml file.
- [ ] Create an api key on the Registry for publishing from Github. [Instructions](https://docs.comfy.org/registry/publishing#create-an-api-key-for-publishing).
- [ ] Add it to your Github Repository Secrets as `REGISTRY_ACCESS_TOKEN`.
A Github action will run on every git push. You can also run the Github action manually. Full instructions [here](https://docs.comfy.org/registry/publishing). Join our [discord](https://discord.com/invite/comfyorg) if you have any questions!
## Acknowledgments
- Thanks to the authors of the original paper for their research
- ComfyUI team for the excellent framework
- [tt-scale-flux](https://github.com/sayakpaul/tt-scale-flux) repository for the `generate_neighbors` function
- Qwen team for their powerful VLM model