jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def generate_gaussian_pyramid(feature_map, valid_mask, return_mask, num_photo_level, gauss_kernel):
        feature_map_pyramid = list()
        mask_pyramid = list()
        feature_map_pyramid.append(feature_map)
        mask_pyramid.append(valid_mask)

        _, feat_channel, height, width = feature_map.shape
        feature_map = feature_map.reshape(feat_channel, 1, height, width)
        for _ in range(num_photo_level - 1):
            raw_feature_map = F.conv2d(feature_map * valid_mask, weight=gauss_kernel, bias=None,
                                       stride=2, padding=1, dilation=1, groups=1)
            height, width = raw_feature_map.shape[2:]
            raw_mask = F.conv2d(valid_mask,
                                weight=gauss_kernel, bias=None, stride=2,
                                padding=1, dilation=1, groups=1)
            valid_mask = t.nn.functional.upsample(
                valid_mask, size=raw_feature_map.shape[2:], mode='nearest')

            if return_mask:
                mask_pyramid.append(valid_mask.reshape(1, 1, height, width))
            # correct the feature value around the edges and corners
            feature_map = raw_feature_map / (raw_mask + 1.0e-8)
            feature_map_pyramid.append(
                feature_map.reshape(1, feat_channel, height, width))
        if return_mask:
            return list(reversed(feature_map_pyramid)), list(reversed(mask_pyramid))
        else:
            return list(reversed(feature_map_pyramid))
//  [python]
def generate_gaussian_pyramid(feature_map, valid_mask, return_mask, num_photo_level, gauss_kernel):
        feature_map_pyramid = list()
        mask_pyramid = list()
        feature_map_pyramid.append(feature_map)
        mask_pyramid.append(valid_mask)

        _, feat_channel, height, width = feature_map.shape
        feature_map = feature_map.reshape(feat_channel, 1, height, width)
        for _ in range(num_photo_level - 1):
            raw_feature_map = F.conv2d(feature_map * valid_mask, weight=gauss_kernel, bias=None,
                                       stride=2, padding=1, dilation=1, groups=1)
            height, width = raw_feature_map.shape[2:]
            raw_mask = F.conv2d(valid_mask,
                                weight=gauss_kernel, bias=None, stride=2,
                                padding=1, dilation=1, groups=1)
            valid_mask = F.interpolate(
                valid_mask, size=raw_feature_map.shape[2:], mode='nearest')

            if return_mask:
                mask_pyramid.append(valid_mask.reshape(1, 1, height, width))
            # correct the feature value around the edges and corners
            feature_map = raw_feature_map / (raw_mask + 1.0e-8)
            feature_map_pyramid.append(
                feature_map.reshape(1, feat_channel, height, width))
        if return_mask:
            return list(reversed(feature_map_pyramid)), list(reversed(mask_pyramid))
        else:
            return list(reversed(feature_map_pyramid))
//  [python]
def sample_unipc(model, noise, image, sigmas, sampling_function, max_denoise, extra_args=None, callback=None, disable=None, noise_mask=None, variant='bh1'):
        to_zero = False
        if sigmas[-1] == 0:
            timesteps = t.nn.functional.interpolate(sigmas[None,None,:-1], size=(len(sigmas),), mode='linear')[0][0]
            to_zero = True
        else:
            timesteps = sigmas.clone()

        for s in range(timesteps.shape[0]):
            timesteps[s] = (model.sigma_to_t(timesteps[s]) / 1000) + (1 / len(model.sigmas))

        ns = NoiseScheduleVP('discrete', alphas_cumprod=model.inner_model.alphas_cumprod)

        if image is not None:
            img = image * ns.marginal_alpha(timesteps[0])
            if max_denoise:
                noise_mult = 1.0
            else:
                noise_mult = ns.marginal_std(timesteps[0])
            img += noise * noise_mult
        else:
            img = noise

        if to_zero:
            timesteps[-1] = (1 / len(model.sigmas))

        device = noise.device

        if model.parameterization == "v":
            model_type = "v"
        else:
            model_type = "noise"

        model_fn = model_wrapper(
            model.inner_model.inner_model.apply_model,
            sampling_function,
            ns,
            model_type=model_type,
            guidance_type="uncond",
            model_kwargs=extra_args,
        )

        uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, noise_mask=noise_mask, masked_image=image, noise=noise, variant=variant)
        x = uni_pc.sample(img, timesteps=timesteps, skip_type="time_uniform", method="multistep", order=3, lower_order_final=True)
        if not to_zero:
            x /= ns.marginal_alpha(timesteps[-1])
        return x
//  [python]
def sample_unipc(model, noise, image, sigmas, sampling_function, max_denoise, extra_args=None, callback=None, disable=None, noise_mask=None, variant='bh1'):
        to_zero = False
        if sigmas[-1] == 0:
            timesteps = torch.nn.functional.interpolate(sigmas[None,None,:-1], size=(len(sigmas),), mode='linear')[0][0]
            to_zero = True
        else:
            timesteps = sigmas.clone()

        for s in range(timesteps.shape[0]):
            timesteps[s] = (model.sigma_to_t(timesteps[s]) / 1000) + (1 / len(model.sigmas))

        ns = NoiseScheduleVP('discrete', alphas_cumprod=model.inner_model.alphas_cumprod)

        if image is not None:
            img = image * ns.marginal_alpha(timesteps[0])
            if max_denoise:
                noise_mult = 1.0
            else:
                noise_mult = ns.marginal_std(timesteps[0])
            img += noise * noise_mult
        else:
            img = noise

        if to_zero:
            timesteps[-1] = (1 / len(model.sigmas))

        device = noise.device

        if model.parameterization == "v":
            model_type = "v"
        else:
            model_type = "noise"

        model_fn = model_wrapper(
            model.inner_model.inner_model.apply_model,
            sampling_function,
            ns,
            model_type=model_type,
            guidance_type="uncond",
            model_kwargs=extra_args,
        )

        uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, noise_mask=noise_mask, masked_image=image, noise=noise, variant=variant)
        x = uni_pc.sample(img, timesteps=timesteps, skip_type="time_uniform", method="multistep", order=3, lower_order_final=True)
        if not to_zero:
            x /= ns.marginal_alpha(timesteps[-1])
        return x
