feat(domain): add validation and encapsulation to domain models

src/domain/model/__init__.py

@@ -1,9 +1,13 @@
 from ._autoencoder import Autoencoder
+from ._cfg_scale import CFGScale
 from ._clip import Clip
 from ._conditioning import Conditioning
 from ._image import Image
+from ._image_size import ImageSize
 from ._pipeline import Pipeline
 from ._scheduler import Scheduler
+from ._seed import Seed
+from ._steps import Steps
 
 __all__ = [
     "Pipeline",
@@ -12,4 +16,8 @@ __all__ = [
     "Conditioning",
     "Image",
     "Scheduler",
+    "ImageSize",
+    "Steps",
+    "CFGScale",
+    "Seed",
 ]

src/domain/model/_autoencoder.py

@@ -3,7 +3,17 @@ from dataclasses import dataclass
 from diffusers import AutoencoderKL
 
 
-@dataclass
+@dataclass(frozen=True)
 class Autoencoder:
     autoencoder: AutoencoderKL
     path: str
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.autoencoder, AutoencoderKL):
+            raise ValueError(
+                f"Autoencoder must be an instance of AutoencoderKL, got {type(self.autoencoder).__name__}"
+            )
+        if not isinstance(self.path, str):
+            raise ValueError(f"Path must be a string, got {type(self.path).__name__}")
+        if not self.path.strip():
+            raise ValueError("Path cannot be empty or whitespace only")

src/domain/model/_cfg_scale.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class CFGScale:
+    value: float
+
+    def __post_init__(self) -> None:
+        if isinstance(self.value, bool) or not isinstance(self.value, float):
+            raise ValueError(f"CFG scale must be a float, got {type(self.value).__name__}")
+        if self.value < 0.0:
+            raise ValueError(f"CFG scale must be non-negative, got {self.value}")

src/domain/model/_clip.py

@@ -3,8 +3,22 @@ from dataclasses import dataclass
 from transformers import CLIPTextModel, CLIPTokenizer
 
 
-@dataclass
+@dataclass(frozen=True)
 class Clip:
     tokenizer: CLIPTokenizer
     text_encoder: CLIPTextModel
     path: str
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.tokenizer, CLIPTokenizer):
+            raise ValueError(
+                f"Tokenizer must be an instance of CLIPTokenizer, got {type(self.tokenizer).__name__}"
+            )
+        if not isinstance(self.text_encoder, CLIPTextModel):
+            raise ValueError(
+                f"Text encoder must be an instance of CLIPTextModel, got {type(self.text_encoder).__name__}"
+            )
+        if not isinstance(self.path, str):
+            raise ValueError(f"Path must be a string, got {type(self.path).__name__}")
+        if not self.path.strip():
+            raise ValueError("Path cannot be empty or whitespace only")

src/domain/model/_conditioning.py

@@ -3,7 +3,21 @@ from dataclasses import dataclass
 import torch
 
 
-@dataclass
+@dataclass(frozen=True)
 class Conditioning:
     conditioning: torch.Tensor
     path: str
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.conditioning, torch.Tensor):
+            raise ValueError(
+                f"Conditioning must be a torch.Tensor, got {type(self.conditioning).__name__}"
+            )
+        if not isinstance(self.path, str):
+            raise ValueError(f"Path must be a string, got {type(self.path).__name__}")
+        if torch.isnan(self.conditioning).any():
+            raise ValueError("Conditioning tensor contains NaN values")
+        if torch.isinf(self.conditioning).any():
+            raise ValueError("Conditioning tensor contains Inf values")
+        if not self.path.strip():
+            raise ValueError("Path cannot be empty or whitespace only")

src/domain/model/_image.py

@@ -1,8 +1,12 @@
 from dataclasses import dataclass
 
-import PIL.Image
+from PIL import Image as PilImage
 
 
-@dataclass
+@dataclass(frozen=True)
 class Image:
