From 6a9b33c848281cb02f38764e4f91ef767f5e3edd Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Wed, 7 Sep 2022 12:32:28 +0300
Subject: codeformer support

---
 modules/shared.py | 6 ++++++
 1 file changed, 6 insertions(+)

(limited to 'modules/shared.py')

diff --git a/modules/shared.py b/modules/shared.py
index fb5bbd43..12082895 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -1,11 +1,13 @@
 import argparse
 import json
 import os
+
 import gradio as gr
 import torch
 
 import modules.artists
 from modules.paths import script_path, sd_path
+import modules.codeformer_model
 
 config_filename = "config.json"
 
@@ -40,6 +42,7 @@ device = gpu if torch.cuda.is_available() else cpu
 batch_cond_uncond = cmd_opts.always_batch_cond_uncond or not (cmd_opts.lowvram or cmd_opts.medvram)
 parallel_processing_allowed = not cmd_opts.lowvram and not cmd_opts.medvram
 
+
 class State:
     interrupted = False
     job = ""
@@ -65,6 +68,7 @@ state = State()
 
 artist_db = modules.artists.ArtistsDatabase(os.path.join(script_path, 'artists.csv'))
 
+face_restorers = []
 
 def find_any_font():
     fonts = ['/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf']
@@ -116,6 +120,8 @@ class Options:
         "upscale_at_full_resolution_padding": OptionInfo(16, "Inpainting at full resolution: padding, in pixels, for the masked region.", gr.Slider, {"minimum": 0, "maximum": 128, "step": 4}),
         "show_progressbar": OptionInfo(True, "Show progressbar"),
         "show_progress_every_n_steps": OptionInfo(0, "Show show image creation progress every N sampling steps. Set 0 to disable.", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}),
+        "face_restoration_model": OptionInfo(None, "Face restoration model", gr.Radio, lambda: {"choices": [x.name() for x in face_restorers]}),
+        "code_former_weight": OptionInfo(0.5, "CodeFormer weight parameter", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),
     }
 
     def __init__(self):
-- 
cgit v1.2.1