jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self,input,e0, e1, e2, e3):

        # if self.use_drop:
        #     self.dropblock.step()

        if self.use_att:
            #d4 = self.GAU4(e4, e3)  # 16
            d3 = self.GAU3(e3, e2)  # 32
            d2 = self.GAU2(d3, e1)  # 64
            d1 = self.GAU1(d2, e0)  # 128
        else:
            #d4 = self.decoder4(e4) + e3  # [2,256,16,16]
            d3 = self.decoder3(e3) + e2  # [2,128,32,32]
            d2 = self.decoder2(d3) + e1  # [2,64,64,64]
            d1 = self.decoder1(d2)  # [2,64,128,128]
        # out = self.finaldeconv1(d1)  # [2,64,256,256]
        # out = self.finalrelu1(out)
        # out = self.finalconv2(out)  # [2,32,256,256]
        # out = self.finalrelu2(out)
        # out=self.dropblock(out)
        # final = self.finalconv3(out)  # [2,1,256,256]
        #
        # if self.RefNet:
        #     final = self.RefNet(final)
        if self.use_rfnet:
            out = self.final_conv(d1)#[8,1,64,64]
            _, _, H0, W0 = e0.size()
            _, _, H, W = input.size()
            fea0 = F.interpolate(input, (H0, W0), mode='bilinear', align_corners=True)
            fea = torch.cat([fea0, e0], dim=1)#[8,67,64,64]
            out = self.RefNet(out, fea)
            final = F.interpolate(out, (H, W), mode='bilinear', align_corners=True)

        else:
            out = self.finaldeconv1(d1)  # [2,64,256,256]
            out = self.finalrelu1(out)
            out = self.finalconv2(out)  # [2,32,256,256]
            out = self.finalrelu2(out)
            final = self.finalconv3(out)  # [2,1,256,256]




        if self.n_class == 1:
            return F.sigmoid(final)  # for binary only
        else:
            return final
//  [python]
def forward(self, f0, upp):
        """sine_tensor, uv = forward(f0)
        input F0: tensor(batchsize=1, length, dim=1)
                  f0 for unvoiced steps should be 0
        output sine_tensor: tensor(batchsize=1, length, dim)
        output uv: tensor(batchsize=1, length, 1)
        """
        f0 = f0.unsqueeze(-1)
        fn = torch.multiply(f0, torch.arange(1, self.dim + 1, device=f0.device).reshape((1, 1, -1)))
        rad_values = (fn / self.sampling_rate) % 1  # %1意味着n_har的乘积无法后处理优化
        rand_ini = torch.rand(fn.shape[0], fn.shape[2], device=fn.device)
        rand_ini[:, 0] = 0
        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
        is_half = rad_values.dtype is not torch.float32
        tmp_over_one = torch.cumsum(rad_values.double(), 1)  # % 1  #####%1意味着后面的cumsum无法再优化
        if is_half:
            tmp_over_one = tmp_over_one.half()
        else:
            tmp_over_one = tmp_over_one.float()
        tmp_over_one *= upp
        tmp_over_one = t.nn.functional.interpolate(tmp_over_one.transpose(2, 1), scale_factor=upp, mode="linear", align_corners=True).transpose(2, 1)
        rad_values = t.nn.functional.interpolate(rad_values.transpose(2, 1), scale_factor=upp, mode="nearest").transpose(2, 1)
        tmp_over_one %= 1
        tmp_over_one_idx = (tmp_over_one[:, 1:, :] - tmp_over_one[:, :-1, :]) < 0
        cumsum_shift = torch.zeros_like(rad_values)
        cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
        rad_values = rad_values.double()
        cumsum_shift = cumsum_shift.double()
        sine_waves = torch.sin(torch.cumsum(rad_values + cumsum_shift, dim=1) * 2 * np.pi)
        if is_half:
            sine_waves = sine_waves.half()
        else:
            sine_waves = sine_waves.float()
        sine_waves = sine_waves * self.sine_amp
        uv = self._f02uv(f0)
        uv = t.nn.functional.interpolate(uv.transpose(2, 1), scale_factor=upp, mode="nearest").transpose(2, 1)
        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
        noise = noise_amp * torch.randn_like(sine_waves)
        sine_waves = sine_waves * uv + noise
        return sine_waves, uv, noise
//  [python]
def forward(self, f0, upp):
        """sine_tensor, uv = forward(f0)
        input F0: tensor(batchsize=1, length, dim=1)
                  f0 for unvoiced steps should be 0
        output sine_tensor: tensor(batchsize=1, length, dim)
        output uv: tensor(batchsize=1, length, 1)
        """
        f0 = f0.unsqueeze(-1)
        fn = torch.multiply(f0, torch.arange(1, self.dim + 1, device=f0.device).reshape((1, 1, -1)))
        rad_values = (fn / self.sampling_rate) % 1  # %1意味着n_har的乘积无法后处理优化
        rand_ini = torch.rand(fn.shape[0], fn.shape[2], device=fn.device)
        rand_ini[:, 0] = 0
        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
        is_half = rad_values.dtype is not torch.float32
        tmp_over_one = torch.cumsum(rad_values.double(), 1)  # % 1  #####%1意味着后面的cumsum无法再优化
        if is_half:
            tmp_over_one = tmp_over_one.half()
        else:
            tmp_over_one = tmp_over_one.float()
        tmp_over_one *= upp
        tmp_over_one = F.interpolate(tmp_over_one.transpose(2, 1), scale_factor=upp, mode="linear", align_corners=True).transpose(2, 1)
        rad_values = F.interpolate(rad_values.transpose(2, 1), scale_factor=upp, mode="nearest").transpose(2, 1)
        tmp_over_one %= 1
        tmp_over_one_idx = (tmp_over_one[:, 1:, :] - tmp_over_one[:, :-1, :]) < 0
        cumsum_shift = torch.zeros_like(rad_values)
        cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
        rad_values = rad_values.double()
        cumsum_shift = cumsum_shift.double()
        sine_waves = torch.sin(torch.cumsum(rad_values + cumsum_shift, dim=1) * 2 * np.pi)
        if is_half:
            sine_waves = sine_waves.half()
        else:
            sine_waves = sine_waves.float()
        sine_waves = sine_waves * self.sine_amp
        uv = self._f02uv(f0)
        uv = F.interpolate(uv.transpose(2, 1), scale_factor=upp, mode="nearest").transpose(2, 1)
        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
        noise = noise_amp * torch.randn_like(sine_waves)
        sine_waves = sine_waves * uv + noise
        return sine_waves, uv, noise
