a little bit of rework for extras tab

This commit is contained in:
AUTOMATIC 2022-09-07 13:35:02 +03:00
parent 5e12c23a67
commit 8dcd3e57e7
3 changed files with 31 additions and 11 deletions

View File

@ -16,6 +16,7 @@ from importlib import reload
pretrain_model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
have_codeformer = False
codeformer = None
def setup_codeformer():
path = modules.paths.paths.get("CodeFormer", None)
@ -64,7 +65,7 @@ def setup_codeformer():
return net, face_helper
def restore(self, np_image):
def restore(self, np_image, w=None):
np_image = np_image[:, :, ::-1]
net, face_helper = self.create_models()
@ -80,7 +81,7 @@ def setup_codeformer():
try:
with torch.no_grad():
output = net(cropped_face_t, w=shared.opts.code_former_weight, adain=True)[0]
output = net(cropped_face_t, w=w if w is not None else shared.opts.code_former_weight, adain=True)[0]
restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
del output
torch.cuda.empty_cache()
@ -99,7 +100,10 @@ def setup_codeformer():
global have_codeformer
have_codeformer = True
shared.face_restorers.append(FaceRestorerCodeFormer())
global codeformer
codeformer = FaceRestorerCodeFormer()
shared.face_restorers.append(codeformer)
except Exception:
print("Error setting up CodeFormer:", file=sys.stderr)

View File

@ -21,9 +21,10 @@ from modules.paths import script_path
from modules.shared import opts, cmd_opts
import modules.shared as shared
from modules.sd_samplers import samplers, samplers_for_img2img
import modules.gfpgan_model as gfpgan
import modules.realesrgan_model as realesrgan
import modules.scripts
import modules.gfpgan_model
import modules.codeformer_model
# this is a fix for Windows users. Without it, javascript files will be served with text/html content-type and the bowser will not show any UI
mimetypes.init()
@ -521,7 +522,11 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
extras_upscaler_2_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Upscaler 2 visibility", value=1)
with gr.Group():
face_restoration_blending = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Faces restoration visibility", value=0, interactive=len(shared.face_restorers) > 1)
gfpgan_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="GFPGAN visibility", value=0, interactive=modules.gfpgan_model.have_gfpgan)
with gr.Group():
codeformer_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="CodeFormer visibility", value=0, interactive=modules.codeformer_model.have_codeformer)
codeformer_weight = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="CodeFormer weight (0=max, 1=none)", value=0, interactive=modules.codeformer_model.have_codeformer)
submit = gr.Button('Generate', elem_id="extras_generate", variant='primary')
@ -534,7 +539,9 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
fn=run_extras,
inputs=[
image,
face_restoration_blending,
gfpgan_visibility,
codeformer_visibility,
codeformer_weight,
upscaling_resize,
extras_upscaler_1,
extras_upscaler_2,

View File

@ -58,19 +58,28 @@ def load_model_from_config(config, ckpt, verbose=False):
cached_images = {}
def run_extras(image, face_restoration_blending, upscaling_resize, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):
def run_extras(image, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):
processing.torch_gc()
image = image.convert("RGB")
outpath = opts.outdir_samples or opts.outdir_extras_samples
if face_restoration_blending > 0:
restored_img = modules.face_restoration.restore_faces(np.array(image, dtype=np.uint8))
if gfpgan_visibility > 0:
restored_img = modules.gfpgan_model.gfpgan_fix_faces(np.array(image, dtype=np.uint8))
res = Image.fromarray(restored_img)
if face_restoration_blending < 1.0:
res = Image.blend(image, res, face_restoration_blending)
if gfpgan_visibility < 1.0:
res = Image.blend(image, res, gfpgan_visibility)
image = res
if codeformer_visibility > 0:
restored_img = modules.codeformer_model.codeformer.restore(np.array(image, dtype=np.uint8), w=codeformer_weight)
res = Image.fromarray(restored_img)
if codeformer_visibility < 1.0:
res = Image.blend(image, res, codeformer_visibility)
image = res