diff --git a/iopaint/model/__init__.py b/iopaint/model/__init__.py index 799e2ec..5914dc4 100644 --- a/iopaint/model/__init__.py +++ b/iopaint/model/__init__.py @@ -3,7 +3,7 @@ from .controlnet import ControlNet from .fcf import FcF from .instruct_pix2pix import InstructPix2Pix from .kandinsky import Kandinsky22 -from .lama import LaMa +from .lama import LaMa, AnimeLaMa from .ldm import LDM from .manga import Manga from .mat import MAT @@ -17,6 +17,7 @@ from .zits import ZITS models = { LaMa.name: LaMa, + AnimeLaMa.name: AnimeLaMa, LDM.name: LDM, ZITS.name: ZITS, MAT.name: MAT, diff --git a/iopaint/model/lama.py b/iopaint/model/lama.py index 7aba242..56fd5eb 100644 --- a/iopaint/model/lama.py +++ b/iopaint/model/lama.py @@ -19,6 +19,14 @@ LAMA_MODEL_URL = os.environ.get( ) LAMA_MODEL_MD5 = os.environ.get("LAMA_MODEL_MD5", "e3aa4aaa15225a33ec84f9f4bc47e500") +ANIME_LAMA_MODEL_URL = os.environ.get( + "ANIME_LAMA_MODEL_URL", + "https://github.com/Sanster/models/releases/download/AnimeMangaInpainting/anime-manga-big-lama.pt", +) +ANIME_LAMA_MODEL_MD5 = os.environ.get( + "ANIME_LAMA_MODEL_MD5", "29f284f36a0a510bcacf39ecf4c4d54f" +) + class LaMa(InpaintModel): name = "lama" @@ -55,3 +63,20 @@ class LaMa(InpaintModel): cur_res = np.clip(cur_res * 255, 0, 255).astype("uint8") cur_res = cv2.cvtColor(cur_res, cv2.COLOR_RGB2BGR) return cur_res + + +class AnimeLaMa(LaMa): + name = "anime-lama" + + @staticmethod + def download(): + download_model(ANIME_LAMA_MODEL_URL, ANIME_LAMA_MODEL_MD5) + + def init_model(self, device, **kwargs): + self.model = load_jit_model( + ANIME_LAMA_MODEL_URL, device, ANIME_LAMA_MODEL_MD5 + ).eval() + + @staticmethod + def is_downloaded() -> bool: + return os.path.exists(get_cache_path_by_url(ANIME_LAMA_MODEL_URL))