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

src/infra/diffusers/_sampler_repository.py

@@ -4,6 +4,7 @@ from diffusers import AutoencoderKL, StableDiffusionPipeline
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from PIL import Image
 
+from ...domain.model import CFGScale, ImageSize, Seed, Steps
 from ...domain.repositories import SamplerRepository
 
 
@@ -18,22 +19,21 @@ class DiffusersSamplerRepository(SamplerRepository):
         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]:
         result = pipeline(  # type: ignore[operator]
             prompt_embeds=positive_embeds,
             vae=vae,
             scheduler=scheduler,
-            height=height,
-            width=width,
-            num_inference_steps=steps,
-            guidance_scale=cfg,
+            height=image_size.height,
+            width=image_size.width,
+            num_inference_steps=steps.value,
+            guidance_scale=cfg.value,
             negative_prompt_embeds=negative_embeds,
-            generator=torch.Generator(self.device).manual_seed(seed),
+            generator=torch.Generator(self.device).manual_seed(seed.value),
         )
         images: list[Image.Image] = result.images
         return images

src/nodes/_diffusers_clip_text_encode.py

@@ -31,7 +31,10 @@ class DiffusersClipTextEncode:
     ) -> tuple[ComfyUIConditioningDTO]:
         clip_domain = ComfyUIClipDTO.to_domain(clip)
         conditioning_model = usecase.execute(
-            clip_domain.tokenizer, clip_domain.text_encoder, clip_domain.path, text
+            clip_domain.tokenizer,
+            clip_domain.text_encoder,
+            clip_domain.path,
+            text,
         )
         conditioning_dto = ComfyUIConditioningDTO.from_domain(conditioning_model)
         return (conditioning_dto,)

src/nodes/_diffusers_sampler.py

@@ -1,6 +1,7 @@
 from dependency_injector.wiring import Provide, inject
 
 from ..di import Container
+from ..domain.model import CFGScale, ImageSize, Seed, Steps
 from ..usecase import SamplerUsecase
 from .dto import (
     ComfyUIAutoencoderDTO,
@@ -62,17 +63,22 @@ class DiffusersSampler:
         scheduler_domain = ComfyUISchedulerDTO.to_domain(scheduler)
         positive_embeds_domain = ComfyUIConditioningDTO.to_domain(positive_embeds)
         negative_embeds_domain = ComfyUIConditioningDTO.to_domain(negative_embeds)
+
+        image_size = ImageSize(width=width, height=height)
+        steps_vo = Steps(value=steps)
+        cfg_vo = CFGScale(value=cfg)
+        seed_vo = Seed(value=seed)
+
         images_model = usecase.execute(
             pipeline_domain,
             vae_domain,
             scheduler_domain,
             positive_embeds_domain,
             negative_embeds_domain,
-            width,
-            height,
-            steps,
-            cfg,
-            seed,
+            image_size,
+            steps_vo,
+            cfg_vo,
+            seed_vo,
         )
         sampler_dto = ComfyUIImage.from_domains(images_model)
         return (sampler_dto,)

src/nodes/dto/_clip.py

@@ -57,7 +57,11 @@ class ComfyUIClipDTO:
         Returns:
             ComfyUIClipDTO instance
         """
-        return cls(tokenizer=clip.tokenizer, text_encoder=clip.text_encoder, path=clip.path)
+        return cls(
+            tokenizer=clip.tokenizer,
+            text_encoder=clip.text_encoder,
+            path=clip.path,
+        )
 
     @classmethod
     def to_domain(cls, dto: "ComfyUIClipDTO") -> "Clip":

src/nodes/dto/_conditioning_dto.py

@@ -54,7 +54,10 @@ class ComfyUIConditioningDTO:
         Returns:
             ComfyUIConditioningDTO instance
         """
-        return cls(conditioning=conditioning.conditioning, path=conditioning.path)
+        return cls(
+            conditioning=conditioning.conditioning,
+            path=conditioning.path,
+        )
 
     @classmethod
     def to_domain(cls, dto: "ComfyUIConditioningDTO") -> Conditioning:

src/usecase/_clip_text_encode_usecase.py

@@ -11,7 +11,11 @@ class ClipTextEncodeUsecase:
         self.text_encoder_repo = text_encoder_repo
 
     def execute(
-        self, tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel, path: str, text: str
+        self,
+        tokenizer: CLIPTokenizer,
+        text_encoder: CLIPTextModel,
+        path: str,
+        text: str,
     ) -> Conditioning:
         if not text.strip():
             raise ValueError("Text input is empty or whitespace only.")

src/usecase/_sampler_usecase.py

@@ -1,6 +1,18 @@
+from __future__ import annotations
+
 import torch
 
-from ..domain.model import Autoencoder, Conditioning, Image, Pipeline, Scheduler
+from ..domain.model import (
+    Autoencoder,
+    CFGScale,
+    Conditioning,
+    Image,
+    ImageSize,
+    Pipeline,
+    Scheduler,
+    Seed,
+    Steps,
+)
 from ..domain.repositories import SamplerRepository
 
 
@@ -16,35 +28,24 @@ class SamplerUsecase:
         scheduler: Scheduler,
         positive_embeds: Conditioning,
         negative_embeds: Conditioning,
-        width: int,
-        height: int,
-        steps: int,
-        cfg: float,
-        seed: int,
+        image_size: ImageSize,
+        steps: Steps,
+        cfg: CFGScale,
+        seed: Seed,
     ) -> list[Image]:
-        if width <= 0 or height <= 0:
-            raise ValueError(f"Invalid dimensions: {width}x{height}")
-        if steps <= 0:
-            raise ValueError(f"Steps must be positive: {steps}")
-        if cfg < 0.0:
-            raise ValueError(f"CFG must be non-negative: {cfg}")
-        if seed < 0:
-            raise ValueError(f"Seed must be non-negative: {seed}")
-
         images = self.sampler_repo.sample(
             pipeline.pipeline,
             vae.autoencoder,
             scheduler.scheduler,
             positive_embeds.conditioning,
             negative_embeds.conditioning,
-            width,
-            height,
+            image_size,
             steps,
             cfg,
             seed,
         )
-        if images is None:
-            raise RuntimeError("Sampler repository returned no images.")
-        images_domain: list[Image] = [Image(image=img) for img in images]
 
-        return images_domain
+        if not images:
+            raise RuntimeError("Sampler repository returned no images.")
+
+        return [Image(image=img) for img in images]

src/usecase/_scheduler_usecase.py

@@ -11,13 +11,10 @@ class SchedulerUsecase:
         self.dtype = torch.float32
 
     def execute(self, pipeline: Pipeline, scheduler_type: Scheduler.Type) -> Scheduler:
-        if pipeline is None or pipeline.pipeline is None:
-            raise ValueError("Pipeline is None.")
-        if not pipeline.path:
-            raise ValueError("Invalid pipeline to create scheduler.")
-
         scheduler_obj: SchedulerMixin = self.scheduler_repo.create_scheduler(
-            pipeline.path, self.dtype, scheduler_type
+            pipeline.path,
+            self.dtype,
+            scheduler_type,
         )
 
         return Scheduler(scheduler=scheduler_obj, path=pipeline.path)