jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def encode(self, positive, negative, pixels, vae, mask):
        x = (pixels.shape[1] // 8) * 8
        y = (pixels.shape[2] // 8) * 8
        mask = t.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")

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

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

        out_latent = {}

        out_latent["samples"] = orig_latent
        out_latent["noise_mask"] = mask

        out = []
        for conditioning in [positive, negative]:
            c = []
            for t in conditioning:
                d = t[1].copy()
                d["concat_latent_image"] = concat_latent
                d["concat_mask"] = mask
                n = [t[0], d]
                c.append(n)
            out.append(c)
        return (out[0], out[1], out_latent)
//  [python]
def encode(self, positive, negative, pixels, vae, mask):
        x = (pixels.shape[1] // 8) * 8
        y = (pixels.shape[2] // 8) * 8
        mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")

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

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

        out_latent = {}

        out_latent["samples"] = orig_latent
        out_latent["noise_mask"] = mask

        out = []
        for conditioning in [positive, negative]:
            c = []
            for t in conditioning:
                d = t[1].copy()
                d["concat_latent_image"] = concat_latent
                d["concat_mask"] = mask
                n = [t[0], d]
                c.append(n)
            out.append(c)
        return (out[0], out[1], out_latent)
//  [python]
def forward(self, x1, x2):
        """Forward function.

        Args:
            x1 (Tensor): Input tensor with shape (n, c, h, w).
            x2 (Tensor): Input tensor with shape (n, c, 2h, 2w).

        Returns:
            x1 (Tensor): Output tensor with shape (n, c, h, w).
            x2 (Tensor): Output tensor with shape (n, c, 2h, 2w).
        """

        x12 = F.upsample(
            x1, scale_factor=2, mode='bicubic', align_corners=False)
        x12 = F.relu(self.conv1to2(x12))
        x21 = F.relu(self.conv2to1(x2))

        x1 = F.relu(self.conv_merge1(torch.cat((x1, x21), dim=1)))
        x2 = F.relu(self.conv_merge2(torch.cat((x2, x12), dim=1)))

        return x1, x2
//  [python]
def forward(self, x1, x2):
        """Forward function.

        Args:
            x1 (Tensor): Input tensor with shape (n, c, h, w).
            x2 (Tensor): Input tensor with shape (n, c, 2h, 2w).

        Returns:
            x1 (Tensor): Output tensor with shape (n, c, h, w).
            x2 (Tensor): Output tensor with shape (n, c, 2h, 2w).
        """

        x12 = F.interpolate(
            x1, scale_factor=2, mode='bicubic', align_corners=False)
        x12 = F.relu(self.conv1to2(x12))
        x21 = F.relu(self.conv2to1(x2))

        x1 = F.relu(self.conv_merge1(torch.cat((x1, x21), dim=1)))
        x2 = F.relu(self.conv_merge2(torch.cat((x2, x12), dim=1)))

        return x1, x2
//  [python]
def forward(self, x1, x2, x4):
        """Forward function.

        Args:
            x1 (Tensor): Input tensor with shape (n, c, h, w).
            x2 (Tensor): Input tensor with shape (n, c, 2h, 2w).
            x4 (Tensor): Input tensor with shape (n, c, 4h, 4w).

        Returns:
            x (Tensor): Output tensor with shape (n, c_out, 4h, 4w).
        """

        x14 = F.upsample(
            x1, scale_factor=4, mode='bicubic', align_corners=False)
        x14 = F.relu(self.conv1to4(x14))
        x24 = F.upsample(
            x2, scale_factor=2, mode='bicubic', align_corners=False)
        x24 = F.relu(self.conv2to4(x24))

        x = F.relu(self.conv_merge(torch.cat((x4, x14, x24), dim=1)))
        x = self.conv_last1(x)
        x = self.conv_last2(x)
        x = torch.clamp(x, -1, 1)

        return x
//  [python]
def forward(self, x1, x2, x4):
        """Forward function.

        Args:
            x1 (Tensor): Input tensor with shape (n, c, h, w).
            x2 (Tensor): Input tensor with shape (n, c, 2h, 2w).
            x4 (Tensor): Input tensor with shape (n, c, 4h, 4w).

        Returns:
            x (Tensor): Output tensor with shape (n, c_out, 4h, 4w).
        """

        x14 = F.interpolate(
            x1, scale_factor=4, mode='bicubic', align_corners=False)
        x14 = F.relu(self.conv1to4(x14))
        x24 = F.interpolate(
            x2, scale_factor=2, mode='bicubic', align_corners=False)
        x24 = F.relu(self.conv2to4(x24))

        x = F.relu(self.conv_merge(torch.cat((x4, x14, x24), dim=1)))
        x = self.conv_last1(x)
        x = self.conv_last2(x)
        x = torch.clamp(x, -1, 1)

        return x
//  [python]
def forward(self, x, targets=None):
        img_size = x.size(2), x.size(3)
        # feed to backbone network
        backbone_outputs = self.transforminputs(self.backbone_net(x), selected_indices=self.cfg['backbone'].get('selected_indices'))
        # feed to norm layer
        assert len(backbone_outputs) == len(self.norm_layers)
        for idx in range(len(backbone_outputs)):
            backbone_outputs[idx] = self.norm(backbone_outputs[idx], self.norm_layers[idx])
        # feed to decoder
        seg_logits = self.decoder(backbone_outputs[-1])
        # forward according to the mode
        if self.mode == 'TRAIN':
            seg_logits = t.interpolate(seg_logits, size=img_size, mode='bilinear', align_corners=self.align_corners)
            outputs_dict = {'loss_cls': seg_logits}
            backbone_outputs = backbone_outputs[:-1]
            for idx, (out, dec) in enumerate(zip(backbone_outputs, self.auxiliary_decoders)):
                seg_logits_aux = dec(out)
                seg_logits_aux = t.interpolate(seg_logits_aux, size=img_size, mode='bilinear', align_corners=self.align_corners)
                outputs_dict[f'loss_aux{idx+1}'] = seg_logits_aux
            return self.calculatelosses(
                predictions=outputs_dict, targets=targets, losses_cfg=self.cfg['losses']
            )
        return seg_logits
//  [python]
def forward(self, x, targets=None):
        img_size = x.size(2), x.size(3)
        # feed to backbone network
        backbone_outputs = self.transforminputs(self.backbone_net(x), selected_indices=self.cfg['backbone'].get('selected_indices'))
        # feed to norm layer
        assert len(backbone_outputs) == len(self.norm_layers)
        for idx in range(len(backbone_outputs)):
            backbone_outputs[idx] = self.norm(backbone_outputs[idx], self.norm_layers[idx])
        # feed to decoder
        seg_logits = self.decoder(backbone_outputs[-1])
        # forward according to the mode
        if self.mode == 'TRAIN':
            seg_logits = F.interpolate(seg_logits, size=img_size, mode='bilinear', align_corners=self.align_corners)
            outputs_dict = {'loss_cls': seg_logits}
            backbone_outputs = backbone_outputs[:-1]
            for idx, (out, dec) in enumerate(zip(backbone_outputs, self.auxiliary_decoders)):
                seg_logits_aux = dec(out)
                seg_logits_aux = F.interpolate(seg_logits_aux, size=img_size, mode='bilinear', align_corners=self.align_corners)
                outputs_dict[f'loss_aux{idx+1}'] = seg_logits_aux
            return self.calculatelosses(
                predictions=outputs_dict, targets=targets, losses_cfg=self.cfg['losses']
            )
        return seg_logits
//  [python]
def resize(clip, target_size, interpolation_mode):
    assert len(target_size) == 2, "target size should be tuple (height, width)"
    return torch.nn.functional.interpolate(
        clip, size=target_size, mode=interpolation_mode
    )
//  [python]
def forward(self, x1, x2, x4):
        """Forward function.

        Args:
            x1 (Tensor): Input tensor with shape (n, c, h, w).
            x2 (Tensor): Input tensor with shape (n, c, 2h, 2w).
            x4 (Tensor): Input tensor with shape (n, c, 4h, 4w).

        Returns:
            x1 (Tensor): Output tensor with shape (n, c, h, w).
            x2 (Tensor): Output tensor with shape (n, c, 2h, 2w).
            x4 (Tensor): Output tensor with shape (n, c, 4h, 4w).
        """

        x12 = F.upsample(
            x1, scale_factor=2, mode='bicubic', align_corners=False)
        x12 = F.relu(self.conv1to2(x12))
        x14 = F.upsample(
            x1, scale_factor=4, mode='bicubic', align_corners=False)
        x14 = F.relu(self.conv1to4(x14))

        x21 = F.relu(self.conv2to1(x2))
        x24 = F.upsample(
            x2, scale_factor=2, mode='bicubic', align_corners=False)
        x24 = F.relu(self.conv2to4(x24))

        x41 = F.relu(self.conv4to1_1(x4))
        x41 = F.relu(self.conv4to1_2(x41))
        x42 = F.relu(self.conv4to2(x4))

        x1 = F.relu(self.conv_merge1(torch.cat((x1, x21, x41), dim=1)))
        x2 = F.relu(self.conv_merge2(torch.cat((x2, x12, x42), dim=1)))
        x4 = F.relu(self.conv_merge4(torch.cat((x4, x14, x24), dim=1)))

        return x1, x2, x4
//  [python]
def forward(self, x1, x2, x4):
        """Forward function.

        Args:
            x1 (Tensor): Input tensor with shape (n, c, h, w).
            x2 (Tensor): Input tensor with shape (n, c, 2h, 2w).
            x4 (Tensor): Input tensor with shape (n, c, 4h, 4w).

        Returns:
            x1 (Tensor): Output tensor with shape (n, c, h, w).
            x2 (Tensor): Output tensor with shape (n, c, 2h, 2w).
            x4 (Tensor): Output tensor with shape (n, c, 4h, 4w).
        """

        x12 = F.interpolate(
            x1, scale_factor=2, mode='bicubic', align_corners=False)
        x12 = F.relu(self.conv1to2(x12))
        x14 = F.interpolate(
            x1, scale_factor=4, mode='bicubic', align_corners=False)
        x14 = F.relu(self.conv1to4(x14))

        x21 = F.relu(self.conv2to1(x2))
        x24 = F.interpolate(
            x2, scale_factor=2, mode='bicubic', align_corners=False)
        x24 = F.relu(self.conv2to4(x24))

        x41 = F.relu(self.conv4to1_1(x4))
        x41 = F.relu(self.conv4to1_2(x41))
        x42 = F.relu(self.conv4to2(x4))

        x1 = F.relu(self.conv_merge1(torch.cat((x1, x21, x41), dim=1)))
        x2 = F.relu(self.conv_merge2(torch.cat((x2, x12, x42), dim=1)))
        x4 = F.relu(self.conv_merge4(torch.cat((x4, x14, x24), dim=1)))

        return x1, x2, x4
//  [python]
def debug_writing(writer, outputs, labels, inputs, iterations):
    tmp_tar = torch.unsqueeze(labels.cpu().data[0], dim=1)
    # tmp_out = torch.unsqueeze(outputs.cpu().data[0], dim=1)

    tmp_inp = inputs.cpu().data[0]
    tmp_inp[0] += 0.406
    tmp_inp[1] += 0.457
    tmp_inp[2] += 0.480

    tmp_inp[0] += torch.sum(t.upsample(tmp_tar, scale_factor=4, mode='bilinear'), dim=0)[0]
    tmp_inp.clamp_(0, 1)

    writer.add_image('Data/input', tmp_inp, iterations)
//  [python]
def debug_writing(writer, outputs, labels, inputs, iterations):
    tmp_tar = torch.unsqueeze(labels.cpu().data[0], dim=1)
    # tmp_out = torch.unsqueeze(outputs.cpu().data[0], dim=1)

    tmp_inp = inputs.cpu().data[0]
    tmp_inp[0] += 0.406
    tmp_inp[1] += 0.457
    tmp_inp[2] += 0.480

    tmp_inp[0] += torch.sum(F.interpolate(tmp_tar, scale_factor=4, mode='bilinear'), dim=0)[0]
    tmp_inp.clamp_(0, 1)

    writer.add_image('Data/input', tmp_inp, iterations)
//  [python]
def forward(self, x, targets=None):
        img_size = x.size(2), x.size(3)
        # feed to backbone network
        backbone_outputs = self.transforminputs(self.backbone_net(x), selected_indices=self.cfg['backbone'].get('selected_indices'))
        # feed to mla neck
        feats_list = self.mla_neck(list(backbone_outputs))
        # feed to decoder
        outputs = []
        assert len(feats_list) == len(self.up_convs)
        for feats, up_conv in zip(feats_list, self.up_convs):
            outputs.append(up_conv(feats))
        outputs = torch.cat(outputs, dim=1)
        seg_logits = self.decoder(outputs)
        # forward according to the mode
        if self.mode == 'TRAIN':
            seg_logits = F.interpolate(seg_logits, size=img_size, mode='bilinear', align_corners=self.align_corners)
            outputs_dict = {'loss_cls': seg_logits}
            feats_list = feats_list[-len(self.auxiliary_decoders):]
            for idx, (out, dec) in enumerate(zip(feats_list, self.auxiliary_decoders)):
                seg_logits_aux = dec(out)
                seg_logits_aux = F.interpolate(seg_logits_aux, size=img_size, mode='bilinear', align_corners=self.align_corners)
                outputs_dict[f'loss_aux{idx+1}'] = seg_logits_aux
            return self.calculatelosses(
                predictions=outputs_dict, targets=targets, losses_cfg=self.cfg['losses']
            )
        return seg_logits
//  [python]
def forward(self, x, soft_attention, textures):
        """Forward function.

        Args:
            x (Tensor): Input tensor with shape (n, c, h, w).
            soft_attention (Tensor): Soft-Attention tensor with shape
                (n, 1, h, w).
            textures (Tuple[Tensor]): Transferred HR texture tensors.
                [(N, C, H, W), (N, C/2, 2H, 2W), ...]

        Returns:
            Tensor: Forward results.
        """

        assert textures[-1].shape[1] == self.texture_channels

        x1 = self.sfe(x)

        # stage 1
        x1_res = torch.cat((x1, textures[0]), dim=1)
        x1_res = self.conv_first1(x1_res)

        # soft-attention
        x1 = x1 + x1_res * soft_attention

        x1_res = self.res_block1(x1)
        x1_res = self.conv_last1(x1_res)

        x1 = x1 + x1_res

        # stage 2
        x21 = x1
        x22 = self.up1(x1)
        x22 = F.relu(x22)

        x22_res = torch.cat((x22, textures[1]), dim=1)
        x22_res = self.conv_first2(x22_res)

        # soft-attention
        x22_res = x22_res * t(
            soft_attention,
            scale_factor=2,
            mode='bicubic',
            align_corners=False)
        x22 = x22 + x22_res

        x21_res, x22_res = self.csfi2(x21, x22)

        x21_res = self.res_block2_1(x21_res)
        x22_res = self.res_block2_2(x22_res)

        x21_res = self.conv_last2_1(x21_res)
        x22_res = self.conv_last2_2(x22_res)

        x21 = x21 + x21_res
        x22 = x22 + x22_res

        # stage 3
        x31 = x21
        x32 = x22
        x33 = self.up2(x22)
        x33 = F.relu(x33)

        x33_res = torch.cat((x33, textures[2]), dim=1)
        x33_res = self.conv_first3(x33_res)

        # soft-attention
        x33_res = x33_res * t(
            soft_attention,
            scale_factor=4,
            mode='bicubic',
            align_corners=False)
        x33 = x33 + x33_res

        x31_res, x32_res, x33_res = self.csfi3(x31, x32, x33)

        x31_res = self.res_block3_1(x31_res)
        x32_res = self.res_block3_2(x32_res)
        x33_res = self.res_block3_3(x33_res)

        x31_res = self.conv_last3_1(x31_res)
        x32_res = self.conv_last3_2(x32_res)
        x33_res = self.conv_last3_3(x33_res)

        x31 = x31 + x31_res
        x32 = x32 + x32_res
        x33 = x33 + x33_res
        x = self.merge_features(x31, x32, x33)

        return x
//  [python]
def forward(self, x, soft_attention, textures):
        """Forward function.

        Args:
            x (Tensor): Input tensor with shape (n, c, h, w).
            soft_attention (Tensor): Soft-Attention tensor with shape
                (n, 1, h, w).
            textures (Tuple[Tensor]): Transferred HR texture tensors.
                [(N, C, H, W), (N, C/2, 2H, 2W), ...]

        Returns:
            Tensor: Forward results.
        """

        assert textures[-1].shape[1] == self.texture_channels

        x1 = self.sfe(x)

        # stage 1
        x1_res = torch.cat((x1, textures[0]), dim=1)
        x1_res = self.conv_first1(x1_res)

        # soft-attention
        x1 = x1 + x1_res * soft_attention

        x1_res = self.res_block1(x1)
        x1_res = self.conv_last1(x1_res)

        x1 = x1 + x1_res

        # stage 2
        x21 = x1
        x22 = self.up1(x1)
        x22 = F.relu(x22)

        x22_res = torch.cat((x22, textures[1]), dim=1)
        x22_res = self.conv_first2(x22_res)

        # soft-attention
        x22_res = x22_res * F.interpolate(
            soft_attention,
            scale_factor=2,
            mode='bicubic',
            align_corners=False)
        x22 = x22 + x22_res

        x21_res, x22_res = self.csfi2(x21, x22)

        x21_res = self.res_block2_1(x21_res)
        x22_res = self.res_block2_2(x22_res)

        x21_res = self.conv_last2_1(x21_res)
        x22_res = self.conv_last2_2(x22_res)

        x21 = x21 + x21_res
        x22 = x22 + x22_res

        # stage 3
        x31 = x21
        x32 = x22
        x33 = self.up2(x22)
        x33 = F.relu(x33)

        x33_res = torch.cat((x33, textures[2]), dim=1)
        x33_res = self.conv_first3(x33_res)

        # soft-attention
        x33_res = x33_res * F.interpolate(
            soft_attention,
            scale_factor=4,
            mode='bicubic',
            align_corners=False)
        x33 = x33 + x33_res

        x31_res, x32_res, x33_res = self.csfi3(x31, x32, x33)

        x31_res = self.res_block3_1(x31_res)
        x32_res = self.res_block3_2(x32_res)
        x33_res = self.res_block3_3(x33_res)

        x31_res = self.conv_last3_1(x31_res)
        x32_res = self.conv_last3_2(x32_res)
        x33_res = self.conv_last3_3(x33_res)

        x31 = x31 + x31_res
        x32 = x32 + x32_res
        x33 = x33 + x33_res
        x = self.merge_features(x31, x32, x33)

        return x
//  [python]
def _down_encode(self, x, r_embed, clip, cnet=None):
        level_outputs = []
        block_group = zip(self.down_blocks, self.down_downscalers, self.down_repeat_mappers)
        for down_block, downscaler, repmap in block_group:
            x = downscaler(x)
            for i in range(len(repmap) + 1):
                for block in down_block:
                    if isinstance(block, ResBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  ResBlock)):
                        if cnet is not None:
                            next_cnet = cnet.pop()
                            if next_cnet is not None:
                                x = x + t.interpolate(next_cnet, size=x.shape[-2:], mode='bilinear',
                                                                  align_corners=True).to(x.dtype)
                        x = block(x)
                    elif isinstance(block, AttnBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  AttnBlock)):
                        x = block(x, clip)
                    elif isinstance(block, TimestepBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  TimestepBlock)):
                        x = block(x, r_embed)
                    else:
                        x = block(x)
                if i < len(repmap):
                    x = repmap[i](x)
            level_outputs.insert(0, x)
        return level_outputs
//  [python]
def _down_encode(self, x, r_embed, clip, cnet=None):
        level_outputs = []
        block_group = zip(self.down_blocks, self.down_downscalers, self.down_repeat_mappers)
        for down_block, downscaler, repmap in block_group:
            x = downscaler(x)
            for i in range(len(repmap) + 1):
                for block in down_block:
                    if isinstance(block, ResBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  ResBlock)):
                        if cnet is not None:
                            next_cnet = cnet.pop()
                            if next_cnet is not None:
                                x = x + nn.functional.interpolate(next_cnet, size=x.shape[-2:], mode='bilinear',
                                                                  align_corners=True).to(x.dtype)
                        x = block(x)
                    elif isinstance(block, AttnBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  AttnBlock)):
                        x = block(x, clip)
                    elif isinstance(block, TimestepBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  TimestepBlock)):
                        x = block(x, r_embed)
                    else:
                        x = block(x)
                if i < len(repmap):
                    x = repmap[i](x)
            level_outputs.insert(0, x)
        return level_outputs
