jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(
        self,
        x,
        mask,
        feats,
        img,
        ws,
        fname=None,
        force_fp32=False,
        fused_modconv=None,
        **layer_kwargs,
    ):
        dtype = torch.float16 if self.use_fp16 and not force_fp32 else torch.float32
        dtype = torch.float32
        memory_format = (
            torch.channels_last
            if self.channels_last and not force_fp32
            else torch.contiguous_format
        )
        if fused_modconv is None:
            fused_modconv = (not self.training) and (
                dtype == torch.float32 or int(x.shape[0]) == 1
            )

        x = x.to(dtype=dtype, memory_format=memory_format)
        x_skip = (
            feats[self.resolution].clone().to(dtype=dtype, memory_format=memory_format)
        )

        # Main layers.
        if self.in_channels == 0:
            x = self.conv1(x, ws[1], fused_modconv=fused_modconv, **layer_kwargs)
        elif self.architecture == "resnet":
            y = self.skip(x, gain=np.sqrt(0.5))
            x = self.conv0(
                x, ws[0].clone(), fused_modconv=fused_modconv, **layer_kwargs
            )
            if len(self.ffc_skip) > 0:
                mask = F.interpolate(
                    mask,
                    size=x_skip.shape[2:],
                )
                z = x + x_skip
                for fres in self.ffc_skip:
                    z = fres(z, mask)
                x = x + z
            else:
                x = x + x_skip
            x = self.conv1(
                x,
                ws[1].clone(),
                fused_modconv=fused_modconv,
                gain=np.sqrt(0.5),
                **layer_kwargs,
            )
            x = y.add_(x)
        else:
            x = self.conv0(
                x, ws[0].clone(), fused_modconv=fused_modconv, **layer_kwargs
            )
            if len(self.ffc_skip) > 0:
                mask = F.interpolate(
                    mask,
                    size=x_skip.shape[2:],
                )
                z = x + x_skip
                for fres in self.ffc_skip:
                    z = fres(z, mask)
                x = x + z
            else:
                x = x + x_skip
            x = self.conv1(
                x, ws[1].clone(), fused_modconv=fused_modconv, **layer_kwargs
            )
        # ToRGB.
        if img is not None:
            img = upsample2d(img, self.resample_filter)
        if self.is_last or self.architecture == "skip":
            y = self.torgb(x, ws[2].clone(), fused_modconv=fused_modconv)
            y = y.to(dtype=torch.float32, memory_format=torch.contiguous_format)
            img = img.add_(y) if img is not None else y

        x = x.to(dtype=dtype)
        assert x.dtype == dtype
        assert img is None or img.dtype == torch.float32
        return x, img
//  [python]
def forward(self, x, **kwargs):
        out = self.decoder(self.encoder(x), **kwargs)

        # Bilinearly upsample the output to match the input resolution
        up_out = t.interpolate(out, size=[x.size(2), x.size(3)], mode='bilinear', align_corners=False)
        
        # L2-normalize the first three channels / ensure positive value for concentration parameters (kappa)
        up_out = norm_normalize(up_out) 
        return up_out
//  [python]
def forward(self, x, **kwargs):
        out = self.decoder(self.encoder(x), **kwargs)

        # Bilinearly upsample the output to match the input resolution
        up_out = F.interpolate(out, size=[x.size(2), x.size(3)], mode='bilinear', align_corners=False)
        
        # L2-normalize the first three channels / ensure positive value for concentration parameters (kappa)
        up_out = norm_normalize(up_out) 
        return up_out
//  [python]
def resize_pos_embed(posemb, posemb_new):
    """
    Rescale the grid of position embeddings when loading from state_dict
    Expected shape of position embeddings is (1, T, N, C), and considers only square images
    """
    _logger.info('Resized position embedding: %s to %s', posemb.shape, posemb_new.shape)
    seq_length_old = posemb.shape[2]
    num_blocks_new, seq_length_new = posemb_new.shape[1:3]
    size_new = int(math.sqrt(num_blocks_new*seq_length_new))
    # First change to (1, C, H, W)
    posemb = deblockify(posemb, int(math.sqrt(seq_length_old))).permute(0, 3, 1, 2)
    posemb = t.nn.functional.interpolate(posemb, size=[size_new, size_new], mode='bicubic', align_corners=False)
    # Now change to new (1, T, N, C)
    posemb = blockify(posemb.permute(0, 2, 3, 1), int(math.sqrt(seq_length_new)))
    return posemb
//  [python]
def resize_pos_embed(posemb, posemb_new):
    """
    Rescale the grid of position embeddings when loading from state_dict
    Expected shape of position embeddings is (1, T, N, C), and considers only square images
    """
    _logger.info('Resized position embedding: %s to %s', posemb.shape, posemb_new.shape)
    seq_length_old = posemb.shape[2]
    num_blocks_new, seq_length_new = posemb_new.shape[1:3]
    size_new = int(math.sqrt(num_blocks_new*seq_length_new))
    # First change to (1, C, H, W)
    posemb = deblockify(posemb, int(math.sqrt(seq_length_old))).permute(0, 3, 1, 2)
    posemb = F.interpolate(posemb, size=[size_new, size_new], mode='bicubic', align_corners=False)
    # Now change to new (1, T, N, C)
    posemb = blockify(posemb.permute(0, 2, 3, 1), int(math.sqrt(seq_length_new)))
    return posemb
