jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(
        self,
        hidden_states: List[Tensor],
        contexts: List[List[Tensor]],
        corr_features: Tensor,
        disparity: Tensor,
        level_processed: List[bool],
    ) -> List[Tensor]:
        # We call it reverse_i because it has a reversed ordering compared to hidden_states
        # see self.grus on the constructor for more detail
        for reverse_i, gru in enumerate(self.grus):
            i = len(self.grus) - 1 - reverse_i
            if level_processed[i]:
                # X is concatenation of 2x downsampled hidden_dim (or motion_features if no bigger dim) with
                # upsampled hidden_dim (or nothing if not exist).
                if i == 0:
                    features = self.motion_encoder(disparity, corr_features)
                else:
                    # 2x downsampled features from larger hidden states
                    features = F.avg_pool2d(hidden_states[i - 1], kernel_size=3, stride=2, padding=1)

                if i < len(self.grus) - 1:
                    # Concat with 2x upsampled features from smaller hidden states
                    _, _, h, w = hidden_states[i + 1].shape
                    features = torch.cat(
                        [
                            features,
                            F.interpolate(
                                hidden_states[i + 1], size=(2 * h, 2 * w), mode="bilinear", align_corners=True
                            ),
                        ],
                        dim=1,
                    )

                hidden_states[i] = gru(hidden_states[i], features, contexts[i])

                # NOTE: For slow-fast gru, we don't always want to calculate delta disparity for every call on UpdateBlock
                # Hence we move the delta disparity calculation to the RAFT-Stereo main forward

        return hidden_states
//  [python]
def p_losses(self, x_start, cond, t, noise=None, mask=None):
        noise = default(noise, lambda: torch.randn_like(x_start))
        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
        model_output = self.apply_model(x_noisy, t, cond)

        loss_dict = {}
        prefix = 'train' if self.training else 'val'

        if self.parameterization == "x0":
            target = x_start
        elif self.parameterization == "eps":
            target = noise
        elif self.parameterization == "v":
            target = self.get_v(x_start, noise, t)
        else:
            raise NotImplementedError()

        if mask is not None:
            mask = F.interpolate(mask, model_output.size()[-2:], mode='bilinear')
            model_output = model_output * mask
            target = target * mask

        loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3])
        loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})

        logvar_t = self.logvar[t].to(self.device)
        loss = loss_simple / torch.exp(logvar_t) + logvar_t
        # loss = loss_simple / torch.exp(self.logvar) + self.logvar
        if self.learn_logvar:
            loss_dict.update({f'{prefix}/loss_gamma': loss.mean()})
            loss_dict.update({'logvar': self.logvar.data.mean()})

        loss = self.l_simple_weight * loss.mean()

        loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3))
        loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean()
        loss_dict.update({f'{prefix}/loss_vlb': loss_vlb})
        loss += (self.original_elbo_weight * loss_vlb)
        loss_dict.update({f'{prefix}/loss': loss})

        return loss, loss_dict
//  [python]
def collate_fn4(batch):
        img, label, path, shapes = zip(*batch)  # transposed
        n = len(shapes) // 4
        im4, label4, path4, shapes4 = [], [], path[:n], shapes[:n]

        ho = torch.tensor([[0.0, 0, 0, 1, 0, 0]])
        wo = torch.tensor([[0.0, 0, 1, 0, 0, 0]])
        s = torch.tensor([[1, 1, 0.5, 0.5, 0.5, 0.5]])  # scale
        for i in range(n):  # zidane torch.zeros(16,3,720,1280)  # BCHW
            i *= 4
            if random.random() < 0.5:
                im = t.interpolate(img[i].unsqueeze(0).float(), scale_factor=2.0, mode='bilinear',
                                   align_corners=False)[0].type(img[i].type())
                lb = label[i]
            else:
                im = torch.cat((torch.cat((img[i], img[i + 1]), 1), torch.cat((img[i + 2], img[i + 3]), 1)), 2)
                lb = torch.cat((label[i], label[i + 1] + ho, label[i + 2] + wo, label[i + 3] + ho + wo), 0) * s
            im4.append(im)
            label4.append(lb)

        for i, lb in enumerate(label4):
            lb[:, 0] = i  # add target image index for build_targets()

        return torch.stack(im4, 0), torch.cat(label4, 0), path4, shapes4
//  [python]
def collate_fn4(batch):
        img, label, path, shapes = zip(*batch)  # transposed
        n = len(shapes) // 4
        im4, label4, path4, shapes4 = [], [], path[:n], shapes[:n]

        ho = torch.tensor([[0.0, 0, 0, 1, 0, 0]])
        wo = torch.tensor([[0.0, 0, 1, 0, 0, 0]])
        s = torch.tensor([[1, 1, 0.5, 0.5, 0.5, 0.5]])  # scale
        for i in range(n):  # zidane torch.zeros(16,3,720,1280)  # BCHW
            i *= 4
            if random.random() < 0.5:
                im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2.0, mode='bilinear',
                                   align_corners=False)[0].type(img[i].type())
                lb = label[i]
            else:
                im = torch.cat((torch.cat((img[i], img[i + 1]), 1), torch.cat((img[i + 2], img[i + 3]), 1)), 2)
                lb = torch.cat((label[i], label[i + 1] + ho, label[i + 2] + wo, label[i + 3] + ho + wo), 0) * s
            im4.append(im)
            label4.append(lb)

        for i, lb in enumerate(label4):
            lb[:, 0] = i  # add target image index for build_targets()

        return torch.stack(im4, 0), torch.cat(label4, 0), path4, shapes4
