#28 Support Flux

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"},
+    },
 }

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
 

py/defs/samplers.py

@@ -7,4 +7,8 @@ SAMPLERS = {
         "positive": "positive",
         "negative": "negative",
     },
+    # Flux - https://comfyanonymous.github.io/ComfyUI_examples/flux/
+    "SamplerCustomAdvanced": {
+        "positive": "guider",
+    },
 }

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"]