jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def neural_warpfield(self, view, seq_length):
        warpfield = view
        for layer in self.layers:
            warpfield = F.pad(warpfield, pad=[1, 0])
            warpfield = F.relu(layer(warpfield))
        warpfield = self.linear(warpfield)
        warpfield = F.interpolate(warpfield, size=seq_length)
        return warpfield
//  [python]
def forward(
        self,
        cameras: CamerasBase,
        evaluation_mode: EvaluationMode,
        mask: Optional[torch.Tensor] = None,
    ) -> ImplicitronRayBundle:
        """

        Args:
            cameras: A batch of `batch_size` cameras from which the rays are emitted.
            evaluation_mode: one of `EvaluationMode.TRAINING` or
                `EvaluationMode.EVALUATION` which determines the sampling mode
                that is used.
            mask: Active for the `RenderSamplingMode.MASK_SAMPLE` sampling mode.
                Defines a non-negative mask of shape
                `(batch_size, image_height, image_width)` where each per-pixel
                value is proportional to the probability of sampling the
                corresponding pixel's ray.

        Returns:
            ray_bundle: A `ImplicitronRayBundle` object containing the parametrizations of the
                sampled rendering rays.
        """
        sample_mask = None
        if (
            self._sampling_mode[evaluation_mode] == RenderSamplingMode.MASK_SAMPLE
            and mask is not None
        ):
            sample_mask = t.nn.functional.interpolate(
                mask,
                size=[self.image_height, self.image_width],
                mode="nearest",
            )[:, 0]

        min_depth, max_depth = self._get_min_max_depth_bounds(cameras)

        raysampler = {
            EvaluationMode.TRAINING: self._training_raysampler,
            EvaluationMode.EVALUATION: self._evaluation_raysampler,
        }[evaluation_mode]

        ray_bundle = raysampler(
            cameras=cameras,
            mask=sample_mask,
            min_depth=min_depth,
            max_depth=max_depth,
        )
        if self.cast_ray_bundle_as_cone and isinstance(
            ray_bundle, HeterogeneousRayBundle
        ):
            # If this error rises it means that raysampler has among
            # its arguments `n_ray_totals`. If it is the case
            # then you should update the radii computation and lengths
            # computation to handle padding and unpadding.
            raise TypeError(
                "Heterogeneous ray bundle is not supported for conical frustum computation yet"
            )
        elif self.cast_ray_bundle_as_cone:
            pixel_hw: Tuple[float, float] = (self.pixel_height, self.pixel_width)
            pixel_radii_2d = compute_radii(cameras, ray_bundle.xys[..., :2], pixel_hw)
            return ImplicitronRayBundle(
                directions=ray_bundle.directions,
                origins=ray_bundle.origins,
                lengths=None,
                xys=ray_bundle.xys,
                bins=ray_bundle.lengths,
                pixel_radii_2d=pixel_radii_2d,
            )

        return ImplicitronRayBundle(
            directions=ray_bundle.directions,
            origins=ray_bundle.origins,
            lengths=ray_bundle.lengths,
            xys=ray_bundle.xys,
            camera_counts=getattr(ray_bundle, "camera_counts", None),
            camera_ids=getattr(ray_bundle, "camera_ids", None),
        )
//  [python]
def forward(
        self,
        cameras: CamerasBase,
        evaluation_mode: EvaluationMode,
        mask: Optional[torch.Tensor] = None,
    ) -> ImplicitronRayBundle:
        """

        Args:
            cameras: A batch of `batch_size` cameras from which the rays are emitted.
            evaluation_mode: one of `EvaluationMode.TRAINING` or
                `EvaluationMode.EVALUATION` which determines the sampling mode
                that is used.
            mask: Active for the `RenderSamplingMode.MASK_SAMPLE` sampling mode.
                Defines a non-negative mask of shape
                `(batch_size, image_height, image_width)` where each per-pixel
                value is proportional to the probability of sampling the
                corresponding pixel's ray.

        Returns:
            ray_bundle: A `ImplicitronRayBundle` object containing the parametrizations of the
                sampled rendering rays.
        """
        sample_mask = None
        if (
            self._sampling_mode[evaluation_mode] == RenderSamplingMode.MASK_SAMPLE
            and mask is not None
        ):
            sample_mask = torch.nn.functional.interpolate(
                mask,
                size=[self.image_height, self.image_width],
                mode="nearest",
            )[:, 0]

        min_depth, max_depth = self._get_min_max_depth_bounds(cameras)

        raysampler = {
            EvaluationMode.TRAINING: self._training_raysampler,
            EvaluationMode.EVALUATION: self._evaluation_raysampler,
        }[evaluation_mode]

        ray_bundle = raysampler(
            cameras=cameras,
            mask=sample_mask,
            min_depth=min_depth,
            max_depth=max_depth,
        )
        if self.cast_ray_bundle_as_cone and isinstance(
            ray_bundle, HeterogeneousRayBundle
        ):
            # If this error rises it means that raysampler has among
            # its arguments `n_ray_totals`. If it is the case
            # then you should update the radii computation and lengths
            # computation to handle padding and unpadding.
            raise TypeError(
                "Heterogeneous ray bundle is not supported for conical frustum computation yet"
            )
        elif self.cast_ray_bundle_as_cone:
            pixel_hw: Tuple[float, float] = (self.pixel_height, self.pixel_width)
            pixel_radii_2d = compute_radii(cameras, ray_bundle.xys[..., :2], pixel_hw)
            return ImplicitronRayBundle(
                directions=ray_bundle.directions,
                origins=ray_bundle.origins,
                lengths=None,
                xys=ray_bundle.xys,
                bins=ray_bundle.lengths,
                pixel_radii_2d=pixel_radii_2d,
            )

        return ImplicitronRayBundle(
            directions=ray_bundle.directions,
            origins=ray_bundle.origins,
            lengths=ray_bundle.lengths,
            xys=ray_bundle.xys,
            camera_counts=getattr(ray_bundle, "camera_counts", None),
            camera_ids=getattr(ray_bundle, "camera_ids", None),
        )
