Load model in inpaint when using free_gpu_mem option (#1938)

* Load model in inpaint when using free_gpu_mem option * Passing free_gpu_mem option to inpaint generator
invoke-ai · Dec 12, 2022 · f4e7383 · f4e7383
1 parent 96a1209
commit f4e7383
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/ldm/generate.py b/ldm/generate.py
@@ -806,6 +806,7 @@ def _make_inpaint(self):
         if not self.generators.get('inpaint'):
             from ldm.invoke.generator.inpaint import Inpaint
             self.generators['inpaint'] = Inpaint(self.model, self.precision)
+            self.generators['inpaint'].free_gpu_mem = self.free_gpu_mem
         return self.generators['inpaint']
 
     # "omnibus" supports the runwayML custom inpainting model, which does

diff --git a/ldm/invoke/generator/inpaint.py b/ldm/invoke/generator/inpaint.py
@@ -59,7 +59,7 @@ def get_tile_images(self, image: np.ndarray, width=8, height=8):
             writeable=False
         )
 
-    def infill_patchmatch(self, im: Image.Image) -> Image:        
+    def infill_patchmatch(self, im: Image.Image) -> Image:
         if im.mode != 'RGBA':
             return im
 
@@ -128,7 +128,7 @@ def mask_edge(self, mask: Image, edge_size: int, edge_blur: int) -> Image:
         # Combine
         npmask = npgradient + npedge
 
-        # Expand 
+        # Expand
         npmask = cv.dilate(npmask, np.ones((3,3), np.uint8), iterations = int(edge_size / 2))
 
         new_mask = Image.fromarray(npmask)
@@ -221,7 +221,7 @@ def get_make_image(self,prompt,sampler,steps,cfg_scale,ddim_eta,
                 init_filled = init_filled.resize((inpaint_width, inpaint_height))
 
             debug_image(init_filled, "init_filled", debug_status=self.enable_image_debugging)
-            
+
             # Create init tensor
             init_image = self._image_to_tensor(init_filled.convert('RGB'))
 
@@ -254,7 +254,7 @@ def get_make_image(self,prompt,sampler,steps,cfg_scale,ddim_eta,
                 f">> Using recommended DDIM sampler for inpainting."
             )
             sampler = DDIMSampler(self.model, device=self.model.device)
-        
+
         sampler.make_schedule(
             ddim_num_steps=steps, ddim_eta=ddim_eta, verbose=False
         )
@@ -291,6 +291,9 @@ def make_image(x_T):
                 masked_region = (1.0-inpaint_replace) * inverted_mask * z_enc + inpaint_replace * inverted_mask * l_noise
                 z_enc   = z_enc * mask_image + masked_region
 
+            if self.free_gpu_mem and self.model.model.device != self.model.device:
+                self.model.model.to(self.model.device)
+
             # decode it
             samples = sampler.decode(
                 z_enc,
@@ -353,7 +356,7 @@ def sample_to_image(self, samples)->Image.Image:
 
         if self.pil_image is None or self.pil_mask is None:
             return gen_result
-        
+
         corrected_result = super().repaste_and_color_correct(gen_result, self.pil_image, self.pil_mask, self.mask_blur_radius)
         debug_image(corrected_result, "corrected_result", debug_status=self.enable_image_debugging)