jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def scale_img(img, ratio=1.0, same_shape=False, gs=32):  # img(16,3,256,416)
    """Scales an image tensor `img` of shape (bs,3,y,x) by `ratio`, optionally maintaining the original shape, padded to
    multiples of `gs`.
    """
    if ratio == 1.0:
        return img
    h, w = img.shape[2:]
    s = (int(h * ratio), int(w * ratio))  # new size
    img = t.interpolate(img, size=s, mode="bilinear", align_corners=False)  # resize
    if not same_shape:  # pad/crop img
        h, w = (math.ceil(x * ratio / gs) * gs for x in (h, w))
    return t.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)
//  [python]
def scale_img(img, ratio=1.0, same_shape=False, gs=32):  # img(16,3,256,416)
    """Scales an image tensor `img` of shape (bs,3,y,x) by `ratio`, optionally maintaining the original shape, padded to
    multiples of `gs`.
    """
    if ratio == 1.0:
        return img
    h, w = img.shape[2:]
    s = (int(h * ratio), int(w * ratio))  # new size
    img = F.interpolate(img, size=s, mode="bilinear", align_corners=False)  # resize
    if not same_shape:  # pad/crop img
        h, w = (math.ceil(x * ratio / gs) * gs for x in (h, w))
    return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)
//  [python]
def loss(self, mask_preds: Tensor, labels: Tensor) -> Tensor:
        """Loss function.

        Args:
            mask_preds (Tensor): Predicted mask logits.
            labels (Tensor): Ground truth.

        Returns:
            Tensor: Semantic segmentation loss.
        """
        labels = t.nn.functional.interpolate(
            labels.float(), scale_factor=self.seg_scale_factor, mode='nearest')
        labels = labels.squeeze(1).long()
        loss_semantic_seg = self.criterion(mask_preds, labels)
        return loss_semantic_seg
//  [python]
def loss(self, mask_preds: Tensor, labels: Tensor) -> Tensor:
        """Loss function.

        Args:
            mask_preds (Tensor): Predicted mask logits.
            labels (Tensor): Ground truth.

        Returns:
            Tensor: Semantic segmentation loss.
        """
        labels = F.interpolate(
            labels.float(), scale_factor=self.seg_scale_factor, mode='nearest')
        labels = labels.squeeze(1).long()
        loss_semantic_seg = self.criterion(mask_preds, labels)
        return loss_semantic_seg
//  [python]
def preprocess(self, input: Input) -> Dict[str, Any]:
        input_img_list = []
        if isinstance(input, str):
            input_img_list = [input]
            input_type = 0
        elif isinstance(input, tuple) and len(input) == 2:
            input_img_list = list(input)
            input_type = 1
        else:
            raise TypeError(
                'modelscope error: Only support "str" or "tuple (img1, img2)" , but got {type(input)}'
            )

        if input_type == 0:
            logger.info('Processing Similar Image Generation mode')
        if input_type == 1:
            logger.info('Processing Interpolation mode')

        img_list = []
        for i, input_img in enumerate(input_img_list):
            img = LoadImage.convert_to_img(input_img)
            logger.info(f'Load {i}-th image done')
            img_list.append(img)

        transforms = T.Compose([
            data.PadToSquare(),
            T.Resize(
                self.cfg.DATA.scale_size,
                interpolation=T.InterpolationMode.BICUBIC),
            T.ToTensor(),
            T.Normalize(mean=self.cfg.DATA.mean, std=self.cfg.DATA.std)
        ])

        y_list = []
        for img in img_list:
            img = transforms(img)
            imgs = torch.unsqueeze(img, 0)
            imgs = imgs.to(self._device)
            imgs_x0 = self.autoencoder.encode(imgs)
            b, c, h, w = imgs_x0.shape
            aug_imgs = TF.normalize(
                T.functional.interpolate(
                    imgs.add(1).div(2), (self.cfg.Params.vit.vit_image_size,
                                         self.cfg.Params.vit.vit_image_size),
                    mode='bilinear',
                    align_corners=True), self.cfg.Params.vit.vit_mean,
                self.cfg.Params.vit.vit_std)
            uy = self.vit(aug_imgs)
            y = F.normalize(uy, p=2, dim=1)
            y_list.append(y)

        if input_type == 0:
            result = {
                'image_data': y_list[0],
                'c': c,
                'h': h,
                'w': w,
                'type': input_type
            }
        elif input_type == 1:
            result = {
                'image_data': y_list[0],
                'image_data_s': y_list[1],
                'c': c,
                'h': h,
                'w': w,
                'type': input_type
            }
        return result