//  [python]
def test_unspec_float_precision(self):
        def fn(image, scale_factor):
            image = t.interpolate(
                image[None],
                size=None,
                scale_factor=scale_factor,
                mode="bilinear",
                recompute_scale_factor=True,
                align_corners=False,
            )[0]

            return image.shape

        x = torch.rand([3, 427, 640])
        scale_factor = 1.873536229133606
        ref = fn(x, scale_factor)
        cnts = torch._dynamo.testing.CompileCounter()
        opt_fn = torch._dynamo.optimize(cnts)(fn)
        res = opt_fn(x, scale_factor)
        self.assertTrue(same(ref, res))
//  [python]
def test_unspec_float_precision(self):
        def fn(image, scale_factor):
            image = torch.nn.functional.interpolate(
                image[None],
                size=None,
                scale_factor=scale_factor,
                mode="bilinear",
                recompute_scale_factor=True,
                align_corners=False,
            )[0]

            return image.shape

        x = torch.rand([3, 427, 640])
        scale_factor = 1.873536229133606
        ref = fn(x, scale_factor)
        cnts = torch._dynamo.testing.CompileCounter()
        opt_fn = torch._dynamo.optimize(cnts)(fn)
        res = opt_fn(x, scale_factor)
        self.assertTrue(same(ref, res))
//  [python]
def _hour_glass_forward(self, n: int, x: Tensor) -> Tensor:
        up1 = self.hg[n - 1][0](x)
        low1 = F.max_pool2d(x, 2, stride=2)
        low1 = self.hg[n - 1][1](low1)

        if n > 1:
            low2 = self._hour_glass_forward(n - 1, low1)
        else:
            low2 = self.hg[n - 1][3](low1)
        low3 = self.hg[n - 1][2](low2)
        up2 = F.upsample(low3, size=up1.shape[2:])
        out = up1 + up2
        return out
//  [python]
def _hour_glass_forward(self, n: int, x: Tensor) -> Tensor:
        up1 = self.hg[n - 1][0](x)
        low1 = F.max_pool2d(x, 2, stride=2)
        low1 = self.hg[n - 1][1](low1)

        if n > 1:
            low2 = self._hour_glass_forward(n - 1, low1)
        else:
            low2 = self.hg[n - 1][3](low1)
        low3 = self.hg[n - 1][2](low2)
        up2 = F.interpolate(low3, size=up1.shape[2:])
        out = up1 + up2
        return out
//  [python]
def __call__(self, imgs, num_workers=0):
        # preprocess
        if isinstance(imgs, Image.Image):
            imgs = [imgs]
        assert isinstance(imgs,
                          (tuple, list)) and isinstance(imgs[0], Image.Image)
        sizes = [u.size for u in imgs]
        imgs = torch.stack(parallel(self.transforms, imgs, num_workers), dim=0)

        # forward
        masks = []
        for batch in imgs.split(self.batch_size, dim=0):
            batch = batch.to(self.device, non_blocking=True)
            masks.append(self.softmax(self.net(batch), dim=1))

        # postprocess
        masks = torch.cat(masks, dim=0).unsqueeze(1)
        masks = [
            t.interpolate(u, v, mode='bilinear', align_corners=False)
            for u, v in zip(masks, sizes)
        ]
        return masks
//  [python]
def __call__(self, imgs, num_workers=0):
        # preprocess
        if isinstance(imgs, Image.Image):
            imgs = [imgs]
        assert isinstance(imgs,
                          (tuple, list)) and isinstance(imgs[0], Image.Image)
        sizes = [u.size for u in imgs]
        imgs = torch.stack(parallel(self.transforms, imgs, num_workers), dim=0)

        # forward
        masks = []
        for batch in imgs.split(self.batch_size, dim=0):
            batch = batch.to(self.device, non_blocking=True)
            masks.append(self.softmax(self.net(batch), dim=1))

        # postprocess
        masks = torch.cat(masks, dim=0).unsqueeze(1)
        masks = [
            F.interpolate(u, v, mode='bilinear', align_corners=False)
            for u, v in zip(masks, sizes)
        ]
        return masks
//  [python]
def forward(self, x: torch.Tensor) ->List[torch.Tensor]:
        """Get Inception feature maps

        Args:
            x: Batch of images with shape (N, 3, H, W). RGB colour order.
                Values are expected to be in range (0, 1) if `normalize_input` is True,
                and in range (-1, 1) otherwise.

        Returns:
            List of torch.autograd.Variable, corresponding to the selected output block, sorted ascending by index.
        """
        outp = []
        input_range = (0, 1) if self.normalize_input else (-1, 1)
        assert x.min() >= input_range[0] and x.max() <= input_range[1], f'Input tensor should be normalized in ({input_range[0]}, {input_range[0]}) range.'
        if self.resize_input:
            x = t.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False)
        if self.normalize_input:
            x = 2 * x - 1
        for idx, block in enumerate(self.blocks):
            x = block(x)
            if idx in self.output_blocks:
                outp.append(x)
            if idx == self.last_needed_block:
                break
        return outp
