Upload Prompt-Depth-Anything-Small (transformers port, original code)

29ebc1d verified 3 days ago

22.7 kB

# ruff: noqa
# fmt: off
"""Prompt Depth Anything, vendored verbatim from DepthAnything/PromptDA @ 9161547ca7151e14fcc8f490ad3cc02fa9437483 (promptda/ and the local torchhub
copy of facebookresearch/dinov2 it loads): the files below concatenated, package-internal imports replaced by
relative imports of the other vendored files (generated by scripts/model_ports/common/vendor_flat.py; do not edit by hand). Code license: Apache-2.0."""


# ==============================================================================================================
# PromptDA/promptda/utils/logger.py
# ==============================================================================================================
import os


class Log:
    log_on = True  # fast switch
    used_tags = dict()  # To keep track of used tags
    _is_main_cached = None  # Cache to store the main process check result

    @staticmethod
    def is_main_process():
        if Log._is_main_cached is not None:
            return Log._is_main_cached
        try:
            from pytorch_lightning.utilities import rank_zero_only
            if rank_zero_only.rank == 0:
                Log._is_main_cached = True
            else:
                Log._is_main_cached = False
        except:
            Log._is_main_cached = True
        return Log._is_main_cached

    @staticmethod
    def _should_log(tag):
        """
        Determine if the log information should be recorded.
        Conditions: log function is enabled, current process is the main process, and the tag has not been used.
        """
        if not Log.log_on:
            return False
        if not Log.is_main_process():
            return False
        if tag is None:
            return True
        if '__' in tag:
            num = int(tag.split('__')[-1])
            tag = tag.split('__')[0]  # can output num same information
        else:
            num = 3  # default 3

        if tag not in Log.used_tags:
            Log.used_tags[tag] = num
        Log.used_tags[tag] -= 1
        if Log.used_tags[tag] >= 0:
            return True
        else:
            return False

    @staticmethod
    def info(*args, tag=None):
        """
        Output INFO level log information.
        """
        if Log._should_log(tag):
            print("\033[1;32m[INFO]\033[0;0m", *args)

    @staticmethod
    def warn(*args, tag=None):
        """
        Output WARN level log information.
        """
        if Log._should_log(tag):
            print("\033[1;35m[WARN]\033[0;0m", *args)

    @staticmethod
    def error(*args, tag=None):
        print("\033[1;31m[ERROR]\033[0;0m", *args)

    @staticmethod
    def debug(*args, tag=None):
        """
        Output DEBUG level log information.
        """
        if Log._should_log(tag) and 'HT_DEBUG' in os.environ and os.environ['HT_DEBUG'] == '1':
            print("\033[1;33m[DEBUG]\033[0;0m", *args)


# ==============================================================================================================
# PromptDA/promptda/model/blocks.py
# ==============================================================================================================
import torch.nn as nn
import torch
import torch.nn.functional as F
import os
import numpy as np


def _make_fusion_block(features, use_bn, size=None):
    return FeatureFusionDepthBlock(
        features,
        nn.ReLU(False),
        deconv=False,
        bn=use_bn,
        expand=False,
        align_corners=True,
        size=size,
    )


def _make_scratch(in_shape, out_shape, groups=1, expand=False):
    scratch = nn.Module()

    out_shape1 = out_shape
    out_shape2 = out_shape
    out_shape3 = out_shape
    if len(in_shape) >= 4:
        out_shape4 = out_shape

    if expand:
        out_shape1 = out_shape
        out_shape2 = out_shape*2
        out_shape3 = out_shape*4
        if len(in_shape) >= 4:
            out_shape4 = out_shape*8

    scratch.layer1_rn = nn.Conv2d(
        in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
    )
    scratch.layer2_rn = nn.Conv2d(
        in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
    )
    scratch.layer3_rn = nn.Conv2d(
        in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
    )
    if len(in_shape) >= 4:
        scratch.layer4_rn = nn.Conv2d(
            in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
        )

    return scratch


class ResidualConvUnit(nn.Module):
    """Residual convolution module.
    """

    def __init__(self, features, activation, bn):
        """Init.
        Args:
            features (int): number of features
        """
        super().__init__()

        self.bn = bn

        self.groups = 1

        self.conv1 = nn.Conv2d(
            features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
        )

        self.conv2 = nn.Conv2d(
            features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
        )

        if self.bn == True:
            self.bn1 = nn.BatchNorm2d(features)
            self.bn2 = nn.BatchNorm2d(features)

        self.activation = activation

        self.skip_add = nn.quantized.FloatFunctional()

    def forward(self, x):
        """Forward pass.
        Args:
            x (tensor): input
        Returns:
            tensor: output
        """

        out = self.activation(x)
        out = self.conv1(out)
        if self.bn == True:
            out = self.bn1(out)

        out = self.activation(out)
        out = self.conv2(out)
        if self.bn == True:
            out = self.bn2(out)

        if self.groups > 1:
            out = self.conv_merge(out)

        return self.skip_add.add(out, x)


