diff --git a/src/domain/repositories/_stream_diffusion_repository.py b/src/domain/repositories/_stream_diffusion_repository.py index 7f54e27..63a2fee 100644 --- a/src/domain/repositories/_stream_diffusion_repository.py +++ b/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 diff --git a/src/infra/diffusers/_stream_diffusion_repository.py b/src/infra/diffusers/_stream_diffusion_repository.py index 627d13f..6db20df 100644 --- a/src/infra/diffusers/_stream_diffusion_repository.py +++ b/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, ) diff --git a/src/usecase/_stream_diffusion_sample_usecase.py b/src/usecase/_stream_diffusion_sample_usecase.py index d5d1fec..cbb8631 100644 --- a/src/usecase/_stream_diffusion_sample_usecase.py +++ b/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, diff --git a/src/usecase/_stream_diffusion_warmup_usecase.py b/src/usecase/_stream_diffusion_warmup_usecase.py index a93451f..0dcf60a 100644 --- a/src/usecase/_stream_diffusion_warmup_usecase.py +++ b/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