jiahuimbzuai/CoQuIR huggingface dataset:
//  [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',
        )
        out3 = in3 + F.interpolate(
            out4,
            scale_factor=2,
            mode='nearest',
        )
        out2 = in2 + F.interpolate(
            out3,
            scale_factor=2,
            mode='nearest',
        )

        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',
        )
        p4 = F.interpolate(
            p4,
            scale_factor=4,
            mode='nearest',
        )
        p3 = F.interpolate(
            p3,
            scale_factor=2,
            mode='nearest',
        )

        fuse = torch.cat([p5, p4, p3, p2], dim=1)
        return fuse
//  [python]
def postprocess(self, feats, curr_sample_rate):
        def resample(x, factor):
            return t.interpolate(x.view(1, 1, -1), scale_factor=factor).squeeze()

        if feats.dim() == 2:
            feats = feats.mean(-1)

        if curr_sample_rate != self.sample_rate:
            factor = self.sample_rate / curr_sample_rate
            feats = resample(feats, factor)

        assert feats.dim() == 1, feats.dim()
        return feats
//  [python]
def postprocess(self, feats, curr_sample_rate):
        def resample(x, factor):
            return F.interpolate(x.view(1, 1, -1), scale_factor=factor).squeeze()

        if feats.dim() == 2:
            feats = feats.mean(-1)

        if curr_sample_rate != self.sample_rate:
            factor = self.sample_rate / curr_sample_rate
            feats = resample(feats, factor)

        assert feats.dim() == 1, feats.dim()
        return feats
//  [python]
def forward(self, hidden_states: List[torch.Tensor], patch_height, patch_width) -> torch.Tensor:
        hidden_states = hidden_states[self.head_in_index]

        predicted_depth = self.conv1(hidden_states)
        predicted_depth = torch.nn.functional.interpolate(
            predicted_depth,
            (int(patch_height * self.patch_size), int(patch_width * self.patch_size)),
            mode="bilinear",
            align_corners=True,
        )
        predicted_depth = self.conv2(predicted_depth)
        predicted_depth = self.activation1(predicted_depth)
        predicted_depth = self.conv3(predicted_depth)
        predicted_depth = self.activation2(predicted_depth)
        predicted_depth = predicted_depth.squeeze(dim=1)  # shape (batch_size, height, width)

        return predicted_depth
//  [python]
def forward(self, hidden_states: List[torch.Tensor], patch_height, patch_width) -> torch.Tensor:
        hidden_states = hidden_states[self.head_in_index]

        predicted_depth = self.conv1(hidden_states)
        predicted_depth = nn.functional.interpolate(
            predicted_depth,
            (int(patch_height * self.patch_size), int(patch_width * self.patch_size)),
            mode="bilinear",
            align_corners=True,
        )
        predicted_depth = self.conv2(predicted_depth)
        predicted_depth = self.activation1(predicted_depth)
        predicted_depth = self.conv3(predicted_depth)
        predicted_depth = self.activation2(predicted_depth)
        predicted_depth = predicted_depth.squeeze(dim=1)  # shape (batch_size, height, width)

        return predicted_depth
//  [python]
def preprocess(self, x):
        x = (x + 1) / 2
        x = (x - self.RGB_mean) / self.RGB_std
        if x.shape[3] < 224:
            x = torch.nn.functional.upsample_bilinear(x, size=(224, 224))
        return x
//  [python]
def preprocess(self, x):
        x = (x + 1) / 2
        x = (x - self.RGB_mean) / self.RGB_std
        if x.shape[3] < 224:
            x = torch.nn.functional.interpolate(x, size=(224, 224), mode='bilinear', align_corners=False)
        return x
//  [python]
def get_codebook_entry(self, indices_list, segm_map, shape):
        # flatten segm_map (b, h, w)
        segm_map = F.upsample(
            segm_map, size=(shape[1], shape[2]), mode='nearest')
        segm_map_flatten = segm_map.view(-1)

        z_q = torch.zeros((shape[0] * shape[1] * shape[2]),
                          self.e_dim).to(segm_map.device)
        for codebook_idx in range(18):
            if torch.sum(segm_map_flatten == codebook_idx) > 0:
                min_encoding_indices_selected = indices_list[
                    codebook_idx].view(-1)[segm_map_flatten == codebook_idx]
                z_q_selected = self.embedding_list[codebook_idx](
                    min_encoding_indices_selected)
                z_q[segm_map_flatten == codebook_idx] = z_q_selected

        z_q = z_q.view(shape)
        # reshape back to match original input shape
        z_q = z_q.permute(0, 3, 1, 2).contiguous()

        return z_q
//  [python]
def get_codebook_entry(self, indices_list, segm_map, shape):
        # flatten segm_map (b, h, w)
        segm_map = F.interpolate(
            segm_map, size=(shape[1], shape[2]), mode='nearest')
        segm_map_flatten = segm_map.view(-1)

        z_q = torch.zeros((shape[0] * shape[1] * shape[2]),
                          self.e_dim).to(segm_map.device)
        for codebook_idx in range(18):
            if torch.sum(segm_map_flatten == codebook_idx) > 0:
                min_encoding_indices_selected = indices_list[
                    codebook_idx].view(-1)[segm_map_flatten == codebook_idx]
                z_q_selected = self.embedding_list[codebook_idx](
                    min_encoding_indices_selected)
                z_q[segm_map_flatten == codebook_idx] = z_q_selected

        z_q = z_q.view(shape)
        # reshape back to match original input shape
        z_q = z_q.permute(0, 3, 1, 2).contiguous()

        return z_q
//  [python]
def forward(self, skip_f, up_f):
        x = self.skip_conv(skip_f)
        x = x + t.upsample_bilinear(
            up_f,
            scale_factor=self.scale_factor)
        x = self.out_conv(x)
        return x
