jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def interpolate_pos_embed(model, checkpoint_model):
    if "pos_embed" in checkpoint_model:
        pos_embed_checkpoint = checkpoint_model["pos_embed"]
        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 forward(self, x, gst=None):
        x = self.token_emb(x)
        feats = [x]

        if gst is not None:
            gst = self.gst_emb(gst)
            gst = rearrange(gst, "b c -> b c 1")
            gst = F.upsample(gst, size=x.shape[1])
            gst = rearrange(gst, "b c t -> b t c")
            feats.append(gst)

        x = torch.cat(feats, dim=-1)

        for i, conv in enumerate(self.conv_layer):
            if i != 0:
                x = conv(x) + x
            else:
                x = conv(x)

        x = self.proj(x)
        x = x.squeeze(-1)

        if self.f0_pred == "mean":
            x = torch.sigmoid(x)
        elif self.f0_pred == "argmax":
            x = torch.softmax(x, dim=-1)
        else:
            raise NotImplementedError
        return x
//  [python]
def forward(self, x, gst=None):
        x = self.token_emb(x)
        feats = [x]

        if gst is not None:
            gst = self.gst_emb(gst)
            gst = rearrange(gst, "b c -> b c 1")
            gst = F.interpolate(gst, x.shape[1])
            gst = rearrange(gst, "b c t -> b t c")
            feats.append(gst)

        x = torch.cat(feats, dim=-1)

        for i, conv in enumerate(self.conv_layer):
            if i != 0:
                x = conv(x) + x
            else:
                x = conv(x)

        x = self.proj(x)
        x = x.squeeze(-1)

        if self.f0_pred == "mean":
            x = torch.sigmoid(x)
        elif self.f0_pred == "argmax":
            x = torch.softmax(x, dim=-1)
        else:
            raise NotImplementedError
        return x
//  [python]
def process_batch(detections, labels, iouv, pred_masks=None, gt_masks=None, overlap=False, masks=False):
    """
    Return correct prediction matrix
    Arguments:
        detections (array[N, 6]), x1, y1, x2, y2, conf, class
        labels (array[M, 5]), class, x1, y1, x2, y2
    Returns:
        correct (array[N, 10]), for 10 IoU levels
    """
    if masks:
        if overlap:
            nl = len(labels)
            index = torch.arange(nl, device=gt_masks.device).view(nl, 1, 1) + 1
            gt_masks = gt_masks.repeat(nl, 1, 1)  # shape(1,640,640) -> (n,640,640)
            gt_masks = torch.where(gt_masks == index, 1.0, 0.0)
        if gt_masks.shape[1:] != pred_masks.shape[1:]:
            gt_masks = t.interpolate(gt_masks[None], pred_masks.shape[1:], mode="bilinear", align_corners=False)[0]
            gt_masks = gt_masks.gt_(0.5)
        iou = mask_iou(gt_masks.view(gt_masks.shape[0], -1), pred_masks.view(pred_masks.shape[0], -1))
    else:  # boxes
        iou = box_iou(labels[:, 1:], detections[:, :4])

    correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool)
    correct_class = labels[:, 0:1] == detections[:, 5]
    for i in range(len(iouv)):
        x = torch.where((iou >= iouv[i]) & correct_class)  # IoU > threshold and classes match
        if x[0].shape[0]:
            matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()  # [label, detect, iou]
            if x[0].shape[0] > 1:
                matches = matches[matches[:, 2].argsort()[::-1]]
                matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
                # matches = matches[matches[:, 2].argsort()[::-1]]
                matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
            correct[matches[:, 1].astype(int), i] = True
    return torch.tensor(correct, dtype=torch.bool, device=iouv.device)
//  [python]
def process_batch(detections, labels, iouv, pred_masks=None, gt_masks=None, overlap=False, masks=False):
    """
    Return correct prediction matrix
    Arguments:
        detections (array[N, 6]), x1, y1, x2, y2, conf, class
        labels (array[M, 5]), class, x1, y1, x2, y2
    Returns:
        correct (array[N, 10]), for 10 IoU levels
    """
    if masks:
        if overlap:
            nl = len(labels)
            index = torch.arange(nl, device=gt_masks.device).view(nl, 1, 1) + 1
            gt_masks = gt_masks.repeat(nl, 1, 1)  # shape(1,640,640) -> (n,640,640)
            gt_masks = torch.where(gt_masks == index, 1.0, 0.0)
        if gt_masks.shape[1:] != pred_masks.shape[1:]:
            gt_masks = F.interpolate(gt_masks[None], pred_masks.shape[1:], mode="bilinear", align_corners=False)[0]
            gt_masks = gt_masks.gt_(0.5)
        iou = mask_iou(gt_masks.view(gt_masks.shape[0], -1), pred_masks.view(pred_masks.shape[0], -1))
    else:  # boxes
        iou = box_iou(labels[:, 1:], detections[:, :4])

    correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool)
    correct_class = labels[:, 0:1] == detections[:, 5]
    for i in range(len(iouv)):
        x = torch.where((iou >= iouv[i]) & correct_class)  # IoU > threshold and classes match
        if x[0].shape[0]:
            matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()  # [label, detect, iou]
            if x[0].shape[0] > 1:
                matches = matches[matches[:, 2].argsort()[::-1]]
                matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
                # matches = matches[matches[:, 2].argsort()[::-1]]
                matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
            correct[matches[:, 1].astype(int), i] = True
    return torch.tensor(correct, dtype=torch.bool, device=iouv.device)
