jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self,
                input: torch.Tensor,
                target_is_real: bool,
                is_disc: bool = False,
                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Args:
            input (Tensor): The input for the loss module, i.e., the network
                prediction.
            target_is_real (bool): Whether the target is real or fake.
            is_disc (bool): Whether the loss for discriminators or not.
                Default: False.
            mask (Tensor): The mask tensor. Default: None.

        Returns:
            Tensor: GAN loss value.
        """

        target_label = self.get_target_label(input, target_is_real)
        if self.gan_type == 'hinge':
            if is_disc:  # for discriminators in hinge-gan
                input = -input if target_is_real else input
                loss = self.loss(1 + input).mean()
            else:  # for generators in hinge-gan
                loss = -input.mean()
        elif self.gan_type == 'smgan':

            input_height, input_width = input.shape[2:]
            mask_height, mask_width = mask.shape[2:]

            # Handle inconsistent size between outputs and masks
            if input_height != mask_height or input_width != mask_width:
                input = F.upsample(
                    input,
                    size=(mask_height, mask_width),
                    mode='bilinear')

                target_label = self.get_target_label(input, target_is_real)

            if is_disc:
                if target_is_real:
                    target_label = target_label
                else:
                    target_label = self.gaussian_blur(mask).detach().cuda(
                    ) if mask.is_cuda else self.gaussian_blur(
                        mask).detach().cpu()
                    # target_label = self.gaussian_blur(mask).detach().cpu()
                loss = self.loss(input, target_label)
            else:
                loss = self.loss(input, target_label) * mask / mask.mean()
                loss = loss.mean()
        else:  # other gan types
            loss = self.loss(input, target_label)

        # loss_weight is always 1.0 for discriminators
        return loss if is_disc else loss * self.loss_weight
//  [python]
def forward(self,
                input: torch.Tensor,
                target_is_real: bool,
                is_disc: bool = False,
                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Args:
            input (Tensor): The input for the loss module, i.e., the network
                prediction.
            target_is_real (bool): Whether the target is real or fake.
            is_disc (bool): Whether the loss for discriminators or not.
                Default: False.
            mask (Tensor): The mask tensor. Default: None.

        Returns:
            Tensor: GAN loss value.
        """

        target_label = self.get_target_label(input, target_is_real)
        if self.gan_type == 'hinge':
            if is_disc:  # for discriminators in hinge-gan
                input = -input if target_is_real else input
                loss = self.loss(1 + input).mean()
            else:  # for generators in hinge-gan
                loss = -input.mean()
        elif self.gan_type == 'smgan':

            input_height, input_width = input.shape[2:]
            mask_height, mask_width = mask.shape[2:]

            # Handle inconsistent size between outputs and masks
            if input_height != mask_height or input_width != mask_width:
                input = F.interpolate(
                    input,
                    size=(mask_height, mask_width),
                    mode='bilinear',
                    align_corners=True)

                target_label = self.get_target_label(input, target_is_real)

            if is_disc:
                if target_is_real:
                    target_label = target_label
                else:
                    target_label = self.gaussian_blur(mask).detach().cuda(
                    ) if mask.is_cuda else self.gaussian_blur(
                        mask).detach().cpu()
                    # target_label = self.gaussian_blur(mask).detach().cpu()
                loss = self.loss(input, target_label)
            else:
                loss = self.loss(input, target_label) * mask / mask.mean()
                loss = loss.mean()
        else:  # other gan types
            loss = self.loss(input, target_label)

        # loss_weight is always 1.0 for discriminators
        return loss if is_disc else loss * self.loss_weight
//  [python]
def smooth_interp(x, size):
        r"""Smooth interpolation of segmentation maps.

        Args:
            x (4D tensor): Segmentation maps.
            size(2D list): Target size (H, W).
        """
        x = t.nn.functional.interpolate(x, size=size, mode='area')
        onehot_idx = torch.argmax(x, dim=-3, keepdims=True)
        x.fill_(0.0)
        x.scatter_(1, onehot_idx, 1.0)
        return x
//  [python]
def smooth_interp(x, size):
        r"""Smooth interpolation of segmentation maps.

        Args:
            x (4D tensor): Segmentation maps.
            size(2D list): Target size (H, W).
        """
        x = F.interpolate(x, size=size, mode='area')
        onehot_idx = torch.argmax(x, dim=-3, keepdims=True)
        x.fill_(0.0)
        x.scatter_(1, onehot_idx, 1.0)
        return x
//  [python]
def resize_abs_pos_embed(self, state_dict):
        if 'pos_embed' in state_dict:
            pos_embed_checkpoint = state_dict['pos_embed']
            embedding_size = pos_embed_checkpoint.shape[-1]
            num_extra_tokens = self.pos_embed.shape[-2] - self.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(self.num_patches**0.5)
            # class_token and dist_token are kept unchanged
            if orig_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)
                state_dict['pos_embed'] = new_pos_embed
        return state_dict
//  [python]
def __call__(self, image):
        # image: np.ndarray, uint8, [H, W, 3]
        H, W = image.shape[:2]
        image = Image.fromarray(image)

        image = self.aug(image).unsqueeze(0).to(self.device)

        if self.task == 'depth':
            depth = self.model(image).clamp(0, 1)
            depth = t.interpolate(depth.unsqueeze(1), size=(H, W), mode='bicubic', align_corners=False)
            depth = depth.squeeze(1).cpu().numpy()
            return depth
        else:
            normal = self.model(image).clamp(0, 1)
            normal = t.interpolate(normal, size=(H, W), mode='bicubic', align_corners=False)
            normal = normal.cpu().numpy()
            return normal
