refactor: any

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:

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]:

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]:

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:

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":