//  [python]
def interpolate_pos_encoding(self, height: int, width: int) -> torch.Tensor:
        patch_pos_embed = self.position_embeddings[:, 1:]
        num_patches = patch_pos_embed.shape[1]
        pretrain_patch_size = int(math.sqrt(num_patches))

        if pretrain_patch_size != height or pretrain_patch_size != width:
            patch_pos_embed = t.interpolate(
                patch_pos_embed.reshape(1, pretrain_patch_size, pretrain_patch_size, -1).permute(0, 3, 1, 2),
                size=(height, width),
                mode="bicubic",
                align_corners=False,
            )

            return patch_pos_embed.permute(0, 2, 3, 1)
        else:
            return patch_pos_embed.reshape(1, height, width, -1)
//  [python]
def interpolate_pos_encoding(self, height: int, width: int) -> torch.Tensor:
        patch_pos_embed = self.position_embeddings[:, 1:]
        num_patches = patch_pos_embed.shape[1]
        pretrain_patch_size = int(math.sqrt(num_patches))

        if pretrain_patch_size != height or pretrain_patch_size != width:
            patch_pos_embed = F.interpolate(
                patch_pos_embed.reshape(1, pretrain_patch_size, pretrain_patch_size, -1).permute(0, 3, 1, 2),
                size=(height, width),
                mode="bicubic",
                align_corners=False,
            )

            return patch_pos_embed.permute(0, 2, 3, 1)
        else:
            return patch_pos_embed.reshape(1, height, width, -1)
//  [python]
def __call__(self, data):
        img, boxes = data
        img = t.nn.functional.interpolate(img.unsqueeze(0), size=self.size, mode="nearest").squeeze(0)
        return img, boxes
//  [python]
def __call__(self, data):
        img, boxes = data
        img = F.interpolate(img.unsqueeze(0), size=self.size, mode="nearest").squeeze(0)
        return img, boxes
//  [python]
def forward(self, concat_x, features_list):
        concat_x = self.conv(concat_x)
        score = self.enhanced_attention(concat_x)
        assert len(features_list) == self.out_features_num
        if self.type not in ['scale_channel_spatial', 'scale_spatial']:
            shape = features_list[0].shape[2:]
            score = F.upsample_bilinear(score, size=shape)
        x = []
        for i in range(self.out_features_num):
            x.append(score[:, i:i+1] * features_list[i])
        return torch.cat(x, dim=1)
//  [python]
def forward(self, concat_x, features_list):
        concat_x = self.conv(concat_x)
        score = self.enhanced_attention(concat_x)
        assert len(features_list) == self.out_features_num
        if self.type not in ['scale_channel_spatial', 'scale_spatial']:
            shape = features_list[0].shape[2:]
            score = F.interpolate(score, size=shape, mode='bilinear')
        x = []
        for i in range(self.out_features_num):
            x.append(score[:, i:i+1] * features_list[i])
        return torch.cat(x, dim=1)
//  [python]
def get_rel_pos(self, q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
        """
        Get relative positional embeddings according to the relative positions of
            query and key sizes.

        Args:
            q_size (int):
                size of the query.
            k_size (int):
                size of key k.
            rel_pos (`torch.Tensor`):
                relative position embeddings (L, channel).

        Returns:
            Extracted positional embeddings according to relative positions.
        """
        max_rel_dist = int(2 * max(q_size, k_size) - 1)
        # Interpolate rel pos.
        rel_pos_resized = F.upsample(
            rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
            scale_factor=max_rel_dist / rel_pos.shape[0],
            mode="linear",
        )
        rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)

        # Scale the coords with short length if shapes for q and k are different.
        q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
        k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
        relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)

        return rel_pos_resized[relative_coords.long()]