//  [python]
def forward(self, source_image, kp_driving, kp_source):
        # Encoding (downsampling) part
        out = self.first(source_image)
        for i in range(len(self.down_blocks)):
            out = self.down_blocks[i](out)
        out = self.second(out)
        bs, c, h, w = out.shape
        # print(out.shape)
        feature_3d = out.view(bs, self.reshape_channel, self.reshape_depth, h ,w) 
        feature_3d = self.resblocks_3d(feature_3d)

        # Transforming feature representation according to deformation and occlusion
        output_dict = {}
        if self.dense_motion_network is not None:
            dense_motion = self.dense_motion_network(feature=feature_3d, kp_driving=kp_driving,
                                                     kp_source=kp_source)
            output_dict['mask'] = dense_motion['mask']

            if 'occlusion_map' in dense_motion:
                occlusion_map = dense_motion['occlusion_map']
                output_dict['occlusion_map'] = occlusion_map
            else:
                occlusion_map = None
            deformation = dense_motion['deformation']
            out = self.deform_input(feature_3d, deformation)

            bs, c, d, h, w = out.shape
            out = out.view(bs, c*d, h, w)
            out = self.third(out)
            out = self.fourth(out)

            if occlusion_map is not None:
                if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]:
                    occlusion_map = t.nn.functional.interpolate(occlusion_map, size=out.shape[2:], mode='bilinear')
                out = out * occlusion_map

            # output_dict["deformed"] = self.deform_input(source_image, deformation)  # 3d deformation cannot deform 2d image

        # Decoding part
        out = self.resblocks_2d(out)
        for i in range(len(self.up_blocks)):
            out = self.up_blocks[i](out)
        out = self.final(out)
        out = F.sigmoid(out)

        output_dict["prediction"] = out

        return output_dict