//  [python]
def forward_prediction_heads(self, output, mask_features, attention_mask_target_size):
        decoder_output = self.decoder_norm(output)
        decoder_output = decoder_output.transpose(0, 1)
        outputs_class = self.class_embed(decoder_output)
        mask_embed = self.mask_embed(decoder_output)
        outputs_mask = torch.einsum("bqc,bchw->bqhw", mask_embed, mask_features)

        attention_mask = t.interpolate(
            outputs_mask, size=attention_mask_target_size, mode="bilinear", align_corners=False
        )

        # must use bool type
        # If a BoolTensor is provided, positions with ``True`` are not allowed to attend while ``False`` values will be unchanged.
        attention_mask = (
            attention_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1).flatten(0, 1) < 0.5
        ).bool()
        attention_mask = attention_mask.detach()

        return outputs_class, outputs_mask, attention_mask
//  [python]
def forward_prediction_heads(self, output, mask_features, attention_mask_target_size):
        decoder_output = self.decoder_norm(output)
        decoder_output = decoder_output.transpose(0, 1)
        outputs_class = self.class_embed(decoder_output)
        mask_embed = self.mask_embed(decoder_output)
        outputs_mask = torch.einsum("bqc,bchw->bqhw", mask_embed, mask_features)

        attention_mask = nn.functional.interpolate(
            outputs_mask, size=attention_mask_target_size, mode="bilinear", align_corners=False
        )

        # must use bool type
        # If a BoolTensor is provided, positions with ``True`` are not allowed to attend while ``False`` values will be unchanged.
        attention_mask = (
            attention_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1).flatten(0, 1) < 0.5
        ).bool()
        attention_mask = attention_mask.detach()

        return outputs_class, outputs_mask, attention_mask
//  [python]
def forward(self, inputs: Tuple[Tensor]) -> Tuple[Tensor]:
        """Forward function."""
        assert len(inputs) == self.num_levels

        # step 1: gather multi-level features by resize and average
        feats = []
        gather_size = inputs[self.refine_level].size()[2:]
        for i in range(self.num_levels):
            if i < self.refine_level:
                gathered = F.adaptive_max_pool2d(
                    inputs[i], output_size=gather_size)
            else:
                gathered = F.upsample(
                    inputs[i], size=gather_size, mode='nearest')
            feats.append(gathered)

        bsf = sum(feats) / len(feats)

        # step 2: refine gathered features
        if self.refine_type is not None:
            bsf = self.refine(bsf)

        # step 3: scatter refined features to multi-levels by a residual path
        outs = []
        for i in range(self.num_levels):
            out_size = inputs[i].size()[2:]
            if i < self.refine_level:
                residual = F.upsample(bsf, size=out_size, mode='nearest')
            else:
                residual = F.adaptive_max_pool2d(bsf, output_size=out_size)
            outs.append(residual + inputs[i])

        return tuple(outs)
//  [python]
def forward(self, inputs: Tuple[Tensor]) -> Tuple[Tensor]:
        """Forward function."""
        assert len(inputs) == self.num_levels

        # step 1: gather multi-level features by resize and average
        feats = []
        gather_size = inputs[self.refine_level].size()[2:]
        for i in range(self.num_levels):
            if i < self.refine_level:
                gathered = F.adaptive_max_pool2d(
                    inputs[i], output_size=gather_size)
            else:
                gathered = F.interpolate(
                    inputs[i], size=gather_size, mode='nearest')
            feats.append(gathered)

        bsf = sum(feats) / len(feats)

        # step 2: refine gathered features
        if self.refine_type is not None:
            bsf = self.refine(bsf)

        # step 3: scatter refined features to multi-levels by a residual path
        outs = []
        for i in range(self.num_levels):
            out_size = inputs[i].size()[2:]
            if i < self.refine_level:
                residual = F.interpolate(bsf, size=out_size, mode='nearest')
            else:
                residual = F.adaptive_max_pool2d(bsf, output_size=out_size)
            outs.append(residual + inputs[i])

        return tuple(outs)
//  [python]
def forward(self, x, **kwargs):

        if self.mode == 'pixel_shuffle':
            return F.pixel_shuffle(x, int(self.scale_factor))
        elif self.mode == 'nearest':
            return F.upsample(x, self.size, self.scale_factor, self.mode)
        else:
            return F.upsample(x, self.size, self.scale_factor, self.mode, self.align_corners)
//  [python]
def forward(self, x, **kwargs):

        if self.mode == 'pixel_shuffle':
            return F.pixel_shuffle(x, int(self.scale_factor))
        elif self.mode == 'nearest':
            return F.interpolate(x, self.size, self.scale_factor, self.mode, None)
        else:
            return F.interpolate(x, self.size, self.scale_factor, self.mode, self.align_corners)
//  [python]
def forward(self, x, **kwargs):

        if self.mode == 'pixel_shuffle':
            return F.pixel_shuffle(concate([x]*int(self.scale_factor)**2,axis=1), int(self.scale_factor))
        elif self.mode == 'nearest':
            return F.interpolate(x, self.size, self.scale_factor, self.mode, None)
        else:
            return F.interpolate(x, self.size, self.scale_factor, mode=self.mode, align_corners=self.align_corners,antialias=self.antialias)