//  [python]
def get_rel_pos(self, q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
        """
        Get relative positional embeddings according to the relative positions of
            query and key sizes.

        Args:
            q_size (int):
                size of the query.
            k_size (int):
                size of key k.
            rel_pos (`torch.Tensor`):
                relative position embeddings (L, channel).

        Returns:
            Extracted positional embeddings according to relative positions.
        """
        max_rel_dist = int(2 * max(q_size, k_size) - 1)
        # Interpolate rel pos.
        rel_pos_resized = F.interpolate(
            rel_pos.reshape(1, rel_pos.shape[0], -1).permute(0, 2, 1),
            size=max_rel_dist,
            mode="linear",
        )
        rel_pos_resized = rel_pos_resized.reshape(-1, max_rel_dist).permute(1, 0)

        # Scale the coords with short length if shapes for q and k are different.
        q_coords = torch.arange(q_size)[:, None] * max(k_size / q_size, 1.0)
        k_coords = torch.arange(k_size)[None, :] * max(q_size / k_size, 1.0)
        relative_coords = (q_coords - k_coords) + (k_size - 1) * max(q_size / k_size, 1.0)

        return rel_pos_resized[relative_coords.long()]
//  [python]
def forward(self, x, shortcut, dec_idx_feat=None):
        """Forward function.

        Args:
            x (Tensor): Input feature map with shape (N, C, H, W).
            shortcut (Tensor): The shortcut connection with shape
                (N, C, H', W').
            dec_idx_feat (Tensor, optional): The decode index feature map with
                shape (N, C, H', W'). Defaults to None.

        Returns:
            Tensor: Output tensor with shape (N, C, H', W').
        """
        if dec_idx_feat is not None:
            assert shortcut.dim() == 4, (
                'shortcut must be tensor with 4 dimensions')
            x = dec_idx_feat * F.interpolate(x, size=shortcut.shape[2:])
        out = torch.cat((x, shortcut), dim=1)
        return self.conv(out)
//  [python]
def resize_mask(self, mask):
        # mask transform is applied AFTER mapper, so we need to post-process it in eval.py
        h, w = mask.shape[-2:]
        min_hw = min(h, w)
        return F.interpolate(mask, (int(h/min_hw*self.size), int(w/min_hw*self.size)), 
                    mode='nearest')
//  [python]
def forward(self, tpv_list, points=None):
    tpv_hw, tpv_zh, tpv_wz = tpv_list
    bs, c, _, _ = tpv_hw.shape

    if self.scale_h != 1 or self.scale_w != 1:
      tpv_hw = t.nn.functional.interpolate(
          tpv_hw,
          size=(self.tpv_h * self.scale_h, self.tpv_w * self.scale_w),
          mode='bilinear',
      )
    if self.scale_z != 1 or self.scale_h != 1:
      tpv_zh = t.nn.functional.interpolate(
          tpv_zh,
          size=(self.tpv_z * self.scale_z, self.tpv_h * self.scale_h),
          mode='bilinear',
      )
    if self.scale_w != 1 or self.scale_z != 1:
      tpv_wz = t.nn.functional.interpolate(
          tpv_wz,
          size=(self.tpv_w * self.scale_w, self.tpv_z * self.scale_z),
          mode='bilinear',
      )

    if points is not None:
      _, n, _ = points.shape
      points = points.reshape(bs, 1, n, 3)
      ## torch grid sample, range [-1, 1]
      points = points * 2 - 1

      sample_loc = points[:, :, :, [1, 0]]
      tpv_hw_pts = F.grid_sample(tpv_hw, sample_loc).squeeze(2)  # bs, c, n
      sample_loc = points[:, :, :, [2, 1]]
      tpv_zh_pts = F.grid_sample(tpv_zh, sample_loc).squeeze(2)
      sample_loc = points[:, :, :, [2, 0]]
      tpv_wz_pts = F.grid_sample(tpv_wz, sample_loc).squeeze(2)

      fused_pts = tpv_hw_pts + tpv_zh_pts + tpv_wz_pts
      fused_pts = fused_pts.permute(0, 2, 1)
      return fused_pts
//  [python]
def forward(self, tpv_list, points=None):
    tpv_hw, tpv_zh, tpv_wz = tpv_list
    bs, c, _, _ = tpv_hw.shape

    if self.scale_h != 1 or self.scale_w != 1:
      tpv_hw = F.interpolate(
          tpv_hw,
          size=(self.tpv_h * self.scale_h, self.tpv_w * self.scale_w),
          mode='bilinear',
      )
    if self.scale_z != 1 or self.scale_h != 1:
      tpv_zh = F.interpolate(
          tpv_zh,
          size=(self.tpv_z * self.scale_z, self.tpv_h * self.scale_h),
          mode='bilinear',
      )
    if self.scale_w != 1 or self.scale_z != 1:
      tpv_wz = F.interpolate(
          tpv_wz,
          size=(self.tpv_w * self.scale_w, self.tpv_z * self.scale_z),
          mode='bilinear',
      )

    if points is not None:
      # points: bs, n, 3
      _, n, _ = points.shape
      points = points.reshape(bs, 1, n, 3)
      ## torch grid sample, range [-1, 1]
      points = points * 2 - 1

      # sample_loc = points[:, :, :, [0, 1]]
      sample_loc = points[:, :, :, [1, 0]]
      tpv_hw_pts = F.grid_sample(tpv_hw, sample_loc).squeeze(2)  # bs, c, n
      # sample_loc = points[:, :, :, [1, 2]]
      sample_loc = points[:, :, :, [2, 1]]
      tpv_zh_pts = F.grid_sample(tpv_zh, sample_loc).squeeze(2)
      # sample_loc = points[:, :, :, [2, 0]]
      sample_loc = points[:, :, :, [2, 0]]
      tpv_wz_pts = F.grid_sample(tpv_wz, sample_loc).squeeze(2)

      fused_pts = tpv_hw_pts + tpv_zh_pts + tpv_wz_pts
      fused_pts = fused_pts.permute(0, 2, 1)
      return fused_pts
//  [python]
def forward(self, tensor_list: NestedTensor):
        xs = self.body(tensor_list.tensors)
        out: Dict[str, NestedTensor] = {}
        for name, x in xs.items():
            m = tensor_list.mask
            assert m is not None
            mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]
            out[name] = NestedTensor(x, mask)
        # import ipdb; ipdb.set_trace()
        return out
//  [python]
def reconstruct_image(recon_model, face_model, renderer, image):
    input_img = image.clone()
    input_img = t.nn.functional.interpolate(input_img, size=224, mode='bilinear', align_corners=True)
    output_coeff = recon_model(input_img)
    pred_coeffs_dict = face_model.split_coeff(output_coeff)
    pred_vertex, pred_tex, pred_color, pred_lm = face_model.compute_for_render(pred_coeffs_dict)
    pred_mask, _, pred_face = renderer(pred_vertex, face_model.face_buf, feat=pred_color)
    
    return pred_coeffs_dict, pred_face
//  [python]
def reconstruct_image(recon_model, face_model, renderer, image):
    input_img = image.clone()
    input_img = F.interpolate(input_img, size=224, mode='bilinear', align_corners=True)
    output_coeff = recon_model(input_img)
    pred_coeffs_dict = face_model.split_coeff(output_coeff)
    pred_vertex, pred_tex, pred_color, pred_lm = face_model.compute_for_render(pred_coeffs_dict)
    pred_mask, _, pred_face = renderer(pred_vertex, face_model.face_buf, feat=pred_color)
    
    return pred_coeffs_dict, pred_face
