jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def convert_dpt_checkpoint(checkpoint_url, pytorch_dump_folder_path, push_to_hub, model_name, show_prediction):
    """
    Copy/paste/tweak model's weights to our DPT structure.
    """

    # define DPT configuration based on URL
    config, expected_shape = get_dpt_config(checkpoint_url)
    # load original state_dict from URL
    # state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")
    state_dict = torch.load(checkpoint_url, map_location="cpu")
    # remove certain keys
    remove_ignore_keys_(state_dict)
    # rename keys
    for key in state_dict.copy().keys():
        val = state_dict.pop(key)
        state_dict[rename_key(key)] = val
    # read in qkv matrices
    read_in_q_k_v(state_dict, config)

    # load HuggingFace model
    model = DPTForSemanticSegmentation(config) if "ade" in checkpoint_url else DPTForDepthEstimation(config)
    model.load_state_dict(state_dict)
    model.eval()

    # Check outputs on an image
    size = 480 if "ade" in checkpoint_url else 384
    image_processor = DPTImageProcessor(size=size)

    image = prepare_img()
    encoding = image_processor(image, return_tensors="pt")

    # forward pass
    outputs = model(**encoding).logits if "ade" in checkpoint_url else model(**encoding).predicted_depth

    if show_prediction:
        prediction = (
            t.nn.functional.interpolate(
                outputs.unsqueeze(1),
                size=(image.size[1], image.size[0]),
                mode="bicubic",
                align_corners=False,
            )
            .squeeze()
            .cpu()
            .numpy()
        )

        Image.fromarray((prediction / prediction.max()) * 255).show()

    if pytorch_dump_folder_path is not None:
        Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
        print(f"Saving model to {pytorch_dump_folder_path}")
        model.save_pretrained(pytorch_dump_folder_path)
        print(f"Saving image processor to {pytorch_dump_folder_path}")
        image_processor.save_pretrained(pytorch_dump_folder_path)

    if push_to_hub:
        model.push_to_hub("ybelkada/dpt-hybrid-midas")
        image_processor.push_to_hub("ybelkada/dpt-hybrid-midas")
//  [python]
def forward(self, features):
        x = features[self.feature_key]

        x = self.aspp(x)

        xs = []
        # build decoder
        for i in range(self.decoder_stage):
            l = features[self.low_level_key[i]]
            l = self.project[i](l)
            x = F.interpolate(x, size=l.size()[2:], mode='bilinear', align_corners=True)
            x = torch.cat((x, l), dim=1)
            x = self.fuse[i](x)

            xs.append(x)

        if self.dsn:
            return xs[-1], xs[-2]
        return xs[-1]
//  [python]
def forward(self, f: torch.FloatTensor, zq: torch.FloatTensor) -> torch.FloatTensor:
        f_size = f.shape[-2:]
        zq = F.upsample(zq, size=f_size, mode="nearest")
        norm_f = self.norm_layer(f)
        new_f = norm_f * self.conv_y(zq) + self.conv_b(zq)
        return new_f
//  [python]
def forward(self, f: torch.FloatTensor, zq: torch.FloatTensor) -> torch.FloatTensor:
        f_size = f.shape[-2:]
        zq = F.interpolate(zq, size=f_size, mode="nearest")
        norm_f = self.norm_layer(f)
        new_f = norm_f * self.conv_y(zq) + self.conv_b(zq)
        return new_f
//  [python]
def forward(self, feats):
        x = self.lateral_convs[self.fusion_level](feats[self.fusion_level])
        fused_size = tuple(x.shape[-2:])
        for i, feat in enumerate(feats):
            if i != self.fusion_level:
                feat = F.upsample(
                    feat, size=fused_size, mode='bilinear', align_corners=True)
                x += self.lateral_convs[i](feat)

        for i in range(self.num_convs):
            x = self.convs[i](x)

        mask_pred = self.conv_logits(x)
        x = self.conv_embedding(x)
        return mask_pred, x
//  [python]
def forward(self, feats):
        x = self.lateral_convs[self.fusion_level](feats[self.fusion_level])
        fused_size = tuple(x.shape[-2:])
        for i, feat in enumerate(feats):
            if i != self.fusion_level:
                feat = F.interpolate(
                    feat, size=fused_size, mode='bilinear', align_corners=True)
                x += self.lateral_convs[i](feat)

        for i in range(self.num_convs):
            x = self.convs[i](x)

        mask_pred = self.conv_logits(x)
        x = self.conv_embedding(x)
        return mask_pred, x
//  [python]
def forward(self, features: Tuple[Tensor, ...]) -> Tensor:
        # multi-level feature fusion
        fusion_feats = [features[0]]
        size = features[0].shape[-2:]
        for i in range(1, self.num_levels):
            f = t.interpolate(features[i], size=size, mode='bilinear')
            fusion_feats.append(f)
        fusion_feats = torch.cat(fusion_feats, dim=1)
        fusion_feats = self.fusion_conv(fusion_feats)
        # pred mask feats
        mask_features = self.stacked_convs(fusion_feats)
        mask_features = self.projection(mask_features)
        return mask_features
//  [python]
def forward(self, features: Tuple[Tensor, ...]) -> Tensor:
        # multi-level feature fusion
        fusion_feats = [features[0]]
        size = features[0].shape[-2:]
        for i in range(1, self.num_levels):
            f = F.interpolate(features[i], size=size, mode='bilinear')
            fusion_feats.append(f)
        fusion_feats = torch.cat(fusion_feats, dim=1)
        fusion_feats = self.fusion_conv(fusion_feats)
        # pred mask feats
        mask_features = self.stacked_convs(fusion_feats)
        mask_features = self.projection(mask_features)
        return mask_features
