From 2999212480f375957e007452442c4247c515684f Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sat, 28 Jun 2025 13:53:25 -0700 Subject: [PATCH 1/3] Moved ui preview-related classes out of io.py and into ui.py, refactored UIImages and related into PreviewImage and related --- comfy_api/v3/io.py | 73 +---------------- comfy_api/v3/ui.py | 143 ++++++++++++++++++++++++++++++++++ comfy_extras/nodes_v3_test.py | 4 +- 3 files changed, 147 insertions(+), 73 deletions(-) create mode 100644 comfy_api/v3/ui.py diff --git a/comfy_api/v3/io.py b/comfy_api/v3/io.py index 2dd9da8ee..268f131b8 100644 --- a/comfy_api/v3/io.py +++ b/comfy_api/v3/io.py @@ -1207,7 +1207,7 @@ class NodeOutput: ''' Standardized output of a node; can pass in any number of args and/or a UIOutput into 'ui' kwarg. ''' - def __init__(self, *args: Any, ui: UIOutput | dict=None, expand: dict=None, block_execution: str=None, **kwargs): + def __init__(self, *args: Any, ui: _UIOutput | dict=None, expand: dict=None, block_execution: str=None, **kwargs): self.args = args self.ui = ui self.expand = expand @@ -1219,21 +1219,7 @@ class NodeOutput: # TODO: use kwargs to refer to outputs by id + organize in proper order return self.args if len(self.args) > 0 else None - -class SavedResult: - def __init__(self, filename: str, subfolder: str, type: FolderType): - self.filename = filename - self.subfolder = subfolder - self.type = type - - def as_dict(self): - return { - "filename": self.filename, - "subfolder": self.subfolder, - "type": self.type.value - } - -class UIOutput(ABC): +class _UIOutput(ABC): def __init__(self): pass @@ -1241,61 +1227,6 @@ class UIOutput(ABC): def as_dict(self) -> dict: ... # TODO: finish -class UIImages(UIOutput): - def __init__(self, values: list[SavedResult | dict], animated=False, **kwargs): - self.values = values - self.animated = animated - - def as_dict(self): - values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values] - return { - "images": values, - "animated": (self.animated,) - } - -class UILatents(UIOutput): - def __init__(self, values: list[SavedResult | dict], **kwargs): - self.values = values - - def as_dict(self): - values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values] - return { - "latents": values, - } - -class UIAudio(UIOutput): - def __init__(self, values: list[SavedResult | dict], **kwargs): - self.values = values - - def as_dict(self): - values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values] - return { - "audio": values, - } - -class UI3D(UIOutput): - def __init__(self, values: list[SavedResult | dict], **kwargs): - self.values = values - - def as_dict(self): - values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values] - return { - "3d": values, - } - -class UIText(UIOutput): - def __init__(self, value: str, **kwargs): - self.value = value - - def as_dict(self): - return {"text": (self.value,)} - - -def create_image_preview(image: Image.Type) -> UIImages: - # TODO: finish, right now is just Cursor's hallucination - return UIImages([SavedResult("preview.png", "comfy_org", FolderType.output)]) - - class TestNode(ComfyNodeV3): @classmethod def DEFINE_SCHEMA(cls): diff --git a/comfy_api/v3/ui.py b/comfy_api/v3/ui.py new file mode 100644 index 000000000..4f1951cb1 --- /dev/null +++ b/comfy_api/v3/ui.py @@ -0,0 +1,143 @@ +from __future__ import annotations +from abc import ABC, abstractmethod + +from comfy_api.v3.io import Image, Mask, FolderType, _UIOutput +# used for image preview +import folder_paths +import random +from PIL import Image as PILImage +import os +import numpy as np + + +class SavedResult: + def __init__(self, filename: str, subfolder: str, type: FolderType): + self.filename = filename + self.subfolder = subfolder + self.type = type + + def as_dict(self): + return { + "filename": self.filename, + "subfolder": self.subfolder, + "type": self.type + } + +class PreviewImage(_UIOutput): + def __init__(self, image: Image.Type, animated: bool=False, **kwargs): + output_dir = folder_paths.get_temp_directory() + type = "temp" + prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) + compress_level = 1 + filename_prefix = "ComfyUI" + + filename_prefix += prefix_append + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, output_dir, image[0].shape[1], image[0].shape[0]) + results = list() + for (batch_number, image) in enumerate(image): + i = 255. * image.cpu().numpy() + img = PILImage.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + metadata = None + # if not args.disable_metadata: + # metadata = PngInfo() + # if prompt is not None: + # metadata.add_text("prompt", json.dumps(prompt)) + # if extra_pnginfo is not None: + # for x in extra_pnginfo: + # metadata.add_text(x, json.dumps(extra_pnginfo[x])) + + filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) + file = f"{filename_with_batch_num}_{counter:05}_.png" + img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=compress_level) + results.append(SavedResult(file, subfolder, type)) + counter += 1 + + self.values = results + self.animated = animated + + def as_dict(self): + values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values] + return { + "images": values, + "animated": (self.animated,) + } + +class PreviewMask(PreviewImage): + def __init__(self, mask: PreviewMask.Type, animated: bool=False, **kwargs): + preview = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) + super().__init__(preview, animated, **kwargs) + +# class UILatent(_UIOutput): +# def __init__(self, values: list[SavedResult | dict], **kwargs): +# output_dir = folder_paths.get_temp_directory() +# type = "temp" +# prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) +# compress_level = 1 +# filename_prefix = "ComfyUI" + + +# full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) + +# # support save metadata for latent sharing +# prompt_info = "" +# if prompt is not None: +# prompt_info = json.dumps(prompt) + +# metadata = None +# if not args.disable_metadata: +# metadata = {"prompt": prompt_info} +# if extra_pnginfo is not None: +# for x in extra_pnginfo: +# metadata[x] = json.dumps(extra_pnginfo[x]) + +# file = f"{filename}_{counter:05}_.latent" + +# results: list[FileLocator] = [] +# results.append({ +# "filename": file, +# "subfolder": subfolder, +# "type": "output" +# }) + +# file = os.path.join(full_output_folder, file) + +# output = {} +# output["latent_tensor"] = samples["samples"].contiguous() +# output["latent_format_version_0"] = torch.tensor([]) + +# comfy.utils.save_torch_file(output, file, metadata=metadata) + +# self.values = values + +# def as_dict(self): +# values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values] +# return { +# "latents": values, +# } + +class PreviewAudio(_UIOutput): + def __init__(self, values: list[SavedResult | dict], **kwargs): + self.values = values + + def as_dict(self): + values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values] + return { + "audio": values, + } + +class PreviewUI3D(_UIOutput): + def __init__(self, values: list[SavedResult | dict], **kwargs): + self.values = values + + def as_dict(self): + values = [x.as_dict() if isinstance(x, SavedResult) else x for x in self.values] + return { + "3d": values, + } + +class PreviewText(_UIOutput): + def __init__(self, value: str, **kwargs): + self.value = value + + def as_dict(self): + return {"text": (self.value,)} diff --git a/comfy_extras/nodes_v3_test.py b/comfy_extras/nodes_v3_test.py index 1bfc8dc37..2f8a062b5 100644 --- a/comfy_extras/nodes_v3_test.py +++ b/comfy_extras/nodes_v3_test.py @@ -1,5 +1,5 @@ import torch -from comfy_api.v3 import io +from comfy_api.v3 import io, ui import logging import folder_paths import comfy.utils @@ -96,7 +96,7 @@ class V3TestNode(io.ComfyNodeV3): if hasattr(cls, "doohickey"): raise Exception("The 'cls' variable leaked state on class properties between runs!") cls.doohickey = "LOLJK" - return io.NodeOutput(some_int, image) + return io.NodeOutput(some_int, image, ui=ui.PreviewImage(image)) class V3LoraLoader(io.ComfyNodeV3): From 0e7ff98e1dd6f2ec91b57780826eb48f1e0cb8a8 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sat, 28 Jun 2025 15:47:02 -0700 Subject: [PATCH 2/3] Introduced Resources to ComfyNodeV3 --- comfy_api/v3/io.py | 55 +++++++++++++++++++++++++++++++++++ comfy_extras/nodes_v3_test.py | 16 +--------- execution.py | 7 ++++- 3 files changed, 62 insertions(+), 16 deletions(-) diff --git a/comfy_api/v3/io.py b/comfy_api/v3/io.py index 268f131b8..3d36793a4 100644 --- a/comfy_api/v3/io.py +++ b/comfy_api/v3/io.py @@ -4,6 +4,9 @@ from enum import Enum from abc import ABC, abstractmethod from dataclasses import dataclass, asdict from collections import Counter +import comfy.utils +import folder_paths +import logging # used for type hinting import torch from spandrel import ImageModelDescriptor @@ -281,6 +284,56 @@ class NodeStateLocal(NodeState): def __delitem__(self, key: str): del self.local_state[key] + +class ResourceKey(ABC): + def __init__(self): + ... + +class ResourceKeyFolderFilename(ResourceKey): + def __init__(self, folder_name: str, file_name: str): + self.folder_name = folder_name + self.file_name = file_name + + def __hash__(self): + return hash((self.folder_name, self.file_name)) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, ResourceKeyFolderFilename): + return False + return self.folder_name == other.folder_name and self.file_name == other.file_name + + def __str__(self): + return f"{self.folder_name} -> {self.file_name}" + +class Resources(ABC): + def __init__(self): + ... + + @abstractmethod + def get_torch_dict(self, key: ResourceKey) -> dict[str, torch.Tensor]: + pass + +class ResourcesLocal(Resources): + def __init__(self): + super().__init__() + self.local_resources: dict[ResourceKey, dict[str, torch.Tensor]] = {} + + def get_torch_dict(self, key: ResourceKey) -> dict[str, torch.Tensor]: + cached = self.local_resources.get(key, None) + if cached is not None: + logging.info(f"Using cached resource '{key}'") + return cached + logging.info(f"Loading resource '{key}'") + to_return = None + if isinstance(key, ResourceKeyFolderFilename): + to_return = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise(key.folder_name, key.file_name), safe_load=True) + + if to_return is not None: + self.local_resources[key] = to_return + return to_return + raise Exception(f"Unsupported resource key type: {type(key)}") + + @comfytype(io_type="BOOLEAN") class Boolean: Type = bool @@ -966,6 +1019,7 @@ class ComfyNodeV3: # filled in during execution state: NodeState = None + resources: Resources = None hidden: HiddenHolder = None @classmethod @@ -995,6 +1049,7 @@ class ComfyNodeV3: def __init__(self): self.local_state: NodeStateLocal = None + self.local_resources: ResourcesLocal = None self.__class__.VALIDATE_CLASS() @classmethod diff --git a/comfy_extras/nodes_v3_test.py b/comfy_extras/nodes_v3_test.py index 2f8a062b5..e84fdaa87 100644 --- a/comfy_extras/nodes_v3_test.py +++ b/comfy_extras/nodes_v3_test.py @@ -100,10 +100,6 @@ class V3TestNode(io.ComfyNodeV3): class V3LoraLoader(io.ComfyNodeV3): - class State(io.NodeState): - loaded_lora: tuple[str, Any] | None = None - state: State - @classmethod def DEFINE_SCHEMA(cls): return io.SchemaV3( @@ -147,17 +143,7 @@ class V3LoraLoader(io.ComfyNodeV3): if strength_model == 0 and strength_clip == 0: return io.NodeOutput(model, clip) - lora_path = folder_paths.get_full_path_or_raise("loras", lora_name) - lora = None - if cls.state.loaded_lora is not None: - if cls.state.loaded_lora[0] == lora_path: - lora = cls.state.loaded_lora[1] - else: - cls.state.loaded_lora = None - - if lora is None: - lora = comfy.utils.load_torch_file(lora_path, safe_load=True) - cls.state.loaded_lora = (lora_path, lora) + lora = cls.resources.get_torch_dict(io.ResourceKeyFolderFilename("loras", lora_name)) model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip) return io.NodeOutput(model_lora, clip_lora) diff --git a/execution.py b/execution.py index 8d5f708e3..5ab7f1fe0 100644 --- a/execution.py +++ b/execution.py @@ -28,7 +28,7 @@ from comfy_execution.graph import ( ) from comfy_execution.graph_utils import GraphBuilder, is_link from comfy_execution.validation import validate_node_input -from comfy_api.v3.io import NodeOutput, ComfyNodeV3, Hidden, NodeStateLocal +from comfy_api.v3.io import NodeOutput, ComfyNodeV3, Hidden, NodeStateLocal, ResourcesLocal class ExecutionResult(Enum): @@ -224,6 +224,11 @@ def _map_node_over_list(obj, input_data_all, func, allow_interrupt=False, execut if obj.local_state is None: obj.local_state = NodeStateLocal(class_clone.hidden.unique_id) class_clone.state = obj.local_state + # NOTE: this is a mock of resource management; for local, just stores ResourcesLocal on node instance + if hasattr(obj, "local_resources"): + if obj.local_resources is None: + obj.local_resources = ResourcesLocal() + class_clone.resources = obj.local_resources results.append(getattr(type(obj), func).__func__(class_clone, **inputs)) # V1 else: From 3758c65107d83f94da93efe3bccb635b16235f11 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Sat, 28 Jun 2025 16:46:45 -0700 Subject: [PATCH 3/3] Extracted resources to separate file --- comfy_api/v3/io.py | 53 +--------------------------- comfy_api/v3/resources.py | 65 +++++++++++++++++++++++++++++++++++ comfy_extras/nodes_v3_test.py | 4 +-- 3 files changed, 68 insertions(+), 54 deletions(-) create mode 100644 comfy_api/v3/resources.py diff --git a/comfy_api/v3/io.py b/comfy_api/v3/io.py index 3d36793a4..4cc758562 100644 --- a/comfy_api/v3/io.py +++ b/comfy_api/v3/io.py @@ -4,9 +4,7 @@ from enum import Enum from abc import ABC, abstractmethod from dataclasses import dataclass, asdict from collections import Counter -import comfy.utils -import folder_paths -import logging +from comfy_api.v3.resources import Resources, ResourcesLocal # used for type hinting import torch from spandrel import ImageModelDescriptor @@ -285,55 +283,6 @@ class NodeStateLocal(NodeState): del self.local_state[key] -class ResourceKey(ABC): - def __init__(self): - ... - -class ResourceKeyFolderFilename(ResourceKey): - def __init__(self, folder_name: str, file_name: str): - self.folder_name = folder_name - self.file_name = file_name - - def __hash__(self): - return hash((self.folder_name, self.file_name)) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, ResourceKeyFolderFilename): - return False - return self.folder_name == other.folder_name and self.file_name == other.file_name - - def __str__(self): - return f"{self.folder_name} -> {self.file_name}" - -class Resources(ABC): - def __init__(self): - ... - - @abstractmethod - def get_torch_dict(self, key: ResourceKey) -> dict[str, torch.Tensor]: - pass - -class ResourcesLocal(Resources): - def __init__(self): - super().__init__() - self.local_resources: dict[ResourceKey, dict[str, torch.Tensor]] = {} - - def get_torch_dict(self, key: ResourceKey) -> dict[str, torch.Tensor]: - cached = self.local_resources.get(key, None) - if cached is not None: - logging.info(f"Using cached resource '{key}'") - return cached - logging.info(f"Loading resource '{key}'") - to_return = None - if isinstance(key, ResourceKeyFolderFilename): - to_return = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise(key.folder_name, key.file_name), safe_load=True) - - if to_return is not None: - self.local_resources[key] = to_return - return to_return - raise Exception(f"Unsupported resource key type: {type(key)}") - - @comfytype(io_type="BOOLEAN") class Boolean: Type = bool diff --git a/comfy_api/v3/resources.py b/comfy_api/v3/resources.py new file mode 100644 index 000000000..6ff59d6ae --- /dev/null +++ b/comfy_api/v3/resources.py @@ -0,0 +1,65 @@ +from __future__ import annotations +import comfy.utils +import folder_paths +import logging +from abc import ABC, abstractmethod +from typing import Any +import torch + +class ResourceKey(ABC): + Type = Any + def __init__(self): + ... + +class TorchDictFolderFilename(ResourceKey): + '''Key for requesting a torch file via file_name from a folder category.''' + Type = dict[str, torch.Tensor] + def __init__(self, folder_name: str, file_name: str): + self.folder_name = folder_name + self.file_name = file_name + + def __hash__(self): + return hash((self.folder_name, self.file_name)) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, TorchDictFolderFilename): + return False + return self.folder_name == other.folder_name and self.file_name == other.file_name + + def __str__(self): + return f"{self.folder_name} -> {self.file_name}" + +class Resources(ABC): + def __init__(self): + ... + + @abstractmethod + def get(self, key: ResourceKey, default: Any=...) -> Any: + pass + +class ResourcesLocal(Resources): + def __init__(self): + super().__init__() + self.local_resources: dict[ResourceKey, Any] = {} + + def get(self, key: ResourceKey, default: Any=...) -> Any: + cached = self.local_resources.get(key, None) + if cached is not None: + logging.info(f"Using cached resource '{key}'") + return cached + logging.info(f"Loading resource '{key}'") + to_return = None + if isinstance(key, TorchDictFolderFilename): + if default is ...: + to_return = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise(key.folder_name, key.file_name), safe_load=True) + else: + full_path = folder_paths.get_full_path(key.folder_name, key.file_name) + if full_path is not None: + to_return = comfy.utils.load_torch_file(full_path, safe_load=True) + + if to_return is not None: + self.local_resources[key] = to_return + return to_return + if default is not ...: + return default + raise Exception(f"Unsupported resource key type: {type(key)}") diff --git a/comfy_extras/nodes_v3_test.py b/comfy_extras/nodes_v3_test.py index e84fdaa87..9120d8b8c 100644 --- a/comfy_extras/nodes_v3_test.py +++ b/comfy_extras/nodes_v3_test.py @@ -1,5 +1,5 @@ import torch -from comfy_api.v3 import io, ui +from comfy_api.v3 import io, ui, resources import logging import folder_paths import comfy.utils @@ -143,7 +143,7 @@ class V3LoraLoader(io.ComfyNodeV3): if strength_model == 0 and strength_clip == 0: return io.NodeOutput(model, clip) - lora = cls.resources.get_torch_dict(io.ResourceKeyFolderFilename("loras", lora_name)) + lora = cls.resources.get(resources.TorchDictFolderFilename("loras", lora_name)) model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip) return io.NodeOutput(model_lora, clip_lora)