jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def random_resized_crop(
    images,
    target_height,
    target_width,
    scale=(0.8, 1.0),
    ratio=(3.0 / 4.0, 4.0 / 3.0),
):
    """
    Crop the given images to random size and aspect ratio. A crop of random
    size (default: of 0.08 to 1.0) of the original size and a random aspect
    ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This
    crop is finally resized to given size. This is popularly used to train the
    Inception networks.

    Args:
        images: Images to perform resizing and cropping.
        target_height: Desired height after cropping.
        target_width: Desired width after cropping.
        scale: Scale range of Inception-style area based random resizing.
        ratio: Aspect ratio range of Inception-style area based random resizing.
    """

    height = images.shape[2]
    width = images.shape[3]

    i, j, h, w = _get_param_spatial_crop(scale, ratio, height, width)
    cropped = images[:, :, i : i + h, j : j + w]
    return t.interpolate(
        cropped,
        size=(target_height, target_width),
        mode="bilinear",
        align_corners=False,
    )
//  [python]
def random_resized_crop(
    images,
    target_height,
    target_width,
    scale=(0.8, 1.0),
    ratio=(3.0 / 4.0, 4.0 / 3.0),
):
    """
    Crop the given images to random size and aspect ratio. A crop of random
    size (default: of 0.08 to 1.0) of the original size and a random aspect
    ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This
    crop is finally resized to given size. This is popularly used to train the
    Inception networks.

    Args:
        images: Images to perform resizing and cropping.
        target_height: Desired height after cropping.
        target_width: Desired width after cropping.
        scale: Scale range of Inception-style area based random resizing.
        ratio: Aspect ratio range of Inception-style area based random resizing.
    """

    height = images.shape[2]
    width = images.shape[3]

    i, j, h, w = _get_param_spatial_crop(scale, ratio, height, width)
    cropped = images[:, :, i : i + h, j : j + w]
    return torch.nn.functional.interpolate(
        cropped,
        size=(target_height, target_width),
        mode="bilinear",
        align_corners=False,
    )
//  [python]
def forward(self, 
                input_tensor, 
                guide_plus, 
                guide_cat, 
                bbox=None, 
                fine_feat_crop=None, 
                coarse_feat_whole=None, 
                coarse_feat_whole_hack=None, 
                coarse_feat_crop=None):

        # apply unscaled feat to swin
        if coarse_feat_whole_hack is not None:
            coarse_feat_whole = coarse_feat_whole_hack

        feat_list = []
        
        x = self.inc(input_tensor)
        feat_list.append(x)
        
        for layer in self.down_conv_list:
            x = layer(x)
            feat_list.append(x)
        
        output = []
        feat_inv_list = feat_list[::-1]
        for idx, (feat_enc, feat_c) in enumerate(zip(feat_inv_list, coarse_feat_whole)):
            
            # in case for depth-anything
            _, _, h, w = feat_enc.shape
            if h != feat_c.shape[-2] or w != feat_c.shape[-1]:
                feat_enc = t.interpolate(feat_enc, size=feat_c.shape[-2:], mode='bilinear', align_corners=True)
                
            if idx == 0:
                pass
            else:
                feat_enc = self.up_conv_list[idx-1](torch.cat([temp_feat, guide_cat[idx-1]], dim=1), feat_enc)
            
            _, _, h, w = feat_c.shape
            feat_c = self.g2l_list[idx](feat_c, None)
            feat_c = torch_roi_align(feat_c, bbox, (h, w), h/self.patch_process_shape[0], aligned=True)
            x = self.convs[idx](torch.cat([feat_enc, feat_c], dim=1))
            temp_feat = x
            output.append(x)

        return output[::-1]
//  [python]
def forward(self, 
                input_tensor, 
                guide_plus, 
                guide_cat, 
                bbox=None, 
                fine_feat_crop=None, 
                coarse_feat_whole=None, 
                coarse_feat_whole_hack=None, 
                coarse_feat_crop=None):

        # apply unscaled feat to swin
        if coarse_feat_whole_hack is not None:
            coarse_feat_whole = coarse_feat_whole_hack

        feat_list = []
        
        x = self.inc(input_tensor)
        feat_list.append(x)
        
        for layer in self.down_conv_list:
            x = layer(x)
            feat_list.append(x)
        
        output = []
        feat_inv_list = feat_list[::-1]
        for idx, (feat_enc, feat_c) in enumerate(zip(feat_inv_list, coarse_feat_whole)):
            
            # in case for depth-anything
            _, _, h, w = feat_enc.shape
            if h != feat_c.shape[-2] or w != feat_c.shape[-1]:
                feat_enc = F.interpolate(feat_enc, size=feat_c.shape[-2:], mode='bilinear', align_corners=True)
                
            if idx == 0:
                pass
            else:
                feat_enc = self.up_conv_list[idx-1](torch.cat([temp_feat, guide_cat[idx-1]], dim=1), feat_enc)
            
            _, _, h, w = feat_c.shape
            feat_c = self.g2l_list[idx](feat_c, None)
            feat_c = torch_roi_align(feat_c, bbox, (h, w), h/self.patch_process_shape[0], aligned=True)
            x = self.convs[idx](torch.cat([feat_enc, feat_c], dim=1))
            temp_feat = x
            output.append(x)

        return output[::-1]
//  [python]
def upsample(self, source, guidance_unscaled, upsampler, scale):
        _, _, H, W = source.shape
        guidance = t.interpolate(guidance_unscaled, size=(H * scale, W * scale), mode='bilinear')
        return upsampler(source, guidance)
//  [python]
def upsample(self, source, guidance_unscaled, upsampler, scale):
        _, _, H, W = source.shape
        guidance = F.interpolate(guidance_unscaled, size=(H * scale, W * scale), mode='bilinear')
        return upsampler(source, guidance)