//  [python]
def forward(self, source_image, kp_driving, kp_source):
        # Encoding (downsampling) part
        out = self.first(source_image)
        for i in range(len(self.down_blocks)):
            out = self.down_blocks[i](out)
        out = self.second(out)
        bs, c, h, w = out.shape
        # print(out.shape)
        feature_3d = out.view(bs, self.reshape_channel, self.reshape_depth, h ,w) 
        feature_3d = self.resblocks_3d(feature_3d)

        # Transforming feature representation according to deformation and occlusion
        output_dict = {}
        if self.dense_motion_network is not None:
            dense_motion = self.dense_motion_network(feature=feature_3d, kp_driving=kp_driving,
                                                     kp_source=kp_source)
            output_dict['mask'] = dense_motion['mask']

            if 'occlusion_map' in dense_motion:
                occlusion_map = dense_motion['occlusion_map']
                output_dict['occlusion_map'] = occlusion_map
            else:
                occlusion_map = None
            deformation = dense_motion['deformation']
            out = self.deform_input(feature_3d, deformation)

            bs, c, d, h, w = out.shape
            out = out.view(bs, c*d, h, w)
            out = self.third(out)
            out = self.fourth(out)

            if occlusion_map is not None:
                if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]:
                    occlusion_map = F.interpolate(occlusion_map, size=out.shape[2:], mode='bilinear')
                out = out * occlusion_map

            # output_dict["deformed"] = self.deform_input(source_image, deformation)  # 3d deformation cannot deform 2d image

        # Decoding part
        out = self.resblocks_2d(out)
        for i in range(len(self.up_blocks)):
            out = self.up_blocks[i](out)
        out = self.final(out)
        out = F.sigmoid(out)

        output_dict["prediction"] = out

        return output_dict