//  [python]
def forward(self, x):

        # conv1
        x = self.conv1_1(x)
        x = self.conv1_2(x)
        #print(x.shape)
        #x = nn.MaxPool2d(x)


        # conv2
        x = self.conv2_1(x)
        x = self.conv2_2(x)
        #print(x.shape)
        #x = nn.MaxPool2d(x)


        # conv3
        x = self.conv3_1(x)
        x = self.conv3_2(x)
        x = self.conv3_3(x)
        #print(x.shape)
        #x = nn.MaxPool2d(x)

        # conv4
        x = self.conv4_1(x)
        x = self.conv4_2(x)
        x = self.conv4_3(x)
        #print(x.shape)
        #x = nn.MaxPool2d(x)

        # conv5
        x = self.conv5_1(x)
        x = self.conv5_2(x)
        x = self.conv5_3(x)
        #print(x.shape)
        #x = nn.MaxPool2d(x)

        # conv6
        x = self.conv6_1(x)
        #print(x.shape)

        # conv7
        x = nn.Upsample(scale_factor=2, mode='nearest')(x)
        x = self.conv7_1(x)
        #print(x.shape)
        x = F.dropout2d(x)

        # conv8
        x = nn.Upsample(scale_factor=2, mode='nearest')(x)
        x = self.conv8_1(x)
        #print(x.shape)

        # conv9
        x = nn.Upsample(scale_factor=2, mode='nearest')(x)
        x = self.conv9_1(x)
        #print(x.shape)

        # conv10
        x = nn.Upsample(scale_factor=2, mode='nearest')(x)
        x = self.conv10_1(x)
        #print(x.shape)

        # conv11
        x = nn.Upsample(scale_factor=2, mode='nearest')(x)
        x = self.conv11_1(x)
        #print(x.shape)

        # conv12
        #x = nn.Upsample(scale_factor=2, mode='nearest')(x)
        x = self.conv12_1(x)
        #print(x.shape)

        # h_out
        x_out = self.h_out(x)
        #print(x_out.shape)

        return x_out
//  [python]
def forward(self, x):

        # conv1
        x = self.conv1_1(x)
        x = self.conv1_2(x)
        #print(x.shape)
        #x = nn.MaxPool2d(x)


        # conv2
        x = self.conv2_1(x)
        x = self.conv2_2(x)
        #print(x.shape)
        #x = nn.MaxPool2d(x)


        # conv3
        x = self.conv3_1(x)
        x = self.conv3_2(x)
        x = self.conv3_3(x)
        #print(x.shape)
        #x = nn.MaxPool2d(x)

        # conv4
        x = self.conv4_1(x)
        x = self.conv4_2(x)
        x = self.conv4_3(x)
        #print(x.shape)
        #x = nn.MaxPool2d(x)

        # conv5
        x = self.conv5_1(x)
        x = self.conv5_2(x)
        x = self.conv5_3(x)
        #print(x.shape)
        #x = nn.MaxPool2d(x)

        # conv6
        x = self.conv6_1(x)
        #print(x.shape)

        # conv7
        x = F.interpolate(x, scale_factor=2, mode='nearest')
        x = self.conv7_1(x)
        #print(x.shape)
        x = F.dropout2d(x)

        # conv8
        x = F.interpolate(x, scale_factor=2, mode='nearest')
        x = self.conv8_1(x)
        #print(x.shape)

        # conv9
        x = F.interpolate(x, scale_factor=2, mode='nearest')
        x = self.conv9_1(x)
        #print(x.shape)

        # conv10
        x = F.interpolate(x, scale_factor=2, mode='nearest')
        x = self.conv10_1(x)
        #print(x.shape)

        # conv11
        x = F.interpolate(x, scale_factor=2, mode='nearest')
        x = self.conv11_1(x)
        #print(x.shape)

        # conv12
        #x = F.interpolate(x, scale_factor=2, mode='nearest')
        x = self.conv12_1(x)
        #print(x.shape)

        # h_out
        x_out = self.h_out(x)
        #print(x_out.shape)

        return x_out
//  [python]
def forward(self, x: Tuple[Tensor]) -> tuple:
        """Forward features from the upstream network.

        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.

        Returns:
            tuple: A tuple of classification scores and mask prediction.

                - mlvl_mask_preds (list[Tensor]): Multi-level mask prediction.
                  Each element in the list has shape
                  (batch_size, num_grids**2 ,h ,w).
                - mlvl_cls_preds (list[Tensor]): Multi-level scores.
                  Each element in the list has shape
                  (batch_size, num_classes, num_grids ,num_grids).
        """
        assert len(x) == self.num_levels
        feats = self.resize_feats(x)
        mlvl_mask_preds = []
        mlvl_cls_preds = []
        for i in range(self.num_levels):
            x = feats[i]
            mask_feat = x
            cls_feat = x
            # generate and concat the coordinate
            coord_feat = generate_coordinate(mask_feat.size(),
                                             mask_feat.device)
            mask_feat = torch.cat([mask_feat, coord_feat], 1)

            for mask_layer in (self.mask_convs):
                mask_feat = mask_layer(mask_feat)

            mask_feat = F.upsample(
                mask_feat, scale_factor=2, mode='bilinear')
            mask_preds = self.conv_mask_list[i](mask_feat)

            # cls branch
            for j, cls_layer in enumerate(self.cls_convs):
                if j == self.cls_down_index:
                    num_grid = self.num_grids[i]
                    cls_feat = F.upsample(
                        cls_feat, size=num_grid, mode='bilinear')
                cls_feat = cls_layer(cls_feat)

            cls_pred = self.conv_cls(cls_feat)

            if not self.training:
                feat_wh = feats[0].size()[-2:]
                upsampled_size = (feat_wh[0] * 2, feat_wh[1] * 2)
                mask_preds = F.upsample(
                    mask_preds.sigmoid(), size=upsampled_size, mode='bilinear')
                cls_pred = cls_pred.sigmoid()
                # get local maximum
                local_max = F.max_pool2d(cls_pred, 2, stride=1, padding=1)
                keep_mask = local_max[:, :, :-1, :-1] == cls_pred
                cls_pred = cls_pred * keep_mask

            mlvl_mask_preds.append(mask_preds)
            mlvl_cls_preds.append(cls_pred)
        return mlvl_mask_preds, mlvl_cls_preds