//  [python]
def preprocess(self, input: Input) -> Dict[str, Any]:
        input_img_list = []
        if isinstance(input, str):
            input_img_list = [input]
            input_type = 0
        elif isinstance(input, tuple) and len(input) == 2:
            input_img_list = list(input)
            input_type = 1
        else:
            raise TypeError(
                'modelscope error: Only support "str" or "tuple (img1, img2)" , but got {type(input)}'
            )

        if input_type == 0:
            logger.info('Processing Similar Image Generation mode')
        if input_type == 1:
            logger.info('Processing Interpolation mode')

        img_list = []
        for i, input_img in enumerate(input_img_list):
            img = LoadImage.convert_to_img(input_img)
            logger.info(f'Load {i}-th image done')
            img_list.append(img)

        transforms = T.Compose([
            data.PadToSquare(),
            T.Resize(
                self.cfg.DATA.scale_size,
                interpolation=T.InterpolationMode.BICUBIC),
            T.ToTensor(),
            T.Normalize(mean=self.cfg.DATA.mean, std=self.cfg.DATA.std)
        ])

        y_list = []
        for img in img_list:
            img = transforms(img)
            imgs = torch.unsqueeze(img, 0)
            imgs = imgs.to(self._device)
            imgs_x0 = self.autoencoder.encode(imgs)
            b, c, h, w = imgs_x0.shape
            aug_imgs = TF.normalize(
                F.interpolate(
                    imgs.add(1).div(2), (self.cfg.Params.vit.vit_image_size,
                                         self.cfg.Params.vit.vit_image_size),
                    mode='bilinear',
                    align_corners=True), self.cfg.Params.vit.vit_mean,
                self.cfg.Params.vit.vit_std)
            uy = self.vit(aug_imgs)
            y = F.normalize(uy, p=2, dim=1)
            y_list.append(y)

        if input_type == 0:
            result = {
                'image_data': y_list[0],
                'c': c,
                'h': h,
                'w': w,
                'type': input_type
            }
        elif input_type == 1:
            result = {
                'image_data': y_list[0],
                'image_data_s': y_list[1],
                'c': c,
                'h': h,
                'w': w,
                'type': input_type
            }
        return result
//  [python]
def forward(self, feats: Tuple[Tensor]) -> Tuple[Tensor]:
        """Forward function.

        Args:
            feats (tuple[Tensor]): Multi scale feature maps.

        Returns:
            tuple[Tensor]:

                - mask_preds (Tensor): Predicted mask logits.
                - x (Tensor): Fused feature.
        """
        x = self.lateral_convs[self.fusion_level](feats[self.fusion_level])
        fused_size = tuple(x.shape[-2:])
        for i, feat in enumerate(feats):
            if i != self.fusion_level:
                feat = t.interpolate(
                    feat, size=fused_size, mode='bilinear', align_corners=True)
                # fix runtime error of "+=" inplace operation in PyTorch 1.10
                x = x + self.lateral_convs[i](feat)

        for i in range(self.num_convs):
            x = self.convs[i](x)

        mask_preds = self.conv_logits(x)
        x = self.conv_embedding(x)
        return mask_preds, x
//  [python]
def forward(self, feats: Tuple[Tensor]) -> Tuple[Tensor]:
        """Forward function.

        Args:
            feats (tuple[Tensor]): Multi scale feature maps.

        Returns:
            tuple[Tensor]:

                - mask_preds (Tensor): Predicted mask logits.
                - x (Tensor): Fused feature.
        """
        x = self.lateral_convs[self.fusion_level](feats[self.fusion_level])
        fused_size = tuple(x.shape[-2:])
        for i, feat in enumerate(feats):
            if i != self.fusion_level:
                feat = F.interpolate(
                    feat, size=fused_size, mode='bilinear', align_corners=True)
                # fix runtime error of "+=" inplace operation in PyTorch 1.10
                x = x + self.lateral_convs[i](feat)

        for i in range(self.num_convs):
            x = self.convs[i](x)

        mask_preds = self.conv_logits(x)
        x = self.conv_embedding(x)
        return mask_preds, x
//  [python]
def forward(self, x, x_mask=None, g=None, **kwargs):  # pylint: disable=unused-argument
        output = torch.zeros_like(x)
        n_channels_tensor = torch.IntTensor([self.hidden_channels])
        x_mask = 1.0 if x_mask is None else x_mask
        if g is not None:
            g = self.cond_layer(g)

        g = t.interpolate(g.unsqueeze(0).unsqueeze(0), (g.shape[0], g.shape[1], self.in_layers[0](x).shape[2]))[0][0]

        for i in range(self.num_layers):
            x_in = self.in_layers[i](x)
            x_in = self.dropout(x_in)
            if g is not None:
                cond_offset = i * 2 * self.hidden_channels
                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
            else:
                g_l = torch.zeros_like(x_in)
            acts = fused_add_tanh_sigmoid_multiply(x_in, g_l, n_channels_tensor)
            res_skip_acts = self.res_skip_layers[i](acts)
            if i < self.num_layers - 1:
                x = (x + res_skip_acts[:, : self.hidden_channels, :]) * x_mask
                output = output + res_skip_acts[:, self.hidden_channels :, :]
            else:
                output = output + res_skip_acts
        return output * x_mask
//  [python]
def forward(self, x, x_mask=None, g=None, **kwargs):  # pylint: disable=unused-argument
        output = torch.zeros_like(x)
        n_channels_tensor = torch.IntTensor([self.hidden_channels])
        x_mask = 1.0 if x_mask is None else x_mask
        if g is not None:
            g = self.cond_layer(g)

        g = torch.nn.functional.interpolate(g.unsqueeze(0).unsqueeze(0), (g.shape[0], g.shape[1], self.in_layers[0](x).shape[2]))[0][0]

        for i in range(self.num_layers):
            x_in = self.in_layers[i](x)
            x_in = self.dropout(x_in)
            if g is not None:
                cond_offset = i * 2 * self.hidden_channels
                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
            else:
                g_l = torch.zeros_like(x_in)
            acts = fused_add_tanh_sigmoid_multiply(x_in, g_l, n_channels_tensor)
            res_skip_acts = self.res_skip_layers[i](acts)
            if i < self.num_layers - 1:
                x = (x + res_skip_acts[:, : self.hidden_channels, :]) * x_mask
                output = output + res_skip_acts[:, self.hidden_channels :, :]
            else:
                output = output + res_skip_acts
        return output * x_mask
