diff options
| author | InvincibleDude <81354513+InvincibleDude@users.noreply.github.com> | 2023-01-24 15:44:09 +0300 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-01-24 15:44:09 +0300 |
| commit | 44c0e6b993d00bb2f441f0fde409bcb79136f034 (patch) | |
| tree | e27a45d1a3ceb8aab884631c7a806c5fe2c8386d /scripts/postprocessing_codeformer.py | |
| parent | 3bc8ee998db5f461b8011a72e6f167012ccb8bc1 (diff) | |
| parent | 602a1864b05075ca4283986e6f5c7d5bce864e11 (diff) | |
Merge branch 'AUTOMATIC1111:master' into master
Diffstat (limited to 'scripts/postprocessing_codeformer.py')
| -rw-r--r-- | scripts/postprocessing_codeformer.py | 36 |
1 files changed, 36 insertions, 0 deletions
diff --git a/scripts/postprocessing_codeformer.py b/scripts/postprocessing_codeformer.py new file mode 100644 index 00000000..a7d80d40 --- /dev/null +++ b/scripts/postprocessing_codeformer.py @@ -0,0 +1,36 @@ +from PIL import Image
+import numpy as np
+
+from modules import scripts_postprocessing, codeformer_model
+import gradio as gr
+
+from modules.ui_components import FormRow
+
+
+class ScriptPostprocessingCodeFormer(scripts_postprocessing.ScriptPostprocessing):
+ name = "CodeFormer"
+ order = 3000
+
+ def ui(self):
+ with FormRow():
+ codeformer_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="CodeFormer visibility", value=0, elem_id="extras_codeformer_visibility")
+ codeformer_weight = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="CodeFormer weight (0 = maximum effect, 1 = minimum effect)", value=0, elem_id="extras_codeformer_weight")
+
+ return {
+ "codeformer_visibility": codeformer_visibility,
+ "codeformer_weight": codeformer_weight,
+ }
+
+ def process(self, pp: scripts_postprocessing.PostprocessedImage, codeformer_visibility, codeformer_weight):
+ if codeformer_visibility == 0:
+ return
+
+ restored_img = codeformer_model.codeformer.restore(np.array(pp.image, dtype=np.uint8), w=codeformer_weight)
+ res = Image.fromarray(restored_img)
+
+ if codeformer_visibility < 1.0:
+ res = Image.blend(pp.image, res, codeformer_visibility)
+
+ pp.image = res
+ pp.info["CodeFormer visibility"] = round(codeformer_visibility, 3)
+ pp.info["CodeFormer weight"] = round(codeformer_weight, 3)
|