//  [python]
def forward(self, x: Tuple[Tensor]) -> tuple:
        """Forward features from the upstream network.

        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.

        Returns:
            tuple: A tuple of classification scores and mask prediction.

                - mlvl_mask_preds (list[Tensor]): Multi-level mask prediction.
                  Each element in the list has shape
                  (batch_size, num_grids**2 ,h ,w).
                - mlvl_cls_preds (list[Tensor]): Multi-level scores.
                  Each element in the list has shape
                  (batch_size, num_classes, num_grids ,num_grids).
        """
        assert len(x) == self.num_levels
        feats = self.resize_feats(x)
        mlvl_mask_preds = []
        mlvl_cls_preds = []
        for i in range(self.num_levels):
            x = feats[i]
            mask_feat = x
            cls_feat = x
            # generate and concat the coordinate
            coord_feat = generate_coordinate(mask_feat.size(),
                                             mask_feat.device)
            mask_feat = torch.cat([mask_feat, coord_feat], 1)

            for mask_layer in (self.mask_convs):
                mask_feat = mask_layer(mask_feat)

            mask_feat = F.interpolate(
                mask_feat, scale_factor=2, mode='bilinear')
            mask_preds = self.conv_mask_list[i](mask_feat)

            # cls branch
            for j, cls_layer in enumerate(self.cls_convs):
                if j == self.cls_down_index:
                    num_grid = self.num_grids[i]
                    cls_feat = F.interpolate(
                        cls_feat, size=num_grid, mode='bilinear')
                cls_feat = cls_layer(cls_feat)

            cls_pred = self.conv_cls(cls_feat)

            if not self.training:
                feat_wh = feats[0].size()[-2:]
                upsampled_size = (feat_wh[0] * 2, feat_wh[1] * 2)
                mask_preds = F.interpolate(
                    mask_preds.sigmoid(), size=upsampled_size, mode='bilinear')
                cls_pred = cls_pred.sigmoid()
                # get local maximum
                local_max = F.max_pool2d(cls_pred, 2, stride=1, padding=1)
                keep_mask = local_max[:, :, :-1, :-1] == cls_pred
                cls_pred = cls_pred * keep_mask

            mlvl_mask_preds.append(mask_preds)
            mlvl_cls_preds.append(cls_pred)
        return mlvl_mask_preds, mlvl_cls_preds
//  [python]
def forward(self, x, gst=None):
        x = self.token_emb(x)
        feats = [x]
        if gst is not None:
            gst = self.gst_emb(gst)
            gst = rearrange(gst, 'b c -> b c 1')
            gst = F.upsample(gst, size=x.shape[1])
            gst = rearrange(gst, 'b c t -> b t c')
            feats.append(gst)
        x = torch.cat(feats, dim=-1)
        for i, conv in enumerate(self.conv_layer):
            if i != 0:
                x = conv(x) + x
            else:
                x = conv(x)
        x = self.proj(x)
        x = x.squeeze(-1)
        if self.f0_pred == 'mean':
            x = torch.sigmoid(x)
        elif self.f0_pred == 'argmax':
            x = torch.softmax(x, dim=-1)
        else:
            raise NotImplementedError
        return x
//  [python]
def forward(self, x, gst=None):
        x = self.token_emb(x)
        feats = [x]
        if gst is not None:
            gst = self.gst_emb(gst)
            gst = rearrange(gst, 'b c -> b c 1')
            gst = F.interpolate(gst, x.shape[1])
            gst = rearrange(gst, 'b c t -> b t c')
            feats.append(gst)
        x = torch.cat(feats, dim=-1)
        for i, conv in enumerate(self.conv_layer):
            if i != 0:
                x = conv(x) + x
            else:
                x = conv(x)
        x = self.proj(x)
        x = x.squeeze(-1)
        if self.f0_pred == 'mean':
            x = torch.sigmoid(x)
        elif self.f0_pred == 'argmax':
            x = torch.softmax(x, dim=-1)
        else:
            raise NotImplementedError
        return x
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor): The tensor to upsample.

        Returns:
            torch.Tensor: The upsample results.
        """
        assert x.shape[1] == self.channels
        if self.dims == 3:
            x = t.upsample(
                x, size=(x.shape[2], x.shape[3] * 2, x.shape[4] * 2),
                mode='nearest')
        else:
            x = t.upsample(x, scale_factor=2, mode='nearest')
        if self.use_conv:
            x = self.conv(x)
        return x
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (torch.Tensor): The tensor to upsample.

        Returns:
            torch.Tensor: The upsample results.
        """
        assert x.shape[1] == self.channels
        if self.dims == 3:
            x = F.interpolate(
                x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2),
                mode='nearest')
        else:
            x = F.interpolate(x, scale_factor=2, mode='nearest')
        if self.use_conv:
            x = self.conv(x)
        return x
