kamp/IOPaint
1
0
Fork 0
Code Issues Pull Requests Releases 1 Activity

Clear model from memory before switch, fixes out of memory

Browse Source
This commit is contained in:
David Bejar Caceres 2022-12-16 16:14:15 +01:00
parent ea422ca68a
commit 0c4de4776e
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
lama_cleaner/model_manager.py
View File

@ -1,4 +1,5 @@
import torch import torch
import gc
from lama_cleaner.model.fcf import FcF from lama_cleaner.model.fcf import FcF
from lama_cleaner.model.lama import LaMa from lama_cleaner.model.lama import LaMa
@ -42,6 +43,12 @@ class ModelManager:
if new_name == self.name: if new_name == self.name:
return return
try: try:
if (torch.cuda.memory_allocated() > 0):
# Clear current loaded model from memory
torch.cuda.empty_cache()
gc.collect()
del self.model
self.model = self.init_model(new_name, self.device, **self.kwargs) self.model = self.init_model(new_name, self.device, **self.kwargs)
self.name = new_name self.name = new_name
except NotImplementedError as e: except NotImplementedError as e: