jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, features):
        features = features[::-1]
        prev_features = self.fpn_laterals[0](features[0])
        if self.with_ppm:
            prev_features = self.ppm(prev_features)
        outputs = [self.fpn_outputs[0](prev_features)]
        for feature, lat_conv, output_conv in zip(features[1:],
                                                  self.fpn_laterals[1:],
                                                  self.fpn_outputs[1:]):
            lat_features = lat_conv(feature)
            top_down_features = F.upsample(
                prev_features, scale_factor=2.0, mode='nearest')
            prev_features = lat_features + top_down_features
            outputs.insert(0, output_conv(prev_features))
        size = outputs[0].shape[2:]
        features = [outputs[0]] + [
            F.upsample(x, size, mode='bilinear', align_corners=False)
            for x in outputs[1:]
        ]
        features = self.fusion(torch.cat(features, dim=1))
        return features
//  [python]
def forward(self, features):
        features = features[::-1]
        prev_features = self.fpn_laterals[0](features[0])
        if self.with_ppm:
            prev_features = self.ppm(prev_features)
        outputs = [self.fpn_outputs[0](prev_features)]
        for feature, lat_conv, output_conv in zip(features[1:],
                                                  self.fpn_laterals[1:],
                                                  self.fpn_outputs[1:]):
            lat_features = lat_conv(feature)
            top_down_features = F.interpolate(
                prev_features, scale_factor=2.0, mode='nearest')
            prev_features = lat_features + top_down_features
            outputs.insert(0, output_conv(prev_features))
        size = outputs[0].shape[2:]
        features = [outputs[0]] + [
            F.interpolate(x, size, mode='bilinear', align_corners=False)
            for x in outputs[1:]
        ]
        features = self.fusion(torch.cat(features, dim=1))
        return features
//  [python]
def resize(input,
           size=None,
           scale_factor=None,
           mode='nearest',
           align_corners=None,
           warning=True):
    if warning:
        if size is not None and align_corners:
            input_h, input_w = tuple(int(x) for x in input.shape[2:])
            output_h, output_w = tuple(int(x) for x in size)
            if output_h > input_h or output_w > output_h:
                if ((output_h > 1 and output_w > 1 and input_h > 1
                     and input_w > 1) and (output_h - 1) % (input_h - 1)
                        and (output_w - 1) % (input_w - 1)):
                    warnings.warn(
                        f'When align_corners={align_corners}, '
                        'the output would more aligned if '
                        f'input size {(input_h, input_w)} is `x+1` and '
                        f'out size {(output_h, output_w)} is `nx+1`')
    return t.interpolate(input, size, scale_factor, mode, align_corners)
//  [python]
def resize(input,
           size=None,
           scale_factor=None,
           mode='nearest',
           align_corners=None,
           warning=True):
    if warning:
        if size is not None and align_corners:
            input_h, input_w = tuple(int(x) for x in input.shape[2:])
            output_h, output_w = tuple(int(x) for x in size)
            if output_h > input_h or output_w > output_h:
                if ((output_h > 1 and output_w > 1 and input_h > 1
                     and input_w > 1) and (output_h - 1) % (input_h - 1)
                        and (output_w - 1) % (input_w - 1)):
                    warnings.warn(
                        f'When align_corners={align_corners}, '
                        'the output would more aligned if '
                        f'input size {(input_h, input_w)} is `x+1` and '
                        f'out size {(output_h, output_w)} is `nx+1`')
    return F.interpolate(input, size, scale_factor, mode, align_corners)
//  [python]
def postprocess(self, input: Dict[str, Any]) -> Dict[str, Any]:
        outputs = input['outputs']
        batched_inputs = input['batched_inputs']
        images = input['images']
        if self.training:
            raise NotImplementedError
        else:
            mask_cls_results = outputs['pred_logits']  # (B, Q, C+1)
            mask_pred_results = outputs['pred_masks']  # (B, Q, H, W)
            # upsample masks
            mask_pred_results = t.interpolate(
                mask_pred_results,
                size=(images.tensor.shape[-2], images.tensor.shape[-1]),
                mode='bilinear',
                align_corners=False,
            )

            del outputs

            processed_results = []
            for mask_cls_result, mask_pred_result, input_per_image, image_size in zip(
                    mask_cls_results, mask_pred_results, batched_inputs,
                    images.image_sizes):
                height = input_per_image.get('height', image_size[0])
                width = input_per_image.get('width', image_size[1])
                processed_results.append({})  # for each image

                if self.sem_seg_postprocess_before_inference:
                    if not self.single_human:
                        mask_pred_result = self.sem_seg_postprocess(
                            mask_pred_result, image_size, height, width)
                    else:
                        mask_pred_result = self.single_human_sem_seg_postprocess(
                            mask_pred_result, image_size,
                            input_per_image['crop_box'], height, width)
                    mask_cls_result = mask_cls_result.to(mask_pred_result)

                # semantic segmentation inference
                if self.semantic_on:
                    r = self.semantic_inference(mask_cls_result,
                                                mask_pred_result)
                    if not self.sem_seg_postprocess_before_inference:
                        if not self.single_human:
                            r = self.sem_seg_postprocess(
                                r, image_size, height, width)
                        else:
                            r = self.single_human_sem_seg_postprocess(
                                r, image_size, input_per_image['crop_box'],
                                height, width)
                        processed_results[-1]['sem_seg'] = r

                # parsing inference
                if self.parsing_on:
                    parsing_r = self.instance_parsing_inference(
                        mask_cls_result, mask_pred_result)
                    processed_results[-1]['parsing'] = parsing_r

        return dict(eval_result=processed_results)
