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

move tokenizer into cpu cache as well

Browse Source
This commit is contained in:
Lincoln Stein 2022-10-12 03:03:29 -04:00
parent 7c06849c4d
commit b537e92789
2 changed files with 9 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
ldm/generate.py
View File

@ -678,8 +678,6 @@ class Generate:
self.embedding_path, self.precision == 'float32' or self.precision == 'autocast'
)
# model.to doesn't change the cond_stage_model.device used to move the tokenizer output, so set it here
self.model.cond_stage_model.device = self.device
self._set_sampler()
for m in self.model.modules():

10
ldm/invoke/model_cache.py
View File

@ -114,7 +114,11 @@ class ModelCache(object):
'''
models = self.list_models()
for name in models:
print(f'{name:20s} {models[name]["status"]:>10s} {models[name]["description"]}')
line = f'{name:25s} {models[name]["status"]:>10s} {models[name]["description"]}'
if models[name]['status'] == 'active':
print(f'\033[1m{line}\033[0m')
else:
print(line)
def _check_memory(self):
avail_memory = psutil.virtual_memory()[1]
@ -164,6 +168,8 @@ class ModelCache(object):
print('>> Using more accurate float32 precision')
model.to(self.device)
# model.to doesn't change the cond_stage_model.device used to move the tokenizer output, so set it here
model.cond_stage_model.device = self.device
model.eval()
# usage statistics
@ -190,6 +196,7 @@ class ModelCache(object):
def _model_to_cpu(self,model):
if self._has_cuda():
model.cond_stage_model.device = 'cpu'
model.first_stage_model.to('cpu')
model.cond_stage_model.to('cpu')
model.model.to('cpu')
@ -200,6 +207,7 @@ class ModelCache(object):
model.to(self.device)
model.first_stage_model.to(self.device)
model.cond_stage_model.to(self.device)
model.cond_stage_model.device = self.device
return model
def _pop_oldest_model(self):