cleanup: Lama

This commit is contained in:
blessedcoolant 2023-08-24 08:51:48 +12:00 committed by Kent Keirsey
parent f45ea1145d
commit fa3fcd7820

View File

@ -1,5 +1,4 @@
import gc
import pathlib
from typing import Any
import numpy as np
@ -44,6 +43,7 @@ class LaMA:
image = torch.from_numpy(image).unsqueeze(0).to(device)
mask = torch.from_numpy(mask).unsqueeze(0).to(device)
with torch.inference_mode():
infilled_image = model(image, mask)
infilled_image = infilled_image[0].permute(1, 2, 0).detach().cpu().numpy()