//  [python]
def postprocess(self, input: Dict[str, Any]) -> Dict[str, Any]:
        outputs = input['outputs']
        batched_inputs = input['batched_inputs']
        images = input['images']
        if self.training:
            raise NotImplementedError
        else:
            mask_cls_results = outputs['pred_logits']  # (B, Q, C+1)
            mask_pred_results = outputs['pred_masks']  # (B, Q, H, W)
            # upsample masks
            mask_pred_results = F.interpolate(
                mask_pred_results,
                size=(images.tensor.shape[-2], images.tensor.shape[-1]),
                mode='bilinear',
                align_corners=False,
            )

            del outputs

            processed_results = []
            for mask_cls_result, mask_pred_result, input_per_image, image_size in zip(
                    mask_cls_results, mask_pred_results, batched_inputs,
                    images.image_sizes):
                height = input_per_image.get('height', image_size[0])
                width = input_per_image.get('width', image_size[1])
                processed_results.append({})  # for each image

                if self.sem_seg_postprocess_before_inference:
                    if not self.single_human:
                        mask_pred_result = self.sem_seg_postprocess(
                            mask_pred_result, image_size, height, width)
                    else:
                        mask_pred_result = self.single_human_sem_seg_postprocess(
                            mask_pred_result, image_size,
                            input_per_image['crop_box'], height, width)
                    mask_cls_result = mask_cls_result.to(mask_pred_result)

                # semantic segmentation inference
                if self.semantic_on:
                    r = self.semantic_inference(mask_cls_result,
                                                mask_pred_result)
                    if not self.sem_seg_postprocess_before_inference:
                        if not self.single_human:
                            r = self.sem_seg_postprocess(
                                r, image_size, height, width)
                        else:
                            r = self.single_human_sem_seg_postprocess(
                                r, image_size, input_per_image['crop_box'],
                                height, width)
                        processed_results[-1]['sem_seg'] = r

                # parsing inference
                if self.parsing_on:
                    parsing_r = self.instance_parsing_inference(
                        mask_cls_result, mask_pred_result)
                    processed_results[-1]['parsing'] = parsing_r

        return dict(eval_result=processed_results)
//  [python]
def forward(self, x):
        c2, c3, c4, c5 = x

        in5 = self.in5_conv(c5)
        in4 = self.in4_conv(c4)
        in3 = self.in3_conv(c3)
        in2 = self.in2_conv(c2)

        out4 = in4 + t.interpolate(
            in5, scale_factor=2, mode="nearest", )#align_mode=1)  # 1/16
        out3 = in3 + t.interpolate(
            out4, scale_factor=2, mode="nearest", )#align_mode=1)  # 1/8
        out2 = in2 + t.interpolate(
            out3, scale_factor=2, mode="nearest", )#align_mode=1)  # 1/4

        p5 = self.p5_conv(in5)
        p4 = self.p4_conv(out4)
        p3 = self.p3_conv(out3)
        p2 = self.p2_conv(out2)
        p5 = t.interpolate(p5, scale_factor=8, mode="nearest", )#align_mode=1)
        p4 = t.interpolate(p4, scale_factor=4, mode="nearest", )#align_mode=1)
        p3 = t.interpolate(p3, scale_factor=2, mode="nearest", )#align_mode=1)

        fuse = torch.cat([p5, p4, p3, p2], dim=1)

        if self.use_asf is True:
            fuse = self.asf(fuse, [p5, p4, p3, p2])

        return fuse
//  [python]
def forward(self, x):
        c2, c3, c4, c5 = x

        in5 = self.in5_conv(c5)
        in4 = self.in4_conv(c4)
        in3 = self.in3_conv(c3)
        in2 = self.in2_conv(c2)

        out4 = in4 + F.interpolate(
            in5, scale_factor=2, mode="nearest", )#align_mode=1)  # 1/16
        out3 = in3 + F.interpolate(
            out4, scale_factor=2, mode="nearest", )#align_mode=1)  # 1/8
        out2 = in2 + F.interpolate(
            out3, scale_factor=2, mode="nearest", )#align_mode=1)  # 1/4

        p5 = self.p5_conv(in5)
        p4 = self.p4_conv(out4)
        p3 = self.p3_conv(out3)
        p2 = self.p2_conv(out2)
        p5 = F.interpolate(p5, scale_factor=8, mode="nearest", )#align_mode=1)
        p4 = F.interpolate(p4, scale_factor=4, mode="nearest", )#align_mode=1)
        p3 = F.interpolate(p3, scale_factor=2, mode="nearest", )#align_mode=1)

        fuse = torch.cat([p5, p4, p3, p2], dim=1)

        if self.use_asf is True:
            fuse = self.asf(fuse, [p5, p4, p3, p2])

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

        for attn, net in zip(self.attns, self.nets):
            res_x = xs[-1]
            xs = xs[:-1]
            x = torch.cat([x, res_x], dim=1)
            x = net(x)
            if attn:
                x = attn(x)
            
        if self.upsample:
            x = F.interpolate(x, scale_factor=2.0, mode='nearest')
            x = self.upsample(x)

        return x
//  [python]
def downscale(x):
    return t.interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=False)
//  [python]
def downscale(x):
    return interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=False)