//  [python]
def forward(self, features):
        res2 = features['res2']
        res3 = features['res3']
        res4 = features['res4']
        res5 = features['res5']

        pred = OrderedDict()

        res5_feature = self.res5_fuse(res5)
        for key in self.class_key:
            pred[key + '-5'] = self.classifiers[key + '-5'](res5_feature)

        res4_feature = self.res4_fuse(res4, res5_feature)
        res4_feature = self.res4_ocr(res4_feature, pred['semantic-5'])
        for key in self.class_key:
            if key == 'semantic':
                pred[key + '-4'] = self.classifiers[key + '-4'](res4_feature)
            else:
                offset_pred = t.interpolate(pred['offset-5'],
                                            size=res4_feature.shape[2:4],
                                            mode='bilinear')
                offset_pred *= 2
                residual = self.classifiers[key + '-4'](torch.cat([res4_feature, offset_pred], dim=1))
                pred[key + '-4'] = offset_pred + residual

        res3_feature = self.res3_fuse(res3, res4_feature)
        res3_feature = self.res3_ocr(res3_feature, pred['semantic-4'])
        for key in self.class_key:
            if key == 'semantic':
                pred[key + '-3'] = self.classifiers[key + '-3'](res3_feature)
            else:
                offset_pred = t.interpolate(pred['offset-4'],
                                            size=res3_feature.shape[2:4],
                                            mode='bilinear')
                offset_pred *= 2
                residual = self.classifiers[key + '-3'](torch.cat([res3_feature, offset_pred], dim=1))
                pred[key + '-3'] = offset_pred + residual

        res2_feature = self.res2_fuse(res2, res3_feature)
        res2_feature = self.res2_ocr(res2_feature, pred['semantic-3'])
        for key in self.class_key:
            if key == 'semantic':
                pred[key + '-2'] = self.classifiers[key + '-2'](res2_feature)
            else:
                offset_pred = t.interpolate(pred['offset-3'],
                                            size=res2_feature.shape[2:4],
                                            mode='bilinear')
                offset_pred *= 2
                residual = self.classifiers[key + '-2'](torch.cat([res2_feature, offset_pred], dim=1))
                pred[key + '-2'] = offset_pred + residual

        return pred
//  [python]
def forward(self, features):
        res2 = features['res2']
        res3 = features['res3']
        res4 = features['res4']
        res5 = features['res5']

        pred = OrderedDict()

        res5_feature = self.res5_fuse(res5)
        for key in self.class_key:
            pred[key + '-5'] = self.classifiers[key + '-5'](res5_feature)

        res4_feature = self.res4_fuse(res4, res5_feature)
        res4_feature = self.res4_ocr(res4_feature, pred['semantic-5'])
        for key in self.class_key:
            if key == 'semantic':
                pred[key + '-4'] = self.classifiers[key + '-4'](res4_feature)
            else:
                offset_pred = F.interpolate(pred['offset-5'],
                                            size=res4_feature.shape[2:4],
                                            mode='bilinear')
                offset_pred *= 2
                residual = self.classifiers[key + '-4'](torch.cat([res4_feature, offset_pred], dim=1))
                pred[key + '-4'] = offset_pred + residual

        res3_feature = self.res3_fuse(res3, res4_feature)
        res3_feature = self.res3_ocr(res3_feature, pred['semantic-4'])
        for key in self.class_key:
            if key == 'semantic':
                pred[key + '-3'] = self.classifiers[key + '-3'](res3_feature)
            else:
                offset_pred = F.interpolate(pred['offset-4'],
                                            size=res3_feature.shape[2:4],
                                            mode='bilinear')
                offset_pred *= 2
                residual = self.classifiers[key + '-3'](torch.cat([res3_feature, offset_pred], dim=1))
                pred[key + '-3'] = offset_pred + residual

        res2_feature = self.res2_fuse(res2, res3_feature)
        res2_feature = self.res2_ocr(res2_feature, pred['semantic-3'])
        for key in self.class_key:
            if key == 'semantic':
                pred[key + '-2'] = self.classifiers[key + '-2'](res2_feature)
            else:
                offset_pred = F.interpolate(pred['offset-3'],
                                            size=res2_feature.shape[2:4],
                                            mode='bilinear')
                offset_pred *= 2
                residual = self.classifiers[key + '-2'](torch.cat([res2_feature, offset_pred], dim=1))
                pred[key + '-2'] = offset_pred + residual

        return pred
//  [python]
def inference(self, image):
        """
        image should be tensor, 1 * 3 * 1024 * 1024
        """
        with torch.no_grad():
            if self.device_id == -1:
                output = self.model(image)
            else:
                device = torch.device('cuda', self.device_id)
                output = self.model(image.to(device))
            output = F.upsample(output, size=(1024, 1024), mode='bilinear')
            output = F.softmax(output, dim=1)
            output = torch.argmax(output, dim=1)
            output = output[0]
            if self.device_id == -1:
                pred = output.data.numpy()
            else:
                pred = output.data.cpu().numpy()

            del output
        return pred
