2022-04-15 18:11:51 +02:00
|
|
|
from lama_cleaner.model.lama import LaMa
|
|
|
|
from lama_cleaner.model.ldm import LDM
|
2022-08-22 17:24:02 +02:00
|
|
|
from lama_cleaner.model.mat import MAT
|
2022-07-14 10:49:03 +02:00
|
|
|
from lama_cleaner.model.zits import ZITS
|
2022-04-15 18:11:51 +02:00
|
|
|
from lama_cleaner.schema import Config
|
|
|
|
|
2022-07-14 10:49:03 +02:00
|
|
|
models = {
|
|
|
|
'lama': LaMa,
|
|
|
|
'ldm': LDM,
|
2022-08-22 17:24:02 +02:00
|
|
|
'zits': ZITS,
|
|
|
|
'mat': MAT
|
2022-07-14 10:49:03 +02:00
|
|
|
}
|
2022-04-15 18:11:51 +02:00
|
|
|
|
|
|
|
|
2022-07-14 10:49:03 +02:00
|
|
|
class ModelManager:
|
2022-04-15 18:11:51 +02:00
|
|
|
def __init__(self, name: str, device):
|
|
|
|
self.name = name
|
|
|
|
self.device = device
|
|
|
|
self.model = self.init_model(name, device)
|
|
|
|
|
|
|
|
def init_model(self, name: str, device):
|
2022-07-14 10:49:03 +02:00
|
|
|
if name in models:
|
|
|
|
model = models[name](device)
|
2022-04-15 18:11:51 +02:00
|
|
|
else:
|
|
|
|
raise NotImplementedError(f"Not supported model: {name}")
|
|
|
|
return model
|
|
|
|
|
2022-04-17 17:31:12 +02:00
|
|
|
def is_downloaded(self, name: str) -> bool:
|
2022-07-14 10:49:03 +02:00
|
|
|
if name in models:
|
|
|
|
return models[name].is_downloaded()
|
2022-04-17 17:31:12 +02:00
|
|
|
else:
|
|
|
|
raise NotImplementedError(f"Not supported model: {name}")
|
|
|
|
|
2022-04-15 18:11:51 +02:00
|
|
|
def __call__(self, image, mask, config: Config):
|
|
|
|
return self.model(image, mask, config)
|
|
|
|
|
|
|
|
def switch(self, new_name: str):
|
|
|
|
if new_name == self.name:
|
|
|
|
return
|
|
|
|
try:
|
|
|
|
self.model = self.init_model(new_name, self.device)
|
|
|
|
self.name = new_name
|
|
|
|
except NotImplementedError as e:
|
|
|
|
raise e
|