From 053b142a87cd7c0e3b2db7abfa8bdc22cc8322e6 Mon Sep 17 00:00:00 2001
From: hodanov <1031hoda@gmail.com>
Date: Mon, 5 Jun 2023 10:16:09 +0900
Subject: [PATCH] Refactor some codes.

---
 sd_cli.py | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/sd_cli.py b/sd_cli.py
index 7263ccc..1275609 100644
--- a/sd_cli.py
+++ b/sd_cli.py
@@ -162,18 +162,13 @@ def entrypoint(
 
     inputs["max_embeddings_multiples"] = util.count_token(p=prompt, n=n_prompt)
     directory = util.make_directory()
-    util.save_prompts(inputs)
 
     sd = StableDiffusion()
     for i in range(samples):
         start_time = time.time()
         images = sd.run_inference.call(inputs)
-        for j, image_bytes in enumerate(images):
-            formatted_time = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
-            output_path = directory / f"{formatted_time}_{i}_{j}.png"
-            print(f"Saving it to {output_path}")
-            with open(output_path, "wb") as file:
-                file.write(image_bytes)
-
+        util.save_images(directory, images, i)
         total_time = time.time() - start_time
         print(f"Sample {i} took {total_time:.3f}s ({(total_time)/len(images):.3f}s / image).")
+
+    util.save_prompts(inputs)