diff --git a/comfy_api/latest/_io.py b/comfy_api/latest/_io.py index 764fa8b2b..deb0ddcab 100644 --- a/comfy_api/latest/_io.py +++ b/comfy_api/latest/_io.py @@ -24,6 +24,7 @@ if TYPE_CHECKING: from comfy.sd import CLIP, VAE from comfy.sd import StyleModel as StyleModel_ from comfy_api.input import VideoInput + from ._ui_types import UIOutputDict from comfy_api.internal import (_ComfyNodeInternal, _NodeOutputInternal, classproperty, copy_class, first_real_override, is_class, prune_dict, shallow_clone_class) from comfy_execution.graph_utils import ExecutionBlocker @@ -1873,7 +1874,7 @@ class _UIOutput(ABC): pass @abstractmethod - def as_dict(self) -> dict: + def as_dict(self) -> "UIOutputDict": ... diff --git a/comfy_api/latest/_ui.py b/comfy_api/latest/_ui.py index e238cdf3c..7a46a57de 100644 --- a/comfy_api/latest/_ui.py +++ b/comfy_api/latest/_ui.py @@ -5,6 +5,7 @@ import os import random import uuid from io import BytesIO +from typing import TYPE_CHECKING, Any import av import numpy as np @@ -23,6 +24,15 @@ import folder_paths from comfy.cli_args import args from ._io import ComfyNode, FolderType, Image, _UIOutput +if TYPE_CHECKING: + from ._ui_types import ( + ImagesUIOutput, + AudioUIOutput, + VideoUIOutput, + TextUIOutput, + UI3DUIOutput, + ) + class SavedResult(dict): def __init__(self, filename: str, subfolder: str, type: FolderType): @@ -48,11 +58,11 @@ class SavedImages(_UIOutput): self.results = results self.is_animated = is_animated - def as_dict(self) -> dict: - data = {"images": self.results} + def as_dict(self) -> "ImagesUIOutput": + data: dict[str, Any] = {"images": self.results} if self.is_animated: data["animated"] = (True,) - return data + return data # type: ignore[return-value] class SavedAudios(_UIOutput): @@ -61,8 +71,8 @@ class SavedAudios(_UIOutput): super().__init__() self.results = results - def as_dict(self) -> dict: - return {"audio": self.results} + def as_dict(self) -> "AudioUIOutput": + return {"audio": self.results} # type: ignore[return-value] def _get_directory_by_folder_type(folder_type: FolderType) -> str: @@ -397,9 +407,9 @@ class PreviewImage(_UIOutput): ) self.animated = animated - def as_dict(self): + def as_dict(self) -> "ImagesUIOutput": return { - "images": self.values, + "images": self.values, # type: ignore[typeddict-item] "animated": (self.animated,) } @@ -421,16 +431,16 @@ class PreviewAudio(_UIOutput): quality="128k", ) - def as_dict(self) -> dict: - return {"audio": self.values} + def as_dict(self) -> "AudioUIOutput": + return {"audio": self.values} # type: ignore[return-value] class PreviewVideo(_UIOutput): def __init__(self, values: list[SavedResult | dict], **kwargs): self.values = values - def as_dict(self): - return {"images": self.values, "animated": (True,)} + def as_dict(self) -> "VideoUIOutput": + return {"images": self.values, "animated": (True,)} # type: ignore[return-value] class PreviewUI3D(_UIOutput): @@ -448,7 +458,7 @@ class PreviewUI3D(_UIOutput): img.save(bg_image_path, compress_level=1) self.bg_image_path = f"temp/{filename}" - def as_dict(self): + def as_dict(self) -> "UI3DUIOutput": return {"result": [self.model_file, self.camera_info, self.bg_image_path]} @@ -456,7 +466,7 @@ class PreviewText(_UIOutput): def __init__(self, value: str, **kwargs): self.value = value - def as_dict(self): + def as_dict(self) -> "TextUIOutput": return {"text": (self.value,)} diff --git a/comfy_api/latest/_ui_public.py b/comfy_api/latest/_ui_public.py index 85b11d78b..916e95b65 100644 --- a/comfy_api/latest/_ui_public.py +++ b/comfy_api/latest/_ui_public.py @@ -1 +1,2 @@ from ._ui import * # noqa: F403 +from ._ui_types import * # noqa: F403 diff --git a/comfy_api/latest/_ui_types.py b/comfy_api/latest/_ui_types.py new file mode 100644 index 000000000..363e50594 --- /dev/null +++ b/comfy_api/latest/_ui_types.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +from typing import Any, Literal, Union +from typing_extensions import NotRequired, TypedDict + + +FolderTypeLiteral = Literal["input", "output", "temp"] + + +class FileResultDict(TypedDict): + filename: str + subfolder: str + type: FolderTypeLiteral + + +class ImagesUIOutput(TypedDict): + images: list[FileResultDict] + animated: NotRequired[tuple[bool]] + + +class AudioUIOutput(TypedDict): + audio: list[FileResultDict] + + +class VideoUIOutput(TypedDict): + images: list[FileResultDict | dict[str, Any]] + animated: tuple[Literal[True]] + + +class TextUIOutput(TypedDict): + text: tuple[str, ...] + + +class CameraInfoDict(TypedDict, total=False): + position: dict[str, float | int] + target: dict[str, float | int] + zoom: int + cameraType: str + + +class UI3DUIOutput(TypedDict): + result: list[str | CameraInfoDict | None] + + +class LatentsUIOutput(TypedDict): + latents: list[FileResultDict] + + +UIOutputDict = Union[ + ImagesUIOutput, + AudioUIOutput, + VideoUIOutput, + TextUIOutput, + UI3DUIOutput, + LatentsUIOutput, +]