From 48eaf60c514a3fa327ac9e4e0621e3dfd9298636 Mon Sep 17 00:00:00 2001 From: nolbert82 Date: Wed, 22 Oct 2025 21:27:39 +0200 Subject: [PATCH] added Apple's "Depth Pro" preprocessor --- modules/control/proc/depth_pro/__init__.py | 95 ++++++++++++++++++++++ modules/control/processor.py | 1 + modules/control/processors.py | 4 + modules/ui_control_elements.py | 2 + 4 files changed, 102 insertions(+) create mode 100644 modules/control/proc/depth_pro/__init__.py diff --git a/modules/control/proc/depth_pro/__init__.py b/modules/control/proc/depth_pro/__init__.py new file mode 100644 index 000000000..43c2458a0 --- /dev/null +++ b/modules/control/proc/depth_pro/__init__.py @@ -0,0 +1,95 @@ +import cv2 +import numpy as np +import torch +from PIL import Image + +from modules import devices, masking +from modules.shared import opts + + +class DepthProDetector: + """Wrapper around Apple's DepthPro depth estimation model.""" + + def __init__(self, model, processor): + self.model = model + self.processor = processor + + @classmethod + def from_pretrained(cls, pretrained_model_or_path: str, cache_dir: str, use_fast_processor: bool = False, **kwargs): + from transformers import AutoImageProcessor, DepthProForDepthEstimation + + processor_kwargs = {"cache_dir": cache_dir} + processor_kwargs.update(kwargs) + if use_fast_processor: + from transformers.models.depth_pro.image_processing_depth_pro_fast import DepthProImageProcessorFast + + processor = DepthProImageProcessorFast.from_pretrained( + pretrained_model_or_path, + **processor_kwargs, + ) + else: + processor = AutoImageProcessor.from_pretrained( + pretrained_model_or_path, + **processor_kwargs, + ) + + model = DepthProForDepthEstimation.from_pretrained( + pretrained_model_or_path, + cache_dir=cache_dir, + ) + model = model.to(device=devices.device).eval() + return cls(model, processor) + + def _prepare_inputs(self, image: Image.Image) -> dict: + inputs = self.processor(images=image, return_tensors="pt") + tensor_inputs = {} + for key, value in inputs.items(): + if isinstance(value, torch.Tensor): + tensor_inputs[key] = value.to(device=devices.device) + else: + tensor_inputs[key] = value + return tensor_inputs + + def __call__( + self, + image, + color_map: str = "inferno", + output_type: str = "pil", + ): + if isinstance(image, list): + image = image[0] + if image is None: + return image + if not isinstance(image, Image.Image): + image = Image.fromarray(np.array(image)) + + original_size = (image.height, image.width) + inputs = self._prepare_inputs(image) + with devices.inference_context(): + outputs = self.model(**inputs) + results = self.processor.post_process_depth_estimation(outputs, target_sizes=[original_size]) + depth_tensor = results[0]["predicted_depth"].to(torch.float32) + if opts.control_move_processor: + self.model.to("cpu") + + # Invert to align with other depth processors that render near as bright + depth_tensor = 1.0 / torch.clamp(depth_tensor, min=1e-6) + depth_tensor -= depth_tensor.min() + max_val = depth_tensor.max() + if max_val > 0: + depth_tensor /= max_val + depth_tensor = (depth_tensor * 255.0).clamp(0, 255).to(torch.uint8) + depth = depth_tensor.cpu().numpy() + + if color_map and color_map.lower() != "none": + color = color_map.lower() + if color not in masking.COLORMAP: + color = "inferno" + processed = cv2.applyColorMap(depth, masking.COLORMAP.index(color))[:, :, ::-1] + else: + processed = depth + + if output_type == "pil": + mode = "RGB" if processed.ndim == 3 else "L" + processed = Image.fromarray(processed, mode=mode) + return processed diff --git a/modules/control/processor.py b/modules/control/processor.py index 80ca18cd9..2022c7959 100644 --- a/modules/control/processor.py +++ b/modules/control/processor.py @@ -34,6 +34,7 @@ processors = [ 'DPT Depth Hybrid', 'GLPN Depth', 'Depth Anything', + 'Depth Pro', ] diff --git a/modules/control/processors.py b/modules/control/processors.py index 24002a4fb..edb02bd09 100644 --- a/modules/control/processors.py +++ b/modules/control/processors.py @@ -39,6 +39,7 @@ config = { 'DPT Depth Hybrid': {'class': None, 'checkpoint': False, 'params': {}}, 'GLPN Depth': {'class': None, 'checkpoint': False, 'params': {}}, 'Depth Anything': {'class': None, 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'LiheYoung/depth_anything_vitl14' }, 'params': { 'color_map': 'inferno' }}, + 'Depth Pro': {'class': None, 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'apple/DepthPro-hf'}, 'params': {'color_map': 'inferno'}}, # 'Midas Depth Large': {'class': MidasDetector, 'checkpoint': True, 'params': {'bg_th': 0.1, 'depth_and_normal': False}, 'load_config': {'pretrained_model_or_path': 'Intel/dpt-large', 'model_type': "dpt_large", 'filename': ''}}, # 'Zoe Depth Zoe': {'class': ZoeDetector, 'checkpoint': True, 'params': {}}, # 'Zoe Depth NK': {'class': ZoeDetector, 'checkpoint': True, 'params': {}, 'load_config': {'pretrained_model_or_path': 'halffried/gyre_zoedepth', 'filename': 'ZoeD_M12_NK.safetensors', 'model_type': "zoedepth_nk"}}, @@ -67,6 +68,7 @@ def delay_load_config(): from modules.control.proc.dpt import DPTDetector from modules.control.proc.glpn import GLPNDetector from modules.control.proc.depth_anything import DepthAnythingDetector + from modules.control.proc.depth_pro import DepthProDetector config = { # placeholder 'None': {}, @@ -95,6 +97,7 @@ def delay_load_config(): 'DPT Depth Hybrid': {'class': DPTDetector, 'checkpoint': False, 'params': {}}, 'GLPN Depth': {'class': GLPNDetector, 'checkpoint': False, 'params': {}}, 'Depth Anything': {'class': DepthAnythingDetector, 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'LiheYoung/depth_anything_vitl14' }, 'params': { 'color_map': 'inferno' }}, + 'Depth Pro': {'class': DepthProDetector, 'checkpoint': True, 'load_config': {'pretrained_model_or_path': 'apple/DepthPro-hf'}, 'params': {'color_map': 'inferno'}}, # 'Midas Depth Large': {'class': MidasDetector, 'checkpoint': True, 'params': {'bg_th': 0.1, 'depth_and_normal': False}, 'load_config': {'pretrained_model_or_path': 'Intel/dpt-large', 'model_type': "dpt_large", 'filename': ''}}, # 'Zoe Depth Zoe': {'class': ZoeDetector, 'checkpoint': True, 'params': {}}, # 'Zoe Depth NK': {'class': ZoeDetector, 'checkpoint': True, 'params': {}, 'load_config': {'pretrained_model_or_path': 'halffried/gyre_zoedepth', 'filename': 'ZoeD_M12_NK.safetensors', 'model_type': "zoedepth_nk"}}, @@ -155,6 +158,7 @@ def update_settings(*settings): update(['Marigold Depth', 'params', 'denoising_steps'], settings[25]) update(['Marigold Depth', 'params', 'ensemble_size'], settings[26]) update(['Depth Anything', 'params', 'color_map'], settings[27]) + update(['Depth Pro', 'params', 'color_map'], settings[28]) class Processor(): diff --git a/modules/ui_control_elements.py b/modules/ui_control_elements.py index 21f7b4ca5..4788b4b99 100644 --- a/modules/ui_control_elements.py +++ b/modules/ui_control_elements.py @@ -317,5 +317,7 @@ def create_ui_elements(units, result_txt, preview_process): settings.append(gr.Slider(label="Ensemble size", minimum=1, maximum=99, step=1, value=10)) with gr.Accordion('Depth Anything', open=True, elem_classes=['processor-settings']): settings.append(gr.Dropdown(label="Depth map", choices=['none'] + masking.COLORMAP, value='inferno')) + with gr.Accordion('Depth Pro', open=True, elem_classes=['processor-settings']): + settings.append(gr.Dropdown(label="Depth map", choices=['none'] + masking.COLORMAP, value='inferno')) for setting in settings: setting.change(fn=processors.update_settings, inputs=settings, outputs=[])