jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def resize_image_to(self, images, resolution):
        return t.interpolate(images, resolution, mode = self.resize_mode)
//  [python]
def resize_image_to(self, images, resolution):
        return F.interpolate(images, resolution, mode = self.resize_mode)
//  [python]
def sample(self, model, dl_iter, batch_size):
        G_kwargs, maybe_text_kwargs = self.generate_kwargs(dl_iter, batch_size)

        with self.accelerator.autocast():
            generator_output = model(**G_kwargs, **maybe_text_kwargs)

        if not self.train_upsampler:
            return generator_output

        output_size = generator_output.shape[-1]
        lowres_image = G_kwargs['lowres_image']
        lowres_image = F.interpolate(lowres_image, (output_size, output_size))

        return torch.cat([lowres_image, generator_output])
//  [python]
def resize_rel_pos_bias_table_levit(
        position_bias_table,
        new_size,
        interpolation: str = 'bicubic',
        antialias: bool = True,
):
    """
    Resample relative position bias table suggested in LeVit
    Adapted from: https://github.com/microsoft/Cream/blob/main/TinyViT/utils.py
    """
    L1, nH1 = position_bias_table.size()
    L2, nH2 = new_size
    assert nH1 == nH2
    if L1 != L2:
        orig_dtype = position_bias_table.dtype
        position_bias_table = position_bias_table.float()
        # bicubic interpolate relative_position_bias_table if not match
        S1 = int(L1 ** 0.5)
        S2 = int(L2 ** 0.5)
        relative_position_bias_table_resized = t.interpolate(
            position_bias_table.permute(1, 0).view(1, nH1, S1, S1),
            size=(S2, S2),
            mode=interpolation,
            align_corners=False if interpolation == 'linear' else None)
        relative_position_bias_table_resized = \
            relative_position_bias_table_resized.view(nH2, L2).permute(1, 0)
        relative_position_bias_table_resized = relative_position_bias_table_resized.to(orig_dtype)
        return relative_position_bias_table_resized
    else:
        return position_bias_table
//  [python]
def resize_rel_pos_bias_table_levit(
        position_bias_table,
        new_size,
        interpolation: str = 'bicubic',
        antialias: bool = True,
):
    """
    Resample relative position bias table suggested in LeVit
    Adapted from: https://github.com/microsoft/Cream/blob/main/TinyViT/utils.py
    """
    L1, nH1 = position_bias_table.size()
    L2, nH2 = new_size
    assert nH1 == nH2
    if L1 != L2:
        orig_dtype = position_bias_table.dtype
        position_bias_table = position_bias_table.float()
        # bicubic interpolate relative_position_bias_table if not match
        S1 = int(L1 ** 0.5)
        S2 = int(L2 ** 0.5)
        relative_position_bias_table_resized = F.interpolate(
            position_bias_table.permute(1, 0).view(1, nH1, S1, S1),
            size=(S2, S2),
            mode=interpolation,
            antialias=antialias)
        relative_position_bias_table_resized = \
            relative_position_bias_table_resized.view(nH2, L2).permute(1, 0)
        relative_position_bias_table_resized.to(orig_dtype)
        return relative_position_bias_table_resized
    else:
        return position_bias_table
//  [python]
def __call__(self, x_t, attention_store):
        k = 1
        maps = attention_store['down_cross'][2:4] + attention_store[
            'up_cross'][:3]
        maps = [
            item.reshape(self.alpha_layers.shape[0], -1, 1, 16, 16,
                         MAX_NUM_WORDS) for item in maps
        ]
        maps = torch.cat(maps, dim=1)
        maps = (maps * self.alpha_layers).sum(-1).mean(1)
        mask = nnf.max_pool2d(
            maps, (k * 2 + 1, k * 2 + 1), (1, 1), padding=(k, k))
        mask = nnf.upsample(mask, size=(x_t.shape[2:]))
        mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]
        mask = mask.gt(self.threshold)
        mask = (mask[:1] + mask[1:]).float()
        x_t = x_t[:1] + mask * (x_t - x_t[:1])
        return x_t
//  [python]
def __call__(self, x_t, attention_store):
        k = 1
        maps = attention_store['down_cross'][2:4] + attention_store[
            'up_cross'][:3]
        maps = [
            item.reshape(self.alpha_layers.shape[0], -1, 1, 16, 16,
                         MAX_NUM_WORDS) for item in maps
        ]
        maps = torch.cat(maps, dim=1)
        maps = (maps * self.alpha_layers).sum(-1).mean(1)
        mask = nnf.max_pool2d(
            maps, (k * 2 + 1, k * 2 + 1), (1, 1), padding=(k, k))
        mask = nnf.interpolate(mask, size=(x_t.shape[2:]))
        mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]
        mask = mask.gt(self.threshold)
        mask = (mask[:1] + mask[1:]).float()
        x_t = x_t[:1] + mask * (x_t - x_t[:1])
        return x_t
//  [python]
def generate_kwargs(self, dl_iter, batch_size):
        # what to pass into the generator
        # depends on whether training upsampler or not

        maybe_text_kwargs = dict()
        if self.train_upsampler or not self.unconditional:
            assert exists(dl_iter)

            if self.unconditional:
                real_images = next(dl_iter)
            else:
                result = next(dl_iter)
                assert isinstance(result, tuple), 'dataset should return a tuple of two items for text conditioned training, (images: Tensor, texts: List[str])'
                real_images, texts = result

                maybe_text_kwargs['texts'] = texts[:batch_size]

            real_images = real_images.to(self.device)

        # if training upsample generator, need to downsample real images

        if self.train_upsampler:
            size = self.unwrapped_G.input_image_size
            lowres_real_images = F.interpolate(real_images, (size, size), mode='bilinear', align_corners=False)

            G_kwargs = dict(lowres_image = lowres_real_images)
        else:
            assert exists(batch_size)

            G_kwargs = dict(batch_size = batch_size)

        # create noise

        noise = torch.randn(batch_size, self.unwrapped_G.style_network.dim, device = self.device)

        G_kwargs.update(noise = noise)

        return G_kwargs, maybe_text_kwargs
