tmp

.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/

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "pixeldetector"]
+	path = pixeldetector
+	url = https://github.com/Astropulse/pixeldetector

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.

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)

__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']    

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)        

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
+}

requirements.txt

@@ -0,0 +1,2 @@
+opencv-python-headless
+diffusers