//  [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))+v_max
        c3 = self.relu(self.conv3(v_range))+v_range
        c3 = self.conv3_(c3)+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))+v_max
        c3 = self.relu(self.conv3(v_range))+v_range
        c3 = self.conv3_(c3)+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 forward(self, x: Tuple[Tensor]) -> Tuple:
        """Forward features from the upstream network.

        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.

        Returns:
            tuple: A tuple of classification scores and mask prediction.

                - mlvl_mask_preds_x (list[Tensor]): Multi-level mask prediction
                  from x branch. Each element in the list has shape
                  (batch_size, num_grids ,h ,w).
                - mlvl_mask_preds_y (list[Tensor]): Multi-level mask prediction
                  from y branch. Each element in the list has shape
                  (batch_size, num_grids ,h ,w).
                - mlvl_cls_preds (list[Tensor]): Multi-level scores.
                  Each element in the list has shape
                  (batch_size, num_classes, num_grids ,num_grids).
        """
        assert len(x) == self.num_levels
        feats = self.resize_feats(x)
        mask_preds_x = []
        mask_preds_y = []
        cls_preds = []
        for i in range(self.num_levels):
            x = feats[i]
            mask_feat = x
            cls_feat = x
            # generate and concat the coordinate
            coord_feat = generate_coordinate(mask_feat.size(),
                                             mask_feat.device)
            mask_feat = torch.cat([mask_feat, coord_feat], 1)

            for mask_layer in self.mask_convs:
                mask_feat = mask_layer(mask_feat)

            mask_feat = F.upsample(
                mask_feat, scale_factor=2, mode='bilinear')

            mask_pred_x = self.conv_mask_list_x[i](mask_feat)
            mask_pred_y = self.conv_mask_list_y[i](mask_feat)

            # cls branch
            for j, cls_layer in enumerate(self.cls_convs):
                if j == self.cls_down_index:
                    num_grid = self.num_grids[i]
                    cls_feat = F.upsample(
                        cls_feat, size=num_grid, mode='bilinear')
                cls_feat = cls_layer(cls_feat)

            cls_pred = self.conv_cls(cls_feat)

            if not self.training:
                feat_wh = feats[0].size()[-2:]
                upsampled_size = (feat_wh[0] * 2, feat_wh[1] * 2)
                mask_pred_x = F.upsample(
                    mask_pred_x.sigmoid(),
                    size=upsampled_size,
                    mode='bilinear')
                mask_pred_y = F.upsample(
                    mask_pred_y.sigmoid(),
                    size=upsampled_size,
                    mode='bilinear')
                cls_pred = cls_pred.sigmoid()
                # get local maximum
                local_max = F.max_pool2d(cls_pred, 2, stride=1, padding=1)
                keep_mask = local_max[:, :, :-1, :-1] == cls_pred
                cls_pred = cls_pred * keep_mask

            mask_preds_x.append(mask_pred_x)
            mask_preds_y.append(mask_pred_y)
            cls_preds.append(cls_pred)
        return mask_preds_x, mask_preds_y, cls_preds
//  [python]
def forward(self, x: Tuple[Tensor]) -> Tuple:
        """Forward features from the upstream network.

        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.

        Returns:
            tuple: A tuple of classification scores and mask prediction.

                - mlvl_mask_preds_x (list[Tensor]): Multi-level mask prediction
                  from x branch. Each element in the list has shape
                  (batch_size, num_grids ,h ,w).
                - mlvl_mask_preds_y (list[Tensor]): Multi-level mask prediction
                  from y branch. Each element in the list has shape
                  (batch_size, num_grids ,h ,w).
                - mlvl_cls_preds (list[Tensor]): Multi-level scores.
                  Each element in the list has shape
                  (batch_size, num_classes, num_grids ,num_grids).
        """
        assert len(x) == self.num_levels
        feats = self.resize_feats(x)
        mask_preds_x = []
        mask_preds_y = []
        cls_preds = []
        for i in range(self.num_levels):
            x = feats[i]
            mask_feat = x
            cls_feat = x
            # generate and concat the coordinate
            coord_feat = generate_coordinate(mask_feat.size(),
                                             mask_feat.device)
            mask_feat = torch.cat([mask_feat, coord_feat], 1)

            for mask_layer in self.mask_convs:
                mask_feat = mask_layer(mask_feat)

            mask_feat = F.interpolate(
                mask_feat, scale_factor=2, mode='bilinear')

            mask_pred_x = self.conv_mask_list_x[i](mask_feat)
            mask_pred_y = self.conv_mask_list_y[i](mask_feat)

            # cls branch
            for j, cls_layer in enumerate(self.cls_convs):
                if j == self.cls_down_index:
                    num_grid = self.num_grids[i]
                    cls_feat = F.interpolate(
                        cls_feat, size=num_grid, mode='bilinear')
                cls_feat = cls_layer(cls_feat)

            cls_pred = self.conv_cls(cls_feat)

            if not self.training:
                feat_wh = feats[0].size()[-2:]
                upsampled_size = (feat_wh[0] * 2, feat_wh[1] * 2)
                mask_pred_x = F.interpolate(
                    mask_pred_x.sigmoid(),
                    size=upsampled_size,
                    mode='bilinear')
                mask_pred_y = F.interpolate(
                    mask_pred_y.sigmoid(),
                    size=upsampled_size,
                    mode='bilinear')
                cls_pred = cls_pred.sigmoid()
                # get local maximum
                local_max = F.max_pool2d(cls_pred, 2, stride=1, padding=1)
                keep_mask = local_max[:, :, :-1, :-1] == cls_pred
                cls_pred = cls_pred * keep_mask

            mask_preds_x.append(mask_pred_x)
            mask_preds_y.append(mask_pred_y)
            cls_preds.append(cls_pred)
        return mask_preds_x, mask_preds_y, cls_preds
