94 lines
2.7 KiB
Python
94 lines
2.7 KiB
Python
import numpy as np
|
|
from PIL import Image
|
|
from torchvision.transforms.functional import to_pil_image, to_tensor
|
|
|
|
from .humanparsing.aigc_run_parsing import Parsing
|
|
from .inference_ootd import OOTDiffusion
|
|
from .ootd_utils import get_mask_location
|
|
from .openpose.run_openpose import OpenPose
|
|
|
|
|
|
class LoadOOTDPipeline:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {}}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("pipe",)
|
|
FUNCTION = "load"
|
|
|
|
CATEGORY = "OOTD"
|
|
|
|
def load(self):
|
|
return (OOTDiffusion(hg_root="models/OOTDiffusion"),)
|
|
|
|
|
|
class OOTDGenerate:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pipe": ("MODEL",),
|
|
"cloth_image": ("IMAGE",),
|
|
"model_image": ("IMAGE",),
|
|
# Openpose from comfyui-controlnet-aux not work
|
|
# "keypoints": ("POSE_KEYPOINT",),
|
|
# "seed": ("INT",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("image", "image_masked")
|
|
FUNCTION = "generate"
|
|
|
|
CATEGORY = "OOTD"
|
|
|
|
def generate(self, pipe, cloth_image, model_image):
|
|
# (1,H,W,3) -> (3,H,W)
|
|
model_image = model_image.squeeze(0)
|
|
model_image = model_image.permute((2, 0, 1))
|
|
model_image = to_pil_image(model_image)
|
|
cloth_image = cloth_image.squeeze(0)
|
|
cloth_image = cloth_image.permute((2, 0, 1))
|
|
cloth_image = to_pil_image(cloth_image)
|
|
model_parse, _ = Parsing(pipe.device)(model_image)
|
|
keypoints = OpenPose()(model_image)
|
|
mask, mask_gray = get_mask_location(
|
|
"hd",
|
|
"upper_body",
|
|
model_parse,
|
|
keypoints,
|
|
width=model_image.size[0],
|
|
height=model_image.size[1],
|
|
)
|
|
masked_vton_img = Image.composite(mask_gray, model_image, mask)
|
|
images = pipe(
|
|
model_type="hd",
|
|
category="upper_body",
|
|
image_garm=cloth_image,
|
|
image_vton=masked_vton_img,
|
|
mask=mask,
|
|
image_ori=model_image,
|
|
num_samples=1,
|
|
num_steps=20,
|
|
image_scale=1,
|
|
seed=-1,
|
|
)
|
|
output_image = to_tensor(images[0])
|
|
output_image = output_image.permute((1, 2, 0))
|
|
masked_vton_img = masked_vton_img.convert("RGB")
|
|
masked_vton_img = to_tensor(masked_vton_img)
|
|
masked_vton_img = masked_vton_img.permute((1, 2, 0))
|
|
return ([output_image], [masked_vton_img])
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LoadOOTDPipeline": LoadOOTDPipeline,
|
|
"OOTDGenerate": OOTDGenerate,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LoadOOTDPipeline": "Load OOTDiffusion",
|
|
"OOTDGenerate": "OOTDiffusion Generate",
|
|
}
|