//  [python]
def apply_masked_loss(loss, batch):
    # mask image is -1 to 1. we need to convert it to 0 to 1
    mask_image = batch["conditioning_images"].to(dtype=loss.dtype)[:, 0].unsqueeze(1)  # use R channel

    # resize to the same size as the loss
    mask_image = t.nn.functional.interpolate(mask_image, size=loss.shape[2:], mode="area")
    mask_image = mask_image / 2 + 0.5
    loss = loss * mask_image
    return loss
//  [python]
def apply_masked_loss(loss, batch):
    # mask image is -1 to 1. we need to convert it to 0 to 1
    mask_image = batch["conditioning_images"].to(dtype=loss.dtype)[:, 0].unsqueeze(1)  # use R channel

    # resize to the same size as the loss
    mask_image = torch.nn.functional.interpolate(mask_image, size=loss.shape[2:], mode="area")
    mask_image = mask_image / 2 + 0.5
    loss = loss * mask_image
    return loss
//  [python]
def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor):
        """
        Outputs feature maps of latter stages C_3 through C_5 in ResNet if `config.num_feature_levels > 1`, otherwise
        outputs feature maps of C_5.
        """
        # send pixel_values through the model to get list of feature maps
        features = self.model(pixel_values)

        out = []
        for feature_map in features:
            # downsample pixel_mask to match shape of corresponding feature_map
            mask = t.interpolate(pixel_mask[None].float(), size=feature_map.shape[-2:]).to(torch.bool)[0]
            out.append((feature_map, mask))
        return out
//  [python]
def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor):
        """
        Outputs feature maps of latter stages C_3 through C_5 in ResNet if `config.num_feature_levels > 1`, otherwise
        outputs feature maps of C_5.
        """
        # send pixel_values through the model to get list of feature maps
        features = self.model(pixel_values)

        out = []
        for feature_map in features:
            # downsample pixel_mask to match shape of corresponding feature_map
            mask = nn.functional.interpolate(pixel_mask[None].float(), size=feature_map.shape[-2:]).to(torch.bool)[0]
            out.append((feature_map, mask))
        return out
//  [python]
def forward(self, x, **kwargs):

        if self.mode == 'pixel_shuffle':
            batch_size, channels, in_depth, in_height, in_width = x.size()
            channels //= self.upscale_factor ** 3

            out_depth = in_depth * self.upscale_factor
            out_height = in_height * self.upscale_factor
            out_width = in_width * self.upscale_factor
            input_view = x.contiguous().view(
                batch_size, channels, self.upscale_factor, self.upscale_factor, self.upscale_factor,
                in_depth, in_height, in_width)
            shuffle_out = input_view.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
            return shuffle_out.view(batch_size, channels, out_depth, out_height, out_width)
        elif self.mode == 'nearest':
            return t.upsample(x, self.size, self.scale_factor, self.mode)
        else:
            return t.upsample(x, self.size, self.scale_factor, self.mode, self.align_corners)
//  [python]
def forward(self, x, **kwargs):

        if self.mode == 'pixel_shuffle':
            batch_size, channels, in_depth, in_height, in_width = x.size()
            channels //= self.upscale_factor ** 3

            out_depth = in_depth * self.upscale_factor
            out_height = in_height * self.upscale_factor
            out_width = in_width * self.upscale_factor
            input_view = x.contiguous().view(
                batch_size, channels, self.upscale_factor, self.upscale_factor, self.upscale_factor,
                in_depth, in_height, in_width)
            shuffle_out = input_view.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
            return shuffle_out.view(batch_size, channels, out_depth, out_height, out_width)
        elif self.mode == 'nearest':
            return F.interpolate(x, self.size, self.scale_factor, self.mode, None)
        else:
            return F.interpolate(x, self.size, self.scale_factor, self.mode, self.align_corners)
