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

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

        def _get_torch_upsample_prediction(
            x, scale_factor=(2, 2), align_corners=False, recompute_scale_factor=True
        ):
            x = torch.from_numpy(x)
            out = torch.nn.functional.interpolate(
                x,
                scale_factor=scale_factor,
                mode="bilinear",
                align_corners=align_corners,
                recompute_scale_factor=recompute_scale_factor,
            )
            return out.numpy()

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

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

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

        expected_output_type = torch_pred.shape + (types.fp32,)
        run_compare_builder(
            build_upsample,
            input_placeholder_dict,
            input_value_dict,
            expected_output_type,
            torch_pred,
            compute_unit=compute_unit,
            backend=backend,
            rtol=0.5,
        )
//  [python]
def forward(self, x_):
        h = x_.size(2)
        w = x_.size(3)
        x0, x1, x2, x3, x4 = self.backbone(x_)
        up4 = F.upsample(x4, size=(x3.size(2), x3.size(3)), mode="bilinear")
        x = torch.cat([up4, x3], dim=1)       
        x = self.dconv_up4(x)
        up3 = F.upsample(x, size=(x2.size(2), x2.size(3)), mode="bilinear")
        x = torch.cat([up3, x2], dim=1)       
        x = self.dconv_up3(x)
        up2 = F.upsample(x, size=(x1.size(2), x1.size(3)), mode="bilinear")
        x = torch.cat([up2, x1], dim=1)       
        x = self.dconv_up2(x)
        up1 = F.upsample(x, size=(x0.size(2), x0.size(3)), mode="bilinear")
        x = torch.cat([up1, x0], dim=1)       
        x = self.dconv_up1(x)
        x = F.upsample(x, size=(x_.size(2), x_.size(3)), mode="bilinear")
        out = self.APNB(x)
#         out = self.conv_last(x)
 
        ## side output
        d1 = F.upsample(self.dsn1(x0), size=(h,w), mode='bilinear')
        d2 = F.upsample(self.dsn2(x1), size=(h,w), mode='bilinear')
        d3 = F.upsample(self.dsn3(x2), size=(h,w), mode='bilinear')
        d4 = F.upsample(self.dsn4(x3), size=(h,w), mode='bilinear')
        d5 = F.upsample(self.dsn5(x4), size=(h,w), mode='bilinear')

        ###########sigmoid ver
        d1_out = F.sigmoid(d1)
        d2_out = F.sigmoid(d2)
        d3_out = F.sigmoid(d3)
        d4_out = F.sigmoid(d4)
        d5_out = F.sigmoid(d5)

        concat = torch.cat((d1_out, d2_out, d3_out, d4_out, d5_out), 1)

              
        fuse_box = self.fuse(concat)
        GAP = F.adaptive_avg_pool2d(fuse_box,(1,1))
        GAP = GAP.view(-1, 64)        
        se_like = self.SE_mimic(GAP)        
        se_like = torch.unsqueeze(se_like, 2)
        se_like = torch.unsqueeze(se_like, 3)

        feat_se = concat * se_like.expand_as(concat)
        boundary = self.final_boundary(feat_se)
        boundary_out = torch.unsqueeze(boundary[:,1,:,:],1)        
        bd_sftmax = F.softmax(boundary, dim=1)
        boundary_scale = torch.unsqueeze(bd_sftmax[:,1,:,:],1)        
        
        feat_concat = torch.cat( [out, fuse_box], 1)
        feat_concat_conv = self.final_conv(feat_concat)
        mask = self.final_mask(feat_concat_conv)
        mask_sftmax = F.softmax(mask,dim=1)        
        mask_scale = torch.unsqueeze(mask_sftmax[:,1,:,:],1)

        if self.mode == 'Train':
            scalefactor = torch.clamp(mask_scale+boundary_scale,0,1)            
        elif self.mode == 'Infer':
            scalefactor = torch.clamp(mask_scale+5*boundary_scale,0,1)
        
        
        mask_out = torch.unsqueeze(mask[:,1,:,:],1)
        relu = self.relu(mask_out)        
        scalar = relu.cpu().detach().numpy()
        if np.sum(scalar) == 0:
            average = 0
        else : 
            average = scalar[np.nonzero(scalar)].mean()
        mask_out = mask_out-relu + (average*scalefactor)
        
        if self.mode == 'Train':
            mask_out = F.sigmoid(mask_out)
            boundary_out = F.sigmoid(boundary_out)

            return d1_out, d2_out, d3_out, d4_out, d5_out, boundary_out, mask_out
        elif self.mode =='Infer':
            return mask_out
//  [python]
def forward(self, feature, shape):
        x0 = feature[0]
        x1 = feature[1]
        x2 = feature[2]
        x3 = feature[3]
        x4 = feature[4]
        x = self.up1(x4, x3)
        if self.training:
            dp3_out_seg = self.out_conv_dp3(Dropout(x, p=0.5))
        else:
            dp3_out_seg = self.out_conv_dp3(x)
        dp3_out_seg = t.interpolate(dp3_out_seg, shape)

        x = self.up2(x, x2)
        if self.training:
            dp2_out_seg = self.out_conv_dp2(FeatureDropout(x))
        else:
            dp2_out_seg = self.out_conv_dp2(x)
        dp2_out_seg = t.interpolate(dp2_out_seg, shape)

        x = self.up3(x, x1)
        if self.training:
            dp1_out_seg = self.out_conv_dp1(self.feature_noise(x))
        else:
            dp1_out_seg = self.out_conv_dp1(x)
        dp1_out_seg = t.interpolate(dp1_out_seg, shape)

        x = self.up4(x, x0)
        dp0_out_seg = self.out_conv(x)
        return dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg
