Fix ComfyUI prompt extraction for ImpactWildcardProcessor
osirigunso
parent
fcc23b5337
commit
97f760aa54
|
|
@ -487,48 +487,196 @@ def is_img_created_by_comfyui(img: Image):
|
|||
def is_img_created_by_comfyui_with_webui_gen_info(img: Image):
|
||||
return is_img_created_by_comfyui(img) and img.info.get('parameters')
|
||||
|
||||
|
||||
def get_comfyui_exif_data(img: Image):
|
||||
prompt = None
|
||||
if img.format == "PNG":
|
||||
prompt = img.info.get('prompt')
|
||||
elif img.format == "WEBP":
|
||||
exif = img.info.get("exif")
|
||||
split = [x.decode("utf-8", errors="ignore") for x in exif.split(b"\x00")]
|
||||
prompt_str = find(split, lambda x: x.lower().startswith("prompt:"))
|
||||
if prompt_str:
|
||||
prompt = prompt_str.split(":", 1)[1] if prompt_str else None
|
||||
if exif:
|
||||
split = [x.decode("utf-8", errors="ignore") for x in exif.split(b"\x00")]
|
||||
prompt_str = find(split, lambda x: x.lower().startswith("prompt:"))
|
||||
if prompt_str:
|
||||
prompt = prompt_str.split(":", 1)[1]
|
||||
|
||||
if not prompt:
|
||||
return {}
|
||||
|
||||
data: Dict[str, Any] = json.loads(prompt)
|
||||
|
||||
meta_key = '3'
|
||||
data: Dict[str, any] = json.loads(prompt)
|
||||
for i in data.keys():
|
||||
try:
|
||||
if data[i]["class_type"].startswith("KSampler"):
|
||||
meta_key = i
|
||||
break
|
||||
except:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if meta_key not in data:
|
||||
return {}
|
||||
|
||||
meta = {}
|
||||
KSampler_entry = data[meta_key]["inputs"]
|
||||
#print(KSampler_entry) # for testing
|
||||
|
||||
# As a workaround to bypass parsing errors in the parser.
|
||||
# https://github.com/jiw0220/stable-diffusion-image-metadata/blob/00b8d42d4d1a536862bba0b07c332bdebb2a0ce5/src/index.ts#L130
|
||||
meta["Steps"] = KSampler_entry.get("steps", "Unknown")
|
||||
meta["Sampler"] = KSampler_entry["sampler_name"]
|
||||
meta["Model"] = data[KSampler_entry["model"][0]]["inputs"].get("ckpt_name")
|
||||
meta["Sampler"] = KSampler_entry.get("sampler_name", "Unknown")
|
||||
try:
|
||||
meta["Model"] = data[KSampler_entry["model"][0]]["inputs"].get("ckpt_name")
|
||||
except Exception:
|
||||
meta["Model"] = None
|
||||
meta["Source Identifier"] = "ComfyUI"
|
||||
def get_text_from_clip(idx: str) :
|
||||
try:
|
||||
inputs = data[idx]["inputs"]
|
||||
if "text" in inputs:
|
||||
text = inputs["text"]
|
||||
elif "t5xxl" in inputs:
|
||||
text = inputs["t5xxl"]
|
||||
|
||||
text_key_priority = [
|
||||
"populated_text", # ImpactWildcardProcessor の最終展開結果
|
||||
"text",
|
||||
"prompt",
|
||||
"positive",
|
||||
"negative",
|
||||
"string",
|
||||
"value",
|
||||
"t5xxl",
|
||||
]
|
||||
text_key_blacklist = {
|
||||
"wildcard_text", # 展開前テンプレート
|
||||
"Select to add Wildcard", # UI 用
|
||||
"select_to_add_wildcard",
|
||||
"template",
|
||||
"pattern",
|
||||
}
|
||||
|
||||
wildcard_patterns = [
|
||||
re.compile(r"__[^_\n]+__"),
|
||||
re.compile(r"\{[^{}\n]*\|[^{}\n]*\}"),
|
||||
]
|
||||
|
||||
def normalize_text(value):
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
value = value.strip()
|
||||
return value if value else None
|
||||
|
||||
def looks_unexpanded_wildcard(text: str) -> bool:
|
||||
return any(p.search(text) for p in wildcard_patterns)
|
||||
|
||||
def get_node(node_id):
|
||||
key = str(node_id)
|
||||
return data.get(key) or data.get(node_id)
|
||||
|
||||
def get_best_text_from_inputs(inputs: dict):
|
||||
if not isinstance(inputs, dict):
|
||||
return ""
|
||||
|
||||
clean_candidates = []
|
||||
wildcard_candidates = []
|
||||
|
||||
def add_candidate(value):
|
||||
text = normalize_text(value)
|
||||
if not text:
|
||||
return
|
||||
if looks_unexpanded_wildcard(text):
|
||||
wildcard_candidates.append(text)
|
||||
else:
|
||||
return ""
|
||||
if isinstance(text, list): # type:CLIPTextEncode (NSP) mode:Wildcards
|
||||
text = data[text[0]]["inputs"]["text"]
|
||||
return text.strip()
|
||||
clean_candidates.append(text)
|
||||
|
||||
for key in text_key_priority:
|
||||
if key in text_key_blacklist:
|
||||
continue
|
||||
add_candidate(inputs.get(key))
|
||||
|
||||
if clean_candidates:
|
||||
return clean_candidates[0]
|
||||
if wildcard_candidates:
|
||||
return wildcard_candidates[0]
|
||||
|
||||
for key, value in inputs.items():
|
||||
if key in text_key_blacklist:
|
||||
continue
|
||||
add_candidate(value)
|
||||
|
||||
if clean_candidates:
|
||||
return clean_candidates[0]
|
||||
if wildcard_candidates:
|
||||
return wildcard_candidates[0]
|
||||
return ""
|
||||
|
||||
def extract_text_from_node(node: dict):
|
||||
if not isinstance(node, dict):
|
||||
return ""
|
||||
|
||||
class_type = node.get("class_type", "")
|
||||
inputs = node.get("inputs", {}) or {}
|
||||
|
||||
# 明示的な特例: ImpactWildcardProcessor は populated_text を最優先
|
||||
if class_type == "ImpactWildcardProcessor":
|
||||
populated = normalize_text(inputs.get("populated_text"))
|
||||
if populated:
|
||||
return populated
|
||||
# wildcard_text は意図的に返さない
|
||||
return ""
|
||||
|
||||
return get_best_text_from_inputs(inputs)
|
||||
|
||||
def resolve_text_from_ref(ref, visited=None):
|
||||
if visited is None:
|
||||
visited = set()
|
||||
|
||||
node_id = ref[0] if isinstance(ref, list) and len(ref) >= 1 else ref
|
||||
node_key = str(node_id)
|
||||
if node_key in visited:
|
||||
return ""
|
||||
visited.add(node_key)
|
||||
|
||||
node = get_node(node_id)
|
||||
if not isinstance(node, dict):
|
||||
return ""
|
||||
|
||||
direct_text = extract_text_from_node(node)
|
||||
if direct_text:
|
||||
return direct_text
|
||||
|
||||
inputs = node.get("inputs", {}) or {}
|
||||
|
||||
# FluxGuidance の場合は conditioning を優先して辿る
|
||||
if node.get("class_type") == "FluxGuidance":
|
||||
conditioning = inputs.get("conditioning")
|
||||
if isinstance(conditioning, list) and len(conditioning) >= 1:
|
||||
resolved = resolve_text_from_ref(conditioning, visited)
|
||||
if resolved:
|
||||
return resolved
|
||||
|
||||
# よく使う接続キーを優先
|
||||
preferred_link_keys = [
|
||||
"text",
|
||||
"conditioning",
|
||||
"positive",
|
||||
"negative",
|
||||
"prompt",
|
||||
"string",
|
||||
"value",
|
||||
]
|
||||
for key in preferred_link_keys:
|
||||
value = inputs.get(key)
|
||||
if isinstance(value, list) and len(value) >= 1:
|
||||
resolved = resolve_text_from_ref(value, visited)
|
||||
if resolved:
|
||||
return resolved
|
||||
|
||||
# fallback: 全 input を走査
|
||||
for value in inputs.values():
|
||||
if isinstance(value, list) and len(value) >= 1:
|
||||
resolved = resolve_text_from_ref(value, visited)
|
||||
if resolved:
|
||||
return resolved
|
||||
|
||||
return ""
|
||||
|
||||
def get_text_from_clip(idx):
|
||||
try:
|
||||
return resolve_text_from_ref(idx)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return ""
|
||||
|
|
@ -542,36 +690,30 @@ def get_comfyui_exif_data(img: Image):
|
|||
for node_id, node_data in data.items():
|
||||
try:
|
||||
class_type = node_data.get("class_type", "")
|
||||
inputs = node_data.get("inputs", {})
|
||||
|
||||
if "CLIPTextEncode" in class_type:
|
||||
text = inputs.get("text", "")
|
||||
if isinstance(text, list):
|
||||
text = data[text[0]]["inputs"].get("text", "")
|
||||
if "CLIPTextEncode" in class_type or class_type == "ImpactWildcardProcessor":
|
||||
text = resolve_text_from_ref(node_id)
|
||||
if text:
|
||||
all_prompts.append(text.strip())
|
||||
except Exception as e:
|
||||
print(e)
|
||||
pass
|
||||
|
||||
all_prompts_str = "\nBREAK\n".join(all_prompts) if all_prompts else ""
|
||||
all_prompts_str = "\nBREAK\n".join(unique_by(all_prompts)) if all_prompts else ""
|
||||
pos_prompt = all_prompts_str
|
||||
neg_prompt = ""
|
||||
else:
|
||||
in_node = data[str(KSampler_entry["positive"][0])]
|
||||
if in_node["class_type"] != "FluxGuidance":
|
||||
pos_prompt = get_text_from_clip(KSampler_entry["positive"][0])
|
||||
else:
|
||||
pos_prompt = get_text_from_clip(in_node["inputs"]["conditioning"][0])
|
||||
positive_ref = KSampler_entry.get("positive")
|
||||
negative_ref = KSampler_entry.get("negative")
|
||||
|
||||
neg_prompt = get_text_from_clip(KSampler_entry["negative"][0])
|
||||
pos_prompt = get_text_from_clip(positive_ref) if positive_ref else ""
|
||||
neg_prompt = get_text_from_clip(negative_ref) if negative_ref else ""
|
||||
|
||||
pos_prompt_arr = unique_by(parse_prompt(pos_prompt)["pos_prompt"])
|
||||
return {
|
||||
"meta": meta,
|
||||
"pos_prompt": pos_prompt_arr,
|
||||
"pos_prompt_raw": pos_prompt,
|
||||
"neg_prompt_raw" : neg_prompt
|
||||
"neg_prompt_raw": neg_prompt
|
||||
}
|
||||
|
||||
def comfyui_exif_data_to_str(data):
|
||||
|
|
|
|||
Loading…
Reference in New Issue