42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
import folder_paths
|
|
import comfy.sd
|
|
|
|
class CheckpointLoaderIterator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"batch_index": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "STRING")
|
|
RETURN_NAMES = ("MODEL", "CLIP", "VAE", "CKPT_NAME")
|
|
FUNCTION = "load_checkpoint"
|
|
CATEGORY = "loaders"
|
|
|
|
def load_checkpoint(self, batch_index):
|
|
checkpoints = folder_paths.get_filename_list("checkpoints")
|
|
if not checkpoints:
|
|
raise ValueError("No checkpoints found.")
|
|
|
|
# Calculate the index to use based on batch_index
|
|
index = batch_index % len(checkpoints)
|
|
ckpt_name = checkpoints[index]
|
|
|
|
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
|
|
out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
|
|
|
return (out[0], out[1], out[2], ckpt_name)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, batch_index):
|
|
# We want to re-evaluate on every execution if the batch_index changes
|
|
return batch_index
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CheckpointLoaderIterator": CheckpointLoaderIterator
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"CheckpointLoaderIterator": "Checkpoint Loader Iterator"
|
|
}
|