diff --git a/src/domain/model/_stream_diffusion_stream.py b/src/domain/model/_stream_diffusion_stream.py index 2db40c1..4c39b44 100644 --- a/src/domain/model/_stream_diffusion_stream.py +++ b/src/domain/model/_stream_diffusion_stream.py @@ -1,12 +1,13 @@ from dataclasses import dataclass -from typing import Any + +from streamdiffusion import StreamDiffusion # type: ignore[import-untyped] @dataclass(frozen=True) class StreamDiffusionStream: """Stream Diffusion stream instance.""" - stream: Any # StreamDiffusion instance from streamdiffusion library + stream: StreamDiffusion def __post_init__(self) -> None: if self.stream is None: diff --git a/src/domain/repositories/_stream_diffusion_repository.py b/src/domain/repositories/_stream_diffusion_repository.py index 63a2fee..7b34b63 100644 --- a/src/domain/repositories/_stream_diffusion_repository.py +++ b/src/domain/repositories/_stream_diffusion_repository.py @@ -1,9 +1,10 @@ from abc import ABC, abstractmethod -from typing import Any +import torch from diffusers import StableDiffusionPipeline from diffusers.schedulers.scheduling_utils import SchedulerMixin from PIL import Image +from streamdiffusion import StreamDiffusion # type: ignore[import-untyped] from ..model import ( CFGScale, @@ -31,10 +32,10 @@ class StreamDiffusionRepository(ABC): use_denoising_batch: bool, frame_buffer_size: FrameBufferSize, cfg_type: CFGType, - lcm_lora_weights: dict[str, Any], + lcm_lora_weights: dict[str, torch.Tensor], tiny_vae_name: str, enable_xformers: bool, - ) -> Any: + ) -> StreamDiffusion: """Create a Stream Diffusion stream instance. Args: @@ -57,7 +58,7 @@ class StreamDiffusionRepository(ABC): @abstractmethod def warmup_stream( self, - stream: Any, + stream: StreamDiffusion, warmup_count: WarmupCount, input_image: Image.Image | None = None, ) -> None: @@ -72,7 +73,7 @@ class StreamDiffusionRepository(ABC): @abstractmethod def prepare_stream( self, - stream: Any, + stream: StreamDiffusion, prompt: str, negative_prompt: str, steps: Steps, @@ -93,7 +94,7 @@ class StreamDiffusionRepository(ABC): """ @abstractmethod - def update_prompt(self, stream: Any, prompt: str) -> None: + def update_prompt(self, stream: StreamDiffusion, prompt: str) -> None: """Update the prompt for the stream. Args: @@ -104,7 +105,7 @@ class StreamDiffusionRepository(ABC): @abstractmethod def sample_txt2img( self, - stream: Any, + stream: StreamDiffusion, num_samples: NumSamples, ) -> list[Image.Image]: """Generate images using txt2img. @@ -120,7 +121,7 @@ class StreamDiffusionRepository(ABC): @abstractmethod def sample_with_images( self, - stream: Any, + stream: StreamDiffusion, num_samples: NumSamples, input_images: list[Image.Image] | None, ) -> list[Image.Image]: diff --git a/src/infra/diffusers/_stream_diffusion_repository.py b/src/infra/diffusers/_stream_diffusion_repository.py index 6db20df..33224b9 100644 --- a/src/infra/diffusers/_stream_diffusion_repository.py +++ b/src/infra/diffusers/_stream_diffusion_repository.py @@ -1,5 +1,4 @@ import copy -from typing import Any import torch from comfy.model_management import get_torch_device # pyright: ignore[reportMissingImports] @@ -41,10 +40,10 @@ class DiffusersStreamDiffusionRepository(StreamDiffusionRepository): use_denoising_batch: bool, frame_buffer_size: FrameBufferSize, cfg_type: CFGType, - lcm_lora_weights: dict[str, Any], + lcm_lora_weights: dict[str, torch.Tensor], tiny_vae_name: str, enable_xformers: bool, - ) -> Any: + ) -> StreamDiffusion: """Create a Stream Diffusion stream instance.""" # Deep copy to avoid modifying the original pipeline # Note: load_lcm_lora() and fuse_lora() modify the pipeline, @@ -90,7 +89,7 @@ class DiffusersStreamDiffusionRepository(StreamDiffusionRepository): def warmup_stream( self, - stream: Any, + stream: StreamDiffusion, warmup_count: WarmupCount, input_image: Image.Image | None = None, ) -> None: @@ -106,7 +105,7 @@ class DiffusersStreamDiffusionRepository(StreamDiffusionRepository): def prepare_stream( self, - stream: Any, + stream: StreamDiffusion, prompt: str, negative_prompt: str, steps: Steps, @@ -125,13 +124,13 @@ class DiffusersStreamDiffusionRepository(StreamDiffusionRepository): seed=seed.value, ) - def update_prompt(self, stream: Any, prompt: str) -> None: + def update_prompt(self, stream: StreamDiffusion, prompt: str) -> None: """Update the prompt for the stream.""" stream.update_prompt(prompt) def sample_txt2img( self, - stream: Any, + stream: StreamDiffusion, num_samples: NumSamples, ) -> list[Image.Image]: """Generate images using txt2img.""" @@ -144,7 +143,7 @@ class DiffusersStreamDiffusionRepository(StreamDiffusionRepository): def sample_with_images( self, - stream: Any, + stream: StreamDiffusion, num_samples: NumSamples, input_images: list[Image.Image] | None, ) -> list[Image.Image]: diff --git a/src/infra/diffusers/_text_encoder_repository.py b/src/infra/diffusers/_text_encoder_repository.py index 0457d6d..47ba918 100644 --- a/src/infra/diffusers/_text_encoder_repository.py +++ b/src/infra/diffusers/_text_encoder_repository.py @@ -1,7 +1,5 @@ from __future__ import annotations -from typing import Any - import torch from comfy.model_management import get_torch_device # pyright: ignore[reportMissingImports] from transformers import CLIPTextModel, CLIPTokenizer @@ -35,7 +33,7 @@ class DiffusersTextEncoderRepository(TextEncoderRepository): with torch.no_grad(): for i in range(0, text_ids.shape[-1], max_length): segment_ids = text_ids[:, i : i + max_length] - outputs: Any = text_encoder(segment_ids) + outputs = text_encoder(segment_ids) if hasattr(outputs, "last_hidden_state"): embeds = outputs.last_hidden_state else: diff --git a/src/nodes/dto/_stream_diffusion_stream_dto.py b/src/nodes/dto/_stream_diffusion_stream_dto.py index 1ceee2f..8147c8e 100644 --- a/src/nodes/dto/_stream_diffusion_stream_dto.py +++ b/src/nodes/dto/_stream_diffusion_stream_dto.py @@ -1,5 +1,7 @@ from dataclasses import dataclass -from typing import Any, ClassVar +from typing import ClassVar + +from streamdiffusion import StreamDiffusion # type: ignore[import-untyped] from ...domain.model import StreamDiffusionStream @@ -17,7 +19,7 @@ class ComfyUIStreamDTO: COMFY_TYPE: ClassVar[str] = "DIFFUSERS_STREAM" - stream: Any + stream: StreamDiffusion @classmethod def from_domain(cls, stream: StreamDiffusionStream) -> "ComfyUIStreamDTO": @@ -59,7 +61,7 @@ class ComfyUIWarmupStreamDTO: COMFY_TYPE: ClassVar[str] = "DIFFUSERS_WARMUP_STREAM" - stream: Any + stream: StreamDiffusion @classmethod def from_domain(cls, stream: StreamDiffusionStream) -> "ComfyUIWarmupStreamDTO":