//  [python]
def forward(self, x: torch.Tensor) ->List[torch.Tensor]:
        """Get Inception feature maps

        Args:
            x: Batch of images with shape (N, 3, H, W). RGB colour order.
                Values are expected to be in range (0, 1) if `normalize_input` is True,
                and in range (-1, 1) otherwise.

        Returns:
            List of torch.autograd.Variable, corresponding to the selected output block, sorted ascending by index.
        """
        outp = []
        input_range = (0, 1) if self.normalize_input else (-1, 1)
        assert x.min() >= input_range[0] and x.max() <= input_range[1], f'Input tensor should be normalized in ({input_range[0]}, {input_range[0]}) range.'
        if self.resize_input:
            x = F.interpolate(x, size=(299, 299), mode='bilinear', align_corners=False)
        if self.normalize_input:
            x = 2 * x - 1
        for idx, block in enumerate(self.blocks):
            x = block(x)
            if idx in self.output_blocks:
                outp.append(x)
            if idx == self.last_needed_block:
                break
        return outp
//  [python]
def _image_enlarge(x: torch.Tensor) ->torch.Tensor:
    """Custom bilinear upscaling of an image.
    The function upscales an input image with upscaling factor 4x-3, adds padding on boundaries as difference
    and downscaled by the factor of 2.

    Args:
        x: An input tensor. Shape :math:`(N, C, H, W)`.

    Returns:
        Upscaled tensor.
    """
    t1 = t.imresize(x, size=(int(4 * x.size(-2) - 3), int(4 * x.size(-1) - 3)), mode='bilinear')
    t2 = torch.zeros([x.size(0), 1, 4 * x.size(-2) - 1, 4 * x.size(-1) - 1])
    t2[:, :, 1:-1, 1:-1] = t1
    t2[:, :, 0, :] = 2 * t2[:, :, 1, :] - t2[:, :, 2, :]
    t2[:, :, -1, :] = 2 * t2[:, :, -2, :] - t2[:, :, -3, :]
    t2[:, :, :, 0] = 2 * t2[:, :, :, 1] - t2[:, :, :, 2]
    t2[:, :, :, -1] = 2 * t2[:, :, :, -2] - t2[:, :, :, -3]
    out = t2[:, :, ::2, ::2]
    return out
//  [python]
def _image_enlarge(x: torch.Tensor) ->torch.Tensor:
    """Custom bilinear upscaling of an image.
    The function upscales an input image with upscaling factor 4x-3, adds padding on boundaries as difference
    and downscaled by the factor of 2.

    Args:
        x: An input tensor. Shape :math:`(N, C, H, W)`.

    Returns:
        Upscaled tensor.
    """
    t1 = F.interpolate(x, size=(int(4 * x.size(-2) - 3), int(4 * x.size(-1) - 3)), mode='bilinear', align_corners=False)
    t2 = torch.zeros([x.size(0), 1, 4 * x.size(-2) - 1, 4 * x.size(-1) - 1])
    t2[:, :, 1:-1, 1:-1] = t1
    t2[:, :, 0, :] = 2 * t2[:, :, 1, :] - t2[:, :, 2, :]
    t2[:, :, -1, :] = 2 * t2[:, :, -2, :] - t2[:, :, -3, :]
    t2[:, :, :, 0] = 2 * t2[:, :, :, 1] - t2[:, :, :, 2]
    t2[:, :, :, -1] = 2 * t2[:, :, :, -2] - t2[:, :, :, -3]
    out = t2[:, :, ::2, ::2]
    return out
//  [python]
def forward(self, x: torch.Tensor, y: torch.Tensor) ->torch.Tensor:
        """

        Args:
            x: An input tensor. Shape :math:`(N, C, H, W)`.
            y: A target tensor. Shape :math:`(N, C, H, W)`.

        Returns:
            Deep Image Structure and Texture Similarity loss, i.e. ``1-DISTS`` in range [0, 1].
        """
        _, _, H, W = x.shape
        if min(H, W) > 256:
            x = t.interpolate(x, scale_factor=256 / min(H, W), recompute_scale_factor=False, mode='bilinear')
            y = t.interpolate(y, scale_factor=256 / min(H, W), recompute_scale_factor=False, mode='bilinear')
        loss = super().forward(x, y)
        return 1 - loss
//  [python]
def forward(self, x: torch.Tensor, y: torch.Tensor) ->torch.Tensor:
        """

        Args:
            x: An input tensor. Shape :math:`(N, C, H, W)`.
            y: A target tensor. Shape :math:`(N, C, H, W)`.

        Returns:
            Deep Image Structure and Texture Similarity loss, i.e. ``1-DISTS`` in range [0, 1].
        """
        _, _, H, W = x.shape
        if min(H, W) > 256:
            x = torch.nn.functional.interpolate(x, scale_factor=256 / min(H, W), recompute_scale_factor=False, mode='bilinear')
            y = torch.nn.functional.interpolate(y, scale_factor=256 / min(H, W), recompute_scale_factor=False, mode='bilinear')
        loss = super().forward(x, y)
        return 1 - loss
//  [python]
def forward(self, hidden_state, residual=None):
        if residual is not None:
            if hidden_state.shape != residual.shape:
                residual = nn.functional.interpolate(
                    residual, size=(hidden_state.shape[2], hidden_state.shape[3]), mode="bilinear", align_corners=False
                )
            hidden_state = hidden_state + self.residual_layer1(residual)

        hidden_state = self.residual_layer2(hidden_state)
        hidden_state = nn.functional.interpolate(
            hidden_state, scale_factor=2, mode="bilinear", align_corners=self.align_corners
        )
        hidden_state = self.projection(hidden_state)

        return hidden_state
//  [python]
def _resize_pos_embed(self, posemb, grid_size_height, grid_size_width, start_index=1):
        posemb_tok = posemb[:, :start_index]
        posemb_grid = posemb[0, start_index:]

        old_grid_size = int(math.sqrt(len(posemb_grid)))

        posemb_grid = posemb_grid.reshape(1, old_grid_size, old_grid_size, -1).permute(0, 3, 1, 2)
        posemb_grid = nn.functional.interpolate(posemb_grid, size=(grid_size_height, grid_size_width), mode="bilinear")
        posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, grid_size_height * grid_size_width, -1)

        posemb = torch.cat([posemb_tok, posemb_grid], dim=1)

        return posemb
