diff --git a/src/di/di.py b/src/di/di.py index c93fa67..3185a49 100644 --- a/src/di/di.py +++ b/src/di/di.py @@ -5,23 +5,16 @@ from dependency_injector import containers, providers from ..infra.diffusers import ( DiffusersAutoencoderRepository, DiffusersPipelineRepository, - DiffusersSamplingRepository, + DiffusersSamplerRepository, DiffusersSchedulerRepository, DiffusersTextEncoderRepository, ) -from ..service import ( - AutoencoderService, - ClipTextEncodeService, - PipelineService, - SamplingService, - SchedulerService, -) -from ..ui import ( - AutoencoderHandler, - ClipTextEncodeHandler, - PipelineHandler, - SamplerHandler, - SchedulerHandler, +from ..usecase import ( + AutoencoderUsecase, + ClipTextEncodeUsecase, + PipelineUsecase, + SamplerUsecase, + SchedulerUsecase, ) @@ -32,49 +25,27 @@ class Container(containers.DeclarativeContainer): pipeline_repository = providers.Factory(DiffusersPipelineRepository) autoencoder_repository = providers.Factory(DiffusersAutoencoderRepository) text_encoder_repository = providers.Factory(DiffusersTextEncoderRepository) - sampling_repository = providers.Factory(DiffusersSamplingRepository) + sampler_repository = providers.Factory(DiffusersSamplerRepository) scheduler_repository = providers.Factory(DiffusersSchedulerRepository) - # 2. Services - pipeline_service = providers.Factory( - PipelineService, + # 2. Usecases + pipeline_usecase = providers.Factory( + PipelineUsecase, pipeline_repo=pipeline_repository, ) - autoencoder_service = providers.Factory( - AutoencoderService, + autoencoder_usecase = providers.Factory( + AutoencoderUsecase, autoencoder_repo=autoencoder_repository, ) - clip_text_encode_service = providers.Factory( - ClipTextEncodeService, + clip_text_encode_usecase = providers.Factory( + ClipTextEncodeUsecase, text_encoder_repo=text_encoder_repository, ) - sampling_service = providers.Factory( - SamplingService, - sampling_repo=sampling_repository, + sampler_usecase = providers.Factory( + SamplerUsecase, + sampler_repo=sampler_repository, ) - scheduler_service = providers.Factory( - SchedulerService, + scheduler_usecase = providers.Factory( + SchedulerUsecase, scheduler_repo=scheduler_repository, ) - - # 3. Handler - pipeline_handler = providers.Factory( - PipelineHandler, - pipeline_service=pipeline_service, - ) - autoencoder_handler = providers.Factory( - AutoencoderHandler, - autoencoder_service=autoencoder_service, - ) - clip_text_encode_handler = providers.Factory( - ClipTextEncodeHandler, - clip_text_encode_service=clip_text_encode_service, - ) - sampler_handler = providers.Factory( - SamplerHandler, - sampling_service=sampling_service, - ) - scheduler_handler = providers.Factory( - SchedulerHandler, - scheduler_service=scheduler_service, - ) diff --git a/src/domain/repositories/__init__.py b/src/domain/repositories/__init__.py index e0268b9..0ce3074 100644 --- a/src/domain/repositories/__init__.py +++ b/src/domain/repositories/__init__.py @@ -1,6 +1,6 @@ from ._autoencoder_repository import AutoencoderRepository from ._pipeline_repository import PipelineRepository -from ._sampling_repository import SamplingRepository +from ._sampler_repository import SamplerRepository from ._scheduler_repository import SchedulerRepository from ._text_encoder_repository import TextEncoderRepository @@ -8,6 +8,6 @@ __all__ = [ "PipelineRepository", "AutoencoderRepository", "TextEncoderRepository", - "SamplingRepository", + "SamplerRepository", "SchedulerRepository", ] diff --git a/src/domain/repositories/_sampling_repository.py b/src/domain/repositories/_sampler_repository.py similarity index 94% rename from src/domain/repositories/_sampling_repository.py rename to src/domain/repositories/_sampler_repository.py index d9c1204..14d627c 100644 --- a/src/domain/repositories/_sampling_repository.py +++ b/src/domain/repositories/_sampler_repository.py @@ -6,7 +6,7 @@ from diffusers.schedulers.scheduling_utils import SchedulerMixin from PIL import Image -class SamplingRepository(ABC): +class SamplerRepository(ABC): @abstractmethod def sample( self, diff --git a/src/infra/diffusers/__init__.py b/src/infra/diffusers/__init__.py index 3f20a8d..3617bf9 100644 --- a/src/infra/diffusers/__init__.py +++ b/src/infra/diffusers/__init__.py @@ -1,6 +1,6 @@ from ._autoencoder_repository import DiffusersAutoencoderRepository from ._pipeline_repository import DiffusersPipelineRepository -from ._sampling_repository import DiffusersSamplingRepository +from ._sampler_repository import DiffusersSamplerRepository from ._scheduler_repository import DiffusersSchedulerRepository from ._text_encoder_repository import DiffusersTextEncoderRepository @@ -8,6 +8,6 @@ __all__ = [ "DiffusersPipelineRepository", "DiffusersAutoencoderRepository", "DiffusersTextEncoderRepository", - "DiffusersSamplingRepository", + "DiffusersSamplerRepository", "DiffusersSchedulerRepository", ] diff --git a/src/infra/diffusers/_sampling_repository.py b/src/infra/diffusers/_sampler_repository.py similarity index 91% rename from src/infra/diffusers/_sampling_repository.py rename to src/infra/diffusers/_sampler_repository.py index 96dcc09..ef23525 100644 --- a/src/infra/diffusers/_sampling_repository.py +++ b/src/infra/diffusers/_sampler_repository.py @@ -4,10 +4,10 @@ from diffusers import AutoencoderKL, StableDiffusionPipeline from diffusers.schedulers.scheduling_utils import SchedulerMixin from PIL import Image -from ...domain.repositories import SamplingRepository +from ...domain.repositories import SamplerRepository -class DiffusersSamplingRepository(SamplingRepository): +class DiffusersSamplerRepository(SamplerRepository): def __init__(self) -> None: self.device = get_torch_device() diff --git a/src/nodes/_diffusers_clip_text_encode.py b/src/nodes/_diffusers_clip_text_encode.py index 621b3ed..ffb0756 100644 --- a/src/nodes/_diffusers_clip_text_encode.py +++ b/src/nodes/_diffusers_clip_text_encode.py @@ -1,7 +1,7 @@ from dependency_injector.wiring import Provide, inject from ..di import Container -from ..ui import ClipTextEncodeHandler +from ..usecase import ClipTextEncodeUsecase from .dto import ComfyUIClipDTO, ComfyUIConditioningDTO @@ -27,10 +27,10 @@ class DiffusersClipTextEncode: self, clip: ComfyUIClipDTO, text: str, - handler: ClipTextEncodeHandler = Provide[Container.clip_text_encode_handler], + usecase: ClipTextEncodeUsecase = Provide[Container.clip_text_encode_usecase], ) -> tuple[ComfyUIConditioningDTO]: clip_domain = ComfyUIClipDTO.to_domain(clip) - conditioning_model = handler.encode( + conditioning_model = usecase.execute( clip_domain.tokenizer, clip_domain.text_encoder, clip_domain.path, text ) conditioning_dto = ComfyUIConditioningDTO.from_domain(conditioning_model) diff --git a/src/nodes/_diffusers_pipeline_loader.py b/src/nodes/_diffusers_pipeline_loader.py index 4bfba40..db90319 100644 --- a/src/nodes/_diffusers_pipeline_loader.py +++ b/src/nodes/_diffusers_pipeline_loader.py @@ -2,7 +2,7 @@ import folder_paths # pyright: ignore[reportMissingImports] from dependency_injector.wiring import Provide, inject from ..di import Container -from ..ui import PipelineHandler +from ..usecase import PipelineUsecase from .dto import ComfyUIClipDTO, ComfyUIPipelineDTO @@ -29,10 +29,10 @@ class DiffusersPipelineLoader: def execute( self, checkpoint_name: str, - handler: PipelineHandler = Provide[Container.pipeline_handler], + usecase: PipelineUsecase = Provide[Container.pipeline_usecase], ) -> tuple[ComfyUIPipelineDTO, ComfyUIClipDTO]: checkpoint_path = folder_paths.get_full_path("checkpoints", checkpoint_name) - pipeline_model, clip_model = handler.create(checkpoint_path) + pipeline_model, clip_model = usecase.execute(checkpoint_path) pipeline_dto = ComfyUIPipelineDTO.from_domain(pipeline_model) clip_dto = ComfyUIClipDTO.from_domain(clip_model) return ( diff --git a/src/nodes/_diffusers_sampler.py b/src/nodes/_diffusers_sampler.py index 7302e89..5c87135 100644 --- a/src/nodes/_diffusers_sampler.py +++ b/src/nodes/_diffusers_sampler.py @@ -1,7 +1,7 @@ from dependency_injector.wiring import Provide, inject from ..di import Container -from ..ui import SamplerHandler +from ..usecase import SamplerUsecase from .dto import ( ComfyUIAutoencoderDTO, ComfyUIConditioningDTO, @@ -55,14 +55,14 @@ class DiffusersSampler: steps: int, cfg: float, seed: int, - handler: SamplerHandler = Provide[Container.sampler_handler], + usecase: SamplerUsecase = Provide[Container.sampler_usecase], ) -> 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( + images_model = usecase.execute( pipeline_domain, vae_domain, scheduler_domain, diff --git a/src/nodes/_diffusers_scheduler_loader.py b/src/nodes/_diffusers_scheduler_loader.py index d518d78..2746c7d 100644 --- a/src/nodes/_diffusers_scheduler_loader.py +++ b/src/nodes/_diffusers_scheduler_loader.py @@ -1,7 +1,7 @@ from dependency_injector.wiring import Provide, inject from ..di import Container -from ..ui import SchedulerHandler +from ..usecase import SchedulerUsecase from .dto import ComfyUIPipelineDTO, ComfyUISchedulerDTO from .type import ComfyUISchedulerType @@ -28,10 +28,10 @@ class DiffusersSchedulerLoader: self, pipeline: ComfyUIPipelineDTO, scheduler_name: str, - handler: SchedulerHandler = Provide[Container.scheduler_handler], + usecase: SchedulerUsecase = Provide[Container.scheduler_usecase], ) -> tuple[ComfyUISchedulerDTO]: domain = ComfyUIPipelineDTO.to_domain(pipeline) scheduler_type = ComfyUISchedulerType.to_domain(scheduler_name) - scheduler_domain = handler.create(domain, scheduler_type) + scheduler_domain = usecase.execute(domain, scheduler_type) scheduler_dto = ComfyUISchedulerDTO.from_domain(scheduler_domain) return (scheduler_dto,) diff --git a/src/nodes/_diffusers_vae_loader.py b/src/nodes/_diffusers_vae_loader.py index 019917c..8aa18e0 100644 --- a/src/nodes/_diffusers_vae_loader.py +++ b/src/nodes/_diffusers_vae_loader.py @@ -2,7 +2,7 @@ import folder_paths # pyright: ignore[reportMissingImports] from dependency_injector.wiring import Provide, inject from ..di import Container -from ..ui import AutoencoderHandler +from ..usecase import AutoencoderUsecase from .dto import ComfyUIAutoencoderDTO @@ -26,9 +26,9 @@ class DiffusersVaeLoader: def execute( self, vae_name: str, - handler: AutoencoderHandler = Provide[Container.autoencoder_handler], + usecase: AutoencoderUsecase = Provide[Container.autoencoder_usecase], ) -> tuple[ComfyUIAutoencoderDTO]: vae_path = folder_paths.get_full_path("vae", vae_name) - autoencoder_model = handler.create(vae_path) + autoencoder_model = usecase.execute(vae_path) autoencoder_dto = ComfyUIAutoencoderDTO.from_domain(autoencoder_model) return (autoencoder_dto,) diff --git a/src/nodes/dto/_scheduler_dto.py b/src/nodes/dto/_scheduler_dto.py index aaec1b4..8d8cfd6 100644 --- a/src/nodes/dto/_scheduler_dto.py +++ b/src/nodes/dto/_scheduler_dto.py @@ -30,7 +30,7 @@ class ComfyUISchedulerDTO: return {"required": {"scheduler": (ComfyUISchedulerDTO.COMFY_TYPE,)}} def execute(self, scheduler: ComfyUISchedulerDTO): - # Use the scheduler for sampling operations + # Use the scheduler for sampler operations pass ``` """ diff --git a/src/service/__init__.py b/src/service/__init__.py deleted file mode 100644 index 6d4bff8..0000000 --- a/src/service/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -from ._autoencoder_service import AutoencoderService -from ._clip_text_encode_service import ClipTextEncodeService -from ._pipeline_service import PipelineService -from ._sampling_service import SamplingService -from ._scheduler_service import SchedulerService - -__all__ = [ - "PipelineService", - "AutoencoderService", - "ClipTextEncodeService", - "SamplingService", - "SchedulerService", -] diff --git a/src/service/_sampling_service.py b/src/service/_sampling_service.py deleted file mode 100644 index 961e3d4..0000000 --- a/src/service/_sampling_service.py +++ /dev/null @@ -1,40 +0,0 @@ -import torch - -from ..domain.model import Autoencoder, Conditioning, Image, Pipeline, Scheduler -from ..domain.repositories import SamplingRepository - - -class SamplingService: - def __init__(self, sampling_repo: SamplingRepository) -> None: - self.sampling_repo = sampling_repo - self.dtype = torch.float32 - - def sample( - self, - pipeline: Pipeline, - vae: Autoencoder, - scheduler: Scheduler, - positive_embeds: Conditioning, - negative_embeds: Conditioning, - width: int, - height: int, - steps: int, - cfg: float, - seed: int, - ) -> list[Image]: - images = self.sampling_repo.sample( - 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.") - images_domain: list[Image] = [Image(image=img) for img in images] - return images_domain diff --git a/src/ui/__init__.py b/src/ui/__init__.py deleted file mode 100644 index af0a7c9..0000000 --- a/src/ui/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -from ._autoencoder_handler import AutoencoderHandler -from ._clip_text_encode_handler import ClipTextEncodeHandler -from ._pipeline_handler import PipelineHandler -from ._sampler_handler import SamplerHandler -from ._scheduler_handler import SchedulerHandler - -__all__ = [ - "PipelineHandler", - "AutoencoderHandler", - "ClipTextEncodeHandler", - "SamplerHandler", - "SchedulerHandler", -] diff --git a/src/ui/_autoencoder_handler.py b/src/ui/_autoencoder_handler.py deleted file mode 100644 index bbc09a8..0000000 --- a/src/ui/_autoencoder_handler.py +++ /dev/null @@ -1,17 +0,0 @@ -import os - -from ..domain.model import Autoencoder -from ..service import AutoencoderService - - -class AutoencoderHandler: - def __init__(self, autoencoder_service: AutoencoderService) -> None: - self.create_autoencoder_service = autoencoder_service - - def create(self, vae_path: str) -> Autoencoder: - if not os.path.exists(vae_path): - raise FileNotFoundError(f"VAE file not found: {vae_path}") - autoencoder = self.create_autoencoder_service.create_from_checkpoint(vae_path) - if autoencoder is None: - raise RuntimeError(f"Failed to create autoencoder from checkpoint: {vae_path}") - return autoencoder diff --git a/src/ui/_clip_text_encode_handler.py b/src/ui/_clip_text_encode_handler.py deleted file mode 100644 index bd83181..0000000 --- a/src/ui/_clip_text_encode_handler.py +++ /dev/null @@ -1,27 +0,0 @@ -from __future__ import annotations - -from transformers import CLIPTextModel, CLIPTokenizer - -from ..domain.model import Conditioning -from ..service._clip_text_encode_service import ClipTextEncodeService - - -class ClipTextEncodeHandler: - def __init__(self, clip_text_encode_service: ClipTextEncodeService) -> None: - self.clip_text_encode_service = clip_text_encode_service - - def encode( - self, tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel, path: str, text: str - ) -> Conditioning: - if not text: - raise ValueError("Text input is empty.") - if path is None: - raise ValueError("Clip path is None.") - if tokenizer is None or text_encoder is None: - raise ValueError("Tokenizer or Text Encoder is None.") - if not isinstance(text, str): - raise TypeError("Text input must be a string.") - conditioning = self.clip_text_encode_service.encode(tokenizer, text_encoder, path, text) - if conditioning is None: - raise RuntimeError("Failed to encode text into conditioning.") - return conditioning diff --git a/src/ui/_pipeline_handler.py b/src/ui/_pipeline_handler.py deleted file mode 100644 index 1e123df..0000000 --- a/src/ui/_pipeline_handler.py +++ /dev/null @@ -1,19 +0,0 @@ -import os - -from ..domain.model import Clip, Pipeline -from ..service import PipelineService - - -class PipelineHandler: - def __init__(self, pipeline_service: PipelineService) -> None: - self.create_pipeline_service = pipeline_service - - def create(self, checkpoint_path: str) -> tuple[Pipeline, Clip]: - if not os.path.exists(checkpoint_path): - raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}") - pipeline, clip = self.create_pipeline_service.create_from_checkpoint(checkpoint_path) - if pipeline is None: - raise RuntimeError(f"Failed to create pipeline from checkpoint: {checkpoint_path}") - if clip is None: - raise RuntimeError(f"Failed to extract clip from checkpoint: {checkpoint_path}") - return pipeline, clip diff --git a/src/ui/_scheduler_handler.py b/src/ui/_scheduler_handler.py deleted file mode 100644 index aade076..0000000 --- a/src/ui/_scheduler_handler.py +++ /dev/null @@ -1,12 +0,0 @@ -from ..domain.model import Pipeline, Scheduler -from ..service import SchedulerService - - -class SchedulerHandler: - def __init__(self, scheduler_service: SchedulerService) -> None: - self.scheduler_service = scheduler_service - - def create(self, pipeline: Pipeline, scheduler_type: Scheduler.Type) -> Scheduler: - if pipeline is None or pipeline.pipeline is None: - raise ValueError("Pipeline is None.") - return self.scheduler_service.create_scheduler(pipeline, scheduler_type) diff --git a/src/usecase/__init__.py b/src/usecase/__init__.py new file mode 100644 index 0000000..63ae7cc --- /dev/null +++ b/src/usecase/__init__.py @@ -0,0 +1,13 @@ +from ._autoencoder_usecase import AutoencoderUsecase +from ._clip_text_encode_usecase import ClipTextEncodeUsecase +from ._pipeline_usecase import PipelineUsecase +from ._sampler_usecase import SamplerUsecase +from ._scheduler_usecase import SchedulerUsecase + +__all__ = [ + "PipelineUsecase", + "AutoencoderUsecase", + "ClipTextEncodeUsecase", + "SamplerUsecase", + "SchedulerUsecase", +] diff --git a/src/service/_autoencoder_service.py b/src/usecase/_autoencoder_usecase.py similarity index 64% rename from src/service/_autoencoder_service.py rename to src/usecase/_autoencoder_usecase.py index 301f4e6..a0b3cbe 100644 --- a/src/service/_autoencoder_service.py +++ b/src/usecase/_autoencoder_usecase.py @@ -6,22 +6,21 @@ from ..domain.model import Autoencoder from ..domain.repositories import AutoencoderRepository -class AutoencoderService: +class AutoencoderUsecase: def __init__(self, autoencoder_repo: AutoencoderRepository) -> None: self.autoencoder_repo = autoencoder_repo self.dtype = torch.float32 - def create_from_checkpoint(self, checkpoint_full_path: str) -> Autoencoder: - path = self.autoencoder_repo.convert_and_save_from_single_file( - checkpoint_full_path, self.dtype - ) + def execute(self, vae_path: str) -> Autoencoder: + if not os.path.exists(vae_path): + raise FileNotFoundError(f"VAE file not found: {vae_path}") + + path = self.autoencoder_repo.convert_and_save_from_single_file(vae_path, self.dtype) if not os.path.exists(path): - raise RuntimeError(f"Failed to convert VAE checkpoint: {checkpoint_full_path}") + raise RuntimeError(f"Failed to convert VAE checkpoint: {vae_path}") vae = self.autoencoder_repo.load_autoencoder_from_path(path, self.dtype) if vae is None: - raise RuntimeError( - f"Failed to create autoencoder from checkpoint: {checkpoint_full_path}" - ) + raise RuntimeError(f"Failed to create autoencoder from checkpoint: {vae_path}") return Autoencoder(autoencoder=vae, path=path) diff --git a/src/service/_clip_text_encode_service.py b/src/usecase/_clip_text_encode_usecase.py similarity index 62% rename from src/service/_clip_text_encode_service.py rename to src/usecase/_clip_text_encode_usecase.py index 71bd446..715d3ec 100644 --- a/src/service/_clip_text_encode_service.py +++ b/src/usecase/_clip_text_encode_usecase.py @@ -6,13 +6,21 @@ from ..domain.model import Conditioning from ..domain.repositories._text_encoder_repository import TextEncoderRepository -class ClipTextEncodeService: +class ClipTextEncodeUsecase: def __init__(self, text_encoder_repo: TextEncoderRepository) -> None: self.text_encoder_repo = text_encoder_repo - def encode( + def execute( self, tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel, path: str, text: str ) -> Conditioning: + if not text: + raise ValueError("Text input is empty.") + if path is None: + raise ValueError("Clip path is None.") + if tokenizer is None or text_encoder is None: + raise ValueError("Tokenizer or Text Encoder is None.") + if not isinstance(text, str): + raise TypeError("Text input must be a string.") embeddings = self.text_encoder_repo.encode(tokenizer, text_encoder, text) if embeddings is None: raise RuntimeError("Failed to generate embeddings.") diff --git a/src/service/_pipeline_service.py b/src/usecase/_pipeline_usecase.py similarity index 70% rename from src/service/_pipeline_service.py rename to src/usecase/_pipeline_usecase.py index 1ca6c27..92063a4 100644 --- a/src/service/_pipeline_service.py +++ b/src/usecase/_pipeline_usecase.py @@ -6,28 +6,27 @@ from ..domain.model import Clip, Pipeline from ..domain.repositories import PipelineRepository -class PipelineService: +class PipelineUsecase: def __init__(self, pipeline_repo: PipelineRepository) -> None: self.pipeline_repo = pipeline_repo self.dtype = torch.float32 - def create_from_checkpoint(self, checkpoint_full_path: str) -> tuple[Pipeline, Clip]: - path = self.pipeline_repo.convert_and_save_from_single_file( - checkpoint_full_path, self.dtype - ) + def execute(self, checkpoint_path: str) -> tuple[Pipeline, Clip]: + if not os.path.exists(checkpoint_path): + raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}") + + path = self.pipeline_repo.convert_and_save_from_single_file(checkpoint_path, self.dtype) if not os.path.exists(path): - raise RuntimeError(f"Failed to convert pipeline checkpoint: {checkpoint_full_path}") + raise RuntimeError(f"Failed to convert pipeline checkpoint: {checkpoint_path}") pipe = self.pipeline_repo.load_pipeline_from_path(path, self.dtype) if pipe is None: - raise RuntimeError(f"Failed to create pipeline from checkpoint: {checkpoint_full_path}") + raise RuntimeError(f"Failed to create pipeline from checkpoint: {checkpoint_path}") pipeline = Pipeline(pipeline=pipe, path=path) - clip = Clip( tokenizer=pipe.tokenizer, # type: ignore[attr-defined] text_encoder=pipe.text_encoder, # type: ignore[attr-defined] - path=checkpoint_full_path, + path=checkpoint_path, ) - return pipeline, clip diff --git a/src/ui/_sampler_handler.py b/src/usecase/_sampler_usecase.py similarity index 64% rename from src/ui/_sampler_handler.py rename to src/usecase/_sampler_usecase.py index fe9247a..b576db0 100644 --- a/src/ui/_sampler_handler.py +++ b/src/usecase/_sampler_usecase.py @@ -1,12 +1,15 @@ +import torch + from ..domain.model import Autoencoder, Conditioning, Image, Pipeline, Scheduler -from ..service import SamplingService +from ..domain.repositories import SamplerRepository -class SamplerHandler: - def __init__(self, sampling_service: SamplingService) -> None: - self.sampling_service = sampling_service +class SamplerUsecase: + def __init__(self, sampler_repo: SamplerRepository) -> None: + self.sampler_repo = sampler_repo + self.dtype = torch.float32 - def sample( + def execute( self, pipeline: Pipeline, vae: Autoencoder, @@ -35,18 +38,20 @@ class SamplerHandler: raise ValueError("CFG must be a non-negative float.") if seed < 0: raise ValueError("Seed must be a non-negative integer.") - domain_images = self.sampling_service.sample( - pipeline, - vae, - scheduler, - positive_embeds, - negative_embeds, + + images = self.sampler_repo.sample( + pipeline.pipeline, + vae.autoencoder, + scheduler.scheduler, + positive_embeds.conditioning, + negative_embeds.conditioning, width, height, steps, cfg, seed, ) - if domain_images is None: - raise RuntimeError("Failed to sample images from pipeline.") - return domain_images + 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 diff --git a/src/service/_scheduler_service.py b/src/usecase/_scheduler_usecase.py similarity index 71% rename from src/service/_scheduler_service.py rename to src/usecase/_scheduler_usecase.py index 2ff4da8..8f6339b 100644 --- a/src/service/_scheduler_service.py +++ b/src/usecase/_scheduler_usecase.py @@ -5,13 +5,15 @@ from ..domain.model import Pipeline, Scheduler from ..domain.repositories import SchedulerRepository -class SchedulerService: +class SchedulerUsecase: def __init__(self, scheduler_repo: SchedulerRepository) -> None: self.scheduler_repo = scheduler_repo self.dtype = torch.float32 - def create_scheduler(self, pipeline: Pipeline, scheduler_type: Scheduler.Type) -> Scheduler: - if pipeline is None or pipeline.pipeline is None or not pipeline.path: + 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