diff --git a/comfy/comfy_types/__init__.py b/comfy/comfy_types/__init__.py index 19ec33f98..7640fbe3f 100644 --- a/comfy/comfy_types/__init__.py +++ b/comfy/comfy_types/__init__.py @@ -1,6 +1,6 @@ import torch from typing import Callable, Protocol, TypedDict, Optional, List -from .node_typing import IO, InputTypeDict, ComfyNodeABC, CheckLazyMixin +from .node_typing import IO, InputTypeDict, ComfyNodeABC, CheckLazyMixin, FileLocator class UnetApplyFunction(Protocol): @@ -42,4 +42,5 @@ __all__ = [ InputTypeDict.__name__, ComfyNodeABC.__name__, CheckLazyMixin.__name__, + FileLocator.__name__, ] diff --git a/comfy/comfy_types/node_typing.py b/comfy/comfy_types/node_typing.py index 0696dbe5e..acf912819 100644 --- a/comfy/comfy_types/node_typing.py +++ b/comfy/comfy_types/node_typing.py @@ -293,3 +293,14 @@ class CheckLazyMixin: need = [name for name in kwargs if kwargs[name] is None] return need + + +class FileLocator(TypedDict): + """Provides type hinting for the file location""" + + filename: str + """The filename of the file.""" + subfolder: str + """The subfolder of the file.""" + type: Literal["input", "output", "temp"] + """The root folder of the file.""" diff --git a/comfy_extras/nodes_audio.py b/comfy_extras/nodes_audio.py index 3cb918e09..7609b35cd 100644 --- a/comfy_extras/nodes_audio.py +++ b/comfy_extras/nodes_audio.py @@ -10,6 +10,7 @@ import random import hashlib import node_helpers from comfy.cli_args import args +from comfy.comfy_types import FileLocator class EmptyLatentAudio: def __init__(self): @@ -164,7 +165,7 @@ class SaveAudio: def save_audio(self, audio, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): filename_prefix += self.prefix_append full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) - results = list() + results: list[FileLocator] = [] metadata = {} if not args.disable_metadata: diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index af37666b2..3503fb996 100644 --- a/comfy_extras/nodes_images.py +++ b/comfy_extras/nodes_images.py @@ -9,6 +9,8 @@ import numpy as np import json import os +from comfy.comfy_types import FileLocator + MAX_RESOLUTION = nodes.MAX_RESOLUTION class ImageCrop: @@ -99,7 +101,7 @@ class SaveAnimatedWEBP: method = self.methods.get(method) filename_prefix += self.prefix_append full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) - results = list() + results: list[FileLocator] = [] pil_images = [] for image in images: i = 255. * image.cpu().numpy() diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py index f3922e03d..8e7f6b97a 100644 --- a/comfy_extras/nodes_video.py +++ b/comfy_extras/nodes_video.py @@ -4,6 +4,7 @@ import torch import folder_paths import json from fractions import Fraction +from comfy.comfy_types import FileLocator class SaveWEBM: @@ -61,7 +62,7 @@ class SaveWEBM: container.mux(packet) container.close() - results = [{ + results: list[FileLocator] = [{ "filename": file, "subfolder": subfolder, "type": self.type diff --git a/nodes.py b/nodes.py index 4e68af79d..030360694 100644 --- a/nodes.py +++ b/nodes.py @@ -25,7 +25,7 @@ import comfy.sample import comfy.sd import comfy.utils import comfy.controlnet -from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict +from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict, FileLocator import comfy.clip_vision @@ -479,7 +479,7 @@ class SaveLatent: file = f"{filename}_{counter:05}_.latent" - results = list() + results: list[FileLocator] = [] results.append({ "filename": file, "subfolder": subfolder,