From fd2e8912a46107d33b29cc442967b67821afe460 Mon Sep 17 00:00:00 2001
From: hodanov <1031hoda@gmail.com>
Date: Sun, 3 Nov 2024 19:31:31 +0900
Subject: [PATCH] Rename util.py to domain.py. Refactor some functions.

---
 cmd/domain.py | 123 ++++++++++++++++++++++++++++++++++++++++++++++++++
 cmd/util.py   |  55 ----------------------
 2 files changed, 123 insertions(+), 55 deletions(-)
 create mode 100644 cmd/domain.py
 delete mode 100644 cmd/util.py

diff --git a/cmd/domain.py b/cmd/domain.py
new file mode 100644
index 0000000..ec6c9ee
--- /dev/null
+++ b/cmd/domain.py
@@ -0,0 +1,123 @@
+"""Utility functions for the script."""
+
+from __future__ import annotations
+
+import secrets
+import time
+from datetime import date
+from pathlib import Path
+
+
+class Seed:
+    def __init__(self, seed: int) -> None:
+        if seed != -1:
+            self.__value = seed
+            return
+
+        self.__value = self.__generate_seed()
+
+    def __generate_seed(self) -> int:
+        max_limit_value = 4294967295
+        return secrets.randbelow(max_limit_value)
+
+    @property
+    def value(self) -> int:
+        return self.__value
+
+
+class Prompts:
+    def __init__(
+        self,
+        prompt: str,
+        n_prompt: str,
+        height: int,
+        width: int,
+        samples: int,
+        steps: int,
+    ) -> None:
+        if prompt == "":
+            msg = "prompt should not be empty."
+            raise ValueError(msg)
+
+        if n_prompt == "":
+            msg = "n_prompt should not be empty."
+            raise ValueError(msg)
+
+        if height <= 0:
+            msg = "height should be positive."
+            raise ValueError(msg)
+
+        if width <= 0:
+            msg = "width should be positive."
+            raise ValueError(msg)
+
+        if samples <= 0:
+            msg = "samples should be positive."
+            raise ValueError(msg)
+
+        if steps <= 0:
+            msg = "steps should be positive."
+            raise ValueError(msg)
+
+        self.__dict: dict[str, int | str] = {
+            "prompt": prompt,
+            "n_prompt": n_prompt,
+            "height": height,
+            "width": width,
+            "samples": samples,
+            "steps": steps,
+        }
+
+    @property
+    def dict(self) -> dict[str, int | str]:
+        return self.__dict
+
+
+class OutputDirectory:
+    def __init__(self) -> None:
+        self.__output_directory_name = "outputs"
+        self.__date_today = date.today().strftime("%Y-%m-%d")
+        self.__make_path()
+
+    def __make_path(self) -> None:
+        self.__path = Path(f"{self.__output_directory_name}/{self.__date_today}")
+
+    def make_directory(self) -> Path:
+        """Make a directory for saving outputs."""
+        if not self.__path.exists():
+            self.__path.mkdir(exist_ok=True, parents=True)
+
+        return self.__path
+
+
+class StableDiffusionOutputManger:
+    def __init__(self, prompts: Prompts, output_directory: Path) -> None:
+        self.__prompts = prompts
+        self.__output_directory = output_directory
+
+    def save_prompts(self) -> str:
+        """Save prompts to a file."""
+        prompts_filename = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
+        output_path = f"{self.__output_directory}/prompts_{prompts_filename}.txt"
+        with Path(output_path).open("wb") as file:
+            for name, value in self.__prompts.dict.items():
+                file.write(f"{name} = {value!r}\n".encode())
+
+        return output_path
+
+    def save_image(
+        self,
+        image: bytes,
+        seed: int,
+        i: int,
+        j: int,
+        output_format: str = "png",
+    ) -> str:
+        """Save image to a file."""
+        formatted_time = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
+        filename = f"{formatted_time}_{seed}_{i}_{j}.{output_format}"
+        output_path = f"{self.__output_directory}/{filename}"
+        with Path(output_path).open("wb") as file:
+            file.write(image)
+
+        return output_path
diff --git a/cmd/util.py b/cmd/util.py
deleted file mode 100644
index 643be2e..0000000
--- a/cmd/util.py
+++ /dev/null
@@ -1,55 +0,0 @@
-""" Utility functions for the script. """
-import random
-import time
-from datetime import date
-from pathlib import Path
-
-OUTPUT_DIRECTORY = "outputs"
-DATE_TODAY = date.today().strftime("%Y-%m-%d")
-
-
-def generate_seed() -> int:
-    """
-    Generate a random seed.
-    """
-    seed = random.randint(0, 4294967295)
-    print(f"Generate a random seed: {seed}")
-
-    return seed
-
-
-def make_directory() -> Path:
-    """
-    Make a directory for saving outputs.
-    """
-    directory = Path(f"{OUTPUT_DIRECTORY}/{DATE_TODAY}")
-    if not directory.exists():
-        directory.mkdir(exist_ok=True, parents=True)
-        print(f"Make a directory: {directory}")
-
-    return directory
-
-
-def save_prompts(inputs: dict):
-    """
-    Save prompts to a file.
-    """
-    prompts_filename = time.strftime("%Y%m%d%H%M%S", time.localtime(time.time()))
-    with open(
-        file=f"{OUTPUT_DIRECTORY}/{DATE_TODAY}/prompts_{prompts_filename}.txt", mode="w", encoding="utf-8"
-    ) as file:
-        for name, value in inputs.items():
-            file.write(f"{name} = {repr(value)}\n")
-        print(f"Save prompts: {prompts_filename}.txt")
-
-
-def save_images(directory: Path, images: list[bytes], seed: int, i: int, output_format: str = "png"):
-    """
-    Save images to a file.
-    """
-    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}_{seed}_{i}_{j}.{output_format}"
-        print(f"Saving it to {output_path}")
-        with open(output_path, "wb") as file:
-            file.write(image_bytes)