diff --git a/src/domain/repositories/_sampling_repository.py b/src/domain/repositories/_sampling_repository.py index 612f69e..d9c1204 100644 --- a/src/domain/repositories/_sampling_repository.py +++ b/src/domain/repositories/_sampling_repository.py @@ -1,7 +1,8 @@ from abc import ABC, abstractmethod import torch -from diffusers import StableDiffusionPipeline +from diffusers import AutoencoderKL, StableDiffusionPipeline +from diffusers.schedulers.scheduling_utils import SchedulerMixin from PIL import Image @@ -10,6 +11,8 @@ class SamplingRepository(ABC): def sample( self, pipeline: StableDiffusionPipeline, + vae: AutoencoderKL, + scheduler: SchedulerMixin, positive_embeds: torch.Tensor, negative_embeds: torch.Tensor, width: int, diff --git a/src/infra/diffusers/_sampling_repository.py b/src/infra/diffusers/_sampling_repository.py index 3554758..96dcc09 100644 --- a/src/infra/diffusers/_sampling_repository.py +++ b/src/infra/diffusers/_sampling_repository.py @@ -1,6 +1,7 @@ import torch from comfy.model_management import get_torch_device # pyright: ignore[reportMissingImports] -from diffusers import StableDiffusionPipeline +from diffusers import AutoencoderKL, StableDiffusionPipeline +from diffusers.schedulers.scheduling_utils import SchedulerMixin from PIL import Image from ...domain.repositories import SamplingRepository @@ -13,6 +14,8 @@ class DiffusersSamplingRepository(SamplingRepository): def sample( self, pipeline: StableDiffusionPipeline, + vae: AutoencoderKL, + scheduler: SchedulerMixin, positive_embeds: torch.Tensor, negative_embeds: torch.Tensor, width: int, @@ -23,6 +26,8 @@ class DiffusersSamplingRepository(SamplingRepository): ) -> list[Image.Image]: result = pipeline( # type: ignore[operator] prompt_embeds=positive_embeds, + vae=vae, + scheduler=scheduler, height=height, width=width, num_inference_steps=steps, diff --git a/src/nodes/_diffusers_sampler.py b/src/nodes/_diffusers_sampler.py index ab38c68..7302e89 100644 --- a/src/nodes/_diffusers_sampler.py +++ b/src/nodes/_diffusers_sampler.py @@ -2,7 +2,14 @@ from dependency_injector.wiring import Provide, inject from ..di import Container from ..ui import SamplerHandler -from .dto import ComfyUIConditioningDTO, ComfyUIImage, ComfyUIImageDTO, ComfyUIPipelineDTO +from .dto import ( + ComfyUIAutoencoderDTO, + ComfyUIConditioningDTO, + ComfyUIImage, + ComfyUIImageDTO, + ComfyUIPipelineDTO, + ComfyUISchedulerDTO, +) class DiffusersSampler: @@ -16,6 +23,8 @@ class DiffusersSampler: return { "required": { "pipeline": (ComfyUIPipelineDTO.COMFY_TYPE,), + "vae": (ComfyUIAutoencoderDTO.COMFY_TYPE,), + "scheduler": (ComfyUISchedulerDTO.COMFY_TYPE,), "positive_embeds": (ComfyUIConditioningDTO.COMFY_TYPE,), "negative_embeds": (ComfyUIConditioningDTO.COMFY_TYPE,), "width": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1.0}), @@ -37,6 +46,8 @@ class DiffusersSampler: def execute( self, pipeline: ComfyUIPipelineDTO, + vae: ComfyUIAutoencoderDTO, + scheduler: ComfyUISchedulerDTO, positive_embeds: ComfyUIConditioningDTO, negative_embeds: ComfyUIConditioningDTO, width: int, @@ -46,10 +57,17 @@ class DiffusersSampler: seed: int, handler: SamplerHandler = Provide[Container.sampler_handler], ) -> tuple[ComfyUIImageDTO]: + pipeline_domain = ComfyUIPipelineDTO.to_domain(pipeline) + vae_domain = ComfyUIAutoencoderDTO.to_domain(vae) + scheduler_domain = ComfyUISchedulerDTO.to_domain(scheduler) + positive_embeds_domain = ComfyUIConditioningDTO.to_domain(positive_embeds) + negative_embeds_domain = ComfyUIConditioningDTO.to_domain(negative_embeds) images_model = handler.sample( - pipeline.pipeline, - positive_embeds.conditioning, - negative_embeds.conditioning, + pipeline_domain, + vae_domain, + scheduler_domain, + positive_embeds_domain, + negative_embeds_domain, width, height, steps, diff --git a/src/service/_sampling_service.py b/src/service/_sampling_service.py index 2167ec7..961e3d4 100644 --- a/src/service/_sampling_service.py +++ b/src/service/_sampling_service.py @@ -1,7 +1,6 @@ import torch -from diffusers import StableDiffusionPipeline -from ..domain.model import Image +from ..domain.model import Autoencoder, Conditioning, Image, Pipeline, Scheduler from ..domain.repositories import SamplingRepository @@ -12,9 +11,11 @@ class SamplingService: def sample( self, - pipeline: StableDiffusionPipeline, - positive_embeds: torch.Tensor, - negative_embeds: torch.Tensor, + pipeline: Pipeline, + vae: Autoencoder, + scheduler: Scheduler, + positive_embeds: Conditioning, + negative_embeds: Conditioning, width: int, height: int, steps: int, @@ -22,7 +23,16 @@ class SamplingService: seed: int, ) -> list[Image]: images = self.sampling_repo.sample( - pipeline, positive_embeds, negative_embeds, width, height, steps, cfg, seed + pipeline.pipeline, + vae.autoencoder, + scheduler.scheduler, + positive_embeds.conditioning, + negative_embeds.conditioning, + width, + height, + steps, + cfg, + seed, ) if images is None: raise RuntimeError("Sampling repository returned no images.") diff --git a/src/ui/_sampler_handler.py b/src/ui/_sampler_handler.py index 1ba56a7..fe9247a 100644 --- a/src/ui/_sampler_handler.py +++ b/src/ui/_sampler_handler.py @@ -1,7 +1,4 @@ -import torch -from diffusers import StableDiffusionPipeline - -from ..domain.model import Image +from ..domain.model import Autoencoder, Conditioning, Image, Pipeline, Scheduler from ..service import SamplingService @@ -11,9 +8,11 @@ class SamplerHandler: def sample( self, - pipeline: StableDiffusionPipeline, - positive_embeds: torch.Tensor, - negative_embeds: torch.Tensor, + pipeline: Pipeline, + vae: Autoencoder, + scheduler: Scheduler, + positive_embeds: Conditioning, + negative_embeds: Conditioning, width: int, height: int, steps: int, @@ -22,6 +21,10 @@ class SamplerHandler: ) -> list[Image]: if pipeline is None: raise ValueError("Pipeline is None.") + if vae is None: + raise ValueError("VAE is None.") + if scheduler is None: + raise ValueError("Scheduler is None.") if positive_embeds is None or negative_embeds is None: raise ValueError("Conditioning embeddings are None.") if width <= 0 or height <= 0: @@ -33,7 +36,16 @@ class SamplerHandler: if seed < 0: raise ValueError("Seed must be a non-negative integer.") domain_images = self.sampling_service.sample( - pipeline, positive_embeds, negative_embeds, width, height, steps, cfg, seed + pipeline, + vae, + scheduler, + positive_embeds, + negative_embeds, + width, + height, + steps, + cfg, + seed, ) if domain_images is None: raise RuntimeError("Failed to sample images from pipeline.")