//  [python]
def generate_kwargs(self, dl_iter, batch_size):
        # what to pass into the generator
        # depends on whether training upsampler or not

        maybe_text_kwargs = dict()
        if self.train_upsampler or not self.unconditional:
            assert exists(dl_iter)

            if self.unconditional:
                real_images = next(dl_iter)
            else:
                result = next(dl_iter)
                assert isinstance(result, tuple), 'dataset should return a tuple of two items for text conditioned training, (images: Tensor, texts: List[str])'
                real_images, texts = result

                maybe_text_kwargs['texts'] = texts[:batch_size]

            real_images = real_images.to(self.device)

        # if training upsample generator, need to downsample real images

        if self.train_upsampler:
            size = self.unwrapped_G.input_image_size
            lowres_real_images = F.interpolate(real_images, (size, size))

            G_kwargs = dict(lowres_image = lowres_real_images)
        else:
            assert exists(batch_size)

            G_kwargs = dict(batch_size = batch_size)

        # create noise

        noise = torch.randn(batch_size, self.unwrapped_G.style_network.dim, device = self.device)

        G_kwargs.update(noise = noise)

        return G_kwargs, maybe_text_kwargs
//  [python]
def resize_rel_pos_bias_table_simple(
        rel_pos_bias,
        new_window_size: Tuple[int, int],
        new_bias_shape: Tuple[int, ...],
):
    dst_size = (new_window_size[0] * 2 - 1, new_window_size[1] * 2 - 1)
    if rel_pos_bias.ndim == 3:
        # TF maxvit style (num_heads, H, W) bias shape, no extra tokens currently supported
        _, dst_h, dst_w = new_bias_shape
        num_attn_heads, src_h, src_w = rel_pos_bias.shape
        assert dst_h == dst_size[0] and dst_w == dst_size[1]
        if src_h != dst_h or src_w != dst_w:
            rel_pos_bias = t.nn.functional.interpolate(
                rel_pos_bias.unsqueeze(0),
                size=dst_size,
                mode="bicubic",
                align_corners=False,
            ).squeeze(0)
    else:
        assert rel_pos_bias.ndim == 2
        # (num_pos, num_heads) (aka flat) bias shape
        dst_num_pos, _ = new_bias_shape
        src_num_pos, num_attn_heads = rel_pos_bias.shape
        num_extra_tokens = dst_num_pos - (dst_size[0] * dst_size[1])
        src_size = int((src_num_pos - num_extra_tokens) ** 0.5)
        src_size = (src_size, src_size)  # FIXME could support non-equal src if argument passed

        if src_size[0] != dst_size[0] or src_size[1] != dst_size[1]:
            if num_extra_tokens:
                extra_tokens = rel_pos_bias[-num_extra_tokens:, :]
                rel_pos_bias = rel_pos_bias[:-num_extra_tokens, :]
            else:
                extra_tokens = None

            rel_pos_bias = t.nn.functional.interpolate(
                rel_pos_bias.transpose(1, 0).reshape((1, -1, src_size[0], src_size[1])),
                size=dst_size,
                mode="bicubic",
                align_corners=False,
            ).view(-1, dst_num_pos - num_extra_tokens).transpose(0, 1)

            if extra_tokens is not None:
                rel_pos_bias = torch.cat((rel_pos_bias, extra_tokens), dim=0)

    return rel_pos_bias
//  [python]
def resize_rel_pos_bias_table_simple(
        rel_pos_bias,
        new_window_size: Tuple[int, int],
        new_bias_shape: Tuple[int, ...],
):
    dst_size = (new_window_size[0] * 2 - 1, new_window_size[1] * 2 - 1)
    if rel_pos_bias.ndim == 3:
        # TF maxvit style (num_heads, H, W) bias shape, no extra tokens currently supported
        _, dst_h, dst_w = new_bias_shape
        num_attn_heads, src_h, src_w = rel_pos_bias.shape
        assert dst_h == dst_size[0] and dst_w == dst_size[1]
        if src_h != dst_h or src_w != dst_w:
            rel_pos_bias = torch.nn.functional.interpolate(
                rel_pos_bias.unsqueeze(0),
                size=dst_size,
                mode="bicubic",
                align_corners=False,
            ).squeeze(0)
    else:
        assert rel_pos_bias.ndim == 2
        # (num_pos, num_heads) (aka flat) bias shape
        dst_num_pos, _ = new_bias_shape
        src_num_pos, num_attn_heads = rel_pos_bias.shape
        num_extra_tokens = dst_num_pos - (dst_size[0] * dst_size[1])
        src_size = int((src_num_pos - num_extra_tokens) ** 0.5)
        src_size = (src_size, src_size)  # FIXME could support non-equal src if argument passed

        if src_size[0] != dst_size[0] or src_size[1] != dst_size[1]:
            if num_extra_tokens:
                extra_tokens = rel_pos_bias[-num_extra_tokens:, :]
                rel_pos_bias = rel_pos_bias[:-num_extra_tokens, :]
            else:
                extra_tokens = None

            rel_pos_bias = torch.nn.functional.interpolate(
                rel_pos_bias.transpose(1, 0).reshape((1, -1, src_size[0], src_size[1])),
                size=dst_size,
                mode="bicubic",
                align_corners=False,
            ).view(-1, dst_num_pos - num_extra_tokens).transpose(0, 1)

            if extra_tokens is not None:
                rel_pos_bias = torch.cat((rel_pos_bias, extra_tokens), dim=0)

    return rel_pos_bias
