feat: scheduler node

src/di/di.py

@@ -6,10 +6,23 @@ from ..infra.diffusers import (
     DiffusersAutoencoderRepository,
     DiffusersPipelineRepository,
     DiffusersSamplingRepository,
+    DiffusersSchedulerRepository,
     DiffusersTextEncoderRepository,
 )
-from ..service import AutoencoderService, ClipTextEncodeService, PipelineService, SamplingService
-from ..ui import AutoencoderHandler, ClipTextEncodeHandler, PipelineHandler, SamplerHandler
+from ..service import (
+    AutoencoderService,
+    ClipTextEncodeService,
+    PipelineService,
+    SamplingService,
+    SchedulerService,
+)
+from ..ui import (
+    AutoencoderHandler,
+    ClipTextEncodeHandler,
+    PipelineHandler,
+    SamplerHandler,
+    SchedulerHandler,
+)
 
 
 class Container(containers.DeclarativeContainer):
@@ -20,6 +33,7 @@ class Container(containers.DeclarativeContainer):
     autoencoder_repository = providers.Factory(DiffusersAutoencoderRepository)
     text_encoder_repository = providers.Factory(DiffusersTextEncoderRepository)
     sampling_repository = providers.Factory(DiffusersSamplingRepository)
+    scheduler_repository = providers.Factory(DiffusersSchedulerRepository)
 
     # 2. Services
     pipeline_service = providers.Factory(
@@ -38,6 +52,10 @@ class Container(containers.DeclarativeContainer):
         SamplingService,
         sampling_repo=sampling_repository,
     )
+    scheduler_service = providers.Factory(
+        SchedulerService,
+        scheduler_repo=scheduler_repository,
+    )
 
     # 3. Handler
     pipeline_handler = providers.Factory(
@@ -56,3 +74,7 @@ class Container(containers.DeclarativeContainer):
         SamplerHandler,
         sampling_service=sampling_service,
     )
+    scheduler_handler = providers.Factory(
+        SchedulerHandler,
+        scheduler_service=scheduler_service,
+    )

src/domain/model/__init__.py

@@ -3,6 +3,7 @@ from ._clip import Clip
 from ._conditioning import Conditioning
 from ._image import Image
 from ._pipeline import Pipeline
