jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def reshape_mel2img(self, normalized_input_features):
        """
        The input is 4 normalized log mel spectrograms. It is reshape to the common shape of images. Each channel
        should represent 1 of the 4 crops of the spectrogram. For more details, refer to the [`ClapFeatureExtractor`].
        """
        _, _, time_length, freq_length = normalized_input_features.shape

        spec_width = int(self.spec_size * self.freq_ratio)
        spec_heigth = self.spec_size // self.freq_ratio

        if time_length > spec_width or freq_length > spec_heigth:
            raise ValueError("the wav size should be less than or equal to the swin input size")

        # to avoid bicubic zero error
        if time_length < spec_width:
            normalized_input_features = t.nn.functional.interpolate(
                normalized_input_features, (spec_width, freq_length), mode="bicubic", align_corners=True
            )
        if freq_length < spec_heigth:
            normalized_input_features = t.nn.functional.interpolate(
                normalized_input_features, (time_length, spec_heigth), mode="bicubic", align_corners=True
            )

        batch, channels, time, freq = normalized_input_features.shape

        # batch_size, channels, spec_width, spec_heigth --> batch_size, channels, spec_heigth * freq_ratio, spec_width // freq_ratio
        normalized_input_features = normalized_input_features.reshape(
            batch, channels * self.freq_ratio, time // self.freq_ratio, freq
        )
        normalized_input_features = normalized_input_features.permute(0, 1, 3, 2).contiguous()
        normalized_input_features = normalized_input_features.reshape(
            batch, channels, freq * self.freq_ratio, time // self.freq_ratio
        )

        return normalized_input_features
//  [python]
def reshape_mel2img(self, normalized_input_features):
        """
        The input is 4 normalized log mel spectrograms. It is reshape to the common shape of images. Each channel
        should represent 1 of the 4 crops of the spectrogram. For more details, refer to the [`ClapFeatureExtractor`].
        """
        _, _, time_length, freq_length = normalized_input_features.shape

        spec_width = int(self.spec_size * self.freq_ratio)
        spec_heigth = self.spec_size // self.freq_ratio

        if time_length > spec_width or freq_length > spec_heigth:
            raise ValueError("the wav size should be less than or equal to the swin input size")

        # to avoid bicubic zero error
        if time_length < spec_width:
            normalized_input_features = nn.functional.interpolate(
                normalized_input_features, (spec_width, freq_length), mode="bicubic", align_corners=True
            )
        if freq_length < spec_heigth:
            normalized_input_features = nn.functional.interpolate(
                normalized_input_features, (time_length, spec_heigth), mode="bicubic", align_corners=True
            )

        batch, channels, time, freq = normalized_input_features.shape

        # batch_size, channels, spec_width, spec_heigth --> batch_size, channels, spec_heigth * freq_ratio, spec_width // freq_ratio
        normalized_input_features = normalized_input_features.reshape(
            batch, channels * self.freq_ratio, time // self.freq_ratio, freq
        )
        normalized_input_features = normalized_input_features.permute(0, 1, 3, 2).contiguous()
        normalized_input_features = normalized_input_features.reshape(
            batch, channels, freq * self.freq_ratio, time // self.freq_ratio
        )

        return normalized_input_features
//  [python]
def forward(self, x, y):
        """
        Args:
            x: low level feature
            y: high level feature
        """
        h, w = x.size(2), x.size(3)
        y_up = F.upsample(y, size=(h, w), mode=self.upscale_mode)
        x = self.conv2(x)
        y = self.conv1(y)
        z = torch.mul(x, y)
        return y_up + z
//  [python]
def forward(self, x, y):
        """
        Args:
            x: low level feature
            y: high level feature
        """
        h, w = x.size(2), x.size(3)
        y_up = F.interpolate(y, size=(h, w), mode=self.upscale_mode, align_corners=self.align_corners)
        x = self.conv2(x)
        y = self.conv1(y)
        z = torch.mul(x, y)
        return y_up + z