//  [python]
def inference(self, image):
        """
        image should be tensor, 1 * 3 * 1024 * 1024
        """
        with torch.no_grad():
            if self.device_id == -1:
                output = self.model(image)
            else:
                device = torch.device('cuda', self.device_id)
                output = self.model(image.to(device))
            output = F.interpolate(output, size=(1024, 1024), mode='bilinear')
            output = F.softmax(output, dim=1)
            output = torch.argmax(output, dim=1)
            output = output[0]
            if self.device_id == -1:
                pred = output.data.numpy()
            else:
                pred = output.data.cpu().numpy()

            del output
        return pred
//  [python]
def forward(self, x_v: torch.Tensor):
        """
        x: shape [B,C,T,V_v]
        """
        B, C, T, V = x_v.shape
        # data normalization
        if self.data_bn is not None:
            if hasattr(self.cfg, 'PJN') and self.cfg.PJN:
                x_v = self.data_bn(x_v.permute(0, 1, 3, 2).contiguous().view(B, -1, T)).view(B, C, V, T) \
                    .contiguous().permute(0, 1, 3, 2)
            else:
                x_v = self.data_bn(x_v)

        x_e = self.get_edge_fea(x_v)

        # forward
        feature = []
        for idx, hoa_block in enumerate(self.downsample):
            x_v, x_e = hoa_block(x_v, x_e)
            if idx == 0 or idx == 2 or idx == 4 or idx == 6:
                feature.append((x_v, x_e))

        feature.append((x_v, x_e))
        feature = feature[::-1]

        x_v, x_e = feature[0]
        identity_feature = feature[1:]

        ushape_feature = []
        ushape_feature.append((x_v, x_e))
        for idx, (hoa_block, id) in \
                enumerate(zip(self.upsample, identity_feature)):
            x_v, x_e = hoa_block(x_v, x_e)
            if hasattr(self.cfg, 'deterministic') and self.cfg.deterministic:
                x_v = t.nn.functional.interpolate(x_v, scale_factor=(2, 1), mode='nearest')
            else:
                x_v = t.nn.functional.interpolate(
                    x_v,
                    scale_factor=(2, 1),
                    mode='bilinear',
                    align_corners=False)
            x_v += id[0]
            ushape_feature.append((x_v, x_e))

        ushape_feature = ushape_feature[:-1]
        for idx, (hoa_block, u) in \
                enumerate(zip(self.merge, ushape_feature)):
            x_v2, x_e2 = hoa_block(*u)
            if hasattr(self.cfg, 'deterministic') and self.cfg.deterministic:
                x_v += t.nn.functional.interpolate(
                    x_v2, scale_factor=(2**(4 - idx), 1), mode='nearest')
            else:
                x_v += t.nn.functional.interpolate(
                    x_v2,
                    scale_factor=(2**(4 - idx), 1),
                    mode='bilinear',
                    align_corners=False)
        return x_v, x_e
//  [python]
def forward(self, x_v: torch.Tensor):
        """
        x: shape [B,C,T,V_v]
        """
        B, C, T, V = x_v.shape
        # data normalization
        if self.data_bn is not None:
            if hasattr(self.cfg, 'PJN') and self.cfg.PJN:
                x_v = self.data_bn(x_v.permute(0, 1, 3, 2).contiguous().view(B, -1, T)).view(B, C, V, T) \
                    .contiguous().permute(0, 1, 3, 2)
            else:
                x_v = self.data_bn(x_v)

        x_e = self.get_edge_fea(x_v)

        # forward
        feature = []
        for idx, hoa_block in enumerate(self.downsample):
            x_v, x_e = hoa_block(x_v, x_e)
            if idx == 0 or idx == 2 or idx == 4 or idx == 6:
                feature.append((x_v, x_e))

        feature.append((x_v, x_e))
        feature = feature[::-1]

        x_v, x_e = feature[0]
        identity_feature = feature[1:]

        ushape_feature = []
        ushape_feature.append((x_v, x_e))
        for idx, (hoa_block, id) in \
                enumerate(zip(self.upsample, identity_feature)):
            x_v, x_e = hoa_block(x_v, x_e)
            if hasattr(self.cfg, 'deterministic') and self.cfg.deterministic:
                x_v = F.interpolate(x_v, scale_factor=(2, 1), mode='nearest')
            else:
                x_v = F.interpolate(
                    x_v,
                    scale_factor=(2, 1),
                    mode='bilinear',
                    align_corners=False)
            x_v += id[0]
            ushape_feature.append((x_v, x_e))

        ushape_feature = ushape_feature[:-1]
        for idx, (hoa_block, u) in \
                enumerate(zip(self.merge, ushape_feature)):
            x_v2, x_e2 = hoa_block(*u)
            if hasattr(self.cfg, 'deterministic') and self.cfg.deterministic:
                x_v += F.interpolate(
                    x_v2, scale_factor=(2**(4 - idx), 1), mode='nearest')
            else:
                x_v += F.interpolate(
                    x_v2,
                    scale_factor=(2**(4 - idx), 1),
                    mode='bilinear',
                    align_corners=False)
        return x_v, x_e
//  [python]
def forward(self, x):
        c1_ = (self.conv1(x))
        c1 = self.conv2(c1_)
        v_max = F.max_pool2d(c1, kernel_size=7, stride=3)
        v_range = self.relu(self.conv_max(v_max))
        c3 = self.relu(self.conv3(v_range))
        c3 = self.conv3_(c3)
        c3 = F.upsample_bilinear(c3, size=(x.size(2), x.size(3)))
        cf = self.conv_f(c1_)
        c4 = self.conv4(c3+cf)
        m = self.sigmoid(c4)
        
        return x * m