//  [python]
def forward(self, input_dict):
        """Forward Function.

        Args:
            input_dict (dict | torch.Tensor): Input dict with middle features
                or torch.Tensor.

        Returns:
            torch.Tensor: Output tensor with shape of (n, c, h, w).
        """
        hidden_feats = input_dict['hidden_feats']
        hidden_masks = input_dict['hidden_masks']
        h_key = 'h{:d}'.format(self.num_layers)
        h, h_mask = hidden_feats[h_key], hidden_masks[h_key]

        for i in range(self.num_layers, 0, -1):
            enc_h_key = f'h{i-1}'
            dec_l_key = f'dec{i}'

            h = t.interpolate(h, scale_factor=2, mode=self.interpolation)
            h_mask = t.interpolate(
                h_mask, scale_factor=2, mode=self.interpolation)

            h = torch.cat([h, hidden_feats[enc_h_key]], dim=1)
            h_mask = torch.cat([h_mask, hidden_masks[enc_h_key]], dim=1)

            h, h_mask = getattr(self, dec_l_key)(h, h_mask)

        return h, h_mask
//  [python]
def forward(self, input_dict):
        """Forward Function.

        Args:
            input_dict (dict | torch.Tensor): Input dict with middle features
                or torch.Tensor.

        Returns:
            torch.Tensor: Output tensor with shape of (n, c, h, w).
        """
        hidden_feats = input_dict['hidden_feats']
        hidden_masks = input_dict['hidden_masks']
        h_key = 'h{:d}'.format(self.num_layers)
        h, h_mask = hidden_feats[h_key], hidden_masks[h_key]

        for i in range(self.num_layers, 0, -1):
            enc_h_key = f'h{i-1}'
            dec_l_key = f'dec{i}'

            h = F.interpolate(h, scale_factor=2, mode=self.interpolation)
            h_mask = F.interpolate(
                h_mask, scale_factor=2, mode=self.interpolation)

            h = torch.cat([h, hidden_feats[enc_h_key]], dim=1)
            h_mask = torch.cat([h_mask, hidden_masks[enc_h_key]], dim=1)

            h, h_mask = getattr(self, dec_l_key)(h, h_mask)

        return h, h_mask
//  [python]
def forward(self, x):# [B, C, T]
        if self.residual:
            scale = np.product(self.scales)
            x_interp = t.interpolate(x, scale_factor=scale, mode='linear' if self.residual_linear else 'nearest', align_corners=False)
        
        for layer in self.t_convs:
            x = layer(x)
        
        if self.residual:
            if self.res_weight:
                x *= self.res_weight
            x[:, :self.res_channels] += x_interp[:, :self.res_channels]
        return x
//  [python]
def forward(self, x):# [B, C, T]
        if self.residual:
            scale = np.product(self.scales)
            x_interp = F.interpolate(x, scale_factor=scale, mode='linear' if self.residual_linear else 'nearest', align_corners=False)
        
        for layer in self.t_convs:
            x = layer(x)
        
        if self.residual:
            if self.res_weight:
                x *= self.res_weight
            x[:, :self.res_channels] += x_interp[:, :self.res_channels]
        return x
//  [python]
def post_process_prediction(self, data, outputs, ret_type):
        """
        Post-process segmentation results to match the size of original input images.

        Parameters
        ----------
        data
            Pandas data frame or dict containing the image information.
        outputs
            A list of segmentation output results.
        ret_type
            What kind of information to extract from model outputs.

        Returns
        -------
        A list of the post-processed segmentation results.
        """
        logits = [ele[ret_type] for ele in outputs]
        image_column_name = self.get_image_column_name(data)
        for idx in range(data.shape[0]):
            ori_image_size = Image.open(data[image_column_name][idx]).size  # width, height
            logits[idx] = t.interpolate(
                logits[idx].float(), (ori_image_size[1], ori_image_size[0]), mode="bilinear", align_corners=False
            )
        return logits
//  [python]
def post_process_prediction(self, data, outputs, ret_type):
        """
        Post-process segmentation results to match the size of original input images.

        Parameters
        ----------
        data
            Pandas data frame or dict containing the image information.
        outputs
            A list of segmentation output results.
        ret_type
            What kind of information to extract from model outputs.

        Returns
        -------
        A list of the post-processed segmentation results.
        """
        logits = [ele[ret_type] for ele in outputs]
        image_column_name = self.get_image_column_name(data)
        for idx in range(data.shape[0]):
            ori_image_size = Image.open(data[image_column_name][idx]).size  # width, height
            logits[idx] = F.interpolate(
                logits[idx].float(), (ori_image_size[1], ori_image_size[0]), mode="bilinear", align_corners=False
            )
        return logits
//  [python]
def test_interpolate_function_substitution(self):
        class ScriptModel(torch.jit.ScriptModule):
            @torch.jit.script_method
            def forward(self, x):
                return t.interpolate(
                    x, mode="nearest", scale_factor=2.0
                )

        class ScriptModule(torch.jit.ScriptModule):
            def __init__(self):
                super().__init__()
                self.submodule = ScriptModel()

            @torch.jit.script_method
            def forward(self, input):
                return self.submodule(input)

        x = torch.randn(1, 2, 4, 4, 6)
        self.run_test(ScriptModule(), (x,))

        @torch.jit.script
        def script_method(x):
            return t.interpolate(x, mode="nearest", scale_factor=2.0)

        class TracingModule(torch.nn.Module):
            def forward(self, x):
                return script_method(x)

        self.run_test(TracingModule(), (x,))
