Merge pull request #4294 from nolbert82/dev

Added Depth pro + Fixed guidance end

modules/control/proc/depth_pro/__init__.py

@@ -0,0 +1,62 @@
+import cv2
+import torch
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+
+from modules import devices, masking
+from modules.shared import opts
+
+
+class DepthProDetector:
+    """Apple DepthPro detector (aligned with Depth Anything style)."""
+
+    def __init__(self, model, processor):
+        self.model = model
+        self.processor = processor
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_or_path: str = "apple/DepthPro-hf", cache_dir: str | None = None) -> "DepthProDetector":
+        from transformers import AutoImageProcessor, DepthProForDepthEstimation
+
+        processor = AutoImageProcessor.from_pretrained(pretrained_model_or_path, cache_dir=cache_dir)
+        model = DepthProForDepthEstimation.from_pretrained(
+            pretrained_model_or_path,
+            cache_dir=cache_dir,
+        ).to(devices.device).eval()
+        return cls(model, processor)
+
+    def __call__(self, image, color_map: str = "none", output_type: str = "pil"):
+        self.model.to(devices.device)
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        h, w = image.shape[:2]
+        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        pil_image = Image.fromarray(image_rgb)
+
+        inputs = self.processor(images=pil_image, return_tensors="pt")
+        inputs = {k: v.to(devices.device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
+
+        with devices.inference_context():
+            outputs = self.model(**inputs)
+        results = self.processor.post_process_depth_estimation(outputs, target_sizes=[(h, w)])
+        depth_tensor = results[0]["predicted_depth"].to(devices.device, dtype=torch.float32)
+
+        if opts.control_move_processor:
+            self.model.to("cpu")
+
+        depth_tensor = F.interpolate(depth_tensor[None, None], size=(h, w), mode="bilinear", align_corners=False)[0, 0]
+        depth_tensor = 1.0 / torch.clamp(depth_tensor, min=1e-6)
+        depth_tensor -= depth_tensor.min()
+        depth_max = depth_tensor.max()
+        if depth_max > 0:
+            depth_tensor /= depth_max
+        depth = (depth_tensor * 255.0).clamp(0, 255).to(torch.uint8).cpu().numpy()
+
+        if color_map != "none":
+            colormap_key = color_map if color_map in masking.COLORMAP else "inferno"
+            depth = cv2.applyColorMap(depth, masking.COLORMAP.index(colormap_key))[:, :, ::-1]
+        if output_type == "pil":
+            mode = "RGB" if depth.ndim == 3 else "L"
+            depth = Image.fromarray(depth, mode=mode)
+        return depth

modules/control/processor.py

@@ -34,6 +34,7 @@ processors = [
     'DPT Depth Hybrid',
     'GLPN Depth',
     'Depth Anything',
+    'Depth Pro',
 ]
 
 

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():

modules/processing_callbacks.py

@@ -93,15 +93,28 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {}
     if step != getattr(pipe, 'num_timesteps', 0):
         kwargs = processing_correction.correction_callback(p, timestep, kwargs, initial=step == 0)
     kwargs = prompt_callback(step, kwargs)  # monkey patch for diffusers callback issues
-    if step == int(getattr(pipe, 'num_timesteps', 100) * p.cfg_end) and 'prompt_embeds' in kwargs and 'negative_prompt_embeds' in kwargs:
+
+    if step == 0:
+        setattr(pipe, "_cfg_end_applied", False)
+
+    cfg_end = getattr(p, "cfg_end", 1.0) or 1.0
+    total_steps = getattr(pipe, "num_timesteps", 0)
+    target_step = int(total_steps * cfg_end) if total_steps else 0
+    if (
+        cfg_end < 1.0
+        and not getattr(pipe, "_cfg_end_applied", False)
+        and step >= target_step
+    ):
+        setattr(pipe, "_cfg_end_applied", True)
         if "PAG" in shared.sd_model.__class__.__name__:
             pipe._guidance_scale = 1.001 if pipe._guidance_scale > 1 else pipe._guidance_scale  # pylint: disable=protected-access
             pipe._pag_scale = 0.001  # pylint: disable=protected-access
         else:
             pipe._guidance_scale = 0.0  # pylint: disable=protected-access
-            for key in {"prompt_embeds", "negative_prompt_embeds", "add_text_embeds", "add_time_ids"} & set(kwargs):
-                if kwargs[key] is not None:
-                    kwargs[key] = kwargs[key].chunk(2)[-1]
+            for key in {"prompt_embeds", "negative_prompt_embeds", "add_text_embeds", "add_time_ids"}:
+                tensor = kwargs.get(key, None)
+                if tensor is not None and hasattr(tensor, "chunk") and tensor.shape[0] % 2 == 0:
+                    kwargs[key] = tensor.chunk(2)[-1]
     try:
         current_noise_pred = kwargs.get("noise_pred", None)
         if current_noise_pred is None:

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=[])