//  [python]
def forward(self, x):
        """Forward Function.

        Args:
            x (Tensor): Input tensor with shape of (n, c, h, w).

        Returns:
            Tensor: Output tensor with shape of (n, c, h', w').
        """
        for i in range(0, len(self.decoder)):
            if i <= 1:
                x = F.upsample(
                    x, scale_factor=2, mode='bilinear', align_corners=True)
            x = self.decoder[i](x)

        return self.output_act(x)
//  [python]
def forward(self, x):
        """Forward Function.

        Args:
            x (Tensor): Input tensor with shape of (n, c, h, w).

        Returns:
            Tensor: Output tensor with shape of (n, c, h', w').
        """
        for i in range(0, len(self.decoder)):
            if i <= 1:
                x = F.interpolate(
                    x, scale_factor=2, mode='bilinear', align_corners=True)
            x = self.decoder[i](x)

        return self.output_act(x)
//  [python]
def random_sized_crop_img(
    im,
    size,
    jitter_scale=(0.08, 1.0),
    jitter_aspect=(3.0 / 4.0, 4.0 / 3.0),
    max_iter=10,
):
    """
    Performs Inception-style cropping (used for training).
    """
    assert (
        len(im.shape) == 3
    ), "Currently only support image for random_sized_crop"
    h, w = im.shape[1:3]
    i, j, h, w = _get_param_spatial_crop(
        scale=jitter_scale,
        ratio=jitter_aspect,
        height=h,
        width=w,
        num_repeat=max_iter,
        log_scale=False,
        switch_hw=True,
    )
    cropped = im[:, i : i + h, j : j + w]
    return t.nn.functional.interpolate(
        cropped.unsqueeze(0),
        size=(size, size),
        mode="bilinear",
        align_corners=False,
    ).squeeze(0)
//  [python]
def random_sized_crop_img(
    im,
    size,
    jitter_scale=(0.08, 1.0),
    jitter_aspect=(3.0 / 4.0, 4.0 / 3.0),
    max_iter=10,
):
    """
    Performs Inception-style cropping (used for training).
    """
    assert (
        len(im.shape) == 3
    ), "Currently only support image for random_sized_crop"
    h, w = im.shape[1:3]
    i, j, h, w = _get_param_spatial_crop(
        scale=jitter_scale,
        ratio=jitter_aspect,
        height=h,
        width=w,
        num_repeat=max_iter,
        log_scale=False,
        switch_hw=True,
    )
    cropped = im[:, i : i + h, j : j + w]
    return torch.nn.functional.interpolate(
        cropped.unsqueeze(0),
        size=(size, size),
        mode="bilinear",
        align_corners=False,
    ).squeeze(0)
//  [python]
def random_resized_crop_with_shift(
    images,
    target_height,
    target_width,
    scale=(0.8, 1.0),
    ratio=(3.0 / 4.0, 4.0 / 3.0),
):
    """
    This is similar to random_resized_crop. However, it samples two different
    boxes (for cropping) for the first and last frame. It then linearly
    interpolates the two boxes for other frames.

    Args:
        images: Images to perform resizing and cropping.
        target_height: Desired height after cropping.
        target_width: Desired width after cropping.
        scale: Scale range of Inception-style area based random resizing.
        ratio: Aspect ratio range of Inception-style area based random resizing.
    """
    t = images.shape[1]
    height = images.shape[2]
    width = images.shape[3]

    i, j, h, w = _get_param_spatial_crop(scale, ratio, height, width)
    i_, j_, h_, w_ = _get_param_spatial_crop(scale, ratio, height, width)
    i_s = [int(i) for i in torch.linspace(i, i_, steps=t).tolist()]
    j_s = [int(i) for i in torch.linspace(j, j_, steps=t).tolist()]
    h_s = [int(i) for i in torch.linspace(h, h_, steps=t).tolist()]
    w_s = [int(i) for i in torch.linspace(w, w_, steps=t).tolist()]
    out = torch.zeros((3, t, target_height, target_width))
    for ind in range(t):
        out[:, ind : ind + 1, :, :] = t(
            images[
                :,
                ind : ind + 1,
                i_s[ind] : i_s[ind] + h_s[ind],
                j_s[ind] : j_s[ind] + w_s[ind],
            ],
            size=(target_height, target_width),
            mode="bilinear",
            align_corners=False,
        )
    return out
//  [python]
def random_resized_crop_with_shift(
    images,
    target_height,
    target_width,
    scale=(0.8, 1.0),
    ratio=(3.0 / 4.0, 4.0 / 3.0),
):
    """
    This is similar to random_resized_crop. However, it samples two different
    boxes (for cropping) for the first and last frame. It then linearly
    interpolates the two boxes for other frames.

    Args:
        images: Images to perform resizing and cropping.
        target_height: Desired height after cropping.
        target_width: Desired width after cropping.
        scale: Scale range of Inception-style area based random resizing.
        ratio: Aspect ratio range of Inception-style area based random resizing.
    """
    t = images.shape[1]
    height = images.shape[2]
    width = images.shape[3]

    i, j, h, w = _get_param_spatial_crop(scale, ratio, height, width)
    i_, j_, h_, w_ = _get_param_spatial_crop(scale, ratio, height, width)
    i_s = [int(i) for i in torch.linspace(i, i_, steps=t).tolist()]
    j_s = [int(i) for i in torch.linspace(j, j_, steps=t).tolist()]
    h_s = [int(i) for i in torch.linspace(h, h_, steps=t).tolist()]
    w_s = [int(i) for i in torch.linspace(w, w_, steps=t).tolist()]
    out = torch.zeros((3, t, target_height, target_width))
    for ind in range(t):
        out[:, ind : ind + 1, :, :] = torch.nn.functional.interpolate(
            images[
                :,
                ind : ind + 1,
                i_s[ind] : i_s[ind] + h_s[ind],
                j_s[ind] : j_s[ind] + w_s[ind],
            ],
            size=(target_height, target_width),
            mode="bilinear",
            align_corners=False,
        )
    return out