//  [python]
def forward(self,
              input,
              alpha,
              options=None,
              **kwargs):

    start = self.img_size_to_layer[input.shape[-1]]
    x = self.fromRGB[start](input)

    if kwargs.get('instance_noise', 0) > 0:
      x = x + torch.randn_like(x) * kwargs['instance_noise']

    for i, layer in enumerate(self.layers[start:]):
      if i == 1 and alpha < 1:
        x = alpha * x + (1 - alpha) * self.fromRGB[start + 1](t.interpolate(input, scale_factor=0.5, mode='nearest'))

      x = layer(x)

    x = self.final_layer(x).reshape(x.shape[0], -1)

    prediction = x[..., 0:1]
    latent = x[..., 1:1 + self.dim_z]
    position = x[..., 1 + self.dim_z:]

    return prediction, latent, position
//  [python]
def forward(self,
              input,
              alpha,
              options=None,
              **kwargs):

    start = self.img_size_to_layer[input.shape[-1]]
    x = self.fromRGB[start](input)

    if kwargs.get('instance_noise', 0) > 0:
      x = x + torch.randn_like(x) * kwargs['instance_noise']

    for i, layer in enumerate(self.layers[start:]):
      if i == 1 and alpha < 1:
        x = alpha * x + (1 - alpha) * self.fromRGB[start + 1](F.interpolate(input, scale_factor=0.5, mode='nearest'))

      x = layer(x)

    x = self.final_layer(x).reshape(x.shape[0], -1)

    prediction = x[..., 0:1]
    latent = x[..., 1:1 + self.dim_z]
    position = x[..., 1 + self.dim_z:]

    return prediction, latent, position
//  [python]
def forward(self, inputs):
        assert len(inputs) == len(self.in_channels)

        # build laterals
        laterals = [
            lateral_conv(inputs[i + self.start_level])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]

        # build top-down path
        used_backbone_levels = len(laterals)
        for i in range(used_backbone_levels - 1, 0, -1):
            # In some cases, fixing `scale factor` (e.g. 2) is preferred, but
            #  it cannot co-exist with `size` in `F.interpolate`.
            if 'scale_factor' in self.upsample_cfg:
                laterals[i - 1] += t.interpolate(laterals[i],
                                                 **self.upsample_cfg)
            else:
                prev_shape = laterals[i - 1].shape[2:]
                laterals[i - 1] += t.interpolate(
                    laterals[i], size=prev_shape, **self.upsample_cfg)

        # build outputs
        # part 1: from original levels
        outs = [
            self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
        ]
        # part 2: add extra levels
        if self.num_outs > len(outs):
            # use max pool to get more levels on top of outputs
            # (e.g., Faster R-CNN, Mask R-CNN)
            if not self.add_extra_convs:
                for i in range(self.num_outs - used_backbone_levels):
                    outs.append(F.max_pool2d(outs[-1], 1, stride=2))
            # add conv layers on top of original feature maps (RetinaNet)
            else:
                if self.add_extra_convs == 'on_input':
                    extra_source = inputs[self.backbone_end_level - 1]
                elif self.add_extra_convs == 'on_lateral':
                    extra_source = laterals[-1]
                elif self.add_extra_convs == 'on_output':
                    extra_source = outs[-1]
                else:
                    raise NotImplementedError
                outs.append(self.fpn_convs[used_backbone_levels](extra_source))
                for i in range(used_backbone_levels + 1, self.num_outs):
                    if self.relu_before_extra_convs:
                        outs.append(self.fpn_convs[i](F.relu(outs[-1])))
                    else:
                        outs.append(self.fpn_convs[i](outs[-1]))
        return tuple(outs)
//  [python]
def forward(self, inputs):
        assert len(inputs) == len(self.in_channels)

        # build laterals
        laterals = [
            lateral_conv(inputs[i + self.start_level])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]

        # build top-down path
        used_backbone_levels = len(laterals)
        for i in range(used_backbone_levels - 1, 0, -1):
            # In some cases, fixing `scale factor` (e.g. 2) is preferred, but
            #  it cannot co-exist with `size` in `F.interpolate`.
            if 'scale_factor' in self.upsample_cfg:
                laterals[i - 1] += F.interpolate(laterals[i],
                                                 **self.upsample_cfg)
            else:
                prev_shape = laterals[i - 1].shape[2:]
                laterals[i - 1] += F.interpolate(
                    laterals[i], size=prev_shape, **self.upsample_cfg)

        # build outputs
        # part 1: from original levels
        outs = [
            self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
        ]
        # part 2: add extra levels
        if self.num_outs > len(outs):
            # use max pool to get more levels on top of outputs
            # (e.g., Faster R-CNN, Mask R-CNN)
            if not self.add_extra_convs:
                for i in range(self.num_outs - used_backbone_levels):
                    outs.append(F.max_pool2d(outs[-1], 1, stride=2))
            # add conv layers on top of original feature maps (RetinaNet)
            else:
                if self.add_extra_convs == 'on_input':
                    extra_source = inputs[self.backbone_end_level - 1]
                elif self.add_extra_convs == 'on_lateral':
                    extra_source = laterals[-1]
                elif self.add_extra_convs == 'on_output':
                    extra_source = outs[-1]
                else:
                    raise NotImplementedError
                outs.append(self.fpn_convs[used_backbone_levels](extra_source))
                for i in range(used_backbone_levels + 1, self.num_outs):
                    if self.relu_before_extra_convs:
                        outs.append(self.fpn_convs[i](F.relu(outs[-1])))
                    else:
                        outs.append(self.fpn_convs[i](outs[-1]))
        return tuple(outs)
