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

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

@ -48,6 +48,10 @@ class Img2Img(Generator):
torch.tensor([t_enc]).to(self.model.device), torch.tensor([t_enc]).to(self.model.device),
noise=x_T 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 # decode it
samples = sampler.decode( samples = sampler.decode(
z_enc, z_enc,
@ -61,6 +65,9 @@ class Img2Img(Generator):
all_timesteps_count = steps all_timesteps_count = steps
) )
if self.free_gpu_mem:
self.model.model.to("cpu")
return self.sample_to_image(samples) return self.sample_to_image(samples)
return make_image return make_image