//  [python]
def scale_img_nhwc(x  : torch.Tensor, size, mag='bilinear', min='area') -> torch.Tensor:
    # assert (x.shape[1] >= size[0] and x.shape[2] >= size[1]) or (x.shape[1] < size[0] and x.shape[2] < size[1]), "Trying to magnify image in one dimension and minify in the other"
    assert (x.shape[1] >= size[0] and x.shape[2] >= size[1]) or (x.shape[1] <= size[0] and x.shape[2] <= size[1]), "Trying to magnify image in one dimension and minify in the other"
    y = x.permute(0, 3, 1, 2) # NHWC -> NCHW
    if x.shape[1] > size[0] and x.shape[2] > size[1]: # Minification, previous size was bigger
        y = t.nn.functional.interpolate(y, size, mode=min)
    else: # Magnification
        if mag == 'bilinear' or mag == 'bicubic':
            y = t.nn.functional.interpolate(y, size, mode=mag, align_corners=True)
        else:
            y = t.nn.functional.interpolate(y, size, mode=mag)
    return y.permute(0, 2, 3, 1).contiguous()
//  [python]
def scale_img_nhwc(x  : torch.Tensor, size, mag='bilinear', min='area') -> torch.Tensor:
    # assert (x.shape[1] >= size[0] and x.shape[2] >= size[1]) or (x.shape[1] < size[0] and x.shape[2] < size[1]), "Trying to magnify image in one dimension and minify in the other"
    assert (x.shape[1] >= size[0] and x.shape[2] >= size[1]) or (x.shape[1] <= size[0] and x.shape[2] <= size[1]), "Trying to magnify image in one dimension and minify in the other"
    y = x.permute(0, 3, 1, 2) # NHWC -> NCHW
    if x.shape[1] > size[0] and x.shape[2] > size[1]: # Minification, previous size was bigger
        y = torch.nn.functional.interpolate(y, size, mode=min)
    else: # Magnification
        if mag == 'bilinear' or mag == 'bicubic':
            y = torch.nn.functional.interpolate(y, size, mode=mag, align_corners=True)
        else:
            y = torch.nn.functional.interpolate(y, size, mode=mag)
    return y.permute(0, 2, 3, 1).contiguous()
//  [python]
def resize(input,
           size=None,
           scale_factor=None,
           mode='nearest',
           align_corners=None,
           warning=True):
    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`')
    if isinstance(size, torch.Size):
        size = tuple(int(x) for x in size)
    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=True):
    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`')
    if isinstance(size, torch.Size):
        size = tuple(int(x) for x in size)
    return F.interpolate(input, size, scale_factor, mode, align_corners)
//  [python]
def interpolate_pos_encoding(self, x, w, h):
        npatch = x.shape[1] - 1
        N = self.pos_embed.shape[1] - 1
        if npatch == N and w == h:
            return self.pos_embed
        class_pos_embed = self.pos_embed[:, 0]
        patch_pos_embed = self.pos_embed[:, 1:]
        dim = x.shape[-1]
        w0 = w // self.patch_embed.patch_size[0]
        h0 = h // self.patch_embed.patch_size[0]
        # we add a small number to avoid floating point error in the interpolation
        # see discussion at https://github.com/facebookresearch/dino/issues/8
        w0, h0 = w0 + 0.1, h0 + 0.1
        patch_pos_embed = nn.functional.interpolate(
            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
            size=(int(w0), int(h0)),
            mode='bicubic',
        )
        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
//  [python]
def interpolate_pos_encoding(self, x, w, h):
        npatch = x.shape[1] - 1
        N = self.pos_embed.shape[1] - 1
        if npatch == N and w == h:
            return self.pos_embed
        class_pos_embed = self.pos_embed[:, 0]
        patch_pos_embed = self.pos_embed[:, 1:]
        dim = x.shape[-1]
        w0 = w // self.patch_embed.patch_size[0]
        h0 = h // self.patch_embed.patch_size[0]
        # we add a small number to avoid floating point error in the interpolation
        # see discussion at https://github.com/facebookresearch/dino/issues/8
        w0, h0 = w0 + 0.1, h0 + 0.1
        patch_pos_embed = nn.functional.interpolate(
            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
            mode='bicubic',
        )
        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
//  [python]
def _upsample_mels(self, cond, audio_size, force_interpolate=False):
        if hasattr(self, 'upsample_net') and self.upsample_net is not None:
            cond = self.upsample_net(cond)
        if force_interpolate or (self.interpolation_required and not (cond.shape[2] == audio_size[3])):
            cond = t.interpolate(cond, size=audio_size[3], mode=self.upsample_mode, align_corners=True if self.upsample_mode == 'linear' else None)
        else:
            pad = (cond.shape[2] - audio_size[3])//2
            remainder = pad%2
            cond = cond[:, :, pad:-(pad+remainder)]
        return cond
//  [python]
def _upsample_mels(self, cond, audio_size, force_interpolate=False):
        if hasattr(self, 'upsample_net') and self.upsample_net is not None:
            cond = self.upsample_net(cond)
        if force_interpolate or (self.interpolation_required and not (cond.shape[2] == audio_size[3])):
            cond = F.interpolate(cond, size=audio_size[3], mode=self.upsample_mode, align_corners=True if self.upsample_mode == 'linear' else None)
        else:
            pad = (cond.shape[2] - audio_size[3])//2
            remainder = pad%2
            cond = cond[:, :, pad:-(pad+remainder)]
        return cond