//  [python]
def forward(self, input, degraded_image, z=None):
        seg = input

        if self.opt.use_vae:
            if z is None:
                z = torch.randn(input.size(0), self.opt.z_dim, dtype=torch.float32, device=input.get_device())
            x = self.fc(z)
            x = x.view(-1, 16 * self.opt.ngf, self.sh, self.sw)
        else:
            if self.opt.no_parsing_map:
                x = t.interpolate(degraded_image, size=(self.sh, self.sw), mode="bilinear")
            else:
                x = t.interpolate(seg, size=(self.sh, self.sw), mode="nearest")
            x = self.fc(x)

        x = self.head_0(x, seg, degraded_image)

        x = self.up(x)
        x = self.G_middle_0(x, seg, degraded_image)

        if self.opt.num_upsampling_layers == "more" or self.opt.num_upsampling_layers == "most":
            x = self.up(x)

        x = self.G_middle_1(x, seg, degraded_image)

        x = self.up(x)
        x = self.up_0(x, seg, degraded_image)
        x = self.up(x)
        x = self.up_1(x, seg, degraded_image)
        x = self.up(x)
        x = self.up_2(x, seg, degraded_image)
        x = self.up(x)
        x = self.up_3(x, seg, degraded_image)

        if self.opt.num_upsampling_layers == "most":
            x = self.up(x)
            x = self.up_4(x, seg, degraded_image)

        x = self.conv_img(F.leaky_relu(x, 2e-1))
        x = F.tanh(x)

        return x
//  [python]
def forward(self, input, degraded_image, z=None):
        seg = input

        if self.opt.use_vae:
            # we sample z from unit normal and reshape the tensor
            if z is None:
                z = torch.randn(input.size(0), self.opt.z_dim, dtype=torch.float32, device=input.get_device())
            x = self.fc(z)
            x = x.view(-1, 16 * self.opt.ngf, self.sh, self.sw)
        else:
            # we downsample segmap and run convolution
            if self.opt.no_parsing_map:
                x = F.interpolate(degraded_image, size=(self.sh, self.sw), mode="bilinear")
            else:
                x = F.interpolate(seg, size=(self.sh, self.sw), mode="nearest")
            x = self.fc(x)

        x = self.head_0(x, seg, degraded_image)

        x = self.up(x)
        x = self.G_middle_0(x, seg, degraded_image)

        if self.opt.num_upsampling_layers == "more" or self.opt.num_upsampling_layers == "most":
            x = self.up(x)

        x = self.G_middle_1(x, seg, degraded_image)

        x = self.up(x)
        x = self.up_0(x, seg, degraded_image)
        x = self.up(x)
        x = self.up_1(x, seg, degraded_image)
        x = self.up(x)
        x = self.up_2(x, seg, degraded_image)
        x = self.up(x)
        x = self.up_3(x, seg, degraded_image)

        if self.opt.num_upsampling_layers == "most":
            x = self.up(x)
            x = self.up_4(x, seg, degraded_image)

        x = self.conv_img(F.leaky_relu(x, 2e-1))
        x = F.tanh(x)

        return x
//  [python]
def timestep_independent(self, aligned_conditioning, conditioning_latent, expected_seq_len, return_code_pred):
        # Shuffle aligned_latent to BxCxS format
        if is_latent(aligned_conditioning):
            aligned_conditioning = aligned_conditioning.permute(0, 2, 1)

        cond_scale, cond_shift = torch.chunk(conditioning_latent, 2, dim=1)
        if is_latent(aligned_conditioning):
            code_emb = self.latent_conditioner(aligned_conditioning)
        else:
            code_emb = self.code_embedding(aligned_conditioning).permute(0, 2, 1)
            code_emb = self.code_converter(code_emb)
        code_emb = self.code_norm(code_emb) * (1 + cond_scale.unsqueeze(-1)) + cond_shift.unsqueeze(-1)

        unconditioned_batches = torch.zeros((code_emb.shape[0], 1, 1), device=code_emb.device)
        # Mask out the conditioning branch for whole batch elements, implementing something similar to classifier-free guidance.
        if self.training and self.unconditioned_percentage > 0:
            unconditioned_batches = torch.rand((code_emb.shape[0], 1, 1),
                                               device=code_emb.device) < self.unconditioned_percentage
            code_emb = torch.where(unconditioned_batches, self.unconditioned_embedding.repeat(aligned_conditioning.shape[0], 1, 1),
                                   code_emb)
        expanded_code_emb = t.interpolate(code_emb, size=expected_seq_len, mode='nearest')

        if not return_code_pred:
            return expanded_code_emb
        else:
            mel_pred = self.mel_head(expanded_code_emb)
            # Multiply mel_pred by !unconditioned_branches, which drops the gradient on unconditioned branches. This is because we don't want that gradient being used to train parameters through the codes_embedder as it unbalances contributions to that network from the MSE loss.
            mel_pred = mel_pred * unconditioned_batches.logical_not()
            return expanded_code_emb, mel_pred
