804 lines
30 KiB
Python
804 lines
30 KiB
Python
"""
|
|
Workflow Extractor for MetaHub Save Node
|
|
Auto-extracts sampler params, prompts, model, VAE, and LoRAs from ComfyUI prompt data.
|
|
"""
|
|
|
|
from typing import Any, Dict, List, Optional, Set, Tuple
|
|
|
|
|
|
class WorkflowExtractor:
|
|
SAMPLER_NODES = [
|
|
"KSampler",
|
|
"KSamplerAdvanced",
|
|
"SamplerCustom",
|
|
"SamplerCustomAdvanced",
|
|
"KSamplerSelect",
|
|
"KSampler (Efficient)",
|
|
"ImpactKSamplerBasicPipe",
|
|
"ImpactKSamplerAdvancedPipe",
|
|
"ImpactKSampler",
|
|
]
|
|
|
|
CHECKPOINT_NODES = [
|
|
"CheckpointLoaderSimple",
|
|
"CheckpointLoader",
|
|
"UNETLoader",
|
|
"UnetLoaderGGUF",
|
|
"DualDiffusionLoader",
|
|
"DiffusionModelLoader",
|
|
]
|
|
|
|
VAE_NODES = [
|
|
"VAELoader",
|
|
"CheckpointLoaderSimple",
|
|
"CheckpointLoader",
|
|
]
|
|
|
|
LORA_NODES = [
|
|
"LoraLoader",
|
|
"LoraLoaderModelOnly",
|
|
"LoraLoaderModelOnlyAdvanced",
|
|
"LoraLoaderAdvanced",
|
|
"LoraLoaderAny",
|
|
]
|
|
|
|
CLIP_NODES = [
|
|
"CLIPTextEncode",
|
|
"CLIPTextEncodeSDXL",
|
|
"CLIPTextEncodeSDXLPlus",
|
|
"CLIPTextEncodeSDXLRefiner",
|
|
"CLIPTextEncodeSD3",
|
|
"TextEncodeQwenImage21",
|
|
]
|
|
|
|
VAE_DECODE_NODES = [
|
|
"VAEDecode",
|
|
"VAEDecodeTiled",
|
|
"VAEDecode (Tiled)",
|
|
]
|
|
|
|
LOAD_IMAGE_NODES = [
|
|
"LoadImage",
|
|
"LoadImageMask",
|
|
]
|
|
|
|
def __init__(self, prompt: Optional[Dict[str, Any]]):
|
|
self.prompt: Dict[str, Any] = {}
|
|
if isinstance(prompt, dict):
|
|
for node_id, node_data in prompt.items():
|
|
self.prompt[str(node_id)] = node_data
|
|
|
|
def extract(self, save_node_id: Optional[str] = None) -> Tuple[Dict[str, Any], Set[str]]:
|
|
data: Dict[str, Any] = {"lora_list": []}
|
|
missing: Set[str] = set()
|
|
|
|
sampler_node_id = self.find_sampler_for_save_node(save_node_id)
|
|
if sampler_node_id:
|
|
sampler_params = self.extract_sampler_params(sampler_node_id)
|
|
for key in ("seed", "steps", "cfg", "sampler_name", "scheduler", "denoise"):
|
|
value = sampler_params.get(key)
|
|
if value is None:
|
|
missing.add(key)
|
|
else:
|
|
data[key] = value
|
|
|
|
positive, negative = self.extract_prompts(sampler_node_id)
|
|
if positive is None:
|
|
missing.add("positive")
|
|
else:
|
|
data["positive"] = positive
|
|
if negative is None:
|
|
missing.add("negative")
|
|
else:
|
|
data["negative"] = negative
|
|
|
|
model_name = self.extract_model_name(sampler_node_id)
|
|
if model_name is None:
|
|
missing.add("model_name")
|
|
else:
|
|
data["model_name"] = model_name
|
|
else:
|
|
missing.update(
|
|
{
|
|
"seed",
|
|
"steps",
|
|
"cfg",
|
|
"sampler_name",
|
|
"scheduler",
|
|
"denoise",
|
|
"positive",
|
|
"negative",
|
|
"model_name",
|
|
}
|
|
)
|
|
|
|
vae_name = self.extract_vae_name(save_node_id)
|
|
if vae_name is None:
|
|
missing.add("vae_name")
|
|
else:
|
|
data["vae_name"] = vae_name
|
|
|
|
lora_list, has_lora_nodes = self.extract_loras()
|
|
if lora_list:
|
|
data["lora_list"] = lora_list
|
|
elif has_lora_nodes:
|
|
missing.add("loras")
|
|
|
|
lineage = self.extract_model_3d_lineage(save_node_id) or self.extract_lineage(sampler_node_id)
|
|
if lineage:
|
|
data.update(lineage)
|
|
|
|
return data, missing
|
|
|
|
def find_sampler_for_save_node(self, save_node_id: Optional[str]) -> Optional[str]:
|
|
sampler_nodes = self._find_nodes_by_type(self.SAMPLER_NODES)
|
|
if not sampler_nodes:
|
|
return None
|
|
|
|
if save_node_id:
|
|
save_node = self._get_node(save_node_id)
|
|
if save_node:
|
|
inputs = save_node.get("inputs", {})
|
|
for input_name in ("images", "model_3d"):
|
|
start_node_id = self._get_connection_node_id(inputs.get(input_name))
|
|
if start_node_id:
|
|
sampler_id = self._find_sampler_from_images_source(start_node_id)
|
|
if sampler_id:
|
|
return sampler_id
|
|
|
|
return sampler_nodes[0]
|
|
|
|
def extract_sampler_params(self, sampler_node_id: str) -> Dict[str, Any]:
|
|
node = self._get_node(sampler_node_id)
|
|
if not node:
|
|
return {}
|
|
inputs = node.get("inputs", {})
|
|
|
|
seed = self._coerce_int(self._get_literal_input(inputs, "seed"))
|
|
if seed is None:
|
|
seed = self._coerce_int(
|
|
self._resolve_scalar_from_connection(inputs.get("seed"), ("seed", "value"))
|
|
)
|
|
if seed is None:
|
|
seed = self._coerce_int(
|
|
self._resolve_scalar_from_connection(inputs.get("noise"), ("noise_seed", "seed", "value"))
|
|
)
|
|
if seed is None:
|
|
seed = self._coerce_int(
|
|
self._resolve_scalar_from_connection(inputs.get("sampler"), ("seed", "noise_seed", "value"))
|
|
)
|
|
steps = self._coerce_int(self._get_literal_input(inputs, "steps"))
|
|
if steps is None:
|
|
steps = self._coerce_int(
|
|
self._resolve_scalar_from_connection(inputs.get("sigmas"), ("steps", "value"))
|
|
)
|
|
cfg = self._coerce_float(self._get_first_literal(inputs, ("cfg", "cfg_scale")))
|
|
if cfg is None:
|
|
cfg = self._coerce_float(
|
|
self._resolve_scalar_from_connection(inputs.get("guider"), ("cfg", "cfg_scale", "scale"))
|
|
)
|
|
sampler_name = self._get_first_literal(inputs, ("sampler_name", "sampler"))
|
|
if sampler_name is None:
|
|
sampler_name = self._resolve_scalar_from_connection(
|
|
inputs.get("sampler"), ("sampler_name", "sampler", "name")
|
|
)
|
|
scheduler = self._get_first_literal(inputs, ("scheduler", "scheduler_name"))
|
|
if scheduler is None:
|
|
scheduler = self._resolve_scalar_from_connection(
|
|
inputs.get("sigmas"), ("scheduler", "scheduler_name", "name")
|
|
)
|
|
denoise = self._coerce_float(self._get_first_literal(inputs, ("denoise", "denoise_strength")))
|
|
if denoise is None:
|
|
denoise = self._coerce_float(
|
|
self._resolve_scalar_from_connection(inputs.get("sigmas"), ("denoise", "denoise_strength"))
|
|
)
|
|
|
|
return {
|
|
"seed": seed,
|
|
"steps": steps,
|
|
"cfg": cfg,
|
|
"sampler_name": sampler_name,
|
|
"scheduler": scheduler,
|
|
"denoise": denoise,
|
|
}
|
|
|
|
def extract_model_name(self, sampler_node_id: str) -> Optional[str]:
|
|
node = self._get_node(sampler_node_id)
|
|
if not node:
|
|
return None
|
|
inputs = node.get("inputs", {})
|
|
start_node_ids = []
|
|
for input_name in ("model", "guider", "sigmas"):
|
|
start_node_id = self._get_connection_node_id(inputs.get(input_name))
|
|
if start_node_id:
|
|
start_node_ids.append(start_node_id)
|
|
if not start_node_ids:
|
|
return None
|
|
checkpoint_node_id = self._bfs_upstream(
|
|
start_node_ids,
|
|
lambda n: self._class_type(n) in self.CHECKPOINT_NODES,
|
|
)
|
|
if not checkpoint_node_id:
|
|
return None
|
|
checkpoint_node = self._get_node(checkpoint_node_id)
|
|
if not checkpoint_node:
|
|
return None
|
|
return self._get_checkpoint_name(checkpoint_node)
|
|
|
|
def extract_vae_name(self, save_node_id: Optional[str]) -> Optional[str]:
|
|
vae_decode_id = self._find_vae_decode_node(save_node_id)
|
|
if not vae_decode_id:
|
|
return None
|
|
decode_node = self._get_node(vae_decode_id)
|
|
if not decode_node:
|
|
return None
|
|
vae_conn = decode_node.get("inputs", {}).get("vae")
|
|
start_node_id = self._get_connection_node_id(vae_conn)
|
|
if not start_node_id:
|
|
return None
|
|
vae_node_id = self._bfs_upstream(
|
|
[start_node_id],
|
|
lambda n: self._class_type(n) in self.VAE_NODES,
|
|
)
|
|
if not vae_node_id:
|
|
return None
|
|
vae_node = self._get_node(vae_node_id)
|
|
if not vae_node:
|
|
return None
|
|
return self._get_vae_name(vae_node)
|
|
|
|
def extract_prompts(self, sampler_node_id: str) -> Tuple[Optional[str], Optional[str]]:
|
|
node = self._get_node(sampler_node_id)
|
|
if not node:
|
|
return None, None
|
|
inputs = node.get("inputs", {})
|
|
positive_conn = inputs.get("positive") or inputs.get("positive_cond")
|
|
negative_conn = inputs.get("negative") or inputs.get("negative_cond")
|
|
positive = self._extract_text_from_connection(positive_conn)
|
|
negative = self._extract_text_from_connection(negative_conn)
|
|
if positive is None:
|
|
positive = self._extract_text_from_connection(inputs.get("guider"))
|
|
if negative is None:
|
|
negative = self._extract_text_from_connection(inputs.get("negative_conditioning"))
|
|
return positive, negative
|
|
|
|
def extract_loras(self) -> Tuple[List[Dict[str, Any]], bool]:
|
|
loras: List[Dict[str, Any]] = []
|
|
has_lora_nodes = False
|
|
|
|
for _, node in self.prompt.items():
|
|
class_type = self._class_type(node)
|
|
if not self._is_lora_node(class_type):
|
|
continue
|
|
has_lora_nodes = True
|
|
inputs = node.get("inputs", {})
|
|
stacked_loras = self._extract_lora_manager_entries(inputs)
|
|
if stacked_loras:
|
|
loras.extend(stacked_loras)
|
|
continue
|
|
dynamic_loras = self._extract_dynamic_lora_entries(inputs)
|
|
if dynamic_loras:
|
|
loras.extend(dynamic_loras)
|
|
continue
|
|
lora_name = self._get_first_literal(inputs, ("lora_name", "lora", "lora_name_1", "lora_1"))
|
|
if not isinstance(lora_name, str) or not lora_name.strip():
|
|
continue
|
|
strength_model = self._get_first_literal(
|
|
inputs,
|
|
("strength_model", "strength", "strength_unet"),
|
|
)
|
|
weight = self._coerce_float(strength_model)
|
|
if weight is None:
|
|
strength_clip = self._get_first_literal(inputs, ("strength_clip",))
|
|
weight = self._coerce_float(strength_clip)
|
|
if weight is None:
|
|
weight = 1.0
|
|
loras.append({"name": lora_name, "weight": float(weight)})
|
|
|
|
return loras, has_lora_nodes
|
|
|
|
def _extract_dynamic_lora_entries(self, inputs: Dict[str, Any]) -> List[Dict[str, Any]]:
|
|
loras: List[Dict[str, Any]] = []
|
|
for key, raw_lora in inputs.items():
|
|
if not key.lower().startswith("lora_") or not isinstance(raw_lora, dict):
|
|
continue
|
|
if raw_lora.get("on") is False or raw_lora.get("active") is False:
|
|
continue
|
|
name = raw_lora.get("lora") or raw_lora.get("lora_name") or raw_lora.get("name")
|
|
if not isinstance(name, str) or not name.strip():
|
|
continue
|
|
weight = self._coerce_float(
|
|
raw_lora.get("strength")
|
|
if raw_lora.get("strength") is not None
|
|
else raw_lora.get("modelStrength")
|
|
)
|
|
if weight is None:
|
|
weight = self._coerce_float(raw_lora.get("strength_model"))
|
|
if weight is None:
|
|
weight = 1.0
|
|
loras.append({"name": name, "weight": float(weight)})
|
|
return loras
|
|
|
|
def _extract_lora_manager_entries(self, inputs: Dict[str, Any]) -> List[Dict[str, Any]]:
|
|
raw_loras = inputs.get("loras")
|
|
if isinstance(raw_loras, dict) and isinstance(raw_loras.get("__value__"), list):
|
|
raw_loras = raw_loras["__value__"]
|
|
if not isinstance(raw_loras, list):
|
|
return []
|
|
|
|
loras: List[Dict[str, Any]] = []
|
|
for raw_lora in raw_loras:
|
|
if not isinstance(raw_lora, dict):
|
|
continue
|
|
if raw_lora.get("active") is False:
|
|
continue
|
|
name = raw_lora.get("name") or raw_lora.get("lora_name")
|
|
if not name:
|
|
continue
|
|
weight = self._coerce_float(
|
|
raw_lora.get("strength")
|
|
if raw_lora.get("strength") is not None
|
|
else raw_lora.get("modelStrength")
|
|
)
|
|
if weight is None:
|
|
weight = self._coerce_float(raw_lora.get("clipStrength"))
|
|
if weight is None:
|
|
weight = 1.0
|
|
loras.append({"name": str(name), "weight": float(weight)})
|
|
return loras
|
|
|
|
def extract_lineage(self, sampler_node_id: Optional[str]) -> Dict[str, Any]:
|
|
if not sampler_node_id:
|
|
return {}
|
|
|
|
node = self._get_node(sampler_node_id)
|
|
if not node:
|
|
return {}
|
|
|
|
inputs = node.get("inputs", {})
|
|
latent_conn = inputs.get("latent_image") or inputs.get("latent") or inputs.get("samples")
|
|
start_node_id = self._get_connection_node_id(latent_conn)
|
|
if not start_node_id:
|
|
return {}
|
|
|
|
queue = [start_node_id]
|
|
visited: Set[str] = set()
|
|
source_image: Optional[Dict[str, Any]] = None
|
|
has_load_image = False
|
|
has_inpaint = False
|
|
has_outpaint = False
|
|
|
|
while queue:
|
|
current_id = queue.pop(0)
|
|
if current_id in visited:
|
|
continue
|
|
visited.add(current_id)
|
|
|
|
current = self._get_node(current_id)
|
|
if not current:
|
|
continue
|
|
|
|
class_type = self._class_type(current)
|
|
class_lower = class_type.lower()
|
|
|
|
if class_type in self.LOAD_IMAGE_NODES or class_lower in {"loadimage", "loadimagemask"}:
|
|
has_load_image = True
|
|
if not source_image:
|
|
image_value = current.get("inputs", {}).get("image")
|
|
if isinstance(image_value, str) and image_value.strip():
|
|
normalized = image_value.replace("\\", "/")
|
|
source_image = {
|
|
"fileName": normalized.split("/")[-1],
|
|
"relativePath": normalized,
|
|
"nodeId": current_id,
|
|
"nodeType": class_type,
|
|
}
|
|
|
|
if "outpaint" in class_lower or class_lower in {"padforoutpaint", "imagepadforoutpaint"}:
|
|
has_outpaint = True
|
|
|
|
if "inpaint" in class_lower or class_lower == "setlatentnoisemask" or "mask" in class_lower:
|
|
has_inpaint = True
|
|
|
|
for input_value in current.get("inputs", {}).values():
|
|
upstream_id = self._get_connection_node_id(input_value)
|
|
if upstream_id and upstream_id not in visited:
|
|
queue.append(upstream_id)
|
|
|
|
if not has_load_image:
|
|
return {}
|
|
|
|
generation_type = "outpaint" if has_outpaint else "inpaint" if has_inpaint else "img2img"
|
|
return {
|
|
"generation_type": generation_type,
|
|
"source_image": source_image,
|
|
}
|
|
|
|
def extract_model_3d_lineage(self, save_node_id: Optional[str]) -> Dict[str, Any]:
|
|
if not save_node_id:
|
|
return {}
|
|
|
|
save_node = self._get_node(save_node_id)
|
|
if not save_node:
|
|
return {}
|
|
|
|
class_type = self._class_type(save_node).lower().replace("_", "").replace(" ", "")
|
|
if "save3d" not in class_type and "saveglb" not in class_type:
|
|
return {}
|
|
|
|
inputs = save_node.get("inputs", {})
|
|
model_conn = inputs.get("model_3d") or inputs.get("mesh") or inputs.get("model")
|
|
start_node_id = self._get_connection_node_id(model_conn)
|
|
if not start_node_id:
|
|
return {}
|
|
|
|
queue = [start_node_id]
|
|
visited: Set[str] = set()
|
|
while queue:
|
|
current_id = queue.pop(0)
|
|
if current_id in visited:
|
|
continue
|
|
visited.add(current_id)
|
|
|
|
current = self._get_node(current_id)
|
|
if not current:
|
|
continue
|
|
|
|
current_type = self._class_type(current)
|
|
if current_type in self.LOAD_IMAGE_NODES or current_type.lower() == "loadimage":
|
|
image_value = current.get("inputs", {}).get("image")
|
|
if isinstance(image_value, str) and image_value.strip():
|
|
normalized = image_value.replace("\\", "/")
|
|
return {
|
|
"generation_type": "image2model3d",
|
|
"source_image": {
|
|
"fileName": normalized.split("/")[-1],
|
|
"relativePath": normalized,
|
|
"nodeId": current_id,
|
|
"nodeType": current_type,
|
|
},
|
|
}
|
|
|
|
for input_value in current.get("inputs", {}).values():
|
|
upstream_id = self._get_connection_node_id(input_value)
|
|
if upstream_id and upstream_id not in visited:
|
|
queue.append(upstream_id)
|
|
|
|
return {}
|
|
|
|
def _find_sampler_from_images_source(self, start_node_id: str) -> Optional[str]:
|
|
start_node = self._get_node(start_node_id)
|
|
if not start_node:
|
|
return None
|
|
if self._class_type(start_node) in self.SAMPLER_NODES:
|
|
return start_node_id
|
|
|
|
if self._class_type(start_node) in self.VAE_DECODE_NODES:
|
|
samples_conn = start_node.get("inputs", {}).get("samples")
|
|
samples_node_id = self._get_connection_node_id(samples_conn)
|
|
if samples_node_id:
|
|
sampler_id = self._bfs_upstream(
|
|
[samples_node_id],
|
|
lambda n: self._class_type(n) in self.SAMPLER_NODES,
|
|
)
|
|
if sampler_id:
|
|
return sampler_id
|
|
|
|
return self._bfs_upstream(
|
|
[start_node_id],
|
|
lambda n: self._class_type(n) in self.SAMPLER_NODES,
|
|
)
|
|
|
|
def _find_vae_decode_node(self, save_node_id: Optional[str]) -> Optional[str]:
|
|
if not save_node_id:
|
|
return None
|
|
save_node = self._get_node(save_node_id)
|
|
if not save_node:
|
|
return None
|
|
inputs = save_node.get("inputs", {})
|
|
for input_name in ("images", "model_3d"):
|
|
start_node_id = self._get_connection_node_id(inputs.get(input_name))
|
|
if not start_node_id:
|
|
continue
|
|
vae_decode_id = self._bfs_upstream(
|
|
[start_node_id],
|
|
lambda n: self._class_type(n) in self.VAE_DECODE_NODES,
|
|
)
|
|
if vae_decode_id:
|
|
return vae_decode_id
|
|
return None
|
|
|
|
def _find_nodes_by_type(self, types: List[str]) -> List[str]:
|
|
matched = []
|
|
for node_id, node in self.prompt.items():
|
|
if self._class_type(node) in types:
|
|
matched.append(node_id)
|
|
return matched
|
|
|
|
def _extract_text_from_connection(self, conn: Any) -> Optional[str]:
|
|
node_id = self._get_connection_node_id(conn)
|
|
if not node_id:
|
|
return None
|
|
return self._extract_text_from_node(node_id, conn[1])
|
|
|
|
def _extract_text_from_node(self, start_node_id: str, output_slot: int = 0) -> Optional[str]:
|
|
texts: List[str] = []
|
|
queue = [(start_node_id, output_slot)]
|
|
visited: Set[Tuple[str, int]] = set()
|
|
|
|
while queue:
|
|
node_id, current_slot = queue.pop(0)
|
|
if (node_id, current_slot) in visited:
|
|
continue
|
|
visited.add((node_id, current_slot))
|
|
node = self._get_node(node_id)
|
|
if not node:
|
|
continue
|
|
class_type = self._class_type(node)
|
|
if class_type == "ConditioningZeroOut":
|
|
return ""
|
|
if class_type == "TextEncodeQwenImage21":
|
|
text = self._get_qwen_image_21_text(node, current_slot)
|
|
if text is not None:
|
|
texts.append(text)
|
|
continue
|
|
if self._is_prompt_encoder(class_type):
|
|
text = self._get_clip_text(node)
|
|
if text is not None:
|
|
texts.append(text)
|
|
continue
|
|
for input_val in node.get("inputs", {}).values():
|
|
conn_id = self._get_connection_node_id(input_val)
|
|
if conn_id and (conn_id, input_val[1]) not in visited:
|
|
queue.append((conn_id, input_val[1]))
|
|
|
|
if texts:
|
|
return "\n".join(texts)
|
|
return None
|
|
|
|
def _get_qwen_image_21_text(self, node: Dict[str, Any], output_slot: int) -> Optional[str]:
|
|
if output_slot not in (0, 1):
|
|
return None
|
|
key = "prompt" if output_slot == 0 else "negative_prompt"
|
|
inputs = node.get("inputs", {})
|
|
value = self._get_literal_input(inputs, key)
|
|
if isinstance(value, str):
|
|
return value
|
|
conn = inputs.get(key)
|
|
source_id = self._get_connection_node_id(conn)
|
|
if source_id:
|
|
source = self._get_node(source_id)
|
|
if source:
|
|
for source_key in ("text", "string", "value", "prompt", "negative_prompt"):
|
|
source_value = self._get_literal_input(source.get("inputs", {}), source_key)
|
|
if isinstance(source_value, str):
|
|
return source_value
|
|
return self._resolve_string_from_connection(conn)
|
|
return None
|
|
|
|
def _get_clip_text(self, node: Dict[str, Any]) -> Optional[str]:
|
|
inputs = node.get("inputs", {})
|
|
text = self._get_literal_input(inputs, "text")
|
|
if text is not None and not isinstance(text, bool) and str(text).strip():
|
|
return str(text)
|
|
text_from_conn = self._resolve_string_from_connection(inputs.get("text"))
|
|
if text_from_conn is not None:
|
|
return text_from_conn
|
|
parts: List[str] = []
|
|
for key in ("text_g", "text_l"):
|
|
value = self._get_literal_input(inputs, key)
|
|
if value:
|
|
parts.append(str(value))
|
|
if parts:
|
|
return "\n".join(parts)
|
|
return None
|
|
|
|
def _resolve_string_from_connection(self, conn: Any) -> Optional[str]:
|
|
node_id = self._get_connection_node_id(conn)
|
|
if not node_id:
|
|
return None
|
|
return self._resolve_string_from_node(node_id, set())
|
|
|
|
def _resolve_string_from_node(self, node_id: str, visited: Set[str]) -> Optional[str]:
|
|
if node_id in visited:
|
|
return None
|
|
visited.add(node_id)
|
|
|
|
node = self._get_node(node_id)
|
|
if not node:
|
|
return None
|
|
|
|
inputs = node.get("inputs", {})
|
|
class_type = self._class_type(node)
|
|
class_lower = class_type.lower()
|
|
|
|
if class_type == "ConditioningZeroOut":
|
|
return ""
|
|
|
|
if class_type in ("PrimitiveStringMultiline", "String Literal"):
|
|
for key in ("value", "string", "text"):
|
|
value = self._get_literal_input(inputs, key)
|
|
if isinstance(value, str):
|
|
return value
|
|
|
|
if self._is_prompt_encoder(class_type):
|
|
return self._get_clip_text(node)
|
|
|
|
if "switch" in class_lower and "on_false" in inputs and "on_true" in inputs:
|
|
switch_value = self._resolve_boolean_input(inputs.get("switch"))
|
|
branch = "on_true" if switch_value is True else "on_false"
|
|
value = self._resolve_string_from_connection_with_visited(inputs.get(branch), visited)
|
|
if value is not None:
|
|
return value
|
|
|
|
if class_lower == "joinstrings" or "join strings" in class_lower:
|
|
delimiter = self._get_literal_input(inputs, "delimiter")
|
|
joiner = str(delimiter) if delimiter is not None else " "
|
|
parts: List[str] = []
|
|
for key in ("string1", "string2", "string3", "string4"):
|
|
value = self._get_literal_input(inputs, key)
|
|
if value is None:
|
|
value = self._resolve_string_from_connection_with_visited(inputs.get(key), visited)
|
|
if value is not None and str(value).strip():
|
|
parts.append(str(value).strip())
|
|
return joiner.join(parts) if parts else ""
|
|
|
|
for key in ("text", "positive", "negative", "prompt", "string", "value"):
|
|
value = self._get_literal_input(inputs, key)
|
|
if value is not None and not isinstance(value, bool) and str(value).strip():
|
|
return str(value)
|
|
|
|
for key in ("positive", "text", "string1", "string2", "conditioning"):
|
|
value = self._resolve_string_from_connection_with_visited(inputs.get(key), visited)
|
|
if value is not None and str(value).strip():
|
|
return value
|
|
|
|
for input_val in inputs.values():
|
|
value = self._resolve_string_from_connection_with_visited(input_val, visited)
|
|
if value is not None and str(value).strip():
|
|
return value
|
|
|
|
return None
|
|
|
|
def _resolve_boolean_input(self, value: Any) -> Optional[bool]:
|
|
if isinstance(value, bool):
|
|
return value
|
|
node_id = self._get_connection_node_id(value)
|
|
if not node_id:
|
|
return None
|
|
node = self._get_node(node_id)
|
|
literal = node.get("inputs", {}).get("value") if node else None
|
|
return literal if isinstance(literal, bool) else None
|
|
|
|
def _resolve_string_from_connection_with_visited(self, conn: Any, visited: Set[str]) -> Optional[str]:
|
|
node_id = self._get_connection_node_id(conn)
|
|
if not node_id:
|
|
return None
|
|
return self._resolve_string_from_node(node_id, visited)
|
|
|
|
def _resolve_scalar_from_connection(self, conn: Any, keys: Tuple[str, ...]) -> Any:
|
|
node_id = self._get_connection_node_id(conn)
|
|
if not node_id:
|
|
return None
|
|
|
|
queue = [node_id]
|
|
visited: Set[str] = set()
|
|
while queue:
|
|
current_id = queue.pop(0)
|
|
if current_id in visited:
|
|
continue
|
|
visited.add(current_id)
|
|
|
|
node = self._get_node(current_id)
|
|
if not node:
|
|
continue
|
|
|
|
inputs = node.get("inputs", {})
|
|
for key in keys:
|
|
value = self._get_literal_input(inputs, key)
|
|
if value is not None:
|
|
return value
|
|
|
|
for input_val in inputs.values():
|
|
conn_id = self._get_connection_node_id(input_val)
|
|
if conn_id and conn_id not in visited:
|
|
queue.append(conn_id)
|
|
|
|
return None
|
|
|
|
def _get_checkpoint_name(self, node: Dict[str, Any]) -> Optional[str]:
|
|
inputs = node.get("inputs", {})
|
|
for key in ("ckpt_name", "checkpoint", "model_name", "unet_name", "diffusion_model_name", "model"):
|
|
value = self._get_literal_input(inputs, key)
|
|
if value:
|
|
return str(value)
|
|
return None
|
|
|
|
def _get_vae_name(self, node: Dict[str, Any]) -> Optional[str]:
|
|
inputs = node.get("inputs", {})
|
|
for key in ("vae_name", "ckpt_name", "model_name"):
|
|
value = self._get_literal_input(inputs, key)
|
|
if value:
|
|
return str(value)
|
|
return None
|
|
|
|
def _get_node(self, node_id: Optional[str]) -> Optional[Dict[str, Any]]:
|
|
if node_id is None:
|
|
return None
|
|
return self.prompt.get(str(node_id))
|
|
|
|
@staticmethod
|
|
def _class_type(node: Dict[str, Any]) -> str:
|
|
return str(node.get("class_type", ""))
|
|
|
|
@staticmethod
|
|
def _is_connection(value: Any) -> bool:
|
|
return isinstance(value, (list, tuple)) and len(value) >= 2
|
|
|
|
def _get_connection_node_id(self, value: Any) -> Optional[str]:
|
|
if not self._is_connection(value):
|
|
return None
|
|
return str(value[0])
|
|
|
|
def _get_literal_input(self, inputs: Dict[str, Any], key: str) -> Any:
|
|
value = inputs.get(key)
|
|
if self._is_connection(value):
|
|
return None
|
|
return value
|
|
|
|
def _get_first_literal(self, inputs: Dict[str, Any], keys: Tuple[str, ...]) -> Any:
|
|
for key in keys:
|
|
value = self._get_literal_input(inputs, key)
|
|
if value is not None:
|
|
return value
|
|
return None
|
|
|
|
def _bfs_upstream(self, start_ids: List[str], match_fn) -> Optional[str]:
|
|
queue = list(start_ids)
|
|
visited: Set[str] = set()
|
|
while queue:
|
|
node_id = queue.pop(0)
|
|
if node_id in visited:
|
|
continue
|
|
visited.add(node_id)
|
|
node = self._get_node(node_id)
|
|
if not node:
|
|
continue
|
|
if match_fn(node):
|
|
return node_id
|
|
for input_val in node.get("inputs", {}).values():
|
|
conn_id = self._get_connection_node_id(input_val)
|
|
if conn_id and conn_id not in visited:
|
|
queue.append(conn_id)
|
|
return None
|
|
|
|
@staticmethod
|
|
def _coerce_int(value: Any) -> Optional[int]:
|
|
if value is None:
|
|
return None
|
|
try:
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
@staticmethod
|
|
def _coerce_float(value: Any) -> Optional[float]:
|
|
if value is None:
|
|
return None
|
|
try:
|
|
return float(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
@classmethod
|
|
def _is_lora_node(cls, class_type: str) -> bool:
|
|
if class_type in cls.LORA_NODES:
|
|
return True
|
|
return "lora" in class_type.lower()
|
|
|
|
@classmethod
|
|
def _is_prompt_encoder(cls, class_type: str) -> bool:
|
|
if class_type in cls.CLIP_NODES:
|
|
return True
|
|
return "promptencoder" in class_type.lower().replace(" ", "")
|