Skip to content

Commit b0f9e32

Browse files
alexisrollandalexisrolland
authored
Add output socket to save nodes (Comfy-Org#13866)
1 parent 0d8b751 commit b0f9e32

4 files changed

Lines changed: 52 additions & 58 deletions

File tree

‎comfy_extras/nodes_audio.py‎

Lines changed: 17 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -158,36 +158,36 @@ def define_schema(cls):
158158
return IO.Schema(
159159
node_id="SaveAudio",
160160
search_aliases=["export flac"],
161-
display_name="Save Audio (FLAC) (Deprecated)",
161+
display_name="Save Audio (FLAC) (DEPRECATED)",
162162
category="audio",
163163
essentials_category="Audio",
164164
inputs=[
165165
IO.Audio.Input("audio"),
166166
IO.String.Input("filename_prefix", default="audio/ComfyUI"),
167167
],
168168
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
169-
is_output_node=True,
170169
is_deprecated=True,
170+
is_output_node=True,
171+
outputs=[IO.Audio.Output("audio")]
171172
)
172173

173174
@classmethod
174175
def execute(cls, audio, filename_prefix="ComfyUI", format="flac") -> IO.NodeOutput:
175176
if audio is None:
176177
raise ValueError("SaveAudio: input audio is None (source video may have no audio track).")
177178
return IO.NodeOutput(
179+
audio,
178180
ui=UI.AudioSaveHelper.get_save_audio_ui(audio, filename_prefix=filename_prefix, cls=cls, format=format)
179181
)
180182

181-
save_flac = execute # TODO: remove
182-
183183

184184
class SaveAudioMP3(IO.ComfyNode):
185185
@classmethod
186186
def define_schema(cls):
187187
return IO.Schema(
188188
node_id="SaveAudioMP3",
189189
search_aliases=["export mp3"],
190-
display_name="Save Audio (MP3) (Deprecated)",
190+
display_name="Save Audio (MP3) (DEPRECATED)",
191191
category="audio",
192192
essentials_category="Audio",
193193
inputs=[
@@ -196,53 +196,53 @@ def define_schema(cls):
196196
IO.Combo.Input("quality", options=["V0", "128k", "320k"], default="V0"),
197197
],
198198
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
199-
is_output_node=True,
200199
is_deprecated=True,
200+
is_output_node=True,
201+
outputs=[IO.Audio.Output("audio")]
201202
)
202203

203204
@classmethod
204205
def execute(cls, audio, filename_prefix="ComfyUI", format="mp3", quality="128k") -> IO.NodeOutput:
205206
if audio is None:
206207
raise ValueError("SaveAudioMP3: input audio is None (source video may have no audio track).")
207208
return IO.NodeOutput(
209+
audio,
208210
ui=UI.AudioSaveHelper.get_save_audio_ui(
209211
audio, filename_prefix=filename_prefix, cls=cls, format=format, quality=quality
210212
)
211213
)
212214

213-
save_mp3 = execute # TODO: remove
214-
215215

216216
class SaveAudioOpus(IO.ComfyNode):
217217
@classmethod
218218
def define_schema(cls):
219219
return IO.Schema(
220220
node_id="SaveAudioOpus",
221221
search_aliases=["export opus"],
222-
display_name="Save Audio (Opus) (Deprecated)",
222+
display_name="Save Audio (Opus) (DEPRECATED)",
223223
category="audio",
224224
inputs=[
225225
IO.Audio.Input("audio"),
226226
IO.String.Input("filename_prefix", default="audio/ComfyUI"),
227227
IO.Combo.Input("quality", options=["64k", "96k", "128k", "192k", "320k"], default="128k"),
228228
],
229229
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
230-
is_output_node=True,
231230
is_deprecated=True,
231+
is_output_node=True,
232+
outputs=[IO.Audio.Output("audio")]
232233
)
233234

234235
@classmethod
235236
def execute(cls, audio, filename_prefix="ComfyUI", format="opus", quality="V3") -> IO.NodeOutput:
236237
if audio is None:
237238
raise ValueError("SaveAudioOpus: input audio is None (source video may have no audio track).")
238239
return IO.NodeOutput(
240+
audio,
239241
ui=UI.AudioSaveHelper.get_save_audio_ui(
240242
audio, filename_prefix=filename_prefix, cls=cls, format=format, quality=quality
241243
)
242244
)
243245

244-
save_opus = execute # TODO: remove
245-
246246

247247
class SaveAudioAdvanced(IO.ComfyNode):
248248
@classmethod
@@ -258,10 +258,7 @@ def define_schema(cls):
258258
IO.String.Input(
259259
"filename_prefix",
260260
default="audio/ComfyUI",
261-
tooltip=(
262-
"The prefix for the file to save. May include formatting tokens "
263-
"such as %date:yyyy-MM-dd%."
264-
),
261+
tooltip=("The prefix for the file to save. May include formatting tokens such as %date:yyyy-MM-dd%."),
265262
),
266263
IO.DynamicCombo.Input(
267264
"format",
@@ -279,6 +276,7 @@ def define_schema(cls):
279276
],
280277
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
281278
is_output_node=True,
279+
outputs=[IO.Audio.Output("audio")],
282280
)
283281

