fix slow sd test

2022-11-14 18:19:50 +08:00
parent eec41734c3
commit c2005786d7
2 changed files with 8 additions and 5 deletions
--- a/lama_cleaner/model_manager.py
+++ b/lama_cleaner/model_manager.py
@@ -1,3 +1,5 @@
+import torch
+
 from lama_cleaner.model.fcf import FcF
 from lama_cleaner.model.lama import LaMa
 from lama_cleaner.model.ldm import LDM
@@ -11,7 +13,7 @@ models = {"lama": LaMa, "ldm": LDM, "zits": ZITS, "mat": MAT, "fcf": FcF, "sd1.5


 class ModelManager:
-    def __init__(self, name: str, device, **kwargs):
+    def __init__(self, name: str, device: torch.device, **kwargs):
        self.name = name
        self.device = device
        self.kwargs = kwargs