diff --git a/py/defs/captures.py b/py/defs/captures.py index d842adb..51eaccb 100644 --- a/py/defs/captures.py +++ b/py/defs/captures.py @@ -4,6 +4,7 @@ from .formatters import ( calc_model_hash, calc_vae_hash, calc_lora_hash, + calc_unet_hash, convert_skip_clip, get_scaled_width, get_scaled_height, @@ -92,4 +93,19 @@ CAPTURE_FIELD_LIST = { MetaField.LORA_STRENGTH_MODEL: {"field_name": "strength_model"}, MetaField.LORA_STRENGTH_CLIP: {"value": 0}, }, + # Flux - https://comfyanonymous.github.io/ComfyUI_examples/flux/ + "UNETLoader": { + MetaField.MODEL_NAME: {"field_name": "unet_name"}, + MetaField.MODEL_HASH: {"field_name": "unet_name", "format": calc_unet_hash}, + }, + "RandomNoise": { + MetaField.SEED: {"field_name": "noise_seed"}, + }, + "BasicScheduler": { + MetaField.STEPS: {"field_name": "steps"}, + MetaField.SCHEDULER: {"field_name": "scheduler"}, + }, + "KSamplerSelect": { + MetaField.SAMPLER_NAME: {"field_name": "sampler_name"}, + }, } diff --git a/py/defs/formatters.py b/py/defs/formatters.py index 1a9a4f2..7c68d66 100644 --- a/py/defs/formatters.py +++ b/py/defs/formatters.py @@ -28,6 +28,11 @@ def calc_lora_hash(model_name, input_data): return calc_hash(filename) +def calc_unet_hash(model_name, input_data): + filename = folder_paths.get_full_path("unet", model_name) + return calc_hash(filename) + + def convert_skip_clip(stop_at_clip_layer, input_data): return stop_at_clip_layer * -1 diff --git a/py/defs/samplers.py b/py/defs/samplers.py index 60e3530..a5632e9 100644 --- a/py/defs/samplers.py +++ b/py/defs/samplers.py @@ -7,4 +7,8 @@ SAMPLERS = { "positive": "positive", "negative": "negative", }, + # Flux - https://comfyanonymous.github.io/ComfyUI_examples/flux/ + "SamplerCustomAdvanced": { + "positive": "guider", + }, } diff --git a/py/defs/validators.py b/py/defs/validators.py index 746abda..0bd19c7 100644 --- a/py/defs/validators.py +++ b/py/defs/validators.py @@ -18,7 +18,8 @@ def _get_node_id_list(prompt, field_name): if node["class_type"] == sampler_type: # There are nodes between "KSampler" and "CLIP Text Encode" in the SD3 workflow d = deque() - d.append(node["inputs"][field_map[field_name]][0]) + if field_name in field_map and field_map[field_name] in node["inputs"]: + d.append(node["inputs"][field_map[field_name]][0]) while len(d) > 0: nid2 = d.popleft() class_type = prompt[nid2]["class_type"]