//  [python]
def forward(self, x):
        c1_ = (self.conv1(x))
        c1 = self.conv2(c1_)
        v_max = F.max_pool2d(c1, kernel_size=7, stride=3)
        v_range = self.relu(self.conv_max(v_max))
        c3 = self.relu(self.conv3(v_range))
        c3 = self.conv3_(c3)
        c3 = F.interpolate(c3, (x.size(2), x.size(3)), mode='bilinear', align_corners=False) 
        cf = self.conv_f(c1_)
        c4 = self.conv4(c3+cf)
        m = self.sigmoid(c4)
        
        return x * m
//  [python]
def gen(self, z=None, layout=None, ema=False, norm_img=False, ret_layout=False, ret_latents=False, noise=None,
            **kwargs):
        assert (z is not None) or (layout is not None)
        kwargs['return_metadata'] = ret_layout
        layout = self.generate_layout(z, metadata=layout, ema=ema, **kwargs)
        # interpolate scores_pyramid and feature grid in layout
        G = self.generator_ema if ema else self.generator
        layout['feature_grid'] = t.interpolate(layout['feature_grid'], size=(G.size_in,int(self.aspect_ratio*G.size_in)))
        layout['scores_pyramid'] = {k: t.interpolate(v, size=(k,int(self.aspect_ratio*k))) for k, v in layout['scores_pyramid'].items()}
        gen_input = {
            'input':  layout['feature_grid'],
            'styles': {k: layout[k] for k in SPLAT_KEYS} if self.spatial_style else z,
            'return_image_only': not ret_latents,
            'return_latents': ret_latents,
            'noise': noise
        }
        out = G(**gen_input)
        if norm_img:
            img = out[0] if ret_latents else out
            img.add_(1).div_(2).mul_(255)
        if ret_layout:
            if not ret_latents:
                out = [out]
            return [layout, *out]
        else:
            return out
//  [python]
def gen(self, z=None, layout=None, ema=False, norm_img=False, ret_layout=False, ret_latents=False, noise=None,
            **kwargs):
        assert (z is not None) or (layout is not None)
        kwargs['return_metadata'] = ret_layout
        layout = self.generate_layout(z, metadata=layout, ema=ema, **kwargs)
        # interpolate scores_pyramid and feature grid in layout
        G = self.generator_ema if ema else self.generator
        layout['feature_grid'] = F.interpolate(layout['feature_grid'], size=(G.size_in,int(self.aspect_ratio*G.size_in)))
        layout['scores_pyramid'] = {k: F.interpolate(v, size=(k,int(self.aspect_ratio*k))) for k, v in layout['scores_pyramid'].items()}
        gen_input = {
            'input':  layout['feature_grid'],
            'styles': {k: layout[k] for k in SPLAT_KEYS} if self.spatial_style else z,
            'return_image_only': not ret_latents,
            'return_latents': ret_latents,
            'noise': noise
        }
        out = G(**gen_input)
        if norm_img:
            img = out[0] if ret_latents else out
            img.add_(1).div_(2).mul_(255)
        if ret_layout:
            if not ret_latents:
                out = [out]
            return [layout, *out]
        else:
            return out
//  [python]
def forward(self, x, crop=False):
        # Not sure of the image range here
        if crop:
            x = t.nn.functional.interpolate(x, (256, 256), mode="area")
            x = x[:, :, 35:223, 32:220]
        x = self.face_pool(x)
        x_feats = self.facenet(x)
        return x_feats
//  [python]
def forward(self, x, crop=False):
        # Not sure of the image range here
        if crop:
            x = torch.nn.functional.interpolate(x, (256, 256), mode="area")
            x = x[:, :, 35:223, 32:220]
        x = self.face_pool(x)
        x_feats = self.facenet(x)
        return x_feats
//  [python]
def test_torch_upsample(self):
        def run_upsample(
            N,
            C,
            idim,
            scale_factor,
            mode,
            align_corner,
            dtype="float32",
            atol=1e-5,
        ):
            with self.subTest(
                N=N,
                C=C,
                idim=idim,
                scale_factor=scale_factor,
                mode=mode,
                align_corner=align_corner,
            ):
                np_dtype = getattr(np, dtype)
                np.random.seed(0)
                iH, iW = idim
                in_np = np.random.normal(-1.0, 1.0, (N, iH, iW, C)).astype(np_dtype)

                in_mx = mx.array(in_np)
                in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)).to("cpu")

                out_mx = nn.Upsample(
                    scale_factor=scale_factor,
                    mode=mode,
                    align_corners=align_corner,
                )(in_mx)
                mode_pt = {
                    "linear": "bilinear",
                    "cubic": "bicubic",
                }[mode]
                out_pt = t.interpolate(
                    in_pt,
                    scale_factor=scale_factor,
                    mode=mode_pt,
                    align_corners=align_corner,
                )
                out_pt = torch.permute(out_pt, (0, 2, 3, 1)).numpy(force=True)
                self.assertEqual(out_pt.shape, out_mx.shape)
                self.assertTrue(np.allclose(out_pt, out_mx, atol=atol))

        for dtype in ("float32",):
            for N, C in ((1, 1), (2, 3)):
                # only test cases in which target sizes are intergers
                # if not, there will be numerical difference between mlx
                # and torch due to different indices selection.
                for idim, scale_factor in (
                    ((2, 2), (1.0, 1.0)),
                    ((2, 2), (1.5, 1.5)),
                    ((2, 2), (2.0, 2.0)),
                    ((4, 4), (0.5, 0.5)),
                    ((7, 7), (2.0, 2.0)),
                    ((10, 10), (0.2, 0.2)),
                    ((11, 21), (3.0, 3.0)),
                    ((11, 21), (3.0, 2.0)),
                ):
                    # only test linear and cubic interpolation
                    # there will be numerical difference in nearest
                    # due to different indices selection.
                    for mode in ("cubic", "linear"):
                        for align_corner in (False, True):
                            run_upsample(
                                N,
                                C,
                                idim,
                                scale_factor,
                                mode,
                                align_corner,
                                dtype=dtype,
                            )
