279 lines
10 KiB
Python
279 lines
10 KiB
Python
# Indexed variants (__LoadImagePositive_1__, _2__, …) pick the Nth LoadImage
|
|
# node sorted by node id; the un-indexed form picks the one upstream of the
|
|
# KSampler's latent_image input — without that distinction, branches that mix
|
|
# multiple LoadImages couldn't address them deterministically.
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
|
|
_LOAD_IMAGE_TYPES = {"LoadImage", "LoadImageOutput"}
|
|
_SAMPLER_TYPES = {"KSampler", "KSamplerAdvanced", "FaceDetailer",
|
|
"FaceDetailerPipe", "SamplerCustomAdvanced", "UltimateSDUpscale"}
|
|
|
|
_KEYWORD_RE = re.compile(r"__LoadImage(Positive|Negative)(?:_(\d+))?__")
|
|
|
|
|
|
def resolve_load_image_keywords(prompt_text: str, execution_prompt: dict | None) -> dict | None:
|
|
if not prompt_text or not execution_prompt:
|
|
return None
|
|
if "__LoadImage" not in prompt_text:
|
|
return None
|
|
|
|
matches = list(_KEYWORD_RE.finditer(prompt_text))
|
|
if not matches:
|
|
return None
|
|
|
|
# collect all unique indices referenced (None = un-indexed)
|
|
indices = set()
|
|
for m in matches:
|
|
indices.add(int(m.group(2)) if m.group(2) else None)
|
|
|
|
load_images = _find_all_load_images(execution_prompt)
|
|
if not load_images:
|
|
return {m.group(0): "" for m in matches}
|
|
|
|
# cache extracted prompts per image path to avoid re-reading
|
|
_cache: dict[str, tuple[str, str]] = {}
|
|
|
|
def _get_prompts(path: str) -> tuple[str, str]:
|
|
if path not in _cache:
|
|
_cache[path] = extract_prompts_from_file(path) if path else ("", "")
|
|
return _cache[path]
|
|
|
|
replacements = {}
|
|
for m in matches:
|
|
keyword = m.group(0)
|
|
polarity = m.group(1) # "Positive" or "Negative"
|
|
idx = int(m.group(2)) if m.group(2) else None
|
|
|
|
if idx is not None:
|
|
# indexed: pick the Nth LoadImage (1-based)
|
|
image_path = load_images[idx - 1] if 1 <= idx <= len(load_images) else None
|
|
else:
|
|
# un-indexed: pick the first (primary) LoadImage
|
|
image_path = load_images[0] if load_images else None
|
|
|
|
pos, neg = _get_prompts(image_path) if image_path else ("", "")
|
|
replacements[keyword] = pos if polarity == "Positive" else neg
|
|
|
|
return replacements
|
|
|
|
|
|
def _find_all_load_images(execution_prompt: dict) -> list[str]:
|
|
# KSampler.latent_image's upstream LoadImage is the "primary" — keyword
|
|
# resolution puts it first so an un-indexed __LoadImagePositive__ refers
|
|
# to the image the user is actually sampling from, not whichever node
|
|
# happens to have the lowest id.
|
|
primary_path = None
|
|
primary_nid = None
|
|
for nid, node in execution_prompt.items():
|
|
if not isinstance(node, dict):
|
|
continue
|
|
if node.get("class_type", "") not in _SAMPLER_TYPES:
|
|
continue
|
|
latent_ref = node.get("inputs", {}).get("latent_image")
|
|
if isinstance(latent_ref, list) and len(latent_ref) == 2:
|
|
primary_nid, primary_path = _trace_upstream_load_image(
|
|
execution_prompt, str(latent_ref[0]))
|
|
break
|
|
|
|
# collect all LoadImage nodes, sorted by node id for stable ordering
|
|
all_nodes = []
|
|
for nid, node in execution_prompt.items():
|
|
if not isinstance(node, dict):
|
|
continue
|
|
if node.get("class_type", "") not in _LOAD_IMAGE_TYPES:
|
|
continue
|
|
image_name = node.get("inputs", {}).get("image", "")
|
|
if not image_name:
|
|
continue
|
|
path = _resolve_image_path(image_name)
|
|
if path:
|
|
all_nodes.append((nid, path))
|
|
|
|
# Sort by leading digit prefix, then alphabetically. Without the
|
|
# string tiebreaker, all non-numeric IDs hashed to 0 and their order
|
|
# was arbitrary — which mattered for "which image is primary" logic
|
|
# below.
|
|
def _sort_key(nid):
|
|
digits = "".join(c for c in nid if c.isdigit())
|
|
return (int(digits) if digits else 0, nid)
|
|
all_nodes.sort(key=lambda x: _sort_key(x[0]))
|
|
|
|
# put primary first if found
|
|
if primary_path and primary_nid:
|
|
result = [primary_path]
|
|
for nid, path in all_nodes:
|
|
if nid != primary_nid:
|
|
result.append(path)
|
|
return result
|
|
|
|
return [path for _, path in all_nodes]
|
|
|
|
|
|
def _trace_upstream_load_image(prompt: dict, node_id: str) -> tuple[str | None, str | None]:
|
|
visited = set()
|
|
queue = [node_id]
|
|
while queue:
|
|
nid = queue.pop(0)
|
|
if nid in visited:
|
|
continue
|
|
visited.add(nid)
|
|
node = prompt.get(nid)
|
|
if not isinstance(node, dict):
|
|
continue
|
|
if node.get("class_type", "") in _LOAD_IMAGE_TYPES:
|
|
image_name = node.get("inputs", {}).get("image", "")
|
|
if image_name:
|
|
return nid, _resolve_image_path(image_name)
|
|
for val in node.get("inputs", {}).values():
|
|
if isinstance(val, list) and len(val) == 2:
|
|
queue.append(str(val[0]))
|
|
return None, None
|
|
|
|
|
|
def _resolve_image_path(image_name: str) -> str | None:
|
|
try:
|
|
import folder_paths
|
|
return folder_paths.get_annotated_filepath(image_name)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def extract_prompts_from_image(img) -> tuple[str, str]:
|
|
"""Read pos/neg prompts from an open PIL.Image's ComfyUI metadata."""
|
|
raw_prompt = img.info.get("prompt", "")
|
|
raw_workflow = img.info.get("workflow", "")
|
|
if not raw_prompt:
|
|
return "", ""
|
|
try:
|
|
prompt_data = json.loads(raw_prompt)
|
|
workflow_data = json.loads(raw_workflow) if raw_workflow else None
|
|
except Exception:
|
|
return "", ""
|
|
return _trace_prompts(prompt_data, workflow_data)
|
|
|
|
|
|
def extract_prompts_from_file(image_path: str) -> tuple[str, str]:
|
|
try:
|
|
from PIL import Image
|
|
with Image.open(image_path) as img:
|
|
return extract_prompts_from_image(img)
|
|
except Exception:
|
|
return "", ""
|
|
|
|
|
|
def extract_prompts_from_bytes(raw: bytes) -> tuple[str, str]:
|
|
if not raw:
|
|
return "", ""
|
|
try:
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
with Image.open(BytesIO(raw)) as img:
|
|
return extract_prompts_from_image(img)
|
|
except Exception:
|
|
return "", ""
|
|
|
|
|
|
def _trace_prompts(prompt_data: dict, workflow_data: dict | None) -> tuple[str, str]:
|
|
wf_props = {}
|
|
if workflow_data:
|
|
for wf_node in workflow_data.get("nodes", []):
|
|
wf_props[str(wf_node.get("id", ""))] = wf_node.get("properties", {})
|
|
|
|
sampler = None
|
|
for nid, node in prompt_data.items():
|
|
ct = node.get("class_type", "")
|
|
if "Sampler" in ct or "sampler" in ct:
|
|
sampler = node
|
|
break
|
|
|
|
if not sampler:
|
|
return "", ""
|
|
|
|
def resolve_text(node_id, slot, depth=0):
|
|
if depth > 20:
|
|
return ""
|
|
node_id = str(node_id)
|
|
node = prompt_data.get(node_id)
|
|
if not node:
|
|
return ""
|
|
|
|
ct = node.get("class_type", "")
|
|
|
|
if "PromptChain" in ct:
|
|
props = wf_props.get(node_id, {})
|
|
if slot == 1 and props.get("pcrCompiledOutput"):
|
|
return props["pcrCompiledOutput"]
|
|
elif slot == 2 and props.get("pcrCompiledNegOutput"):
|
|
return props["pcrCompiledNegOutput"]
|
|
# fallback for older images without compiled outputs:
|
|
# use the raw prompt text and split on Negative Prompt:
|
|
raw = node.get("inputs", {}).get("prompt", "")
|
|
if raw:
|
|
import re
|
|
parts = re.split(r"Negative Prompt:", raw, maxsplit=1, flags=re.IGNORECASE)
|
|
pos_part = parts[0].strip() if parts else ""
|
|
neg_part = parts[1].strip() if len(parts) > 1 else ""
|
|
return pos_part if slot == 1 else neg_part
|
|
|
|
# CLIP-like encode nodes — check "text" and "prompt" fields.
|
|
# Covers CLIPTextEncode, TextEncodeQwenImageEditPlus, etc.
|
|
if "Encode" in ct or "CLIP" in ct:
|
|
for field in ("text", "prompt"):
|
|
text_val = node.get("inputs", {}).get(field)
|
|
if text_val is None:
|
|
continue
|
|
if isinstance(text_val, str):
|
|
return text_val
|
|
if isinstance(text_val, list) and len(text_val) == 2:
|
|
return resolve_text(text_val[0], text_val[1], depth + 1)
|
|
return ""
|
|
|
|
# Unknown-class node between sampler and encode (FluxGuidance,
|
|
# ConditioningZeroOut, ConditioningCombine, etc.). Recurse
|
|
# through link-typed inputs, preferring conditioning-ish names,
|
|
# before falling back to string scrape. Without this, a bare
|
|
# string input on a passthrough node (e.g. a scheduler's
|
|
# "index_timestep_zero" mode) gets returned as the prompt.
|
|
inputs = node.get("inputs", {})
|
|
preferred, other = [], []
|
|
for name, val in inputs.items():
|
|
if not (isinstance(val, list) and len(val) == 2):
|
|
continue
|
|
low = name.lower()
|
|
if "cond" in low or low in ("positive", "negative", "prompt", "text"):
|
|
preferred.append(val)
|
|
else:
|
|
other.append(val)
|
|
for ref in preferred + other:
|
|
text = resolve_text(ref[0], ref[1], depth + 1)
|
|
if text:
|
|
return text
|
|
|
|
# Bare-string fallback for nodes that don't expose conditioning
|
|
# links but still carry a raw prompt string (custom text-primitive
|
|
# nodes, etc.). Require a space so we reject mode/parameter
|
|
# identifiers like "index_timestep_zero" that would otherwise
|
|
# surface as the negative prompt when the graph goes through a
|
|
# non-text conditioning node.
|
|
for val in inputs.values():
|
|
if isinstance(val, str) and len(val) > 5 and " " in val:
|
|
return val
|
|
|
|
return ""
|
|
|
|
inputs = sampler.get("inputs", {})
|
|
positive = negative = ""
|
|
|
|
pos_ref = inputs.get("positive")
|
|
neg_ref = inputs.get("negative")
|
|
|
|
if isinstance(pos_ref, list) and len(pos_ref) == 2:
|
|
positive = resolve_text(pos_ref[0], pos_ref[1])
|
|
if isinstance(neg_ref, list) and len(neg_ref) == 2:
|
|
negative = resolve_text(neg_ref[0], neg_ref[1])
|
|
|
|
return positive, negative
|