From f97bf1311dd7fe4eb7b3655ebf35b566adfb031e Mon Sep 17 00:00:00 2001
From: Bingsu <ks2515@naver.com>
Date: Thu, 27 Apr 2023 00:17:24 +0900
Subject: [PATCH] feat: save preview

---
 pyproject.toml        |  4 ++--
 scripts/!adetailer.py | 25 ++++++++++++++++++++++++-
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index f09b2fc..84ac568 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -13,11 +13,11 @@ repository = "https://github.com/Bing-su/adetailer"
 
 [tool.isort]
 profile = "black"
-known_first_party = ["modules", "launch"]
+known_first_party = ["launch", "modules"]
 
 [tool.ruff]
 select = ["A", "B", "C4", "E", "F", "I001", "ISC", "N", "PIE", "PT", "RET", "SIM", "UP", "W"]
 ignore = ["B008", "B905", "E501", "UP007", "F401"]
 
 [tool.ruff.isort]
-known-first-party = ["modules", "launch"]
+known-first-party = ["launch", "modules"]
diff --git a/scripts/!adetailer.py b/scripts/!adetailer.py
index 1603b72..f835e46 100644
--- a/scripts/!adetailer.py
+++ b/scripts/!adetailer.py
@@ -8,7 +8,7 @@ import torch
 import modules
 from adetailer import __version__, get_models, mediapipe_predict, ultralytics_predict
 from adetailer.common import dilate_erode, is_all_black, offset
-from modules import devices, scripts
+from modules import devices, images, script_callbacks, scripts, shared
 from modules.paths import data_path, models_path
 from modules.processing import StableDiffusionProcessingImg2Img, process_images
 from modules.safe import load, unsafe_torch_load
@@ -367,6 +367,18 @@ class AfterDetailerScript(scripts.Script):
         steps = len(masks)
         processed = None
 
+        if opts.data.get("ad_save_previews", False):
+            images.save_image(
+                pred.preview,
+                p.outpath_samples,
+                "",
+                seed,
+                p.all_prompts[i],
+                opts.samples_format,
+                p=p,
+                suffix="-ad-preview",
+            )
+
         for j in range(steps):
             mask = masks[j]
             mask = dilate_erode(mask, args.ad_dilate_erode)
@@ -387,3 +399,14 @@ class AfterDetailerScript(scripts.Script):
         if original_params:
             with params_txt.open("w", encoding="utf-8") as f:
                 f.write(original_params)
+
+
+def on_ui_settings():
+    section = ("ADetailer", AFTER_DETAILER)
+    shared.opts.add_option(
+        "ad_save_previews",
+        shared.OptionInfo(False, "Save mask previews", section=section),
+    )
+
+
+script_callbacks.on_ui_settings(on_ui_settings)