//  [python]
def test_interpolate_function_substitution(self):
        class ScriptModel(torch.jit.ScriptModule):
            @torch.jit.script_method
            def forward(self, x):
                return torch.nn.functional.interpolate(
                    x, mode="nearest", scale_factor=2.0
                )

        class ScriptModule(torch.jit.ScriptModule):
            def __init__(self):
                super().__init__()
                self.submodule = ScriptModel()

            @torch.jit.script_method
            def forward(self, input):
                return self.submodule(input)

        x = torch.randn(1, 2, 4, 4, 6)
        self.run_test(ScriptModule(), (x,))

        @torch.jit.script
        def script_method(x):
            return torch.nn.functional.interpolate(x, mode="nearest", scale_factor=2.0)

        class TracingModule(torch.nn.Module):
            def forward(self, x):
                return script_method(x)

        self.run_test(TracingModule(), (x,))
//  [python]
def __call__(self, images: torch.Tensor) -> torch.Tensor:
        """
        Args:
            images (torch.Tensor): tensor of size [N, 3, H, W] that contains
                BGR data (typically in uint8)
        Returns:
            images (torch.Tensor): tensor of size [N, 3, H1, W1] where
                H1 and W1 are chosen to respect the specified min and max sizes
                and preserve the original aspect ratio, the data channels
                follow BGR order and the data type is `torch.float32`
        """
        # resize with min size
        images = images.float()
        min_size = min(images.shape[-2:])
        max_size = max(images.shape[-2:])
        scale = min(self.min_size / min_size, self.max_size / max_size)
        images = t.interpolate(
            images,
            scale_factor=scale,
            mode="bilinear",
            align_corners=False,
        )
        return images
//  [python]
def __call__(self, images: torch.Tensor) -> torch.Tensor:
        """
        Args:
            images (torch.Tensor): tensor of size [N, 3, H, W] that contains
                BGR data (typically in uint8)
        Returns:
            images (torch.Tensor): tensor of size [N, 3, H1, W1] where
                H1 and W1 are chosen to respect the specified min and max sizes
                and preserve the original aspect ratio, the data channels
                follow BGR order and the data type is `torch.float32`
        """
        # resize with min size
        images = images.float()
        min_size = min(images.shape[-2:])
        max_size = max(images.shape[-2:])
        scale = min(self.min_size / min_size, self.max_size / max_size)
        images = torch.nn.functional.interpolate(
            images,
            scale_factor=scale,
            mode="bilinear",
            align_corners=False,
        )
        return images
//  [python]
def test_interpolate_half_pixel(self):
        # testing whether it uses "half_pixel" or "pytorch_half_pixel"
        # see https://github.com/onnx/onnx/blob/main/docs/Operators.md#Resize

        class MyModel(torch.nn.Module):
            def __init__(self, mode, size):
                super().__init__()
                self.mode = mode
                self.size = size

            def forward(self, x):
                return t.interpolate(
                    x, mode=self.mode, size=self.size
                )

        modes = ["linear", "bicubic"]
        x = [
            torch.randn(1, 2, 6, requires_grad=True),
            torch.randn(1, 2, 4, 6, requires_grad=True),
            torch.randn(1, 2, 4, 4, 6, requires_grad=True),
        ]
        for mode in modes:
            for xi in x:
                mode_i = mode
                if mode == "bicubic" and xi.dim() != 4:
                    continue
                elif mode == "linear":
                    if xi.dim() == 4:
                        mode_i = "bilinear"
                    elif xi.dim() == 5:
                        mode_i = "trilinear"
                for i in range(xi.dim() - 2):
                    size = list(xi.shape[2:])
                    size[i] = 1
                    self.run_test(MyModel(mode_i, size), xi)
//  [python]
def test_interpolate_half_pixel(self):
        # testing whether it uses "half_pixel" or "pytorch_half_pixel"
        # see https://github.com/onnx/onnx/blob/main/docs/Operators.md#Resize

        class MyModel(torch.nn.Module):
            def __init__(self, mode, size):
                super().__init__()
                self.mode = mode
                self.size = size

            def forward(self, x):
                return torch.nn.functional.interpolate(
                    x, mode=self.mode, size=self.size
                )

        modes = ["linear", "bicubic"]
        x = [
            torch.randn(1, 2, 6, requires_grad=True),
            torch.randn(1, 2, 4, 6, requires_grad=True),
            torch.randn(1, 2, 4, 4, 6, requires_grad=True),
        ]
        for mode in modes:
            for xi in x:
                mode_i = mode
                if mode == "bicubic" and xi.dim() != 4:
                    continue
                elif mode == "linear":
                    if xi.dim() == 4:
                        mode_i = "bilinear"
                    elif xi.dim() == 5:
                        mode_i = "trilinear"
                for i in range(xi.dim() - 2):
                    size = list(xi.shape[2:])
                    size[i] = 1
                    self.run_test(MyModel(mode_i, size), xi)
//  [python]
def post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple] = None):
        """
        Converts the output of [`BeitForSemanticSegmentation`] into semantic segmentation maps. Only supports PyTorch.

        Args:
            outputs ([`BeitForSemanticSegmentation`]):
                Raw outputs of the model.
            target_sizes (`List[Tuple]` of length `batch_size`, *optional*):
                List of tuples corresponding to the requested final size (height, width) of each prediction. If left to
                None, predictions will not be resized.
        Returns:
            semantic_segmentation: `List[torch.Tensor]` of length `batch_size`, where each item is a semantic
            segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is
            specified). Each entry of each `torch.Tensor` correspond to a semantic class id.
        """
        logits = outputs.logits

        # Resize logits and compute semantic segmentation maps
        if target_sizes is not None:
            if len(logits) != len(target_sizes):
                raise ValueError(
                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"
                )

            if is_torch_tensor(target_sizes):
                target_sizes = target_sizes.numpy()

            semantic_segmentation = []

            for idx in range(len(logits)):
                resized_logits = t.nn.functional.interpolate(
                    logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False
                )
                semantic_map = resized_logits[0].argmax(dim=0)
                semantic_segmentation.append(semantic_map)
        else:
            semantic_segmentation = logits.argmax(dim=1)
            semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]

        return semantic_segmentation
