From 03e232f3f401d743f39f4748c854cd384db6575f Mon Sep 17 00:00:00 2001 From: MaximClouser Date: Fri, 21 Feb 2025 21:29:40 +0000 Subject: [PATCH] Initial commit --- .editorconfig | 21 +++ .github/ISSUE_TEMPLATE.md | 15 ++ .github/workflows/build-pipeline.yml | 34 ++++ .github/workflows/publish_node.yml | 21 +++ .github/workflows/validate.yml | 13 ++ .gitignore | 100 ++++++++++++ .pre-commit-config.yaml | 10 ++ .vscode/settings.json | 8 + LICENSE | 22 +++ MANIFEST.in | 9 ++ README.md | 67 ++++++++ __init__.py | 16 ++ pyproject.toml | 79 ++++++++++ src/inferencescale/__init__.py | 0 src/inferencescale/nodes.py | 224 +++++++++++++++++++++++++++ tests/__init__.py | 1 + tests/conftest.py | 6 + tests/pytest.ini | 4 + tests/test_inferencescale.py | 21 +++ 19 files changed, 671 insertions(+) create mode 100644 .editorconfig create mode 100644 .github/ISSUE_TEMPLATE.md create mode 100644 .github/workflows/build-pipeline.yml create mode 100644 .github/workflows/publish_node.yml create mode 100644 .github/workflows/validate.yml create mode 100644 .gitignore create mode 100644 .pre-commit-config.yaml create mode 100644 .vscode/settings.json create mode 100644 LICENSE create mode 100644 MANIFEST.in create mode 100644 README.md create mode 100644 __init__.py create mode 100644 pyproject.toml create mode 100644 src/inferencescale/__init__.py create mode 100644 src/inferencescale/nodes.py create mode 100644 tests/__init__.py create mode 100644 tests/conftest.py create mode 100644 tests/pytest.ini create mode 100644 tests/test_inferencescale.py diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..d4a2c44 --- /dev/null +++ b/.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 diff --git a/.github/ISSUE_TEMPLATE.md b/.github/ISSUE_TEMPLATE.md new file mode 100644 index 0000000..d4832d9 --- /dev/null +++ b/.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. +``` diff --git a/.github/workflows/build-pipeline.yml b/.github/workflows/build-pipeline.yml new file mode 100644 index 0000000..2feee4a --- /dev/null +++ b/.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/ diff --git a/.github/workflows/publish_node.yml b/.github/workflows/publish_node.yml new file mode 100644 index 0000000..0118542 --- /dev/null +++ b/.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 }} diff --git a/.github/workflows/validate.yml b/.github/workflows/validate.yml new file mode 100644 index 0000000..8d3b5d4 --- /dev/null +++ b/.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 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d2d9a89 --- /dev/null +++ b/.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 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..2539924 --- /dev/null +++ b/.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 diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..50056f0 --- /dev/null +++ b/.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/" + ], +} diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..d9fc82a --- /dev/null +++ b/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. + diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..475df85 --- /dev/null +++ b/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 diff --git a/README.md b/README.md new file mode 100644 index 0000000..2a8fc06 --- /dev/null +++ b/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! + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..9e0d390 --- /dev/null +++ b/__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" diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..745d4a2 --- /dev/null +++ b/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" diff --git a/src/inferencescale/__init__.py b/src/inferencescale/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/inferencescale/nodes.py b/src/inferencescale/nodes.py new file mode 100644 index 0000000..757d1d3 --- /dev/null +++ b/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" +} + + diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..21de109 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Unit test package for inferencescale.""" diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..310609c --- /dev/null +++ b/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__), '..'))) diff --git a/tests/pytest.ini b/tests/pytest.ini new file mode 100644 index 0000000..95c76f1 --- /dev/null +++ b/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 diff --git a/tests/test_inferencescale.py b/tests/test_inferencescale.py new file mode 100644 index 0000000..b2cf230 --- /dev/null +++ b/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"