jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_pyramid_layer(image, gaussian_k, direction='down'):
    """Get the pyramid layer.

    Args:
        image (Tensor): Input image.
        gaussian_k (Tensor): Gaussian kernel
        direction (str, optional): The direction of pyramid. Defaults to
            'down'.

    Returns:
        Tensor: The output of the pyramid.
    """
    gaussian_k = gaussian_k.to(image.device)
    if direction == 'up':
        image = t.nn.functional.upsample(image, scale_factor=2)
    multiband = [
        F.conv2d(
            image[:, i:i + 1, :, :],
            gaussian_k,
            padding=2,
            stride=1 if direction == 'up' else 2) for i in range(3)
    ]
    image = torch.cat(multiband, dim=1)
    return image
//  [python]
def get_pyramid_layer(image, gaussian_k, direction='down'):
    """Get the pyramid layer.

    Args:
        image (Tensor): Input image.
        gaussian_k (Tensor): Gaussian kernel
        direction (str, optional): The direction of pyramid. Defaults to
            'down'.

    Returns:
        Tensor: The output of the pyramid.
    """
    gaussian_k = gaussian_k.to(image.device)
    if direction == 'up':
        image = F.interpolate(image, scale_factor=2)
    multiband = [
        F.conv2d(
            image[:, i:i + 1, :, :],
            gaussian_k,
            padding=2,
            stride=1 if direction == 'up' else 2) for i in range(3)
    ]
    image = torch.cat(multiband, dim=1)
    return image
//  [python]
def forward(self, x):
        """Computes the upsampling pass

        Arguments
        ---------
        x: torch.Tensor
            layer inputs

        Returns
        -------
        result: torch.Tensor
            upsampled outputs"""
        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):
        """Computes the upsampling pass

        Arguments
        ---------
        x: torch.Tensor
            layer inputs

        Returns
        -------
        result: torch.Tensor
            upsampled outputs"""
        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 feat_kl_loss(feats, labels, feats_mem):
    B, C, H, W = feats.shape
    CLS = 19
    _, H_org, W_org = labels.shape
    labels = t.nn.functional.interpolate(labels.unsqueeze(1).float(), (H, W), mode='nearest')

    select_feat = torch.clone(feats)
    feats = feats.permute(0, 2, 3, 1).contiguous().view(-1, C)
    select_feat = select_feat.permute(0, 2, 3, 1).contiguous().view(-1, C)
    labels = labels.view(-1)

    feats_mem = feats_mem.squeeze(1)
    batch_feats_mem = torch.zeros_like(feats_mem)

    ignore_index = 255
    for c in labels.unique():
        if c == ignore_index: continue
        c = c.item()
        feats_cls = feats[labels == c].mean(0)
        batch_feats_mem[int(c)] = feats_cls

        m = labels == c
        m = m[..., None].repeat(1, C)
        feat_temp = feats_mem[int(c)][None, ...].expand(labels.shape[0], -1)
        select_feat = torch.where(m, feat_temp, select_feat)
    feats = feats.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()
    select_feat = select_feat.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()

    T = 20
    alpha = 0.9
    loss_kl = nn.KLDivLoss()(
        F.log_softmax(feats / T, dim=1),
        F.softmax(select_feat / T, dim=1)) * (alpha * T * T) + F.cross_entropy(feats, torch.argmax(select_feat, dim=1).long()) * (1. - alpha)


    return loss_kl, batch_feats_mem, select_feat
//  [python]
def feat_kl_loss(feats, labels, feats_mem):
    B, C, H, W = feats.shape
    CLS = 19
    _, H_org, W_org = labels.shape
    labels = F.interpolate(labels.unsqueeze(1).float(), (H, W), mode='nearest')

    select_feat = torch.clone(feats)
    feats = feats.permute(0, 2, 3, 1).contiguous().view(-1, C)
    select_feat = select_feat.permute(0, 2, 3, 1).contiguous().view(-1, C)
    labels = labels.view(-1)

    feats_mem = feats_mem.squeeze(1)
    batch_feats_mem = torch.zeros_like(feats_mem)

    ignore_index = 255
    for c in labels.unique():
        if c == ignore_index: continue
        c = c.item()
        feats_cls = feats[labels == c].mean(0)
        batch_feats_mem[int(c)] = feats_cls

        m = labels == c
        m = m[..., None].repeat(1, C)
        feat_temp = feats_mem[int(c)][None, ...].expand(labels.shape[0], -1)
        select_feat = torch.where(m, feat_temp, select_feat)
    feats = feats.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()
    select_feat = select_feat.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()

    T = 20
    alpha = 0.9
    loss_kl = nn.KLDivLoss()(
        F.log_softmax(feats / T, dim=1),
        F.softmax(select_feat / T, dim=1)) * (alpha * T * T) + F.cross_entropy(feats, torch.argmax(select_feat, dim=1).long()) * (1. - alpha)


    return loss_kl, batch_feats_mem, select_feat