//  [python]
def test_torch_upsample(self):
        def run_upsample(
            N,
            C,
            idim,
            scale_factor,
            mode,
            align_corner,
            dtype="float32",
            atol=1e-5,
        ):
            with self.subTest(
                N=N,
                C=C,
                idim=idim,
                scale_factor=scale_factor,
                mode=mode,
                align_corner=align_corner,
            ):
                np_dtype = getattr(np, dtype)
                np.random.seed(0)
                iH, iW = idim
                in_np = np.random.normal(-1.0, 1.0, (N, iH, iW, C)).astype(np_dtype)

                in_mx = mx.array(in_np)
                in_pt = torch.from_numpy(in_np.transpose(0, 3, 1, 2)).to("cpu")

                out_mx = nn.Upsample(
                    scale_factor=scale_factor,
                    mode=mode,
                    align_corners=align_corner,
                )(in_mx)
                mode_pt = {
                    "linear": "bilinear",
                    "cubic": "bicubic",
                }[mode]
                out_pt = F.interpolate(
                    in_pt,
                    scale_factor=scale_factor,
                    mode=mode_pt,
                    align_corners=align_corner,
                )
                out_pt = torch.permute(out_pt, (0, 2, 3, 1)).numpy(force=True)
                self.assertEqual(out_pt.shape, out_mx.shape)
                self.assertTrue(np.allclose(out_pt, out_mx, atol=atol))

        for dtype in ("float32",):
            for N, C in ((1, 1), (2, 3)):
                # only test cases in which target sizes are intergers
                # if not, there will be numerical difference between mlx
                # and torch due to different indices selection.
                for idim, scale_factor in (
                    ((2, 2), (1.0, 1.0)),
                    ((2, 2), (1.5, 1.5)),
                    ((2, 2), (2.0, 2.0)),
                    ((4, 4), (0.5, 0.5)),
                    ((7, 7), (2.0, 2.0)),
                    ((10, 10), (0.2, 0.2)),
                    ((11, 21), (3.0, 3.0)),
                    ((11, 21), (3.0, 2.0)),
                ):
                    # only test linear and cubic interpolation
                    # there will be numerical difference in nearest
                    # due to different indices selection.
                    for mode in ("cubic", "linear"):
                        for align_corner in (False, True):
                            run_upsample(
                                N,
                                C,
                                idim,
                                scale_factor,
                                mode,
                                align_corner,
                                dtype=dtype,
                            )
//  [python]
def run_midas_depth_estimation(image):
    global midas_depth_estimation
    global midas_transform
    if image is None:
        return None
    if midas_depth_estimation is None or midas_transform is None:
        try_loading_midas_depth_estimation()
    if midas_depth_estimation is None or midas_transform is None:
        return None

    image = image.convert("RGB")
    image = np.asarray(image)

    device = "cpu"
    input_batch = midas_transform(image).to(device)
    with torch.no_grad():
        prediction = midas_depth_estimation(input_batch)

        prediction = t.nn.functional.interpolate(
            prediction.unsqueeze(1),
            size=image.shape[:2],
            mode="bicubic",
            align_corners=False,
        ).squeeze()

    output = prediction.cpu().numpy()
    depth = 1 - output / np.max(output)
    return depth
//  [python]
def run_midas_depth_estimation(image):
    global midas_depth_estimation
    global midas_transform
    if image is None:
        return None
    if midas_depth_estimation is None or midas_transform is None:
        try_loading_midas_depth_estimation()
    if midas_depth_estimation is None or midas_transform is None:
        return None

    image = image.convert("RGB")
    image = np.asarray(image)

    device = "cpu"
    input_batch = midas_transform(image).to(device)
    with torch.no_grad():
        prediction = midas_depth_estimation(input_batch)

        prediction = torch.nn.functional.interpolate(
            prediction.unsqueeze(1),
            size=image.shape[:2],
            mode="bicubic",
            align_corners=False,
        ).squeeze()

    output = prediction.cpu().numpy()
    depth = 1 - output / np.max(output)
    return depth
