diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..8560b79 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,2 @@ +# Declare files that always have LF line endings on checkout +* text eol=lf diff --git a/.github/CODE_OF_CONDUCT.md b/.github/CODE_OF_CONDUCT.md new file mode 100644 index 0000000..00769f2 --- /dev/null +++ b/.github/CODE_OF_CONDUCT.md @@ -0,0 +1,9 @@ +# Code of Conduct + +Be openness, as well as friendly and didactic in discussions. + +Treat everybody equally, and value their contributions. + +Decisions are made based on technical merit and consensus. + +Try to follow most principles described here: https://nextcloud.com/code-of-conduct/ diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ee33183 --- /dev/null +++ b/.gitignore @@ -0,0 +1,150 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST +converted/ + +# 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/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# Pycharm settings +.idea/ + +# VSCode workspace settings +.vscode + +.DS_Store + +/out +/dev/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..e84f372 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,32 @@ +exclude: ^(screenshots)/ +repos: +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: check-yaml + - id: check-toml + - id: end-of-file-fixer + - id: trailing-whitespace + - id: mixed-line-ending + +- repo: https://github.com/PyCQA/isort + rev: 5.13.2 + hooks: + - id: isort + files: . + +- repo: https://github.com/psf/black + rev: 24.10.0 + hooks: + - id: black + files: . + +- repo: https://github.com/tox-dev/pyproject-fmt + rev: 2.3.1 + hooks: + - id: pyproject-fmt + +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.6.9 + hooks: + - id: ruff diff --git a/AUTHORS b/AUTHORS new file mode 100644 index 0000000..cbbe3d9 --- /dev/null +++ b/AUTHORS @@ -0,0 +1,14 @@ +Here is an inevitably incomplete list of MUCH-APPRECIATED CONTRIBUTORS -- +people who have submitted patches, reported bugs, added translations, helped +answer newbie questions, and contributed in other ways: + + Alexander Piskun + Andrey Borysenko + + +A big THANK YOU goes to: + + All ComfyUI community. + All Python community. + + Guido van Rossum for creating Python. diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..71e952e --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,9 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +## [1.0.0 - 2024-10-20] + +### Added + +- First release. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..39a8c6b --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..030e368 --- /dev/null +++ b/nodes.py @@ -0,0 +1,88 @@ +import logging + +import google.generativeai as genai +from torch import Tensor + +from .utils import images_to_pillow, temporary_env_var + + +class GeminiNode: + + @classmethod + def INPUT_TYPES(cls): # noqa + return { + "required": { + "prompt": ("STRING", {"default": "Why number 42 is important?", "multiline": True}), + "safety_settings": (["BLOCK_NONE", "BLOCK_ONLY_HIGH", "BLOCK_MEDIUM_AND_ABOVE"],), + "response_type": (["text", "json"],), + "model": (["gemini-1.5-flash-002", "gemini-1.5-pro-002"],), + }, + "optional": { + "api_key": ("STRING",), + "proxy": ("STRING",), + "image_1": ("IMAGE",), + "image_2": ("IMAGE",), + "image_3": ("IMAGE",), + "system_instruction": ("STRING",), + "error_fallback_value": ("STRING", {"lazy": True}), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("text",) + FUNCTION = "ask_gemini" + + CATEGORY = "Gemini" + + def __init__(self): + self.text_output: str | None = None + + def ask_gemini(self, **kwargs): + return (kwargs["error_fallback_value"] if self.text_output is None else self.text_output,) + + def check_lazy_status( + self, + prompt: str, + safety_settings: str, + response_type: str, + model: str, + api_key: str | None = None, + proxy: str | None = None, + image_1: Tensor | list[Tensor] | None = None, + image_2: Tensor | list[Tensor] | None = None, + image_3: Tensor | list[Tensor] | None = None, + system_instruction: str | None = None, + error_fallback_value: str | None = None, + ): + self.text_output = None + if not system_instruction: + system_instruction = None + images_to_send = [] + for image in [image_1, image_2, image_3]: + if image is not None: + images_to_send.extend(images_to_pillow(image)) + genai.configure(api_key=api_key, transport="rest") + model = genai.GenerativeModel(model, safety_settings=safety_settings, system_instruction=system_instruction) + generation_config = genai.GenerationConfig( + response_mime_type="application/json" if response_type == "json" else "text/plain" + ) + try: + with temporary_env_var("HTTP_PROXY", proxy), temporary_env_var("HTTPS_PROXY", proxy): + response = model.generate_content([prompt, *images_to_send], generation_config=generation_config) + self.text_output = response.text + except Exception: + if error_fallback_value is None: + logging.getLogger("ComfyUI-Gemini").debug("ComfyUI-Gemini: exception occurred:", exc_info=True) + return ["error_fallback_value"] + if error_fallback_value == "": + raise + return [] + + +NODE_CLASS_MAPPINGS = { + "Ask_Gemini": GeminiNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "Ask_Gemini": "Ask Gemini", +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..7921e1d --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,57 @@ +[project] +name = "comfyui-gemini" +version = "1.0.0" +description = "Nodes for Google Gemini API, focusing on backward compatibility and stability within ComfyUI." +license = "LICENSE" +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", +] +dependencies = [ "google-generativeai>=0.8.3" ] +urls.Changelog = "https://github.com/Visionatrix/ComfyUI-Gemini/blob/main/CHANGELOG.md" +urls.Repository = "https://github.com/Visionatrix/ComfyUI-Gemini" +urls.Source = "https://github.com/Visionatrix/ComfyUI-Gemini" + +[tool.black] +line-length = 120 +preview = true + +[tool.ruff] +target-version = "py310" +line-length = 120 +lint.select = [ + "A", + "B", + "C", + "E", + "F", + "G", + "I", + "PIE", + "Q", + "RET", + "RUF", + "S", + "SIM", + "UP", + "W", +] +lint.extend-ignore = [ + "I001", + "RUF100", + "S311", + "S603", +] +lint.mccabe.max-complexity = 20 + +[tool.isort] +profile = "black" + +[tool.comfy] +PublisherId = "" +DisplayName = "ComfyUI-Gemini" +Icon = "" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..498470a --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +torch +pillow +numpy +google-generativeai>=0.8.3 diff --git a/screenshots/race_gender.jpg b/screenshots/race_gender.jpg new file mode 100644 index 0000000..0038370 Binary files /dev/null and b/screenshots/race_gender.jpg differ diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..7219b07 --- /dev/null +++ b/utils.py @@ -0,0 +1,28 @@ +import os +from contextlib import contextmanager + +import numpy as np +from PIL import Image +from torch import Tensor + + +@contextmanager +def temporary_env_var(key: str, new_value: str | None): + old_value = os.environ.get(key) + if new_value: + os.environ[key] = new_value + try: + yield + finally: + if old_value is not None: + os.environ[key] = old_value + elif key in os.environ: + del os.environ[key] + + +def images_to_pillow(images: Tensor | list[Tensor]) -> list[Image]: + pillow_images = [] + for _bn, image in enumerate(images): + i = 255.0 * image.cpu().numpy() + pillow_images.append(Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))) + return pillow_images