//  [python]
def _pyramid_noise(self, epsilon, size_range=None, levels=10, scale_mode='nearest'):
        epsilon = epsilon.clone()
        multipliers = [1]
        for i in range(1, levels):
            m = 0.75 ** i
            h, w = epsilon.size(-2) // (2 ** i), epsilon.size(-2) // (2 ** i)
            if size_range is None or (size_range[0] <= h <= size_range[1] or size_range[0] <= w <= size_range[1]):
                offset = torch.randn(epsilon.size(0), epsilon.size(1), h, w, device=self.device)
                epsilon = epsilon + t.interpolate(offset, size=epsilon.shape[-2:],
                                                                    mode=scale_mode) * m
                multipliers.append(m)
            if h <= 1 or w <= 1:
                break
        epsilon = epsilon / sum([m ** 2 for m in multipliers]) ** 0.5
        # epsilon = epsilon / epsilon.std()
        return epsilon
//  [python]
def _pyramid_noise(self, epsilon, size_range=None, levels=10, scale_mode='nearest'):
        epsilon = epsilon.clone()
        multipliers = [1]
        for i in range(1, levels):
            m = 0.75 ** i
            h, w = epsilon.size(-2) // (2 ** i), epsilon.size(-2) // (2 ** i)
            if size_range is None or (size_range[0] <= h <= size_range[1] or size_range[0] <= w <= size_range[1]):
                offset = torch.randn(epsilon.size(0), epsilon.size(1), h, w, device=self.device)
                epsilon = epsilon + torch.nn.functional.interpolate(offset, size=epsilon.shape[-2:],
                                                                    mode=scale_mode) * m
                multipliers.append(m)
            if h <= 1 or w <= 1:
                break
        epsilon = epsilon / sum([m ** 2 for m in multipliers]) ** 0.5
        # epsilon = epsilon / epsilon.std()
        return epsilon
//  [python]
def forward(self, sequence_output):
        sequence_output = self.conv_before_upsample(sequence_output)
        sequence_output = self.activation(sequence_output)
        sequence_output = self.lrelu(
            self.conv_up1(torch.nn.functional.upsample(sequence_output, scale_factor=2, mode="nearest"))
        )
        sequence_output = self.lrelu(
            self.conv_up2(torch.nn.functional.upsample(sequence_output, scale_factor=2, mode="nearest"))
        )
        reconstruction = self.final_convolution(self.lrelu(self.conv_hr(sequence_output)))
        return reconstruction
//  [python]
def forward(self, sequence_output):
        sequence_output = self.conv_before_upsample(sequence_output)
        sequence_output = self.activation(sequence_output)
        sequence_output = self.lrelu(
            self.conv_up1(torch.nn.functional.interpolate(sequence_output, scale_factor=2, mode="nearest"))
        )
        sequence_output = self.lrelu(
            self.conv_up2(torch.nn.functional.interpolate(sequence_output, scale_factor=2, mode="nearest"))
        )
        reconstruction = self.final_convolution(self.lrelu(self.conv_hr(sequence_output)))
        return reconstruction
//  [python]
def deform_heatmap(inp, deformation):
    _, d_old, h_old, w_old, _ = deformation.shape
    _, _, d, h, w = inp.shape
    if h_old != h or w_old != w or d_old != d:
        deformation = deformation.permute(0, 4, 1, 2, 3)
        deformation = t.nn.functional.interpolate(deformation, size=(d, h, w), mode='trilinear')
        deformation = deformation.permute(0, 2, 3, 4, 1)
    return F.grid_sample(inp, deformation)
//  [python]
def deform_heatmap(inp, deformation):
    _, d_old, h_old, w_old, _ = deformation.shape
    _, _, d, h, w = inp.shape
    if h_old != h or w_old != w or d_old != d:
        deformation = deformation.permute(0, 4, 1, 2, 3)
        deformation = F.interpolate(deformation, size=(d, h, w), mode='trilinear')
        deformation = deformation.permute(0, 2, 3, 4, 1)
    return F.grid_sample(inp, deformation)
//  [python]
def forward(self, x):
         f=x
         c1_ = (self.conv1(f))
        #  c1 = self.conv2(c1_)
         v_max = F.max_pool2d(c1_, kernel_size=7, stride=3)
         v_range = self.relu(self.conv_max(v_max))
        #  c3 = self.relu(self.conv3(v_range))
         c3 = self.conv3_(v_range)
         c3 = F.upsample(c3, size=(x.size(2), x.size(3)), mode='bilinear') 
        #  cf = self.conv_f(c1_)
         c4 = self.conv4(c3+c1_)
         m = self.sigmoid(c4)
         
         return x * m
//  [python]
def forward(self, x):
         f=x
         c1_ = (self.conv1(f))
        #  c1 = self.conv2(c1_)
         v_max = F.max_pool2d(c1_, kernel_size=7, stride=3)
         v_range = self.relu(self.conv_max(v_max))
        #  c3 = self.relu(self.conv3(v_range))
         c3 = self.conv3_(v_range)
         c3 = F.interpolate(c3, (x.size(2), x.size(3)), mode='bilinear') 
        #  cf = self.conv_f(c1_)
         c4 = self.conv4(c3+c1_)
         m = self.sigmoid(c4)
         
         return x * m