//  [python]
def timestep_independent(self, aligned_conditioning, conditioning_latent, expected_seq_len, return_code_pred):
        # Shuffle aligned_latent to BxCxS format
        if is_latent(aligned_conditioning):
            aligned_conditioning = aligned_conditioning.permute(0, 2, 1)

        cond_scale, cond_shift = torch.chunk(conditioning_latent, 2, dim=1)
        if is_latent(aligned_conditioning):
            code_emb = self.latent_conditioner(aligned_conditioning)
        else:
            code_emb = self.code_embedding(aligned_conditioning).permute(0, 2, 1)
            code_emb = self.code_converter(code_emb)
        code_emb = self.code_norm(code_emb) * (1 + cond_scale.unsqueeze(-1)) + cond_shift.unsqueeze(-1)

        unconditioned_batches = torch.zeros((code_emb.shape[0], 1, 1), device=code_emb.device)
        # Mask out the conditioning branch for whole batch elements, implementing something similar to classifier-free guidance.
        if self.training and self.unconditioned_percentage > 0:
            unconditioned_batches = torch.rand((code_emb.shape[0], 1, 1),
                                               device=code_emb.device) < self.unconditioned_percentage
            code_emb = torch.where(unconditioned_batches, self.unconditioned_embedding.repeat(aligned_conditioning.shape[0], 1, 1),
                                   code_emb)
        expanded_code_emb = F.interpolate(code_emb, size=expected_seq_len, mode='nearest')

        if not return_code_pred:
            return expanded_code_emb
        else:
            mel_pred = self.mel_head(expanded_code_emb)
            # Multiply mel_pred by !unconditioned_branches, which drops the gradient on unconditioned branches. This is because we don't want that gradient being used to train parameters through the codes_embedder as it unbalances contributions to that network from the MSE loss.
            mel_pred = mel_pred * unconditioned_batches.logical_not()
            return expanded_code_emb, mel_pred
//  [python]
def forward(self, x):
        conv_ftr_list = x

        feature_list = []
        last_feature = None
        for i, conv_ftr in enumerate(list(reversed(conv_ftr_list))):
            # adaptive
            feature = self.adaptive[i](conv_ftr)
            # fuse
            if last_feature is not None:
                feature = feature + t.upsample(
                    last_feature, scale_factor=2, mode='nearest')
            # smooth
            feature = self.smooth[i](feature)
            last_feature = feature
            feature_list.append(feature)

        return tuple(reversed(feature_list))
//  [python]
def forward(self, x):
        conv_ftr_list = x

        feature_list = []
        last_feature = None
        for i, conv_ftr in enumerate(list(reversed(conv_ftr_list))):
            # adaptive
            feature = self.adaptive[i](conv_ftr)
            # fuse
            if last_feature is not None:
                feature = feature + F.interpolate(
                    last_feature, scale_factor=2, mode='nearest')
            # smooth
            feature = self.smooth[i](feature)
            last_feature = feature
            feature_list.append(feature)

        return tuple(reversed(feature_list))
//  [python]
def forward(self, x):
        out = self.skip_layers(x)
        out = self.output_block(out)
        if self.training and self.deep_supervision:
            out_all = [out]
            for feature_map in self.heads:
                out_all.append(F.interpolate(feature_map, out.shape[2:]))
            return torch.stack(out_all, dim=1)
        return out
//  [python]
def forward(self, x):
        out = self.skip_layers(x)
        out = self.output_block(out)
        if self.training and self.deep_supervision:
            out_all = [out]
            for feature_map in self.heads:
                out_all.append(interpolate(feature_map, out.shape[2:]))
            return torch.stack(out_all, dim=1)
        return out
//  [python]
def encode(
        self,
        input_cameras,
        input_images,
        input_bbox=None
    ):
        """
        Encode cameras and images
        """
        #! ENCODE IMAGES
        if input_images.shape[1] != self.in_dim:
            input_images = input_images.permute(0,3,1,2)
        self.input_images = input_images
        self.input_cameras = input_cameras
        self.input_bbox = input_bbox
        #' input images (B, 3, H, W)

        if self.encoder == 'lite':
            self.encoder_latent = self.encoder_model(input_images)
            self.feat_size = self.encoder_latent.shape[1]
        elif self.encoder == 'resnet18':
            x = self.encoder_model.conv1(input_images)
            x = self.encoder_model.bn1(x)
            x = self.encoder_model.relu(x)

            latents = [x]
            if self.encoder_num_layers > 1:
                x = self.encoder_model.maxpool(x)
                x = self.encoder_model.layer1(x)
                latents.append(x)
            if self.encoder_num_layers > 2:
                x = self.encoder_model.layer2(x)
                latents.append(x)
            if self.encoder_num_layers > 3:
                x = self.encoder_model.layer3(x)
                latents.append(x)
            if self.encoder_num_layers > 4:
                x = self.encoder_model.layer4(x)
                latents.append(x)

            latent_sz = latents[0].shape[-2:]
            for i in range(len(latents)):
                latents[i] = F.interpolate(
                    latents[i],
                    latent_sz,
                    mode='bilinear',
                    align_corners=True,
                )
            
            self.encoder_latent = torch.cat(latents, dim=1)
            self.feat_size = self.encoder_latent.shape[1]
        return self.input_images, self.encoder_latent
