diff --git a/__init__.py b/__init__.py index 7bd32dc..06e218d 100644 --- a/__init__.py +++ b/__init__.py @@ -1,13 +1,40 @@ +import subprocess import sys +import threading + import requests from tqdm import tqdm from .config import * -# import pydevd_pycharm -# pydevd_pycharm.settrace('49.7.62.197', port=10090, stdoutToServer=True, stderrToServer=True) +import pydevd_pycharm + +pydevd_pycharm.settrace('49.7.62.197', port=10090, stdoutToServer=True, stderrToServer=True) sys.path.append(utils_path) +def handle_stream(stream, prefix): + for line in stream: + print(prefix, line, end="") + +def run_script(cmd, cwd='.'): + process = subprocess.Popen(cmd, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1) + + stdout_thread = threading.Thread(target=handle_stream, args=(process.stdout, "")) + stderr_thread = threading.Thread(target=handle_stream, args=(process.stderr, "[!]")) + + stdout_thread.start() + stderr_thread.start() + + stdout_thread.join() + stderr_thread.join() + + return process.wait() + +print("## installing dependencies") + +requirements_path = os.path.join(root_path, "requirements.txt") +run_script([sys.executable, '-s', '-m', 'pip', 'install', '-q', '-r', requirements_path]) + from .node import * def urldownload_progressbar(url, file_path): @@ -26,7 +53,7 @@ print("Start Setting weights") for url, filename in zip(urls, filenames): if os.path.exists(filename): continue - print(f"Start Downloading: {url}") + print("Start Downloading: {url}") os.makedirs(os.path.dirname(filename), exist_ok=True) urldownload_progressbar(url, filename) @@ -43,24 +70,26 @@ NODE_CLASS_MAPPINGS = { "MaskDilateErode": MaskDilateErode, "SkinRetouching": SkinRetouching, "PortraitEnhancement": PortraitEnhancement, - "ResizeImage": ResizeImage, + "ImageScaleShort": ImageScaleShort, + "ImageResizeTarget": ImageResizeTarget, "GetImageInfo": GetImageInfo, } NODE_DISPLAY_NAME_MAPPINGS = { - "RetainFace": "RetainFace", - "FaceFusion": "FaceFusion", - "RatioMerge2Image": "RatioMerge2Image", - "MaskMerge2Image": "MaskMerge2Image", - "ReplaceBoxImg": "ReplaceBoxImg", - "ExpandMaskBox": "ExpandMaskBox", - "BoxCropImage": "BoxCropImage", - "ColorTransfer": "ColorTransfer", - "FaceSkin": "FaceSkin", - "MaskDilateErode": "MaskDilateErode", - "SkinRetouching": "SkinRetouching", - "PortraitEnhancement": "PortraitEnhancement", - "ResizeImage": "ResizeImage", - "GetImageInfo": "GetImageInfo", + "RetainFace": "RetainFace PM", + "FaceFusion": "FaceFusion PM", + "RatioMerge2Image": "RatioMerge2Image PM", + "MaskMerge2Image": "MaskMerge2Image PM", + "ReplaceBoxImg": "ReplaceBoxImg PM", + "ExpandMaskBox": "ExpandMaskBox PM", + "BoxCropImage": "BoxCropImage PM", + "ColorTransfer": "ColorTransfer PM", + "FaceSkin": "FaceSkin PM", + "MaskDilateErode": "MaskDilateErode PM", + "SkinRetouching": "SkinRetouching PM", + "PortraitEnhancement": "PortraitEnhancement PM", + "ImageScaleShort": "ImageScaleShort PM", + "ImageResizeTarget": "ImageResizeTarget PM", + "GetImageInfo": "GetImageInfo PM", } __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/config.py b/config.py index c082052..e6b59d0 100644 --- a/config.py +++ b/config.py @@ -1,4 +1,4 @@ -import os, glob +import os from folder_paths import folder_names_and_paths root_path = os.path.dirname(__file__) @@ -6,6 +6,7 @@ utils_path = os.path.join(os.path.dirname(__file__), "utils") models_path = os.path.join(os.path.dirname(__file__), "models") # save_dirs urls = [ + "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/ChilloutMix-ni-fp16.safetensors", "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/control_v11p_sd15_openpose.pth", "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/control_v11p_sd15_canny.pth", "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/control_v11f1e_sd15_tile.pth", @@ -18,6 +19,7 @@ urls = [ "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/webui/face_skin.pth", ] filenames = [ + os.path.join(folder_names_and_paths['checkpoints'][0][0], "Chilloutmix-Ni-pruned-fp16-fix.safetensors"), os.path.join(folder_names_and_paths['controlnet'][0][0], "control_v11p_sd15_openpose.pth"), os.path.join(folder_names_and_paths['controlnet'][0][0], "control_v11p_sd15_canny.pth"), os.path.join(folder_names_and_paths['controlnet'][0][0], "control_v11f1e_sd15_tile.pth"), diff --git a/node.py b/node.py index 4a05256..e469aaa 100644 --- a/node.py +++ b/node.py @@ -1,16 +1,13 @@ -import copy import os import cv2 import numpy as np -import torch from PIL import Image from modelscope.outputs import OutputKeys from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks from .face_process_utils import call_face_crop, color_transfer, Face_Skin from protrait.img_utils import img_to_tensor, tensor_to_img, tensor_to_np, np_to_tensor, np_to_mask, img_to_mask -import torch from .config import models_path import pydevd_pycharm @@ -114,9 +111,6 @@ class MaskMerge2Image: "image2": ("IMAGE",), "mask": ("MASK",), }, - "optional": { - "box": ("BOX",), - } } RETURN_TYPES = ("IMAGE",) @@ -126,10 +120,7 @@ class MaskMerge2Image: def image_mask_merge(self, image1, image2, mask, box=None): mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) - if box is None: - image1 = image1 * mask + image2 * (1 - mask) - else: - image1[:, box[1]:box[3], box[0]:box[2], :] = image1[:, box[1]:box[3], box[0]:box[2], :] * mask + image2[:, box[1]:box[3], box[0]:box[2], :] * (1 - mask) + image1 = image1 * mask + image2 * (1 - mask) return (image1,) class ExpandMaskFaceWidth: @@ -274,7 +265,7 @@ class PortraitEnhancement: output_image = cv2.cvtColor(self.portrait_enhancement(tensor_to_img(image))[OutputKeys.OUTPUT_IMG], cv2.COLOR_BGR2RGB) return (np_to_tensor(output_image),) -class ResizeImage: +class ImageScaleShort: @classmethod def INPUT_TYPES(s): @@ -285,11 +276,11 @@ class ResizeImage: }} RETURN_TYPES = ("IMAGE",) - FUNCTION = "resize_image" + FUNCTION = "image_scale_short" CATEGORY = "protrait/model" - def resize_image(self, image, size, crop_face): + def image_scale_short(self, image, size, crop_face): input_image = tensor_to_img(image) short_side = min(input_image.width, input_image.height) resize = float(short_side / size) @@ -301,6 +292,26 @@ class ResizeImage: input_image = input_image.resize([new_width, new_height], Image.Resampling.LANCZOS) return (img_to_tensor(input_image),) +class ImageResizeTarget: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image": ("IMAGE",), + "width": ("INT", {"default": 512, "min": 0, "max": 2048, "step": 1}), + "height": ("INT", {"default": 512, "min": 0, "max": 2048, "step": 1}), + }} + + RETURN_TYPES = ("IMAGE",) + + FUNCTION = "image_resize_target" + + CATEGORY = "protrait/model" + + def image_resize_target(self, image, width, height): + imagepi = tensor_to_img(image) + out = imagepi.resize([width, height], Image.Resampling.LANCZOS) + return (img_to_tensor(out),) + class GetImageInfo: @classmethod def INPUT_TYPES(s): diff --git a/requirements.txt b/requirements.txt index 79fc31c..ec4e96f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,6 @@ tensorflow onnx onnxruntime modelscope +scikit-image +matplotlib diffusers==0.18.2 \ No newline at end of file diff --git a/utils/models b/utils/models deleted file mode 100644 index e69de29..0000000