//  [python]
def _up_decode(self, level_outputs, r_embed, clip, cnet=None):
        x = level_outputs[0]
        block_group = zip(self.up_blocks, self.up_upscalers, self.up_repeat_mappers)
        for i, (up_block, upscaler, repmap) in enumerate(block_group):
            for j in range(len(repmap) + 1):
                for k, block in enumerate(up_block):
                    if isinstance(block, ResBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  ResBlock)):
                        skip = level_outputs[i] if k == 0 and i > 0 else None
                        if skip is not None and (x.size(-1) != skip.size(-1) or x.size(-2) != skip.size(-2)):
                            x = t.interpolate(x, skip.shape[-2:], mode='bilinear',
                                                                align_corners=True)
                        if cnet is not None:
                            next_cnet = cnet.pop()
                            if next_cnet is not None:
                                x = x + t.interpolate(next_cnet, size=x.shape[-2:], mode='bilinear',
                                                                  align_corners=True).to(x.dtype)
                        x = block(x, skip)
                    elif isinstance(block, AttnBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  AttnBlock)):
                        x = block(x, clip)
                    elif isinstance(block, TimestepBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  TimestepBlock)):
                        x = block(x, r_embed)
                    else:
                        x = block(x)
                if j < len(repmap):
                    x = repmap[j](x)
            x = upscaler(x)
        return x
