feat: scheduler node

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