//  [python]
def forward(self, x):
        size = x.shape[-2:]
        if self.gather is not None:
            x_ge = self.gather(x)
        else:
            if self.extent == 0:
                # global extent
                x_ge = x.mean(dim=(2, 3), keepdims=True)
                if self.add_maxpool:
                    # experimental codepath, may remove or change
                    x_ge = 0.5 * x_ge + 0.5 * x.amax((2, 3), keepdim=True)
            else:
                x_ge = F.avg_pool2d(
                    x, kernel_size=self.gk, stride=self.gs, padding=self.gk // 2, count_include_pad=False)
                if self.add_maxpool:
                    # experimental codepath, may remove or change
                    x_ge = 0.5 * x_ge + 0.5 * F.max_pool2d(x, kernel_size=self.gk, stride=self.gs, padding=self.gk // 2)
        x_ge = self.mlp(x_ge)
        if x_ge.shape[-1] != 1 or x_ge.shape[-2] != 1:
            x_ge = F.upsample_bilinear(x_ge, size=size)
        return x * self.gate(x_ge)
//  [python]
def forward(self, x):
        size = x.shape[-2:]
        if self.gather is not None:
            x_ge = self.gather(x)
        else:
            if self.extent == 0:
                # global extent
                x_ge = x.mean(dim=(2, 3), keepdims=True)
                if self.add_maxpool:
                    # experimental codepath, may remove or change
                    x_ge = 0.5 * x_ge + 0.5 * x.amax((2, 3), keepdim=True)
            else:
                x_ge = F.avg_pool2d(
                    x, kernel_size=self.gk, stride=self.gs, padding=self.gk // 2, count_include_pad=False)
                if self.add_maxpool:
                    # experimental codepath, may remove or change
                    x_ge = 0.5 * x_ge + 0.5 * F.max_pool2d(x, kernel_size=self.gk, stride=self.gs, padding=self.gk // 2)
        x_ge = self.mlp(x_ge)
        if x_ge.shape[-1] != 1 or x_ge.shape[-2] != 1:
            x_ge = F.interpolate(x_ge, size=size)
        return x * self.gate(x_ge)
//  [python]
def forward(self,
                img: Tensor,
                img_raw: Optional[Tensor] = None,
                cond: Optional[Tensor] = None):
        """Forward function.

        Args:
            img (torch.Tensor): Input high resoluation image tensor.
            img_raw (torch.Tensor): Input raw (low resolution) image tensor.
                Defaults to None.
            cond (torch.Tensor): The conditional input (camera-to-world matrix
                and intrinsics matrix). Defaults to None.

        Returns:
            torch.Tensor: Predict score for the input image.
        """
        if self.use_dual_disc:
            assert img_raw is not None, (
                '\'img_raw\' must be passed when \'use_dual_disc\' is True.')

        # This setting was used to finetune on converted weights
        if self.input_bgr2rgb:
            img = img[:, [2, 1, 0], ...]
            if img_raw is not None:
                img_raw = img_raw[:, [2, 1, 0], ...]

        if img_raw is not None:
            # the official implementation only use 'antialiased' upsampline,
            # therefore we only support 'antialiased' for torch >= 1.11.0
            interpolation_kwargs = dict(
                size=(img.shape[-1], img.shape[-1]),
                mode='bilinear',
                align_corners=False)
            if digit_version(TORCH_VERSION) >= digit_version('1.11.0'):
                interpolation_kwargs['antialias'] = True
            img_raw_sr = t.interpolate(img_raw, **interpolation_kwargs)
            img = torch.cat([img, img_raw_sr], dim=1)

        # convs has own fp-16 controller, do not wrap here
        x = self.convs(img)
        x = self.mbstd_layer(x)

        fp16_enabled = (
            self.final_conv.fp16_enabled or not self.convert_input_fp32)
        with autocast(enabled=fp16_enabled):
            if not fp16_enabled:
                x = x.to(torch.float32)
            x = self.final_conv(x)
            x = x.view(x.shape[0], -1)
            x = self.final_linear(x)

            # conditioning
            if cond is not None:
                assert self.mapping is not None, (
                    '\'mapping\' network must not be None when conditional '
                    'input is passed.')

                # if self.disc_c_noise is not None and self.disc_c_noise > 0:
                if self.disc_c_noise is not None:
                    cond = cond + torch.randn_like(
                        cond) * cond.std() * self.disc_c_noise
                cmap = self.mapping(None, cond)
                x = (x * cmap).sum(
                    dim=1, keepdim=True) * (1 / np.sqrt(cmap.shape[1]))
        return x
//  [python]
def forward(self,
                img: Tensor,
                img_raw: Optional[Tensor] = None,
                cond: Optional[Tensor] = None):
        """Forward function.

        Args:
            img (torch.Tensor): Input high resoluation image tensor.
            img_raw (torch.Tensor): Input raw (low resolution) image tensor.
                Defaults to None.
            cond (torch.Tensor): The conditional input (camera-to-world matrix
                and intrinsics matrix). Defaults to None.

        Returns:
            torch.Tensor: Predict score for the input image.
        """
        if self.use_dual_disc:
            assert img_raw is not None, (
                '\'img_raw\' must be passed when \'use_dual_disc\' is True.')

        # This setting was used to finetune on converted weights
        if self.input_bgr2rgb:
            img = img[:, [2, 1, 0], ...]
            if img_raw is not None:
                img_raw = img_raw[:, [2, 1, 0], ...]

        if img_raw is not None:
            # the official implementation only use 'antialiased' upsampline,
            # therefore we only support 'antialiased' for torch >= 1.11.0
            interpolation_kwargs = dict(
                size=(img.shape[-1], img.shape[-1]),
                mode='bilinear',
                align_corners=False)
            if digit_version(TORCH_VERSION) >= digit_version('1.11.0'):
                interpolation_kwargs['antialias'] = True
            img_raw_sr = F.interpolate(img_raw, **interpolation_kwargs)
            img = torch.cat([img, img_raw_sr], dim=1)

        # convs has own fp-16 controller, do not wrap here
        x = self.convs(img)
        x = self.mbstd_layer(x)

        fp16_enabled = (
            self.final_conv.fp16_enabled or not self.convert_input_fp32)
        with autocast(enabled=fp16_enabled):
            if not fp16_enabled:
                x = x.to(torch.float32)
            x = self.final_conv(x)
            x = x.view(x.shape[0], -1)
            x = self.final_linear(x)

            # conditioning
            if cond is not None:
                assert self.mapping is not None, (
                    '\'mapping\' network must not be None when conditional '
                    'input is passed.')

                # if self.disc_c_noise is not None and self.disc_c_noise > 0:
                if self.disc_c_noise is not None:
                    cond = cond + torch.randn_like(
                        cond) * cond.std() * self.disc_c_noise
                cmap = self.mapping(None, cond)
                x = (x * cmap).sum(
                    dim=1, keepdim=True) * (1 / np.sqrt(cmap.shape[1]))
        return x
//  [python]
def get_hook(G_synth, layer, size, strategy):
    size = np.flip(size)
    if strategy == 'stretch':

        def hook(module, input, output):
            return t.interpolate(output, tuple(size), mode='bicubic', align_corners=False)
        return hook
    elif strategy == 'pad-zero':
        original_size = getattr(G_synth, G_synth.layer_names[max(layer - 1, 0)]).out_size
        pad_h, pad_w = (size - original_size).astype(int) // 2
        padding = pad_w, pad_w, pad_h, pad_h

        def hook(module, input, output):
            return pad(output, padding, mode='constant', value=0)
        return hook
    else:
        raise Exception(f'Resize strategy not found: {strategy}')
//  [python]
def get_hook(G_synth, layer, size, strategy):
    size = np.flip(size)
    if strategy == 'stretch':

        def hook(module, input, output):
            return interpolate(output, tuple(size), mode='bicubic', align_corners=False)
        return hook
    elif strategy == 'pad-zero':
        original_size = getattr(G_synth, G_synth.layer_names[max(layer - 1, 0)]).out_size
        pad_h, pad_w = (size - original_size).astype(int) // 2
        padding = pad_w, pad_w, pad_h, pad_h

        def hook(module, input, output):
            return pad(output, padding, mode='constant', value=0)
        return hook
    else:
        raise Exception(f'Resize strategy not found: {strategy}')
//  [python]
def make_noise_pyramid(self, noise, layer_limit=8):
        noises = {}
        for l, layer in enumerate(self.layer_names[1:]):
            if l > layer_limit:
                continue
            _, block, conv = layer.split('.')
            synth_layer = getattr(self.G_synth.bs[int(block)], conv)
            h, w = synth_layer.noise_const.shape[-2], synth_layer.noise_const.shape[-1]
            try:
                noises[f'noise{l}'] = t.interpolate(noise, (h, w), mode='bicubic', align_corners=False).cpu()
            except RuntimeError:
                noises[f'noise{l}'] = t.interpolate(noise.cpu(), (h, w), mode='bicubic', align_corners=False)
            noises[f'noise{l}'] /= noises[f'noise{l}'].std((1, 2, 3), keepdim=True)
        return noises
//  [python]
def make_noise_pyramid(self, noise, layer_limit=8):
        noises = {}
        for l, layer in enumerate(self.layer_names[1:]):
            if l > layer_limit:
                continue
            _, block, conv = layer.split('.')
            synth_layer = getattr(self.G_synth.bs[int(block)], conv)
            h, w = synth_layer.noise_const.shape[-2], synth_layer.noise_const.shape[-1]
            try:
                noises[f'noise{l}'] = torch.nn.functional.interpolate(noise, (h, w), mode='bicubic', align_corners=False).cpu()
            except RuntimeError:
                noises[f'noise{l}'] = torch.nn.functional.interpolate(noise.cpu(), (h, w), mode='bicubic', align_corners=False)
            noises[f'noise{l}'] /= noises[f'noise{l}'].std((1, 2, 3), keepdim=True)
        return noises
//  [python]
def compute_segments(
    mask_probs,
    pred_scores,
    pred_labels,
    mask_threshold: float = 0.5,
    overlap_mask_area_threshold: float = 0.8,
    label_ids_to_fuse: Optional[Set[int]] = None,
    target_size: Tuple[int, int] = None,
):
    height = mask_probs.shape[1] if target_size is None else target_size[0]
    width = mask_probs.shape[2] if target_size is None else target_size[1]

    segmentation = torch.zeros((height, width), dtype=torch.int32, device=mask_probs.device)
    segments: List[Dict] = []

    if target_size is not None:
        mask_probs = t.nn.functional.interpolate(
            mask_probs.unsqueeze(0), size=target_size, mode="bilinear", align_corners=False
        )[0]

    current_segment_id = 0

    # Weigh each mask by its prediction score
    mask_probs *= pred_scores.view(-1, 1, 1)
    mask_labels = mask_probs.argmax(0)  # [height, width]

    # Keep track of instances of each class
    stuff_memory_list: Dict[str, int] = {}
    for k in range(pred_labels.shape[0]):
        pred_class = pred_labels[k].item()
        should_fuse = pred_class in label_ids_to_fuse

        # Check if mask exists and large enough to be a segment
        mask_exists, mask_k = check_segment_validity(
            mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold
        )

        if mask_exists:
            if pred_class in stuff_memory_list:
                current_segment_id = stuff_memory_list[pred_class]
            else:
                current_segment_id += 1

            # Add current object segment to final segmentation map
            segmentation[mask_k] = current_segment_id
            segment_score = round(pred_scores[k].item(), 6)
            segments.append(
                {
                    "id": current_segment_id,
                    "label_id": pred_class,
                    "was_fused": should_fuse,
                    "score": segment_score,
                }
            )
            if should_fuse:
                stuff_memory_list[pred_class] = current_segment_id

    return segmentation, segments
//  [python]
def compute_segments(
    mask_probs,
    pred_scores,
    pred_labels,
    mask_threshold: float = 0.5,
    overlap_mask_area_threshold: float = 0.8,
    label_ids_to_fuse: Optional[Set[int]] = None,
    target_size: Tuple[int, int] = None,
):
    height = mask_probs.shape[1] if target_size is None else target_size[0]
    width = mask_probs.shape[2] if target_size is None else target_size[1]

    segmentation = torch.zeros((height, width), dtype=torch.int32, device=mask_probs.device)
    segments: List[Dict] = []

    if target_size is not None:
        mask_probs = nn.functional.interpolate(
            mask_probs.unsqueeze(0), size=target_size, mode="bilinear", align_corners=False
        )[0]

    current_segment_id = 0

    # Weigh each mask by its prediction score
    mask_probs *= pred_scores.view(-1, 1, 1)
    mask_labels = mask_probs.argmax(0)  # [height, width]

    # Keep track of instances of each class
    stuff_memory_list: Dict[str, int] = {}
    for k in range(pred_labels.shape[0]):
        pred_class = pred_labels[k].item()
        should_fuse = pred_class in label_ids_to_fuse

        # Check if mask exists and large enough to be a segment
        mask_exists, mask_k = check_segment_validity(
            mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold
        )

        if mask_exists:
            if pred_class in stuff_memory_list:
                current_segment_id = stuff_memory_list[pred_class]
            else:
                current_segment_id += 1

            # Add current object segment to final segmentation map
            segmentation[mask_k] = current_segment_id
            segment_score = round(pred_scores[k].item(), 6)
            segments.append(
                {
                    "id": current_segment_id,
                    "label_id": pred_class,
                    "was_fused": should_fuse,
                    "score": segment_score,
                }
            )
            if should_fuse:
                stuff_memory_list[pred_class] = current_segment_id

    return segmentation, segments
//  [python]
def forward(self, latents: Tensor, translation: Optional[Tensor]=None, translation_layer: int=7, zoom: Optional[Tensor]=None, zoom_layer: int=7, zoom_center: Optional[int]=None, rotation: Optional[Tensor]=None, rotation_layer: int=7, rotation_center: Optional[int]=None, **noise) ->Tensor:
        if translation is not None:
            self.apply_translation(translation_layer, translation)
        if zoom is not None:
            self.apply_zoom(zoom_layer, zoom, zoom_center)
        if rotation is not None:
            self.apply_rotation(rotation_layer, rotation, rotation_center)
        if noise is not None:
            noises, l = list(noise.values()), 0
            for block in self.G_synth.bs:
                if l >= len(noises):
                    continue
                for c in (([block.conv0] if hasattr(block, 'conv0') else []) + [block.conv1]):
                    noise_l = noises[l]
                    if (noise_l.shape[-2], noise_l.shape[-1]) != (c.noise_const.shape[-2], c.noise_const.shape[-1]):
                        warnings.warn(f'Supplied noise for SynthesisLayer {l} has shape {noise_l.shape} while the expected shape is {c.noise_const.shape}. Resizing the supplied noise to match...')
                        h, w = c.noise_const.shape[-2], c.noise_const.shape[-1]
                        noise_l = t.interpolate(noise_l, (h, w), mode='bicubic', align_corners=False)
                    setattr(c, 'noise_const', noise_l)
                    l += 1
        return self.G_synth.forward(latents, noise_mode='const')
//  [python]
def forward(self, latents: Tensor, translation: Optional[Tensor]=None, translation_layer: int=7, zoom: Optional[Tensor]=None, zoom_layer: int=7, zoom_center: Optional[int]=None, rotation: Optional[Tensor]=None, rotation_layer: int=7, rotation_center: Optional[int]=None, **noise) ->Tensor:
        if translation is not None:
            self.apply_translation(translation_layer, translation)
        if zoom is not None:
            self.apply_zoom(zoom_layer, zoom, zoom_center)
        if rotation is not None:
            self.apply_rotation(rotation_layer, rotation, rotation_center)
        if noise is not None:
            noises, l = list(noise.values()), 0
            for block in self.G_synth.bs:
                if l >= len(noises):
                    continue
                for c in (([block.conv0] if hasattr(block, 'conv0') else []) + [block.conv1]):
                    noise_l = noises[l]
                    if (noise_l.shape[-2], noise_l.shape[-1]) != (c.noise_const.shape[-2], c.noise_const.shape[-1]):
                        warnings.warn(f'Supplied noise for SynthesisLayer {l} has shape {noise_l.shape} while the expected shape is {c.noise_const.shape}. Resizing the supplied noise to match...')
                        h, w = c.noise_const.shape[-2], c.noise_const.shape[-1]
                        noise_l = torch.nn.functional.interpolate(noise_l, (h, w), mode='bicubic', align_corners=False)
                    setattr(c, 'noise_const', noise_l)
                    l += 1
        return self.G_synth.forward(latents, noise_mode='const')
//  [python]
def post_process_segmentation(
        self, outputs: "MaskFormerForInstanceSegmentationOutput", target_size: Tuple[int, int] = None
    ) -> "torch.Tensor":
        """
        Converts the output of [`MaskFormerForInstanceSegmentationOutput`] into image segmentation predictions. Only
        supports PyTorch.

        Args:
            outputs ([`MaskFormerForInstanceSegmentationOutput`]):
                The outputs from [`MaskFormerForInstanceSegmentation`].

            target_size (`Tuple[int, int]`, *optional*):
                If set, the `masks_queries_logits` will be resized to `target_size`.

        Returns:
            `torch.Tensor`:
                A tensor of shape (`batch_size, num_class_labels, height, width`).
        """
        logger.warning(
            "`post_process_segmentation` is deprecated and will be removed in v5 of Transformers, please use"
            " `post_process_instance_segmentation`",
            FutureWarning,
        )

        # class_queries_logits has shape [BATCH, QUERIES, CLASSES + 1]
        class_queries_logits = outputs.class_queries_logits
        # masks_queries_logits has shape [BATCH, QUERIES, HEIGHT, WIDTH]
        masks_queries_logits = outputs.masks_queries_logits
        if target_size is not None:
            masks_queries_logits = t.nn.functional.interpolate(
                masks_queries_logits,
                size=target_size,
                mode="bilinear",
                align_corners=False,
            )
        # remove the null class `[..., :-1]`
        masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
        # mask probs has shape [BATCH, QUERIES, HEIGHT, WIDTH]
        masks_probs = masks_queries_logits.sigmoid()
        # now we want to sum over the queries,
        # $ out_{c,h,w} =  \sum_q p_{q,c} * m_{q,h,w} $
        # where $ softmax(p) \in R^{q, c} $ is the mask classes
        # and $ sigmoid(m) \in R^{q, h, w}$ is the mask probabilities
        # b(atch)q(uery)c(lasses), b(atch)q(uery)h(eight)w(idth)
        segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)

        return segmentation
//  [python]
def post_process_segmentation(
        self, outputs: "MaskFormerForInstanceSegmentationOutput", target_size: Tuple[int, int] = None
    ) -> "torch.Tensor":
        """
        Converts the output of [`MaskFormerForInstanceSegmentationOutput`] into image segmentation predictions. Only
        supports PyTorch.

        Args:
            outputs ([`MaskFormerForInstanceSegmentationOutput`]):
                The outputs from [`MaskFormerForInstanceSegmentation`].

            target_size (`Tuple[int, int]`, *optional*):
                If set, the `masks_queries_logits` will be resized to `target_size`.

        Returns:
            `torch.Tensor`:
                A tensor of shape (`batch_size, num_class_labels, height, width`).
        """
        logger.warning(
            "`post_process_segmentation` is deprecated and will be removed in v5 of Transformers, please use"
            " `post_process_instance_segmentation`",
            FutureWarning,
        )

        # class_queries_logits has shape [BATCH, QUERIES, CLASSES + 1]
        class_queries_logits = outputs.class_queries_logits
        # masks_queries_logits has shape [BATCH, QUERIES, HEIGHT, WIDTH]
        masks_queries_logits = outputs.masks_queries_logits
        if target_size is not None:
            masks_queries_logits = torch.nn.functional.interpolate(
                masks_queries_logits,
                size=target_size,
                mode="bilinear",
                align_corners=False,
            )
        # remove the null class `[..., :-1]`
        masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
        # mask probs has shape [BATCH, QUERIES, HEIGHT, WIDTH]
        masks_probs = masks_queries_logits.sigmoid()
        # now we want to sum over the queries,
        # $ out_{c,h,w} =  \sum_q p_{q,c} * m_{q,h,w} $
        # where $ softmax(p) \in R^{q, c} $ is the mask classes
        # and $ sigmoid(m) \in R^{q, h, w}$ is the mask probabilities
        # b(atch)q(uery)c(lasses), b(atch)q(uery)h(eight)w(idth)
        segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)

        return segmentation
//  [python]
def initialize_optical_flow(cache, init, consistency_trust, width, height, device):
    flow_model = get_flow_model()
    frames = VideoFrames(init, height=min(height, 512), width=min(width, 512), device=device)
    N = len(frames)
    if len(cache.flow) == N and tuple(cache.flow[0].shape[1:3]) == (height, width):
        None
        return
    else:
        cache.flow.length = cache.consistency.length = 0
    for f_n in trange(N, desc='Calculating optical flow...'):
        prev = frames[(f_n - 1) % N].add(1).div(2)
        curr = frames[f_n].add(1).div(2)
        forward = flow_model(curr, prev)
        backward = flow_model(prev, curr)
        maxflow = max(forward.shape[0], forward.shape[1])
        forward, backward = forward.clamp(-maxflow, maxflow), backward.clamp(-maxflow, maxflow)
        if consistency_trust > 0:
            consistency = get_consistency_map(forward, backward)
            consistency = t(consistency.unsqueeze(1), (height, width), mode='bilinear')
            cache.consistency.insert(consistency)
        forward *= np.mean((width / forward.shape[1], height / forward.shape[2]))
        forward = t(forward.permute(0, 3, 1, 2), (height, width), mode='bilinear').permute(0, 2, 3, 1)
        cache.flow.insert(forward)
//  [python]
def initialize_optical_flow(cache, init, consistency_trust, width, height, device):
    flow_model = get_flow_model()
    frames = VideoFrames(init, height=min(height, 512), width=min(width, 512), device=device)
    N = len(frames)
    if len(cache.flow) == N and tuple(cache.flow[0].shape[1:3]) == (height, width):
        None
        return
    else:
        cache.flow.length = cache.consistency.length = 0
    for f_n in trange(N, desc='Calculating optical flow...'):
        prev = frames[(f_n - 1) % N].add(1).div(2)
        curr = frames[f_n].add(1).div(2)
        forward = flow_model(curr, prev)
        backward = flow_model(prev, curr)
        maxflow = max(forward.shape[0], forward.shape[1])
        forward, backward = forward.clamp(-maxflow, maxflow), backward.clamp(-maxflow, maxflow)
        if consistency_trust > 0:
            consistency = get_consistency_map(forward, backward)
            consistency = interpolate(consistency.unsqueeze(1), (height, width), mode='bilinear')
            cache.consistency.insert(consistency)
        forward *= np.mean((width / forward.shape[1], height / forward.shape[2]))
        forward = interpolate(forward.permute(0, 3, 1, 2), (height, width), mode='bilinear').permute(0, 2, 3, 1)
        cache.flow.insert(forward)
//  [python]
def decode_tensor(self):
        width, height = self.image_shape
        pallet = self.sort_pallet()
        values = self.value.clamp(0, 1) * (self.pallet_size - 1)
        value_floors, value_ceils, value_rounds, value_fracs = break_tensor(values)
        value_fracs = value_fracs.unsqueeze(-1).unsqueeze(-1)
        pallet_weights = self.tensor.movedim(0, 2)
        pallets = F.one_hot(pallet_weights.argmax(dim=2), num_classes=self.n_pallets)
        pallet_weights = pallet_weights.softmax(dim=2).unsqueeze(-1)
        pallets = pallets.unsqueeze(-1)
        colors_disc = pallet[value_rounds]
        colors_disc = (colors_disc * pallets).sum(dim=2)
        colors_disc = t(colors_disc.movedim(2, 0).unsqueeze(0), (height, width), mode='nearest')
        colors_cont = pallet[value_floors] * (1 - value_fracs) + pallet[value_ceils] * value_fracs
        colors_cont = (colors_cont * pallet_weights).sum(dim=2)
        colors_cont = t(colors_cont.movedim(2, 0).unsqueeze(0), (height, width), mode='nearest')
        return replace_grad(colors_disc, colors_cont * 0.5 + colors_disc * 0.5)
//  [python]
def decode_tensor(self):
        width, height = self.image_shape
        pallet = self.sort_pallet()
        values = self.value.clamp(0, 1) * (self.pallet_size - 1)
        value_floors, value_ceils, value_rounds, value_fracs = break_tensor(values)
        value_fracs = value_fracs.unsqueeze(-1).unsqueeze(-1)
        pallet_weights = self.tensor.movedim(0, 2)
        pallets = F.one_hot(pallet_weights.argmax(dim=2), num_classes=self.n_pallets)
        pallet_weights = pallet_weights.softmax(dim=2).unsqueeze(-1)
        pallets = pallets.unsqueeze(-1)
        colors_disc = pallet[value_rounds]
        colors_disc = (colors_disc * pallets).sum(dim=2)
        colors_disc = F.interpolate(colors_disc.movedim(2, 0).unsqueeze(0), (height, width), mode='nearest')
        colors_cont = pallet[value_floors] * (1 - value_fracs) + pallet[value_ceils] * value_fracs
        colors_cont = (colors_cont * pallet_weights).sum(dim=2)
        colors_cont = F.interpolate(colors_cont.movedim(2, 0).unsqueeze(0), (height, width), mode='nearest')
        return replace_grad(colors_disc, colors_cont * 0.5 + colors_disc * 0.5)
//  [python]
def interpolate_pos_encoding_2d(target_spatial_size, pos_embed):
        N = pos_embed.shape[1]
        if N == target_spatial_size:
            return pos_embed
        dim = pos_embed.shape[-1]
        pos_embed = t.nn.functional.interpolate(pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2), scale_factor=math.sqrt(target_spatial_size / N), mode='bicubic')
        pos_embed = pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return pos_embed
//  [python]
def interpolate_pos_encoding_2d(target_spatial_size, pos_embed):
        N = pos_embed.shape[1]
        if N == target_spatial_size:
            return pos_embed
        dim = pos_embed.shape[-1]
        pos_embed = nn.functional.interpolate(pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2), scale_factor=math.sqrt(target_spatial_size / N), mode='bicubic')
        pos_embed = pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return pos_embed