//  [python]
def _up_decode(self, level_outputs, r_embed, clip, cnet=None):
        x = level_outputs[0]
        block_group = zip(self.up_blocks, self.up_upscalers, self.up_repeat_mappers)
        for i, (up_block, upscaler, repmap) in enumerate(block_group):
            for j in range(len(repmap) + 1):
                for k, block in enumerate(up_block):
                    if isinstance(block, ResBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  ResBlock)):
                        skip = level_outputs[i] if k == 0 and i > 0 else None
                        if skip is not None and (x.size(-1) != skip.size(-1) or x.size(-2) != skip.size(-2)):
                            x = torch.nn.functional.interpolate(x, skip.shape[-2:], mode='bilinear',
                                                                align_corners=True)
                        if cnet is not None:
                            next_cnet = cnet.pop()
                            if next_cnet is not None:
                                x = x + nn.functional.interpolate(next_cnet, size=x.shape[-2:], mode='bilinear',
                                                                  align_corners=True).to(x.dtype)
                        x = block(x, skip)
                    elif isinstance(block, AttnBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  AttnBlock)):
                        x = block(x, clip)
                    elif isinstance(block, TimestepBlock) or (
                            hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
                                                                                  TimestepBlock)):
                        x = block(x, r_embed)
                    else:
                        x = block(x)
                if j < len(repmap):
                    x = repmap[j](x)
            x = upscaler(x)
        return x