//  [python]
def forward(self, feature, shape):
        x0 = feature[0]
        x1 = feature[1]
        x2 = feature[2]
        x3 = feature[3]
        x4 = feature[4]
        x = self.up1(x4, x3)
        if self.training:
            dp3_out_seg = self.out_conv_dp3(Dropout(x, p=0.5))
        else:
            dp3_out_seg = self.out_conv_dp3(x)
        dp3_out_seg = torch.nn.functional.interpolate(dp3_out_seg, shape)

        x = self.up2(x, x2)
        if self.training:
            dp2_out_seg = self.out_conv_dp2(FeatureDropout(x))
        else:
            dp2_out_seg = self.out_conv_dp2(x)
        dp2_out_seg = torch.nn.functional.interpolate(dp2_out_seg, shape)

        x = self.up3(x, x1)
        if self.training:
            dp1_out_seg = self.out_conv_dp1(self.feature_noise(x))
        else:
            dp1_out_seg = self.out_conv_dp1(x)
        dp1_out_seg = torch.nn.functional.interpolate(dp1_out_seg, shape)

        x = self.up4(x, x0)
        dp0_out_seg = self.out_conv(x)
        return dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg
//  [python]
def _resize_tensor(x: Tensor, size: int = 64) -> Tensor:
    """https://github.com/toshas/torch-fidelity/blob/master/torch_fidelity/sample_similarity_lpips.py#L127C22-L132."""
    if x.shape[-1] > size and x.shape[-2] > size:
        return t.nn.functional.interpolate(x, (size, size), mode="area")
    return t.nn.functional.interpolate(x, (size, size), mode="bilinear", align_corners=False)
//  [python]
def _resize_tensor(x: Tensor, size: int = 64) -> Tensor:
    """https://github.com/toshas/torch-fidelity/blob/master/torch_fidelity/sample_similarity_lpips.py#L127C22-L132."""
    if x.shape[-1] > size and x.shape[-2] > size:
        return torch.nn.functional.interpolate(x, (size, size), mode="area")
    return torch.nn.functional.interpolate(x, (size, size), mode="bilinear", align_corners=False)
//  [python]
def forward(self, conv_out, segSize=None):
        conv5 = conv_out[-1]
        x = self.cbr(conv5)
        x = self.conv_last(x)

        if self.use_softmax: # is True during inference
            x = nn.functional.upsample(
                x, size=segSize, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)
        else:
            x = nn.functional.log_softmax(x, dim=1)

        return x
//  [python]
def forward(self, conv_out, segSize=None):
        conv5 = conv_out[-1]
        x = self.cbr(conv5)
        x = self.conv_last(x)

        if self.use_softmax: # is True during inference
            x = nn.functional.interpolate(
                x, size=segSize, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)
        else:
            x = nn.functional.log_softmax(x, dim=1)

        return x
//  [python]
def forward(self, conv_out, segSize=None):
        conv5 = conv_out[-1]

        input_size = conv5.size()
        ppm_out = [conv5]
        for pool_scale in self.ppm:
            ppm_out.append(t.interpolate(
                pool_scale(conv5),
                (input_size[2], input_size[3]),
                mode='bilinear', align_corners=False))
        ppm_out = torch.cat(ppm_out, 1)

        x = self.conv_last(ppm_out)

        if self.use_softmax:  # is True during inference
            x = t.interpolate(
                x, size=segSize, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)
        else:
            x = nn.functional.log_softmax(x, dim=1)
        return x
//  [python]
def forward(self, conv_out, segSize=None):
        conv5 = conv_out[-1]

        input_size = conv5.size()
        ppm_out = [conv5]
        for pool_scale in self.ppm:
            ppm_out.append(nn.functional.interpolate(
                pool_scale(conv5),
                (input_size[2], input_size[3]),
                mode='bilinear', align_corners=False))
        ppm_out = torch.cat(ppm_out, 1)

        x = self.conv_last(ppm_out)

        if self.use_softmax:  # is True during inference
            x = nn.functional.interpolate(
                x, size=segSize, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)
        else:
            x = nn.functional.log_softmax(x, dim=1)
        return x
//  [python]
def interpolate_image(image, shape, mode='bilinear', align_corners=True):
    """
    Interpolate an image to a different resolution

    Parameters
    ----------
    image : torch.Tensor [B,?,h,w]
        Image to be interpolated
    shape : tuple (H, W)
        Output shape
    mode : str
        Interpolation mode
    align_corners : bool
        True if corners will be aligned after interpolation

    Returns
    -------
    image : torch.Tensor [B,?,H,W]
        Interpolated image
    """
    # Take last two dimensions as shape
    if len(shape) > 2:
        shape = shape[-2:]
    # If the shapes are the same, do nothing
    if same_shape(image.shape[-2:], shape):
        return image
    else:
        # Interpolate image to match the shape
        return t.interpolate(
            image, size=shape, mode=mode, align_corners=align_corners)
//  [python]
def interpolate_image(image, shape, mode='bilinear', align_corners=True):
    """
    Interpolate an image to a different resolution

    Parameters
    ----------
    image : torch.Tensor [B,?,h,w]
        Image to be interpolated
    shape : tuple (H, W)
        Output shape
    mode : str
        Interpolation mode
    align_corners : bool
        True if corners will be aligned after interpolation

    Returns
    -------
    image : torch.Tensor [B,?,H,W]
        Interpolated image
    """
    # Take last two dimensions as shape
    if len(shape) > 2:
        shape = shape[-2:]
    # If the shapes are the same, do nothing
    if same_shape(image.shape[-2:], shape):
        return image
    else:
        # Interpolate image to match the shape
        return funct.interpolate(
            image, size=shape, mode=mode, align_corners=align_corners)
