From 06c9f21ffb5e0a2f56fa0982dcf0f2a1b4545349 Mon Sep 17 00:00:00 2001 From: Limitex <76650151+Limitex@users.noreply.github.com> Date: Sat, 22 Nov 2025 08:23:16 +0900 Subject: [PATCH] feat: stream diffusion impl --- __init__.py | 6 + src/di/di.py | 29 +++ src/domain/model/__init__.py | 17 ++ src/domain/model/_cfg_type.py | 22 +++ src/domain/model/_delta.py | 14 ++ src/domain/model/_frame_buffer_size.py | 14 ++ src/domain/model/_lcm_lora.py | 15 ++ src/domain/model/_num_samples.py | 14 ++ src/domain/model/_stream_diffusion_stream.py | 15 ++ src/domain/model/_t_index_list.py | 19 ++ src/domain/model/_warmup_count.py | 14 ++ src/domain/repositories/__init__.py | 4 + .../repositories/_lcm_lora_repository.py | 15 ++ .../_stream_diffusion_repository.py | 132 ++++++++++++++ src/infra/diffusers/__init__.py | 4 + src/infra/diffusers/_lcm_lora_repository.py | 18 ++ .../diffusers/_stream_diffusion_repository.py | 172 ++++++++++++++++++ src/nodes/__init__.py | 24 +++ src/nodes/_create_int_list_node.py | 36 ++++ src/nodes/_lcm_lora_loader.py | 45 +++++ src/nodes/_stream_diffusion_create_stream.py | 102 +++++++++++ src/nodes/_stream_diffusion_fast_sampler.py | 66 +++++++ src/nodes/_stream_diffusion_sampler.py | 114 ++++++++++++ src/nodes/_stream_diffusion_warmup.py | 85 +++++++++ src/nodes/dto/__init__.py | 5 + src/nodes/dto/_lcm_lora_dto.py | 44 +++++ src/nodes/dto/_stream_diffusion_stream_dto.py | 86 +++++++++ src/usecase/__init__.py | 10 + src/usecase/_lcm_lora_usecase.py | 26 +++ ..._stream_diffusion_create_stream_usecase.py | 66 +++++++ .../_stream_diffusion_fast_sample_usecase.py | 40 ++++ .../_stream_diffusion_sample_usecase.py | 72 ++++++++ .../_stream_diffusion_warmup_usecase.py | 43 +++++ 33 files changed, 1388 insertions(+) create mode 100644 src/domain/model/_cfg_type.py create mode 100644 src/domain/model/_delta.py create mode 100644 src/domain/model/_frame_buffer_size.py create mode 100644 src/domain/model/_lcm_lora.py create mode 100644 src/domain/model/_num_samples.py create mode 100644 src/domain/model/_stream_diffusion_stream.py create mode 100644 src/domain/model/_t_index_list.py create mode 100644 src/domain/model/_warmup_count.py create mode 100644 src/domain/repositories/_lcm_lora_repository.py create mode 100644 src/domain/repositories/_stream_diffusion_repository.py create mode 100644 src/infra/diffusers/_lcm_lora_repository.py create mode 100644 src/infra/diffusers/_stream_diffusion_repository.py create mode 100644 src/nodes/_create_int_list_node.py create mode 100644 src/nodes/_lcm_lora_loader.py create mode 100644 src/nodes/_stream_diffusion_create_stream.py create mode 100644 src/nodes/_stream_diffusion_fast_sampler.py create mode 100644 src/nodes/_stream_diffusion_sampler.py create mode 100644 src/nodes/_stream_diffusion_warmup.py create mode 100644 src/nodes/dto/_lcm_lora_dto.py create mode 100644 src/nodes/dto/_stream_diffusion_stream_dto.py create mode 100644 src/usecase/_lcm_lora_usecase.py create mode 100644 src/usecase/_stream_diffusion_create_stream_usecase.py create mode 100644 src/usecase/_stream_diffusion_fast_sample_usecase.py create mode 100644 src/usecase/_stream_diffusion_sample_usecase.py create mode 100644 src/usecase/_stream_diffusion_warmup_usecase.py diff --git a/__init__.py b/__init__.py index 6f4e249..0463501 100644 --- a/__init__.py +++ b/__init__.py @@ -5,8 +5,11 @@ from .src.nodes import ( ComfyUIAutoencoderDTO, ComfyUIClipDTO, ComfyUIConditioningDTO, + ComfyUILcmLoraDTO, ComfyUIPipelineDTO, ComfyUISchedulerDTO, + ComfyUIStreamDTO, + ComfyUIWarmupStreamDTO, ) load_envs() @@ -19,4 +22,7 @@ __all__ = [ "ComfyUIClipDTO", "ComfyUIConditioningDTO", "ComfyUISchedulerDTO", + "ComfyUILcmLoraDTO", + "ComfyUIStreamDTO", + "ComfyUIWarmupStreamDTO", ] diff --git a/src/di/di.py b/src/di/di.py index 3185a49..4e51912 100644 --- a/src/di/di.py +++ b/src/di/di.py @@ -4,17 +4,24 @@ from dependency_injector import containers, providers from ..infra.diffusers import ( DiffusersAutoencoderRepository, + DiffusersLcmLoraRepository, DiffusersPipelineRepository, DiffusersSamplerRepository, DiffusersSchedulerRepository, + DiffusersStreamDiffusionRepository, DiffusersTextEncoderRepository, ) from ..usecase import ( AutoencoderUsecase, ClipTextEncodeUsecase, + LcmLoraUsecase, PipelineUsecase, SamplerUsecase, SchedulerUsecase, + StreamDiffusionCreateStreamUsecase, + StreamDiffusionFastSampleUsecase, + StreamDiffusionSampleUsecase, + StreamDiffusionWarmupUsecase, ) @@ -27,6 +34,8 @@ class Container(containers.DeclarativeContainer): text_encoder_repository = providers.Factory(DiffusersTextEncoderRepository) sampler_repository = providers.Factory(DiffusersSamplerRepository) scheduler_repository = providers.Factory(DiffusersSchedulerRepository) + lcm_lora_repository = providers.Factory(DiffusersLcmLoraRepository) + stream_diffusion_repository = providers.Factory(DiffusersStreamDiffusionRepository) # 2. Usecases pipeline_usecase = providers.Factory( @@ -49,3 +58,23 @@ class Container(containers.DeclarativeContainer): SchedulerUsecase, scheduler_repo=scheduler_repository, ) + lcm_lora_usecase = providers.Factory( + LcmLoraUsecase, + lcm_lora_repo=lcm_lora_repository, + ) + stream_diffusion_create_stream_usecase = providers.Factory( + StreamDiffusionCreateStreamUsecase, + stream_diffusion_repo=stream_diffusion_repository, + ) + stream_diffusion_warmup_usecase = providers.Factory( + StreamDiffusionWarmupUsecase, + stream_diffusion_repo=stream_diffusion_repository, + ) + stream_diffusion_sample_usecase = providers.Factory( + StreamDiffusionSampleUsecase, + stream_diffusion_repo=stream_diffusion_repository, + ) + stream_diffusion_fast_sample_usecase = providers.Factory( + StreamDiffusionFastSampleUsecase, + stream_diffusion_repo=stream_diffusion_repository, + ) diff --git a/src/domain/model/__init__.py b/src/domain/model/__init__.py index f3b54e9..629f9c2 100644 --- a/src/domain/model/__init__.py +++ b/src/domain/model/__init__.py @@ -1,13 +1,21 @@ from ._autoencoder import Autoencoder from ._cfg_scale import CFGScale +from ._cfg_type import CFGType, CFGTypeEnum from ._clip import Clip from ._conditioning import Conditioning +from ._delta import Delta +from ._frame_buffer_size import FrameBufferSize from ._image import Image from ._image_size import ImageSize +from ._lcm_lora import LcmLora +from ._num_samples import NumSamples from ._pipeline import Pipeline from ._scheduler import Scheduler from ._seed import Seed from ._steps import Steps +from ._stream_diffusion_stream import StreamDiffusionStream +from ._t_index_list import TIndexList +from ._warmup_count import WarmupCount __all__ = [ "Pipeline", @@ -20,4 +28,13 @@ __all__ = [ "Steps", "CFGScale", "Seed", + "LcmLora", + "TIndexList", + "FrameBufferSize", + "CFGType", + "CFGTypeEnum", + "Delta", + "WarmupCount", + "NumSamples", + "StreamDiffusionStream", ] diff --git a/src/domain/model/_cfg_type.py b/src/domain/model/_cfg_type.py new file mode 100644 index 0000000..8a6c896 --- /dev/null +++ b/src/domain/model/_cfg_type.py @@ -0,0 +1,22 @@ +from dataclasses import dataclass +from enum import Enum + + +class CFGTypeEnum(str, Enum): + """CFG type options for Stream Diffusion.""" + + NONE = "none" + FULL = "full" + SELF = "self" + INITIALIZE = "initialize" + + +@dataclass(frozen=True) +class CFGType: + """CFG type configuration for Stream Diffusion.""" + + value: CFGTypeEnum + + def __post_init__(self) -> None: + if not isinstance(self.value, CFGTypeEnum): + raise ValueError(f"CFG type must be a CFGTypeEnum, got {type(self.value).__name__}") diff --git a/src/domain/model/_delta.py b/src/domain/model/_delta.py new file mode 100644 index 0000000..00f917e --- /dev/null +++ b/src/domain/model/_delta.py @@ -0,0 +1,14 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class Delta: + """Delta value for Stream Diffusion.""" + + value: float + + def __post_init__(self) -> None: + if isinstance(self.value, bool) or not isinstance(self.value, float): + raise ValueError(f"Delta must be a float, got {type(self.value).__name__}") + if not (0.0 <= self.value <= 1.0): + raise ValueError(f"Delta must be between 0.0 and 1.0, got {self.value}") diff --git a/src/domain/model/_frame_buffer_size.py b/src/domain/model/_frame_buffer_size.py new file mode 100644 index 0000000..9fb9c7b --- /dev/null +++ b/src/domain/model/_frame_buffer_size.py @@ -0,0 +1,14 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class FrameBufferSize: + """Frame buffer size for Stream Diffusion.""" + + value: int + + def __post_init__(self) -> None: + if isinstance(self.value, bool) or not isinstance(self.value, int): + raise ValueError(f"Frame buffer size must be an int, got {type(self.value).__name__}") + if self.value <= 0: + raise ValueError(f"Frame buffer size must be positive, got {self.value}") diff --git a/src/domain/model/_lcm_lora.py b/src/domain/model/_lcm_lora.py new file mode 100644 index 0000000..af42755 --- /dev/null +++ b/src/domain/model/_lcm_lora.py @@ -0,0 +1,15 @@ +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class LcmLora: + """LCM LoRA weights for Stream Diffusion.""" + + weights: dict[str, Any] + + def __post_init__(self) -> None: + if not isinstance(self.weights, dict): + raise ValueError(f"LCM LoRA weights must be a dict, got {type(self.weights).__name__}") + if not self.weights: + raise ValueError("LCM LoRA weights cannot be empty") diff --git a/src/domain/model/_num_samples.py b/src/domain/model/_num_samples.py new file mode 100644 index 0000000..bfa004d --- /dev/null +++ b/src/domain/model/_num_samples.py @@ -0,0 +1,14 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class NumSamples: + """Number of samples to generate.""" + + value: int + + def __post_init__(self) -> None: + if isinstance(self.value, bool) or not isinstance(self.value, int): + raise ValueError(f"Number of samples must be an int, got {type(self.value).__name__}") + if self.value <= 0: + raise ValueError(f"Number of samples must be positive, got {self.value}") diff --git a/src/domain/model/_stream_diffusion_stream.py b/src/domain/model/_stream_diffusion_stream.py new file mode 100644 index 0000000..2db40c1 --- /dev/null +++ b/src/domain/model/_stream_diffusion_stream.py @@ -0,0 +1,15 @@ +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class StreamDiffusionStream: + """Stream Diffusion stream instance.""" + + stream: Any # StreamDiffusion instance from streamdiffusion library + + def __post_init__(self) -> None: + if self.stream is None: + raise ValueError("Stream cannot be None") + # We can't check for exact type here as it would create circular dependency + # The actual type checking will be done at runtime in the repository layer diff --git a/src/domain/model/_t_index_list.py b/src/domain/model/_t_index_list.py new file mode 100644 index 0000000..20858ea --- /dev/null +++ b/src/domain/model/_t_index_list.py @@ -0,0 +1,19 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class TIndexList: + """List of timestep indices for Stream Diffusion.""" + + indices: list[int] + + def __post_init__(self) -> None: + if not isinstance(self.indices, list): + raise ValueError(f"Indices must be a list, got {type(self.indices).__name__}") + if not self.indices: + raise ValueError("Indices list cannot be empty") + for idx, value in enumerate(self.indices): + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError( + f"Index at position {idx} must be an int, got {type(value).__name__}" + ) diff --git a/src/domain/model/_warmup_count.py b/src/domain/model/_warmup_count.py new file mode 100644 index 0000000..2bc2f25 --- /dev/null +++ b/src/domain/model/_warmup_count.py @@ -0,0 +1,14 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class WarmupCount: + """Number of warmup iterations for Stream Diffusion.""" + + value: int + + def __post_init__(self) -> None: + if isinstance(self.value, bool) or not isinstance(self.value, int): + raise ValueError(f"Warmup count must be an int, got {type(self.value).__name__}") + if self.value < 0: + raise ValueError(f"Warmup count must be non-negative, got {self.value}") diff --git a/src/domain/repositories/__init__.py b/src/domain/repositories/__init__.py index 0ce3074..de2e025 100644 --- a/src/domain/repositories/__init__.py +++ b/src/domain/repositories/__init__.py @@ -1,7 +1,9 @@ from ._autoencoder_repository import AutoencoderRepository +from ._lcm_lora_repository import LcmLoraRepository from ._pipeline_repository import PipelineRepository from ._sampler_repository import SamplerRepository from ._scheduler_repository import SchedulerRepository +from ._stream_diffusion_repository import StreamDiffusionRepository from ._text_encoder_repository import TextEncoderRepository __all__ = [ @@ -10,4 +12,6 @@ __all__ = [ "TextEncoderRepository", "SamplerRepository", "SchedulerRepository", + "LcmLoraRepository", + "StreamDiffusionRepository", ] diff --git a/src/domain/repositories/_lcm_lora_repository.py b/src/domain/repositories/_lcm_lora_repository.py new file mode 100644 index 0000000..51739ef --- /dev/null +++ b/src/domain/repositories/_lcm_lora_repository.py @@ -0,0 +1,15 @@ +from abc import ABC, abstractmethod +from typing import Any + + +class LcmLoraRepository(ABC): + @abstractmethod + def load_lcm_lora(self, lora_path: str) -> dict[str, Any]: + """Load LCM LoRA weights from file. + + Args: + lora_path: Path to the LCM LoRA file + + Returns: + Dictionary containing the loaded weights + """ diff --git a/src/domain/repositories/_stream_diffusion_repository.py b/src/domain/repositories/_stream_diffusion_repository.py new file mode 100644 index 0000000..87508eb --- /dev/null +++ b/src/domain/repositories/_stream_diffusion_repository.py @@ -0,0 +1,132 @@ +from abc import ABC, abstractmethod +from typing import Any + +from diffusers import StableDiffusionPipeline +from PIL import Image + +from ..model import ( + CFGScale, + CFGType, + Delta, + FrameBufferSize, + ImageSize, + NumSamples, + Seed, + Steps, + TIndexList, + WarmupCount, +) + + +class StreamDiffusionRepository(ABC): + @abstractmethod + def create_stream( + self, + pipeline: StableDiffusionPipeline, + t_index_list: TIndexList, + image_size: ImageSize, + do_add_noise: bool, + use_denoising_batch: bool, + frame_buffer_size: FrameBufferSize, + cfg_type: CFGType, + lcm_lora_weights: dict[str, Any], + tiny_vae_name: str, + enable_xformers: bool, + ) -> Any: + """Create a Stream Diffusion stream instance. + + Args: + pipeline: StableDiffusionPipeline instance + t_index_list: List of timestep indices + image_size: Image dimensions + do_add_noise: Whether to add noise + use_denoising_batch: Whether to use denoising batch + frame_buffer_size: Size of frame buffer + cfg_type: CFG type configuration + lcm_lora_weights: LCM LoRA weights + tiny_vae_name: Name of tiny VAE model + enable_xformers: Whether to enable xformers memory efficient attention + + Returns: + StreamDiffusion instance + """ + + @abstractmethod + def warmup_stream( + self, + stream: Any, + negative_prompt: str, + steps: Steps, + cfg: CFGScale, + delta: Delta, + seed: Seed, + warmup_count: WarmupCount, + ) -> None: + """Warm up the stream with given parameters. + + Args: + stream: StreamDiffusion instance + negative_prompt: Negative prompt text + steps: Number of inference steps + cfg: CFG scale value + delta: Delta value + seed: Random seed + warmup_count: Number of warmup iterations + """ + + @abstractmethod + def update_prompt(self, stream: Any, prompt: str) -> None: + """Update the prompt for the stream. + + Args: + stream: StreamDiffusion instance + prompt: New prompt text + """ + + @abstractmethod + def sample_txt2img( + self, + stream: Any, + num_samples: NumSamples, + ) -> list[Image.Image]: + """Generate images using txt2img. + + Args: + stream: StreamDiffusion instance + num_samples: Number of images to generate + + Returns: + List of generated PIL images + """ + + @abstractmethod + def sample_with_images( + self, + stream: Any, + prompt: str, + negative_prompt: str, + steps: Steps, + cfg: CFGScale, + delta: Delta, + seed: Seed, + num_samples: NumSamples, + warmup_count: WarmupCount, + input_images: list[Image.Image] | None, + ) -> list[Image.Image]: + """Generate images with optional input images. + + Args: + stream: StreamDiffusion instance + prompt: Prompt text + negative_prompt: Negative prompt text + steps: Number of inference steps + cfg: CFG scale value + delta: Delta value + seed: Random seed + num_samples: Number of images to generate + warmup_count: Number of warmup iterations + input_images: Optional list of input images + + Returns: + List of generated PIL images + """ diff --git a/src/infra/diffusers/__init__.py b/src/infra/diffusers/__init__.py index 3617bf9..6163882 100644 --- a/src/infra/diffusers/__init__.py +++ b/src/infra/diffusers/__init__.py @@ -1,7 +1,9 @@ from ._autoencoder_repository import DiffusersAutoencoderRepository +from ._lcm_lora_repository import DiffusersLcmLoraRepository from ._pipeline_repository import DiffusersPipelineRepository from ._sampler_repository import DiffusersSamplerRepository from ._scheduler_repository import DiffusersSchedulerRepository +from ._stream_diffusion_repository import DiffusersStreamDiffusionRepository from ._text_encoder_repository import DiffusersTextEncoderRepository __all__ = [ @@ -10,4 +12,6 @@ __all__ = [ "DiffusersTextEncoderRepository", "DiffusersSamplerRepository", "DiffusersSchedulerRepository", + "DiffusersLcmLoraRepository", + "DiffusersStreamDiffusionRepository", ] diff --git a/src/infra/diffusers/_lcm_lora_repository.py b/src/infra/diffusers/_lcm_lora_repository.py new file mode 100644 index 0000000..44654a9 --- /dev/null +++ b/src/infra/diffusers/_lcm_lora_repository.py @@ -0,0 +1,18 @@ +from typing import Any + +from safetensors.torch import load_file + +from ...domain.repositories import LcmLoraRepository + + +class DiffusersLcmLoraRepository(LcmLoraRepository): + def load_lcm_lora(self, lora_path: str) -> dict[str, Any]: + """Load LCM LoRA weights from safetensors file. + + Args: + lora_path: Path to the LCM LoRA safetensors file + + Returns: + Dictionary containing the loaded weights + """ + return load_file(lora_path) diff --git a/src/infra/diffusers/_stream_diffusion_repository.py b/src/infra/diffusers/_stream_diffusion_repository.py new file mode 100644 index 0000000..048d13b --- /dev/null +++ b/src/infra/diffusers/_stream_diffusion_repository.py @@ -0,0 +1,172 @@ +import copy +from typing import Any + +import torch +from comfy.model_management import get_torch_device # pyright: ignore[reportMissingImports] +from diffusers import AutoencoderTiny, StableDiffusionPipeline +from PIL import Image +from streamdiffusion import StreamDiffusion # type: ignore[import-untyped] +from streamdiffusion.image_utils import postprocess_image # type: ignore[import-untyped] + +from ...domain.model import ( + CFGScale, + CFGType, + Delta, + FrameBufferSize, + ImageSize, + NumSamples, + Seed, + Steps, + TIndexList, + WarmupCount, +) +from ...domain.repositories import StreamDiffusionRepository +from ._cache import get_cache_path + + +class DiffusersStreamDiffusionRepository(StreamDiffusionRepository): + def __init__(self) -> None: + self.device = get_torch_device() + self.cache_dir = get_cache_path() + self.dtype = torch.float16 + + def create_stream( + self, + pipeline: StableDiffusionPipeline, + t_index_list: TIndexList, + image_size: ImageSize, + do_add_noise: bool, + use_denoising_batch: bool, + frame_buffer_size: FrameBufferSize, + cfg_type: CFGType, + lcm_lora_weights: dict[str, Any], + tiny_vae_name: str, + enable_xformers: bool, + ) -> Any: + """Create a Stream Diffusion stream instance.""" + # Deep copy to avoid modifying the original pipeline + # Note: load_lcm_lora() and fuse_lora() modify the pipeline, + # so we need to copy it to allow reusing the original pipeline in ComfyUI + pipeline_copy = copy.deepcopy(pipeline) + lora_weights_copy = copy.deepcopy(lcm_lora_weights) + + # Create stream + stream = StreamDiffusion( + pipe=pipeline_copy, + t_index_list=t_index_list.indices, + torch_dtype=self.dtype, + width=image_size.width, + height=image_size.height, + do_add_noise=do_add_noise, + use_denoising_batch=use_denoising_batch, + frame_buffer_size=frame_buffer_size.value, + cfg_type=cfg_type.value.value, # Get the string value from enum + ) + + # Load and fuse LCM LoRA + stream.load_lcm_lora(lora_weights_copy) + stream.fuse_lora() + + # Load tiny VAE + stream.vae = AutoencoderTiny.from_pretrained( # type: ignore[no-untyped-call] + pretrained_model_name_or_path=tiny_vae_name, + torch_dtype=self.dtype, + cache_dir=self.cache_dir, + ).to(device=pipeline_copy.device, dtype=pipeline_copy.dtype) # type: ignore[attr-defined] + + # Enable xformers if requested + if enable_xformers: + pipeline_copy.enable_xformers_memory_efficient_attention() # type: ignore[attr-defined] + + return stream + + def warmup_stream( + self, + stream: Any, + negative_prompt: str, + steps: Steps, + cfg: CFGScale, + delta: Delta, + seed: Seed, + warmup_count: WarmupCount, + ) -> None: + """Warm up the stream with given parameters.""" + stream.prepare( + prompt="", + negative_prompt=negative_prompt, + num_inference_steps=steps.value, + guidance_scale=cfg.value, + delta=delta.value, + seed=seed.value, + ) + + for _ in range(warmup_count.value): + stream() + + def update_prompt(self, stream: Any, prompt: str) -> None: + """Update the prompt for the stream.""" + stream.update_prompt(prompt) + + def sample_txt2img( + self, + stream: Any, + num_samples: NumSamples, + ) -> list[Image.Image]: + """Generate images using txt2img.""" + result: list[Image.Image] = [] + for _ in range(num_samples.value): + x_output = stream.txt2img() + image = postprocess_image(x_output, output_type="pil")[0] + result.append(image) + return result + + def sample_with_images( + self, + stream: Any, + prompt: str, + negative_prompt: str, + steps: Steps, + cfg: CFGScale, + delta: Delta, + seed: Seed, + num_samples: NumSamples, + warmup_count: WarmupCount, + input_images: list[Image.Image] | None, + ) -> list[Image.Image]: + """Generate images with optional input images.""" + # Prepare stream + stream.prepare( + prompt=prompt, + negative_prompt=negative_prompt, + num_inference_steps=steps.value, + guidance_scale=cfg.value, + delta=delta.value, + seed=seed.value, + ) + + # Resize input images if provided + if input_images is not None: + resized_images = [img.resize((stream.width, stream.height)) for img in input_images] + else: + resized_images = None + + # Warmup + for _ in range(warmup_count.value): + stream() + + # Generate images + result: list[Image.Image] = [] + for _ in range(num_samples.value): + x_outputs: list[Any] = [] + if resized_images is None: + x_outputs.append(stream.txt2img()) + else: + stream(resized_images[0]) + for img in resized_images[1:] + resized_images[-1:]: + x_outputs.append(stream(img)) + + for x_output in x_outputs: + image = postprocess_image(x_output, output_type="pil")[0] + result.append(image) + + return result diff --git a/src/nodes/__init__.py b/src/nodes/__init__.py index 9c6b4ca..429536a 100644 --- a/src/nodes/__init__.py +++ b/src/nodes/__init__.py @@ -1,15 +1,24 @@ from ..di import Container +from ._create_int_list_node import CreateIntListNode from ._diffusers_clip_text_encode import DiffusersClipTextEncode from ._diffusers_pipeline_loader import DiffusersPipelineLoader from ._diffusers_sampler import DiffusersSampler from ._diffusers_scheduler_loader import DiffusersSchedulerLoader from ._diffusers_vae_loader import DiffusersVaeLoader +from ._lcm_lora_loader import LcmLoraLoader +from ._stream_diffusion_create_stream import StreamDiffusionCreateStream +from ._stream_diffusion_fast_sampler import StreamDiffusionFastSampler +from ._stream_diffusion_sampler import StreamDiffusionSampler +from ._stream_diffusion_warmup import StreamDiffusionWarmup from .dto import ( ComfyUIAutoencoderDTO, ComfyUIClipDTO, ComfyUIConditioningDTO, + ComfyUILcmLoraDTO, ComfyUIPipelineDTO, ComfyUISchedulerDTO, + ComfyUIStreamDTO, + ComfyUIWarmupStreamDTO, ) container = Container() @@ -21,6 +30,12 @@ NODE_CLASS_MAPPINGS = { DiffusersClipTextEncode.__name__: DiffusersClipTextEncode, DiffusersSampler.__name__: DiffusersSampler, DiffusersSchedulerLoader.__name__: DiffusersSchedulerLoader, + CreateIntListNode.__name__: CreateIntListNode, + LcmLoraLoader.__name__: LcmLoraLoader, + StreamDiffusionCreateStream.__name__: StreamDiffusionCreateStream, + StreamDiffusionWarmup.__name__: StreamDiffusionWarmup, + StreamDiffusionSampler.__name__: StreamDiffusionSampler, + StreamDiffusionFastSampler.__name__: StreamDiffusionFastSampler, } NODE_DISPLAY_NAME_MAPPINGS = { DiffusersPipelineLoader.__name__: "Diffusers Pipeline Loader", @@ -28,6 +43,12 @@ NODE_DISPLAY_NAME_MAPPINGS = { DiffusersClipTextEncode.__name__: "Diffusers CLIP Text Encode", DiffusersSampler.__name__: "Diffusers Sampler", DiffusersSchedulerLoader.__name__: "Diffusers Scheduler Loader", + CreateIntListNode.__name__: "Create Int List", + LcmLoraLoader.__name__: "LCM LoRA Loader", + StreamDiffusionCreateStream.__name__: "StreamDiffusion Create Stream", + StreamDiffusionWarmup.__name__: "StreamDiffusion Warmup", + StreamDiffusionSampler.__name__: "StreamDiffusion Sampler", + StreamDiffusionFastSampler.__name__: "StreamDiffusion Fast Sampler", } __all__ = [ @@ -38,4 +59,7 @@ __all__ = [ "ComfyUIClipDTO", "ComfyUIConditioningDTO", "ComfyUISchedulerDTO", + "ComfyUILcmLoraDTO", + "ComfyUIStreamDTO", + "ComfyUIWarmupStreamDTO", ] diff --git a/src/nodes/_create_int_list_node.py b/src/nodes/_create_int_list_node.py new file mode 100644 index 0000000..12446c0 --- /dev/null +++ b/src/nodes/_create_int_list_node.py @@ -0,0 +1,36 @@ +class CreateIntListNode: + """Utility node to create a list of integers for Stream Diffusion t_index_list.""" + + MAX_ELEMENTS = 10 + + @classmethod + def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str, dict[str, int]] | tuple[str, ...]]]: + inputs: dict[str, dict[str, tuple[str, dict[str, int]] | tuple[str, ...]]] = { + "required": { + "elements_count": ( + "INT", + {"default": 2, "min": 1, "max": cls.MAX_ELEMENTS, "step": 1}, + ), + }, + "optional": {}, + } + for i in range(1, cls.MAX_ELEMENTS): + inputs["optional"][f"element_{i}"] = ("INT", {"default": 0}) + return inputs + + RETURN_TYPES = ("LIST",) + FUNCTION = "execute" + CATEGORY = "Diffusers/StreamDiffusion" + + def execute(self, elements_count: int, **kwargs: int) -> tuple[list[int]]: + """Create a list of integers from the input parameters. + + Args: + elements_count: Number of elements to include in the list + **kwargs: Optional element values + + Returns: + Tuple containing the list of integers + """ + result = [value for key, value in kwargs.items()][:elements_count] + return (result,) diff --git a/src/nodes/_lcm_lora_loader.py b/src/nodes/_lcm_lora_loader.py new file mode 100644 index 0000000..cc138a5 --- /dev/null +++ b/src/nodes/_lcm_lora_loader.py @@ -0,0 +1,45 @@ +import folder_paths # pyright: ignore[reportMissingImports] +from dependency_injector.wiring import Provide, inject + +from ..di import Container +from ..usecase import LcmLoraUsecase +from .dto import ComfyUILcmLoraDTO + + +class LcmLoraLoader: + """Node to load LCM LoRA weights for Stream Diffusion.""" + + def __init__(self) -> None: + pass + + @classmethod + def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[tuple[str, ...], ...]]]: + return { + "required": { + "lora_name": (folder_paths.get_filename_list("loras"),), + } + } + + RETURN_TYPES = (ComfyUILcmLoraDTO.COMFY_TYPE,) + FUNCTION = "execute" + CATEGORY = "Diffusers/StreamDiffusion" + + @inject + def execute( + self, + lora_name: str, + usecase: LcmLoraUsecase = Provide[Container.lcm_lora_usecase], + ) -> tuple[ComfyUILcmLoraDTO]: + """Load LCM LoRA weights. + + Args: + lora_name: Name of the LoRA file + usecase: Injected LcmLoraUsecase + + Returns: + Tuple containing ComfyUILcmLoraDTO + """ + lora_path = folder_paths.get_full_path("loras", lora_name) + lcm_lora = usecase.execute(lora_path) + dto = ComfyUILcmLoraDTO.from_domain(lcm_lora) + return (dto,) diff --git a/src/nodes/_stream_diffusion_create_stream.py b/src/nodes/_stream_diffusion_create_stream.py new file mode 100644 index 0000000..5570691 --- /dev/null +++ b/src/nodes/_stream_diffusion_create_stream.py @@ -0,0 +1,102 @@ +from typing import Any + +from dependency_injector.wiring import Provide, inject + +from ..di import Container +from ..domain.model import CFGType, CFGTypeEnum, FrameBufferSize, ImageSize, TIndexList +from ..usecase import StreamDiffusionCreateStreamUsecase +from .dto import ComfyUILcmLoraDTO, ComfyUIPipelineDTO, ComfyUIStreamDTO + +NodeInputMap = dict[str, dict[str, tuple[str, ...] | tuple[str | list[str], dict[str, Any]]]] + + +class StreamDiffusionCreateStream: + """Node to create a Stream Diffusion stream.""" + + def __init__(self) -> None: + pass + + @classmethod + def INPUT_TYPES( + cls, + ) -> NodeInputMap: + return { + "required": { + "pipeline": (ComfyUIPipelineDTO.COMFY_TYPE,), + "t_index_list": ("LIST",), + "width": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}), + "height": ("INT", {"default": 512, "min": 1, "max": 8192, "step": 1}), + "do_add_noise": ("BOOLEAN", {"default": True}), + "use_denoising_batch": ("BOOLEAN", {"default": True}), + "frame_buffer_size": ("INT", {"default": 1, "min": 1, "max": 10000}), + "cfg_type": (["none", "full", "self", "initialize"], {"default": "none"}), + "lcm_lora": (ComfyUILcmLoraDTO.COMFY_TYPE,), + "tiny_vae": ("STRING", {"default": "madebyollin/taesd"}), + "enable_xformers": ("BOOLEAN", {"default": False}), + }, + } + + RETURN_TYPES = (ComfyUIStreamDTO.COMFY_TYPE,) + FUNCTION = "execute" + CATEGORY = "Diffusers/StreamDiffusion" + + @inject + def execute( + self, + pipeline: ComfyUIPipelineDTO, + t_index_list: list[int], + width: int, + height: int, + do_add_noise: bool, + use_denoising_batch: bool, + frame_buffer_size: int, + cfg_type: str, + lcm_lora: ComfyUILcmLoraDTO, + tiny_vae: str, + enable_xformers: bool, + usecase: StreamDiffusionCreateStreamUsecase = Provide[ + Container.stream_diffusion_create_stream_usecase + ], + ) -> tuple[ComfyUIStreamDTO]: + """Create a Stream Diffusion stream. + + Args: + pipeline: Pipeline DTO + t_index_list: List of timestep indices + width: Image width + height: Image height + do_add_noise: Whether to add noise + use_denoising_batch: Whether to use denoising batch + frame_buffer_size: Size of frame buffer + cfg_type: CFG type string + lcm_lora: LCM LoRA DTO + tiny_vae: Name of tiny VAE model + enable_xformers: Whether to enable xformers + usecase: Injected StreamDiffusionCreateStreamUsecase + + Returns: + Tuple containing ComfyUIStreamDTO + """ + pipeline_domain = ComfyUIPipelineDTO.to_domain(pipeline) + lcm_lora_domain = ComfyUILcmLoraDTO.to_domain(lcm_lora) + + t_index_list_vo = TIndexList(indices=t_index_list) + image_size = ImageSize(width=width, height=height) + frame_buffer_size_vo = FrameBufferSize(value=frame_buffer_size) + cfg_type_vo = CFGType(value=CFGTypeEnum(cfg_type)) + + stream = usecase.execute( + pipeline=pipeline_domain, + t_index_list=t_index_list_vo, + image_size=image_size, + do_add_noise=do_add_noise, + use_denoising_batch=use_denoising_batch, + frame_buffer_size=frame_buffer_size_vo, + cfg_type=cfg_type_vo, + lcm_lora=lcm_lora_domain, + tiny_vae_name=tiny_vae, + enable_xformers=enable_xformers, + ) + + dto = ComfyUIStreamDTO.from_domain(stream) + return (dto,) diff --git a/src/nodes/_stream_diffusion_fast_sampler.py b/src/nodes/_stream_diffusion_fast_sampler.py new file mode 100644 index 0000000..2a8f611 --- /dev/null +++ b/src/nodes/_stream_diffusion_fast_sampler.py @@ -0,0 +1,66 @@ +from typing import Any + +from dependency_injector.wiring import Provide, inject + +from ..di import Container +from ..domain.model import NumSamples +from ..usecase import StreamDiffusionFastSampleUsecase +from .dto import ComfyUIImage, ComfyUIImageDTO, ComfyUIWarmupStreamDTO + +NodeInputMap = dict[str, dict[str, tuple[str, ...] | tuple[str | list[str], dict[str, Any]]]] + + +class StreamDiffusionFastSampler: + """Node for fast sampling from a warmed-up Stream Diffusion stream.""" + + def __init__(self) -> None: + pass + + @classmethod + def INPUT_TYPES( + cls, + ) -> NodeInputMap: + return { + "required": { + "warmup_stream": (ComfyUIWarmupStreamDTO.COMFY_TYPE,), + "positive_prompt": ("STRING", {"multiline": True}), + "num": ("INT", {"default": 1, "min": 1, "max": 10000}), + }, + } + + RETURN_TYPES = (ComfyUIImage.COMFY_TYPE,) + FUNCTION = "execute" + CATEGORY = "Diffusers/StreamDiffusion" + + @inject + def execute( + self, + warmup_stream: ComfyUIWarmupStreamDTO, + positive_prompt: str, + num: int, + usecase: StreamDiffusionFastSampleUsecase = Provide[ + Container.stream_diffusion_fast_sample_usecase + ], + ) -> tuple[ComfyUIImageDTO]: + """Fast sample images from a warmed-up stream. + + Args: + warmup_stream: Warmed-up stream DTO + positive_prompt: Prompt text + num: Number of images to generate + usecase: Injected StreamDiffusionFastSampleUsecase + + Returns: + Tuple containing image tensor + """ + stream_domain = ComfyUIWarmupStreamDTO.to_domain(warmup_stream) + num_vo = NumSamples(value=num) + + images = usecase.execute( + stream=stream_domain, + prompt=positive_prompt, + num_samples=num_vo, + ) + + result_dto = ComfyUIImage.from_domains(images) + return (result_dto,) diff --git a/src/nodes/_stream_diffusion_sampler.py b/src/nodes/_stream_diffusion_sampler.py new file mode 100644 index 0000000..56483c9 --- /dev/null +++ b/src/nodes/_stream_diffusion_sampler.py @@ -0,0 +1,114 @@ +from typing import Any + +import numpy as np +from dependency_injector.wiring import Provide, inject +from PIL import Image as PilImage + +from ..di import Container +from ..domain.model import CFGScale, Delta, Image, NumSamples, Seed, Steps, WarmupCount +from ..usecase import StreamDiffusionSampleUsecase +from .dto import ComfyUIImage, ComfyUIImageDTO, ComfyUIStreamDTO + +NodeInputMap = dict[str, dict[str, tuple[str, ...] | tuple[str | list[str], dict[str, Any]]]] + + +class StreamDiffusionSampler: + """Node to sample images from a Stream Diffusion stream.""" + + def __init__(self) -> None: + pass + + @classmethod + def INPUT_TYPES( + cls, + ) -> NodeInputMap: + return { + "required": { + "stream": (ComfyUIStreamDTO.COMFY_TYPE,), + "positive_prompt": ("STRING", {"multiline": True}), + "negative_prompt": ("STRING", {"multiline": True, "default": ""}), + "steps": ("INT", {"default": 50, "min": 1, "max": 10000}), + "cfg": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 100.0}), + "delta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), + "num": ("INT", {"default": 1, "min": 1, "max": 10000}), + "warmup": ("INT", {"default": 1, "min": 0, "max": 10000}), + }, + "optional": { + "image": (ComfyUIImage.COMFY_TYPE,), + }, + } + + RETURN_TYPES = (ComfyUIImage.COMFY_TYPE,) + FUNCTION = "execute" + CATEGORY = "Diffusers/StreamDiffusion" + + @inject + def execute( + self, + stream: ComfyUIStreamDTO, + positive_prompt: str, + negative_prompt: str, + steps: int, + cfg: float, + delta: float, + seed: int, + num: int, + warmup: int, + image: ComfyUIImageDTO | None = None, + usecase: StreamDiffusionSampleUsecase = Provide[Container.stream_diffusion_sample_usecase], + ) -> tuple[ComfyUIImageDTO]: + """Sample images from a Stream Diffusion stream. + + Args: + stream: Stream DTO + positive_prompt: Prompt text + negative_prompt: Negative prompt text + steps: Number of inference steps + cfg: CFG scale value + delta: Delta value + seed: Random seed + num: Number of images to generate + warmup: Number of warmup iterations + image: Optional input image tensor + usecase: Injected StreamDiffusionSampleUsecase + + Returns: + Tuple containing image tensor + """ + stream_domain = ComfyUIStreamDTO.to_domain(stream) + + steps_vo = Steps(value=steps) + cfg_vo = CFGScale(value=cfg) + delta_vo = Delta(value=delta) + seed_vo = Seed(value=seed) + num_vo = NumSamples(value=num) + warmup_vo = WarmupCount(value=warmup) + + # Convert input images if provided + input_images_domain: list[Image] | None = None + if image is not None: + # ComfyUI images are in format [B, H, W, C] with values in [0, 1] + images_np = image.cpu().numpy() + input_images_domain = [] + for img_np in images_np: + # Convert to uint8 and create PIL image + img_uint8 = (np.clip(img_np, 0, 1) * 255).astype(np.uint8) + pil_img = PilImage.fromarray(img_uint8) + input_images_domain.append(Image(image=pil_img)) + + images = usecase.execute( + stream=stream_domain, + prompt=positive_prompt, + negative_prompt=negative_prompt, + steps=steps_vo, + cfg=cfg_vo, + delta=delta_vo, + seed=seed_vo, + num_samples=num_vo, + warmup_count=warmup_vo, + input_images=input_images_domain, + ) + + result_dto = ComfyUIImage.from_domains(images) + return (result_dto,) diff --git a/src/nodes/_stream_diffusion_warmup.py b/src/nodes/_stream_diffusion_warmup.py new file mode 100644 index 0000000..7b9fb37 --- /dev/null +++ b/src/nodes/_stream_diffusion_warmup.py @@ -0,0 +1,85 @@ +from typing import Any + +from dependency_injector.wiring import Provide, inject + +from ..di import Container +from ..domain.model import CFGScale, Delta, Seed, Steps, WarmupCount +from ..usecase import StreamDiffusionWarmupUsecase +from .dto import ComfyUIStreamDTO, ComfyUIWarmupStreamDTO + +NodeInputMap = dict[str, dict[str, tuple[str, ...] | tuple[str | list[str], dict[str, Any]]]] + + +class StreamDiffusionWarmup: + """Node to warm up a Stream Diffusion stream.""" + + def __init__(self) -> None: + pass + + @classmethod + def INPUT_TYPES( + cls, + ) -> NodeInputMap: + return { + "required": { + "stream": (ComfyUIStreamDTO.COMFY_TYPE,), + "negative_prompt": ("STRING", {"multiline": True, "default": ""}), + "steps": ("INT", {"default": 50, "min": 1, "max": 10000}), + "cfg": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 100.0}), + "delta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}), + "warmup": ("INT", {"default": 1, "min": 0, "max": 10000}), + }, + } + + RETURN_TYPES = (ComfyUIWarmupStreamDTO.COMFY_TYPE,) + FUNCTION = "execute" + CATEGORY = "Diffusers/StreamDiffusion" + + @inject + def execute( + self, + stream: ComfyUIStreamDTO, + negative_prompt: str, + steps: int, + cfg: float, + delta: float, + seed: int, + warmup: int, + usecase: StreamDiffusionWarmupUsecase = Provide[Container.stream_diffusion_warmup_usecase], + ) -> tuple[ComfyUIWarmupStreamDTO]: + """Warm up a Stream Diffusion stream. + + Args: + stream: Stream DTO + negative_prompt: Negative prompt text + steps: Number of inference steps + cfg: CFG scale value + delta: Delta value + seed: Random seed + warmup: Number of warmup iterations + usecase: Injected StreamDiffusionWarmupUsecase + + Returns: + Tuple containing ComfyUIWarmupStreamDTO + """ + stream_domain = ComfyUIStreamDTO.to_domain(stream) + + steps_vo = Steps(value=steps) + cfg_vo = CFGScale(value=cfg) + delta_vo = Delta(value=delta) + seed_vo = Seed(value=seed) + warmup_vo = WarmupCount(value=warmup) + + warmed_stream = usecase.execute( + stream=stream_domain, + negative_prompt=negative_prompt, + steps=steps_vo, + cfg=cfg_vo, + delta=delta_vo, + seed=seed_vo, + warmup_count=warmup_vo, + ) + + dto = ComfyUIWarmupStreamDTO.from_domain(warmed_stream) + return (dto,) diff --git a/src/nodes/dto/__init__.py b/src/nodes/dto/__init__.py index c5c67b1..6f9165a 100644 --- a/src/nodes/dto/__init__.py +++ b/src/nodes/dto/__init__.py @@ -2,8 +2,10 @@ from ._autoencoder_dto import ComfyUIAutoencoderDTO from ._clip import ComfyUIClipDTO from ._comfyui import ComfyUIImage, ComfyUIImageDTO from ._conditioning_dto import ComfyUIConditioningDTO +from ._lcm_lora_dto import ComfyUILcmLoraDTO from ._pipeline_dto import ComfyUIPipelineDTO from ._scheduler_dto import ComfyUISchedulerDTO +from ._stream_diffusion_stream_dto import ComfyUIStreamDTO, ComfyUIWarmupStreamDTO __all__ = [ "ComfyUIPipelineDTO", @@ -13,4 +15,7 @@ __all__ = [ "ComfyUIImage", "ComfyUIImageDTO", "ComfyUISchedulerDTO", + "ComfyUILcmLoraDTO", + "ComfyUIStreamDTO", + "ComfyUIWarmupStreamDTO", ] diff --git a/src/nodes/dto/_lcm_lora_dto.py b/src/nodes/dto/_lcm_lora_dto.py new file mode 100644 index 0000000..0b1b9c0 --- /dev/null +++ b/src/nodes/dto/_lcm_lora_dto.py @@ -0,0 +1,44 @@ +from dataclasses import dataclass +from typing import Any, ClassVar + +from ...domain.model import LcmLora + + +@dataclass +class ComfyUILcmLoraDTO: + """Data Transfer Object for LCM LoRA in ComfyUI. + + Attributes: + weights: Dictionary containing the LCM LoRA weights + + Class Attributes: + COMFY_TYPE: Type name used in ComfyUI's RETURN_TYPES and INPUT_TYPES + """ + + COMFY_TYPE: ClassVar[str] = "DIFFUSERS_LCM_LORA" + + weights: dict[str, Any] + + @classmethod + def from_domain(cls, lcm_lora: LcmLora) -> "ComfyUILcmLoraDTO": + """Create DTO from domain model. + + Args: + lcm_lora: LcmLora from domain layer + + Returns: + ComfyUILcmLoraDTO instance + """ + return cls(weights=lcm_lora.weights) + + @classmethod + def to_domain(cls, dto: "ComfyUILcmLoraDTO") -> LcmLora: + """Convert DTO back to domain model. + + Args: + dto: ComfyUILcmLoraDTO instance + + Returns: + LcmLora domain model instance + """ + return LcmLora(weights=dto.weights) diff --git a/src/nodes/dto/_stream_diffusion_stream_dto.py b/src/nodes/dto/_stream_diffusion_stream_dto.py new file mode 100644 index 0000000..1ceee2f --- /dev/null +++ b/src/nodes/dto/_stream_diffusion_stream_dto.py @@ -0,0 +1,86 @@ +from dataclasses import dataclass +from typing import Any, ClassVar + +from ...domain.model import StreamDiffusionStream + + +@dataclass +class ComfyUIStreamDTO: + """Data Transfer Object for Stream Diffusion Stream in ComfyUI. + + Attributes: + stream: StreamDiffusion instance from streamdiffusion library + + Class Attributes: + COMFY_TYPE: Type name used in ComfyUI's RETURN_TYPES and INPUT_TYPES + """ + + COMFY_TYPE: ClassVar[str] = "DIFFUSERS_STREAM" + + stream: Any + + @classmethod + def from_domain(cls, stream: StreamDiffusionStream) -> "ComfyUIStreamDTO": + """Create DTO from domain model. + + Args: + stream: StreamDiffusionStream from domain layer + + Returns: + ComfyUIStreamDTO instance + """ + return cls(stream=stream.stream) + + @classmethod + def to_domain(cls, dto: "ComfyUIStreamDTO") -> StreamDiffusionStream: + """Convert DTO back to domain model. + + Args: + dto: ComfyUIStreamDTO instance + + Returns: + StreamDiffusionStream domain model instance + """ + return StreamDiffusionStream(stream=dto.stream) + + +@dataclass +class ComfyUIWarmupStreamDTO: + """Data Transfer Object for Warmed-up Stream Diffusion Stream in ComfyUI. + + This is a separate type to distinguish warmed-up streams from fresh streams. + + Attributes: + stream: StreamDiffusion instance from streamdiffusion library + + Class Attributes: + COMFY_TYPE: Type name used in ComfyUI's RETURN_TYPES and INPUT_TYPES + """ + + COMFY_TYPE: ClassVar[str] = "DIFFUSERS_WARMUP_STREAM" + + stream: Any + + @classmethod + def from_domain(cls, stream: StreamDiffusionStream) -> "ComfyUIWarmupStreamDTO": + """Create DTO from domain model. + + Args: + stream: StreamDiffusionStream from domain layer + + Returns: + ComfyUIWarmupStreamDTO instance + """ + return cls(stream=stream.stream) + + @classmethod + def to_domain(cls, dto: "ComfyUIWarmupStreamDTO") -> StreamDiffusionStream: + """Convert DTO back to domain model. + + Args: + dto: ComfyUIWarmupStreamDTO instance + + Returns: + StreamDiffusionStream domain model instance + """ + return StreamDiffusionStream(stream=dto.stream) diff --git a/src/usecase/__init__.py b/src/usecase/__init__.py index 63ae7cc..a19f2ec 100644 --- a/src/usecase/__init__.py +++ b/src/usecase/__init__.py @@ -1,8 +1,13 @@ from ._autoencoder_usecase import AutoencoderUsecase from ._clip_text_encode_usecase import ClipTextEncodeUsecase +from ._lcm_lora_usecase import LcmLoraUsecase from ._pipeline_usecase import PipelineUsecase from ._sampler_usecase import SamplerUsecase from ._scheduler_usecase import SchedulerUsecase +from ._stream_diffusion_create_stream_usecase import StreamDiffusionCreateStreamUsecase +from ._stream_diffusion_fast_sample_usecase import StreamDiffusionFastSampleUsecase +from ._stream_diffusion_sample_usecase import StreamDiffusionSampleUsecase +from ._stream_diffusion_warmup_usecase import StreamDiffusionWarmupUsecase __all__ = [ "PipelineUsecase", @@ -10,4 +15,9 @@ __all__ = [ "ClipTextEncodeUsecase", "SamplerUsecase", "SchedulerUsecase", + "LcmLoraUsecase", + "StreamDiffusionCreateStreamUsecase", + "StreamDiffusionWarmupUsecase", + "StreamDiffusionSampleUsecase", + "StreamDiffusionFastSampleUsecase", ] diff --git a/src/usecase/_lcm_lora_usecase.py b/src/usecase/_lcm_lora_usecase.py new file mode 100644 index 0000000..c07771a --- /dev/null +++ b/src/usecase/_lcm_lora_usecase.py @@ -0,0 +1,26 @@ +from ..domain.model import LcmLora +from ..domain.repositories import LcmLoraRepository + + +class LcmLoraUsecase: + def __init__(self, lcm_lora_repo: LcmLoraRepository) -> None: + self.lcm_lora_repo = lcm_lora_repo + + def execute(self, lora_path: str) -> LcmLora: + """Load LCM LoRA weights from file. + + Args: + lora_path: Path to the LCM LoRA file + + Returns: + LcmLora domain model + + Raises: + FileNotFoundError: If the lora file doesn't exist + RuntimeError: If loading fails + """ + weights = self.lcm_lora_repo.load_lcm_lora(lora_path) + if not weights: + raise RuntimeError(f"Failed to load LCM LoRA weights from: {lora_path}") + + return LcmLora(weights=weights) diff --git a/src/usecase/_stream_diffusion_create_stream_usecase.py b/src/usecase/_stream_diffusion_create_stream_usecase.py new file mode 100644 index 0000000..df5564b --- /dev/null +++ b/src/usecase/_stream_diffusion_create_stream_usecase.py @@ -0,0 +1,66 @@ +from ..domain.model import ( + CFGType, + FrameBufferSize, + ImageSize, + LcmLora, + Pipeline, + StreamDiffusionStream, + TIndexList, +) +from ..domain.repositories import StreamDiffusionRepository + + +class StreamDiffusionCreateStreamUsecase: + def __init__(self, stream_diffusion_repo: StreamDiffusionRepository) -> None: + self.stream_diffusion_repo = stream_diffusion_repo + + def execute( + self, + pipeline: Pipeline, + t_index_list: TIndexList, + image_size: ImageSize, + do_add_noise: bool, + use_denoising_batch: bool, + frame_buffer_size: FrameBufferSize, + cfg_type: CFGType, + lcm_lora: LcmLora, + tiny_vae_name: str, + enable_xformers: bool, + ) -> StreamDiffusionStream: + """Create a Stream Diffusion stream. + + Args: + pipeline: Pipeline model + t_index_list: List of timestep indices + image_size: Image dimensions + do_add_noise: Whether to add noise + use_denoising_batch: Whether to use denoising batch + frame_buffer_size: Size of frame buffer + cfg_type: CFG type configuration + lcm_lora: LCM LoRA model + tiny_vae_name: Name of tiny VAE model + enable_xformers: Whether to enable xformers + + Returns: + StreamDiffusionStream domain model + + Raises: + RuntimeError: If stream creation fails + """ + stream = self.stream_diffusion_repo.create_stream( + pipeline=pipeline.pipeline, + t_index_list=t_index_list, + image_size=image_size, + do_add_noise=do_add_noise, + use_denoising_batch=use_denoising_batch, + frame_buffer_size=frame_buffer_size, + cfg_type=cfg_type, + lcm_lora_weights=lcm_lora.weights, + tiny_vae_name=tiny_vae_name, + enable_xformers=enable_xformers, + ) + + if stream is None: + raise RuntimeError("Failed to create Stream Diffusion stream") + + return StreamDiffusionStream(stream=stream) diff --git a/src/usecase/_stream_diffusion_fast_sample_usecase.py b/src/usecase/_stream_diffusion_fast_sample_usecase.py new file mode 100644 index 0000000..789a03b --- /dev/null +++ b/src/usecase/_stream_diffusion_fast_sample_usecase.py @@ -0,0 +1,40 @@ +from ..domain.model import Image, NumSamples, StreamDiffusionStream +from ..domain.repositories import StreamDiffusionRepository + + +class StreamDiffusionFastSampleUsecase: + def __init__(self, stream_diffusion_repo: StreamDiffusionRepository) -> None: + self.stream_diffusion_repo = stream_diffusion_repo + + def execute( + self, + stream: StreamDiffusionStream, + prompt: str, + num_samples: NumSamples, + ) -> list[Image]: + """Fast sample images from a warmed-up Stream Diffusion stream. + + Args: + stream: Warmed-up stream to sample from + prompt: Prompt text + num_samples: Number of images to generate + + Returns: + List of generated images as domain models + + Raises: + RuntimeError: If sampling fails + """ + # Update prompt + self.stream_diffusion_repo.update_prompt(stream=stream.stream, prompt=prompt) + + # Sample images + images = self.stream_diffusion_repo.sample_txt2img( + stream=stream.stream, + num_samples=num_samples, + ) + + if not images: + raise RuntimeError("Stream Diffusion repository returned no images") + + return [Image(image=img) for img in images] diff --git a/src/usecase/_stream_diffusion_sample_usecase.py b/src/usecase/_stream_diffusion_sample_usecase.py new file mode 100644 index 0000000..5736c77 --- /dev/null +++ b/src/usecase/_stream_diffusion_sample_usecase.py @@ -0,0 +1,72 @@ +from ..domain.model import ( + CFGScale, + Delta, + Image, + NumSamples, + Seed, + Steps, + StreamDiffusionStream, + WarmupCount, +) +from ..domain.repositories import StreamDiffusionRepository + + +class StreamDiffusionSampleUsecase: + def __init__(self, stream_diffusion_repo: StreamDiffusionRepository) -> None: + self.stream_diffusion_repo = stream_diffusion_repo + + def execute( + self, + stream: StreamDiffusionStream, + prompt: str, + negative_prompt: str, + steps: Steps, + cfg: CFGScale, + delta: Delta, + seed: Seed, + num_samples: NumSamples, + warmup_count: WarmupCount, + input_images: list[Image] | None, + ) -> list[Image]: + """Sample images from a Stream Diffusion stream. + + Args: + stream: Stream to sample from + prompt: Prompt text + negative_prompt: Negative prompt text + steps: Number of inference steps + cfg: CFG scale value + delta: Delta value + seed: Random seed + num_samples: Number of images to generate + warmup_count: Number of warmup iterations + input_images: Optional list of input images + + Returns: + List of generated images as domain models + + Raises: + RuntimeError: If sampling fails + """ + # Convert domain images to PIL images if provided + pil_images = None + if input_images is not None: + pil_images = [img.image for img in input_images] + + images = self.stream_diffusion_repo.sample_with_images( + stream=stream.stream, + prompt=prompt, + negative_prompt=negative_prompt, + steps=steps, + cfg=cfg, + delta=delta, + seed=seed, + num_samples=num_samples, + warmup_count=warmup_count, + input_images=pil_images, + ) + + if not images: + raise RuntimeError("Stream Diffusion repository returned no images") + + return [Image(image=img) for img in images] diff --git a/src/usecase/_stream_diffusion_warmup_usecase.py b/src/usecase/_stream_diffusion_warmup_usecase.py new file mode 100644 index 0000000..9bdebd4 --- /dev/null +++ b/src/usecase/_stream_diffusion_warmup_usecase.py @@ -0,0 +1,43 @@ +from ..domain.model import CFGScale, Delta, Seed, Steps, StreamDiffusionStream, WarmupCount +from ..domain.repositories import StreamDiffusionRepository + + +class StreamDiffusionWarmupUsecase: + def __init__(self, stream_diffusion_repo: StreamDiffusionRepository) -> None: + self.stream_diffusion_repo = stream_diffusion_repo + + def execute( + self, + stream: StreamDiffusionStream, + negative_prompt: str, + steps: Steps, + cfg: CFGScale, + delta: Delta, + seed: Seed, + warmup_count: WarmupCount, + ) -> StreamDiffusionStream: + """Warm up a Stream Diffusion stream. + + Args: + stream: Stream to warm up + negative_prompt: Negative prompt text + steps: Number of inference steps + cfg: CFG scale value + delta: Delta value + seed: Random seed + warmup_count: Number of warmup iterations + + Returns: + Warmed up StreamDiffusionStream (same instance) + """ + self.stream_diffusion_repo.warmup_stream( + stream=stream.stream, + negative_prompt=negative_prompt, + steps=steps, + cfg=cfg, + delta=delta, + seed=seed, + warmup_count=warmup_count, + ) + + return stream