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

Add workaround low memory crash during image saving

Browse Source 
Sample model moved to cpu before processing samples to image
This commit is contained in:
ArDiouscuros 2022-09-24 10:29:06 +02:00 committed by Lincoln Stein
parent b40bfb5116
commit 485fe67c92
4 changed files with 17 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
ldm/dream/args.py
View File

@ -339,6 +339,12 @@ class Args(object):
action='store_true',
help='Deprecated way to set --precision=float32',
)
model_group.add_argument(
'--free_gpu_mem',
dest='free_gpu_mem',
action='store_true',
help='Force free gpu memory before final decoding',
)
model_group.add_argument(
'--precision',
dest='precision',

8
ldm/dream/generator/txt2img.py
View File

@ -27,6 +27,10 @@ class Txt2Img(Generator):
height // self.downsampling_factor,
width // self.downsampling_factor,
]
if self.free_gpu_mem and self.model.model.device != self.model.device:
self.model.model.to(self.model.device)
samples, _ = sampler.sample(
batch_size = 1,
S = steps,
@ -39,6 +43,10 @@ class Txt2Img(Generator):
eta = ddim_eta,
img_callback = step_callback
)
if self.free_gpu_mem:
self.model.model.to("cpu")
return self.sample_to_image(samples)
return make_image

1
ldm/generate.py
View File

@ -655,6 +655,7 @@ class Generate:
if not self.generators.get('txt2img'):
from ldm.dream.generator.txt2img import Txt2Img
self.generators['txt2img'] = Txt2Img(self.model, self.precision)
self.generators['txt2img'].free_gpu_mem = self.free_gpu_mem
return self.generators['txt2img']
def _make_inpaint(self):

2
scripts/dream.py
View File

@ -108,6 +108,8 @@ def main():
# preload the model
gen.load_model()
#set additional option
gen.free_gpu_mem = opt.free_gpu_mem
if not infile:
print(