//  [python]
def __init__(self, in_channel, out_channel=1, depth=64, base_size=None, stage=None):
        super(PAA_d, self).__init__()
        self.conv1 = Conv2d(in_channel ,depth, 3)
        self.conv2 = Conv2d(depth, depth, 3)
        self.conv3 = Conv2d(depth, depth, 3)
        self.conv4 = Conv2d(depth, depth, 3)
        self.conv5 = Conv2d(depth, out_channel, 3, bn=False)
        
        self.base_size = base_size
        self.stage = stage
        
        if base_size is not None and stage is not None:
            self.stage_size = (base_size[0] // (2 ** stage), base_size[1] // (2 ** stage))
        else:
            self.stage_size = [None, None]

        self.Hattn = SelfAttention(depth, 'h', self.stage_size[0])
        self.Wattn = SelfAttention(depth, 'w', self.stage_size[1])

        self.upsample = lambda img, size: t.interpolate(img, size=size, mode='bilinear', align_corners=True)
//  [python]
def __init__(self, in_channel, out_channel=1, depth=64, base_size=None, stage=None):
        super(PAA_d, self).__init__()
        self.conv1 = Conv2d(in_channel ,depth, 3)
        self.conv2 = Conv2d(depth, depth, 3)
        self.conv3 = Conv2d(depth, depth, 3)
        self.conv4 = Conv2d(depth, depth, 3)
        self.conv5 = Conv2d(depth, out_channel, 3, bn=False)
        
        self.base_size = base_size
        self.stage = stage
        
        if base_size is not None and stage is not None:
            self.stage_size = (base_size[0] // (2 ** stage), base_size[1] // (2 ** stage))
        else:
            self.stage_size = [None, None]

        self.Hattn = SelfAttention(depth, 'h', self.stage_size[0])
        self.Wattn = SelfAttention(depth, 'w', self.stage_size[1])

        self.upsample = lambda img, size: F.interpolate(img, size=size, mode='bilinear', align_corners=True)
//  [python]
def forward(self, x):
        extra = 11
        x = t.interpolate(x, (x.shape[2] * 2, x.shape[3] * 2))
        x = torch.nn.functional.pad(x, (extra, extra, extra, extra))
        for layer in [self.conv1, self.conv2, self.conv3, self.conv4, self.conv5, self.conv6, self.conv7, self.conv8]:
            x = layer(x)
            x = torch.nn.functional.leaky_relu(x, 0.1)
        return x
//  [python]
def forward(self, x):
        extra = 11
        x = torch.nn.functional.interpolate(x, (x.shape[2] * 2, x.shape[3] * 2))
        x = torch.nn.functional.pad(x, (extra, extra, extra, extra))
        for layer in [self.conv1, self.conv2, self.conv3, self.conv4, self.conv5, self.conv6, self.conv7, self.conv8]:
            x = layer(x)
            x = torch.nn.functional.leaky_relu(x, 0.1)
        return x
//  [python]
def encode_vae_inpaint(vae, pixels, mask):
    assert mask.ndim == 3 and pixels.ndim == 4
    assert mask.shape[-1] == pixels.shape[-2]
    assert mask.shape[-2] == pixels.shape[-3]

    w = mask.round()[..., None]
    pixels = pixels * (1 - w) + 0.5 * w

    latent = vae.encode(pixels)
    B, C, H, W = latent.shape

    latent_mask = mask[:, None, :, :]
    latent_mask = t.nn.functional.interpolate(latent_mask, size=(H * 8, W * 8), mode="bilinear").round()
    latent_mask = t.nn.functional.max_pool2d(latent_mask, (8, 8)).round().to(latent)

    return latent, latent_mask
//  [python]
def encode_vae_inpaint(vae, pixels, mask):
    assert mask.ndim == 3 and pixels.ndim == 4
    assert mask.shape[-1] == pixels.shape[-2]
    assert mask.shape[-2] == pixels.shape[-3]

    w = mask.round()[..., None]
    pixels = pixels * (1 - w) + 0.5 * w

    latent = vae.encode(pixels)
    B, C, H, W = latent.shape

    latent_mask = mask[:, None, :, :]
    latent_mask = torch.nn.functional.interpolate(latent_mask, size=(H * 8, W * 8), mode="bilinear").round()
    latent_mask = torch.nn.functional.max_pool2d(latent_mask, (8, 8)).round().to(latent)

    return latent, latent_mask
//  [python]
def forward(self, x):
        if self.layer_type == 'none':
            return x
        elif self.layer_type == 'timepreserve':
            return t.upsample(x, scale_factor=(2, 1), mode='nearest')
        elif self.layer_type == 'half':
            return t.upsample(x, scale_factor=2, mode='nearest')
        else:
            raise RuntimeError('Got unexpected upsampletype %s, expected is [none, timepreserve, half]' % self.layer_type)
//  [python]
def forward(self, x):
        if self.layer_type == 'none':
            return x
        elif self.layer_type == 'timepreserve':
            return F.interpolate(x, scale_factor=(2, 1), mode='nearest')
        elif self.layer_type == 'half':
            return F.interpolate(x, scale_factor=2, mode='nearest')
        else:
            raise RuntimeError('Got unexpected upsampletype %s, expected is [none, timepreserve, half]' % self.layer_type)
//  [python]
def fasterrcnn_reshape_transform(x):
    target_size = x['pool'].size()[-2:]
    activations = []
    for key, value in x.items():
        activations.append(
            t.interpolate(
                torch.abs(value),
                target_size,
                mode='bilinear'))
    activations = torch.cat(activations, axis=1)
    return activations
//  [python]
def fasterrcnn_reshape_transform(x):
    target_size = x['pool'].size()[-2:]
    activations = []
    for key, value in x.items():
        activations.append(
            torch.nn.functional.interpolate(
                torch.abs(value),
                target_size,
                mode='bilinear'))
    activations = torch.cat(activations, axis=1)
    return activations
//  [python]
def forward(self, x):
        out_lr, in_hr = x
        # low resolution
        out_lr = self.m_lr(out_lr).sigmoid_()
        # super resolution
        out_lr_sr = t.interpolate(out_lr, mode='bilinear', size=in_hr.shape[2:], align_corners=False)
        # concat
        out_hr = torch.cat((in_hr, out_lr_sr), dim=1)
        # high resolution
        out_hr = self.m_hr(out_hr).sigmoid_()
        
        return out_lr, out_hr
//  [python]
def forward(self, x):
        out_lr, in_hr = x
        # low resolution
        out_lr = self.m_lr(out_lr).sigmoid_()
        # super resolution
        out_lr_sr = F.interpolate(out_lr, mode='bilinear', size=in_hr.shape[2:], align_corners=False)
        # concat
        out_hr = torch.cat((in_hr, out_lr_sr), dim=1)
        # high resolution
        out_hr = self.m_hr(out_hr).sigmoid_()
        
        return out_lr, out_hr
//  [python]
def resize(
        x: torch.Tensor,
        size: any or None = None,
        scale_factor=None,
        mode: str = "bicubic",
        align_corners: bool or None = False,
) -> torch.Tensor:
    if mode in ["bilinear", "bicubic"]:
        return t.interpolate(
            x,
            size=size,
            scale_factor=scale_factor,
            mode=mode,
            align_corners=align_corners,
        )
    elif mode in ["nearest", "area"]:
        return t.interpolate(x, size=size, scale_factor=scale_factor, mode=mode)
    else:
        raise NotImplementedError(f"resize(mode={mode}) not implemented.")
//  [python]
def resize(
        x: torch.Tensor,
        size: any or None = None,
        scale_factor=None,
        mode: str = "bicubic",
        align_corners: bool or None = False,
) -> torch.Tensor:
    if mode in ["bilinear", "bicubic"]:
        return F.interpolate(
            x,
            size=size,
            scale_factor=scale_factor,
            mode=mode,
            align_corners=align_corners,
        )
    elif mode in ["nearest", "area"]:
        return F.interpolate(x, size=size, scale_factor=scale_factor, mode=mode)
    else:
        raise NotImplementedError(f"resize(mode={mode}) not implemented.")
//  [python]
def forward(self, input_video):
        if len(input_video.shape) == 3:
            input_video = input_video.unsqueeze(1)
        input_video = torch.transpose(input_video, 1, 3)
        output_layer = F.relu(self.batch_norm1(self.conv1(input_video)))
        output_layer = F.relu(self.batch_norm2(self.conv2(output_layer)))
        output_layer = F.relu(self.batch_norm3(self.conv3(output_layer)))
        output_layer = F.relu(self.batch_norm4(self.conv4(output_layer)))
        output_layer = F.relu(self.batch_norm5(self.conv5(output_layer)))
        output_layer = F.relu(self.batch_norm6(self.conv6(output_layer)))
        output_layer = F.interpolate(output_layer, size=(298, 1), mode='nearest')
        return output_layer
//  [python]
def forward(self, input_video):
        if len(input_video.shape) == 3:
            input_video = input_video.unsqueeze(1)
        input_video = torch.transpose(input_video, 1, 3)
        output_layer = F.relu(self.batch_norm1(self.conv1(input_video)))
        output_layer = F.relu(self.batch_norm2(self.conv2(output_layer)))
        output_layer = F.relu(self.batch_norm3(self.conv3(output_layer)))
        output_layer = F.relu(self.batch_norm4(self.conv4(output_layer)))
        output_layer = F.relu(self.batch_norm5(self.conv5(output_layer)))
        output_layer = F.relu(self.batch_norm6(self.conv6(output_layer)))
        output_layer = nn.functional.interpolate(output_layer, size=(298, 1), mode='nearest')
        return output_layer
//  [python]
def get_depth_image(image):
    """
    Create depth map for SDXL depth control net.
    """
    from transformers import DPTFeatureExtractor, DPTForDepthEstimation

    depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda")
    feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-hybrid-midas")

    image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
    with torch.no_grad(), torch.autocast("cuda"):
        depth_map = depth_estimator(image).predicted_depth

    # The depth map is 384x384 by default, here we interpolate to the default output size.
    # Note that it will be resized to output image size later. May change the size here to avoid interpolate twice.
    depth_map = torch.nn.functional.interpolate(
        depth_map.unsqueeze(1),
        size=(1024, 1024),
        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 = (depth_map - depth_min) / (depth_max - depth_min)
    image = torch.cat([depth_map] * 3, dim=1)

    image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
    image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
    return image
//  [python]
def forward(self, x: torch.Tensor):
        resolution = x.shape[-2]
        height, width = x.shape[-2] // self.patch_size, x.shape[
            -1] // self.patch_size
        x = self.conv1(x)  # shape = [*, width, grid, grid]
        x = x.reshape(x.shape[0], x.shape[1],
                      -1)  # shape = [*, width, grid ** 2]
        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]

        if resolution != self.input_resolution:
            old_pe = self.positional_embedding[1:]
            patch_num = self.input_resolution // self.patch_size
            old_pe = old_pe.reshape(1, patch_num, patch_num,
                                    -1).permute(0, 3, 1, 2)
            new_pe = F.upsample_bilinear(
                old_pe, size=(height, width))
            new_pe = new_pe.permute(0, 2, 3, 1).reshape(height * width, -1)
            x = x + new_pe.to(x.dtype)
        else:
            x = x + self.positional_embedding[1:].to(x.dtype)
        x = self.ln_pre(x)

        x = x.permute(1, 0, 2)  # NLD -> LND
        x = self.transformer(x)
        x = x.permute(1, 0, 2)  # LND -> NLD

        bz, seq, hidden = x.shape
        x = x.transpose(1, 2).reshape(bz, hidden, height, width)

        return x
//  [python]
def forward(self, x: torch.Tensor):
        resolution = x.shape[-2]
        height, width = x.shape[-2] // self.patch_size, x.shape[
            -1] // self.patch_size
        x = self.conv1(x)  # shape = [*, width, grid, grid]
        x = x.reshape(x.shape[0], x.shape[1],
                      -1)  # shape = [*, width, grid ** 2]
        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]

        if resolution != self.input_resolution:
            old_pe = self.positional_embedding[1:]
            patch_num = self.input_resolution // self.patch_size
            old_pe = old_pe.reshape(1, patch_num, patch_num,
                                    -1).permute(0, 3, 1, 2)
            new_pe = F.interpolate(
                old_pe, size=(height, width), mode='bilinear')
            new_pe = new_pe.permute(0, 2, 3, 1).reshape(height * width, -1)
            x = x + new_pe.to(x.dtype)
        else:
            x = x + self.positional_embedding[1:].to(x.dtype)
        x = self.ln_pre(x)

        x = x.permute(1, 0, 2)  # NLD -> LND
        x = self.transformer(x)
        x = x.permute(1, 0, 2)  # LND -> NLD

        bz, seq, hidden = x.shape
        x = x.transpose(1, 2).reshape(bz, hidden, height, width)

        return x
