chore: connect nodes to sampler

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,

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,

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,

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.")

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.")