google-labs-jules[bot]
b6f8c16b0b
- Implemented a text-based table formatter for batches. - Updated Load CSV node to show a table preview in the workflow. - Added a new Preview Batch node for visualizing any batch data. - Added a JavaScript extension to handle node visualization. - Updated English and Spanish documentation.
297 lines
8.1 KiB
Python
297 lines
8.1 KiB
Python
import collections
|
|
import itertools
|
|
import csv
|
|
import os
|
|
|
|
from . import register_node
|
|
|
|
|
|
def format_batch_as_table(batch, current_index=-1, max_rows=20):
|
|
if not batch:
|
|
return "(Empty Batch)"
|
|
|
|
# Get all unique keys across all rows
|
|
keys = []
|
|
for row in batch:
|
|
for k in row.keys():
|
|
if k not in keys:
|
|
keys.append(k)
|
|
|
|
if not keys:
|
|
return "(Empty rows)"
|
|
|
|
# Truncate values for display
|
|
def trunc(v, l=20):
|
|
s = str(v).replace('\n', ' ')
|
|
return (s[:l-3] + "...") if len(s) > l else s
|
|
|
|
# Calculate widths
|
|
widths = {k: len(str(k)) for k in keys}
|
|
rows_to_show = batch[:max_rows]
|
|
for row in rows_to_show:
|
|
for k in keys:
|
|
widths[k] = max(widths[k], len(trunc(row.get(k, ""))))
|
|
|
|
# Build header
|
|
header = "| " + " | ".join(str(k).ljust(widths[k]) for k in keys) + " |"
|
|
sep = "|-" + "-|-".join("-" * widths[k] for k in keys) + "-|"
|
|
|
|
idx_to_highlight = (current_index % len(batch)) if (len(batch) > 0 and current_index != -1) else (0 if len(batch) > 0 else -1)
|
|
|
|
lines = [header, sep]
|
|
for i, row in enumerate(rows_to_show):
|
|
prefix = ">" if i == idx_to_highlight else " "
|
|
line = f"{prefix}| " + " | ".join(trunc(row.get(k, "")).ljust(widths[k]) for k in keys) + " |"
|
|
lines.append(line)
|
|
|
|
if len(batch) > max_rows:
|
|
lines.append(f"... and {len(batch) - max_rows} more rows")
|
|
|
|
return "\n".join(lines)
|
|
|
|
|
|
@register_node
|
|
class MakeBatch:
|
|
"""Turns a sequence into a batch with one attribute."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"sequence": ("SEQUENCE", ),
|
|
"name": ("STRING", {"default": ''}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("BATCH", "INT")
|
|
RETURN_NAMES = ("batch", "count")
|
|
FUNCTION = "go"
|
|
CATEGORY = "🔁 Sequential Batcher/Batch"
|
|
|
|
def merge_dicts(self, *dicts):
|
|
#return collections.ChainMap(*reversed(dicts))
|
|
return dict(itertools.chain.from_iterable(d.items() for d in dicts))
|
|
|
|
def go(self, sequence, name):
|
|
result = [{name: value} for value in sequence]
|
|
return (result, len(result))
|
|
|
|
|
|
@register_node
|
|
class CombineBatches(MakeBatch):
|
|
"""Combines multiple batches."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"a": ("BATCH", ),
|
|
"method": (("zip", "product"), {"default": "zip"}),
|
|
},
|
|
"optional": {
|
|
x: ("BATCH", ) for x in ('b', 'c', 'd', 'e')
|
|
}
|
|
}
|
|
|
|
def go(self, method, **kwargs):
|
|
method = {'product': itertools.product, 'zip': zip}[method]
|
|
result = [self.merge_dicts(*steps) for steps in method(*kwargs.values())]
|
|
return (result, len(result))
|
|
|
|
|
|
@register_node
|
|
class EnumerateBatch(MakeBatch):
|
|
"""Combines multiple batches."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"batch": ("BATCH", ),
|
|
"name": ("STRING", {"default": ''}),
|
|
},
|
|
}
|
|
|
|
def go(self, batch, name):
|
|
result = [self.merge_dicts(step, {name: n}) for n, step in enumerate(batch)]
|
|
return (result, len(result))
|
|
|
|
|
|
@register_node
|
|
class GetBatchStep:
|
|
"""Gets the batch step by number."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"batch": ("BATCH", ),
|
|
"step": ("INT", {"default": 0}),
|
|
"wrap": (("repeat", "clamp"), {"default": "repeat"}),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("ATTRIBUTES", )
|
|
RETURN_NAMES = ("attributes", )
|
|
FUNCTION = "go"
|
|
CATEGORY = "🔁 Sequential Batcher/Batch"
|
|
|
|
def go(self, batch, step, wrap):
|
|
if wrap == 'repeat':
|
|
while step < 0:
|
|
step += len(batch)
|
|
step = step % len(batch)
|
|
elif wrap == 'clamp':
|
|
step = max(min(step, len(batch)-1), 0)
|
|
return (batch[step], )
|
|
|
|
|
|
@register_node
|
|
class FormatAttributes:
|
|
"""Applies attributes to a format string."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"attributes": ("ATTRIBUTES",),
|
|
"format": ("STRING", {'default': '', 'multiline': True, "dynamicPrompts": False})
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", )
|
|
RETURN_NAMES = ("string", )
|
|
FUNCTION = "go"
|
|
CATEGORY = "🔁 Sequential Batcher/Batch"
|
|
|
|
def go(self, attributes, format):
|
|
return (format.format(**attributes), )
|
|
|
|
|
|
class AnyType(str):
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
|
|
@register_node
|
|
class GetAttribute:
|
|
"""Gets a named attribute from a step."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"attributes": ("ATTRIBUTES", ),
|
|
"name": ("STRING", {"default": ''}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (AnyType("*"), )
|
|
RETURN_NAMES = ("value", )
|
|
FUNCTION = "go"
|
|
CATEGORY = "🔁 Sequential Batcher/Batch"
|
|
|
|
def go(self, attributes, name):
|
|
return (attributes[name], )
|
|
|
|
|
|
# Dynamically register typed attribute getters to avoid wildcard bug
|
|
# https://github.com/comfyanonymous/ComfyUI/pull/770
|
|
for t in ('INT', 'FLOAT', 'STRING'):
|
|
register_node(type(
|
|
'GetAttribute'+t.title(),
|
|
(GetAttribute, ),
|
|
{
|
|
'RETURN_TYPES': (t, ),
|
|
}
|
|
))
|
|
|
|
|
|
@register_node
|
|
class LoadCSV:
|
|
"""Loads a CSV file into a batch."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"path": ("STRING", {"default": "prompt.csv"}),
|
|
"delimiter": (["comma", "tab", "semicolon"], {"default": "comma"}),
|
|
"quotechar": (["\"", "'"], {"default": "\""}),
|
|
},
|
|
"optional": {
|
|
"index": ("INT", {"default": -1, "min": -1, "max": 999999}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BATCH", "ATTRIBUTES", "INT")
|
|
RETURN_NAMES = ("batch", "current_attributes", "count")
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "go"
|
|
CATEGORY = "🔁 Sequential Batcher/Batch"
|
|
|
|
def go(self, path, delimiter, quotechar, index=-1):
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError(f"CSV file not found: {path}")
|
|
|
|
delim = {
|
|
"comma": ",",
|
|
"tab": "\t",
|
|
"semicolon": ";",
|
|
}[delimiter]
|
|
|
|
batch = []
|
|
with open(path, mode='r', encoding='utf-8-sig') as f:
|
|
reader = csv.DictReader(f, delimiter=delim, quotechar=quotechar)
|
|
for row in reader:
|
|
# Clean up keys and values
|
|
clean_row = {k.strip() if k else k: v.strip() if v else v for k, v in row.items()}
|
|
batch.append(clean_row)
|
|
|
|
count = len(batch)
|
|
current_attributes = batch[index % count] if count > 0 and index != -1 else (batch[0] if count > 0 else {})
|
|
|
|
table = format_batch_as_table(batch, index)
|
|
|
|
return {"ui": {"text": [table]}, "result": (batch, current_attributes, count)}
|
|
|
|
|
|
@register_node
|
|
class PreviewBatch:
|
|
"""Shows a table of the batch."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"batch": ("BATCH",),
|
|
},
|
|
"optional": {
|
|
"index": ("INT", {"default": -1, "min": -1, "max": 999999}),
|
|
"max_rows": ("INT", {"default": 20, "min": 1, "max": 1000}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BATCH",)
|
|
RETURN_NAMES = ("batch",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "go"
|
|
CATEGORY = "🔁 Sequential Batcher/Batch"
|
|
|
|
def go(self, batch, index=-1, max_rows=20):
|
|
table = format_batch_as_table(batch, index, max_rows)
|
|
return {"ui": {"text": [table]}, "result": (batch,)}
|
|
|
|
|
|
@register_node
|
|
class BatchToList:
|
|
"""Converts a batch into a list."""
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"batch": ("BATCH",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("ATTRIBUTES",)
|
|
RETURN_NAMES = ("attributes",)
|
|
OUTPUT_IS_LIST = (True,)
|
|
FUNCTION = "go"
|
|
CATEGORY = "🔁 Sequential Batcher/Batch"
|
|
|
|
def go(self, batch):
|
|
return (batch,)
|