From 4be7637b5184c56ad598ebb81a1cd6db561bf4e1 Mon Sep 17 00:00:00 2001 From: Limitex <76650151+Limitex@users.noreply.github.com> Date: Wed, 5 Nov 2025 21:16:15 +0900 Subject: [PATCH] feat: scheduler node --- src/di/di.py | 26 ++++++- src/domain/model/__init__.py | 2 + src/domain/model/_scheduler.py | 36 +++++++++ src/domain/repositories/__init__.py | 2 + .../repositories/_scheduler_repository.py | 14 ++++ src/infra/diffusers/__init__.py | 2 + src/infra/diffusers/_scheduler_repository.py | 22 ++++++ src/nodes/__init__.py | 12 ++- src/nodes/_diffusers_scheduler_loader.py | 37 +++++++++ src/nodes/dto/__init__.py | 2 + src/nodes/dto/_pipeline_dto.py | 12 +++ src/nodes/dto/_scheduler_dto.py | 54 +++++++++++++ src/nodes/dto/type/__init__.py | 5 ++ src/nodes/dto/type/_scheduler_type.py | 75 +++++++++++++++++++ src/service/__init__.py | 9 ++- src/service/_scheduler_service.py | 19 +++++ src/ui/__init__.py | 9 ++- src/ui/_scheduler_handler.py | 12 +++ 18 files changed, 345 insertions(+), 5 deletions(-) create mode 100644 src/domain/model/_scheduler.py create mode 100644 src/domain/repositories/_scheduler_repository.py create mode 100644 src/infra/diffusers/_scheduler_repository.py create mode 100644 src/nodes/_diffusers_scheduler_loader.py create mode 100644 src/nodes/dto/_scheduler_dto.py create mode 100644 src/nodes/dto/type/__init__.py create mode 100644 src/nodes/dto/type/_scheduler_type.py create mode 100644 src/service/_scheduler_service.py create mode 100644 src/ui/_scheduler_handler.py diff --git a/src/di/di.py b/src/di/di.py index aa93669..c93fa67 100644 --- a/src/di/di.py +++ b/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, + ) diff --git a/src/domain/model/__init__.py b/src/domain/model/__init__.py index 0fd96dd..dc4f4e9 100644 --- a/src/domain/model/__init__.py +++ b/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", ] diff --git a/src/domain/model/_scheduler.py b/src/domain/model/_scheduler.py new file mode 100644 index 0000000..cbcb5f6 --- /dev/null +++ b/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 diff --git a/src/domain/repositories/__init__.py b/src/domain/repositories/__init__.py index 4604ffa..e0268b9 100644 --- a/src/domain/repositories/__init__.py +++ b/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", ] diff --git a/src/domain/repositories/_scheduler_repository.py b/src/domain/repositories/_scheduler_repository.py new file mode 100644 index 0000000..2b84f66 --- /dev/null +++ b/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 diff --git a/src/infra/diffusers/__init__.py b/src/infra/diffusers/__init__.py index a1e7a55..3f20a8d 100644 --- a/src/infra/diffusers/__init__.py +++ b/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", ] diff --git a/src/infra/diffusers/_scheduler_repository.py b/src/infra/diffusers/_scheduler_repository.py new file mode 100644 index 0000000..4ccc834 --- /dev/null +++ b/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 diff --git a/src/nodes/__init__.py b/src/nodes/__init__.py index da18a27..9c6b4ca 100644 --- a/src/nodes/__init__.py +++ b/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", ] diff --git a/src/nodes/_diffusers_scheduler_loader.py b/src/nodes/_diffusers_scheduler_loader.py new file mode 100644 index 0000000..9795352 --- /dev/null +++ b/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,) diff --git a/src/nodes/dto/__init__.py b/src/nodes/dto/__init__.py index 44a6c56..c5c67b1 100644 --- a/src/nodes/dto/__init__.py +++ b/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", ] diff --git a/src/nodes/dto/_pipeline_dto.py b/src/nodes/dto/_pipeline_dto.py index f43a6f3..5dd18d3 100644 --- a/src/nodes/dto/_pipeline_dto.py +++ b/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) diff --git a/src/nodes/dto/_scheduler_dto.py b/src/nodes/dto/_scheduler_dto.py new file mode 100644 index 0000000..51ec5fc --- /dev/null +++ b/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) diff --git a/src/nodes/dto/type/__init__.py b/src/nodes/dto/type/__init__.py new file mode 100644 index 0000000..1f890b0 --- /dev/null +++ b/src/nodes/dto/type/__init__.py @@ -0,0 +1,5 @@ +from ._scheduler_type import ComfyUISchedulerType + +__all__ = [ + "ComfyUISchedulerType", +] diff --git a/src/nodes/dto/type/_scheduler_type.py b/src/nodes/dto/type/_scheduler_type.py new file mode 100644 index 0000000..2bd8586 --- /dev/null +++ b/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] diff --git a/src/service/__init__.py b/src/service/__init__.py index b06377c..6d4bff8 100644 --- a/src/service/__init__.py +++ b/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", +] diff --git a/src/service/_scheduler_service.py b/src/service/_scheduler_service.py new file mode 100644 index 0000000..2ff4da8 --- /dev/null +++ b/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) diff --git a/src/ui/__init__.py b/src/ui/__init__.py index f14e125..af0a7c9 100644 --- a/src/ui/__init__.py +++ b/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", +] diff --git a/src/ui/_scheduler_handler.py b/src/ui/_scheduler_handler.py new file mode 100644 index 0000000..aade076 --- /dev/null +++ b/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)