//  [python]
def forward(self, x, label):
        x0 = x
        if self.conditional_bn:
            x = self.bn1(x, label)
        else:
            x = self.bn1(x)
        x = self.activation(x)
        x = t.upsample(x, scale_factor=2, mode='nearest')
        x = self.conv2d1(x)
        if self.conditional_bn:
            x = self.bn2(x, label)
        else:
            x = self.bn2(x)
        x = self.activation(x)
        x = self.conv2d2(x)

        x0 = t.upsample(x0, scale_factor=2, mode='nearest')
        x0 = self.conv2d0(x0)

        out = x + x0
        return out
//  [python]
def forward(self, x, label):
        x0 = x
        if self.conditional_bn:
            x = self.bn1(x, label)
        else:
            x = self.bn1(x)
        x = self.activation(x)
        x = F.interpolate(x, scale_factor=2, mode='nearest')
        x = self.conv2d1(x)
        if self.conditional_bn:
            x = self.bn2(x, label)
        else:
            x = self.bn2(x)
        x = self.activation(x)
        x = self.conv2d2(x)

        x0 = F.interpolate(x0, scale_factor=2, mode='nearest')
        x0 = self.conv2d0(x0)

        out = x + x0
        return out
//  [python]
def forward(self, pitches, base_lengths, new_lengths):
        new_Max_Length = torch.max(new_lengths)

        pitches = [
            torch.nn.functional.upsample(
                input= pitch[:base_Length].unsqueeze(0).unsqueeze(0),
                size= new_Length,
                mode= 'linear'
                ).squeeze(0).squeeze(0)
            for pitch, base_Length, new_Length in zip(pitches, base_lengths, new_lengths)
            ]
        pitches = torch.stack([
            torch.nn.functional.pad(pitch, [0, new_Max_Length - pitch.size(0)])
            for pitch in pitches
            ])
        
        return pitches
//  [python]
def forward(self, pitches, base_lengths, new_lengths):
        new_Max_Length = torch.max(new_lengths)

        pitches = [
            torch.nn.functional.interpolate(
                input= pitch[:base_Length].unsqueeze(0).unsqueeze(0),
                size= new_Length,
                mode= 'linear',
                align_corners= True
                ).squeeze(0).squeeze(0)
            for pitch, base_Length, new_Length in zip(pitches, base_lengths, new_lengths)
            ]
        pitches = torch.stack([
            torch.nn.functional.pad(pitch, [0, new_Max_Length - pitch.size(0)])
            for pitch in pitches
            ])
        
        return pitches
//  [python]
def post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple] = None):
        """
        Converts the output of [`SegformerForSemanticSegmentation`] into semantic segmentation maps. Only supports PyTorch.

        Args:
            outputs ([`SegformerForSemanticSegmentation`]):
                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 = t.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 post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple] = None):
        """
        Converts the output of [`SegformerForSemanticSegmentation`] into semantic segmentation maps. Only supports PyTorch.

        Args:
            outputs ([`SegformerForSemanticSegmentation`]):
                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 interpolate(
    x: torch.Tensor,
    size: Union[Sequence[int], int],
    /,
    *,
    mode: Literal[
        "linear",
        "bilinear",
        "trilinear",
        "nd",
        "nearest",
        "area",
        "nearest_exact",
        "tf_area",
        "tf_bicubic",
        "bicubic",
        "mitchellcubic",
        "lanczos3",
        "lanczos5",
        "gaussian",
    ] = "linear",
    scale_factor: Optional[Union[Sequence[int], int]] = None,
    recompute_scale_factor: Optional[bool] = None,
    align_corners: bool = False,
    antialias: bool = False,
    out: Optional[torch.Tensor] = None,
):
    if mode not in ["linear", "bilinear", "bicubic", "trilinear"]:
        align_corners = None
    return t.interpolate(
        x,
        size=size,
        mode=mode,
        align_corners=align_corners,
        antialias=antialias,
        scale_factor=scale_factor,
        recompute_scale_factor=recompute_scale_factor,
    )
//  [python]
def interpolate(
    x: torch.Tensor,
    size: Union[Sequence[int], int],
    /,
    *,
    mode: Literal[
        "linear",
        "bilinear",
        "trilinear",
        "nd",
        "nearest",
        "area",
        "nearest_exact",
        "tf_area",
        "tf_bicubic",
        "bicubic",
        "mitchellcubic",
        "lanczos3",
        "lanczos5",
        "gaussian",
    ] = "linear",
    scale_factor: Optional[Union[Sequence[int], int]] = None,
    recompute_scale_factor: Optional[bool] = None,
    align_corners: bool = False,
    antialias: bool = False,
    out: Optional[torch.Tensor] = None,
):
    if mode not in ["linear", "bilinear", "bicubic", "trilinear"]:
        align_corners = None
    return torch.nn.functional.interpolate(
        x,
        size=size,
        mode=mode,
        align_corners=align_corners,
        antialias=antialias,
        scale_factor=scale_factor,
        recompute_scale_factor=recompute_scale_factor,
    )