284282
@classmethod
@@ -289,7 +287,7 @@ def execute(cls, audio, filename_prefix: str, format: dict) -> IO.NodeOutput:
289287
ui=UI.AudioSaveHelper.get_save_audio_ui(audio, filename_prefix=filename_prefix, cls=cls, format=file_format, quality=quality)
290288
else:
291289
ui=UI.AudioSaveHelper.get_save_audio_ui(audio, filename_prefix=filename_prefix, cls=cls, format=file_format)
292-
return IO.NodeOutput(ui=ui)
290+
return IO.NodeOutput(audio, ui=ui)
293291

294292

295293
class PreviewAudio(IO.ComfyNode):
@@ -305,13 +303,14 @@ def define_schema(cls):
305303
],
306304
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
307305
is_output_node=True,
306+
outputs=[IO.Audio.Output("audio")]
308307
)
309308

310309
@classmethod
311310
def execute(cls, audio) -> IO.NodeOutput:
312311
if audio is None:
313312
raise ValueError("PreviewAudio: input audio is None (source video may have no audio track).")
314-
return IO.NodeOutput(ui=UI.PreviewAudio(audio, cls=cls))
313+
return IO.NodeOutput(audio, ui=UI.PreviewAudio(audio, cls=cls))
315314

316315
save_flac = execute # TODO: remove
317316

‎comfy_extras/nodes_images.py‎

Lines changed: 16 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -214,11 +214,13 @@ def define_schema(cls):
214214
],
215215
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
216216
is_output_node=True,
217+
outputs=[IO.Image.Output(display_name="images")]
217218
)
218219

219220
@classmethod
220221
def execute(cls, images, fps, filename_prefix, lossless, quality, method, num_frames=0) -> IO.NodeOutput:
221222
return IO.NodeOutput(
223+
images,
222224
ui=UI.ImageSaveHelper.get_save_animated_webp_ui(
223225
images=images,
224226
filename_prefix=filename_prefix,
@@ -230,8 +232,6 @@ def execute(cls, images, fps, filename_prefix, lossless, quality, method, num_fr
230232
)
231233
)
232234

233-
save_images = execute # TODO: remove
234-
235235

236236
class SaveAnimatedPNG(IO.ComfyNode):
237237

@@ -249,11 +249,13 @@ def define_schema(cls):
249249
],
250250
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
251251
is_output_node=True,
252+
outputs=[IO.Image.Output(display_name="images")]
252253
)
253254

254255
@classmethod
255256
def execute(cls, images, fps, compress_level, filename_prefix="ComfyUI") -> IO.NodeOutput:
256257
return IO.NodeOutput(
258+
images,
257259
ui=UI.ImageSaveHelper.get_save_animated_png_ui(
258260
images=images,
259261
filename_prefix=filename_prefix,
@@ -263,8 +265,6 @@ def execute(cls, images, fps, compress_level, filename_prefix="ComfyUI") -> IO.N
263265
)
264266
)
265267