//  [python]
def forward(self, x):
        h, w = x.size(2), x.size(3)
        b1 = self.branch1(x)
        upscale_parameters = dict(mode=self.upscale_mode, align_corners=self.align_corners)
        b1 = F.upsample(b1, size=(h, w), **upscale_parameters)

        mid = self.mid(x)
        x1 = self.down1(x)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x3 = F.upsample(x3, size=(h // 4, w // 4), **upscale_parameters)

        x2 = self.conv2(x2)
        x = x2 + x3
        x = F.upsample(x, size=(h // 2, w // 2), **upscale_parameters)

        x1 = self.conv1(x1)
        x = x + x1
        x = F.upsample(x, size=(h, w), **upscale_parameters)

        x = torch.mul(x, mid)
        x = x + b1
        return x
//  [python]
def forward(self, x):
        h, w = x.size(2), x.size(3)
        b1 = self.branch1(x)
        upscale_parameters = dict(mode=self.upscale_mode, align_corners=self.align_corners)
        b1 = F.interpolate(b1, size=(h, w), **upscale_parameters)

        mid = self.mid(x)
        x1 = self.down1(x)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x3 = F.interpolate(x3, size=(h // 4, w // 4), **upscale_parameters)

        x2 = self.conv2(x2)
        x = x2 + x3
        x = F.interpolate(x, size=(h // 2, w // 2), **upscale_parameters)

        x1 = self.conv1(x1)
        x = x + x1
        x = F.interpolate(x, size=(h, w), **upscale_parameters)

        x = torch.mul(x, mid)
        x = x + b1
        return x
//  [python]
def forward(self, *xs):
        """Forward pass.
        Returns:
            tensor: output
        """
        output = xs[0]

        if len(xs) == 2:
            output += self.resConfUnit1(xs[1])

        output = self.resConfUnit2(output)

        output = t.interpolate(
            output, scale_factor=2, mode="bilinear", align_corners=True
        )

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

        if len(xs) == 2:
            output += self.resConfUnit1(xs[1])

        output = self.resConfUnit2(output)

        output = nn.functional.interpolate(
            output, scale_factor=2, mode="bilinear", align_corners=True
        )

        return output
//  [python]
def _resize_pos_embed(self, posemb, gs_h, gs_w):
    posemb_tok, posemb_grid = (
        posemb[:, : self.start_index],
        posemb[0, self.start_index :],
    )

    gs_old = int(math.sqrt(posemb_grid.shape[0]))

    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
    posemb_grid = t.interpolate(posemb_grid, size=(gs_h, gs_w), mode="bilinear")
    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_h * gs_w, -1)

    posemb = torch.cat([posemb_tok, posemb_grid], dim=1)

    return posemb
//  [python]
def _resize_pos_embed(self, posemb, gs_h, gs_w):
    posemb_tok, posemb_grid = (
        posemb[:, : self.start_index],
        posemb[0, self.start_index :],
    )

    gs_old = int(math.sqrt(posemb_grid.shape[0]))

    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
    posemb_grid = F.interpolate(posemb_grid, size=(gs_h, gs_w), mode="bilinear")
    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_h * gs_w, -1)

    posemb = torch.cat([posemb_tok, posemb_grid], dim=1)

    return posemb
//  [python]
def forward(self, *xs):
        """Forward pass.
        Returns:
            tensor: output
        """
        output = xs[0]

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

        output = self.resConfUnit2(output)

        output = nn.functional.upsample(
            output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
        )

        output = self.out_conv(output)

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

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

        output = self.resConfUnit2(output)

        output = nn.functional.interpolate(
            output, scale_factor=2, mode="bilinear", align_corners=self.align_corners
        )

        output = self.out_conv(output)

        return output
//  [python]
def forward(self, points, K, T,scale=0):
        K[:,:2,:] = (K[:,:2,:]/(2 ** scale)).trunc()
        b,n,n = K.shape
        b,w,h,n,n = T.shape
        pad = torch.tensor([0.0,0.0,0.0]).view(1,3,1).expand(b,3,1).cuda()
        K = torch.cat([K,pad],-1)
        pad = torch.tensor([0.0,0.0,0.0,1.0]).view(1,1,4).expand(b,1,4).cuda()
        K = torch.cat([K,pad],1)
        K = K.unsqueeze(1).unsqueeze(1).expand_as(T).contiguous().view(b*w*h,n,n).contiguous()
        T = T.view(b*w*h,n,n).contiguous()
        P = torch.matmul(K, T)[:, :3, :]
        P = P.view(b,w,h,3,4).contiguous().view(b,w,h,12).contiguous().permute(0,3,1,2).contiguous()
        P = F.upsample(P,scale_factor=32,mode='area')
        P = P.view(b,12,256*256).contiguous().permute(0,2,1).contiguous().view(-1,3,4).contiguous()
        points = points.permute(0,2,1).contiguous().view(-1,4).contiguous().unsqueeze(-1)
        cam_points = torch.matmul(P, points)
        cam_points = cam_points.view(b,256*256,3,1).contiguous().permute(0,2,3,1).squeeze(2)

        pix_coords = cam_points[:, :2, :] / (cam_points[:, 2, :].unsqueeze(1) + self.eps)
        pix_coords = pix_coords.view(self.batch_size, 2, self.height, self.width)
        pix_coords = pix_coords.permute(0, 2, 3, 1)
        pix_coords[..., 0] /= self.width - 1
        pix_coords[..., 1] /= self.height - 1
        pix_coords = (pix_coords - 0.5) * 2
        return pix_coords
//  [python]
def forward(self, points, K, T,scale=0):
        # K[0, :] *= self.width // (2 ** scale)
        # K[1, :] *= self.height // (2 ** scale)
        K[:,:2,:] = (K[:,:2,:]/(2 ** scale)).trunc()
        b,n,n = K.shape
        b,w,h,n,n = T.shape
        pad = torch.tensor([0.0,0.0,0.0]).view(1,3,1).expand(b,3,1).cuda()
        K = torch.cat([K,pad],-1)
        pad = torch.tensor([0.0,0.0,0.0,1.0]).view(1,1,4).expand(b,1,4).cuda()
        K = torch.cat([K,pad],1)
        K = K.unsqueeze(1).unsqueeze(1).expand_as(T).contiguous().view(b*w*h,n,n).contiguous()
        T = T.view(b*w*h,n,n).contiguous()
        P = torch.matmul(K, T)[:, :3, :]
        P = P.view(b,w,h,3,4).contiguous().view(b,w,h,12).contiguous().permute(0,3,1,2).contiguous()
        P = F.interpolate(P,scale_factor=32,mode='area')
        P = P.view(b,12,256*256).contiguous().permute(0,2,1).contiguous().view(-1,3,4).contiguous()
        points = points.permute(0,2,1).contiguous().view(-1,4).contiguous().unsqueeze(-1)
        cam_points = torch.matmul(P, points)
        cam_points = cam_points.view(b,256*256,3,1).contiguous().permute(0,2,3,1).squeeze(2)

        pix_coords = cam_points[:, :2, :] / (cam_points[:, 2, :].unsqueeze(1) + self.eps)
        pix_coords = pix_coords.view(self.batch_size, 2, self.height, self.width)
        pix_coords = pix_coords.permute(0, 2, 3, 1)
        pix_coords[..., 0] /= self.width - 1
        pix_coords[..., 1] /= self.height - 1
        pix_coords = (pix_coords - 0.5) * 2
        return pix_coords
//  [python]
def forward(self, x):
        """Get Inception feature maps.

        Args:
            x (Tensor): Input tensor of shape (b, 3, h, w).
                Values are expected to be in range (-1, 1). You can also input
                (0, 1) with setting normalize_input = True.

        Returns:
            list[Tensor]: Corresponding to the selected output block, sorted
            ascending by index.
        """
        output = []

        if self.resize_input:
            x = t.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False)

        if self.normalize_input:
            x = 2 * x - 1  # Scale from range (0, 1) to range (-1, 1)

        for idx, block in enumerate(self.blocks):
            x = block(x)
            if idx in self.output_blocks:
                output.append(x)

            if idx == self.last_needed_block:
                break

        return output
//  [python]
def forward(self, x):
        """Get Inception feature maps.

        Args:
            x (Tensor): Input tensor of shape (b, 3, h, w).
                Values are expected to be in range (-1, 1). You can also input
                (0, 1) with setting normalize_input = True.

        Returns:
            list[Tensor]: Corresponding to the selected output block, sorted
            ascending by index.
        """
        output = []

        if self.resize_input:
            x = F.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False)

        if self.normalize_input:
            x = 2 * x - 1  # Scale from range (0, 1) to range (-1, 1)

        for idx, block in enumerate(self.blocks):
            x = block(x)
            if idx in self.output_blocks:
                output.append(x)

            if idx == self.last_needed_block:
                break

        return output
//  [python]
def forward(self, x):
        """
        Arguments:
            x (list[Tensor]): feature maps for each feature level.
        Returns:
            results (tuple[Tensor]): feature maps after FPN layers.
                They are ordered from highest resolution first.
        """
        if type(x) is tuple:
            # for the case of VL backbone
            x, x_text = x[0], x[1]
        # print([v.shape for v in x])
        swint_feature_c4 = None
        if self.return_swint_feature_before_fusion:
            # TODO: here we only return last single scale feature map before the backbone fusion, should be more flexible
            swint_feature_c4 = x[-2]

        if self.spp_block:
            last_inner = getattr(self, self.inner_blocks[-1])(self.spp_block(x[-1]))
        else:
            last_inner = getattr(self, self.inner_blocks[-1])(x[-1])
        results = []
        results.append(getattr(self, self.layer_blocks[-1])(last_inner))
        for feature, inner_block, layer_block in zip(
            x[:-1][::-1], self.inner_blocks[:-1][::-1], self.layer_blocks[:-1][::-1]
        ):
            if not inner_block:
                continue
            inner_lateral = getattr(self, inner_block)(feature)

            if inner_lateral.shape[-2:] != last_inner.shape[-2:]:
                # TODO: could also give size instead of
                inner_top_down = F.upsample(last_inner, size=inner_lateral.shape[-2:], mode="nearest")
            else:
                inner_top_down = last_inner

            # TODO use size instead of scale to make it robust to different sizes
            # inner_top_down = F.upsample(last_inner, size=inner_lateral.shape[-2:],
            # mode='bilinear', align_corners=False)
            last_inner = inner_lateral + inner_top_down
            if self.drop_block and self.training:
                results.insert(0, getattr(self, layer_block)(self.drop_block(last_inner)))
            else:
                results.insert(0, getattr(self, layer_block)(last_inner))

        if self.pan_blocks:
            pan_results = []
            last_outer = results[0]
            pan_results.append(last_outer)
            for outer_top_down, pan_block in zip(results[1:], self.pan_blocks):

                if self.drop_block and self.training:
                    pan_lateral = getattr(self, pan_block[0])(self.drop_block(last_outer))
                else:
                    pan_lateral = getattr(self, pan_block[0])(last_outer)

                last_outer = getattr(self, pan_block[1])(pan_lateral + outer_top_down)
                pan_results.append(last_outer)
            results = pan_results

        if isinstance(self.top_blocks, LastLevelP6P7):
            last_results = self.top_blocks(x[-1], results[-1])
            results.extend(last_results)
        elif isinstance(self.top_blocks, LastLevelMaxPool):
            last_results = self.top_blocks(results[-1])
            results.extend(last_results)

        try:
            return tuple(results), x_text, swint_feature_c4
        except NameError as e:
            return tuple(results)
//  [python]
def forward(self, x):
        """
        Arguments:
            x (list[Tensor]): feature maps for each feature level.
        Returns:
            results (tuple[Tensor]): feature maps after FPN layers.
                They are ordered from highest resolution first.
        """
        if type(x) is tuple:
            # for the case of VL backbone
            x, x_text = x[0], x[1]
        # print([v.shape for v in x])
        swint_feature_c4 = None
        if self.return_swint_feature_before_fusion:
            # TODO: here we only return last single scale feature map before the backbone fusion, should be more flexible
            swint_feature_c4 = x[-2]

        if self.spp_block:
            last_inner = getattr(self, self.inner_blocks[-1])(self.spp_block(x[-1]))
        else:
            last_inner = getattr(self, self.inner_blocks[-1])(x[-1])
        results = []
        results.append(getattr(self, self.layer_blocks[-1])(last_inner))
        for feature, inner_block, layer_block in zip(
            x[:-1][::-1], self.inner_blocks[:-1][::-1], self.layer_blocks[:-1][::-1]
        ):
            if not inner_block:
                continue
            inner_lateral = getattr(self, inner_block)(feature)

            if inner_lateral.shape[-2:] != last_inner.shape[-2:]:
                # TODO: could also give size instead of
                inner_top_down = F.interpolate(last_inner, size=inner_lateral.shape[-2:], mode="nearest")
            else:
                inner_top_down = last_inner

            # TODO use size instead of scale to make it robust to different sizes
            # inner_top_down = F.upsample(last_inner, size=inner_lateral.shape[-2:],
            # mode='bilinear', align_corners=False)
            last_inner = inner_lateral + inner_top_down
            if self.drop_block and self.training:
                results.insert(0, getattr(self, layer_block)(self.drop_block(last_inner)))
            else:
                results.insert(0, getattr(self, layer_block)(last_inner))

        if self.pan_blocks:
            pan_results = []
            last_outer = results[0]
            pan_results.append(last_outer)
            for outer_top_down, pan_block in zip(results[1:], self.pan_blocks):

                if self.drop_block and self.training:
                    pan_lateral = getattr(self, pan_block[0])(self.drop_block(last_outer))
                else:
                    pan_lateral = getattr(self, pan_block[0])(last_outer)

                last_outer = getattr(self, pan_block[1])(pan_lateral + outer_top_down)
                pan_results.append(last_outer)
            results = pan_results

        if isinstance(self.top_blocks, LastLevelP6P7):
            last_results = self.top_blocks(x[-1], results[-1])
            results.extend(last_results)
        elif isinstance(self.top_blocks, LastLevelMaxPool):
            last_results = self.top_blocks(results[-1])
            results.extend(last_results)

        try:
            return tuple(results), x_text, swint_feature_c4
        except NameError as e:
            return tuple(results)
//  [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:257]
        position = x[..., 257:259]
        return prediction, latent, position
//  [python]
def forward(self, input, alpha, instance_noise=0, **kwargs):
        start = self.img_size_to_layer[input.shape[-1]]
        x = self.fromRGB[start](input)
        for i, layer in enumerate(self.layers[start:]):
            if i == 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)
        return x
//  [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, x, mask):
        """ Forward function.

            Args:
                x: input rgb images
                mask: input padding masks [0: rgb values, 1: padded values]

            Returns:
                patch_outs: multi-scale [PATCH] tokens (four scales are used)
                    these tokens are the first input of the neck decoder
                det_tgt: final [DET] tokens obtained at the last stage
                    this tokens are the second input of the neck decoder
                det_pos: the learnable pos encoding for [DET] tokens.
                    these encodings are used to generate reference points in deformable attention
        """

        # original input shape
        B, _, _ = x.shape[0], x.shape[2], x.shape[3]

        # patch embedding
        x = self.patch_embed(x)

        Wh, Ww = x.size(2), x.size(3)
        x = x.flatten(2).transpose(1, 2)
        x = self.pos_drop(x)

        # expand det_token for all examples in the batch
        det_token = self.det_token.expand(B, -1, -1)

        # det pos encoding -> will be projected in each block
        det_pos = self.det_pos_embed

        # prepare a mask for cross attention
        mask = t.nn.functional.interpolate(
            mask[None].float(),
            size=(Wh // self.mask_divisor,
                  Ww // self.mask_divisor)).to(torch.bool)[0]

        patch_outs = []
        for stage in range(self.num_layers):
            layer = self.layers[stage]

            # whether to use cross-attention
            cross_attn = True if stage in self.cross_indices else False

            # concat input
            x = torch.cat([x, det_token], dim=1)

            # inference
            x_out, H, W, x, Wh, Ww = layer(
                x,
                Wh,
                Ww,
                # additional input for VIDT
                input_mask=mask,
                det_pos=det_pos,
                cross_attn=cross_attn)

            x, det_token = x[:, :-self.det_token_num, :], x[:, -self.
                                                            det_token_num:, :]

            # Aggregate intermediate outputs
            if stage > 0:
                patch_out = x_out[:, :-self.det_token_num, :].view(
                    B, H, W, -1).permute(0, 3, 1, 2)
                patch_outs.append(patch_out)

        # patch token reduced from last stage output
        patch_outs.append(x.view(B, Wh, Ww, -1).permute(0, 3, 1, 2))

        # det token
        det_tgt = x_out[:, -self.det_token_num:, :].permute(0, 2, 1)

        # det token pos encoding
        det_pos = det_pos.permute(0, 2, 1)

        features_0, features_1, features_2, features_3 = patch_outs
        return features_0, features_1, features_2, features_3, det_tgt, det_pos
//  [python]
def forward(self, x, mask):
        """ Forward function.

            Args:
                x: input rgb images
                mask: input padding masks [0: rgb values, 1: padded values]

            Returns:
                patch_outs: multi-scale [PATCH] tokens (four scales are used)
                    these tokens are the first input of the neck decoder
                det_tgt: final [DET] tokens obtained at the last stage
                    this tokens are the second input of the neck decoder
                det_pos: the learnable pos encoding for [DET] tokens.
                    these encodings are used to generate reference points in deformable attention
        """

        # original input shape
        B, _, _ = x.shape[0], x.shape[2], x.shape[3]

        # patch embedding
        x = self.patch_embed(x)

        Wh, Ww = x.size(2), x.size(3)
        x = x.flatten(2).transpose(1, 2)
        x = self.pos_drop(x)

        # expand det_token for all examples in the batch
        det_token = self.det_token.expand(B, -1, -1)

        # det pos encoding -> will be projected in each block
        det_pos = self.det_pos_embed

        # prepare a mask for cross attention
        mask = F.interpolate(
            mask[None].float(),
            size=(Wh // self.mask_divisor,
                  Ww // self.mask_divisor)).to(torch.bool)[0]

        patch_outs = []
        for stage in range(self.num_layers):
            layer = self.layers[stage]

            # whether to use cross-attention
            cross_attn = True if stage in self.cross_indices else False

            # concat input
            x = torch.cat([x, det_token], dim=1)

            # inference
            x_out, H, W, x, Wh, Ww = layer(
                x,
                Wh,
                Ww,
                # additional input for VIDT
                input_mask=mask,
                det_pos=det_pos,
                cross_attn=cross_attn)

            x, det_token = x[:, :-self.det_token_num, :], x[:, -self.
                                                            det_token_num:, :]

            # Aggregate intermediate outputs
            if stage > 0:
                patch_out = x_out[:, :-self.det_token_num, :].view(
                    B, H, W, -1).permute(0, 3, 1, 2)
                patch_outs.append(patch_out)

        # patch token reduced from last stage output
        patch_outs.append(x.view(B, Wh, Ww, -1).permute(0, 3, 1, 2))

        # det token
        det_tgt = x_out[:, -self.det_token_num:, :].permute(0, 2, 1)

        # det token pos encoding
        det_pos = det_pos.permute(0, 2, 1)

        features_0, features_1, features_2, features_3 = patch_outs
        return features_0, features_1, features_2, features_3, det_tgt, det_pos
//  [python]
def forward(self, input, alpha, instance_noise=0, **kwargs):
        if instance_noise > 0:
            input = input + torch.randn_like(input) * instance_noise
        start = self.img_size_to_layer[input.shape[-1]]
        x = self.fromRGB[start](input)
        for i, layer in enumerate(self.layers[start:]):
            if i == 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:257]
        position = x[..., 257:259]
        return prediction, latent, position
//  [python]
def forward(self, input, alpha, **kwargs):
        size = input.shape[-1]
        log_size = int(math.log(size, 2))
        if global_cfg.tl_debug:
            VerboseModel.forward_verbose(self.conv_in[f'{2 ** log_size}'], inputs_args=(input,), name_prefix=f'conv_in[{2 ** log_size}].')
        cur_size_out = self.conv_in[f'{2 ** log_size}'](input)
        if global_cfg.tl_debug:
            VerboseModel.forward_verbose(self.convs[f'{2 ** log_size}'], inputs_args=(cur_size_out,), name_prefix=f'convs[{2 ** log_size}].')
        cur_size_out = self.convs[f'{2 ** log_size}'](cur_size_out)
        if alpha < 1:
            down_input = t.interpolate(input, scale_factor=0.5, mode='bilinear')
            if global_cfg.tl_debug:
                VerboseModel.forward_verbose(self.conv_in[f'{2 ** (log_size - 1)}'], inputs_args=(down_input,), name_prefix=f'conv_in[{2 ** (log_size - 1)}].')
            down_size_out = self.conv_in[f'{2 ** (log_size - 1)}'](down_input)
            out = alpha * cur_size_out + (1 - alpha) * down_size_out
        else:
            out = cur_size_out
        for i in range(log_size - 1, 2, -1):
            if global_cfg.tl_debug:
                VerboseModel.forward_verbose(self.convs[f'{2 ** i}'], inputs_args=(out,), name_prefix=f'convs[{2 ** i}].')
            out = self.convs[f'{2 ** i}'](out)
        batch, channel, height, width = out.shape
        group = min(batch, self.stddev_group)
        stddev = out.view(group, -1, self.stddev_feat, channel // self.stddev_feat, height, width)
        stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-08)
        stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
        stddev = stddev.repeat(group, 1, height, width)
        out = torch.cat([out, stddev], 1)
        if global_cfg.tl_debug:
            VerboseModel.forward_verbose(self.final_conv, inputs_args=(out,), name_prefix='final_conv.')
        out = self.final_conv(out)
        out = out.view(batch, -1)
        if global_cfg.tl_debug:
            VerboseModel.forward_verbose(self.final_linear, inputs_args=(out,), name_prefix='final_linear.')
        out = self.final_linear(out)
        latent, position = None, None
        return out, latent, position
//  [python]
def forward(self, input, alpha, **kwargs):
        size = input.shape[-1]
        log_size = int(math.log(size, 2))
        if global_cfg.tl_debug:
            VerboseModel.forward_verbose(self.conv_in[f'{2 ** log_size}'], inputs_args=(input,), name_prefix=f'conv_in[{2 ** log_size}].')
        cur_size_out = self.conv_in[f'{2 ** log_size}'](input)
        if global_cfg.tl_debug:
            VerboseModel.forward_verbose(self.convs[f'{2 ** log_size}'], inputs_args=(cur_size_out,), name_prefix=f'convs[{2 ** log_size}].')
        cur_size_out = self.convs[f'{2 ** log_size}'](cur_size_out)
        if alpha < 1:
            down_input = F.interpolate(input, scale_factor=0.5, mode='bilinear')
            if global_cfg.tl_debug:
                VerboseModel.forward_verbose(self.conv_in[f'{2 ** (log_size - 1)}'], inputs_args=(down_input,), name_prefix=f'conv_in[{2 ** (log_size - 1)}].')
            down_size_out = self.conv_in[f'{2 ** (log_size - 1)}'](down_input)
            out = alpha * cur_size_out + (1 - alpha) * down_size_out
        else:
            out = cur_size_out
        for i in range(log_size - 1, 2, -1):
            if global_cfg.tl_debug:
                VerboseModel.forward_verbose(self.convs[f'{2 ** i}'], inputs_args=(out,), name_prefix=f'convs[{2 ** i}].')
            out = self.convs[f'{2 ** i}'](out)
        batch, channel, height, width = out.shape
        group = min(batch, self.stddev_group)
        stddev = out.view(group, -1, self.stddev_feat, channel // self.stddev_feat, height, width)
        stddev = torch.sqrt(stddev.var(0, unbiased=False) + 1e-08)
        stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
        stddev = stddev.repeat(group, 1, height, width)
        out = torch.cat([out, stddev], 1)
        if global_cfg.tl_debug:
            VerboseModel.forward_verbose(self.final_conv, inputs_args=(out,), name_prefix='final_conv.')
        out = self.final_conv(out)
        out = out.view(batch, -1)
        if global_cfg.tl_debug:
            VerboseModel.forward_verbose(self.final_linear, inputs_args=(out,), name_prefix='final_linear.')
        out = self.final_linear(out)
        latent, position = None, None
        return out, latent, position
//  [python]
def resize_attention_map(attentions, height, width, align_corners=False):
    """
    Args:
        attentions (`torch.Tensor`): attention map of shape [batch_size, groups, feat_height*feat_width]
        height (`int`): height of the output attention map
        width (`int`): width of the output attention map
        align_corners (`bool`, *optional*): the `align_corner` argument for `nn.functional.interpolate`.

    Returns:
        `torch.Tensor`: resized attention map of shape [batch_size, groups, height, width]
    """

    scale = (height * width // attentions.shape[2]) ** 0.5
    if height > width:
        feat_width = int(np.round(width / scale))
        feat_height = attentions.shape[2] // feat_width
    else:
        feat_height = int(np.round(height / scale))
        feat_width = attentions.shape[2] // feat_height

    batch_size = attentions.shape[0]
    groups = attentions.shape[1]  # number of group token
    # [batch_size, groups, height*width, groups] -> [batch_size, groups, height, width]
    attentions = attentions.reshape(batch_size, groups, feat_height, feat_width)
    attentions = t.interpolate(
        attentions, size=(height, width), mode="bilinear", align_corners=align_corners
    )
    return attentions
//  [python]
def resize_attention_map(attentions, height, width, align_corners=False):
    """
    Args:
        attentions (`torch.Tensor`): attention map of shape [batch_size, groups, feat_height*feat_width]
        height (`int`): height of the output attention map
        width (`int`): width of the output attention map
        align_corners (`bool`, *optional*): the `align_corner` argument for `nn.functional.interpolate`.

    Returns:
        `torch.Tensor`: resized attention map of shape [batch_size, groups, height, width]
    """

    scale = (height * width // attentions.shape[2]) ** 0.5
    if height > width:
        feat_width = int(np.round(width / scale))
        feat_height = attentions.shape[2] // feat_width
    else:
        feat_height = int(np.round(height / scale))
        feat_width = attentions.shape[2] // feat_height

    batch_size = attentions.shape[0]
    groups = attentions.shape[1]  # number of group token
    # [batch_size, groups, height*width, groups] -> [batch_size, groups, height, width]
    attentions = attentions.reshape(batch_size, groups, feat_height, feat_width)
    attentions = nn.functional.interpolate(
        attentions, size=(height, width), mode="bilinear", align_corners=align_corners
    )
    return attentions
//  [python]
def forward(self, x, segmap, degraded_image):

        # Part 1. generate parameter-free normalized activations
        normalized = self.param_free_norm(x)

        # Part 2. produce scaling and bias conditioned on semantic map
        segmap = t.interpolate(segmap, size=x.size()[2:], mode="nearest")
        degraded_face = t.interpolate(degraded_image, size=x.size()[2:], mode="bilinear")

        if self.opt.no_parsing_map:
            actv = self.mlp_shared(degraded_face)
        else:
            actv = self.mlp_shared(torch.cat((segmap, degraded_face), dim=1))
        gamma = self.mlp_gamma(actv)
        beta = self.mlp_beta(actv)

        # apply scale and bias
        out = normalized * (1 + gamma) + beta

        return out
//  [python]
def forward(self, x, segmap, degraded_image):

        # Part 1. generate parameter-free normalized activations
        normalized = self.param_free_norm(x)

        # Part 2. produce scaling and bias conditioned on semantic map
        segmap = F.interpolate(segmap, size=x.size()[2:], mode="nearest")
        degraded_face = F.interpolate(degraded_image, size=x.size()[2:], mode="bilinear")

        if self.opt.no_parsing_map:
            actv = self.mlp_shared(degraded_face)
        else:
            actv = self.mlp_shared(torch.cat((segmap, degraded_face), dim=1))
        gamma = self.mlp_gamma(actv)
        beta = self.mlp_beta(actv)

        # apply scale and bias
        out = normalized * (1 + gamma) + beta

        return out
//  [python]
def upsample(signal: torch.Tensor, factor: int) -> torch.Tensor:
    signal = signal.permute(0, 2, 1)
    signal = t.nn.functional.interpolate(torch.cat((signal, signal[:, :, -1:]), 2), size=signal.shape[-1] * factor + 1, mode='linear', align_corners=True)
    signal = signal[:, :, :-1]
    return signal.permute(0, 2, 1)
//  [python]
def upsample(signal: torch.Tensor, factor: int) -> torch.Tensor:
    signal = signal.permute(0, 2, 1)
    signal = nn.functional.interpolate(torch.cat((signal, signal[:, :, -1:]), 2), size=signal.shape[-1] * factor + 1, mode='linear', align_corners=True)
    signal = signal[:, :, :-1]
    return signal.permute(0, 2, 1)
//  [python]
def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
        """
        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
        resolution images.

        Source:
        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
        """

        npatch = embeddings.shape[1]
        if npatch == self.position_embeddings.shape[1] and height == width:
            return self.position_embeddings
        patch_pos_embed = self.position_embeddings
        num_original_pos_embed = patch_pos_embed.shape[1]
        dim = embeddings.shape[-1]
        feat_height = height // self.config.patch_size
        feat_width = width // self.config.patch_size
        # we add a small number to avoid floating point error in the interpolation
        # see discussion at https://github.com/facebookresearch/dino/issues/8
        feat_height, feat_width = feat_height + 0.1, feat_width + 0.1
        original_height = original_width = math.sqrt(num_original_pos_embed)
        reshaped_patch_pos_embed = patch_pos_embed.reshape(1, int(original_height), int(original_width), dim).permute(
            0, 3, 1, 2
        )
        scale_factor = (feat_height / original_height, feat_width / original_width)
        patch_pos_embed = nn.functional.interpolate(
            reshaped_patch_pos_embed,
            scale_factor=scale_factor,
            mode="bicubic",
            align_corners=False,
        )
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return patch_pos_embed
//  [python]
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        if self.up_flag:
            x1 = F.upsample(x1, size=x2.shape[2:], mode='bilinear')
        return self.relu(self.bn(self.conv(torch.cat([x1, x2], dim=1))))
//  [python]
def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        if self.up_flag:
            x1 = F.interpolate(x1, size=x2.shape[2:], mode='bilinear', align_corners=False)
        return self.relu(self.bn(self.conv(torch.cat([x1, x2], dim=1))))
//  [python]
def interpolate_rel_pos_embed(state_dict_origin,
                              state_dict_model,
                              temporal=True,
                              verbose=False):
    rel_pos_embed_types = ['rel_pos_h', 'rel_pos_w']
    if temporal:
        rel_pos_embed_types += ['rel_pos_t']

    state_dict_inflated = state_dict_origin.copy()
    for k, v2d in state_dict_origin.items():
        if any([x in k for x in rel_pos_embed_types]):
            v3d = state_dict_model[k]
            if v2d.shape[0] != v3d.shape[0]:
                rel_pos_resized = t.nn.functional.interpolate(
                    v2d.reshape(1, v2d.shape[0], -1).permute(0, 2, 1),
                    size=v3d.shape[0],
                    mode='linear',
                )
                v3d = rel_pos_resized.reshape(-1, v3d.shape[0]).permute(1, 0)
                if verbose:
                    print('Inflate {}: {} -> {}: {}'.format(
                        k, v2d.shape, k, v3d.shape))
            else:
                v3d = v2d
            state_dict_inflated[k] = v3d.clone()
    return state_dict_inflated
//  [python]
def interpolate_rel_pos_embed(state_dict_origin,
                              state_dict_model,
                              temporal=True,
                              verbose=False):
    rel_pos_embed_types = ['rel_pos_h', 'rel_pos_w']
    if temporal:
        rel_pos_embed_types += ['rel_pos_t']

    state_dict_inflated = state_dict_origin.copy()
    for k, v2d in state_dict_origin.items():
        if any([x in k for x in rel_pos_embed_types]):
            v3d = state_dict_model[k]
            if v2d.shape[0] != v3d.shape[0]:
                rel_pos_resized = F.interpolate(
                    v2d.reshape(1, v2d.shape[0], -1).permute(0, 2, 1),
                    size=v3d.shape[0],
                    mode='linear',
                )
                v3d = rel_pos_resized.reshape(-1, v3d.shape[0]).permute(1, 0)
                if verbose:
                    print('Inflate {}: {} -> {}: {}'.format(
                        k, v2d.shape, k, v3d.shape))
            else:
                v3d = v2d
            state_dict_inflated[k] = v3d.clone()
    return state_dict_inflated
//  [python]
def _get_pos_embed(self, pos_embed, bcthw):
        t, h, w = bcthw[-3], bcthw[-2], bcthw[-1]
        if self.cls_embed_on:
            cls_pos_embed = pos_embed[:, 0:1, :]
            pos_embed = pos_embed[:, 1:]
        txy_num = pos_embed.shape[1]
        p_t, p_h, p_w = self.patch_dims
        assert p_t * p_h * p_w == txy_num

        if (p_t, p_h, p_w) != (t, h, w):
            new_pos_embed = t.upsample(
                pos_embed[:, :, :].reshape(1, p_t, p_h, p_w,
                                           -1).permute(0, 4, 1, 2, 3),
                scale_factor=(t / p_t, h / p_h, w / p_w),
                mode='trilinear',
            )
            pos_embed = new_pos_embed.reshape(1, -1,
                                              t * h * w).permute(0, 2, 1)

        if self.cls_embed_on:
            pos_embed = torch.cat((cls_pos_embed, pos_embed), dim=1)

        return pos_embed
//  [python]
def _get_pos_embed(self, pos_embed, bcthw):
        t, h, w = bcthw[-3], bcthw[-2], bcthw[-1]
        if self.cls_embed_on:
            cls_pos_embed = pos_embed[:, 0:1, :]
            pos_embed = pos_embed[:, 1:]
        txy_num = pos_embed.shape[1]
        p_t, p_h, p_w = self.patch_dims
        assert p_t * p_h * p_w == txy_num

        if (p_t, p_h, p_w) != (t, h, w):
            new_pos_embed = F.interpolate(
                pos_embed[:, :, :].reshape(1, p_t, p_h, p_w,
                                           -1).permute(0, 4, 1, 2, 3),
                size=(t, h, w),
                mode='trilinear',
            )
            pos_embed = new_pos_embed.reshape(1, -1,
                                              t * h * w).permute(0, 2, 1)

        if self.cls_embed_on:
            pos_embed = torch.cat((cls_pos_embed, pos_embed), dim=1)

        return pos_embed
//  [python]
def resize(input, size=None, scale_factor=None, mode="nearest", align_corners=None, warning=False):
    if warning:
        if size is not None and align_corners:
            input_h, input_w = tuple(int(x) for x in input.shape[2:])
            output_h, output_w = tuple(int(x) for x in size)
            if output_h > input_h or output_w > output_h:
                if (
                    (output_h > 1 and output_w > 1 and input_h > 1 and input_w > 1)
                    and (output_h - 1) % (input_h - 1)
                    and (output_w - 1) % (input_w - 1)
                ):
                    warnings.warn(
                        f"When align_corners={align_corners}, "
                        "the output would more aligned if "
                        f"input size {(input_h, input_w)} is `x+1` and "
                        f"out size {(output_h, output_w)} is `nx+1`"
                    )
    return t.interpolate(input, size, scale_factor, mode, align_corners)
//  [python]
def resize(input, size=None, scale_factor=None, mode="nearest", align_corners=None, warning=False):
    if warning:
        if size is not None and align_corners:
            input_h, input_w = tuple(int(x) for x in input.shape[2:])
            output_h, output_w = tuple(int(x) for x in size)
            if output_h > input_h or output_w > output_h:
                if (
                    (output_h > 1 and output_w > 1 and input_h > 1 and input_w > 1)
                    and (output_h - 1) % (input_h - 1)
                    and (output_w - 1) % (input_w - 1)
                ):
                    warnings.warn(
                        f"When align_corners={align_corners}, "
                        "the output would more aligned if "
                        f"input size {(input_h, input_w)} is `x+1` and "
                        f"out size {(output_h, output_w)} is `nx+1`"
                    )
    return F.interpolate(input, size, scale_factor, mode, align_corners)
//  [python]
def forward(self, x: torch.Tensor) -> Union[torch.Tensor, List[torch.Tensor]]:
        _, _, h, w = x.shape

        # collect encode outputs
        encode_outputs = []
        for i, m in enumerate(self.encode_modules):
            x = m(x)
            encode_outputs.append(x)
            if i != self.encode_num - 1:
                x = F.max_pool2d(x, kernel_size=2, stride=2, ceil_mode=True)

        # collect decode outputs
        x = encode_outputs.pop()
        decode_outputs = [x]
        for m in self.decode_modules:
            x2 = encode_outputs.pop()
            x = F.upsample(x, size=x2.shape[2:], mode='bilinear')
            x = m(torch.concat([x, x2], dim=1))
            decode_outputs.insert(0, x)

        # collect side outputs
        side_outputs = []
        for m in self.side_modules:
            x = decode_outputs.pop()
            x = F.upsample(m(x), size=[h, w], mode='bilinear')
            side_outputs.insert(0, x)

        x = self.out_conv(torch.concat(side_outputs, dim=1))

        if self.training:
            # do not use torch.sigmoid for amp safe
            return [x] + side_outputs
        else:
            return torch.sigmoid(x)