//  [python]
def forward(self, x):
        if self.num_branches == 1:
            return [self.branches[0](x[0])]
        for i in range(self.num_branches):
            x[i] = self.branches[i](x[i])
        x_fuse = []
        for i in range(len(self.fuse_layers)):
            y = 0
            for j in range(self.num_branches):
                if i == j:
                    y += x[j]
                elif j > i:
                    y = y + t.interpolate(self.fuse_layers[i][j](x[j]), size=x[i].shape[2:], mode='bilinear', align_corners=False)
                else:
                    y += self.fuse_layers[i][j](x[j])
            x_fuse.append(self.relu(y))
        return x_fuse
//  [python]
def forward(self, x):
        if self.num_branches == 1:
            return [self.branches[0](x[0])]
        for i in range(self.num_branches):
            x[i] = self.branches[i](x[i])
        x_fuse = []
        for i in range(len(self.fuse_layers)):
            y = 0
            for j in range(self.num_branches):
                if i == j:
                    y += x[j]
                elif j > i:
                    y = y + F.interpolate(self.fuse_layers[i][j](x[j]), size=x[i].shape[2:], mode='bilinear', align_corners=False)
                else:
                    y += self.fuse_layers[i][j](x[j])
            x_fuse.append(self.relu(y))
        return x_fuse
//  [python]
def forward(self, hidden_states, output_size=None):
        assert hidden_states.shape[1] == self.channels

        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
        # TODO(Suraj): Remove this cast once the issue is fixed in PyTorch
        # https://github.com/pytorch/pytorch/issues/86679
        dtype = hidden_states.dtype
        if dtype == torch.bfloat16:
            hidden_states = hidden_states.to(torch.float32)

        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
        if hidden_states.shape[0] >= 64:
            hidden_states = hidden_states.contiguous()

        b = hidden_states.shape[0]
        is_video = hidden_states.ndim == 5
        if is_video:
            hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")

        # if `output_size` is passed we force the interpolation output
        # size and do not make use of `scale_factor=2`
        if output_size is None:
            hidden_states = F.upsample(hidden_states, scale_factor=2.0, mode="nearest")
        else:
            hidden_states = F.upsample(hidden_states, size=output_size, mode="nearest")
        
        if is_video:
            td =  ('temporal_downsample' in self.model_config and self.model_config['temporal_downsample'] is True)

                
            if td:
                hidden_states = rearrange(hidden_states, " (b f) c h w -> b c h w f ", b=b)
                t_b, t_c, t_h, t_w, t_f = hidden_states.shape
                hidden_states = rearrange(hidden_states, " b c h w f -> (b c) (h w) f ", b=b)
                
                hidden_states = F.upsample(hidden_states, scale_factor=2.0, mode="linear")
                hidden_states = rearrange(hidden_states, " (b c) (h w) f  ->  (b f) c h w ", b=t_b, h=t_h)
            
        # If the input is bfloat16, we cast back to bfloat16
        if dtype == torch.bfloat16:
            hidden_states = hidden_states.to(dtype)

        if is_video:
            hidden_states = rearrange(hidden_states, "(b f) c h w -> b c f h w", b=b)

        # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed
        if self.use_conv:
            if self.name == "conv":
                hidden_states = self.conv(hidden_states)
            else:
                hidden_states = self.Conv2d_0(hidden_states)

        return hidden_states
//  [python]
def forward(self, hidden_states, output_size=None):
        assert hidden_states.shape[1] == self.channels

        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
        # TODO(Suraj): Remove this cast once the issue is fixed in PyTorch
        # https://github.com/pytorch/pytorch/issues/86679
        dtype = hidden_states.dtype
        if dtype == torch.bfloat16:
            hidden_states = hidden_states.to(torch.float32)

        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
        if hidden_states.shape[0] >= 64:
            hidden_states = hidden_states.contiguous()

        b = hidden_states.shape[0]
        is_video = hidden_states.ndim == 5
        if is_video:
            hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")

        # if `output_size` is passed we force the interpolation output
        # size and do not make use of `scale_factor=2`
        if output_size is None:
            hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
        else:
            hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
        
        if is_video:
            td =  ('temporal_downsample' in self.model_config and self.model_config['temporal_downsample'] is True)

                
            if td:
                hidden_states = rearrange(hidden_states, " (b f) c h w -> b c h w f ", b=b)
                t_b, t_c, t_h, t_w, t_f = hidden_states.shape
                hidden_states = rearrange(hidden_states, " b c h w f -> (b c) (h w) f ", b=b)
                
                hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="linear")
                hidden_states = rearrange(hidden_states, " (b c) (h w) f  ->  (b f) c h w ", b=t_b, h=t_h)
            
        # If the input is bfloat16, we cast back to bfloat16
        if dtype == torch.bfloat16:
            hidden_states = hidden_states.to(dtype)

        if is_video:
            hidden_states = rearrange(hidden_states, "(b f) c h w -> b c f h w", b=b)

        # TODO(Suraj, Patrick) - clean up after weight dicts are correctly renamed
        if self.use_conv:
            if self.name == "conv":
                hidden_states = self.conv(hidden_states)
            else:
                hidden_states = self.Conv2d_0(hidden_states)

        return hidden_states