//  [python]
def __init__(self, latent_selection, loop_len, type='spline', smooth=10):
        super().__init__()
        latent_selection = latent_selection.cpu()
        if loop_len == 1 or type == 'constant':
            latents = latent_selection[[0]]
            loop_len = 1
        elif type == 'spline':
            latent_selection = np.concatenate([latent_selection, latent_selection[[0]]])
            x = np.linspace(0, 1, loop_len)
            latents = np.zeros((loop_len, *latent_selection.shape[1:]))
            for lay in range(latent_selection.shape[1]):
                for lat in range(latent_selection.shape[2]):
                    tck = splrep(np.linspace(0, 1, latent_selection.shape[0]), latent_selection[:, lay, lat])
                    latents[:, lay, lat] = splev(x, tck)
            latents = torch.from_numpy(latents)
        elif type == 'slerp':
            latent_selection = np.concatenate([latent_selection, latent_selection[[0]]])
            latents = []
            for n in range(len(latent_selection)):
                for val in np.linspace(0.0, 1.0, int(ceil(loop_len / len(latent_selection)))):
                    latents.append(torch.from_numpy(slerp(latent_selection[n % len(latent_selection)][0], latent_selection[(n + 1) % len(latent_selection)][0], val)))
            latents = torch.stack(latents).unsqueeze(1).tile(1, 18, 1)
            latents = t(latents.permute(2, 1, 0), size=(loop_len,), mode='linear', align_corners=False).permute(2, 1, 0)
            latents = gaussian_filter(latents, 1)
        elif type == 'gaussian':
            latents = torch.cat([lat.tile(round(loop_len / len(latent_selection)), 1, 1) for lat in latent_selection])
            latents = t(latents.permute(2, 1, 0), size=(loop_len,), mode='linear', align_corners=False).permute(2, 1, 0)
            latents = gaussian_filter(latents, smooth)
        self.register_buffer('latents', latents)
        self.index, self.length = 0, loop_len
