wraith-executioner-year-2-comfyui-misc/py/any_split.py

"""
Any Split (misc)

LIST / BATCH 入力 * を any_01, any_02 … と length (INT) に分解する。型チェックなし。
"""

try:
    import torch
except ImportError:
    torch = None

from .utils import ByPassNamesTuple, SplitReturnTypes, any_type, get_category, get_name

MIN_DATA_OUTPUTS = 1
SEND_TO_CPU_KEY = "Send to CPU"
_LEGACY_SEND_TO_CUDA_KEY = "Send to CUDA"
_CPU_DEVICE = "cpu"


def _unwrap_singleton(value):
    """INPUT_IS_LIST 時の hidden 入力ラップを外す。"""
    if isinstance(value, list) and len(value) == 1:
        return value[0]
    return value


def _unwrap_bool(value, default=False):
    """INPUT_IS_LIST 時の BOOLEAN ウィジェット値を bool にする。"""
    value = _unwrap_singleton(value)
    if value is None:
        return default
    return bool(value)


def _send_value_to_cpu(value):
    """テンソル・.to 可能オブジェクト・入れ子 dict/list を CPU へ送る。"""
    if torch is None or value is None:
        return value
    if isinstance(value, torch.Tensor):
        return value.to(_CPU_DEVICE)
    if isinstance(value, dict):
        return {key: _send_value_to_cpu(item) for key, item in value.items()}
    if isinstance(value, list):
        return [_send_value_to_cpu(item) for item in value]
    if isinstance(value, tuple):
        return tuple(_send_value_to_cpu(item) for item in value)
    to_device = getattr(value, "to", None)
    if callable(to_device):
        try:
            return to_device(_CPU_DEVICE)
        except (RuntimeError, TypeError, AttributeError):
            return value
    return value


def _normalize_list_input(value):
    """INPUT_IS_LIST で渡る * を平坦な要素列にする。"""
    if value is None:
        return []
    if isinstance(value, (list, tuple)):
        out = []
        for item in value:
            if isinstance(item, (list, tuple)):
                out.extend(_normalize_list_input(item))
            else:
                out.append(item)
        return out
    return [value]


def _canonicalize_batch_tensor(value):
    """[B, 1, H, W] 形式の MASK バッチを ComfyUI 標準の [B, H, W] に揃える。"""
    if torch is None or not isinstance(value, torch.Tensor):
        return value
    if value.ndim == 4 and int(value.shape[1]) == 1:
        return value.squeeze(1)
    return value


def _batch_size(value):
    """BATCH 入力の要素数。判定不能なら None。"""
    value = _canonicalize_batch_tensor(value)

    if torch is not None and isinstance(value, torch.Tensor):
        if value.ndim == 2:
            return 1
        if value.ndim >= 1:
            return int(value.shape[0])

    if isinstance(value, dict):
        samples = value.get("samples")
        if torch is not None and isinstance(samples, torch.Tensor) and samples.ndim >= 1:
            return int(samples.shape[0])
        batch_index = value.get("batch_index")
        if isinstance(batch_index, list) and batch_index:
            return len(batch_index)

    return None


def _slice_batch(value, index):
    """BATCH 値の index 番目（バッチ次元 1 を維持）。"""
    value = _canonicalize_batch_tensor(value)

    if torch is not None and isinstance(value, torch.Tensor):
        if value.ndim == 2:
            return value.unsqueeze(0)
        return value[index : index + 1]

    if isinstance(value, dict) and "samples" in value:
        out = value.copy()
        samples = value["samples"]
        if torch is not None and isinstance(samples, torch.Tensor):
            out["samples"] = samples[index : index + 1].clone()
        if "noise_mask" in value and torch is not None:
            masks = value["noise_mask"]
            if isinstance(masks, torch.Tensor):
                if masks.shape[0] == 1:
                    out["noise_mask"] = masks.clone()
                elif masks.shape[0] >= samples.shape[0]:
                    out["noise_mask"] = masks[index : index + 1].clone()
        if "batch_index" in value and isinstance(value["batch_index"], list):
            out["batch_index"] = value["batch_index"][index : index + 1]
        return out

    return value


def _expand_batch_input(items):
    """LIST はそのまま。単一 BATCH 値は要素列へ展開。"""
    if not items or len(items) != 1:
        return items

    size = _batch_size(items[0])
    if size is None:
        return items
    if size <= 1:
        return [_slice_batch(items[0], 0)]
    return [_slice_batch(items[0], i) for i in range(size)]


def _linked_output_slots(prompt, unique_id):
    """ワークフロー JSON から接続済み出力スロット番号を集める。"""
    uid = str(unique_id)
    slots = []
    for node in prompt.values():
        for value in node.get("inputs", {}).values():
            if isinstance(value, list) and len(value) == 2 and str(value[0]) == uid:
                slots.append(int(value[1]))
    return slots


def _compute_split_num_data(linked_slots, min_data=MIN_DATA_OUTPUTS):
    """接続状況から data 出力本数を決める（web/logic/any-split-layout.js と同等）。"""
    if not linked_slots:
        return min_data

    slots = sorted(set(int(s) for s in linked_slots))

    if len(slots) == 1 and slots[0] == min_data:
        return min_data

    length_candidates = [slot for slot in slots if any(s < slot for s in slots)]
    length_slot = max(length_candidates) if length_candidates else -1
    data_slots = [slot for slot in slots if slot != length_slot] if length_slot >= 0 else slots

    max_data = max(data_slots) if data_slots else -1
    return min_data if max_data < 0 else max(min_data, max_data + 2)


class MiscAnySplit:
    """LIST / BATCH 入力を可変 any_* 出力に分解する。"""

    NAME = get_name("Any Split")
    CATEGORY = get_category()

    @classmethod
    def INPUT_TYPES(cls):  # pylint: disable=invalid-name
        return {
            "required": {
                "*": (any_type,),
                SEND_TO_CPU_KEY: ("BOOLEAN", {"default": False}),
            },
            "hidden": {
                "prompt": "PROMPT",
                "unique_id": "UNIQUE_ID",
            },
        }

    RETURN_TYPES = SplitReturnTypes()
    RETURN_NAMES = ByPassNamesTuple(("any_01", "length"))
    INPUT_IS_LIST = True
    FUNCTION = "split"

    def split(self, *, prompt=None, unique_id=None, **kwargs):
        prompt = _unwrap_singleton(prompt)
        unique_id = _unwrap_singleton(unique_id)
        send_to_cpu = _unwrap_bool(
            kwargs.get(SEND_TO_CPU_KEY, kwargs.get(_LEGACY_SEND_TO_CUDA_KEY, False)),
        )

        items = _expand_batch_input(_normalize_list_input(kwargs.get("*")))
        if send_to_cpu:
            items = [_send_value_to_cpu(item) for item in items]
        length = len(items)

        if not items:
            return (length,)

        linked = (
            _linked_output_slots(prompt, unique_id)
            if prompt is not None and unique_id is not None
            else []
        )
        num_data = _compute_split_num_data(linked, MIN_DATA_OUTPUTS)

        result = [items[i] if i < length else None for i in range(num_data)]
        result.append(length)
        return tuple(result)


NODE_CLASS_MAPPINGS = {
    MiscAnySplit.NAME: MiscAnySplit,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    MiscAnySplit.NAME: MiscAnySplit.NAME,
}