89 lines
2.5 KiB
Python
89 lines
2.5 KiB
Python
import random
|
|
import torch
|
|
from .videos_nodes import check_license_shared
|
|
|
|
|
|
def _rand_hw():
|
|
return random.randint(513, 1023), random.randint(513, 1023)
|
|
|
|
|
|
def _black_image(batch: int, h: int, w: int):
|
|
return torch.zeros((batch, h, w, 3), dtype=torch.float32)
|
|
|
|
|
|
def _white_mask(batch: int, h: int, w: int):
|
|
return torch.ones((batch, h, w), dtype=torch.float32)
|
|
|
|
|
|
class VF9_SetValue:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"license_key": ("STRING", {"default": "", "multiline": False}),
|
|
"enable_background": ("BOOLEAN", {"default": True}),
|
|
"enable_mask": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"background": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
RETURN_NAMES = ("background", "mask")
|
|
FUNCTION = "run"
|
|
CATEGORY = "comfyui-sortlist"
|
|
|
|
def run(
|
|
self,
|
|
license_key: str,
|
|
enable_background: bool,
|
|
enable_mask: bool,
|
|
background=None,
|
|
mask=None,
|
|
):
|
|
|
|
batch = 1
|
|
if isinstance(background, torch.Tensor) and background.ndim == 4:
|
|
batch = int(background.shape[0])
|
|
elif isinstance(mask, torch.Tensor) and mask.ndim == 3:
|
|
batch = int(mask.shape[0])
|
|
|
|
def fb_img():
|
|
return _black_image(batch, 512, 512)
|
|
|
|
def fb_mask():
|
|
return _white_mask(batch, 512, 512)
|
|
|
|
|
|
try:
|
|
res = check_license_shared(license_key)
|
|
valid = (res.get("valid") is True)
|
|
except Exception:
|
|
valid = False
|
|
|
|
if valid:
|
|
out_bg = background if (enable_background and background is not None) else fb_img()
|
|
out_mask = mask if (enable_mask and mask is not None) else fb_mask()
|
|
return (out_bg, out_mask)
|
|
|
|
|
|
if enable_background:
|
|
h, w = _rand_hw()
|
|
out_bg = _black_image(batch, h, w)
|
|
else:
|
|
out_bg = background if background is not None else fb_img()
|
|
|
|
if enable_mask:
|
|
h, w = _rand_hw()
|
|
out_mask = _white_mask(batch, h, w) # full trắng
|
|
else:
|
|
out_mask = mask if mask is not None else fb_mask()
|
|
|
|
return (out_bg, out_mask)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"VF9_SetValue": VF9_SetValue}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"VF9_SetValue": "Set Value"}
|