//  [python]
def compute_depth_metrics(self, data):
        # TODO: This is only correct for batchsize 1!
        depth_gt = data["depths"][0]
        depth_pred = data["fine"][0]["depth"][:, :1]

        depth_pred = t.nn.functional.interpolate(depth_pred, depth_gt.shape[-2:])

        if self.depth_scaling == "median":
            mask = depth_gt > 0
            scaling = torch.median(depth_gt[mask]) / torch.median(depth_pred[mask])
            depth_pred = scaling * depth_pred
        elif self.depth_scaling == "l2":
            mask = depth_gt > 0
            depth_pred = depth_pred
            depth_gt_ = depth_gt[mask]
            depth_pred_ = depth_pred[mask]
            depth_pred_ = torch.stack((depth_pred_, torch.ones_like(depth_pred_)), dim=-1)
            x = torch.linalg.lstsq(depth_pred_.to(torch.float32), depth_gt_.unsqueeze(-1).to(torch.float32)).solution.squeeze()
            depth_pred = depth_pred * x[0] + x[1]

        depth_pred = torch.clamp(depth_pred, 1e-3, 80)
        mask = depth_gt != 0

        depth_gt = depth_gt[mask]
        depth_pred = depth_pred[mask]

        thresh = torch.maximum((depth_gt / depth_pred), (depth_pred / depth_gt))
        a1 = (thresh < 1.25).to(torch.float)
        a2 = (thresh < 1.25 ** 2).to(torch.float)
        a3 = (thresh < 1.25 ** 3).to(torch.float)
        a1 = a1.mean()
        a2 = a2.mean()
        a3 = a3.mean()

        rmse = (depth_gt - depth_pred) ** 2
        rmse = rmse.mean() ** .5

        rmse_log = (torch.log(depth_gt) - torch.log(depth_pred)) ** 2
        rmse_log = rmse_log.mean() ** .5

        abs_rel = torch.abs(depth_gt - depth_pred) / depth_gt
        abs_rel = abs_rel.mean()

        sq_rel = ((depth_gt - depth_pred) ** 2) / depth_gt
        sq_rel = sq_rel.mean()

        metrics_dict = {
            "abs_rel": abs_rel,
            "sq_rel": sq_rel,
            "rmse": rmse,
            "rmse_log": rmse_log,
            "a1": a1,
            "a2": a2,
            "a3": a3
        }
        return metrics_dict
//  [python]
def compute_depth_metrics(self, data):
        # TODO: This is only correct for batchsize 1!
        depth_gt = data["depths"][0]
        depth_pred = data["fine"][0]["depth"][:, :1]

        depth_pred = F.interpolate(depth_pred, depth_gt.shape[-2:])

        if self.depth_scaling == "median":
            mask = depth_gt > 0
            scaling = torch.median(depth_gt[mask]) / torch.median(depth_pred[mask])
            depth_pred = scaling * depth_pred
        elif self.depth_scaling == "l2":
            mask = depth_gt > 0
            depth_pred = depth_pred
            depth_gt_ = depth_gt[mask]
            depth_pred_ = depth_pred[mask]
            depth_pred_ = torch.stack((depth_pred_, torch.ones_like(depth_pred_)), dim=-1)
            x = torch.linalg.lstsq(depth_pred_.to(torch.float32), depth_gt_.unsqueeze(-1).to(torch.float32)).solution.squeeze()
            depth_pred = depth_pred * x[0] + x[1]

        depth_pred = torch.clamp(depth_pred, 1e-3, 80)
        mask = depth_gt != 0

        depth_gt = depth_gt[mask]
        depth_pred = depth_pred[mask]

        thresh = torch.maximum((depth_gt / depth_pred), (depth_pred / depth_gt))
        a1 = (thresh < 1.25).to(torch.float)
        a2 = (thresh < 1.25 ** 2).to(torch.float)
        a3 = (thresh < 1.25 ** 3).to(torch.float)
        a1 = a1.mean()
        a2 = a2.mean()
        a3 = a3.mean()

        rmse = (depth_gt - depth_pred) ** 2
        rmse = rmse.mean() ** .5

        rmse_log = (torch.log(depth_gt) - torch.log(depth_pred)) ** 2
        rmse_log = rmse_log.mean() ** .5

        abs_rel = torch.abs(depth_gt - depth_pred) / depth_gt
        abs_rel = abs_rel.mean()

        sq_rel = ((depth_gt - depth_pred) ** 2) / depth_gt
        sq_rel = sq_rel.mean()

        metrics_dict = {
            "abs_rel": abs_rel,
            "sq_rel": sq_rel,
            "rmse": rmse,
            "rmse_log": rmse_log,
            "a1": a1,
            "a2": a2,
            "a3": a3
        }
        return metrics_dict
//  [python]
def forward(self,
              input,
              alpha,
              options=None,
              **kwargs):

    start = self.img_size_to_layer[input.shape[-1]]

    if global_cfg.tl_debug:
      VerboseModel.forward_verbose(self.fromRGB[start],
                                   inputs_args=(input, ),
                                   submodels=['model', ],
                                   name_prefix=f"fromRGB[{start}].")
    x = self.fromRGB[start](input)

    if kwargs.get('instance_noise', 0) > 0:
      x = x + torch.randn_like(x) * kwargs['instance_noise']

    for i, layer in enumerate(self.layers[start:]):
      if i == 1 and alpha < 1:
        x = alpha * x + (1 - alpha) * self.fromRGB[start + 1](t.interpolate(input, scale_factor=0.5, mode='nearest'))

      if global_cfg.tl_debug:
        VerboseModel.forward_verbose(layer,
                                     inputs_args=(x,),
                                     submodels=['network', 'network.0', 'network.2'],
                                     name_prefix=f"layers[{start + i}].")
      x = layer(x)

    if global_cfg.tl_debug:
      VerboseModel.forward_verbose(self.final_layer,
                                   inputs_args=(x, ),
                                   name_prefix=f"final_layer.")
    x = self.final_layer(x).reshape(x.shape[0], -1)

    prediction = x[..., 0:1]
    latent = x[..., 1:1 + self.dim_z]
    position = x[..., 1 + self.dim_z:]

    return prediction, latent, position