//  [python]
def smooth_border_mg(diffuse_mask, mg):
    mg = mg - 0.5
    diffuse_mask = t.nn.functional.interpolate(
        diffuse_mask, mg.shape[:2], mode='bilinear')[0].permute(1, 2, 0)
    mg = mg * diffuse_mask
    mg = mg + 0.5
    return mg
//  [python]
def smooth_border_mg(diffuse_mask, mg):
    mg = mg - 0.5
    diffuse_mask = F.interpolate(
        diffuse_mask, mg.shape[:2], mode='bilinear')[0].permute(1, 2, 0)
    mg = mg * diffuse_mask
    mg = mg + 0.5
    return mg
//  [python]
def mask_correlation_map(self, correlation_map, mask):
        """Add mask weight for correlation map.

        Add a negative infinity number to the masked regions so that softmax
        function will result in 'zero' in those regions.

        Args:
            correlation_map (torch.Tensor): Correlation map with shape of
                (n, h_unfold*w_unfold, h_map, w_map).
            mask (torch.Tensor): Mask tensor with shape of (n, c, h, w). '1'
                in the mask indicates masked region while '0' indicates valid
                region.

        Returns:
            torch.Tensor: Updated correlation map with mask.
        """
        if mask is not None:
            mask = t.nn.functional.interpolate(mask, scale_factor=self.scale)
            # if any pixel is masked in patch, the patch is considered to be
            # masked
            mask_cols = self.im2col(
                mask,
                kernel_size=self.unfold_corr_kernel_size,
                stride=self.unfold_corr_stride,
                padding=self.unfold_corr_padding,
                dilation=self.unfold_corr_dilation)
            mask_cols = (mask_cols.sum(dim=1, keepdim=True) > 0).float()
            mask_cols = mask_cols.permute(0, 2,
                                          1).reshape(mask.size(0), -1, 1, 1)
            # add negative inf will bring zero in softmax
            mask_cols[mask_cols == 1] = -float('inf')
            correlation_map += mask_cols
        return correlation_map
//  [python]
def mask_correlation_map(self, correlation_map, mask):
        """Add mask weight for correlation map.

        Add a negative infinity number to the masked regions so that softmax
        function will result in 'zero' in those regions.

        Args:
            correlation_map (torch.Tensor): Correlation map with shape of
                (n, h_unfold*w_unfold, h_map, w_map).
            mask (torch.Tensor): Mask tensor with shape of (n, c, h, w). '1'
                in the mask indicates masked region while '0' indicates valid
                region.

        Returns:
            torch.Tensor: Updated correlation map with mask.
        """
        if mask is not None:
            mask = F.interpolate(mask, scale_factor=self.scale)
            # if any pixel is masked in patch, the patch is considered to be
            # masked
            mask_cols = self.im2col(
                mask,
                kernel_size=self.unfold_corr_kernel_size,
                stride=self.unfold_corr_stride,
                padding=self.unfold_corr_padding,
                dilation=self.unfold_corr_dilation)
            mask_cols = (mask_cols.sum(dim=1, keepdim=True) > 0).float()
            mask_cols = mask_cols.permute(0, 2,
                                          1).reshape(mask.size(0), -1, 1, 1)
            # add negative inf will bring zero in softmax
            mask_cols[mask_cols == 1] = -float('inf')
            correlation_map += mask_cols
        return correlation_map
//  [python]
def forward(self, x, flow=None, scale=1):
        x = t.interpolate(
            x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False
        )
        if flow is not None:
            flow = (
                t.interpolate(
                    flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False
                )
                * 1.0
                / scale
            )
            x = torch.cat((x, flow), 1)
        feat = self.conv0(x)
        feat = self.convblock(feat)
        tmp = self.lastconv(feat)
        tmp = t.interpolate(
            tmp, scale_factor=scale, mode="bilinear", align_corners=False
        )
        flow = tmp[:, :4] * scale
        mask = tmp[:, 4:5]
        return flow, mask
//  [python]
def forward(self, x, flow=None, scale=1):
        x = F.interpolate(
            x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False
        )
        if flow is not None:
            flow = (
                F.interpolate(
                    flow, scale_factor=1.0 / scale, mode="bilinear", align_corners=False
                )
                * 1.0
                / scale
            )
            x = torch.cat((x, flow), 1)
        feat = self.conv0(x)
        feat = self.convblock(feat)
        tmp = self.lastconv(feat)
        tmp = F.interpolate(
            tmp, scale_factor=scale, mode="bilinear", align_corners=False
        )
        flow = tmp[:, :4] * scale
        mask = tmp[:, 4:5]
        return flow, mask
