41 lines
1.1 KiB
Python
41 lines
1.1 KiB
Python
import torch
|
|
import gc
|
|
|
|
class Iter_ModelSwitch:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@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": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "switch_model"
|
|
CATEGORY = "CustomNodes/Iterators"
|
|
|
|
def switch_model(self, output_index, input_count, use_last_input, **kwargs):
|
|
# 1. 確定目標索引 (嚴格對應編號)
|
|
target_idx = output_index % input_count
|
|
|
|
# 2. 往回搜尋有效模型
|
|
res_model = None
|
|
for i in range(target_idx, -1, -1):
|
|
val = kwargs.get(f"model_{i}")
|
|
if val is not None or not use_last_input:
|
|
res_model = val
|
|
break
|
|
|
|
|
|
# 4. VRAM 管理
|
|
#if force_gc == "enable":
|
|
#gc.collect()
|
|
#torch.cuda.empty_cache()
|
|
|
|
return (res_model,)
|