fix: warmup input image

src/domain/repositories/_stream_diffusion_repository.py

@@ -59,12 +59,14 @@ class StreamDiffusionRepository(ABC):
         self,
         stream: Any,
         warmup_count: WarmupCount,
+        input_image: Image.Image | None = None,
     ) -> None:
         """Warm up the stream with given parameters.
 
         Args:
             stream: StreamDiffusion instance
             warmup_count: Number of warmup iterations
+            input_image: Optional input image for img2img warmup
         """
 
     @abstractmethod

src/infra/diffusers/_stream_diffusion_repository.py

@@ -92,10 +92,17 @@ class DiffusersStreamDiffusionRepository(StreamDiffusionRepository):
         self,
         stream: Any,
         warmup_count: WarmupCount,
+        input_image: Image.Image | None = None,
     ) -> None:
         """Warm up the stream with given parameters."""
-        for _ in range(warmup_count.value):
-            stream()
+        if input_image is not None:
+            # Resize input image to match stream dimensions
+            resized_image = input_image.resize((stream.width, stream.height))
+            for _ in range(warmup_count.value):
+                stream(resized_image)
+        else:
+            for _ in range(warmup_count.value):
+                stream()
 
     def prepare_stream(
         self,
@@ -114,6 +121,7 @@ class DiffusersStreamDiffusionRepository(StreamDiffusionRepository):
             num_inference_steps=steps.value,
             guidance_scale=cfg.value,
             delta=delta.value,
+            generator=torch.Generator().manual_seed(seed.value),
             seed=seed.value,
         )
 

src/usecase/_stream_diffusion_sample_usecase.py

@@ -59,17 +59,20 @@ class StreamDiffusionSampleUsecase:
             seed=seed,
         )
 
-        # Warmup
+        # Convert domain images to PIL images if provided
+        pil_images = None
+        warmup_image = None
+        if input_images is not None:
+            pil_images = [img.image for img in input_images]
+            # Use first image for warmup (img2img mode)
+            warmup_image = pil_images[0] if pil_images else None
+
         self.stream_diffusion_repo.warmup_stream(
             stream=stream.stream,
             warmup_count=warmup_count,
+            input_image=warmup_image,
         )
 
-        # 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,

src/usecase/_stream_diffusion_warmup_usecase.py

@@ -41,10 +41,11 @@ class StreamDiffusionWarmupUsecase:
             seed=seed,
         )
 
-        # Warmup
+        # Warmup (txt2img mode - no input image)
         self.stream_diffusion_repo.warmup_stream(
             stream=stream.stream,
             warmup_count=warmup_count,
+            input_image=None,
         )
 
         return stream