//  [python]
def get_pose_map(self, items):
        pt_mats = torch.einsum('nj,jxy->nxy', self.lbs, items['cano2live_jnt_mats_woRoot'])
        live_pts = torch.einsum('nxy,ny->nx', pt_mats[..., :3, :3], self.init_points) + pt_mats[..., :3, 3]
        live_pos_map = torch.zeros_like(self.cano_smpl_map)
        live_pos_map[self.cano_smpl_mask] = live_pts
        live_pos_map = t.upsample(live_pos_map.permute(2, 0, 1)[None], scale_factor=0.5, mode='nearest')[0]
        live_pos_map = torch.cat(torch.split(live_pos_map, [512, 512], 2), 0)
        items.update({
            'smpl_pos_map': live_pos_map
        })
        return live_pos_map
//  [python]
def get_pose_map(self, items):
        pt_mats = torch.einsum('nj,jxy->nxy', self.lbs, items['cano2live_jnt_mats_woRoot'])
        live_pts = torch.einsum('nxy,ny->nx', pt_mats[..., :3, :3], self.init_points) + pt_mats[..., :3, 3]
        live_pos_map = torch.zeros_like(self.cano_smpl_map)
        live_pos_map[self.cano_smpl_mask] = live_pts
        live_pos_map = F.interpolate(live_pos_map.permute(2, 0, 1)[None], None, [0.5, 0.5], mode = 'nearest')[0]
        live_pos_map = torch.cat(torch.split(live_pos_map, [512, 512], 2), 0)
        items.update({
            'smpl_pos_map': live_pos_map
        })
        return live_pos_map