//  [python]
def forward(self, skip_f, up_f):
        x = self.skip_conv(skip_f)
        x = x + F.interpolate(
            up_f,
            scale_factor=self.scale_factor,
            mode='bilinear',
            align_corners=False)
        x = self.out_conv(x)
        return x
//  [python]
def downsample(gt_s1, map_sizes):
    gt_ms = []
    gt = gt_s1
    for size in map_sizes:
        gt = t.nn.functional.interpolate(gt, size=size, mode="bilinear")
        gt_ms.append(gt)

    return gt_ms
//  [python]
def downsample(gt_s1, map_sizes):
    gt_ms = []
    gt = gt_s1
    for size in map_sizes:
        gt = F.interpolate(gt, size=size, mode="bilinear")
        gt_ms.append(gt)

    return gt_ms
//  [python]
def prepare_depth_map(self, image, depth_map, batch_size, do_classifier_free_guidance, dtype, device):
        if isinstance(image, PIL.Image.Image):
            image = [image]
        else:
            image = list(image)

        if isinstance(image[0], PIL.Image.Image):
            width, height = image[0].size
        elif isinstance(image[0], np.ndarray):
            width, height = image[0].shape[:-1]
        else:
            height, width = image[0].shape[-2:]

        if depth_map is None:
            pixel_values = self.feature_extractor(images=image, return_tensors="pt").pixel_values
            pixel_values = pixel_values.to(device=device)
            # The DPT-Hybrid model uses batch-norm layers which are not compatible with fp16.
            # So we use `torch.autocast` here for half precision inference.
            if torch.backends.mps.is_available():
                autocast_ctx = contextlib.nullcontext()
                logger.warning(
                    "The DPT-Hybrid model uses batch-norm layers which are not compatible with fp16, but autocast is not yet supported on MPS."
                )
            else:
                autocast_ctx = torch.autocast(device.type, dtype=dtype)

            with autocast_ctx:
                depth_map = self.depth_estimator(pixel_values).predicted_depth
        else:
            depth_map = depth_map.to(device=device, dtype=dtype)

        depth_map = t.nn.functional.interpolate(
            depth_map.unsqueeze(1),
            size=(height // self.vae_scale_factor, width // self.vae_scale_factor),
            mode="bicubic",
            align_corners=False,
        )

        depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
        depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
        depth_map = 2.0 * (depth_map - depth_min) / (depth_max - depth_min) - 1.0
        depth_map = depth_map.to(dtype)

        # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method
        if depth_map.shape[0] < batch_size:
            repeat_by = batch_size // depth_map.shape[0]
            depth_map = depth_map.repeat(repeat_by, 1, 1, 1)

        depth_map = torch.cat([depth_map] * 2) if do_classifier_free_guidance else depth_map
        return depth_map
//  [python]
def prepare_depth_map(self, image, depth_map, batch_size, do_classifier_free_guidance, dtype, device):
        if isinstance(image, PIL.Image.Image):
            image = [image]
        else:
            image = list(image)

        if isinstance(image[0], PIL.Image.Image):
            width, height = image[0].size
        elif isinstance(image[0], np.ndarray):
            width, height = image[0].shape[:-1]
        else:
            height, width = image[0].shape[-2:]

        if depth_map is None:
            pixel_values = self.feature_extractor(images=image, return_tensors="pt").pixel_values
            pixel_values = pixel_values.to(device=device)
            # The DPT-Hybrid model uses batch-norm layers which are not compatible with fp16.
            # So we use `torch.autocast` here for half precision inference.
            if torch.backends.mps.is_available():
                autocast_ctx = contextlib.nullcontext()
                logger.warning(
                    "The DPT-Hybrid model uses batch-norm layers which are not compatible with fp16, but autocast is not yet supported on MPS."
                )
            else:
                autocast_ctx = torch.autocast(device.type, dtype=dtype)

            with autocast_ctx:
                depth_map = self.depth_estimator(pixel_values).predicted_depth
        else:
            depth_map = depth_map.to(device=device, dtype=dtype)

        depth_map = torch.nn.functional.interpolate(
            depth_map.unsqueeze(1),
            size=(height // self.vae_scale_factor, width // self.vae_scale_factor),
            mode="bicubic",
            align_corners=False,
        )

        depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
        depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
        depth_map = 2.0 * (depth_map - depth_min) / (depth_max - depth_min) - 1.0
        depth_map = depth_map.to(dtype)

        # duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method
        if depth_map.shape[0] < batch_size:
            repeat_by = batch_size // depth_map.shape[0]
            depth_map = depth_map.repeat(repeat_by, 1, 1, 1)

        depth_map = torch.cat([depth_map] * 2) if do_classifier_free_guidance else depth_map
        return depth_map
//  [python]
def downsample_gt(gt_s1, map_sizes, threshold=0, batch_list=False):
    gt_ms = []
    gt = gt_s1

    if batch_list:

        for batch in gt_s1:
            batch_ms = []
            for size in map_sizes:
                gt = t.interpolate(batch.unsqueeze(0), size=size, mode="bilinear")
                batch_ms.append((gt.squeeze(0) > threshold).float())
            gt_ms.append(batch_ms)

        return gt_ms
    else:
        ndims_gt = len(gt.shape)
        for size in map_sizes:
            if ndims_gt == 5:
                batch_size, t, c, _, _ = gt.shape

                # flatten to 4D
                gt = gt.flatten(start_dim=0, end_dim=1)

                gt = t.interpolate(gt, size=size, mode="bilinear")

                # partition back to original dims if originally 5D
                gt = gt.unsqueeze(1).view(batch_size, t, c, size[0], size[1])

            else:
                gt = t.interpolate(gt, size=size, mode="bilinear")

            gt_ms.append(gt)

        return [(gt > threshold).float() for gt in gt_ms]
//  [python]
def downsample_gt(gt_s1, map_sizes, threshold=0, batch_list=False):
    gt_ms = []
    gt = gt_s1

    if batch_list:

        for batch in gt_s1:
            batch_ms = []
            for size in map_sizes:
                gt = F.interpolate(batch.unsqueeze(0), size=size, mode="bilinear")
                batch_ms.append((gt.squeeze(0) > threshold).float())
            gt_ms.append(batch_ms)

        return gt_ms
    else:
        ndims_gt = len(gt.shape)
        for size in map_sizes:
            if ndims_gt == 5:
                batch_size, t, c, _, _ = gt.shape

                # flatten to 4D
                gt = gt.flatten(start_dim=0, end_dim=1)

                gt = F.interpolate(gt, size=size, mode="bilinear")

                # partition back to original dims if originally 5D
                gt = gt.unsqueeze(1).view(batch_size, t, c, size[0], size[1])

            else:
                gt = F.interpolate(gt, size=size, mode="bilinear")

            gt_ms.append(gt)

        return [(gt > threshold).float() for gt in gt_ms]
//  [python]
def forward(self, x):
        input_size = x.size()
        ppm_out = [x]
        for pool_scale in self.ppm:
            ppm_out.append(F.upsample(pool_scale(x), size=(input_size[2], input_size[3]),
                                         mode='bilinear'))

        ppm_out = torch.cat(ppm_out, 1)

        return ppm_out
//  [python]
def forward(self, x):
        input_size = x.size()
        ppm_out = [x]
        for pool_scale in self.ppm:
            ppm_out.append(F.interpolate(pool_scale(x), (input_size[2], input_size[3]),
                                         mode='bilinear', align_corners=False))

        ppm_out = torch.cat(ppm_out, 1)

        return ppm_out
//  [python]
def train_step(self, text_embeddings, pred_rgb, guidance_scale=100, grad_scale=1):

        images = t.interpolate(pred_rgb, (64, 64), mode='bilinear', align_corners=False) * 2 - 1

        t = torch.randint(self.min_step, self.max_step + 1, (images.shape[0],), dtype=torch.long, device=self.device)

        with torch.no_grad():
            noise = torch.randn_like(images)
            images_noisy = self.scheduler.add_noise(images, noise, t)

            model_input = torch.cat([images_noisy] * 2)
            model_input = self.scheduler.scale_model_input(model_input, t)
            tt = torch.cat([t] * 2)
            noise_pred = self.unet(model_input, tt, encoder_hidden_states=text_embeddings).sample
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred_uncond, _ = noise_pred_uncond.split(model_input.shape[1], dim=1)
            noise_pred_text, predicted_variance = noise_pred_text.split(model_input.shape[1], dim=1)
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

            # TODO: how to use the variance here?
            # noise_pred = torch.cat([noise_pred, predicted_variance], dim=1)

        w = (1 - self.alphas[t])
        grad = grad_scale * w[:, None, None, None] * (noise_pred - noise)
        grad = torch.nan_to_num(grad)

        targets = (images - grad).detach()
        loss = 0.5 * F.mse_loss(images.float(), targets, reduction='sum') / images.shape[0]

        return loss
//  [python]
def train_step(self, text_embeddings, pred_rgb, guidance_scale=100, grad_scale=1):

        # [0, 1] to [-1, 1] and make sure shape is [64, 64]
        images = F.interpolate(pred_rgb, (64, 64), mode='bilinear', align_corners=False) * 2 - 1

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

        # predict the noise residual with unet, NO grad!
        with torch.no_grad():
            # add noise
            noise = torch.randn_like(images)
            images_noisy = self.scheduler.add_noise(images, noise, t)

            # pred noise
            model_input = torch.cat([images_noisy] * 2)
            model_input = self.scheduler.scale_model_input(model_input, t)
            tt = torch.cat([t] * 2)
            noise_pred = self.unet(model_input, tt, encoder_hidden_states=text_embeddings).sample
            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
            noise_pred_uncond, _ = noise_pred_uncond.split(model_input.shape[1], dim=1)
            noise_pred_text, predicted_variance = noise_pred_text.split(model_input.shape[1], dim=1)
            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

            # TODO: how to use the variance here?
            # noise_pred = torch.cat([noise_pred, predicted_variance], dim=1)

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

        targets = (images - grad).detach()
        loss = 0.5 * F.mse_loss(images.float(), targets, reduction='sum') / images.shape[0]

        return loss
//  [python]
def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        input_shape = x.shape[-2:]
        backbone_out = self.backbone(x)
        x = self.up1(backbone_out['stage4'], backbone_out['stage3'])
        x = self.up2(x, backbone_out['stage2'])
        x = self.up3(x, backbone_out['stage1'])
        x = self.up4(x, backbone_out['stage0'])
        x = self.conv(x)
        x = F.upsample(x, size=input_shape, mode="bilinear")

        return {"out": x}
//  [python]
def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        input_shape = x.shape[-2:]
        backbone_out = self.backbone(x)
        x = self.up1(backbone_out['stage4'], backbone_out['stage3'])
        x = self.up2(x, backbone_out['stage2'])
        x = self.up3(x, backbone_out['stage1'])
        x = self.up4(x, backbone_out['stage0'])
        x = self.conv(x)
        x = F.interpolate(x, size=input_shape, mode="bilinear", align_corners=False)

        return {"out": x}
//  [python]
def loss_by_feat(self, seg_logits: Tensor,
                     batch_data_samples: SampleList) -> dict:
        """Compute Detail Aggregation Loss."""
        # Note: The paper claims `fusion_kernel` is a trainable 1x1 conv
        # parameters. However, it is a constant in original repo and other
        # codebase because it would not be added into computation graph
        # after threshold operation.
        seg_label = self._stack_batch_gt(batch_data_samples).to(
            self.laplacian_kernel)
        boundary_targets = F.conv2d(
            seg_label, self.laplacian_kernel, padding=1)
        boundary_targets = boundary_targets.clamp(min=0)
        boundary_targets[boundary_targets > self.boundary_threshold] = 1
        boundary_targets[boundary_targets <= self.boundary_threshold] = 0

        boundary_targets_x2 = F.conv2d(
            seg_label, self.laplacian_kernel, stride=2, padding=1)
        boundary_targets_x2 = boundary_targets_x2.clamp(min=0)

        boundary_targets_x4 = F.conv2d(
            seg_label, self.laplacian_kernel, stride=4, padding=1)
        boundary_targets_x4 = boundary_targets_x4.clamp(min=0)

        boundary_targets_x4_up = t.nn.functional.upsample(
            boundary_targets_x4, scale_factor=2, mode='nearest')
        boundary_targets_x2_up = t.nn.functional.upsample(
            boundary_targets_x2, scale_factor=2, mode='nearest')

        boundary_targets_x2_up[
            boundary_targets_x2_up > self.boundary_threshold] = 1
        boundary_targets_x2_up[
            boundary_targets_x2_up <= self.boundary_threshold] = 0

        boundary_targets_x4_up[
            boundary_targets_x4_up > self.boundary_threshold] = 1
        boundary_targets_x4_up[
            boundary_targets_x4_up <= self.boundary_threshold] = 0

        boundary_targets_pyramids = torch.stack(
            (boundary_targets, boundary_targets_x2_up, boundary_targets_x4_up),
            dim=1)

        boundary_targets_pyramids = boundary_targets_pyramids.squeeze(2)
        boudary_targets_pyramid = F.conv2d(boundary_targets_pyramids,
                                           self.fusion_kernel)

        boudary_targets_pyramid[
            boudary_targets_pyramid > self.boundary_threshold] = 1
        boudary_targets_pyramid[
            boudary_targets_pyramid <= self.boundary_threshold] = 0

        seg_labels = boudary_targets_pyramid.long()
        batch_sample_list = []
        for label in seg_labels:
            seg_data_sample = SegDataSample()
            seg_data_sample.gt_sem_seg = PixelData(data=label)
            batch_sample_list.append(seg_data_sample)

        loss = super().loss_by_feat(seg_logits, batch_sample_list)
        return loss
//  [python]
def loss_by_feat(self, seg_logits: Tensor,
                     batch_data_samples: SampleList) -> dict:
        """Compute Detail Aggregation Loss."""
        # Note: The paper claims `fusion_kernel` is a trainable 1x1 conv
        # parameters. However, it is a constant in original repo and other
        # codebase because it would not be added into computation graph
        # after threshold operation.
        seg_label = self._stack_batch_gt(batch_data_samples).to(
            self.laplacian_kernel)
        boundary_targets = F.conv2d(
            seg_label, self.laplacian_kernel, padding=1)
        boundary_targets = boundary_targets.clamp(min=0)
        boundary_targets[boundary_targets > self.boundary_threshold] = 1
        boundary_targets[boundary_targets <= self.boundary_threshold] = 0

        boundary_targets_x2 = F.conv2d(
            seg_label, self.laplacian_kernel, stride=2, padding=1)
        boundary_targets_x2 = boundary_targets_x2.clamp(min=0)

        boundary_targets_x4 = F.conv2d(
            seg_label, self.laplacian_kernel, stride=4, padding=1)
        boundary_targets_x4 = boundary_targets_x4.clamp(min=0)

        boundary_targets_x4_up = F.interpolate(
            boundary_targets_x4, boundary_targets.shape[2:], mode='nearest')
        boundary_targets_x2_up = F.interpolate(
            boundary_targets_x2, boundary_targets.shape[2:], mode='nearest')

        boundary_targets_x2_up[
            boundary_targets_x2_up > self.boundary_threshold] = 1
        boundary_targets_x2_up[
            boundary_targets_x2_up <= self.boundary_threshold] = 0

        boundary_targets_x4_up[
            boundary_targets_x4_up > self.boundary_threshold] = 1
        boundary_targets_x4_up[
            boundary_targets_x4_up <= self.boundary_threshold] = 0

        boundary_targets_pyramids = torch.stack(
            (boundary_targets, boundary_targets_x2_up, boundary_targets_x4_up),
            dim=1)

        boundary_targets_pyramids = boundary_targets_pyramids.squeeze(2)
        boudary_targets_pyramid = F.conv2d(boundary_targets_pyramids,
                                           self.fusion_kernel)

        boudary_targets_pyramid[
            boudary_targets_pyramid > self.boundary_threshold] = 1
        boudary_targets_pyramid[
            boudary_targets_pyramid <= self.boundary_threshold] = 0

        seg_labels = boudary_targets_pyramid.long()
        batch_sample_list = []
        for label in seg_labels:
            seg_data_sample = SegDataSample()
            seg_data_sample.gt_sem_seg = PixelData(data=label)
            batch_sample_list.append(seg_data_sample)

        loss = super().loss_by_feat(seg_logits, batch_sample_list)
        return loss
//  [python]
def forward(self, object_pixel_values):
        """Forward function.

        Args:
            object_pixel_values (torch.Tensor ):
                You can directly input a ``torch.Tensor``.

        Returns:
            torch.Tensor : ``torch.tensor`` will be returned.
        """

        b, num_objects, c, h, w = object_pixel_values.shape

        object_pixel_values = object_pixel_values.view(b * num_objects, c, h,
                                                       w)

        if h != self.image_size or w != self.image_size:
            h, w = self.image_size, self.image_size
            object_pixel_values = F.interpolate(
                object_pixel_values, (h, w), mode='bilinear')

        object_pixel_values = self.vision_processor(object_pixel_values)
        object_embeds = self.vision_model(object_pixel_values)[1]
        object_embeds = self.visual_projection(object_embeds)
        object_embeds = object_embeds.view(b, num_objects, 1, -1)
        return object_embeds
//  [python]
def forward(self, data_dict):
        x = self.stage1(data_dict['img'])
        aux_x = self.dsn(x)
        x = self.stage2(x)
        x = self.ppm(x)
        x = self.cls(x)
        x_dsn = t.interpolate(aux_x, size=(data_dict['img'].size(2), data_dict['img'].size(3)),
                              mode="bilinear", align_corners=False)
        x = t.interpolate(x, size=(data_dict['img'].size(2), data_dict['img'].size(3)),
                          mode="bilinear", align_corners=False)
        out_dict = dict(dsn_out=x_dsn, out=x)
        if self.configer.get('phase') == 'test':
            return out_dict

        loss_dict = dict()
        if 'dsn_ce_loss' in self.valid_loss_dict:
            loss_dict['dsn_ce_loss'] = dict(
                params=[x_dsn, data_dict['labelmap']],
                type=torch.cuda.LongTensor([BASE_LOSS_DICT['ce_loss']]),
                weight=torch.cuda.FloatTensor([self.valid_loss_dict['dsn_ce_loss']])
            )

        if 'ce_loss' in self.valid_loss_dict:
            loss_dict['ce_loss'] = dict(
                params=[x, data_dict['labelmap']],
                type=torch.cuda.LongTensor([BASE_LOSS_DICT['ce_loss']]),
                weight=torch.cuda.FloatTensor([self.valid_loss_dict['ce_loss']])
            )

        if 'ohem_ce_loss' in self.valid_loss_dict:
            loss_dict['ohem_ce_loss'] = dict(
                params=[x, data_dict['labelmap']],
                type=torch.cuda.LongTensor([BASE_LOSS_DICT['ohem_ce_loss']]),
                weight=torch.cuda.FloatTensor([self.valid_loss_dict['ohem_ce_loss']])
            )
        return out_dict, loss_dict
//  [python]
def forward(self, data_dict):
        x = self.stage1(data_dict['img'])
        aux_x = self.dsn(x)
        x = self.stage2(x)
        x = self.ppm(x)
        x = self.cls(x)
        x_dsn = F.interpolate(aux_x, size=(data_dict['img'].size(2), data_dict['img'].size(3)),
                              mode="bilinear", align_corners=False)
        x = F.interpolate(x, size=(data_dict['img'].size(2), data_dict['img'].size(3)),
                          mode="bilinear", align_corners=False)
        out_dict = dict(dsn_out=x_dsn, out=x)
        if self.configer.get('phase') == 'test':
            return out_dict

        loss_dict = dict()
        if 'dsn_ce_loss' in self.valid_loss_dict:
            loss_dict['dsn_ce_loss'] = dict(
                params=[x_dsn, data_dict['labelmap']],
                type=torch.cuda.LongTensor([BASE_LOSS_DICT['ce_loss']]),
                weight=torch.cuda.FloatTensor([self.valid_loss_dict['dsn_ce_loss']])
            )

        if 'ce_loss' in self.valid_loss_dict:
            loss_dict['ce_loss'] = dict(
                params=[x, data_dict['labelmap']],
                type=torch.cuda.LongTensor([BASE_LOSS_DICT['ce_loss']]),
                weight=torch.cuda.FloatTensor([self.valid_loss_dict['ce_loss']])
            )

        if 'ohem_ce_loss' in self.valid_loss_dict:
            loss_dict['ohem_ce_loss'] = dict(
                params=[x, data_dict['labelmap']],
                type=torch.cuda.LongTensor([BASE_LOSS_DICT['ohem_ce_loss']]),
                weight=torch.cuda.FloatTensor([self.valid_loss_dict['ohem_ce_loss']])
            )
        return out_dict, loss_dict
//  [python]
def forward(self, x):
        assert x.shape[1] == self.channels
        x = F.upsample(x, scale_factor=self.factor, mode="nearest")
        if self.use_conv:
            x = self.conv(x)
        return x
//  [python]
def forward(self, x):
        assert x.shape[1] == self.channels
        x = F.interpolate(x, scale_factor=self.factor, mode="nearest")
        if self.use_conv:
            x = self.conv(x)
        return x
//  [python]
def train_step_perpneg(self, text_embeddings, weights, pred_rgb, guidance_scale=100, grad_scale=1):

        B = pred_rgb.shape[0]
        K = (text_embeddings.shape[0] // B) - 1 # maximum number of prompts        

        # [0, 1] to [-1, 1] and make sure shape is [64, 64]
        images = t(pred_rgb, (64, 64), mode='bilinear', align_corners=False) * 2 - 1

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

        # predict the noise residual with unet, NO grad!
        with torch.no_grad():
            # add noise
            noise = torch.randn_like(images)
            images_noisy = self.scheduler.add_noise(images, noise, t)

            # pred noise
            model_input = torch.cat([images_noisy] * (1 + K))
            model_input = self.scheduler.scale_model_input(model_input, t)
            tt = torch.cat([t] * (1 + K))
            unet_output = self.unet(model_input, tt, encoder_hidden_states=text_embeddings).sample
            noise_pred_uncond, noise_pred_text = unet_output[:B], unet_output[B:]
            noise_pred_uncond, _ = noise_pred_uncond.split(model_input.shape[1], dim=1)
            noise_pred_text, predicted_variance = noise_pred_text.split(model_input.shape[1], dim=1)
            # noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
            delta_noise_preds = noise_pred_text - noise_pred_uncond.repeat(K, 1, 1, 1)
            noise_pred = noise_pred_uncond + guidance_scale * weighted_perpendicular_aggregator(delta_noise_preds, weights, B)



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

        targets = (images - grad).detach()
        loss = 0.5 * F.mse_loss(images.float(), targets, reduction='sum') / images.shape[0]

        return loss
//  [python]
def train_step_perpneg(self, text_embeddings, weights, pred_rgb, guidance_scale=100, grad_scale=1):

        B = pred_rgb.shape[0]
        K = (text_embeddings.shape[0] // B) - 1 # maximum number of prompts        

        # [0, 1] to [-1, 1] and make sure shape is [64, 64]
        images = F.interpolate(pred_rgb, (64, 64), mode='bilinear', align_corners=False) * 2 - 1

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

        # predict the noise residual with unet, NO grad!
        with torch.no_grad():
            # add noise
            noise = torch.randn_like(images)
            images_noisy = self.scheduler.add_noise(images, noise, t)

            # pred noise
            model_input = torch.cat([images_noisy] * (1 + K))
            model_input = self.scheduler.scale_model_input(model_input, t)
            tt = torch.cat([t] * (1 + K))
            unet_output = self.unet(model_input, tt, encoder_hidden_states=text_embeddings).sample
            noise_pred_uncond, noise_pred_text = unet_output[:B], unet_output[B:]
            noise_pred_uncond, _ = noise_pred_uncond.split(model_input.shape[1], dim=1)
            noise_pred_text, predicted_variance = noise_pred_text.split(model_input.shape[1], dim=1)
            # noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
            delta_noise_preds = noise_pred_text - noise_pred_uncond.repeat(K, 1, 1, 1)
            noise_pred = noise_pred_uncond + guidance_scale * weighted_perpendicular_aggregator(delta_noise_preds, weights, B)



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

        targets = (images - grad).detach()
        loss = 0.5 * F.mse_loss(images.float(), targets, reduction='sum') / images.shape[0]

        return loss
//  [python]
def _get_rel_pos_bias(self, window_size):
    """
    Modification of timm.models.beit.py: Attention._get_rel_pos_bias to support arbitrary window sizes.
    """
    old_height = 2 * self.window_size[0] - 1
    old_width = 2 * self.window_size[1] - 1

    new_height = 2 * window_size[0] - 1
    new_width = 2 * window_size[1] - 1

    old_relative_position_bias_table = self.relative_position_bias_table

    old_num_relative_distance = self.num_relative_distance
    new_num_relative_distance = new_height * new_width + 3

    old_sub_table = old_relative_position_bias_table[:old_num_relative_distance - 3]

    old_sub_table = old_sub_table.reshape(1, old_width, old_height, -1).permute(0, 3, 1, 2)
    new_sub_table = t.nn.functional.interpolate(old_sub_table, size=(int(new_height), int(new_width)), mode="bilinear")
    new_sub_table = new_sub_table.permute(0, 2, 3, 1).reshape(new_num_relative_distance - 3, -1)

    new_relative_position_bias_table = torch.cat(
        [new_sub_table, old_relative_position_bias_table[old_num_relative_distance - 3:]])

    key = str(window_size[1]) + "," + str(window_size[0])
    if key not in self.relative_position_indices.keys():
        self.relative_position_indices[key] = gen_relative_position_index(window_size)

    relative_position_bias = new_relative_position_bias_table[
        self.relative_position_indices[key].view(-1)].view(
        window_size[0] * window_size[1] + 1,
        window_size[0] * window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
    relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
    return relative_position_bias.unsqueeze(0)
//  [python]
def _get_rel_pos_bias(self, window_size):
    """
    Modification of timm.models.beit.py: Attention._get_rel_pos_bias to support arbitrary window sizes.
    """
    old_height = 2 * self.window_size[0] - 1
    old_width = 2 * self.window_size[1] - 1

    new_height = 2 * window_size[0] - 1
    new_width = 2 * window_size[1] - 1

    old_relative_position_bias_table = self.relative_position_bias_table

    old_num_relative_distance = self.num_relative_distance
    new_num_relative_distance = new_height * new_width + 3

    old_sub_table = old_relative_position_bias_table[:old_num_relative_distance - 3]

    old_sub_table = old_sub_table.reshape(1, old_width, old_height, -1).permute(0, 3, 1, 2)
    new_sub_table = F.interpolate(old_sub_table, size=(int(new_height), int(new_width)), mode="bilinear")
    new_sub_table = new_sub_table.permute(0, 2, 3, 1).reshape(new_num_relative_distance - 3, -1)

    new_relative_position_bias_table = torch.cat(
        [new_sub_table, old_relative_position_bias_table[old_num_relative_distance - 3:]])

    key = str(window_size[1]) + "," + str(window_size[0])
    if key not in self.relative_position_indices.keys():
        self.relative_position_indices[key] = gen_relative_position_index(window_size)

    relative_position_bias = new_relative_position_bias_table[
        self.relative_position_indices[key].view(-1)].view(
        window_size[0] * window_size[1] + 1,
        window_size[0] * window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
    relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
    return relative_position_bias.unsqueeze(0)
//  [python]
def test_upsample(self):
        class MyModule(Module):
            def forward(self, x):
                size = [v * 2 for v in x.size()[2:]]
                size = [int(i) for i in size]
                return torch.nn.functional.interpolate(x, size=size, mode="nearest")

        module = MyModule()
        ops8 = [
            {
                "op_name": "Upsample",
                "attributes": [
                    {"name": "mode", "s": (b"nearest"), "type": 3},
                    {"name": "scales", "floats": [1.0, 1.0, 2.0, 2.0], "type": 6},
                ],
            }
        ]
        ops9 = [
            {"op_name": "Constant"},
            {
                "op_name": "Upsample",
                "attributes": [{"name": "mode", "s": (b"nearest"), "type": 3}],
            },
        ]
        ops = {8: ops8, 9: ops9}
        x = torch.randn(2, 2, 2, 2)
        check_onnx_opsets_operator(module, x, ops, opset_versions=[8, 9])
//  [python]
def partial_losses(self, input, target, mask=None):
        check_and_warn_input_range(target, 0, 1, 'PerceptualLoss target in partial_losses')

        # we expect input and target to be in [0, 1] range
        losses = []

        if self.normalize_inputs:
            features_input = self.do_normalize_inputs(input)
            features_target = self.do_normalize_inputs(target)
        else:
            features_input = input
            features_target = target

        for layer in self.vgg[:30]:

            features_input = layer(features_input)
            features_target = layer(features_target)

            if layer.__class__.__name__ == 'ReLU':
                loss = F.mse_loss(features_input, features_target, reduction='none')

                if mask is not None:
                    cur_mask = t.interpolate(mask, size=features_input.shape[-2:],
                                             mode='bilinear', align_corners=False)
                    loss = loss * (1 - cur_mask)

                loss = loss.mean(dim=tuple(range(1, len(loss.shape))))
                losses.append(loss)

        return losses
//  [python]
def partial_losses(self, input, target, mask=None):
        check_and_warn_input_range(target, 0, 1, 'PerceptualLoss target in partial_losses')

        # we expect input and target to be in [0, 1] range
        losses = []

        if self.normalize_inputs:
            features_input = self.do_normalize_inputs(input)
            features_target = self.do_normalize_inputs(target)
        else:
            features_input = input
            features_target = target

        for layer in self.vgg[:30]:

            features_input = layer(features_input)
            features_target = layer(features_target)

            if layer.__class__.__name__ == 'ReLU':
                loss = F.mse_loss(features_input, features_target, reduction='none')

                if mask is not None:
                    cur_mask = F.interpolate(mask, size=features_input.shape[-2:],
                                             mode='bilinear', align_corners=False)
                    loss = loss * (1 - cur_mask)

                loss = loss.mean(dim=tuple(range(1, len(loss.shape))))
                losses.append(loss)

        return losses
//  [python]
def test_interpolate():
    """test torch translator for interpolate"""

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

    input_info = [([1, 3, 10, 10], "float32")]
    for via_relax in [True, False]:
        verify_model(Interpolate(), input_info, via_relax)
//  [python]
def test_interpolate():
    """test torch 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")]
    for via_relax in [True, False]:
        verify_model(Interpolate(), input_info, via_relax)
//  [python]
def get_input(self, batch, k):
        x = batch[k]
        if len(x.shape) == 3:
            x = x[..., None]
        x = x.permute(0, 3, 1,
                      2).to(memory_format=torch.contiguous_format).float()
        if self.batch_resize_range is not None:
            lower_size = self.batch_resize_range[0]
            upper_size = self.batch_resize_range[1]
            if self.global_step <= 4:
                # do the first few batches with max size to avoid later oom
                new_resize = upper_size
            else:
                new_resize = np.random.choice(
                    np.arange(lower_size, upper_size + 16, 16))
            if new_resize != x.shape[2]:
                x = t.interpolate(x, size=new_resize, mode='bicubic')
            x = x.detach()
        return x
//  [python]
def get_input(self, batch, k):
        x = batch[k]
        if len(x.shape) == 3:
            x = x[..., None]
        x = x.permute(0, 3, 1,
                      2).to(memory_format=torch.contiguous_format).float()
        if self.batch_resize_range is not None:
            lower_size = self.batch_resize_range[0]
            upper_size = self.batch_resize_range[1]
            if self.global_step <= 4:
                # do the first few batches with max size to avoid later oom
                new_resize = upper_size
            else:
                new_resize = np.random.choice(
                    np.arange(lower_size, upper_size + 16, 16))
            if new_resize != x.shape[2]:
                x = F.interpolate(x, size=new_resize, mode='bicubic')
            x = x.detach()
        return x
//  [python]
def DynP_extract_ms(net, input, ms, msp):
    
    v = torch.zeros(net.meta['outputdim'])
    dyn_p_ms = torch.zeros(len(ms))
    
    for i_s,s in enumerate(ms): 
        if s == 1:
            input_t = input.clone()
        else:    
            input_t = nn.functional.upsample_bilinear(input, scale_factor=s)
            # input_t = nn.functional.interpolate(input, scale_factor=s, mode='bilinear', align_corners=False)
        
        v_temp = net.features(input_t)
        v_temp, dyn_p = net.DAME3(v_temp)
        dyn_p_ms[i_s] = dyn_p
        v_temp = L2N()(v_temp)
        v_temp = v_temp.permute(1,0)
        v += v_temp.pow(msp).cpu().data.squeeze()
        
    v /= len(ms)
    v = v.pow(1./msp)
    v /= v.norm()

    return v,dyn_p_ms
//  [python]
def DynP_extract_ms(net, input, ms, msp):
    
    v = torch.zeros(net.meta['outputdim'])
    dyn_p_ms = torch.zeros(len(ms))
    
    for i_s,s in enumerate(ms): 
        if s == 1:
            input_t = input.clone()
        else:    
            input_t = nn.functional.interpolate(input, scale_factor=s, mode='bilinear', align_corners=False)
            # input_t = nn.functional.upsample_bilinear(input, scale_factor=s)
        
        v_temp = net.features(input_t)
        v_temp, dyn_p = net.DAME3(v_temp)
        dyn_p_ms[i_s] = dyn_p
        v_temp = L2N()(v_temp)
        v_temp = v_temp.permute(1,0)
        v += v_temp.pow(msp).cpu().data.squeeze()
        
    v /= len(ms)
    v = v.pow(1./msp)
    v /= v.norm()

    return v,dyn_p_ms
//  [python]
def get_object_localization_loss_for_one_layer(
    cross_attention_scores,
    object_segmaps,
    object_token_idx,
    object_token_idx_mask,
    loss_fn,
):
    """Get object localization loss for one layer."""

    bxh, num_noise_latents, num_text_tokens = cross_attention_scores.shape
    b, max_num_objects, _, _ = object_segmaps.shape
    size = int(num_noise_latents**0.5)

    # Resize the object segmentation maps to
    # the size of the cross attention scores
    object_segmaps = t.interpolate(
        object_segmaps, size=(size, size), mode='bilinear')
    # (b, max_num_objects, size, size)

    object_segmaps = object_segmaps.view(
        b, max_num_objects, -1)  # (b, max_num_objects, num_noise_latents)

    num_heads = bxh // b

    cross_attention_scores = cross_attention_scores.view(
        b, num_heads, num_noise_latents, num_text_tokens)

    # Gather object_token_attn_prob
    object_token_attn_prob = torch.gather(
        cross_attention_scores,
        dim=3,
        index=object_token_idx.view(b, 1, 1, max_num_objects).expand(
            b, num_heads, num_noise_latents, max_num_objects),
    )  # (b, num_heads, num_noise_latents, max_num_objects)

    object_segmaps = (
        object_segmaps.permute(0, 2,
                               1).unsqueeze(1).expand(b, num_heads,
                                                      num_noise_latents,
                                                      max_num_objects))

    loss = loss_fn(object_token_attn_prob, object_segmaps)

    loss = loss * object_token_idx_mask.view(b, 1, max_num_objects)
    object_token_cnt = object_token_idx_mask.sum(dim=1).view(b, 1) + 1e-5
    loss = (loss.sum(dim=2) / object_token_cnt).mean()

    return loss
//  [python]
def get_object_localization_loss_for_one_layer(
    cross_attention_scores,
    object_segmaps,
    object_token_idx,
    object_token_idx_mask,
    loss_fn,
):
    """Get object localization loss for one layer."""

    bxh, num_noise_latents, num_text_tokens = cross_attention_scores.shape
    b, max_num_objects, _, _ = object_segmaps.shape
    size = int(num_noise_latents**0.5)

    # Resize the object segmentation maps to
    # the size of the cross attention scores
    object_segmaps = F.interpolate(
        object_segmaps, size=(size, size), mode='bilinear')
    # (b, max_num_objects, size, size)

    object_segmaps = object_segmaps.view(
        b, max_num_objects, -1)  # (b, max_num_objects, num_noise_latents)

    num_heads = bxh // b

    cross_attention_scores = cross_attention_scores.view(
        b, num_heads, num_noise_latents, num_text_tokens)

    # Gather object_token_attn_prob
    object_token_attn_prob = torch.gather(
        cross_attention_scores,
        dim=3,
        index=object_token_idx.view(b, 1, 1, max_num_objects).expand(
            b, num_heads, num_noise_latents, max_num_objects),
    )  # (b, num_heads, num_noise_latents, max_num_objects)

    object_segmaps = (
        object_segmaps.permute(0, 2,
                               1).unsqueeze(1).expand(b, num_heads,
                                                      num_noise_latents,
                                                      max_num_objects))

    loss = loss_fn(object_token_attn_prob, object_segmaps)

    loss = loss * object_token_idx_mask.view(b, 1, max_num_objects)
    object_token_cnt = object_token_idx_mask.sum(dim=1).view(b, 1) + 1e-5
    loss = (loss.sum(dim=2) / object_token_cnt).mean()

    return loss
//  [python]
def forward(self, x, t, concat):
        # embedding
        if concat is not None:
            if concat.shape[-2:] != x.shape[-2:]:
                concat = F.upsample(
                    concat, size=x.shape[-2:], mode='bilinear')
            x = torch.cat([x, concat], dim=1)
        e = self.time_embedding(sinusoidal_embedding(t, self.dim))

        # encoder
        xs = []
        for block in self.encoder:
            x = self._forward_single(block, x, e)
            xs.append(x)

        # middle
        for block in self.middle:
            x = self._forward_single(block, x, e)

        # decoder
        for block in self.decoder:
            x = torch.cat([x, xs.pop()], dim=1)
            x = self._forward_single(block, x, e)

        # head
        x = self.head(x)
        return x
//  [python]
def forward(self, x, t, concat):
        # embedding
        if concat is not None:
            if concat.shape[-2:] != x.shape[-2:]:
                concat = F.interpolate(
                    concat, x.shape[-2:], mode='bilinear', align_corners=False)
            x = torch.cat([x, concat], dim=1)
        e = self.time_embedding(sinusoidal_embedding(t, self.dim))

        # encoder
        xs = []
        for block in self.encoder:
            x = self._forward_single(block, x, e)
            xs.append(x)

        # middle
        for block in self.middle:
            x = self._forward_single(block, x, e)

        # decoder
        for block in self.decoder:
            x = torch.cat([x, xs.pop()], dim=1)
            x = self._forward_single(block, x, e)

        # head
        x = self.head(x)
        return x