From 92242f43b127bba67da3eefd3a13d5609b5c8cd5 Mon Sep 17 00:00:00 2001 From: Brian Jemilo II Date: Sat, 12 Apr 2025 22:51:14 -0500 Subject: [PATCH] Tensor Batch Fixes --- CONTRIBUTORS.md | 44 +++++++ LICENSE | 214 +++------------------------------ __init__.py | 35 +++--- nodes/bbox_insert_image.py | 91 ++++++++++++++ nodes/bbox_split.py | 24 ++++ nodes/cascades.py | 7 +- nodes/image_filters.py | 226 +++++++++++++++++++++++++++++++++++ nodes/luts.py | 161 +++++++++++++++++++++---- nodes/remove_area_by_mask.py | 60 ++++++++++ nodes/remove_background.py | 130 ++++++++++++++++++++ nodes/seed.py | 27 +++++ nodes/waifu2x.py | 93 ++++++++------ requirements.txt | 2 + utils/convert.py | 10 ++ web/index.js | 24 ++++ 15 files changed, 866 insertions(+), 282 deletions(-) create mode 100644 CONTRIBUTORS.md create mode 100644 nodes/bbox_insert_image.py create mode 100644 nodes/bbox_split.py create mode 100644 nodes/image_filters.py create mode 100644 nodes/remove_area_by_mask.py create mode 100644 nodes/remove_background.py create mode 100644 nodes/seed.py create mode 100644 requirements.txt create mode 100644 web/index.js diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md new file mode 100644 index 0000000..605035e --- /dev/null +++ b/CONTRIBUTORS.md @@ -0,0 +1,44 @@ +# Contributors and Attributions + +This project integrates and builds upon work from the following developers and repositories. + +--- + +### Third-Party Comfy UI Nodes + +- **WD14-Tagger** + - **Author:** pythongosssss + - **Source:** https://github.com/pythongosssss/stable-diffusion-webui-wd14-tagger + - **License:** MIT + - **Usage:** + - `WD14Tagger`: Tagging images with anime-style concept labels + - **Note:** Switched from my own WD14-Tagger implementation due to its extra convenient features. + +- **Essentials** + - **Author:** Matteo Spinelli + - **Source:** https://github.com/mat-sz/ComfyUI-Utilities + - **License:** MIT + - **Usage:** + - `ImageApplyLUT`: Applying LUTs to images + - **Note:** I adopted this implementation over my own original node due to better customization. + +- **WAS NODE Suite** + - **Author:** Jordan Thompson (WASasquatch) + - **Source:** https://github.com/WASasquatch/was-node-suite-comfyui + - **License:** MIT + - **Usage:** + - `WAS_Image_Filters`: Used for advanced image filter processing + - `WAS_Remove_Rembg`: Used for removing backgrounds efficiently + - **Note:** Replaced my earlier attempts with WAS nodes because I preferred their feature set and reliability. + +--- + +### Special Thanks + +- These developers significantly improved the quality and capability of this project through their publicly shared work. + +--- + +### Original Work + +All other code and systems are original to this repository’s authors. diff --git a/LICENSE b/LICENSE index 8bd9a50..874f59d 100644 --- a/LICENSE +++ b/LICENSE @@ -1,201 +1,21 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ +MIT License - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION +Copyright (c) 2025 Brian Jemilo II - 1. Definitions. +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: - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright 2024 Brian Jemilo II - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. +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/__init__.py b/__init__.py index c67163b..eb010eb 100644 --- a/__init__.py +++ b/__init__.py @@ -1,28 +1,27 @@ -from .nodes.cascades import NODE_CLASS_MAPPINGS as CASCADES_NODE_CLASS_MAPPINGS -from .nodes.cascades import NODE_DISPLAY_NAME_MAPPINGS as CASCADES_NODE_DISPLAY_NAME_MAPPINGS -from .nodes.luts import NODE_CLASS_MAPPINGS as LUTS_NODE_CLASS_MAPPINGS -from .nodes.luts import NODE_DISPLAY_NAME_MAPPINGS as LUTS_NODE_DISPLAY_NAME_MAPPINGS -from .nodes.waifu2x import NODE_CLASS_MAPPINGS as WAIFU2X_NODE_CLASS_MAPPINGS -from .nodes.waifu2x import NODE_DISPLAY_NAME_MAPPINGS as WAIFU2X_NODE_DISPLAY_NAME_MAPPINGS +import os +import importlib WEB_DIRECTORY = "./web" NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} -# Cascades -NODE_CLASS_MAPPINGS.update(CASCADES_NODE_CLASS_MAPPINGS) -NODE_DISPLAY_NAME_MAPPINGS.update(CASCADES_NODE_DISPLAY_NAME_MAPPINGS) -print(f'\033[93m[Shibiko AI] \033[31mCascades \033[92mLoaded\033[0m') -# Luts -NODE_CLASS_MAPPINGS.update(LUTS_NODE_CLASS_MAPPINGS) -NODE_DISPLAY_NAME_MAPPINGS.update(LUTS_NODE_DISPLAY_NAME_MAPPINGS) -print(f'\033[93m[Shibiko AI] \033[31mLuts \033[92mLoaded\033[0m') +def load_module(module_name): + try: + module = importlib.import_module(f'.nodes.{module_name}', package=__package__) + NODE_CLASS_MAPPINGS.update(getattr(module, 'NODE_CLASS_MAPPINGS', {})) + NODE_DISPLAY_NAME_MAPPINGS.update(getattr(module, 'NODE_DISPLAY_NAME_MAPPINGS', {})) + print(f'\033[93m[Shibiko AI] \033[31m{module_name.capitalize()} \033[92mloaded\033[0m') + except ImportError as e: + print(f'\033[93m[Shibiko AI] \033[31mFailed to load {module_name}: {str(e)}\033[0m') -# Waifu2x -NODE_CLASS_MAPPINGS.update(WAIFU2X_NODE_CLASS_MAPPINGS) -NODE_DISPLAY_NAME_MAPPINGS.update(WAIFU2X_NODE_DISPLAY_NAME_MAPPINGS) -print(f'\033[93m[Shibiko AI] \033[31mWaifu2x \033[92mLoaded\033[0m') + +current_dir = os.path.dirname(os.path.abspath(__file__)) +nodes_dir = os.path.join(current_dir, 'nodes') + +for filename in os.listdir(nodes_dir): + if filename.endswith('.py') and not filename.startswith('__'): + load_module(filename[:-3]) __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] print(f'\033[93m[Shibiko AI] \033[92mLoading Complete!\033[0m') diff --git a/nodes/bbox_insert_image.py b/nodes/bbox_insert_image.py new file mode 100644 index 0000000..4162d2b --- /dev/null +++ b/nodes/bbox_insert_image.py @@ -0,0 +1,91 @@ +import numpy as np +import torch +import torch.nn.functional as F + + +class BboxInsertImage: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "bbox_image": ("IMAGE",), + "bbox": ("BBOX",), + "blend_factor": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.1}), + } + } + + CATEGORY = "Shibiko AI" + DESCRIPTION = "Insert Image into Bounding Box Area with Smooth Blending." + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "__call__" + + def __call__(self, image, bbox_image, bbox, blend_factor=0.5): + device = image.device + bbox_image = bbox_image.to(device) + + if isinstance(bbox, (np.ndarray, torch.Tensor)): + bbox = bbox.tolist() + + if len(bbox) == 1 and isinstance(bbox[0], (list, tuple)): + bbox = bbox[0] + + if len(bbox) != 4: + raise ValueError(f"Invalid bbox format. Expected 4 values, but got {len(bbox)}") + + x, y, w, h = map(int, bbox) + + # Ensure bbox_image has 4 dimensions (N, H, W, C) + if bbox_image.dim() == 3: + bbox_image = bbox_image.unsqueeze(0) + + # If bbox_image doesn't have a channel dimension, add it + if bbox_image.shape[-1] != image.shape[-1]: + bbox_image = bbox_image.unsqueeze(-1).repeat(1, 1, 1, image.shape[-1]) + + # Get the dimensions of the original region + original_region = image[:, y:y + h, x:x + w, :] + original_h, original_w = original_region.shape[1:3] + + # Resize bbox_image to match the original region size + bbox_image = F.interpolate(bbox_image.permute(0, 3, 1, 2), + size=(original_h, original_w), + mode='bilinear', + align_corners=False) + bbox_image = bbox_image.permute(0, 2, 3, 1) + + # Ensure bbox_image has the same batch size as image + if bbox_image.shape[0] != image.shape[0]: + bbox_image = bbox_image.repeat(image.shape[0], 1, 1, 1) + + # Create a Gaussian blending mask + blend_mask = self.create_gaussian_mask(original_h, original_w, device) + + # Apply blending + blended_image = blend_factor * bbox_image + (1 - blend_factor) * original_region + blended_image = blend_mask * bbox_image + (1 - blend_mask) * original_region + + # Insert blended image into original image + image[:, y:y + h, x:x + w, :] = blended_image + + return (image,) + + def create_gaussian_mask(self, height, width, device, sigma=10): + """ + Creates a Gaussian blend mask that smoothly fades from the center outwards. + """ + y, x = torch.meshgrid(torch.linspace(-1, 1, height, device=device), + torch.linspace(-1, 1, width, device=device), indexing='ij') + + distance = torch.sqrt(x ** 2 + y ** 2) + mask = torch.exp(-distance ** 2 / (2 * sigma ** 2)) + mask = (mask - mask.min()) / (mask.max() - mask.min()) # Normalize to 0-1 + + return mask.unsqueeze(0).unsqueeze(-1) # Add batch and channel dims + + +NODE_CLASS_MAPPINGS = {"BboxInsertImage": BboxInsertImage} +NODE_DISPLAY_NAME_MAPPINGS = {"BboxInsertImage": "Shibiko AI - BBOX Insert Image"} diff --git a/nodes/bbox_split.py b/nodes/bbox_split.py new file mode 100644 index 0000000..17b88c5 --- /dev/null +++ b/nodes/bbox_split.py @@ -0,0 +1,24 @@ +class BboxSplit: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "bbox": ("BBOX",), + } + } + + CATEGORY = "Shibiko AI" + DESCRIPTION = "Split a BBOX into x, y, width, height." + + RETURN_TYPES = ("INT", "INT", "INT", "INT") + RETURN_NAMES = ("x", "y", "width", "height") + + FUNCTION = "__call__" + + def __call__(self, bbox): + x, y, width, height = bbox[0] + return x, y, width, height, + + +NODE_CLASS_MAPPINGS = {"BboxSplit": BboxSplit} +NODE_DISPLAY_NAME_MAPPINGS = {"BboxSplit": "Shibiko AI - BBOX Split"} diff --git a/nodes/cascades.py b/nodes/cascades.py index 52eb6a0..3941eb8 100644 --- a/nodes/cascades.py +++ b/nodes/cascades.py @@ -31,7 +31,7 @@ class Cascade: }, } - CATEGORY = "Shibiko" + CATEGORY = "Shibiko AI" RETURN_TYPES = ("IMAGE", "MASK", "BBOX") RETURN_NAMES = ("image", "mask", "bbox") @@ -71,7 +71,7 @@ class Cascade: gray = self.image # Detect faces in the image - self.bbox = self.haar_cascade_face.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(24, 24)) + self.bbox = self.haar_cascade_face.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(32, 32)) return self.bbox def draw(self, image=None, padding=None): @@ -101,7 +101,7 @@ class Cascade: for (x, y, w, h) in self.bbox: mask[y - self.padding:y + h + self.padding, x - self.padding:x + w + self.padding] = 255 - if blur is not None: + if blur is not None and blur > 0: if blur_type == 'gaussian': blur = blur if blur % 2 == 1 else blur + 1 dilation = dilation if dilation % 2 == 1 else dilation + 1 @@ -145,6 +145,7 @@ class Cascade: cascade='frontalface_default', dilation=4, padding=50, + **kwargs ): if cascade != self.cascade: self.load(cascade) diff --git a/nodes/image_filters.py b/nodes/image_filters.py new file mode 100644 index 0000000..e147762 --- /dev/null +++ b/nodes/image_filters.py @@ -0,0 +1,226 @@ +import numpy as np +import torch +from PIL import Image, ImageEnhance, ImageFilter +from ..utils.convert import tensor2pil, pil2tensor + + +# From the original code of WAS Node Suite with some additions and modifications +class ImageFilters: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "brightness": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), + "contrast": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 2.0, "step": 0.01}), + "saturation": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01}), + "sharpness": ("FLOAT", {"default": 1.0, "min": -5.0, "max": 5.0, "step": 0.01}), + "blur": ("INT", {"default": 0, "min": 0, "max": 16, "step": 1}), + "sharpen": ("INT", {"default": 0, "min": 0, "max": 16, "step": 1}), + "smooth": ("INT", {"default": 0, "min": 0, "max": 16, "step": 1}), + "smooth_more": ("INT", {"default": 0, "min": 0, "max": 16, "step": 1}), + "gaussian_blur": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1024.0, "step": 0.1}), + "edge_enhance": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "detail_enhance": (["false", "true"],), + }, + } + + CATEGORY = "Shibiko AI" + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "__call__" + + def __call__(self, + image, + brightness=0.0, + contrast=1.0, + saturation=1.0, + sharpness=1.0, + blur=0, + sharpen=0, + smooth=0, + smooth_more=0, + gaussian_blur=0.0, + edge_enhance=0.0, + detail_enhance=False, + **kwargs): + tensors = [] + if len(image) > 1: + for img in image: + + pil_image = None + + # Apply NP Adjustments + if brightness > 0.0 or brightness < 0.0: + # Apply brightness + img = np.clip(img + brightness, 0.0, 1.0) + + if contrast > 1.0 or contrast < 1.0: + # Apply contrast + img = np.clip(img * contrast, 0.0, 1.0) + + # Apply PIL Adjustments + if saturation > 1.0 or saturation < 1.0: + # PIL Image + pil_image = tensor2pil(img) + # Apply saturation + pil_image = ImageEnhance.Color(pil_image).enhance(saturation) + + if sharpness > 1.0 or sharpness < 1.0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Apply sharpness + pil_image = ImageEnhance.Sharpness(pil_image).enhance(sharpness) + + if blur > 0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Apply blur + for _ in range(blur): + pil_image = pil_image.filter(ImageFilter.BLUR) + + if sharpen > 0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Apply blur + for _ in range(sharpen): + pil_image = pil_image.filter(ImageFilter.SHARPEN) + + if smooth > 0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Apply blur + for _ in range(smooth): + pil_image = pil_image.filter(ImageFilter.SMOOTH) + + if smooth_more > 0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Apply blur + for _ in range(smooth_more): + pil_image = pil_image.filter(ImageFilter.SMOOTH_MORE) + + if gaussian_blur > 0.0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Apply Gaussian blur + pil_image = pil_image.filter(ImageFilter.GaussianBlur(radius=gaussian_blur)) + + if edge_enhance > 0.0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Edge Enhancement + edge_enhanced_img = pil_image.filter(ImageFilter.EDGE_ENHANCE_MORE) + # Blend Mask + blend_mask = Image.new(mode="L", size=pil_image.size, color=(round(edge_enhance * 255))) + # Composite Original and Enhanced Version + pil_image = Image.composite(edge_enhanced_img, pil_image, blend_mask) + # Clean-up + del blend_mask, edge_enhanced_img + + if detail_enhance == "true" or detail_enhance is True: + pil_image = pil_image if pil_image else tensor2pil(img) + pil_image = pil_image.filter(ImageFilter.DETAIL) + + # Output image + out_image = (pil2tensor(pil_image) if pil_image else img.unsqueeze(0)) + + tensors.append(out_image) + + tensors = torch.cat(tensors, dim=0) + + else: + + pil_image = None + img = image + + # Apply NP Adjustments + if brightness > 0.0 or brightness < 0.0: + # Apply brightness + img = np.clip(img + brightness, 0.0, 1.0) + + if contrast > 1.0 or contrast < 1.0: + # Apply contrast + img = np.clip(img * contrast, 0.0, 1.0) + + # Apply PIL Adjustments + if saturation > 1.0 or saturation < 1.0: + # PIL Image + pil_image = tensor2pil(img) + # Apply saturation + pil_image = ImageEnhance.Color(pil_image).enhance(saturation) + + if sharpness > 1.0 or sharpness < 1.0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Apply sharpness + pil_image = ImageEnhance.Sharpness(pil_image).enhance(sharpness) + + if blur > 0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Apply blur + for _ in range(blur): + pil_image = pil_image.filter(ImageFilter.BLUR) + + if sharpen > 0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Apply blur + for _ in range(sharpen): + pil_image = pil_image.filter(ImageFilter.SHARPEN) + + if smooth > 0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Apply blur + for _ in range(smooth): + pil_image = pil_image.filter(ImageFilter.SMOOTH) + + if smooth_more > 0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Apply blur + for _ in range(smooth_more): + pil_image = pil_image.filter(ImageFilter.SMOOTH_MORE) + + if gaussian_blur > 0.0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Apply Gaussian blur + pil_image = pil_image.filter( + ImageFilter.GaussianBlur(radius=gaussian_blur)) + + if edge_enhance > 0.0: + # Assign or create PIL Image + pil_image = pil_image if pil_image else tensor2pil(img) + # Edge Enhancement + edge_enhanced_img = pil_image.filter(ImageFilter.EDGE_ENHANCE_MORE) + # Blend Mask + blend_mask = Image.new( + mode="L", size=pil_image.size, color=(round(edge_enhance * 255))) + # Composite Original and Enhanced Version + pil_image = Image.composite( + edge_enhanced_img, pil_image, blend_mask) + # Clean-up + del blend_mask, edge_enhanced_img + + if detail_enhance == "true": + pil_image = pil_image if pil_image else tensor2pil(img) + pil_image = pil_image.filter(ImageFilter.DETAIL) + + # Output image + out_image = (pil2tensor(pil_image) if pil_image else img) + + tensors = out_image + + return (tensors, ) + + +NODE_CLASS_MAPPINGS = {"ImageFilters": ImageFilters} +NODE_DISPLAY_NAME_MAPPINGS = {"Cascade": "Shibiko AI - Image Filters"} diff --git a/nodes/luts.py b/nodes/luts.py index 1ac65a7..1152b38 100644 --- a/nodes/luts.py +++ b/nodes/luts.py @@ -1,16 +1,22 @@ import os + +import folder_paths +import numpy as np +import torch +from pathlib import Path from PIL import ImageFilter +from colour.io.luts.iridas_cube import read_LUT_IridasCube from typing import List, Union -from ..utils.convert import convert +from ..utils.convert import convert, tensor2pil, pil2tensor from ..utils.directory import initialize_directory class Luts: luts_directory = initialize_directory('luts') - + @classmethod - def INPUT_TYPES(self): - luts = [lut.split(".")[0] for lut in os.listdir(self.luts_directory) if lut.endswith(".cube")] + def INPUT_TYPES(cls): + luts = [lut.split(".")[0] for lut in os.listdir(cls.luts_directory) if lut.endswith(".cube")] return { "required": { @@ -19,7 +25,7 @@ class Luts: }, } - CATEGORY = "Shibiko" + CATEGORY = "Shibiko AI" RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) @@ -27,18 +33,19 @@ class Luts: FUNCTION = "__call__" def apply_lut(self, image, lut): - print("Applying LUT...") if not lut.endswith(".cube"): lut += ".cube" lut_file_name = lut lut_directory = initialize_directory('luts') - lut_path = os.path.join(lut_directory, f"{lut_file_name}") - lut = self.read_lut(lut_path) + lut_path = os.path.join(lut_directory, lut_file_name) + lut_filter = self.read_lut(lut_path) - if isinstance(image, list): - image = [img.filter(lut) for img in image] - else: - image = image.filter(lut) + # Convert tensor to PIL, apply LUT, then convert back to tensor + if isinstance(image, torch.Tensor): + image = tensor2pil(image) + + # Apply the LUT filter + image = image.filter(lut_filter) return image @@ -49,35 +56,139 @@ class Luts: size = round(len(lut_raw) ** (1 / 3)) lut_table = [ - self.row2val(row.split(" ")) for row in lut_raw if self.is_3dlut_row(row.split(" ")) + self.row2val(row.split(" ")) for row in lut_raw if self.is_3D_lut_row(row.split(" ")) ] return ImageFilter.Color3DLUT(size, lut_table, num_channels) @staticmethod - def is_3dlut_row(row: List) -> bool: + def is_3D_lut_row(row: List) -> bool: """Check if one line in the file has exactly 3 values""" row_values = [] for val in row: try: row_values.append(float(val)) - except: + except ValueError: return False - if len(row_values) == 3: - return True - return False + return len(row_values) == 3 @staticmethod def row2val(row): return tuple([float(val) for val in row]) - def __call__(self, image, lut, unique_id): - image = convert(image) - image = self.apply_lut(image, lut) - image = convert(image) + def __call__(self, image, lut, **kwargs): + # Handle single image or batch + if image.dim() == 3: # Single image (H, W, C) + images = [image] + else: # Batch of images (B, H, W, C) + images = [img for img in image] - return (image,) + processed_images = [] + for img in images: + # Convert to PIL + img_pil = tensor2pil(img) + + # Apply LUT + img_processed = self.apply_lut(img_pil, lut) + + # Convert back to tensor + img_tensor = pil2tensor(img_processed) + processed_images.append(img_tensor) + + # Stack all tensors if we have multiple images + if len(processed_images) > 1: + output = torch.cat(processed_images, dim=0) + else: + output = processed_images[0] + + # Ensure proper output format (B, H, W, C) with float32 dtype + if output.dtype != torch.float32: + output = output.float() + if output.max() > 1.0: + output = output / 255.0 + + return (output,) -NODE_CLASS_MAPPINGS = {"Luts": Luts} -NODE_DISPLAY_NAME_MAPPINGS = {"Luts": "Shibiko AI - Luts"} +# From https://github.com/yoonsikp/pycubelut/blob/master/pycubelut.py (MIT license) +class LutsAdvanced: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "lut_file": (folder_paths.get_filename_list("luts"), {"default": "Cinematic.cube"}), + "gamma_correction": ("BOOLEAN", {"default": True}), + "clip_values": ("BOOLEAN", {"default": True}), + "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1}), + }} + + CATEGORY = "Shibiko AI" + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "__call__" + + # TODO: check if we can do without numpy + def __call__(self, image, lut_file="cinematic", gamma_correction=True, clip_values=True, strength=1.0): + lut_file_path = folder_paths.get_full_path("luts", lut_file + '.cube') + if not lut_file_path or not Path(lut_file_path).exists(): + print(f"Could not find LUT file: {lut_file_path}") + return (image,) + + device = image.device + lut = read_LUT_IridasCube(lut_file_path) + lut.name = lut_file + + if clip_values: + if lut.domain[0].max() == lut.domain[0].min() and lut.domain[1].max() == lut.domain[1].min(): + lut.table = np.clip(lut.table, lut.domain[0, 0], lut.domain[1, 0]) + else: + for dim in range(3): + domain0 = lut.domain[0, dim] + domain1 = lut.domain[1, dim] + + if len(lut.table.shape) == 2: + lut.table[:, dim] = np.clip(lut.table[:, dim], domain0, domain1) + else: # 3D + lut.table[:, :, :, dim] = np.clip(lut.table[:, :, :, dim], domain0, domain1) + + out = [] + for img in image: # TODO: is this more resource efficient? should we use a batch instead? + lut_img = img.cpu().numpy().copy() + + is_non_default_domain = not np.array_equal(lut.domain, np.array([[0., 0., 0.], [1., 1., 1.]])) + dom_scale = None + if is_non_default_domain: + dom_scale = lut.domain[1] - lut.domain[0] + lut_img = lut_img * dom_scale + lut.domain[0] + + if gamma_correction: + lut_img = lut_img ** (1 / 2.2) + + lut_img = lut.apply(lut_img) + if gamma_correction: + lut_img = lut_img ** 2.2 + + if is_non_default_domain: + lut_img = (lut_img - lut.domain[0]) / dom_scale + + lut_img = torch.from_numpy(lut_img).to(device) + if strength < 1.0: + lut_img = strength * lut_img + (1 - strength) * img + out.append(lut_img) + + out = torch.stack(out) + + return (out,) + + +NODE_CLASS_MAPPINGS = { + "Luts": Luts, + "LutsAdvanced": LutsAdvanced +} +NODE_DISPLAY_NAME_MAPPINGS = { + "Luts": "Shibiko AI - Luts", + "LutsAdvanced": "Shibiko AI - Luts (Advanced)" +} diff --git a/nodes/remove_area_by_mask.py b/nodes/remove_area_by_mask.py new file mode 100644 index 0000000..0e3439f --- /dev/null +++ b/nodes/remove_area_by_mask.py @@ -0,0 +1,60 @@ +import cv2 +import numpy as np +import torch + + +class RemoveAreaByMask: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "mask": ("MASK",), + }, + "optional": { + "invert_mask": ("BOOLEAN", {"default": False}), + }, + } + + CATEGORY = "Shibiko AI" + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "__call__" + + @staticmethod + def NHWC_to_CV2(image): + image_numpy = image.cpu().numpy() + return (image_numpy[0] * 255).astype(np.uint8) # Convert to uint8, keep NHWC + + @staticmethod + def CV2_to_NHWC(image): + image_rgb = image.astype(np.float32) / 255.0 + return torch.from_numpy(image_rgb).unsqueeze(0) # (1, H, W, C) + + def __call__(self, image, mask, invert_mask=False): + image_cv2 = self.NHWC_to_CV2(image) + mask_cv2 = self.NHWC_to_CV2(mask) + + if mask_cv2.ndim == 3 and mask_cv2.shape[2] == 3: + mask_cv2 = cv2.cvtColor(mask_cv2, cv2.COLOR_BGR2GRAY) + + if mask_cv2.shape[:2] != image_cv2.shape[:2]: + mask_cv2 = cv2.resize(mask_cv2, (image_cv2.shape[1], image_cv2.shape[0]), interpolation=cv2.INTER_NEAREST) + + mask_float = mask_cv2.astype(np.float32) / 255.0 + + if mask_float.ndim == 2: + mask_float = np.expand_dims(mask_float, axis=-1) + + if invert_mask: + mask_float = 1.0 - mask_float + + masked_image = (image_cv2 * (1 - mask_float)).astype(np.uint8) + + return (self.CV2_to_NHWC(masked_image),) # Return (1, H, W, C) + + +NODE_CLASS_MAPPINGS = {"RemoveAreaByMask": RemoveAreaByMask} +NODE_DISPLAY_NAME_MAPPINGS = {"RemoveAreaByMask": "Shibiko AI - Remove Area By Mask"} diff --git a/nodes/remove_background.py b/nodes/remove_background.py new file mode 100644 index 0000000..0c8ae40 --- /dev/null +++ b/nodes/remove_background.py @@ -0,0 +1,130 @@ +import os +import torch +import folder_paths +from rembg import remove, new_session +from ..utils.convert import pil2tensor, tensor2pil + + +class RemoveBackground: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "transparency": ("BOOLEAN", {"default": True},), + "model": (["u2net", "u2netp", "u2net_human_seg", "silueta", "isnet-general-use", "isnet-anime"], {"default": "isnet-anime"}), + "post_processing": ("BOOLEAN", {"default": False}), + "only_mask": ("BOOLEAN", {"default": False},), + "alpha_matting": ("BOOLEAN", {"default": False},), + "alpha_matting_foreground_threshold": ("INT", {"default": 240, "min": 0, "max": 255}), + "alpha_matting_background_threshold": ("INT", {"default": 10, "min": 0, "max": 255}), + "alpha_matting_erode_size": ("INT", {"default": 10, "min": 0, "max": 255}), + "background_color": (["none", "black", "white", "magenta", "chroma green", "chroma blue"],), + "putalpha": ("BOOLEAN", {"default": False},), + }, + } + + CATEGORY = "Shibiko AI" + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "__call__" + + # A helper function to convert from strings to logical boolean + # Conforms to https://docs.python.org/3/library/stdtypes.html#truth-value-testing + # With the addition of evaluating string representations of Falsy types + @staticmethod + def __convertToBool(x): + + # Evaluate string representation of False types + if type(x) is str: + x = x.strip() + if (x.lower() == 'false' + or x.lower() == 'none' + or x == '0' + or x == '0.0' + or x == '0j' + or x == "''" + or x == '""' + or x == "()" + or x == "[]" + or x == "{}" + or x.lower() == "decimal(0)" + or x.lower() == "fraction(0,1)" + or x.lower() == "set()" + or x.lower() == "range(0)"): + return False + else: + return True + + return bool(x) + + def __call__( + self, + images, + transparency=True, + model="isnet-anime", + alpha_matting=False, + alpha_matting_foreground_threshold=240, + alpha_matting_background_threshold=10, + alpha_matting_erode_size=10, + post_processing=False, + only_mask=False, + background_color="none", + putalpha=False, + **kwargs + ): + # ComfyUI will allow strings in place of booleans, validate the input. + transparency = transparency if type(transparency) is bool else self.__convertToBool(transparency) + alpha_matting = alpha_matting if type(alpha_matting) is bool else self.__convertToBool(alpha_matting) + post_processing = post_processing if type(post_processing) is bool else self.__convertToBool(post_processing) + only_mask = only_mask if type(only_mask) is bool else self.__convertToBool(only_mask) + + os.environ['U2NET_HOME'] = os.path.join(folder_paths.models_dir, 'rembg') + os.makedirs(os.environ['U2NET_HOME'], exist_ok=True) + + # Set background color + if background_color == "black": + bgrgba = [0, 0, 0, 255] + elif background_color == "white": + bgrgba = [255, 255, 255, 255] + elif background_color == "magenta": + bgrgba = [255, 0, 255, 255] + elif background_color == "chroma green": + bgrgba = [0, 177, 64, 255] + elif background_color == "chroma blue": + bgrgba = [0, 71, 187, 255] + else: + bgrgba = None + + if transparency and bgrgba is not None: + bgrgba[3] = 0 + + batch_tensor = [] + for image in images: + image = tensor2pil(image) + batch_tensor.append(pil2tensor( + remove( + image, + session=new_session(model), + post_process_mask=post_processing, + alpha_matting=alpha_matting, + alpha_matting_foreground_threshold=alpha_matting_foreground_threshold, + alpha_matting_background_threshold=alpha_matting_background_threshold, + alpha_matting_erode_size=alpha_matting_erode_size, + only_mask=only_mask, + bgcolor=bgrgba, + putalpha=putalpha, + ) + .convert(('RGBA' if transparency else 'RGB')))) + batch_tensor = torch.cat(batch_tensor, dim=0) + + return (batch_tensor,) + + +NODE_CLASS_MAPPINGS = {"RemoveBackground": RemoveBackground} +NODE_DISPLAY_NAME_MAPPINGS = {"RemoveBackground": "Shibiko AI - Remove Background"} diff --git a/nodes/seed.py b/nodes/seed.py new file mode 100644 index 0000000..f7bafba --- /dev/null +++ b/nodes/seed.py @@ -0,0 +1,27 @@ +class SeedGenerator: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "seed": ("INT", {"default": 0, "min": 0, "max": 18446744073709551615}), + } + } + + @classmethod + def IS_CHANGED(cls, seed=0): + return seed + + CATEGORY = "Shibiko AI" + DESCRIPTION = "Share Your Seed Across Multiple Nodes." + + RETURN_TYPES = ("INT",) + RETURN_NAMES = ("SEED",) + + FUNCTION = "__call__" + + def __call__(self, seed=0): + return (seed,) + + +NODE_CLASS_MAPPINGS = {"SeedGenerator": SeedGenerator} +NODE_DISPLAY_NAME_MAPPINGS = {"SeedGenerator": "Shibiko AI - Seed Generator"} diff --git a/nodes/waifu2x.py b/nodes/waifu2x.py index fdc1c89..33ece4e 100644 --- a/nodes/waifu2x.py +++ b/nodes/waifu2x.py @@ -1,6 +1,6 @@ import torch from typing import Optional -from ..utils.convert import convert +from ..utils.convert import convert, tensor2pil, pil2tensor class Waifu2x: @@ -17,7 +17,7 @@ class Waifu2x: }, } - CATEGORY = "Shibiko" + CATEGORY = "Shibiko AI" RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) @@ -35,7 +35,6 @@ class Waifu2x: self.output_type = kwargs.get('output_type', 'pil') self.scale = kwargs.get('scale', 1) self.tile_size = kwargs.get('tile_size', 256) - self.load() def load(self, model_type: Optional[str] = 'art'): @@ -47,22 +46,21 @@ class Waifu2x: keep_alpha=self.keep_alpha, amp=self.amp, trust_repo=True, - ).to('cuda') - + ) self.method = self.waifu2x_method(self.scale, self.noise_level) self.model.set_mode(method=self.method, noise_level=self.noise_level) def set( - self, - amp: Optional[bool] = None, - batch_size: Optional[int] = None, - enabled: Optional[bool] = None, - keep_alpha: Optional[bool] = None, - model_type: Optional[str] = None, - noise_level: Optional[int] = None, - output_type: Optional[str] = None, - scale: Optional[int] = None, - tile_size: Optional[int] = None, + self, + amp: Optional[bool] = None, + batch_size: Optional[int] = None, + enabled: Optional[bool] = None, + keep_alpha: Optional[bool] = None, + model_type: Optional[str] = None, + noise_level: Optional[int] = None, + output_type: Optional[str] = None, + scale: Optional[int] = None, + tile_size: Optional[int] = None, ): settings = locals() settings.pop('self') @@ -86,36 +84,53 @@ class Waifu2x: return 'noise_scale4x' def __call__( - self, - image, - scale: Optional[int] = 1, - noise_level: Optional[int] = 3, - model_type: Optional[str] = 'art', + self, + image, + scale: Optional[int] = 1, + noise_level: Optional[int] = 3, + model_type: Optional[str] = 'art', + **kwargs ): - print(f'Waifu2x Upscaling image with noise_level:{noise_level} scale:{scale}x model_type:{model_type}...') - print(f'Type of image: {type(image)}') - + print('WAIFU2X CALL:', scale, noise_level, model_type) if self.model_type != model_type: self.load(model_type) - image = convert(image) - scale = scale if scale is not None else self.scale - noise_level = noise_level if noise_level is not None else self.noise_level + if isinstance(image, torch.Tensor): + if image.dim() == 3: + image = [image] + elif image.dim() == 4: + image = [img for img in image] - if self.scale != scale or self.noise_level != noise_level: - self.method = self.waifu2x_method(scale, noise_level) - self.model.set_mode(self.method, noise_level) - image = self.model.infer(image, method=self.method, noise_level=noise_level, output_type='pil') - image = convert(image) - return (image,) + tensors = [] + for img in image: + img = img.cpu() + img_pil = tensor2pil(img) - # @classmethod - # def IS_CHANGED(self, video, **kwargs): - # return hash_path(video) - # - # @classmethod - # def VALIDATE_INPUTS(self, video, **kwargs): - # return validate_path(video, allow_none=True) + scale = scale if scale is not None else self.scale + noise_level = noise_level if noise_level is not None else self.noise_level + + if self.scale != scale or self.noise_level != noise_level: + self.method = self.waifu2x_method(scale, noise_level) + self.model.set_mode(self.method, noise_level) + + print('METHOD:', self.method) + + img_pil = self.model.infer(img_pil, method=self.method, noise_level=noise_level, output_type='pil') + + img_tensor = pil2tensor(img_pil) + tensors.append(img_tensor) + + if len(tensors) > 1: + output = torch.cat(tensors, dim=0) + else: + output = tensors[0] + + if output.dtype != torch.float32: + output = output.float() + if output.max() > 1.0: + output = output / 255.0 + + return (output,) NODE_CLASS_MAPPINGS = {"Waifu2x": Waifu2x} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..a89eca2 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +colour-science +rembg diff --git a/utils/convert.py b/utils/convert.py index 34d0276..ad199ac 100644 --- a/utils/convert.py +++ b/utils/convert.py @@ -16,3 +16,13 @@ def convert(image): image = Image.fromarray(np.clip(img_array, 0, 255).astype(np.uint8)) return image + + +# Tensor to PIL +def tensor2pil(image): + return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + + +# PIL to Tensor +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) diff --git a/web/index.js b/web/index.js new file mode 100644 index 0000000..2351979 --- /dev/null +++ b/web/index.js @@ -0,0 +1,24 @@ +import { app } from '../../scripts/app.js'; +import { loadCSS, loadScript, sleep } from './utils.js'; +import { applyStyles } from './node-styles.js'; +import { showCode } from './show-code.js'; +import { storeAnyNodeOutputs } from './store-outputs.js'; + +app.registerExtension({ + name: 'ShibikoAI', + async init() { + const CDN_BASE_URL = 'https://cdnjs.cloudflare.com/ajax/libs'; + await loadScript(`${CDN_BASE_URL}/highlight.js/11.9.0/highlight.min.js`); + await loadScript(`${CDN_BASE_URL}/highlightjs-line-numbers.js/2.8.0/highlightjs-line-numbers.min.js`); + await loadCSS(`${CDN_BASE_URL}/highlight.js/11.9.0/styles/atom-one-dark.min.css`); + await loadCSS(`/extensions/anynode/css/show-code.css`); + hljs.configure({ ignoreUnescapedHTML: true }); + }, + async nodeCreated(node) { + await sleep(0); // Wait for node object to be updated + if (node.type.includes('AnyNode')) { + applyStyles(node); + node.type === 'AnyNodeShowCode' ? showCode(node) : storeAnyNodeOutputs(node); + } + }, +});