//  [python]
def _preprocess_fast_guided_blur(
    guidance: Tensor, input: Tensor, kernel_size: tuple[int, int] | int, subsample: int = 1
) -> tuple[Tensor, Tensor, tuple[int, int]]:
    ky, kx = _unpack_2d_ks(kernel_size)
    if subsample > 1:
        s = 1 / subsample
        guidance_sub = t.interpolate(guidance, scale_factor=s, mode="nearest")
        input_sub = guidance_sub if input is guidance else t.interpolate(input, scale_factor=s, mode="nearest")
        ky, kx = ((k - 1) // subsample + 1 for k in (ky, kx))
    else:
        guidance_sub = guidance
        input_sub = input
    return guidance_sub, input_sub, (ky, kx)
//  [python]
def _preprocess_fast_guided_blur(
    guidance: Tensor, input: Tensor, kernel_size: tuple[int, int] | int, subsample: int = 1
) -> tuple[Tensor, Tensor, tuple[int, int]]:
    ky, kx = _unpack_2d_ks(kernel_size)
    if subsample > 1:
        s = 1 / subsample
        guidance_sub = interpolate(guidance, scale_factor=s, mode="nearest")
        input_sub = guidance_sub if input is guidance else interpolate(input, scale_factor=s, mode="nearest")
        ky, kx = ((k - 1) // subsample + 1 for k in (ky, kx))
    else:
        guidance_sub = guidance
        input_sub = input
    return guidance_sub, input_sub, (ky, kx)
//  [python]
def forward(self, inputs):
        """Forward function."""
        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):
            prev_shape = laterals[i - 1].shape[2:]
            laterals[i - 1] += t.interpolate(
                laterals[i], size=prev_shape, mode='nearest')

        # build outputs
        # part 1: from original levels
        inter_outs = [
            self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
        ]

        # part 2: add bottom-up path
        for i in range(0, used_backbone_levels - 1):
            inter_outs[i + 1] += self.downsample_convs[i](inter_outs[i])

        outs = []
        outs.append(inter_outs[0])
        outs.extend([
            self.pafpn_convs[i - 1](inter_outs[i])
            for i in range(1, used_backbone_levels)
        ])

        # part 3: 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':
                    orig = inputs[self.backbone_end_level - 1]
                    outs.append(self.fpn_convs[used_backbone_levels](orig))
                elif self.add_extra_convs == 'on_lateral':
                    outs.append(self.fpn_convs[used_backbone_levels](
                        laterals[-1]))
                elif self.add_extra_convs == 'on_output':
                    outs.append(self.fpn_convs[used_backbone_levels](outs[-1]))
                else:
                    raise NotImplementedError
                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):
        """Forward function."""
        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):
            prev_shape = laterals[i - 1].shape[2:]
            laterals[i - 1] += F.interpolate(
                laterals[i], size=prev_shape, mode='nearest')

        # build outputs
        # part 1: from original levels
        inter_outs = [
            self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
        ]

        # part 2: add bottom-up path
        for i in range(0, used_backbone_levels - 1):
            inter_outs[i + 1] += self.downsample_convs[i](inter_outs[i])

        outs = []
        outs.append(inter_outs[0])
        outs.extend([
            self.pafpn_convs[i - 1](inter_outs[i])
            for i in range(1, used_backbone_levels)
        ])

        # part 3: 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':
                    orig = inputs[self.backbone_end_level - 1]
                    outs.append(self.fpn_convs[used_backbone_levels](orig))
                elif self.add_extra_convs == 'on_lateral':
                    outs.append(self.fpn_convs[used_backbone_levels](
                        laterals[-1]))
                elif self.add_extra_convs == 'on_output':
                    outs.append(self.fpn_convs[used_backbone_levels](outs[-1]))
                else:
                    raise NotImplementedError
                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 test_Interp(self):
        """Interp"""

        def _test_interp(input_shape, size=None, scale=None, mode=None):

            class Model(nn.Module):

                def __init__(self):
                    super(Model, self).__init__()

                def forward(self, x):
                    y = t.interpolate(x,
                                                  size,
                                                  scale_factor=scale,
                                                  mode=mode,
                                                  align_corners=False)
                    return y

            self.trace_and_test([input_shape], Model())

        _test_interp((1, 3, 100, 100), None, 4, 'bilinear')
        _test_interp((1, 1, 32, 32), None, 10,'bilinear')
        _test_interp((2, 32, 16, 16), None, 16, 'bilinear')
        if self.chip in ["bm1684x", "bm1688"]:
            _test_interp((2, 3, 224), None, 2, 'linear')
            _test_interp((2, 3, 224), (100), None, 'linear')
//  [python]
def test_Interp(self):
        """Interp"""

        def _test_interp(input_shape, size=None, scale=None, mode=None):

            class Model(nn.Module):

                def __init__(self):
                    super(Model, self).__init__()

                def forward(self, x):
                    y = nn.functional.interpolate(x,
                                                  size,
                                                  scale,
                                                  mode=mode,
                                                  align_corners=False)
                    return y

            self.trace_and_test([input_shape], Model())

        _test_interp((1, 3, 100, 100), None, 4, 'bilinear')
        _test_interp((1, 1, 32, 32), None, 10,'bilinear')
        _test_interp((2, 32, 16, 16), None, 16, 'bilinear')
        if self.chip in ["bm1684x", "bm1688"]:
            _test_interp((2, 3, 224), None, 2, 'linear')
            _test_interp((2, 3, 224), (100), None, 'linear')
//  [python]
def set_source_image(self, source_image):
        self.source = to_tensor(source_image).to(self.device)
        self.kp_source = self.kp_detector(self.source)

        if self.enc_downscale > 1:
            h, w = int(self.source.shape[2] / self.enc_downscale), int(self.source.shape[3] / self.enc_downscale)
            source_enc = torch.nn.functional.interpolate(self.source, size=(h, w), mode='bilinear')
        else:
            source_enc = self.source

        self.generator.encode_source(source_enc)
