36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
import torch
|
|
import gc
|
|
|
|
|
|
class Iter_ImageSwitch:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"output_index": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"input_count": ("INT", {"default": 2, "min": 1, "max": 64, "step": 1}),
|
|
"use_last_input": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "switch_image"
|
|
CATEGORY = "CustomNodes/Iterators"
|
|
|
|
def switch_image(self, output_index, input_count, use_last_input, **kwargs):
|
|
# 1. 取得目標索引（對 input_count 取模，確保不溢出）
|
|
# 這樣 index 0~n 就會嚴格對應到 image_0 ~ image_n
|
|
target_idx = output_index % input_count
|
|
|
|
# 2. 從目標索引開始「往回」搜尋有效接點
|
|
res_image = None
|
|
for i in range(target_idx, -1, -1):
|
|
val = kwargs.get(f"image_{i}")
|
|
if val is not None or not use_last_input:
|
|
res_image = val
|
|
break
|
|
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return (res_image,)
|