//  [python]
def forward(
        self,
        x,
        emb = None
    ):
        if self.downsample:
            t, h, w = x.shape[-3:]
            resize_factors = tuple((2 if downsample else 1) for downsample in self.downsample_config)
            interpolate_shape = tuple(shape // factor for shape, factor in zip((t, h, w), resize_factors))

            x = F.interpolate(x, interpolate_shape, mode = 'trilinear')
            x = self.downsample_conv(x)

        x = self.pixel_norm(x)

        res = x.clone()

        x = self.block1(x)

        if exists(emb):
            scale = self.to_emb(emb) + 1
            x = x * rearrange(scale, 'b c -> b c 1 1 1')

        x = self.block2(x)

        x = self.res_mp_add(x, res)

        if exists(self.attn):
            if self.factorized_attn:
                attn_space, attn_time = self.attn
                x = attn_space(x)
                x = attn_time(x)

            else:
                x = self.attn(x)

        return x
//  [python]
def forward(
        self,
        x,
        emb = None
    ):
        if self.upsample:
            t, h, w = x.shape[-3:]
            resize_factors = tuple((2 if upsample else 1) for upsample in self.upsample_config)
            interpolate_shape = tuple(shape * factor for shape, factor in zip((t, h, w), resize_factors))

            x = F.interpolate(x, interpolate_shape, mode = 'trilinear')

        res = self.res_conv(x)

        x = self.block1(x)

        if exists(emb):
            scale = self.to_emb(emb) + 1
            x = x * rearrange(scale, 'b c -> b c 1 1 1')

        x = self.block2(x)

        x = self.res_mp_add(x, res)

        if exists(self.attn):
            if self.factorized_attn:
                attn_space, attn_time = self.attn
                x = attn_space(x)
                x = attn_time(x)

            else:
                x = self.attn(x)

        return x
//  [python]
def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None):
    # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor
    """
    Equivalent to nn.functional.interpolate, but with support for empty batch sizes.
    This will eventually be supported natively by PyTorch, and this
    class can go away.
    """
    if __torchvision_need_compat_flag < 0.7:
        if input.numel() > 0:
            return t.interpolate(input, size, scale_factor, mode, align_corners)

        output_shape = _output_size(2, input, size, scale_factor)
        output_shape = list(input.shape[:-2]) + list(output_shape)
        return _new_empty_tensor(input, output_shape)
    else:
        return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners)
//  [python]
def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None):
    # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor
    """
    Equivalent to nn.functional.interpolate, but with support for empty batch sizes.
    This will eventually be supported natively by PyTorch, and this
    class can go away.
    """
    if __torchvision_need_compat_flag < 0.7:
        if input.numel() > 0:
            return torch.nn.functional.interpolate(input, size, scale_factor, mode, align_corners)

        output_shape = _output_size(2, input, size, scale_factor)
        output_shape = list(input.shape[:-2]) + list(output_shape)
        return _new_empty_tensor(input, output_shape)
    else:
        return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners)
//  [python]
def PatchGaussian(self, img, patch_scale=0.125):
        batch_size, _, img_size, img_size = img.size()
        patch_size = int(img_size*patch_scale)
        noise = (torch.rand((batch_size, 1, patch_size, patch_size)).to(img.device) - 0.5) * 0.1
        mask = img.lt(0.99)
        img = img + mask * t(noise, size=(img_size, img_size))
        return img
//  [python]
def PatchGaussian(self, img, patch_scale=0.125):
        batch_size, _, img_size, img_size = img.size()
        patch_size = int(img_size*patch_scale)
        noise = (torch.rand((batch_size, 1, patch_size, patch_size)).to(img.device) - 0.5) * 0.1
        mask = img.lt(0.99)
        img = img + mask * F.interpolate(noise, scale_factor=1/patch_scale)
        return img
//  [python]
def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.layer1(x)
        x_list = []
        for i in range(self.stage2_cfg['num_branches']):
            if self.transition1[i] is not None:
                x_list.append(self.transition1[i](x))
            else:
                x_list.append(x)
        y_list = self.stage2(x_list)
        x_list = []
        for i in range(self.stage3_cfg['num_branches']):
            if self.transition2[i] is not None:
                x_list.append(self.transition2[i](y_list[-1]))
            else:
                x_list.append(y_list[i])
        y_list = self.stage3(x_list)
        x_list = []
        for i in range(self.stage4_cfg['num_branches']):
            if self.transition3[i] is not None:
                x_list.append(self.transition3[i](y_list[-1]))
            else:
                x_list.append(y_list[i])
        y_list = self.stage4(x_list)
        h, w = max([y.shape[2] for y in y_list]), max([y.shape[3] for y in y_list])
        out = torch.cat([t.interpolate(y, size=(h, w), mode='bilinear', align_corners=False) for y in y_list], dim=1)
        outs = [out]
        return tuple(outs)
//  [python]
def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.layer1(x)
        x_list = []
        for i in range(self.stage2_cfg['num_branches']):
            if self.transition1[i] is not None:
                x_list.append(self.transition1[i](x))
            else:
                x_list.append(x)
        y_list = self.stage2(x_list)
        x_list = []
        for i in range(self.stage3_cfg['num_branches']):
            if self.transition2[i] is not None:
                x_list.append(self.transition2[i](y_list[-1]))
            else:
                x_list.append(y_list[i])
        y_list = self.stage3(x_list)
        x_list = []
        for i in range(self.stage4_cfg['num_branches']):
            if self.transition3[i] is not None:
                x_list.append(self.transition3[i](y_list[-1]))
            else:
                x_list.append(y_list[i])
        y_list = self.stage4(x_list)
        h, w = max([y.shape[2] for y in y_list]), max([y.shape[3] for y in y_list])
        out = torch.cat([F.interpolate(y, size=(h, w), mode='bilinear', align_corners=False) for y in y_list], dim=1)
        outs = [out]
        return tuple(outs)