//  [python]
def post_process_semantic_segmentation(self, outputs, target_sizes: List[Tuple] = None):
        """
        Converts the output of [`BeitForSemanticSegmentation`] into semantic segmentation maps. Only supports PyTorch.

        Args:
            outputs ([`BeitForSemanticSegmentation`]):
                Raw outputs of the model.
            target_sizes (`List[Tuple]` of length `batch_size`, *optional*):
                List of tuples corresponding to the requested final size (height, width) of each prediction. If left to
                None, predictions will not be resized.
        Returns:
            semantic_segmentation: `List[torch.Tensor]` of length `batch_size`, where each item is a semantic
            segmentation map of shape (height, width) corresponding to the target_sizes entry (if `target_sizes` is
            specified). Each entry of each `torch.Tensor` correspond to a semantic class id.
        """
        logits = outputs.logits

        # Resize logits and compute semantic segmentation maps
        if target_sizes is not None:
            if len(logits) != len(target_sizes):
                raise ValueError(
                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"
                )

            if is_torch_tensor(target_sizes):
                target_sizes = target_sizes.numpy()

            semantic_segmentation = []

            for idx in range(len(logits)):
                resized_logits = torch.nn.functional.interpolate(
                    logits[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False
                )
                semantic_map = resized_logits[0].argmax(dim=0)
                semantic_segmentation.append(semantic_map)
        else:
            semantic_segmentation = logits.argmax(dim=1)
            semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]

        return semantic_segmentation
//  [python]
def optimization_loss(feature_map, x, C, opts, mask_classes=[]):
	b, c, h, w = feature_map.shape
	hw = h*w
	feature_map_reshaped = feature_map.reshape(b,c,hw)
	x_reshaped = t.nn.functional.interpolate(x.float(), size=(h,w)).reshape(b,opts.label_nc,hw).permute(0,2,1)
	xC = torch.bmm(x_reshaped, C).permute(0,2,1)
	x_reshaped_normalized = F.normalize(x_reshaped, p=1, dim=1)
	xF = torch.bmm(x_reshaped_normalized.permute(0,2,1), feature_map_reshaped.permute(0,2,1))

	mask = torch.ones_like(feature_map_reshaped)
	for cid in mask_classes:
		mask.permute(0,2,1)[x_reshaped[:,:,cid]==1] = 0.
		
	loss = 1.-F.cosine_similarity(mask*feature_map_reshaped, mask*xC, dim=1).mean()

	return loss
//  [python]
def optimization_loss(feature_map, x, C, opts, mask_classes=[]):
	b, c, h, w = feature_map.shape
	hw = h*w
	feature_map_reshaped = feature_map.reshape(b,c,hw)
	x_reshaped = F.interpolate(x.float(), size=(h,w)).reshape(b,opts.label_nc,hw).permute(0,2,1)
	xC = torch.bmm(x_reshaped, C).permute(0,2,1)
	x_reshaped_normalized = F.normalize(x_reshaped, p=1, dim=1)
	xF = torch.bmm(x_reshaped_normalized.permute(0,2,1), feature_map_reshaped.permute(0,2,1))

	mask = torch.ones_like(feature_map_reshaped)
	for cid in mask_classes:
		mask.permute(0,2,1)[x_reshaped[:,:,cid]==1] = 0.
		
	loss = 1.-F.cosine_similarity(mask*feature_map_reshaped, mask*xC, dim=1).mean()

	return loss
//  [python]
def forward(self, f0, upp):
        """ sine_tensor, uv = forward(f0)
        input F0: tensor(batchsize=1, length, dim=1)
                  f0 for unvoiced steps should be 0
        output sine_tensor: tensor(batchsize=1, length, dim)
        output uv: tensor(batchsize=1, length, 1)
        """
        f0 = f0.unsqueeze(-1)
        fn = torch.multiply(f0, torch.arange(1, self.dim + 1, device=f0.device).reshape((1, 1, -1)))
        sine_waves = self._f02sine(fn, upp) * self.sine_amp
        uv = (f0 > self.voiced_threshold).float()
        uv = F.upsample(uv.transpose(2, 1), scale_factor=upp, mode='nearest').transpose(2, 1)
        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
        noise = noise_amp * torch.randn_like(sine_waves)
        sine_waves = sine_waves * uv + noise
        return sine_waves
//  [python]
def forward(self, f0, upp):
        """ sine_tensor, uv = forward(f0)
        input F0: tensor(batchsize=1, length, dim=1)
                  f0 for unvoiced steps should be 0
        output sine_tensor: tensor(batchsize=1, length, dim)
        output uv: tensor(batchsize=1, length, 1)
        """
        f0 = f0.unsqueeze(-1)
        fn = torch.multiply(f0, torch.arange(1, self.dim + 1, device=f0.device).reshape((1, 1, -1)))
        sine_waves = self._f02sine(fn, upp) * self.sine_amp
        uv = (f0 > self.voiced_threshold).float()
        uv = F.interpolate(uv.transpose(2, 1), scale_factor=upp, mode='nearest').transpose(2, 1)
        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
        noise = noise_amp * torch.randn_like(sine_waves)
        sine_waves = sine_waves * uv + noise
        return sine_waves
//  [python]
def forward(self, x, skip):
        x = self.conv(x)

        if self.up is not None:
            up = self.up(skip)
            up = F.upsample(up, size=x.shape[-2:])

            x = x + up

        return self.relu(x)
//  [python]
def forward(self, x, skip):
        x = self.conv(x)

        if self.up is not None:
            up = self.up(skip)
            up = F.interpolate(up, x.shape[-2:])

            x = x + up

        return self.relu(x)
