86 lines
2.3 KiB
Python
86 lines
2.3 KiB
Python
import checkbin
|
|
import os
|
|
|
|
|
|
class CheckbinStartRunNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"token": ("STRING",),
|
|
},
|
|
"optional": {
|
|
"app_key": ("STRING",),
|
|
"directory": ("STRING",),
|
|
"checkin_id": ("STRING",),
|
|
"set_id": ("STRING",),
|
|
"run_id": ("STRING",),
|
|
"sample_size": ("INT",),
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN")
|
|
|
|
RETURN_TYPES = ("CHECKBIN",)
|
|
RETURN_NAMES = ("bin",)
|
|
OUTPUT_IS_LIST = (True,)
|
|
CATEGORY = "comparison"
|
|
FUNCTION = "start_run"
|
|
|
|
def start_run(
|
|
self,
|
|
token,
|
|
app_key="comfyui",
|
|
directory=None,
|
|
checkin_id=None,
|
|
set_id=None,
|
|
run_id=None,
|
|
sample_size=None,
|
|
):
|
|
checkbin.authenticate(token=token)
|
|
checkbin_app = checkbin.App(app_key=app_key)
|
|
|
|
if directory is not None and directory != "":
|
|
if not os.path.isdir(directory):
|
|
raise FileNotFoundError(f"Directory '{directory} cannot be found.'")
|
|
dir_files = os.listdir(directory)
|
|
if len(dir_files) == 0:
|
|
raise FileNotFoundError(f"No files in directory '{directory}'.")
|
|
|
|
# Filter files by extension
|
|
valid_extensions = [".jpg", ".jpeg", ".png", ".webp"]
|
|
dir_files = [
|
|
f
|
|
for f in dir_files
|
|
if any(f.lower().endswith(ext) for ext in valid_extensions)
|
|
]
|
|
|
|
input_set = checkbin_app.create_input_set(
|
|
"ComfyUI auto generated input set"
|
|
)
|
|
|
|
for file in dir_files:
|
|
checkin = input_set.add_input()
|
|
checkin.upload_file("image", f"{directory}/{file}", "image")
|
|
|
|
set_id = input_set.submit()
|
|
|
|
if checkin_id == "":
|
|
checkin_id = None
|
|
if set_id == "":
|
|
set_id = None
|
|
if run_id == "":
|
|
run_id = None
|
|
if sample_size == 0:
|
|
sample_size = None
|
|
|
|
with checkbin_app.start_run(
|
|
checkin_id=checkin_id,
|
|
set_id=set_id,
|
|
run_id=run_id,
|
|
sample_size=sample_size,
|
|
) as bins:
|
|
return (list(bins),)
|