//  [python]
def forward(self, x):
        batch_size = x.shape[0]

        # pooling
        x = x.unsqueeze(1)
        x = self.pooling_layer(x)
        x = x.squeeze(1)

        # fully connected layer stack
        x = self.layers(x)

        # forked linear layers producing waveform generator parameters
        theta_backcast = self.backcast_linear_layer(x)
        theta_forecast = self.forecast_linear_layer(x)

        # set the expansion coefs in last dimension for the forecasts
        theta_forecast = theta_forecast.view(batch_size, self.nr_params, -1)

        # interpolate function expects (batch, "channels", time)
        theta_backcast = theta_backcast.unsqueeze(1)

        # interpolate both backcast and forecast from the thetas
        x_hat = F.interpolate(
            theta_backcast, size=self.input_chunk_length, mode="linear"
        )
        y_hat = F.interpolate(
            theta_forecast, size=self.output_chunk_length, mode="linear"
        )

        x_hat = x_hat.squeeze(1)  # remove 2nd dim we added before interpolation

        # Set the distribution parameters as the last dimension
        y_hat = y_hat.reshape(x.shape[0], self.output_chunk_length, self.nr_params)

        return x_hat, y_hat
//  [python]
def _guided_blur_grayscale_guidance(
    guidance: Tensor,
    input: Tensor,
    kernel_size: tuple[int, int] | int,
    eps: float | Tensor,
    border_type: str = "reflect",
    subsample: int = 1,
) -> Tensor:
    guidance_sub, input_sub, kernel_size = _preprocess_fast_guided_blur(guidance, input, kernel_size, subsample)

    mean_I = box_blur(guidance_sub, kernel_size, border_type)
    corr_I = box_blur(guidance_sub.square(), kernel_size, border_type)
    var_I = corr_I - mean_I.square()

    if input is guidance:
        mean_p = mean_I
        cov_Ip = var_I

    else:
        mean_p = box_blur(input_sub, kernel_size, border_type)
        corr_Ip = box_blur(guidance_sub * input_sub, kernel_size, border_type)
        cov_Ip = corr_Ip - mean_I * mean_p

    if isinstance(eps, Tensor):
        eps = eps.view(-1, 1, 1, 1)  # N -> NCHW

    a = cov_Ip / (var_I + eps)
    b = mean_p - a * mean_I

    mean_a = box_blur(a, kernel_size, border_type)
    mean_b = box_blur(b, kernel_size, border_type)

    if subsample > 1:
        mean_a = t.nn.functional.interpolate(mean_a, scale_factor=subsample, mode="bilinear")
        mean_b = t.nn.functional.interpolate(mean_b, scale_factor=subsample, mode="bilinear")

    return mean_a * guidance + mean_b
//  [python]
def _guided_blur_grayscale_guidance(
    guidance: Tensor,
    input: Tensor,
    kernel_size: tuple[int, int] | int,
    eps: float | Tensor,
    border_type: str = "reflect",
    subsample: int = 1,
) -> Tensor:
    guidance_sub, input_sub, kernel_size = _preprocess_fast_guided_blur(guidance, input, kernel_size, subsample)

    mean_I = box_blur(guidance_sub, kernel_size, border_type)
    corr_I = box_blur(guidance_sub.square(), kernel_size, border_type)
    var_I = corr_I - mean_I.square()

    if input is guidance:
        mean_p = mean_I
        cov_Ip = var_I

    else:
        mean_p = box_blur(input_sub, kernel_size, border_type)
        corr_Ip = box_blur(guidance_sub * input_sub, kernel_size, border_type)
        cov_Ip = corr_Ip - mean_I * mean_p

    if isinstance(eps, Tensor):
        eps = eps.view(-1, 1, 1, 1)  # N -> NCHW

    a = cov_Ip / (var_I + eps)
    b = mean_p - a * mean_I

    mean_a = box_blur(a, kernel_size, border_type)
    mean_b = box_blur(b, kernel_size, border_type)

    if subsample > 1:
        mean_a = interpolate(mean_a, scale_factor=subsample, mode="bilinear")
        mean_b = interpolate(mean_b, scale_factor=subsample, mode="bilinear")

    return mean_a * guidance + mean_b
//  [python]
def forward(self, f16, f8, f4):
        x = self.compress(f16)
        x = self.up_16_8(f8, x)
        x = self.up_8_4(f4, x)

        x = self.pred(F.relu(x))

        x = t.interpolate(
            x, scale_factor=4, mode='bilinear', align_corners=False)
        return x
//  [python]
def forward(self, f16, f8, f4):
        x = self.compress(f16)
        x = self.up_16_8(f8, x)
        x = self.up_8_4(f4, x)

        x = self.pred(F.relu(x))

        x = F.interpolate(
            x, scale_factor=4, mode='bilinear', align_corners=False)
        return x
//  [python]
def forward(self, x):
        up = t.upsample_nearest(x, scale_factor=self.stride)
        return self.conv(up)
//  [python]
def forward(self, x):
        up = nn.functional.interpolate(x, scale_factor=self.stride, mode="nearest")
        return self.conv(up)
