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

check that model name matches format

Browse Source
This commit is contained in:
Lincoln Stein 2023-05-09 23:46:59 -04:00
parent 3d85e769ce
commit 99c692f397
2 changed files with 12 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
invokeai/backend/model_management/model_cache.py
View File

@ -289,19 +289,26 @@ class ModelCache(object):
cache = self.cache cache = self.cache
key = self.key key = self.key
model = self.model model = self.model
# NOTE that the model has to have the to() method in order for this
# code to move it into GPU!
if self.gpu_load and hasattr(model,'to'): if self.gpu_load and hasattr(model,'to'):
cache.loaded_models.add(key) cache.loaded_models.add(key)
cache.locked_models[key] += 1 cache.locked_models[key] += 1
if cache.lazy_offloading: if cache.lazy_offloading:
cache._offload_unlocked_models() cache._offload_unlocked_models()
if model.device != cache.execution_device: if model.device != cache.execution_device:
cache.logger.debug(f'Moving {key} into {cache.execution_device}') cache.logger.debug(f'Moving {key} into {cache.execution_device}')
with VRAMUsage() as mem: with VRAMUsage() as mem:
model.to(cache.execution_device) # move into GPU model.to(cache.execution_device) # move into GPU
cache.logger.debug(f'GPU VRAM used for load: {(mem.vram_used/GIG):.2f} GB') cache.logger.debug(f'GPU VRAM used for load: {(mem.vram_used/GIG):.2f} GB')
cache.model_sizes[key] = mem.vram_used # more accurate size cache.model_sizes[key] = mem.vram_used # more accurate size
cache.logger.debug(f'Locking {key} in {cache.execution_device}') cache.logger.debug(f'Locking {key} in {cache.execution_device}')
cache._print_cuda_stats() cache._print_cuda_stats()
else: else:
# in the event that the caller wants the model in RAM, we # in the event that the caller wants the model in RAM, we
# move it into CPU if it is in GPU and not locked # move it into CPU if it is in GPU and not locked

7
invokeai/backend/model_management/model_manager.py
View File

@ -1054,10 +1054,13 @@ class ModelManager(object):
full_name = f"{model_name}/{model_type.name}" full_name = f"{model_name}/{model_type.name}"
if full_name in self.config: if full_name in self.config:
return full_name return full_name
if model_name in self.config: # special case - if diffusers requested, then allow name without type appended
if model_type==SDModelType.diffusers \
and model_name in self.config \
and self.config[model_name].format=='diffusers':
return model_name return model_name
raise InvalidModelError( raise InvalidModelError(
f'Neither "{model_name}" nor "{full_name}" are known model names. Please check your models.yaml file' f'"{full_name}" is not a known model name. Please check your models.yaml file'
) )