//  [python]
def uniform_crop(images, size, spatial_idx, boxes=None, scale_size=None):
    """
    Perform uniform spatial sampling on the images and corresponding boxes.
    Args:
        images (tensor): images to perform uniform crop. The dimension is
            `num frames` x `channel` x `height` x `width`.
        size (int): size of height and weight to crop the images.
        spatial_idx (int): 0, 1, or 2 for left, center, and right crop if width
            is larger than height. Or 0, 1, or 2 for top, center, and bottom
            crop if height is larger than width.
        boxes (ndarray or None): optional. Corresponding boxes to images.
            Dimension is `num boxes` x 4.
        scale_size (int): optinal. If not None, resize the images to scale_size before
            performing any crop.
    Returns:
        cropped (tensor): images with dimension of
            `num frames` x `channel` x `size` x `size`.
        cropped_boxes (ndarray or None): the cropped boxes with dimension of
            `num boxes` x 4.
    """
    assert spatial_idx in [0, 1, 2]
    ndim = len(images.shape)
    if ndim == 3:
        images = images.unsqueeze(0)
    height = images.shape[2]
    width = images.shape[3]

    if scale_size is not None:
        if width <= height:
            width, height = scale_size, int(height / width * scale_size)
        else:
            width, height = int(width / height * scale_size), scale_size
        images = t.nn.functional.interpolate(
            images,
            size=(height, width),
            mode="bilinear",
            align_corners=False,
        )

    y_offset = int(math.ceil((height - size) / 2))
    x_offset = int(math.ceil((width - size) / 2))

    if height > width:
        if spatial_idx == 0:
            y_offset = 0
        elif spatial_idx == 2:
            y_offset = height - size
    else:
        if spatial_idx == 0:
            x_offset = 0
        elif spatial_idx == 2:
            x_offset = width - size
    cropped = images[
        :, :, y_offset : y_offset + size, x_offset : x_offset + size
    ]
    cropped_boxes = (
        crop_boxes(boxes, x_offset, y_offset) if boxes is not None else None
    )
    if ndim == 3:
        cropped = cropped.squeeze(0)
    return cropped, cropped_boxes
//  [python]
def uniform_crop(images, size, spatial_idx, boxes=None, scale_size=None):
    """
    Perform uniform spatial sampling on the images and corresponding boxes.
    Args:
        images (tensor): images to perform uniform crop. The dimension is
            `num frames` x `channel` x `height` x `width`.
        size (int): size of height and weight to crop the images.
        spatial_idx (int): 0, 1, or 2 for left, center, and right crop if width
            is larger than height. Or 0, 1, or 2 for top, center, and bottom
            crop if height is larger than width.
        boxes (ndarray or None): optional. Corresponding boxes to images.
            Dimension is `num boxes` x 4.
        scale_size (int): optinal. If not None, resize the images to scale_size before
            performing any crop.
    Returns:
        cropped (tensor): images with dimension of
            `num frames` x `channel` x `size` x `size`.
        cropped_boxes (ndarray or None): the cropped boxes with dimension of
            `num boxes` x 4.
    """
    assert spatial_idx in [0, 1, 2]
    ndim = len(images.shape)
    if ndim == 3:
        images = images.unsqueeze(0)
    height = images.shape[2]
    width = images.shape[3]

    if scale_size is not None:
        if width <= height:
            width, height = scale_size, int(height / width * scale_size)
        else:
            width, height = int(width / height * scale_size), scale_size
        images = torch.nn.functional.interpolate(
            images,
            size=(height, width),
            mode="bilinear",
            align_corners=False,
        )

    y_offset = int(math.ceil((height - size) / 2))
    x_offset = int(math.ceil((width - size) / 2))

    if height > width:
        if spatial_idx == 0:
            y_offset = 0
        elif spatial_idx == 2:
            y_offset = height - size
    else:
        if spatial_idx == 0:
            x_offset = 0
        elif spatial_idx == 2:
            x_offset = width - size
    cropped = images[
        :, :, y_offset : y_offset + size, x_offset : x_offset + size
    ]
    cropped_boxes = (
        crop_boxes(boxes, x_offset, y_offset) if boxes is not None else None
    )
    if ndim == 3:
        cropped = cropped.squeeze(0)
    return cropped, cropped_boxes
//  [python]
def forward(self, x):

        # Encoder 128 -> 128 -> 256 -> 512  (50, 50, 25, 12)
        skips = []
        out_sizes = []
        for resnet_layer in self.resnet_layers:
            x = resnet_layer(x)
            skips.append(x)
            out_sizes.append(x[0, 0, :, :].shape)

        # Decoder
        out_sizes.pop()
        skips.pop()
        for up_layer in self.up_layers:
            x = t.interpolate(x, size=out_sizes.pop(), mode='bilinear', align_corners=True)
            x = torch.cat([x, skips.pop()], dim=1)
            x = up_layer(x)

        return x
//  [python]
def forward(self, x):

        # Encoder 128 -> 128 -> 256 -> 512  (50, 50, 25, 12)
        skips = []
        out_sizes = []
        for resnet_layer in self.resnet_layers:
            x = resnet_layer(x)
            skips.append(x)
            out_sizes.append(x[0, 0, :, :].shape)

        # Decoder
        out_sizes.pop()
        skips.pop()
        for up_layer in self.up_layers:
            x = nn.functional.interpolate(x, size=out_sizes.pop(), mode='bilinear', align_corners=True)
            x = torch.cat([x, skips.pop()], dim=1)
            x = up_layer(x)

        return x
//  [python]
def gray_resize_for_identity(self, out, size=128):
        out_gray = (0.2989 * out[:, 0, :, :] + 0.5870 * out[:, 1, :, :] + 0.1140 * out[:, 2, :, :])
        out_gray = out_gray.unsqueeze(1)
        out_gray = t.nn.functional.interpolate(out_gray, (size, size), mode='bilinear', align_corners=False)
        return out_gray
//  [python]
def gray_resize_for_identity(self, out, size=128):
        out_gray = (0.2989 * out[:, 0, :, :] + 0.5870 * out[:, 1, :, :] + 0.1140 * out[:, 2, :, :])
        out_gray = out_gray.unsqueeze(1)
        out_gray = F.interpolate(out_gray, (size, size), mode='bilinear', align_corners=False)
        return out_gray