//  [python]
def forward(self,
              input,
              alpha,
              options=None,
              **kwargs):

    start = self.img_size_to_layer[input.shape[-1]]

    if global_cfg.tl_debug:
      VerboseModel.forward_verbose(self.fromRGB[start],
                                   inputs_args=(input, ),
                                   submodels=['model', ],
                                   name_prefix=f"fromRGB[{start}].")
    x = self.fromRGB[start](input)

    if kwargs.get('instance_noise', 0) > 0:
      x = x + torch.randn_like(x) * kwargs['instance_noise']

    for i, layer in enumerate(self.layers[start:]):
      if i == 1 and alpha < 1:
        x = alpha * x + (1 - alpha) * self.fromRGB[start + 1](F.interpolate(input, scale_factor=0.5, mode='nearest'))

      if global_cfg.tl_debug:
        VerboseModel.forward_verbose(layer,
                                     inputs_args=(x,),
                                     submodels=['network', 'network.0', 'network.2'],
                                     name_prefix=f"layers[{start + i}].")
      x = layer(x)

    if global_cfg.tl_debug:
      VerboseModel.forward_verbose(self.final_layer,
                                   inputs_args=(x, ),
                                   name_prefix=f"final_layer.")
    x = self.final_layer(x).reshape(x.shape[0], -1)

    prediction = x[..., 0:1]
    latent = x[..., 1:1 + self.dim_z]
    position = x[..., 1 + self.dim_z:]

    return prediction, latent, position
//  [python]
def forward(self, features: torch.Tensor) -> torch.Tensor:
        spatial_size = features.shape[-2:]
        features = self.global_pool(features)
        features = self.conv_1x1(features)
        features = torch.nn.functional.interpolate(features, size=spatial_size, mode="bilinear", align_corners=False)
        return features
//  [python]
def forward(self, features: torch.Tensor) -> torch.Tensor:
        spatial_size = features.shape[-2:]
        features = self.global_pool(features)
        features = self.conv_1x1(features)
        features = nn.functional.interpolate(features, size=spatial_size, mode="bilinear", align_corners=False)
        return features
//  [python]
def _preprocess(self, image: Tensor) -> Tensor:
        """Preprocess image before pass to the Inception. Preprocess operations
        contain channel conversion and resize.

        Args:
            image (Tensor): Image tensor before preprocess.

        Returns:
            Tensor: Image tensor after resize and channel conversion
                (if need.)
        """
        # image must passed in 'bgr'
        image = image[:, [2, 1, 0]]
        if not self.resize:
            return image
        if self.use_pillow_resize:
            image = image.to(torch.uint8)
            x_np = [x_.permute(1, 2, 0).detach().cpu().numpy() for x_ in image]

            # use bicubic resize as default
            x_pil = [
                Image.fromarray(x_).resize((299, 299),
                                           resample=self.resize_method)
                for x_ in x_np
            ]
            x_ten = torch.cat(
                [torch.FloatTensor(np.array(x_)[None, ...]) for x_ in x_pil])
            return x_ten.permute(0, 3, 1, 2)
        else:
            return t.interpolate(
                image, size=(299, 299), mode=self.resize_method)
//  [python]
def _preprocess(self, image: Tensor) -> Tensor:
        """Preprocess image before pass to the Inception. Preprocess operations
        contain channel conversion and resize.

        Args:
            image (Tensor): Image tensor before preprocess.

        Returns:
            Tensor: Image tensor after resize and channel conversion
                (if need.)
        """
        # image must passed in 'bgr'
        image = image[:, [2, 1, 0]]
        if not self.resize:
            return image
        if self.use_pillow_resize:
            image = image.to(torch.uint8)
            x_np = [x_.permute(1, 2, 0).detach().cpu().numpy() for x_ in image]

            # use bicubic resize as default
            x_pil = [
                Image.fromarray(x_).resize((299, 299),
                                           resample=self.resize_method)
                for x_ in x_np
            ]
            x_ten = torch.cat(
                [torch.FloatTensor(np.array(x_)[None, ...]) for x_ in x_pil])
            return x_ten.permute(0, 3, 1, 2)
        else:
            return F.interpolate(
                image, size=(299, 299), mode=self.resize_method)
//  [python]
def process_mask(self,
                     mask_proto: Tensor,
                     mask_coeff_pred: Tensor,
                     bboxes: Tensor,
                     shape: Tuple[int, int],
                     upsample: bool = False) -> Tensor:
        """Generate mask logits results.

        Args:
            mask_proto (Tensor): Mask prototype features.
                Has shape (num_instance, mask_channels).
            mask_coeff_pred (Tensor): Mask coefficients prediction for
                single image. Has shape (mask_channels, H, W)
            bboxes (Tensor): Tensor of the bbox. Has shape (num_instance, 4).
            shape (Tuple): Batch input shape of image.
            upsample (bool): Whether upsample masks results to batch input
                shape. Default to False.
        Return:
            Tensor: Instance segmentation masks for each instance.
                Has shape (num_instance, H, W).
        """
        c, mh, mw = mask_proto.shape  # CHW
        masks = (
            mask_coeff_pred @ mask_proto.float().view(c, -1)).sigmoid().view(
                -1, mh, mw)[None]
        if upsample:
            masks = t.nn.functional.interpolate(
                masks, shape, mode='bilinear', align_corners=False)  # 1CHW
        masks = self.crop_mask(masks, bboxes)
        return masks
