100 lines
2.8 KiB
Python
100 lines
2.8 KiB
Python
import gc as _iutils_gc
|
|
import torch as _iutils_torch
|
|
|
|
|
|
def _iutils_ensure_bhwc_rgb(image):
|
|
t = image
|
|
if not isinstance(t, _iutils_torch.Tensor):
|
|
t = _iutils_torch.tensor(t)
|
|
t = t.to(_iutils_torch.float32)
|
|
if t.dim() == 3:
|
|
t = t.unsqueeze(0)
|
|
# Expect [B,H,W,C]
|
|
if t.shape[-1] == 3:
|
|
return t.clamp(0.0, 1.0)
|
|
elif t.shape[-1] > 3:
|
|
return t[..., :3].clamp(0.0, 1.0)
|
|
elif t.shape[-1] == 1:
|
|
return t.repeat(1, 1, 1, 3).clamp(0.0, 1.0)
|
|
else:
|
|
return t.clamp(0.0, 1.0)
|
|
|
|
|
|
class GpuSyncCacheClear:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"image": ("IMAGE",)}}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "apply"
|
|
CATEGORY = "image/utils"
|
|
|
|
def apply(self, image):
|
|
try:
|
|
if _iutils_torch.cuda.is_available():
|
|
_iutils_torch.cuda.synchronize()
|
|
except Exception:
|
|
pass
|
|
_iutils_gc.collect()
|
|
try:
|
|
if _iutils_torch.cuda.is_available():
|
|
_iutils_torch.cuda.empty_cache()
|
|
except Exception:
|
|
pass
|
|
return (image,)
|
|
|
|
|
|
class SeedVR2Prep:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"ensure_even_dims": ("BOOLEAN", {"default": True}),
|
|
"to_fp16_on_gpu": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "apply"
|
|
CATEGORY = "image/utils"
|
|
|
|
def apply(self, image, ensure_even_dims=True, to_fp16_on_gpu=True):
|
|
x = _iutils_ensure_bhwc_rgb(image)
|
|
if ensure_even_dims:
|
|
H, W = x.shape[1], x.shape[2]
|
|
# Width
|
|
if (W % 2) == 1:
|
|
if W > 1:
|
|
x = x[:, :, : W - 1, :]
|
|
else:
|
|
# pad a column by repeating the last column
|
|
x = _iutils_torch.cat([x, x[:, :, -1:, :]], dim=2)
|
|
# Height
|
|
H, W = x.shape[1], x.shape[2]
|
|
if (H % 2) == 1:
|
|
if H > 1:
|
|
x = x[:, : H - 1, :, :]
|
|
else:
|
|
# pad a row by repeating last row
|
|
x = _iutils_torch.cat([x, x[:, -1:, :, :]], dim=1)
|
|
|
|
if to_fp16_on_gpu and _iutils_torch.cuda.is_available():
|
|
if x.device.type != "cuda":
|
|
x = x.to("cuda", non_blocking=True)
|
|
x = x.to(dtype=_iutils_torch.float16)
|
|
return (x,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GpuSyncCacheClear": GpuSyncCacheClear,
|
|
"SeedVR2Prep": SeedVR2Prep,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"GpuSyncCacheClear": "GPU Sync & Cache Clear",
|
|
"SeedVR2Prep": "SeedVR2 Prep (sanitize)",
|
|
}
|