//  [python]
def forward(self, img, lr8x, hr2x):
        lr4x = t.interpolate(lr8x, scale_factor=2, mode='bilinear', align_corners=False)
        lr4x = self.conv_lr4x(lr4x)
        lr2x = t.interpolate(lr4x, scale_factor=2, mode='bilinear', align_corners=False)

        f2x = self.conv_f2x(torch.cat((lr2x, hr2x), dim=1))
        f = t.interpolate(f2x, scale_factor=2, mode='bilinear', align_corners=False)
        f = self.conv_f(torch.cat((f, img), dim=1))
        pred_matte = torch.sigmoid(f)

        return pred_matte
//  [python]
def forward(self, img, lr8x, hr2x):
        lr4x = F.interpolate(lr8x, scale_factor=2, mode='bilinear', align_corners=False)
        lr4x = self.conv_lr4x(lr4x)
        lr2x = F.interpolate(lr4x, scale_factor=2, mode='bilinear', align_corners=False)

        f2x = self.conv_f2x(torch.cat((lr2x, hr2x), dim=1))
        f = F.interpolate(f2x, scale_factor=2, mode='bilinear', align_corners=False)
        f = self.conv_f(torch.cat((f, img), dim=1))
        pred_matte = torch.sigmoid(f)

        return pred_matte
//  [python]
def forward(self, img, inference):
        enc_features = self.backbone.forward(img)
        enc2x, enc4x, enc32x = enc_features[0], enc_features[1], enc_features[4]

        enc32x = self.se_block(enc32x)
        lr16x = t.interpolate(enc32x, scale_factor=2, mode='bilinear', align_corners=False)
        lr16x = self.conv_lr16x(lr16x)
        lr8x = t.interpolate(lr16x, scale_factor=2, mode='bilinear', align_corners=False)
        lr8x = self.conv_lr8x(lr8x)

        pred_semantic = None
        if not inference:
            lr = self.conv_lr(lr8x)
            pred_semantic = torch.sigmoid(lr)

        return pred_semantic, lr8x, [enc2x, enc4x]
//  [python]
def forward(self, img, inference):
        enc_features = self.backbone.forward(img)
        enc2x, enc4x, enc32x = enc_features[0], enc_features[1], enc_features[4]

        enc32x = self.se_block(enc32x)
        lr16x = F.interpolate(enc32x, scale_factor=2, mode='bilinear', align_corners=False)
        lr16x = self.conv_lr16x(lr16x)
        lr8x = F.interpolate(lr16x, scale_factor=2, mode='bilinear', align_corners=False)
        lr8x = self.conv_lr8x(lr8x)

        pred_semantic = None
        if not inference:
            lr = self.conv_lr(lr8x)
            pred_semantic = torch.sigmoid(lr)

        return pred_semantic, lr8x, [enc2x, enc4x]