//  [python]
def __call__(self, image):
        # image: np.ndarray, uint8, [H, W, 3]
        H, W = image.shape[:2]
        image = Image.fromarray(image)

        image = self.aug(image).unsqueeze(0).to(self.device)

        if self.task == 'depth':
            depth = self.model(image).clamp(0, 1)
            depth = F.interpolate(depth.unsqueeze(1), size=(H, W), mode='bicubic', align_corners=False)
            depth = depth.squeeze(1).cpu().numpy()
            return depth
        else:
            normal = self.model(image).clamp(0, 1)
            normal = F.interpolate(normal, size=(H, W), mode='bicubic', align_corners=False)
            normal = normal.cpu().numpy()
            return normal
//  [python]
def resize(input,
           size=None,
           scale_factor=None,
           mode='nearest',
           align_corners=None,
           warning=True):

    return t.interpolate(input, size, scale_factor, mode, align_corners)
//  [python]
def resize(input,
           size=None,
           scale_factor=None,
           mode='nearest',
           align_corners=None,
           warning=True):

    return F.interpolate(input, size, scale_factor, mode, align_corners)
//  [python]
def forward(self, latent_inputs, latent_predictions, image_inputs, split="train"):
        log = {}
        loss = (latent_inputs - latent_predictions) ** 2
        log[f"{split}/latent_l2_loss"] = loss.mean().detach()
        image_reconstructions = None
        if self.perceptual_weight > 0.0:
            image_reconstructions = self.decoder.decode(latent_predictions)
            image_targets = self.decoder.decode(latent_inputs)
            perceptual_loss = self.perceptual_loss(
                image_targets.contiguous(), image_reconstructions.contiguous()
            )
            loss = (
                self.latent_weight * loss.mean()
                + self.perceptual_weight * perceptual_loss.mean()
            )
            log[f"{split}/perceptual_loss"] = perceptual_loss.mean().detach()

        if self.perceptual_weight_on_inputs > 0.0:
            image_reconstructions = default(
                image_reconstructions, self.decoder.decode(latent_predictions)
            )
            if self.scale_input_to_tgt_size:
                image_inputs = t.interpolate(
                    image_inputs,
                    image_reconstructions.shape[2:],
                    mode="bicubic",
                    antialias=True,
                )
            elif self.scale_tgt_to_input_size:
                image_reconstructions = t.interpolate(
                    image_reconstructions,
                    image_inputs.shape[2:],
                    mode="bicubic",
                    antialias=True,
                )

            perceptual_loss2 = self.perceptual_loss(
                image_inputs.contiguous(), image_reconstructions.contiguous()
            )
            loss = loss + self.perceptual_weight_on_inputs * perceptual_loss2.mean()
            log[f"{split}/perceptual_loss_on_inputs"] = perceptual_loss2.mean().detach()
        return loss, log
//  [python]
def forward(self, latent_inputs, latent_predictions, image_inputs, split="train"):
        log = {}
        loss = (latent_inputs - latent_predictions) ** 2
        log[f"{split}/latent_l2_loss"] = loss.mean().detach()
        image_reconstructions = None
        if self.perceptual_weight > 0.0:
            image_reconstructions = self.decoder.decode(latent_predictions)
            image_targets = self.decoder.decode(latent_inputs)
            perceptual_loss = self.perceptual_loss(
                image_targets.contiguous(), image_reconstructions.contiguous()
            )
            loss = (
                self.latent_weight * loss.mean()
                + self.perceptual_weight * perceptual_loss.mean()
            )
            log[f"{split}/perceptual_loss"] = perceptual_loss.mean().detach()

        if self.perceptual_weight_on_inputs > 0.0:
            image_reconstructions = default(
                image_reconstructions, self.decoder.decode(latent_predictions)
            )
            if self.scale_input_to_tgt_size:
                image_inputs = torch.nn.functional.interpolate(
                    image_inputs,
                    image_reconstructions.shape[2:],
                    mode="bicubic",
                    antialias=True,
                )
            elif self.scale_tgt_to_input_size:
                image_reconstructions = torch.nn.functional.interpolate(
                    image_reconstructions,
                    image_inputs.shape[2:],
                    mode="bicubic",
                    antialias=True,
                )

            perceptual_loss2 = self.perceptual_loss(
                image_inputs.contiguous(), image_reconstructions.contiguous()
            )
            loss = loss + self.perceptual_weight_on_inputs * perceptual_loss2.mean()
            log[f"{split}/perceptual_loss_on_inputs"] = perceptual_loss2.mean().detach()
        return loss, log
//  [python]
def test_interpolate():
    """test tensorrt translator for interpolate"""

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

    input_info = [([1, 3, 10, 10], "float32")]
    verify_model(Interpolate(), input_info)
//  [python]
def test_interpolate():
    """test tensorrt translator for interpolate"""

    class Interpolate(Module):
        def forward(self, data):
            return torch.nn.functional.interpolate(data, (5, 5))

    input_info = [([1, 3, 10, 10], "float32")]
    verify_model(Interpolate(), input_info)
