diff --git a/src/nodes/_diffusers_scheduler_loader.py b/src/nodes/_diffusers_scheduler_loader.py index 2746c7d..ff74972 100644 --- a/src/nodes/_diffusers_scheduler_loader.py +++ b/src/nodes/_diffusers_scheduler_loader.py @@ -3,7 +3,7 @@ from dependency_injector.wiring import Provide, inject from ..di import Container from ..usecase import SchedulerUsecase from .dto import ComfyUIPipelineDTO, ComfyUISchedulerDTO -from .type import ComfyUISchedulerType +from .map import SchedulerMap class DiffusersSchedulerLoader: @@ -15,7 +15,7 @@ class DiffusersSchedulerLoader: return { "required": { "pipeline": (ComfyUIPipelineDTO.COMFY_TYPE,), - "scheduler_name": (list(ComfyUISchedulerType.SCHEDULERS),), + "scheduler_name": (SchedulerMap.SCHEDULERS,), } } @@ -31,7 +31,7 @@ class DiffusersSchedulerLoader: usecase: SchedulerUsecase = Provide[Container.scheduler_usecase], ) -> tuple[ComfyUISchedulerDTO]: domain = ComfyUIPipelineDTO.to_domain(pipeline) - scheduler_type = ComfyUISchedulerType.to_domain(scheduler_name) + scheduler_type = SchedulerMap.to_domain(scheduler_name) scheduler_domain = usecase.execute(domain, scheduler_type) scheduler_dto = ComfyUISchedulerDTO.from_domain(scheduler_domain) return (scheduler_dto,) diff --git a/src/nodes/map/__init__.py b/src/nodes/map/__init__.py new file mode 100644 index 0000000..14a8cf1 --- /dev/null +++ b/src/nodes/map/__init__.py @@ -0,0 +1,5 @@ +from ._scheduler_map import SchedulerMap + +__all__ = [ + "SchedulerMap", +] diff --git a/src/nodes/map/_scheduler_map.py b/src/nodes/map/_scheduler_map.py new file mode 100644 index 0000000..3c0a38a --- /dev/null +++ b/src/nodes/map/_scheduler_map.py @@ -0,0 +1,76 @@ +from typing import Final + +from ...domain.model import Scheduler + + +class SchedulerMap: + """Maps between domain Scheduler.Type enum and ComfyUI string representation. + + This map provides bidirectional conversion between the domain layer's + type-safe enum and ComfyUI's string-based interface requirements. + """ + + SCHEDULERS: Final[list[str]] = [ + "DDIM", + "DDPM", + "DEISMultistep", + "DPMSolverMultistep", + "DPMSolverSinglestep", + "EulerAncestralDiscrete", + "EulerDiscrete", + "HeunDiscrete", + "KDPM2AncestralDiscrete", + "KDPM2Discrete", + "UniPCMultistep", + ] + + _DOMAIN_TO_STRING: 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: dict[str, Scheduler.Type] = {v: k for k, v in _DOMAIN_TO_STRING.items()} + + @classmethod + def from_domain(cls, scheduler_type: Scheduler.Type) -> str: + """Convert domain Scheduler.Type enum to ComfyUI string representation. + + Args: + scheduler_type: Scheduler.Type enum value from domain layer + + Returns: + String representation for ComfyUI interface + + Raises: + KeyError: If the scheduler_type is not in the mapping + """ + return cls._DOMAIN_TO_STRING[scheduler_type] + + @classmethod + def to_domain(cls, scheduler_name: str) -> Scheduler.Type: + """Convert ComfyUI string to domain Scheduler.Type enum. + + Args: + scheduler_name: String representation from ComfyUI interface + + Returns: + Corresponding Scheduler.Type enum value from domain layer + + Raises: + ValueError: If the scheduler_name is not recognized + """ + if scheduler_name not in cls._STRING_TO_DOMAIN: + raise ValueError( + f"Unknown scheduler type: '{scheduler_name}'. " + f"Available types: {', '.join(cls.SCHEDULERS)}" + ) + return cls._STRING_TO_DOMAIN[scheduler_name] diff --git a/src/nodes/type/__init__.py b/src/nodes/type/__init__.py deleted file mode 100644 index 1f890b0..0000000 --- a/src/nodes/type/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from ._scheduler_type import ComfyUISchedulerType - -__all__ = [ - "ComfyUISchedulerType", -] diff --git a/src/nodes/type/_scheduler_type.py b/src/nodes/type/_scheduler_type.py deleted file mode 100644 index 3ba1dbc..0000000 --- a/src/nodes/type/_scheduler_type.py +++ /dev/null @@ -1,75 +0,0 @@ -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]