From d7ba041007ec3a23f85eea8c380adef5b780e501 Mon Sep 17 00:00:00 2001 From: Daya Adianto Date: Sat, 3 Dec 2022 15:12:15 +0700 Subject: [PATCH] Enable force free GPU memory in img2img --- ldm/invoke/generator/img2img.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/ldm/invoke/generator/img2img.py b/ldm/invoke/generator/img2img.py index 1981b4eacb..f188d9b23d 100644 --- a/ldm/invoke/generator/img2img.py +++ b/ldm/invoke/generator/img2img.py @@ -48,6 +48,10 @@ class Img2Img(Generator): torch.tensor([t_enc]).to(self.model.device), noise=x_T ) + + 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, @@ -61,6 +65,9 @@ class Img2Img(Generator): all_timesteps_count = steps ) + if self.free_gpu_mem: + self.model.model.to("cpu") + return self.sample_to_image(samples) return make_image @@ -87,4 +94,4 @@ class Img2Img(Generator): image = torch.from_numpy(image) if normalize: image = 2.0 * image - 1.0 - return image.to(self.model.device) + return image.to(self.model.device)