//  [python]
def train_step(
        self,
        text_embeddings: torch.Tensor,
        pred_rgb: torch.Tensor,
        guidance_scale: float = 100.0,
        as_latent: bool = False,
    ) -> float:
        """
        Train step function for StableDiffusion.

        Parameters:
            text_embeddings (torch.Tensor): Embeddings tensor [B, 512].
            pred_rgb (torch.Tensor): Predicted RGB tensor [B, 3, 512, 512].
            guidance_scale (float): Guidance scaling factor.
            as_latent (bool): If True, considers pred_rgb as latent.

        Returns:
            float: Loss value.
        """
        if as_latent:
            latents = F.interpolate(pred_rgb, (64, 64), mode='bilinear', align_corners=False) * 2 - 1
        else:
            pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode='bilinear', align_corners=False)
            latents = self.encode_imgs(pred_rgb_512)

        t = torch.randint(self.min_step, self.max_step + 1, (latents.shape[0],), dtype=torch.long, device=self.device)

        with torch.no_grad():
            # add noise
            noise = torch.randn_like(latents)
            latents_noisy = self.scheduler.add_noise(latents, noise, t)
            # pred noise
            latent_model_input = torch.cat([latents_noisy] * 2)
            td = torch.cat([t] * 2)
            noise_pred = self.unet(latent_model_input, td, encoder_hidden_states=text_embeddings).sample

            noise_pred_uncond, noise_pred_pos = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_pos - noise_pred_uncond)

        w = 1 - self.alphas[t]
        grad = w[:, None, None, None] * (noise_pred - noise)
        grad = torch.nan_to_num(grad)

        targets = (latents - grad).detach()
        loss = 0.5 * F.mse_loss(latents.float(), targets, reduction='sum') / latents.shape[0]
        return loss
//  [python]
def _down_encode(self, x, r_embed, clip, cnet=None):
        level_outputs = []
        block_group = zip(self.down_blocks, self.down_downscalers, self.down_repeat_mappers)
        for down_block, downscaler, repmap in block_group:
            x = downscaler(x)
            for i in range(len(repmap) + 1):
                for block in down_block:
                    if isinstance(block, ResBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  ResBlock)):
                        if cnet is not None:
                            next_cnet = cnet()
                            if next_cnet is not None:
                                x = x + t.interpolate(next_cnet, size=x.shape[-2:], mode='bilinear',
                                                                  align_corners=True)
                        x = block(x)
                    elif isinstance(block, AttnBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  AttnBlock)):
                        x = block(x, clip)
                    elif isinstance(block, TimestepBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  TimestepBlock)):
                        x = block(x, r_embed)
                    else:
                        x = block(x)
                if i < len(repmap):
                    x = repmap[i](x)
            level_outputs.insert(0, x)
        return level_outputs
//  [python]
def _down_encode(self, x, r_embed, clip, cnet=None):
        level_outputs = []
        block_group = zip(self.down_blocks, self.down_downscalers, self.down_repeat_mappers)
        for down_block, downscaler, repmap in block_group:
            x = downscaler(x)
            for i in range(len(repmap) + 1):
                for block in down_block:
                    if isinstance(block, ResBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  ResBlock)):
                        if cnet is not None:
                            next_cnet = cnet()
                            if next_cnet is not None:
                                x = x + nn.functional.interpolate(next_cnet, size=x.shape[-2:], mode='bilinear',
                                                                  align_corners=True)
                        x = block(x)
                    elif isinstance(block, AttnBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  AttnBlock)):
                        x = block(x, clip)
                    elif isinstance(block, TimestepBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  TimestepBlock)):
                        x = block(x, r_embed)
                    else:
                        x = block(x)
                if i < len(repmap):
                    x = repmap[i](x)
            level_outputs.insert(0, x)
        return level_outputs
//  [python]
def process_unknown_mask(self, unknown, img_feat, softmax_scale):
        """Process unknown mask.

        Args:
            unknown (Tensor, optional): Unknown area map generated by trimap of
                shape (N, 1, ori_h, ori_w)
            img_feat (Tensor): The interpolated image feature map of shape
                (N, img_c, img_h, img_w).
            softmax_scale (float, optional): The softmax scale of the attention
                if unknown area is not provided in forward. Default: 1.

        Returns:
            tuple: 2-tuple of

                ``Tensor``: Interpolated unknown area map of shape \
                    (N, img_h*img_w, img_h, img_w).

                ``Tensor``: Softmax scale tensor of known and unknown area of \
                    shape (N, 2).
        """
        n, _, h, w = img_feat.shape

        if unknown is not None:
            unknown = unknown.clone()
            unknown = t.interpolate(
                unknown, scale_factor=1 / self.rate, mode=self.interpolation)
            unknown_mean = unknown.mean(dim=[2, 3])
            known_mean = 1 - unknown_mean
            unknown_scale = torch.clamp(
                torch.sqrt(unknown_mean / known_mean), 0.1, 10).to(img_feat)
            known_scale = torch.clamp(
                torch.sqrt(known_mean / unknown_mean), 0.1, 10).to(img_feat)
            softmax_scale = torch.cat([unknown_scale, known_scale], dim=1)
        else:
            unknown = torch.ones((n, 1, h, w)).to(img_feat)
            softmax_scale = torch.FloatTensor(
                [softmax_scale,
                 softmax_scale]).view(1, 2).repeat(n, 1).to(img_feat)

        return unknown, softmax_scale