+from ._scheduler import Scheduler
 
 __all__ = [
     "Pipeline",
@@ -10,4 +11,5 @@ __all__ = [
     "Clip",
     "Conditioning",
     "Image",
+    "Scheduler",
 ]

src/domain/model/_scheduler.py

@@ -0,0 +1,36 @@
+from dataclasses import dataclass
+from enum import Enum
+
+from diffusers import (
+    DDIMScheduler,
+    DDPMScheduler,
+    DEISMultistepScheduler,
+    DPMSolverMultistepScheduler,
+    DPMSolverSinglestepScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    HeunDiscreteScheduler,
+    KDPM2AncestralDiscreteScheduler,
+    KDPM2DiscreteScheduler,
+    UniPCMultistepScheduler,
+)
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+
+@dataclass
+class Scheduler:
+    class Type(Enum):
+        DDIM = DDIMScheduler
+        DDPM = DDPMScheduler
+        DEISMultistep = DEISMultistepScheduler
+        DPMSolverMultistep = DPMSolverMultistepScheduler
+        DPMSolverSinglestep = DPMSolverSinglestepScheduler
+        EulerAncestralDiscrete = EulerAncestralDiscreteScheduler
+        EulerDiscrete = EulerDiscreteScheduler
+        HeunDiscrete = HeunDiscreteScheduler
+        KDPM2AncestralDiscrete = KDPM2AncestralDiscreteScheduler
+        KDPM2Discrete = KDPM2DiscreteScheduler
+        UniPCMultistep = UniPCMultistepScheduler
+
+    scheduler: SchedulerMixin
+    path: str

src/domain/repositories/__init__.py

@@ -1,6 +1,7 @@
 from ._autoencoder_repository import AutoencoderRepository
 from ._pipeline_repository import PipelineRepository
 from ._sampling_repository import SamplingRepository
+from ._scheduler_repository import SchedulerRepository
 from ._text_encoder_repository import TextEncoderRepository
 
 __all__ = [
@@ -8,4 +9,5 @@ __all__ = [
     "AutoencoderRepository",
     "TextEncoderRepository",
     "SamplingRepository",
+    "SchedulerRepository",
 ]

src/domain/repositories/_scheduler_repository.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+
+import torch
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+from ...domain.model import Scheduler
+
+
+class SchedulerRepository(ABC):
+    @abstractmethod
+    def create_scheduler(
+        self, model_path: str, dtype: torch.dtype, scheduler_type: Scheduler.Type
+    ) -> SchedulerMixin:
+        pass

src/infra/diffusers/__init__.py

@@ -1,6 +1,7 @@
 from ._autoencoder_repository import DiffusersAutoencoderRepository
 from ._pipeline_repository import DiffusersPipelineRepository
 from ._sampling_repository import DiffusersSamplingRepository
+from ._scheduler_repository import DiffusersSchedulerRepository
 from ._text_encoder_repository import DiffusersTextEncoderRepository
 
 __all__ = [
@@ -8,4 +9,5 @@ __all__ = [
     "DiffusersAutoencoderRepository",
     "DiffusersTextEncoderRepository",
     "DiffusersSamplingRepository",
+    "DiffusersSchedulerRepository",
 ]

src/infra/diffusers/_scheduler_repository.py

@@ -0,0 +1,22 @@
+import folder_paths  # pyright: ignore[reportMissingImports]
+import torch
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+from ...domain.model import Scheduler
+from ...domain.repositories import SchedulerRepository
+
+
+class DiffusersSchedulerRepository(SchedulerRepository):
+    def __init__(self) -> None:
+        self.cache_dir = folder_paths.get_temp_directory()
+
+    def create_scheduler(
+        self, model_path: str, dtype: torch.dtype, scheduler_type: Scheduler.Type
+    ) -> SchedulerMixin:
+        scheduler: SchedulerMixin = scheduler_type.value.from_pretrained(
+            pretrained_model_name_or_path=model_path,
+            subfolder="scheduler",
+            torch_dtype=dtype,
+            cache_dir=self.cache_dir,
+        )
+        return scheduler

src/nodes/__init__.py

@@ -2,8 +2,15 @@ from ..di import Container
 from ._diffusers_clip_text_encode import DiffusersClipTextEncode
 from ._diffusers_pipeline_loader import DiffusersPipelineLoader
 from ._diffusers_sampler import DiffusersSampler
+from ._diffusers_scheduler_loader import DiffusersSchedulerLoader
 from ._diffusers_vae_loader import DiffusersVaeLoader
-from .dto import ComfyUIAutoencoderDTO, ComfyUIClipDTO, ComfyUIConditioningDTO, ComfyUIPipelineDTO
+from .dto import (
+    ComfyUIAutoencoderDTO,
+    ComfyUIClipDTO,
+    ComfyUIConditioningDTO,
+    ComfyUIPipelineDTO,
+    ComfyUISchedulerDTO,
+)
 
 container = Container()
 container.wire(modules=[__name__])
@@ -13,12 +20,14 @@ NODE_CLASS_MAPPINGS = {
     DiffusersVaeLoader.__name__: DiffusersVaeLoader,
     DiffusersClipTextEncode.__name__: DiffusersClipTextEncode,
     DiffusersSampler.__name__: DiffusersSampler,
+    DiffusersSchedulerLoader.__name__: DiffusersSchedulerLoader,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     DiffusersPipelineLoader.__name__: "Diffusers Pipeline Loader",
     DiffusersVaeLoader.__name__: "Diffusers VAE Loader",
     DiffusersClipTextEncode.__name__: "Diffusers CLIP Text Encode",
     DiffusersSampler.__name__: "Diffusers Sampler",
+    DiffusersSchedulerLoader.__name__: "Diffusers Scheduler Loader",
 }
 
 __all__ = [
@@ -28,4 +37,5 @@ __all__ = [
     "ComfyUIAutoencoderDTO",
     "ComfyUIClipDTO",
     "ComfyUIConditioningDTO",
+    "ComfyUISchedulerDTO",
 ]

src/nodes/_diffusers_scheduler_loader.py

@@ -0,0 +1,37 @@
+from dependency_injector.wiring import Provide, inject
+
+from ..di import Container
+from ..ui import SchedulerHandler
+from .dto import ComfyUIPipelineDTO, ComfyUISchedulerDTO
+from .dto.type import ComfyUISchedulerType
+
+
+class DiffusersSchedulerLoader:
+    def __init__(self) -> None:
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str] | tuple[list[str]]]]:
+        return {
+            "required": {
+                "pipeline": (ComfyUIPipelineDTO.COMFY_TYPE,),
+                "scheduler_name": (list(ComfyUISchedulerType.SCHEDULERS),),
+            }
+        }
+
+    RETURN_TYPES = (ComfyUISchedulerDTO.COMFY_TYPE,)
+    FUNCTION = "execute"
+    CATEGORY = "Diffusers"
+
+    @inject
+    def execute(
+        self,
+        pipeline: ComfyUIPipelineDTO,
+        scheduler_name: str,
+        handler: SchedulerHandler = Provide[Container.scheduler_handler],
+    ) -> tuple[ComfyUISchedulerDTO]:
+        domain = ComfyUIPipelineDTO.to_domain(pipeline)
+        scheduler_type = ComfyUISchedulerType.to_domain(scheduler_name)
+        scheduler_domain = handler.create(domain, scheduler_type)
+        scheduler_dto = ComfyUISchedulerDTO.from_domain(scheduler_domain)
+        return (scheduler_dto,)

src/nodes/dto/__init__.py

@@ -3,6 +3,7 @@ from ._clip import ComfyUIClipDTO
 from ._comfyui import ComfyUIImage, ComfyUIImageDTO
 from ._conditioning_dto import ComfyUIConditioningDTO
 from ._pipeline_dto import ComfyUIPipelineDTO
+from ._scheduler_dto import ComfyUISchedulerDTO
 
 __all__ = [
     "ComfyUIPipelineDTO",
@@ -11,4 +12,5 @@ __all__ = [
     "ComfyUIConditioningDTO",
     "ComfyUIImage",
     "ComfyUIImageDTO",
+    "ComfyUISchedulerDTO",
 ]

src/nodes/dto/_pipeline_dto.py

@@ -54,3 +54,15 @@ class ComfyUIPipelineDTO:
             ComfyUIPipelineDTO instance
         """
         return cls(pipeline=pipeline.pipeline, path=pipeline.path)
+
+    @classmethod
+    def to_domain(cls, dto: "ComfyUIPipelineDTO") -> Pipeline:
+        """Convert DTO back to domain model.
+
+        Args:
+            dto: ComfyUIPipelineDTO instance
+
+        Returns:
+            Pipeline domain model instance
+        """
+        return Pipeline(pipeline=dto.pipeline, path=dto.path)

src/nodes/dto/_scheduler_dto.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass
+from typing import ClassVar
+
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+from ...domain.model import Scheduler
+
+
+@dataclass
+class ComfyUISchedulerDTO:
+    """Data Transfer Object for Diffusers Scheduler in ComfyUI.
+
+    This DTO is designed to be used across different custom nodes.
+    Other node developers can import and use this type for scheduler operations.
+
+    Attributes:
+        scheduler: SchedulerMixin instance from diffusers library
+        path: Path to the source 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 ComfyUISchedulerDTO
+
+        class YourCustomNode:
+            @classmethod
+            def INPUT_TYPES(cls):
+                return {"required": {"scheduler": (ComfyUISchedulerDTO.COMFY_TYPE,)}}
+
+            def execute(self, scheduler: ComfyUISchedulerDTO):
+                # Use the scheduler for sampling operations
+                pass
+        ```
+    """
+
+    # ClassVar to avoid being treated as a dataclass field
+    COMFY_TYPE: ClassVar[str] = "DIFFUSERS_SCHEDULER"
+
+    scheduler: SchedulerMixin
+    path: str
+
+    @classmethod
+    def from_domain(cls, scheduler: Scheduler) -> "ComfyUISchedulerDTO":
+        """Create DTO from domain model.
+
+        Args:
+            scheduler: Scheduler from domain layer
+
+        Returns:
+            ComfyUISchedulerDTO instance
+        """
+        return cls(scheduler=scheduler.scheduler, path=scheduler.path)

src/nodes/dto/type/__init__.py

@@ -0,0 +1,5 @@
+from ._scheduler_type import ComfyUISchedulerType
+
+__all__ = [
+    "ComfyUISchedulerType",
+]

src/nodes/dto/type/_scheduler_type.py

@@ -0,0 +1,75 @@
+from ....domain.model import Scheduler
+
+
+class ComfyUISchedulerType:
+    SCHEDULERS: list[str] = [
+        "DDIM",
+        "DDPM",
+        "DEISMultistep",
+        "DPMSolverMultistep",
+        "DPMSolverSinglestep",
+        "EulerAncestralDiscrete",
+        "EulerDiscrete",
+        "HeunDiscrete",
+        "KDPM2AncestralDiscrete",
+        "KDPM2Discrete",
+        "UniPCMultistep",
+    ]
+
+    _DOMAIN_TO_STRING_MAP: dict[Scheduler.Type, str] = {
+        Scheduler.Type.DDIM: "DDIM",
+        Scheduler.Type.DDPM: "DDPM",
+        Scheduler.Type.DEISMultistep: "DEISMultistep",
+        Scheduler.Type.DPMSolverMultistep: "DPMSolverMultistep",
+        Scheduler.Type.DPMSolverSinglestep: "DPMSolverSinglestep",
+        Scheduler.Type.EulerAncestralDiscrete: "EulerAncestralDiscrete",
+        Scheduler.Type.EulerDiscrete: "EulerDiscrete",
+        Scheduler.Type.HeunDiscrete: "HeunDiscrete",
+        Scheduler.Type.KDPM2AncestralDiscrete: "KDPM2AncestralDiscrete",
+        Scheduler.Type.KDPM2Discrete: "KDPM2Discrete",
+        Scheduler.Type.UniPCMultistep: "UniPCMultistep",
+    }
+
+    _STRING_TO_DOMAIN_MAP: dict[str, Scheduler.Type] = {
+        "DDIM": Scheduler.Type.DDIM,
+        "DDPM": Scheduler.Type.DDPM,
+        "DEISMultistep": Scheduler.Type.DEISMultistep,
+        "DPMSolverMultistep": Scheduler.Type.DPMSolverMultistep,
+        "DPMSolverSinglestep": Scheduler.Type.DPMSolverSinglestep,
+        "EulerAncestralDiscrete": Scheduler.Type.EulerAncestralDiscrete,
+        "EulerDiscrete": Scheduler.Type.EulerDiscrete,
+        "HeunDiscrete": Scheduler.Type.HeunDiscrete,
+        "KDPM2AncestralDiscrete": Scheduler.Type.KDPM2AncestralDiscrete,
+        "KDPM2Discrete": Scheduler.Type.KDPM2Discrete,
+        "UniPCMultistep": Scheduler.Type.UniPCMultistep,
+    }
+
+    @classmethod
+    def from_domain(cls, scheduler: Scheduler.Type) -> str:
+        """Convert Scheduler.Type enum to string representation.
+
+        Args:
+            scheduler: Scheduler.Type enum value
+
+        Returns:
+            String representation of the scheduler type
+        """
+        return cls._DOMAIN_TO_STRING_MAP[scheduler]
+
+    @classmethod
+    def to_domain(cls, scheduler_type: str) -> Scheduler.Type:
+        """Convert string type to Scheduler.Type enum.
+
+        Args:
+            scheduler_type: String representation of scheduler type
+
+        Returns:
+            Corresponding Scheduler.Type enum value
+
+        Raises:
+            ValueError: If the scheduler_type is not recognized
+        """
+        if scheduler_type not in cls._STRING_TO_DOMAIN_MAP:
+            msg = f"Unknown scheduler type: {scheduler_type}"
+            raise ValueError(msg)
+        return cls._STRING_TO_DOMAIN_MAP[scheduler_type]

src/service/__init__.py

@@ -2,5 +2,12 @@ from ._autoencoder_service import AutoencoderService
 from ._clip_text_encode_service import ClipTextEncodeService
 from ._pipeline_service import PipelineService
 from ._sampling_service import SamplingService
+from ._scheduler_service import SchedulerService
 
-__all__ = ["PipelineService", "AutoencoderService", "ClipTextEncodeService", "SamplingService"]
+__all__ = [
+    "PipelineService",
+    "AutoencoderService",
+    "ClipTextEncodeService",
+    "SamplingService",
+    "SchedulerService",
+]

src/service/_scheduler_service.py

@@ -0,0 +1,19 @@
+import torch
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+from ..domain.model import Pipeline, Scheduler
+from ..domain.repositories import SchedulerRepository
+
+
+class SchedulerService:
+    def __init__(self, scheduler_repo: SchedulerRepository) -> None:
+        self.scheduler_repo = scheduler_repo
+        self.dtype = torch.float32
+
+    def create_scheduler(self, pipeline: Pipeline, scheduler_type: Scheduler.Type) -> Scheduler:
+        if pipeline is None or pipeline.pipeline is None or not pipeline.path:
+            raise ValueError("Invalid pipeline to create scheduler.")
+        scheduler_obj: SchedulerMixin = self.scheduler_repo.create_scheduler(
+            pipeline.path, self.dtype, scheduler_type
+        )
+        return Scheduler(scheduler=scheduler_obj, path=pipeline.path)

src/ui/__init__.py

@@ -2,5 +2,12 @@ from ._autoencoder_handler import AutoencoderHandler
 from ._clip_text_encode_handler import ClipTextEncodeHandler
 from ._pipeline_handler import PipelineHandler
 from ._sampler_handler import SamplerHandler
+from ._scheduler_handler import SchedulerHandler
 
-__all__ = ["PipelineHandler", "AutoencoderHandler", "ClipTextEncodeHandler", "SamplerHandler"]
+__all__ = [
+    "PipelineHandler",
+    "AutoencoderHandler",
+    "ClipTextEncodeHandler",
+    "SamplerHandler",
+    "SchedulerHandler",
+]

src/ui/_scheduler_handler.py

@@ -0,0 +1,12 @@
+from ..domain.model import Pipeline, Scheduler
+from ..service import SchedulerService
+
+
+class SchedulerHandler:
+    def __init__(self, scheduler_service: SchedulerService) -> None:
+        self.scheduler_service = scheduler_service
+
+    def create(self, pipeline: Pipeline, scheduler_type: Scheduler.Type) -> Scheduler:
+        if pipeline is None or pipeline.pipeline is None:
+            raise ValueError("Pipeline is None.")
+        return self.scheduler_service.create_scheduler(pipeline, scheduler_type)