Make nodes map over input lists (#579) · code/app-python-comfyui@1201d2e · GitHub
Skip to content

Commit 1201d2e

Browse files
BlenderNekoBlenderNeko
authored
Make nodes map over input lists (#579)
* allow nodes to map over lists * make work with IS_CHANGED and VALIDATE_INPUTS * give list outputs distinct socket shape * add rebatch node * add batch index logic * add repeat latent batch * deal with noise mask edge cases in latentfrombatch
1 parent 997dd1b commit 1201d2e

6 files changed

Lines changed: 250 additions & 26 deletions

File tree

‎comfy/sample.py‎

Lines changed: 13 additions & 4 deletions

‎comfy_extras/nodes_rebatch.py‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
import torch
2+
3+
class LatentRebatch:
4+
@classmethod
5+
def INPUT_TYPES(s):
6+
return {"required": { "latents": ("LATENT",),
7+
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
8+
}}
9+
RETURN_TYPES = ("LATENT",)
10+
INPUT_IS_LIST = True
11+
OUTPUT_IS_LIST = (True, )
12+
13+
FUNCTION = "rebatch"
14+
15+
CATEGORY = "latent/batch"
16+
17+
@staticmethod
18+
def get_batch(latents, list_ind, offset):
19+
'''prepare a batch out of the list of latents'''
20+
samples = latents[list_ind]['samples']
21+
shape = samples.shape
22+
mask = latents[list_ind]['noise_mask'] if 'noise_mask' in latents[list_ind] else torch.ones((shape[0], 1, shape[2]*8, shape[3]*8), device='cpu')
23+
if mask.shape[-1] != shape[-1] * 8 or mask.shape[-2] != shape[-2]:
24+
torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[-2]*8, shape[-1]*8), mode="bilinear")
25+
if mask.shape[0] < samples.shape[0]:
26+
mask = mask.repeat((shape[0] - 1) // mask.shape[0] + 1, 1, 1, 1)[:shape[0]]
27+
if 'batch_index' in latents[list_ind]:
28+
batch_inds = latents[list_ind]['batch_index']
29+
else:
30+
batch_inds = [x+offset for x in range(shape[0])]
31+
return samples, mask, batch_inds
32+
33+
@staticmethod
34+
def get_slices(indexable, num, batch_size):
35+
'''divides an indexable object into num slices of length batch_size, and a remainder'''
36+
slices = []
37+
for i in range(num):
38+
slices.append(indexable[i*batch_size:(i+1)*batch_size])
39+
if num * batch_size < len(indexable):
40+
return slices, indexable[num * batch_size:]
41+
else:
42+
return slices, None
43+
44+
@staticmethod
45+
def slice_batch(batch, num, batch_size):
46+
result = [LatentRebatch.get_slices(x, num, batch_size) for x in batch]
47+
return list(zip(*result))
48+
49+
@staticmethod
50+
def cat_batch(batch1, batch2):
51+
if batch1[0] is None:
52+
return batch2
53+
result = [torch.cat((b1, b2)) if torch.is_tensor(b1) else b1 + b2 for b1, b2 in zip(batch1, batch2)]
54+
return result
55+
56+
def rebatch(self, latents, batch_size):
57+
batch_size = batch_size[0]
58+
59+
output_list = []
60+
current_batch = (None, None, None)
61+
processed = 0
62+
63+
for i in range(len(latents)):
64+
# fetch new entry of list
65+
#samples, masks, indices = self.get_batch(latents, i)
66+
next_batch = self.get_batch(latents, i, processed)
67+
processed += len(next_batch[2])
68+
# set to current if current is None
69+
if current_batch[0] is None:
70+
current_batch = next_batch
71+
# add previous to list if dimensions do not match
72+
elif next_batch[0].shape[-1] != current_batch[0].shape[-1] or next_batch[0].shape[-2] != current_batch[0].shape[-2]:
73+
sliced, _ = self.slice_batch(current_batch, 1, batch_size)
74+
output_list.append({'samples': sliced[0][0], 'noise_mask': sliced[1][0], 'batch_index': sliced[2][0]})
75+
current_batch = next_batch
76+
# cat if everything checks out
77+
else:
78+
current_batch = self.cat_batch(current_batch, next_batch)
79+
80+
# add to list if dimensions gone above target batch size
81+
if current_batch[0].shape[0] > batch_size:
82+
num = current_batch[0].shape[0] // batch_size
83+
sliced, remainder = self.slice_batch(current_batch, num, batch_size)
84+
85+
for i in range(num):
86+
output_list.append({'samples': sliced[0][i], 'noise_mask': sliced[1][i], 'batch_index': sliced[2][i]})
87+
88+
current_batch = remainder
89+
90+
#add remainder
91+
if current_batch[0] is not None:
92+
sliced, _ = self.slice_batch(current_batch, 1, batch_size)
93+
output_list.append({'samples': sliced[0][0], 'noise_mask': sliced[1][0], 'batch_index': sliced[2][0]})
94+
95+
#get rid of empty masks
96+
for s in output_list:
97+
if s['noise_mask'].mean() == 1.0:
98+
del s['noise_mask']
99+
100+
return (output_list,)
101+
102+
NODE_CLASS_MAPPINGS = {
103+
"RebatchLatents": LatentRebatch,
104+
}
105+
106+
NODE_DISPLAY_NAME_MAPPINGS = {
107+
"RebatchLatents": "Rebatch Latents",
108+
}

‎execution.py‎

Lines changed: 76 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -26,20 +26,81 @@ def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_da
2626
input_data_all[x] = obj
2727
else:
2828
if ("required" in valid_inputs and x in valid_inputs["required"]) or ("optional" in valid_inputs and x in valid_inputs["optional"]):
29-
input_data_all[x] = input_data
29+
input_data_all[x] = [input_data]
3030

3131
if "hidden" in valid_inputs:
3232
h = valid_inputs["hidden"]
3333
for x in h:
3434
if h[x] == "PROMPT":
35-
input_data_all[x] = prompt
35+
input_data_all[x] = [prompt]
3636
if h[x] == "EXTRA_PNGINFO":
3737
if "extra_pnginfo" in extra_data:
38-
input_data_all[x] = extra_data['extra_pnginfo']
38+
input_data_all[x] = [extra_data['extra_pnginfo']]
3939
if h[x] == "UNIQUE_ID":
40-
input_data_all[x] = unique_id
40+
input_data_all[x] = [unique_id]
4141
return input_data_all
4242

43+
def map_node_over_list(obj, input_data_all, func, allow_interrupt=False):
44+
# check if node wants the lists
45+
intput_is_list = False
46+
if hasattr(obj, "INPUT_IS_LIST"):
47+
intput_is_list = obj.INPUT_IS_LIST
48+
49+
max_len_input = max([len(x) for x in input_data_all.values()])
50+
51+
# get a slice of inputs, repeat last input when list isn't long enough
52+
def slice_dict(d, i):
53+
d_new = dict()
54+
for k,v in d.items():
55+
d_new[k] = v[i if len(v) > i else -1]
56+
return d_new
57+
58+
results = []
59+
if intput_is_list:
60+
if allow_interrupt:
61+
nodes.before_node_execution()
62+
results.append(getattr(obj, func)(**input_data_all))
63+
else:
64+
for i in range(max_len_input):
65+
if allow_interrupt:
66+
nodes.before_node_execution()
67+
results.append(getattr(obj, func)(**slice_dict(input_data_all, i)))
68+
return results
69+
70+
def get_output_data(obj, input_data_all):
71+
72+
results = []
73+
uis = []
74+
return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True)
75+
76+
for r in return_values:
77+
if isinstance(r, dict):
78+
if 'ui' in r:
79+
uis.append(r['ui'])
80+
if 'result' in r:
81+
results.append(r['result'])
82+
else:
83+
results.append(r)
84+
85+
output = []
86+
if len(results) > 0:
87+
# check which outputs need concatenating
88+
output_is_list = [False] * len(results[0])
89+
if hasattr(obj, "OUTPUT_IS_LIST"):
90+
output_is_list = obj.OUTPUT_IS_LIST
91+
92+
# merge node execution results
93+
for i, is_list in zip(range(len(results[0])), output_is_list):
94+
if is_list:
95+
output.append([x for o in results for x in o[i]])
96+
else:
97+
output.append([o[i] for o in results])
98+
99+
ui = dict()
100+
if len(uis) > 0:
101+
ui = {k: [y for x in uis for y in x[k]] for k in uis[0].keys()}
102+
return output, ui
103+
43104
def recursive_execute(server, prompt, outputs, current_item, extra_data, executed):
44105
unique_id = current_item
45106
inputs = prompt[unique_id]['inputs']
@@ -63,13 +124,11 @@ def recursive_execute(server, prompt, outputs, current_item, extra_data, execute
63124
server.send_sync("executing", { "node": unique_id }, server.client_id)
64125
obj = class_def()
65126

66-
nodes.before_node_execution()
67-
outputs[unique_id] = getattr(obj, obj.FUNCTION)(**input_data_all)
68-
if "ui" in outputs[unique_id]:
127+
output_data, output_ui = get_output_data(obj, input_data_all)
128+
outputs[unique_id] = output_data
129+
if len(output_ui) > 0:
69130
if server.client_id is not None:
70-
server.send_sync("executed", { "node": unique_id, "output": outputs[unique_id]["ui"] }, server.client_id)
71-
if "result" in outputs[unique_id]:
72-
outputs[unique_id] = outputs[unique_id]["result"]
131+
server.send_sync("executed", { "node": unique_id, "output": output_ui }, server.client_id)
73132
executed.add(unique_id)
74133

75134
def recursive_will_execute(prompt, outputs, current_item):
@@ -105,7 +164,8 @@ def recursive_output_delete_if_changed(prompt, old_prompt, outputs, current_item
105164
input_data_all = get_input_data(inputs, class_def, unique_id, outputs)
106165
if input_data_all is not None:
107166
try:
108-
is_changed = class_def.IS_CHANGED(**input_data_all)
167+
#is_changed = class_def.IS_CHANGED(**input_data_all)
168+
is_changed = map_node_over_list(class_def, input_data_all, "IS_CHANGED")
109169
prompt[unique_id]['is_changed'] = is_changed
110170
except:
111171
to_delete = True
@@ -261,9 +321,11 @@ def validate_inputs(prompt, item, validated):
261321

262322
if hasattr(obj_class, "VALIDATE_INPUTS"):
263323
input_data_all = get_input_data(inputs, obj_class, unique_id)
264-
ret = obj_class.VALIDATE_INPUTS(**input_data_all)
265-
if ret != True:
266-
return (False, "{}, {}".format(class_type, ret))
324+
#ret = obj_class.VALIDATE_INPUTS(**input_data_all)
325+
ret = map_node_over_list(obj_class, input_data_all, "VALIDATE_INPUTS")
326+
for r in ret:
327+
if r != True:
328+
return (False, "{}, {}".format(class_type, r))
267329
else:
268330
if isinstance(type_input, list):
269331
if val not in type_input:

‎nodes.py‎

Lines changed: 50 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -629,18 +629,57 @@ class LatentFromBatch:
629629
def INPUT_TYPES(s):
630630
return {"required": { "samples": ("LATENT",),
631631
"batch_index": ("INT", {"default": 0, "min": 0, "max": 63}),
632+
"length": ("INT", {"default": 1, "min": 1, "max": 64}),
632633
}}
633634
RETURN_TYPES = ("LATENT",)
634-
FUNCTION = "rotate"
635+
FUNCTION = "frombatch"
635636

636-
CATEGORY = "latent"
637+
CATEGORY = "latent/batch"
637638

638-
def rotate(self, samples, batch_index):
639+
def frombatch(self, samples, batch_index, length):
639640
s = samples.copy()
640641
s_in = samples["samples"]
641642
batch_index = min(s_in.shape[0] - 1, batch_index)
642-
s["samples"] = s_in[batch_index:batch_index + 1].clone()
643-
s["batch_index"] = batch_index
643+
length = min(s_in.shape[0] - batch_index, length)
644+
s["samples"] = s_in[batch_index:batch_index + length].clone()
645+
if "noise_mask" in samples:
646+
masks = samples["noise_mask"]
647+
if masks.shape[0] == 1:
648+
s["noise_mask"] = masks.clone()
649+
else:
650+
if masks.shape[0] < s_in.shape[0]:
651+
masks = masks.repeat(math.ceil(s_in.shape[0] / masks.shape[0]), 1, 1, 1)[:s_in.shape[0]]
652+
s["noise_mask"] = masks[batch_index:batch_index + length].clone()
653+
if "batch_index" not in s:
654+
s["batch_index"] = [x for x in range(batch_index, batch_index+length)]
655+
else:
656+
s["batch_index"] = samples["batch_index"][batch_index:batch_index + length]
657+
return (s,)
658+
659+
class RepeatLatentBatch:
660+
@classmethod
661+
def INPUT_TYPES(s):
662+
return {"required": { "samples": ("LATENT",),
663+
"amount": ("INT", {"default": 1, "min": 1, "max": 64}),
664+
}}
665+
RETURN_TYPES = ("LATENT",)
666+
FUNCTION = "repeat"
667+
668+
CATEGORY = "latent/batch"
669+
670+
def repeat(self, samples, amount):
671+
s = samples.copy()
672+
s_in = samples["samples"]
673+
674+
s["samples"] = s_in.repeat((amount, 1,1,1))
675+
if "noise_mask" in samples and samples["noise_mask"].shape[0] > 1:
676+
masks = samples["noise_mask"]
677+
if masks.shape[0] < s_in.shape[0]:
678+
masks = masks.repeat(math.ceil(s_in.shape[0] / masks.shape[0]), 1, 1, 1)[:s_in.shape[0]]
679+
s["noise_mask"] = samples["noise_mask"].repeat((amount, 1,1,1))
680+
if "batch_index" in s:
681+
offset = max(s["batch_index"]) - min(s["batch_index"]) + 1
682+
s["batch_index"] = s["batch_index"] + [x + (i * offset) for i in range(1, amount) for x in s["batch_index"]]
644683
return (s,)
645684

646685
class LatentUpscale:
@@ -805,8 +844,8 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
805844
if disable_noise:
806845
noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
807846
else:
808-
skip = latent["batch_index"] if "batch_index" in latent else 0
809-
noise = comfy.sample.prepare_noise(latent_image, seed, skip)
847+
batch_inds = latent["batch_index"] if "batch_index" in latent else None
848+
noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
810849

811850
noise_mask = None
812851
if "noise_mask" in latent:
@@ -1170,6 +1209,7 @@ def expand_image(self, image, left, top, right, bottom, feathering):
11701209
"EmptyLatentImage": EmptyLatentImage,
11711210
"LatentUpscale": LatentUpscale,
11721211
"LatentFromBatch": LatentFromBatch,
1212+
"RepeatLatentBatch": RepeatLatentBatch,
11731213
"SaveImage": SaveImage,
11741214
"PreviewImage": PreviewImage,
11751215
"LoadImage": LoadImage,
@@ -1244,6 +1284,8 @@ def expand_image(self, image, left, top, right, bottom, feathering):
12441284
"EmptyLatentImage": "Empty Latent Image",
12451285
"LatentUpscale": "Upscale Latent",
12461286
"LatentComposite": "Latent Composite",
1287+
"LatentFromBatch" : "Latent From Batch",
1288+
"RepeatLatentBatch": "Repeat Latent Batch",
12471289
# Image
12481290
"SaveImage": "Save Image",
12491291
"PreviewImage": "Preview Image",
@@ -1299,3 +1341,4 @@ def init_custom_nodes():
12991341
load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_upscale_model.py"))
13001342
load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_post_processing.py"))
13011343
load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_mask.py"))
1344+
load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_rebatch.py"))

‎server.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,7 @@ async def get_object_info(request):
268268
info = {}
269269
info['input'] = obj_class.INPUT_TYPES()
270270
info['output'] = obj_class.RETURN_TYPES
271+
info['output_is_list'] = obj_class.OUTPUT_IS_LIST if hasattr(obj_class, 'OUTPUT_IS_LIST') else [False] * len(obj_class.RETURN_TYPES)
271272
info['output_name'] = obj_class.RETURN_NAMES if hasattr(obj_class, 'RETURN_NAMES') else info['output']
272273
info['name'] = x
273274
info['display_name'] = nodes.NODE_DISPLAY_NAME_MAPPINGS[x] if x in nodes.NODE_DISPLAY_NAME_MAPPINGS.keys() else x

‎web/scripts/app.js‎

Lines changed: 2 additions & 1 deletion

0 commit comments

Comments
 (0)