import os
import shutil
import numpy as np
import torchvision.transforms as transforms
import cv2
from omegaconf import OmegaConf
from torch.utils.data import DataLoader
import torch
from importlib import import_module
from .cldm.model import create_model
from .cldm.plms_hacked import PLMSSampler
from .utils.utils import *
from .utils.file_util import *
vition_path = node_path("ComfyUI_Seg_VITON")
cache_dir = os.path.join(vition_path,"cache")
model_load_path = os.path.join( vition_path,"checkpoints/VITONHD.ckpt")
yaml_path = os.path.join(vition_path,"configs/VITON512_COMFYUI.yaml")
def tensor2img_seg(x):
'''
x : [BS x c x H x W] or [c x H x W]
'''
if x.ndim == 3:
x = x.unsqueeze(0)
BS, C, H, W = x.shape
x = x.permute(0,2,3,1).reshape(-1, W, C).detach().cpu().numpy()
x = np.clip(x, -1, 1)
x = (x+1)/2
x = np.uint8(x*255.0)
if x.shape[-1] == 1:
x = np.concatenate([x,x,x], axis=-1)
return x
def imread(p, h, w, is_mask=False, in_inverse_mask=False, img=None):
if img is None:
img = cv2.imread(p)
if not is_mask:
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (w,h))
img = (img.astype(np.float32) / 127.5) - 1.0 # [-1, 1]
else:
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img = cv2.resize(img, (w,h))
img = (img >= 128).astype(np.float32) # 0 or 1
img = img[:,:,None]
if in_inverse_mask:
img = 1-img
return img
class stabel_vition:
def __init__(self):
self.model = None
self.sampler = None
@classmethod
def INPUT_TYPES(cls):
return {"required":
{
"agn":("IMAGE", {"default": "","multiline": False}),
"agn_mask":("MASK", {"default": "","multiline": False}),
"cloth":("IMAGE", {"default": "","multiline": False}),
"image":("IMAGE", {"default": "","multiline": False}),
"image_densepose":("IMAGE", {"default": "","multiline": False}),
"img_H": ("INT", {"default": 512, "min": 268, "max": 2048}),
"img_W": ("INT", {"default": 384, "min": 268, "max": 2048}),
"denoise_steps": ("INT", {"default": 20, "min": 5, "max": 200}),
"batch_size": ("INT", {"default": 16, "min": 0, "max": 32, "step": 16}),
"eta": ("INT", {"default": 0, "min": 0, "max": 200}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"cache": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}),
"repaint": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}),
}
}
RETURN_TYPES = ("IMAGE","BOOLEAN")
RETURN_NAMES = ("image","open")
OUTPUT_NODE = True
FUNCTION = "sample"
CATEGORY = "CXH"
def sample(self,agn,agn_mask,cloth,image,image_densepose,img_H,img_W,denoise_steps,batch_size,eta,seed,cache,repaint):
seed = str(seed)
img_fn = seed+"_img.jpg"
cloth_fn = seed+"_cloth.jpg"
#创建缓存文件夹 +缓存本地（待优化直接tensor转cv2）
mkdir(cache_dir)
agnostic_v3_2_dir = os.path.join(cache_dir,seed,"agnostic_v3_2")
mkdir(agnostic_v3_2_dir)
agnostic_v3_2_img_path = os.path.join(agnostic_v3_2_dir,img_fn)
save_tensor_image(agn,agnostic_v3_2_img_path)
agnostic_mask_dir = os.path.join(cache_dir,seed,"agnostic_mask")
mkdir(agnostic_mask_dir)
agnostic_mask_img_path = os.path.join(agnostic_mask_dir,img_fn)
save_tensor_image(agn_mask,agnostic_mask_img_path)
cloth_dir = os.path.join(cache_dir,seed,"cloth")
mkdir(cloth_dir)
cloth_img_path = os.path.join(cloth_dir,img_fn)
save_tensor_image(cloth,cloth_img_path)
image_dir = os.path.join(cache_dir,seed,"image")
mkdir(image_dir)
image_img_path = os.path.join(image_dir,img_fn)
save_tensor_image(image,image_img_path)
image_densepose_dir = os.path.join(cache_dir,seed,"image_densepose")
mkdir(image_densepose_dir)
image_densepose_img_path = os.path.join(image_densepose_dir,img_fn)
save_tensor_image(image_densepose,image_densepose_img_path)
agn = imread(agnostic_v3_2_img_path, img_H, img_W)
agn_mask = imread(agnostic_mask_img_path, img_H, img_W, is_mask=True, in_inverse_mask=True)
cloth = imread(cloth_img_path, img_H, img_W)
image = imread(image_img_path, img_H, img_W)
image_densepose = imread(image_densepose_img_path, img_H, img_W)
config = OmegaConf.load(yaml_path)
config.model.params.img_H = img_H
config.model.params.img_W = img_W
params = config.model.params
if self.model == None:
self.model = create_model(config_path=None, config=config)
self.model.load_state_dict(torch.load(model_load_path, map_location="cpu"))
self.model = self.model.cuda()
self.model.eval()
if self.sampler == None:
self.sampler = PLMSSampler(self.model)
dataset = getattr(import_module("comyui_dataset"), config.dataset_name)(
img_fn,
cloth_fn,
agn,
agn_mask,
cloth,
image,
image_densepose,
)
dataloader = DataLoader(dataset, num_workers=4, shuffle=False, batch_size=batch_size, pin_memory=True)
shape = (4, img_H//8, img_W//8)
x_sample_list =[]
for batch_idx, batch in enumerate(dataloader):
print(f"{batch_idx}/{len(dataloader)}")
z, c = self.model.get_input(batch, params.first_stage_key)
bs = z.shape[0]
c_crossattn = c["c_crossattn"][0][:bs]
if c_crossattn.ndim == 4:
c_crossattn = self.model.get_learned_conditioning(c_crossattn)
c["c_crossattn"] = [c_crossattn]
uc_cross = self.model.get_unconditional_conditioning(bs)
uc_full = {"c_concat": c["c_concat"], "c_crossattn": [uc_cross]}
uc_full["first_stage_cond"] = c["first_stage_cond"]
for k, v in batch.items():
if isinstance(v, torch.Tensor):
batch[k] = v.cuda()
self.sampler.model.batch = batch
ts = torch.full((1,), 999, device=z.device, dtype=torch.long)
start_code = self.model.q_sample(z, ts)
samples, _, _ = self.sampler.sample(
denoise_steps,
bs,
shape,
c,
x_T=start_code,
verbose=False,
eta=eta,
unconditional_conditioning=uc_full,
)
x_samples = self.model.decode_first_stage(samples)
for sample_idx, (x_sample, fn, cloth_fn) in enumerate(zip(x_samples, batch['img_fn'], batch["cloth_fn"])):
x_sample_img = tensor2img_seg(x_sample)
x_sample_list.append(x_sample_img)
if repaint:
repaint_agn_img = np.uint8((batch["image"][sample_idx].cpu().numpy()+1)/2 * 255) # [0,255]
repaint_agn_mask_img = batch["agn_mask"][sample_idx].cpu().numpy() # 0 or 1
x_sample_img = repaint_agn_img * repaint_agn_mask_img + x_sample_img * (1-repaint_agn_mask_img)
x_sample_img = np.uint8(x_sample_img)
to_path = os.path.join(cache_dir,seed,"result_"+str(sample_idx)+".jpg")
cv2.imwrite(to_path, x_sample_img[:,:,::-1])
if not cache:
shutil.rmtree(os.path.join(cache_dir,seed))
return pil2tensor(x_sample_list[0]),True