//  [python]
def interpolate_pos_embed(model, checkpoint_model):
    if 'pos_embed' in checkpoint_model:
        pos_embed_checkpoint = checkpoint_model['pos_embed'].float()
        embedding_size = pos_embed_checkpoint.shape[-1]
        num_patches = model.patch_embed.num_patches
        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
        # height (== width) for the checkpoint position embedding
        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
        # height (== width) for the new position embedding
        new_size = int(num_patches ** 0.5)
        # class_token and dist_token are kept unchanged
        if orig_size != new_size:
            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
            # only the position tokens are interpolated
            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
            pos_tokens = torch.nn.functional.interpolate(
                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
            checkpoint_model['pos_embed'] = new_pos_embed
//  [python]
def _extract_mask(self, dt: Dict[str, Any]) -> np.ndarray:
        if "densepose" in dt:
            densepose_results_quantized = dt["densepose"]
            return densepose_results_quantized.labels_uv_uint8[0].numpy()
        elif "cse_mask" in dt:
            return dt["cse_mask"]
        elif "coarse_segm" in dt:
            dy = max(int(dt["bbox"][3]), 1)
            dx = max(int(dt["bbox"][2]), 1)
            return (
                t.interpolate(
                    dt["coarse_segm"].unsqueeze(0),
                    (dy, dx),
                    mode="bilinear",
                    align_corners=False,
                )
                .squeeze(0)
                .argmax(0)
                .numpy()
                .astype(np.uint8)
            )
        elif "record_id" in dt:
            assert (
                self.multi_storage is not None
            ), f"Storage record id encountered in a detection {dt}, but no storage provided!"
            record = self.multi_storage.get(dt["rank"], dt["record_id"])
            coarse_segm = record["coarse_segm"]
            dy = max(int(dt["bbox"][3]), 1)
            dx = max(int(dt["bbox"][2]), 1)
            return (
                t.interpolate(
                    coarse_segm.unsqueeze(0),
                    (dy, dx),
                    mode="bilinear",
                    align_corners=False,
                )
                .squeeze(0)
                .argmax(0)
                .numpy()
                .astype(np.uint8)
            )
        else:
            raise Exception(f"No mask data in the detection: {dt}")
        raise ValueError('The prediction dict needs to contain either "densepose" or "cse_mask"')
//  [python]
def _extract_mask(self, dt: Dict[str, Any]) -> np.ndarray:
        if "densepose" in dt:
            densepose_results_quantized = dt["densepose"]
            return densepose_results_quantized.labels_uv_uint8[0].numpy()
        elif "cse_mask" in dt:
            return dt["cse_mask"]
        elif "coarse_segm" in dt:
            dy = max(int(dt["bbox"][3]), 1)
            dx = max(int(dt["bbox"][2]), 1)
            return (
                F.interpolate(
                    dt["coarse_segm"].unsqueeze(0),
                    (dy, dx),
                    mode="bilinear",
                    align_corners=False,
                )
                .squeeze(0)
                .argmax(0)
                .numpy()
                .astype(np.uint8)
            )
        elif "record_id" in dt:
            assert (
                self.multi_storage is not None
            ), f"Storage record id encountered in a detection {dt}, but no storage provided!"
            record = self.multi_storage.get(dt["rank"], dt["record_id"])
            coarse_segm = record["coarse_segm"]
            dy = max(int(dt["bbox"][3]), 1)
            dx = max(int(dt["bbox"][2]), 1)
            return (
                F.interpolate(
                    coarse_segm.unsqueeze(0),
                    (dy, dx),
                    mode="bilinear",
                    align_corners=False,
                )
                .squeeze(0)
                .argmax(0)
                .numpy()
                .astype(np.uint8)
            )
        else:
            raise Exception(f"No mask data in the detection: {dt}")
        raise ValueError('The prediction dict needs to contain either "densepose" or "cse_mask"')
//  [python]
def computeOgps_single_pair_cse(
        self, dt, gt, py, px, pt_mask, coarse_segm, embedding, bbox_xywh_abs
    ):
        # 0-based mesh vertex indices
        cVertsGT = torch.as_tensor(gt["dp_vertex"], dtype=torch.int64)
        # label for each pixel of the bbox, [H, W] tensor of long
        labels_dt = resample_coarse_segm_tensor_to_bbox(
            coarse_segm.unsqueeze(0), bbox_xywh_abs
        ).squeeze(0)
        x, y, w, h = bbox_xywh_abs
        # embedding for each pixel of the bbox, [D, H, W] tensor of float32
        embedding = t.nn.functional.interpolate(
            embedding.unsqueeze(0), (int(h), int(w)), mode="bilinear", align_corners=False
        ).squeeze(0)
        # valid locations py, px
        py_pt = torch.from_numpy(py[pt_mask > -1])
        px_pt = torch.from_numpy(px[pt_mask > -1])
        cVerts = torch.ones_like(cVertsGT) * -1
        cVerts[pt_mask > -1] = self.findClosestVertsCse(
            embedding, py_pt, px_pt, labels_dt, gt["ref_model"]
        )
        # Get pairwise geodesic distances between gt and estimated mesh points.
        dist = self.getDistancesCse(cVertsGT, cVerts, gt["ref_model"])
        # normalize distances
        if (gt["ref_model"] == "smpl_27554") and ("dp_I" in gt):
            Current_Mean_Distances = self.Mean_Distances[
                self.CoarseParts[np.array(gt["dp_I"], dtype=int)]
            ]
        else:
            Current_Mean_Distances = 0.255
        return dist, Current_Mean_Distances
//  [python]
def computeOgps_single_pair_cse(
        self, dt, gt, py, px, pt_mask, coarse_segm, embedding, bbox_xywh_abs
    ):
        # 0-based mesh vertex indices
        cVertsGT = torch.as_tensor(gt["dp_vertex"], dtype=torch.int64)
        # label for each pixel of the bbox, [H, W] tensor of long
        labels_dt = resample_coarse_segm_tensor_to_bbox(
            coarse_segm.unsqueeze(0), bbox_xywh_abs
        ).squeeze(0)
        x, y, w, h = bbox_xywh_abs
        # embedding for each pixel of the bbox, [D, H, W] tensor of float32
        embedding = F.interpolate(
            embedding.unsqueeze(0), (int(h), int(w)), mode="bilinear", align_corners=False
        ).squeeze(0)
        # valid locations py, px
        py_pt = torch.from_numpy(py[pt_mask > -1])
        px_pt = torch.from_numpy(px[pt_mask > -1])
        cVerts = torch.ones_like(cVertsGT) * -1
        cVerts[pt_mask > -1] = self.findClosestVertsCse(
            embedding, py_pt, px_pt, labels_dt, gt["ref_model"]
        )
        # Get pairwise geodesic distances between gt and estimated mesh points.
        dist = self.getDistancesCse(cVertsGT, cVerts, gt["ref_model"])
        # normalize distances
        if (gt["ref_model"] == "smpl_27554") and ("dp_I" in gt):
            Current_Mean_Distances = self.Mean_Distances[
                self.CoarseParts[np.array(gt["dp_I"], dtype=int)]
            ]
        else:
            Current_Mean_Distances = 0.255
        return dist, Current_Mean_Distances
//  [python]
def forward(self, x):
        x1 = self.aspp1(x)
        x2 = self.aspp2(x)
        x3 = self.aspp3(x)
        x4 = self.aspp4(x)
        x5 = self.global_avg_pool(x)
        x5 = F.upsample(
            x5, size=x4.size()[2:], mode='bilinear', align_corners=True)
        x = torch.cat((x1, x2, x3, x4, x5), dim=1)

        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        return self.dropout(x)
//  [python]
def forward(self, x):
        x1 = self.aspp1(x)
        x2 = self.aspp2(x)
        x3 = self.aspp3(x)
        x4 = self.aspp4(x)
        x5 = self.global_avg_pool(x)
        x5 = F.interpolate(
            x5, size=x4.size()[2:], mode='bilinear', align_corners=True)
        x = torch.cat((x1, x2, x3, x4, x5), dim=1)

        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)

        return self.dropout(x)