//  [python]
def forward(self, x):
        res = self.short_conv(F.avg_pool2d(x, kernel_size=2, stride=2))
        x1 = self.primary_conv(x)
        x2 = self.cheap_operation(x1)
        out = torch.cat([x1, x2], dim=1)
        return out[:, :self.out_chs, :, :] * F.upsample(
            self.gate_fn(res), size=(out.shape[-2], out.shape[-1]), mode='nearest')
//  [python]
def forward(self, x):
        res = self.short_conv(F.avg_pool2d(x, kernel_size=2, stride=2))
        x1 = self.primary_conv(x)
        x2 = self.cheap_operation(x1)
        out = torch.cat([x1, x2], dim=1)
        return out[:, :self.out_chs, :, :] * F.interpolate(
            self.gate_fn(res), size=(out.shape[-2], out.shape[-1]), mode='nearest')
//  [python]
def photometric_masked_loss(prediction: torch.Tensor,
                            target: torch.Tensor,
                            mask: Optional[torch.BoolTensor] = None,
                            multiplier: Optional[torch.Tensor] = None,
                            mode: str = 'l1',
                            eps: float = 1e-6,
                            drop_zeros: bool = False,
                            allow_different_size_of_gt_and_pred=False,
                            ) -> torch.Tensor:
    """
    Args:
        prediction: B x C x H x W
        target: B x C x H x W
        mask: B x 1 x h x w
        multiplier: B x 1 x h x w or B x C x H x W
        mode: l1 | l2
        eps: const for numerical stability

    Returns:
        loss
    """
    if allow_different_size_of_gt_and_pred:
        assert prediction.shape[:2] == target.shape[:2], \
            f'prediction.shape[:2] - {prediction.shape[:2]}, target.shape[:2] - {target.shape[:2]}'
        if prediction.shape != target.shape:
            target = t.interpolate(target, size=[prediction.shape[2], prediction.shape[3]])
    else:
        assert prediction.shape == target.shape

    if mode == 'l1':
        func = lambda x, y: torch.abs(x - y)
    elif mode == 'l2':
        func = lambda x, y: (x - y) ** 2
    else:
        raise ValueError(f'Unknown mode={mode}')

    if multiplier is not None and multiplier.shape[-2:] != target.shape[-2:]:
        multiplier = t.interpolate(multiplier.expand(-1, target.shape[1], -1, -1),
                                   size=target.shape[-2:], mode='bilinear')

    if mask is None:
        pixelwise = func(prediction, target)
        if multiplier is None:
            return pixelwise.mean()
        else:
            return torch.mean(pixelwise * multiplier)
    else:
        mask = t.interpolate(mask.float(), size=target.shape[-2:], mode='bilinear')
        pixelwise = func(prediction, target) * mask

        if drop_zeros:
            mask_not_zeros = ~torch.all(torch.isclose(pixelwise, torch.zeros_like(pixelwise)).bool(), dim=1, keepdim=True)
            mask_not_zeros = mask_not_zeros.detach()

            if multiplier is None:
                return torch.mean(
                    torch.sum(pixelwise, dim=[-1, -2]) / mask_not_zeros.sum(dim=[-1, -2]).add(eps)
                )
            else:
                return torch.mean(
                    torch.sum(pixelwise * multiplier, dim=[-1, -2]) / mask_not_zeros.sum(dim=[-1, -2]).add(eps)
                )
                
        if multiplier is None:
            return torch.mean(
                torch.sum(pixelwise, dim=[-1, -2]) / mask.sum(dim=[-1, -2]).add(eps)
            )
        else:
            return torch.mean(
                torch.sum(pixelwise * multiplier, dim=[-1, -2]) / mask.sum(dim=[-1, -2]).add(eps)
            )
//  [python]
def photometric_masked_loss(prediction: torch.Tensor,
                            target: torch.Tensor,
                            mask: Optional[torch.BoolTensor] = None,
                            multiplier: Optional[torch.Tensor] = None,
                            mode: str = 'l1',
                            eps: float = 1e-6,
                            drop_zeros: bool = False,
                            allow_different_size_of_gt_and_pred=False,
                            ) -> torch.Tensor:
    """
    Args:
        prediction: B x C x H x W
        target: B x C x H x W
        mask: B x 1 x h x w
        multiplier: B x 1 x h x w or B x C x H x W
        mode: l1 | l2
        eps: const for numerical stability

    Returns:
        loss
    """
    if allow_different_size_of_gt_and_pred:
        assert prediction.shape[:2] == target.shape[:2], \
            f'prediction.shape[:2] - {prediction.shape[:2]}, target.shape[:2] - {target.shape[:2]}'
        if prediction.shape != target.shape:
            target = F.interpolate(target, size=[prediction.shape[2], prediction.shape[3]])
    else:
        assert prediction.shape == target.shape

    if mode == 'l1':
        func = lambda x, y: torch.abs(x - y)
    elif mode == 'l2':
        func = lambda x, y: (x - y) ** 2
    else:
        raise ValueError(f'Unknown mode={mode}')

    if multiplier is not None and multiplier.shape[-2:] != target.shape[-2:]:
        multiplier = F.interpolate(multiplier.expand(-1, target.shape[1], -1, -1),
                                   size=target.shape[-2:], mode='bilinear')

    if mask is None:
        pixelwise = func(prediction, target)
        if multiplier is None:
            return pixelwise.mean()
        else:
            return torch.mean(pixelwise * multiplier)
    else:
        mask = F.interpolate(mask.float(), size=target.shape[-2:], mode='bilinear')
        pixelwise = func(prediction, target) * mask

        if drop_zeros:
            mask_not_zeros = ~torch.all(torch.isclose(pixelwise, torch.zeros_like(pixelwise)).bool(), dim=1, keepdim=True)
            mask_not_zeros = mask_not_zeros.detach()

            if multiplier is None:
                return torch.mean(
                    torch.sum(pixelwise, dim=[-1, -2]) / mask_not_zeros.sum(dim=[-1, -2]).add(eps)
                )
            else:
                return torch.mean(
                    torch.sum(pixelwise * multiplier, dim=[-1, -2]) / mask_not_zeros.sum(dim=[-1, -2]).add(eps)
                )
                
        if multiplier is None:
            return torch.mean(
                torch.sum(pixelwise, dim=[-1, -2]) / mask.sum(dim=[-1, -2]).add(eps)
            )
        else:
            return torch.mean(
                torch.sum(pixelwise * multiplier, dim=[-1, -2]) / mask.sum(dim=[-1, -2]).add(eps)
            )