//  [python]
def __init__(self, latent_selection, loop_len, type='spline', smooth=10):
        super().__init__()
        latent_selection = latent_selection.cpu()
        if loop_len == 1 or type == 'constant':
            latents = latent_selection[[0]]
            loop_len = 1
        elif type == 'spline':
            latent_selection = np.concatenate([latent_selection, latent_selection[[0]]])
            x = np.linspace(0, 1, loop_len)
            latents = np.zeros((loop_len, *latent_selection.shape[1:]))
            for lay in range(latent_selection.shape[1]):
                for lat in range(latent_selection.shape[2]):
                    tck = splrep(np.linspace(0, 1, latent_selection.shape[0]), latent_selection[:, lay, lat])
                    latents[:, lay, lat] = splev(x, tck)
            latents = torch.from_numpy(latents)
        elif type == 'slerp':
            latent_selection = np.concatenate([latent_selection, latent_selection[[0]]])
            latents = []
            for n in range(len(latent_selection)):
                for val in np.linspace(0.0, 1.0, int(ceil(loop_len / len(latent_selection)))):
                    latents.append(torch.from_numpy(slerp(latent_selection[n % len(latent_selection)][0], latent_selection[(n + 1) % len(latent_selection)][0], val)))
            latents = torch.stack(latents).unsqueeze(1).tile(1, 18, 1)
            latents = F.interpolate(latents.permute(2, 1, 0), loop_len, mode='linear', align_corners=False).permute(2, 1, 0)
            latents = gaussian_filter(latents, 1)
        elif type == 'gaussian':
            latents = torch.cat([lat.tile(round(loop_len / len(latent_selection)), 1, 1) for lat in latent_selection])
            latents = F.interpolate(latents.permute(2, 1, 0), loop_len, mode='linear', align_corners=False).permute(2, 1, 0)
            latents = gaussian_filter(latents, smooth)
        self.register_buffer('latents', latents)
        self.index, self.length = 0, loop_len
