add anime-lama
This commit is contained in:
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user