Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Enable force free GPU memory in img2img

Browse Source
This commit is contained in:
Daya Adianto 2022-12-03 15:12:15 +07:00 committed by Lincoln Stein
parent 3859302f1c
commit d7ba041007
1 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
ldm/invoke/generator/img2img.py
View File

@ -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)