//  [python]
def interpolate(self, x, output_size, size):
        """ Feature map interpolation. """
        B, N, C = x.shape
        H, W = size
        assert N == 1 + H * W or 2 + H * W

        diff = N - H * W
        other_token  = x[:, :diff, :]
        img_tokens = x[:, diff:, :]
        
        img_tokens = img_tokens.transpose(1, 2).reshape(B, C, H, W)
        img_tokens = F.interpolate(img_tokens, size=output_size, mode='bilinear')  # FIXME: May have alignment issue.
        img_tokens = img_tokens.reshape(B, C, -1).transpose(1, 2)
        
        out = torch.cat((other_token, img_tokens), dim=1)

        return out
//  [python]
def forward(self, out_features, patch_h, patch_w):
        out = []
        for i, x in enumerate(out_features):
            if self.use_clstoken:
                x, cls_token = x[0], x[1]
                readout = cls_token.unsqueeze(1).expand_as(x)
                x = self.readout_projects[i](torch.cat((x, readout), -1))
            else:
                x = x[0]

            x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))

            x = self.projects[i](x)
            x = self.resize_layers[i](x)

            out.append(x)

        layer_1, layer_2, layer_3, layer_4 = out

        layer_1_rn = self.scratch.layer1_rn(layer_1)
        layer_2_rn = self.scratch.layer2_rn(layer_2)
        layer_3_rn = self.scratch.layer3_rn(layer_3)
        layer_4_rn = self.scratch.layer4_rn(layer_4)

        path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
        path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])
        path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])
        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)

        out = self.scratch.output_conv1(path_1)
        out = t.interpolate(out, (int(patch_h * 14), int(patch_w * 14)), mode="bilinear", align_corners=True)
        out = self.scratch.output_conv2(out)

        return out
//  [python]
def forward(self, out_features, patch_h, patch_w):
        out = []
        for i, x in enumerate(out_features):
            if self.use_clstoken:
                x, cls_token = x[0], x[1]
                readout = cls_token.unsqueeze(1).expand_as(x)
                x = self.readout_projects[i](torch.cat((x, readout), -1))
            else:
                x = x[0]

            x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))

            x = self.projects[i](x)
            x = self.resize_layers[i](x)

            out.append(x)

        layer_1, layer_2, layer_3, layer_4 = out

        layer_1_rn = self.scratch.layer1_rn(layer_1)
        layer_2_rn = self.scratch.layer2_rn(layer_2)
        layer_3_rn = self.scratch.layer3_rn(layer_3)
        layer_4_rn = self.scratch.layer4_rn(layer_4)

        path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
        path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])
        path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])
        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)

        out = self.scratch.output_conv1(path_1)
        out = F.interpolate(out, (int(patch_h * 14), int(patch_w * 14)), mode="bilinear", align_corners=True)
        out = self.scratch.output_conv2(out)

        return out
//  [python]
def test_nondeterministic_alert_interpolate_linear(self, device):
        input = torch.randn(1, 2, 4, device=device, requires_grad=True)
        res = t.nn.functional.interpolate(
            input,
            size=12,
            mode='linear',
            align_corners=False)
        grad = torch.ones_like(res)

        @expectedAlertNondeterministic('upsample_linear1d_backward_out_cuda', 'cuda')
        def backward_func(slf, device):
            res.backward(grad)

        backward_func(self, device)
//  [python]
def test_nondeterministic_alert_interpolate_linear(self, device):
        input = torch.randn(1, 2, 4, device=device, requires_grad=True)
        res = torch.nn.functional.interpolate(
            input,
            size=12,
            mode='linear',
            align_corners=False)
        grad = torch.ones_like(res)

        @expectedAlertNondeterministic('upsample_linear1d_backward_out_cuda', 'cuda')
        def backward_func(slf, device):
            res.backward(grad)

        backward_func(self, device)
//  [python]
def test_nondeterministic_alert_interpolate_bicubic(self, device):
        input = torch.randn(1, 2, 4, 4, device=device, requires_grad=True)
        res = t.interpolate(
            input,
            size=12,
            mode='bicubic',
            align_corners=False)
        grad = torch.ones_like(res)

        @expectedAlertNondeterministic('upsample_bicubic2d_backward_out_cuda', 'cuda')
        def backward_func(slf, device):
            res.backward(grad)

        backward_func(self, device)
//  [python]
def test_nondeterministic_alert_interpolate_bicubic(self, device):
        input = torch.randn(1, 2, 4, 4, device=device, requires_grad=True)
        res = torch.nn.functional.interpolate(
            input,
            size=12,
            mode='bicubic',
            align_corners=False)
        grad = torch.ones_like(res)

        @expectedAlertNondeterministic('upsample_bicubic2d_backward_out_cuda', 'cuda')
        def backward_func(slf, device):
            res.backward(grad)

        backward_func(self, device)
//  [python]
def test_nondeterministic_alert_interpolate_bilinear(self, device):
        input = torch.randn(1, 2, 4, 4, device=device, requires_grad=True)
        res = t.interpolate(
            input,
            size=12,
            mode='bilinear',
            align_corners=False)
        grad = torch.ones_like(res)

        @expectedAlertNondeterministic('upsample_bilinear2d_backward_out_cuda', 'cuda')
        def backward_func(slf, device):
            res.backward(grad)

        backward_func(self, device)