//  [python]
def test_interpolate(self):
        input1 = torch.empty(3, 3, 4, 2)
        npu_input = input1.npu()

        cpu_output = t.interpolate(input1, scale_factor=4)
        npu_output = t.interpolate(npu_input, scale_factor=4)

        self.assertRtolEqual(cpu_output.detach().numpy(), npu_output.detach().cpu().numpy())
//  [python]
def test_interpolate(self):
        input1 = torch.empty(3, 3, 4, 2)
        npu_input = input1.npu()

        cpu_output = F.interpolate(input1, 4)
        npu_output = F.interpolate(npu_input, 4)

        self.assertRtolEqual(cpu_output.detach().numpy(), npu_output.detach().cpu().numpy())
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        outputs = []
        if self.upsample_mode == "transpose":
            for project_module, pool_module, up_module in zip(self.project_modules, self.pool_modules, self.up_modules):
                output = up_module(project_module(pool_module(x)))
                outputs.append(output)
        else:
            for project_module, pool_module in zip(self.project_modules, self.pool_modules):
                interpolate_size = x.shape[2:]
                align_corners: Union[bool, None] = None
                if self.upsample_mode in ["trilinear", "bilinear"]:
                    align_corners = True
                output = t.interpolate(
                    project_module(pool_module(x)),
                    size=interpolate_size,
                    mode=self.upsample_mode,
                    align_corners=align_corners,
                )
                outputs.append(output)
        x = torch.cat(outputs, dim=1)
        return x
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        outputs = []
        if self.upsample_mode == "transpose":
            for project_module, pool_module, up_module in zip(self.project_modules, self.pool_modules, self.up_modules):
                output = up_module(project_module(pool_module(x)))
                outputs.append(output)
        else:
            for project_module, pool_module in zip(self.project_modules, self.pool_modules):
                interpolate_size = x.shape[2:]
                align_corners: Union[bool, None] = None
                if self.upsample_mode in ["trilinear", "bilinear"]:
                    align_corners = True
                output = F.interpolate(
                    project_module(pool_module(x)),
                    size=interpolate_size,
                    mode=self.upsample_mode,
                    align_corners=align_corners,
                )
                outputs.append(output)
        x = torch.cat(outputs, dim=1)
        return x
//  [python]
def forward(
        self,
        images: Tensor,
    ) -> Tensor:
        features, sizes = self.encoder(images)
        context = None
        scales = self.decoder.scales
        for idx, (feature_map, scale) in enumerate(zip(reversed(features), scales)):
            if idx == 0:
                descriptions, context = self.decoder(feature_map, scale=scale, context=context)
            else:
                delta_descriptions, context = self.decoder(feature_map, scale=scale, context=context)
                descriptions = descriptions + delta_descriptions
            if idx < len(scales) - 1:
                size = sizes[-(idx + 2)]
                descriptions = F.interpolate(descriptions, size=size, mode="bilinear", align_corners=False)
                context = F.interpolate(context, size=size, mode="bilinear", align_corners=False)
        return descriptions