class FeatureFusionBlock(nn.Module):
    """Feature fusion block.
    """

    def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, size=None):
        """Init.
        Args:
            features (int): number of features
        """
        super(FeatureFusionBlock, self).__init__()

        self.deconv = deconv
        self.align_corners = align_corners

        self.groups = 1

        self.expand = expand
        out_features = features
        if self.expand == True:
            out_features = features//2

        self.out_conv = nn.Conv2d(
            features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)

        self.resConfUnit1 = ResidualConvUnit(features, activation, bn)
        self.resConfUnit2 = ResidualConvUnit(features, activation, bn)

        self.skip_add = nn.quantized.FloatFunctional()

        self.size = size

    def forward(self, *xs, size=None):
        """Forward pass.
        Returns:
            tensor: output
        """
        output = xs[0]

        if len(xs) == 2:
            res = self.resConfUnit1(xs[1])
            output = self.skip_add.add(output, res)

        output = self.resConfUnit2(output)

        if (size is None) and (self.size is None):
            modifier = {"scale_factor": 2}
        elif size is None:
            modifier = {"size": self.size}
        else:
            modifier = {"size": size}

        output = nn.functional.interpolate(
            output, **modifier, mode="bilinear", align_corners=self.align_corners
        )

        output = self.out_conv(output)

        return output


class FeatureFusionControlBlock(FeatureFusionBlock):
    """Feature fusion block.
    """

    def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, size=None):
        """Init.
        Args:
            features (int): number of features
        """
        super.__init__(features, activation, deconv,
                       bn, expand, align_corners, size)
        self.copy_block = FeatureFusionBlock(
            features, activation, deconv, bn, expand, align_corners, size)

    def forward(self, *xs, size=None):
        """Forward pass.
        Returns:
            tensor: output
        """
        output = xs[0]

        if len(xs) == 2:
            res = self.resConfUnit1(xs[1])
            output = self.skip_add.add(output, res)

        output = self.resConfUnit2(output)

        if (size is None) and (self.size is None):
            modifier = {"scale_factor": 2}
        elif size is None:
            modifier = {"size": self.size}
        else:
            modifier = {"size": size}

        output = nn.functional.interpolate(
            output, **modifier, mode="bilinear", align_corners=self.align_corners
        )

        output = self.out_conv(output)

        return output


def zero_module(module):
    """
    Zero out the parameters of a module and return it.
    """
    for p in module.parameters():
        p.detach().zero_()
    return module


class FeatureFusionDepthBlock(nn.Module):
    """Feature fusion block.
    """

    def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, size=None):
        """Init.
        Args:
            features (int): number of features
        """
        super(FeatureFusionDepthBlock, self).__init__()

        self.deconv = deconv
        self.align_corners = align_corners

        self.groups = 1

        self.expand = expand
        out_features = features
        if self.expand == True:
            out_features = features//2

        self.out_conv = nn.Conv2d(
            features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)

        self.resConfUnit1 = ResidualConvUnit(features, activation, bn)
        self.resConfUnit2 = ResidualConvUnit(features, activation, bn)
        self.resConfUnit_depth = nn.Sequential(
            nn.Conv2d(1, features, kernel_size=3, stride=1,
                      padding=1, bias=True, groups=1),
            activation,
            nn.Conv2d(features, features, kernel_size=3,
                      stride=1, padding=1, bias=True, groups=1),
            activation,
            zero_module(
                nn.Conv2d(features, features, kernel_size=3,
                          stride=1, padding=1, bias=True, groups=1)
            )
        )
        self.skip_add = nn.quantized.FloatFunctional()
        self.size = size

    def forward(self, *xs, prompt_depth=None, size=None):
        """Forward pass.
        Returns:
            tensor: output
        """
        output = xs[0]

        if len(xs) == 2:
            res = self.resConfUnit1(xs[1])
            output = self.skip_add.add(output, res)

        output = self.resConfUnit2(output)

        if prompt_depth is not None:
            prompt_depth = F.interpolate(
                prompt_depth, output.shape[2:], mode='bilinear', align_corners=False)
            res = self.resConfUnit_depth(prompt_depth)
            output = self.skip_add.add(output, res)

        if (size is None) and (self.size is None):
            modifier = {"scale_factor": 2}
        elif size is None:
            modifier = {"size": self.size}
        else:
            modifier = {"size": size}

        output = nn.functional.interpolate(
            output, **modifier, mode="bilinear", align_corners=self.align_corners
        )

        output = self.out_conv(output)

        return output


# ==============================================================================================================
# PromptDA/promptda/model/dpt.py
# ==============================================================================================================
# Copyright (c) 2024, Depth Anything V2
# https://github.com/DepthAnything/Depth-Anything-V2/blob/main/depth_anything_v2/dpt.py
import torch
import torch.nn as nn
import torch.nn.functional as F