//  [python]
def test_nondeterministic_alert_interpolate_bilinear(self, device):
        input = torch.randn(1, 2, 4, 4, device=device, requires_grad=True)
        res = torch.nn.functional.interpolate(
            input,
            size=12,
            mode='bilinear',
            align_corners=False)
        grad = torch.ones_like(res)

        @expectedAlertNondeterministic('upsample_bilinear2d_backward_out_cuda', 'cuda')
        def backward_func(slf, device):
            res.backward(grad)

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

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

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

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

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

            semantic_segmentation = []

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

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

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

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

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

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

            semantic_segmentation = []

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

        return semantic_segmentation
//  [python]
def test_nondeterministic_alert_interpolate_trilinear(self, device):
        input = torch.randn(1, 2, 4, 4, 4, device=device, requires_grad=True)
        res = t.upsample_trilinear3d(
            input,
            size=12,
            align_corners=False)
        grad = torch.ones_like(res)

        @expectedAlertNondeterministic('upsample_trilinear3d_backward_out_cuda', 'cuda')
        def backward_func(slf, device):
            res.backward(grad)

        backward_func(self, device)
//  [python]
def test_nondeterministic_alert_interpolate_trilinear(self, device):
        input = torch.randn(1, 2, 4, 4, 4, device=device, requires_grad=True)
        res = torch.nn.functional.interpolate(
            input,
            size=12,
            mode='trilinear',
            align_corners=False)
        grad = torch.ones_like(res)

        @expectedAlertNondeterministic('upsample_trilinear3d_backward_out_cuda', 'cuda')
        def backward_func(slf, device):
            res.backward(grad)

        backward_func(self, device)
//  [python]
def forward(self, x: NestedTensor):
        """
        :param x: NestedTensor
        :return: list containing feature descriptors at different spatial resolution
                0: [2N, 3, H, W]
                1: [2N, C0, H//4, W//4]
                2: [2N, C1, H//8, W//8]
                3: [2N, C2, H//16, W//16]
        """
        _, _, h, w = x.left.shape

        src_stereo = torch.cat([x.left, x.right], dim=0)  # 2NxCxHxW

        # in conv
        output = self.in_conv(src_stereo)  # 1/2

        # res blocks
        output_1 = self.resblock_1(output)  # 1/4
        output_2 = self.resblock_2(output_1)  # 1/8

        # spp
        h_spp, w_spp = math.ceil(h / 16), math.ceil(w / 16)
        spp_1 = self.branch1(output_2)
        spp_1 = t.upsample(spp_1, size=(h_spp, w_spp), mode='bilinear')
        spp_2 = self.branch2(output_2)
        spp_2 = t.upsample(spp_2, size=(h_spp, w_spp), mode='bilinear')
        spp_3 = self.branch3(output_2)
        spp_3 = t.upsample(spp_3, size=(h_spp, w_spp), mode='bilinear')
        spp_4 = self.branch4(output_2)
        spp_4 = t.upsample(spp_4, size=(h_spp, w_spp), mode='bilinear')
        output_3 = torch.cat([spp_1, spp_2, spp_3, spp_4], dim=1)  # 1/16

        return [src_stereo, output_1, output_2, output_3]
//  [python]
def forward(self, x: NestedTensor):
        """
        :param x: NestedTensor
        :return: list containing feature descriptors at different spatial resolution
                0: [2N, 3, H, W]
                1: [2N, C0, H//4, W//4]
                2: [2N, C1, H//8, W//8]
                3: [2N, C2, H//16, W//16]
        """
        _, _, h, w = x.left.shape

        src_stereo = torch.cat([x.left, x.right], dim=0)  # 2NxCxHxW

        # in conv
        output = self.in_conv(src_stereo)  # 1/2

        # res blocks
        output_1 = self.resblock_1(output)  # 1/4
        output_2 = self.resblock_2(output_1)  # 1/8

        # spp
        h_spp, w_spp = math.ceil(h / 16), math.ceil(w / 16)
        spp_1 = self.branch1(output_2)
        spp_1 = F.interpolate(spp_1, size=(h_spp, w_spp), mode='bilinear', align_corners=False)
        spp_2 = self.branch2(output_2)
        spp_2 = F.interpolate(spp_2, size=(h_spp, w_spp), mode='bilinear', align_corners=False)
        spp_3 = self.branch3(output_2)
        spp_3 = F.interpolate(spp_3, size=(h_spp, w_spp), mode='bilinear', align_corners=False)
        spp_4 = self.branch4(output_2)
        spp_4 = F.interpolate(spp_4, size=(h_spp, w_spp), mode='bilinear', align_corners=False)
        output_3 = torch.cat([spp_1, spp_2, spp_3, spp_4], dim=1)  # 1/16

        return [src_stereo, output_1, output_2, output_3]
//  [python]
def forward(self, x,target):
        x1=self.conv1(torch.cat([x,target],1))
        x1 = self.trans_block1(x1)
        x2=self.conv2(x1)
        x2 = self.trans_block2(x2)
        x3=self.conv3(x2)
        x3 = self.trans_block3(x3)
        x4=self.conv3(x3)
        x4 = self.trans_block4(x4)
        #print(x4.size())
        residual = self.sig(self.conv_refin(self.sig(F.adaptive_avg_pool2d(x4,(1,1)))))
        #print(residual)
        residual = F.upsample(residual, size=x.size()[2:])
        #print(residual.size())
        return residual