//  [python]
def forward(self, feats):
        assert len(feats) == self.num_scales

        # processed from bottom (high-lvl) to top (low-lvl)
        outs = []
        out = self.detect1(feats[-1])
        outs.append(out)

        for i, x in enumerate(reversed(feats[:-1])):
            conv = getattr(self, f'conv{i+1}')
            tmp = conv(out)

            # Cat with low-lvl feats
            tmp = t.upsample(tmp, scale_factor=2)
            tmp = torch.cat((tmp, x), 1)

            detect = getattr(self, f'detect{i+2}')
            out = detect(tmp)
            outs.append(out)

        return tuple(outs)
//  [python]
def forward(self, feats):
        assert len(feats) == self.num_scales

        # processed from bottom (high-lvl) to top (low-lvl)
        outs = []
        out = self.detect1(feats[-1])
        outs.append(out)

        for i, x in enumerate(reversed(feats[:-1])):
            conv = getattr(self, f'conv{i+1}')
            tmp = conv(out)

            # Cat with low-lvl feats
            tmp = F.interpolate(tmp, scale_factor=2)
            tmp = torch.cat((tmp, x), 1)

            detect = getattr(self, f'detect{i+2}')
            out = detect(tmp)
            outs.append(out)

        return tuple(outs)
//  [python]
def set_input_train(self, input):
        self.outer = input['data_outer'].to(self.device)
        self.outer = t.interpolate(self.outer,(1024,1024),mode='bilinear',align_corners=False)

        self.inner = input['data_inner'].to(self.device)
        self.inner = t.interpolate(self.inner,(1024,1024),mode='bilinear',align_corners=False)

        self.image_paths = input['image_path']

        if self.isTrain:
            self.gtfake = input['data_gtfake'].to(self.device)
            self.gtfake = t.interpolate(self.gtfake, (1024, 1024), mode='bilinear', align_corners=False)
            self.real_B = self.gtfake

        self.real_A = torch.cat((self.outer, self.inner), 1)
//  [python]
def set_input_train(self, input):
        self.outer = input['data_outer'].to(self.device)
        self.outer = torch.nn.functional.interpolate(self.outer,(1024,1024),mode='bilinear',align_corners=False)

        self.inner = input['data_inner'].to(self.device)
        self.inner = torch.nn.functional.interpolate(self.inner,(1024,1024),mode='bilinear',align_corners=False)

        self.image_paths = input['image_path']

        if self.isTrain:
            self.gtfake = input['data_gtfake'].to(self.device)
            self.gtfake = torch.nn.functional.interpolate(self.gtfake, (1024, 1024), mode='bilinear', align_corners=False)
            self.real_B = self.gtfake

        self.real_A = torch.cat((self.outer, self.inner), 1)
//  [python]
def get_input(self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False):
        # note: restricted to non-trainable encoders currently
        assert not self.cond_stage_trainable, 'trainable cond stages not yet supported for depth2img'
        z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True,
                                              force_c_encode=True, return_original_cond=True, bs=bs)

        assert exists(self.concat_keys)
        assert len(self.concat_keys) == 1
        c_cat = list()
        for ck in self.concat_keys:
            cc = batch[ck]
            if bs is not None:
                cc = cc[:bs]
                cc = cc.to(self.device)
            cc = self.depth_model(cc)
            cc = t.interpolate(
                cc,
                size=z.shape[2:],
                mode="bicubic",
                align_corners=False,
            )

            depth_min, depth_max = torch.amin(cc, dim=[1, 2, 3], keepdim=True), torch.amax(cc, dim=[1, 2, 3],
                                                                                           keepdim=True)
            cc = 2. * (cc - depth_min) / (depth_max - depth_min + 0.001) - 1.
            c_cat.append(cc)
        c_cat = torch.cat(c_cat, dim=1)
        all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
        if return_first_stage_outputs:
            return z, all_conds, x, xrec, xc
        return z, all_conds
//  [python]
def get_input(self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False):
        # note: restricted to non-trainable encoders currently
        assert not self.cond_stage_trainable, 'trainable cond stages not yet supported for depth2img'
        z, c, x, xrec, xc = super().get_input(batch, self.first_stage_key, return_first_stage_outputs=True,
                                              force_c_encode=True, return_original_cond=True, bs=bs)

        assert exists(self.concat_keys)
        assert len(self.concat_keys) == 1
        c_cat = list()
        for ck in self.concat_keys:
            cc = batch[ck]
            if bs is not None:
                cc = cc[:bs]
                cc = cc.to(self.device)
            cc = self.depth_model(cc)
            cc = torch.nn.functional.interpolate(
                cc,
                size=z.shape[2:],
                mode="bicubic",
                align_corners=False,
            )

            depth_min, depth_max = torch.amin(cc, dim=[1, 2, 3], keepdim=True), torch.amax(cc, dim=[1, 2, 3],
                                                                                           keepdim=True)
            cc = 2. * (cc - depth_min) / (depth_max - depth_min + 0.001) - 1.
            c_cat.append(cc)
        c_cat = torch.cat(c_cat, dim=1)
        all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
        if return_first_stage_outputs:
            return z, all_conds, x, xrec, xc
        return z, all_conds
