From 008f3310a5d096363bd895d35f4248f3ef38dd40 Mon Sep 17 00:00:00 2001
From: Limitex <76650151+Limitex@users.noreply.github.com>
Date: Sun, 23 Nov 2025 00:00:02 +0900
Subject: [PATCH] refactor: simplify image generation loop in StreamDiffusion

---
 .../diffusers/_stream_diffusion_repository.py | 21 ++++++++++---------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/infra/diffusers/_stream_diffusion_repository.py b/src/infra/diffusers/_stream_diffusion_repository.py
index 048d13b..77e0e53 100644
--- a/src/infra/diffusers/_stream_diffusion_repository.py
+++ b/src/infra/diffusers/_stream_diffusion_repository.py
@@ -156,17 +156,18 @@ class DiffusersStreamDiffusionRepository(StreamDiffusionRepository):
 
         # Generate images
         result: list[Image.Image] = []
-        for _ in range(num_samples.value):
-            x_outputs: list[Any] = []
-            if resized_images is None:
-                x_outputs.append(stream.txt2img())
-            else:
-                stream(resized_images[0])
-                for img in resized_images[1:] + resized_images[-1:]:
-                    x_outputs.append(stream(img))
-
-            for x_output in x_outputs:
+        if resized_images is None:
+            # Text-to-image: Generate num_samples images
+            for _ in range(num_samples.value):
+                x_output = stream.txt2img()
                 image = postprocess_image(x_output, output_type="pil")[0]
                 result.append(image)
+        else:
+            # Image-to-image: Process each input image num_samples times
+            for _ in range(num_samples.value):
+                for img in resized_images:
+                    x_output = stream(img)
+                    image = postprocess_image(x_output, output_type="pil")[0]
+                    result.append(image)
 
         return result