//  [python]
def forward(self, x,target):
        x1=self.conv1(torch.cat([x,target],1))
        x1 = self.trans_block1(x1)
        x2=self.conv2(x1)
        x2 = self.trans_block2(x2)
        x3=self.conv3(x2)
        x3 = self.trans_block3(x3)
        x4=self.conv3(x3)
        x4 = self.trans_block4(x4)
        #print(x4.size())
        residual = self.sig(self.conv_refin(self.sig(F.adaptive_avg_pool2d(x4,(1,1)))))
        #print(residual)
        residual = F.interpolate(residual, size=x.size()[2:])
        #print(residual.size())
        return residual
//  [python]
def prepare_masked_condition(
        self,
        image,
        batch_size,
        num_channels_latents,
        num_frames,
        height,
        width,
        dtype,
        device,
        generator,
        motion_scale=0,
    ):
        shape = (
            batch_size,
            num_channels_latents,
            num_frames,
            height // self.vae_scale_factor,
            width // self.vae_scale_factor,
        )
        _, _, _, scaled_height, scaled_width = shape

        image = self.image_processor.preprocess(image)
        image = image.to(device, dtype)

        if isinstance(generator, list):
            image_latent = [
                self.vae.encode(image[k : k + 1]).latent_dist.sample(generator[k]) for k in range(batch_size)
            ]
            image_latent = torch.cat(image_latent, dim=0)
        else:
            image_latent = self.vae.encode(image).latent_dist.sample(generator)

        image_latent = image_latent.to(device=device, dtype=dtype)
        image_latent = t.nn.functional.interpolate(image_latent, size=[scaled_height, scaled_width])
        image_latent_padding = image_latent.clone() * self.vae.config.scaling_factor

        mask = torch.zeros((batch_size, 1, num_frames, scaled_height, scaled_width)).to(device=device, dtype=dtype)
        mask_coef = prepare_mask_coef_by_statistics(num_frames, 0, motion_scale)
        masked_image = torch.zeros(batch_size, 4, num_frames, scaled_height, scaled_width).to(
            device=device, dtype=self.unet.dtype
        )
        for f in range(num_frames):
            mask[:, :, f, :, :] = mask_coef[f]
            masked_image[:, :, f, :, :] = image_latent_padding.clone()

        mask = torch.cat([mask] * 2) if self.do_classifier_free_guidance else mask
        masked_image = torch.cat([masked_image] * 2) if self.do_classifier_free_guidance else masked_image

        return mask, masked_image
//  [python]
def prepare_masked_condition(
        self,
        image,
        batch_size,
        num_channels_latents,
        num_frames,
        height,
        width,
        dtype,
        device,
        generator,
        motion_scale=0,
    ):
        shape = (
            batch_size,
            num_channels_latents,
            num_frames,
            height // self.vae_scale_factor,
            width // self.vae_scale_factor,
        )
        _, _, _, scaled_height, scaled_width = shape

        image = self.image_processor.preprocess(image)
        image = image.to(device, dtype)

        if isinstance(generator, list):
            image_latent = [
                self.vae.encode(image[k : k + 1]).latent_dist.sample(generator[k]) for k in range(batch_size)
            ]
            image_latent = torch.cat(image_latent, dim=0)
        else:
            image_latent = self.vae.encode(image).latent_dist.sample(generator)

        image_latent = image_latent.to(device=device, dtype=dtype)
        image_latent = torch.nn.functional.interpolate(image_latent, size=[scaled_height, scaled_width])
        image_latent_padding = image_latent.clone() * self.vae.config.scaling_factor

        mask = torch.zeros((batch_size, 1, num_frames, scaled_height, scaled_width)).to(device=device, dtype=dtype)
        mask_coef = prepare_mask_coef_by_statistics(num_frames, 0, motion_scale)
        masked_image = torch.zeros(batch_size, 4, num_frames, scaled_height, scaled_width).to(
            device=device, dtype=self.unet.dtype
        )
        for f in range(num_frames):
            mask[:, :, f, :, :] = mask_coef[f]
            masked_image[:, :, f, :, :] = image_latent_padding.clone()

        mask = torch.cat([mask] * 2) if self.do_classifier_free_guidance else mask
        masked_image = torch.cat([masked_image] * 2) if self.do_classifier_free_guidance else masked_image

        return mask, masked_image