//  [python]
def forward(
        self,
        imgs: Optional[torch.Tensor],
        masks: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> Dict[Any, torch.Tensor]:
        """
        Args:
            imgs: A batch of input images of shape `(B, 3, H, W)`.
            masks: A batch of input masks of shape `(B, 3, H, W)`.

        Returns:
            out_feats: A dict `{f_i: t_i}` keyed by predicted feature names `f_i`
                and their corresponding tensors `t_i` of shape `(B, dim_i, H_i, W_i)`.
        """

        out_feats = {}

        imgs_input = imgs
        if self.image_rescale != 1.0 and imgs_input is not None:
            imgs_resized = t.interpolate(
                imgs_input,
                scale_factor=self.image_rescale,
                mode="bilinear",
            )
        else:
            imgs_resized = imgs_input

        if len(self.stages) > 0:
            assert imgs_resized is not None

            if self.normalize_image:
                imgs_normed = self._resnet_normalize_image(imgs_resized)
            else:
                imgs_normed = imgs_resized
            #  is not a function.
            feats = self.stem(imgs_normed)
            for stage, (layer, proj) in enumerate(zip(self.layers, self.proj_layers)):
                feats = layer(feats)
                # just a sanity check below
                assert feats.shape[1] == _FEAT_DIMS[self.name][stage]
                if (stage + 1) in self.stages:
                    f = proj(feats)
                    if self.global_average_pool:
                        f = f.mean(dims=(2, 3))
                    if self.l2_norm:
                        normfac = 1.0 / math.sqrt(len(self.stages))
                        f = Fu.normalize(f, dim=1) * normfac
                    feature_name = self._get_resnet_stage_feature_name(stage)
                    out_feats[feature_name] = f

        if self.add_masks:
            assert masks is not None
            out_feats[MASK_FEATURE_NAME] = masks

        if self.add_images:
            assert imgs_resized is not None
            out_feats[IMAGE_FEATURE_NAME] = imgs_resized

        if self.feature_rescale != 1.0:
            out_feats = {k: self.feature_rescale * f for k, f in out_feats.items()}

        # pyre-fixme[7]: Incompatible return type, expected `Dict[typing.Any, Tensor]`
        # but got `Dict[typing.Any, float]`
        return out_feats
//  [python]
def forward(
        self,
        imgs: Optional[torch.Tensor],
        masks: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> Dict[Any, torch.Tensor]:
        """
        Args:
            imgs: A batch of input images of shape `(B, 3, H, W)`.
            masks: A batch of input masks of shape `(B, 3, H, W)`.

        Returns:
            out_feats: A dict `{f_i: t_i}` keyed by predicted feature names `f_i`
                and their corresponding tensors `t_i` of shape `(B, dim_i, H_i, W_i)`.
        """

        out_feats = {}

        imgs_input = imgs
        if self.image_rescale != 1.0 and imgs_input is not None:
            imgs_resized = Fu.interpolate(
                imgs_input,
                scale_factor=self.image_rescale,
                mode="bilinear",
            )
        else:
            imgs_resized = imgs_input

        if len(self.stages) > 0:
            assert imgs_resized is not None

            if self.normalize_image:
                imgs_normed = self._resnet_normalize_image(imgs_resized)
            else:
                imgs_normed = imgs_resized
            #  is not a function.
            feats = self.stem(imgs_normed)
            for stage, (layer, proj) in enumerate(zip(self.layers, self.proj_layers)):
                feats = layer(feats)
                # just a sanity check below
                assert feats.shape[1] == _FEAT_DIMS[self.name][stage]
                if (stage + 1) in self.stages:
                    f = proj(feats)
                    if self.global_average_pool:
                        f = f.mean(dims=(2, 3))
                    if self.l2_norm:
                        normfac = 1.0 / math.sqrt(len(self.stages))
                        f = Fu.normalize(f, dim=1) * normfac
                    feature_name = self._get_resnet_stage_feature_name(stage)
                    out_feats[feature_name] = f

        if self.add_masks:
            assert masks is not None
            out_feats[MASK_FEATURE_NAME] = masks

        if self.add_images:
            assert imgs_resized is not None
            out_feats[IMAGE_FEATURE_NAME] = imgs_resized

        if self.feature_rescale != 1.0:
            out_feats = {k: self.feature_rescale * f for k, f in out_feats.items()}

        # pyre-fixme[7]: Incompatible return type, expected `Dict[typing.Any, Tensor]`
        # but got `Dict[typing.Any, float]`
        return out_feats
//  [python]
def forward(self, x):
        h, w = x.shape[-2:]

        features = self.pretrained.get_intermediate_layers(x, 4, return_class_token=True)

        patch_h, patch_w = h // 14, w // 14

        depth = self.depth_head(features, patch_h, patch_w)
        depth = F.upsample(depth, size=(h, w), mode="bilinear")
        depth = F.relu(depth)

        return depth.squeeze(1)
//  [python]
def forward(self, x):
        h, w = x.shape[-2:]

        features = self.pretrained.get_intermediate_layers(x, 4, return_class_token=True)

        patch_h, patch_w = h // 14, w // 14

        depth = self.depth_head(features, patch_h, patch_w)
        depth = F.interpolate(depth, size=(h, w), mode="bilinear", align_corners=True)
        depth = F.relu(depth)

        return depth.squeeze(1)