diff --git a/modules/control/proc/depth_pro/__init__.py b/modules/control/proc/depth_pro/__init__.py new file mode 100644 index 000000000..ac2075632 --- /dev/null +++ b/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 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 1f5b0494b..4d66bfac6 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/processing_callbacks.py b/modules/processing_callbacks.py index 3aec83fca..e36bd0e75 100644 --- a/modules/processing_callbacks.py +++ b/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: 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=[])