//  [python]
def score_distillation_sampling(self,
                                    pred_rgb: torch.Tensor,
                                    crop_size: int,
                                    augments: str,
                                    prompt: Union[List, str],
                                    negative_prompt: Union[List, str] = None,
                                    guidance_scale: float = 100,
                                    as_latent: bool = False,
                                    grad_scale: float = 1,
                                    t_range: Union[List[float], Tuple[float]] = (0.02, 0.98)):
        num_train_timesteps = self.scheduler.config.num_train_timesteps
        min_step = int(num_train_timesteps * t_range[0])
        max_step = int(num_train_timesteps * t_range[1])
        alphas = self.scheduler.alphas_cumprod.to(self.device)  # for convenience

        # sketch augmentation
        pred_rgb_a = self.S_aug(pred_rgb, crop_size, augments)

        # interp to crop_size x crop_size to be fed into vae.
        if as_latent:
            latents = F.interpolate(pred_rgb_a, (64, 64), mode='bilinear', align_corners=False) * 2 - 1
        else:
            # encode image into latents with vae, requires grad!
            latents = self.encode_(pred_rgb_a)

        #  Encode input prompt
        num_images_per_prompt = 1  # the number of images to generate per prompt
        do_classifier_free_guidance = guidance_scale > 1.0
        text_embeddings = self._encode_prompt(
            prompt, self.device, num_images_per_prompt, do_classifier_free_guidance,
            negative_prompt=negative_prompt,
        )

        # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
        t = torch.randint(min_step, max_step + 1, [1], dtype=torch.long, device=self.device)

        # predict the noise residual with unet, stop gradient
        with torch.no_grad():
            # add noise
            noise = torch.randn_like(latents)
            latents_noisy = self.scheduler.add_noise(latents, noise, t)
            # pred noise
            latent_model_input = torch.cat([latents_noisy] * 2) if do_classifier_free_guidance else latents_noisy
            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

        # perform guidance (high scale from paper!)
        if do_classifier_free_guidance:
            noise_pred_uncond, noise_pred_pos = noise_pred.chunk(2)
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_pos - noise_pred_uncond)

        # w(t), sigma_t^2
        w = (1 - alphas[t])
        grad = grad_scale * w * (noise_pred - noise)
        grad = torch.nan_to_num(grad)

        # since we omitted an item in grad, we need to use the custom function to specify the gradient
        loss = SpecifyGradient.apply(latents, grad)

        return loss, grad.mean()
//  [python]
def scale_img(img, ratio=1.0, same_shape=False, gs=32):  # img(16,3,256,416)
    # scales img(bs,3,y,x) by ratio constrained to gs-multiple
    if ratio == 1.0:
        return img
    else:
        h, w = img.shape[2:]
        s = (int(h * ratio), int(w * ratio))  # new size
        img = t.nn.functional.interpolate(img, size=s, mode='bilinear', align_corners=False)  # resize
        if not same_shape:  # pad/crop img
            h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]
        return t.nn.functional.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)
//  [python]
def scale_img(img, ratio=1.0, same_shape=False, gs=32):  # img(16,3,256,416)
    # scales img(bs,3,y,x) by ratio constrained to gs-multiple
    if ratio == 1.0:
        return img
    else:
        h, w = img.shape[2:]
        s = (int(h * ratio), int(w * ratio))  # new size
        img = F.interpolate(img, size=s, mode='bilinear', align_corners=False)  # resize
        if not same_shape:  # pad/crop img
            h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]
        return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447)
//  [python]
def test_interpolate():
    """test relax translator for interpolate"""

    class Interpolate(Module):
        def forward(self, data):
            return t.interpolate(data, size=(5, 5))

    input_info = [([1, 3, 10, 10], "float32")]
    _verify_model(Interpolate(), input_info)
//  [python]
def test_interpolate():
    """test relax translator for interpolate"""

    class Interpolate(Module):
        def forward(self, data):
            return torch.nn.functional.interpolate(data, (5, 5))

    input_info = [([1, 3, 10, 10], "float32")]
    _verify_model(Interpolate(), input_info)
//  [python]
def _f02sine(self, f0_values, upp):
        """ f0_values: (batchsize, length, dim)
            where dim indicates fundamental tone and overtones
        """
        rad_values = (f0_values / self.sampling_rate).fmod(1.)  # %1意味着n_har的乘积无法后处理优化
        rand_ini = torch.rand(1, self.dim, device=f0_values.device)
        rand_ini[:, 0] = 0
        rad_values[:, 0, :] += rand_ini
        is_half = rad_values.dtype is not torch.float32
        tmp_over_one = torch.cumsum(rad_values.double(), 1)  # % 1  #####%1意味着后面的cumsum无法再优化
        if is_half:
            tmp_over_one = tmp_over_one.half()
        else:
            tmp_over_one = tmp_over_one.float()
        tmp_over_one *= upp
        tmp_over_one = t.nn.functional.interpolate(
            tmp_over_one.transpose(2, 1), scale_factor=upp,
            mode='linear', align_corners=True
        ).transpose(2, 1)
        rad_values = t.nn.functional.interpolate(rad_values.transpose(2, 1), scale_factor=upp, mode='nearest').transpose(2, 1)
        tmp_over_one = tmp_over_one.fmod(1.)
        diff = F.conv2d(
            tmp_over_one.unsqueeze(1), torch.FloatTensor([[[[-1.], [1.]]]]).to(tmp_over_one.device),
            stride=(1, 1), padding=0, dilation=(1, 1)
        ).squeeze(1)  # Equivalent to torch.diff, but able to export ONNX
        cumsum_shift = (diff < 0).double()
        cumsum_shift = torch.cat((
            torch.zeros((f0_values.size()[0], 1, self.dim), dtype=torch.double).to(f0_values.device),
            cumsum_shift
        ), dim=1)
        sines = torch.sin(torch.cumsum(rad_values.double() + cumsum_shift, dim=1) * 2 * np.pi)
        if is_half:
            sines = sines.half()
        else:
            sines = sines.float()
        return sines
