refactor: warmup

src/domain/repositories/_stream_diffusion_repository.py

@@ -58,23 +58,36 @@ class StreamDiffusionRepository(ABC):
     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
+            warmup_count: Number of warmup iterations
+        """
+
+    @abstractmethod
+    def prepare_stream(
+        self,
+        stream: Any,
+        prompt: str,
+        negative_prompt: str,
+        steps: Steps,
+        cfg: CFGScale,
+        delta: Delta,
+        seed: Seed,
+    ) -> None:
+        """Prepare the stream with given parameters.
+
+        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
-            warmup_count: Number of warmup iterations
         """
 
     @abstractmethod
@@ -106,28 +119,14 @@ class StreamDiffusionRepository(ABC):
     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:

src/infra/diffusers/_stream_diffusion_repository.py

@@ -91,16 +91,25 @@ class DiffusersStreamDiffusionRepository(StreamDiffusionRepository):
     def warmup_stream(
         self,
         stream: Any,
+        warmup_count: WarmupCount,
+    ) -> None:
+        """Warm up the stream with given parameters."""
+        for _ in range(warmup_count.value):
+            stream()
+
+    def prepare_stream(
+        self,
+        stream: Any,
+        prompt: str,
         negative_prompt: str,
         steps: Steps,
         cfg: CFGScale,
         delta: Delta,
         seed: Seed,
-        warmup_count: WarmupCount,
     ) -> None:
-        """Warm up the stream with given parameters."""
+        """Prepare the stream with given parameters."""
         stream.prepare(
-            prompt="",
+            prompt=prompt,
             negative_prompt=negative_prompt,
             num_inference_steps=steps.value,
             guidance_scale=cfg.value,
@@ -108,9 +117,6 @@ class DiffusersStreamDiffusionRepository(StreamDiffusionRepository):
             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)
@@ -131,37 +137,16 @@ class DiffusersStreamDiffusionRepository(StreamDiffusionRepository):
     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] = []
         if resized_images is None:

src/usecase/_stream_diffusion_sample_usecase.py

@@ -48,12 +48,8 @@ class StreamDiffusionSampleUsecase:
         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(
+        # Prepare stream with parameters
+        self.stream_diffusion_repo.prepare_stream(
             stream=stream.stream,
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -61,11 +57,31 @@ class StreamDiffusionSampleUsecase:
             cfg=cfg,
             delta=delta,
             seed=seed,
-            num_samples=num_samples,
-            warmup_count=warmup_count,
-            input_images=pil_images,
         )
 
+        # Warmup
+        self.stream_diffusion_repo.warmup_stream(
+            stream=stream.stream,
+            warmup_count=warmup_count,
+        )
+
+        # 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]
+
+        if pil_images:
+            images = self.stream_diffusion_repo.sample_with_images(
+                stream=stream.stream,
+                num_samples=num_samples,
+                input_images=pil_images,
+            )
+        else:
+            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")
 

src/usecase/_stream_diffusion_warmup_usecase.py

@@ -30,13 +30,20 @@ class StreamDiffusionWarmupUsecase:
         Returns:
             Warmed up StreamDiffusionStream (same instance)
         """
-        self.stream_diffusion_repo.warmup_stream(
+        # Prepare stream with parameters
+        self.stream_diffusion_repo.prepare_stream(
             stream=stream.stream,
+            prompt="",
             negative_prompt=negative_prompt,
             steps=steps,
             cfg=cfg,
             delta=delta,
             seed=seed,
+        )
+
+        # Warmup
+        self.stream_diffusion_repo.warmup_stream(
+            stream=stream.stream,
             warmup_count=warmup_count,
         )