update nodes

This commit is contained in:
toto
2023-11-20 18:00:24 +08:00
parent 4e531256df
commit cd0cc7ea2f
5 changed files with 138 additions and 53 deletions
+4 -2
View File
@@ -10,19 +10,21 @@ from .comfyui.nodes import *
from .comfyui.style_loader_node import *
NODE_CLASS_MAPPINGS = {
# "FC_LoraMerge": FCLoraMerge,
"FC_FaceFusion": FCFaceFusion,
"FC_StyleLoraLoad": FCStyleLoraLoad,
"FC_FaceDetection": FCFaceDetection,
"FC_CropMask": FCCropMask,
"FC_Segment": FCSegment,
"FC_ReplaceImage": FCReplaceImage,
"FC_CropBottom": FCCropBottom,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FC_FaceFusion": "FC FaceFusion",
"FC_StyleLoraLoad": "FC StyleLoraLoad",
"FC_FaceDetection": "FC FaceDetection",
"FC_CropMask": "FC CropMask",
"FC_Segment": "FC Segment",
"FC_ReplaceImage": "FC ReplaceImage",
"FC_CropBottom": "FC CropBottom",
}
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
+65
View File
@@ -15,10 +15,12 @@ from torch import multiprocessing
from transformers import pipeline as tpipeline
import pydevd_pycharm
pydevd_pycharm.settrace('49.7.62.197', port=10090, stdoutToServer=True, stderrToServer=True)
from .model_holder import *
from .utils.img_utils import *
from .utils.convert_utils import *
class FCLoraMerge:
@classmethod
@@ -222,3 +224,66 @@ class FCSegment:
source_image = tensor_to_img(source_image)
mask = self.segment(get_segmentation(), source_image, ksize=0.1)
return (mask_np2_to_mask_tensor(mask),)
class FCReplaceImage:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"source_image": ("IMAGE",),
"replace_image": ("IMAGE",),
"face_box": ("BOX",),
"mask": ("MASK",)
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "replace_image"
CATEGORY = "facechain/model"
def replace_image(self, source_image, replace_image, face_box, mask):
face_ratio = 0.45
h, w, _ = replace_image.shape
cropl = int(max(face_box[3] - face_box[1], face_box[2] - face_box[0]) / face_ratio / 2)
cx = int((face_box[2] + face_box[0]) / 2)
cy = int((face_box[1] + face_box[3]) / 2)
cropup = min(cy, cropl)
cropbo = min(h - cy, cropl)
crople = min(cx, cropl)
cropri = min(w - cx, cropl)
ksize = int(10 * cropl / 256)
rst_gen = cv2.resize(replace_image, (cropl * 2, cropl * 2))
rst_crop = rst_gen[cropl - cropup:cropl + cropbo, cropl - crople:cropl + cropri]
print(rst_crop.shape)
inpaint_img_rst = np.zeros_like(source_image)
print('Start pasting.')
inpaint_img_rst[cy - cropup:cy + cropbo, cx - crople:cx + cropri] = rst_crop
print('Fininsh pasting.')
print(inpaint_img_rst.shape, mask.shape, source_image.shape)
mask_large = mask.astype(np.float32)
kernel = np.ones((ksize * 2, ksize * 2))
mask_large1 = cv2.erode(mask_large, kernel, iterations=1)
mask_large1 = cv2.GaussianBlur(mask_large1, (int(ksize * 1.8) * 2 + 1, int(ksize * 1.8) * 2 + 1), 0)
mask_large1[face_box[1]:face_box[3], face_box[0]:face_box[2]] = 1
mask_large = mask_large * mask_large1
final_inpaint_rst = (inpaint_img_rst.astype(np.float32) * mask_large.astype(np.float32) + source_image.astype(np.float32) * (1.0 - mask_large.astype(np.float32))).astype(np.uint8)
return (final_inpaint_rst,)
class FCCropBottom:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"source_image": ("IMAGE",),
"face_index": ("INT", {"default": 0, "min": 0, "max": 10, "step": 1})
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "crop_bottom"
CATEGORY = "facechain/crop"
def crop_bottom(self, source_image, width):
source_image = tensor_to_img(source_image)
crop_result = crop_bottom(source_image, width)
return (img_to_tensor(crop_result),)
+1
View File
@@ -1,3 +1,4 @@
class FCStyleLoraLoad:
@classmethod
def INPUT_TYPES(s):
+54
View File
@@ -0,0 +1,54 @@
import numpy as np
import torch
from PIL import ImageOps
from PIL import Image
def img_to_tensor(input):
i = ImageOps.exif_transpose(input)
image = i.convert("RGB")
image = np.array(image).astype(np.float32) / 255.0
tensor = torch.from_numpy(image)[None,]
return tensor
def img_to_np(input):
i = ImageOps.exif_transpose(input)
image = i.convert("RGB")
image_np = np.array(image).astype(np.float32)
return image_np
def img_to_mask(input):
i = ImageOps.exif_transpose(input)
image = i.convert("RGB")
new_np = np.array(image).astype(np.float32) / 255.0
mask_tensor = torch.from_numpy(new_np).permute(2, 0, 1)[0:1, :, :]
return mask_tensor
def image_np2_to_mask_tensor(input):
image = input.astype(np.float32) / 255.0
tensor = torch.from_numpy(image)[None,]
return tensor
def mask_np2_to_mask_tensor(input):
image = input.astype(np.float32)
tensor = torch.from_numpy(image)[None,]
return tensor
def mask_np3_to_mask_tensor(input):
image = input.astype(np.float32)
tensor = torch.from_numpy(image).permute(2, 0, 1)[0:1, :, :]
return tensor
def tensor_to_img(image):
image = image[0]
i = 255. * image.cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)).convert("RGB")
return img
def tensor_to_np(image):
image = image[0]
i = 255. * image.cpu().numpy()
result = np.clip(i, 0, 255).astype(np.uint8)
return result
def image_np_to_mask(input):
new_np = input.astype(np.float32) / 255.0
tensor = torch.from_numpy(new_np).permute(2, 0, 1)[0:1, :, :]
return tensor
+14 -51
View File
@@ -1,54 +1,17 @@
import numpy as np
import torch
from PIL import ImageOps
from PIL import Image
def img_to_tensor(input):
i = ImageOps.exif_transpose(input)
image = i.convert("RGB")
image = np.array(image).astype(np.float32) / 255.0
tensor = torch.from_numpy(image)[None,]
return tensor
def img_to_np(input):
i = ImageOps.exif_transpose(input)
image = i.convert("RGB")
image_np = np.array(image).astype(np.float32)
return image_np
def img_to_mask(input):
i = ImageOps.exif_transpose(input)
image = i.convert("RGB")
new_np = np.array(image).astype(np.float32) / 255.0
mask_tensor = torch.from_numpy(new_np).permute(2, 0, 1)[0:1, :, :]
return mask_tensor
def image_np2_to_mask_tensor(input):
image = input.astype(np.float32) / 255.0
tensor = torch.from_numpy(image)[None,]
return tensor
def mask_np2_to_mask_tensor(input):
image = input.astype(np.float32)
tensor = torch.from_numpy(image)[None,]
return tensor
def mask_np3_to_mask_tensor(input):
image = input.astype(np.float32)
tensor = torch.from_numpy(image).permute(2, 0, 1)[0:1, :, :]
return tensor
def tensor_to_img(image):
image = image[0]
i = 255. * image.cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)).convert("RGB")
return img
def tensor_to_np(image):
image = image[0]
i = 255. * image.cpu().numpy()
result = np.clip(i, 0, 255).astype(np.uint8)
return result
def image_np_to_mask(input):
new_np = input.astype(np.float32) / 255.0
tensor = torch.from_numpy(new_np).permute(2, 0, 1)[0:1, :, :]
return tensor
def crop_bottom(pil_file, width):
if width == 512:
height = 768
else:
height = 1152
w, h = pil_file.size
factor = w / width
new_h = int(h / factor)
pil_file = pil_file.resize((width, new_h))
crop_h = min(int(new_h / 32) * 32, height)
array_file = np.array(pil_file)
array_file = array_file[:crop_h, :, :]
output_file = Image.fromarray(array_file)
return output_file