82 lines
3.3 KiB
Python
82 lines
3.3 KiB
Python
import torch
|
|
|
|
class AlternatingOutputB:
|
|
"""
|
|
Implements a custom alternating output pattern:
|
|
- If first_output is "input_a": Outputs A, B, B, ..., B, A, B, B, ..., B, A, ... (A, then n-1 times B, then A, and repeat)
|
|
- If first_output is "input_b": Outputs B, A, A, ..., A, B, A, A, ..., A, B, ... (B, then n-1 times A, then B, and repeat)
|
|
Handles only IMAGE type input.
|
|
Adds integer output ports for image width and height.
|
|
Adds a first_output option to select the first image to output.
|
|
"""
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"input_a": (
|
|
"IMAGE",
|
|
{"forceInput": True},
|
|
),
|
|
"input_b": (
|
|
"IMAGE",
|
|
{"forceInput": True},
|
|
),
|
|
"n": ("INT", {"default": 3, "min": 2, "max": 10000}), # n should be at least 2
|
|
"first_output": (["input_a", "input_b"], {"default": "input_a"}),
|
|
"counter": ("INT", {"default": 0, "min": 0, "max": 10000}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "INT", "INT")
|
|
RETURN_NAMES = ("image", "width", "height", "counter")
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "alternate"
|
|
CATEGORY = "Utils"
|
|
|
|
|
|
def alternate(self, input_a, input_b, n, first_output, counter):
|
|
# 检查和处理 input_a
|
|
if torch.is_tensor(input_a) and input_a.ndim == 4:
|
|
image_a = input_a
|
|
elif isinstance(input_a, tuple) and len(input_a) > 0 and torch.is_tensor(input_a[0]) and input_a[0].ndim == 4:
|
|
image_a = input_a[0]
|
|
print(f"[AlternatingOutputB] WARNING: input_a is a tuple. Extracting the first element as the IMAGE tensor.")
|
|
else:
|
|
raise ValueError(f"[AlternatingOutputB] ERROR: input_a is not an IMAGE tensor or a valid tuple. Received type: {type(input_a)}")
|
|
|
|
# 检查和处理 input_b
|
|
if torch.is_tensor(input_b) and input_b.ndim == 4:
|
|
image_b = input_b
|
|
elif isinstance(input_b, tuple) and len(input_b) > 0 and torch.is_tensor(input_b[0]) and input_b[0].ndim == 4:
|
|
image_b = input_b[0]
|
|
print(f"[AlternatingOutputB] WARNING: input_b is a tuple. Extracting the first element as the IMAGE tensor.")
|
|
else:
|
|
raise ValueError(f"[AlternatingOutputB] ERROR: input_b is not an IMAGE tensor or a valid tuple. Received type: {type(input_b)}")
|
|
|
|
# 根据 first_output 和 counter 决定输出哪个图像
|
|
if counter == 0:
|
|
output = image_a if first_output == "input_a" else image_b
|
|
else:
|
|
cycle_position = (counter - 1) % n
|
|
if first_output == "input_a":
|
|
if cycle_position == 0:
|
|
output = image_a
|
|
else:
|
|
output = image_b
|
|
else: # first_output == "input_b"
|
|
if cycle_position == 0:
|
|
output = image_b
|
|
else:
|
|
output = image_a
|
|
|
|
# 获取图像的宽度和高度
|
|
_, height, width, _ = output.shape
|
|
|
|
# 递增 counter
|
|
counter += 1
|
|
|
|
return (output, width, height, counter)
|