Skip to content

Commit

Permalink
metadata readability
Browse files Browse the repository at this point in the history
  • Loading branch information
@d3x-at
d3x-at committed Jul 7, 2024
1 parent 43a80a8 commit 33b8cb8
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 31 deletions.
57 changes: 44 additions & 13 deletions src/sd_parsers/parsers/_comfyui.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
POSITIVE_PROMPT_KEYS = ["text", "positive"]
NEGATIVE_PROMPT_KEYS = ["text", "negative"]
IGNORE_LINK_TYPES_PROMPT = ["CLIP"]
IGNORE_CLASS_TYPES = ["ConditioningCombine"]


class ComfyUIParser(Parser):
Expand Down Expand Up @@ -158,11 +159,38 @@ def _get_input_values(self, inputs):
except Exception:
return {}

def _get_trace_metadata(self, trace: list[int]):
metadata = {}

for node_id in trace:
try:
node = self.prompt[node_id]
class_type = node["class_type"]

if class_type in IGNORE_CLASS_TYPES:
continue

value = self._get_input_values(node["inputs"])

except KeyError:
continue

try:
entry = metadata[class_type]
if isinstance(entry, list):
entry.append(value)
else:
metadata[class_type] = [entry, value]
except KeyError:
metadata[class_type] = value

return metadata

def _get_model(self, initial_node_id: int) -> Optional[Model]:
"""Get the first model reached from the given node_id"""
logger.debug("looking for model: #%s", initial_node_id)

for node_id, node, _ in self._traverse(initial_node_id):
for node_id, node, trace in self._traverse(initial_node_id):
try:
inputs = dict(node["inputs"])
ckpt_name = inputs.pop("ckpt_name")
Expand All @@ -173,6 +201,8 @@ def _get_model(self, initial_node_id: int) -> Optional[Model]:
logger.debug("found model #%d: %s", node_id, ckpt_name)

metadata = self._get_input_values(inputs)
metadata.update(self._get_trace_metadata(trace))

model = Model(model_id=node_id, name=ckpt_name, metadata=metadata)
return model

Expand All @@ -188,26 +218,27 @@ def check_inputs(node_id: int, inputs: Dict, trace: List[int]) -> bool:
found_prompt = False
for key in text_keys:
try:
text = inputs[key]
text = inputs.pop(key)
except KeyError:
continue

if isinstance(text, str):
logger.debug("found prompt %s#%d: %s", key, node_id, text)

metadata = {}
for tmp_id in trace[:-1]:
with suppress(KeyError):
tmp_node = self.prompt[tmp_id]
metadata[tmp_node["class_type"], tmp_id] = self._get_input_values(
tmp_node["inputs"]
)
metadata = self._get_input_values(inputs)
metadata.update(self._get_trace_metadata(trace))

prompts.append(Prompt(value=text.strip(), prompt_id=node_id, metadata=metadata))
prompts.append(
Prompt(
value=text.strip(),
prompt_id=node_id,
metadata=metadata,
)
)
found_prompt = True

if found_prompt:
self.processed_nodes.update(trace)
self.processed_nodes.update([node_id], trace)
return False
return True

Expand All @@ -221,7 +252,7 @@ def check_inputs(node_id: int, inputs: Dict, trace: List[int]) -> bool:
except KeyError:
pass
else:
recurse = check_inputs(node_id, inputs, trace)
recurse = check_inputs(node_id, dict(inputs), trace)

node_id, node, trace = prompt_iterator.send(recurse)

Expand All @@ -238,7 +269,7 @@ def traverse_inner(node_id: int, trace: List[int]):
visited.add(node_id)

with suppress(KeyError):
recurse = yield node_id, self.prompt[node_id], trace
recurse = yield node_id, self.prompt[node_id], trace[:-1]
if recurse is False:
return

Expand Down
26 changes: 8 additions & 18 deletions tests/resources/parsers/ComfyUI/night_evening_day_morning_cropped.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,7 @@
value="(best quality) (daytime:1.2) sky (blue)",
prompt_id=17,
metadata={
("ConditioningCombine", 19): {},
("ConditioningCombine", 12): {},
("ConditioningSetArea", 11): {
"ConditioningSetArea": {
"width": 704,
"height": 384,
"x": 0,
Expand All @@ -37,10 +35,7 @@
"(galaxy:1.2) (space) (universe)",
prompt_id=14,
metadata={
("ConditioningCombine", 19): {},
("ConditioningCombine", 12): {},
("ConditioningCombine", 35): {},
("ConditioningSetArea", 34): {
"ConditioningSetArea": {
"width": 704,
"height": 384,
"x": 0,
Expand All @@ -54,10 +49,7 @@
"(HDR:1.2) (sunset:1.3)",
prompt_id=13,
metadata={
("ConditioningCombine", 19): {},
("ConditioningCombine", 12): {},
("ConditioningCombine", 35): {},
("ConditioningSetArea", 18): {
"ConditioningSetArea": {
"width": 704,
"height": 384,
"x": 0,
Expand All @@ -70,9 +62,7 @@
value="(masterpiece) (best quality) morning sky",
prompt_id=33,
metadata={
("ConditioningCombine", 19): {},
("ConditioningCombine", 10): {},
("ConditioningSetArea", 15): {
"ConditioningSetArea": {
"width": 704,
"height": 384,
"x": 0,
Expand All @@ -86,10 +76,7 @@
"amazing view nature photograph forest mountains ocean (sky) "
"national park scenery",
prompt_id=6,
metadata={
("ConditioningCombine", 19): {},
("ConditioningCombine", 10): {},
},
metadata={},
),
],
negative_prompts=[
Expand All @@ -104,6 +91,7 @@
"(forehead mark) (depth of field) "
"(emotionless) (penis) (pumpkin)",
prompt_id=7,
metadata={},
)
],
),
Expand All @@ -130,6 +118,7 @@
"mountains ocean daytime night evening morning, "
"(sky:1.2)",
prompt_id=26,
metadata={},
)
],
negative_prompts=[
Expand All @@ -144,6 +133,7 @@
"(forehead mark) (depth of field) "
"(emotionless) (penis) (pumpkin)",
prompt_id=27,
metadata={},
)
],
),
Expand Down

0 comments on commit 33b8cb8

Please sign in to comment.