//  [python]
def resolve_areas_and_cond_masks(conditions, h, w, device):
    # We need to decide on an area outside the sampling loop in order to properly generate opposite areas of equal sizes.
    # While we're doing this, we can also resolve the mask device and scaling for performance reasons
    for i in range(len(conditions)):
        c = conditions[i]
        if 'area' in c:
            area = c['area']
            if area[0] == "percentage":
                modified = c.copy()
                area = (max(1, round(area[1] * h)), max(1, round(area[2] * w)), round(area[3] * h), round(area[4] * w))
                modified['area'] = area
                c = modified
                conditions[i] = c

        if 'mask' in c:
            mask = c['mask']
            mask = mask.to(device=device)
            modified = c.copy()
            if len(mask.shape) == 2:
                mask = mask.unsqueeze(0)
            if mask.shape[1] != h or mask.shape[2] != w:
                mask = t.nn.functional.interpolate(mask.unsqueeze(1), size=(h, w), mode='bilinear', align_corners=False).squeeze(1)

            if modified.get("set_area_to_bounds", False):
                bounds = torch.max(torch.abs(mask),dim=0).values.unsqueeze(0)
                boxes, is_empty = get_mask_aabb(bounds)
                if is_empty[0]:
                    # Use the minimum possible size for efficiency reasons. (Since the mask is all-0, this becomes a noop anyway)
                    modified['area'] = (8, 8, 0, 0)
                else:
                    box = boxes[0]
                    H, W, Y, X = (box[3] - box[1] + 1, box[2] - box[0] + 1, box[1], box[0])
                    H = max(8, H)
                    W = max(8, W)
                    area = (int(H), int(W), int(Y), int(X))
                    modified['area'] = area

            modified['mask'] = mask
            conditions[i] = modified
//  [python]
def resolve_areas_and_cond_masks(conditions, h, w, device):
    # We need to decide on an area outside the sampling loop in order to properly generate opposite areas of equal sizes.
    # While we're doing this, we can also resolve the mask device and scaling for performance reasons
    for i in range(len(conditions)):
        c = conditions[i]
        if 'area' in c:
            area = c['area']
            if area[0] == "percentage":
                modified = c.copy()
                area = (max(1, round(area[1] * h)), max(1, round(area[2] * w)), round(area[3] * h), round(area[4] * w))
                modified['area'] = area
                c = modified
                conditions[i] = c

        if 'mask' in c:
            mask = c['mask']
            mask = mask.to(device=device)
            modified = c.copy()
            if len(mask.shape) == 2:
                mask = mask.unsqueeze(0)
            if mask.shape[1] != h or mask.shape[2] != w:
                mask = torch.nn.functional.interpolate(mask.unsqueeze(1), size=(h, w), mode='bilinear', align_corners=False).squeeze(1)

            if modified.get("set_area_to_bounds", False):
                bounds = torch.max(torch.abs(mask),dim=0).values.unsqueeze(0)
                boxes, is_empty = get_mask_aabb(bounds)
                if is_empty[0]:
                    # Use the minimum possible size for efficiency reasons. (Since the mask is all-0, this becomes a noop anyway)
                    modified['area'] = (8, 8, 0, 0)
                else:
                    box = boxes[0]
                    H, W, Y, X = (box[3] - box[1] + 1, box[2] - box[0] + 1, box[1], box[0])
                    H = max(8, H)
                    W = max(8, W)
                    area = (int(H), int(W), int(Y), int(X))
                    modified['area'] = area

            modified['mask'] = mask
            conditions[i] = modified
//  [python]
def forward(self, x):
        """Forward function.

        Args:
            x (Tensor): Input tensor with shape (N, C, H, W).

        Returns:
            dict: Contains the output tensor, shortcut feature and \
                intermediate feature.
        """
        out = self.conv1(x)
        x1 = self.conv2(out)
        out = self.conv3(x1)

        img_feat = self.guidance_head(x[:, :3, ...])
        if self.trimap_channels == 3:
            unknown = x[:, 4:5, ...]
        else:
            unknown = x[:, 3:, ...].eq(1).float()
        # same as img_feat, downsample to 1/8
        unknown = F.interpolate(unknown, scale_factor=1 / 8, mode='nearest')

        x2 = self.layer1(out)
        x3 = self.layer2(x2)
        x3 = self.gca(img_feat, x3, unknown)
        x4 = self.layer3(x3)
        out = self.layer4(x4)

        # shortcut block
        feat1 = self.shortcut[0](x)
        feat2 = self.shortcut[1](x1)
        feat3 = self.shortcut[2](x2)
        feat4 = self.shortcut[3](x3)
        feat5 = self.shortcut[4](x4)

        return {
            'out': out,
            'feat1': feat1,
            'feat2': feat2,
            'feat3': feat3,
            'feat4': feat4,
            'feat5': feat5,
            'img_feat': img_feat,
            'unknown': unknown
        }
//  [python]
def forward(self, x: Tensor) -> Tensor:
        """
        x - input image
        """
        shape_im = x.shape
        feats: List[Tensor] = [self.feature_extractor(x)]
        for i in range(1, self.num_levels):
            x = pyrdown(x, factor=1.2)
            feats_i = self.feature_extractor(x)
            feats_i = t.interpolate(feats_i, size=(shape_im[2], shape_im[3]), mode="bilinear")
            feats.append(feats_i)
        scores = self.last_conv(concatenate(feats, 1))
        return scores
//  [python]
def forward(self, x: Tensor) -> Tensor:
        """
        x - input image
        """
        shape_im = x.shape
        feats: List[Tensor] = [self.feature_extractor(x)]
        for i in range(1, self.num_levels):
            x = pyrdown(x, factor=1.2)
            feats_i = self.feature_extractor(x)
            feats_i = F.interpolate(feats_i, size=(shape_im[2], shape_im[3]), mode="bilinear")
            feats.append(feats_i)
        scores = self.last_conv(concatenate(feats, 1))
        return scores
