Add files via upload
update
This commit is contained in:
smthemex
@@ -0,0 +1,177 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
import einops
|
||||
import torch
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
import cv2
|
||||
from folder_paths import base_path
|
||||
sys.path.append(os.path.join(base_path,"custom_nodes","ComfyUI_CustomNet"))
|
||||
try:
|
||||
from pytorch_lightning.utilities.rank_zero import rank_zero_only
|
||||
except:
|
||||
try:
|
||||
from pytorch_lightning.utilities.distributed import rank_zero_only
|
||||
except:
|
||||
raise "import pytorch_lightning rank_zero_only error"
|
||||
import numpy as np
|
||||
from torch.optim.lr_scheduler import LambdaLR
|
||||
|
||||
from .util import (
|
||||
conv_nd,
|
||||
linear,
|
||||
zero_module,
|
||||
timestep_embedding,
|
||||
)
|
||||
|
||||
from einops import rearrange, repeat
|
||||
from torchvision.utils import make_grid
|
||||
from .attention import SpatialTransformer
|
||||
from .openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample, AttentionBlock
|
||||
from .ddpm import LatentDiffusion
|
||||
from .customnet_util import log_txt_as_img, exists, instantiate_from_config
|
||||
|
||||
from .dddim import DDIMSampler
|
||||
from .customnet_util import load_state_dict
|
||||
|
||||
|
||||
class CustomNet(LatentDiffusion):
|
||||
def __init__(self,
|
||||
text_encoder_config,
|
||||
sd_15_ckpt=None,
|
||||
use_cond_concat=False,
|
||||
use_bbox_mask=False,
|
||||
use_bg_inpainting=False,
|
||||
learning_rate_scale=10,
|
||||
*args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
|
||||
self.text_encoder = instantiate_from_config(text_encoder_config)
|
||||
|
||||
if sd_15_ckpt is not None:
|
||||
self.load_model_from_ckpt(ckpt=sd_15_ckpt)
|
||||
|
||||
self.use_cond_concat = use_cond_concat
|
||||
self.use_bbox_mask = use_bbox_mask
|
||||
self.use_bg_inpainting = use_bg_inpainting
|
||||
self.learning_rate_scale = learning_rate_scale
|
||||
|
||||
|
||||
def load_model_from_ckpt(self, ckpt, verbose=True):
|
||||
print(" =========================== init Stable Diffusion pretrained checkpoint =========================== ")
|
||||
print(f"Loading model from {ckpt}")
|
||||
pl_sd = torch.load(ckpt, map_location="cpu")
|
||||
if "global_step" in pl_sd:
|
||||
print(f"Global Step: {pl_sd['global_step']}")
|
||||
sd = pl_sd["state_dict"]
|
||||
sd_keys = sd.keys()
|
||||
|
||||
|
||||
missing = []
|
||||
text_encoder_sd = self.text_encoder.state_dict()
|
||||
for k in text_encoder_sd.keys():
|
||||
sd_k = "cond_stage_model."+ k
|
||||
if sd_k in sd_keys:
|
||||
text_encoder_sd[k] = sd[sd_k]
|
||||
else:
|
||||
missing.append(k)
|
||||
|
||||
self.text_encoder.load_state_dict(text_encoder_sd)
|
||||
|
||||
|
||||
|
||||
|
||||
def configure_optimizers(self):
|
||||
lr = self.learning_rate
|
||||
params = []
|
||||
params += list(self.cc_projection.parameters())
|
||||
|
||||
|
||||
params_dualattn = []
|
||||
for k, v in self.model.named_parameters():
|
||||
if "to_k_text" in k or "to_v_text" in k:
|
||||
params_dualattn.append(v)
|
||||
print("training weight: ", k)
|
||||
else:
|
||||
params.append(v)
|
||||
|
||||
|
||||
opt = torch.optim.AdamW([
|
||||
{'params':params_dualattn, 'lr': lr*self.learning_rate_scale},
|
||||
{'params': params, 'lr': lr}
|
||||
])
|
||||
|
||||
|
||||
if self.use_scheduler:
|
||||
assert 'target' in self.scheduler_config
|
||||
scheduler = instantiate_from_config(self.scheduler_config)
|
||||
|
||||
print("Setting up LambdaLR scheduler...")
|
||||
scheduler = [
|
||||
{
|
||||
'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule),
|
||||
'interval': 'step',
|
||||
'frequency': 1
|
||||
}]
|
||||
return [opt], scheduler
|
||||
return opt
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
loss, loss_dict = self.shared_step(batch)
|
||||
|
||||
self.log_dict(loss_dict, prog_bar=True,
|
||||
logger=True, on_step=True, on_epoch=True)
|
||||
|
||||
self.log("global_step", self.global_step,
|
||||
prog_bar=True, logger=True, on_step=True, on_epoch=False)
|
||||
|
||||
if self.use_scheduler:
|
||||
lr = self.optimizers().param_groups[0]['lr']
|
||||
self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
|
||||
|
||||
def shared_step(self, batch, **kwargs):
|
||||
|
||||
if 'txt' in self.ucg_training:
|
||||
k = 'txt'
|
||||
p = self.ucg_training[k]
|
||||
for i in range(len(batch[k])):
|
||||
if self.ucg_prng.choice(2, p=[1 - p, p]):
|
||||
if isinstance(batch[k], list):
|
||||
batch[k][i] = ""
|
||||
|
||||
with torch.no_grad():
|
||||
text = batch['txt']
|
||||
text_embedding = self.text_encoder(text)
|
||||
|
||||
|
||||
x, c = self.get_input(batch, self.first_stage_key)
|
||||
|
||||
c["c_crossattn"].append(text_embedding)
|
||||
loss = self(x, c,)
|
||||
return loss
|
||||
|
||||
|
||||
def apply_model(self, x_noisy, t, cond, return_ids=False, **kwargs):
|
||||
|
||||
if isinstance(cond, dict):
|
||||
# hybrid case, cond is exptected to be a dict
|
||||
pass
|
||||
else:
|
||||
if not isinstance(cond, list):
|
||||
cond = [cond]
|
||||
key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn'
|
||||
cond = {key: cond}
|
||||
|
||||
x_recon = self.model(x_noisy, t, **cond)
|
||||
|
||||
if isinstance(x_recon, tuple) and not return_ids:
|
||||
return x_recon[0]
|
||||
else:
|
||||
return x_recon
|
||||
@@ -0,0 +1,321 @@
|
||||
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__))
|
||||
from folder_paths import base_path
|
||||
paths = os.path.join(base_path,"custom_nodes","ComfyUI_CustomNet")
|
||||
# print(paths)
|
||||
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)
|
||||
Reference in New Issue
Block a user