From ef8e613b01ab56f463d716a2493c4275ba3e443e Mon Sep 17 00:00:00 2001
From: hodanov <1031hoda@gmail.com>
Date: Wed, 28 Jun 2023 20:18:56 +0900
Subject: [PATCH] Modify a directory structure.

---
 Makefile                          | 13 +++++++++++--
 sdcli/__init__.py                 |  0
 entrypoint.py => sdcli/txt2img.py |  6 ++++++
 util.py => sdcli/util.py          |  0
 setup.py                          | 25 +++++++++++++++----------
 5 files changed, 32 insertions(+), 12 deletions(-)
 create mode 100644 sdcli/__init__.py
 rename entrypoint.py => sdcli/txt2img.py (86%)
 rename util.py => sdcli/util.py (100%)

diff --git a/Makefile b/Makefile
index 6709073..67716ae 100644
--- a/Makefile
+++ b/Makefile
@@ -1,11 +1,20 @@
 deploy:
 	modal deploy setup.py
 
+# `--upscaler` is a name of upscaler you want to use.
+# You can use upscalers the below:
+#   - `RealESRGAN_x4plus`
+#   - `RealESRNet_x4plus`
+#   - `RealESRGAN_x4plus_anime_6B`
+#   - `RealESRGAN_x2plus`
 run:
-	modal run entrypoint.py \
+	cd ./sdcli && modal run txt2img.py \
 	--prompt "a photograph of an astronaut riding a horse" \
 	--n-prompt "" \
 	--height 512 \
 	--width 512 \
 	--samples 1 \
-	--steps 50
+	--steps 50 \
+	--upscaler "" \
+	--use-face-enhancer "False" \
+	--use-hires-fix "False"
diff --git a/sdcli/__init__.py b/sdcli/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/entrypoint.py b/sdcli/txt2img.py
similarity index 86%
rename from entrypoint.py
rename to sdcli/txt2img.py
index 2763317..46d37c1 100644
--- a/entrypoint.py
+++ b/sdcli/txt2img.py
@@ -16,6 +16,9 @@ def main(
     batch_size: int = 1,
     steps: int = 20,
     seed: int = -1,
+    upscaler: str = "",
+    use_face_enhancer: str = "False",
+    use_hires_fix: str = "False",
 ):
     """
     This function is the entrypoint for the Runway CLI.
@@ -38,6 +41,9 @@ def main(
             batch_size=batch_size,
             steps=steps,
             seed=seed_generated,
+            upscaler=upscaler,
+            use_face_enhancer=use_face_enhancer == "True",
+            use_hires_fix=use_hires_fix == "True",
         )
         util.save_images(directory, images, seed_generated, i)
         total_time = time.time() - start_time
diff --git a/util.py b/sdcli/util.py
similarity index 100%
rename from util.py
rename to sdcli/util.py
diff --git a/setup.py b/setup.py
index 406d2c3..33f905b 100644
--- a/setup.py
+++ b/setup.py
@@ -97,10 +97,6 @@ class StableDiffusion(ClsMixin):
         import diffusers
         import torch
 
-        self.use_vae = os.environ["USE_VAE"] == "true"
-        self.upscaler = os.environ["UPSCALER"]
-        self.use_face_enhancer = os.environ["USE_FACE_ENHANCER"] == "true"
-        self.use_hires_fix = os.environ["USE_HIRES_FIX"] == "true"
         self.cache_path = os.path.join(BASE_CACHE_PATH, os.environ["MODEL_NAME"])
         if os.path.exists(self.cache_path):
             print(f"The directory '{self.cache_path}' exists.")
@@ -123,7 +119,7 @@ class StableDiffusion(ClsMixin):
             subfolder="scheduler",
         )
 
-        if self.use_vae:
+        if os.environ["USE_VAE"] == "true":
             self.pipe.vae = diffusers.AutoencoderKL.from_pretrained(
                 self.cache_path,
                 subfolder="vae",
@@ -194,6 +190,9 @@ class StableDiffusion(ClsMixin):
         batch_size: int = 1,
         steps: int = 30,
         seed: int = 1,
+        upscaler: str = "",
+        use_face_enhancer: bool = False,
+        use_hires_fix: bool = False,
     ) -> list[bytes]:
         """
         Runs the Stable Diffusion pipeline on the given prompt and outputs images.
@@ -215,14 +214,17 @@ class StableDiffusion(ClsMixin):
                     generator=generator,
                 ).images
 
-        if self.upscaler != "":
+        if upscaler != "":
             upscaled = self.upscale(
                 base_images=base_images,
                 half_precision=False,
                 tile=700,
+                upscaler=upscaler,
+                use_face_enhancer=use_face_enhancer,
+                use_hires_fix=use_hires_fix,
             )
             base_images.extend(upscaled)
-            if self.use_hires_fix:
+            if use_hires_fix:
                 torch.cuda.empty_cache()
                 for img in upscaled:
                     with torch.inference_mode():
@@ -256,6 +258,9 @@ class StableDiffusion(ClsMixin):
         tile: int = 0,
         tile_pad: int = 10,
         pre_pad: int = 0,
+        upscaler: str = "",
+        use_face_enhancer: bool = False,
+        use_hires_fix: bool = False,
     ) -> list[Image.Image]:
         """
         Upscales the given images using the given model.
@@ -268,7 +273,7 @@ class StableDiffusion(ClsMixin):
         from realesrgan import RealESRGANer
         from tqdm import tqdm
 
-        model_name = self.upscaler
+        model_name = upscaler
         if model_name == "RealESRGAN_x4plus":
             upscale_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
             netscale = 4
@@ -298,7 +303,7 @@ class StableDiffusion(ClsMixin):
 
         from gfpgan import GFPGANer
 
-        if self.use_face_enhancer:
+        if use_face_enhancer:
             face_enhancer = GFPGANer(
                 model_path=os.path.join(BASE_CACHE_PATH, "esrgan", "GFPGANv1.3.pth"),
                 upscale=netscale,
@@ -312,7 +317,7 @@ class StableDiffusion(ClsMixin):
         with tqdm(total=len(base_images)) as progress_bar:
             for img in base_images:
                 img = numpy.array(img)
-                if self.use_face_enhancer:
+                if use_face_enhancer:
                     _, _, enhance_result = face_enhancer.enhance(
                         img,
                         has_aligned=False,