266-
save_images = execute # TODO: remove
267-
268268

269269
class ImageStitch(IO.ComfyNode):
270270
"""Upstreamed from https://github.com/kijai/ComfyUI-KJNodes"""
@@ -513,6 +513,7 @@ def define_schema(cls):
513513
],
514514
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
515515
is_output_node=True,
516+
outputs=[IO.SVG.Output("svg")],
516517
)
517518

518519
@classmethod
@@ -562,9 +563,7 @@ def replacement(match):
562563

563564
results.append(UI.SavedResult(filename=file, subfolder=subfolder, type=IO.FolderType.output))
564565
counter += 1
565-
return IO.NodeOutput(ui={"images": results})
566-
567-
save_svg = execute # TODO: remove
566+
return IO.NodeOutput(svg, ui={"images": results})
568567

569568

570569
class GetImageSize(IO.ComfyNode):
@@ -1157,40 +1156,27 @@ def define_schema(cls):
11571156
IO.String.Input(
11581157
"filename_prefix",
11591158
default="ComfyUI",
1160-
tooltip=(
1161-
"The prefix for the file to save. May include formatting tokens "
1162-
"such as %date:yyyy-MM-dd% or %Empty Latent Image.width%."
1163-
),
1159+
tooltip=("The prefix for the file to save. May include formatting tokens such as %date:yyyy-MM-dd% or %Empty Latent Image.width%."),
11641160
),
11651161
IO.DynamicCombo.Input(
11661162
"format",
11671163
options=[
11681164
IO.DynamicCombo.Option("png", [
1169-
IO.Combo.Input("bit_depth", options=["8-bit", "16-bit"],
1170-
default="8-bit", advanced=True),
1171-
IO.Combo.Input("input_color_space", options=["sRGB"],
1172-
default="sRGB", advanced=True),
1165+
IO.Combo.Input("bit_depth", options=["8-bit", "16-bit"], default="8-bit", advanced=True),
1166+
IO.Combo.Input("input_color_space", options=["sRGB"], default="sRGB", advanced=True),
11731167
]),
11741168
IO.DynamicCombo.Option("exr", [
1175-
IO.Combo.Input("bit_depth", options=["32-bit float"],
1176-
default="32-bit float", advanced=True),
1169+
IO.Combo.Input("bit_depth", options=["32-bit float"], default="32-bit float", advanced=True),
11771170
IO.Combo.Input(
11781171
"input_color_space",
11791172
options=["sRGB", "HDR", "linear"],
11801173
default="sRGB",
11811174
advanced=True,
11821175
tooltip=(
1183-
"Colorspace of the input tensor. The EXR is "
1184-
"always written as scene-linear in the matching "
1185-
"gamut.\n"
1186-
" 'sRGB' — input is sRGB-encoded Rec.709; "
1187-
"the inverse sRGB EOTF is applied.\n"
1188-
" 'HDR' — input is HLG-encoded Rec.2020 "
1189-
"(BT.2100); the inverse HLG OETF is applied "
1190-
"to get scene-linear light.\n"
1191-
" 'linear' — input is already scene-linear "
1192-
"(Rec.709 primaries); written through unchanged. "
1193-
"Use this for renderer/compositor output."
1176+
"Colorspace of the input tensor. The EXR is always written as scene-linear in the matching gamut.\n"
1177+
"sRGB — input is sRGB-encoded Rec.709; the inverse sRGB EOTF is applied.\n"
1178+
"HDR — input is HLG-encoded Rec.2020 (BT.2100); the inverse HLG OETF is applied to get scene-linear light.\n"
1179+
"linear — input is already scene-linear (Rec.709 primaries); written through unchanged. Use this for renderer/compositor output."
11941180
),
11951181
),
11961182
]),
@@ -1200,6 +1186,7 @@ def define_schema(cls):
12001186
],
12011187
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
12021188
is_output_node=True,
1189+
outputs=[IO.Image.Output(display_name="images")]
12031190
)
12041191