//  [python]
def render_channel(self, pallet_i):
        width, height = self.image_shape
        pallet = self.sort_pallet()
        pallet[:, :pallet_i, :] = 0.5
        pallet[:, pallet_i + 1:, :] = 0.5
        values = self.value.clamp(0, 1) * (self.pallet_size - 1)
        value_floors, value_ceils, value_rounds, value_fracs = break_tensor(values)
        value_fracs = value_fracs.unsqueeze(-1).unsqueeze(-1)
        pallet_weights = self.tensor.movedim(0, 2)
        pallets = F.one_hot(pallet_weights.argmax(dim=2), num_classes=self.n_pallets)
        pallet_weights = pallet_weights.softmax(dim=2).unsqueeze(-1)
        colors_cont = pallet[value_floors] * (1 - value_fracs) + pallet[value_ceils] * value_fracs
        colors_cont = (colors_cont * pallet_weights).sum(dim=2)
        colors_cont = t.imresize(colors_cont.movedim(2, 0).unsqueeze(0), (height, width), mode='nearest')
        tensor = named_rearrange(colors_cont, self.output_axes, ('y', 'x', 's'))
        array = np.array(tensor.mul(255).clamp(0, 255).cpu().detach().numpy().astype(np.uint8))[:, :, :]
        return Image.fromarray(array)
