diff --git a/Comfyui_CustomNet_node.py b/Comfyui_CustomNet_node.py index 48ec99e..111ebeb 100644 --- a/Comfyui_CustomNet_node.py +++ b/Comfyui_CustomNet_node.py @@ -2,12 +2,6 @@ # -*- coding: UTF-8 -*- import sys import os -from folder_paths import base_path -sys.path.append(os.path.join(base_path,"custom_nodes","ComfyUI_CustomNet","custom_net")) -# print(base_path) -path = os.path.dirname(os.path.abspath(__file__)) -path_dir = os.path.dirname(path) -file_path = os.path.dirname(path_dir) import numpy as np import torch from omegaconf import OmegaConf @@ -17,9 +11,13 @@ from .custom_net.ddim import DDIMSampler from einops import rearrange import math from PIL import Image, ImageDraw -from .customnet_util import instantiate_from_config, img2tensor +from .custom_net.customnet_util import instantiate_from_config, img2tensor from .gradio_utils import load_preprocess_model, preprocess_image - +from comfy.utils import common_upscale +import folder_paths +from folder_paths import base_path +sys.path.append(os.path.join(base_path,"custom_nodes","ComfyUI_CustomNet","custom_net")) +cur_path = os.path.dirname(os.path.abspath(__file__)) def get_instance_path(path): os_path = os.path.normpath(path) @@ -27,33 +25,137 @@ def get_instance_path(path): os_path = os_path.replace('\\', "/") return os_path +def tensor_to_image( tensor): + image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() + image = Image.fromarray(image_np, mode='RGB') + return image -class CustomNet_Plus: +def upscale_to_pil(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_image(samples) + return img_pil +def preprocess_input(preprocess_model, input_image): # reg background + processed_image = preprocess_image(preprocess_model, input_image) + # input_img = (processed_image / 255.0).astype(np.float32) + return processed_image + # return processed_image, processed_image + + +def prepare_data(device, input_image, x0, y0, x1, y1, polar, azimuth, prompt,bg_image,use_inpaint): + + # if input_image.size[0] != 256 or input_image.size[1] != 256: + # input_image = input_image.resize((256, 256)) + # input_image = np.array(input_image) + # img_cond = img2tensor(input_image, bgr2rgb=False, float32=True).unsqueeze(0) / 255. + input_image = np.array(input_image) + img_cond = img2tensor(input_image, bgr2rgb=False, float32=True).unsqueeze(0) / 255. + img_cond = img_cond * 2 - 1 + + img_location = copy.deepcopy(img_cond) + input_im_padding = torch.ones_like(img_location) + + x_0 = min(x0, x1) + x_1 = max(x0, x1) + y_0 = min(y0, y1) + y_1 = max(y0, y1) + + # print(x0, y0, x1, y1) + # print(x_0, y_0, x_1, y_1) + img_location = torch.nn.functional.interpolate(img_location, (y_1 - y_0, x_1 - x_0), mode="bilinear") + input_im_padding[:, :, y_0:y_1, x_0:x_1] = img_location + img_location = input_im_padding + + if use_inpaint: + bg_image = np.array(bg_image) + bg_cond=img2tensor( bg_image, bgr2rgb=False, float32=True) / 255. + bg_cond =bg_cond* 2 - 1 + bg_cond =bg_cond.unsqueeze(0) + bg_cond[:, :, y_0:y_1, x_0:x_1] = 1 + else: + bg_cond=img_cond #no use + + T = torch.tensor( + [[math.radians(polar), math.sin(math.radians(azimuth)), math.cos(math.radians(azimuth)), 0.0]]).unsqueeze(1) + + + batch = { + "image_cond": img_cond.to(device), + "image_location": img_location.to(device), + "bg_cond":bg_cond.to(device), + 'T': T.to(device), + 'text': [prompt], + } + return batch + + +device = torch.device("cuda") +class CustomNet_LoadModel: def __init__(self): pass - def tensor_to_image(self, tensor): - tensor = tensor.cpu() - image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() - image = Image.fromarray(image_np, mode='RGB') - return image - @classmethod def INPUT_TYPES(cls): return { "required": { - "input_image": ("IMAGE",), + "ckpt_name": (["none"] + folder_paths.get_filename_list("checkpoints"),), + } + } + + RETURN_TYPES = ("MODEL","DICT",) + RETURN_NAMES = ("model","info") + FUNCTION = "main_loader" + CATEGORY = "CustomNet_Plus" + + def main_loader(self,ckpt_name,): + # load model + ckpt = folder_paths.get_full_path("checkpoints", ckpt_name) + if "inpaint"in ckpt_name: + path_yaml = os.path.join(cur_path, "configs", "config_customnet_inpaint.yaml") + use_inpaint=True + else: + path_yaml = os.path.join(cur_path, "configs", "config_customnet.yaml") + use_inpaint = False + path_yaml = get_instance_path(path_yaml) + config = OmegaConf.load(path_yaml) + model = instantiate_from_config(config.model) + ckpt_load = torch.load(ckpt, map_location="cpu") + model.load_state_dict(ckpt_load, strict=False) + del ckpt_load + model = model.to(device) + info={"inpaint":use_inpaint,} + return (model,info,) + +class CustomNet_Sampler: + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model":("MODEL",), + "info":("DICT",), + "image": ("IMAGE",), "prompt": ("STRING", {"multiline": True, "default": "on the seaside"}), "neg_prompt": ("STRING", {"multiline": True,"default": ""}), + "steps": ("INT", {"default": 50, "min": 1, "max": 1024,"step": 1,"display": "number"}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - "x0": ("INT", {"default": 0, "min": 0, "max": 256, "step": 1, "display": "number"}), - "y0": ("INT", {"default": 0, "min": 0, "max": 256, "step": 1, "display": "number"}), - "x1": ("INT", {"default": 256, "min": 0, "max": 256, "step": 1, "display": "number"}), - "y1": ("INT", {"default": 256, "min": 0, "max": 256, "step": 1, "display": "number"}), + "width": ("INT", {"default": 256, "min": 128, "max": 512, "step": 64, "display": "number"}), + "height": ("INT", {"default": 256, "min": 128, "max": 512, "step": 64, "display": "number"}), + "obj_x": ("INT", {"default": 50, "min": 0, "max": 256, "step": 1, "display": "number"}), + "obj_y": ("INT", {"default": 50, "min": 0, "max": 256, "step": 1, "display": "number"}), + "bg_x": ("INT", {"default": 200, "min": 0, "max": 256, "step": 1, "display": "number"}), + "bg_y": ("INT", {"default": 200, "min": 0, "max": 256, "step": 1, "display": "number"}), "polar": ("FLOAT", {"default": 0, "min": -30, "max": 30, "step": -0.5, "display": "number"}), - "azimuth": ("FLOAT", {"default": 0, "min": -60, "max": 30, "step": -0.5, "display": "number"}) - } + "azimuth": ("FLOAT", {"default": 0, "min": -60, "max": 30, "step": -0.5, "display": "number"}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 256, "step": 1, "display": "number"}), }, + "optional": { + "bg_image": ("IMAGE",), } + } RETURN_TYPES = ("IMAGE",) @@ -61,71 +163,30 @@ class CustomNet_Plus: FUNCTION = "customnet_main" CATEGORY = "CustomNet_Plus" - def preprocess_input(self,preprocess_model, input_image): # reg background - # global input_img - processed_image = preprocess_image(preprocess_model, input_image) - # input_img = (processed_image / 255.0).astype(np.float32) - return processed_image - # return processed_image, processed_image - - def prepare_data(self,device, input_image, x0, y0, x1, y1, polar, azimuth, prompt): - if input_image.size[0] != 256 or input_image.size[1] != 256: - input_image = input_image.resize((256, 256)) - input_image = np.array(input_image) - - img_cond = img2tensor(input_image, bgr2rgb=False, float32=True).unsqueeze(0) / 255. - img_cond = img_cond * 2 - 1 - - img_location = copy.deepcopy(img_cond) - input_im_padding = torch.ones_like(img_location) - - x_0 = min(x0, x1) - x_1 = max(x0, x1) - y_0 = min(y0, y1) - y_1 = max(y0, y1) - # print(x0, y0, x1, y1) - # print(x_0, y_0, x_1, y_1) - img_location = torch.nn.functional.interpolate(img_location, (y_1 - y_0, x_1 - x_0), mode="bilinear") - input_im_padding[:, :, y_0:y_1, x_0:x_1] = img_location - img_location = input_im_padding - - T = torch.tensor( - [[math.radians(polar), math.sin(math.radians(azimuth)), math.cos(math.radians(azimuth)), 0.0]]).unsqueeze(1) - batch = { - "image_cond": img_cond.to(device), - "image_location": img_location.to(device), - 'T': T.to(device), - 'text': [prompt], - } - return batch @torch.no_grad() - def customnet_main(self,input_image,neg_prompt,x0, y0, x1, y1, polar, azimuth, prompt, seed): - # load model - device = torch.device("cuda") + def customnet_main(self,model,info,image, prompt,neg_prompt,steps,seed,width,height,obj_x, obj_y, bg_x, bg_y, polar, azimuth,batch_size,**kwargs): preprocess_model = load_preprocess_model() - - path_yaml = os.path.join(path,"configs","config_customnet.yaml") - path_yaml = get_instance_path(path_yaml) - config = OmegaConf.load(path_yaml) - config = config.model - model = instantiate_from_config(config) - path_ckpt = os.path.join(path,"pretrain","customnet_v1.pt") - path_ckpt = get_instance_path(path_ckpt) - ckpt = torch.load(path_ckpt, map_location="cpu") - model.load_state_dict(ckpt, strict=False) - del ckpt - model = model.to(device) sampler = DDIMSampler(model, device=device) - - input_image = self.tensor_to_image(input_image) # comfy img2pil_img - input_image = self.preprocess_input(preprocess_model, input_image) # reg img - + use_inpaint=info["inpaint"] + print(f"inpaint is {use_inpaint}") + input_image = upscale_to_pil(image, width, height) # comfy upscale tensor2pil + if use_inpaint: + bg_image=kwargs.get("bg_image") + bg_image=upscale_to_pil(bg_image, width, height) + else: + bg_image=input_image + + input_image = preprocess_input(preprocess_model, input_image) # using interface reg img seed_everything(seed) - batch = self.prepare_data(device, input_image, x0, y0, x1, y1, polar, azimuth, prompt) + + batch = prepare_data(device, input_image, obj_x, obj_y, bg_x, bg_y, polar, azimuth, prompt,bg_image,use_inpaint) + c = model.get_learned_conditioning(batch["image_cond"]) c = torch.cat([c, batch["T"]], dim=-1) c = model.cc_projection(c) + if use_inpaint: + bg_concat = model.encode_first_stage(batch["bg_cond"]).mode().detach() ## condition cond = {} @@ -133,20 +194,24 @@ class CustomNet_Plus: cond['c_crossattn'] = [c] text_embedding = model.text_encoder(batch["text"]) cond["c_crossattn"].append(text_embedding) + if use_inpaint: + cond['c_concat'].append(bg_concat) ## null-condition uc = {} - uc['c_concat'] = [torch.zeros(1, 4, 32, 32).to(c.device)] uc['c_crossattn'] = [torch.zeros_like(c).to(c.device)] uc_text_embedding = model.text_encoder([neg_prompt]) uc['c_crossattn'].append(uc_text_embedding) + if use_inpaint: + uc["c_concat"].append(bg_concat) + ## sample shape = [4, 32, 32] samples_latents, _ = sampler.sample( - S=50, - batch_size=1, + S=steps, + batch_size=batch_size, shape=shape, verbose=False, unconditional_guidance_scale=999, # useless @@ -165,9 +230,11 @@ class CustomNet_Plus: NODE_CLASS_MAPPINGS = { - "CustomNet_Plus": CustomNet_Plus + "CustomNet_LoadModel":CustomNet_LoadModel, + "CustomNet_Sampler": CustomNet_Sampler } NODE_DISPLAY_NAME_MAPPINGS = { - "CustomNet_Plus": "CustomNet_Plus" + "CustomNet_LoadModel":"CustomNet_LoadModel", + "CustomNet_Sampler": "CustomNet_Sampler" }