//  [python]
def interpolate_scales(images,
                       shape=None,
                       mode='bilinear',
                       align_corners=False):
    """
    Interpolate list of images to the same shape

    Parameters
    ----------
    images : list of torch.Tensor [B,?,?,?]
        Images to be interpolated, with different resolutions
    shape : tuple (H, W)
        Output shape
    mode : str
        Interpolation mode
    align_corners : bool
        True if corners will be aligned after interpolation

    Returns
    -------
    images : list of torch.Tensor [B,?,H,W]
        Interpolated images, with the same resolution
    """
    # If no shape is provided, interpolate to highest resolution
    if shape is None:
        shape = images[0].shape
    # Take last two dimensions as shape
    if len(shape) > 2:
        shape = shape[-2:]
    # Interpolate all images
    return [
        t.interpolate(
            image, shape, mode=mode, align_corners=align_corners)
        for image in images
    ]
//  [python]
def interpolate_scales(images,
                       shape=None,
                       mode='bilinear',
                       align_corners=False):
    """
    Interpolate list of images to the same shape

    Parameters
    ----------
    images : list of torch.Tensor [B,?,?,?]
        Images to be interpolated, with different resolutions
    shape : tuple (H, W)
        Output shape
    mode : str
        Interpolation mode
    align_corners : bool
        True if corners will be aligned after interpolation

    Returns
    -------
    images : list of torch.Tensor [B,?,H,W]
        Interpolated images, with the same resolution
    """
    # If no shape is provided, interpolate to highest resolution
    if shape is None:
        shape = images[0].shape
    # Take last two dimensions as shape
    if len(shape) > 2:
        shape = shape[-2:]
    # Interpolate all images
    return [
        funct.interpolate(
            image, shape, mode=mode, align_corners=align_corners)
        for image in images
    ]
//  [python]
def predict(self, tensor, imgSizes=(-1,),  # (300, 375, 450, 525, 600)
                segSize=None):
        """Entry-point for segmentation. Use this methods instead of forward
        Arguments:
            tensor {torch.Tensor} -- BCHW
        Keyword Arguments:
            imgSizes {tuple or list} -- imgSizes for segmentation input.
                default: (300, 450)
                original implementation: (300, 375, 450, 525, 600)

        """
        if segSize is None:
            segSize = tensor.shape[-2:]
        segSize = (tensor.shape[2], tensor.shape[3])
        with torch.no_grad():
            if self.use_default_normalization:
                tensor = self.normalize_input(tensor)
            scores = torch.zeros(1, NUM_CLASS, segSize[0], segSize[1]).to(self.device)
            features = torch.zeros(1, self.feature_maps_channels, segSize[0], segSize[1]).to(self.device)

            result = []
            for img_size in imgSizes:
                if img_size != -1:
                    img_data = t.interpolate(tensor.clone(), size=img_size)
                else:
                    img_data = tensor.clone()

                if self.return_feature_maps:
                    pred_current, fmaps = self.forward(img_data, segSize=segSize)
                else:
                    pred_current = self.forward(img_data, segSize=segSize)


                result.append(pred_current)
                scores = scores + pred_current / len(imgSizes)

                # Disclaimer: We use and aggregate only last fmaps: fmaps[3]
                if self.return_feature_maps:
                    features = features + t.interpolate(fmaps[self.return_feature_maps_level], size=segSize) / len(imgSizes)

            _, pred = torch.max(scores, dim=1)

            if self.return_feature_maps:
                return features

            return pred, result
//  [python]
def predict(self, tensor, imgSizes=(-1,),  # (300, 375, 450, 525, 600)
                segSize=None):
        """Entry-point for segmentation. Use this methods instead of forward
        Arguments:
            tensor {torch.Tensor} -- BCHW
        Keyword Arguments:
            imgSizes {tuple or list} -- imgSizes for segmentation input.
                default: (300, 450)
                original implementation: (300, 375, 450, 525, 600)

        """
        if segSize is None:
            segSize = tensor.shape[-2:]
        segSize = (tensor.shape[2], tensor.shape[3])
        with torch.no_grad():
            if self.use_default_normalization:
                tensor = self.normalize_input(tensor)
            scores = torch.zeros(1, NUM_CLASS, segSize[0], segSize[1]).to(self.device)
            features = torch.zeros(1, self.feature_maps_channels, segSize[0], segSize[1]).to(self.device)

            result = []
            for img_size in imgSizes:
                if img_size != -1:
                    img_data = F.interpolate(tensor.clone(), size=img_size)
                else:
                    img_data = tensor.clone()

                if self.return_feature_maps:
                    pred_current, fmaps = self.forward(img_data, segSize=segSize)
                else:
                    pred_current = self.forward(img_data, segSize=segSize)


                result.append(pred_current)
                scores = scores + pred_current / len(imgSizes)

                # Disclaimer: We use and aggregate only last fmaps: fmaps[3]
                if self.return_feature_maps:
                    features = features + F.interpolate(fmaps[self.return_feature_maps_level], size=segSize) / len(imgSizes)

            _, pred = torch.max(scores, dim=1)

            if self.return_feature_maps:
                return features

            return pred, result
//  [python]
def forward(self, x, alpha, steps):
        out = self.initial(x)

        if steps == 0:
            return self.initial_rgb(out)

        for step in range(steps):
            upscaled = t.interpolate(out, scale_factor=2, mode="nearest")
            out = self.prog_blocks[step](upscaled)

        # The number of channels in upscale will stay the same, while
        # out which has moved through prog_blocks might change. To ensure
        # we can convert both to rgb we use different rgb_layers
        # (steps-1) and steps for upscaled, out respectively
        final_upscaled = self.rgb_layers[steps - 1](upscaled)
        final_out = self.rgb_layers[steps](out)
        return self.fade_in(alpha, final_upscaled, final_out)
