Add files via upload

main

Comfyui_CustomNet_node.py

@@ -0,0 +1,171 @@
+# !/usr/bin/env python
+# -*- coding: UTF-8 -*-
+import sys
+import os
+path = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(os.path.join(path,"custom_net"))
+path_dir = os.path.dirname(path)
+file_path = os.path.dirname(path_dir)
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from pytorch_lightning import seed_everything
+import copy
+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 .gradio_utils import load_preprocess_model, preprocess_image
+
+
+def get_instance_path(path):
+    os_path = os.path.normpath(path)
+    if sys.platform.startswith('win32'):
+        os_path = os_path.replace('\\', "/")
+    return os_path
+
+
+class CustomNet_Plus:
+
+    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",),
+                "prompt": ("STRING", {"multiline": True, "default": "on the seaside"}),
+                "neg_prompt": ("STRING", {"multiline": True,"default": ""}),
+                "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"}),
+                "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"})
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("output_image", )
+    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")
+        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
+
+        seed_everything(seed)
+        batch = self.prepare_data(device, input_image, x0, y0, x1, y1, polar, azimuth, prompt)
+        c = model.get_learned_conditioning(batch["image_cond"])
+        c = torch.cat([c, batch["T"]], dim=-1)
+        c = model.cc_projection(c)
+
+        ## condition
+        cond = {}
+        cond['c_concat'] = [model.encode_first_stage((batch["image_location"])).mode().detach()]
+        cond['c_crossattn'] = [c]
+        text_embedding = model.text_encoder(batch["text"])
+        cond["c_crossattn"].append(text_embedding)
+
+        ## 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)
+
+        ## sample
+        shape = [4, 32, 32]
+        samples_latents, _ = sampler.sample(
+            S=50,
+            batch_size=1,
+            shape=shape,
+            verbose=False,
+            unconditional_guidance_scale=999,  # useless
+            conditioning=cond,
+            unconditional_conditioning=uc,
+            cfg_type=0,
+            cfg_scale_dict={"img": 0., "text": 0., "all": 3.0}
+        )
+
+        x_samples = model.decode_first_stage(samples_latents)
+        x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0).cpu().numpy()
+        x_samples = rearrange(255.0 * x_samples[0], 'c h w -> h w c').astype(np.uint8)
+        image = Image.fromarray(x_samples)
+        image = torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+        return (image,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "CustomNet_Plus": CustomNet_Plus
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "CustomNet_Plus": "CustomNet_Plus"
+}

__init__.py

@@ -0,0 +1,5 @@
+import sys
+from .custom_net.customnet import CustomNet
+from .Comfyui_CustomNet_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

customnet_util.py