//  [python]
def process_mask(self,
                     mask_proto: Tensor,
                     mask_coeff_pred: Tensor,
                     bboxes: Tensor,
                     shape: Tuple[int, int],
                     upsample: bool = False) -> Tensor:
        """Generate mask logits results.

        Args:
            mask_proto (Tensor): Mask prototype features.
                Has shape (num_instance, mask_channels).
            mask_coeff_pred (Tensor): Mask coefficients prediction for
                single image. Has shape (mask_channels, H, W)
            bboxes (Tensor): Tensor of the bbox. Has shape (num_instance, 4).
            shape (Tuple): Batch input shape of image.
            upsample (bool): Whether upsample masks results to batch input
                shape. Default to False.
        Return:
            Tensor: Instance segmentation masks for each instance.
                Has shape (num_instance, H, W).
        """
        c, mh, mw = mask_proto.shape  # CHW
        masks = (
            mask_coeff_pred @ mask_proto.float().view(c, -1)).sigmoid().view(
                -1, mh, mw)[None]
        if upsample:
            masks = F.interpolate(
                masks, shape, mode='bilinear', align_corners=False)  # 1CHW
        masks = self.crop_mask(masks, bboxes)
        return masks
//  [python]
def interpolate(self, x):
        upsampled = F.interpolate(
            x.unsqueeze(1), self.output_size, mode="linear", align_corners=True
        ).squeeze(1)
        if self.trainable:
            upsampled = upsampled * self.gate(self.mask.unsqueeze(0)) * 2.0
        return upsampled
//  [python]
def forward(self, masks: torch.Tensor) -> torch.Tensor:
    masks = F.interpolate(
        masks,
        size=(500, 500),
        mode="bilinear",
        align_corners=False,
    )
    return masks
//  [python]
def post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple] = None):
        """
        Converts the output of [`DPTForSemanticSegmentation`] into semantic segmentation maps. Only supports PyTorch.

        Args:
            outputs ([`DPTForSemanticSegmentation`]):
                Raw outputs of the model.
            target_sizes (`List[Tuple]` of length `batch_size`, *optional*):
                List of tuples corresponding to the requested final size (height, width) of each prediction. If unset,
                predictions will not be resized.

        Returns:
            semantic_segmentation: `List[torch.Tensor]` of length `batch_size`, where each item is a semantic
            segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is
            specified). Each entry of each `torch.Tensor` correspond to a semantic class id.
        """
        # TODO: add support for other frameworks
        logits = outputs.logits

        # Resize logits and compute semantic segmentation maps
        if target_sizes is not None:
            if len(logits) != len(target_sizes):
                raise ValueError(
                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"
                )

            if is_torch_tensor(target_sizes):
                target_sizes = target_sizes.numpy()

            semantic_segmentation = []

            for idx in range(len(logits)):
                resized_logits = torch.nn.functional.upsample(
                    logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear"
                )
                semantic_map = resized_logits[0].argmax(dim=0)
                semantic_segmentation.append(semantic_map)
        else:
            semantic_segmentation = logits.argmax(dim=1)
            semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]

        return semantic_segmentation
//  [python]
def post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple] = None):
        """
        Converts the output of [`DPTForSemanticSegmentation`] into semantic segmentation maps. Only supports PyTorch.

        Args:
            outputs ([`DPTForSemanticSegmentation`]):
                Raw outputs of the model.
            target_sizes (`List[Tuple]` of length `batch_size`, *optional*):
                List of tuples corresponding to the requested final size (height, width) of each prediction. If unset,
                predictions will not be resized.

        Returns:
            semantic_segmentation: `List[torch.Tensor]` of length `batch_size`, where each item is a semantic
            segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is
            specified). Each entry of each `torch.Tensor` correspond to a semantic class id.
        """
        # TODO: add support for other frameworks
        logits = outputs.logits

        # Resize logits and compute semantic segmentation maps
        if target_sizes is not None:
            if len(logits) != len(target_sizes):
                raise ValueError(
                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"
                )

            if is_torch_tensor(target_sizes):
                target_sizes = target_sizes.numpy()

            semantic_segmentation = []

            for idx in range(len(logits)):
                resized_logits = torch.nn.functional.interpolate(
                    logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False
                )
                semantic_map = resized_logits[0].argmax(dim=0)
                semantic_segmentation.append(semantic_map)
        else:
            semantic_segmentation = logits.argmax(dim=1)
            semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]

        return semantic_segmentation
