jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _sample_points_strided_grid(self, enc_images, enc_segmentations):
        """Downsample images by a factor of self.stride and pick those points
        """
        B, N, C, H, W = enc_segmentations.size()
        B, N, D, H, W = enc_images.size()
        stride = self.stride
        Hs, Ws = int(H//stride), int(W//stride)
        enc_segmentations = t.nn.functional.interpolate(enc_segmentations.view(B*N, C, H, W), size=(Hs, Ws), mode="nearest")
        enc_segmentations = enc_segmentations.reshape(B, N, C, Hs, Ws)
        enc_images = t.nn.functional.interpolate(enc_images.view(B*N, D, H, W), size=(Hs, Ws), mode="nearest")
        enc_images = enc_images.reshape(B, N, D, Hs, Ws)
        x_s = enc_images.permute(0,1,3,4,2).reshape(B, N*Hs*Ws, D)
        y_s = enc_segmentations.permute(0,1,3,4,2).reshape(B, N*Hs*Ws, C)
        return x_s, y_s
//  [python]
def _sample_points_strided_grid(self, enc_images, enc_segmentations):
        """Downsample images by a factor of self.stride and pick those points
        """
        B, N, C, H, W = enc_segmentations.size()
        B, N, D, H, W = enc_images.size()
        stride = self.stride
        Hs, Ws = int(H//stride), int(W//stride)
        enc_segmentations = F.interpolate(enc_segmentations.view(B*N, C, H, W), size=(Hs, Ws), mode="nearest")
        enc_segmentations = enc_segmentations.reshape(B, N, C, Hs, Ws)
        enc_images = F.interpolate(enc_images.view(B*N, D, H, W), size=(Hs, Ws), mode="nearest")
        enc_images = enc_images.reshape(B, N, D, Hs, Ws)
        x_s = enc_images.permute(0,1,3,4,2).reshape(B, N*Hs*Ws, D)
        y_s = enc_segmentations.permute(0,1,3,4,2).reshape(B, N*Hs*Ws, C)
        return x_s, y_s
//  [python]
def _upsample_add(x, y):
    """Upsample and add two feature maps.
    Args:
      x: (Variable) top feature map to be upsampled.
      y: (Variable) lateral feature map.
    Returns:
      (Variable) added feature map.
    Note in PyTorch, when input size is odd, the upsampled feature map
    with `F.upsample(..., scale_factor=2, mode='nearest')`
    maybe not equal to the lateral feature map size.
    e.g.
    original input size: [N,_,15,15] ->
    conv2d feature map size: [N,_,8,8] ->
    upsampled feature map size: [N,_,16,16]
    So we choose bilinear upsample which supports arbitrary output sizes.
    """
    _, _, H, W = y.size()
    return F.upsample(x, size=(H, W), mode='bilinear') + y
//  [python]
def _upsample_add(x, y):
    """Upsample and add two feature maps.
    Args:
      x: (Variable) top feature map to be upsampled.
      y: (Variable) lateral feature map.
    Returns:
      (Variable) added feature map.
    Note in PyTorch, when input size is odd, the upsampled feature map
    with `F.upsample(..., scale_factor=2, mode='nearest')`
    maybe not equal to the lateral feature map size.
    e.g.
    original input size: [N,_,15,15] ->
    conv2d feature map size: [N,_,8,8] ->
    upsampled feature map size: [N,_,16,16]
    So we choose bilinear upsample which supports arbitrary output sizes.
    """
    _, _, H, W = y.size()
    return F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True) + y
//  [python]
def adaptive_downsample256(img, mode='bilinear'):
  img = img.clamp(-1, 1)
  if img.shape[-1] > 256:
    return t.interpolate(img, size=(256, 256), mode=mode)
  else:
    return img
//  [python]
def adaptive_downsample256(img, mode='bilinear'):
  img = img.clamp(-1, 1)
  if img.shape[-1] > 256:
    return F.interpolate(img, size=(256, 256), mode=mode)
  else:
    return img
//  [python]
def create_depth_map(
    img: "Tensor", model_type="dpt_hybrid_384", max_size=512
) -> "Tensor":
    import torch

    orig_size = img.shape[2:]

    depth_pt = _create_depth_map_raw(img, max_size=max_size, model_type=model_type)
    # copy the depth map to the other channels
    depth_pt = torch.cat([depth_pt, depth_pt, depth_pt], dim=0)

    depth_pt -= torch.min(depth_pt)
    depth_pt /= torch.max(depth_pt)
    depth_pt = depth_pt.unsqueeze(0)
    # depth_pt = depth_pt.cpu().numpy()

    depth_pt = torch.nn.functional.interpolate(
        depth_pt,
        size=orig_size,
        mode="bilinear",
    )

    return depth_pt
//  [python]
def _create_depth_map_raw(
    img: "Tensor", max_size=512, model_type="dpt_large_384"
) -> "Tensor":
    import torch

    from imaginairy.modules.midas.api import MiDaSInference, midas_device

    model = MiDaSInference(model_type=model_type).to(midas_device())
    img = img.to(midas_device())

    # calculate new size such that image fits within 512x512 but keeps aspect ratio
    if img.shape[2] > img.shape[3]:
        new_size = (max_size, int(max_size * img.shape[3] / img.shape[2]))
    else:
        new_size = (int(max_size * img.shape[2] / img.shape[3]), max_size)

    # resize torch image to be multiple of 32
    img = torch.nn.functional.interpolate(
        img,
        size=(new_size[0] // 32 * 32, new_size[1] // 32 * 32),
        mode="bilinear",
        align_corners=False,
    )

    depth_pt = model(img)[0]
    return depth_pt
//  [python]
def forward(self, features):
        """
        Arguments:
            features (list[Tensor]): FPN feature map tensors in high to low resolution.
                Each tensor in the list correspond to different feature levels.

        Returns:
            pred_logits (list[Tensor]): #lvl tensors, each has shape (N, AxK, Hi, Wi).
                The tensor predicts the classification probability
                at each spatial position for each of the A anchors and K object
                classes.
            pred_deltas (list[Tensor]): #lvl tensors, each has shape (N, Ax4, Hi, Wi).
                The tensor predicts 4-vector (dx,dy,dw,dh) box
                regression values for every anchor. These values are the
                relative offset between the anchor and the ground truth box.
            pred_masks (list(list[Tensor])): #lvl list of tensors, each is a list of
                A tensors of shape (N, M_{i,a}, Hi, Wi).
                The tensor predicts a dense set of M_ixM_i masks at every location.
        """
        pred_logits = [self.cls_score(self.cls_subnet(x)) for x in features]
        pred_deltas = [self.bbox_pred(self.bbox_subnet(x)) for x in features]

        pred_masks = None
        if self.mask_on:
            mask_feats = [self.mask_subnet(x) for x in features]

            if self.bipyramid_on:
                mask_feat_high_res = mask_feats[0]
                H, W = mask_feat_high_res.shape[-2:]
                mask_feats_up = []
                for lvl, mask_feat in enumerate(mask_feats):
                    lambda_val = 2.0**lvl
                    mask_feat_up = mask_feat
                    if lvl > 0:
                        mask_feat_up = t.nn.functional.interpolate(
                            mask_feat, scale_factor=lambda_val, mode="bilinear", align_corners=False
                        )
                    mask_feats_up.append(
                        self.mask_fuse(mask_feat_up[:, :, :H, :W] + mask_feat_high_res)
                    )
                mask_feats = mask_feats_up

            pred_masks = []
            for lvl, mask_feat in enumerate(mask_feats):
                cur_masks = []
                for mask_size in self.mask_sizes:
                    cur_mask_module = getattr(self, "mask_pred_%02d" % mask_size)
                    cur_mask = cur_mask_module(mask_feat)
                    if self.align_on:
                        if self.bipyramid_on:
                            cur_mask_module = getattr(self, "align2nat_%02d" % lvl)
                            cur_mask = cur_mask_module(cur_mask)
                        else:
                            cur_mask = self.align2nat(cur_mask)
                    cur_masks.append(cur_mask)
                pred_masks.append(cur_masks)
        return pred_logits, pred_deltas, pred_masks
//  [python]
def forward(self, features):
        """
        Arguments:
            features (list[Tensor]): FPN feature map tensors in high to low resolution.
                Each tensor in the list correspond to different feature levels.

        Returns:
            pred_logits (list[Tensor]): #lvl tensors, each has shape (N, AxK, Hi, Wi).
                The tensor predicts the classification probability
                at each spatial position for each of the A anchors and K object
                classes.
            pred_deltas (list[Tensor]): #lvl tensors, each has shape (N, Ax4, Hi, Wi).
                The tensor predicts 4-vector (dx,dy,dw,dh) box
                regression values for every anchor. These values are the
                relative offset between the anchor and the ground truth box.
            pred_masks (list(list[Tensor])): #lvl list of tensors, each is a list of
                A tensors of shape (N, M_{i,a}, Hi, Wi).
                The tensor predicts a dense set of M_ixM_i masks at every location.
        """
        pred_logits = [self.cls_score(self.cls_subnet(x)) for x in features]
        pred_deltas = [self.bbox_pred(self.bbox_subnet(x)) for x in features]

        pred_masks = None
        if self.mask_on:
            mask_feats = [self.mask_subnet(x) for x in features]

            if self.bipyramid_on:
                mask_feat_high_res = mask_feats[0]
                H, W = mask_feat_high_res.shape[-2:]
                mask_feats_up = []
                for lvl, mask_feat in enumerate(mask_feats):
                    lambda_val = 2.0**lvl
                    mask_feat_up = mask_feat
                    if lvl > 0:
                        mask_feat_up = F.interpolate(
                            mask_feat, scale_factor=lambda_val, mode="bilinear", align_corners=False
                        )
                    mask_feats_up.append(
                        self.mask_fuse(mask_feat_up[:, :, :H, :W] + mask_feat_high_res)
                    )
                mask_feats = mask_feats_up

            pred_masks = []
            for lvl, mask_feat in enumerate(mask_feats):
                cur_masks = []
                for mask_size in self.mask_sizes:
                    cur_mask_module = getattr(self, "mask_pred_%02d" % mask_size)
                    cur_mask = cur_mask_module(mask_feat)
                    if self.align_on:
                        if self.bipyramid_on:
                            cur_mask_module = getattr(self, "align2nat_%02d" % lvl)
                            cur_mask = cur_mask_module(cur_mask)
                        else:
                            cur_mask = self.align2nat(cur_mask)
                    cur_masks.append(cur_mask)
                pred_masks.append(cur_masks)
        return pred_logits, pred_deltas, pred_masks
//  [python]
def test_upsamplingNearest1d(self):
        m = nn.Upsample(size=4, mode='nearest')
        in_t = torch.ones(1, 1, 2)
        in_uint8_t = torch.ones(1, 1, 2, dtype=torch.uint8)
        with warnings.catch_warnings(record=True) as w:
            out_t = m(in_t)
            out_uint8_t = m(in_uint8_t)
        self.assertEqual(torch.ones(1, 1, 4), out_t.data)
        self.assertEqual(torch.ones(1, 1, 4, dtype=torch.uint8), out_uint8_t.data)

        input = torch.randn(1, 1, 2, requires_grad=True)
        gradcheck(lambda x: nn.functional.upsample(x, scale_factor=4, mode='nearest'), [input])
//  [python]
def test_upsamplingNearest1d(self):
        m = nn.Upsample(size=4, mode='nearest')
        in_t = torch.ones(1, 1, 2)
        in_uint8_t = torch.ones(1, 1, 2, dtype=torch.uint8)
        with warnings.catch_warnings(record=True) as w:
            out_t = m(in_t)
            out_uint8_t = m(in_uint8_t)
        self.assertEqual(torch.ones(1, 1, 4), out_t.data)
        self.assertEqual(torch.ones(1, 1, 4, dtype=torch.uint8), out_uint8_t.data)

        input = torch.randn(1, 1, 2, requires_grad=True)
        gradcheck(lambda x: F.interpolate(x, 4, mode='nearest'), [input])
//  [python]
def pyrdown(input: Tensor, border_type: str = "reflect", align_corners: bool = False, factor: float = 2.0) -> Tensor:
    r"""Blur a tensor and downsamples it.

    .. image:: _static/img/pyrdown.png

    Args:
        input: the tensor to be downsampled.
        border_type: the padding mode to be applied before convolving.
          The expected modes are: ``'constant'``, ``'reflect'``,
          ``'replicate'`` or ``'circular'``.
        align_corners: interpolation flag.
        factor: the downsampling factor

    Return:
        the downsampled tensor.

    Examples:
        >>> input = torch.arange(16, dtype=torch.float32).reshape(1, 1, 4, 4)
        >>> pyrdown(input, align_corners=True)
        tensor([[[[ 3.7500,  5.2500],
                  [ 9.7500, 11.2500]]]])
    """
    KORNIA_CHECK_SHAPE(input, ["B", "C", "H", "W"])

    kernel: Tensor = _get_pyramid_gaussian_kernel()
    _, _, height, width = input.shape
    # blur image
    x_blur: Tensor = filter2d(input, kernel, border_type)

    # TODO: use kornia.geometry.resize/rescale
    # downsample.
    out: Tensor = t.interpolate(
        x_blur,
        size=(int(float(height) / factor), int(float(width) // factor)),
        mode="bilinear",
        align_corners=align_corners,
    )
    return out
//  [python]
def pyrdown(input: Tensor, border_type: str = "reflect", align_corners: bool = False, factor: float = 2.0) -> Tensor:
    r"""Blur a tensor and downsamples it.

    .. image:: _static/img/pyrdown.png

    Args:
        input: the tensor to be downsampled.
        border_type: the padding mode to be applied before convolving.
          The expected modes are: ``'constant'``, ``'reflect'``,
          ``'replicate'`` or ``'circular'``.
        align_corners: interpolation flag.
        factor: the downsampling factor

    Return:
        the downsampled tensor.

    Examples:
        >>> input = torch.arange(16, dtype=torch.float32).reshape(1, 1, 4, 4)
        >>> pyrdown(input, align_corners=True)
        tensor([[[[ 3.7500,  5.2500],
                  [ 9.7500, 11.2500]]]])
    """
    KORNIA_CHECK_SHAPE(input, ["B", "C", "H", "W"])

    kernel: Tensor = _get_pyramid_gaussian_kernel()
    _, _, height, width = input.shape
    # blur image
    x_blur: Tensor = filter2d(input, kernel, border_type)

    # TODO: use kornia.geometry.resize/rescale
    # downsample.
    out: Tensor = F.interpolate(
        x_blur,
        size=(int(float(height) / factor), int(float(width) // factor)),
        mode="bilinear",
        align_corners=align_corners,
    )
    return out
//  [python]
def pyrup(input: Tensor, border_type: str = "reflect", align_corners: bool = False) -> Tensor:
    r"""Upsample a tensor and then blurs it.

    .. image:: _static/img/pyrup.png

    Args:
        input: the tensor to be downsampled.
        border_type: the padding mode to be applied before convolving.
          The expected modes are: ``'constant'``, ``'reflect'``, ``'replicate'`` or ``'circular'``.
        align_corners: interpolation flag.

    Return:
        the downsampled tensor.

    Examples:
        >>> input = torch.arange(4, dtype=torch.float32).reshape(1, 1, 2, 2)
        >>> pyrup(input, align_corners=True)
        tensor([[[[0.7500, 0.8750, 1.1250, 1.2500],
                  [1.0000, 1.1250, 1.3750, 1.5000],
                  [1.5000, 1.6250, 1.8750, 2.0000],
                  [1.7500, 1.8750, 2.1250, 2.2500]]]])
    """
    KORNIA_CHECK_SHAPE(input, ["B", "C", "H", "W"])

    kernel: Tensor = _get_pyramid_gaussian_kernel()
    # upsample tensor
    _, _, height, width = input.shape
    # TODO: use kornia.geometry.resize/rescale
    x_up: Tensor = t.interpolate(input, size=(height * 2, width * 2), mode="bilinear", align_corners=align_corners)

    # blurs upsampled tensor
    x_blur: Tensor = filter2d(x_up, kernel, border_type)
    return x_blur
//  [python]
def pyrup(input: Tensor, border_type: str = "reflect", align_corners: bool = False) -> Tensor:
    r"""Upsample a tensor and then blurs it.

    .. image:: _static/img/pyrup.png

    Args:
        input: the tensor to be downsampled.
        border_type: the padding mode to be applied before convolving.
          The expected modes are: ``'constant'``, ``'reflect'``, ``'replicate'`` or ``'circular'``.
        align_corners: interpolation flag.

    Return:
        the downsampled tensor.

    Examples:
        >>> input = torch.arange(4, dtype=torch.float32).reshape(1, 1, 2, 2)
        >>> pyrup(input, align_corners=True)
        tensor([[[[0.7500, 0.8750, 1.1250, 1.2500],
                  [1.0000, 1.1250, 1.3750, 1.5000],
                  [1.5000, 1.6250, 1.8750, 2.0000],
                  [1.7500, 1.8750, 2.1250, 2.2500]]]])
    """
    KORNIA_CHECK_SHAPE(input, ["B", "C", "H", "W"])

    kernel: Tensor = _get_pyramid_gaussian_kernel()
    # upsample tensor
    _, _, height, width = input.shape
    # TODO: use kornia.geometry.resize/rescale
    x_up: Tensor = F.interpolate(input, size=(height * 2, width * 2), mode="bilinear", align_corners=align_corners)

    # blurs upsampled tensor
    x_blur: Tensor = filter2d(x_up, kernel, border_type)
    return x_blur
//  [python]
def forward(
        self,
        input: torch.Tensor,
        conditions: tp.Dict[str, ConditionType]
    ) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
        """Fuse the conditions to the provided model input.

        Args:
            input (torch.Tensor): Transformer input.
            conditions (dict[str, ConditionType]): Dict of conditions.
        Returns:
            tuple[torch.Tensor, torch.Tensor]: The first tensor is the transformer input
                after the conditions have been fused. The second output tensor is the tensor
                used for cross-attention or None if no cross attention inputs exist.
        """
        B, T, _ = input.shape

        if 'offsets' in self._streaming_state:
            first_step = False
            offsets = self._streaming_state['offsets']
        else:
            first_step = True
            offsets = torch.zeros(input.shape[0], dtype=torch.long, device=input.device)

        assert set(conditions.keys()).issubset(set(self.cond2fuse.keys())), \
            f"given conditions contain unknown attributes for fuser, " \
            f"expected {self.cond2fuse.keys()}, got {conditions.keys()}"
        cross_attention_output = None
        for cond_type, (cond, cond_mask) in conditions.items():
            op = self.cond2fuse[cond_type]
            if op == 'sum':
                input += cond
            elif op == 'input_interpolate':
                cond = einops.rearrange(cond, "b t d -> b d t")
                cond = F.interpolate(cond, size=input.shape[1], mode='linear', align_corners=False)
                input += einops.rearrange(cond, "b d t -> b t d")
            elif op == 'prepend':
                if first_step:
                    input = torch.cat([cond, input], dim=1)
            elif op == 'cross':
                if cross_attention_output is not None:
                    cross_attention_output = torch.cat([cross_attention_output, cond], dim=1)
                else:
                    cross_attention_output = cond
            else:
                raise ValueError(f"unknown op ({op})")

        if self.cross_attention_pos_emb and cross_attention_output is not None:
            positions = torch.arange(
                cross_attention_output.shape[1],
                device=cross_attention_output.device
            ).view(1, -1, 1)
            pos_emb = create_sin_embedding(positions, cross_attention_output.shape[-1])
            cross_attention_output = cross_attention_output + self.cross_attention_pos_emb_scale * pos_emb

        if self._is_streaming:
            self._streaming_state['offsets'] = offsets + T

        return input, cross_attention_output
//  [python]
def forward(
        self,
        input: torch.Tensor,
        conditions: tp.Dict[str, ConditionType]
    ) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
        """Fuse the conditions to the provided model input.

        Args:
            input (torch.Tensor): Transformer input.
            conditions (dict[str, ConditionType]): Dict of conditions.
        Returns:
            tuple[torch.Tensor, torch.Tensor]: The first tensor is the transformer input
                after the conditions have been fused. The second output tensor is the tensor
                used for cross-attention or None if no cross attention inputs exist.
        """
        B, T, _ = input.shape

        if 'offsets' in self._streaming_state:
            first_step = False
            offsets = self._streaming_state['offsets']
        else:
            first_step = True
            offsets = torch.zeros(input.shape[0], dtype=torch.long, device=input.device)

        assert set(conditions.keys()).issubset(set(self.cond2fuse.keys())), \
            f"given conditions contain unknown attributes for fuser, " \
            f"expected {self.cond2fuse.keys()}, got {conditions.keys()}"
        cross_attention_output = None
        for cond_type, (cond, cond_mask) in conditions.items():
            op = self.cond2fuse[cond_type]
            if op == 'sum':
                input += cond
            elif op == 'input_interpolate':
                cond = einops.rearrange(cond, "b t d -> b d t")
                cond = F.interpolate(cond, size=input.shape[1])
                input += einops.rearrange(cond, "b d t -> b t d")
            elif op == 'prepend':
                if first_step:
                    input = torch.cat([cond, input], dim=1)
            elif op == 'cross':
                if cross_attention_output is not None:
                    cross_attention_output = torch.cat([cross_attention_output, cond], dim=1)
                else:
                    cross_attention_output = cond
            else:
                raise ValueError(f"unknown op ({op})")

        if self.cross_attention_pos_emb and cross_attention_output is not None:
            positions = torch.arange(
                cross_attention_output.shape[1],
                device=cross_attention_output.device
            ).view(1, -1, 1)
            pos_emb = create_sin_embedding(positions, cross_attention_output.shape[-1])
            cross_attention_output = cross_attention_output + self.cross_attention_pos_emb_scale * pos_emb

        if self._is_streaming:
            self._streaming_state['offsets'] = offsets + T

        return input, cross_attention_output
//  [python]
def forward(self, x):
        b, _, h, w, p = *x.size(), self.patch_size
        assert h % p == 0 and w % p == 0, f'Image size ({w}, {h}) is not divisible by patch size ({p}, {p})'
        hp, wp, grid = h // p, w // p, self.image_size // p

        # embeddings
        pos_embedding = torch.cat([
            self.pos_embedding[:, :1],
            t.interpolate(
                self.pos_embedding[:, 1:].reshape(1, grid, grid, -1).permute(
                    0, 3, 1, 2),
                size=(hp, wp),
                mode='bilinear',
                align_corners=False).permute(0, 2, 3, 1).reshape(
                    1, hp * wp, -1)
        ],
                                  dim=1)  # noqa
        x = self.patch_embedding(x).flatten(2).permute(0, 2, 1)
        x = torch.cat([self.cls_embedding.repeat(b, 1, 1), x], dim=1)
        x = x + pos_embedding

        # stage1
        x = self.blocks[self.slices[0]](x)
        x1 = torch.cat([x[:, 1:], x[:, :1].expand_as(x[:, 1:])], dim=-1)
        x1 = self.fc1(x1).permute(0, 2, 1).unflatten(2, (hp, wp))
        x1 = self.conv1(x1)

        # stage2
        x = self.blocks[self.slices[1]](x)
        x2 = torch.cat([x[:, 1:], x[:, :1].expand_as(x[:, 1:])], dim=-1)
        x2 = self.fc2(x2).permute(0, 2, 1).unflatten(2, (hp, wp))
        x2 = self.conv2(x2)

        # stage3
        x = self.blocks[self.slices[2]](x)
        x3 = torch.cat([x[:, 1:], x[:, :1].expand_as(x[:, 1:])], dim=-1)
        x3 = self.fc3(x3).permute(0, 2, 1).unflatten(2, (hp, wp))
        x3 = self.conv3(x3)

        # stage4
        x = self.blocks[self.slices[3]](x)
        x4 = torch.cat([x[:, 1:], x[:, :1].expand_as(x[:, 1:])], dim=-1)
        x4 = self.fc4(x4).permute(0, 2, 1).unflatten(2, (hp, wp))
        x4 = self.conv4(x4)

        # fusion
        x4 = self.fusion4(x4)
        x3 = self.fusion3(x4, x3)
        x2 = self.fusion2(x3, x2)
        x1 = self.fusion1(x2, x1)

        # head
        x = self.head(x1)
        return x
//  [python]
def forward(self, x):
        b, _, h, w, p = *x.size(), self.patch_size
        assert h % p == 0 and w % p == 0, f'Image size ({w}, {h}) is not divisible by patch size ({p}, {p})'
        hp, wp, grid = h // p, w // p, self.image_size // p

        # embeddings
        pos_embedding = torch.cat([
            self.pos_embedding[:, :1],
            F.interpolate(
                self.pos_embedding[:, 1:].reshape(1, grid, grid, -1).permute(
                    0, 3, 1, 2),
                size=(hp, wp),
                mode='bilinear',
                align_corners=False).permute(0, 2, 3, 1).reshape(
                    1, hp * wp, -1)
        ],
                                  dim=1)  # noqa
        x = self.patch_embedding(x).flatten(2).permute(0, 2, 1)
        x = torch.cat([self.cls_embedding.repeat(b, 1, 1), x], dim=1)
        x = x + pos_embedding

        # stage1
        x = self.blocks[self.slices[0]](x)
        x1 = torch.cat([x[:, 1:], x[:, :1].expand_as(x[:, 1:])], dim=-1)
        x1 = self.fc1(x1).permute(0, 2, 1).unflatten(2, (hp, wp))
        x1 = self.conv1(x1)

        # stage2
        x = self.blocks[self.slices[1]](x)
        x2 = torch.cat([x[:, 1:], x[:, :1].expand_as(x[:, 1:])], dim=-1)
        x2 = self.fc2(x2).permute(0, 2, 1).unflatten(2, (hp, wp))
        x2 = self.conv2(x2)

        # stage3
        x = self.blocks[self.slices[2]](x)
        x3 = torch.cat([x[:, 1:], x[:, :1].expand_as(x[:, 1:])], dim=-1)
        x3 = self.fc3(x3).permute(0, 2, 1).unflatten(2, (hp, wp))
        x3 = self.conv3(x3)

        # stage4
        x = self.blocks[self.slices[3]](x)
        x4 = torch.cat([x[:, 1:], x[:, :1].expand_as(x[:, 1:])], dim=-1)
        x4 = self.fc4(x4).permute(0, 2, 1).unflatten(2, (hp, wp))
        x4 = self.conv4(x4)

        # fusion
        x4 = self.fusion4(x4)
        x3 = self.fusion3(x4, x3)
        x2 = self.fusion2(x3, x2)
        x1 = self.fusion1(x2, x1)

        # head
        x = self.head(x1)
        return x
//  [python]
def resize_video_to(
    video,
    target_image_size,
    target_frames = None,
    clamp_range = None,
    mode = 'nearest'
):
    orig_video_size = video.shape[-1]

    frames = video.shape[2]
    target_frames = default(target_frames, frames)

    target_shape = (target_frames, target_image_size, target_image_size)

    if tuple(video.shape[-3:]) == target_shape:
        return video

    out = t.interpolate(video, target_shape, mode = mode)

    if exists(clamp_range):
        out = out.clamp(*clamp_range)
        
    return out
//  [python]
def resize_video_to(
    video,
    target_image_size,
    target_frames = None,
    clamp_range = None,
    mode = 'nearest'
):
    orig_video_size = video.shape[-1]

    frames = video.shape[2]
    target_frames = default(target_frames, frames)

    target_shape = (target_frames, target_image_size, target_image_size)

    if tuple(video.shape[-3:]) == target_shape:
        return video

    out = F.interpolate(video, target_shape, mode = mode)

    if exists(clamp_range):
        out = out.clamp(*clamp_range)
        
    return out
//  [python]
def forward(self, xt, xl):
        _, _, h, w = xl.size()
        xt = F.upsample(self.conv1(xt), size=(h, w), mode="bilinear")
        xl = self.conv2(xl)
        x = torch.cat([xt, xl], dim=1)
        x = self.conv3(x)
        seg = self.conv4(x)
        return seg, x
//  [python]
def forward(self, xt, xl):
        _, _, h, w = xl.size()
        xt = F.interpolate(self.conv1(xt), size=(h, w), mode="bilinear", align_corners=True)
        xl = self.conv2(xl)
        x = torch.cat([xt, xl], dim=1)
        x = self.conv3(x)
        seg = self.conv4(x)
        return seg, x
//  [python]
def test_upsamplingBicubic2d(self):
        # test output against known input: align_corners=False result must match opencv
        in_t = torch.arange(8.).view(1, 2, 2, 2)
        expected_out_t = torch.tensor(
            [[[[-0.31641, 0.01562, 0.56250, 0.89453],
              [0.34766, 0.67969, 1.22656, 1.55859],
              [1.44141, 1.77344, 2.32031, 2.65234],
              [2.10547, 2.43750, 2.98438, 3.31641]],

             [[3.68359, 4.01562, 4.56250, 4.89453],
              [4.34766, 4.67969, 5.22656, 5.55859],
              [5.44141, 5.77344, 6.32031, 6.65234],
              [6.10547, 6.43750, 6.98438, 7.31641]]]])
        out_t = t.nn.functional.interpolate(in_t, scale_factor=2, mode='bicubic', align_corners=False)
        torch.set_printoptions(precision=5)
        self.assertEqual(out_t, expected_out_t, atol=1e-5, rtol=0)


        device_list = ['cpu']
        if TEST_CUDA:
            device_list.append('cuda')

        for align_corners in [True, False]:
            kwargs = dict(mode='bicubic', align_corners=align_corners)
            # test float scale factor up & downsampling
            for device in device_list:
                for scale_factor in [0.5, 1, 1.5, 2]:
                    in_t = torch.ones(2, 2, 2, 2).to(device)
                    out_t = t.nn.functional.interpolate(in_t, scale_factor=scale_factor, **kwargs)
                    out_size = int(math.floor(in_t.shape[-1] * scale_factor))
                    self.assertEqual(torch.ones(2, 2, out_size, out_size), out_t.data,
                                     atol=1e-5, rtol=0)

                    input = torch.randn(2, 2, 2, 2, requires_grad=True)
                    gradcheck(lambda x: t.nn.functional.interpolate(x, out_size, **kwargs), [input])
//  [python]
def test_upsamplingBicubic2d(self):
        # test output against known input: align_corners=False result must match opencv
        in_t = torch.arange(8.).view(1, 2, 2, 2)
        expected_out_t = torch.tensor(
            [[[[-0.31641, 0.01562, 0.56250, 0.89453],
              [0.34766, 0.67969, 1.22656, 1.55859],
              [1.44141, 1.77344, 2.32031, 2.65234],
              [2.10547, 2.43750, 2.98438, 3.31641]],

             [[3.68359, 4.01562, 4.56250, 4.89453],
              [4.34766, 4.67969, 5.22656, 5.55859],
              [5.44141, 5.77344, 6.32031, 6.65234],
              [6.10547, 6.43750, 6.98438, 7.31641]]]])
        out_t = F.interpolate(in_t, scale_factor=2, mode='bicubic', align_corners=False)
        torch.set_printoptions(precision=5)
        self.assertEqual(out_t, expected_out_t, atol=1e-5, rtol=0)


        device_list = ['cpu']
        if TEST_CUDA:
            device_list.append('cuda')

        for align_corners in [True, False]:
            kwargs = dict(mode='bicubic', align_corners=align_corners)
            # test float scale factor up & downsampling
            for device in device_list:
                for scale_factor in [0.5, 1, 1.5, 2]:
                    in_t = torch.ones(2, 2, 2, 2).to(device)
                    out_t = F.interpolate(in_t, scale_factor=scale_factor, **kwargs)
                    out_size = int(math.floor(in_t.shape[-1] * scale_factor))
                    self.assertEqual(torch.ones(2, 2, out_size, out_size), out_t.data,
                                     atol=1e-5, rtol=0)

                    input = torch.randn(2, 2, 2, 2, requires_grad=True)
                    gradcheck(lambda x: F.interpolate(x, out_size, **kwargs), [input])
//  [python]
def forward(self, x1, x2, x3):
        _, _, h, w = x1.size()

        edge1_fea = self.conv1(x1)
        edge1 = self.conv4(edge1_fea)
        edge2_fea = self.conv2(x2)
        edge2 = self.conv4(edge2_fea)
        edge3_fea = self.conv3(x3)
        edge3 = self.conv4(edge3_fea)

        edge2_fea = F.upsample(edge2_fea, size=(h, w), mode="bilinear")
        edge3_fea = F.upsample(edge3_fea, size=(h, w), mode="bilinear")
        edge2 = F.upsample(edge2, size=(h, w), mode="bilinear")
        edge3 = F.upsample(edge3, size=(h, w), mode="bilinear")

        edge = torch.cat([edge1, edge2, edge3], dim=1)
        edge_fea = torch.cat([edge1_fea, edge2_fea, edge3_fea], dim=1)
        edge = self.conv5(edge)

        return edge, edge_fea
//  [python]
def forward(self, x1, x2, x3):
        _, _, h, w = x1.size()

        edge1_fea = self.conv1(x1)
        edge1 = self.conv4(edge1_fea)
        edge2_fea = self.conv2(x2)
        edge2 = self.conv4(edge2_fea)
        edge3_fea = self.conv3(x3)
        edge3 = self.conv4(edge3_fea)

        edge2_fea = F.interpolate(edge2_fea, size=(h, w), mode="bilinear", align_corners=True)
        edge3_fea = F.interpolate(edge3_fea, size=(h, w), mode="bilinear", align_corners=True)
        edge2 = F.interpolate(edge2, size=(h, w), mode="bilinear", align_corners=True)
        edge3 = F.interpolate(edge3, size=(h, w), mode="bilinear", align_corners=True)

        edge = torch.cat([edge1, edge2, edge3], dim=1)
        edge_fea = torch.cat([edge1_fea, edge2_fea, edge3_fea], dim=1)
        edge = self.conv5(edge)

        return edge, edge_fea
//  [python]
def muti_loss_fusion(preds, target):
    loss0 = 0.0
    loss = 0.0

    for i in range(0,len(preds)):
        # print("i: ", i, preds[i].shape)
        if(preds[i].shape[2]!=target.shape[2] or preds[i].shape[3]!=target.shape[3]):
            # tmp_target = _upsample_like(target,preds[i])
            tmp_target = t.interpolate(target, size=preds[i].size()[2:], mode='bilinear', align_corners=True)
            loss = loss + bce_loss(preds[i],tmp_target)
        else:
            loss = loss + bce_loss(preds[i],target)
        if(i==0):
            loss0 = loss
    return loss0, loss
//  [python]
def muti_loss_fusion(preds, target):
    loss0 = 0.0
    loss = 0.0

    for i in range(0,len(preds)):
        # print("i: ", i, preds[i].shape)
        if(preds[i].shape[2]!=target.shape[2] or preds[i].shape[3]!=target.shape[3]):
            # tmp_target = _upsample_like(target,preds[i])
            tmp_target = F.interpolate(target, size=preds[i].size()[2:], mode='bilinear', align_corners=True)
            loss = loss + bce_loss(preds[i],tmp_target)
        else:
            loss = loss + bce_loss(preds[i],target)
        if(i==0):
            loss0 = loss
    return loss0, loss
//  [python]
def preprocess_single(video, resolution, sequence_length=None):
    # video: THWC, {0, ..., 255}
    video = video.permute(0, 3, 1, 2).float() / 255. # TCHW
    t, c, h, w = video.shape

    # temporal crop
    if sequence_length is not None:
        assert sequence_length <= t
        video = video[:sequence_length]

    # scale shorter side to resolution
    scale = resolution / min(h, w)
    if h < w:
        target_size = (resolution, math.ceil(w * scale))
    else:
        target_size = (math.ceil(h * scale), resolution)
    video = t(video, size=target_size, mode='bilinear',
                          align_corners=False)

    # center crop
    t, c, h, w = video.shape
    w_start = (w - resolution) // 2
    h_start = (h - resolution) // 2
    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]
    video = video.permute(1, 0, 2, 3).contiguous() # CTHW

    video -= 0.5

    return video
//  [python]
def preprocess_single(video, resolution, sequence_length=None):
    # video: THWC, {0, ..., 255}
    video = video.permute(0, 3, 1, 2).float() / 255. # TCHW
    t, c, h, w = video.shape

    # temporal crop
    if sequence_length is not None:
        assert sequence_length <= t
        video = video[:sequence_length]

    # scale shorter side to resolution
    scale = resolution / min(h, w)
    if h < w:
        target_size = (resolution, math.ceil(w * scale))
    else:
        target_size = (math.ceil(h * scale), resolution)
    video = F.interpolate(video, size=target_size, mode='bilinear',
                          align_corners=False)

    # center crop
    t, c, h, w = video.shape
    w_start = (w - resolution) // 2
    h_start = (h - resolution) // 2
    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]
    video = video.permute(1, 0, 2, 3).contiguous() # CTHW

    video -= 0.5

    return video
//  [python]
def test_upsamplingBilinear2d(self, device):
        for align_corners in [True, False]:
            kwargs = dict(mode='bilinear', align_corners=align_corners)
            for memory_format in [torch.contiguous_format, torch.channels_last]:
                # test float scale factor up & downsampling
                for scale_factor in [0.5, 1.5, 2]:
                    in_t = torch.ones(1, 2, 2, 2, device=device).contiguous(memory_format=memory_format).requires_grad_()
                    out_size = int(math.floor(in_t.shape[-1] * scale_factor))
                    with warnings.catch_warnings(record=True) as w:
                        out_t = t.upsample_bilinear(in_t, scale_factor=scale_factor, **kwargs)
                    self.assertEqual(torch.ones(1, 2, out_size, out_size, device=device), out_t.data)
                    # Assert that memory format is carried through to the output
                    self.assertTrue(out_t.is_contiguous(memory_format=memory_format))
                    out_t.backward(torch.randn_like(out_t))
                    self.assertTrue(in_t.grad.is_contiguous(memory_format=memory_format))

                    input = torch.randn(1, 2, 2, 2, device=device).contiguous(memory_format=memory_format).requires_grad_()
                    gradcheck(lambda x: t.upsample_bilinear(x, out_size, **kwargs), [input])

                    # Assert that cpu and cuda give same results
                    if torch.device(device).type == 'cuda':
                        for shapes in [
                            (2, 2, 3, 4), (2, 3, 4, 5), (3, 1, 2, 2), (1, 5, 3, 2)
                        ]:
                            a_cuda = torch.randn(*shapes, device=device).contiguous(memory_format=memory_format).requires_grad_()
                            a_cpu = a_cuda.detach().cpu().requires_grad_()

                            with warnings.catch_warnings(record=True):
                                out_cuda = t.upsample_bilinear(a_cuda, scale_factor=scale_factor, **kwargs)
                                out_cpu = t.upsample_bilinear(a_cpu, scale_factor=scale_factor, **kwargs)

                            self.assertEqual(out_cpu.cuda(), out_cuda)

                            g_cuda = torch.randn_like(out_cuda)
                            g_cpu = g_cuda.cpu()

                            out_cuda.backward(g_cuda)
                            out_cpu.backward(g_cpu)

                            self.assertEqual(a_cuda.grad, a_cpu.grad)
//  [python]
def test_upsamplingBilinear2d(self, device):
        for align_corners in [True, False]:
            kwargs = dict(mode='bilinear', align_corners=align_corners)
            for memory_format in [torch.contiguous_format, torch.channels_last]:
                # test float scale factor up & downsampling
                for scale_factor in [0.5, 1.5, 2]:
                    in_t = torch.ones(1, 2, 2, 2, device=device).contiguous(memory_format=memory_format).requires_grad_()
                    out_size = int(math.floor(in_t.shape[-1] * scale_factor))
                    with warnings.catch_warnings(record=True) as w:
                        out_t = F.interpolate(in_t, scale_factor=scale_factor, **kwargs)
                    self.assertEqual(torch.ones(1, 2, out_size, out_size, device=device), out_t.data)
                    # Assert that memory format is carried through to the output
                    self.assertTrue(out_t.is_contiguous(memory_format=memory_format))
                    out_t.backward(torch.randn_like(out_t))
                    self.assertTrue(in_t.grad.is_contiguous(memory_format=memory_format))

                    input = torch.randn(1, 2, 2, 2, device=device).contiguous(memory_format=memory_format).requires_grad_()
                    gradcheck(lambda x: F.interpolate(x, out_size, **kwargs), [input])

                    # Assert that cpu and cuda give same results
                    if torch.device(device).type == 'cuda':
                        for shapes in [
                            (2, 2, 3, 4), (2, 3, 4, 5), (3, 1, 2, 2), (1, 5, 3, 2)
                        ]:
                            a_cuda = torch.randn(*shapes, device=device).contiguous(memory_format=memory_format).requires_grad_()
                            a_cpu = a_cuda.detach().cpu().requires_grad_()

                            with warnings.catch_warnings(record=True):
                                out_cuda = F.interpolate(a_cuda, scale_factor=scale_factor, **kwargs)
                                out_cpu = F.interpolate(a_cpu, scale_factor=scale_factor, **kwargs)

                            self.assertEqual(out_cpu.cuda(), out_cuda)

                            g_cuda = torch.randn_like(out_cuda)
                            g_cpu = g_cuda.cpu()

                            out_cuda.backward(g_cuda)
                            out_cpu.backward(g_cpu)

                            self.assertEqual(a_cuda.grad, a_cpu.grad)
//  [python]
def _upsample_like(src, tar):
    src = t.interpolate(src, size=tar.shape[2:], mode="bilinear")

    return src
//  [python]
def _upsample_like(src, tar):
    src = F.interpolate(src, size=tar.shape[2:], mode="bilinear")

    return src
//  [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 inpainting"
        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)
        c_cat = list()
        for ck in self.concat_keys:
            if self.use_fp16:
                cc = rearrange(batch[ck], "b h w c -> b c h w").to(memory_format=torch.contiguous_format).half()
            else:
                cc = rearrange(batch[ck], "b h w c -> b c h w").to(memory_format=torch.contiguous_format).float()
            if bs is not None:
                cc = cc[:bs]
                cc = cc.to(self.device)
            bchw = z.shape
            if ck != self.masked_image_key:
                cc = t.interpolate(cc, size=bchw[-2:])
            else:
                cc = self.get_first_stage_encoding(self.encode_first_stage(cc))
            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 inpainting"
        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)
        c_cat = list()
        for ck in self.concat_keys:
            if self.use_fp16:
                cc = rearrange(batch[ck], "b h w c -> b c h w").to(memory_format=torch.contiguous_format).half()
            else:
                cc = rearrange(batch[ck], "b h w c -> b c h w").to(memory_format=torch.contiguous_format).float()
            if bs is not None:
                cc = cc[:bs]
                cc = cc.to(self.device)
            bchw = z.shape
            if ck != self.masked_image_key:
                cc = torch.nn.functional.interpolate(cc, size=bchw[-2:])
            else:
                cc = self.get_first_stage_encoding(self.encode_first_stage(cc))
            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 muti_loss_fusion_kl(preds, target, dfs, fs, mode='MSE'):
    loss0 = 0.0
    loss = 0.0

    for i in range(0,len(preds)):
        # print("i: ", i, preds[i].shape)
        if(preds[i].shape[2]!=target.shape[2] or preds[i].shape[3]!=target.shape[3]):
            # tmp_target = _upsample_like(target,preds[i])
            tmp_target = t.nn.functional.interpolate(target, size=preds[i].size()[2:], mode='bilinear', align_corners=True)
            loss = loss + bce_loss(preds[i],tmp_target)
        else:
            loss = loss + bce_loss(preds[i],target)
        if(i==0):
            loss0 = loss

    for i in range(0,len(dfs)):
        if(mode=='MSE'):
            loss = loss + fea_loss(dfs[i],fs[i]) ### add the mse loss of features as additional constraints
            # print("fea_loss: ", fea_loss(dfs[i],fs[i]).item())
        elif(mode=='KL'):
            loss = loss + kl_loss(F.log_softmax(dfs[i],dim=1),F.softmax(fs[i],dim=1))
            # print("kl_loss: ", kl_loss(F.log_softmax(dfs[i],dim=1),F.softmax(fs[i],dim=1)).item())
        elif(mode=='MAE'):
            loss = loss + l1_loss(dfs[i],fs[i])
            # print("ls_loss: ", l1_loss(dfs[i],fs[i]))
        elif(mode=='SmoothL1'):
            loss = loss + smooth_l1_loss(dfs[i],fs[i])
            # print("SmoothL1: ", smooth_l1_loss(dfs[i],fs[i]).item())

    return loss0, loss
//  [python]
def muti_loss_fusion_kl(preds, target, dfs, fs, mode='MSE'):
    loss0 = 0.0
    loss = 0.0

    for i in range(0,len(preds)):
        # print("i: ", i, preds[i].shape)
        if(preds[i].shape[2]!=target.shape[2] or preds[i].shape[3]!=target.shape[3]):
            # tmp_target = _upsample_like(target,preds[i])
            tmp_target = F.interpolate(target, size=preds[i].size()[2:], mode='bilinear', align_corners=True)
            loss = loss + bce_loss(preds[i],tmp_target)
        else:
            loss = loss + bce_loss(preds[i],target)
        if(i==0):
            loss0 = loss

    for i in range(0,len(dfs)):
        if(mode=='MSE'):
            loss = loss + fea_loss(dfs[i],fs[i]) ### add the mse loss of features as additional constraints
            # print("fea_loss: ", fea_loss(dfs[i],fs[i]).item())
        elif(mode=='KL'):
            loss = loss + kl_loss(F.log_softmax(dfs[i],dim=1),F.softmax(fs[i],dim=1))
            # print("kl_loss: ", kl_loss(F.log_softmax(dfs[i],dim=1),F.softmax(fs[i],dim=1)).item())
        elif(mode=='MAE'):
            loss = loss + l1_loss(dfs[i],fs[i])
            # print("ls_loss: ", l1_loss(dfs[i],fs[i]))
        elif(mode=='SmoothL1'):
            loss = loss + smooth_l1_loss(dfs[i],fs[i])
            # print("SmoothL1: ", smooth_l1_loss(dfs[i],fs[i]).item())

    return loss0, loss
//  [python]
def forward(self, x, sampleidx=None):
        trimapmask = x[:, 3:, :, :]
        # set outputs
        outs = []

        # get outs[0]
        # outs.append(self.shortcut[0](x))

        # forward patch-embedding layer
        x = self.patch_embed(x)
        _, _, Wh, Ww = x.shape
        trimapmask = F.upsample(trimapmask, scale_factor=1/4, mode='nearest')

        # get outs[1]
        # outs.append(self.shortcut[1](F.upsample_bilinear(x, scale_factor=2.0)))

        # dropout
        x = x.flatten(2).transpose(1, 2)
        x = self.pos_drop(x)

        # get outs[2~5]
        for i in range(self.num_layers):
            layer = self.layers[i]

            trimapmask_ = F.upsample(trimapmask, scale_factor=1/(pow(2, i)), mode='nearest')
            area_fg = trimapmask_[:, 0, :, :].unsqueeze(-1)  # background area
            area_bg = trimapmask_[:, 2, :, :].unsqueeze(-1)  # foreground area
            area_uk = trimapmask_[:, 1, :, :].unsqueeze(-1)  # unknown area

            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww, area_fg=area_fg, area_bg=area_bg, area_uk=area_uk)

            # if i in self.out_indices:
            #     norm_layer = getattr(self, f'norm{i}')
            #     x_out = norm_layer(x_out)
            #     out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
                # out = self.shortcut[i+2](out)
                # outs.append(out)
        out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
        # return tuple(outs) # 增加一层upsample
        return out