class DPTHead(nn.Module):
    def __init__(self,
                 nclass,
                 in_channels,
                 features=256,
                 out_channels=[256, 512, 1024, 1024],
                 use_bn=False,
                 use_clstoken=False,
                 output_act='sigmoid'):
        super(DPTHead, self).__init__()

        self.nclass = nclass
        self.use_clstoken = use_clstoken

        self.projects = nn.ModuleList([
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=out_channel,
                kernel_size=1,
                stride=1,
                padding=0,
            ) for out_channel in out_channels
        ])

        self.resize_layers = nn.ModuleList([
            nn.ConvTranspose2d(
                in_channels=out_channels[0],
                out_channels=out_channels[0],
                kernel_size=4,
                stride=4,
                padding=0),
            nn.ConvTranspose2d(
                in_channels=out_channels[1],
                out_channels=out_channels[1],
                kernel_size=2,
                stride=2,
                padding=0),
            nn.Identity(),
            nn.Conv2d(
                in_channels=out_channels[3],
                out_channels=out_channels[3],
                kernel_size=3,
                stride=2,
                padding=1)
        ])

        if use_clstoken:
            self.readout_projects = nn.ModuleList()
            for _ in range(len(self.projects)):
                self.readout_projects.append(
                    nn.Sequential(
                        nn.Linear(2 * in_channels, in_channels),
                        nn.GELU()))

        self.scratch = _make_scratch(
            out_channels,
            features,
            groups=1,
            expand=False,
        )

        self.scratch.stem_transpose = None

        self.scratch.refinenet1 = _make_fusion_block(
            features, use_bn)
        self.scratch.refinenet2 = _make_fusion_block(
            features, use_bn)
        self.scratch.refinenet3 = _make_fusion_block(
            features, use_bn)
        self.scratch.refinenet4 = _make_fusion_block(
            features, use_bn)

        head_features_1 = features
        head_features_2 = 32

        act_func = nn.Sigmoid() if output_act == 'sigmoid' else nn.Identity()

        if nclass > 1:
            self.scratch.output_conv = nn.Sequential(
                nn.Conv2d(head_features_1, head_features_1,
                          kernel_size=3, stride=1, padding=1),
                nn.ReLU(True),
                nn.Conv2d(head_features_1, nclass,
                          kernel_size=1, stride=1, padding=0),
            )
        else:
            self.scratch.output_conv1 = nn.Conv2d(
                head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1)

            self.scratch.output_conv2 = nn.Sequential(
                nn.Conv2d(head_features_1 // 2, head_features_2,
                          kernel_size=3, stride=1, padding=1),
                nn.ReLU(True),
                nn.Conv2d(head_features_2, 1, kernel_size=1,
                          stride=1, padding=0),
                act_func,
            )

    def forward(self, out_features, patch_h, patch_w, prompt_depth=None):
        out = []
        for i, x in enumerate(out_features):
            if self.use_clstoken:
                x, cls_token = x[0], x[1]
                readout = cls_token.unsqueeze(1).expand_as(x)
                x = self.readout_projects[i](torch.cat((x, readout), -1))
            else:
                x = x[0]

            x = x.permute(0, 2, 1).reshape(
                (x.shape[0], x.shape[-1], patch_h, patch_w))

            x = self.projects[i](x)
            x = self.resize_layers[i](x)

            out.append(x)

        layer_1, layer_2, layer_3, layer_4 = out

        layer_1_rn = self.scratch.layer1_rn(layer_1)
        layer_2_rn = self.scratch.layer2_rn(layer_2)
        layer_3_rn = self.scratch.layer3_rn(layer_3)
        layer_4_rn = self.scratch.layer4_rn(layer_4)

        path_4 = self.scratch.refinenet4(
            layer_4_rn, size=layer_3_rn.shape[2:], prompt_depth=prompt_depth)
        path_3 = self.scratch.refinenet3(
            path_4, layer_3_rn, size=layer_2_rn.shape[2:], prompt_depth=prompt_depth)
        path_2 = self.scratch.refinenet2(
            path_3, layer_2_rn, size=layer_1_rn.shape[2:], prompt_depth=prompt_depth)
        path_1 = self.scratch.refinenet1(
            path_2, layer_1_rn, prompt_depth=prompt_depth)
        out = self.scratch.output_conv1(path_1)
        out_feat = F.interpolate(
            out, (int(patch_h * 14), int(patch_w * 14)),
            mode="bilinear", align_corners=True)
        out = self.scratch.output_conv2(out_feat)
        return out


# ==============================================================================================================
# PromptDA/promptda/model/config.py
# ==============================================================================================================
model_configs = {
    'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384], 'layer_idxs': [2, 5, 8, 11]},
    'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768], 'layer_idxs': [2, 5, 8, 11]},
    'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024], 'layer_idxs': [4, 11, 17, 23]},
    'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536], 'layer_idxs': [9, 19, 29, 39]}
}