//  [python]
def process_unknown_mask(self, unknown, img_feat, softmax_scale):
        """Process unknown mask.

        Args:
            unknown (Tensor, optional): Unknown area map generated by trimap of
                shape (N, 1, ori_h, ori_w)
            img_feat (Tensor): The interpolated image feature map of shape
                (N, img_c, img_h, img_w).
            softmax_scale (float, optional): The softmax scale of the attention
                if unknown area is not provided in forward. Default: 1.

        Returns:
            tuple: 2-tuple of

                ``Tensor``: Interpolated unknown area map of shape \
                    (N, img_h*img_w, img_h, img_w).

                ``Tensor``: Softmax scale tensor of known and unknown area of \
                    shape (N, 2).
        """
        n, _, h, w = img_feat.shape

        if unknown is not None:
            unknown = unknown.clone()
            unknown = F.interpolate(
                unknown, scale_factor=1 / self.rate, mode=self.interpolation)
            unknown_mean = unknown.mean(dim=[2, 3])
            known_mean = 1 - unknown_mean
            unknown_scale = torch.clamp(
                torch.sqrt(unknown_mean / known_mean), 0.1, 10).to(img_feat)
            known_scale = torch.clamp(
                torch.sqrt(known_mean / unknown_mean), 0.1, 10).to(img_feat)
            softmax_scale = torch.cat([unknown_scale, known_scale], dim=1)
        else:
            unknown = torch.ones((n, 1, h, w)).to(img_feat)
            softmax_scale = torch.FloatTensor(
                [softmax_scale,
                 softmax_scale]).view(1, 2).repeat(n, 1).to(img_feat)

        return unknown, softmax_scale
//  [python]
def resize_norm_img(self, img, max_wh_ratio):
        imgC, imgH, imgW = self.rec_image_shape
        assert imgC == img.shape[0]
        imgW = int((imgH * max_wh_ratio))

        h, w = img.shape[1:]
        ratio = w / float(h)
        if math.ceil(imgH * ratio) > imgW:
            resized_w = imgW
        else:
            resized_w = int(math.ceil(imgH * ratio))
        resized_image = t.nn.functional.interpolate(
            img.unsqueeze(0),
            size=(imgH, resized_w),
            mode="bilinear",
            align_corners=True,
        )
        resized_image /= 255.0
        resized_image -= 0.5
        resized_image /= 0.5
        padding_im = torch.zeros((imgC, imgH, imgW), dtype=torch.float32).to(img.device)
        padding_im[:, :, 0:resized_w] = resized_image[0]
        return padding_im
//  [python]
def resize_norm_img(self, img, max_wh_ratio):
        imgC, imgH, imgW = self.rec_image_shape
        assert imgC == img.shape[0]
        imgW = int((imgH * max_wh_ratio))

        h, w = img.shape[1:]
        ratio = w / float(h)
        if math.ceil(imgH * ratio) > imgW:
            resized_w = imgW
        else:
            resized_w = int(math.ceil(imgH * ratio))
        resized_image = torch.nn.functional.interpolate(
            img.unsqueeze(0),
            size=(imgH, resized_w),
            mode="bilinear",
            align_corners=True,
        )
        resized_image /= 255.0
        resized_image -= 0.5
        resized_image /= 0.5
        padding_im = torch.zeros((imgC, imgH, imgW), dtype=torch.float32).to(img.device)
        padding_im[:, :, 0:resized_w] = resized_image[0]
        return padding_im
//  [python]
def forward(self, img_feat, alpha_feat, unknown=None, softmax_scale=1.):
        """Forward function of GCAModule.

        Args:
            img_feat (Tensor): Image feature map of shape
                (N, ori_c, ori_h, ori_w).
            alpha_feat (Tensor): Alpha feature map of shape
                (N, alpha_c, ori_h, ori_w).
            unknown (Tensor, optional): Unknown area map generated by trimap.
                If specified, this tensor should have shape
                (N, 1, ori_h, ori_w).
            softmax_scale (float, optional): The softmax scale of the attention
                if unknown area is not provided in forward. Default: 1.

        Returns:
            Tensor: The augmented alpha feature.
        """

        if alpha_feat.shape[2:4] != img_feat.shape[2:4]:
            raise ValueError(
                'image feature size does not align with alpha feature size: '
                f'image feature size {img_feat.shape[2:4]}, '
                f'alpha feature size {alpha_feat.shape[2:4]}')

        if unknown is not None and unknown.shape[2:4] != img_feat.shape[2:4]:
            raise ValueError(
                'image feature size does not align with unknown mask size: '
                f'image feature size {img_feat.shape[2:4]}, '
                f'unknown mask size {unknown.shape[2:4]}')

        # preprocess image feature
        img_feat = self.guidance_conv(img_feat)
        img_feat = t.interpolate(
            img_feat, scale_factor=1 / self.rate, mode=self.interpolation)

        # preprocess unknown mask
        unknown, softmax_scale = self.process_unknown_mask(
            unknown, img_feat, softmax_scale)

        img_ps, alpha_ps, unknown_ps = self.extract_feature_maps_patches(
            img_feat, alpha_feat, unknown)

        # create self correlation mask with shape:
        # (N, img_h*img_w, img_h, img_w)
        self_mask = self.get_self_correlation_mask(img_feat)

        # split tensors by batch dimension; tuple is returned
        img_groups = torch.split(img_feat, 1, dim=0)
        img_ps_groups = torch.split(img_ps, 1, dim=0)
        alpha_ps_groups = torch.split(alpha_ps, 1, dim=0)
        unknown_ps_groups = torch.split(unknown_ps, 1, dim=0)
        scale_groups = torch.split(softmax_scale, 1, dim=0)
        groups = (img_groups, img_ps_groups, alpha_ps_groups,
                  unknown_ps_groups, scale_groups)

        out = []
        # i is the virtual index of the sample in the current batch
        for img_i, img_ps_i, alpha_ps_i, unknown_ps_i, scale_i in zip(*groups):
            similarity_map = self.compute_similarity_map(img_i, img_ps_i)

            gca_score = self.compute_guided_attention_score(
                similarity_map, unknown_ps_i, scale_i, self_mask)

            out_i = self.propagate_alpha_feature(gca_score, alpha_ps_i)

            out.append(out_i)

        out = torch.cat(out, dim=0)
        out.reshape_as(alpha_feat)

        out = self.out_conv(out) + alpha_feat
        return out
