commit fa9601b9263dcf2ab9aa83b050bb4568bbd8bf37 Author: Pfaeff Date: Wed Aug 2 20:49:03 2023 +0200 tmp diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..68bc17f --- /dev/null +++ b/.gitignore @@ -0,0 +1,160 @@ +# 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 + +# 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 + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__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 +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..8c68db9 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "pixeldetector"] + path = pixeldetector + url = https://github.com/Astropulse/pixeldetector diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..927fe62 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 Pfaeff + +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/README.md b/README.md new file mode 100644 index 0000000..9af96b4 --- /dev/null +++ b/README.md @@ -0,0 +1,21 @@ +# pfaeff-comfyui + +Work in progress + +## Installation (not tested, yet) + +1. `git clone https://github.com/Pfaeff/pfaeff-comfyui.git` in `ComfyUI/custom_nodes/` +2. Restart ComfyUI, installation of dependencies should happen automatically + +## SDXL pixel art example workflows + +Simply drag one of the images into ComfyUI after installation. + +### SDXL Pixelart workflow with background removal + +![SDXL Pixelart workflow](workflows/SDXL_pixelart.png) + + +### Aspect ratio preserving SDXL Pixelart workflow + +![Aspect ratio preserving SDXL Pixelart workflow](workflows/SDXL_pixelart_preserve_aspect_ratio.png) \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..610069c --- /dev/null +++ b/__init__.py @@ -0,0 +1,17 @@ +import os +import importlib + +def install_dependencies(): + spec = importlib.util.spec_from_file_location('pfaeff_install', os.path.join(os.path.dirname(__file__), 'install.py')) + pfaeff_install = importlib.util.module_from_spec(spec) + spec.loader.exec_module(pfaeff_install) + +try: + import diffusers + import cv2 +except: + install_dependencies() + +from .pfaeff import NODE_CLASS_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS'] \ No newline at end of file diff --git a/install.py b/install.py new file mode 100644 index 0000000..b9e38c5 --- /dev/null +++ b/install.py @@ -0,0 +1,19 @@ +import sys +import os +import subprocess + +if sys.argv[0] == 'install.py': + sys.path.append('.') + +pfaeff_root_path = os.path.join(os.path.dirname(__file__)) + +if "python_embed" in sys.executable: + pip_install = [sys.executable, '-s', '-m', 'pip', 'install'] +else: + pip_install = [sys.executable, '-m', 'pip', 'install'] + +requirements_txt = os.path.join(pfaeff_root_path, "requirements.txt") +if os.path.exists(requirements_txt): + subprocess.run(pip_install + ['-r', 'requirements.txt'], cwd=pfaeff_root_path) + +subprocess.run(['git', 'submodule', 'init', '--init', '--recursive'], cwd=pfaeff_root_path) \ No newline at end of file diff --git a/pfaeff.py b/pfaeff.py new file mode 100644 index 0000000..291a7a4 --- /dev/null +++ b/pfaeff.py @@ -0,0 +1,279 @@ +import os +import torch +import nodes +from nodes import MAX_RESOLUTION +import cv2 +import numpy as np +from PIL import Image +import tempfile +from collections import Counter + +from diffusers import StableDiffusionInpaintPipeline +from diffusers.models import AutoencoderKL +import subprocess + + +class AstropulsePixelDetector: + @classmethod + def INPUT_TYPES(s): + return {"required": { "image": ("IMAGE", ), + "max_colors": ("INT", {"default": 128}) + }} + RETURN_TYPES = ("IMAGE",) + FUNCTION = "run" + + CATEGORY = "Pfaeff/image" + + def __init__(self): + current_script_dir = os.path.dirname(os.path.realpath(__file__)) + self.pixeldetector_path = os.path.join(current_script_dir, 'pixeldetector', 'pixeldetector.py') + + + def run(self, image, max_colors): + + if image.shape[0] > 1: + raise Exception("Batches are not supported since output images can have different resolutions!") + + numpy_image = (image[0, :, :, :].cpu().detach().numpy() * 255.0).astype(np.uint8) + + # Create a temporary file for the input image + temp_input_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False) + temp_input_path = temp_input_file.name + Image.fromarray(numpy_image).save(temp_input_path) + temp_input_file.close() + + # Create a temporary file for the output image + temp_output_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False) + temp_output_path = temp_output_file.name + temp_output_file.close() + + # Construct the command for calling the pixeldetector script + command = [ + "python", self.pixeldetector_path, + "--input", temp_input_path, + "--output", temp_output_path + ] + if max_colors: + command.extend(["--max", str(max_colors)]) + + command.extend(["--palette"]) # TODO make this a parameter + + # Run the command + subprocess.run(command) + + # Read the processed image + processed_image = Image.open(temp_output_path) + processed_image = np.array(processed_image) + # Close and delete the temporary files + os.unlink(temp_input_path) + os.unlink(temp_output_path) + + # Convert back to torch + image_result = torch.from_numpy(processed_image).float() / 255.0 + image_result = image_result[None, :, :, :] + image_result = image_result.to(image.device) + + return (image_result, ) + + +class BackgroundRemover: + @classmethod + def INPUT_TYPES(s): + return {"required": { "image": ("IMAGE", ), + }} + RETURN_TYPES = ("IMAGE",) + FUNCTION = "run" + + CATEGORY = "Pfaeff/image" + + def detect_most_common_color(self, image): + # Reshape the image to be a list of pixels + pixels = image.reshape(-1, 3) + + # Convert pixels to tuples so they can be hashed + pixel_tuples = [tuple(pixel) for pixel in pixels] + + # Find the most common color using a Counter + most_common_color = Counter(pixel_tuples).most_common(1)[0][0] + + return np.array(most_common_color) + + + def run(self, image): + batch_size, height, width, channels = image.size() + + if channels != 3: + raise Exception("Input must be a 3-channel image") + + image_result = torch.zeros((batch_size, height, width, 4), dtype=torch.float32) + + for b in range(image.shape[0]): + numpy_image = (image[b, :, :, :].cpu().detach().numpy() * 255.0).astype(np.uint8) + + # Find the most common color in the image + most_common_color = self.detect_most_common_color(numpy_image) + + # Create an output image with an additional alpha channel (shape will be (height, width, 4)) + output_image = np.zeros((numpy_image.shape[0], numpy_image.shape[1], 4)) + + # Iterate through the image and set the RGB channels to match the original image + # If the pixel matches the most common color, set the alpha channel to 0; otherwise, set it to 255 + for i in range(numpy_image.shape[0]): + for j in range(numpy_image.shape[1]): + if np.array_equal(numpy_image[i, j], most_common_color): + output_image[i, j, :3] = numpy_image[i, j] + output_image[i, j, 3] = 0 + else: + output_image[i, j, :3] = numpy_image[i, j] + output_image[i, j, 3] = 255 + + image_result[b, :, :, :] = torch.from_numpy(output_image).float() / 255.0 + image_result = image_result.to(image.device) + + return (image_result, ) + + +class InpaintingPipelineLoader: + @classmethod + def INPUT_TYPES(s): + return {"required": { "model_name": ("STRING", {"default": "stabilityai/stable-diffusion-2-inpainting"}), + "vae_name": ("STRING", {"default": "stabilityai/sd-vae-ft-mse"}), + }} + RETURN_TYPES = ("INPAINT_PIPELINE",) + FUNCTION = "load_inpainting_pipeline" + + CATEGORY = "Pfaeff/loaders" + + def load_inpainting_pipeline(self, model_name, vae_name): + vae = AutoencoderKL.from_pretrained(vae_name) + inpaint = StableDiffusionInpaintPipeline.from_pretrained( + model_name, vae=vae + ) + + if torch.cuda.is_available(): + inpaint.to("cuda") + + return (inpaint,) + + +class Inpainting: + @classmethod + def INPUT_TYPES(s): + return {"required": { "inpainting_pipeline": ("INPAINT_PIPELINE", ), + "image": ("IMAGE", ), + "mask": ("MASK", ), + "text": ("STRING", {"multiline": True}), + "num_inference_steps": ("INT", {"default": 20}), + "guidance_scale": ("FLOAT", {"default": 8.0}), + }} + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "inpaint" + + CATEGORY = "Pfaeff/inpainting" + + def inpaint(self, inpainting_pipeline, image, mask, text, num_inference_steps, guidance_scale): + + extra_kwargs = { + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + } + + image_result = torch.zeros_like(image) + + for i in range(image.shape[0]): + + numpy_image = (image[i, :, :, :].cpu().detach().numpy() * 255.0).astype(np.uint8) + init_image = Image.fromarray(numpy_image) + + numpy_mask = (mask.cpu().detach().numpy() * 255.0).astype(np.uint8) + mask_image = Image.fromarray(numpy_mask) + + inpainted = inpainting_pipeline( + prompt=text, + image=init_image, + mask_image=mask_image, + width=image.shape[2], + height=image.shape[1], + **extra_kwargs, + )["images"] + + image_result[i, :, :, :] = torch.from_numpy(np.array(inpainted[0])).float() / 255.0 + + image_result = image_result.to(image.device) + + return (image_result,) + + +class ImagePadForBetterOutpaint: + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "left": ("INT", {"default": 256, "min": 0, "max": MAX_RESOLUTION, "step": 8}), + "top": ("INT", {"default": 256, "min": 0, "max": MAX_RESOLUTION, "step": 8}), + "right": ("INT", {"default": 256, "min": 0, "max": MAX_RESOLUTION, "step": 8}), + "bottom": ("INT", {"default": 256, "min": 0, "max": MAX_RESOLUTION, "step": 8}), + "inpaint_radius": ("INT", {"default": 5, "min": 3, "max": 128, "step": 8}), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + FUNCTION = "expand_image" + + CATEGORY = "Pfaeff/outpainting" + + def add_padding_and_create_mask(self, image, left, top, right, bottom): + # Add padding around the image + padded_image = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0) + + # Create a mask with same size as padded image + h, w = padded_image.shape[:2] + mask = np.zeros((h, w), dtype=np.uint8) + + # Set the padded region in the mask to 255 + mask[:top, :] = 255 + mask[-bottom:, :] = 255 + mask[:, :left] = 255 + mask[:, -right:] = 255 + + return padded_image, mask + + + def expand_image(self, image, left, top, right, bottom, inpaint_radius): + batch_size, height, width, channels = image.size() + + image_result = torch.zeros((batch_size, height + top + bottom, width + left + right, channels), dtype=torch.float32) + mask_result = torch.zeros((height + top + bottom, width + left + right), dtype=torch.float32) + + for i in range(batch_size): + img = (image[i, :, :, :].cpu().detach().numpy() * 255.0).astype(np.uint8) + + img, mask = self.add_padding_and_create_mask(img, left, top, right, bottom) + + inpainted = cv2.inpaint(img, mask, inpaint_radius, cv2.INPAINT_NS) + + image_result[i, :, :, :] = torch.from_numpy(inpainted).float() / 255.0 + + if i == 0: + mask_result[:, :] = torch.from_numpy(mask).float() / 255.0 + + image_result = image_result.to(image.device) + mask_result = mask_result.to(image.device) + + masked_image = torch.zeros_like(image_result) + for i in range(image_result.shape[-1]): + masked_image[:, :, :, i] = image_result[:, :, :, i] * (mask < 0.5) + + return (image_result, mask_result, masked_image) + + +NODE_CLASS_MAPPINGS = { + "AstropulsePixelDetector": AstropulsePixelDetector, + "BackgroundRemover": BackgroundRemover, + "ImagePadForBetterOutpaint": ImagePadForBetterOutpaint, + "InpaintingPipelineLoader": InpaintingPipelineLoader, + "Inpainting": Inpainting +} \ No newline at end of file diff --git a/pixeldetector b/pixeldetector new file mode 160000 index 0000000..6267db4 --- /dev/null +++ b/pixeldetector @@ -0,0 +1 @@ +Subproject commit 6267db449241135515a6252f46db119f446683b0 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..36c27aa --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +opencv-python-headless +diffusers \ No newline at end of file diff --git a/workflows/SDXL_pixelart.png b/workflows/SDXL_pixelart.png new file mode 100644 index 0000000..e0b4b86 Binary files /dev/null and b/workflows/SDXL_pixelart.png differ diff --git a/workflows/SDXL_pixelart_preserve_aspect_ratio.png b/workflows/SDXL_pixelart_preserve_aspect_ratio.png new file mode 100644 index 0000000..3560368 Binary files /dev/null and b/workflows/SDXL_pixelart_preserve_aspect_ratio.png differ