From cca8d14c796964188272794f596a8080054ac2b8 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Tue, 20 Dec 2022 18:32:35 -0500 Subject: [PATCH] defer patchmatch loading (#2039) * defer patchmatch loading Because of the way that patchmatch was loaded early at import time, it was impossible to turn off the attempted loading with --no-patchmatch. In addition, the patchmatch loading messages appear early on during initialization, interfering with ability to print out the version cleanly when --version provided to invoke script. This commit creates a thin wrapper class for patch_match that is only loaded when needed, solving both problems. * create a singleton patchmatch object for use in inpainting This creates a thin wrapper to patchmatch which loads the module on demand, respecting the global "trypatchmatch" option. * address 2d round of issues in PR 2039 comments * Patchmatch->PatchMatch and misc cleanup --- backend/invoke_ai_web_server.py | 2 +- ldm/generate.py | 4 ++- ldm/invoke/generator/inpaint.py | 33 +++++++++++-------------- ldm/invoke/globals.py | 3 +-- ldm/invoke/patchmatch.py | 44 +++++++++++++++++++++++++++++++++ 5 files changed, 63 insertions(+), 23 deletions(-) create mode 100644 ldm/invoke/patchmatch.py diff --git a/backend/invoke_ai_web_server.py b/backend/invoke_ai_web_server.py index d91d66e5be..4a087f720f 100644 --- a/backend/invoke_ai_web_server.py +++ b/backend/invoke_ai_web_server.py @@ -294,7 +294,7 @@ class InvokeAIWebServer: print(f">> System config requested") config = self.get_system_config() config["model_list"] = self.generate.model_cache.list_models() - config["infill_methods"] = infill_methods + config["infill_methods"] = infill_methods() socketio.emit("systemConfig", config) @socketio.on("requestModelChange") diff --git a/ldm/generate.py b/ldm/generate.py index 5feb221c2a..e6821245cb 100644 --- a/ldm/generate.py +++ b/ldm/generate.py @@ -190,6 +190,7 @@ class Generate: self.txt2mask = None self.safety_checker = None self.karras_max = None + self.infill_method = None # Note that in previous versions, there was an option to pass the # device to Generate(). However the device was then ignored, so @@ -326,7 +327,7 @@ class Generate: seam_strength: float = 0.7, seam_steps: int = 10, tile_size: int = 32, - infill_method = infill_methods[0], # The infill method to use + infill_method = None, force_outpaint: bool = False, enable_image_debugging = False, @@ -392,6 +393,7 @@ class Generate: self.log_tokenization = log_tokenization self.step_callback = step_callback self.karras_max = karras_max + self.infill_method = infill_method or infill_methods()[0], # The infill method to use with_variations = [] if with_variations is None else with_variations # will instantiate the model or return it from cache diff --git a/ldm/invoke/generator/inpaint.py b/ldm/invoke/generator/inpaint.py index a2613f48ea..dfa64c571e 100644 --- a/ldm/invoke/generator/inpaint.py +++ b/ldm/invoke/generator/inpaint.py @@ -17,22 +17,15 @@ from ldm.models.diffusion.ddim import DDIMSampler from ldm.models.diffusion.ksampler import KSampler from ldm.invoke.generator.base import downsampling from ldm.util import debug_image +from ldm.invoke.patchmatch import PatchMatch from ldm.invoke.globals import Globals -infill_methods: list[str] = list() - -if Globals.try_patchmatch: - from patchmatch import patch_match - if patch_match.patchmatch_available: - print('>> Patchmatch initialized') - infill_methods.append('patchmatch') - else: - print('>> Patchmatch not loaded (nonfatal)') -else: - print('>> Patchmatch loading disabled') - -infill_methods.append('tile') - +def infill_methods()->list[str]: + methods = list() + if PatchMatch.patchmatch_available(): + methods.append('patchmatch') + methods.append('tile') + return methods class Inpaint(Img2Img): def __init__(self, model, precision): @@ -40,6 +33,7 @@ class Inpaint(Img2Img): self.pil_image = None self.pil_mask = None self.mask_blur_radius = 0 + self.infill_method = None super().__init__(model, precision) # Outpaint support code @@ -64,11 +58,11 @@ class Inpaint(Img2Img): return im # Skip patchmatch if patchmatch isn't available - if not patch_match.patchmatch_available: + if not PatchMatch.patchmatch_available(): return im # Patchmatch (note, we may want to expose patch_size? Increasing it significantly impacts performance though) - im_patched_np = patch_match.inpaint(im.convert('RGB'), ImageOps.invert(im.split()[-1]), patch_size = 3) + im_patched_np = PatchMatch.inpaint(im.convert('RGB'), ImageOps.invert(im.split()[-1]), patch_size = 3) im_patched = Image.fromarray(im_patched_np, mode = 'RGB') return im_patched @@ -187,7 +181,7 @@ class Inpaint(Img2Img): tile_size: int = 32, step_callback=None, inpaint_replace=False, enable_image_debugging=False, - infill_method = infill_methods[0], # The infill method to use + infill_method = None, inpaint_width=None, inpaint_height=None, **kwargs): @@ -198,7 +192,8 @@ class Inpaint(Img2Img): """ self.enable_image_debugging = enable_image_debugging - + self.infill_method = infill_method or infill_methods()[0], # The infill method to use + self.inpaint_width = inpaint_width self.inpaint_height = inpaint_height @@ -206,7 +201,7 @@ class Inpaint(Img2Img): self.pil_image = init_image.copy() # Do infill - if infill_method == 'patchmatch' and patch_match.patchmatch_available: + if infill_method == 'patchmatch' and PatchMatch.patchmatch_available(): init_filled = self.infill_patchmatch(self.pil_image.copy()) else: # if infill_method == 'tile': # Only two methods right now, so always use 'tile' if not patchmatch init_filled = self.tile_fill_missing( diff --git a/ldm/invoke/globals.py b/ldm/invoke/globals.py index b15f579c9d..806861a355 100644 --- a/ldm/invoke/globals.py +++ b/ldm/invoke/globals.py @@ -27,6 +27,5 @@ else: # Where to look for the initialization file Globals.initfile = 'invokeai.init' -# Awkward workaround to disable attempted loading of pypatchmatch -# which is causing CI tests to error out. +# Try loading patchmatch Globals.try_patchmatch = True diff --git a/ldm/invoke/patchmatch.py b/ldm/invoke/patchmatch.py new file mode 100644 index 0000000000..925fabe1bf --- /dev/null +++ b/ldm/invoke/patchmatch.py @@ -0,0 +1,44 @@ +''' +This module defines a singleton object, "patchmatch" that +wraps the actual patchmatch object. It respects the global +"try_patchmatch" attribute, so that patchmatch loading can +be suppressed or deferred +''' +from ldm.invoke.globals import Globals +import numpy as np + +class PatchMatch: + ''' + Thin class wrapper around the patchmatch function. + ''' + + patch_match = None + tried_load:bool = False + + def __init__(self): + super().__init__() + + @classmethod + def _load_patch_match(self): + if self.tried_load: + return + if Globals.try_patchmatch: + from patchmatch import patch_match as pm + if pm.patchmatch_available: + print('>> Patchmatch initialized') + else: + print('>> Patchmatch not loaded (nonfatal)') + self.patch_match = pm + else: + print('>> Patchmatch loading disabled') + self.tried_load = True + + @classmethod + def patchmatch_available(self)->bool: + self._load_patch_match() + return self.patch_match and self.patch_match.patchmatch_available + + @classmethod + def inpaint(self,*args,**kwargs)->np.ndarray: + if self.patchmatch_available(): + return self.patch_match.inpaint(*args,**kwargs)