//  [python]
def test_interpolate(self, X, size, mode, scale_factor, align_corners, nhwc_layout):
        """
        This test cover upsample_nearest2d and upsample_bilinear2d
        """
        X, (scale, zero_point, torch_type) = X
        H, W = X.shape[-2:]

        if scale_factor is not None:
            size = None
        if mode == "nearest":
            align_corners = None

        if nhwc_layout:
            if X.shape[1] < 176:
                X = np.repeat(X, 176 / X.shape[1], 1)

            X_nchw = np.ascontiguousarray(X.transpose([0, 2, 3, 1]))
            X = torch.from_numpy(X_nchw).permute([0, 3, 1, 2])

            qX = torch.quantize_per_tensor(X, scale=scale, zero_point=zero_point,
                                           dtype=torch_type).permute([0, 3, 1, 2])
        else:
            X = torch.from_numpy(X)
            qX = torch.quantize_per_tensor(X, scale=scale, zero_point=zero_point,
                                           dtype=torch_type)

        X_ref = t.nn.functional.interpolate(
            qX.int_repr().to(torch.float), size=size, scale_factor=scale_factor,
            mode=mode, align_corners=align_corners)

        ops_under_test = {
            "nn.functional": t.nn.functional.interpolate,
            "nn.quantized.functional": torch.nn.quantized.functional.interpolate
        }
        error_message = r"Results are off for {}:\n\tExpected:\n{}\n\tGot:\n{}"
        for name, op in ops_under_test.items():
            qX_hat = op(qX, size=size, scale_factor=scale_factor,
                        mode=mode, align_corners=align_corners)
            # TODO(#38095): Replace assertEqualIgnoreType. See issue #38095
            self.assertEqualIgnoreType(X_ref, qX_hat.int_repr(), atol=1.0, rtol=0,
                                       msg="{} results are off: qX_hat={} X_ref={}"
                                           .format(name, qX_hat.int_repr(), X_ref))
            self.assertEqual(scale, qX_hat.q_scale(),
                             msg=error_message.format(name + '.scale', scale, qX_hat.q_scale()))
            self.assertEqual(zero_point, qX_hat.q_zero_point(),
                             msg=error_message.format(name + '.zero_point', scale,
                                                      qX_hat.q_zero_point()))
//  [python]
def test_interpolate(self, X, size, mode, scale_factor, align_corners, nhwc_layout):
        """
        This test cover upsample_nearest2d and upsample_bilinear2d
        """
        X, (scale, zero_point, torch_type) = X
        H, W = X.shape[-2:]

        if scale_factor is not None:
            size = None
        if mode == "nearest":
            align_corners = None

        if nhwc_layout:
            if X.shape[1] < 176:
                X = np.repeat(X, 176 / X.shape[1], 1)

            X_nchw = np.ascontiguousarray(X.transpose([0, 2, 3, 1]))
            X = torch.from_numpy(X_nchw).permute([0, 3, 1, 2])

            qX = torch.quantize_per_tensor(X, scale=scale, zero_point=zero_point,
                                           dtype=torch_type).permute([0, 3, 1, 2])
        else:
            X = torch.from_numpy(X)
            qX = torch.quantize_per_tensor(X, scale=scale, zero_point=zero_point,
                                           dtype=torch_type)

        X_ref = torch.nn.functional.interpolate(
            qX.int_repr().to(torch.float), size=size, scale_factor=scale_factor,
            mode=mode, align_corners=align_corners)

        ops_under_test = {
            "nn.functional": torch.nn.functional.interpolate,
            "nn.quantized.functional": torch.nn.quantized.functional.interpolate
        }
        error_message = r"Results are off for {}:\n\tExpected:\n{}\n\tGot:\n{}"
        for name, op in ops_under_test.items():
            qX_hat = op(qX, size=size, scale_factor=scale_factor,
                        mode=mode, align_corners=align_corners)
            # TODO(#38095): Replace assertEqualIgnoreType. See issue #38095
            self.assertEqualIgnoreType(X_ref, qX_hat.int_repr(), atol=1.0, rtol=0,
                                       msg="{} results are off: qX_hat={} X_ref={}"
                                           .format(name, qX_hat.int_repr(), X_ref))
            self.assertEqual(scale, qX_hat.q_scale(),
                             msg=error_message.format(name + '.scale', scale, qX_hat.q_scale()))
            self.assertEqual(zero_point, qX_hat.q_zero_point(),
                             msg=error_message.format(name + '.zero_point', scale,
                                                      qX_hat.q_zero_point()))
//  [python]
def forward(self, latents, g=None):
        """
        Args:
            x (Tensor): feature input tensor (GPT latent).
            g (Tensor): global conditioning input tensor.

        Returns:
            Tensor: output waveform.

        Shapes:
            x: [B, C, T]
            Tensor: [B, 1, T]
        """

        z = t.interpolate(
            latents.transpose(1, 2),
            scale_factor=[self.ar_mel_length_compression / self.output_hop_length],
            mode="linear",
        ).squeeze(1)
        # upsample to the right sr
        if self.output_sample_rate != self.input_sample_rate:
            z = t.interpolate(
                z,
                scale_factor=[self.output_sample_rate / self.input_sample_rate],
                mode="linear",
            ).squeeze(0)
        o = self.waveform_decoder(z, g=g)
        return o
//  [python]
def forward(self, latents, g=None):
        """
        Args:
            x (Tensor): feature input tensor (GPT latent).
            g (Tensor): global conditioning input tensor.

        Returns:
            Tensor: output waveform.

        Shapes:
            x: [B, C, T]
            Tensor: [B, 1, T]
        """

        z = torch.nn.functional.interpolate(
            latents.transpose(1, 2),
            scale_factor=[self.ar_mel_length_compression / self.output_hop_length],
            mode="linear",
        ).squeeze(1)
        # upsample to the right sr
        if self.output_sample_rate != self.input_sample_rate:
            z = torch.nn.functional.interpolate(
                z,
                scale_factor=[self.output_sample_rate / self.input_sample_rate],
                mode="linear",
            ).squeeze(0)
        o = self.waveform_decoder(z, g=g)
        return o