//  [python]
def extract_id(self, np_source, f5p):
        Xs = warp_and_crop_face(
            np_source,
            f5p,
            reference_pts=get_reference_facial_points(default_square=True),
            crop_size=(256, 256))

        Xs = Image.fromarray(Xs)
        Xs = self.test_transform(Xs)
        Xs = Xs.unsqueeze(0).to(self.device)
        with torch.no_grad():
            embeds, Xs_feats = self.arcface(
                F.upsample(
                    Xs, (112, 112), mode='bilinear'))
        return embeds, Xs
//  [python]
def extract_id(self, np_source, f5p):
        Xs = warp_and_crop_face(
            np_source,
            f5p,
            reference_pts=get_reference_facial_points(default_square=True),
            crop_size=(256, 256))

        Xs = Image.fromarray(Xs)
        Xs = self.test_transform(Xs)
        Xs = Xs.unsqueeze(0).to(self.device)
        with torch.no_grad():
            embeds, Xs_feats = self.arcface(
                F.interpolate(
                    Xs, (112, 112), mode='bilinear', align_corners=True))
        return embeds, Xs
//  [python]
def encode(self, vae, pixels, mask, grow_mask_by=6):
        x = (pixels.shape[1] // vae.downscale_ratio) * vae.downscale_ratio
        y = (pixels.shape[2] // vae.downscale_ratio) * vae.downscale_ratio
        mask = t.nn.functional.upsample(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")

        pixels = pixels.clone()
        if pixels.shape[1] != x or pixels.shape[2] != y:
            x_offset = (pixels.shape[1] % vae.downscale_ratio) // 2
            y_offset = (pixels.shape[2] % vae.downscale_ratio) // 2
            pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
            mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset]

        #grow mask by a few pixels to keep things seamless in latent space
        if grow_mask_by == 0:
            mask_erosion = mask
        else:
            kernel_tensor = torch.ones((1, 1, grow_mask_by, grow_mask_by))
            padding = math.ceil((grow_mask_by - 1) / 2)

            mask_erosion = torch.clamp(torch.nn.functional.conv2d(mask.round(), kernel_tensor, padding=padding), 0, 1)

        m = (1.0 - mask.round()).squeeze(1)
        for i in range(3):
            pixels[:,:,:,i] -= 0.5
            pixels[:,:,:,i] *= m
            pixels[:,:,:,i] += 0.5
        t = vae.encode(pixels)

        return ({"samples":t, "noise_mask": (mask_erosion[:,:,:x,:y].round())}, )
//  [python]
def encode(self, vae, pixels, mask, grow_mask_by=6):
        x = (pixels.shape[1] // vae.downscale_ratio) * vae.downscale_ratio
        y = (pixels.shape[2] // vae.downscale_ratio) * vae.downscale_ratio
        mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")

        pixels = pixels.clone()
        if pixels.shape[1] != x or pixels.shape[2] != y:
            x_offset = (pixels.shape[1] % vae.downscale_ratio) // 2
            y_offset = (pixels.shape[2] % vae.downscale_ratio) // 2
            pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
            mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset]

        #grow mask by a few pixels to keep things seamless in latent space
        if grow_mask_by == 0:
            mask_erosion = mask
        else:
            kernel_tensor = torch.ones((1, 1, grow_mask_by, grow_mask_by))
            padding = math.ceil((grow_mask_by - 1) / 2)

            mask_erosion = torch.clamp(torch.nn.functional.conv2d(mask.round(), kernel_tensor, padding=padding), 0, 1)

        m = (1.0 - mask.round()).squeeze(1)
        for i in range(3):
            pixels[:,:,:,i] -= 0.5
            pixels[:,:,:,i] *= m
            pixels[:,:,:,i] += 0.5
        t = vae.encode(pixels)

        return ({"samples":t, "noise_mask": (mask_erosion[:,:,:x,:y].round())}, )
//  [python]
def common_upscale(samples, width, height, upscale_method, crop):
        if crop == "center":
            old_width = samples.shape[3]
            old_height = samples.shape[2]
            old_aspect = old_width / old_height
            new_aspect = width / height
            x = 0
            y = 0
            if old_aspect > new_aspect:
                x = round((old_width - old_width * (new_aspect / old_aspect)) / 2)
            elif old_aspect < new_aspect:
                y = round((old_height - old_height * (old_aspect / new_aspect)) / 2)
            s = samples[:,:,y:old_height-y,x:old_width-x]
        else:
            s = samples

        if upscale_method == "bislerp":
            return bislerp(s, width, height)
        elif upscale_method == "lanczos":
            return lanczos(s, width, height)
        else:
            return t.interpolate(s, size=(height, width), mode=upscale_method)
//  [python]
def common_upscale(samples, width, height, upscale_method, crop):
        if crop == "center":
            old_width = samples.shape[3]
            old_height = samples.shape[2]
            old_aspect = old_width / old_height
            new_aspect = width / height
            x = 0
            y = 0
            if old_aspect > new_aspect:
                x = round((old_width - old_width * (new_aspect / old_aspect)) / 2)
            elif old_aspect < new_aspect:
                y = round((old_height - old_height * (old_aspect / new_aspect)) / 2)
            s = samples[:,:,y:old_height-y,x:old_width-x]
        else:
            s = samples

        if upscale_method == "bislerp":
            return bislerp(s, width, height)
        elif upscale_method == "lanczos":
            return lanczos(s, width, height)
        else:
            return torch.nn.functional.interpolate(s, size=(height, width), mode=upscale_method)
//  [python]
def thin_plate_dense(x1, y1, shape, step, lambd=.0, unroll_step_size=2 ** 12):
    device = x1.device
    D, H, W = shape
    D1, H1, W1 = D // step, H // step, W // step

    x2 = F.affine_grid(torch.eye(3, 4, device=device).unsqueeze(0), (1, 1, D1, H1, W1), align_corners=True).view(-1, 3)
    tps = TPS()
    theta = tps.fit(x1[0], y1[0], lambd)

    y2 = torch.zeros((1, D1 * H1 * W1, 3), device=device)
    N = D1 * H1 * W1
    n = math.ceil(N / unroll_step_size)
    for j in range(n):
        j1 = j * unroll_step_size
        j2 = min((j + 1) * unroll_step_size, N)
        y2[0, j1:j2, :] = tps.z(x2[j1:j2], x1[0], theta)

    y2 = y2.view(1, D1, H1, W1, 3).permute(0, 4, 1, 2, 3)
    y2 = t.upsample(y2, size=(D, H, W), mode='trilinear', align_corners=True).permute(0, 2, 3, 4, 1)

    return y2
//  [python]
def thin_plate_dense(x1, y1, shape, step, lambd=.0, unroll_step_size=2 ** 12):
    device = x1.device
    D, H, W = shape
    D1, H1, W1 = D // step, H // step, W // step

    x2 = F.affine_grid(torch.eye(3, 4, device=device).unsqueeze(0), (1, 1, D1, H1, W1), align_corners=True).view(-1, 3)
    tps = TPS()
    theta = tps.fit(x1[0], y1[0], lambd)

    y2 = torch.zeros((1, D1 * H1 * W1, 3), device=device)
    N = D1 * H1 * W1
    n = math.ceil(N / unroll_step_size)
    for j in range(n):
        j1 = j * unroll_step_size
        j2 = min((j + 1) * unroll_step_size, N)
        y2[0, j1:j2, :] = tps.z(x2[j1:j2], x1[0], theta)

    y2 = y2.view(1, D1, H1, W1, 3).permute(0, 4, 1, 2, 3)
    y2 = F.interpolate(y2, (D, H, W), mode='trilinear', align_corners=True).permute(0, 2, 3, 4, 1)

    return y2
//  [python]
def forward(
        self,
        batch,
    ):
        """
        Parameters
        ----------
        batch
            A dictionary containing the input mini-batch data.
            We need to use the keys with the model prefix to index required data.

        Returns
        -------
            A dictionary with mask predictions.
        """
        # binary
        if self.num_classes == 1:
            rets = self.model(batch[self.image_key], multimask_output=False, output_moe_loss=self.output_moe_loss)
            pred_masks = rets.pred_masks[:, 0, :, :, :]
            pred_masks = t.interpolate(
                pred_masks, (self.image_size, self.image_size), mode="bilinear", align_corners=False
            )
            if self.training:
                rets_dict = {self.prefix: {LOGITS: pred_masks}}
            else:
                rets_dict = {self.prefix: {LOGITS: pred_masks, LABEL: batch[self.label_key]}}
        # multi-class
        else:
            rets = self.model(batch[self.image_key], multimask_output=False, output_moe_loss=self.output_moe_loss)
            rets, class_predictions = rets
            pred_masks = rets.pred_masks[:, 0, :, :, :]
            pred_classes = class_predictions[:, 0, :, :]
            pred_masks = t.interpolate(
                pred_masks, (self.image_size, self.image_size), mode="bilinear", align_corners=False
            )
            if self.training:
                rets_dict = {self.prefix: {LOGITS: pred_masks, CLASS_LOGITS: pred_classes}}
            else:
                rets_dict = {
                    self.prefix: {
                        LOGITS: pred_masks,
                        CLASS_LOGITS: pred_classes,
                        LABEL: batch[self.label_key],
                    }
                }
        if self.output_moe_loss:
            rets_dict[self.prefix].update({MOE_LOSS: rets.vision_moe_loss})
        return rets_dict
//  [python]
def forward(
        self,
        batch,
    ):
        """
        Parameters
        ----------
        batch
            A dictionary containing the input mini-batch data.
            We need to use the keys with the model prefix to index required data.

        Returns
        -------
            A dictionary with mask predictions.
        """
        # binary
        if self.num_classes == 1:
            rets = self.model(batch[self.image_key], multimask_output=False, output_moe_loss=self.output_moe_loss)
            pred_masks = rets.pred_masks[:, 0, :, :, :]
            pred_masks = F.interpolate(
                pred_masks, (self.image_size, self.image_size), mode="bilinear", align_corners=False
            )
            if self.training:
                rets_dict = {self.prefix: {LOGITS: pred_masks}}
            else:
                rets_dict = {self.prefix: {LOGITS: pred_masks, LABEL: batch[self.label_key]}}
        # multi-class
        else:
            rets = self.model(batch[self.image_key], multimask_output=False, output_moe_loss=self.output_moe_loss)
            rets, class_predictions = rets
            pred_masks = rets.pred_masks[:, 0, :, :, :]
            pred_classes = class_predictions[:, 0, :, :]
            pred_masks = F.interpolate(
                pred_masks, (self.image_size, self.image_size), mode="bilinear", align_corners=False
            )
            if self.training:
                rets_dict = {self.prefix: {LOGITS: pred_masks, CLASS_LOGITS: pred_classes}}
            else:
                rets_dict = {
                    self.prefix: {
                        LOGITS: pred_masks,
                        CLASS_LOGITS: pred_classes,
                        LABEL: batch[self.label_key],
                    }
                }
        if self.output_moe_loss:
            rets_dict[self.prefix].update({MOE_LOSS: rets.vision_moe_loss})
        return rets_dict
//  [python]
def decode(
        self,
        z: torch.FloatTensor,
        generator: Optional[torch.Generator] = None,
        return_dict: bool = True,
        num_inference_steps: int = 2,
    ) -> Union[DecoderOutput, Tuple[torch.FloatTensor]]:
        """
        Decodes the input latent vector `z` using the consistency decoder VAE model.

        Args:
            z (torch.FloatTensor): The input latent vector.
            generator (Optional[torch.Generator]): The random number generator. Default is None.
            return_dict (bool): Whether to return the output as a dictionary. Default is True.
            num_inference_steps (int): The number of inference steps. Default is 2.

        Returns:
            Union[DecoderOutput, Tuple[torch.FloatTensor]]: The decoded output.

        """
        z = (z * self.config.scaling_factor - self.means) / self.stds

        scale_factor = 2 ** (len(self.config.block_out_channels) - 1)
        z = F.interpolate(z, mode="nearest", scale_factor=scale_factor)

        batch_size, _, height, width = z.shape

        self.decoder_scheduler.set_timesteps(num_inference_steps, device=self.device)

        x_t = self.decoder_scheduler.init_noise_sigma * randn_tensor(
            (batch_size, 3, height, width), generator=generator, dtype=z.dtype, device=z.device
        )

        for t in self.decoder_scheduler.timesteps:
            model_input = torch.concat([self.decoder_scheduler.scale_model_input(x_t, t), z], dim=1)
            model_output = self.decoder_unet(model_input, t).sample[:, :3, :, :]
            prev_sample = self.decoder_scheduler.step(model_output, t, x_t, generator).prev_sample
            x_t = prev_sample

        x_0 = x_t

        if not return_dict:
            return (x_0,)

        return DecoderOutput(sample=x_0)
//  [python]
def sdsp(x: torch.Tensor, data_range: Union[int, float]=255, omega_0: float=0.021, sigma_f: float=1.34, sigma_d: float=145.0, sigma_c: float=0.001) ->torch.Tensor:
    """SDSP algorithm for salient region detection from a given image.

    Supports only colour images with RGB channel order.

    Args:
        x: Tensor. Shape :math:`(N, 3, H, W)`.
        data_range: Maximum value range of images (usually 1.0 or 255).
        omega_0: coefficient for log Gabor filter
        sigma_f: coefficient for log Gabor filter
        sigma_d: coefficient for the central areas, which have a bias towards attention
        sigma_c: coefficient for the warm colors, which have a bias towards attention

    Returns:
        torch.Tensor: Visual saliency map
    """
    x = x / float(data_range) * 255
    size = x.size()
    size_to_use = 256, 256
    x = t.interpolate(input=x, size=size_to_use, mode='bilinear', align_corners=False)
    x_lab = rgb2lab(x, data_range=255)
    lg = _log_gabor(size_to_use, omega_0, sigma_f).view(1, 1, *size_to_use)
    recommended_torch_version = _parse_version('1.8.0')
    torch_version = _parse_version(torch.__version__)
    if len(torch_version) != 0 and torch_version >= recommended_torch_version:
        x_fft = torch.fft.fft2(x_lab)
        x_ifft_real = torch.fft.ifft2(x_fft * lg).real
    else:
        x_fft = torch.rfft(x_lab, 2, onesided=False)
        x_ifft_real = torch.ifft(x_fft * lg.unsqueeze(-1), 2)[..., 0]
    s_f = x_ifft_real.pow(2).sum(dim=1, keepdim=True).sqrt()
    coordinates = torch.stack(get_meshgrid(size_to_use), dim=0)
    coordinates = coordinates * size_to_use[0] + 1
    s_d = torch.exp(-torch.sum(coordinates ** 2, dim=0) / sigma_d ** 2).view(1, 1, *size_to_use)
    eps = torch.finfo(x_lab.dtype).eps
    min_x = x_lab.min(dim=-1, keepdim=True).values.min(dim=-2, keepdim=True).values
    max_x = x_lab.max(dim=-1, keepdim=True).values.max(dim=-2, keepdim=True).values
    normalized = (x_lab - min_x) / (max_x - min_x + eps)
    norm = normalized[:, 1:].pow(2).sum(dim=1, keepdim=True)
    s_c = 1 - torch.exp(-norm / sigma_c ** 2)
    vs_m = s_f * s_d * s_c
    vs_m = t.interpolate(vs_m, size[-2:], mode='bilinear', align_corners=True)
    min_vs_m = vs_m.min(dim=-1, keepdim=True).values.min(dim=-2, keepdim=True).values
    max_vs_m = vs_m.max(dim=-1, keepdim=True).values.max(dim=-2, keepdim=True).values
    return (vs_m - min_vs_m) / (max_vs_m - min_vs_m + eps)
//  [python]
def sdsp(x: torch.Tensor, data_range: Union[int, float]=255, omega_0: float=0.021, sigma_f: float=1.34, sigma_d: float=145.0, sigma_c: float=0.001) ->torch.Tensor:
    """SDSP algorithm for salient region detection from a given image.

    Supports only colour images with RGB channel order.

    Args:
        x: Tensor. Shape :math:`(N, 3, H, W)`.
        data_range: Maximum value range of images (usually 1.0 or 255).
        omega_0: coefficient for log Gabor filter
        sigma_f: coefficient for log Gabor filter
        sigma_d: coefficient for the central areas, which have a bias towards attention
        sigma_c: coefficient for the warm colors, which have a bias towards attention

    Returns:
        torch.Tensor: Visual saliency map
    """
    x = x / float(data_range) * 255
    size = x.size()
    size_to_use = 256, 256
    x = interpolate(input=x, size=size_to_use, mode='bilinear', align_corners=False)
    x_lab = rgb2lab(x, data_range=255)
    lg = _log_gabor(size_to_use, omega_0, sigma_f).view(1, 1, *size_to_use)
    recommended_torch_version = _parse_version('1.8.0')
    torch_version = _parse_version(torch.__version__)
    if len(torch_version) != 0 and torch_version >= recommended_torch_version:
        x_fft = torch.fft.fft2(x_lab)
        x_ifft_real = torch.fft.ifft2(x_fft * lg).real
    else:
        x_fft = torch.rfft(x_lab, 2, onesided=False)
        x_ifft_real = torch.ifft(x_fft * lg.unsqueeze(-1), 2)[..., 0]
    s_f = x_ifft_real.pow(2).sum(dim=1, keepdim=True).sqrt()
    coordinates = torch.stack(get_meshgrid(size_to_use), dim=0)
    coordinates = coordinates * size_to_use[0] + 1
    s_d = torch.exp(-torch.sum(coordinates ** 2, dim=0) / sigma_d ** 2).view(1, 1, *size_to_use)
    eps = torch.finfo(x_lab.dtype).eps
    min_x = x_lab.min(dim=-1, keepdim=True).values.min(dim=-2, keepdim=True).values
    max_x = x_lab.max(dim=-1, keepdim=True).values.max(dim=-2, keepdim=True).values
    normalized = (x_lab - min_x) / (max_x - min_x + eps)
    norm = normalized[:, 1:].pow(2).sum(dim=1, keepdim=True)
    s_c = 1 - torch.exp(-norm / sigma_c ** 2)
    vs_m = s_f * s_d * s_c
    vs_m = interpolate(vs_m, size[-2:], mode='bilinear', align_corners=True)
    min_vs_m = vs_m.min(dim=-1, keepdim=True).values.min(dim=-2, keepdim=True).values
    max_vs_m = vs_m.max(dim=-1, keepdim=True).values.max(dim=-2, keepdim=True).values
    return (vs_m - min_vs_m) / (max_vs_m - min_vs_m + eps)
//  [python]
def forward(self, feats=Tuple[Tensor]) -> Tuple[Tensor]:
        assert len(feats) == self.num_scales

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

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

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

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

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

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

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

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

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

        return tuple(outs)
//  [python]
def forward(self, x, skip=None, upscale=True):
        if self.upscale and upscale:
            if self.use_pixel_shuffle:
                x = self.pixel_shuffle(x)
            else:
                x = t.interpolate(x, scale_factor=2, mode="nearest")
        if skip is not None:
            x = torch.cat([x, skip], dim=1)
            x = self.attention1(x)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.attention2(x)
        return x
//  [python]
def forward(self, x, skip=None, upscale=True):
        if self.upscale and upscale:
            if self.use_pixel_shuffle:
                x = self.pixel_shuffle(x)
            else:
                x = F.interpolate(x, scale_factor=2, mode="nearest")
        if skip is not None:
            x = torch.cat([x, skip], dim=1)
            x = self.attention1(x)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.attention2(x)
        return x
//  [python]
def process(self, im):
        preds = []
        imt = self.img2tensor(im)
        imt = t.interpolate(imt, (self.resolution, self.resolution))

        with torch.no_grad():
            img_out, __ = self.model(imt)

        face = self.tensor2img(img_out)

        return face, preds
//  [python]
def process(self, im):
        preds = []
        imt = self.img2tensor(im)
        imt = F.interpolate(imt, (self.resolution, self.resolution))

        with torch.no_grad():
            img_out, __ = self.model(imt)

        face = self.tensor2img(img_out)

        return face, preds
//  [python]
def compute_3d_params(self, Xs, Xt):
        kp_fuse = {}
        kp_t = {}

        c_s = self.f_3d(
            t.interpolate(Xs * 0.5 + 0.5, size=224, mode='bilinear'))
        c_t = self.f_3d(
            t.interpolate(Xt * 0.5 + 0.5, size=224, mode='bilinear'))
        c_fuse = torch.cat(((c_s[:, :80] + c_t[:, :80]) / 2, c_t[:, 80:]),
                           dim=1)
        _, _, _, q_fuse = self.face_model.compute_for_render(c_fuse)
        q_fuse = q_fuse / 224
        q_fuse[..., 1] = 1 - q_fuse[..., 1]
        q_fuse = q_fuse * 2 - 1
        delta = int((17 - self.num_kp) / 2)

        _, _, _, q_t = self.face_model.compute_for_render(c_t)
        q_t = q_t / 224
        q_t[..., 1] = 1 - q_t[..., 1]
        q_t = q_t * 2 - 1

        kp_fuse['value'] = q_fuse[:, delta:17 - delta, :]
        kp_t['value'] = q_t[:, delta:17 - delta, :]

        return kp_fuse, kp_t
//  [python]
def compute_3d_params(self, Xs, Xt):
        kp_fuse = {}
        kp_t = {}

        c_s = self.f_3d(
            F.interpolate(Xs * 0.5 + 0.5, size=224, mode='bilinear'))
        c_t = self.f_3d(
            F.interpolate(Xt * 0.5 + 0.5, size=224, mode='bilinear'))
        c_fuse = torch.cat(((c_s[:, :80] + c_t[:, :80]) / 2, c_t[:, 80:]),
                           dim=1)
        _, _, _, q_fuse = self.face_model.compute_for_render(c_fuse)
        q_fuse = q_fuse / 224
        q_fuse[..., 1] = 1 - q_fuse[..., 1]
        q_fuse = q_fuse * 2 - 1
        delta = int((17 - self.num_kp) / 2)

        _, _, _, q_t = self.face_model.compute_for_render(c_t)
        q_t = q_t / 224
        q_t[..., 1] = 1 - q_t[..., 1]
        q_t = q_t * 2 - 1

        kp_fuse['value'] = q_fuse[:, delta:17 - delta, :]
        kp_t['value'] = q_t[:, delta:17 - delta, :]

        return kp_fuse, kp_t
//  [python]
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        assert inputs.shape[1] == self.channels
        if self.use_conv_transpose:
            return self.conv(inputs)

        outputs = F.upsample(inputs, scale_factor=2.0, mode="nearest")

        if self.use_conv:
            outputs = self.conv(outputs)

        return outputs
//  [python]
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        assert inputs.shape[1] == self.channels
        if self.use_conv_transpose:
            return self.conv(inputs)

        outputs = F.interpolate(inputs, scale_factor=2.0, mode="nearest")

        if self.use_conv:
            outputs = self.conv(outputs)

        return outputs
//  [python]
def _predict(self, inputs, target_sizes):
        with torch.no_grad():
            outputs = self.model(**inputs)

        predicted_depth = outputs.predicted_depth
        prediction = torch.nn.functional.interpolate(
            predicted_depth.unsqueeze(1),
            size=target_sizes[0],
            mode="bicubic",
            align_corners=False,
        )
        prediction = prediction.squeeze(1).cpu().numpy()

        return to_heatmap(prediction)
//  [python]
def test_upsamplingBilinear2d(self):
        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]:
                    m = nn.Upsample(scale_factor=scale_factor, **kwargs)
                    in_t = torch.ones(1, 2, 2, 2).contiguous(memory_format=memory_format)
                    out_size = int(math.floor(in_t.shape[-1] * scale_factor))
                    with warnings.catch_warnings(record=True) as w:
                        out_t = m(in_t)
                    self.assertEqual(torch.ones(1, 2, out_size, out_size), out_t.data)
                    # Assert that memory format is carried through to the output
                    self.assertTrue(out_t.is_contiguous(memory_format=memory_format))

                    input = torch.randn(1, 2, 2, 2, requires_grad=True)
                    gradcheck(lambda x: t.interpolate(x, out_size, **kwargs), [input])
//  [python]
def test_upsamplingBilinear2d(self):
        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]:
                    m = nn.Upsample(scale_factor=scale_factor, **kwargs)
                    in_t = torch.ones(1, 2, 2, 2).contiguous(memory_format=memory_format)
                    out_size = int(math.floor(in_t.shape[-1] * scale_factor))
                    with warnings.catch_warnings(record=True) as w:
                        out_t = m(in_t)
                    self.assertEqual(torch.ones(1, 2, out_size, out_size), out_t.data)
                    # Assert that memory format is carried through to the output
                    self.assertTrue(out_t.is_contiguous(memory_format=memory_format))

                    input = torch.randn(1, 2, 2, 2, requires_grad=True)
                    gradcheck(lambda x: F.interpolate(x, out_size, **kwargs), [input])
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (Tensor): Input tensor with shape (B, C, H, W).

        Returns:
            Tensor: Forward results.
        """
        H, W = x.shape[2:]
        x = self.check_image_size(x)

        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        if self.upsampler == 'pixelshuffle':
            # for classical SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))
        elif self.upsampler == 'pixelshuffledirect':
            # for lightweight SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.upsample(x)
        elif self.upsampler == 'nearest+conv':
            # for real-world SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.lrelu(
                self.conv_up1(
                    t.nn.functional.interpolate(
                        x, scale_factor=2, mode='nearest')))
            if self.upscale == 4:
                x = self.lrelu(
                    self.conv_up2(
                        t.nn.functional.interpolate(
                            x, scale_factor=2, mode='nearest')))
            x = self.conv_last(self.lrelu(self.conv_hr(x)))
        else:
            # for image denoising and JPEG compression artifact reduction
            x_first = self.conv_first(x)
            res = self.conv_after_body(
                self.forward_features(x_first)) + x_first
            x = x + self.conv_last(res)

        x = x / self.img_range + self.mean

        return x[:, :, :H * self.upscale, :W * self.upscale]
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (Tensor): Input tensor with shape (B, C, H, W).

        Returns:
            Tensor: Forward results.
        """
        H, W = x.shape[2:]
        x = self.check_image_size(x)

        self.mean = self.mean.type_as(x)
        x = (x - self.mean) * self.img_range

        if self.upsampler == 'pixelshuffle':
            # for classical SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.conv_last(self.upsample(x))
        elif self.upsampler == 'pixelshuffledirect':
            # for lightweight SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.upsample(x)
        elif self.upsampler == 'nearest+conv':
            # for real-world SR
            x = self.conv_first(x)
            x = self.conv_after_body(self.forward_features(x)) + x
            x = self.conv_before_upsample(x)
            x = self.lrelu(
                self.conv_up1(
                    torch.nn.functional.interpolate(
                        x, scale_factor=2, mode='nearest')))
            if self.upscale == 4:
                x = self.lrelu(
                    self.conv_up2(
                        torch.nn.functional.interpolate(
                            x, scale_factor=2, mode='nearest')))
            x = self.conv_last(self.lrelu(self.conv_hr(x)))
        else:
            # for image denoising and JPEG compression artifact reduction
            x_first = self.conv_first(x)
            res = self.conv_after_body(
                self.forward_features(x_first)) + x_first
            x = x + self.conv_last(res)

        x = x / self.img_range + self.mean

        return x[:, :, :H * self.upscale, :W * self.upscale]
//  [python]
def get_mean_attention_map(attn, token, shape):
    attn = attn[:, :, token, 1:]
    attn = attn.unflatten(2, torch.Size([shape[2] // 16,
                                         shape[3] // 16])).float()
    attn = torch.nn.functional.upsample(
        attn, size=shape[2:], mode='bicubic').squeeze(0)

    all_attn = torch.mean(attn, 0)

    return all_attn
//  [python]
def get_mean_attention_map(attn, token, shape):
    attn = attn[:, :, token, 1:]
    attn = attn.unflatten(2, torch.Size([shape[2] // 16,
                                         shape[3] // 16])).float()
    attn = torch.nn.functional.interpolate(
        attn, size=shape[2:], mode='bicubic', align_corners=False).squeeze(0)

    all_attn = torch.mean(attn, 0)

    return all_attn