//  [python]
def forward(self, pixel_values):
        batch_size = pixel_values.shape[0]

        # create patch embeddings
        embeddings, num_frames, patch_width = self.patch_embeddings(pixel_values)

        cls_tokens = self.cls_token.expand(embeddings.size(0), -1, -1)
        embeddings = torch.cat((cls_tokens, embeddings), dim=1)

        # resizing the positional embeddings in case they don't match the input at inference
        if embeddings.size(1) != self.position_embeddings.size(1):
            position_embeddings = self.position_embeddings
            cls_pos_embed = position_embeddings[0, 0, :].unsqueeze(0).unsqueeze(1)
            other_pos_embed = position_embeddings[0, 1:, :].unsqueeze(0).transpose(1, 2)
            patch_num = int(other_pos_embed.size(2) ** 0.5)
            patch_height = embeddings.size(1) // patch_width
            other_pos_embed = other_pos_embed.reshape(1, embeddings.size(2), patch_num, patch_num)
            new_pos_embed = nn.functional.upsample(
                other_pos_embed, size=(patch_height, patch_width), mode="nearest"
            )
            new_pos_embed = new_pos_embed.flatten(2)
            new_pos_embed = new_pos_embed.transpose(1, 2)
            new_pos_embed = torch.cat((cls_pos_embed, new_pos_embed), 1)
            embeddings = embeddings + new_pos_embed
        else:
            embeddings = embeddings + self.position_embeddings
        embeddings = self.pos_drop(embeddings)

        # Time Embeddings
        if self.attention_type != "space_only":
            cls_tokens = embeddings[:batch_size, 0, :].unsqueeze(1)
            embeddings = embeddings[:, 1:]
            _, patch_height, patch_width = embeddings.shape
            embeddings = (
                embeddings.reshape(batch_size, num_frames, patch_height, patch_width)
                .permute(0, 2, 1, 3)
                .reshape(batch_size * patch_height, num_frames, patch_width)
            )
            # Resizing time embeddings in case they don't match
            if num_frames != self.time_embeddings.size(1):
                time_embeddings = self.time_embeddings.transpose(1, 2)
                new_time_embeddings = nn.functional.upsample(time_embeddings, size=(num_frames), mode="nearest")
                new_time_embeddings = new_time_embeddings.transpose(1, 2)
                embeddings = embeddings + new_time_embeddings
            else:
                embeddings = embeddings + self.time_embeddings
            embeddings = self.time_drop(embeddings)
            embeddings = embeddings.view(batch_size, patch_height, num_frames, patch_width).reshape(
                batch_size, patch_height * num_frames, patch_width
            )
            embeddings = torch.cat((cls_tokens, embeddings), dim=1)

        return embeddings
//  [python]
def forward(self, pixel_values):
        batch_size = pixel_values.shape[0]

        # create patch embeddings
        embeddings, num_frames, patch_width = self.patch_embeddings(pixel_values)

        cls_tokens = self.cls_token.expand(embeddings.size(0), -1, -1)
        embeddings = torch.cat((cls_tokens, embeddings), dim=1)

        # resizing the positional embeddings in case they don't match the input at inference
        if embeddings.size(1) != self.position_embeddings.size(1):
            position_embeddings = self.position_embeddings
            cls_pos_embed = position_embeddings[0, 0, :].unsqueeze(0).unsqueeze(1)
            other_pos_embed = position_embeddings[0, 1:, :].unsqueeze(0).transpose(1, 2)
            patch_num = int(other_pos_embed.size(2) ** 0.5)
            patch_height = embeddings.size(1) // patch_width
            other_pos_embed = other_pos_embed.reshape(1, embeddings.size(2), patch_num, patch_num)
            new_pos_embed = nn.functional.interpolate(
                other_pos_embed, size=(patch_height, patch_width), mode="nearest"
            )
            new_pos_embed = new_pos_embed.flatten(2)
            new_pos_embed = new_pos_embed.transpose(1, 2)
            new_pos_embed = torch.cat((cls_pos_embed, new_pos_embed), 1)
            embeddings = embeddings + new_pos_embed
        else:
            embeddings = embeddings + self.position_embeddings
        embeddings = self.pos_drop(embeddings)

        # Time Embeddings
        if self.attention_type != "space_only":
            cls_tokens = embeddings[:batch_size, 0, :].unsqueeze(1)
            embeddings = embeddings[:, 1:]
            _, patch_height, patch_width = embeddings.shape
            embeddings = (
                embeddings.reshape(batch_size, num_frames, patch_height, patch_width)
                .permute(0, 2, 1, 3)
                .reshape(batch_size * patch_height, num_frames, patch_width)
            )
            # Resizing time embeddings in case they don't match
            if num_frames != self.time_embeddings.size(1):
                time_embeddings = self.time_embeddings.transpose(1, 2)
                new_time_embeddings = nn.functional.interpolate(time_embeddings, size=(num_frames), mode="nearest")
                new_time_embeddings = new_time_embeddings.transpose(1, 2)
                embeddings = embeddings + new_time_embeddings
            else:
                embeddings = embeddings + self.time_embeddings
            embeddings = self.time_drop(embeddings)
            embeddings = embeddings.view(batch_size, patch_height, num_frames, patch_width).reshape(
                batch_size, patch_height * num_frames, patch_width
            )
            embeddings = torch.cat((cls_tokens, embeddings), dim=1)

        return embeddings
