From 1f852778ec98140d51024cd6bfbb2f3ca7a78153 Mon Sep 17 00:00:00 2001
From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com>
Date: Tue, 11 Nov 2025 16:23:24 +0800
Subject: [PATCH] Add outputs' id/name

---
 comfy_extras/nodes_dataset.py | 48 +++++++++++++++++++++++++++--------
 comfy_extras/nodes_train.py   | 10 ++++----
 2 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py
index 820d8a4e9..282010b23 100644
--- a/comfy_extras/nodes_dataset.py
+++ b/comfy_extras/nodes_dataset.py
@@ -59,7 +59,9 @@ class LoadImageDataSetFromFolderNode(io.ComfyNode):
                 )
             ],
             outputs=[
-                io.Image.Output(is_output_list=True, tooltip="List of loaded images")
+                io.Image.Output(
+                    "images", is_output_list=True, tooltip="List of loaded images"
+                )
             ],
         )
 
@@ -92,8 +94,12 @@ class LoadImageTextDataSetFromFolderNode(io.ComfyNode):
                 )
             ],
             outputs=[
-                io.Image.Output(is_output_list=True, tooltip="List of loaded images"),
-                io.String.Output(is_output_list=True, tooltip="List of text captions"),
+                io.Image.Output(
+                    "images", is_output_list=True, tooltip="List of loaded images"
+                ),
+                io.String.Output(
+                    "texts", is_output_list=True, tooltip="List of text captions"
+                ),
             ],
         )
 
@@ -329,7 +335,11 @@ class ImageProcessingNode(io.ComfyNode):
             is_experimental=True,
             is_input_list=True,
             inputs=inputs,
-            outputs=[io.Image.Output(is_output_list=True, tooltip="Processed images")],
+            outputs=[
+                io.Image.Output(
+                    "images", is_output_list=True, tooltip="Processed images"
+                )
+            ],
         )
 
     @classmethod
@@ -393,7 +403,11 @@ class TextProcessingNode(io.ComfyNode):
             is_experimental=True,
             is_input_list=True,
             inputs=inputs,
-            outputs=[io.String.Output(is_output_list=True, tooltip="Processed texts")],
+            outputs=[
+                io.String.Output(
+                    "texts", is_output_list=True, tooltip="Processed texts"
+                )
+            ],
         )
 
     @classmethod
@@ -752,8 +766,12 @@ class ShuffleImageTextDatasetNode(io.ComfyNode):
                 ),
             ],
             outputs=[
-                io.Image.Output(is_output_list=True, tooltip="Shuffled images"),
-                io.String.Output(is_output_list=True, tooltip="Shuffled texts"),
+                io.Image.Output(
+                    "images", is_output_list=True, tooltip="Shuffled images"
+                ),
+                io.String.Output(
+                    "texts", is_output_list=True, tooltip="Shuffled texts"
+                ),
             ],
         )
 
@@ -884,9 +902,13 @@ class MakeTrainingDataset(io.ComfyNode):
                 ),
             ],
             outputs=[
-                io.Latent.Output(is_output_list=True, tooltip="List of latent dicts"),
+                io.Latent.Output(
+                    "latents", is_output_list=True, tooltip="List of latent dicts"
+                ),
                 io.Conditioning.Output(
-                    is_output_list=True, tooltip="List of conditioning lists"
+                    "conditioning",
+                    is_output_list=True,
+                    tooltip="List of conditioning lists",
                 ),
             ],
         )
@@ -1059,9 +1081,13 @@ class LoadTrainingDataset(io.ComfyNode):
                 ),
             ],
             outputs=[
-                io.Latent.Output(is_output_list=True, tooltip="List of latent dicts"),
+                io.Latent.Output(
+                    "latents", is_output_list=True, tooltip="List of latent dicts"
+                ),
                 io.Conditioning.Output(
-                    is_output_list=True, tooltip="List of conditioning lists"
+                    "conditioning",
+                    is_output_list=True,
+                    tooltip="List of conditioning lists",
                 ),
             ],
         )
diff --git a/comfy_extras/nodes_train.py b/comfy_extras/nodes_train.py
index 636b74b7f..cdb26332a 100644
--- a/comfy_extras/nodes_train.py
+++ b/comfy_extras/nodes_train.py
@@ -387,10 +387,10 @@ class TrainLoraNode(io.ComfyNode):
                 ),
             ],
             outputs=[
-                io.Model.Output(tooltip="Model with LoRA applied"),
-                io.Custom("LORA_MODEL").Output(tooltip="LoRA weights"),
-                io.Custom("LOSS_MAP").Output(tooltip="Loss history"),
-                io.Int.Output(tooltip="Total training steps"),
+                io.Model.Output("model", tooltip="Model with LoRA applied"),
+                io.Custom("LORA_MODEL").Output("lora", tooltip="LoRA weights"),
+                io.Custom("LOSS_MAP").Output("loss_map", tooltip="Loss history"),
+                io.Int.Output("steps", tooltip="Total training steps"),
             ],
         )
 
@@ -663,7 +663,7 @@ class LoraModelLoader(io.ComfyNode):
                 ),
             ],
             outputs=[
-                io.Model.Output(tooltip="The modified diffusion model."),
+                io.Model.Output("model", tooltip="The modified diffusion model."),
             ],
         )