247 lines
9.5 KiB
Python
247 lines
9.5 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
|
|
import folder_paths
|
|
import os
|
|
import torch
|
|
|
|
|
|
from diffusers import (DDIMScheduler, ControlNetModel,
|
|
KDPM2AncestralDiscreteScheduler, LMSDiscreteScheduler, DPMSolverMultistepScheduler,
|
|
DPMSolverSinglestepScheduler,
|
|
EulerDiscreteScheduler, HeunDiscreteScheduler, KDPM2DiscreteScheduler,
|
|
EulerAncestralDiscreteScheduler, UniPCMultistepScheduler,
|
|
DDPMScheduler, LCMScheduler, StableDiffusionPipeline, )
|
|
from .pipeline_sd15 import StableDiffusionControlNetPipeline
|
|
from .detail_encoder.encoder_plus import detail_encoder
|
|
from .import spiga_draw
|
|
from PIL import Image
|
|
from .facelib import FaceDetector
|
|
from comfy.utils import common_upscale
|
|
import numpy as np
|
|
|
|
makeup_current_path = os.path.dirname(os.path.abspath(__file__))
|
|
weigths_current_path = os.path.join(folder_paths.models_dir, "stable_makeup")
|
|
if not os.path.exists(weigths_current_path):
|
|
os.makedirs(weigths_current_path)
|
|
|
|
scheduler_list = ["DDIM",
|
|
"Euler",
|
|
"Euler a",
|
|
"DDPM",
|
|
"DPM++ 2M",
|
|
"DPM++ 2M Karras",
|
|
"DPM++ 2M SDE",
|
|
"DPM++ 2M SDE Karras",
|
|
"DPM++ SDE",
|
|
"DPM++ SDE Karras",
|
|
"DPM2",
|
|
"DPM2 Karras",
|
|
"DPM2 a",
|
|
"DPM2 a Karras",
|
|
"Heun",
|
|
"LCM",
|
|
"LMS",
|
|
"LMS Karras",
|
|
"UniPC",
|
|
]
|
|
|
|
def get_sheduler(name):
|
|
scheduler = False
|
|
if name == "Euler":
|
|
scheduler = EulerDiscreteScheduler()
|
|
elif name == "Euler a":
|
|
scheduler = EulerAncestralDiscreteScheduler()
|
|
elif name == "DDIM":
|
|
scheduler = DDIMScheduler()
|
|
elif name == "DDPM":
|
|
scheduler = DDPMScheduler()
|
|
elif name == "DPM++ 2M":
|
|
scheduler = DPMSolverMultistepScheduler()
|
|
elif name == "DPM++ 2M Karras":
|
|
scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True)
|
|
elif name == "DPM++ 2M SDE":
|
|
scheduler = DPMSolverMultistepScheduler(algorithm_type="sde-dpmsolver++")
|
|
elif name == "DPM++ 2M SDE Karras":
|
|
scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True, algorithm_type="sde-dpmsolver++")
|
|
elif name == "DPM++ SDE":
|
|
scheduler = DPMSolverSinglestepScheduler()
|
|
elif name == "DPM++ SDE Karras":
|
|
scheduler = DPMSolverSinglestepScheduler(use_karras_sigmas=True)
|
|
elif name == "DPM2":
|
|
scheduler = KDPM2DiscreteScheduler()
|
|
elif name == "DPM2 Karras":
|
|
scheduler = KDPM2DiscreteScheduler(use_karras_sigmas=True)
|
|
elif name == "DPM2 a":
|
|
scheduler = KDPM2AncestralDiscreteScheduler()
|
|
elif name == "DPM2 a Karras":
|
|
scheduler = KDPM2AncestralDiscreteScheduler(use_karras_sigmas=True)
|
|
elif name == "Heun":
|
|
scheduler = HeunDiscreteScheduler()
|
|
elif name == "LCM":
|
|
scheduler = LCMScheduler()
|
|
elif name == "LMS":
|
|
scheduler = LMSDiscreteScheduler()
|
|
elif name == "LMS Karras":
|
|
scheduler = LMSDiscreteScheduler(use_karras_sigmas=True)
|
|
elif name == "UniPC":
|
|
scheduler = UniPCMultistepScheduler()
|
|
return scheduler
|
|
|
|
def tensor_to_pil(tensor):
|
|
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
|
image = Image.fromarray(image_np, mode='RGB')
|
|
return image
|
|
|
|
def pil2narry(img):
|
|
img = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
|
|
return img
|
|
|
|
def nomarl_upscale(img_tensor, width, height):
|
|
samples = img_tensor.movedim(-1, 1)
|
|
img = common_upscale(samples, width, height, "nearest-exact", "center")
|
|
samples = img.movedim(1, -1)
|
|
img_pil = tensor_to_pil(samples)
|
|
return img_pil
|
|
|
|
|
|
class StableMakeup_LoadModel:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
|
|
"clip":(folder_paths.get_filename_list("clip"),),
|
|
"lora":(["none"]+folder_paths.get_filename_list("loras"),),
|
|
"lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1, "round": 0.01}),
|
|
"lora_trigger_words":("STRING", {"default": "best"}),
|
|
"scheduler": (scheduler_list,),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MAKEUP_MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "main_loader"
|
|
CATEGORY = "Stable_Makeup"
|
|
|
|
def main_loader(self,ckpt_name,clip,lora,lora_scale,lora_trigger_words,scheduler):
|
|
clip=folder_paths.get_full_path("clip", clip)
|
|
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
|
scheduler_used = get_sheduler(scheduler)
|
|
makeup_encoder_path = os.path.join(weigths_current_path,"pytorch_model.bin")
|
|
id_encoder_path = os.path.join(weigths_current_path,"pytorch_model_1.bin")
|
|
pose_encoder_path = os.path.join(weigths_current_path,"pytorch_model_2.bin")
|
|
original_config_file=os.path.join(makeup_current_path,"datasets/v1-inference.yaml")
|
|
sd15_config=os.path.join(makeup_current_path,"sd15_config")
|
|
try:
|
|
pipe = StableDiffusionPipeline.from_single_file(
|
|
ckpt_path,config=sd15_config, original_config=original_config_file)
|
|
except:
|
|
pipe = StableDiffusionPipeline.from_single_file(
|
|
ckpt_path, config=sd15_config,original_config_file=original_config_file)
|
|
pipe.to("cuda")
|
|
Unet= pipe.unet
|
|
vae=pipe.vae
|
|
text_encoder = pipe.text_encoder
|
|
id_encoder = ControlNetModel.from_unet(Unet)
|
|
pose_encoder = ControlNetModel.from_unet(Unet)
|
|
repo=os.path.join(makeup_current_path,"clip")
|
|
makeup_encoder = detail_encoder(Unet, clip,repo, "cuda", dtype=torch.float32)
|
|
makeup_state_dict = torch.load(makeup_encoder_path)
|
|
id_state_dict = torch.load(id_encoder_path)
|
|
|
|
id_encoder.load_state_dict(id_state_dict, strict=False)
|
|
pose_state_dict = torch.load(pose_encoder_path)
|
|
pose_encoder.load_state_dict(pose_state_dict, strict=False)
|
|
makeup_encoder.load_state_dict(makeup_state_dict, strict=False)
|
|
|
|
del id_state_dict,makeup_state_dict,pose_state_dict
|
|
torch.cuda.empty_cache()
|
|
|
|
id_encoder.to("cuda")
|
|
pose_encoder.to("cuda")
|
|
makeup_encoder.to("cuda")
|
|
if lora!="none":
|
|
lora_path = folder_paths.get_full_path("loras", lora)
|
|
pipe.load_lora_weights(lora_path, adapter_name=lora_trigger_words)
|
|
pipe.fuse_lora(adapter_names=[lora_trigger_words, ], lora_scale=lora_scale)
|
|
pipe = StableDiffusionControlNetPipeline.from_pretrained(
|
|
sd15_config,
|
|
safety_checker=None,
|
|
unet=Unet,
|
|
vae=vae,
|
|
text_encoder=text_encoder,
|
|
controlnet=[id_encoder, pose_encoder],
|
|
torch_dtype=torch.float32).to("cuda")
|
|
pipe.scheduler = scheduler_used.from_config(pipe.scheduler.config)
|
|
|
|
return ({"pipe":pipe,"makeup_encoder":makeup_encoder},)
|
|
|
|
class StableMakeup_Sampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"id_image": ("IMAGE",),
|
|
"makeup_image": ("IMAGE",),
|
|
"model": ("MAKEUP_MODEL",),
|
|
"facedetector": (["mobilenet","resnet"],),
|
|
"dataname": (["300wpublic", "300wprivate","merlrav","wflw"],),
|
|
"cfg": ("FLOAT", {"default": 1.6, "min": 0.0, "max": 30.0, "step": 0.1, "round": 0.01}),
|
|
"steps": ("INT", {"default": 30, "min": 1, "max": 10000}),
|
|
"width": ("INT", {"default": 512, "min": 256, "max": 768, "step": 64, "display": "number"}),
|
|
"height": ("INT", {"default": 512, "min": 256, "max": 768, "step": 64, "display": "number"}),
|
|
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "makeup_main"
|
|
CATEGORY = "Stable_Makeup"
|
|
|
|
|
|
def makeup_main(self, id_image, makeup_image, model,facedetector,dataname,cfg, steps,width, height ):
|
|
|
|
pipe=model.get("pipe")
|
|
makeup_encoder=model.get("makeup_encoder")
|
|
if facedetector=="mobilenet":
|
|
weight_path=os.path.join(weigths_current_path, "mobilenet0.25_Final.pth")
|
|
else:
|
|
weight_path=os.path.join(weigths_current_path, "resnet50.pth")
|
|
detector = FaceDetector(name=facedetector,weight_path=weight_path)
|
|
|
|
def get_draw(pil_img, size,dataname):
|
|
spigas = spiga_draw.spiga_process(pil_img, detector,dataname)
|
|
if spigas == False:
|
|
width, height = pil_img.size
|
|
black_image_pil = Image.new('RGB', (width, height), color=(0, 0, 0))
|
|
return black_image_pil
|
|
else:
|
|
spigas_faces = spiga_draw.spiga_segmentation(spigas, size=size)
|
|
return spigas_faces
|
|
|
|
id_image=nomarl_upscale(id_image, width, height)
|
|
makeup_image = nomarl_upscale(makeup_image, width, height)
|
|
pose_image = get_draw(id_image, size=width,dataname=dataname)
|
|
result_img = makeup_encoder.generate(id_image=[id_image, pose_image], makeup_image=makeup_image,num_inference_steps=steps,
|
|
pipe=pipe, guidance_scale=cfg)
|
|
|
|
image=pil2narry(result_img)
|
|
return (image,)
|
|
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"StableMakeup_LoadModel":StableMakeup_LoadModel,
|
|
"StableMakeup_Sampler": StableMakeup_Sampler
|
|
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"StableMakeup_LoadModel":"StableMakeup_LoadModel",
|
|
"StableMakeup_Sampler": "StableMakeup_Sampler"
|
|
}
|