//  [python]
def resize_and_crop_tensor(samples: torch.Tensor, new_width: int, new_height: int) -> torch.Tensor:
        orig_height, orig_width = samples.shape[2], samples.shape[3]

        # Check if resizing is needed
        if orig_height != new_height or orig_width != new_width:
            ratio = max(new_height / orig_height, new_width / orig_width)
            resized_width = int(orig_width * ratio)
            resized_height = int(orig_height * ratio)

            # Resize
            samples = t.interpolate(
                samples, size=(resized_height, resized_width), mode="bilinear", align_corners=False
            )

            # Center Crop
            start_x = (resized_width - new_width) // 2
            end_x = start_x + new_width
            start_y = (resized_height - new_height) // 2
            end_y = start_y + new_height
            samples = samples[:, :, start_y:end_y, start_x:end_x]

        return samples
//  [python]
def resize_and_crop_tensor(samples: torch.Tensor, new_width: int, new_height: int) -> torch.Tensor:
        orig_height, orig_width = samples.shape[2], samples.shape[3]

        # Check if resizing is needed
        if orig_height != new_height or orig_width != new_width:
            ratio = max(new_height / orig_height, new_width / orig_width)
            resized_width = int(orig_width * ratio)
            resized_height = int(orig_height * ratio)

            # Resize
            samples = F.interpolate(
                samples, size=(resized_height, resized_width), mode="bilinear", align_corners=False
            )

            # Center Crop
            start_x = (resized_width - new_width) // 2
            end_x = start_x + new_width
            start_y = (resized_height - new_height) // 2
            end_y = start_y + new_height
            samples = samples[:, :, start_y:end_y, start_x:end_x]

        return samples
//  [python]
def forward(self, x, P2=None):
        #residual = x
        #x3, x4, x5, x6 = self.backbone(x)
        x3, x4 = self.backbone.forward1(x)
        x4a = 0.1 * x4 + 0.9 * x4.detach()
        x4a = self.bt(x4a)
        dht = self.dht(x4a)
        dht = self.dht_backbone(dht)

        x5 = self.backbone.forward2(x4)
        
        x5a = x5.clone()
        h, w = x5a.shape[2], x5a.shape[3]
        dht = dht.expand(-1, -1, h, w)
        x5a = torch.cat([x5a, dht], 1)
        x5a = self.cord(x5a)

        x6 = self.backbone.forward3(x5a)
        outs = {}
        #x6 = F.relu(x6 + self.fullencoders(x6))
        x = self.up0(x6, x5, P2=P2, scale=32)
        x = self.up1(x, x4, P2=P2, scale=16)
        outs['scale_8'] = self.out_scale_8(x)
        x = self.up2(x, x3)
        outs['scale_4'] = self.out_scale_4(x)
        #x = F.upsample_bilinear(x, scale_factor=4)
        x = F.upsample_bilinear(x, scale_factor=4)
        #x = self.up3(x)
        #x = self.up4(x, residual)
        #x = torch.cat([x, residual], dim=1)
        outs['scale_1'] = self.outc(x)
        return outs
//  [python]
def forward(self, x, P2=None):
        #residual = x
        #x3, x4, x5, x6 = self.backbone(x)
        x3, x4 = self.backbone.forward1(x)
        x4a = 0.1 * x4 + 0.9 * x4.detach()
        x4a = self.bt(x4a)
        dht = self.dht(x4a)
        dht = self.dht_backbone(dht)

        x5 = self.backbone.forward2(x4)
        
        x5a = x5.clone()
        h, w = x5a.shape[2], x5a.shape[3]
        dht = dht.expand(-1, -1, h, w)
        x5a = torch.cat([x5a, dht], 1)
        x5a = self.cord(x5a)

        x6 = self.backbone.forward3(x5a)
        outs = {}
        #x6 = F.relu(x6 + self.fullencoders(x6))
        x = self.up0(x6, x5, P2=P2, scale=32)
        x = self.up1(x, x4, P2=P2, scale=16)
        outs['scale_8'] = self.out_scale_8(x)
        x = self.up2(x, x3)
        outs['scale_4'] = self.out_scale_4(x)
        #x = F.upsample_bilinear(x, scale_factor=4)
        x = F.interpolate(x, scale_factor=4, align_corners=True, mode='bilinear')
        #x = self.up3(x)
        #x = self.up4(x, residual)
        #x = torch.cat([x, residual], dim=1)
        outs['scale_1'] = self.outc(x)
        return outs
//  [python]
def test_upsample():
    """test_upsample"""

    class Upsample(Module):
        def __init__(self, size=None, scale=None, mode="nearest", align_corners=None):
            super().__init__()
            self.size = size
            self.scale = scale
            self.mode = mode
            self.align_corners = align_corners

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

    inp = torch.rand((1, 3, 32, 32))
    verify_model(Upsample(size=(64, 64), mode="nearest"), inp)
    verify_model(Upsample(scale=2, mode="nearest"), inp)
    verify_model(Upsample(size=(50, 50), mode="nearest"), inp)
    verify_model(Upsample(size=(64, 64), mode="bilinear", align_corners=True), inp)
    verify_model(Upsample(scale=2, mode="bilinear", align_corners=True), inp)
    verify_model(Upsample(size=(50, 50), mode="bilinear", align_corners=True), inp)
    verify_model(Upsample(size=(64, 64), mode="bicubic", align_corners=True), inp)
    verify_model(Upsample(scale=2, mode="bicubic", align_corners=True), inp)
    verify_model(Upsample(size=(50, 50), mode="bicubic", align_corners=True), inp)
