IOPaint/lama_cleaner/model_manager.py

119 lines
4.1 KiB
Python
Raw Normal View History

2022-11-14 11:19:50 +01:00
import torch
import gc
2022-11-14 11:19:50 +01:00
2023-05-13 07:45:27 +02:00
from loguru import logger
2023-03-19 15:40:23 +01:00
from lama_cleaner.const import SD15_MODELS
2023-02-11 06:30:09 +01:00
from lama_cleaner.helper import switch_mps_device
2023-03-19 15:40:23 +01:00
from lama_cleaner.model.controlnet import ControlNet
2022-09-02 04:37:30 +02:00
from lama_cleaner.model.fcf import FcF
2022-04-15 18:11:51 +02:00
from lama_cleaner.model.lama import LaMa
from lama_cleaner.model.ldm import LDM
2022-11-18 14:40:12 +01:00
from lama_cleaner.model.manga import Manga
2022-08-22 17:24:02 +02:00
from lama_cleaner.model.mat import MAT
2022-12-10 15:06:15 +01:00
from lama_cleaner.model.paint_by_example import PaintByExample
2023-01-28 14:24:51 +01:00
from lama_cleaner.model.instruct_pix2pix import InstructPix2Pix
2023-03-01 14:44:02 +01:00
from lama_cleaner.model.sd import SD15, SD2, Anything4, RealisticVision14
2023-05-11 15:51:58 +02:00
from lama_cleaner.model.utils import torch_gc
2022-07-14 10:49:03 +02:00
from lama_cleaner.model.zits import ZITS
2022-09-25 15:27:12 +02:00
from lama_cleaner.model.opencv2 import OpenCV2
2022-04-15 18:11:51 +02:00
from lama_cleaner.schema import Config
2023-02-11 06:30:09 +01:00
models = {
"lama": LaMa,
"ldm": LDM,
"zits": ZITS,
"mat": MAT,
"fcf": FcF,
2023-03-19 15:40:23 +01:00
SD15.name: SD15,
2023-03-01 14:44:02 +01:00
Anything4.name: Anything4,
RealisticVision14.name: RealisticVision14,
2023-02-11 06:30:09 +01:00
"cv2": OpenCV2,
"manga": Manga,
"sd2": SD2,
"paint_by_example": PaintByExample,
"instruct_pix2pix": InstructPix2Pix,
}
2022-04-15 18:11:51 +02:00
2022-07-14 10:49:03 +02:00
class ModelManager:
2022-11-14 11:19:50 +01:00
def __init__(self, name: str, device: torch.device, **kwargs):
2022-04-15 18:11:51 +02:00
self.name = name
self.device = device
2022-09-15 16:21:27 +02:00
self.kwargs = kwargs
self.model = self.init_model(name, device, **kwargs)
2022-04-15 18:11:51 +02:00
2022-09-15 16:21:27 +02:00
def init_model(self, name: str, device, **kwargs):
2023-03-19 15:40:23 +01:00
if name in SD15_MODELS and kwargs.get("sd_controlnet", False):
return ControlNet(device, **{**kwargs, "name": name})
2022-07-14 10:49:03 +02:00
if name in models:
2022-09-15 16:21:27 +02:00
model = models[name](device, **kwargs)
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):
2023-05-13 07:45:27 +02:00
self.switch_controlnet_method(control_method=config.controlnet_method)
2022-04-15 18:11:51 +02:00
return self.model(image, mask, config)
2023-05-11 15:51:58 +02:00
def switch(self, new_name: str, **kwargs):
2022-04-15 18:11:51 +02:00
if new_name == self.name:
return
try:
2023-02-11 06:30:09 +01:00
if torch.cuda.memory_allocated() > 0:
# Clear current loaded model from memory
torch.cuda.empty_cache()
del self.model
gc.collect()
2023-02-11 06:30:09 +01:00
self.model = self.init_model(
new_name, switch_mps_device(new_name, self.device), **self.kwargs
)
2022-04-15 18:11:51 +02:00
self.name = new_name
except NotImplementedError as e:
raise e
2023-05-11 15:51:58 +02:00
def switch_controlnet_method(self, control_method: str):
if not self.kwargs.get("sd_controlnet"):
return
if self.kwargs["sd_controlnet_method"] == control_method:
return
if not hasattr(self.model, "is_local_sd_model"):
return
2023-05-18 07:07:12 +02:00
if self.model.is_local_sd_model:
# is_native_control_inpaint 表示加载了普通 SD 模型
if (
self.model.is_native_control_inpaint
and control_method != "control_v11p_sd15_inpaint"
):
raise RuntimeError(
f"--sd-local-model-path load a normal SD model, "
f"to use {control_method} you should load an inpainting SD model"
)
elif (
not self.model.is_native_control_inpaint
and control_method == "control_v11p_sd15_inpaint"
):
raise RuntimeError(
f"--sd-local-model-path load an inpainting SD model, "
f"to use {control_method} you should load a norml SD model"
)
2023-05-11 15:51:58 +02:00
del self.model
torch_gc()
2023-05-13 07:45:27 +02:00
old_method = self.kwargs["sd_controlnet_method"]
2023-05-11 15:51:58 +02:00
self.kwargs["sd_controlnet_method"] = control_method
self.model = self.init_model(
self.name, switch_mps_device(self.name, self.device), **self.kwargs
)
2023-05-13 07:45:27 +02:00
logger.info(f"Switch ControlNet method from {old_method} to {control_method}")