//  [python]
def get_viewdir_feat(self, items):
        with torch.no_grad():
            pt_mats = torch.einsum('nj,jxy->nxy', self.lbs, items['cano2live_jnt_mats'])
            live_pts = torch.einsum('nxy,ny->nx', pt_mats[..., :3, :3], self.init_points) + pt_mats[..., :3, 3]
            live_nmls = torch.einsum('nxy,ny->nx', pt_mats[..., :3, :3], self.cano_nmls)
            cam_pos = -torch.matmul(torch.linalg.inv(items['extr'][:3, :3]), items['extr'][:3, 3])
            viewdirs = F.normalize(cam_pos[None] - live_pts, dim = -1, eps = 1e-3)
            if self.training:
                viewdirs += torch.randn(*viewdirs.shape).to(viewdirs) * 0.1
            viewdirs = F.normalize(viewdirs, dim = -1, eps = 1e-3)
            viewdirs = (live_nmls * viewdirs).sum(-1)

            viewdirs_map = torch.zeros(*self.cano_nml_map.shape[:2]).to(viewdirs)
            viewdirs_map[self.cano_smpl_mask] = viewdirs

            viewdirs_map = viewdirs_map[None, None]
            viewdirs_map = F.interpolate(viewdirs_map, scale_factor=0.5, mode='nearest')
            front_viewdirs, back_viewdirs = torch.split(viewdirs_map, [512, 512], -1)

        front_viewdirs = self.opt.get('weight_viewdirs', 1.) * self.viewdir_net(front_viewdirs)
        back_viewdirs = self.opt.get('weight_viewdirs', 1.) * self.viewdir_net(back_viewdirs)
        return front_viewdirs, back_viewdirs
//  [python]
def get_viewdir_feat(self, items):
        with torch.no_grad():
            pt_mats = torch.einsum('nj,jxy->nxy', self.lbs, items['cano2live_jnt_mats'])
            live_pts = torch.einsum('nxy,ny->nx', pt_mats[..., :3, :3], self.init_points) + pt_mats[..., :3, 3]
            live_nmls = torch.einsum('nxy,ny->nx', pt_mats[..., :3, :3], self.cano_nmls)
            cam_pos = -torch.matmul(torch.linalg.inv(items['extr'][:3, :3]), items['extr'][:3, 3])
            viewdirs = F.normalize(cam_pos[None] - live_pts, dim = -1, eps = 1e-3)
            if self.training:
                viewdirs += torch.randn(*viewdirs.shape).to(viewdirs) * 0.1
            viewdirs = F.normalize(viewdirs, dim = -1, eps = 1e-3)
            viewdirs = (live_nmls * viewdirs).sum(-1)

            viewdirs_map = torch.zeros(*self.cano_nml_map.shape[:2]).to(viewdirs)
            viewdirs_map[self.cano_smpl_mask] = viewdirs

            viewdirs_map = viewdirs_map[None, None]
            viewdirs_map = F.interpolate(viewdirs_map, None, 0.5, 'nearest')
            front_viewdirs, back_viewdirs = torch.split(viewdirs_map, [512, 512], -1)

        front_viewdirs = self.opt.get('weight_viewdirs', 1.) * self.viewdir_net(front_viewdirs)
        back_viewdirs = self.opt.get('weight_viewdirs', 1.) * self.viewdir_net(back_viewdirs)
        return front_viewdirs, back_viewdirs
//  [python]
def test_deterministic_interpolate_bilinear(self, device):
        input = torch.randn(1, 2, 4, 4, device=device, requires_grad=True)
        grad = None
        with DeterministicGuard(True):
            for _ in range(5):
                res = t.interpolate(
                    input,
                    size=12,
                    mode='bilinear',
                    align_corners=False)
                res.backward(torch.ones_like(res))
                if grad is None:
                    grad = input.grad
                else:
                    self.assertEqual(grad, input.grad, atol=0, rtol=0)
                input.grad = None
//  [python]
def test_deterministic_interpolate_bilinear(self, device):
        input = torch.randn(1, 2, 4, 4, device=device, requires_grad=True)
        grad = None
        with DeterministicGuard(True):
            for _ in range(5):
                res = torch.nn.functional.interpolate(
                    input,
                    size=12,
                    mode='bilinear',
                    align_corners=False)
                res.backward(torch.ones_like(res))
                if grad is None:
                    grad = input.grad
                else:
                    self.assertEqual(grad, input.grad, atol=0, rtol=0)
                input.grad = None
//  [python]
def forward(self, x):
        
        input_spatial_dim = x.size()[2:]
        
        x = self.resnet101_8s(x)
        
        x = nn.functional.upsample(input=x, size=input_spatial_dim, mode='bilinear')
        
        return x