//  [python]
def attn_visualization_all(Ic, Is, Fc, Fs, Attn, slide_size=4):
	if Ic.size(3) != Is.size(3):
		Is = t.nn.functional.interpolate(Is, size=(Ic.size(2), Ic.size(3)))
	B, C, S = Attn.size()
	B, c, c_w, c_h = Ic.size()
	B, c, s_w, s_h = Is.size()
	_, _, f_w, f_h = Fs.size()
	attn = torch.zeros(1, Fs.size(2)*Fs.size(3)).cuda()
	mask = torch.zeros(1, Fc.size(2)*Fc.size(3)).cuda()
	for index in range(0, Fc.size(2)*Fc.size(3), slide_size):
		for idx in range(slide_size):
			start_idx = index+Fc.size(3)*idx
			start_idx_ref = index+Fs.size(3)*idx
			mask[0][start_idx:start_idx+slide_size] = 1
			attn += torch.sum(Attn[0][start_idx:start_idx+slide_size], dim=0)
	masked = t.nn.functional.interpolate(mask.view(1, 1, Fc.size(2), Fc.size(3)), size=(c_w, c_h), mode='nearest')
	attn = t.nn.functional.interpolate(attn.view(1, 1, Fs.size(2), Fs.size(3)), size=(s_w, s_h), mode='nearest')
	#imsave_no_norm(Is, attn.repeat(1,3,1,1),'style.png')
	return attn.repeat(1,3,1,1), Ic*(masked+0.2), Is*(attn*(1/attn.max())+0.2), Is*(attn*0.05+0.2)
//  [python]
def attn_visualization_all(Ic, Is, Fc, Fs, Attn, slide_size=4):
	if Ic.size(3) != Is.size(3):
		Is = torch.nn.functional.interpolate(Is, size=(Ic.size(2), Ic.size(3)))
	B, C, S = Attn.size()
	B, c, c_w, c_h = Ic.size()
	B, c, s_w, s_h = Is.size()
	_, _, f_w, f_h = Fs.size()
	attn = torch.zeros(1, Fs.size(2)*Fs.size(3)).cuda()
	mask = torch.zeros(1, Fc.size(2)*Fc.size(3)).cuda()
	for index in range(0, Fc.size(2)*Fc.size(3), slide_size):
		for idx in range(slide_size):
			start_idx = index+Fc.size(3)*idx
			start_idx_ref = index+Fs.size(3)*idx
			mask[0][start_idx:start_idx+slide_size] = 1
			attn += torch.sum(Attn[0][start_idx:start_idx+slide_size], dim=0)
	masked = torch.nn.functional.interpolate(mask.view(1, 1, Fc.size(2), Fc.size(3)), size=(c_w, c_h), mode='nearest')
	attn = torch.nn.functional.interpolate(attn.view(1, 1, Fs.size(2), Fs.size(3)), size=(s_w, s_h), mode='nearest')
	#imsave_no_norm(Is, attn.repeat(1,3,1,1),'style.png')
	return attn.repeat(1,3,1,1), Ic*(masked+0.2), Is*(attn*(1/attn.max())+0.2), Is*(attn*0.05+0.2)
//  [python]
def forward(
        self, inputs: Tensor, data_samples: List[DetDataSample]
    ) -> Tuple[Tensor, List[DetDataSample]]:
        """resize a batch of images and bboxes."""

        batch_height, batch_width = inputs.shape[-2:]
        target_height, target_width, scale = self.get_target_size(
            batch_height, batch_width)

        inputs = t.interpolate(
            inputs,
            size=(target_height, target_width),
            mode='bilinear',
            align_corners=False)

        inputs = self.get_padded_tensor(inputs, self.pad_value)

        if data_samples is not None:
            batch_input_shape = tuple(inputs.size()[-2:])
            for data_sample in data_samples:
                img_shape = [
                    int(scale * _) for _ in list(data_sample.img_shape)
                ]
                data_sample.set_metainfo({
                    'img_shape': tuple(img_shape),
                    'batch_input_shape': batch_input_shape,
                    'pad_shape': batch_input_shape,
                    'scale_factor': (scale, scale)
                })

                data_sample.gt_instances.bboxes *= scale
                data_sample.ignored_instances.bboxes *= scale

        return inputs, data_samples
//  [python]
def forward(
        self, inputs: Tensor, data_samples: List[DetDataSample]
    ) -> Tuple[Tensor, List[DetDataSample]]:
        """resize a batch of images and bboxes."""

        batch_height, batch_width = inputs.shape[-2:]
        target_height, target_width, scale = self.get_target_size(
            batch_height, batch_width)

        inputs = F.interpolate(
            inputs,
            size=(target_height, target_width),
            mode='bilinear',
            align_corners=False)

        inputs = self.get_padded_tensor(inputs, self.pad_value)

        if data_samples is not None:
            batch_input_shape = tuple(inputs.size()[-2:])
            for data_sample in data_samples:
                img_shape = [
                    int(scale * _) for _ in list(data_sample.img_shape)
                ]
                data_sample.set_metainfo({
                    'img_shape': tuple(img_shape),
                    'batch_input_shape': batch_input_shape,
                    'pad_shape': batch_input_shape,
                    'scale_factor': (scale, scale)
                })

                data_sample.gt_instances.bboxes *= scale
                data_sample.ignored_instances.bboxes *= scale

        return inputs, data_samples