github/InvokeAI
1
1
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2026-04-23 03:00:31 -04:00
Code Issues Packages Projects Releases Wiki Activity

code cleanup after @ryand review

Browse Source
This commit is contained in:
Lincoln Stein
2024-06-21 15:36:37 -04:00
parent 787671c2c2
commit 27195b1672
1 changed files with 19 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
invokeai/app/api/routers/model_manager.py
View File

@@ -847,24 +847,37 @@ async def set_cache_size(
"""Set the current RAM or VRAM cache size setting (in GB). ."""
cache = ApiDependencies.invoker.services.model_manager.load.ram_cache
app_config = get_config()
vram_bak, ram_bak = (app_config.vram, app_config.ram)
if cache_type == CacheType.RAM:
cache.max_cache_size = value
app_config.ram = value
elif cache_type == CacheType.VRAM:
cache.max_vram_cache_size = value
app_config.vram = value
else:
raise ValueError(f"Unexpected {cache_type=}.")
if persist:
config_path = app_config.config_file_path
print(f"DEBUG: config_path = {config_path}")
new_config_path = config_path.with_suffix(".yaml.new")
backup_config_path = config_path.with_suffix(".yaml.bak")
shutil.copy(config_path, backup_config_path)
try:
shutil.copy(config_path, config_path.with_suffix(".yaml.bak"))
app_config.write_file(config_path)
app_config.write_file(new_config_path)
shutil.move(new_config_path, config_path)
except Exception as e:
shutil.move(config_path.with_suffix(".yaml.bak"), config_path)
raise RuntimeError(f"Failed to write modified configuration to {config_path}: {e}") from e
shutil.move(backup_config_path, config_path)
app_config.max_vram_cache_size = vram_bak
app_config.max_cache_size = ram_bak
raise RuntimeError(f"Failed to save configuration to {config_path}: {e}") from e
return cache.max_vram_cache_size if cache_type == CacheType.VRAM else cache.max_cache_size
if cache_type == CacheType.VRAM:
return cache.max_vram_cache_size
elif cache_type == CacheType.RAM:
return cache.max_cache_size
else:
raise ValueError(f"Unexpected {cache_type=}.")
@model_manager_router.get(