add anime-lama

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,

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))