//  [python]
def forward(self, x):
        assert x.shape[1] == self.channels
        if self.dims == 3:
            x = F.upsample(
                x, size=(x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
            )
        else:
            x = F.upsample(x, scale_factor=2, mode="nearest")
        if self.use_conv:
            x = self.conv(x)
        return x
//  [python]
def forward(self, x):
        assert x.shape[1] == self.channels
        if self.dims == 3:
            x = F.interpolate(
                x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
            )
        else:
            x = F.interpolate(x, scale_factor=2, mode="nearest")
        if self.use_conv:
            x = self.conv(x)
        return x
//  [python]
def forward(self, img, enc2x, enc4x, lr8x, inference):
        img2x = t.interpolate(img, scale_factor=1 / 2, mode='bilinear', align_corners=False)
        img4x = t.interpolate(img, scale_factor=1 / 4, mode='bilinear', align_corners=False)

        enc2x = self.tohr_enc2x(enc2x)
        hr4x = self.conv_enc2x(torch.cat((img2x, enc2x), dim=1))

        enc4x = self.tohr_enc4x(enc4x)
        hr4x = self.conv_enc4x(torch.cat((hr4x, enc4x), dim=1))

        lr4x = t.interpolate(lr8x, scale_factor=2, mode='bilinear', align_corners=False)
        hr4x = self.conv_hr4x(torch.cat((hr4x, lr4x, img4x), dim=1))

        hr2x = t.interpolate(hr4x, scale_factor=2, mode='bilinear', align_corners=False)
        hr2x = self.conv_hr2x(torch.cat((hr2x, enc2x), dim=1))

        pred_detail = None
        if not inference:
            hr = t.interpolate(hr2x, scale_factor=2, mode='bilinear', align_corners=False)
            hr = self.conv_hr(torch.cat((hr, img), dim=1))
            pred_detail = torch.sigmoid(hr)

        return pred_detail, hr2x
//  [python]
def forward(self, img, enc2x, enc4x, lr8x, inference):
        img2x = F.interpolate(img, scale_factor=1 / 2, mode='bilinear', align_corners=False)
        img4x = F.interpolate(img, scale_factor=1 / 4, mode='bilinear', align_corners=False)

        enc2x = self.tohr_enc2x(enc2x)
        hr4x = self.conv_enc2x(torch.cat((img2x, enc2x), dim=1))

        enc4x = self.tohr_enc4x(enc4x)
        hr4x = self.conv_enc4x(torch.cat((hr4x, enc4x), dim=1))

        lr4x = F.interpolate(lr8x, scale_factor=2, mode='bilinear', align_corners=False)
        hr4x = self.conv_hr4x(torch.cat((hr4x, lr4x, img4x), dim=1))

        hr2x = F.interpolate(hr4x, scale_factor=2, mode='bilinear', align_corners=False)
        hr2x = self.conv_hr2x(torch.cat((hr2x, enc2x), dim=1))

        pred_detail = None
        if not inference:
            hr = F.interpolate(hr2x, scale_factor=2, mode='bilinear', align_corners=False)
            hr = self.conv_hr(torch.cat((hr, img), dim=1))
            pred_detail = torch.sigmoid(hr)

        return pred_detail, hr2x
//  [python]
def upsample2d_flow_as(inputs, target_as, mode="bilinear", if_rate=False):

    _, _, h, w = target_as.size()
    if if_rate:
        _, _, h_, w_ = inputs.size()
        inputs[:, 0, :, :] *= (w / w_)
        inputs[:, 1, :, :] *= (h / h_)
    res = t.nn.functional.interpolate(inputs, [h, w], mode=mode, align_corners=True)
    return res
//  [python]
def upsample2d_flow_as(inputs, target_as, mode="bilinear", if_rate=False):

    _, _, h, w = target_as.size()
    if if_rate:
        _, _, h_, w_ = inputs.size()
        inputs[:, 0, :, :] *= (w / w_)
        inputs[:, 1, :, :] *= (h / h_)
    res = F.interpolate(inputs, [h, w], mode=mode, align_corners=True)
    return res
//  [python]
def forward(
        self,
        x,
        emb = None
    ):
        if self.upsample:
            h, w = x.shape[-2:]
            x = F.interpolate(x, (h * 2, w * 2), mode = 'bilinear')

        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')

        x = self.block2(x)

        x = self.res_mp_add(x, res)

        if exists(self.attn):
            x = self.attn(x)

        return x
//  [python]
def forward_bidirect_flow(self, masked_local_frames):
        b, l_t, c, h, w = masked_local_frames.size()

        # compute forward and backward flows of masked frames
        masked_local_frames = F.interpolate(masked_local_frames.view(
            -1, c, h, w),
                                            scale_factor=1 / 4,
                                            mode='bilinear',
                                            align_corners=True,
                                            recompute_scale_factor=True)
        masked_local_frames = masked_local_frames.view(b, l_t, c, h // 4,
                                                       w // 4)
        mlf_1 = masked_local_frames[:, :-1, :, :, :].reshape(
            -1, c, h // 4, w // 4)
        mlf_2 = masked_local_frames[:, 1:, :, :, :].reshape(
            -1, c, h // 4, w // 4)
        pred_flows_forward = self.update_spynet(mlf_1, mlf_2)
        pred_flows_backward = self.update_spynet(mlf_2, mlf_1)

        pred_flows_forward = pred_flows_forward.view(b, l_t - 1, 2, h // 4,
                                                     w // 4)
        pred_flows_backward = pred_flows_backward.view(b, l_t - 1, 2, h // 4,
                                                       w // 4)

        return pred_flows_forward, pred_flows_backward
//  [python]
def forward(
        self,
        x,
        emb = None
    ):
        if self.downsample:
            h, w = x.shape[-2:]
            x = t.interpolate(x, (h // 2, w // 2), mode = 'bilinear')
            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')

        x = self.block2(x)

        x = self.res_mp_add(x, res)

        if exists(self.attn):
            x = self.attn(x)

        return x
//  [python]
def forward(
        self,
        x,
        emb = None
    ):
        if self.downsample:
            h, w = x.shape[-2:]
            x = F.interpolate(x, (h // 2, w // 2), mode = 'bilinear')
            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')

        x = self.block2(x)

        x = self.res_mp_add(x, res)

        if exists(self.attn):
            x = self.attn(x)

        return x
//  [python]
def sag_masking(self, original_latents, attn_map, map_size, t, eps):
        # Same masking process as in SAG paper: https://arxiv.org/pdf/2210.00939.pdf
        bh, hw1, hw2 = attn_map.shape
        b, latent_channel, latent_h, latent_w = original_latents.shape
        h = self.unet.config.attention_head_dim
        if isinstance(h, list):
            h = h[-1]

        # Produce attention mask
        attn_map = attn_map.reshape(b, h, hw1, hw2)
        attn_mask = attn_map.mean(1, keepdim=False).sum(1, keepdim=False) > 1.0
        attn_mask = (
            attn_mask.reshape(b, map_size[0], map_size[1])
            .unsqueeze(1)
            .repeat(1, latent_channel, 1, 1)
            .type(attn_map.dtype)
        )
        attn_mask = F.interpolate(attn_mask, (latent_h, latent_w), mode='bilinear', align_corners=False)

        # Blur according to the self-attention mask
        degraded_latents = gaussian_blur_2d(original_latents, kernel_size=9, sigma=1.0)
        degraded_latents = degraded_latents * attn_mask + original_latents * (1 - attn_mask)

        # Noise it again to match the noise level
        degraded_latents = self.scheduler.add_noise(degraded_latents, noise=eps, timesteps=t[None])

        return degraded_latents
//  [python]
def sag_masking(self, original_latents, attn_map, map_size, t, eps):
        # Same masking process as in SAG paper: https://arxiv.org/pdf/2210.00939.pdf
        bh, hw1, hw2 = attn_map.shape
        b, latent_channel, latent_h, latent_w = original_latents.shape
        h = self.unet.config.attention_head_dim
        if isinstance(h, list):
            h = h[-1]

        # Produce attention mask
        attn_map = attn_map.reshape(b, h, hw1, hw2)
        attn_mask = attn_map.mean(1, keepdim=False).sum(1, keepdim=False) > 1.0
        attn_mask = (
            attn_mask.reshape(b, map_size[0], map_size[1])
            .unsqueeze(1)
            .repeat(1, latent_channel, 1, 1)
            .type(attn_map.dtype)
        )
        attn_mask = F.interpolate(attn_mask, (latent_h, latent_w))

        # Blur according to the self-attention mask
        degraded_latents = gaussian_blur_2d(original_latents, kernel_size=9, sigma=1.0)
        degraded_latents = degraded_latents * attn_mask + original_latents * (1 - attn_mask)

        # Noise it again to match the noise level
        degraded_latents = self.scheduler.add_noise(degraded_latents, noise=eps, timesteps=t[None])

        return degraded_latents
//  [python]
def interpolate_pos_embed(model, checkpoint_model):
    if 'pos_embed' in checkpoint_model:
        pos_embed_checkpoint = checkpoint_model['pos_embed']
        embedding_size = pos_embed_checkpoint.shape[-1]
        num_patches = model.patch_embed.num_patches
        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
        # height (== width) for the checkpoint position embedding
        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
        # height (== width) for the new position embedding
        new_size = int(num_patches ** 0.5)
        # class_token and dist_token are kept unchanged
        if orig_size != new_size:
            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
            # only the position tokens are interpolated
            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
            pos_tokens = t.nn.functional.interpolate(
                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
            checkpoint_model['pos_embed'] = new_pos_embed
//  [python]
def interpolate_pos_embed(model, checkpoint_model):
    if 'pos_embed' in checkpoint_model:
        pos_embed_checkpoint = checkpoint_model['pos_embed']
        embedding_size = pos_embed_checkpoint.shape[-1]
        num_patches = model.patch_embed.num_patches
        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
        # height (== width) for the checkpoint position embedding
        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
        # height (== width) for the new position embedding
        new_size = int(num_patches ** 0.5)
        # class_token and dist_token are kept unchanged
        if orig_size != new_size:
            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
            # only the position tokens are interpolated
            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
            pos_tokens = torch.nn.functional.interpolate(
                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
            checkpoint_model['pos_embed'] = new_pos_embed
//  [python]
def print_gradimg(gradimg,it,shape=None):
    out=torch.zeros((gradimg.shape[0],gradimg.shape[1],3),requires_grad=False,dtype=torch.float32)
    for x in range(gradimg.shape[0]):
        for y in range(gradimg.shape[1]):
            h=math.atan2(gradimg[x,y,1],gradimg[x,y,0])
            s=math.tanh(np.linalg.norm(gradimg[x,y,:]))
            v=1.
            vec=(gradimg[x,y,:].clip(min=-1,max=1)/2)+.5
            #out[x,y,:]=torch.tensor(colorsys.hsv_to_rgb(h,s,v),dtype=torch.float32)
            out[x,y,:]=torch.tensor([vec[0],vec[1],0])

    if shape is not None:
        out=t.nn.functional.interpolate(out.permute(2,0,1).unsqueeze(0),shape,mode='bilinear').squeeze().permute(1,2,0)
    pydiffvg.imwrite(out.cpu(), 'results/simple_transform_svg/grad_{:04}.png'.format(it))
//  [python]
def print_gradimg(gradimg,it,shape=None):
    out=torch.zeros((gradimg.shape[0],gradimg.shape[1],3),requires_grad=False,dtype=torch.float32)
    for x in range(gradimg.shape[0]):
        for y in range(gradimg.shape[1]):
            h=math.atan2(gradimg[x,y,1],gradimg[x,y,0])
            s=math.tanh(np.linalg.norm(gradimg[x,y,:]))
            v=1.
            vec=(gradimg[x,y,:].clip(min=-1,max=1)/2)+.5
            #out[x,y,:]=torch.tensor(colorsys.hsv_to_rgb(h,s,v),dtype=torch.float32)
            out[x,y,:]=torch.tensor([vec[0],vec[1],0])

    if shape is not None:
        out=torch.nn.functional.interpolate(out.permute(2,0,1).unsqueeze(0),shape,mode='bilinear').squeeze().permute(1,2,0)
    pydiffvg.imwrite(out.cpu(), 'results/simple_transform_svg/grad_{:04}.png'.format(it))
//  [python]
def comp_loss_and_grad(img, tgt, it, sz):
    dif=img-tgt

    loss=dif.pow(2).mean()

    dif=dif.detach()

    cdif=dif.clone().abs()
    cdif[:,:,3]=1.

    resdif=t.nn.functional.interpolate(cdif.permute(2,0,1).unsqueeze(0),sz,mode='bilinear').squeeze().permute(1,2,0).abs()
    pydiffvg.imwrite(resdif[:,:,0:4], 'results/simple_transform_svg/dif_{:04}.png'.format(it))

    dif=dif.numpy()
    padded=np.pad(dif,[(1,1),(1,1),(0,0)],mode='edge')
    #print(padded[:-2,:,:].shape)
    grad_x=(padded[:-2,:,:]-padded[2:,:,:])[:,1:-1,:]
    grad_y=(padded[:,:-2,:]-padded[:,2:,:])[1:-1,:,:]

    resshape=dif.shape
    resshape=(resshape[0],resshape[1],2)
    res=np.zeros(resshape)

    for x in range(resshape[0]):
        for y in range(resshape[1]):
            A=np.concatenate((grad_x[x,y,:][:,np.newaxis],grad_y[x,y,:][:,np.newaxis]),axis=1)
            b=-dif[x,y,:]
            v=np.linalg.lstsq(np.dot(A.T,A),np.dot(A.T,b))
            res[x,y,:]=v[0]

    return loss, res
//  [python]
def comp_loss_and_grad(img, tgt, it, sz):
    dif=img-tgt

    loss=dif.pow(2).mean()

    dif=dif.detach()

    cdif=dif.clone().abs()
    cdif[:,:,3]=1.

    resdif=torch.nn.functional.interpolate(cdif.permute(2,0,1).unsqueeze(0),sz,mode='bilinear').squeeze().permute(1,2,0).abs()
    pydiffvg.imwrite(resdif[:,:,0:4], 'results/simple_transform_svg/dif_{:04}.png'.format(it))

    dif=dif.numpy()
    padded=np.pad(dif,[(1,1),(1,1),(0,0)],mode='edge')
    #print(padded[:-2,:,:].shape)
    grad_x=(padded[:-2,:,:]-padded[2:,:,:])[:,1:-1,:]
    grad_y=(padded[:,:-2,:]-padded[:,2:,:])[1:-1,:,:]

    resshape=dif.shape
    resshape=(resshape[0],resshape[1],2)
    res=np.zeros(resshape)

    for x in range(resshape[0]):
        for y in range(resshape[1]):
            A=np.concatenate((grad_x[x,y,:][:,np.newaxis],grad_y[x,y,:][:,np.newaxis]),axis=1)
            b=-dif[x,y,:]
            v=np.linalg.lstsq(np.dot(A.T,A),np.dot(A.T,b))
            res[x,y,:]=v[0]

    return loss, res
//  [python]
def create_mask(self, row=None, col=None, mask=None):
        """
        Creating coded aperture mask.
        """

        if mask is not None:
            self.mask = mask
            assert row is None and col is None, "Row and col should not be specified"

        elif row is not None:
            assert col is not None, "Both row and col should be specified"
            self.row = row
            self.col = col

        # output product if necessary
        if self.row is not None:
            if self.is_torch:
                self.mask = torch.outer(self.row, self.col)
            else:
                self.mask = np.outer(self.row, self.col)
        assert self.mask is not None, "Mask should be specified"

        # resize to sensor shape
        if np.any(self.resolution != self.mask.shape):

            if self.is_torch:
                self.mask = self.mask.unsqueeze(0).unsqueeze(0)
                self.mask = t.nn.functional.interpolate(
                    self.mask, size=tuple(self.resolution), mode="nearest"
                ).squeeze()
            else:
                # self.mask = resize(self.mask[:, :, np.newaxis], shape=tuple(self.resolution) + (1,))
                self.mask = resize(
                    self.mask[:, :, np.newaxis],
                    shape=tuple(self.resolution) + (1,),
                    interpolation=cv.INTER_NEAREST,
                ).squeeze()
//  [python]
def create_mask(self, row=None, col=None, mask=None):
        """
        Creating coded aperture mask.
        """

        if mask is not None:
            self.mask = mask
            assert row is None and col is None, "Row and col should not be specified"

        elif row is not None:
            assert col is not None, "Both row and col should be specified"
            self.row = row
            self.col = col

        # output product if necessary
        if self.row is not None:
            if self.is_torch:
                self.mask = torch.outer(self.row, self.col)
            else:
                self.mask = np.outer(self.row, self.col)
        assert self.mask is not None, "Mask should be specified"

        # resize to sensor shape
        if np.any(self.resolution != self.mask.shape):

            if self.is_torch:
                self.mask = self.mask.unsqueeze(0).unsqueeze(0)
                self.mask = torch.nn.functional.interpolate(
                    self.mask, size=tuple(self.resolution), mode="nearest"
                ).squeeze()
            else:
                # self.mask = resize(self.mask[:, :, np.newaxis], shape=tuple(self.resolution) + (1,))
                self.mask = resize(
                    self.mask[:, :, np.newaxis],
                    shape=tuple(self.resolution) + (1,),
                    interpolation=cv.INTER_NEAREST,
                ).squeeze()
//  [python]
def forward(self, inputs: Tuple[torch.Tensor, ...]) -> torch.Tensor:
        """

        :param inputs: Tuple with single Tensor of shape (B, T, D)
        :return:
        """
        # Transpose the input to the required format (B, D, T)
        x, = inputs
        x = x.transpose(1, 2)

        # Put signal through TCN. Output-shape: (B, nb_filters, T)
        x = self.tcn_enc(x)
        # Now, adjust the number of channels...
        x = self.conv1d(x)
        x = self.activation(x)

        # Do some average (max) pooling to get a compressed representation of the time series
        # (e.g. a sequence of length 8)
        seq_len = x.shape[-1]
        x = self.pooler(x)
        # x = self.activation(x)

        # Now we should have a short sequence, which we will upsample again and
        # then try to reconstruct the original series
        x = F.upsample(x, scale_factor=seq_len, mode='nearest')
        x = self.tcn_dec(x)
        # Put the filter-outputs through a dense layer finally, to get the reconstructed signal
        x = self.linear(x)

        # Put output dimensions in the correct order again
        x = x.transpose(1, 2)

        return x
//  [python]
def forward(self, inputs: Tuple[torch.Tensor, ...]) -> torch.Tensor:
        """

        :param inputs: Tuple with single Tensor of shape (B, T, D)
        :return:
        """
        # Transpose the input to the required format (B, D, T)
        x, = inputs
        x = x.transpose(1, 2)

        # Put signal through TCN. Output-shape: (B, nb_filters, T)
        x = self.tcn_enc(x)
        # Now, adjust the number of channels...
        x = self.conv1d(x)
        x = self.activation(x)

        # Do some average (max) pooling to get a compressed representation of the time series
        # (e.g. a sequence of length 8)
        seq_len = x.shape[-1]
        x = self.pooler(x)
        # x = self.activation(x)

        # Now we should have a short sequence, which we will upsample again and
        # then try to reconstruct the original series
        x = F.interpolate(x, seq_len, mode='nearest')
        x = self.tcn_dec(x)
        # Put the filter-outputs through a dense layer finally, to get the reconstructed signal
        x = self.linear(x)

        # Put output dimensions in the correct order again
        x = x.transpose(1, 2)

        return x
//  [python]
def forward(self,
                img: torch.Tensor,
                feature: torch.Tensor,
                styles: Union[torch.Tensor, List[torch.Tensor]],
                add_noise: bool = False) -> torch.Tensor:
        """Forward function.

        Args:
            img (torch.Tensor): Image to super resolution.
            x (torch.Tensor): Feature map of the input image.
            styles (torch.Tensor): Style codes in w space.
            add_noise (bool, optional): Whether add noise to image.
                Defaults to False.

        Returns:
            torch.Tensor: Image after super resolution.
        """
        if isinstance(styles, list):
            styles = styles[-1]
        if styles.ndim == 3:
            styles = styles[-1]
        assert styles.ndim == 2 and styles.shape[-1] == self.style_channels
        styles = styles[:, None, :].repeat(1, 3, 1)
        if feature.shape[-1] != self.in_size:
            interpolation_kwargs = dict(
                size=(self.in_size, self.in_size),
                mode='bilinear',
                align_corners=False)
            if digit_version(TORCH_VERSION) >= digit_version('1.11.0'):
                interpolation_kwargs['antialias'] = self.sr_antialias
            feature = F.interpolate(feature, **interpolation_kwargs)
            img = F.interpolate(img, **interpolation_kwargs)
        feature, img = self.block0(feature, img, styles, add_noise)
        feature, img = self.block1(feature, img, styles, add_noise)
        return img
//  [python]
def forward(self, input, z=None):
        seg = input

        if self.opt.use_vae:
            # we sample z from unit normal and reshape the tensor
            if z is None:
                z = torch.randn(input.size(0), self.opt.z_dim,
                                dtype=torch.float32, device=input.get_device())
            x = self.fc(z)
            x = x.view(-1, 16 * self.opt.ngf, self.sh, self.sw)
        else:
            # we downsample segmap and run convolution
            x = t.interpolate(seg, size=(self.sh, self.sw))
            x = self.fc(x)

        x = self.head_0(x, seg)

        x = self.up(x)
        x = self.G_middle_0(x, seg)

        if self.opt.num_upsampling_layers == 'more' or \
           self.opt.num_upsampling_layers == 'most':
            x = self.up(x)

        x = self.G_middle_1(x, seg)

        x = self.up(x)
        x = self.up_0(x, seg)
        x = self.up(x)
        x = self.up_1(x, seg)
        x = self.up(x)
        x = self.up_2(x, seg)
        x = self.up(x)
        x = self.up_3(x, seg)

        if self.opt.num_upsampling_layers == 'most':
            x = self.up(x)
            x = self.up_4(x, seg)

        x = self.conv_img(F.leaky_relu(x, 2e-1))
        x = F.tanh(x)

        return x
//  [python]
def forward(self, input, z=None):
        seg = input

        if self.opt.use_vae:
            # we sample z from unit normal and reshape the tensor
            if z is None:
                z = torch.randn(input.size(0), self.opt.z_dim,
                                dtype=torch.float32, device=input.get_device())
            x = self.fc(z)
            x = x.view(-1, 16 * self.opt.ngf, self.sh, self.sw)
        else:
            # we downsample segmap and run convolution
            x = F.interpolate(seg, size=(self.sh, self.sw))
            x = self.fc(x)

        x = self.head_0(x, seg)

        x = self.up(x)
        x = self.G_middle_0(x, seg)

        if self.opt.num_upsampling_layers == 'more' or \
           self.opt.num_upsampling_layers == 'most':
            x = self.up(x)

        x = self.G_middle_1(x, seg)

        x = self.up(x)
        x = self.up_0(x, seg)
        x = self.up(x)
        x = self.up_1(x, seg)
        x = self.up(x)
        x = self.up_2(x, seg)
        x = self.up(x)
        x = self.up_3(x, seg)

        if self.opt.num_upsampling_layers == 'most':
            x = self.up(x)
            x = self.up_4(x, seg)

        x = self.conv_img(F.leaky_relu(x, 2e-1))
        x = F.tanh(x)

        return x
//  [python]
def _upsample_add(self, x, y):
        _,_,H,W = y.size()
        return F.upsample(x, scale_factor=float(H)/x.size(2), mode='bicubic') + y
//  [python]
def _upsample_add(self, x, y):
        _,_,H,W = y.size()
        return F.interpolate(x, size=(H,W), mode='bicubic') + y
//  [python]
def forward(self, x, noise, mask):
        W,_ = torch.split(self.weight.view(1, -1, 1, 1), self.channels // 2, dim=1)
        B,_ = torch.split(self.bias.view(1, -1, 1, 1), self.channels // 2, dim=1)
        Z = torch.zeros_like(W)
        w = torch.cat([W,Z], dim=1).to(x.device)
        b = torch.cat([B,Z], dim=1).to(x.device)
        # adds = w * noise.to(x.device) + b
        # return x + adds.type_as(x)*(1-mask)
        adds = w * torch.randn_like(x) + b
        mask = F.upsample(mask, size=(x.shape[2],x.shape[3]), mode='nearest')
        return x + adds.type_as(x)*(1-mask)
//  [python]
def forward(self, x, noise, mask):
        W,_ = torch.split(self.weight.view(1, -1, 1, 1), self.channels // 2, dim=1)
        B,_ = torch.split(self.bias.view(1, -1, 1, 1), self.channels // 2, dim=1)
        Z = torch.zeros_like(W)
        w = torch.cat([W,Z], dim=1).to(x.device)
        b = torch.cat([B,Z], dim=1).to(x.device)
        # adds = w * noise.to(x.device) + b
        # return x + adds.type_as(x)*(1-mask)
        adds = w * torch.randn_like(x) + b
        mask = F.interpolate(mask, (x.shape[2],x.shape[3]))
        return x + adds.type_as(x)*(1-mask)
//  [python]
def forward(self, x, gts=None):

        x2, x3, x4, x5 = self.backbone(x)

        outs = {}

        metric = self.mlayer(x5)

        x = self.up_4(x5, x4)
        
        d = self.vlayer(x)

        if self.training:
            var_loss = self.var_loss(x, d, gts)


        x, d  = self.ref_4(x, d)

        d_u4 = t.interpolate(d, scale_factor=16, mode='bilinear', align_corners=True)

        x = self.up_3(x, x3)

        x, d = self.ref_3(x, t.interpolate(d, scale_factor=2, mode='bilinear', align_corners=True))

        d_u3 = t.interpolate(d, scale_factor=8, mode='bilinear', align_corners=True)

        x = self.up_2(x, x2)

        x, d = self.ref_2(x, t.interpolate(d, scale_factor=2, mode='bilinear', align_corners=True))

        d_u2 = t.interpolate(d, scale_factor=4, mode='bilinear', align_corners=True)

        d = d_u2 + d_u3 + d_u4

        d = torch.sigmoid(metric[:, 0:1]) * (self.outc(d) + torch.exp(metric[:, 1:2]))

        outs['scale_1'] = d

        if self.training:
            si_loss = self.si_loss(outs, gts)
            return outs['scale_1'], var_loss + si_loss
        else:
            return outs['scale_1']