feat: comfy dto

__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",
+]

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",
+]

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,)

src/nodes/dto/__init__.py

@@ -0,0 +1,3 @@
+from .pipeline_dto import ComfyUIPipelineDTO
+
+__all__ = ["ComfyUIPipelineDTO"]

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)