174 lines
6.7 KiB
Python
174 lines
6.7 KiB
Python
# !/usr/bin/env python
|
|
# -*- 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
|
|
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"
|
|
}
|