use torch.jit to load Lama model

modules/inpaint/lama/__init__.py

@@ -1,14 +1,12 @@
 # https://github.com/advimman/lama
 import os
-import yaml
 import torch
 import torch.nn.functional as F
 
 import folder_paths
 import comfy.model_management as model_management
 
 from ...model_utils import download_model
-from ...utils import ensure_package
 
 
 lama = None
@@ -35,25 +33,15 @@ def pad_tensor_to_modulo(img, mod):
 def load_model():
     global lama
     if lama is None:
-        ensure_package("omegaconf")
-
-        from omegaconf import OmegaConf
-        from .saicinpainting.training.trainers import load_checkpoint
-
         files = download_model(
             model_path=model_dir,
             model_url=model_url,
             ext_filter=[".pt"],
             download_name="big-lama.pt",
         )
 
-        cfg = yaml.safe_load(open(config_path, "rt"))
-        cfg = OmegaConf.create(cfg)
-        cfg.training_model.predict_only = True
-        cfg.visualizer.kind = "noop"
-
-        lama = load_checkpoint(cfg, files[0], strict=False, map_location="cpu")
-        lama.freeze()
+        lama = torch.jit.load(files[0], map_location="cpu")
+        lama.eval()
 
     return lama
 
@@ -98,13 +86,12 @@ def lama_inpaint(
                 msk = (msk > 0) * 1.0
                 msk = msk.unsqueeze(0).unsqueeze(0)
 
-                batch = {}
-                batch["image"] = pad_tensor_to_modulo(img, 8).to(device)
-                batch["mask"] = pad_tensor_to_modulo(msk, 8).to(device)
+
+                src_image = pad_tensor_to_modulo(img, 8).to(device)
+                src_mask = pad_tensor_to_modulo(msk, 8).to(device)
 
-                res = model(batch)
-                res = batch["inpainted"][0].permute(1, 2, 0)
-                res = res.detach().cpu()
+                res = model(src_image, src_mask)
+                res = res[0].permute(1, 2, 0).detach().cpu()
                 res = res[:orig_h, :orig_w]
 
                 inpainted.append(res)