Files
Amorano-Jovi_GLSL/core/__init__.py
T
2025-01-07 23:56:10 -05:00

650 lines
24 KiB
Python

"""
Jovi_GLSL - http://www.github.com/Amorano/Jovi_GLSL
Core
"""
import os
import re
import json
from enum import Enum
from pathlib import Path
from typing import Any, List, Optional, Tuple, Union, Generator
import cv2
import torch
import numpy as np
from loguru import logger
from .. import JOV_TYPE_ANY, ROOT, \
JOVBaseNode, \
load_file
# ==============================================================================
# === SHADER LOADER ===
# ==============================================================================
ROOT_GLSL = ROOT / 'glsl'
GLSL_PROGRAMS = {
"vertex": { },
"fragment": { }
}
GLSL_PROGRAMS['vertex'].update({str(f.relative_to(ROOT_GLSL).as_posix()):
str(f) for f in Path(ROOT_GLSL).rglob('*.vert')})
USER_GLSL = ROOT / 'user'
USER_GLSL.mkdir(parents=True, exist_ok=True)
logger.debug(f"user shader folder: {USER_GLSL}")
if (USER_GLSL := os.getenv("JOV_GLSL", str(USER_GLSL))) is not None:
GLSL_PROGRAMS['vertex'].update({str(f.relative_to(USER_GLSL).as_posix()):
str(f) for f in Path(USER_GLSL).rglob('*.vert')})
GLSL_PROGRAMS['fragment'].update({str(f.relative_to(ROOT_GLSL).as_posix()):
str(f) for f in Path(ROOT_GLSL).rglob('*.frag')})
if USER_GLSL is not None:
GLSL_PROGRAMS['fragment'].update({str(f.relative_to(USER_GLSL).as_posix()):
str(f) for f in Path(USER_GLSL).rglob('*.frag')})
PROG_VERTEX = None
try:
prog = GLSL_PROGRAMS['vertex'].pop('.lib/_.vert')
PROG_VERTEX = load_file(prog)
except Exception as e:
logger.error(e)
raise Exception("failed load default vertex program .lib/_.vert")
PROG_FRAGMENT = None
try:
prog = GLSL_PROGRAMS['fragment'].pop('.lib/_.frag')
PROG_FRAGMENT = load_file(prog)
except Exception as e:
logger.error(e)
raise Exception("failed load default fragment program .lib/_.frag")
PROG_HEADER = load_file(ROOT_GLSL / '.lib/_.head')
PROG_FOOTER = load_file(ROOT_GLSL / '.lib/_.foot')
logger.info(f" vertex programs: {len(GLSL_PROGRAMS['vertex'])}")
logger.info(f"fragment programs: {len(GLSL_PROGRAMS['fragment'])}")
# ==============================================================================
# === EXCEPTION ===
# ==============================================================================
class CompileException(Exception): pass
# ==============================================================================
# === THERE CAN BE ONLY ONE ===
# ==============================================================================
class Singleton(type):
_instances = {}
def __call__(cls, *arg, **kw) -> Any:
# If the instance does not exist, create and store it
if cls not in cls._instances:
instance = super().__call__(*arg, **kw)
cls._instances[cls] = instance
return cls._instances[cls]
# ==============================================================================
# === CORE NODES ===
# ==============================================================================
class JOVBaseGLSLNode(JOVBaseNode):
NOT_IDEMPOTENT = True
RETURN_TYPES = ("IMAGE", "IMAGE", "MASK")
RETURN_NAMES = ('RGBA', 'RGB', 'MASK')
FUNCTION = "run"
@classmethod
def INPUT_TYPES(cls) -> dict:
d = super().INPUT_TYPES()
d = deep_merge(d, {
"outputs": {
0: ("IMAGE", {"tooltips":"Full channel [RGBA] image. If there is an alpha, the image will be masked out with it when using this output."}),
1: ("IMAGE", {"tooltips":"Three channel [RGB] image. There will be no alpha."}),
2: ("MASK", {"tooltips":"Single channel mask output."}),
}
})
return d
# ==============================================================================
# === CONSTANT ===
# ==============================================================================
RE_VARIABLE = re.compile(r"uniform\s+(\w+)\s+(\w+);\s*(?:\/\/\s*([^;|]*))?\s*(?:;\s*([^;|]*))?\s*(?:;\s*([^;|]*))?\s*(?:;\s*([^;|]*))?\s*(?:;\s*([^;|]*))?\s*(?:\|\s*(.*))?$", re.MULTILINE)
IMAGE_SIZE_DEFAULT: int = 512
IMAGE_SIZE_MIN: int = 64
IMAGE_SIZE_MAX: int = 8192
# ==============================================================================
# === ENUMERATION ===
# ==============================================================================
class EnumConvertType(Enum):
BOOLEAN = 1
FLOAT = 10
INT = 12
VEC2 = 20
VEC2INT = 25
VEC3 = 30
VEC3INT = 35
VEC4 = 40
VEC4INT = 45
COORD2D = 22
STRING = 0
LIST = 2
DICT = 3
IMAGE = 4
LATENT = 5
# ENUM = 6
ANY = 9
MASK = 7
# MIXLAB LAYER
LAYER = 8
PTYPE = {
'bool': EnumConvertType.BOOLEAN,
'int': EnumConvertType.INT,
'ivec2': EnumConvertType.VEC2INT,
'ivec3': EnumConvertType.VEC3INT,
'ivec4': EnumConvertType.VEC4INT,
'float': EnumConvertType.FLOAT,
'vec2': EnumConvertType.VEC2,
'vec3': EnumConvertType.VEC3,
'vec4': EnumConvertType.VEC4,
'sampler2D': EnumConvertType.IMAGE
}
# ==============================================================================
# === TYPE ===
# ==============================================================================
TYPE_fCOORD2D = Tuple[float, float]
TYPE_iRGB = Tuple[int, int, int]
TYPE_iRGBA = Tuple[int, int, int, int]
TYPE_fRGB = Tuple[float, float, float]
TYPE_fRGBA = Tuple[float, float, float, float]
TYPE_PIXEL = Union[int, float, TYPE_iRGB, TYPE_iRGBA, TYPE_fRGB, TYPE_fRGBA]
TYPE_IMAGE = Union[np.ndarray, torch.Tensor]
# want to make explicit entries; comfy only looks for single type
JOV_TYPE_COMFY = "BOOLEAN|FLOAT|INT"
JOV_TYPE_VECTOR = "VEC2|VEC3|VEC4|VEC2INT|VEC3INT|VEC4INT|COORD2D"
JOV_TYPE_NUMBER = f"{JOV_TYPE_COMFY}|{JOV_TYPE_VECTOR}"
JOV_TYPE_IMAGE = "IMAGE|MASK"
JOV_TYPE_FULL = f"{JOV_TYPE_NUMBER}|{JOV_TYPE_IMAGE}"
JOV_TYPE_COMFY = JOV_TYPE_ANY
JOV_TYPE_VECTOR = JOV_TYPE_ANY
JOV_TYPE_NUMBER = JOV_TYPE_ANY
JOV_TYPE_IMAGE = JOV_TYPE_ANY
JOV_TYPE_FULL = JOV_TYPE_ANY
# ==============================================================================
# === SUPPORT ===
# ==============================================================================
def deep_merge(d1: dict, d2: dict) -> dict:
"""
Deep merge multiple dictionaries recursively.
Args:
*dicts: Variable number of dictionaries to be merged.
Returns:
dict: Merged dictionary.
"""
for key in d2:
if key in d1:
if isinstance(d1[key], dict) and isinstance(d2[key], dict):
deep_merge(d1[key], d2[key])
else:
d1[key] = d2[key]
else:
d1[key] = d2[key]
return d1
def zip_longest_fill(*iterables: Any) -> Generator[Tuple[Any, ...], None, None]:
"""
Zip longest with fill value.
This function behaves like itertools.zip_longest, but it fills the values
of exhausted iterators with their own last values instead of None.
"""
try:
iterators = [iter(iterable) for iterable in iterables]
except Exception as e:
logger.error(iterables)
logger.error(str(e))
else:
while True:
values = [next(iterator, None) for iterator in iterators]
# Check if all iterators are exhausted
if all(value is None for value in values):
break
# Fill in the last values of exhausted iterators with their own last values
for i, _ in enumerate(iterators):
if values[i] is None:
iterator_copy = iter(iterables[i])
while True:
current_value = next(iterator_copy, None)
if current_value is None:
break
values[i] = current_value
yield tuple(values)
def parse_value(val:Any, typ:EnumConvertType, default: Any,
clip_min: Optional[float]=None, clip_max: Optional[float]=None,
zero:int=0) -> List[Any]:
"""Convert target value into the new specified type."""
if typ == EnumConvertType.ANY:
return val
if isinstance(default, torch.Tensor) and typ not in [EnumConvertType.IMAGE,
EnumConvertType.MASK,
EnumConvertType.LATENT]:
h, w = default.shape[:2]
cc = default.shape[2] if len(default.shape) > 2 else 1
default = (w, h, cc)
if val is None:
if default is None:
return None
val = default
if isinstance(val, dict):
# old index?
if '0' in val or 0 in val:
val = [val.get(i, val.get(str(i), 0)) for i in range(min(len(val), 4))]
# coord2d?
elif 'x' in val:
val = [val.get(c, 0) for c in 'xyzw']
# wacky color struct?
elif 'r' in val:
val = [val.get(c, 0) for c in 'rgba']
elif isinstance(val, torch.Tensor) and typ not in [EnumConvertType.IMAGE,
EnumConvertType.MASK,
EnumConvertType.LATENT]:
h, w = val.shape[:2]
cc = val.shape[2] if len(val.shape) > 2 else 1
val = (w, h, cc)
new_val = val
if typ in [EnumConvertType.FLOAT, EnumConvertType.INT,
EnumConvertType.VEC2, EnumConvertType.VEC2INT,
EnumConvertType.VEC3, EnumConvertType.VEC3INT,
EnumConvertType.VEC4, EnumConvertType.VEC4INT,
EnumConvertType.COORD2D]:
if not isinstance(val, (list, tuple, torch.Tensor)):
val = [val]
elif isinstance(val, (tuple, )):
val = list(val)
size = max(1, int(typ.value / 10))
new_val = []
for idx in range(size):
try:
d = default[idx] if idx < len(default) else 0
except:
try:
d = default.get(str(idx), 0)
except:
d = default
v = d if val is None else val[idx] if idx < len(val) else d
if isinstance(v, (str, )):
v = v.strip('\n').strip()
if v == '':
v = 0
try:
if typ in [EnumConvertType.FLOAT, EnumConvertType.VEC2, EnumConvertType.VEC3, EnumConvertType.VEC4]:
v = round(float(v or 0), 16)
else:
v = int(v)
if clip_min is not None:
v = max(v, clip_min)
if clip_max is not None:
v = min(v, clip_max)
except Exception as e:
logger.exception(e)
logger.error(f"Error converting value: {val} -- {v}")
v = 0
if v == 0:
v = zero
new_val.append(v)
new_val = new_val[0] if size == 1 else tuple(new_val)
elif typ == EnumConvertType.DICT:
try:
if isinstance(new_val, (str,)):
try:
new_val = json.loads(new_val)
except json.decoder.JSONDecodeError:
new_val = {}
else:
if not isinstance(new_val, (list, tuple,)):
new_val = [new_val]
new_val = {i: v for i, v in enumerate(new_val)}
except Exception as e:
logger.exception(e)
elif typ == EnumConvertType.LIST:
new_val = list(new_val)
elif typ == EnumConvertType.STRING:
if isinstance(new_val, (str, list, int, float,)):
new_val = [new_val]
new_val = ", ".join(map(str, new_val)) if not isinstance(new_val, str) else new_val
elif typ == EnumConvertType.BOOLEAN:
if isinstance(new_val, (torch.Tensor,)):
new_val = True
elif isinstance(new_val, (dict,)):
new_val = len(new_val.keys()) > 0
elif isinstance(new_val, (list, tuple,)) and len(new_val) > 0 and (nv := new_val[0]) is not None:
if isinstance(nv, (bool, str,)):
new_val = bool(nv)
elif isinstance(nv, (int, float,)):
new_val = nv > 0
elif typ == EnumConvertType.LATENT:
# covert image into latent
if isinstance(new_val, (torch.Tensor,)):
new_val = {'samples': new_val.unsqueeze(0)}
else:
# convert whatever into a latent sample...
new_val = torch.empty((4, 64, 64), dtype=torch.uint8).unsqueeze(0)
new_val = {'samples': new_val}
elif typ == EnumConvertType.IMAGE:
# covert image into image? just skip if already an image
if not isinstance(new_val, (torch.Tensor,)):
color = parse_value(new_val, EnumConvertType.VEC4INT, (0,0,0,255), 0, 255)
color = torch.tensor(color, dtype=torch.int32).tolist()
new_val = torch.empty((IMAGE_SIZE_MIN, IMAGE_SIZE_MIN, 4), dtype=torch.uint8)
new_val[0,:,:] = color[0]
new_val[1,:,:] = color[1]
new_val[2,:,:] = color[2]
new_val[3,:,:] = color[3]
elif typ == EnumConvertType.MASK:
# @TODO: FIX FOR MULTI-CHAN?
if not isinstance(new_val, (torch.Tensor,)):
color = parse_value(new_val, EnumConvertType.INT, 0, 0, 255)
color = torch.tensor(color, dtype=torch.int32).tolist()
new_val = torch.empty((IMAGE_SIZE_MIN, IMAGE_SIZE_MIN, 1), dtype=torch.uint8)
new_val[0,:,:] = color
elif issubclass(typ, Enum):
new_val = typ[val]
if typ == EnumConvertType.COORD2D:
new_val = {'x': new_val[0], 'y': new_val[1]}
return new_val
def parse_param(data:dict, key:str, typ:EnumConvertType, default: Any,
clip_min: Optional[float]=None, clip_max: Optional[float]=None,
zero:int=0) -> List[Any]:
"""Convenience because of the dictionary parameters.
Convert list of values into a list of specified type.
"""
val = data.get(key, default)
if typ == EnumConvertType.ANY:
if val is None:
val = [default]
return val
elif isinstance(val, (list,)):
val = val[0]
if isinstance(val, (str,)):
try: val = json.loads(val.replace("'", '"'))
except json.JSONDecodeError: pass
# see if we are a hacked vector blob... {0:x, 1:y, 2:z, 3:w}
elif isinstance(val, dict):
# mixlab layer?
if (image := val.get('image', None)) is not None:
ret = image
if (mask := val.get('mask', None)) is not None:
while len(mask.shape) < len(image.shape):
mask = mask.unsqueeze(-1)
ret = torch.cat((image, mask), dim=-1)
if ret.ndim > 3:
val = [t for t in ret]
elif ret.ndim == 3:
val = [v.unsqueeze(-1) for v in ret]
# vector patch....
elif 'xyzw' in val:
val = tuple(x for x in val["xyzw"])
# latents....
elif 'samples' in val:
val = tuple(x for x in val["samples"])
elif ('0' in val) or (0 in val):
val = tuple(val.get(i, val.get(str(i), 0)) for i in range(min(len(val), 4)))
elif 'x' in val and 'y' in val:
val = tuple(val.get(c, 0) for c in 'xyzw')
elif 'r' in val and 'g' in val:
val = tuple(val.get(c, 0) for c in 'rgba')
elif len(val) == 0:
val = tuple()
elif isinstance(val, (torch.Tensor,)):
# a batch of RGB(A)
if val.ndim > 3:
val = [t for t in val]
# a batch of Grayscale
else:
val = [t.unsqueeze(-1) for t in val]
elif isinstance(val, (list, tuple, set)):
if isinstance(val, (tuple, set,)):
val = list(val)
elif issubclass(type(val), (Enum,)):
val = [str(val.name)]
if not isinstance(val, (list,)):
val = [val]
return [parse_value(v, typ, default, clip_min, clip_max, zero) for v in val]
# ==============================================================================
# === IMAGE SUPPORT ===
# ==============================================================================
def cv2tensor_full(image: TYPE_IMAGE, matte:TYPE_PIXEL=(0,0,0,255)) \
-> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
rgba = image_convert(image, 4)
rgb = image_matte(rgba, matte)[...,:3]
mask = image_mask(image)
rgba = torch.from_numpy(rgba.astype(np.float32) / 255.0)
rgb = torch.from_numpy(rgb.astype(np.float32) / 255.0)
mask = torch.from_numpy(mask.astype(np.float32) / 255.0)
return rgba, rgb, mask
def tensor2cv(tensor: torch.Tensor, invert_mask:bool=True) -> TYPE_IMAGE:
"""Convert a torch Tensor to a numpy ndarray."""
if tensor.ndim > 3:
raise Exception("Tensor is batch of tensors")
if tensor.ndim < 3:
tensor = tensor.unsqueeze(-1)
if tensor.shape[2] == 1 and invert_mask:
tensor = 1. - tensor
tensor = tensor.cpu().numpy()
return np.clip(255.0 * tensor, 0, 255).astype(np.uint8)
def image_convert(image: TYPE_IMAGE, channels: int, width: int=None, height: int=None,
matte: Tuple[int, ...]=(0, 0, 0, 255)) -> TYPE_IMAGE:
"""Force image format to a specific number of channels.
Args:
image (TYPE_IMAGE): Input image.
channels (int): Desired number of channels (1, 3, or 4).
width (int): Desired width. `None` means leave unchanged.
height (int): Desired height. `None` means leave unchanged.
matte (tuple): RGBA color to use as background color for transparent areas.
Returns:
TYPE_IMAGE: Image with the specified number of channels.
"""
if image.ndim == 2:
image = np.expand_dims(image, axis=-1)
if (cc := image.shape[2]) != channels:
if cc == 1 and channels == 3:
image = np.repeat(image, 3, axis=2)
elif cc == 1 and channels == 4:
rgb = np.repeat(image, 3, axis=2)
alpha = np.full(image.shape[:2] + (1,), matte[3], dtype=image.dtype)
image = np.concatenate([rgb, alpha], axis=2)
elif cc == 3 and channels == 1:
image = np.mean(image, axis=2, keepdims=True).astype(image.dtype)
elif cc == 3 and channels == 4:
alpha = np.full(image.shape[:2] + (1,), matte[3], dtype=image.dtype)
image = np.concatenate([image, alpha], axis=2)
elif cc == 4 and channels == 1:
rgb = image[..., :3]
alpha = image[..., 3:4] / 255.0
image = (np.mean(rgb, axis=2, keepdims=True) * alpha).astype(image.dtype)
elif cc == 4 and channels == 3:
image = image[..., :3]
# Resize if width or height is specified
h, w = image.shape[:2]
new_width = width if width is not None else w
new_height = height if height is not None else h
if (new_width, new_height) != (w, h):
# Create a new image with the matte color
new_image = np.full((new_height, new_width, channels), matte[:channels], dtype=image.dtype)
paste_x = (new_width - w) // 2
paste_y = (new_height - h) // 2
new_image[paste_y:paste_y+h, paste_x:paste_x+w] = image[:h, :w]
image = new_image
return image
def image_crop(image: TYPE_IMAGE, width:int=None, height:int=None, offset:Tuple[float, float]=(0, 0)) -> TYPE_IMAGE:
h, w = image.shape[:2]
width = width if width is not None else w
height = height if height is not None else h
x, y = offset
x = max(0, min(width, x))
y = max(0, min(width, y))
x2 = max(0, min(width, x + width))
y2 = max(0, min(height, y + height))
points = [(x, y), (x2, y), (x2, y2), (x, y2)]
return image_crop_polygonal(image, points)
def image_crop_center(image: TYPE_IMAGE, width:int=None, height:int=None) -> TYPE_IMAGE:
"""Helper crop function to find the "center" of the area of interest."""
h, w = image.shape[:2]
cx = w // 2
cy = h // 2
width = w if width is None else width
height = h if height is None else height
x1 = max(0, int(cx - width // 2))
y1 = max(0, int(cy - height // 2))
x2 = min(w, int(cx + width // 2)) - 1
y2 = min(h, int(cy + height // 2)) - 1
points = [(x1, y1), (x2, y1), (x2, y2), (x1, y2)]
return image_crop_polygonal(image, points)
def image_crop_polygonal(image: TYPE_IMAGE, points: List[TYPE_fCOORD2D]) -> TYPE_IMAGE:
cc = image.shape[2] if image.ndim == 3 else 1
height, width = image.shape[:2]
point_mask = np.zeros((height, width), dtype=np.uint8)
points = np.array(points, np.int32).reshape((-1, 1, 2))
point_mask = cv2.fillPoly(point_mask, [points], 255)
x, y, w, h = cv2.boundingRect(point_mask)
cropped_image = cv2.resize(image[y:y+h, x:x+w], (w, h)).astype(np.uint8)
# Apply the mask to the cropped image
point_mask_cropped = cv2.resize(point_mask[y:y+h, x:x+w], (w, h))
if cc == 4:
mask = image_mask(image, 0)
alpha_channel = cv2.resize(mask[y:y+h, x:x+w], (w, h))
cropped_image = cv2.cvtColor(cropped_image, cv2.COLOR_BGRA2BGR)
cropped_image = cv2.bitwise_and(cropped_image, cropped_image, mask=point_mask_cropped)
return image_mask_add(cropped_image, alpha_channel)
elif cc == 1:
cropped_image = cv2.cvtColor(cropped_image, cv2.COLOR_GRAY2BGR)
cropped_image = cv2.bitwise_and(cropped_image, cropped_image, mask=point_mask_cropped)
return image_convert(cropped_image, cc)
return cv2.bitwise_and(cropped_image, cropped_image, mask=point_mask_cropped)
def image_mask(image: TYPE_IMAGE, color: TYPE_PIXEL = 255) -> TYPE_IMAGE:
"""Create a mask from the image, preserving transparency.
Args:
image (TYPE_IMAGE): Input image, assumed to be 2D or 3D (with or without alpha channel).
color (TYPE_PIXEL): Value to fill the mask (default is 255).
Returns:
TYPE_IMAGE: Mask of the image, either the alpha channel or a full mask of the given color.
"""
if image.ndim == 3 and image.shape[2] == 4:
return image[..., 3]
h, w = image.shape[:2]
return np.ones((h, w), dtype=np.uint8) * color
def image_mask_add(image:TYPE_IMAGE, mask:TYPE_IMAGE=None, alpha:float=255) -> TYPE_IMAGE:
"""Put custom mask into an image. If there is no mask, alpha is applied.
Images are expanded to 4 channels.
Existing 4 channel images with no mask input just return themselves.
"""
image = image_convert(image, 4)
mask = image_mask(image, alpha) if mask is None else image_convert(mask, 1)
image[..., 3] = mask if mask.ndim == 2 else mask[:, :, 0]
return image
def image_matte(image: TYPE_IMAGE, color: TYPE_iRGBA=(0,0,0,255), width: int=None, height: int=None) -> TYPE_IMAGE:
"""
Puts an RGBA image atop a colored matte expanding or clipping the image if requested.
Args:
image (TYPE_IMAGE): The input RGBA image.
color (TYPE_iRGBA): The color of the matte as a tuple (R, G, B, A).
width (int, optional): The width of the matte. Defaults to the image width.
height (int, optional): The height of the matte. Defaults to the image height.
Returns:
TYPE_IMAGE: Composited RGBA image on a matte with original alpha channel.
"""
#if image.ndim != 4 or image.shape[2] != 4:
# return image
# Determine the dimensions of the image and the matte
image_height, image_width = image.shape[:2]
width = width or image_width
height = height or image_height
# Create a solid matte with the specified color
matte = np.full((height, width, 4), color, dtype=np.uint8)
# Extract the alpha channel from the image
alpha = None
if image.ndim == 3 and image.shape[2] == 4:
alpha = image[:, :, 3] / 255.0
# Calculate the center position for the image on the matte
x_offset = (width - image_width) // 2
y_offset = (height - image_height) // 2
if alpha is not None:
# Place the image onto the matte using the alpha channel for blending
for c in range(0, 3):
matte[y_offset:y_offset + image_height, x_offset:x_offset + image_width, c] = \
(1 - alpha) * matte[y_offset:y_offset + image_height, x_offset:x_offset + image_width, c] + \
alpha * image[:, :, c]
# Set the alpha channel of the matte to the maximum of the matte's and the image's alpha
matte[y_offset:y_offset + image_height, x_offset:x_offset + image_width, 3] = \
np.maximum(matte[y_offset:y_offset + image_height, x_offset:x_offset + image_width, 3], image[:, :, 3])
else:
image = image[y_offset:y_offset + image_height, x_offset:x_offset + image_width, :]
return matte