# ==============================================================================================================
# PromptDA/promptda/promptda.py
# ==============================================================================================================
import torch
import torch.nn as nn
import os
from pathlib import Path
from huggingface_hub import hf_hub_download


class PromptDA(nn.Module):
    patch_size = 14  # patch size of the pretrained dinov2 model
    use_bn = False
    use_clstoken = False
    output_act = 'sigmoid'

    def __init__(self,
                 encoder='vitl',
                 ckpt_path='checkpoints/promptda_vitl.ckpt'):
        super().__init__()
        model_config = model_configs[encoder]

        self.encoder = encoder
        self.model_config = model_config
        module_path = Path(__file__) # From anywhere else: module_path = Path(inspect.getfile(PromptDA))
        package_base_dir = str(Path(*module_path.parts[:-2])) # extract path to PromptDA module, then resolve to repo base dir for dinov2 load
        self.pretrained = torch.hub.load(
            f'{package_base_dir}/torchhub/facebookresearch_dinov2_main',
            'dinov2_{:}14'.format(encoder),
            source='local',
            pretrained=False)
        dim = self.pretrained.blocks[0].attn.qkv.in_features
        self.depth_head = DPTHead(nclass=1,
                                  in_channels=dim,
                                  features=model_config['features'],
                                  out_channels=model_config['out_channels'],
                                  use_bn=self.use_bn,
                                  use_clstoken=self.use_clstoken,
                                  output_act=self.output_act)

        # mean and std of the pretrained dinov2 model
        self.register_buffer('_mean', torch.tensor(
            [0.485, 0.456, 0.406]).view(1, 3, 1, 1))
        self.register_buffer('_std', torch.tensor(
            [0.229, 0.224, 0.225]).view(1, 3, 1, 1))

        self.load_checkpoint(ckpt_path)
    
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path = None, model_kwargs = None, **hf_kwargs):
        """
        Load a model from a checkpoint file.
        ### Parameters:
        - `pretrained_model_name_or_path`: path to the checkpoint file or repo id.
        - `model_kwargs`: additional keyword arguments to override the parameters in the checkpoint.
        - `hf_kwargs`: additional keyword arguments to pass to the `hf_hub_download` function. Ignored if `pretrained_model_name_or_path` is a local path.
        ### Returns:
        - A new instance of `MoGe` with the parameters loaded from the checkpoint.
        """
        ckpt_path = None
        if Path(pretrained_model_name_or_path).exists():
            ckpt_path = pretrained_model_name_or_path
        else:
            cached_checkpoint_path = hf_hub_download(
                repo_id=pretrained_model_name_or_path,
                repo_type="model",
                filename="model.ckpt",
                **hf_kwargs
            )
            ckpt_path = cached_checkpoint_path
        # model_config = checkpoint['model_config']
        # if model_kwargs is not None:
        #     model_config.update(model_kwargs)
        if model_kwargs is None:
            model_kwargs = {}
        model_kwargs.update({'ckpt_path': ckpt_path})
        model = cls(**model_kwargs)
        return model

    def load_checkpoint(self, ckpt_path):
        if os.path.exists(ckpt_path):
            Log.info(f'Loading checkpoint from {ckpt_path}')
            checkpoint = torch.load(ckpt_path, map_location='cpu')
            self.load_state_dict(
                {k[9:]: v for k, v in checkpoint['state_dict'].items()})
        else:
            Log.warn(f'Checkpoint {ckpt_path} not found')

    def forward(self, x, prompt_depth=None):
        assert prompt_depth is not None, 'prompt_depth is required'
        prompt_depth, min_val, max_val = self.normalize(prompt_depth)
        h, w = x.shape[-2:]
        features = self.pretrained.get_intermediate_layers(
            (x - self._mean) / self._std, self.model_config['layer_idxs'],
            return_class_token=True)
        patch_h, patch_w = h // self.patch_size, w // self.patch_size
        depth = self.depth_head(features, patch_h, patch_w, prompt_depth)
        depth = self.denormalize(depth, min_val, max_val)
        return depth

    @torch.no_grad()
    def predict(self,
                image: torch.Tensor,
                prompt_depth: torch.Tensor):
        return self.forward(image, prompt_depth)

    def normalize(self,
                  prompt_depth: torch.Tensor):
        B, C, H, W = prompt_depth.shape
        min_val = torch.quantile(
            prompt_depth.reshape(B, -1), 0., dim=1, keepdim=True)[:, :, None, None]
        max_val = torch.quantile(
            prompt_depth.reshape(B, -1), 1., dim=1, keepdim=True)[:, :, None, None]
        prompt_depth = (prompt_depth - min_val) / (max_val - min_val)
        return prompt_depth, min_val, max_val

    def denormalize(self,
                    depth: torch.Tensor,
                    min_val: torch.Tensor,
                    max_val: torch.Tensor):
        return depth * (max_val - min_val) + min_val

	# ruff: noqa
	# fmt: off
	"""Prompt Depth Anything, vendored verbatim from DepthAnything/PromptDA @ 9161547ca7151e14fcc8f490ad3cc02fa9437483 (promptda/ and the local torchhub
	copy of facebookresearch/dinov2 it loads): the files below concatenated, package-internal imports replaced by
	relative imports of the other vendored files (generated by scripts/model_ports/common/vendor_flat.py; do not edit by hand). Code license: Apache-2.0."""


	# ==============================================================================================================
	# PromptDA/promptda/utils/logger.py
	# ==============================================================================================================
	import os


	class Log:
	log_on = True # fast switch
	used_tags = dict() # To keep track of used tags
	_is_main_cached = None # Cache to store the main process check result

	@staticmethod
	def is_main_process():
	if Log._is_main_cached is not None:
	return Log._is_main_cached
	try:
	from pytorch_lightning.utilities import rank_zero_only
	if rank_zero_only.rank == 0:
	Log._is_main_cached = True
	else:
	Log._is_main_cached = False
	except:
	Log._is_main_cached = True
	return Log._is_main_cached

	@staticmethod
	def _should_log(tag):
	"""
	Determine if the log information should be recorded.
	Conditions: log function is enabled, current process is the main process, and the tag has not been used.
	"""
	if not Log.log_on:
	return False
	if not Log.is_main_process():
	return False
	if tag is None:
	return True
	if '__' in tag:
	num = int(tag.split('__')[-1])
	tag = tag.split('__')[0] # can output num same information
	else:
	num = 3 # default 3

	if tag not in Log.used_tags:
	Log.used_tags[tag] = num
	Log.used_tags[tag] -= 1
	if Log.used_tags[tag] >= 0:
	return True
	else:
	return False

	@staticmethod
	def info(*args, tag=None):
	"""
	Output INFO level log information.
	"""
	if Log._should_log(tag):
	print("\033[1;32m[INFO]\033[0;0m", *args)

	@staticmethod
	def warn(*args, tag=None):
	"""
	Output WARN level log information.
	"""
	if Log._should_log(tag):
	print("\033[1;35m[WARN]\033[0;0m", *args)

	@staticmethod
	def error(*args, tag=None):
	print("\033[1;31m[ERROR]\033[0;0m", *args)

	@staticmethod
	def debug(*args, tag=None):
	"""
	Output DEBUG level log information.
	"""
	if Log._should_log(tag) and 'HT_DEBUG' in os.environ and os.environ['HT_DEBUG'] == '1':
	print("\033[1;33m[DEBUG]\033[0;0m", *args)


	# ==============================================================================================================
	# PromptDA/promptda/model/blocks.py
	# ==============================================================================================================
	import torch.nn as nn
	import torch
	import torch.nn.functional as F
	import os
	import numpy as np


	def _make_fusion_block(features, use_bn, size=None):
	return FeatureFusionDepthBlock(
	features,
	nn.ReLU(False),
	deconv=False,
	bn=use_bn,
	expand=False,
	align_corners=True,
	size=size,
	)


	def _make_scratch(in_shape, out_shape, groups=1, expand=False):
	scratch = nn.Module()

	out_shape1 = out_shape
	out_shape2 = out_shape
	out_shape3 = out_shape
	if len(in_shape) >= 4:
	out_shape4 = out_shape

	if expand:
	out_shape1 = out_shape
	out_shape2 = out_shape*2
	out_shape3 = out_shape*4
	if len(in_shape) >= 4:
	out_shape4 = out_shape*8

	scratch.layer1_rn = nn.Conv2d(
	in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
	)
	scratch.layer2_rn = nn.Conv2d(
	in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
	)
	scratch.layer3_rn = nn.Conv2d(
	in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
	)
	if len(in_shape) >= 4:
	scratch.layer4_rn = nn.Conv2d(
	in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
	)

	return scratch


	class ResidualConvUnit(nn.Module):
	"""Residual convolution module.
	"""

	def __init__(self, features, activation, bn):
	"""Init.

	Args:
	features (int): number of features
	"""
	super().__init__()

	self.bn = bn

	self.groups = 1

	self.conv1 = nn.Conv2d(
	features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
	)

	self.conv2 = nn.Conv2d(
	features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups
	)

	if self.bn == True:
	self.bn1 = nn.BatchNorm2d(features)
	self.bn2 = nn.BatchNorm2d(features)

	self.activation = activation

	self.skip_add = nn.quantized.FloatFunctional()

	def forward(self, x):
	"""Forward pass.

	Args:
	x (tensor): input

	Returns:
	tensor: output
	"""

	out = self.activation(x)
	out = self.conv1(out)
	if self.bn == True:
	out = self.bn1(out)

	out = self.activation(out)
	out = self.conv2(out)
	if self.bn == True:
	out = self.bn2(out)

	if self.groups > 1:
	out = self.conv_merge(out)

	return self.skip_add.add(out, x)


	class FeatureFusionBlock(nn.Module):
	"""Feature fusion block.
	"""

	def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, size=None):
	"""Init.

	Args:
	features (int): number of features
	"""
	super(FeatureFusionBlock, self).__init__()

	self.deconv = deconv
	self.align_corners = align_corners

	self.groups = 1

	self.expand = expand
	out_features = features
	if self.expand == True:
	out_features = features//2

	self.out_conv = nn.Conv2d(
	features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)

	self.resConfUnit1 = ResidualConvUnit(features, activation, bn)
	self.resConfUnit2 = ResidualConvUnit(features, activation, bn)

	self.skip_add = nn.quantized.FloatFunctional()

	self.size = size

	def forward(self, *xs, size=None):
	"""Forward pass.

	Returns:
	tensor: output
	"""
	output = xs[0]

	if len(xs) == 2:
	res = self.resConfUnit1(xs[1])
	output = self.skip_add.add(output, res)

	output = self.resConfUnit2(output)

	if (size is None) and (self.size is None):
	modifier = {"scale_factor": 2}
	elif size is None:
	modifier = {"size": self.size}
	else:
	modifier = {"size": size}

	output = nn.functional.interpolate(
	output, **modifier, mode="bilinear", align_corners=self.align_corners
	)

	output = self.out_conv(output)

	return output


	class FeatureFusionControlBlock(FeatureFusionBlock):
	"""Feature fusion block.
	"""

	def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, size=None):
	"""Init.

	Args:
	features (int): number of features
	"""
	super.__init__(features, activation, deconv,
	bn, expand, align_corners, size)
	self.copy_block = FeatureFusionBlock(
	features, activation, deconv, bn, expand, align_corners, size)

	def forward(self, *xs, size=None):
	"""Forward pass.

	Returns:
	tensor: output
	"""
	output = xs[0]

	if len(xs) == 2:
	res = self.resConfUnit1(xs[1])
	output = self.skip_add.add(output, res)

	output = self.resConfUnit2(output)

	if (size is None) and (self.size is None):
	modifier = {"scale_factor": 2}
	elif size is None:
	modifier = {"size": self.size}
	else:
	modifier = {"size": size}

	output = nn.functional.interpolate(
	output, **modifier, mode="bilinear", align_corners=self.align_corners
	)

	output = self.out_conv(output)

	return output


	def zero_module(module):
	"""
	Zero out the parameters of a module and return it.
	"""
	for p in module.parameters():
	p.detach().zero_()
	return module


	class FeatureFusionDepthBlock(nn.Module):
	"""Feature fusion block.
	"""

	def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, size=None):
	"""Init.

	Args:
	features (int): number of features
	"""
	super(FeatureFusionDepthBlock, self).__init__()

	self.deconv = deconv
	self.align_corners = align_corners

	self.groups = 1

	self.expand = expand
	out_features = features
	if self.expand == True:
	out_features = features//2

	self.out_conv = nn.Conv2d(
	features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)

	self.resConfUnit1 = ResidualConvUnit(features, activation, bn)
	self.resConfUnit2 = ResidualConvUnit(features, activation, bn)
	self.resConfUnit_depth = nn.Sequential(
	nn.Conv2d(1, features, kernel_size=3, stride=1,
	padding=1, bias=True, groups=1),
	activation,
	nn.Conv2d(features, features, kernel_size=3,
	stride=1, padding=1, bias=True, groups=1),
	activation,
	zero_module(
	nn.Conv2d(features, features, kernel_size=3,
	stride=1, padding=1, bias=True, groups=1)
	)
	)
	self.skip_add = nn.quantized.FloatFunctional()
	self.size = size

	def forward(self, *xs, prompt_depth=None, size=None):
	"""Forward pass.

	Returns:
	tensor: output
	"""
	output = xs[0]

	if len(xs) == 2:
	res = self.resConfUnit1(xs[1])
	output = self.skip_add.add(output, res)

	output = self.resConfUnit2(output)

	if prompt_depth is not None:
	prompt_depth = F.interpolate(
	prompt_depth, output.shape[2:], mode='bilinear', align_corners=False)
	res = self.resConfUnit_depth(prompt_depth)
	output = self.skip_add.add(output, res)

	if (size is None) and (self.size is None):
	modifier = {"scale_factor": 2}
	elif size is None:
	modifier = {"size": self.size}
	else:
	modifier = {"size": size}

	output = nn.functional.interpolate(
	output, **modifier, mode="bilinear", align_corners=self.align_corners
	)

	output = self.out_conv(output)

	return output


	# ==============================================================================================================
	# PromptDA/promptda/model/dpt.py
	# ==============================================================================================================
	# Copyright (c) 2024, Depth Anything V2
	# https://github.com/DepthAnything/Depth-Anything-V2/blob/main/depth_anything_v2/dpt.py
	import torch
	import torch.nn as nn
	import torch.nn.functional as F


	class DPTHead(nn.Module):
	def __init__(self,
	nclass,
	in_channels,
	features=256,
	out_channels=[256, 512, 1024, 1024],
	use_bn=False,
	use_clstoken=False,
	output_act='sigmoid'):
	super(DPTHead, self).__init__()

	self.nclass = nclass
	self.use_clstoken = use_clstoken

	self.projects = nn.ModuleList([
	nn.Conv2d(
	in_channels=in_channels,
	out_channels=out_channel,
	kernel_size=1,
	stride=1,
	padding=0,
	) for out_channel in out_channels
	])

	self.resize_layers = nn.ModuleList([
	nn.ConvTranspose2d(
	in_channels=out_channels[0],
	out_channels=out_channels[0],
	kernel_size=4,
	stride=4,
	padding=0),
	nn.ConvTranspose2d(
	in_channels=out_channels[1],
	out_channels=out_channels[1],
	kernel_size=2,
	stride=2,
	padding=0),
	nn.Identity(),
	nn.Conv2d(
	in_channels=out_channels[3],
	out_channels=out_channels[3],
	kernel_size=3,
	stride=2,
	padding=1)
	])

	if use_clstoken:
	self.readout_projects = nn.ModuleList()
	for _ in range(len(self.projects)):
	self.readout_projects.append(
	nn.Sequential(
	nn.Linear(2 * in_channels, in_channels),
	nn.GELU()))

	self.scratch = _make_scratch(
	out_channels,
	features,
	groups=1,
	expand=False,
	)

	self.scratch.stem_transpose = None

	self.scratch.refinenet1 = _make_fusion_block(
	features, use_bn)
	self.scratch.refinenet2 = _make_fusion_block(
	features, use_bn)
	self.scratch.refinenet3 = _make_fusion_block(
	features, use_bn)
	self.scratch.refinenet4 = _make_fusion_block(
	features, use_bn)

	head_features_1 = features
	head_features_2 = 32

	act_func = nn.Sigmoid() if output_act == 'sigmoid' else nn.Identity()

	if nclass > 1:
	self.scratch.output_conv = nn.Sequential(
	nn.Conv2d(head_features_1, head_features_1,
	kernel_size=3, stride=1, padding=1),
	nn.ReLU(True),
	nn.Conv2d(head_features_1, nclass,
	kernel_size=1, stride=1, padding=0),
	)
	else:
	self.scratch.output_conv1 = nn.Conv2d(
	head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1)

	self.scratch.output_conv2 = nn.Sequential(
	nn.Conv2d(head_features_1 // 2, head_features_2,
	kernel_size=3, stride=1, padding=1),
	nn.ReLU(True),
	nn.Conv2d(head_features_2, 1, kernel_size=1,
	stride=1, padding=0),
	act_func,
	)

	def forward(self, out_features, patch_h, patch_w, prompt_depth=None):
	out = []
	for i, x in enumerate(out_features):
	if self.use_clstoken:
	x, cls_token = x[0], x[1]
	readout = cls_token.unsqueeze(1).expand_as(x)
	x = self.readout_projects[i](torch.cat((x, readout), -1))
	else:
	x = x[0]

	x = x.permute(0, 2, 1).reshape(
	(x.shape[0], x.shape[-1], patch_h, patch_w))

	x = self.projects[i](x)
	x = self.resize_layers[i](x)

	out.append(x)

	layer_1, layer_2, layer_3, layer_4 = out

	layer_1_rn = self.scratch.layer1_rn(layer_1)
	layer_2_rn = self.scratch.layer2_rn(layer_2)
	layer_3_rn = self.scratch.layer3_rn(layer_3)
	layer_4_rn = self.scratch.layer4_rn(layer_4)

	path_4 = self.scratch.refinenet4(
	layer_4_rn, size=layer_3_rn.shape[2:], prompt_depth=prompt_depth)
	path_3 = self.scratch.refinenet3(
	path_4, layer_3_rn, size=layer_2_rn.shape[2:], prompt_depth=prompt_depth)
	path_2 = self.scratch.refinenet2(
	path_3, layer_2_rn, size=layer_1_rn.shape[2:], prompt_depth=prompt_depth)
	path_1 = self.scratch.refinenet1(
	path_2, layer_1_rn, prompt_depth=prompt_depth)
	out = self.scratch.output_conv1(path_1)
	out_feat = F.interpolate(
	out, (int(patch_h * 14), int(patch_w * 14)),
	mode="bilinear", align_corners=True)
	out = self.scratch.output_conv2(out_feat)
	return out


	# ==============================================================================================================
	# PromptDA/promptda/model/config.py
	# ==============================================================================================================
	model_configs = {
	'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384], 'layer_idxs': [2, 5, 8, 11]},
	'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768], 'layer_idxs': [2, 5, 8, 11]},
	'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024], 'layer_idxs': [4, 11, 17, 23]},
	'vitg': {'encoder': 'vitg', 'features': 384, 'out_channels': [1536, 1536, 1536, 1536], 'layer_idxs': [9, 19, 29, 39]}
	}


	# ==============================================================================================================
	# PromptDA/promptda/promptda.py
	# ==============================================================================================================
	import torch
	import torch.nn as nn
	import os
	from pathlib import Path
	from huggingface_hub import hf_hub_download


	class PromptDA(nn.Module):
	patch_size = 14 # patch size of the pretrained dinov2 model
	use_bn = False
	use_clstoken = False
	output_act = 'sigmoid'

	def __init__(self,
	encoder='vitl',
	ckpt_path='checkpoints/promptda_vitl.ckpt'):
	super().__init__()
	model_config = model_configs[encoder]

	self.encoder = encoder
	self.model_config = model_config
	module_path = Path(__file__) # From anywhere else: module_path = Path(inspect.getfile(PromptDA))
	package_base_dir = str(Path(*module_path.parts[:-2])) # extract path to PromptDA module, then resolve to repo base dir for dinov2 load
	self.pretrained = torch.hub.load(
	f'{package_base_dir}/torchhub/facebookresearch_dinov2_main',
	'dinov2_{:}14'.format(encoder),
	source='local',
	pretrained=False)
	dim = self.pretrained.blocks[0].attn.qkv.in_features
	self.depth_head = DPTHead(nclass=1,
	in_channels=dim,
	features=model_config['features'],
	out_channels=model_config['out_channels'],
	use_bn=self.use_bn,
	use_clstoken=self.use_clstoken,
	output_act=self.output_act)

	# mean and std of the pretrained dinov2 model
	self.register_buffer('_mean', torch.tensor(
	[0.485, 0.456, 0.406]).view(1, 3, 1, 1))
	self.register_buffer('_std', torch.tensor(
	[0.229, 0.224, 0.225]).view(1, 3, 1, 1))

	self.load_checkpoint(ckpt_path)

	@classmethod
	def from_pretrained(cls, pretrained_model_name_or_path = None, model_kwargs = None, **hf_kwargs):
	"""
	Load a model from a checkpoint file.
	### Parameters:
	- `pretrained_model_name_or_path`: path to the checkpoint file or repo id.
	- `model_kwargs`: additional keyword arguments to override the parameters in the checkpoint.
	- `hf_kwargs`: additional keyword arguments to pass to the `hf_hub_download` function. Ignored if `pretrained_model_name_or_path` is a local path.
	### Returns:
	- A new instance of `MoGe` with the parameters loaded from the checkpoint.
	"""
	ckpt_path = None
	if Path(pretrained_model_name_or_path).exists():
	ckpt_path = pretrained_model_name_or_path
	else:
	cached_checkpoint_path = hf_hub_download(
	repo_id=pretrained_model_name_or_path,
	repo_type="model",
	filename="model.ckpt",
	**hf_kwargs
	)
	ckpt_path = cached_checkpoint_path
	# model_config = checkpoint['model_config']
	# if model_kwargs is not None:
	# model_config.update(model_kwargs)
	if model_kwargs is None:
	model_kwargs = {}
	model_kwargs.update({'ckpt_path': ckpt_path})
	model = cls(**model_kwargs)
	return model

	def load_checkpoint(self, ckpt_path):
	if os.path.exists(ckpt_path):
	Log.info(f'Loading checkpoint from {ckpt_path}')
	checkpoint = torch.load(ckpt_path, map_location='cpu')
	self.load_state_dict(
	{k[9:]: v for k, v in checkpoint['state_dict'].items()})
	else:
	Log.warn(f'Checkpoint {ckpt_path} not found')

	def forward(self, x, prompt_depth=None):
	assert prompt_depth is not None, 'prompt_depth is required'
	prompt_depth, min_val, max_val = self.normalize(prompt_depth)
	h, w = x.shape[-2:]
	features = self.pretrained.get_intermediate_layers(
	(x - self._mean) / self._std, self.model_config['layer_idxs'],
	return_class_token=True)
	patch_h, patch_w = h // self.patch_size, w // self.patch_size
	depth = self.depth_head(features, patch_h, patch_w, prompt_depth)
	depth = self.denormalize(depth, min_val, max_val)
	return depth

	@torch.no_grad()
	def predict(self,
	image: torch.Tensor,
	prompt_depth: torch.Tensor):
	return self.forward(image, prompt_depth)

	def normalize(self,
	prompt_depth: torch.Tensor):
	B, C, H, W = prompt_depth.shape
	min_val = torch.quantile(
	prompt_depth.reshape(B, -1), 0., dim=1, keepdim=True)[:, :, None, None]
	max_val = torch.quantile(
	prompt_depth.reshape(B, -1), 1., dim=1, keepdim=True)[:, :, None, None]
	prompt_depth = (prompt_depth - min_val) / (max_val - min_val)
	return prompt_depth, min_val, max_val

	def denormalize(self,
	depth: torch.Tensor,
	min_val: torch.Tensor,
	max_val: torch.Tensor):
	return depth * (max_val - min_val) + min_val