//  [python]
def whiten_img(image, skin_mask, whitening_degree, flag_bigKernal=False):
    """
    image: rgb
    """
    dilate_kernalsize = 30
    if flag_bigKernal:
        dilate_kernalsize = 80
    new_kernel1 = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE, (dilate_kernalsize, dilate_kernalsize))
    new_kernel2 = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE, (dilate_kernalsize, dilate_kernalsize))
    if len(skin_mask.shape) == 3:
        skin_mask = skin_mask[:, :, -1]
    skin_mask = cv2.dilate(skin_mask, new_kernel1, 1)
    skin_mask = cv2.erode(skin_mask, new_kernel2, 1)
    skin_mask = cv2.blur(skin_mask, (20, 20)) / 255.0
    skin_mask = skin_mask.squeeze()
    skin_mask = torch.from_numpy(skin_mask).to(image.device)
    skin_mask = torch.stack([skin_mask, skin_mask, skin_mask], dim=0)[None,
                                                                      ...]
    skin_mask[:, 1:, :, :] *= 0.75

    whiten_mg = skin_mask * 0.2 * whitening_degree + 0.5
    assert len(whiten_mg.shape) == 4
    whiten_mg = t.interpolate(
        whiten_mg, image.shape[:2], mode='bilinear')[0].permute(1, 2,
                                                                0).half()
    output_pred = image.half()
    output_pred = output_pred / 255.0
    output_pred = (
        -2 * whiten_mg + 1
    ) * output_pred * output_pred + 2 * whiten_mg * output_pred  # value: 0~1
    output_pred = output_pred * 255.0
    output_pred = output_pred.byte()

    output_pred = output_pred.cpu().numpy()
    return output_pred
//  [python]
def whiten_img(image, skin_mask, whitening_degree, flag_bigKernal=False):
    """
    image: rgb
    """
    dilate_kernalsize = 30
    if flag_bigKernal:
        dilate_kernalsize = 80
    new_kernel1 = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE, (dilate_kernalsize, dilate_kernalsize))
    new_kernel2 = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE, (dilate_kernalsize, dilate_kernalsize))
    if len(skin_mask.shape) == 3:
        skin_mask = skin_mask[:, :, -1]
    skin_mask = cv2.dilate(skin_mask, new_kernel1, 1)
    skin_mask = cv2.erode(skin_mask, new_kernel2, 1)
    skin_mask = cv2.blur(skin_mask, (20, 20)) / 255.0
    skin_mask = skin_mask.squeeze()
    skin_mask = torch.from_numpy(skin_mask).to(image.device)
    skin_mask = torch.stack([skin_mask, skin_mask, skin_mask], dim=0)[None,
                                                                      ...]
    skin_mask[:, 1:, :, :] *= 0.75

    whiten_mg = skin_mask * 0.2 * whitening_degree + 0.5
    assert len(whiten_mg.shape) == 4
    whiten_mg = F.interpolate(
        whiten_mg, image.shape[:2], mode='bilinear')[0].permute(1, 2,
                                                                0).half()
    output_pred = image.half()
    output_pred = output_pred / 255.0
    output_pred = (
        -2 * whiten_mg + 1
    ) * output_pred * output_pred + 2 * whiten_mg * output_pred  # value: 0~1
    output_pred = output_pred * 255.0
    output_pred = output_pred.byte()

    output_pred = output_pred.cpu().numpy()
    return output_pred
//  [python]
def forward(self, feature, shape):
        x0 = feature[0]
        x1 = feature[1]
        x2 = feature[2]
        x3 = feature[3]
        x4 = feature[4]
        x = self.up1(x4, x3)
        dp3_out_seg = self.out_conv_dp3(x)
        dp3_out_seg = torch.nn.functional.interpolate(dp3_out_seg, shape)

        x = self.up2(x, x2)
        dp2_out_seg = self.out_conv_dp2(x)
        dp2_out_seg = torch.nn.functional.interpolate(dp2_out_seg, shape)

        x = self.up3(x, x1)
        dp1_out_seg = self.out_conv_dp1(x)
        dp1_out_seg = torch.nn.functional.interpolate(dp1_out_seg, shape)

        x = self.up4(x, x0)
        dp0_out_seg = self.out_conv(x)
        return dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg
//  [python]
def forward(self, x, context, mask=None):
        """Forward Function.

        Args:
            x (torch.Tensor): Tensor with shape (n, c, h, w).
            context (torch.Tensor): Tensor with shape (n, c, h, w).
            mask (torch.Tensor): Tensor with shape (n, 1, h, w). Default: None.

        Returns:
            tuple(torch.Tensor): Features after contextural attention.
        """
        # raw features to be used in copy (deconv)
        raw_context = context
        raw_context_cols = self.im2col(
            raw_context,
            kernel_size=self.unfold_raw_kernel_size,
            stride=self.unfold_raw_stride,
            padding=self.unfold_raw_padding,
            normalize=False,
            return_cols=True)
        # resize the feature to reduce computational cost
        x = t.nn.functional.interpolate(x, scale_factor=self.scale, mode='bilinear', align_corners=False)
        context = t.nn.functional.interpolate(context, scale_factor=self.scale, mode='bilinear', align_corners=False)

        context_cols = self.im2col(
            context,
            kernel_size=self.unfold_corr_kernel_size,
            stride=self.unfold_corr_stride,
            padding=self.unfold_corr_padding,
            dilation=self.unfold_corr_dilation,
            normalize=True,
            return_cols=True)
        h_unfold, w_unfold = self.calculate_unfold_hw(
            context.size()[-2:],
            kernel_size=self.unfold_corr_kernel_size,
            stride=self.unfold_corr_stride,
            padding=self.unfold_corr_padding,
            dilation=self.unfold_corr_dilation,
        )
        # reshape context_cols to
        # (n*h_unfold*w_unfold, c, unfold_mks, unfold_mks)
        # 'mks' is short for 'mask_kernel_size'
        context_cols = context_cols.reshape(-1, *context_cols.shape[2:])

        # the shape of correlation map should be:
        # (n, h_unfold*w_unfold, h', w')
        correlation_map = self.patch_correlation(x, context_cols)
        # fuse correlation map to enlarge consistent attention region.
        if self.with_fuse_correlation:
            correlation_map = self.fuse_correlation_map(
                correlation_map, h_unfold, w_unfold)

        correlation_map = self.mask_correlation_map(correlation_map, mask=mask)

        attention_score = self.softmax(correlation_map * self.softmax_scale)

        raw_context_filter = raw_context_cols.reshape(
            -1, *raw_context_cols.shape[2:])
        output = self.patch_copy_deconv(attention_score, raw_context_filter)
        # deconv will cause overlap and we need to remove the effects of that
        overlap_factor = self.calculate_overlap_factor(attention_score)
        output /= overlap_factor

        if self.return_attention_score:
            n, _, h_s, w_s = attention_score.size()
            attention_score = attention_score.view(n, h_unfold, w_unfold, h_s,
                                                   w_s)
            return output, attention_score

        return output
