68 lines
2.5 KiB
Python
68 lines
2.5 KiB
Python
import torch
|
|
import numpy as np
|
|
|
|
def _as_batched_tensor(imgs):
|
|
if isinstance(imgs, torch.Tensor):
|
|
return imgs if imgs.ndim == 4 else imgs.unsqueeze(0)
|
|
if isinstance(imgs, (list, tuple)):
|
|
ts = []
|
|
for im in imgs:
|
|
if not isinstance(im, torch.Tensor):
|
|
# convert numpy HWC uint8 or float to torch float32 0..1
|
|
if isinstance(im, np.ndarray):
|
|
arr = im.astype(np.float32)
|
|
if arr.ndim == 4 and arr.shape[0] == 1:
|
|
arr = arr[0]
|
|
if arr.ndim != 3 or arr.shape[-1] < 3:
|
|
raise ValueError("Unsupported numpy image shape")
|
|
if arr.dtype != np.float32:
|
|
arr = arr.astype(np.float32)
|
|
if arr.max() > 1.0: arr = arr / 255.0
|
|
im = torch.from_numpy(arr)
|
|
else:
|
|
raise ValueError("Unsupported IMAGE element type")
|
|
ts.append(im if im.ndim == 3 else im.squeeze(0))
|
|
return torch.stack(ts, dim=0) if ts else torch.zeros((0, 64, 64, 3), dtype=torch.float32)
|
|
raise ValueError("Unsupported IMAGE container")
|
|
|
|
def _is_black_or_empty(img_tensor):
|
|
if not isinstance(img_tensor, torch.Tensor): return True
|
|
if img_tensor.numel() == 0: return True
|
|
return bool(torch.all(img_tensor == 0))
|
|
|
|
class MergeImageBatches:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"batch1": ("IMAGE",),
|
|
"batch2": ("IMAGE",),
|
|
"prefer": (["batch1", "batch2"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("MERGED",)
|
|
FUNCTION = "merge"
|
|
CATEGORY = "image/util"
|
|
|
|
def merge(self, batch1, batch2, prefer):
|
|
b1 = _as_batched_tensor(batch1)
|
|
b2 = _as_batched_tensor(batch2)
|
|
n = min(int(b1.shape[0]), int(b2.shape[0]))
|
|
items = []
|
|
for i in range(n):
|
|
a = b1[i]; b = b2[i]
|
|
a_bad = _is_black_or_empty(a)
|
|
b_bad = _is_black_or_empty(b)
|
|
if a_bad and not b_bad:
|
|
items.append(b)
|
|
elif b_bad and not a_bad:
|
|
items.append(a)
|
|
else:
|
|
items.append(a if prefer == "batch1" else b)
|
|
return (torch.stack(items, dim=0) if items else b1[0:0],)
|
|
|
|
NODE_CLASS_MAPPINGS = {"MergeImageBatches": MergeImageBatches}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"MergeImageBatches": "Merge image batches"}
|