From 26d08193848568b06105a1ee7b76f338ebf0f0ee Mon Sep 17 00:00:00 2001 From: Chris OBryan <13701027+cobryan05@users.noreply.github.com> Date: Fri, 28 Oct 2022 13:24:11 -0500 Subject: extras: Add option to run upscaling before face fixing Face restoration can look much better if ran after upscaling, as it allows the restoration to fix upscaling artifacts. This patch adds an option to choose which order to run upscaling/face fixing in. --- modules/ui.py | 4 ++++ 1 file changed, 4 insertions(+) (limited to 'modules/ui.py') diff --git a/modules/ui.py b/modules/ui.py index 0a63e357..16b6ac49 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -1119,6 +1119,9 @@ def create_ui(wrap_gradio_gpu_call): 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 = maximum effect, 1 = minimum effect)", value=0, interactive=modules.codeformer_model.have_codeformer) + with gr.Group(): + upscale_before_face_fix = gr.Checkbox(label='Upscale Before Restoring Faces', value=False) + submit = gr.Button('Generate', elem_id="extras_generate", variant='primary') with gr.Column(variant='panel'): @@ -1152,6 +1155,7 @@ def create_ui(wrap_gradio_gpu_call): extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, + upscale_before_face_fix, ], outputs=[ result_images, -- cgit v1.2.3 From 1f1b327959b546b5e6f995905a1699c5fe4a0c35 Mon Sep 17 00:00:00 2001 From: Chris OBryan <13701027+cobryan05@users.noreply.github.com> Date: Fri, 28 Oct 2022 16:11:16 -0500 Subject: extras: Make image cache LRU This changes the extras image cache into a Least-Recently-Used cache. This allows more experimentation with different upscalers without missing the cache. Max cache size is increased to 5 and is cleared on source image update. --- modules/extras.py | 67 +++++++++++++++++++++++++++++++------------------------ modules/ui.py | 5 +++++ 2 files changed, 43 insertions(+), 29 deletions(-) (limited to 'modules/ui.py') diff --git a/modules/extras.py b/modules/extras.py index cffe0381..72cc6d1d 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -1,3 +1,4 @@ +from __future__ import annotations import math import os @@ -7,7 +8,7 @@ from PIL import Image import torch import tqdm -from typing import Callable, Dict, List, Tuple +from typing import Callable, List, OrderedDict, Tuple from functools import partial from dataclasses import dataclass @@ -21,18 +22,34 @@ import piexif.helper import gradio as gr -@dataclass(frozen=True) -class CacheKey: - image_hash: int - info_hash: int - args_hash: int +class LruCache(OrderedDict): + @dataclass(frozen=True) + class Key: + image_hash: int + info_hash: int + args_hash: int -@dataclass -class CacheEntry: - image: Image.Image - info: str + @dataclass + class Value: + image: Image.Image + info: str + + def __init__(self, max_size:int = 5, *args, **kwargs): + super().__init__(*args, **kwargs) + self._max_size = max_size + + def get(self, key: LruCache.Key) -> LruCache.Value: + ret = super().get(key) + if ret is not None: + self.move_to_end(key) # Move to end of eviction list + return ret + + def put(self, key: LruCache.Key, value: LruCache.Value) -> None: + self[key] = value + while len(self) > self._max_size: + self.popitem(last=False) -cached_images: Dict[CacheKey, CacheEntry] = {} +cached_images: LruCache = LruCache(max_size = 5) def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, upscale_first: bool ): @@ -121,14 +138,14 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_ blended_result: Image.Image = None for upscaler in params: upscale_args = (upscaler.upscaler_idx, upscaling_resize, resize_mode, upscaling_resize_w, upscaling_resize_h, upscaling_crop) - cache_key = CacheKey( image_hash = hash(np.array(image.getdata()).tobytes()), + cache_key = LruCache.Key( image_hash = hash(np.array(image.getdata()).tobytes()), info_hash = hash(info), - args_hash = hash(upscale_args) ) + args_hash = hash(upscale_args + (upscaler.blend_alpha,)) ) cached_entry = cached_images.get(cache_key) if cached_entry is None: res = upscale(image, *upscale_args) info += f"Upscale: {round(upscaling_resize, 3)}, visibility: {upscaler.blend_alpha}, model:{shared.sd_upscalers[upscaler.upscaler_idx].name}\n" - cached_images[cache_key] = CacheEntry(image=res, info=info) + cached_images.put(cache_key, LruCache.Value(image=res, info=info)) else: res, info = cached_entry.image, cached_entry.info @@ -140,14 +157,11 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_ # Build a list of operations to run facefix_ops: List[Callable] = [] - if gfpgan_visibility > 0: - facefix_ops.append(run_gfpgan) - if codeformer_visibility > 0: - facefix_ops.append(run_codeformer) + facefix_ops += [run_gfpgan] if gfpgan_visibility > 0 else [] + facefix_ops += [run_codeformer] if codeformer_visibility > 0 else [] upscale_ops: List[Callable] = [] - if resize_mode == 1: - upscale_ops.append(run_prepare_crop) + upscale_ops += [run_prepare_crop] if resize_mode == 1 else [] if upscaling_resize != 0: step_params: List[UpscaleParams] = [] @@ -157,12 +171,7 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_ upscale_ops.append( partial(run_upscalers_blend, step_params) ) - - extras_ops: List[Callable] = [] - if upscale_first: - extras_ops = upscale_ops + facefix_ops - else: - extras_ops = facefix_ops + upscale_ops + extras_ops: List[Callable] = (upscale_ops + facefix_ops) if upscale_first else (facefix_ops + upscale_ops) for image, image_name in zip(imageArr, imageNameArr): @@ -176,9 +185,6 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_ for op in extras_ops: image, info = op(image, info) - while len(cached_images) > 2: - del cached_images[next(iter(cached_images.keys()))] - if opts.use_original_name_batch and image_name != None: basename = os.path.splitext(os.path.basename(image_name))[0] else: @@ -198,6 +204,9 @@ def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_ return outputs, plaintext_to_html(info), '' +def clear_cache(): + cached_images.clear() + def run_pnginfo(image): if image is None: diff --git a/modules/ui.py b/modules/ui.py index 16b6ac49..b7c36c55 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -1178,6 +1178,11 @@ def create_ui(wrap_gradio_gpu_call): outputs=[init_img_with_mask], ) + extras_image.change( + fn=modules.extras.clear_cache, + inputs=[], outputs=[] + ) + with gr.Blocks(analytics_enabled=False) as pnginfo_interface: with gr.Row().style(equal_height=False): with gr.Column(variant='panel'): -- cgit v1.2.3