//  [python]
def _get_rel_pos_bias(self, window_size):
    """
    Modification of timm.models.beit.py: Attention._get_rel_pos_bias to support arbitrary window sizes.
    """
    old_height = 2 * self.window_size[0] - 1
    old_width = 2 * self.window_size[1] - 1

    new_height = 2 * window_size[0] - 1
    new_width = 2 * window_size[1] - 1

    old_relative_position_bias_table = self.relative_position_bias_table

    old_num_relative_distance = self.num_relative_distance
    new_num_relative_distance = new_height * new_width + 3

    old_sub_table = old_relative_position_bias_table[: old_num_relative_distance - 3]

    old_sub_table = old_sub_table.reshape(1, old_width, old_height, -1).permute(
        0, 3, 1, 2
    )
    new_sub_table = t.nn.functional.interpolate(
        old_sub_table, size=(int(new_height), int(new_width)), mode="bilinear"
    )
    new_sub_table = new_sub_table.permute(0, 2, 3, 1).reshape(
        new_num_relative_distance - 3, -1
    )

    new_relative_position_bias_table = torch.cat(
        [
            new_sub_table,
            old_relative_position_bias_table[old_num_relative_distance - 3 :],
        ]
    )

    key = str(window_size[1]) + "," + str(window_size[0])
    if key not in self.relative_position_indices:
        self.relative_position_indices[key] = gen_relative_position_index(window_size)

    relative_position_bias = new_relative_position_bias_table[
        self.relative_position_indices[key].view(-1)
    ].view(
        window_size[0] * window_size[1] + 1, window_size[0] * window_size[1] + 1, -1
    )  # Wh*Ww,Wh*Ww,nH
    relative_position_bias = relative_position_bias.permute(
        2, 0, 1
    ).contiguous()  # nH, Wh*Ww, Wh*Ww
    return relative_position_bias.unsqueeze(0)
//  [python]
def _get_rel_pos_bias(self, window_size):
    """
    Modification of timm.models.beit.py: Attention._get_rel_pos_bias to support arbitrary window sizes.
    """
    old_height = 2 * self.window_size[0] - 1
    old_width = 2 * self.window_size[1] - 1

    new_height = 2 * window_size[0] - 1
    new_width = 2 * window_size[1] - 1

    old_relative_position_bias_table = self.relative_position_bias_table

    old_num_relative_distance = self.num_relative_distance
    new_num_relative_distance = new_height * new_width + 3

    old_sub_table = old_relative_position_bias_table[: old_num_relative_distance - 3]

    old_sub_table = old_sub_table.reshape(1, old_width, old_height, -1).permute(
        0, 3, 1, 2
    )
    new_sub_table = F.interpolate(
        old_sub_table, size=(int(new_height), int(new_width)), mode="bilinear"
    )
    new_sub_table = new_sub_table.permute(0, 2, 3, 1).reshape(
        new_num_relative_distance - 3, -1
    )

    new_relative_position_bias_table = torch.cat(
        [
            new_sub_table,
            old_relative_position_bias_table[old_num_relative_distance - 3 :],
        ]
    )

    key = str(window_size[1]) + "," + str(window_size[0])
    if key not in self.relative_position_indices:
        self.relative_position_indices[key] = gen_relative_position_index(window_size)

    relative_position_bias = new_relative_position_bias_table[
        self.relative_position_indices[key].view(-1)
    ].view(
        window_size[0] * window_size[1] + 1, window_size[0] * window_size[1] + 1, -1
    )  # Wh*Ww,Wh*Ww,nH
    relative_position_bias = relative_position_bias.permute(
        2, 0, 1
    ).contiguous()  # nH, Wh*Ww, Wh*Ww
    return relative_position_bias.unsqueeze(0)
//  [python]
def forward(self, hidden_states, output_size=None):
        assert hidden_states.shape[1] == self.channels

        if self.use_conv_transpose:
            return self.conv(hidden_states)

        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
        # Remove this cast once the issue is fixed in PyTorch
        # https://github.com/pytorch/pytorch/issues/86679
        dtype = hidden_states.dtype
        if dtype == torch.bfloat16:
            hidden_states = hidden_states.to(torch.float32)

        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
        if hidden_states.shape[0] >= 64:
            hidden_states = hidden_states.contiguous()

        # if `output_size` is passed we force the interpolation output
        # size and do not make use of `scale_factor=2`
        if output_size is None:
            hidden_states = F.upsample_nearest(
                hidden_states, scale_factor=2.0
            )
        else:
            hidden_states = F.upsample_nearest(
                hidden_states, size=output_size
            )

        # If the input is bfloat16, we cast back to bfloat16
        if dtype == torch.bfloat16:
            hidden_states = hidden_states.to(dtype)

        if self.use_conv:
            if self.name == "conv":
                hidden_states = self.conv(hidden_states)
            else:
                hidden_states = self.Conv2d_0(hidden_states)

        return hidden_states
//  [python]
def forward(self, hidden_states, output_size=None):
        assert hidden_states.shape[1] == self.channels

        if self.use_conv_transpose:
            return self.conv(hidden_states)

        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
        # Remove this cast once the issue is fixed in PyTorch
        # https://github.com/pytorch/pytorch/issues/86679
        dtype = hidden_states.dtype
        if dtype == torch.bfloat16:
            hidden_states = hidden_states.to(torch.float32)

        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
        if hidden_states.shape[0] >= 64:
            hidden_states = hidden_states.contiguous()

        # if `output_size` is passed we force the interpolation output
        # size and do not make use of `scale_factor=2`
        if output_size is None:
            hidden_states = F.interpolate(
                hidden_states, scale_factor=2.0, mode="nearest"
            )
        else:
            hidden_states = F.interpolate(
                hidden_states, size=output_size, mode="nearest"
            )

        # If the input is bfloat16, we cast back to bfloat16
        if dtype == torch.bfloat16:
            hidden_states = hidden_states.to(dtype)

        if self.use_conv:
            if self.name == "conv":
                hidden_states = self.conv(hidden_states)
            else:
                hidden_states = self.Conv2d_0(hidden_states)

        return hidden_states