//  [python]
def render_channel(self, pallet_i):
        width, height = self.image_shape
        pallet = self.sort_pallet()
        pallet[:, :pallet_i, :] = 0.5
        pallet[:, pallet_i + 1:, :] = 0.5
        values = self.value.clamp(0, 1) * (self.pallet_size - 1)
        value_floors, value_ceils, value_rounds, value_fracs = break_tensor(values)
        value_fracs = value_fracs.unsqueeze(-1).unsqueeze(-1)
        pallet_weights = self.tensor.movedim(0, 2)
        pallets = F.one_hot(pallet_weights.argmax(dim=2), num_classes=self.n_pallets)
        pallet_weights = pallet_weights.softmax(dim=2).unsqueeze(-1)
        colors_cont = pallet[value_floors] * (1 - value_fracs) + pallet[value_ceils] * value_fracs
        colors_cont = (colors_cont * pallet_weights).sum(dim=2)
        colors_cont = F.interpolate(colors_cont.movedim(2, 0).unsqueeze(0), (height, width), mode='nearest')
        tensor = named_rearrange(colors_cont, self.output_axes, ('y', 'x', 's'))
        array = np.array(tensor.mul(255).clamp(0, 255).cpu().detach().numpy().astype(np.uint8))[:, :, :]
        return Image.fromarray(array)