//  [python]
def forward(
        self,
        sample,
        timestep_ratio,
        clip_text_pooled,
        clip_text=None,
        clip_img=None,
        effnet=None,
        pixels=None,
        sca=None,
        crp=None,
        return_dict=True,
    ):
        if pixels is None:
            pixels = sample.new_zeros(sample.size(0), 3, 8, 8)

        # Process the conditioning embeddings
        timestep_ratio_embed = self.get_timestep_ratio_embedding(timestep_ratio)
        for c in self.config.timestep_conditioning_type:
            if c == "sca":
                cond = sca
            elif c == "crp":
                cond = crp
            else:
                cond = None
            t_cond = cond or torch.zeros_like(timestep_ratio)
            timestep_ratio_embed = torch.cat([timestep_ratio_embed, self.get_timestep_ratio_embedding(t_cond)], dim=1)
        clip = self.get_clip_embeddings(clip_txt_pooled=clip_text_pooled, clip_txt=clip_text, clip_img=clip_img)

        # Model Blocks
        x = self.embedding(sample)
        if hasattr(self, "effnet_mapper") and effnet is not None:
            x = x + self.effnet_mapper(
                t.interpolate(effnet, size=x.shape[-2:], mode="bilinear", align_corners=True)
            )
        if hasattr(self, "pixels_mapper"):
            x = x + t.interpolate(
                self.pixels_mapper(pixels), size=x.shape[-2:], mode="bilinear", align_corners=True
            )
        level_outputs = self._down_encode(x, timestep_ratio_embed, clip)
        x = self._up_decode(level_outputs, timestep_ratio_embed, clip)
        sample = self.clf(x)

        if not return_dict:
            return (sample,)
        return StableCascadeUNetOutput(sample=sample)
//  [python]
def forward(
        self,
        sample,
        timestep_ratio,
        clip_text_pooled,
        clip_text=None,
        clip_img=None,
        effnet=None,
        pixels=None,
        sca=None,
        crp=None,
        return_dict=True,
    ):
        if pixels is None:
            pixels = sample.new_zeros(sample.size(0), 3, 8, 8)

        # Process the conditioning embeddings
        timestep_ratio_embed = self.get_timestep_ratio_embedding(timestep_ratio)
        for c in self.config.timestep_conditioning_type:
            if c == "sca":
                cond = sca
            elif c == "crp":
                cond = crp
            else:
                cond = None
            t_cond = cond or torch.zeros_like(timestep_ratio)
            timestep_ratio_embed = torch.cat([timestep_ratio_embed, self.get_timestep_ratio_embedding(t_cond)], dim=1)
        clip = self.get_clip_embeddings(clip_txt_pooled=clip_text_pooled, clip_txt=clip_text, clip_img=clip_img)

        # Model Blocks
        x = self.embedding(sample)
        if hasattr(self, "effnet_mapper") and effnet is not None:
            x = x + self.effnet_mapper(
                nn.functional.interpolate(effnet, size=x.shape[-2:], mode="bilinear", align_corners=True)
            )
        if hasattr(self, "pixels_mapper"):
            x = x + nn.functional.interpolate(
                self.pixels_mapper(pixels), size=x.shape[-2:], mode="bilinear", align_corners=True
            )
        level_outputs = self._down_encode(x, timestep_ratio_embed, clip)
        x = self._up_decode(level_outputs, timestep_ratio_embed, clip)
        sample = self.clf(x)

        if not return_dict:
            return (sample,)
        return StableCascadeUNetOutput(sample=sample)