//  [python]
def _f02sine(self, f0_values, upp):
        """ f0_values: (batchsize, length, dim)
            where dim indicates fundamental tone and overtones
        """
        rad_values = (f0_values / self.sampling_rate).fmod(1.)  # %1意味着n_har的乘积无法后处理优化
        rand_ini = torch.rand(1, self.dim, device=f0_values.device)
        rand_ini[:, 0] = 0
        rad_values[:, 0, :] += rand_ini
        is_half = rad_values.dtype is not torch.float32
        tmp_over_one = torch.cumsum(rad_values.double(), 1)  # % 1  #####%1意味着后面的cumsum无法再优化
        if is_half:
            tmp_over_one = tmp_over_one.half()
        else:
            tmp_over_one = tmp_over_one.float()
        tmp_over_one *= upp
        tmp_over_one = F.interpolate(
            tmp_over_one.transpose(2, 1), scale_factor=upp,
            mode='linear', align_corners=True
        ).transpose(2, 1)
        rad_values = F.interpolate(rad_values.transpose(2, 1), scale_factor=upp, mode='nearest').transpose(2, 1)
        tmp_over_one = tmp_over_one.fmod(1.)
        diff = F.conv2d(
            tmp_over_one.unsqueeze(1), torch.FloatTensor([[[[-1.], [1.]]]]).to(tmp_over_one.device),
            stride=(1, 1), padding=0, dilation=(1, 1)
        ).squeeze(1)  # Equivalent to torch.diff, but able to export ONNX
        cumsum_shift = (diff < 0).double()
        cumsum_shift = torch.cat((
            torch.zeros((f0_values.size()[0], 1, self.dim), dtype=torch.double).to(f0_values.device),
            cumsum_shift
        ), dim=1)
        sines = torch.sin(torch.cumsum(rad_values.double() + cumsum_shift, dim=1) * 2 * np.pi)
        if is_half:
            sines = sines.half()
        else:
            sines = sines.float()
        return sines
//  [python]
def approximate_decoded_latents(latents, scale=None):
    """
    Approximate the decoded latents without using the VAE.
    """
    import torch
    import torch.nn.functional as F
    # origingally adapted from code by @erucipe and @keturn here:
    # https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/7

    # these updated numbers for v1.5 are from @torridgristle
    v1_5_latent_rgb_factors = torch.tensor([
        #    R        G        B
        [ 0.3444,  0.1385,  0.0670], # L1
        [ 0.1247,  0.4027,  0.1494], # L2
        [-0.3192,  0.2513,  0.2103], # L3
        [-0.1307, -0.1874, -0.7445]  # L4
    ], dtype=latents.dtype, device=latents.device)

    latent_image = latents[0].permute(1, 2, 0) @ v1_5_latent_rgb_factors
    if scale is not None:
        latent_image = F.interpolate(
            latent_image.permute(2, 0, 1).unsqueeze(0), scale_factor=scale, mode="nearest"
        ).squeeze(0).permute(1, 2, 0)
    latent_image = ((latent_image + 1) / 2).clamp(0, 1).cpu()
    return latent_image.numpy()
//  [python]
def approximate_decoded_latents(latents, scale=None):
    """
    Approximate the decoded latents without using the VAE.
    """
    import torch
    # origingally adapted from code by @erucipe and @keturn here:
    # https://discuss.huggingface.co/t/decoding-latents-to-rgb-without-upscaling/23204/7

    # these updated numbers for v1.5 are from @torridgristle
    v1_5_latent_rgb_factors = torch.tensor([
        #    R        G        B
        [ 0.3444,  0.1385,  0.0670], # L1
        [ 0.1247,  0.4027,  0.1494], # L2
        [-0.3192,  0.2513,  0.2103], # L3
        [-0.1307, -0.1874, -0.7445]  # L4
    ], dtype=latents.dtype, device=latents.device)

    latent_image = latents[0].permute(1, 2, 0) @ v1_5_latent_rgb_factors
    if scale is not None:
        latent_image = torch.nn.functional.interpolate(
            latent_image.permute(2, 0, 1).unsqueeze(0), scale_factor=scale, mode="nearest"
        ).squeeze(0).permute(1, 2, 0)
    latent_image = ((latent_image + 1) / 2).clamp(0, 1).cpu()
    return latent_image.numpy()
//  [python]
def forward(self, x: Dict[str, torch.Tensor]) -> List[torch.Tensor]:
        y = list(x.values())

        output1 = self.output1(y[0])
        output2 = self.output2(y[1])
        output3 = self.output3(y[2])

        up3 = F.upsample(
            output3, size=[output2.size(2), output2.size(3)], mode='nearest')
        output2 = output2 + up3
        output2 = self.merge2(output2)

        up2 = F.upsample(
            output2, size=[output1.size(2), output1.size(3)], mode='nearest')
        output1 = output1 + up2
        output1 = self.merge1(output1)

        return [output1, output2, output3]
//  [python]
def forward(self, x: Dict[str, torch.Tensor]) -> List[torch.Tensor]:
        y = list(x.values())

        output1 = self.output1(y[0])
        output2 = self.output2(y[1])
        output3 = self.output3(y[2])

        up3 = F.interpolate(
            output3, size=[output2.size(2), output2.size(3)], mode='nearest')
        output2 = output2 + up3
        output2 = self.merge2(output2)

        up2 = F.interpolate(
            output2, size=[output1.size(2), output1.size(3)], mode='nearest')
        output1 = output1 + up2
        output1 = self.merge1(output1)

        return [output1, output2, output3]