IOPaint/iopaint/model_manager.py

193 lines
6.9 KiB
Python
Raw Normal View History

2023-12-01 03:15:35 +01:00
from typing import List, Dict
2022-11-14 11:19:50 +01:00
2023-12-01 03:15:35 +01:00
import torch
2023-05-13 07:45:27 +02:00
from loguru import logger
2024-01-04 14:39:59 +01:00
import numpy as np
2023-05-13 07:45:27 +02:00
2024-01-05 08:19:23 +01:00
from iopaint.download import scan_models
from iopaint.helper import switch_mps_device
from iopaint.model import models, ControlNet, SD, SDXL
from iopaint.model.utils import torch_gc, is_local_files_only
2024-01-05 08:19:23 +01:00
from iopaint.model_info import ModelInfo, ModelType
from iopaint.schema import InpaintRequest
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
2023-12-01 03:15:35 +01:00
self.available_models: Dict[str, ModelInfo] = {}
self.scan_models()
2023-12-19 06:16:30 +01:00
2023-12-27 15:00:07 +01:00
self.enable_controlnet = kwargs.get("enable_controlnet", False)
controlnet_method = kwargs.get("controlnet_method", None)
if (
controlnet_method is None
and name in self.available_models
and self.available_models[name].support_controlnet
):
controlnet_method = self.available_models[name].controlnets[0]
self.controlnet_method = controlnet_method
2022-09-15 16:21:27 +02:00
self.model = self.init_model(name, device, **kwargs)
2022-04-15 18:11:51 +02:00
2023-12-24 08:32:27 +01:00
@property
2023-12-30 16:36:44 +01:00
def current_model(self) -> ModelInfo:
return self.available_models[self.name]
2023-12-24 08:32:27 +01:00
def init_model(self, name: str, device, **kwargs):
logger.info(f"Loading model: {name}")
2023-12-01 03:15:35 +01:00
if name not in self.available_models:
2023-12-30 16:36:44 +01:00
raise NotImplementedError(
2024-01-08 16:43:20 +01:00
f"Unsupported model: {name}. Available models: {list(self.available_models.keys())}"
2023-12-30 16:36:44 +01:00
)
2023-12-01 03:15:35 +01:00
model_info = self.available_models[name]
2023-12-19 06:16:30 +01:00
kwargs = {
**kwargs,
"model_info": model_info,
2023-12-27 15:00:07 +01:00
"enable_controlnet": self.enable_controlnet,
"controlnet_method": self.controlnet_method,
2023-12-19 06:16:30 +01:00
}
2023-12-27 15:00:07 +01:00
if model_info.support_controlnet and self.enable_controlnet:
2023-12-15 05:40:29 +01:00
return ControlNet(device, **kwargs)
2023-12-27 15:00:07 +01:00
elif model_info.name in models:
return models[name](device, **kwargs)
2022-04-15 18:11:51 +02:00
else:
2023-12-01 03:15:35 +01:00
if model_info.model_type in [
2023-12-11 15:28:07 +01:00
ModelType.DIFFUSERS_SD_INPAINT,
2023-12-01 03:15:35 +01:00
ModelType.DIFFUSERS_SD,
]:
2023-12-15 05:40:29 +01:00
return SD(device, **kwargs)
2022-04-15 18:11:51 +02:00
2023-12-11 15:28:07 +01:00
if model_info.model_type in [
ModelType.DIFFUSERS_SDXL_INPAINT,
ModelType.DIFFUSERS_SDXL,
]:
2023-12-15 05:40:29 +01:00
return SDXL(device, **kwargs)
2023-12-01 03:15:35 +01:00
raise NotImplementedError(f"Unsupported model: {name}")
2022-04-17 17:31:12 +02:00
2024-01-09 15:54:20 +01:00
@torch.inference_mode()
2023-12-30 16:36:44 +01:00
def __call__(self, image, mask, config: InpaintRequest):
"""
Args:
image: [H, W, C] RGB
mask: [H, W, 1] 255 means area to repaint
config:
Returns:
2024-01-04 14:39:59 +01:00
BGR image
2023-12-30 16:36:44 +01:00
"""
2023-12-19 06:16:30 +01:00
self.switch_controlnet_method(config)
2023-11-14 15:04:16 +01:00
self.enable_disable_freeu(config)
2023-11-15 01:50:35 +01:00
self.enable_disable_lcm_lora(config)
2024-01-04 14:39:59 +01:00
return self.model(image, mask, config).astype(np.uint8)
2022-04-15 18:11:51 +02:00
2023-12-01 03:15:35 +01:00
def scan_models(self) -> List[ModelInfo]:
available_models = scan_models()
self.available_models = {it.name: it for it in available_models}
return available_models
def switch(self, new_name: str):
2022-04-15 18:11:51 +02:00
if new_name == self.name:
return
2023-12-01 03:15:35 +01:00
old_name = self.name
2023-12-27 15:00:07 +01:00
old_controlnet_method = self.controlnet_method
2023-12-01 03:15:35 +01:00
self.name = new_name
2023-12-19 06:16:30 +01:00
if (
self.available_models[new_name].support_controlnet
2023-12-27 15:00:07 +01:00
and self.controlnet_method
2023-12-19 06:16:30 +01:00
not in self.available_models[new_name].controlnets
):
2023-12-27 15:00:07 +01:00
self.controlnet_method = self.available_models[new_name].controlnets[0]
2022-04-15 18:11:51 +02:00
try:
2023-12-24 08:32:27 +01:00
# TODO: enable/disable controlnet without reload model
2023-12-19 06:16:30 +01:00
del self.model
torch_gc()
2023-02-11 06:30:09 +01:00
self.model = self.init_model(
new_name, switch_mps_device(new_name, self.device), **self.kwargs
)
2023-12-01 03:15:35 +01:00
except Exception as e:
self.name = old_name
2023-12-27 15:00:07 +01:00
self.controlnet_method = old_controlnet_method
2023-12-19 06:16:30 +01:00
logger.info(f"Switch model from {old_name} to {new_name} failed, rollback")
self.model = self.init_model(
old_name, switch_mps_device(old_name, self.device), **self.kwargs
)
2022-04-15 18:11:51 +02:00
raise e
2023-05-11 15:51:58 +02:00
2023-12-19 06:16:30 +01:00
def switch_controlnet_method(self, config):
2023-12-11 15:28:07 +01:00
if not self.available_models[self.name].support_controlnet:
2023-12-01 03:15:35 +01:00
return
2023-05-11 15:51:58 +02:00
2023-12-22 07:00:30 +01:00
if (
2023-12-27 15:00:07 +01:00
self.enable_controlnet
2023-12-22 07:00:30 +01:00
and config.controlnet_method
2023-12-27 15:00:07 +01:00
and self.controlnet_method != config.controlnet_method
2023-12-19 06:16:30 +01:00
):
2023-12-27 15:00:07 +01:00
old_controlnet_method = self.controlnet_method
self.controlnet_method = config.controlnet_method
2023-12-22 07:00:30 +01:00
self.model.switch_controlnet_method(config.controlnet_method)
logger.info(
2023-12-27 15:00:07 +01:00
f"Switch Controlnet method from {old_controlnet_method} to {config.controlnet_method}"
2023-12-22 07:00:30 +01:00
)
2023-12-27 15:00:07 +01:00
elif self.enable_controlnet != config.enable_controlnet:
self.enable_controlnet = config.enable_controlnet
self.controlnet_method = config.controlnet_method
2023-05-11 15:51:58 +02:00
pipe_components = {
"vae": self.model.model.vae,
"text_encoder": self.model.model.text_encoder,
"unet": self.model.model.unet,
}
if hasattr(self.model.model, "text_encoder_2"):
pipe_components["text_encoder_2"] = self.model.model.text_encoder_2
2023-12-19 06:16:30 +01:00
self.model = self.init_model(
self.name,
switch_mps_device(self.name, self.device),
pipe_components=pipe_components,
**self.kwargs,
2023-12-19 06:16:30 +01:00
)
2023-12-27 15:00:07 +01:00
if not config.enable_controlnet:
2023-12-19 06:16:30 +01:00
logger.info(f"Disable controlnet")
else:
logger.info(f"Enable controlnet: {config.controlnet_method}")
2023-11-14 15:04:16 +01:00
2023-12-30 16:36:44 +01:00
def enable_disable_freeu(self, config: InpaintRequest):
2023-11-14 15:04:16 +01:00
if str(self.model.device) == "mps":
return
2023-12-11 15:28:07 +01:00
if self.available_models[self.name].support_freeu:
2023-11-14 15:04:16 +01:00
if config.sd_freeu:
freeu_config = config.sd_freeu_config
self.model.model.enable_freeu(
s1=freeu_config.s1,
s2=freeu_config.s2,
b1=freeu_config.b1,
b2=freeu_config.b2,
)
else:
self.model.model.disable_freeu()
2023-11-15 01:50:35 +01:00
2023-12-30 16:36:44 +01:00
def enable_disable_lcm_lora(self, config: InpaintRequest):
2023-12-11 15:28:07 +01:00
if self.available_models[self.name].support_lcm_lora:
2024-01-08 14:49:27 +01:00
# TODO: change this if load other lora is supported
lcm_lora_loaded = bool(self.model.model.get_list_adapters())
2023-11-15 01:50:35 +01:00
if config.sd_lcm_lora:
2024-01-08 14:49:27 +01:00
if not lcm_lora_loaded:
self.model.model.load_lora_weights(
self.model.lcm_lora_id,
weight_name="pytorch_lora_weights.safetensors",
local_files_only=is_local_files_only(),
)
2023-11-15 01:50:35 +01:00
else:
2024-01-08 14:49:27 +01:00
if lcm_lora_loaded:
self.model.model.disable_lora()