//  [python]
def visualize_basics(
    viz: "Visdom",
    preds: Dict[str, Any],
    visdom_env_imgs: str,
    title: str = "",
    visualize_preds_keys: Tuple[str, ...] = (
        "image_rgb",
        "images_render",
        "fg_probability",
        "masks_render",
        "depths_render",
        "depth_map",
    ),
    store_history: bool = False,
) -> None:
    """
    Visualize basic outputs of a `GenericModel` to visdom.

    Args:
        viz: The visdom object.
        preds: A dictionary containing `GenericModel` outputs.
        visdom_env_imgs: Target visdom environment name.
        title: The title of produced visdom window.
        visualize_preds_keys: The list of keys of `preds` for visualization.
        store_history: Store the history buffer in visdom windows.
    """
    imout = {}
    for k in visualize_preds_keys:
        if k not in preds or preds[k] is None:
            logger.info(f"cant show {k}")
            continue
        v = preds[k].cpu().detach().clone()
        if k.startswith("depth"):
            # divide by 95th percentile
            normfac = (
                v.view(v.shape[0], -1)
                .topk(k=int(0.05 * (v.numel() // v.shape[0])), dim=-1)
                .values[:, -1]
            )
            v = v / normfac[:, None, None, None].clamp(1e-4)
        if v.shape[1] == 1:
            v = v.repeat(1, 3, 1, 1)
        v = t.nn.functional.interpolate(
            v,
            size=(
                600.0
                if (
                    "_eval" in visdom_env_imgs
                    and k in ("images_render", "depths_render")
                )
                else 200.0
            )
            / v.shape[2],
            mode="bilinear",
        )
        imout[k] = v

    # TODO: handle errors on the outside
    try:
        imout = {"all": torch.cat(list(imout.values()), dim=2)}
    except RuntimeError as e:
        print("cant cat!", e.args)

    for k, v in imout.items():
        viz.images(
            v.clamp(0.0, 1.0),
            win=k,
            env=visdom_env_imgs,
            opts={"title": title + "_" + k, "store_history": store_history},
        )
//  [python]
def visualize_basics(
    viz: "Visdom",
    preds: Dict[str, Any],
    visdom_env_imgs: str,
    title: str = "",
    visualize_preds_keys: Tuple[str, ...] = (
        "image_rgb",
        "images_render",
        "fg_probability",
        "masks_render",
        "depths_render",
        "depth_map",
    ),
    store_history: bool = False,
) -> None:
    """
    Visualize basic outputs of a `GenericModel` to visdom.

    Args:
        viz: The visdom object.
        preds: A dictionary containing `GenericModel` outputs.
        visdom_env_imgs: Target visdom environment name.
        title: The title of produced visdom window.
        visualize_preds_keys: The list of keys of `preds` for visualization.
        store_history: Store the history buffer in visdom windows.
    """
    imout = {}
    for k in visualize_preds_keys:
        if k not in preds or preds[k] is None:
            logger.info(f"cant show {k}")
            continue
        v = preds[k].cpu().detach().clone()
        if k.startswith("depth"):
            # divide by 95th percentile
            normfac = (
                v.view(v.shape[0], -1)
                .topk(k=int(0.05 * (v.numel() // v.shape[0])), dim=-1)
                .values[:, -1]
            )
            v = v / normfac[:, None, None, None].clamp(1e-4)
        if v.shape[1] == 1:
            v = v.repeat(1, 3, 1, 1)
        v = torch.nn.functional.interpolate(
            v,
            scale_factor=(
                600.0
                if (
                    "_eval" in visdom_env_imgs
                    and k in ("images_render", "depths_render")
                )
                else 200.0
            )
            / v.shape[2],
            mode="bilinear",
        )
        imout[k] = v

    # TODO: handle errors on the outside
    try:
        imout = {"all": torch.cat(list(imout.values()), dim=2)}
    except RuntimeError as e:
        print("cant cat!", e.args)

    for k, v in imout.items():
        viz.images(
            v.clamp(0.0, 1.0),
            win=k,
            env=visdom_env_imgs,
            opts={"title": title + "_" + k, "store_history": store_history},
        )
//  [python]
def _up_decode(self, level_outputs, r_embed, clip):
        x = level_outputs[0]
        block_group = zip(self.up_blocks, self.up_upscalers, self.up_repeat_mappers)

        if self.training and self.gradient_checkpointing:

            def create_custom_forward(module):
                def custom_forward(*inputs):
                    return module(*inputs)

                return custom_forward

            for i, (up_block, upscaler, repmap) in enumerate(block_group):
                for j in range(len(repmap) + 1):
                    for k, block in enumerate(up_block):
                        if isinstance(block, SDCascadeResBlock):
                            skip = level_outputs[i] if k == 0 and i > 0 else None
                            if skip is not None and (x.size(-1) != skip.size(-1) or x.size(-2) != skip.size(-2)):
                                x = t.interpolate(
                                    x.float(), skip.shape[-2:], mode="bilinear", align_corners=True
                                )
                            x = torch.utils.checkpoint.checkpoint(
                                create_custom_forward(block), x, skip, use_reentrant=False
                            )
                        elif isinstance(block, SDCascadeAttnBlock):
                            x = torch.utils.checkpoint.checkpoint(
                                create_custom_forward(block), x, clip, use_reentrant=False
                            )
                        elif isinstance(block, SDCascadeTimestepBlock):
                            x = torch.utils.checkpoint.checkpoint(
                                create_custom_forward(block), x, r_embed, use_reentrant=False
                            )
                        else:
                            x = torch.utils.checkpoint.checkpoint(create_custom_forward(block), x, use_reentrant=False)
                    if j < len(repmap):
                        x = repmap[j](x)
                x = upscaler(x)
        else:
            for i, (up_block, upscaler, repmap) in enumerate(block_group):
                for j in range(len(repmap) + 1):
                    for k, block in enumerate(up_block):
                        if isinstance(block, SDCascadeResBlock):
                            skip = level_outputs[i] if k == 0 and i > 0 else None
                            if skip is not None and (x.size(-1) != skip.size(-1) or x.size(-2) != skip.size(-2)):
                                x = t.interpolate(
                                    x.float(), skip.shape[-2:], mode="bilinear", align_corners=True
                                )
                            x = block(x, skip)
                        elif isinstance(block, SDCascadeAttnBlock):
                            x = block(x, clip)
                        elif isinstance(block, SDCascadeTimestepBlock):
                            x = block(x, r_embed)
                        else:
                            x = block(x)
                    if j < len(repmap):
                        x = repmap[j](x)
                x = upscaler(x)
        return x
//  [python]
def _up_decode(self, level_outputs, r_embed, clip):
        x = level_outputs[0]
        block_group = zip(self.up_blocks, self.up_upscalers, self.up_repeat_mappers)

        if self.training and self.gradient_checkpointing:

            def create_custom_forward(module):
                def custom_forward(*inputs):
                    return module(*inputs)

                return custom_forward

            for i, (up_block, upscaler, repmap) in enumerate(block_group):
                for j in range(len(repmap) + 1):
                    for k, block in enumerate(up_block):
                        if isinstance(block, SDCascadeResBlock):
                            skip = level_outputs[i] if k == 0 and i > 0 else None
                            if skip is not None and (x.size(-1) != skip.size(-1) or x.size(-2) != skip.size(-2)):
                                x = torch.nn.functional.interpolate(
                                    x.float(), skip.shape[-2:], mode="bilinear", align_corners=True
                                )
                            x = torch.utils.checkpoint.checkpoint(
                                create_custom_forward(block), x, skip, use_reentrant=False
                            )
                        elif isinstance(block, SDCascadeAttnBlock):
                            x = torch.utils.checkpoint.checkpoint(
                                create_custom_forward(block), x, clip, use_reentrant=False
                            )
                        elif isinstance(block, SDCascadeTimestepBlock):
                            x = torch.utils.checkpoint.checkpoint(
                                create_custom_forward(block), x, r_embed, use_reentrant=False
                            )
                        else:
                            x = torch.utils.checkpoint.checkpoint(create_custom_forward(block), x, use_reentrant=False)
                    if j < len(repmap):
                        x = repmap[j](x)
                x = upscaler(x)
        else:
            for i, (up_block, upscaler, repmap) in enumerate(block_group):
                for j in range(len(repmap) + 1):
                    for k, block in enumerate(up_block):
                        if isinstance(block, SDCascadeResBlock):
                            skip = level_outputs[i] if k == 0 and i > 0 else None
                            if skip is not None and (x.size(-1) != skip.size(-1) or x.size(-2) != skip.size(-2)):
                                x = torch.nn.functional.interpolate(
                                    x.float(), skip.shape[-2:], mode="bilinear", align_corners=True
                                )
                            x = block(x, skip)
                        elif isinstance(block, SDCascadeAttnBlock):
                            x = block(x, clip)
                        elif isinstance(block, SDCascadeTimestepBlock):
                            x = block(x, r_embed)
                        else:
                            x = block(x)
                    if j < len(repmap):
                        x = repmap[j](x)
                x = upscaler(x)
        return x
//  [python]
def inference_mask(self, img, img_metas, input_size):
        self.eval()
        raw_h, raw_w = img_metas['ori_shape']
        pad_direction = img_metas['pad_direction']
        coarse_input_size = input_size['coarse_input_size']
        refine_input_size = input_size['refine_input_size']
        h, w = img_metas['refine_shape']
        resize_images = t.nn.functional.interpolate(
            img, coarse_input_size, mode='bilinear', align_corners=True)
        # get coarse result
        pred_scores = self.coarse_model(resize_images)
        if isinstance(pred_scores, (tuple, list)):
            pred_scores = pred_scores[-1]
        score = t.nn.functional.interpolate(
            input=pred_scores,
            size=refine_input_size,
            mode='bilinear',
            align_corners=True,
        )
        _, coarse_pred = torch.max(score, dim=1)  # [B, h, w]
        coarse_pred = coarse_pred.unsqueeze(1).type(img.dtype)
        img = torch.cat([img, coarse_pred], dim=1)  # [B, c=4, h, w]
        del resize_images
        del pred_scores
        del score
        del coarse_pred
        torch.cuda.empty_cache()
        cur_scores = self.refine_model(img)
        del img
        torch.cuda.empty_cache()
        cur_scores = torch.clip(cur_scores, 0, 1)
        cur_scores = cur_scores.detach().cpu().numpy()[0]

        # resize if cur_scores shape are not compatible with origin image shape
        ph, pw = cur_scores.shape
        if ph != h or pw != w:
            cur_scores = t.nn.functional.interpolate(
                input=cur_scores,
                size=(h, w),
                mode='nearest',
                align_corners=True)
        # unpad to get valid area and resize to origin size
        valid_cur_pred = cur_scores[pad_direction[1]:refine_input_size[0]
                                    - pad_direction[3],
                                    pad_direction[0]:refine_input_size[1]
                                    - pad_direction[2], ]
        valid_cur_pred = cv2.resize(valid_cur_pred, (raw_w, raw_h))
        del cur_scores
        torch.cuda.empty_cache()
        print('get refine mask done')
        return valid_cur_pred
//  [python]
def inference_mask(self, img, img_metas, input_size):
        self.eval()
        raw_h, raw_w = img_metas['ori_shape']
        pad_direction = img_metas['pad_direction']
        coarse_input_size = input_size['coarse_input_size']
        refine_input_size = input_size['refine_input_size']
        h, w = img_metas['refine_shape']
        resize_images = F.interpolate(
            img, coarse_input_size, mode='bilinear', align_corners=True)
        # get coarse result
        pred_scores = self.coarse_model(resize_images)
        if isinstance(pred_scores, (tuple, list)):
            pred_scores = pred_scores[-1]
        score = F.interpolate(
            input=pred_scores,
            size=refine_input_size,
            mode='bilinear',
            align_corners=True,
        )
        _, coarse_pred = torch.max(score, dim=1)  # [B, h, w]
        coarse_pred = coarse_pred.unsqueeze(1).type(img.dtype)
        img = torch.cat([img, coarse_pred], dim=1)  # [B, c=4, h, w]
        del resize_images
        del pred_scores
        del score
        del coarse_pred
        torch.cuda.empty_cache()
        cur_scores = self.refine_model(img)
        del img
        torch.cuda.empty_cache()
        cur_scores = torch.clip(cur_scores, 0, 1)
        cur_scores = cur_scores.detach().cpu().numpy()[0]

        # resize if cur_scores shape are not compatible with origin image shape
        ph, pw = cur_scores.shape
        if ph != h or pw != w:
            cur_scores = F.interpolate(
                input=cur_scores,
                size=(h, w),
                mode='nearest',
                align_corners=True)
        # unpad to get valid area and resize to origin size
        valid_cur_pred = cur_scores[pad_direction[1]:refine_input_size[0]
                                    - pad_direction[3],
                                    pad_direction[0]:refine_input_size[1]
                                    - pad_direction[2], ]
        valid_cur_pred = cv2.resize(valid_cur_pred, (raw_w, raw_h))
        del cur_scores
        torch.cuda.empty_cache()
        print('get refine mask done')
        return valid_cur_pred
//  [python]
def preprocess_glide(x, order="NHWC", resize=True):
  if order == "NHWC":
    # x is [NHWC]. Reshape to NCHW.
    x = x.movedim(-1, 1)
  x = scale_glide(x)
  if resize and x.shape[-2:] != (64, 64):
    x = t.interpolate(x, (64, 64), mode="bilinear")
  return x
//  [python]
def preprocess_glide(x, order="NHWC", resize=True):
  if order == "NHWC":
    # x is [NHWC]. Reshape to NCHW.
    x = x.movedim(-1, 1)
  x = scale_glide(x)
  if resize and x.shape[-2:] != (64, 64):
    x = F.interpolate(x, (64, 64), mode="bilinear")
  return x
//  [python]
def forward(self, x, y_seg=None):

        base_x = self.generate_flow_x(self.text_features) # 36 x out_cs
        base_y = self.generate_flow_y(self.text_features) # 36 x out_cs
        base_z = self.generate_flow_z(self.text_features) # 36 x out_cs

        y_seg = t.interpolate(y_seg.float(),scale_factor=0.5,mode='nearest') # bs x 1 x h x w x d

        flow_filter_x = F.embedding(y_seg.long().squeeze(dim=1), base_x).permute(0, 4, 1, 2, 3).contiguous() # bs x out_cs x h x w x d
        flow_filter_y = F.embedding(y_seg.long().squeeze(dim=1), base_y).permute(0, 4, 1, 2, 3).contiguous() # bs x out_cs x h x w x d
        flow_filter_z = F.embedding(y_seg.long().squeeze(dim=1), base_z).permute(0, 4, 1, 2, 3).contiguous() # bs x out_cs x h x w x d

        flow_x = (flow_filter_x * x).sum(dim=1, keepdim=True)
        flow_y = (flow_filter_y * x).sum(dim=1, keepdim=True)
        flow_z = (flow_filter_z * x).sum(dim=1, keepdim=True)
        scp_flow = torch.cat([flow_x, flow_y, flow_z], dim=1)

        x = scp_flow + self.flow(x)

        return x
//  [python]
def forward(self, x, y_seg=None):

        base_x = self.generate_flow_x(self.text_features) # 36 x out_cs
        base_y = self.generate_flow_y(self.text_features) # 36 x out_cs
        base_z = self.generate_flow_z(self.text_features) # 36 x out_cs

        y_seg = nn.functional.interpolate(y_seg.float(),scale_factor=0.5,mode='nearest',recompute_scale_factor=0) # bs x 1 x h x w x d

        flow_filter_x = F.embedding(y_seg.long().squeeze(dim=1), base_x).permute(0, 4, 1, 2, 3).contiguous() # bs x out_cs x h x w x d
        flow_filter_y = F.embedding(y_seg.long().squeeze(dim=1), base_y).permute(0, 4, 1, 2, 3).contiguous() # bs x out_cs x h x w x d
        flow_filter_z = F.embedding(y_seg.long().squeeze(dim=1), base_z).permute(0, 4, 1, 2, 3).contiguous() # bs x out_cs x h x w x d

        flow_x = (flow_filter_x * x).sum(dim=1, keepdim=True)
        flow_y = (flow_filter_y * x).sum(dim=1, keepdim=True)
        flow_z = (flow_filter_z * x).sum(dim=1, keepdim=True)
        scp_flow = torch.cat([flow_x, flow_y, flow_z], dim=1)

        x = scp_flow + self.flow(x)

        return x
//  [python]
def forward(self, image, mask, config: InpaintRequest):
        """
        image: [H, W, C] RGB
        mask: [H, W, 1]
        return: BGR IMAGE
        """
        # image [1,3,512,512] float32
        # mask: [1,1,512,512] float32
        # masked_image: [1,3,512,512] float32
        if config.ldm_sampler == LDMSampler.ddim:
            sampler = DDIMSampler(self.model)
        elif config.ldm_sampler == LDMSampler.plms:
            sampler = PLMSSampler(self.model)
        else:
            raise ValueError()

        steps = config.ldm_steps
        image = norm_img(image)
        mask = norm_img(mask)

        mask[mask < 0.5] = 0
        mask[mask >= 0.5] = 1

        image = torch.from_numpy(image).unsqueeze(0).to(self.device)
        mask = torch.from_numpy(mask).unsqueeze(0).to(self.device)
        masked_image = (1 - mask) * image

        mask = self._norm(mask)
        masked_image = self._norm(masked_image)

        c = self.cond_stage_model_encode(masked_image)
        torch.cuda.empty_cache()

        cc = t.nn.functional.interpolate(mask, size=c.shape[-2:])  # 1,1,128,128
        c = torch.cat((c, cc), dim=1)  # 1,4,128,128

        shape = (c.shape[1] - 1,) + c.shape[2:]
        samples_ddim = sampler.sample(
            steps=steps, conditioning=c, batch_size=c.shape[0], shape=shape
        )
        torch.cuda.empty_cache()
        x_samples_ddim = self.cond_stage_model_decode(
            samples_ddim
        )  # samples_ddim: 1, 3, 128, 128 float32
        torch.cuda.empty_cache()

        # image = torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0)
        # mask = torch.clamp((mask + 1.0) / 2.0, min=0.0, max=1.0)
        inpainted_image = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)

        # inpainted = (1 - mask) * image + mask * predicted_image
        inpainted_image = inpainted_image.cpu().numpy().transpose(0, 2, 3, 1)[0] * 255
        inpainted_image = inpainted_image.astype(np.uint8)[:, :, ::-1]
        return inpainted_image
//  [python]
def forward(self, image, mask, config: InpaintRequest):
        """
        image: [H, W, C] RGB
        mask: [H, W, 1]
        return: BGR IMAGE
        """
        # image [1,3,512,512] float32
        # mask: [1,1,512,512] float32
        # masked_image: [1,3,512,512] float32
        if config.ldm_sampler == LDMSampler.ddim:
            sampler = DDIMSampler(self.model)
        elif config.ldm_sampler == LDMSampler.plms:
            sampler = PLMSSampler(self.model)
        else:
            raise ValueError()

        steps = config.ldm_steps
        image = norm_img(image)
        mask = norm_img(mask)

        mask[mask < 0.5] = 0
        mask[mask >= 0.5] = 1

        image = torch.from_numpy(image).unsqueeze(0).to(self.device)
        mask = torch.from_numpy(mask).unsqueeze(0).to(self.device)
        masked_image = (1 - mask) * image

        mask = self._norm(mask)
        masked_image = self._norm(masked_image)

        c = self.cond_stage_model_encode(masked_image)
        torch.cuda.empty_cache()

        cc = torch.nn.functional.interpolate(mask, size=c.shape[-2:])  # 1,1,128,128
        c = torch.cat((c, cc), dim=1)  # 1,4,128,128

        shape = (c.shape[1] - 1,) + c.shape[2:]
        samples_ddim = sampler.sample(
            steps=steps, conditioning=c, batch_size=c.shape[0], shape=shape
        )
        torch.cuda.empty_cache()
        x_samples_ddim = self.cond_stage_model_decode(
            samples_ddim
        )  # samples_ddim: 1, 3, 128, 128 float32
        torch.cuda.empty_cache()

        # image = torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0)
        # mask = torch.clamp((mask + 1.0) / 2.0, min=0.0, max=1.0)
        inpainted_image = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)

        # inpainted = (1 - mask) * image + mask * predicted_image
        inpainted_image = inpainted_image.cpu().numpy().transpose(0, 2, 3, 1)[0] * 255
        inpainted_image = inpainted_image.astype(np.uint8)[:, :, ::-1]
        return inpainted_image
//  [python]
def __call__(self, input_image, fan_landmarks=None, landmarks_validity=None):
        batched = len(input_image.shape) == 4
        if not batched:
            input_image = input_image.unsqueeze(0)
        if self.mode in [True,  'default']:
            mica_image = self._dirty_image_preprocessing(input_image)
        elif self.mode == 'ported_insightface':
            mica_image = self._dirty_image_preprocessing(input_image)
        elif self.mode in [False, 'none']: 
            mica_image = t.interpolate(input_image, (112,112), mode='bilinear', align_corners=False)
        elif self.mode == 'fan':
            if fan_landmarks is not None:
                mica_image = self._fan_image_preprocessing(input_image, fan_landmarks, landmarks_validity=landmarks_validity)
            else: 
                if self.app is None:
                    self._instantiate_face_analysis()
                mica_image = self._dirty_image_preprocessing(input_image)
        elif self.mode == 'mediapipe':
            if fan_landmarks is not None:
                mica_image = self._mediapipe_image_preprocessing(input_image, fan_landmarks, landmarks_validity=landmarks_validity)
            else: 
                if self.app is None:
                    self._instantiate_face_analysis()
                mica_image = self._dirty_image_preprocessing(input_image)
        else: 
            raise ValueError(f"Invalid mica_preprocessing option: '{self.mode}'")
        if not batched:
            mica_image = mica_image.squeeze(0)
        return mica_image
//  [python]
def __call__(self, input_image, fan_landmarks=None, landmarks_validity=None):
        batched = len(input_image.shape) == 4
        if not batched:
            input_image = input_image.unsqueeze(0)
        if self.mode in [True,  'default']:
            mica_image = self._dirty_image_preprocessing(input_image)
        elif self.mode == 'ported_insightface':
            mica_image = self._dirty_image_preprocessing(input_image)
        elif self.mode in [False, 'none']: 
            mica_image = F.interpolate(input_image, (112,112), mode='bilinear', align_corners=False)
        elif self.mode == 'fan':
            if fan_landmarks is not None:
                mica_image = self._fan_image_preprocessing(input_image, fan_landmarks, landmarks_validity=landmarks_validity)
            else: 
                if self.app is None:
                    self._instantiate_face_analysis()
                mica_image = self._dirty_image_preprocessing(input_image)
        elif self.mode == 'mediapipe':
            if fan_landmarks is not None:
                mica_image = self._mediapipe_image_preprocessing(input_image, fan_landmarks, landmarks_validity=landmarks_validity)
            else: 
                if self.app is None:
                    self._instantiate_face_analysis()
                mica_image = self._dirty_image_preprocessing(input_image)
        else: 
            raise ValueError(f"Invalid mica_preprocessing option: '{self.mode}'")
        if not batched:
            mica_image = mica_image.squeeze(0)
        return mica_image