| 38 | global depthmap_device |
| 39 | |
| 40 | class ModelHolder: |
| 41 | def __init__(self): |
| 42 | self.depth_model = None |
| 43 | self.pix2pix_model = None |
| 44 | self.depth_model_type = None |
| 45 | self.device = None # Target device, the model may be swapped from VRAM into RAM. |
| 46 | self.offloaded = False # True means current device is not the target device |
| 47 | |
| 48 | # Extra stuff |
| 49 | self.resize_mode = None |
| 50 | self.normalization = None |
| 51 | self.tiling_mode = False |
| 52 | |
| 53 | |
| 54 | def update_settings(self, **kvargs): |
| 55 | # Opens the pandora box |
| 56 | for k, v in kvargs.items(): |
| 57 | setattr(self, k, v) |
| 58 | |
| 59 | |
| 60 | def ensure_models(self, model_type, device: torch.device, boost: bool, tiling_mode: bool = False): |
| 61 | # TODO: could make it more granular |
| 62 | if model_type == -1 or model_type is None: |
| 63 | self.unload_models() |
| 64 | return |
| 65 | # Certain optimisations are irreversible and not device-agnostic, thus changing device requires reloading |
| 66 | if ( |
| 67 | model_type != self.depth_model_type or |
| 68 | boost != (self.pix2pix_model is not None) or |
| 69 | device != self.device or |
| 70 | tiling_mode != self.tiling_mode |
| 71 | ): |
| 72 | self.unload_models() |
| 73 | self.load_models(model_type, device, boost, tiling_mode) |
| 74 | self.reload() |
| 75 | |
| 76 | def load_models(self, model_type, device: torch.device, boost: bool, tiling_mode: bool = False): |
| 77 | """Ensure that the depth model is loaded""" |
| 78 | |
| 79 | # TODO: we need to at least try to find models downloaded by other plugins (e.g. controlnet) |
| 80 | |
| 81 | # model path and name |
| 82 | # ZoeDepth and Marigold do not use this |
| 83 | model_dir = "./models/midas" |
| 84 | if model_type == 0: |
| 85 | model_dir = "./models/leres" |
| 86 | if model_type == 11: |
| 87 | model_dir = "./models/depth_anything" |
| 88 | if model_type in [12, 13, 14]: |
| 89 | model_dir = "./models/depth_anything_v2" |
| 90 | |
| 91 | # create paths to model if not present |
| 92 | os.makedirs(model_dir, exist_ok=True) |
| 93 | os.makedirs('./models/pix2pix', exist_ok=True) |
| 94 | |
| 95 | print("Loading model weights from ", end=" ") |
| 96 | |
| 97 | resize_mode = "minimal" |