//  [python]
def test_upsample():
    """test_upsample"""

    class Upsample(Module):
        def __init__(self, size=None, scale=None, mode="nearest", align_corners=None):
            super().__init__()
            self.size = size
            self.scale = scale
            self.mode = mode
            self.align_corners = align_corners

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

    inp = torch.rand((1, 3, 32, 32))
    verify_model(Upsample(size=(64, 64), mode="nearest"), inp)
    verify_model(Upsample(scale=2, mode="nearest"), inp)
    verify_model(Upsample(size=(50, 50), mode="nearest"), inp)
    verify_model(Upsample(size=(64, 64), mode="bilinear", align_corners=True), inp)
    verify_model(Upsample(scale=2, mode="bilinear", align_corners=True), inp)
    verify_model(Upsample(size=(50, 50), mode="bilinear", align_corners=True), inp)
    verify_model(Upsample(size=(64, 64), mode="bicubic", align_corners=True), inp)
    verify_model(Upsample(scale=2, mode="bicubic", align_corners=True), inp)
    verify_model(Upsample(size=(50, 50), mode="bicubic", align_corners=True), inp)
//  [python]
def prepare_mask_latents(self, mask_image, batch_size, height, width, dtype, device, do_classifier_free_guidance):
        # resize the mask to latents shape as we concatenate the mask to the latents
        # we do that before converting to dtype to avoid breaking in case we're using cpu_offload
        # and half precision
        mask_image = t.interpolate(mask_image, size=(height // self.vae_scale_factor, width // self.vae_scale_factor))
        mask_image = mask_image.to(device=device, dtype=dtype)

        # duplicate mask for each generation per prompt, using mps friendly method
        if mask_image.shape[0] < batch_size:
            if not batch_size % mask_image.shape[0] == 0:
                raise ValueError(
                    "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"
                    f" a total batch size of {batch_size}, but {mask_image.shape[0]} masks were passed. Make sure the number"
                    " of masks that you pass is divisible by the total requested batch size."
                )
            mask_image = mask_image.repeat(batch_size // mask_image.shape[0], 1, 1, 1)

        mask_image = torch.cat([mask_image] * 2) if do_classifier_free_guidance else mask_image

        mask_image_latents = mask_image

        return mask_image_latents
//  [python]
def prepare_mask_latents(self, mask_image, batch_size, height, width, dtype, device, do_classifier_free_guidance):
        # resize the mask to latents shape as we concatenate the mask to the latents
        # we do that before converting to dtype to avoid breaking in case we're using cpu_offload
        # and half precision
        mask_image = F.interpolate(mask_image, size=(height // self.vae_scale_factor, width // self.vae_scale_factor))
        mask_image = mask_image.to(device=device, dtype=dtype)

        # duplicate mask for each generation per prompt, using mps friendly method
        if mask_image.shape[0] < batch_size:
            if not batch_size % mask_image.shape[0] == 0:
                raise ValueError(
                    "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"
                    f" a total batch size of {batch_size}, but {mask_image.shape[0]} masks were passed. Make sure the number"
                    " of masks that you pass is divisible by the total requested batch size."
                )
            mask_image = mask_image.repeat(batch_size // mask_image.shape[0], 1, 1, 1)

        mask_image = torch.cat([mask_image] * 2) if do_classifier_free_guidance else mask_image

        mask_image_latents = mask_image

        return mask_image_latents
//  [python]
def forward(self, x):
        x_in = x
        if self.upsample:
            x_in = torch.nn.functional.upsample(x_in, scale_factor=self.upsample, mode='nearest')
        out = self.reflection_pad(x_in)
        out = self.conv2d(out)
        return out
//  [python]
def forward(self, x):
        x_in = x
        if self.upsample:
            x_in = torch.nn.functional.interpolate(x_in, mode='nearest', scale_factor=self.upsample)
        out = self.reflection_pad(x_in)
        out = self.conv2d(out)
        return out
//  [python]
def forward(self, ref, supp):
        """
        Args:
            ref (Tensor): Reference image with shape of (b, 3, h, w).
            supp: The supporting image to be warped: (b, 3, h, w).

        Returns:
            Tensor: Estimated optical flow: (b, 2, h, w).
        """

        num_batches, _, h, w = ref.size()
        ref = [ref]
        supp = [supp]

        # generate downsampled frames
        for _ in range(3):
            ref.insert(
                0,
                F.avg_pool2d(
                    input=ref[0],
                    kernel_size=2,
                    stride=2,
                    count_include_pad=False))
            supp.insert(
                0,
                F.avg_pool2d(
                    input=supp[0],
                    kernel_size=2,
                    stride=2,
                    count_include_pad=False))

        # flow computation
        flow = ref[0].new_zeros(num_batches, 2, h // 16, w // 16)
        for i in range(4):
            flow_up = t.upsample(
                input=flow,
                scale_factor=2,
                mode='bilinear') * 2.0
            flow = flow_up + self.basic_module[i](
                torch.cat([
                    ref[i],
                    flow_warp(
                        supp[i],
                        flow_up.permute(0, 2, 3, 1),
                        padding_mode='border'), flow_up
                ], 1))

        return flow
//  [python]
def forward(self, ref, supp):
        """
        Args:
            ref (Tensor): Reference image with shape of (b, 3, h, w).
            supp: The supporting image to be warped: (b, 3, h, w).

        Returns:
            Tensor: Estimated optical flow: (b, 2, h, w).
        """

        num_batches, _, h, w = ref.size()
        ref = [ref]
        supp = [supp]

        # generate downsampled frames
        for _ in range(3):
            ref.insert(
                0,
                F.avg_pool2d(
                    input=ref[0],
                    kernel_size=2,
                    stride=2,
                    count_include_pad=False))
            supp.insert(
                0,
                F.avg_pool2d(
                    input=supp[0],
                    kernel_size=2,
                    stride=2,
                    count_include_pad=False))

        # flow computation
        flow = ref[0].new_zeros(num_batches, 2, h // 16, w // 16)
        for i in range(4):
            flow_up = F.interpolate(
                input=flow,
                scale_factor=2,
                mode='bilinear',
                align_corners=True) * 2.0
            flow = flow_up + self.basic_module[i](
                torch.cat([
                    ref[i],
                    flow_warp(
                        supp[i],
                        flow_up.permute(0, 2, 3, 1),
                        padding_mode='border'), flow_up
                ], 1))

        return flow
//  [python]
def resize_pos_embed(posemb, posemb_new):
    # Rescale the grid of position embeddings when loading from state_dict. Adapted from
    # https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224
    _logger.info('Resized position embedding: %s to %s', posemb.shape, posemb_new.shape)
    ntok_new = posemb_new.shape[1]
    if True:
        posemb_tok, posemb_grid = posemb[:, :1], posemb[0, 1:]
        ntok_new -= 1
    else:
        posemb_tok, posemb_grid = posemb[:, :0], posemb[0]
    gs_old = int(math.sqrt(len(posemb_grid)))
    gs_new = int(math.sqrt(ntok_new))
    _logger.info('Position embedding grid-size from %s to %s', gs_old, gs_new)
    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
    posemb_grid = t.interpolate(posemb_grid, size=(gs_new, gs_new), mode='bilinear')
    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_new * gs_new, -1)
    posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
    return posemb
//  [python]
def resize_pos_embed(posemb, posemb_new):
    # Rescale the grid of position embeddings when loading from state_dict. Adapted from
    # https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224
    _logger.info('Resized position embedding: %s to %s', posemb.shape, posemb_new.shape)
    ntok_new = posemb_new.shape[1]
    if True:
        posemb_tok, posemb_grid = posemb[:, :1], posemb[0, 1:]
        ntok_new -= 1
    else:
        posemb_tok, posemb_grid = posemb[:, :0], posemb[0]
    gs_old = int(math.sqrt(len(posemb_grid)))
    gs_new = int(math.sqrt(ntok_new))
    _logger.info('Position embedding grid-size from %s to %s', gs_old, gs_new)
    posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
    posemb_grid = F.interpolate(posemb_grid, size=(gs_new, gs_new), mode='bilinear')
    posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_new * gs_new, -1)
    posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
    return posemb
//  [python]
def forward(self, x):
        H, W = x.size()[2:]

        x = self.init_block(x)

        x1 = self.block1_1(x)
        x1 = self.block1_2(x1)
        x1 = self.block1_3(x1)

        x2 = self.block2_1(x1)
        x2 = self.block2_2(x2)
        x2 = self.block2_3(x2)
        x2 = self.block2_4(x2)

        x3 = self.block3_1(x2)
        x3 = self.block3_2(x3)
        x3 = self.block3_3(x3)
        x3 = self.block3_4(x3)

        x4 = self.block4_1(x3)
        x4 = self.block4_2(x4)
        x4 = self.block4_3(x4)
        x4 = self.block4_4(x4)

        x_fuses = []
        if self.sa and self.dil is not None:
            for i, xi in enumerate([x1, x2, x3, x4]):
                x_fuses.append(self.attentions[i](self.dilations[i](xi)))
        elif self.sa:
            for i, xi in enumerate([x1, x2, x3, x4]):
                x_fuses.append(self.attentions[i](xi))
        elif self.dil is not None:
            for i, xi in enumerate([x1, x2, x3, x4]):
                x_fuses.append(self.dilations[i](xi))
        else:
            x_fuses = [x1, x2, x3, x4]

        e1 = self.conv_reduces[0](x_fuses[0])
        e1 = t.upsample(e1, size=(H, W), mode="bilinear")

        e2 = self.conv_reduces[1](x_fuses[1])
        e2 = t.upsample(e2, size=(H, W), mode="bilinear")

        e3 = self.conv_reduces[2](x_fuses[2])
        e3 = t.upsample(e3, size=(H, W), mode="bilinear")

        e4 = self.conv_reduces[3](x_fuses[3])
        e4 = t.upsample(e4, size=(H, W), mode="bilinear")

        outputs = [e1, e2, e3, e4]

        output = self.classifier(torch.cat(outputs, dim=1))
        # if not self.training:
        #    return torch.sigmoid(output)

        outputs.append(output)
        outputs = [torch.sigmoid(r) for r in outputs]
        return outputs