12051192
@classmethod
@@ -1237,7 +1224,7 @@ def execute(cls, images, filename_prefix: str, format: dict) -> IO.NodeOutput:
12371224
results.append({"filename": file, "subfolder": subfolder, "type": "output"})
12381225
counter += 1
12391226

1240-
return IO.NodeOutput(ui={"images": results})
1227+
return IO.NodeOutput(images, ui={"images": results})
12411228

12421229

12431230
class ImagesExtension(ComfyExtension):

‎comfy_extras/nodes_video.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ def define_schema(cls):
2727
],
2828
hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
2929
is_output_node=True,
30+
outputs=[io.Image.Output(display_name="images")]
3031
)
3132

3233
@classmethod
@@ -69,7 +70,7 @@ def execute(cls, images, codec, fps, filename_prefix, crf) -> io.NodeOutput:
6970
container.mux(stream.encode())
7071
container.close()
7172

72-
return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
73+
return io.NodeOutput(images, ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
7374

7475
class SaveVideo(io.ComfyNode):
7576
@classmethod
@@ -89,6 +90,7 @@ def define_schema(cls):
8990
],
9091
hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
9192
is_output_node=True,
93+
outputs=[io.Video.Output("video")],
9294
)
9395

9496
@classmethod
@@ -117,7 +119,7 @@ def execute(cls, video: Input.Video, filename_prefix, format: str, codec) -> io.
117119
metadata=saved_metadata
118120
)
119121

120-
return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
122+
return io.NodeOutput(video, ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
121123

122124

123125
class CreateVideo(io.ComfyNode):

‎nodes.py‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -480,11 +480,13 @@ def __init__(self):
480480

481481
@classmethod
482482
def INPUT_TYPES(s):
483-
return {"required": { "samples": ("LATENT", ),
484-
"filename_prefix": ("STRING", {"default": "latents/ComfyUI"})},
485-
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
486-
}
487-
RETURN_TYPES = ()
483+
return { "required": {
484+
"samples": ("LATENT",),
485+
"filename_prefix": ("STRING", {"default": "latents/ComfyUI"})},
486+
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
487+
}
488+
RETURN_TYPES = ("LATENT",)
489+
RETURN_NAMES = ("samples",)
488490
FUNCTION = "save"
489491

490492
OUTPUT_NODE = True
@@ -522,7 +524,7 @@ def save(self, samples, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=No
522524
output["latent_format_version_0"] = torch.tensor([])
523525

524526
comfy.utils.save_torch_file(output, file, metadata=metadata)
525-
return { "ui": { "latents": results } }
527+
return { "ui": { "latents": results }, "result": (samples,) }
526528

527529

528530
class LoadLatent:
@@ -1627,14 +1629,18 @@ def INPUT_TYPES(s):
16271629
return {
16281630
"required": {
16291631
"images": ("IMAGE", {"tooltip": "The images to save."}),
1630-
"filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."})
1632+
"filename_prefix": ("STRING", {
1633+
"default": "ComfyUI",
1634+
"tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."
1635+
})
16311636
},
16321637
"hidden": {
16331638
"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
16341639
},
16351640
}
16361641

1637-
RETURN_TYPES = ()
1642+
RETURN_TYPES = ("IMAGE",)
1643+
RETURN_NAMES = ("images",)
16381644
FUNCTION = "save_images"
16391645

16401646
OUTPUT_NODE = True
@@ -1670,7 +1676,7 @@ def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pngi
16701676
})
16711677
counter += 1
16721678

1673-
return { "ui": { "images": results } }
1679+
return { "ui": { "images": results }, "result" : (images,) }
16741680

16751681
class PreviewImage(SaveImage):
16761682
def __init__(self):

0 commit comments

Comments
 (0)