//  [python]
def forward(self, x):
        
        input_spatial_dim = x.size()[2:]
        
        x = self.resnet18_16s(x)
        
        x = nn.functional.upsample(input=x, size=input_spatial_dim, mode='bilinear')
        
        return x
//  [python]
def forward(self, x, feature_alignment=False):
        
        input_spatial_dim = x.size()[2:]
        
        if feature_alignment:
            
            x = adjust_input_image_size_for_proper_feature_alignment(x, output_stride=8)
        
        x = self.resnet18_8s(x)
        
        x = t.interpolate(input=x, size=input_spatial_dim, mode='bilinear')
        
        #x = nn.functional.interpolate(input=x, size=input_spatial_dim, mode='bilinear')#, align_corners=False)
        
        return x
//  [python]
def forward(self, x, feature_alignment=False):
        
        input_spatial_dim = x.size()[2:]
        
        if feature_alignment:
            
            x = adjust_input_image_size_for_proper_feature_alignment(x, output_stride=8)
        
        x = self.resnet18_8s(x)
        
        x = nn.functional.interpolate(input=x, size=input_spatial_dim, mode='bilinear')
        
        #x = nn.functional.interpolate(input=x, size=input_spatial_dim, mode='bilinear')#, align_corners=False)
        
        return x
//  [python]
def forward(self, x):
        
        input_spatial_dim = x.size()[2:]
        
        x = self.resnet18_32s(x)
        
        x = nn.functional.upsample(input=x, size=input_spatial_dim, mode='bilinear')
        
        return x
//  [python]
def forward(self, x):
        
        input_spatial_dim = x.size()[2:]
        
        x = self.resnet34_32s(x)
        
        x = nn.functional.upsample(input=x, size=input_spatial_dim, mode='bilinear')
        
        return x
//  [python]
def forward(self, x):
        
        input_spatial_dim = x.size()[2:]
        
        x = self.resnet34_16s(x)
        
        x = nn.functional.upsample(input=x, size=input_spatial_dim, mode='bilinear')
        
        return x
//  [python]
def forward(self, x, feature_alignment=False):
        
        input_spatial_dim = x.size()[2:]
        
        if feature_alignment:
            
            x = adjust_input_image_size_for_proper_feature_alignment(x, output_stride=8)
        
        x = self.resnet34_8s(x)
        
        x = t.interpolate(input=x, size=input_spatial_dim, mode='bilinear', align_corners=False)
        
        return x
//  [python]
def forward(self, x, feature_alignment=False):
        
        input_spatial_dim = x.size()[2:]
        
        if feature_alignment:
            
            x = adjust_input_image_size_for_proper_feature_alignment(x, output_stride=8)
        
        x = self.resnet34_8s(x)
        
        x = nn.functional.interpolate(input=x, size=input_spatial_dim, mode='bilinear', align_corners=False)
        
        return x
//  [python]
def forward(self, x):
        
        input_spatial_dim = x.size()[2:]
        
        x = self.resnet50_32s(x)
        
        x = nn.functional.upsample(input=x, size=input_spatial_dim, mode='bilinear')
        
        return x
//  [python]
def test_upsample_nearest1d_meta(self):
        # TODO: this test should be triggered by test_nn.py but right
        # now meta is not enabled (and even if it was, we are probably
        # missing too many meta functions to get through the test unmolested)

        # NB: Can't make the exponent too big, or it will overflow
        # signed 64-bit integer
        x = torch.empty(2 * 10 ** 8, 3, 2 * 10 ** 8, device='meta')
        z = t.nn.functional.interpolate(x, scale_factor=2)
        self.assertEqual(z.size(), (2 * 10 ** 8, 3, 4 * 10 ** 8))
        self.assertRaises(RuntimeError, lambda: z[0][0][0].item())

        # TODO: the out tests cannot be triggered by test_nn.py because
        # we don't actually do out= arguments for nn functions, so there
        # is no public API by which to get the out version

        # interpolate doesn't seem to support out=
        # (not sure why passing None here doesn't work? How strange...)
        z = torch.empty(0, device='meta')
        torch._C._nn.upsample_nearest1d(x, (4 * 10 ** 8,), 2, out=z)
        self.assertEqual(z.size(), (2 * 10 ** 8, 3, 4 * 10 ** 8))
        self.assertRaises(RuntimeError, lambda: z[0][0][0].item())
//  [python]
def test_upsample_nearest1d_meta(self):
        # TODO: this test should be triggered by test_nn.py but right
        # now meta is not enabled (and even if it was, we are probably
        # missing too many meta functions to get through the test unmolested)

        # NB: Can't make the exponent too big, or it will overflow
        # signed 64-bit integer
        x = torch.empty(2 * 10 ** 8, 3, 2 * 10 ** 8, device='meta')
        z = torch.nn.functional.interpolate(x, scale_factor=2)
        self.assertEqual(z.size(), (2 * 10 ** 8, 3, 4 * 10 ** 8))
        self.assertRaises(RuntimeError, lambda: z[0][0][0].item())

        # TODO: the out tests cannot be triggered by test_nn.py because
        # we don't actually do out= arguments for nn functions, so there
        # is no public API by which to get the out version

        # interpolate doesn't seem to support out=
        # (not sure why passing None here doesn't work? How strange...)
        z = torch.empty(0, device='meta')
        torch._C._nn.upsample_nearest1d(x, (4 * 10 ** 8,), 2, out=z)
        self.assertEqual(z.size(), (2 * 10 ** 8, 3, 4 * 10 ** 8))
        self.assertRaises(RuntimeError, lambda: z[0][0][0].item())
//  [python]
def forward(self, x):
        
        input_spatial_dim = x.size()[2:]
        
        x = self.resnet50_8s(x)
        
        x = nn.functional.upsample(input=x, size=input_spatial_dim, mode='bilinear')
        
        return x