//  [python]
def forward(self, x, timesteps, low_res=None, **kwargs):
        _, _, new_height, new_width = x.shape
        upsampled = t.interpolate(low_res, (new_height, new_width), mode="bilinear")
        x = th.cat([x, upsampled], dim=1)
        return super().forward(x, timesteps, **kwargs)
//  [python]
def forward(self, x, timesteps, low_res=None, **kwargs):
        _, _, new_height, new_width = x.shape
        upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear")
        x = th.cat([x, upsampled], dim=1)
        return super().forward(x, timesteps, **kwargs)
//  [python]
def __getitem__(self, i):
        x = self.km[i]
        x = x.split(" ")
        x = list(map(int, x))
        x = torch.LongTensor(x)

        gst = parse_style(self.tsv[i], self.gst)
        gst = self.gst2id[gst]
        spk = parse_speaker(self.tsv[i], self.spk)
        spk = self.spk2id[spk]

        f0_raw = self._load_f0(self.tsv[i])
        f0 = self._preprocess_f0(f0_raw.clone(), spk)

        f0 = t.nn.functional.interpolate(f0.unsqueeze(0).unsqueeze(0), x.shape[0])[0, 0]
        f0_raw = t.nn.functional.interpolate(f0_raw.unsqueeze(0).unsqueeze(0), x.shape[0])[0, 0]

        f0 = freq2bin(f0, f0_min=self.f0_min, f0_max=self.f0_max, bins=self.f0_bins)
        f0 = F.one_hot(f0.long(), num_classes=len(self.f0_bins) + 1).float()
        if self.f0_smoothing > 0:
            f0 = torch.tensor(
                gaussian_filter1d(f0.float().numpy(), sigma=self.f0_smoothing)
            )
        return x, f0, f0_raw, spk, gst
//  [python]
def __getitem__(self, i):
        x = self.km[i]
        x = x.split(" ")
        x = list(map(int, x))
        x = torch.LongTensor(x)

        gst = parse_style(self.tsv[i], self.gst)
        gst = self.gst2id[gst]
        spk = parse_speaker(self.tsv[i], self.spk)
        spk = self.spk2id[spk]

        f0_raw = self._load_f0(self.tsv[i])
        f0 = self._preprocess_f0(f0_raw.clone(), spk)

        f0 = F.interpolate(f0.unsqueeze(0).unsqueeze(0), x.shape[0])[0, 0]
        f0_raw = F.interpolate(f0_raw.unsqueeze(0).unsqueeze(0), x.shape[0])[0, 0]

        f0 = freq2bin(f0, f0_min=self.f0_min, f0_max=self.f0_max, bins=self.f0_bins)
        f0 = F.one_hot(f0.long(), num_classes=len(self.f0_bins) + 1).float()
        if self.f0_smoothing > 0:
            f0 = torch.tensor(
                gaussian_filter1d(f0.float().numpy(), sigma=self.f0_smoothing)
            )
        return x, f0, f0_raw, spk, gst
//  [python]
def forward(self, x):
        c1_ = self.conv1(x)
        c1 = self.conv2(c1_)
        v_max = F.max_pool2d(c1, kernel_size=7, stride=3)
        c3 = self.conv3(v_max)
        c3 = F.upsample_bilinear(c3, size=(x.size(2), x.size(3)))
        cf = self.conv_f(c1_)
        c4 = self.conv4(c3 + cf)
        m = self.sigmoid(c4)
        return x * m
//  [python]
def forward(self, x):
        c1_ = self.conv1(x)
        c1 = self.conv2(c1_)
        v_max = F.max_pool2d(c1, kernel_size=7, stride=3)
        c3 = self.conv3(v_max)
        c3 = F.interpolate(
            c3, (x.size(2), x.size(3)), mode="bilinear", align_corners=False
        )
        cf = self.conv_f(c1_)
        c4 = self.conv4(c3 + cf)
        m = self.sigmoid(c4)
        return x * m