//  [python]
def forward(self, img_feat, alpha_feat, unknown=None, softmax_scale=1.):
        """Forward function of GCAModule.

        Args:
            img_feat (Tensor): Image feature map of shape
                (N, ori_c, ori_h, ori_w).
            alpha_feat (Tensor): Alpha feature map of shape
                (N, alpha_c, ori_h, ori_w).
            unknown (Tensor, optional): Unknown area map generated by trimap.
                If specified, this tensor should have shape
                (N, 1, ori_h, ori_w).
            softmax_scale (float, optional): The softmax scale of the attention
                if unknown area is not provided in forward. Default: 1.

        Returns:
            Tensor: The augmented alpha feature.
        """

        if alpha_feat.shape[2:4] != img_feat.shape[2:4]:
            raise ValueError(
                'image feature size does not align with alpha feature size: '
                f'image feature size {img_feat.shape[2:4]}, '
                f'alpha feature size {alpha_feat.shape[2:4]}')

        if unknown is not None and unknown.shape[2:4] != img_feat.shape[2:4]:
            raise ValueError(
                'image feature size does not align with unknown mask size: '
                f'image feature size {img_feat.shape[2:4]}, '
                f'unknown mask size {unknown.shape[2:4]}')

        # preprocess image feature
        img_feat = self.guidance_conv(img_feat)
        img_feat = F.interpolate(
            img_feat, scale_factor=1 / self.rate, mode=self.interpolation)

        # preprocess unknown mask
        unknown, softmax_scale = self.process_unknown_mask(
            unknown, img_feat, softmax_scale)

        img_ps, alpha_ps, unknown_ps = self.extract_feature_maps_patches(
            img_feat, alpha_feat, unknown)

        # create self correlation mask with shape:
        # (N, img_h*img_w, img_h, img_w)
        self_mask = self.get_self_correlation_mask(img_feat)

        # split tensors by batch dimension; tuple is returned
        img_groups = torch.split(img_feat, 1, dim=0)
        img_ps_groups = torch.split(img_ps, 1, dim=0)
        alpha_ps_groups = torch.split(alpha_ps, 1, dim=0)
        unknown_ps_groups = torch.split(unknown_ps, 1, dim=0)
        scale_groups = torch.split(softmax_scale, 1, dim=0)
        groups = (img_groups, img_ps_groups, alpha_ps_groups,
                  unknown_ps_groups, scale_groups)

        out = []
        # i is the virtual index of the sample in the current batch
        for img_i, img_ps_i, alpha_ps_i, unknown_ps_i, scale_i in zip(*groups):
            similarity_map = self.compute_similarity_map(img_i, img_ps_i)

            gca_score = self.compute_guided_attention_score(
                similarity_map, unknown_ps_i, scale_i, self_mask)

            out_i = self.propagate_alpha_feature(gca_score, alpha_ps_i)

            out.append(out_i)

        out = torch.cat(out, dim=0)
        out.reshape_as(alpha_feat)

        out = self.out_conv(out) + alpha_feat
        return out
//  [python]
def apply_visual_prompts(self, image, mask):
    if torch.sum(mask).item() <= 1:
      return image
    image_array = np.array(image)
    img_h = image_array.shape[0]
    img_w = image_array.shape[1]
    mask = (
        t.interpolate(mask[None][None], size=(img_h, img_w), mode="nearest")
        .squeeze()
        .detach()
        .cpu()
        .numpy()
    )
    mask = (mask > self.mask_threshold).astype(np.uint8)
    prompted_image = apply_visual_prompts(
        image_array, mask, self.visual_prompt_type, self.visualize
    )
    return prompted_image
//  [python]
def apply_visual_prompts(self, image, mask):
    if torch.sum(mask).item() <= 1:
      return image
    image_array = np.array(image)
    img_h = image_array.shape[0]
    img_w = image_array.shape[1]
    mask = (
        F.interpolate(mask[None][None], size=(img_h, img_w), mode="nearest")
        .squeeze()
        .detach()
        .cpu()
        .numpy()
    )
    mask = (mask > self.mask_threshold).astype(np.uint8)
    prompted_image = apply_visual_prompts(
        image_array, mask, self.visual_prompt_type, self.visualize
    )
    return prompted_image
