add
This commit is contained in:
smthemex
@@ -0,0 +1,390 @@
|
||||
# !/usr/bin/env python
|
||||
# -*- coding: UTF-8 -*-
|
||||
|
||||
import random
|
||||
import sys
|
||||
import os
|
||||
from tqdm.auto import trange
|
||||
import einops
|
||||
from . import k_diffusion as K
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from omegaconf import OmegaConf
|
||||
from PIL import Image, ImageFilter
|
||||
from torch import autocast
|
||||
import cv2
|
||||
|
||||
from .stable_diffusion.ldmx.util import instantiate_from_config
|
||||
import folder_paths
|
||||
from comfy.utils import common_upscale,ProgressBar
|
||||
|
||||
cur_path = os.path.dirname(os.path.abspath(__file__))
|
||||
MAX_SEED = np.iinfo(np.int32).max
|
||||
# device = (
|
||||
# "cuda"
|
||||
# if torch.cuda.is_available()
|
||||
# else "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
# )
|
||||
# dtype = torch.float16 if device == "cuda" else torch.float
|
||||
|
||||
def get_instance_path(path):
|
||||
instance_path = os.path.normpath(path)
|
||||
if sys.platform == 'win32':
|
||||
instance_path = instance_path.replace('\\', "/")
|
||||
return instance_path
|
||||
|
||||
def pil2tensor(pil_img):
|
||||
tensor = torch.from_numpy(np.array(pil_img).astype(np.float32) / 255.0).unsqueeze(0)
|
||||
return tensor
|
||||
|
||||
def tensor2phi(tensor):
|
||||
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
||||
pil_img = Image.fromarray(image_np, mode='RGB')
|
||||
return pil_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)
|
||||
return samples
|
||||
|
||||
def nomarl_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 = tensor2phi(samples)
|
||||
return img_pil
|
||||
|
||||
def narry_list_to_pill_ist(narry_list_in):
|
||||
for i in range(len(narry_list_in)):
|
||||
value = narry_list_in[i]
|
||||
modified_value = pil2tensor(value)
|
||||
narry_list_in[i] = modified_value
|
||||
return narry_list_in
|
||||
|
||||
def narry_list(list_in):
|
||||
for i in range(len(list_in)):
|
||||
value = list_in[i]
|
||||
modified_value = pil2tensor(value)
|
||||
list_in[i] = modified_value
|
||||
return list_in
|
||||
|
||||
def phi_list(list_in):
|
||||
for i in range(len(list_in)):
|
||||
value = list_in[i]
|
||||
list_in[i] = value
|
||||
return list_in
|
||||
|
||||
def get_local_path(file_path, model_name,model_dir):
|
||||
path = os.path.join(file_path, "models", model_name, model_dir)
|
||||
model_path = os.path.normpath(path)
|
||||
if sys.platform.startswith('win32'):
|
||||
model_path = model_path.replace('\\', "/")
|
||||
return model_path
|
||||
|
||||
def get_image_path_list(folder_name):
|
||||
image_basename_list = os.listdir(folder_name)
|
||||
image_path_list = sorted(
|
||||
[os.path.join(folder_name, basename) for basename in image_basename_list]
|
||||
)
|
||||
return image_path_list
|
||||
|
||||
class CFGDenoiser(nn.Module):
|
||||
def __init__(self, model):
|
||||
super().__init__()
|
||||
self.inner_model = model
|
||||
|
||||
def forward(self, z_0, z_1, sigma, cond, uncond, text_cfg_scale, image_cfg_scale):
|
||||
cfg_z_0 = einops.repeat(z_0, "1 ... -> n ...", n=3)
|
||||
cfg_z_1 = einops.repeat(z_1, "1 ... -> n ...", n=3)
|
||||
cfg_sigma = einops.repeat(sigma, "1 ... -> n ...", n=3)
|
||||
cfg_cond = {
|
||||
"c_crossattn": [torch.cat([cond["c_crossattn"][0], uncond["c_crossattn"][0], uncond["c_crossattn"][0]])],
|
||||
"c_concat": [torch.cat([cond["c_concat"][0], cond["c_concat"][0], uncond["c_concat"][0]])],
|
||||
}
|
||||
output_0, output_1 = self.inner_model(cfg_z_0, cfg_z_1, cfg_sigma, cond=cfg_cond)
|
||||
out_cond_0, out_img_cond_0, out_uncond_0 = output_0.chunk(3)
|
||||
out_cond_1, _, _ = output_1.chunk(3)
|
||||
del output_0, output_1
|
||||
torch.cuda.empty_cache()
|
||||
return out_uncond_0 + text_cfg_scale * (out_cond_0 - out_img_cond_0) + image_cfg_scale * (out_img_cond_0 - out_uncond_0), \
|
||||
out_cond_1
|
||||
|
||||
def load_model_from_config(config, ckpt, vae_ckpt=None, verbose=False):
|
||||
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"]
|
||||
if vae_ckpt is not None:
|
||||
print(f"Loading VAE from {vae_ckpt}")
|
||||
vae_sd = torch.load(vae_ckpt, map_location="cpu")["state_dict"]
|
||||
sd = {
|
||||
k: vae_sd[k[len("first_stage_model.") :]] if k.startswith("first_stage_model.") else v
|
||||
for k, v in sd.items()
|
||||
}
|
||||
del vae_sd
|
||||
torch.cuda.empty_cache()
|
||||
model = instantiate_from_config(config.model)
|
||||
m, u = model.load_state_dict(sd, strict=False)
|
||||
del pl_sd,sd
|
||||
torch.cuda.empty_cache()
|
||||
if len(m) > 0 and verbose:
|
||||
print("missing keys:")
|
||||
print(m)
|
||||
if len(u) > 0 and verbose:
|
||||
print("unexpected keys:")
|
||||
print(u)
|
||||
return model
|
||||
|
||||
def append_dims(x, target_dims):
|
||||
"""Appends dimensions to the end of a tensor until it has target_dims dimensions."""
|
||||
dims_to_append = target_dims - x.ndim
|
||||
if dims_to_append < 0:
|
||||
raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
|
||||
return x[(...,) + (None,) * dims_to_append]
|
||||
|
||||
class CompVisDenoiser(K.external.CompVisDenoiser):
|
||||
def __init__(self, model, quantize=False, device='cpu'):
|
||||
super().__init__(model, quantize, device)
|
||||
|
||||
def get_eps(self, *args, **kwargs):
|
||||
return self.inner_model.apply_model(*args, **kwargs)
|
||||
|
||||
def forward(self, input_0, input_1, sigma, **kwargs):
|
||||
c_out, c_in = [append_dims(x, input_0.ndim) for x in self.get_scalings(sigma)]
|
||||
#eps_0, eps_1 = self.get_eps(input_0 * c_in, input_1 * c_in, self.sigma_to_t(sigma), **kwargs)
|
||||
#eps_0, eps_1 = self.get_eps(input_0 * c_in, self.sigma_to_t(sigma).cuda(), **kwargs)
|
||||
eps_0, eps_1 = self.get_eps(input_0 * c_in, self.sigma_to_t(sigma), **kwargs)
|
||||
|
||||
return input_0 + eps_0 * c_out, eps_1
|
||||
|
||||
def to_d(x, sigma, denoised):
|
||||
"""Converts a denoiser output to a Karras ODE derivative."""
|
||||
return (x - denoised) / append_dims(sigma, x.ndim)
|
||||
|
||||
|
||||
def default_noise_sampler(x):
|
||||
return lambda sigma, sigma_next: torch.randn_like(x)
|
||||
|
||||
|
||||
def get_ancestral_step(sigma_from, sigma_to, eta=1.):
|
||||
"""Calculates the noise level (sigma_down) to step down to and the amount
|
||||
of noise to add (sigma_up) when doing an ancestral sampling step."""
|
||||
if not eta:
|
||||
return sigma_to, 0.
|
||||
sigma_up = min(sigma_to, eta * (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5)
|
||||
sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
|
||||
return sigma_down, sigma_up
|
||||
|
||||
|
||||
def decode_mask(mask, height=256, width=256):
|
||||
mask = nn.functional.interpolate(mask, size=(height, width), mode="bilinear", align_corners=False)
|
||||
mask = torch.where(mask > 0, 1, -1) # Thresholding step
|
||||
mask = torch.clamp((mask + 1.0) / 2.0, min=0.0, max=1.0)
|
||||
mask = 255.0 * rearrange(mask, "1 c h w -> h w c")
|
||||
mask = torch.cat([mask, mask, mask], dim=-1)
|
||||
mask = mask.type(torch.uint8).cpu().numpy()
|
||||
return mask
|
||||
|
||||
|
||||
def sample_euler_ancestral(model, x_0, x_1, sigmas, height, width, extra_args=None, disable=None, eta=1., s_noise=1.,
|
||||
noise_sampler=None):
|
||||
"""Ancestral sampling with Euler method steps."""
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
noise_sampler = default_noise_sampler(x_0) if noise_sampler is None else noise_sampler
|
||||
s_in = x_0.new_ones([x_0.shape[0]])
|
||||
|
||||
mask_list = []
|
||||
image_list = []
|
||||
for i in trange(len(sigmas) - 1, disable=disable):
|
||||
denoised_0, denoised_1 = model(x_0, x_1, sigmas[i] * s_in, **extra_args)
|
||||
image_list.append(denoised_0)
|
||||
|
||||
sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
|
||||
d_0 = to_d(x_0, sigmas[i], denoised_0)
|
||||
|
||||
# Euler method
|
||||
dt = sigma_down - sigmas[i]
|
||||
x_0 = x_0 + d_0 * dt
|
||||
|
||||
if sigmas[i + 1] > 0:
|
||||
x_0 = x_0 + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
|
||||
|
||||
x_1 = denoised_1
|
||||
mask_list.append(decode_mask(x_1, height, width))
|
||||
|
||||
image_list = torch.cat(image_list, dim=0)
|
||||
|
||||
return x_0, x_1, image_list, mask_list
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def generate(
|
||||
model,
|
||||
model_wrap,
|
||||
model_wrap_cfg,
|
||||
image: Image.Image,
|
||||
instruction: str,
|
||||
width: int,
|
||||
height: int,
|
||||
steps: int,
|
||||
randomize_seed: bool,
|
||||
seed: int,
|
||||
randomize_cfg: bool,
|
||||
text_cfg_scale: float,
|
||||
image_cfg_scale: float,
|
||||
rgb_mode: str,
|
||||
dtype
|
||||
):
|
||||
|
||||
null_token = model.get_learned_conditioning([""])
|
||||
seed = random.randint(0, 100000) if randomize_seed else seed
|
||||
text_cfg_scale = round(random.uniform(6.0, 9.0), ndigits=2) if randomize_cfg else text_cfg_scale
|
||||
image_cfg_scale = round(random.uniform(1.2, 1.8), ndigits=2) if randomize_cfg else image_cfg_scale
|
||||
|
||||
input_image_copy = image
|
||||
|
||||
if instruction == "":
|
||||
instruction="reflective sunglasses" # prompt为空时，返回默认prompt
|
||||
|
||||
with torch.no_grad(), autocast("cuda"), model.ema_scope():
|
||||
cond = {}
|
||||
cond["c_crossattn"] = [model.get_learned_conditioning([instruction]).to(model.device)]
|
||||
#input_image = 2 * torch.tensor(np.array(image))/ 255 - 1
|
||||
input_image = 2 * torch.tensor(np.array(image)).float() / 255 - 1
|
||||
input_image = rearrange(input_image, "h w c -> 1 c h w").to(model.device)
|
||||
cond["c_concat"] = [model.encode_first_stage(input_image).mode().to(model.device)]
|
||||
|
||||
uncond = {}
|
||||
uncond["c_crossattn"] = [null_token.to(model.device)]
|
||||
uncond["c_concat"] = [torch.zeros_like(cond["c_concat"][0])]
|
||||
|
||||
sigmas = model_wrap.get_sigmas(steps).to(model.device)
|
||||
|
||||
extra_args = {
|
||||
"cond": cond,
|
||||
"uncond": uncond,
|
||||
"text_cfg_scale": text_cfg_scale,
|
||||
"image_cfg_scale": image_cfg_scale,
|
||||
}
|
||||
torch.manual_seed(seed)
|
||||
z_0 = torch.randn_like(cond["c_concat"][0]).to(model.device,dtype=dtype) * sigmas[0]
|
||||
z_1 = torch.randn_like(cond["c_concat"][0]).to(model.device,dtype=dtype) * sigmas[0]
|
||||
|
||||
z_0, z_1, image_list, mask_list = sample_euler_ancestral(model_wrap_cfg, z_0, z_1, sigmas, height, width,
|
||||
extra_args=extra_args)
|
||||
|
||||
x_0 = model.decode_first_stage(z_0)
|
||||
|
||||
x_1 = nn.functional.interpolate(z_1, size=(height, width), mode="bilinear", align_corners=False)
|
||||
x_1 = torch.where(x_1 > 0, 1, -1) # Thresholding step
|
||||
x_0 = torch.clamp((x_0 + 1.0) / 2.0, min=0.0, max=1.0)
|
||||
x_1 = torch.clamp((x_1 + 1.0) / 2.0, min=0.0, max=1.0)
|
||||
x_0 = 255.0 * rearrange(x_0, "1 c h w -> h w c")
|
||||
x_1 = 255.0 * rearrange(x_1, "1 c h w -> h w c")
|
||||
x_1 = torch.cat([x_1, x_1, x_1], dim=-1)
|
||||
edited_image = Image.fromarray(x_0.type(torch.uint8).cpu().numpy())
|
||||
edited_mask = Image.fromarray(x_1.type(torch.uint8).cpu().numpy())
|
||||
|
||||
|
||||
# 对edited_mask做膨胀
|
||||
edited_mask_copy = edited_mask.copy()
|
||||
kernel = np.ones((3, 3), np.uint8)
|
||||
edited_mask = cv2.dilate(np.array(edited_mask), kernel, iterations=3)
|
||||
edited_mask = Image.fromarray(edited_mask)
|
||||
|
||||
m_img = edited_mask.filter(ImageFilter.GaussianBlur(radius=3))
|
||||
m_img = np.asarray(m_img).astype('float') / 255.0
|
||||
img_np = np.asarray(input_image_copy).astype('float') / 255.0
|
||||
ours_np = np.asarray(edited_image).astype('float') / 255.0
|
||||
#print(ours_np)
|
||||
mix_image_np = m_img * ours_np + (1 - m_img) * img_np
|
||||
mix_image = Image.fromarray((mix_image_np * 255).astype(np.uint8)).convert('RGB')
|
||||
|
||||
rgb_choice = np.array(mix_image).astype('float') * 1
|
||||
if rgb_mode=="red":
|
||||
rgb_choice[:, :, 0] = 180.0
|
||||
rgb_choice[:, :, 2] = 0
|
||||
rgb_choice[:, :, 1] = 0
|
||||
elif rgb_mode=="blue":
|
||||
rgb_choice[:, :, 0] = 0
|
||||
rgb_choice[:, :, 2] = 180.0
|
||||
rgb_choice[:, :, 1] = 0
|
||||
else:# green
|
||||
rgb_choice[:, :, 0] = 0
|
||||
rgb_choice[:, :, 2] = 0
|
||||
rgb_choice[:, :, 1] = 180.0
|
||||
|
||||
|
||||
mix_result_with_rgb_mask = np.array(mix_image)
|
||||
mix_result_with_rgb_mask = Image.fromarray(
|
||||
(mix_result_with_rgb_mask.astype('float') * (1 - m_img.astype('float') / 2.0) +
|
||||
m_img.astype('float') / 2.0 * rgb_choice).astype('uint8'))
|
||||
|
||||
mask_img=Image.fromarray((m_img * 255).astype(np.uint8)).convert('RGB')
|
||||
|
||||
return mix_image, mask_img,mix_result_with_rgb_mask
|
||||
|
||||
|
||||
def images_generator(img_list: list, ):
|
||||
# get img size
|
||||
sizes = {}
|
||||
for image_ in img_list:
|
||||
if isinstance(image_, Image.Image):
|
||||
count = sizes.get(image_.size, 0)
|
||||
sizes[image_.size] = count + 1
|
||||
elif isinstance(image_, np.ndarray):
|
||||
count = sizes.get(image_.shape[:2][::-1], 0)
|
||||
sizes[image_.shape[:2][::-1]] = count + 1
|
||||
else:
|
||||
raise "unsupport image list,must be pil or cv2!!!"
|
||||
size = max(sizes.items(), key=lambda x: x[1])[0]
|
||||
yield size[0], size[1]
|
||||
|
||||
# any to tensor
|
||||
def load_image(img_in):
|
||||
if isinstance(img_in, Image.Image):
|
||||
img_in = img_in.convert("RGB")
|
||||
i = np.array(img_in, dtype=np.float32)
|
||||
i = torch.from_numpy(i).div_(255)
|
||||
if i.shape[0] != size[1] or i.shape[1] != size[0]:
|
||||
i = torch.from_numpy(i).movedim(-1, 0).unsqueeze(0)
|
||||
i = common_upscale(i, size[0], size[1], "lanczos", "center")
|
||||
i = i.squeeze(0).movedim(0, -1).numpy()
|
||||
return i
|
||||
elif isinstance(img_in, np.ndarray):
|
||||
i = cv2.cvtColor(img_in, cv2.COLOR_BGR2RGB).astype(np.float32)
|
||||
i = torch.from_numpy(i).div_(255)
|
||||
# print(i.shape)
|
||||
return i
|
||||
else:
|
||||
raise "unsupport image list,must be pil,cv2 or tensor!!!"
|
||||
|
||||
total_images = len(img_list)
|
||||
processed_images = 0
|
||||
pbar = ProgressBar(total_images)
|
||||
images = map(load_image, img_list)
|
||||
try:
|
||||
prev_image = next(images)
|
||||
while True:
|
||||
next_image = next(images)
|
||||
yield prev_image
|
||||
processed_images += 1
|
||||
pbar.update_absolute(processed_images, total_images)
|
||||
prev_image = next_image
|
||||
except StopIteration:
|
||||
pass
|
||||
if prev_image is not None:
|
||||
yield prev_image
|
||||
|
||||
|
||||
def load_images(img_list: list, ):
|
||||
gen = images_generator(img_list)
|
||||
(width, height) = next(gen)
|
||||
images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height, width, 3)))))
|
||||
if len(images) == 0:
|
||||
raise FileNotFoundError(f"No images could be loaded .")
|
||||
return images
|
||||
Reference in New Issue
Block a user