//  [python]
def forward(self, x, context, mask=None):
        """Forward Function.

        Args:
            x (torch.Tensor): Tensor with shape (n, c, h, w).
            context (torch.Tensor): Tensor with shape (n, c, h, w).
            mask (torch.Tensor): Tensor with shape (n, 1, h, w). Default: None.

        Returns:
            tuple(torch.Tensor): Features after contextural attention.
        """
        # raw features to be used in copy (deconv)
        raw_context = context
        raw_context_cols = self.im2col(
            raw_context,
            kernel_size=self.unfold_raw_kernel_size,
            stride=self.unfold_raw_stride,
            padding=self.unfold_raw_padding,
            normalize=False,
            return_cols=True)
        # resize the feature to reduce computational cost
        x = F.interpolate(x, scale_factor=self.scale)
        context = F.interpolate(context, scale_factor=self.scale)

        context_cols = self.im2col(
            context,
            kernel_size=self.unfold_corr_kernel_size,
            stride=self.unfold_corr_stride,
            padding=self.unfold_corr_padding,
            dilation=self.unfold_corr_dilation,
            normalize=True,
            return_cols=True)
        h_unfold, w_unfold = self.calculate_unfold_hw(
            context.size()[-2:],
            kernel_size=self.unfold_corr_kernel_size,
            stride=self.unfold_corr_stride,
            padding=self.unfold_corr_padding,
            dilation=self.unfold_corr_dilation,
        )
        # reshape context_cols to
        # (n*h_unfold*w_unfold, c, unfold_mks, unfold_mks)
        # 'mks' is short for 'mask_kernel_size'
        context_cols = context_cols.reshape(-1, *context_cols.shape[2:])

        # the shape of correlation map should be:
        # (n, h_unfold*w_unfold, h', w')
        correlation_map = self.patch_correlation(x, context_cols)
        # fuse correlation map to enlarge consistent attention region.
        if self.with_fuse_correlation:
            correlation_map = self.fuse_correlation_map(
                correlation_map, h_unfold, w_unfold)

        correlation_map = self.mask_correlation_map(correlation_map, mask=mask)

        attention_score = self.softmax(correlation_map * self.softmax_scale)

        raw_context_filter = raw_context_cols.reshape(
            -1, *raw_context_cols.shape[2:])
        output = self.patch_copy_deconv(attention_score, raw_context_filter)
        # deconv will cause overlap and we need to remove the effects of that
        overlap_factor = self.calculate_overlap_factor(attention_score)
        output /= overlap_factor

        if self.return_attention_score:
            n, _, h_s, w_s = attention_score.size()
            attention_score = attention_score.view(n, h_unfold, w_unfold, h_s,
                                                   w_s)
            return output, attention_score

        return output
//  [python]
def test_builder_to_backend_stress(
        self,
        compute_unit,
        backend,
        input_shape,
        scale_factor,
        align_corners,
        recompute_scale_factor,
    ):
        scale_factor_height, scale_factor_width = scale_factor
        _, _, height, width = input_shape
        height = height * scale_factor_height
        width = width * scale_factor_width
        is_h_float = height - np.floor(height) > 0.001
        is_w_float = width - np.floor(width) > 0.001

        # Currently, MIL is not supporting recompute_scale_factor=False + align_corners=False
        # with fractional output size
        if not recompute_scale_factor and not align_corners and (is_h_float or is_w_float):
            pytest.xfail("rdar://81124053 (Support recompute_scale_factor)")

        def _get_torch_upsample_prediction(
            x, scale_factor=(2, 2), align_corners=False, recompute_scale_factor=True
        ):
            x = torch.from_numpy(x)
            out = t.upsample_bilinear(
                x,
                scale_factor=scale_factor,
                align_corners=align_corners,
            )
            return out.numpy()

        x = random_gen(input_shape, rand_min=-100, rand_max=100)
        torch_pred = _get_torch_upsample_prediction(
            x,
            scale_factor=scale_factor,
            align_corners=align_corners,
            recompute_scale_factor=recompute_scale_factor,
        )

        input_placeholder_dict = {"x": mb.placeholder(shape=x.shape)}
        input_value_dict = {"x": x}

        def build_upsample(x):
            return mb.upsample_bilinear(
                x=x,
                scale_factor_height=scale_factor[0],
                scale_factor_width=scale_factor[1],
                align_corners=align_corners,
            )

        expected_output_type = torch_pred.shape + (types.fp32,)
        run_compare_builder(
            build_upsample,
            input_placeholder_dict,
            input_value_dict,
            expected_output_type,
            torch_pred,
            compute_unit=compute_unit,
            backend=backend,
            rtol=0.5,
        )