From a3e6c284db14af43c7ac3466e00e24cc2537a1ca Mon Sep 17 00:00:00 2001 From: Limitex <76650151+Limitex@users.noreply.github.com> Date: Mon, 24 Nov 2025 01:02:01 +0900 Subject: [PATCH] refactor: warmup --- .../_stream_diffusion_repository.py | 39 +++++++++---------- .../diffusers/_stream_diffusion_repository.py | 39 ++++++------------- .../_stream_diffusion_sample_usecase.py | 34 +++++++++++----- .../_stream_diffusion_warmup_usecase.py | 9 ++++- 4 files changed, 64 insertions(+), 57 deletions(-) diff --git a/src/domain/repositories/_stream_diffusion_repository.py b/src/domain/repositories/_stream_diffusion_repository.py index 26bcefa..7f54e27 100644 --- a/src/domain/repositories/_stream_diffusion_repository.py +++ b/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: diff --git a/src/infra/diffusers/_stream_diffusion_repository.py b/src/infra/diffusers/_stream_diffusion_repository.py index 222e651..627d13f 100644 --- a/src/infra/diffusers/_stream_diffusion_repository.py +++ b/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: diff --git a/src/usecase/_stream_diffusion_sample_usecase.py b/src/usecase/_stream_diffusion_sample_usecase.py index 5736c77..d5d1fec 100644 --- a/src/usecase/_stream_diffusion_sample_usecase.py +++ b/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") diff --git a/src/usecase/_stream_diffusion_warmup_usecase.py b/src/usecase/_stream_diffusion_warmup_usecase.py index 9bdebd4..a93451f 100644 --- a/src/usecase/_stream_diffusion_warmup_usecase.py +++ b/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, )