diff --git a/__init__.py b/__init__.py index 20becb1..bd50b46 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,11 @@ -from .src.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .src.nodes import ( + NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS, + ComfyUIPipelineDTO, +) -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "ComfyUIPipelineDTO", +] diff --git a/src/nodes/__init__.py b/src/nodes/__init__.py index 34eb959..b954c04 100644 --- a/src/nodes/__init__.py +++ b/src/nodes/__init__.py @@ -1,5 +1,6 @@ from ..di import Container from .diffusers_pipeline_loader import DiffusersPipelineLoader +from .dto import ComfyUIPipelineDTO container = Container() container.wire(modules=[__name__]) @@ -7,4 +8,8 @@ container.wire(modules=[__name__]) NODE_CLASS_MAPPINGS = {DiffusersPipelineLoader.__name__: DiffusersPipelineLoader} NODE_DISPLAY_NAME_MAPPINGS = {DiffusersPipelineLoader.__name__: "Diffusers Pipeline Loader"} -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "ComfyUIPipelineDTO", +] diff --git a/src/nodes/diffusers_pipeline_loader.py b/src/nodes/diffusers_pipeline_loader.py index df2d114..b89cad8 100644 --- a/src/nodes/diffusers_pipeline_loader.py +++ b/src/nodes/diffusers_pipeline_loader.py @@ -2,8 +2,8 @@ import folder_paths # pyright: ignore[reportMissingImports] from dependency_injector.wiring import Provide, inject from ..di import Container -from ..domain.model import PipelineModel -from ..ui.pipeline_handler import PipelineHandler +from ..ui import PipelineHandler +from .dto import ComfyUIPipelineDTO class DiffusersPipelineLoader: @@ -18,7 +18,7 @@ class DiffusersPipelineLoader: } } - RETURN_TYPES = ("PIPELINE",) + RETURN_TYPES = (ComfyUIPipelineDTO.COMFY_TYPE,) FUNCTION = "execute" CATEGORY = "Diffusers" @@ -27,6 +27,7 @@ class DiffusersPipelineLoader: self, checkpoint_name: str, handler: PipelineHandler = Provide[Container.pipeline_handler_provider], - ) -> tuple[PipelineModel]: + ) -> tuple[ComfyUIPipelineDTO]: pipeline_model = handler.create(checkpoint_name) - return (pipeline_model,) + pipeline_dto = ComfyUIPipelineDTO.from_domain(pipeline_model) + return (pipeline_dto,) diff --git a/src/nodes/dto/__init__.py b/src/nodes/dto/__init__.py new file mode 100644 index 0000000..7325eea --- /dev/null +++ b/src/nodes/dto/__init__.py @@ -0,0 +1,3 @@ +from .pipeline_dto import ComfyUIPipelineDTO + +__all__ = ["ComfyUIPipelineDTO"] diff --git a/src/nodes/dto/pipeline_dto.py b/src/nodes/dto/pipeline_dto.py new file mode 100644 index 0000000..2a71231 --- /dev/null +++ b/src/nodes/dto/pipeline_dto.py @@ -0,0 +1,58 @@ +from dataclasses import dataclass +from typing import ClassVar + +from diffusers import StableDiffusionPipeline + +from ...domain.model import PipelineModel + + +@dataclass +class ComfyUIPipelineDTO: + """ + Data Transfer Object for Diffusers Pipeline in ComfyUI. + + This DTO is designed to be used across different custom nodes. + Other node developers can import and use this type for pipeline operations. + + Attributes: + pipeline: StableDiffusionPipeline instance from diffusers library + path: Path to the loaded model directory + + Class Attributes: + COMFY_TYPE: Type name used in ComfyUI's RETURN_TYPES and INPUT_TYPES + + Example: + ```python + from your_node_package.nodes.dto import ComfyUIPipelineDTO + + class YourCustomNode: + @classmethod + def INPUT_TYPES(cls): + return {"required": {"pipeline": (ComfyUIPipelineDTO.COMFY_TYPE,)}} + + RETURN_TYPES = (ComfyUIPipelineDTO.COMFY_TYPE,) + + def execute(self, pipeline: ComfyUIPipelineDTO): + result = pipeline.pipeline(prompt="...") + return (result,) + ``` + """ + + # ClassVar to avoid being treated as a dataclass field + COMFY_TYPE: ClassVar[str] = "PIPELINE" + + pipeline: StableDiffusionPipeline + path: str + + @classmethod + def from_domain(cls, model: PipelineModel) -> "ComfyUIPipelineDTO": + """ + Create DTO from domain model. + + Args: + model: PipelineModel from domain layer + + Returns: + ComfyUIPipelineDTO instance + """ + return cls(pipeline=model.pipeline, path=model.path)