//  [python]
def forward(self, lqs):
        """Forward function for BasicVSR++.

        Args:
            lqs (tensor): Input low quality (LQ) sequence with
                shape (n, t, c, h, w).

        Returns:
            Tensor: Output HR sequence with shape (n, t, c, 4h, 4w).
        """

        n, t, c, h, w = lqs.size()

        # whether to cache the features in CPU
        self.cpu_cache = True if t > self.cpu_cache_length else False

        if self.is_low_res_input:
            lqs_downsample = lqs.clone()
        else:
            lqs_downsample = t.interpolate(
                lqs.view(-1, c, h, w), scale_factor=0.25, mode='bicubic').view(n, t, c, h // 4, w // 4)

        # check whether the input is an extended sequence
        self.check_if_mirror_extended(lqs)

        feats = {}
        # compute spatial features
        if self.cpu_cache:
            feats['spatial'] = []
            for i in range(0, t):
                feat = self.feat_extract(lqs[:, i, :, :, :]).cpu()
                feats['spatial'].append(feat)
                torch.cuda.empty_cache()
        else:
            feats_ = self.feat_extract(lqs.view(-1, c, h, w))
            h, w = feats_.shape[2:]
            feats_ = feats_.view(n, t, -1, h, w)
            feats['spatial'] = [feats_[:, i, :, :, :] for i in range(0, t)]

        # compute optical flow using the low-res inputs
        assert lqs_downsample.size(3) >= 64 and lqs_downsample.size(4) >= 64, (
            'The height and width of low-res inputs must be at least 64, '
            f'but got {h} and {w}.')
        flows_forward, flows_backward = self.compute_flow(lqs_downsample)

        # feature propgation
        for iter_ in [1, 2]:
            for direction in ['backward', 'forward']:
                module = f'{direction}_{iter_}'

                feats[module] = []

                if direction == 'backward':
                    flows = flows_backward
                elif flows_forward is not None:
                    flows = flows_forward
                else:
                    flows = flows_backward.flip(1)

                feats = self.propagate(feats, flows, module)
                if self.cpu_cache:
                    del flows
                    torch.cuda.empty_cache()

        return self.upsample(lqs, feats)
//  [python]
def forward(self, lqs):
        """Forward function for BasicVSR++.

        Args:
            lqs (tensor): Input low quality (LQ) sequence with
                shape (n, t, c, h, w).

        Returns:
            Tensor: Output HR sequence with shape (n, t, c, 4h, 4w).
        """

        n, t, c, h, w = lqs.size()

        # whether to cache the features in CPU
        self.cpu_cache = True if t > self.cpu_cache_length else False

        if self.is_low_res_input:
            lqs_downsample = lqs.clone()
        else:
            lqs_downsample = F.interpolate(
                lqs.view(-1, c, h, w), scale_factor=0.25, mode='bicubic').view(n, t, c, h // 4, w // 4)

        # check whether the input is an extended sequence
        self.check_if_mirror_extended(lqs)

        feats = {}
        # compute spatial features
        if self.cpu_cache:
            feats['spatial'] = []
            for i in range(0, t):
                feat = self.feat_extract(lqs[:, i, :, :, :]).cpu()
                feats['spatial'].append(feat)
                torch.cuda.empty_cache()
        else:
            feats_ = self.feat_extract(lqs.view(-1, c, h, w))
            h, w = feats_.shape[2:]
            feats_ = feats_.view(n, t, -1, h, w)
            feats['spatial'] = [feats_[:, i, :, :, :] for i in range(0, t)]

        # compute optical flow using the low-res inputs
        assert lqs_downsample.size(3) >= 64 and lqs_downsample.size(4) >= 64, (
            'The height and width of low-res inputs must be at least 64, '
            f'but got {h} and {w}.')
        flows_forward, flows_backward = self.compute_flow(lqs_downsample)

        # feature propgation
        for iter_ in [1, 2]:
            for direction in ['backward', 'forward']:
                module = f'{direction}_{iter_}'

                feats[module] = []

                if direction == 'backward':
                    flows = flows_backward
                elif flows_forward is not None:
                    flows = flows_forward
                else:
                    flows = flows_backward.flip(1)

                feats = self.propagate(feats, flows, module)
                if self.cpu_cache:
                    del flows
                    torch.cuda.empty_cache()

        return self.upsample(lqs, feats)
//  [python]
def encode(self, input_tensor):
        """
        Encode input_tensor into feature maps, can be overridden in derived classes
        Default: nearest downsampling of 2**5 = 32 times
        """
        h, w = input_tensor.size()[-2:]
        sh, sw = h // 2**self.scale_ratio, w // 2**self.scale_ratio
        x = F.upsample(input_tensor, size=(sh, sw), mode='nearest')
        return self.fc(x)
//  [python]
def encode(self, input_tensor):
        """
        Encode input_tensor into feature maps, can be overridden in derived classes
        Default: nearest downsampling of 2**5 = 32 times
        """
        h, w = input_tensor.size()[-2:]
        sh, sw = h // 2**self.scale_ratio, w // 2**self.scale_ratio
        x = F.interpolate(input_tensor, size=(sh, sw))
        return self.fc(x)
//  [python]
def forward_single(self, x, img_metas):
        """"Forward function for a single feature level.

        Args:
            x (Tensor): Input feature from backbone's single stage, shape
                [bs, c, h, w].
            img_metas (list[dict]): List of image information.

        Returns:
            all_cls_scores (Tensor): Outputs from the classification head,
                shape [nb_dec, bs, num_query, cls_out_channels]. Note
                cls_out_channels should includes background.
            all_bbox_preds (Tensor): Sigmoid outputs from the regression
                head with normalized coordinate format (cx, cy, w, h).
                Shape [nb_dec, bs, num_query, 4].
        """
        # construct binary masks which used for the transformer.
        # NOTE following the official DETR repo, non-zero values representing
        # ignored positions, while zero values means valid positions.
        batch_size = x.size(0)
        input_img_h, input_img_w = img_metas[0]['batch_input_shape']
        masks = x.new_ones((batch_size, input_img_h, input_img_w))
        for img_id in range(batch_size):
            img_h, img_w, _ = img_metas[img_id]['img_shape']
            masks[img_id, :img_h, :img_w] = 0

        x = self.input_proj(x)
        # interpolate masks to have the same spatial shape with x
        masks = t.interpolate(
            masks.unsqueeze(1), size=x.shape[-2:]).to(torch.bool).squeeze(1)
        # position encoding
        pos_embed = self.positional_encoding(masks)  # [bs, embed_dim, h, w]
        # outs_dec: [nb_dec, bs, num_query, embed_dim]
        outs_dec, _ = self.transformer(x, masks, self.query_embedding.weight,
                                       pos_embed)

        all_cls_scores = self.fc_cls(outs_dec)
        all_bbox_preds = self.fc_reg(self.activate(
            self.reg_ffn(outs_dec))).sigmoid()
        return all_cls_scores, all_bbox_preds
//  [python]
def forward_single(self, x, img_metas):
        """"Forward function for a single feature level.

        Args:
            x (Tensor): Input feature from backbone's single stage, shape
                [bs, c, h, w].
            img_metas (list[dict]): List of image information.

        Returns:
            all_cls_scores (Tensor): Outputs from the classification head,
                shape [nb_dec, bs, num_query, cls_out_channels]. Note
                cls_out_channels should includes background.
            all_bbox_preds (Tensor): Sigmoid outputs from the regression
                head with normalized coordinate format (cx, cy, w, h).
                Shape [nb_dec, bs, num_query, 4].
        """
        # construct binary masks which used for the transformer.
        # NOTE following the official DETR repo, non-zero values representing
        # ignored positions, while zero values means valid positions.
        batch_size = x.size(0)
        input_img_h, input_img_w = img_metas[0]['batch_input_shape']
        masks = x.new_ones((batch_size, input_img_h, input_img_w))
        for img_id in range(batch_size):
            img_h, img_w, _ = img_metas[img_id]['img_shape']
            masks[img_id, :img_h, :img_w] = 0

        x = self.input_proj(x)
        # interpolate masks to have the same spatial shape with x
        masks = F.interpolate(
            masks.unsqueeze(1), size=x.shape[-2:]).to(torch.bool).squeeze(1)
        # position encoding
        pos_embed = self.positional_encoding(masks)  # [bs, embed_dim, h, w]
        # outs_dec: [nb_dec, bs, num_query, embed_dim]
        outs_dec, _ = self.transformer(x, masks, self.query_embedding.weight,
                                       pos_embed)

        all_cls_scores = self.fc_cls(outs_dec)
        all_bbox_preds = self.fc_reg(self.activate(
            self.reg_ffn(outs_dec))).sigmoid()
        return all_cls_scores, all_bbox_preds
//  [python]
def forward(self, x):
        x1 = self.encoder(x)
        x2 = self.down(x1)
        high = x1 - t.interpolate(x2, size=x.size()
                                  [-2:], mode='bilinear', align_corners=True)
        for i in range(5):
            x2 = self.decoder_low(x2)
        x3 = x2
        high1 = self.decoder_high(high)
        x4 = t.interpolate(
            x3, size=x.size()[-2:], mode='bilinear', align_corners=True)
        return self.alise(self.att(self.alise2(torch.cat([x4, high1], dim=1)))) + x
//  [python]
def forward(self, x):
        x1 = self.encoder(x)
        x2 = self.down(x1)
        high = x1 - F.interpolate(x2, size=x.size()
                                  [-2:], mode='bilinear', align_corners=True)
        for i in range(5):
            x2 = self.decoder_low(x2)
        x3 = x2
        high1 = self.decoder_high(high)
        x4 = F.interpolate(
            x3, size=x.size()[-2:], mode='bilinear', align_corners=True)
        return self.alise(self.att(self.alise2(torch.cat([x4, high1], dim=1)))) + x
//  [python]
def forward(self, x):
        if self.num_branches == 1:
            return [self.branches[0](x[0])]

        for i in range(self.num_branches):
            x[i] = self.branches[i](x[i])

        x_fuse = []
        for i in range(len(self.fuse_layers)):
            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
            for j in range(1, self.num_branches):
                if i == j:
                    y = y + x[j]
                elif j > i:
                    width_output = x[i].shape[-1]
                    height_output = x[i].shape[-2]
                    y = y + t.interpolate(
                        self.fuse_layers[i][j](x[j]),
                        size=[height_output, width_output],
                        mode="bilinear",
                        align_corners=True,
                    )
                else:
                    y = y + self.fuse_layers[i][j](x[j])
            x_fuse.append(self.relu(y))

        return x_fuse
//  [python]
def forward(self, x):
        if self.num_branches == 1:
            return [self.branches[0](x[0])]

        for i in range(self.num_branches):
            x[i] = self.branches[i](x[i])

        x_fuse = []
        for i in range(len(self.fuse_layers)):
            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
            for j in range(1, self.num_branches):
                if i == j:
                    y = y + x[j]
                elif j > i:
                    width_output = x[i].shape[-1]
                    height_output = x[i].shape[-2]
                    y = y + F.interpolate(
                        self.fuse_layers[i][j](x[j]),
                        size=[height_output, width_output],
                        mode="bilinear",
                        align_corners=True,
                    )
                else:
                    y = y + self.fuse_layers[i][j](x[j])
            x_fuse.append(self.relu(y))

        return x_fuse
//  [python]
def forward(self, f16, f8, f4, hidden_state, memory_readout, h_out=True):
        batch_size, num_objects = memory_readout.shape[:2]

        if self.hidden_update is not None:
            g16 = self.fuser(f16, torch.cat([memory_readout, hidden_state], 2))
        else:
            g16 = self.fuser(f16, memory_readout)

        g8 = self.up_16_8(f8, g16)
        g4 = self.up_8_4(f4, g8)
        logits = self.pred(F.relu(g4.flatten(start_dim=0, end_dim=1)))

        if h_out and self.hidden_update is not None:
            g4 = torch.cat([g4, logits.view(batch_size, num_objects, 1, *logits.shape[-2:])], 2)
            hidden_state = self.hidden_update([g16, g8, g4], hidden_state)
        else:
            hidden_state = None
        
        logits = F.upsample(logits, scale_factor=4, mode='bilinear')
        logits = logits.view(batch_size, num_objects, *logits.shape[-2:])

        return hidden_state, logits
//  [python]
def forward(self, f16, f8, f4, hidden_state, memory_readout, h_out=True):
        batch_size, num_objects = memory_readout.shape[:2]

        if self.hidden_update is not None:
            g16 = self.fuser(f16, torch.cat([memory_readout, hidden_state], 2))
        else:
            g16 = self.fuser(f16, memory_readout)

        g8 = self.up_16_8(f8, g16)
        g4 = self.up_8_4(f4, g8)
        logits = self.pred(F.relu(g4.flatten(start_dim=0, end_dim=1)))

        if h_out and self.hidden_update is not None:
            g4 = torch.cat([g4, logits.view(batch_size, num_objects, 1, *logits.shape[-2:])], 2)
            hidden_state = self.hidden_update([g16, g8, g4], hidden_state)
        else:
            hidden_state = None
        
        logits = F.interpolate(logits, scale_factor=4, mode='bilinear', align_corners=False)
        logits = logits.view(batch_size, num_objects, *logits.shape[-2:])

        return hidden_state, logits
//  [python]
def forward(
        self, x, time, text_embed=None, text_mask=None, x_low_res=None, time_low_res=None,
    ):
        if self.low_res_cond:
            assert x_low_res is not None, 'x_low_res cannot be None'
        else:
            assert x_low_res is None, 'x_low_res cannot be presented'
        if self.noise_cond_aug:
            assert time_low_res is not None, 'time_low_res cannot be None when training with noise conditioning aug'
        else:
            assert time_low_res is None, 'time_low_res cannot be presented'
        # Concatenating low resolution images
        if x_low_res is not None:
            if x_low_res.shape != x.shape:
                # Upscale if not done in the trainer
                _, _, new_height, new_width = x.shape
                x_low_res = F.interpolate(x_low_res, (new_height, new_width), mode="bicubic")
            x = torch.cat([x, x_low_res], dim=1)
        batch_size, device = x.shape[0], x.device

        if x.dtype != time.dtype or time.dtype != text_embed.dtype:
            dtype = text_embed.dtype
            x = x.to(dtype=dtype)
            time = time.to(dtype=dtype)
            if x_low_res is not None:
                x_low_res = x_low_res.to(dtype=dtype)
            if time_low_res is not None:
                time_low_res = time_low_res.to(dtype=dtype)
        # Time Conditioning
        t = self.time_embed(time)
        # Add lowres time conditioning
        if self.noise_cond_aug:
            lowres_t = self.lowres_time_embed(time_low_res)
            t += lowres_t
        # Text Conditioning
        text_cond = self.text_to_cond(text_embed)

        # Context Embedding
        # TODO We may want to concat time token here
        if self.use_null_token:
            # Null Context (Helpful when text_embed is drop)
            null_context = self.null_text_embedding.repeat(batch_size, 1, 1)
            context_emb = torch.cat([text_cond, null_context], dim=1)
            context_mask = torch.cat([text_mask, torch.ones(batch_size, 1).to(device)], dim=1)
        else:
            context_emb = text_cond
            context_mask = text_mask

        # Add pooled text embeddings to the diffusion timestep
        # TODO We may only want to calculated the pooled feature based on text token length
        if self.feature_pooling_type == 'mean':
            pooled_text_cond = text_cond.mean(dim=-2)
        elif self.feature_pooling_type == 'attention':
            pooled_text_cond = self.attention_pooling(text_embed)
        text_hiddens = self.to_text_non_attn_cond(pooled_text_cond)
        t += text_hiddens

        h = x
        hs = []
        # UNet Forward
        for module in self.input_blocks:
            h = module(h, t, context_emb, context_mask)
            hs.append(h)
        h = self.middle_block(h, t, context_emb, context_mask)
        for module in self.output_blocks:
            h_prev = hs.pop()
            h = torch.cat([h, h_prev], dim=1)
            h = module(h, t, context_emb, context_mask)
        return self.out(h)
//  [python]
def get_rel_pos(q_size: int, k_size: int, rel_pos: torch.Tensor) -> torch.Tensor:
    """
    Get relative positional embeddings according to the relative positions of query and key sizes.

    Args:
        q_size (int): size of query q.
        k_size (int): size of key k.
        rel_pos (Tensor): relative position embeddings (L, C).

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

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

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

    Args:
        q_size (int): size of query q.
        k_size (int): size of key k.
        rel_pos (Tensor): relative position embeddings (L, C).

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

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

    return rel_pos_resized[relative_coords.long()]
//  [python]
def forward(
        self, x, time, text_embed, text_mask, x_low_res, time_low_res=None,
    ):
        if self.noise_cond_aug:
            assert time_low_res is not None, 'time_low_res cannot be None when training with noise conditioning aug'
        else:
            assert time_low_res is None, 'time_low_res cannot be presented'

        if x.dtype != time.dtype or time.dtype != text_embed.dtype:
            dtype = text_embed.dtype
            x = x.to(dtype=dtype)
            time = time.to(dtype=dtype)
            if x_low_res is not None:
                x_low_res = x_low_res.to(dtype=dtype)
            if time_low_res is not None:
                time_low_res = time_low_res.to(dtype=dtype)

        batch_size, device = x.shape[0], x.device
        # Time Conditioning
        t = self.time_embed(time)
        # Text Conditioning
        text_cond = self.text_to_cond(text_embed)
        # Concatenating low resolution images
        if x_low_res.shape != x.shape:
            # Upscale if not done in the trainer
            _, _, new_height, new_width = x.shape
            x_low_res = F.interpolate(x_low_res, (new_height, new_width), mode="bicubic")
        x = torch.cat([x, x_low_res], dim=1)

        # Add lowres time conditioning
        if self.noise_cond_aug:
            lowres_t = self.lowres_time_embed(time_low_res)
            t += lowres_t
        # Context Embedding
        # TODO We may want to concat time token here
        if self.use_null_token:
            # Null Context (Helpful when text_embed is drop)
            null_context = self.null_text_embedding.repeat(batch_size, 1, 1)
            context_emb = torch.cat([text_cond, null_context], dim=1)
            context_mask = torch.cat([text_mask, torch.ones(batch_size, 1).to(device)], dim=1)
        else:
            context_emb = text_cond
            context_mask = text_mask

        # Add pooled text embeddings to the diffusion timestep
        # TODO We may only want to calculated the pooled feature based on text token length
        if self.feature_pooling_type == 'mean':
            pooled_text_cond = text_cond.mean(dim=-2)
        elif self.feature_pooling_type == 'attention':
            pooled_text_cond = self.attention_pooling(text_embed)
        text_hiddens = self.to_text_non_attn_cond(pooled_text_cond)
        t += text_hiddens

        # UNet forward
        x = self.init_conv(x)
        feats = dict()
        for level in range(self.n_levels):
            level_key = str(level)
            x = self.DBlocks[level_key](x, t, context_emb, context_mask)
            # Save feats for UBlocks
            if level < self.n_levels - 1:
                feats[level_key] = x
        for level in range(self.n_levels - 1, -1, -1):
            level_key = str(level)
            if level < self.n_levels - 1:
                x = x + feats[level_key]
            x = self.UBlocks[level_key](x, t, context_emb, context_mask)
        return self.out(x)
//  [python]
def forward(self, x, get_fea=False):
        feas = []
        x = (x + 1) / 2
        fea_first = fea = self.conv_first(x)
        for l in self.RRDB_trunk:
            fea = l(fea)
            feas.append(fea)
        trunk = self.trunk_conv(fea)
        fea = fea_first + trunk
        feas.append(fea)
        
        fea = self.lrelu(self.upconv1(t.upsample(fea, scale_factor=2, mode='nearest')))
        fea = self.lrelu(self.upconv2(t.upsample(fea, scale_factor=2, mode='nearest')))
        if hparams['sr_scale'] == 8:
            fea = self.lrelu(self.upconv3(t.upsample(fea, scale_factor=2, mode='nearest')))
        fea_hr = self.HRconv(fea)
        out = self.conv_last(self.lrelu(fea_hr))
        out = out.clamp(0, 1)
        out = out * 2 - 1
        if get_fea:
            return out, feas
        else:
            return out
//  [python]
def forward(self, x, get_fea=False):
        feas = []
        x = (x + 1) / 2
        fea_first = fea = self.conv_first(x)
        for l in self.RRDB_trunk:
            fea = l(fea)
            feas.append(fea)
        trunk = self.trunk_conv(fea)
        fea = fea_first + trunk
        feas.append(fea)
        
        fea = self.lrelu(self.upconv1(F.interpolate(fea, scale_factor=2, mode='nearest')))
        fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest')))
        if hparams['sr_scale'] == 8:
            fea = self.lrelu(self.upconv3(F.interpolate(fea, scale_factor=2, mode='nearest')))
        fea_hr = self.HRconv(fea)
        out = self.conv_last(self.lrelu(fea_hr))
        out = out.clamp(0, 1)
        out = out * 2 - 1
        if get_fea:
            return out, feas
        else:
            return out
//  [python]
def get_input(self, batch, k):
        x = batch[k]
        if len(x.shape) == 3:
            x = x[..., None]
        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
        if self.batch_resize_range is not None:
            lower_size = self.batch_resize_range[0]
            upper_size = self.batch_resize_range[1]
            if self.global_step <= 4:
                # do the first few batches with max size to avoid later oom
                new_resize = upper_size
            else:
                new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16))
            if new_resize != x.shape[2]:
                x = t.interpolate(x, size=new_resize, mode="bicubic")
            x = x.detach()
        return x
//  [python]
def get_input(self, batch, k):
        x = batch[k]
        if len(x.shape) == 3:
            x = x[..., None]
        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
        if self.batch_resize_range is not None:
            lower_size = self.batch_resize_range[0]
            upper_size = self.batch_resize_range[1]
            if self.global_step <= 4:
                # do the first few batches with max size to avoid later oom
                new_resize = upper_size
            else:
                new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16))
            if new_resize != x.shape[2]:
                x = F.interpolate(x, size=new_resize, mode="bicubic")
            x = x.detach()
        return x
//  [python]
def get_closest_vertices_mask_from_ES(
    E: torch.Tensor,
    S: torch.Tensor,
    h: int,
    w: int,
    mesh_vertex_embeddings: torch.Tensor,
    device: torch.device,
):
    """
    Interpolate Embeddings and Segmentations to the size of a given bounding box,
    and compute closest vertices and the segmentation mask

    Args:
        E (tensor [1, D, H, W]): D-dimensional embedding vectors for every point of the
            default-sized box
        S (tensor [1, 2, H, W]): 2-dimensional segmentation mask for every point of the
            default-sized box
        h (int): height of the target bounding box
        w (int): width of the target bounding box
        mesh_vertex_embeddings (tensor [N, D]): vertex embeddings for a chosen mesh
            N is the number of vertices in the mesh, D is feature dimensionality
        device (torch.device): device to move the tensors to
    Return:
        Closest Vertices (tensor [h, w]), int, for every point of the resulting box
        Segmentation mask (tensor [h, w]), boolean, for every point of the resulting box
    """
    embedding_resized = t.nn.functional.interpolate(E, size=(h, w), mode="bilinear")[0].to(device)
    coarse_segm_resized = t.nn.functional.interpolate(S, size=(h, w), mode="bilinear")[0].to(device)
    mask = coarse_segm_resized.argmax(0) > 0
    closest_vertices = torch.zeros(mask.shape, dtype=torch.long, device=device)
    all_embeddings = embedding_resized[:, mask].t()
    size_chunk = 10_000  # Chunking to avoid possible OOM
    edm = []
    if len(all_embeddings) == 0:
        return closest_vertices, mask
    for chunk in range((len(all_embeddings) - 1) // size_chunk + 1):
        chunk_embeddings = all_embeddings[size_chunk * chunk : size_chunk * (chunk + 1)]
        edm.append(
            torch.argmin(
                squared_euclidean_distance_matrix(chunk_embeddings, mesh_vertex_embeddings), dim=1
            )
        )
    closest_vertices[mask] = torch.cat(edm)
    return closest_vertices, mask
//  [python]
def get_closest_vertices_mask_from_ES(
    E: torch.Tensor,
    S: torch.Tensor,
    h: int,
    w: int,
    mesh_vertex_embeddings: torch.Tensor,
    device: torch.device,
):
    """
    Interpolate Embeddings and Segmentations to the size of a given bounding box,
    and compute closest vertices and the segmentation mask

    Args:
        E (tensor [1, D, H, W]): D-dimensional embedding vectors for every point of the
            default-sized box
        S (tensor [1, 2, H, W]): 2-dimensional segmentation mask for every point of the
            default-sized box
        h (int): height of the target bounding box
        w (int): width of the target bounding box
        mesh_vertex_embeddings (tensor [N, D]): vertex embeddings for a chosen mesh
            N is the number of vertices in the mesh, D is feature dimensionality
        device (torch.device): device to move the tensors to
    Return:
        Closest Vertices (tensor [h, w]), int, for every point of the resulting box
        Segmentation mask (tensor [h, w]), boolean, for every point of the resulting box
    """
    embedding_resized = F.interpolate(E, size=(h, w), mode="bilinear")[0].to(device)
    coarse_segm_resized = F.interpolate(S, size=(h, w), mode="bilinear")[0].to(device)
    mask = coarse_segm_resized.argmax(0) > 0
    closest_vertices = torch.zeros(mask.shape, dtype=torch.long, device=device)
    all_embeddings = embedding_resized[:, mask].t()
    size_chunk = 10_000  # Chunking to avoid possible OOM
    edm = []
    if len(all_embeddings) == 0:
        return closest_vertices, mask
    for chunk in range((len(all_embeddings) - 1) // size_chunk + 1):
        chunk_embeddings = all_embeddings[size_chunk * chunk : size_chunk * (chunk + 1)]
        edm.append(
            torch.argmin(
                squared_euclidean_distance_matrix(chunk_embeddings, mesh_vertex_embeddings), dim=1
            )
        )
    closest_vertices[mask] = torch.cat(edm)
    return closest_vertices, mask
//  [python]
def forward(self, x, t, lx, lt, y, context, mask):
        assert context.shape[:-1] == mask.shape

        # embeddings
        t = self.time_embedding(sinusoidal_embedding(t, self.dim)) \
            + self.noise_time_embedding(sinusoidal_embedding(lt, self.dim)) \
            + self.pool_embedding(y)

        context = self.text_embedding(context)

        if lx.shape[-2:] != x.shape[-2:]:
            lx = F.upsample(
                lx, x.shape[-2:], mode='bilinear')
        x = torch.cat([x, lx], dim=1)

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

        # middle
        for block in self.middle:
            x = self._forward_single(block, x, t, context, mask)

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

        # head
        x = self.head(x)
        return x
//  [python]
def forward(self, x, t, lx, lt, y, context, mask):
        assert context.shape[:-1] == mask.shape

        # embeddings
        t = self.time_embedding(sinusoidal_embedding(t, self.dim)) \
            + self.noise_time_embedding(sinusoidal_embedding(lt, self.dim)) \
            + self.pool_embedding(y)

        context = self.text_embedding(context)

        if lx.shape[-2:] != x.shape[-2:]:
            lx = F.interpolate(
                lx, x.shape[-2:], mode='bilinear', align_corners=False)
        x = torch.cat([x, lx], dim=1)

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

        # middle
        for block in self.middle:
            x = self._forward_single(block, x, t, context, mask)

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

        # head
        x = self.head(x)
        return x
//  [python]
def test_upsample_nearest1d(self):
        x = torch.randn(2, 2, 4)
        x1 = x.clone().detach().requires_grad_()
        y1 = t.interpolate(x1, scale_factor=2, mode="nearest")
        y1.mean().backward()

        # test bfloat16
        x3 = x.clone().detach().bfloat16().requires_grad_()
        y3 = t.interpolate(x3, scale_factor=2, mode="nearest")
        y3.mean().backward()
        self.assertTrue(y3.dtype == torch.bfloat16)
        self.assertEqual(y1, y3, prec=0.01)
        self.assertTrue(x3.grad.dtype == torch.bfloat16)
        self.assertEqual(x1.grad, x3.grad)

        # test autocast
        with torch.cpu.amp.autocast():
            for datatype in (torch.bfloat16, torch.float32):
                x4 = x.clone().detach().to(datatype).requires_grad_()
                y4 = t.interpolate(x4, scale_factor=2, mode="nearest")
                y4.mean().backward()
                self.assertTrue(y4.dtype == datatype)
                self.assertTrue(x4.grad.dtype == datatype)
//  [python]
def test_upsample_nearest1d(self):
        x = torch.randn(2, 2, 4)
        x1 = x.clone().detach().requires_grad_()
        y1 = F.interpolate(x1, scale_factor=2, mode="nearest")
        y1.mean().backward()

        # test bfloat16
        x3 = x.clone().detach().bfloat16().requires_grad_()
        y3 = F.interpolate(x3, scale_factor=2, mode="nearest")
        y3.mean().backward()
        self.assertTrue(y3.dtype == torch.bfloat16)
        self.assertEqual(y1, y3, prec=0.01)
        self.assertTrue(x3.grad.dtype == torch.bfloat16)
        self.assertEqual(x1.grad, x3.grad)

        # test autocast
        with torch.cpu.amp.autocast():
            for datatype in (torch.bfloat16, torch.float32):
                x4 = x.clone().detach().to(datatype).requires_grad_()
                y4 = F.interpolate(x4, scale_factor=2, mode="nearest")
                y4.mean().backward()
                self.assertTrue(y4.dtype == datatype)
                self.assertTrue(x4.grad.dtype == datatype)
//  [python]
def test_upsample_linear1d(self):
        x = torch.randn(2, 2, 4)
        x1 = x.clone().detach().requires_grad_()
        y1 = t.interpolate(x1, scale_factor=2, mode="linear")
        y1.mean().backward()

        # test bfloat16
        x3 = x.clone().detach().bfloat16().requires_grad_()
        y3 = t.interpolate(x3, scale_factor=2, mode="linear")
        y3.mean().backward()
        self.assertTrue(y3.dtype == torch.bfloat16)
        self.assertEqual(y1, y3, prec=0.01)
        self.assertTrue(x3.grad.dtype == torch.bfloat16)
        self.assertEqual(x1.grad, x3.grad)

        # test autocast
        with torch.cpu.amp.autocast():
            for datatype in (torch.bfloat16, torch.float32):
                x4 = x.clone().detach().to(datatype).requires_grad_()
                y4 = t.interpolate(x4, scale_factor=2, mode="linear")
                y4.mean().backward()
                self.assertTrue(y4.dtype == datatype)
                self.assertTrue(x4.grad.dtype == datatype)
//  [python]
def test_upsample_linear1d(self):
        x = torch.randn(2, 2, 4)
        x1 = x.clone().detach().requires_grad_()
        y1 = F.interpolate(x1, scale_factor=2, mode="linear")
        y1.mean().backward()

        # test bfloat16
        x3 = x.clone().detach().bfloat16().requires_grad_()
        y3 = F.interpolate(x3, scale_factor=2, mode="linear")
        y3.mean().backward()
        self.assertTrue(y3.dtype == torch.bfloat16)
        self.assertEqual(y1, y3, prec=0.01)
        self.assertTrue(x3.grad.dtype == torch.bfloat16)
        self.assertEqual(x1.grad, x3.grad)

        # test autocast
        with torch.cpu.amp.autocast():
            for datatype in (torch.bfloat16, torch.float32):
                x4 = x.clone().detach().to(datatype).requires_grad_()
                y4 = F.interpolate(x4, scale_factor=2, mode="linear")
                y4.mean().backward()
                self.assertTrue(y4.dtype == datatype)
                self.assertTrue(x4.grad.dtype == datatype)
//  [python]
def spatial_average(in_tens, mask=None, keepdim=True):
    if mask is None:
        return in_tens.mean([2,3], keepdim=keepdim)
    else:
        mask_resized = torch.nn.functional.upsample(mask[:, 0:1, :, :], size=[in_tens.size(2), in_tens.size(3)])
        num_valid = torch.sum(mask_resized)
        return torch.sum(in_tens * mask_resized) / (num_valid + 1e-8)
//  [python]
def spatial_average(in_tens, mask=None, keepdim=True):
    if mask is None:
        return in_tens.mean([2,3], keepdim=keepdim)
    else:
        mask_resized = torch.nn.functional.interpolate(mask[:, 0:1, :, :], size=[in_tens.size(2), in_tens.size(3)])
        num_valid = torch.sum(mask_resized)
        return torch.sum(in_tens * mask_resized) / (num_valid + 1e-8)