Initial commit

.editorconfig

@@ -0,0 +1,21 @@
+# http://editorconfig.org
+
+root = true
+
+[*]
+indent_style = space
+indent_size = 4
+trim_trailing_whitespace = true
+insert_final_newline = true
+charset = utf-8
+end_of_line = lf
+
+[*.bat]
+indent_style = tab
+end_of_line = crlf
+
+[LICENSE]
+insert_final_newline = false
+
+[Makefile]
+indent_style = tab

.github/ISSUE_TEMPLATE.md

@@ -0,0 +1,15 @@
+* inferencescale version:
+* Python version:
+* Operating System:
+
+### Description
+
+Describe what you were trying to get done.
+Tell us what happened, what went wrong, and what you expected to happen.
+
+### What I Did
+
+```
+Paste the command(s) you ran and the output.
+If there was a crash, please include the traceback here.
+```

.github/workflows/build-pipeline.yml

@@ -0,0 +1,34 @@
+# GitHub CI build pipeline
+name: inferencescale CI build
+
+on:
+  pull_request:
+    branches:
+      - master
+      - main
+jobs:
+  build:
+    runs-on:  ${{ matrix.os }} 
+    env:
+      PYTHONIOENCODING: "utf8"
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        python-version: ["3.12"]
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version:  ${{ matrix.python-version }} 
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install .[dev]
+      - name: Run Linting
+        run: |
+          ruff check .
+      - name: Run Tests
+        run: |
+          pytest tests/

.github/workflows/publish_node.yml

@@ -0,0 +1,21 @@
+name: 📦 Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    tags:
+      - '*'
+
+permissions:
+  issues: write
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    steps:
+      - name: ♻️ Check out code
+        uses: actions/checkout@v4
+      - name: 📦 Publish Custom Node
+        uses: Comfy-Org/publish-node-action@main
+        with:
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.github/workflows/validate.yml

@@ -0,0 +1,13 @@
+name: Validate backwards compatibility
+
+on:
+  pull_request:
+    branches:
+      - master
+      - main
+
+jobs:
+  validate:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: comfy-org/node-diff@main

.gitignore

@@ -0,0 +1,100 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# OSX useful to ignore
+*.DS_Store
+.AppleDouble
+.LSOverride
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+venv/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+
+# Sphinx documentation
+docs/_build/
+
+# IntelliJ Idea
+.idea
+*.iml
+*.ipr
+*.iws
+
+# PyBuilder
+target/
+
+# Cookiecutter
+output/
+python_boilerplate/
+cookiecutter-pypackage-env/
+
+# vscode settings
+.history/
+*.code-workspace

.pre-commit-config.yaml

@@ -0,0 +1,10 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.4.9
+    hooks:
+      # Run the linter.
+      - id: ruff
+        args: [ --fix ]
+      # Run the formatter.
+      - id: ruff-format

.vscode/settings.json

@@ -0,0 +1,8 @@
+{
+  // Required - change /PATH/TO to the absolute path to ComfyUI.  Windows e.g.: D:/My Folder/ComfyUI/
+  // This pulls in ComfyUI Python types for the extension.
+  "python.analysis.extraPaths": [
+    "/PATH/TO/ComfyUI/",
+    "/PATH/TO/ComfyUI/custom_nodes/"
+  ],
+}

LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2025, Max Clouser
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+

MANIFEST.in

@@ -0,0 +1,9 @@
+include LICENSE
+include README.md
+
+recursive-exclude * __pycache__
+recursive-exclude * *.py[co]
+
+recursive-include docs *.rst conf.py Makefile make.bat *.jpg *.png *.gif
+
+graft src/inferencescale/web

README.md