//  [python]
def forward(self, x, alpha, steps):
        out = self.initial(x)

        if steps == 0:
            return self.initial_rgb(out)

        for step in range(steps):
            upscaled = F.interpolate(out, scale_factor=2, mode="nearest")
            out = self.prog_blocks[step](upscaled)

        # The number of channels in upscale will stay the same, while
        # out which has moved through prog_blocks might change. To ensure
        # we can convert both to rgb we use different rgb_layers
        # (steps-1) and steps for upscaled, out respectively
        final_upscaled = self.rgb_layers[steps - 1](upscaled)
        final_out = self.rgb_layers[steps](out)
        return self.fade_in(alpha, final_upscaled, final_out)
//  [python]
def upflow8(flow, mode='bilinear'):
    new_size = (8 * flow.shape[2], 8 * flow.shape[3])
    return 8 * t.interpolate(
        flow, size=new_size, mode=mode, align_corners=True)
//  [python]
def upflow8(flow, mode='bilinear'):
    new_size = (8 * flow.shape[2], 8 * flow.shape[3])
    return 8 * F.interpolate(
        flow, size=new_size, mode=mode, align_corners=True)
//  [python]
def inpaint(sampler, image, mask, prompt, seed, scale, ddim_steps, num_samples=1, w=512, h=512):
    device = torch.device(
        "cuda") if torch.cuda.is_available() else torch.device("cpu")
    model = sampler.model

    print("Creating invisible watermark encoder (see https://github.com/ShieldMnt/invisible-watermark)...")
    wm = "SDV2"
    wm_encoder = WatermarkEncoder()
    wm_encoder.set_watermark('bytes', wm.encode('utf-8'))

    prng = np.random.RandomState(seed)
    start_code = prng.randn(num_samples, 4, h // 8, w // 8)
    start_code = torch.from_numpy(start_code).to(
        device=device, dtype=torch.float32)

    with torch.no_grad(), \
            torch.autocast("cuda"):
        batch = make_batch_sd(image, mask, txt=prompt,
                              device=device, num_samples=num_samples)

        c = model.cond_stage_model.encode(batch["txt"])

        c_cat = list()
        for ck in model.concat_keys:
            cc = batch[ck].float()
            if ck != model.masked_image_key:
                bchw = [num_samples, 4, h // 8, w // 8]
                cc = t.nn.functional.interpolate(cc, size=bchw[-2:])
            else:
                cc = model.get_first_stage_encoding(
                    model.encode_first_stage(cc))
            c_cat.append(cc)
        c_cat = torch.cat(c_cat, dim=1)

        # cond
        cond = {"c_concat": [c_cat], "c_crossattn": [c]}

        # uncond cond
        uc_cross = model.get_unconditional_conditioning(num_samples, "")
        uc_full = {"c_concat": [c_cat], "c_crossattn": [uc_cross]}

        shape = [model.channels, h // 8, w // 8]
        samples_cfg, intermediates = sampler.sample(
            ddim_steps,
            num_samples,
            shape,
            cond,
            verbose=False,
            eta=1.0,
            unconditional_guidance_scale=scale,
            unconditional_conditioning=uc_full,
            x_T=start_code,
        )
        x_samples_ddim = model.decode_first_stage(samples_cfg)

        result = torch.clamp((x_samples_ddim + 1.0) / 2.0,
                             min=0.0, max=1.0)

        result = result.cpu().numpy().transpose(0, 2, 3, 1) * 255
    return [put_watermark(Image.fromarray(img.astype(np.uint8)), wm_encoder) for img in result]
//  [python]
def inpaint(sampler, image, mask, prompt, seed, scale, ddim_steps, num_samples=1, w=512, h=512):
    device = torch.device(
        "cuda") if torch.cuda.is_available() else torch.device("cpu")
    model = sampler.model

    print("Creating invisible watermark encoder (see https://github.com/ShieldMnt/invisible-watermark)...")
    wm = "SDV2"
    wm_encoder = WatermarkEncoder()
    wm_encoder.set_watermark('bytes', wm.encode('utf-8'))

    prng = np.random.RandomState(seed)
    start_code = prng.randn(num_samples, 4, h // 8, w // 8)
    start_code = torch.from_numpy(start_code).to(
        device=device, dtype=torch.float32)

    with torch.no_grad(), \
            torch.autocast("cuda"):
        batch = make_batch_sd(image, mask, txt=prompt,
                              device=device, num_samples=num_samples)

        c = model.cond_stage_model.encode(batch["txt"])

        c_cat = list()
        for ck in model.concat_keys:
            cc = batch[ck].float()
            if ck != model.masked_image_key:
                bchw = [num_samples, 4, h // 8, w // 8]
                cc = torch.nn.functional.interpolate(cc, size=bchw[-2:])
            else:
                cc = model.get_first_stage_encoding(
                    model.encode_first_stage(cc))
            c_cat.append(cc)
        c_cat = torch.cat(c_cat, dim=1)

        # cond
        cond = {"c_concat": [c_cat], "c_crossattn": [c]}

        # uncond cond
        uc_cross = model.get_unconditional_conditioning(num_samples, "")
        uc_full = {"c_concat": [c_cat], "c_crossattn": [uc_cross]}

        shape = [model.channels, h // 8, w // 8]
        samples_cfg, intermediates = sampler.sample(
            ddim_steps,
            num_samples,
            shape,
            cond,
            verbose=False,
            eta=1.0,
            unconditional_guidance_scale=scale,
            unconditional_conditioning=uc_full,
            x_T=start_code,
        )
        x_samples_ddim = model.decode_first_stage(samples_cfg)

        result = torch.clamp((x_samples_ddim + 1.0) / 2.0,
                             min=0.0, max=1.0)

        result = result.cpu().numpy().transpose(0, 2, 3, 1) * 255
    return [put_watermark(Image.fromarray(img.astype(np.uint8)), wm_encoder) for img in result]
//  [python]
def resample(input, size, align_corners=True):
    """Lanczos resampling image.

    Args:
        input (torch.Tensor): Input image tensor.
        size (Tuple[int, int]): Output image size.
        align_corners (bool): align_corners argument of F.interpolate.
            Defaults to True.

    Returns:
        torch.Tensor: Resampling results.
    """
    n, c, h, w = input.shape
    dh, dw = size

    input = input.reshape([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.reshape([n, c, h, w])
    return t.interpolate(
        input, size, mode='bicubic', align_corners=align_corners)
//  [python]
def resample(input, size, align_corners=True):
    """Lanczos resampling image.

    Args:
        input (torch.Tensor): Input image tensor.
        size (Tuple[int, int]): Output image size.
        align_corners (bool): align_corners argument of F.interpolate.
            Defaults to True.

    Returns:
        torch.Tensor: Resampling results.
    """
    n, c, h, w = input.shape
    dh, dw = size

    input = input.reshape([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.reshape([n, c, h, w])
    return F.interpolate(
        input, size, mode='bicubic', align_corners=align_corners)
//  [python]
def forward(self, x, cs = None, cache = None):
        if cache is None:
            cache = {}
        batch, height, width = x.shape
        y = self._one_hot(x)
       
        horz = self.shift_down(self.horz_conv(y))
        vert = self.shift_right(self.vert_conv(y))
        y = horz + vert

        bg = self.bg[:, :, :height, :].expand(batch, 2, height, width)

        if cs != None and len(cs) > 0:
            if 'condition' in cache:
                cs = cache['condition']
                cs = cs[:, :, :height, :]
            else:
                cs = [self.cond_in_net[i](c) for i, c in enumerate(cs)] # apply smaller resnet where appropriate
                cs = [t.interpolate(c, size=(height, width)) for c in cs] # interpolate conditions to image size
                cs = torch.cat(cs, dim=1) # concatenate conditions
                cs = self.cond_net(cs) # apply larger conditional resnet
                cache['condition'] = cs.detach().clone()
                cs = cs[:, :, :height, :]

        # iterate over pixelblocks
        for blk in self.blks:
            y = blk(y, bg, c=cs)
        y = self.out(y)
        return y, cache
//  [python]
def forward(self, x, cs = None, cache = None):
        if cache is None:
            cache = {}
        batch, height, width = x.shape
        y = self._one_hot(x)
       
        horz = self.shift_down(self.horz_conv(y))
        vert = self.shift_right(self.vert_conv(y))
        y = horz + vert

        bg = self.bg[:, :, :height, :].expand(batch, 2, height, width)

        if cs != None and len(cs) > 0:
            if 'condition' in cache:
                cs = cache['condition']
                cs = cs[:, :, :height, :]
            else:
                cs = [self.cond_in_net[i](c) for i, c in enumerate(cs)] # apply smaller resnet where appropriate
                cs = [F.interpolate(c, size=(height, width)) for c in cs] # interpolate conditions to image size
                cs = torch.cat(cs, dim=1) # concatenate conditions
                cs = self.cond_net(cs) # apply larger conditional resnet
                cache['condition'] = cs.detach().clone()
                cs = cs[:, :, :height, :]

        # iterate over pixelblocks
        for blk in self.blks:
            y = blk(y, bg, c=cs)
        y = self.out(y)
        return y, cache
//  [python]
def forward_features(self, x):
        B = x.shape[0]
        x, T, W = self.patch_embed(x)
        cls_tokens = self.cls_token.expand(x.size(0), -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        # resizing the positional embeddings in case they don't match the input at inference
        if x.size(1) != self.pos_embed.size(1):
            pos_embed = self.pos_embed
            cls_pos_embed = pos_embed[0, 0, :].unsqueeze(0).unsqueeze(1)
            other_pos_embed = pos_embed[0, 1:, :].unsqueeze(0).transpose(1, 2)
            P = int(other_pos_embed.size(2) ** 0.5)
            H = x.size(1) // W
            other_pos_embed = other_pos_embed.reshape(1, x.size(2), P, P)
            new_pos_embed = F.interpolate(other_pos_embed, size=(H, W), 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)
            x = x + new_pos_embed
        else:
            x = x + self.pos_embed
        x = self.pos_drop(x)

        # Time Embeddings
        if self.attention_type != "space_only":
            cls_tokens = x[:B, 0, :].unsqueeze(1)
            x = x[:, 1:]
            x = rearrange(x, "(b t) n m -> (b n) t m", b=B, t=T)
            # Resizing time embeddings in case they don't match
            if T != self.time_embed.size(1):
                time_embed = self.time_embed.transpose(1, 2)
                new_time_embed = F.interpolate(time_embed, size=(T), mode="nearest")
                new_time_embed = new_time_embed.transpose(1, 2)
                x = x + new_time_embed
            else:
                x = x + self.time_embed
            x = self.time_drop(x)
            x = rearrange(x, "(b n) t m -> b (n t) m", b=B, t=T)
            x = torch.cat((cls_tokens, x), dim=1)

        # Attention blocks
        for blk in self.blocks:
            x = blk(x, B, T, W)

        # Predictions for space-only baseline
        if self.attention_type == "space_only":
            x = rearrange(x, "(b t) n m -> b t n m", b=B, t=T)
            x = torch.mean(x, 1)  # averaging predictions for every frame

        x = self.norm(x)

        return x
//  [python]
def _upsample(self, x: NestedTensor, disp_pred: Tensor, occ_pred: Tensor, scale: int):
        """
        Upsample the raw prediction to full resolution

        :param x: input data
        :param disp_pred: predicted disp at low res
        :param occ_pred: predicted occlusion at low res
        :param scale: high-res to low-res disparity scale
        :return: high res disp and occ prediction
        """
        _, _, h, w = x.left.size()

        # scale disparity
        disp_pred_attn = disp_pred * scale

        # upsample
        disp_pred = t.upsample(disp_pred_attn[None,], size=(h, w), mode='nearest')  # N x 1 x H x W
        occ_pred = t.upsample(occ_pred[None,], size=(h, w), mode='nearest')  # N x 1 x H x W

        if self.cal is not None:
            # normalize disparity
            eps = 1e-6
            mean_disp_pred = disp_pred.mean()
            std_disp_pred = disp_pred.std() + eps
            disp_pred_normalized = (disp_pred - mean_disp_pred) / std_disp_pred

            # normalize occlusion mask
            occ_pred_normalized = (occ_pred - 0.5) / 0.5

            disp_pred_normalized, occ_pred = self.cal(disp_pred_normalized, occ_pred_normalized, x.left)  # N x H x W

            disp_pred_final = disp_pred_normalized * std_disp_pred + mean_disp_pred
        else:
            disp_pred_final = disp_pred.squeeze(1)
            disp_pred_attn = disp_pred_attn.squeeze(1)

        return disp_pred_final.squeeze(1), disp_pred_attn.squeeze(1), occ_pred.squeeze(1)
//  [python]
def _upsample(self, x: NestedTensor, disp_pred: Tensor, occ_pred: Tensor, scale: int):
        """
        Upsample the raw prediction to full resolution

        :param x: input data
        :param disp_pred: predicted disp at low res
        :param occ_pred: predicted occlusion at low res
        :param scale: high-res to low-res disparity scale
        :return: high res disp and occ prediction
        """
        _, _, h, w = x.left.size()

        # scale disparity
        disp_pred_attn = disp_pred * scale

        # upsample
        disp_pred = F.interpolate(disp_pred_attn[None,], size=(h, w), mode='nearest')  # N x 1 x H x W
        occ_pred = F.interpolate(occ_pred[None,], size=(h, w), mode='nearest')  # N x 1 x H x W

        if self.cal is not None:
            # normalize disparity
            eps = 1e-6
            mean_disp_pred = disp_pred.mean()
            std_disp_pred = disp_pred.std() + eps
            disp_pred_normalized = (disp_pred - mean_disp_pred) / std_disp_pred

            # normalize occlusion mask
            occ_pred_normalized = (occ_pred - 0.5) / 0.5

            disp_pred_normalized, occ_pred = self.cal(disp_pred_normalized, occ_pred_normalized, x.left)  # N x H x W

            disp_pred_final = disp_pred_normalized * std_disp_pred + mean_disp_pred
        else:
            disp_pred_final = disp_pred.squeeze(1)
            disp_pred_attn = disp_pred_attn.squeeze(1)

        return disp_pred_final.squeeze(1), disp_pred_attn.squeeze(1), occ_pred.squeeze(1)
//  [python]
def interpolate(self, x, output_size, size):
        """ Feature map interpolation. """
        B, N, C = x.shape
        H, W = size
        assert N == 1 + H * W

        cls_token  = x[:, :1, :]
        img_tokens = x[:, 1:, :]
        
        img_tokens = img_tokens.transpose(1, 2).reshape(B, C, H, W)
        img_tokens = F.upsample_bilinear(img_tokens, size=output_size)  # FIXME: May have alignment issue.
        img_tokens = img_tokens.reshape(B, C, -1).transpose(1, 2)
        
        out = torch.cat((cls_token, img_tokens), dim=1)

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

        cls_token  = x[:, :1, :]
        img_tokens = x[:, 1:, :]
        
        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((cls_token, img_tokens), dim=1)

        return out
//  [python]
def forward(self, x: torch.Tensor):
        # Smoothing or blurring
        x = self.smooth(x)
        # Scaled down
        return F.interpolate(x, (x.shape[2] // 2, x.shape[3] // 2), mode='bilinear', align_corners=False)
//  [python]
def forward(self, x):
        H, W = x.size()[2:]  # [48,48]
        feat1 = self.conv1(x)  # [1,256,48,48]
        feat2 = self.conv2(x)  # [1,256,48,48]
        feat3 = self.conv3(x)  # [1,256,48,48]
        feat4 = self.conv4(x)  # [1,256,48,48]  Expanding block
        if self.with_gp:
            avg = self.avg(x)
            feat5 = self.conv1x1(avg)
            feat5 = F.upsample(feat5, size=(H, W), mode='bilinear')  # [1,256,48,48]
            feat = torch.cat([feat1, feat2, feat3, feat4, feat5], 1)  ##[1,256*5,48,48]

        else:
            feat = torch.cat([feat1, feat2, feat3, feat4], 1)

        feat = self.conv_out(feat)
        return feat
//  [python]
def forward(self, x):
        H, W = x.size()[2:]  # [48,48]
        feat1 = self.conv1(x)  # [1,256,48,48]
        feat2 = self.conv2(x)  # [1,256,48,48]
        feat3 = self.conv3(x)  # [1,256,48,48]
        feat4 = self.conv4(x)  # [1,256,48,48]  Expanding block
        if self.with_gp:
            avg = self.avg(x)
            feat5 = self.conv1x1(avg)
            feat5 = F.interpolate(feat5, (H, W), mode='bilinear', align_corners=True)  # [1,256,48,48]
            feat = torch.cat([feat1, feat2, feat3, feat4, feat5], 1)  ##[1,256*5,48,48]

        else:
            feat = torch.cat([feat1, feat2, feat3, feat4], 1)

        feat = self.conv_out(feat)
        return feat
//  [python]
def generate_cam(self, input_image, target_class=None):
        # Full forward pass
        # conv_output is the output of convolutions at specified layer
        # model_output is the final output of the model (1, 1000)
        conv_output, model_output = self.extractor.forward_pass(input_image)
        if target_class is None:
            target_class = np.argmax(model_output.data.numpy())
        # Get convolution outputs
        target = conv_output[0]
        # Create empty numpy array for cam
        cam = np.ones(target.shape[1:], dtype=np.float32)
        # Multiply each weight with its conv output and then, sum
        for i in range(len(target)):
            # Unsqueeze to 4D
            saliency_map = torch.unsqueeze(torch.unsqueeze(target[i, :, :],0),0)
            # Upsampling to input size
            saliency_map = t.functional.interpolate(saliency_map, size=(224, 224), mode='bilinear', align_corners=False)
            if saliency_map.max() == saliency_map.min():
                continue
            # Scale between 0-1
            norm_saliency_map = (saliency_map - saliency_map.min()) / (saliency_map.max() - saliency_map.min())
            # Get the target score
            w = F.softmax(self.extractor.forward_pass(input_image*norm_saliency_map)[1],dim=1)[0][target_class]
            cam += w.data.numpy() * target[i, :, :].data.numpy()
        cam = np.maximum(cam, 0)
        cam = (cam - np.min(cam)) / (np.max(cam) - np.min(cam))  # Normalize between 0-1
        cam = np.uint8(cam * 255)  # Scale between 0-255 to visualize
        cam = np.uint8(Image.fromarray(cam).resize((input_image.shape[2],
                       input_image.shape[3]), Image.ANTIALIAS))/255
        return cam
//  [python]
def generate_cam(self, input_image, target_class=None):
        # Full forward pass
        # conv_output is the output of convolutions at specified layer
        # model_output is the final output of the model (1, 1000)
        conv_output, model_output = self.extractor.forward_pass(input_image)
        if target_class is None:
            target_class = np.argmax(model_output.data.numpy())
        # Get convolution outputs
        target = conv_output[0]
        # Create empty numpy array for cam
        cam = np.ones(target.shape[1:], dtype=np.float32)
        # Multiply each weight with its conv output and then, sum
        for i in range(len(target)):
            # Unsqueeze to 4D
            saliency_map = torch.unsqueeze(torch.unsqueeze(target[i, :, :],0),0)
            # Upsampling to input size
            saliency_map = F.interpolate(saliency_map, size=(224, 224), mode='bilinear', align_corners=False)
            if saliency_map.max() == saliency_map.min():
                continue
            # Scale between 0-1
            norm_saliency_map = (saliency_map - saliency_map.min()) / (saliency_map.max() - saliency_map.min())
            # Get the target score
            w = F.softmax(self.extractor.forward_pass(input_image*norm_saliency_map)[1],dim=1)[0][target_class]
            cam += w.data.numpy() * target[i, :, :].data.numpy()
        cam = np.maximum(cam, 0)
        cam = (cam - np.min(cam)) / (np.max(cam) - np.min(cam))  # Normalize between 0-1
        cam = np.uint8(cam * 255)  # Scale between 0-255 to visualize
        cam = np.uint8(Image.fromarray(cam).resize((input_image.shape[2],
                       input_image.shape[3]), Image.ANTIALIAS))/255
        return cam
//  [python]
def test_interpolate_trace(self):
        class test(nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = nn.Conv2d(1, 32, kernel_size=3, padding=1)

            def forward(self, x):
                y = self.conv(x)
                w = t.interpolate(
                    y, mode="bilinear", align_corners=False, scale_factor=3
                )
                return w

        f = test()
        # no failure
        g = torch.jit.trace(f, (torch.zeros(1, 1, 28, 28),))
        x = torch.zeros(1, 1, 14, 14)
        # constants not baked in
        self.assertEqual(g(x), f(x))
//  [python]
def test_interpolate_trace(self):
        class test(nn.Module):
            def __init__(self):
                super().__init__()
                self.conv = nn.Conv2d(1, 32, kernel_size=3, padding=1)

            def forward(self, x):
                y = self.conv(x)
                w = nn.functional.interpolate(
                    y, mode="bilinear", align_corners=False, scale_factor=3
                )
                return w

        f = test()
        # no failure
        g = torch.jit.trace(f, (torch.zeros(1, 1, 28, 28),))
        x = torch.zeros(1, 1, 14, 14)
        # constants not baked in
        self.assertEqual(g(x), f(x))
//  [python]
def forward(self,input,e0, e1, e2, e3,e4):

        #e4: [4,512,16,16] e3:[4,256,32,32] e2:[4,128,64,64], e1:[4,64,128,128] e0:[2,64,256,256]


        if self.use_att:
            d4 = self.GAU4(e4, e3)  # 16
            d3 = self.GAU3(d4, e2)  # 32
            d2 = self.GAU2(d3, e1)  # 64
            d1 = self.GAU1(d2, e0)  # 128
        else:
            d4 = self.decoder4(e4) + e3  # [2,256,16,16]
            d3 = self.decoder3(e3) + e2  # [2,128,32,32]
            d2 = self.decoder2(d3) + e1  # [2,64,64,64]
            d1 = self.decoder1(d2)  # [2,64,128,128]

        if self.use_rfnet:
            out=self.final_conv(d1)
            _,_,H0, W0 = e0.size()
            _,_,H,W=input.size()
            fea0=t.nn.functional.upsample(input,(H0, W0), mode='bilinear')
            fea=torch.cat([fea0,e0],dim=1)
            out=self.RefNet(out,fea)
            final=t.nn.functional.upsample(out, (H, W), mode='bilinear')

        else:
            out = self.finaldeconv1(d1)  # [2,64,256,256]
            out = self.finalrelu1(out)
            out = self.finalconv2(out)  # [2,32,256,256]
            out = self.finalrelu2(out)
            final = self.finalconv3(out)  # [2,1,256,256]


        #output
        if self.n_class == 1:
            return F.sigmoid(final)  # for binary only
        else:
            return final
//  [python]
def forward(self,input,e0, e1, e2, e3,e4):

        #e4: [4,512,16,16] e3:[4,256,32,32] e2:[4,128,64,64], e1:[4,64,128,128] e0:[2,64,256,256]


        if self.use_att:
            d4 = self.GAU4(e4, e3)  # 16
            d3 = self.GAU3(d4, e2)  # 32
            d2 = self.GAU2(d3, e1)  # 64
            d1 = self.GAU1(d2, e0)  # 128
        else:
            d4 = self.decoder4(e4) + e3  # [2,256,16,16]
            d3 = self.decoder3(e3) + e2  # [2,128,32,32]
            d2 = self.decoder2(d3) + e1  # [2,64,64,64]
            d1 = self.decoder1(d2)  # [2,64,128,128]

        if self.use_rfnet:
            out=self.final_conv(d1)
            _,_,H0, W0 = e0.size()
            _,_,H,W=input.size()
            fea0=F.interpolate(input,(H0, W0), mode='bilinear', align_corners=True)
            fea=torch.cat([fea0,e0],dim=1)
            out=self.RefNet(out,fea)
            final=F.interpolate(out, (H, W), mode='bilinear', align_corners=True)

        else:
            out = self.finaldeconv1(d1)  # [2,64,256,256]
            out = self.finalrelu1(out)
            out = self.finalconv2(out)  # [2,32,256,256]
            out = self.finalrelu2(out)
            final = self.finalconv3(out)  # [2,1,256,256]


        #output
        if self.n_class == 1:
            return F.sigmoid(final)  # for binary only
        else:
            return final
//  [python]
def test_trace_script(self):
        @torch.jit.script
        def func1(x: Tuple[Tensor, Tensor]) -> Tensor:
            return x[0] + x[1]

        @torch.jit.script
        def func2(x: List[Tensor]) -> Tensor:
            return x[0] + x[1]

        a = torch.randn(5)
        b = torch.randn(5)

        self.checkTrace(func1, ((a, b),))
        self.checkTrace(func2, ((a, b),))

        @torch.jit.script
        def func3(
            x: Tensor, method: str = "bilinear", align_corners: bool = True
        ) -> Tensor:
            hw = x.shape[2:4]
            return t.interpolate(x, hw, mode=method, align_corners=align_corners)

        inp = torch.rand(1, 3, 6, 6)
        self.checkTrace(func3, (inp,))

        @torch.jit.script
        def func4(x: Tensor, a: List[Optional[str]]) -> Tensor:
            if len(a) == 2:
                return x + 2
            else:
                return x
//  [python]
def test_trace_script(self):
        @torch.jit.script
        def func1(x: Tuple[Tensor, Tensor]) -> Tensor:
            return x[0] + x[1]

        @torch.jit.script
        def func2(x: List[Tensor]) -> Tensor:
            return x[0] + x[1]

        a = torch.randn(5)
        b = torch.randn(5)

        self.checkTrace(func1, ((a, b),))
        self.checkTrace(func2, ((a, b),))

        @torch.jit.script
        def func3(
            x: Tensor, method: str = "bilinear", align_corners: bool = True
        ) -> Tensor:
            hw = x.shape[2:4]
            return F.interpolate(x, hw, mode=method, align_corners=align_corners)

        inp = torch.rand(1, 3, 6, 6)
        self.checkTrace(func3, (inp,))

        @torch.jit.script
        def func4(x: Tensor, a: List[Optional[str]]) -> Tensor:
            if len(a) == 2:
                return x + 2
            else:
                return x
//  [python]
def _coarse_decoder_pass(self, mel_specs, encoder_outputs, alignments,
                             input_mask):
        """ Double Decoder Consistency """
        T = mel_specs.shape[1]
        if T % self.coarse_decoder.r > 0:
            padding_size = self.coarse_decoder.r - (T % self.coarse_decoder.r)
            mel_specs = torch.nn.functional.pad(mel_specs,
                                                (0, 0, 0, padding_size, 0, 0))
        decoder_outputs_backward, alignments_backward, _ = self.coarse_decoder(
            encoder_outputs.detach(), mel_specs, input_mask)
        # scale_factor = self.decoder.r_init / self.decoder.r
        alignments_backward = torch.nn.functional.upsample(
            alignments_backward.transpose(1, 2),
            size=alignments.shape[1],
            mode='nearest').transpose(1, 2)
        decoder_outputs_backward = decoder_outputs_backward.transpose(1, 2)
        decoder_outputs_backward = decoder_outputs_backward[:, :T, :]
        return decoder_outputs_backward, alignments_backward
//  [python]
def _coarse_decoder_pass(self, mel_specs, encoder_outputs, alignments,
                             input_mask):
        """ Double Decoder Consistency """
        T = mel_specs.shape[1]
        if T % self.coarse_decoder.r > 0:
            padding_size = self.coarse_decoder.r - (T % self.coarse_decoder.r)
            mel_specs = torch.nn.functional.pad(mel_specs,
                                                (0, 0, 0, padding_size, 0, 0))
        decoder_outputs_backward, alignments_backward, _ = self.coarse_decoder(
            encoder_outputs.detach(), mel_specs, input_mask)
        # scale_factor = self.decoder.r_init / self.decoder.r
        alignments_backward = torch.nn.functional.interpolate(
            alignments_backward.transpose(1, 2),
            size=alignments.shape[1],
            mode='nearest').transpose(1, 2)
        decoder_outputs_backward = decoder_outputs_backward.transpose(1, 2)
        decoder_outputs_backward = decoder_outputs_backward[:, :T, :]
        return decoder_outputs_backward, alignments_backward