-    image: PIL.Image.Image
+    image: PilImage.Image
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.image, PilImage.Image):
+            raise ValueError(f"Image must be a PIL.Image.Image, got {type(self.image).__name__}")

src/domain/model/_image_size.py

@@ -0,0 +1,21 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ImageSize:
+    width: int
+    height: int
+
+    def __post_init__(self) -> None:
+        if isinstance(self.width, bool) or not isinstance(self.width, int):
+            raise ValueError(f"Width must be an int, got {type(self.width).__name__}")
+        if isinstance(self.height, bool) or not isinstance(self.height, int):
+            raise ValueError(f"Height must be an int, got {type(self.height).__name__}")
+        if self.width <= 0:
+            raise ValueError(f"Width must be positive, got {self.width}")
+        if self.height <= 0:
+            raise ValueError(f"Height must be positive, got {self.height}")
+        if self.width % 8 != 0:
+            raise ValueError(f"Width must be divisible by 8, got {self.width}")
+        if self.height % 8 != 0:
+            raise ValueError(f"Height must be divisible by 8, got {self.height}")

src/domain/model/_pipeline.py

@@ -3,7 +3,17 @@ from dataclasses import dataclass
 from diffusers import StableDiffusionPipeline
 
 
-@dataclass
+@dataclass(frozen=True)
 class Pipeline:
     pipeline: StableDiffusionPipeline
     path: str
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.pipeline, StableDiffusionPipeline):
+            raise ValueError(
+                f"Pipeline must be an instance of StableDiffusionPipeline, got {type(self.pipeline).__name__}"
+            )
+        if not isinstance(self.path, str):
+            raise ValueError(f"Path must be a string, got {type(self.path).__name__}")
+        if not self.path.strip():
+            raise ValueError("Path cannot be empty or whitespace only")

src/domain/model/_scheduler.py

@@ -17,7 +17,7 @@ from diffusers import (
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
 
 
-@dataclass
+@dataclass(frozen=True)
 class Scheduler:
     class Type(Enum):
         DDIM = DDIMScheduler
@@ -34,3 +34,13 @@ class Scheduler:
 
     scheduler: SchedulerMixin
     path: str
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.scheduler, SchedulerMixin):
+            raise ValueError(
+                f"Scheduler must inherit from SchedulerMixin, got {type(self.scheduler).__name__}"
+            )
+        if not isinstance(self.path, str):
+            raise ValueError(f"Path must be a string, got {type(self.path).__name__}")
+        if not self.path.strip():
+            raise ValueError("Path cannot be empty or whitespace only")

src/domain/model/_seed.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Seed:
+    value: int
+
+    def __post_init__(self) -> None:
+        if isinstance(self.value, bool) or not isinstance(self.value, int):
+            raise ValueError(f"Seed must be an int, got {type(self.value).__name__}")

src/domain/model/_steps.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Steps:
+    value: int
+
+    def __post_init__(self) -> None:
+        if isinstance(self.value, bool) or not isinstance(self.value, int):
+            raise ValueError(f"Steps must be an int, got {type(self.value).__name__}")
+        if self.value <= 0:
+            raise ValueError(f"Steps must be positive, got {self.value}")

src/domain/repositories/_sampler_repository.py

@@ -5,6 +5,8 @@ from diffusers import AutoencoderKL, StableDiffusionPipeline
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from PIL import Image
 
+from ..model import CFGScale, ImageSize, Seed, Steps
+
 
 class SamplerRepository(ABC):
     @abstractmethod
@@ -15,10 +17,9 @@ class SamplerRepository(ABC):
         scheduler: SchedulerMixin,
         positive_embeds: torch.Tensor,
         negative_embeds: torch.Tensor,
-        width: int,
-        height: int,
-        steps: int,
-        cfg: float,
-        seed: int,
+        image_size: ImageSize,
+        steps: Steps,
+        cfg: CFGScale,
+        seed: Seed,
     ) -> list[Image.Image]:
         pass