@@ -0,0 +1,67 @@
+# InferenceScale
+
+inference diffusion scaling
+
+> [!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.
+
+## Quickstart
+
+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.
+
+# Features
+
+- A list of features
+
+## Develop
+
+To install the dev dependencies and pre-commit (will run the ruff hook), do:
+
+```bash
+cd inferencescale
+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.
+
+## Publish to Github
+
+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!
+

__init__.py

@@ -0,0 +1,16 @@
+"""Top-level package for inferencescale."""
+
+__all__ = [
+    "NODE_CLASS_MAPPINGS",
+    "NODE_DISPLAY_NAME_MAPPINGS",
+    "WEB_DIRECTORY",
+]
+
+__author__ = """Max Clouser"""
+__email__ = "maximclouser@gmail.com"
+__version__ = "0.0.1"
+
+from .src.inferencescale.nodes import NODE_CLASS_MAPPINGS
+from .src.inferencescale.nodes import NODE_DISPLAY_NAME_MAPPINGS
+
+WEB_DIRECTORY = "./web"

pyproject.toml

@@ -0,0 +1,79 @@
+[build-system]
+requires = ["setuptools>=70.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "inferencescale"
+version = "0.0.1"
+description = "inference diffusion scaling"
+authors = [
+  {name = "Max Clouser", email = "maximclouser@gmail.com"}
+]
+readme = "README.md"
+license = {text = "MIT license"}
+classifiers = []
+dependencies = [
+
+]
+
+[project.optional-dependencies]
+dev = [
+    "bump-my-version",
+    "coverage",  # testing
+    "mypy",  # linting
+    "pre-commit", # runs linting on commit
+    "pytest",  # testing
+    "ruff",  # linting
+]
+
+[project.urls]
+bugs = "https://github.com/MaximClouser/inferencescale/issues"
+homepage = "https://github.com/MaximClouser/inferencescale"
+
+
+[tool.comfy]
+PublisherId = "MaximClouser"
+DisplayName = "InferenceScale"
+Icon = ""
+
+[tool.setuptools.package-data]
+"*" = ["*.*"]
+
+[tool.pytest.ini_options]
+minversion = "8.0"
+testpaths = [
+    "tests",
+]
+
+[tool.mypy]
+files = "."
+
+# Use strict defaults
+strict = true
+warn_unreachable = true
+warn_no_return = true
+
+[[tool.mypy.overrides]]
+# Don't require test functions to include types
+module = "tests.*"
+allow_untyped_defs = true
+disable_error_code = "attr-defined"
+
+[tool.ruff]
+# extend-exclude = ["static", "ci/templates"]
+line-length = 140
+src = ["src", "tests"]
+target-version = "py39"
+
+# Add rules to ban exec/eval
+[tool.ruff.lint]
+select = [
+    "S102",  # exec-builtin
+    "S307",  # eval-used
+    "W293",
+    "F",  # The "F" series in Ruff stands for "Pyflakes" rules, which catch various Python syntax errors and undefined names.
+    # See all rules here: https://docs.astral.sh/ruff/rules/#pyflakes-f
+]
+
+[tool.ruff.lint.flake8-quotes]
+inline-quotes = "double"

src/inferencescale/nodes.py

@@ -0,0 +1,224 @@
+import torch
+import torch.nn.functional as F
+import comfy.samplers
+import comfy.sample
+import comfy.sd
+import comfy.utils
+import latent_preview
+import json
+from torchvision.utils import make_grid
+from torchvision.transforms import ToPILImage
+
+from transformers import CLIPProcessor, CLIPModel
+import ImageReward as reward
+
+
+def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
+    latent_image = latent["samples"]
+    latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image)
+
+    if disable_noise:
+        noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
+    else:
+        batch_inds = latent["batch_index"] if "batch_index" in latent else None
+        noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
+
+    noise_mask = None
+    if "noise_mask" in latent:
+        noise_mask = latent["noise_mask"]
+
+    callback = latent_preview.prepare_callback(model, steps)
+    disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
+    samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
+                                  denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step,
+                                  force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
+    out = latent.copy()
+    out["samples"] = samples
+    return (out, )
+
+
+def compute_similarity(clip_model, clip_processor, text, image) -> float:
+    # Process inputs: 'text' is a string, and 'image' should be a PIL image
+    inputs = clip_processor(text=[text], images=image, return_tensors="pt", padding=True)
+    
+    device = next(clip_model.parameters()).device
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    
+    # Get text and image embeddings.
+    with torch.no_grad():
+        text_embeddings = clip_model.get_text_features(
+            input_ids=inputs['input_ids'],
+            attention_mask=inputs['attention_mask']
+        )
+        image_embeddings = clip_model.get_image_features(
+            pixel_values=inputs['pixel_values']
+        )
+    
+    text_embeddings = F.normalize(text_embeddings, dim=-1)
+    image_embeddings = F.normalize(image_embeddings, dim=-1)
+    
+    similarity = F.cosine_similarity(text_embeddings, image_embeddings)
+    # rounded_similarity = round(similarity.item(), 4)
+    rounded_similarity = similarity.item()
+    return rounded_similarity
+
+
+def compute_combined_verifier_score(candidate_image, text_prompt, verifiers, weights):
+    """
+    verifiers: a dict with keys like 'clip' and 'image_reward' and values holding the corresponding model(s)/processor(s)
+    weights: a dict with the same keys, providing a float weight for each verifier
+    Returns a combined score and a dict of individual scores.
+    """
+    scores = {}
+
+    if 'clip' in verifiers and verifiers['clip'] is not None:
+        clip_model, clip_processor = verifiers['clip']
+        clip_score = compute_similarity(clip_model, clip_processor, text_prompt, candidate_image)
+        scores['clip'] = clip_score
+
+    if 'image_reward' in verifiers and verifiers['image_reward'] is not None:
+        image_reward_model = verifiers['image_reward']
+        image_reward_score = image_reward_model.score(text_prompt, candidate_image)
+        scores['image_reward'] = image_reward_score
+
+    combined_score = 0.0
+    for key, score in scores.items():
+        # TODO Implement algo from paper to combine scores
+        combined_score += weights.get(key, 1.0) * score
+        combined_score = round(combined_score, 4)
+    return combined_score, scores
+
+
+class MaximKSampler:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed used for creating the noise."}),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}),
+                "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt; however, too high values may negatively impact quality."}),
+                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The sampling algorithm used. This can affect the quality, speed, and style of the generated output."}),
+                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}),
+                "positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes to include in the image."}),
+                "negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes to exclude from the image."}),
+                "latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}),
+                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The amount of denoising applied; lower values maintain more of the structure of the initial image."}),
+                "text_prompt_to_compare": ("STRING", {"default": "", "multiline": True, "tooltip": "A text prompt used for comparing against generated images."}),
+                "search_budget": ("INT", {"default": 5, "min": 1, "max": 10000, "tooltip": "The number of random noise samples to generate (i.e., the search budget)."}),
+                "vae": ("VAE", {"tooltip": "The VAE model used for decoding the latent image."}),
+                "top_k": ("INT", {"default": 3, "min": 1, "max": 100, "tooltip": "The top k image/score pairs to return."}),
+            },
+            "optional": {
+                "clip_score_verifier": ("STRING", {"default": "openai/clip-vit-base-patch32", "tooltip": "Optional Hugging Face CLIP model identifier used for verifying the generated image quality."}),
+                "clip_score_weight": ("FLOAT", {"default": 1, "min": 0, "max": 1, "tooltip": "The weight of the clip verifier to apply "}),
+                "image_reward_verifier": ("STRING", {"default": "ImageReward-v1.0", "tooltip": "Optional Image Reward model used for verifying the generated image quality."}),
+                "image_reward_weight": ("FLOAT", {"default": 1, "min": 0, "max": 1, "tooltip": "The weight of the iamge rewqrd to apply "})
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE", "IMAGE", "STRING")
+    RETURN_NAMES = ("Best Image", "Top-k Grid", "Top-k Score(s)")
+
+    OUTPUT_TOOLTIPS = ("A grid of the top k images from the random search, and a JSON string with their verifier scores.",)
+    FUNCTION = "execute"
+
+    CATEGORY = "maxtest/sampling"
+    DESCRIPTION = ("Uses the provided model, conditioning, and VAE to denoise the latent image. "
+                   "Performs random search over noise candidates by sampling different noises, decodes each candidate image, "
+                   "evaluates its quality using a verifier (if provided), and returns the top k candidates based on the verifier score.")
+
+
+    def execute(self, 
+                model, 
+                vae, 
+                seed, 
+                steps, 
+                cfg, 
+                sampler_name, 
+                scheduler, 
+                positive, 
+                negative, 
+                latent_image, 
+                denoise, 
+                text_prompt_to_compare, 
+                search_budget,
+                top_k,
+                clip_score_verifier=None,
+                image_reward_verifier=None,
+                clip_score_weight=1.0,
+                image_reward_weight=1.0
+    ):
+        # TODO Move to innit
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        candidate_results = []
+
+        verifiers = {}
+        weights = {}
+        
+        if clip_score_verifier:
+            clip_model = CLIPModel.from_pretrained(clip_score_verifier)
+            clip_processor = CLIPProcessor.from_pretrained(clip_score_verifier)
+            clip_model.to(device)
+            verifiers['clip'] = (clip_model, clip_processor)
+            weights['clip'] = clip_score_weight
+
+        if image_reward_verifier:
+            image_reward_model = reward.load(image_reward_verifier)
+            verifiers['image_reward'] = image_reward_model
+            weights['image_reward'] = image_reward_weight
+
+        for i in range(search_budget):
+            new_seed = seed + i + 1
+            samples, = common_ksampler(model, new_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
+            
+            candidate_images = vae.decode(samples["samples"])
+            if len(candidate_images.shape) == 5:
+                candidate_images = candidate_images.reshape(-1, candidate_images.shape[-3], candidate_images.shape[-2], candidate_images.shape[-1])
+            
+            candidate_image = candidate_images[0:1]  # Using the first image of the batch.
+
+            # Convert from tensor to PIL Image
+            # Remove the batch dimension: [1, H, W, 3] -> [H, W, 3]
+            image_tensor = candidate_image.squeeze(0)
+            # Permute to [C, H, W]
+            image_tensor = image_tensor.permute(2, 0, 1)
+            # Convert to PIL Image
+            candidate_pil_image = ToPILImage()(image_tensor)
+
+            # Score the candidate
+            combined_score, individual_scores = compute_combined_verifier_score(candidate_pil_image, text_prompt_to_compare, verifiers, weights)
+            
+            candidate_results.append((candidate_image, combined_score))
+        
+        candidate_results.sort(key=lambda x: x[1], reverse=True)
+        top_candidates = candidate_results[:min(top_k, len(candidate_results))]
+        
+        images_list = [img for (img, score) in top_candidates]
+        scores_list = [score for (img, score) in top_candidates]
+        
+        # Prepare the images for grid display.
+        images_for_grid = [img.squeeze(0).permute(2, 0, 1) for img in images_list]
+        grid_image = make_grid(torch.stack(images_for_grid), nrow=top_k)
+        grid_image = grid_image.permute(1, 2, 0).unsqueeze(0)
+        
+        scores_json = json.dumps({"results": [{"score": score} for score in scores_list]})
+        best_candidate = images_list[0]
+        return (best_candidate, grid_image, scores_json)
+
+
+
+
+
+
+# Node registration
+NODE_CLASS_MAPPINGS = {
+    "MaximKSampler": MaximKSampler
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "MaximKSampler": "MaximKSampler"
+}
+
+

tests/__init__.py

@@ -0,0 +1 @@
+"""Unit test package for inferencescale."""

tests/conftest.py

@@ -0,0 +1,6 @@
+import os
+import sys
+
+# Add the project root directory to Python path
+# This allows the tests to import the project
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

tests/pytest.ini

@@ -0,0 +1,4 @@
+[pytest]
+testpaths = . # Run tests in the current directory
+python_files = test_*.py # Run tests in files that start with "test_"
+norecursedirs = .. # Don't run tests in the parent directory

tests/test_inferencescale.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python
+
+"""Tests for `inferencescale` package."""
+
+import pytest
+from src.inferencescale.nodes import Example
+
+@pytest.fixture
+def example_node():
+    """Fixture to create an Example node instance."""
+    return Example()
+
+def test_example_node_initialization(example_node):
+    """Test that the node can be instantiated."""
+    assert isinstance(example_node, Example)
+
+def test_return_types():
+    """Test the node's metadata."""
+    assert Example.RETURN_TYPES == ("IMAGE",)
+    assert Example.FUNCTION == "test"
+    assert Example.CATEGORY == "Example"