nodes/Limitex-ComfyUI-Diffusers
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: stream diffusion impl

Browse Source
This commit is contained in:
Limitex
2025-11-23 22:59:05 +09:00
parent 2aecbc841e
commit 06c9f21ffb
33 changed files with 1388 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+6
View File
@@ -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",
]
src/di/di.py
+29
View File
@@ -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,
)
src/domain/model/__init__.py
+17
View File
@@ -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",
]
src/domain/model/_cfg_type.py
+22
View File
@@ -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__}")
src/domain/model/_delta.py
+14
View File
@@ -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}")
src/domain/model/_frame_buffer_size.py
+14
View File
@@ -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}")
src/domain/model/_lcm_lora.py
+15
View File
@@ -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")
src/domain/model/_num_samples.py
+14
View File
@@ -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}")
src/domain/model/_stream_diffusion_stream.py
+15
View File
@@ -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
src/domain/model/_t_index_list.py
+19
View File
@@ -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__}"
)
src/domain/model/_warmup_count.py
+14
View File
@@ -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}")
src/domain/repositories/__init__.py
+4
View File
@@ -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",
]
src/domain/repositories/_lcm_lora_repository.py
+15
View File
@@ -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
"""
src/domain/repositories/_stream_diffusion_repository.py
+132
View File
@@ -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
"""
src/infra/diffusers/__init__.py
+4
View File
@@ -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",
]
src/infra/diffusers/_lcm_lora_repository.py
+18
View File
@@ -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)
src/infra/diffusers/_stream_diffusion_repository.py
+172
View File
@@ -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
src/nodes/__init__.py
+24
View File
@@ -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",
]
src/nodes/_create_int_list_node.py
+36
View File
@@ -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,)
src/nodes/_lcm_lora_loader.py
+45
View File
@@ -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,)
src/nodes/_stream_diffusion_create_stream.py
+102
View File
@@ -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,)
src/nodes/_stream_diffusion_fast_sampler.py
+66
View File
@@ -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,)
src/nodes/_stream_diffusion_sampler.py
+114
View File
@@ -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,)
src/nodes/_stream_diffusion_warmup.py
+85
View File
@@ -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,)
src/nodes/dto/__init__.py
+5
View File
@@ -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",
]
src/nodes/dto/_lcm_lora_dto.py
+44
View File
@@ -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)
src/nodes/dto/_stream_diffusion_stream_dto.py
+86
View File
@@ -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)
src/usecase/__init__.py
+10
View File
@@ -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",
]
src/usecase/_lcm_lora_usecase.py
+26
View File
@@ -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)
src/usecase/_stream_diffusion_create_stream_usecase.py
+66
View File
@@ -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)
src/usecase/_stream_diffusion_fast_sample_usecase.py
+40
View File
@@ -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]
src/usecase/_stream_diffusion_sample_usecase.py
+72
View File
@@ -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]
src/usecase/_stream_diffusion_warmup_usecase.py
+43
View File
@@ -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