79 lines
3.5 KiB
Python
79 lines
3.5 KiB
Python
# Copyright (C) 2023 Deforum LLC
|
|
#
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published by
|
|
# the Free Software Foundation, version 3 of the License.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
# Contact the authors: https://deforum.github.io/
|
|
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torchvision.transforms.functional as TF
|
|
from .general_utils import download_file_with_checksum
|
|
from infer import InferenceHelper
|
|
|
|
class AdaBinsModel:
|
|
_instance = None
|
|
|
|
def __new__(cls, *args, **kwargs):
|
|
keep_in_vram = kwargs.get('keep_in_vram', False)
|
|
if cls._instance is None:
|
|
cls._instance = super().__new__(cls)
|
|
cls._instance._initialize(*args, keep_in_vram=keep_in_vram)
|
|
return cls._instance
|
|
|
|
def _initialize(self, models_path, keep_in_vram=False):
|
|
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
self.keep_in_vram = keep_in_vram
|
|
self.adabins_helper = None
|
|
|
|
download_file_with_checksum(url='https://github.com/hithereai/deforum-for-automatic1111-webui/releases/download/AdaBins/AdaBins_nyu.pt', expected_checksum='643db9785c663aca72f66739427642726b03acc6c4c1d3755a4587aa2239962746410d63722d87b49fc73581dbc98ed8e3f7e996ff7b9c0d56d0fbc98e23e41a', dest_folder=models_path, dest_filename='AdaBins_nyu.pt')
|
|
|
|
self.adabins_helper = InferenceHelper(models_path=models_path, dataset='nyu', device=self.device)
|
|
|
|
def predict(self, img_pil, prev_img_cv2):
|
|
w, h = prev_img_cv2.shape[1], prev_img_cv2.shape[0]
|
|
adabins_depth = np.array([])
|
|
use_adabins = True
|
|
MAX_ADABINS_AREA, MIN_ADABINS_AREA = 500000, 448 * 448
|
|
|
|
image_pil_area, resized = w * h, False
|
|
|
|
if image_pil_area not in range(MIN_ADABINS_AREA, MAX_ADABINS_AREA + 1):
|
|
scale = ((MAX_ADABINS_AREA if image_pil_area > MAX_ADABINS_AREA else MIN_ADABINS_AREA) / image_pil_area) ** 0.5
|
|
depth_input = img_pil.resize((int(w * scale), int(h * scale)), Image.LANCZOS if image_pil_area > MAX_ADABINS_AREA else Image.BICUBIC)
|
|
print(f"AdaBins depth resized to {depth_input.width}x{depth_input.height}")
|
|
resized = True
|
|
else:
|
|
depth_input = img_pil
|
|
|
|
try:
|
|
with torch.no_grad():
|
|
_, adabins_depth = self.adabins_helper.predict_pil(depth_input)
|
|
if resized:
|
|
adabins_depth = TF.resize(torch.from_numpy(adabins_depth), torch.Size([h, w]), interpolation=TF.InterpolationMode.BICUBIC).cpu().numpy()
|
|
adabins_depth = adabins_depth.squeeze()
|
|
except Exception as e:
|
|
print("AdaBins exception encountered. Falling back to pure MiDaS/Zoe (only if running in Legacy Midas/Zoe+AdaBins mode)")
|
|
use_adabins = False
|
|
torch.cuda.empty_cache()
|
|
|
|
return use_adabins, adabins_depth
|
|
|
|
def to(self, device):
|
|
self.device = device
|
|
if self.adabins_helper is not None:
|
|
self.adabins_helper.to(device)
|
|
|
|
def delete_model(self):
|
|
del self.adabins_helper
|