@@ -0,0 +1,320 @@
+import importlib
+
+import torchvision
+from torch import optim
+
+from inspect import isfunction
+from PIL import Image, ImageDraw, ImageFont
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+import time
+import cv2
+from carvekit.api.high import HiInterface
+import PIL
+import sys
+import os
+
+path = os.path.dirname(os.path.abspath(__file__))
+current_path = os.getcwd()
+paths = os.path.join(current_path,"custom_nodes","Comfyui_CustomNet","customnet")
+sys.path.append(paths)
+
+def pil_rectangle_crop(im):
+    width, height = im.size   # Get dimensions
+    
+    if width <= height:
+        left = 0
+        right = width
+        top = (height - width)/2
+        bottom = (height + width)/2
+    else:
+        
+        top = 0
+        bottom = height
+        left = (width - height) / 2
+        bottom = (width + height) / 2
+
+    # Crop the center of the image
+    im = im.crop((left, top, right, bottom))
+    return im
+
+def add_margin(pil_img, color, size=256):
+    width, height = pil_img.size
+    result = Image.new(pil_img.mode, (size, size), color)
+    result.paste(pil_img, ((size - width) // 2, (size - height) // 2))
+    return result
+
+
+def create_carvekit_interface():
+    # Check doc strings for more information
+    interface = HiInterface(object_type="object",  # Can be "object" or "hairs-like".
+                            batch_size_seg=5,
+                            batch_size_matting=1,
+                            device='cuda' if torch.cuda.is_available() else 'cpu',
+                            seg_mask_size=640,  # Use 640 for Tracer B7 and 320 for U2Net
+                            matting_mask_size=2048,
+                            trimap_prob_threshold=231,
+                            trimap_dilation=30,
+                            trimap_erosion_iters=5,
+                            fp16=False)
+
+    return interface
+
+
+def load_and_preprocess(interface, input_im):
+    '''
+    :param input_im (PIL Image).
+    :return image (H, W, 3) array in [0, 1].
+    '''
+    # See https://github.com/Ir1d/image-background-remove-tool
+    image = input_im.convert('RGB')
+
+    image_without_background = interface([image])[0]
+    image_without_background = np.array(image_without_background)
+    est_seg = image_without_background > 127
+    image = np.array(image)
+    foreground = est_seg[:, : , -1].astype(np.bool_)
+    image[~foreground] = [255., 255., 255.]
+    x, y, w, h = cv2.boundingRect(foreground.astype(np.uint8))
+    image = image[y:y+h, x:x+w, :]
+    image = PIL.Image.fromarray(np.array(image))
+    # resize image such that long edge is 512
+    image.thumbnail([200, 200], Image.Resampling.LANCZOS)
+    image = add_margin(image, (255, 255, 255), size=256)
+    #image = np.array(image)
+    return image
+
+
+def log_txt_as_img(wh, xc, size=10):
+    # wh a tuple of (width, height)
+    # xc a list of captions to plot
+    b = len(xc)
+    txts = list()
+    for bi in range(b):
+        txt = Image.new("RGB", wh, color="white")
+        draw = ImageDraw.Draw(txt)
+        font = ImageFont.truetype('data/DejaVuSans.ttf', size=size)
+        nc = int(40 * (wh[0] / 256))
+        lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))
+
+        try:
+            draw.text((0, 0), lines, fill="black", font=font)
+        except UnicodeEncodeError:
+            print("Cant encode string for logging. Skipping.")
+
+        txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
+        txts.append(txt)
+    txts = np.stack(txts)
+    txts = torch.tensor(txts)
+    return txts
+
+
+def ismap(x):
+    if not isinstance(x, torch.Tensor):
+        return False
+    return (len(x.shape) == 4) and (x.shape[1] > 3)
+
+
+def isimage(x):
+    if not isinstance(x,torch.Tensor):
+        return False
+    return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)
+
+
+def exists(x):
+    return x is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def mean_flat(tensor):
+    """
+    https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+def count_params(model, verbose=False):
+    total_params = sum(p.numel() for p in model.parameters())
+    if verbose:
+        print(f"{model.__class__.__name__} has {total_params*1.e-6:.2f} M params.")
+    return total_params
+
+
+def instantiate_from_config(config):
+    if not "target" in config:
+        if config == '__is_first_stage__':
+            return None
+        elif config == "__is_unconditional__":
+            return None
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+def get_obj_from_str(string, reload=False):
+    module, cls = string.rsplit(".", 1)
+    print(module,cls)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+class AdamWwithEMAandWings(optim.Optimizer):
+    # credit to https://gist.github.com/crowsonkb/65f7265353f403714fce3b2595e0b298
+    def __init__(self, params, lr=1.e-3, betas=(0.9, 0.999), eps=1.e-8,  # TODO: check hyperparameters before using
+                 weight_decay=1.e-2, amsgrad=False, ema_decay=0.9999,   # ema decay to match previous code
+                 ema_power=1., param_names=()):
+        """AdamW that saves EMA versions of the parameters."""
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        if not 0.0 <= weight_decay:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+        if not 0.0 <= ema_decay <= 1.0:
+            raise ValueError("Invalid ema_decay value: {}".format(ema_decay))
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                        weight_decay=weight_decay, amsgrad=amsgrad, ema_decay=ema_decay,
+                        ema_power=ema_power, param_names=param_names)
+        super().__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('amsgrad', False)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Args:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            ema_params_with_grad = []
+            state_sums = []
+            max_exp_avg_sqs = []
+            state_steps = []
+            amsgrad = group['amsgrad']
+            beta1, beta2 = group['betas']
+            ema_decay = group['ema_decay']
+            ema_power = group['ema_power']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                params_with_grad.append(p)
+                if p.grad.is_sparse:
+                    raise RuntimeError('AdamW does not support sparse gradients')
+                grads.append(p.grad)
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
+                    # Exponential moving average of parameter values
+                    state['param_exp_avg'] = p.detach().float().clone()
+
+                exp_avgs.append(state['exp_avg'])
+                exp_avg_sqs.append(state['exp_avg_sq'])
+                ema_params_with_grad.append(state['param_exp_avg'])
+
+                if amsgrad:
+                    max_exp_avg_sqs.append(state['max_exp_avg_sq'])
+
+                # update the steps for each param group update
+                state['step'] += 1
+                # record the step after step update
+                state_steps.append(state['step'])
+
+            optim._functional.adamw(params_with_grad,
+                    grads,
+                    exp_avgs,
+                    exp_avg_sqs,
+                    max_exp_avg_sqs,
+                    state_steps,
+                    amsgrad=amsgrad,
+                    beta1=beta1,
+                    beta2=beta2,
+                    lr=group['lr'],
+                    weight_decay=group['weight_decay'],
+                    eps=group['eps'],
+                    maximize=False)
+
+            cur_ema_decay = min(ema_decay, 1 - state['step'] ** -ema_power)
+            for param, ema_param in zip(params_with_grad, ema_params_with_grad):
+                ema_param.mul_(cur_ema_decay).add_(param.float(), alpha=1 - cur_ema_decay)
+
+        return loss
+
+def get_state_dict(d):
+    return d.get('state_dict', d)
+    
+def load_state_dict(ckpt_path, location='cpu'):
+    _, extension = os.path.splitext(ckpt_path)
+    if extension.lower() == ".safetensors":
+        import safetensors.torch
+        state_dict = safetensors.torch.load_file(ckpt_path, device=location)
+    else:
+        state_dict = get_state_dict(torch.load(ckpt_path, map_location=torch.device(location)))
+    state_dict = get_state_dict(state_dict)
+    print(f'Loaded state_dict from [{ckpt_path}]')
+    return state_dict
+
+
+def img2tensor(imgs, bgr2rgb=True, float32=True):
+    """Numpy array to tensor.
+
+    Args:
+        imgs (list[ndarray] | ndarray): Input images.
+        bgr2rgb (bool): Whether to change bgr to rgb.
+        float32 (bool): Whether to change to float32.
+
+    Returns:
+        list[tensor] | tensor: Tensor images. If returned results only have
+            one element, just return tensor.
+    """
+
+    def _totensor(img, bgr2rgb, float32):
+        if img.shape[2] == 3 and bgr2rgb:
+            if img.dtype == 'float64':
+                img = img.astype('float32')
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = torch.from_numpy(img.transpose(2, 0, 1))
+        if float32:
+            img = img.float()
+        return img
+
+    if isinstance(imgs, list):
+        return [_totensor(img, bgr2rgb, float32) for img in imgs]
+    else:
+        return _totensor(imgs, bgr2rgb, float32)

gradio_utils.py

@@ -0,0 +1,13 @@
+from .customnet_util import create_carvekit_interface, load_and_preprocess
+
+def load_preprocess_model():
+    carvekit = create_carvekit_interface()
+    return carvekit
+
+def preprocess_image(models, input_im):
+    '''
+    :param input_im (PIL Image).
+    :return input_im (H, W, 3) array.
+    '''
+    input_im = load_and_preprocess(models, input_im)
+    return input_im

lr_scheduler.py

@@ -0,0 +1,98 @@
+import numpy as np
+
+
+class LambdaWarmUpCosineScheduler:
+    """
+    note: use with a base_lr of 1.0
+    """
+    def __init__(self, warm_up_steps, lr_min, lr_max, lr_start, max_decay_steps, verbosity_interval=0):
+        self.lr_warm_up_steps = warm_up_steps
+        self.lr_start = lr_start
+        self.lr_min = lr_min
+        self.lr_max = lr_max
+        self.lr_max_decay_steps = max_decay_steps
+        self.last_lr = 0.
+        self.verbosity_interval = verbosity_interval
+
+    def schedule(self, n, **kwargs):
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
+        if n < self.lr_warm_up_steps:
+            lr = (self.lr_max - self.lr_start) / self.lr_warm_up_steps * n + self.lr_start
+            self.last_lr = lr
+            return lr
+        else:
+            t = (n - self.lr_warm_up_steps) / (self.lr_max_decay_steps - self.lr_warm_up_steps)
+            t = min(t, 1.0)
+            lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
+                    1 + np.cos(t * np.pi))
+            self.last_lr = lr
+            return lr
+
+    def __call__(self, n, **kwargs):
+        return self.schedule(n,**kwargs)
+
+
+class LambdaWarmUpCosineScheduler2:
+    """
+    supports repeated iterations, configurable via lists
+    note: use with a base_lr of 1.0.
+    """
+    def __init__(self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0):
+        assert len(warm_up_steps) == len(f_min) == len(f_max) == len(f_start) == len(cycle_lengths)
+        self.lr_warm_up_steps = warm_up_steps
+        self.f_start = f_start
+        self.f_min = f_min
+        self.f_max = f_max
+        self.cycle_lengths = cycle_lengths
+        self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths))
+        self.last_f = 0.
+        self.verbosity_interval = verbosity_interval
+
+    def find_in_interval(self, n):
+        interval = 0
+        for cl in self.cum_cycles[1:]:
+            if n <= cl:
+                return interval
+            interval += 1
+
+    def schedule(self, n, **kwargs):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+                                                       f"current cycle {cycle}")
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            t = (n - self.lr_warm_up_steps[cycle]) / (self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle])
+            t = min(t, 1.0)
+            f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (
+                    1 + np.cos(t * np.pi))
+            self.last_f = f
+            return f
+
+    def __call__(self, n, **kwargs):
+        return self.schedule(n, **kwargs)
+
+
+class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2):
+
+    def schedule(self, n, **kwargs):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+                                                       f"current cycle {cycle}")
+
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (self.cycle_lengths[cycle] - n) / (self.cycle_lengths[cycle])
+            self.last_f = f
+            return f
+