From 333d942f303b170d9a778959b80fe0c3246a3cf6 Mon Sep 17 00:00:00 2001 From: bigcat88 Date: Wed, 23 Jul 2025 06:54:15 +0300 Subject: [PATCH] refactored Preview/Save of images --- comfy_api/v3/ui.py | 72 ++++++++++++++++++++++++++++++--- comfy_extras/v3/nodes_images.py | 45 +++++++++------------ 2 files changed, 86 insertions(+), 31 deletions(-) diff --git a/comfy_api/v3/ui.py b/comfy_api/v3/ui.py index 8a276cf0c..caaf894f6 100644 --- a/comfy_api/v3/ui.py +++ b/comfy_api/v3/ui.py @@ -37,6 +37,20 @@ class SavedResult(dict): return FolderType(self["type"]) +class SavedImages(_UIOutput): + """A UI output class to represent one or more saved images, potentially animated.""" + def __init__(self, results: list[SavedResult], is_animated: bool = False): + super().__init__() + self.results = results + self.is_animated = is_animated + + def as_dict(self) -> dict: + data = {"images": self.results} + if self.is_animated: + data["animated"] = (True,) + return data + + def _get_directory_by_folder_type(folder_type: FolderType) -> str: if folder_type == FolderType.input: return folder_paths.get_input_directory() @@ -125,14 +139,22 @@ class ImageSaveHelper: counter += 1 return results + @staticmethod + def get_save_images_ui(images, filename_prefix: str, cls: Type[ComfyNodeV3] | None, compress_level=4) -> SavedImages: + """Saves a batch of images and returns a UI object for the node output.""" + return SavedImages( + ImageSaveHelper.save_images( + images, + filename_prefix=filename_prefix, + folder_type=FolderType.output, + cls=cls, + compress_level=compress_level, + ) + ) + @staticmethod def save_animated_png( - images, - filename_prefix: str, - folder_type: FolderType, - cls: Type[ComfyNodeV3] | None, - fps: float, - compress_level: int + images, filename_prefix: str, folder_type: FolderType, cls: Type[ComfyNodeV3] | None, fps: float, compress_level: int ) -> SavedResult: """Saves a batch of images as a single animated PNG.""" full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path( @@ -152,6 +174,21 @@ class ImageSaveHelper: ) return SavedResult(file, subfolder, folder_type) + @staticmethod + def get_save_animated_png_ui( + images, filename_prefix: str, cls: Type[ComfyNodeV3] | None, fps: float, compress_level: int + ) -> SavedImages: + """Saves an animated PNG and returns a UI object for the node output.""" + result = ImageSaveHelper.save_animated_png( + images, + filename_prefix=filename_prefix, + folder_type=FolderType.output, + cls=cls, + fps=fps, + compress_level=compress_level, + ) + return SavedImages([result], is_animated=len(images) > 1) + @staticmethod def save_animated_webp( images, @@ -182,6 +219,29 @@ class ImageSaveHelper: ) return SavedResult(file, subfolder, folder_type) + @staticmethod + def get_save_animated_webp_ui( + images, + filename_prefix: str, + cls: Type[ComfyNodeV3] | None, + fps: float, + lossless: bool, + quality: int, + method: int, + ) -> SavedImages: + """Saves an animated WebP and returns a UI object for the node output.""" + result = ImageSaveHelper.save_animated_webp( + images, + filename_prefix=filename_prefix, + folder_type=FolderType.output, + cls=cls, + fps=fps, + lossless=lossless, + quality=quality, + method=method, + ) + return SavedImages([result], is_animated=len(images) > 1) + class PreviewImage(_UIOutput): def __init__(self, image: Image.Type, animated: bool=False, cls: ComfyNodeV3=None, **kwargs): diff --git a/comfy_extras/v3/nodes_images.py b/comfy_extras/v3/nodes_images.py index 16e77de4a..a1e26032b 100644 --- a/comfy_extras/v3/nodes_images.py +++ b/comfy_extras/v3/nodes_images.py @@ -630,15 +630,15 @@ class SaveAnimatedPNG(io.ComfyNodeV3): @classmethod def execute(cls, images, fps, compress_level, filename_prefix="ComfyUI") -> io.NodeOutput: - result = ui.ImageSaveHelper.save_animated_png( - images=images, - filename_prefix=filename_prefix, - folder_type=io.FolderType.output, - cls=cls, - fps=fps, - compress_level=compress_level, + return io.NodeOutput( + ui=ui.ImageSaveHelper.get_save_animated_png_ui( + images=images, + filename_prefix=filename_prefix, + cls=cls, + fps=fps, + compress_level=compress_level, + ) ) - return io.NodeOutput(ui={"images": [result], "animated": (len(images) != 1,)}) class SaveAnimatedWEBP(io.ComfyNodeV3): @@ -664,17 +664,17 @@ class SaveAnimatedWEBP(io.ComfyNodeV3): @classmethod def execute(cls, images, fps, filename_prefix, lossless, quality, method) -> io.NodeOutput: - result = ui.ImageSaveHelper.save_animated_webp( - images=images, - filename_prefix=filename_prefix, - folder_type=io.FolderType.output, - cls=cls, - fps=fps, - lossless=lossless, - quality=quality, - method=cls.COMPRESS_METHODS.get(method) + return io.NodeOutput( + ui=ui.ImageSaveHelper.get_save_animated_webp_ui( + images=images, + filename_prefix=filename_prefix, + cls=cls, + fps=fps, + lossless=lossless, + quality=quality, + method=cls.COMPRESS_METHODS.get(method) + ) ) - return io.NodeOutput(ui={"images": [result], "animated": (len(images) != 1,)}) class SaveImage(io.ComfyNodeV3): @@ -703,14 +703,9 @@ class SaveImage(io.ComfyNodeV3): @classmethod def execute(cls, images, filename_prefix="ComfyUI") -> io.NodeOutput: - results = ui.ImageSaveHelper.save_images( - images, - filename_prefix=filename_prefix, - folder_type=io.FolderType.output, - cls=cls, - compress_level=4, + return io.NodeOutput( + ui=ui.ImageSaveHelper.get_save_images_ui(images, filename_prefix=filename_prefix, cls=cls, compress_level=4) ) - return io.NodeOutput(ui={"images": results}) NODES_LIST: list[type[io.ComfyNodeV3]] = [