diff --git a/src/domain/model/__init__.py b/src/domain/model/__init__.py index dc4f4e9..f3b54e9 100644 --- a/src/domain/model/__init__.py +++ b/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", ] diff --git a/src/domain/model/_autoencoder.py b/src/domain/model/_autoencoder.py index e3b703f..2469a58 100644 --- a/src/domain/model/_autoencoder.py +++ b/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") diff --git a/src/domain/model/_cfg_scale.py b/src/domain/model/_cfg_scale.py new file mode 100644 index 0000000..c34c97c --- /dev/null +++ b/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}") diff --git a/src/domain/model/_clip.py b/src/domain/model/_clip.py index 35bd864..cd518f6 100644 --- a/src/domain/model/_clip.py +++ b/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") diff --git a/src/domain/model/_conditioning.py b/src/domain/model/_conditioning.py index f827afe..658be4f 100644 --- a/src/domain/model/_conditioning.py +++ b/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") diff --git a/src/domain/model/_image.py b/src/domain/model/_image.py index 8b7f6ec..e6d6200 100644 --- a/src/domain/model/_image.py +++ b/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__}") diff --git a/src/domain/model/_image_size.py b/src/domain/model/_image_size.py new file mode 100644 index 0000000..f00ee12 --- /dev/null +++ b/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}") diff --git a/src/domain/model/_pipeline.py b/src/domain/model/_pipeline.py index 3641a9c..2da682a 100644 --- a/src/domain/model/_pipeline.py +++ b/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") diff --git a/src/domain/model/_scheduler.py b/src/domain/model/_scheduler.py index cbcb5f6..ee1d5ea 100644 --- a/src/domain/model/_scheduler.py +++ b/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") diff --git a/src/domain/model/_seed.py b/src/domain/model/_seed.py new file mode 100644 index 0000000..09800de --- /dev/null +++ b/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__}") diff --git a/src/domain/model/_steps.py b/src/domain/model/_steps.py new file mode 100644 index 0000000..3b04174 --- /dev/null +++ b/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}") diff --git a/src/domain/repositories/_sampler_repository.py b/src/domain/repositories/_sampler_repository.py index 14d627c..ff0f41e 100644 --- a/src/domain/repositories/_sampler_repository.py +++ b/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 diff --git a/src/infra/diffusers/_sampler_repository.py b/src/infra/diffusers/_sampler_repository.py index ef23525..fe6db34 100644 --- a/src/infra/diffusers/_sampler_repository.py +++ b/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 diff --git a/src/nodes/_diffusers_clip_text_encode.py b/src/nodes/_diffusers_clip_text_encode.py index ffb0756..7ccd910 100644 --- a/src/nodes/_diffusers_clip_text_encode.py +++ b/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,) diff --git a/src/nodes/_diffusers_sampler.py b/src/nodes/_diffusers_sampler.py index 5c87135..cf91d76 100644 --- a/src/nodes/_diffusers_sampler.py +++ b/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,) diff --git a/src/nodes/dto/_clip.py b/src/nodes/dto/_clip.py index 13471a6..813cf4a 100644 --- a/src/nodes/dto/_clip.py +++ b/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": diff --git a/src/nodes/dto/_conditioning_dto.py b/src/nodes/dto/_conditioning_dto.py index 3d7e09c..09fc821 100644 --- a/src/nodes/dto/_conditioning_dto.py +++ b/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: diff --git a/src/usecase/_clip_text_encode_usecase.py b/src/usecase/_clip_text_encode_usecase.py index 84578ae..47f9379 100644 --- a/src/usecase/_clip_text_encode_usecase.py +++ b/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.") diff --git a/src/usecase/_sampler_usecase.py b/src/usecase/_sampler_usecase.py index cccdbe2..780e03c 100644 --- a/src/usecase/_sampler_usecase.py +++ b/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] diff --git a/src/usecase/_scheduler_usecase.py b/src/usecase/_scheduler_usecase.py index 8a9e67f..4c9d369 100644 --- a/src/usecase/_scheduler_usecase.py +++ b/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)