1047 lines
37 KiB
Python
1047 lines
37 KiB
Python
import torch
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
import numpy as np
|
|
import sys, os
|
|
# facerestore
|
|
from comfy_extras.chainner_models import model_loading
|
|
from comfy import model_management
|
|
from custom_nodes.facerestore_cf.basicsr.utils.registry import ARCH_REGISTRY
|
|
import cv2
|
|
from custom_nodes.facerestore_cf.facelib.utils.face_restoration_helper import FaceRestoreHelper
|
|
from torchvision.transforms.functional import normalize
|
|
import math
|
|
# Latent
|
|
import comfy.model_management # CUDA detect
|
|
# LLM
|
|
from llama_cpp import Llama
|
|
from llama_cpp.llama_chat_format import Llava15ChatHandler
|
|
import folder_paths
|
|
from io import BytesIO
|
|
import base64
|
|
|
|
MAX_RESOLUTION = 8192
|
|
|
|
# LLM Special Load
|
|
models_base_path = os.path.join(folder_paths.models_dir, "GPTcheckpoints")
|
|
_folders_whitelist = ["moondream", "joytag"] # ,"internlm"]
|
|
|
|
|
|
def get_model_list(models_base_path, supported_gpt_extensions):
|
|
all_models = []
|
|
try:
|
|
for file in os.listdir(models_base_path):
|
|
|
|
if os.path.isdir(os.path.join(models_base_path, file)):
|
|
if file in _folders_whitelist:
|
|
all_models.append(os.path.join(models_base_path, file))
|
|
|
|
else:
|
|
if file.endswith(tuple(supported_gpt_extensions)):
|
|
all_models.append(os.path.join(models_base_path, file))
|
|
except:
|
|
print(f"Path {models_base_path} not valid.")
|
|
return all_models
|
|
|
|
|
|
def get_model_path(folder_list, model_name):
|
|
for folder_path in folder_list:
|
|
if folder_path.endswith(model_name):
|
|
return folder_path
|
|
|
|
|
|
supported_gpt_extensions = set(['.gguf'])
|
|
supported_clip_extensions = set(['.gguf', '.bin'])
|
|
model_external_path = None
|
|
all_models = []
|
|
try:
|
|
model_external_path = folder_paths.folder_names_and_paths["GPTcheckpoints"][0][0]
|
|
except:
|
|
# no external folder
|
|
pass
|
|
|
|
all_llava_models = get_model_list(os.path.join(folder_paths.models_dir, "GPTcheckpoints", "llava", "models"),
|
|
supported_gpt_extensions)
|
|
all_llava_clips = get_model_list(os.path.join(folder_paths.models_dir, "GPTcheckpoints", "llava", "clips"),
|
|
supported_clip_extensions)
|
|
|
|
all_models = get_model_list(models_base_path, supported_gpt_extensions)
|
|
if model_external_path is not None:
|
|
all_models += get_model_list(model_external_path, supported_gpt_extensions)
|
|
all_models += all_llava_models
|
|
# extract only names
|
|
all_models_names = [os.path.basename(model) for model in all_models]
|
|
|
|
all_clips_names = [os.path.basename(model) for model in all_llava_clips]
|
|
|
|
|
|
class Den_GPTLoaderSimple_llama:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{
|
|
"ckpt_name": (all_models_names,),
|
|
"clip_name": (all_clips_names,),
|
|
"gpu_layers": ("INT", {"default": 27, "min": 0, "max": 100, "step": 1}),
|
|
"n_threads": ("INT", {"default": 8, "min": 1, "max": 100, "step": 1}),
|
|
"max_ctx": ("INT", {"default": 2048, "min": 300, "max": 100000, "step": 4}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CUSTOM",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_gpt_checkpoint"
|
|
|
|
CATEGORY = "LLM"
|
|
|
|
def load_gpt_checkpoint(self, ckpt_name, clip_name, gpu_layers, n_threads, max_ctx, llava_clip=None):
|
|
ckpt_path = get_model_path(all_models, ckpt_name)
|
|
llm = None
|
|
clip_path = get_model_path(all_llava_clips, clip_name)
|
|
llava_clip = Llava15ChatHandler(clip_model_path=clip_path, verbose=False)
|
|
# if is path
|
|
if os.path.isfile(ckpt_path):
|
|
print("GPT MODEL DETECTED")
|
|
if "llava" in ckpt_path:
|
|
if llava_clip is None:
|
|
raise ValueError("Please provide a llava clip")
|
|
llm = Llama(model_path=ckpt_path, n_gpu_layers=gpu_layers, verbose=False, n_threads=n_threads,
|
|
n_ctx=max_ctx, logits_all=True, chat_handler=llava_clip)
|
|
else:
|
|
llm = Llama(model_path=ckpt_path, n_gpu_layers=gpu_layers, verbose=False, n_threads=n_threads,
|
|
n_ctx=max_ctx)
|
|
return ([llm, ckpt_name, ckpt_path],)
|
|
|
|
|
|
class Den_GPTSampler_llama:
|
|
"""
|
|
A custom node by Den for text generation using GPT
|
|
|
|
Attributes
|
|
----------
|
|
max_tokens (`int`): Maximum number of tokens in the generated text.
|
|
temperature (`float`): Temperature parameter for controlling randomness (0.2 to 1.0).
|
|
top_p (`float`): Top-p probability for nucleus sampling.
|
|
logprobs (`int`|`None`): Number of log probabilities to output alongside the generated text.
|
|
echo (`bool`): Whether to print the input prompt alongside the generated text.
|
|
stop (`str`|`List[str]`|`None`): Tokens at which to stop generation.
|
|
frequency_penalty (`float`): Frequency penalty for word repetition.
|
|
presence_penalty (`float`): Presence penalty for word diversity.
|
|
repeat_penalty (`float`): Penalty for repeating a prompt's output.
|
|
top_k (`int`): Top-k tokens to consider during generation.
|
|
stream (`bool`): Whether to generate the text in a streaming fashion.
|
|
tfs_z (`float`): Temperature scaling factor for top frequent samples.
|
|
model (`str`): The GPT model to use for text generation.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.temp_prompt = ""
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
|
|
"model": ("CUSTOM", {"default": ""}),
|
|
"max_tokens": ("INT", {"default": 2048}),
|
|
"temperature": ("FLOAT", {"default": 0.7, "min": 0.2, "max": 1.0}),
|
|
"top_p": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0}),
|
|
"logprobs": ("INT", {"default": 0}),
|
|
"echo": (["enable", "disable"], {"default": "disable"}),
|
|
"stop_token": ("STRING", {"default": "STOPTOKEN"}),
|
|
"frequency_penalty": ("FLOAT", {"default": 0.0}),
|
|
"presence_penalty": ("FLOAT", {"default": 0.0}),
|
|
"repeat_penalty": ("FLOAT", {"default": 1.17647}),
|
|
"top_k": ("INT", {"default": 40}),
|
|
"tfs_z": ("FLOAT", {"default": 1.0}),
|
|
"print_output": (["enable", "disable"], {"default": "disable"}),
|
|
"cached": (["YES", "NO"], {"default": "NO"}),
|
|
"prefix": ("STRING", {"default": "### Instruction: "}),
|
|
"suffix": ("STRING", {"default": "### Response: "}),
|
|
"max_tags": ("INT", {"default": 50}),
|
|
|
|
},
|
|
"optional": {
|
|
"prompt": ("STRING", {"forceInput": True}),
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
OUTPUT_IS_LIST = (True,)
|
|
FUNCTION = "generate_text"
|
|
CATEGORY = "LLM"
|
|
|
|
def generate_text(self, max_tokens, temperature, top_p, logprobs, echo, stop_token, frequency_penalty,
|
|
presence_penalty, repeat_penalty, top_k, tfs_z, model, print_output, cached, prefix, suffix,
|
|
max_tags, image=None, prompt=None):
|
|
model_funct = model[0]
|
|
model_name = model[1]
|
|
model_path = model[2]
|
|
|
|
if cached == "NO":
|
|
if "llava" in model_path:
|
|
print(f"Den Used Llama Model!: {prompt}\n")
|
|
cont = llava_inference(model_funct, prompt, image, max_tokens, stop_token, frequency_penalty,
|
|
presence_penalty, repeat_penalty, temperature, top_k, top_p)
|
|
self.temp_prompt = cont
|
|
else:
|
|
print(f"Den Used Wrong!: {prompt}\n")
|
|
# Call your GPT generation function here using the provided parameters
|
|
composed_prompt = f"{prefix} {prompt} {suffix}"
|
|
cont = ""
|
|
stream = model_funct(max_tokens=max_tokens, stop=[stop_token], stream=False,
|
|
frequency_penalty=frequency_penalty, presence_penalty=presence_penalty,
|
|
repeat_penalty=repeat_penalty, temperature=temperature, top_k=top_k,
|
|
top_p=top_p, model=model_path, prompt=composed_prompt)
|
|
cont = [stream["choices"][0]["text"]]
|
|
self.temp_prompt = cont
|
|
else:
|
|
cont = self.temp_prompt
|
|
# remove fist 30 characters of cont
|
|
try:
|
|
if print_output == "enable":
|
|
print(f"Input: {prompt}\nGenerated Text: {cont}")
|
|
return {"ui": {"text": cont}, "result": (cont,)}
|
|
|
|
except:
|
|
if print_output == "enable":
|
|
print(f"Input: {prompt}\nGenerated Text: ")
|
|
return {"ui": {"text": " "}, "result": (" ",)}
|
|
|
|
|
|
def tensor2pil(image):
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
|
|
# Convert PIL to Tensor
|
|
def pil2tensor(image):
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
|
|
def detect_device():
|
|
"""
|
|
Detects the appropriate device to run on, and return the device and dtype.
|
|
"""
|
|
if torch.cuda.is_available():
|
|
return torch.device("cuda"), torch.float16
|
|
elif torch.backends.mps.is_available():
|
|
return torch.device("mps"), torch.float16
|
|
else:
|
|
return torch.device("cpu"), torch.float32
|
|
|
|
|
|
def llava_inference(model_funct, prompt, images, max_tokens, stop_token, frequency_penalty, presence_penalty,
|
|
repeat_penalty, temperature, top_k, top_p):
|
|
list_descriptions = []
|
|
for image in images:
|
|
pil_image = tensor2pil(image)
|
|
# Convert the PIL image to a bytes buffer
|
|
buffer = BytesIO()
|
|
pil_image.save(buffer, format="JPEG") # You can change the format if needed
|
|
image_bytes = buffer.getvalue()
|
|
base64_string = f"data:image/jpeg;base64,{base64.b64encode(image_bytes).decode('utf-8')}"
|
|
|
|
response = model_funct.create_chat_completion(max_tokens=max_tokens, stop=[stop_token], stream=False,
|
|
frequency_penalty=frequency_penalty,
|
|
presence_penalty=presence_penalty,
|
|
repeat_penalty=repeat_penalty,
|
|
temperature=temperature, top_k=top_k, top_p=top_p,
|
|
messages=[
|
|
{"role": "system",
|
|
"content": "You are an assistant who perfectly describes images."},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "image_url",
|
|
"image_url": {"url": base64_string}},
|
|
{"type": "text", "text": prompt}
|
|
]
|
|
}
|
|
]
|
|
)
|
|
list_descriptions.append(response['choices'][0]['message']['content'])
|
|
return list_descriptions
|
|
|
|
|
|
# SVD from size Image
|
|
class Den_SVD_img2vid:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"clip_vision": ("CLIP_VISION",),
|
|
"init_image": ("IMAGE",),
|
|
"vae": ("VAE",),
|
|
"video_frames": ("INT", {"default": 14, "min": 1, "max": 4096}),
|
|
"motion_bucket_id": ("INT", {"default": 127, "min": 1, "max": 1023}),
|
|
"fps": ("INT", {"default": 6, "min": 1, "max": 1024}),
|
|
"augmentation_level": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.01})
|
|
}}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
|
|
RETURN_NAMES = ("positive", "negative", "latent")
|
|
|
|
FUNCTION = "encode"
|
|
|
|
CATEGORY = "conditioning/video_models"
|
|
|
|
def encode(self, clip_vision, init_image, vae, video_frames, motion_bucket_id, fps, augmentation_level):
|
|
height = init_image.shape[1]
|
|
width = init_image.shape[2]
|
|
output = clip_vision.encode_image(init_image)
|
|
pooled = output.image_embeds.unsqueeze(0)
|
|
pixels = comfy.utils.common_upscale(init_image.movedim(-1, 1), width, height, "bilinear", "center").movedim(1,
|
|
-1)
|
|
encode_pixels = pixels[:, :, :, :3]
|
|
if augmentation_level > 0:
|
|
encode_pixels += torch.randn_like(pixels) * augmentation_level
|
|
t = vae.encode(encode_pixels)
|
|
positive = [[pooled,
|
|
{"motion_bucket_id": motion_bucket_id, "fps": fps, "augmentation_level": augmentation_level,
|
|
"concat_latent_image": t}]]
|
|
negative = [[torch.zeros_like(pooled),
|
|
{"motion_bucket_id": motion_bucket_id, "fps": fps, "augmentation_level": augmentation_level,
|
|
"concat_latent_image": torch.zeros_like(t)}]]
|
|
latent = torch.zeros([video_frames, 4, height // 8, width // 8])
|
|
return (positive, negative, {"samples": latent})
|
|
|
|
|
|
# Latent space from size Image
|
|
class Den_ImageToLatentSpace:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return \
|
|
{"required":
|
|
{
|
|
"image": ("IMAGE",),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
|
"upscale_factor": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 8.0, "step": 0.1})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "image")
|
|
RETURN_TYPES = ("LATENT", "IMAGE")
|
|
FUNCTION = "generate"
|
|
|
|
CATEGORY = "latent"
|
|
|
|
def generate(self, image, batch_size=1, upscale_factor=1):
|
|
upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
|
|
new_image = image.movedim(-1, 1)
|
|
width = round(new_image.shape[3] * upscale_factor)
|
|
height = round(new_image.shape[2] * upscale_factor)
|
|
image2 = comfy.utils.common_upscale(new_image, width, height, upscale_methods[0], "disabled")
|
|
image = image2.movedim(1, -1)
|
|
height = image.shape[1]
|
|
width = image.shape[2]
|
|
latent = torch.zeros([batch_size, 4, height // 8, width // 8],
|
|
device=comfy.model_management.intermediate_device())
|
|
return ({"samples": latent}, image,)
|
|
|
|
|
|
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)
|
|
|
|
|
|
def tensor2img(tensor, rgb2bgr=True, out_type=np.uint8, min_max=(0, 1)):
|
|
"""Convert torch Tensors into image numpy arrays.
|
|
|
|
After clamping to [min, max], values will be normalized to [0, 1].
|
|
|
|
Args:
|
|
tensor (Tensor or list[Tensor]): Accept shapes:
|
|
1) 4D mini-batch Tensor of shape (B x 3/1 x H x W);
|
|
2) 3D Tensor of shape (3/1 x H x W);
|
|
3) 2D Tensor of shape (H x W).
|
|
Tensor channel should be in RGB order.
|
|
rgb2bgr (bool): Whether to change rgb to bgr.
|
|
out_type (numpy type): output types. If ``np.uint8``, transform outputs
|
|
to uint8 type with range [0, 255]; otherwise, float type with
|
|
range [0, 1]. Default: ``np.uint8``.
|
|
min_max (tuple[int]): min and max values for clamp.
|
|
|
|
Returns:
|
|
(Tensor or list): 3D ndarray of shape (H x W x C) OR 2D ndarray of
|
|
shape (H x W). The channel order is BGR.
|
|
"""
|
|
if not (torch.is_tensor(tensor) or (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
|
|
raise TypeError(f'tensor or list of tensors expected, got {type(tensor)}')
|
|
|
|
if torch.is_tensor(tensor):
|
|
tensor = [tensor]
|
|
result = []
|
|
for _tensor in tensor:
|
|
_tensor = _tensor.squeeze(0).float().detach().cpu().clamp_(*min_max)
|
|
_tensor = (_tensor - min_max[0]) / (min_max[1] - min_max[0])
|
|
|
|
n_dim = _tensor.dim()
|
|
if n_dim == 4:
|
|
img_np = make_grid(_tensor, nrow=int(math.sqrt(_tensor.size(0))), normalize=False).numpy()
|
|
img_np = img_np.transpose(1, 2, 0)
|
|
if rgb2bgr:
|
|
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
|
|
elif n_dim == 3:
|
|
img_np = _tensor.numpy()
|
|
img_np = img_np.transpose(1, 2, 0)
|
|
if img_np.shape[2] == 1: # gray image
|
|
img_np = np.squeeze(img_np, axis=2)
|
|
else:
|
|
if rgb2bgr:
|
|
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
|
|
elif n_dim == 2:
|
|
img_np = _tensor.numpy()
|
|
else:
|
|
raise TypeError('Only support 4D, 3D or 2D tensor. ' f'But received with dimension: {n_dim}')
|
|
if out_type == np.uint8:
|
|
# Unlike MATLAB, numpy.unit8() WILL NOT round by default.
|
|
img_np = (img_np * 255.0).round()
|
|
img_np = img_np.astype(out_type)
|
|
result.append(img_np)
|
|
if len(result) == 1:
|
|
result = result[0]
|
|
return result
|
|
|
|
|
|
class Den_FaceRestoreCFWithModel:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model_name": (folder_paths.get_filename_list("facerestore_models"),),
|
|
"image": ("IMAGE",),
|
|
"facedetection": (["retinaface_resnet50", "retinaface_mobile0.25", "YOLOv5l", "YOLOv5n"],),
|
|
"codeformer_fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1, "step": 0.05})
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
FUNCTION = "restore_face"
|
|
|
|
CATEGORY = "facerestore_cf"
|
|
|
|
def __init__(self):
|
|
self.face_helper = None
|
|
|
|
def restore_face(self, model_name, image, facedetection, codeformer_fidelity):
|
|
if "codeformer" in model_name.lower():
|
|
print(f'\tLoading CodeFormer: {model_name}')
|
|
model_path = folder_paths.get_full_path("facerestore_models", model_name)
|
|
device = model_management.get_torch_device()
|
|
codeformer_net = ARCH_REGISTRY.get("CodeFormer")(
|
|
dim_embd=512,
|
|
codebook_size=1024,
|
|
n_head=8,
|
|
n_layers=9,
|
|
connect_list=["32", "64", "128", "256"],
|
|
).to(device)
|
|
checkpoint = torch.load(model_path)["params_ema"]
|
|
codeformer_net.load_state_dict(checkpoint)
|
|
out = codeformer_net.eval()
|
|
else:
|
|
model_path = folder_paths.get_full_path("facerestore_models", model_name)
|
|
sd = comfy.utils.load_torch_file(model_path, safe_load=True)
|
|
out = model_loading.load_state_dict(sd).eval()
|
|
facerestore_model = out
|
|
print(f'\tStarting restore_face with codeformer_fidelity: {codeformer_fidelity}')
|
|
device = model_management.get_torch_device()
|
|
facerestore_model.to(device)
|
|
if self.face_helper is None:
|
|
self.face_helper = FaceRestoreHelper(1, face_size=512, crop_ratio=(1, 1), det_model=facedetection,
|
|
save_ext='png', use_parse=True, device=device)
|
|
|
|
image_np = 255. * image.cpu().numpy()
|
|
|
|
total_images = image_np.shape[0]
|
|
out_images = np.ndarray(shape=image_np.shape)
|
|
|
|
for i in range(total_images):
|
|
cur_image_np = image_np[i, :, :, ::-1]
|
|
|
|
original_resolution = cur_image_np.shape[0:2]
|
|
|
|
if facerestore_model is None or self.face_helper is None:
|
|
return image
|
|
|
|
self.face_helper.clean_all()
|
|
self.face_helper.read_image(cur_image_np)
|
|
self.face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)
|
|
self.face_helper.align_warp_face()
|
|
|
|
restored_face = None
|
|
for idx, cropped_face in enumerate(self.face_helper.cropped_faces):
|
|
cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
|
|
normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
|
|
cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
|
|
|
|
try:
|
|
with torch.no_grad():
|
|
# output = facerestore_model(cropped_face_t, w=strength, adain=True)[0]
|
|
# output = facerestore_model(cropped_face_t)[0]
|
|
output = facerestore_model(cropped_face_t, w=codeformer_fidelity)[0]
|
|
restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
|
|
del output
|
|
torch.cuda.empty_cache()
|
|
except Exception as error:
|
|
print(f'\tFailed inference for CodeFormer: {error}', file=sys.stderr)
|
|
restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))
|
|
|
|
restored_face = restored_face.astype('uint8')
|
|
self.face_helper.add_restored_face(restored_face)
|
|
|
|
self.face_helper.get_inverse_affine(None)
|
|
|
|
restored_img = self.face_helper.paste_faces_to_input_image()
|
|
restored_img = restored_img[:, :, ::-1]
|
|
|
|
if original_resolution != restored_img.shape[0:2]:
|
|
restored_img = cv2.resize(restored_img, (0, 0), fx=original_resolution[1] / restored_img.shape[1],
|
|
fy=original_resolution[0] / restored_img.shape[0],
|
|
interpolation=cv2.INTER_LINEAR)
|
|
|
|
self.face_helper.clean_all()
|
|
|
|
# restored_img = cv2.cvtColor(restored_face, cv2.COLOR_BGR2RGB)
|
|
|
|
out_images[i] = restored_img
|
|
|
|
restored_img_np = np.array(out_images).astype(np.float32) / 255.0
|
|
restored_img_tensor = torch.from_numpy(restored_img_np)
|
|
return (restored_img_tensor,)
|
|
|
|
|
|
class Den_MaskToImage_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"mask": ("MASK",), }}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def convert(self, mask):
|
|
d2, d3 = mask.size()
|
|
print("MASK SIZE:", mask.size())
|
|
|
|
new_image = torch.zeros(
|
|
(1, d2, d3, 3),
|
|
dtype=torch.float32,
|
|
)
|
|
new_image[0, :, :, 0] = mask
|
|
new_image[0, :, :, 1] = mask
|
|
new_image[0, :, :, 2] = mask
|
|
|
|
print("MaskSize", mask.size())
|
|
print("Tyep New img", type(new_image))
|
|
|
|
return (new_image,)
|
|
|
|
|
|
class Den_ImageToMask_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"image": ("IMAGE",), }}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def convert(self, image):
|
|
return (image.squeeze().mean(2),)
|
|
|
|
|
|
class Den_LatentMix_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"samples_to": ("LATENT",),
|
|
"samples_from": ("LATENT",),
|
|
"blend": ("FLOAT", {"default": 0, "min": 0, "max": 100, "step": 1}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "composite"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def composite(self, samples_to, samples_from, blend):
|
|
samples_out = samples_to.copy()
|
|
s_to = samples_to["samples"].clone()
|
|
s_from = samples_from["samples"].clone()
|
|
samples_out["samples"] = s_to * blend / 100 + s_from * (100 - blend) / 100
|
|
return (samples_out,)
|
|
|
|
|
|
class Den_LatentAdd_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"samples_to": ("LATENT",),
|
|
"samples_from": ("LATENT",),
|
|
}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "composite"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def composite(self, samples_to, samples_from):
|
|
samples_out = samples_to.copy()
|
|
s_to = samples_to["samples"].clone()
|
|
s_from = samples_from["samples"].clone()
|
|
samples_out["samples"] = (s_to + s_from)
|
|
return (samples_out,)
|
|
|
|
|
|
class Den_SaveLatent_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"latent_in": ("LATENT",), }}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "doStuff"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def doStuff(self, latent_in):
|
|
torch.save(latent_in, 'latent.pt')
|
|
return (latent_in,)
|
|
|
|
|
|
# a = torch.load("e:/portables/ComfyUI_windows_portable/latent.pt")
|
|
# for idx in range(a['samples'].shape[1]):
|
|
# plt.figure()
|
|
# plt.imshow(a['samples'][0,idx,:,:])
|
|
# plt.show()
|
|
|
|
class Den_LoadLatent_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "doStuff"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def doStuff(self, ):
|
|
latent_out = torch.load('latent.pt')
|
|
return (latent_out,)
|
|
|
|
|
|
class Den_LatentToImages_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"latent_in": ("LATENT",), }}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "doStuff"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def doStuff(self, latent_in):
|
|
s = latent_in['samples']
|
|
s = (s - s.min()) / (s.max() - s.min())
|
|
d1, d2, d3, d4 = s.shape
|
|
images_out = torch.zeros(d2, d3, d4, 3)
|
|
|
|
for idx in range(s.shape[1]):
|
|
for chan in range(3):
|
|
images_out[idx, :, :, chan] = s[0, idx, :, :]
|
|
return (images_out,)
|
|
|
|
|
|
class Den_LatentMixMasked_As:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"samples_to": ("LATENT",),
|
|
"samples_from": ("LATENT",),
|
|
"mask": ("MASK",),
|
|
}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "composite"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def composite(self, samples_to, samples_from, mask):
|
|
print(samples_to["samples"].size())
|
|
samples_out = samples_to.copy()
|
|
s_to = samples_to["samples"].clone()
|
|
s_from = samples_from["samples"].clone()
|
|
samples_out["samples"] = s_to * mask + s_from * (1 - mask)
|
|
return (samples_out,)
|
|
|
|
|
|
class Den_ImageMixMasked_As:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"image_to": ("IMAGE",),
|
|
"image_from": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "composite"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def composite(self, image_to, image_from, mask):
|
|
image_out = image_to.clone()
|
|
image_out[0, :, :, 0] = image_to[0, :, :, 0] * mask + image_from[0, :, :, 0] * (1 - mask)
|
|
image_out[0, :, :, 1] = image_to[0, :, :, 1] * mask + image_from[0, :, :, 1] * (1 - mask)
|
|
image_out[0, :, :, 2] = image_to[0, :, :, 2] * mask + image_from[0, :, :, 2] * (1 - mask)
|
|
return (image_out,)
|
|
|
|
|
|
class Den_TextToImage_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
fonts = os.listdir('C:\Windows\Fonts')
|
|
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"multiline": True}),
|
|
"font": (fonts,),
|
|
"size": ("INT", {"default": 20, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
|
|
"width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
|
|
"height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "doStuff"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def doStuff(self, text, font, size, width, height):
|
|
PIL_image = Image.new("RGB", (width, height), (0, 0, 0))
|
|
|
|
draw = ImageDraw.Draw(PIL_image)
|
|
|
|
# Set the font and size
|
|
font = ImageFont.truetype(font, size)
|
|
|
|
# Get the size of the text
|
|
text_size = draw.textsize(text, font)
|
|
|
|
# Calculate the position of the text
|
|
x = (PIL_image.width - text_size[0]) / 2
|
|
y = (PIL_image.height - text_size[1]) / 2
|
|
|
|
# Draw the text on the image
|
|
draw.text((x, y), text, font=font, fill=(255, 255, 255))
|
|
|
|
np_image = np.array(PIL_image)
|
|
|
|
new_image = torch.zeros(
|
|
(1, height, width, 3),
|
|
dtype=torch.float32,
|
|
)
|
|
new_image[0, :, :, :] = torch.from_numpy(np_image) / 256
|
|
|
|
return (new_image,)
|
|
|
|
|
|
class Den_BatchIndex_AS:
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"batch_index": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT",)
|
|
FUNCTION = "doStuff"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def doStuff(self, batch_index):
|
|
return (batch_index,)
|
|
|
|
|
|
class Den_MapRange_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"value": ("FLOAT", {"default": 0, "min": -sys.float_info.max, "max": sys.float_info.max, "step": 0.01}),
|
|
"in_0": ("FLOAT", {"default": 0, "min": -sys.float_info.max, "max": sys.float_info.max, "step": 0.01}),
|
|
"in_1": ("FLOAT", {"default": 1, "min": -sys.float_info.max, "max": sys.float_info.max, "step": 0.01}),
|
|
"out_0": ("FLOAT", {"default": 0, "min": -sys.float_info.max, "max": sys.float_info.max, "step": 0.01}),
|
|
"out_1": ("FLOAT", {"default": 1, "min": -sys.float_info.max, "max": sys.float_info.max, "step": 0.01}),
|
|
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT", "INT")
|
|
FUNCTION = "mapRange"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def mapRange(self, value, in_0, in_1, out_0, out_1):
|
|
if (in_0 == in_1):
|
|
raise ValueError("MapRange_AS: in_0 and in_1 are equal")
|
|
|
|
run_param = (value - in_0) / (in_1 - in_0)
|
|
result = out_0 + run_param * (out_1 - out_0)
|
|
return (result, round(result))
|
|
|
|
|
|
class Den_Number2Float_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"value": ("number", {}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def convert(self, value):
|
|
return (value,)
|
|
|
|
|
|
class Den_Int2Any_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"value": ("INT", {"default": 0, }),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("*",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def convert(self, value):
|
|
return (value,)
|
|
|
|
|
|
class Den_Number2Int_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"value": ("number", {}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def convert(self, value):
|
|
return (round(value),)
|
|
|
|
|
|
class Den_Eval_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"optional": {
|
|
"i1": ("INT", {"default": 0, "step": 1}),
|
|
"i2": ("INT", {"default": 0, "step": 1}),
|
|
"f1": ("FLOAT", {"default": 0.0, "step": 0.1, "round": 0.01}),
|
|
"f2": ("FLOAT", {"default": 0.0, "step": 0.1, "round": 0.01}),
|
|
"s1": ("STRING", {"multiline": True, "default": ""}),
|
|
"s2": ("STRING", {"multiline": True, "default": ""}),
|
|
},
|
|
"required": {
|
|
"int_prc": ("STRING", {"multiline": False, "default": "0"}),
|
|
"float_prc": ("STRING", {"multiline": False, "default": "0"}),
|
|
"str_prc": ("STRING", {"multiline": False, "default": "0"}),
|
|
},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("INT", "FLOAT", "STRING",)
|
|
FUNCTION = "do_stuff"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def do_stuff(self, int_prc, float_prc, str_prc, i1, i2, f1, f2, s1, s2):
|
|
int_out = eval(int_prc)
|
|
float_out = eval(float_prc)
|
|
str_out = eval(str_prc)
|
|
|
|
return (int_out, float_out, str_out)
|
|
|
|
|
|
class Den_Number_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"value": ("FLOAT", {"default": 0, "min": -sys.float_info.max, "max": sys.float_info.max, }),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("number",)
|
|
FUNCTION = "doStuff"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def doStuff(self, value):
|
|
return (value,)
|
|
|
|
|
|
class Den_Math_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"do": (["add", "subtract", "multiply", "divide", "power"],),
|
|
"in_0": ("FLOAT", {"default": 0, "step": 0.01, "round": 0.01, }),
|
|
"in_1": ("FLOAT", {"default": 0, "step": 0.01, "round": 0.01, }),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT", "INT")
|
|
FUNCTION = "calculate"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def calculate(self, do, in_0, in_1):
|
|
if do == "add":
|
|
result = in_0 + in_1
|
|
if do == "subtract":
|
|
result = in_0 - in_1
|
|
if do == "multiply":
|
|
result = in_0 * in_1
|
|
if do == "divide":
|
|
result = in_0 / in_1
|
|
if do == "power":
|
|
result = in_0 ** in_1
|
|
|
|
return (result, round(result))
|
|
|
|
|
|
class Den_Increment_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"value": ("INT", {"default": 10, "min": 0, "max": 20, "step": 3}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("number",)
|
|
FUNCTION = "doStuff"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def doStuff(self, value):
|
|
return (value,)
|
|
|
|
|
|
class Den_CropImage_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"image": ("IMAGE",),
|
|
"width": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
|
|
"height": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
|
|
"x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
|
|
"y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def process(self, image, width, height, x, y):
|
|
image_out = image.clone()
|
|
print(image.shape)
|
|
image_out = image_out[:, y:y + height, x:x + width]
|
|
return (image_out,)
|
|
|
|
|
|
class Den_TextWildcardList_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"text": ("STRING", {"default": "$list", "multiline": True}),
|
|
"strings": ("STRING", {"multiline": True}),
|
|
"idx": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "encode"
|
|
|
|
CATEGORY = "conditioning"
|
|
|
|
def encode(self, text, strings, idx):
|
|
string_list = strings.split(",")
|
|
wildcard = string_list[idx % len(string_list)].strip()
|
|
return (text.replace("$list", wildcard),)
|
|
|
|
|
|
class Den_NoiseImage_AS:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
|
|
"height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 64}),
|
|
"idx": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "doStuff"
|
|
CATEGORY = "ASNodes"
|
|
|
|
def doStuff(self, width, height, idx):
|
|
new_image = torch.rand(
|
|
(1, height, width, 3),
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
return (new_image,)
|
|
|
|
|
|
# A dictionary that contains all nodes you want to export with their names
|
|
# NOTE: names should be globally unique
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Den_GPTLoaderSimple_llama": Den_GPTLoaderSimple_llama,
|
|
"Den_GPTSampler_llama": Den_GPTSampler_llama,
|
|
"Den_SVD_img2vid": Den_SVD_img2vid,
|
|
"Den_ImageToLatentSpace": Den_ImageToLatentSpace,
|
|
"Den_FaceRestoreCFWithModel": Den_FaceRestoreCFWithModel,
|
|
"Den_MaskToImage_AS": Den_MaskToImage_AS,
|
|
"Den_ImageToMask_AS": Den_ImageToMask_AS,
|
|
"Den_LatentMix_AS": Den_LatentMix_AS,
|
|
"Den_LatentAdd_AS": Den_LatentAdd_AS,
|
|
"Den_SaveLatent_AS": Den_SaveLatent_AS,
|
|
"Den_LoadLatent_AS": Den_LoadLatent_AS,
|
|
"Den_LatentToImages_AS": Den_LatentToImages_AS,
|
|
"Den_LatentMixMasked_As": Den_LatentMixMasked_As,
|
|
"Den_ImageMixMasked_As": Den_ImageMixMasked_As,
|
|
"Den_TextToImage_AS": Den_TextToImage_AS,
|
|
# "Den_BatchIndex_AS": Den_BatchIndex_AS,
|
|
"Den_MapRange_AS": Den_MapRange_AS,
|
|
"Den_Number_AS": Den_Number_AS,
|
|
"Den_Int2Any_AS": Den_Int2Any_AS,
|
|
"Den_Number2Int_AS": Den_Number2Int_AS,
|
|
"Den_Number2Float_AS": Den_Number2Float_AS,
|
|
"Den_Math_AS": Den_Math_AS,
|
|
# "Increment_AS": Increment_AS,
|
|
"Den_CropImage_AS": Den_CropImage_AS,
|
|
"Den_TextWildcardList_AS": Den_TextWildcardList_AS,
|
|
"Den_NoiseImage_AS": Den_NoiseImage_AS,
|
|
"Den_Eval_AS": Den_Eval_AS,
|
|
}
|