//  [python]
def forward(self, x):
        x = self.block(x)
        if self.upsample:
            x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
        return x
//  [python]
def forward(self, *in_x):
        features = self.encoder(*in_x)
        features = features[1:]
        x = features[-1]
        # len(features)-2, decrease to 0, corresponding to the settings above
        for i in range(len(features) - 2, -1, -1):
            # print(i, features[i].shape)
            conv = getattr(self, 'shrink%d' % i)
            # x = F.interpolate(x, scale_factor=2)
            x = F.upsample(x, scale_factor=2)
            # print('conv',conv)
            # print('x',x.shape)
            x = conv(x)
            # print('features[i].shape',features[i].shape)
            if features[i].shape[-1] != x.shape[-1]:
                x2 = F.upsample(features[i], scale_factor=4)
            else:
                x2 = features[i]
            # print(x.shape, ':',x2.shape)
            x = torch.cat([x, x2], 1)
            conv2 = getattr(self, 'shrink2%d' % i)
            # print('conv2',conv2)
            x = conv2(x)
            bn = getattr(self, 'batchnorm%d' % i)
            x = bn(x)
            x = F.relu(x)
            x = self.dropout(x)
        x = F.upsample(x, scale_factor=2)
        logits = self.outc(x)
        return logits
//  [python]
def forward(self, *in_x):
        features = self.encoder(*in_x)
        features = features[1:]
        x = features[-1]
        # len(features)-2, decrease to 0, corresponding to the settings above
        for i in range(len(features) - 2, -1, -1):
            # print(i, features[i].shape)
            conv = getattr(self, 'shrink%d' % i)
            # x = F.interpolate(x, scale_factor=2)
            x = F.interpolate(x, scale_factor=2)
            # print('conv',conv)
            # print('x',x.shape)
            x = conv(x)
            # print('features[i].shape',features[i].shape)
            if features[i].shape[-1] != x.shape[-1]:
                x2 = F.interpolate(features[i], scale_factor=4)
            else:
                x2 = features[i]
            # print(x.shape, ':',x2.shape)
            x = torch.cat([x, x2], 1)
            conv2 = getattr(self, 'shrink2%d' % i)
            # print('conv2',conv2)
            x = conv2(x)
            bn = getattr(self, 'batchnorm%d' % i)
            x = bn(x)
            x = F.relu(x)
            x = self.dropout(x)
        x = F.interpolate(x, scale_factor=2)
        logits = self.outc(x)
        return logits