diff --git a/src/domain/model/_t_index_list.py b/src/domain/model/_t_index_list.py index 20858ea..4d6eca5 100644 --- a/src/domain/model/_t_index_list.py +++ b/src/domain/model/_t_index_list.py @@ -1,19 +1,35 @@ -from dataclasses import dataclass +from __future__ import annotations + +import json +from dataclasses import dataclass, field @dataclass(frozen=True) class TIndexList: - """List of timestep indices for Stream Diffusion.""" + """Value object that validates and stores the t-index sequence.""" - indices: list[int] + raw_value: str + value: tuple[int, ...] = field(init=False) 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__}" - ) + if not isinstance(self.raw_value, str): + raise ValueError("t_index_list must be provided as a string.") + + try: + parsed = json.loads(self.raw_value) + except json.JSONDecodeError as exc: + raise ValueError("t_index_list must be a JSON array formatted string.") from exc + + if not isinstance(parsed, list) or not parsed: + raise ValueError("t_index_list must be a non-empty list of integers.") + + validated: list[int] = [] + for idx, item in enumerate(parsed): + if isinstance(item, bool) or not isinstance(item, int): + raise ValueError(f"t_index_list[{idx}] must be an int, got {type(item).__name__}.") + validated.append(item) + + object.__setattr__(self, "value", tuple(validated)) + + def as_list(self) -> list[int]: + return list(self.value) diff --git a/src/infra/diffusers/_stream_diffusion_repository.py b/src/infra/diffusers/_stream_diffusion_repository.py index 33224b9..302f457 100644 --- a/src/infra/diffusers/_stream_diffusion_repository.py +++ b/src/infra/diffusers/_stream_diffusion_repository.py @@ -55,7 +55,7 @@ class DiffusersStreamDiffusionRepository(StreamDiffusionRepository): # Create stream stream = StreamDiffusion( pipe=pipeline_copy, - t_index_list=t_index_list.indices, + t_index_list=t_index_list.as_list(), torch_dtype=self.dtype, width=image_size.width, height=image_size.height, diff --git a/src/nodes/__init__.py b/src/nodes/__init__.py index 429536a..fb09025 100644 --- a/src/nodes/__init__.py +++ b/src/nodes/__init__.py @@ -1,5 +1,4 @@ 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 @@ -30,7 +29,6 @@ NODE_CLASS_MAPPINGS = { DiffusersClipTextEncode.__name__: DiffusersClipTextEncode, DiffusersSampler.__name__: DiffusersSampler, DiffusersSchedulerLoader.__name__: DiffusersSchedulerLoader, - CreateIntListNode.__name__: CreateIntListNode, LcmLoraLoader.__name__: LcmLoraLoader, StreamDiffusionCreateStream.__name__: StreamDiffusionCreateStream, StreamDiffusionWarmup.__name__: StreamDiffusionWarmup, @@ -43,7 +41,6 @@ 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", diff --git a/src/nodes/_create_int_list_node.py b/src/nodes/_create_int_list_node.py deleted file mode 100644 index 12446c0..0000000 --- a/src/nodes/_create_int_list_node.py +++ /dev/null @@ -1,36 +0,0 @@ -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/_stream_diffusion_create_stream.py b/src/nodes/_stream_diffusion_create_stream.py index 393b042..673dc08 100644 --- a/src/nodes/_stream_diffusion_create_stream.py +++ b/src/nodes/_stream_diffusion_create_stream.py @@ -29,7 +29,7 @@ class StreamDiffusionCreateStream: "required": { "pipeline": (ComfyUIPipelineDTO.COMFY_TYPE,), "scheduler": (ComfyUISchedulerDTO.COMFY_TYPE,), - "t_index_list": ("LIST",), + "t_index_list": ("STRING", {"default": "[0, 16, 32, 45]"}), "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}), @@ -51,7 +51,7 @@ class StreamDiffusionCreateStream: self, pipeline: ComfyUIPipelineDTO, scheduler: ComfyUISchedulerDTO, - t_index_list: list[int], + t_index_list: str, width: int, height: int, do_add_noise: bool, @@ -70,7 +70,7 @@ class StreamDiffusionCreateStream: Args: pipeline: Pipeline DTO scheduler: Scheduler DTO - t_index_list: List of timestep indices + t_index_list: JSON array string of timestep indices width: Image width height: Image height do_add_noise: Whether to add noise @@ -89,7 +89,7 @@ class StreamDiffusionCreateStream: scheduler_domain = ComfyUISchedulerDTO.to_domain(scheduler) lcm_lora_domain = ComfyUILcmLoraDTO.to_domain(lcm_lora) - t_index_list_vo = TIndexList(indices=t_index_list) + t_index_list_vo = TIndexList(raw_value=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))