| 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 |
|
|
