in src/controlnet_aux/zoe/zoedepth/models/zoedepth_nk/zoedepth_nk_v1.py [0:0]
def __init__(self, core, bin_conf, bin_centers_type="softplus", bin_embedding_dim=128,
n_attractors=[16, 8, 4, 1], attractor_alpha=300, attractor_gamma=2, attractor_kind='sum', attractor_type='exp',
min_temp=5, max_temp=50,
memory_efficient=False, train_midas=True,
is_midas_pretrained=True, midas_lr_factor=1, encoder_lr_factor=10, pos_enc_lr_factor=10, inverse_midas=False, **kwargs):
"""ZoeDepthNK model. This is the version of ZoeDepth that has two metric heads and uses a learned router to route to experts.
Args:
core (models.base_models.midas.MidasCore): The base midas model that is used for extraction of "relative" features
bin_conf (List[dict]): A list of dictionaries that contain the bin configuration for each metric head. Each dictionary should contain the following keys:
"name" (str, typically same as the dataset name), "n_bins" (int), "min_depth" (float), "max_depth" (float)
The length of this list determines the number of metric heads.
bin_centers_type (str, optional): "normed" or "softplus". Activation type used for bin centers. For "normed" bin centers, linear normalization trick is applied. This results in bounded bin centers.
For "softplus", softplus activation is used and thus are unbounded. Defaults to "normed".
bin_embedding_dim (int, optional): bin embedding dimension. Defaults to 128.
n_attractors (List[int], optional): Number of bin attractors at decoder layers. Defaults to [16, 8, 4, 1].
attractor_alpha (int, optional): Proportional attractor strength. Refer to models.layers.attractor for more details. Defaults to 300.
attractor_gamma (int, optional): Exponential attractor strength. Refer to models.layers.attractor for more details. Defaults to 2.
attractor_kind (str, optional): Attraction aggregation "sum" or "mean". Defaults to 'sum'.
attractor_type (str, optional): Type of attractor to use; "inv" (Inverse attractor) or "exp" (Exponential attractor). Defaults to 'exp'.
min_temp (int, optional): Lower bound for temperature of output probability distribution. Defaults to 5.
max_temp (int, optional): Upper bound for temperature of output probability distribution. Defaults to 50.
memory_efficient (bool, optional): Whether to use memory efficient version of attractor layers. Memory efficient version is slower but is recommended incase of multiple metric heads in order save GPU memory. Defaults to False.
train_midas (bool, optional): Whether to train "core", the base midas model. Defaults to True.
is_midas_pretrained (bool, optional): Is "core" pretrained? Defaults to True.
midas_lr_factor (int, optional): Learning rate reduction factor for base midas model except its encoder and positional encodings. Defaults to 10.
encoder_lr_factor (int, optional): Learning rate reduction factor for the encoder in midas model. Defaults to 10.
pos_enc_lr_factor (int, optional): Learning rate reduction factor for positional encodings in the base midas model. Defaults to 10.
"""
super().__init__()
self.core = core
self.bin_conf = bin_conf
self.min_temp = min_temp
self.max_temp = max_temp
self.memory_efficient = memory_efficient
self.train_midas = train_midas
self.is_midas_pretrained = is_midas_pretrained
self.midas_lr_factor = midas_lr_factor
self.encoder_lr_factor = encoder_lr_factor
self.pos_enc_lr_factor = pos_enc_lr_factor
self.inverse_midas = inverse_midas
N_MIDAS_OUT = 32
btlnck_features = self.core.output_channels[0]
num_out_features = self.core.output_channels[1:]
# self.scales = [16, 8, 4, 2] # spatial scale factors
self.conv2 = nn.Conv2d(
btlnck_features, btlnck_features, kernel_size=1, stride=1, padding=0)
# Transformer classifier on the bottleneck
self.patch_transformer = PatchTransformerEncoder(
btlnck_features, 1, 128, use_class_token=True)
self.mlp_classifier = nn.Sequential(
nn.Linear(128, 128),
nn.ReLU(),
nn.Linear(128, 2)
)
if bin_centers_type == "normed":
SeedBinRegressorLayer = SeedBinRegressor
Attractor = AttractorLayer
elif bin_centers_type == "softplus":
SeedBinRegressorLayer = SeedBinRegressorUnnormed
Attractor = AttractorLayerUnnormed
elif bin_centers_type == "hybrid1":
SeedBinRegressorLayer = SeedBinRegressor
Attractor = AttractorLayerUnnormed
elif bin_centers_type == "hybrid2":
SeedBinRegressorLayer = SeedBinRegressorUnnormed
Attractor = AttractorLayer
else:
raise ValueError(
"bin_centers_type should be one of 'normed', 'softplus', 'hybrid1', 'hybrid2'")
self.bin_centers_type = bin_centers_type
# We have bins for each bin conf.
# Create a map (ModuleDict) of 'name' -> seed_bin_regressor
self.seed_bin_regressors = nn.ModuleDict(
{conf['name']: SeedBinRegressorLayer(btlnck_features, conf["n_bins"], mlp_dim=bin_embedding_dim//2, min_depth=conf["min_depth"], max_depth=conf["max_depth"])
for conf in bin_conf}
)
self.seed_projector = Projector(
btlnck_features, bin_embedding_dim, mlp_dim=bin_embedding_dim//2)
self.projectors = nn.ModuleList([
Projector(num_out, bin_embedding_dim, mlp_dim=bin_embedding_dim//2)
for num_out in num_out_features
])
# Create a map (ModuleDict) of 'name' -> attractors (ModuleList)
self.attractors = nn.ModuleDict(
{conf['name']: nn.ModuleList([
Attractor(bin_embedding_dim, n_attractors[i],
mlp_dim=bin_embedding_dim, alpha=attractor_alpha,
gamma=attractor_gamma, kind=attractor_kind,
attractor_type=attractor_type, memory_efficient=memory_efficient,
min_depth=conf["min_depth"], max_depth=conf["max_depth"])
for i in range(len(n_attractors))
])
for conf in bin_conf}
)
last_in = N_MIDAS_OUT
# conditional log binomial for each bin conf
self.conditional_log_binomial = nn.ModuleDict(
{conf['name']: ConditionalLogBinomial(last_in, bin_embedding_dim, conf['n_bins'], bottleneck_factor=4, min_temp=self.min_temp, max_temp=self.max_temp)
for conf in bin_conf}
)