//  [python]
def preprocess_mask(mask, batch_size, scale_factor=8):
    if not isinstance(mask, torch.FloatTensor):
        mask = mask.convert("L")
        w, h = mask.size
        w, h = (x - x % 8 for x in (w, h))  # resize to integer multiple of 8
        mask = mask.resize((w // scale_factor, h // scale_factor), resample=PIL_INTERPOLATION["nearest"])
        mask = np.array(mask).astype(np.float32) / 255.0
        mask = np.tile(mask, (4, 1, 1))
        mask = np.vstack([mask[None]] * batch_size)
        mask = 1 - mask  # repaint white, keep black
        mask = torch.from_numpy(mask)
        return mask

    else:
        valid_mask_channel_sizes = [1, 3]
        # if mask channel is fourth tensor dimension, permute dimensions to pytorch standard (B, C, H, W)
        if mask.shape[3] in valid_mask_channel_sizes:
            mask = mask.permute(0, 3, 1, 2)
        elif mask.shape[1] not in valid_mask_channel_sizes:
            raise ValueError(
                f"Mask channel dimension of size in {valid_mask_channel_sizes} should be second or fourth dimension,"
                f" but received mask of shape {tuple(mask.shape)}"
            )
        # (potentially) reduce mask channel dimension from 3 to 1 for broadcasting to latent shape
        mask = mask.mean(dim=1, keepdim=True)
        h, w = mask.shape[-2:]
        h, w = (x - x % 8 for x in (h, w))  # resize to integer multiple of 8
        mask = t.nn.functional.interpolate(mask, (h // scale_factor, w // scale_factor))
        return mask
//  [python]
def preprocess_mask(mask, batch_size, scale_factor=8):
    if not isinstance(mask, torch.FloatTensor):
        mask = mask.convert("L")
        w, h = mask.size
        w, h = (x - x % 8 for x in (w, h))  # resize to integer multiple of 8
        mask = mask.resize((w // scale_factor, h // scale_factor), resample=PIL_INTERPOLATION["nearest"])
        mask = np.array(mask).astype(np.float32) / 255.0
        mask = np.tile(mask, (4, 1, 1))
        mask = np.vstack([mask[None]] * batch_size)
        mask = 1 - mask  # repaint white, keep black
        mask = torch.from_numpy(mask)
        return mask

    else:
        valid_mask_channel_sizes = [1, 3]
        # if mask channel is fourth tensor dimension, permute dimensions to pytorch standard (B, C, H, W)
        if mask.shape[3] in valid_mask_channel_sizes:
            mask = mask.permute(0, 3, 1, 2)
        elif mask.shape[1] not in valid_mask_channel_sizes:
            raise ValueError(
                f"Mask channel dimension of size in {valid_mask_channel_sizes} should be second or fourth dimension,"
                f" but received mask of shape {tuple(mask.shape)}"
            )
        # (potentially) reduce mask channel dimension from 3 to 1 for broadcasting to latent shape
        mask = mask.mean(dim=1, keepdim=True)
        h, w = mask.shape[-2:]
        h, w = (x - x % 8 for x in (h, w))  # resize to integer multiple of 8
        mask = torch.nn.functional.interpolate(mask, (h // scale_factor, w // scale_factor))
        return mask
//  [python]
def forward(self,
                mlvl_feats,
                img_metas,
                dn_label_query=None,
                dn_bbox_query=None,
                attn_mask=None):
        batch_size = mlvl_feats[0].size(0)
        input_img_h, input_img_w = img_metas[0]['batch_input_shape']
        img_masks = mlvl_feats[0].new_ones(
            (batch_size, input_img_h, input_img_w))
        for img_id in range(batch_size):
            img_h, img_w = img_metas[img_id]['img_shape']
            img_masks[img_id, :img_h, :img_w] = 0

        mlvl_masks = []
        mlvl_positional_encodings = []
        for feat in mlvl_feats:
            mlvl_masks.append(
                F.interpolate(img_masks[None],
                              size=feat.shape[-2:]).to(torch.bool).squeeze(0))
            mlvl_positional_encodings.append(
                self.positional_encoding(mlvl_masks[-1]))

        query_embeds = None
        hs, inter_references, topk_score, topk_anchor, enc_outputs = \
            self.transformer(
                mlvl_feats,
                mlvl_masks,
                query_embeds,
                mlvl_positional_encodings,
                dn_label_query,
                dn_bbox_query,
                attn_mask,
                reg_branches=self.reg_branches if self.with_box_refine else None,  # noqa:E501
                cls_branches=self.cls_branches if self.as_two_stage else None  # noqa:E501
            )
        outs = []
        num_level = len(mlvl_feats)
        start = 0
        for lvl in range(num_level):
            bs, c, h, w = mlvl_feats[lvl].shape
            end = start + h * w
            feat = enc_outputs[start:end].permute(1, 2, 0).contiguous()
            start = end
            outs.append(feat.reshape(bs, c, h, w))
        outs.append(self.downsample(outs[-1]))

        hs = hs.permute(0, 2, 1, 3)

        if dn_label_query is not None and dn_label_query.size(1) == 0:
            # NOTE: If there is no target in the image, the parameters of
            # label_embedding won't be used in producing loss, which raises
            # RuntimeError when using distributed mode.
            hs[0] += self.dn_generator.label_embedding.weight[0, 0] * 0.0

        outputs_classes = []
        outputs_coords = []

        for lvl in range(hs.shape[0]):
            reference = inter_references[lvl]
            reference = inverse_sigmoid(reference, eps=1e-3)
            outputs_class = self.cls_branches[lvl](hs[lvl])
            tmp = self.reg_branches[lvl](hs[lvl])
            if reference.shape[-1] == 4:
                tmp += reference
            else:
                assert reference.shape[-1] == 2
                tmp[..., :2] += reference
            outputs_coord = tmp.sigmoid()
            outputs_classes.append(outputs_class)
            outputs_coords.append(outputs_coord)

        outputs_classes = torch.stack(outputs_classes)
        outputs_coords = torch.stack(outputs_coords)

        return outputs_classes, outputs_coords, topk_score, topk_anchor, outs
//  [python]
def postprocess(self, outputs, inputs, task_name):
        """
        Post-process of ViM, based on task_name.

        Args:
            inputs: batched input image (B, 3, H, W)
            outputs: batched output (format based on task_name)
            task_name (str): task name
        """

        _, in_channels, img_height, img_width = inputs.size()

        if 'seg' in task_name:
            # outputs in shape of [1, C, H, W]
            seg = F.softmax(outputs, dim=1)
            seg = F.upsample_bilinear(seg, size=(img_height, img_width))
            seg = seg[0].detach().cpu()
            pred = torch.argmax(seg, dim=0)

            labels = sorted(list(set(pred.reshape(-1).numpy())))

            masks, scores = [], []
            for label in labels:
                mask = (pred == label)
                masks.append(mask.long().numpy())
                scores.append(((mask.float() * seg[label]).sum()
                               / mask.float().sum()).item())

            label_names = [
                self.label_maps[task_name][label] for label in labels
            ]

            return {
                OutputKeys.MASKS: masks,
                OutputKeys.LABELS: label_names,
                OutputKeys.SCORES: scores
            }
        else:
            raise NotImplementedError(
                'Only segmentation task is currently supported in pipeline')
//  [python]
def postprocess(self, outputs, inputs, task_name):
        """
        Post-process of ViM, based on task_name.

        Args:
            inputs: batched input image (B, 3, H, W)
            outputs: batched output (format based on task_name)
            task_name (str): task name
        """

        _, in_channels, img_height, img_width = inputs.size()

        if 'seg' in task_name:
            # outputs in shape of [1, C, H, W]
            seg = F.softmax(outputs, dim=1)
            seg = F.interpolate(seg, (img_height, img_width), None, 'bilinear',
                                True)
            seg = seg[0].detach().cpu()
            pred = torch.argmax(seg, dim=0)

            labels = sorted(list(set(pred.reshape(-1).numpy())))

            masks, scores = [], []
            for label in labels:
                mask = (pred == label)
                masks.append(mask.long().numpy())
                scores.append(((mask.float() * seg[label]).sum()
                               / mask.float().sum()).item())

            label_names = [
                self.label_maps[task_name][label] for label in labels
            ]

            return {
                OutputKeys.MASKS: masks,
                OutputKeys.LABELS: label_names,
                OutputKeys.SCORES: scores
            }
        else:
            raise NotImplementedError(
                'Only segmentation task is currently supported in pipeline')
//  [python]
def resample(input, size, align_corners=True):
    n, c, h, w = input.shape

    if isinstance(size, (int, float)):
        short, long = (w, h) if w <= h else (h, w)
        requested_new_short = size
        new_short, new_long = round(requested_new_short), round(requested_new_short * long / short)
        dw, dh = (new_short, new_long) if w <= h else (new_long, new_short)
    else:
        dh, dw = size

    input = input.view([n * c, 1, h, w])

    if dh < h:
        kernel_h = lanczos(ramp(dh / h, 2), 2).to(input.device, input.dtype)
        pad_h = (kernel_h.shape[0] - 1) // 2
        input = F.pad(input, (0, 0, pad_h, pad_h), "reflect")
        input = F.conv2d(input, kernel_h[None, None, :, None])

    if dw < w:
        kernel_w = lanczos(ramp(dw / w, 2), 2).to(input.device, input.dtype)
        pad_w = (kernel_w.shape[0] - 1) // 2
        input = F.pad(input, (pad_w, pad_w, 0, 0), "reflect")
        input = F.conv2d(input, kernel_w[None, None, None, :])

    input = input.view([n, c, h, w])
    return t.nn.functional.interpolate(input, (dh, dw), mode="bicubic", align_corners=align_corners)
//  [python]
def resample(input, size, align_corners=True):
    n, c, h, w = input.shape

    if isinstance(size, (int, float)):
        short, long = (w, h) if w <= h else (h, w)
        requested_new_short = size
        new_short, new_long = round(requested_new_short), round(requested_new_short * long / short)
        dw, dh = (new_short, new_long) if w <= h else (new_long, new_short)
    else:
        dh, dw = size

    input = input.view([n * c, 1, h, w])

    if dh < h:
        kernel_h = lanczos(ramp(dh / h, 2), 2).to(input.device, input.dtype)
        pad_h = (kernel_h.shape[0] - 1) // 2
        input = F.pad(input, (0, 0, pad_h, pad_h), "reflect")
        input = F.conv2d(input, kernel_h[None, None, :, None])

    if dw < w:
        kernel_w = lanczos(ramp(dw / w, 2), 2).to(input.device, input.dtype)
        pad_w = (kernel_w.shape[0] - 1) // 2
        input = F.pad(input, (pad_w, pad_w, 0, 0), "reflect")
        input = F.conv2d(input, kernel_w[None, None, None, :])

    input = input.view([n, c, h, w])
    return F.interpolate(input, (dh, dw), mode="bicubic", align_corners=align_corners)
//  [python]
def test_interpolate_upsample_dynamic_sizes(self):
        class MyModel(torch.nn.Module):
            def forward(self, x):
                size = [v * 2 for v in x.size()[2:]]
                return nn.functional.upsample(x, size=size, mode="nearest")

        x = torch.randn(1, 2, 3, 4, requires_grad=True)
        model = MyModel()
        self.run_model_test(
            model, train=False, input=(x), batch_size=BATCH_SIZE, use_gpu=False
        )
//  [python]
def test_interpolate_upsample_dynamic_sizes(self):
        class MyModel(torch.nn.Module):
            def forward(self, x):
                size = [v * 2 for v in x.size()[2:]]
                return nn.functional.interpolate(x, size=size, mode="nearest")

        x = torch.randn(1, 2, 3, 4, requires_grad=True)
        model = MyModel()
        self.run_model_test(
            model, train=False, input=(x), batch_size=BATCH_SIZE, use_gpu=False
        )
//  [python]
def test_interpolate_upsample(self):
        class MyModel(torch.nn.Module):
            def forward(self, x):
                size = [v * 2 for v in x.size()[2:]]
                # work around for now: turn the dynamic sizes into constant
                size = [int(i) for i in size]
                return nn.functional.interpolate(x, size=size, mode="nearest")

        x = torch.randn(1, 2, 3, 4, requires_grad=True)
        model = MyModel()
        self.run_model_test(
            model, train=False, input=(x), batch_size=BATCH_SIZE, use_gpu=False
        )
//  [python]
def forward(self, imgs: list[Tensor]) -> tuple[Tensor, list[ImageSize]]:
        # TODO: support other input formats e.g. file path, numpy
        resized_imgs, original_sizes = [], []
        for i in range(len(imgs)):
            img = imgs[i]
            # NOTE: assume that image layout is CHW
            original_sizes.append(ImageSize(height=img.shape[1], width=img.shape[2]))
            resized_imgs.append(
                t.interpolate(img.unsqueeze(0), size=self.size, mode=self.interpolation_mode)
            )
        return concatenate(resized_imgs), original_sizes
//  [python]
def forward(self, imgs: list[Tensor]) -> tuple[Tensor, list[ImageSize]]:
        # TODO: support other input formats e.g. file path, numpy
        resized_imgs, original_sizes = [], []
        for i in range(len(imgs)):
            img = imgs[i]
            # NOTE: assume that image layout is CHW
            original_sizes.append(ImageSize(height=img.shape[1], width=img.shape[2]))
            resized_imgs.append(
                # TODO: fix kornia resize to support onnx
                torch.nn.functional.interpolate(img.unsqueeze(0), size=self.size, mode=self.interpolation_mode)
            )
        return concatenate(resized_imgs), original_sizes
//  [python]
def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        if self.add_relu:
            x = self.activation(x)
        if self.interpolate:
            x = F.upsample(x, scale_factor=2, mode="bilinear")
        return x
//  [python]
def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        if self.add_relu:
            x = self.activation(x)
        if self.interpolate:
            x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
        return x
//  [python]
def sample_keypoints(scoremap, num_samples = 8192, device = get_best_device(), use_nms = True, 
                     sample_topk = False, return_scoremap = False, sharpen = False, upsample = False,
                     increase_coverage = False, remove_borders = False):
    #scoremap = scoremap**2
    log_scoremap = (scoremap+1e-10).log()
    if upsample:
        log_scoremap = t.nn.functional.interpolate(log_scoremap[:,None], scale_factor = 3, mode = "bicubic", align_corners = False)[:,0]#.clamp(min = 0)
        scoremap = log_scoremap.exp()
    B,H,W = scoremap.shape
    if increase_coverage:
        weights = (-torch.linspace(-2, 2, steps = 51, device = device)**2).exp()[None,None]
        # 10000 is just some number for maybe numerical stability, who knows. :), result is invariant anyway
        local_density_x = F.conv2d((scoremap[:,None]+1e-6)*10000,weights[...,None,:], padding = (0,51//2))
        local_density = F.conv2d(local_density_x, weights[...,None], padding = (51//2,0))[:,0]
        scoremap = scoremap * (local_density+1e-8)**(-1/2)
    grid = get_grid(B,H,W, device=device).reshape(B,H*W,2)
    if sharpen:
        laplace_operator = torch.tensor([[[[0,1,0],[1,-4,1],[0,1,0]]]], device = device)/4
        scoremap = scoremap[:,None] - 0.5 * F.conv2d(scoremap[:,None], weight = laplace_operator, padding = 1)
        scoremap = scoremap[:,0].clamp(min = 0)
    if use_nms:
        scoremap = scoremap * (scoremap == F.max_pool2d(scoremap, (3, 3), stride = 1, padding = 1))
    if remove_borders:
        frame = torch.zeros_like(scoremap)
        # we hardcode 4px, could do it nicer, but whatever
        frame[...,4:-4, 4:-4] = 1
        scoremap = scoremap * frame
    if sample_topk:
        inds = torch.topk(scoremap.reshape(B,H*W), k = num_samples).indices
    else:
        inds = torch.multinomial(scoremap.reshape(B,H*W), num_samples = num_samples, replacement=False)
    kps = torch.gather(grid, dim = 1, index = inds[...,None].expand(B,num_samples,2))
    if return_scoremap:
        return kps, torch.gather(scoremap.reshape(B,H*W), dim = 1, index = inds)
    return kps
//  [python]
def sample_keypoints(scoremap, num_samples = 8192, device = get_best_device(), use_nms = True, 
                     sample_topk = False, return_scoremap = False, sharpen = False, upsample = False,
                     increase_coverage = False, remove_borders = False):
    #scoremap = scoremap**2
    log_scoremap = (scoremap+1e-10).log()
    if upsample:
        log_scoremap = F.interpolate(log_scoremap[:,None], scale_factor = 3, mode = "bicubic", align_corners = False)[:,0]#.clamp(min = 0)
        scoremap = log_scoremap.exp()
    B,H,W = scoremap.shape
    if increase_coverage:
        weights = (-torch.linspace(-2, 2, steps = 51, device = device)**2).exp()[None,None]
        # 10000 is just some number for maybe numerical stability, who knows. :), result is invariant anyway
        local_density_x = F.conv2d((scoremap[:,None]+1e-6)*10000,weights[...,None,:], padding = (0,51//2))
        local_density = F.conv2d(local_density_x, weights[...,None], padding = (51//2,0))[:,0]
        scoremap = scoremap * (local_density+1e-8)**(-1/2)
    grid = get_grid(B,H,W, device=device).reshape(B,H*W,2)
    if sharpen:
        laplace_operator = torch.tensor([[[[0,1,0],[1,-4,1],[0,1,0]]]], device = device)/4
        scoremap = scoremap[:,None] - 0.5 * F.conv2d(scoremap[:,None], weight = laplace_operator, padding = 1)
        scoremap = scoremap[:,0].clamp(min = 0)
    if use_nms:
        scoremap = scoremap * (scoremap == F.max_pool2d(scoremap, (3, 3), stride = 1, padding = 1))
    if remove_borders:
        frame = torch.zeros_like(scoremap)
        # we hardcode 4px, could do it nicer, but whatever
        frame[...,4:-4, 4:-4] = 1
        scoremap = scoremap * frame
    if sample_topk:
        inds = torch.topk(scoremap.reshape(B,H*W), k = num_samples).indices
    else:
        inds = torch.multinomial(scoremap.reshape(B,H*W), num_samples = num_samples, replacement=False)
    kps = torch.gather(grid, dim = 1, index = inds[...,None].expand(B,num_samples,2))
    if return_scoremap:
        return kps, torch.gather(scoremap.reshape(B,H*W), dim = 1, index = inds)
    return kps