Marigold Computer Vision

Marigold was proposed in Repurposing Diffusion-Based Image Generators for Monocular Depth Estimation, a CVPR 2024 Oral paper by Bingxin Ke, Anton Obukhov, Shengyu Huang, Nando Metzger, Rodrigo Caye Daudt, and Konrad Schindler. The core idea is to repurpose the generative prior of Text-to-Image Latent Diffusion Models (LDMs) for traditional computer vision tasks. This approach was explored by fine-tuning Stable Diffusion for Monocular Depth Estimation, as demonstrated in the teaser above.

Marigold was later extended in the follow-up paper, Marigold: Affordable Adaptation of Diffusion-Based Image Generators for Image Analysis, authored by Bingxin Ke, Kevin Qu, Tianfu Wang, Nando Metzger, Shengyu Huang, Bo Li, Anton Obukhov, and Konrad Schindler. This work expanded Marigold to support new modalities such as Surface Normals and Intrinsic Image Decomposition (IID), introduced a training protocol for Latent Consistency Models (LCM), and demonstrated High-Resolution (HR) processing capability.

Tip

The early Marigold models (v1-0 and earlier) were optimized for best results with at least 10 inference steps. LCM models were later developed to enable high-quality inference in just 1 to 4 steps. Marigold models v1-1 and later use the DDIM scheduler to achieve optimal results in as few as 1 to 4 steps.

Available Pipelines

Each pipeline is tailored for a specific computer vision task, processing an input RGB image and generating a corresponding prediction. Currently, the following computer vision tasks are implemented:

Pipeline	Recommended Model Checkpoints	Spaces (Interactive Apps)	Predicted Modalities
MarigoldDepthPipeline	prs-eth/marigold-depth-v1-1	Depth Estimation	Depth, Disparity
MarigoldNormalsPipeline	prs-eth/marigold-normals-v1-1	Surface Normals Estimation	Surface normals
MarigoldIntrinsicsPipeline	prs-eth/marigold-iid-appearance-v1-1, prs-eth/marigold-iid-lighting-v1-1	Intrinsic Image Decomposition	Albedo, Materials, Lighting

Available Checkpoints

All original checkpoints are available under the PRS-ETH organization on Hugging Face. They are designed for use with diffusers pipelines and the original codebase, which can also be used to train new model checkpoints. The following is a summary of the recommended checkpoints, all of which produce reliable results with 1 to 4 steps.

Checkpoint	Modality	Comment
prs-eth/marigold-depth-v1-1	Depth	Affine-invariant depth prediction assigns each pixel a value between 0 (near plane) and 1 (far plane), with both planes determined by the model during inference.
prs-eth/marigold-normals-v0-1	Normals	The surface normals predictions are unit-length 3D vectors in the screen space camera, with values in the range from -1 to 1.
prs-eth/marigold-iid-appearance-v1-1	Intrinsics	InteriorVerse decomposition is comprised of Albedo and two BRDF material properties: Roughness and Metallicity.
prs-eth/marigold-iid-lighting-v1-1	Intrinsics	HyperSim decomposition of an image &nbsp\(I\)&nbsp is comprised of Albedo &nbsp\(A\), Diffuse shading &nbsp\(S\), and Non-diffuse residual &nbsp\(R\): &nbsp\(I = A*S+R\).

Tip

Make sure to check out the Schedulers guide to learn how to explore the tradeoff between scheduler speed and quality, and see the reuse components across pipelines section to learn how to efficiently load the same components into multiple pipelines. Also, to know more about reducing the memory usage of this pipeline, refer to the ["Reduce memory usage"] section here.

Warning

Marigold pipelines were designed and tested with the scheduler embedded in the model checkpoint. The optimal number of inference steps varies by scheduler, with no universal value that works best across all cases. To accommodate this, the num_inference_steps parameter in the pipeline's __call__ method defaults to None (see the API reference). Unless set explicitly, it inherits the value from the default_denoising_steps field in the checkpoint configuration file (model_index.json). This ensures high-quality predictions when invoking the pipeline with only the image argument.

See also Marigold usage examples.

Marigold Depth Prediction API

mindone.diffusers.MarigoldDepthPipeline

Bases: DiffusionPipeline

Pipeline for monocular depth estimation using the Marigold method: https://marigoldmonodepth.github.io.

This model inherits from [DiffusionPipeline]. Check the superclass documentation for the generic methods the library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)

PARAMETER	DESCRIPTION
unet	Conditional U-Net to denoise the depth latent, conditioned on image latent. TYPE: `UNet2DConditionModel`
vae	Variational Auto-Encoder (VAE) Model to encode and decode images and predictions to and from latent representations. TYPE: `AutoencoderKL`
scheduler	A scheduler to be used in combination with unet to denoise the encoded image latents. TYPE: `DDIMScheduler` or `LCMScheduler`
text_encoder	Text-encoder, for empty text embedding. TYPE: `CLIPTextModel`
tokenizer	CLIP tokenizer. TYPE: `CLIPTokenizer`
prediction_type	Type of predictions made by the model. TYPE: `str`, *optional* DEFAULT: None
scale_invariant	A model property specifying whether the predicted depth maps are scale-invariant. This value must be set in the model config. When used together with the shift_invariant=True flag, the model is also called "affine-invariant". NB: overriding this value is not supported. TYPE: `bool`, *optional* DEFAULT: True
shift_invariant	A model property specifying whether the predicted depth maps are shift-invariant. This value must be set in the model config. When used together with the scale_invariant=True flag, the model is also called "affine-invariant". NB: overriding this value is not supported. TYPE: `bool`, *optional* DEFAULT: True
default_denoising_steps	The minimum number of denoising diffusion steps that are required to produce a prediction of reasonable quality with the given model. This value must be set in the model config. When the pipeline is called without explicitly setting num_inference_steps, the default value is used. This is required to ensure reasonable results with various model flavors compatible with the pipeline, such as those relying on very short denoising schedules (LCMScheduler) and those with full diffusion schedules (DDIMScheduler). TYPE: `int`, *optional* DEFAULT: None
default_processing_resolution	The recommended value of the processing_resolution parameter of the pipeline. This value must be set in the model config. When the pipeline is called without explicitly setting processing_resolution, the default value is used. This is required to ensure reasonable results with various model flavors trained with varying optimal processing resolution values. TYPE: `int`, *optional* DEFAULT: None

Source code in mindone/diffusers/pipelines/marigold/pipeline_marigold_depth.py

class MarigoldDepthPipeline(DiffusionPipeline):
    """
    Pipeline for monocular depth estimation using the Marigold method: https://marigoldmonodepth.github.io.

    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)

    Args:
        unet (`UNet2DConditionModel`):
            Conditional U-Net to denoise the depth latent, conditioned on image latent.
        vae (`AutoencoderKL`):
            Variational Auto-Encoder (VAE) Model to encode and decode images and predictions to and from latent
            representations.
        scheduler (`DDIMScheduler` or `LCMScheduler`):
            A scheduler to be used in combination with `unet` to denoise the encoded image latents.
        text_encoder (`CLIPTextModel`):
            Text-encoder, for empty text embedding.
        tokenizer (`CLIPTokenizer`):
            CLIP tokenizer.
        prediction_type (`str`, *optional*):
            Type of predictions made by the model.
        scale_invariant (`bool`, *optional*):
            A model property specifying whether the predicted depth maps are scale-invariant. This value must be set in
            the model config. When used together with the `shift_invariant=True` flag, the model is also called
            "affine-invariant". NB: overriding this value is not supported.
        shift_invariant (`bool`, *optional*):
            A model property specifying whether the predicted depth maps are shift-invariant. This value must be set in
            the model config. When used together with the `scale_invariant=True` flag, the model is also called
            "affine-invariant". NB: overriding this value is not supported.
        default_denoising_steps (`int`, *optional*):
            The minimum number of denoising diffusion steps that are required to produce a prediction of reasonable
            quality with the given model. This value must be set in the model config. When the pipeline is called
            without explicitly setting `num_inference_steps`, the default value is used. This is required to ensure
            reasonable results with various model flavors compatible with the pipeline, such as those relying on very
            short denoising schedules (`LCMScheduler`) and those with full diffusion schedules (`DDIMScheduler`).
        default_processing_resolution (`int`, *optional*):
            The recommended value of the `processing_resolution` parameter of the pipeline. This value must be set in
            the model config. When the pipeline is called without explicitly setting `processing_resolution`, the
            default value is used. This is required to ensure reasonable results with various model flavors trained
            with varying optimal processing resolution values.
    """

    model_cpu_offload_seq = "text_encoder->unet->vae"
    supported_prediction_types = ("depth", "disparity")

    def __init__(
        self,
        unet: UNet2DConditionModel,
        vae: AutoencoderKL,
        scheduler: Union[DDIMScheduler, LCMScheduler],
        text_encoder: CLIPTextModel,
        tokenizer: CLIPTokenizer,
        prediction_type: Optional[str] = None,
        scale_invariant: Optional[bool] = True,
        shift_invariant: Optional[bool] = True,
        default_denoising_steps: Optional[int] = None,
        default_processing_resolution: Optional[int] = None,
    ):
        super().__init__()

        if prediction_type not in self.supported_prediction_types:
            logger.warning(
                f"Potentially unsupported `prediction_type='{prediction_type}'`; values supported by the pipeline: "
                f"{self.supported_prediction_types}."
            )

        self.register_modules(
            unet=unet,
            vae=vae,
            scheduler=scheduler,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
        )
        self.register_to_config(
            prediction_type=prediction_type,
            scale_invariant=scale_invariant,
            shift_invariant=shift_invariant,
            default_denoising_steps=default_denoising_steps,
            default_processing_resolution=default_processing_resolution,
        )

        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8

        self.scale_invariant = scale_invariant
        self.shift_invariant = shift_invariant
        self.default_denoising_steps = default_denoising_steps
        self.default_processing_resolution = default_processing_resolution

        self.empty_text_embedding = None

        self.image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor)

    def check_inputs(
        self,
        image: PipelineImageInput,
        num_inference_steps: int,
        ensemble_size: int,
        processing_resolution: int,
        resample_method_input: str,
        resample_method_output: str,
        batch_size: int,
        ensembling_kwargs: Optional[Dict[str, Any]],
        latents: Optional[ms.Tensor],
        generator: Optional[Union[np.random.Generator, List[np.random.Generator]]],
        output_type: str,
        output_uncertainty: bool,
    ) -> int:
        actual_vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        if actual_vae_scale_factor != self.vae_scale_factor:
            raise ValueError(
                f"`vae_scale_factor` computed at initialization ({self.vae_scale_factor}) differs from the actual one ({actual_vae_scale_factor})."
            )
        if num_inference_steps is None:
            raise ValueError("`num_inference_steps` is not specified and could not be resolved from the model config.")
        if num_inference_steps < 1:
            raise ValueError("`num_inference_steps` must be positive.")
        if ensemble_size < 1:
            raise ValueError("`ensemble_size` must be positive.")
        if ensemble_size == 2:
            logger.warning(
                "`ensemble_size` == 2 results are similar to no ensembling (1); "
                "consider increasing the value to at least 3."
            )
        if ensemble_size > 1 and (self.scale_invariant or self.shift_invariant) and not is_scipy_available():
            raise ImportError("Make sure to install scipy if you want to use ensembling.")
        if ensemble_size == 1 and output_uncertainty:
            raise ValueError(
                "Computing uncertainty by setting `output_uncertainty=True` also requires setting `ensemble_size` "
                "greater than 1."
            )
        if processing_resolution is None:
            raise ValueError(
                "`processing_resolution` is not specified and could not be resolved from the model config."
            )
        if processing_resolution < 0:
            raise ValueError(
                "`processing_resolution` must be non-negative: 0 for native resolution, or any positive value for "
                "downsampled processing."
            )
        if processing_resolution % self.vae_scale_factor != 0:
            raise ValueError(f"`processing_resolution` must be a multiple of {self.vae_scale_factor}.")
        if resample_method_input not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
            raise ValueError(
                "`resample_method_input` takes string values compatible with PIL library: "
                "nearest, nearest-exact, bilinear, bicubic, area."
            )
        if resample_method_output not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
            raise ValueError(
                "`resample_method_output` takes string values compatible with PIL library: "
                "nearest, nearest-exact, bilinear, bicubic, area."
            )
        if batch_size < 1:
            raise ValueError("`batch_size` must be positive.")
        if output_type not in ["pt", "np"]:
            raise ValueError("`output_type` must be one of `pt` or `np`.")
        if latents is not None and generator is not None:
            raise ValueError("`latents` and `generator` cannot be used together.")
        if ensembling_kwargs is not None:
            if not isinstance(ensembling_kwargs, dict):
                raise ValueError("`ensembling_kwargs` must be a dictionary.")
            if "reduction" in ensembling_kwargs and ensembling_kwargs["reduction"] not in ("mean", "median"):
                raise ValueError("`ensembling_kwargs['reduction']` can be either `'mean'` or `'median'`.")

        # image checks
        num_images = 0
        W, H = None, None
        if not isinstance(image, list):
            image = [image]
        for i, img in enumerate(image):
            if isinstance(img, np.ndarray) or ops.is_tensor(img):
                if img.ndim not in (2, 3, 4):
                    raise ValueError(f"`image[{i}]` has unsupported dimensions or shape: {img.shape}.")
                H_i, W_i = img.shape[-2:]
                N_i = 1
                if img.ndim == 4:
                    N_i = img.shape[0]
            elif isinstance(img, Image.Image):
                W_i, H_i = img.size
                N_i = 1
            else:
                raise ValueError(f"Unsupported `image[{i}]` type: {type(img)}.")
            if W is None:
                W, H = W_i, H_i
            elif (W, H) != (W_i, H_i):
                raise ValueError(
                    f"Input `image[{i}]` has incompatible dimensions {(W_i, H_i)} with the previous images {(W, H)}"
                )
            num_images += N_i

        # latents checks
        if latents is not None:
            if not ops.is_tensor(latents):
                raise ValueError("`latents` must be a ms.Tensor.")
            if latents.ndim != 4:
                raise ValueError(f"`latents` has unsupported dimensions or shape: {latents.shape}.")

            if processing_resolution > 0:
                max_orig = max(H, W)
                new_H = H * processing_resolution // max_orig
                new_W = W * processing_resolution // max_orig
                if new_H == 0 or new_W == 0:
                    raise ValueError(f"Extreme aspect ratio of the input image: [{W} x {H}]")
                W, H = new_W, new_H
            w = (W + self.vae_scale_factor - 1) // self.vae_scale_factor
            h = (H + self.vae_scale_factor - 1) // self.vae_scale_factor
            shape_expected = (num_images * ensemble_size, self.vae.config.latent_channels, h, w)

            if latents.shape != shape_expected:
                raise ValueError(f"`latents` has unexpected shape={latents.shape} expected={shape_expected}.")

        # generator checks
        if generator is not None:
            if isinstance(generator, list):
                if len(generator) != num_images * ensemble_size:
                    raise ValueError(
                        "The number of generators must match the total number of ensemble members for all input images."
                    )
            elif not isinstance(generator, np.random.Generator):
                raise ValueError(f"Unsupported generator type: {type(generator)}.")

        return num_images

    def progress_bar(self, iterable=None, total=None, desc=None, leave=True):
        if not hasattr(self, "_progress_bar_config"):
            self._progress_bar_config = {}
        elif not isinstance(self._progress_bar_config, dict):
            raise ValueError(
                f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}."
            )

        progress_bar_config = dict(**self._progress_bar_config)
        progress_bar_config["desc"] = progress_bar_config.get("desc", desc)
        progress_bar_config["leave"] = progress_bar_config.get("leave", leave)
        if iterable is not None:
            return tqdm(iterable, **progress_bar_config)
        elif total is not None:
            return tqdm(total=total, **progress_bar_config)
        else:
            raise ValueError("Either `total` or `iterable` has to be defined.")

    def __call__(
        self,
        image: PipelineImageInput,
        num_inference_steps: Optional[int] = None,
        ensemble_size: int = 1,
        processing_resolution: Optional[int] = None,
        match_input_resolution: bool = True,
        resample_method_input: str = "bilinear",
        resample_method_output: str = "bilinear",
        batch_size: int = 1,
        ensembling_kwargs: Optional[Dict[str, Any]] = None,
        latents: Optional[Union[ms.Tensor, List[ms.Tensor]]] = None,
        generator: Optional[Union[np.random.Generator, List[np.random.Generator]]] = None,
        output_type: str = "np",
        output_uncertainty: bool = False,
        output_latent: bool = False,
        return_dict: bool = False,
    ):
        """
        Function invoked when calling the pipeline.

        Args:
            image (`PIL.Image.Image`, `np.ndarray`, `ms.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`),
                `List[ms.Tensor]`: An input image or images used as an input for the depth estimation task. For
                arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible
                by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or
                three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the
                same width and height.
            num_inference_steps (`int`, *optional*, defaults to `None`):
                Number of denoising diffusion steps during inference. The default value `None` results in automatic
                selection.
            ensemble_size (`int`, defaults to `1`):
                Number of ensemble predictions. Higher values result in measurable improvements and visual degradation.
            processing_resolution (`int`, *optional*, defaults to `None`):
                Effective processing resolution. When set to `0`, matches the larger input image dimension. This
                produces crisper predictions, but may also lead to the overall loss of global context. The default
                value `None` resolves to the optimal value from the model config.
            match_input_resolution (`bool`, *optional*, defaults to `True`):
                When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer
                side of the output will equal to `processing_resolution`.
            resample_method_input (`str`, *optional*, defaults to `"bilinear"`):
                Resampling method used to resize input images to `processing_resolution`. The accepted values are:
                `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
            resample_method_output (`str`, *optional*, defaults to `"bilinear"`):
                Resampling method used to resize output predictions to match the input resolution. The accepted values
                are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
            batch_size (`int`, *optional*, defaults to `1`):
                Batch size; only matters when setting `ensemble_size` or passing a tensor of images.
            ensembling_kwargs (`dict`, *optional*, defaults to `None`)
                Extra dictionary with arguments for precise ensembling control. The following options are available:
                - reduction (`str`, *optional*, defaults to `"median"`): Defines the ensembling function applied in
                  every pixel location, can be either `"median"` or `"mean"`.
                - regularizer_strength (`float`, *optional*, defaults to `0.02`): Strength of the regularizer that
                  pulls the aligned predictions to the unit range from 0 to 1.
                - max_iter (`int`, *optional*, defaults to `2`): Maximum number of the alignment solver steps. Refer to
                  `scipy.optimize.minimize` function, `options` argument.
                - tol (`float`, *optional*, defaults to `1e-3`): Alignment solver tolerance. The solver stops when the
                  tolerance is reached.
                - max_res (`int`, *optional*, defaults to `None`): Resolution at which the alignment is performed;
                  `None` matches the `processing_resolution`.
            latents (`ms.Tensor`, or `List[ms.Tensor]`, *optional*, defaults to `None`):
                Latent noise tensors to replace the random initialization. These can be taken from the previous
                function call's output.
            generator (`np.random.Generator`, or `List[np.random.Generator]`, *optional*, defaults to `None`):
                Random number generator object to ensure reproducibility.
            output_type (`str`, *optional*, defaults to `"np"`):
                Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted
                values are: `"np"` (numpy array) or `"pt"` (mindspore tensor).
            output_uncertainty (`bool`, *optional*, defaults to `False`):
                When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that
                the `ensemble_size` argument is set to a value above 2.
            output_latent (`bool`, *optional*, defaults to `False`):
                When enabled, the output's `latent` field contains the latent codes corresponding to the predictions
                within the ensemble. These codes can be saved, modified, and used for subsequent calls with the
                `latents` argument.
            return_dict (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~pipelines.marigold.MarigoldDepthOutput`] instead of a plain tuple.

        Examples:

        Returns:
            [`~pipelines.marigold.MarigoldDepthOutput`] or `tuple`:
                If `return_dict` is `True`, [`~pipelines.marigold.MarigoldDepthOutput`] is returned, otherwise a
                `tuple` is returned where the first element is the prediction, the second element is the uncertainty
                (or `None`), and the third is the latent (or `None`).
        """

        # 0. Resolving variables.
        dtype = self.dtype

        # Model-specific optimal default values leading to fast and reasonable results.
        if num_inference_steps is None:
            num_inference_steps = self.default_denoising_steps
        if processing_resolution is None:
            processing_resolution = self.default_processing_resolution

        # 1. Check inputs.
        num_images = self.check_inputs(
            image,
            num_inference_steps,
            ensemble_size,
            processing_resolution,
            resample_method_input,
            resample_method_output,
            batch_size,
            ensembling_kwargs,
            latents,
            generator,
            output_type,
            output_uncertainty,
        )

        # 2. Prepare empty text conditioning.
        # Model invocation: self.tokenizer, self.text_encoder.
        if self.empty_text_embedding is None:
            prompt = ""
            text_inputs = self.tokenizer(
                prompt,
                padding="do_not_pad",
                max_length=self.tokenizer.model_max_length,
                truncation=True,
                return_tensors="np",
            )
            text_input_ids = ms.Tensor.from_numpy(text_inputs.input_ids)
            self.empty_text_embedding = self.text_encoder(text_input_ids)[0]  # [1,2,1024]

        # 3. Preprocess input images. This function loads input image or images of compatible dimensions `(H, W)`,
        # optionally downsamples them to the `processing_resolution` `(PH, PW)`, where
        # `max(PH, PW) == processing_resolution`, and pads the dimensions to `(PPH, PPW)` such that these values are
        # divisible by the latent space downscaling factor (typically 8 in Stable Diffusion). The default value `None`
        # of `processing_resolution` resolves to the optimal value from the model config. It is a recommended mode of
        # operation and leads to the most reasonable results. Using the native image resolution or any other processing
        # resolution can lead to loss of either fine details or global context in the output predictions.
        image, padding, original_resolution = self.image_processor.preprocess(
            image, processing_resolution, resample_method_input, dtype
        )  # [N,3,PPH,PPW]

        # 4. Encode input image into latent space. At this step, each of the `N` input images is represented with `E`
        # ensemble members. Each ensemble member is an independent diffused prediction, just initialized independently.
        # Latents of each such predictions across all input images and all ensemble members are represented in the
        # `pred_latent` variable. The variable `image_latent` is of the same shape: it contains each input image encoded
        # into latent space and replicated `E` times. The latents can be either generated (see `generator` to ensure
        # reproducibility), or passed explicitly via the `latents` argument. The latter can be set outside the pipeline
        # code. This behavior can be achieved by setting the `output_latent` argument to `True`. The latent space
        # dimensions are `(h, w)`. Encoding into latent space happens in batches of size `batch_size`.
        # Model invocation: self.vae.encoder.
        image_latent, pred_latent = self.prepare_latents(
            image, latents, generator, ensemble_size, batch_size
        )  # [N*E,4,h,w], [N*E,4,h,w]

        del image

        batch_empty_text_embedding = self.empty_text_embedding.to(dtype=dtype).tile((batch_size, 1, 1))  # [B,1024,2]

        # 5. Process the denoising loop. All `N * E` latents are processed sequentially in batches of size `batch_size`.
        # The unet model takes concatenated latent spaces of the input image and the predicted modality as an input, and
        # outputs noise for the predicted modality's latent space. The number of denoising diffusion steps is defined by
        # `num_inference_steps`. It is either set directly, or resolves to the optimal value specific to the loaded
        # model.
        # Model invocation: self.unet.
        pred_latents = []

        for i in self.progress_bar(
            range(0, num_images * ensemble_size, batch_size), leave=True, desc="Marigold predictions..."
        ):
            batch_image_latent = image_latent[i : i + batch_size]  # [B,4,h,w]
            batch_pred_latent = pred_latent[i : i + batch_size]  # [B,4,h,w]
            effective_batch_size = batch_image_latent.shape[0]
            text = batch_empty_text_embedding[:effective_batch_size]  # [B,2,1024]

            self.scheduler.set_timesteps(num_inference_steps)
            for t in self.progress_bar(self.scheduler.timesteps, leave=False, desc="Diffusion steps..."):
                batch_latent = mint.cat([batch_image_latent, batch_pred_latent], dim=1)  # [B,8,h,w]
                noise = self.unet(batch_latent, t, encoder_hidden_states=text, return_dict=False)[0]  # [B,4,h,w]
                batch_pred_latent = self.scheduler.step(noise, t, batch_pred_latent, generator=generator)[
                    0
                ]  # [B,4,h,w]

            pred_latents.append(batch_pred_latent)

        pred_latent = mint.cat(pred_latents, dim=0)  # [N*E,4,h,w]

        del (
            pred_latents,
            image_latent,
            batch_empty_text_embedding,
            batch_image_latent,
            batch_pred_latent,
            text,
            batch_latent,
            noise,
        )

        # 6. Decode predictions from latent into pixel space. The resulting `N * E` predictions have shape `(PPH, PPW)`,
        # which requires slight postprocessing. Decoding into pixel space happens in batches of size `batch_size`.
        # Model invocation: self.vae.decoder.
        prediction = mint.cat(
            [
                self.decode_prediction(pred_latent[i : i + batch_size])
                for i in range(0, pred_latent.shape[0], batch_size)
            ],
            dim=0,
        )  # [N*E,1,PPH,PPW]

        if not output_latent:
            pred_latent = None

        # 7. Remove padding. The output shape is (PH, PW).
        prediction = self.image_processor.unpad_image(prediction, padding)  # [N*E,1,PH,PW]

        # 8. Ensemble and compute uncertainty (when `output_uncertainty` is set). This code treats each of the `N`
        # groups of `E` ensemble predictions independently. For each group it computes an ensembled prediction of shape
        # `(PH, PW)` and an optional uncertainty map of the same dimensions. After computing this pair of outputs for
        # each group independently, it stacks them respectively into batches of `N` almost final predictions and
        # uncertainty maps.
        uncertainty = None
        if ensemble_size > 1:
            prediction = prediction.reshape(num_images, ensemble_size, *prediction.shape[1:])  # [N,E,1,PH,PW]
            prediction = [
                self.ensemble_depth(
                    prediction[i],
                    self.scale_invariant,
                    self.shift_invariant,
                    output_uncertainty,
                    **(ensembling_kwargs or {}),
                )
                for i in range(num_images)
            ]  # [ [[1,1,PH,PW], [1,1,PH,PW]], ... ]
            prediction, uncertainty = zip(*prediction)  # [[1,1,PH,PW], ... ], [[1,1,PH,PW], ... ]
            prediction = mint.cat(prediction, dim=0)  # [N,1,PH,PW]
            if output_uncertainty:
                uncertainty = mint.cat(uncertainty, dim=0)  # [N,1,PH,PW]
            else:
                uncertainty = None

        # 9. If `match_input_resolution` is set, the output prediction and the uncertainty are upsampled to match the
        # input resolution `(H, W)`. This step may introduce upsampling artifacts, and therefore can be disabled.
        # Depending on the downstream use-case, upsampling can be also chosen based on the tolerated artifacts by
        # setting the `resample_method_output` parameter (e.g., to `"nearest"`).
        if match_input_resolution:
            prediction = self.image_processor.resize_antialias(
                prediction, original_resolution, resample_method_output, is_aa=False
            )  # [N,1,H,W]
            if uncertainty is not None and output_uncertainty:
                uncertainty = self.image_processor.resize_antialias(
                    uncertainty, original_resolution, resample_method_output, is_aa=False
                )  # [N,1,H,W]

        # 10. Prepare the final outputs.
        if output_type == "np":
            prediction = self.image_processor.ms_to_numpy(prediction)  # [N,H,W,1]
            if uncertainty is not None and output_uncertainty:
                uncertainty = self.image_processor.ms_to_numpy(uncertainty)  # [N,H,W,1]

        if not return_dict:
            return (prediction, uncertainty, pred_latent)

        return MarigoldDepthOutput(
            prediction=prediction,
            uncertainty=uncertainty,
            latent=pred_latent,
        )

    def prepare_latents(
        self,
        image: ms.Tensor,
        latents: Optional[ms.Tensor],
        generator: Optional[np.random.Generator],
        ensemble_size: int,
        batch_size: int,
    ) -> Tuple[ms.Tensor, ms.Tensor]:
        def retrieve_latents(encoder_output):
            assert ops.is_tensor(
                encoder_output
            ), "Could not access latents of provided encoder_output which is not a tensor"
            if hasattr(self.vae, "diag_gauss_dist"):
                return self.vae.diag_gauss_dist.mode(encoder_output)
            else:
                return encoder_output

        image_latent = mint.cat(
            [
                retrieve_latents(self.vae.encode(image[i : i + batch_size])[0])
                for i in range(0, image.shape[0], batch_size)
            ],
            dim=0,
        )  # [N,4,h,w]
        image_latent = image_latent * self.vae.config.scaling_factor
        image_latent = image_latent.repeat_interleave(ensemble_size, dim=0)  # [N*E,4,h,w]

        pred_latent = latents
        if pred_latent is None:
            pred_latent = randn_tensor(
                image_latent.shape,
                generator=generator,
                dtype=image_latent.dtype,
            )  # [N*E,4,h,w]

        return image_latent, pred_latent

    def decode_prediction(self, pred_latent: ms.Tensor) -> ms.Tensor:
        if pred_latent.ndim != 4 or pred_latent.shape[1] != self.vae.config.latent_channels:
            raise ValueError(
                f"Expecting 4D tensor of shape [B,{self.vae.config.latent_channels},H,W]; got {pred_latent.shape}."
            )

        prediction = self.vae.decode(pred_latent / self.vae.config.scaling_factor, return_dict=False)[0]  # [B,3,H,W]

        prediction = prediction.mean(dim=1, keepdim=True)  # [B,1,H,W]
        prediction = mint.clip(prediction, -1.0, 1.0)  # [B,1,H,W]
        prediction = (prediction + 1.0) / 2.0

        return prediction  # [B,1,H,W]

    @staticmethod
    def ensemble_depth(
        depth: ms.Tensor,
        scale_invariant: bool = True,
        shift_invariant: bool = True,
        output_uncertainty: bool = False,
        reduction: str = "median",
        regularizer_strength: float = 0.02,
        max_iter: int = 2,
        tol: float = 1e-3,
        max_res: int = 1024,
    ) -> Tuple[ms.Tensor, Optional[ms.Tensor]]:
        """
        Ensembles the depth maps represented by the `depth` tensor with expected shape `(B, 1, H, W)`, where B is the
        number of ensemble members for a given prediction of size `(H x W)`. Even though the function is designed for
        depth maps, it can also be used with disparity maps as long as the input tensor values are non-negative. The
        alignment happens when the predictions have one or more degrees of freedom, that is when they are either
        affine-invariant (`scale_invariant=True` and `shift_invariant=True`), or just scale-invariant (only
        `scale_invariant=True`). For absolute predictions (`scale_invariant=False` and `shift_invariant=False`)
        alignment is skipped and only ensembling is performed.

        Args:
            depth (`ms.Tensor`):
                Input ensemble depth maps.
            scale_invariant (`bool`, *optional*, defaults to `True`):
                Whether to treat predictions as scale-invariant.
            shift_invariant (`bool`, *optional*, defaults to `True`):
                Whether to treat predictions as shift-invariant.
            output_uncertainty (`bool`, *optional*, defaults to `False`):
                Whether to output uncertainty map.
            reduction (`str`, *optional*, defaults to `"median"`):
                Reduction method used to ensemble aligned predictions. The accepted values are: `"mean"` and
                `"median"`.
            regularizer_strength (`float`, *optional*, defaults to `0.02`):
                Strength of the regularizer that pulls the aligned predictions to the unit range from 0 to 1.
            max_iter (`int`, *optional*, defaults to `2`):
                Maximum number of the alignment solver steps. Refer to `scipy.optimize.minimize` function, `options`
                argument.
            tol (`float`, *optional*, defaults to `1e-3`):
                Alignment solver tolerance. The solver stops when the tolerance is reached.
            max_res (`int`, *optional*, defaults to `1024`):
                Resolution at which the alignment is performed; `None` matches the `processing_resolution`.
        Returns:
            A tensor of aligned and ensembled depth maps and optionally a tensor of uncertainties of the same shape:
            `(1, 1, H, W)`.
        """
        if depth.ndim != 4 or depth.shape[1] != 1:
            raise ValueError(f"Expecting 4D tensor of shape [B,1,H,W]; got {depth.shape}.")
        if reduction not in ("mean", "median"):
            raise ValueError(f"Unrecognized reduction method: {reduction}.")
        if not scale_invariant and shift_invariant:
            raise ValueError("Pure shift-invariant ensembling is not supported.")

        def init_param(depth: ms.Tensor):
            init_min = depth.reshape(ensemble_size, -1).min(dim=1)
            init_max = depth.reshape(ensemble_size, -1).max(dim=1)

            if scale_invariant and shift_invariant:
                init_s = 1.0 / (init_max - init_min).clamp(min=1e-6)
                init_t = -init_s * init_min
                param = mint.cat((init_s, init_t)).numpy()
            elif scale_invariant:
                init_s = 1.0 / init_max.clamp(min=1e-6)
                param = init_s.numpy()
            else:
                raise ValueError("Unrecognized alignment.")

            return param

        def align(depth: ms.Tensor, param: np.ndarray) -> ms.Tensor:
            if scale_invariant and shift_invariant:
                s, t = np.split(param, 2)
                s = ms.Tensor.from_numpy(s).to(depth.dtype).view(ensemble_size, 1, 1, 1)
                t = ms.Tensor.from_numpy(t).to(depth.dtype).view(ensemble_size, 1, 1, 1)
                out = depth * s + t
            elif scale_invariant:
                s = ms.Tensor.from_numpy(param).to(depth.dtype).view(ensemble_size, 1, 1, 1)
                out = depth * s
            else:
                raise ValueError("Unrecognized alignment.")
            return out

        def ensemble(
            depth_aligned: ms.Tensor, return_uncertainty: bool = False
        ) -> Tuple[ms.Tensor, Optional[ms.Tensor]]:
            uncertainty = None
            if reduction == "mean":
                prediction = mint.mean(depth_aligned, dim=0, keepdim=True)
                if return_uncertainty:
                    uncertainty = mint.std(depth_aligned, dim=0, keepdim=True)
            elif reduction == "median":
                # ops.median has two return values and does not supported some data-type
                prediction = mint.median(depth_aligned.float(), dim=0, keepdim=True)[0]
                prediction = prediction.to(depth_aligned.dtype)
                if return_uncertainty:
                    uncertainty = mint.median(mint.abs(depth_aligned - prediction).float(), dim=0, keepdim=True)[0]
                    uncertainty = uncertainty.to(depth_aligned.dtype)
            else:
                raise ValueError(f"Unrecognized reduction method: {reduction}.")
            return prediction, uncertainty

        def cost_fn(param: np.ndarray, depth: ms.Tensor) -> float:
            cost = 0.0
            depth_aligned = align(depth, param)

            # todo: unavailable mint interface combinations
            for i, j in ops.combinations(mint.arange(ensemble_size)):
                diff = depth_aligned[i] - depth_aligned[j]
                cost += (diff**2).mean().sqrt().item()

            if regularizer_strength > 0:
                prediction, _ = ensemble(depth_aligned, return_uncertainty=False)
                err_near = (prediction.min()).abs().item()
                err_far = (1.0 - prediction.max()).abs().item()
                cost += (err_near + err_far) * regularizer_strength

            return cost

        def compute_param(depth: ms.Tensor):
            import scipy

            depth_to_align = depth.to(ms.float32)
            if max_res is not None and max(depth_to_align.shape[2:]) > max_res:
                depth_to_align = MarigoldImageProcessor.resize_to_max_edge(depth_to_align, max_res, "nearest-exact")

            param = init_param(depth_to_align)

            res = scipy.optimize.minimize(
                partial(cost_fn, depth=depth_to_align),
                param,
                method="BFGS",
                tol=tol,
                options={"maxiter": max_iter, "disp": False},
            )

            return res.x

        requires_aligning = scale_invariant or shift_invariant
        ensemble_size = depth.shape[0]

        if requires_aligning:
            param = compute_param(depth)
            depth = align(depth, param)

        depth, uncertainty = ensemble(depth, return_uncertainty=output_uncertainty)

        depth_max = depth.max()
        if scale_invariant and shift_invariant:
            depth_min = mint.min(depth)
        elif scale_invariant:
            depth_min = 0
        else:
            raise ValueError("Unrecognized alignment.")
        depth_range = (depth_max - depth_min).clamp(min=1e-6)
        depth = (depth - depth_min) / depth_range
        if output_uncertainty:
            uncertainty /= depth_range

        return depth, uncertainty  # [1,1,H,W], [1,1,H,W]

mindone.diffusers.MarigoldDepthPipeline.__call__(image, num_inference_steps=None, ensemble_size=1, processing_resolution=None, match_input_resolution=True, resample_method_input='bilinear', resample_method_output='bilinear', batch_size=1, ensembling_kwargs=None, latents=None, generator=None, output_type='np', output_uncertainty=False, output_latent=False, return_dict=False)

Function invoked when calling the pipeline.

PARAMETER	DESCRIPTION
num_inference_steps	Number of denoising diffusion steps during inference. The default value None results in automatic selection. TYPE: `int`, *optional*, defaults to `None` DEFAULT: None
ensemble_size	Number of ensemble predictions. Higher values result in measurable improvements and visual degradation. TYPE: `int`, defaults to `1` DEFAULT: 1
processing_resolution	Effective processing resolution. When set to 0, matches the larger input image dimension. This produces crisper predictions, but may also lead to the overall loss of global context. The default value None resolves to the optimal value from the model config. TYPE: `int`, *optional*, defaults to `None` DEFAULT: None
match_input_resolution	When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer side of the output will equal to processing_resolution. TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True
resample_method_input	Resampling method used to resize input images to processing_resolution. The accepted values are: "nearest", "nearest-exact", "bilinear", "bicubic", or "area". TYPE: `str`, *optional*, defaults to `"bilinear"` DEFAULT: 'bilinear'
resample_method_output	Resampling method used to resize output predictions to match the input resolution. The accepted values are "nearest", "nearest-exact", "bilinear", "bicubic", or "area". TYPE: `str`, *optional*, defaults to `"bilinear"` DEFAULT: 'bilinear'
batch_size	Batch size; only matters when setting ensemble_size or passing a tensor of images. TYPE: `int`, *optional*, defaults to `1` DEFAULT: 1
latents	Latent noise tensors to replace the random initialization. These can be taken from the previous function call's output. TYPE: `ms.Tensor`, or `List[ms.Tensor]`, *optional*, defaults to `None` DEFAULT: None
generator	Random number generator object to ensure reproducibility. TYPE: `np.random.Generator`, or `List[np.random.Generator]`, *optional*, defaults to `None` DEFAULT: None
output_type	Preferred format of the output's prediction and the optional uncertainty fields. The accepted values are: "np" (numpy array) or "pt" (mindspore tensor). TYPE: `str`, *optional*, defaults to `"np"` DEFAULT: 'np'
output_uncertainty	When enabled, the output's uncertainty field contains the predictive uncertainty map, provided that the ensemble_size argument is set to a value above 2. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
output_latent	When enabled, the output's latent field contains the latent codes corresponding to the predictions within the ensemble. These codes can be saved, modified, and used for subsequent calls with the latents argument. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
return_dict	Whether or not to return a [~pipelines.marigold.MarigoldDepthOutput] instead of a plain tuple. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False

RETURNS	DESCRIPTION
	[~pipelines.marigold.MarigoldDepthOutput] or tuple: If return_dict is True, [~pipelines.marigold.MarigoldDepthOutput] is returned, otherwise a tuple is returned where the first element is the prediction, the second element is the uncertainty (or None), and the third is the latent (or None).

Source code in mindone/diffusers/pipelines/marigold/pipeline_marigold_depth.py

def __call__(
    self,
    image: PipelineImageInput,
    num_inference_steps: Optional[int] = None,
    ensemble_size: int = 1,
    processing_resolution: Optional[int] = None,
    match_input_resolution: bool = True,
    resample_method_input: str = "bilinear",
    resample_method_output: str = "bilinear",
    batch_size: int = 1,
    ensembling_kwargs: Optional[Dict[str, Any]] = None,
    latents: Optional[Union[ms.Tensor, List[ms.Tensor]]] = None,
    generator: Optional[Union[np.random.Generator, List[np.random.Generator]]] = None,
    output_type: str = "np",
    output_uncertainty: bool = False,
    output_latent: bool = False,
    return_dict: bool = False,
):
    """
    Function invoked when calling the pipeline.

    Args:
        image (`PIL.Image.Image`, `np.ndarray`, `ms.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`),
            `List[ms.Tensor]`: An input image or images used as an input for the depth estimation task. For
            arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible
            by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or
            three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the
            same width and height.
        num_inference_steps (`int`, *optional*, defaults to `None`):
            Number of denoising diffusion steps during inference. The default value `None` results in automatic
            selection.
        ensemble_size (`int`, defaults to `1`):
            Number of ensemble predictions. Higher values result in measurable improvements and visual degradation.
        processing_resolution (`int`, *optional*, defaults to `None`):
            Effective processing resolution. When set to `0`, matches the larger input image dimension. This
            produces crisper predictions, but may also lead to the overall loss of global context. The default
            value `None` resolves to the optimal value from the model config.
        match_input_resolution (`bool`, *optional*, defaults to `True`):
            When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer
            side of the output will equal to `processing_resolution`.
        resample_method_input (`str`, *optional*, defaults to `"bilinear"`):
            Resampling method used to resize input images to `processing_resolution`. The accepted values are:
            `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
        resample_method_output (`str`, *optional*, defaults to `"bilinear"`):
            Resampling method used to resize output predictions to match the input resolution. The accepted values
            are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
        batch_size (`int`, *optional*, defaults to `1`):
            Batch size; only matters when setting `ensemble_size` or passing a tensor of images.
        ensembling_kwargs (`dict`, *optional*, defaults to `None`)
            Extra dictionary with arguments for precise ensembling control. The following options are available:
            - reduction (`str`, *optional*, defaults to `"median"`): Defines the ensembling function applied in
              every pixel location, can be either `"median"` or `"mean"`.
            - regularizer_strength (`float`, *optional*, defaults to `0.02`): Strength of the regularizer that
              pulls the aligned predictions to the unit range from 0 to 1.
            - max_iter (`int`, *optional*, defaults to `2`): Maximum number of the alignment solver steps. Refer to
              `scipy.optimize.minimize` function, `options` argument.
            - tol (`float`, *optional*, defaults to `1e-3`): Alignment solver tolerance. The solver stops when the
              tolerance is reached.
            - max_res (`int`, *optional*, defaults to `None`): Resolution at which the alignment is performed;
              `None` matches the `processing_resolution`.
        latents (`ms.Tensor`, or `List[ms.Tensor]`, *optional*, defaults to `None`):
            Latent noise tensors to replace the random initialization. These can be taken from the previous
            function call's output.
        generator (`np.random.Generator`, or `List[np.random.Generator]`, *optional*, defaults to `None`):
            Random number generator object to ensure reproducibility.
        output_type (`str`, *optional*, defaults to `"np"`):
            Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted
            values are: `"np"` (numpy array) or `"pt"` (mindspore tensor).
        output_uncertainty (`bool`, *optional*, defaults to `False`):
            When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that
            the `ensemble_size` argument is set to a value above 2.
        output_latent (`bool`, *optional*, defaults to `False`):
            When enabled, the output's `latent` field contains the latent codes corresponding to the predictions
            within the ensemble. These codes can be saved, modified, and used for subsequent calls with the
            `latents` argument.
        return_dict (`bool`, *optional*, defaults to `False`):
            Whether or not to return a [`~pipelines.marigold.MarigoldDepthOutput`] instead of a plain tuple.

    Examples:

    Returns:
        [`~pipelines.marigold.MarigoldDepthOutput`] or `tuple`:
            If `return_dict` is `True`, [`~pipelines.marigold.MarigoldDepthOutput`] is returned, otherwise a
            `tuple` is returned where the first element is the prediction, the second element is the uncertainty
            (or `None`), and the third is the latent (or `None`).
    """

    # 0. Resolving variables.
    dtype = self.dtype

    # Model-specific optimal default values leading to fast and reasonable results.
    if num_inference_steps is None:
        num_inference_steps = self.default_denoising_steps
    if processing_resolution is None:
        processing_resolution = self.default_processing_resolution

    # 1. Check inputs.
    num_images = self.check_inputs(
        image,
        num_inference_steps,
        ensemble_size,
        processing_resolution,
        resample_method_input,
        resample_method_output,
        batch_size,
        ensembling_kwargs,
        latents,
        generator,
        output_type,
        output_uncertainty,
    )

    # 2. Prepare empty text conditioning.
    # Model invocation: self.tokenizer, self.text_encoder.
    if self.empty_text_embedding is None:
        prompt = ""
        text_inputs = self.tokenizer(
            prompt,
            padding="do_not_pad",
            max_length=self.tokenizer.model_max_length,
            truncation=True,
            return_tensors="np",
        )
        text_input_ids = ms.Tensor.from_numpy(text_inputs.input_ids)
        self.empty_text_embedding = self.text_encoder(text_input_ids)[0]  # [1,2,1024]

    # 3. Preprocess input images. This function loads input image or images of compatible dimensions `(H, W)`,
    # optionally downsamples them to the `processing_resolution` `(PH, PW)`, where
    # `max(PH, PW) == processing_resolution`, and pads the dimensions to `(PPH, PPW)` such that these values are
    # divisible by the latent space downscaling factor (typically 8 in Stable Diffusion). The default value `None`
    # of `processing_resolution` resolves to the optimal value from the model config. It is a recommended mode of
    # operation and leads to the most reasonable results. Using the native image resolution or any other processing
    # resolution can lead to loss of either fine details or global context in the output predictions.
    image, padding, original_resolution = self.image_processor.preprocess(
        image, processing_resolution, resample_method_input, dtype
    )  # [N,3,PPH,PPW]

    # 4. Encode input image into latent space. At this step, each of the `N` input images is represented with `E`
    # ensemble members. Each ensemble member is an independent diffused prediction, just initialized independently.
    # Latents of each such predictions across all input images and all ensemble members are represented in the
    # `pred_latent` variable. The variable `image_latent` is of the same shape: it contains each input image encoded
    # into latent space and replicated `E` times. The latents can be either generated (see `generator` to ensure
    # reproducibility), or passed explicitly via the `latents` argument. The latter can be set outside the pipeline
    # code. This behavior can be achieved by setting the `output_latent` argument to `True`. The latent space
    # dimensions are `(h, w)`. Encoding into latent space happens in batches of size `batch_size`.
    # Model invocation: self.vae.encoder.
    image_latent, pred_latent = self.prepare_latents(
        image, latents, generator, ensemble_size, batch_size
    )  # [N*E,4,h,w], [N*E,4,h,w]

    del image

    batch_empty_text_embedding = self.empty_text_embedding.to(dtype=dtype).tile((batch_size, 1, 1))  # [B,1024,2]

    # 5. Process the denoising loop. All `N * E` latents are processed sequentially in batches of size `batch_size`.
    # The unet model takes concatenated latent spaces of the input image and the predicted modality as an input, and
    # outputs noise for the predicted modality's latent space. The number of denoising diffusion steps is defined by
    # `num_inference_steps`. It is either set directly, or resolves to the optimal value specific to the loaded
    # model.
    # Model invocation: self.unet.
    pred_latents = []

    for i in self.progress_bar(
        range(0, num_images * ensemble_size, batch_size), leave=True, desc="Marigold predictions..."
    ):
        batch_image_latent = image_latent[i : i + batch_size]  # [B,4,h,w]
        batch_pred_latent = pred_latent[i : i + batch_size]  # [B,4,h,w]
        effective_batch_size = batch_image_latent.shape[0]
        text = batch_empty_text_embedding[:effective_batch_size]  # [B,2,1024]

        self.scheduler.set_timesteps(num_inference_steps)
        for t in self.progress_bar(self.scheduler.timesteps, leave=False, desc="Diffusion steps..."):
            batch_latent = mint.cat([batch_image_latent, batch_pred_latent], dim=1)  # [B,8,h,w]
            noise = self.unet(batch_latent, t, encoder_hidden_states=text, return_dict=False)[0]  # [B,4,h,w]
            batch_pred_latent = self.scheduler.step(noise, t, batch_pred_latent, generator=generator)[
                0
            ]  # [B,4,h,w]

        pred_latents.append(batch_pred_latent)

    pred_latent = mint.cat(pred_latents, dim=0)  # [N*E,4,h,w]

    del (
        pred_latents,
        image_latent,
        batch_empty_text_embedding,
        batch_image_latent,
        batch_pred_latent,
        text,
        batch_latent,
        noise,
    )

    # 6. Decode predictions from latent into pixel space. The resulting `N * E` predictions have shape `(PPH, PPW)`,
    # which requires slight postprocessing. Decoding into pixel space happens in batches of size `batch_size`.
    # Model invocation: self.vae.decoder.
    prediction = mint.cat(
        [
            self.decode_prediction(pred_latent[i : i + batch_size])
            for i in range(0, pred_latent.shape[0], batch_size)
        ],
        dim=0,
    )  # [N*E,1,PPH,PPW]

    if not output_latent:
        pred_latent = None

    # 7. Remove padding. The output shape is (PH, PW).
    prediction = self.image_processor.unpad_image(prediction, padding)  # [N*E,1,PH,PW]

    # 8. Ensemble and compute uncertainty (when `output_uncertainty` is set). This code treats each of the `N`
    # groups of `E` ensemble predictions independently. For each group it computes an ensembled prediction of shape
    # `(PH, PW)` and an optional uncertainty map of the same dimensions. After computing this pair of outputs for
    # each group independently, it stacks them respectively into batches of `N` almost final predictions and
    # uncertainty maps.
    uncertainty = None
    if ensemble_size > 1:
        prediction = prediction.reshape(num_images, ensemble_size, *prediction.shape[1:])  # [N,E,1,PH,PW]
        prediction = [
            self.ensemble_depth(
                prediction[i],
                self.scale_invariant,
                self.shift_invariant,
                output_uncertainty,
                **(ensembling_kwargs or {}),
            )
            for i in range(num_images)
        ]  # [ [[1,1,PH,PW], [1,1,PH,PW]], ... ]
        prediction, uncertainty = zip(*prediction)  # [[1,1,PH,PW], ... ], [[1,1,PH,PW], ... ]
        prediction = mint.cat(prediction, dim=0)  # [N,1,PH,PW]
        if output_uncertainty:
            uncertainty = mint.cat(uncertainty, dim=0)  # [N,1,PH,PW]
        else:
            uncertainty = None

    # 9. If `match_input_resolution` is set, the output prediction and the uncertainty are upsampled to match the
    # input resolution `(H, W)`. This step may introduce upsampling artifacts, and therefore can be disabled.
    # Depending on the downstream use-case, upsampling can be also chosen based on the tolerated artifacts by
    # setting the `resample_method_output` parameter (e.g., to `"nearest"`).
    if match_input_resolution:
        prediction = self.image_processor.resize_antialias(
            prediction, original_resolution, resample_method_output, is_aa=False
        )  # [N,1,H,W]
        if uncertainty is not None and output_uncertainty:
            uncertainty = self.image_processor.resize_antialias(
                uncertainty, original_resolution, resample_method_output, is_aa=False
            )  # [N,1,H,W]

    # 10. Prepare the final outputs.
    if output_type == "np":
        prediction = self.image_processor.ms_to_numpy(prediction)  # [N,H,W,1]
        if uncertainty is not None and output_uncertainty:
            uncertainty = self.image_processor.ms_to_numpy(uncertainty)  # [N,H,W,1]

    if not return_dict:
        return (prediction, uncertainty, pred_latent)

    return MarigoldDepthOutput(
        prediction=prediction,
        uncertainty=uncertainty,
        latent=pred_latent,
    )

mindone.diffusers.MarigoldDepthPipeline.ensemble_depth(depth, scale_invariant=True, shift_invariant=True, output_uncertainty=False, reduction='median', regularizer_strength=0.02, max_iter=2, tol=0.001, max_res=1024) staticmethod

Ensembles the depth maps represented by the depth tensor with expected shape (B, 1, H, W), where B is the number of ensemble members for a given prediction of size (H x W). Even though the function is designed for depth maps, it can also be used with disparity maps as long as the input tensor values are non-negative. The alignment happens when the predictions have one or more degrees of freedom, that is when they are either affine-invariant (scale_invariant=True and shift_invariant=True), or just scale-invariant (only scale_invariant=True). For absolute predictions (scale_invariant=False and shift_invariant=False) alignment is skipped and only ensembling is performed.

PARAMETER	DESCRIPTION
depth	Input ensemble depth maps. TYPE: `ms.Tensor`
scale_invariant	Whether to treat predictions as scale-invariant. TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True
shift_invariant	Whether to treat predictions as shift-invariant. TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True
output_uncertainty	Whether to output uncertainty map. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
reduction	Reduction method used to ensemble aligned predictions. The accepted values are: "mean" and "median". TYPE: `str`, *optional*, defaults to `"median"` DEFAULT: 'median'
regularizer_strength	Strength of the regularizer that pulls the aligned predictions to the unit range from 0 to 1. TYPE: `float`, *optional*, defaults to `0.02` DEFAULT: 0.02
max_iter	Maximum number of the alignment solver steps. Refer to scipy.optimize.minimize function, options argument. TYPE: `int`, *optional*, defaults to `2` DEFAULT: 2
tol	Alignment solver tolerance. The solver stops when the tolerance is reached. TYPE: `float`, *optional*, defaults to `1e-3` DEFAULT: 0.001
max_res	Resolution at which the alignment is performed; None matches the processing_resolution. TYPE: `int`, *optional*, defaults to `1024` DEFAULT: 1024

Source code in mindone/diffusers/pipelines/marigold/pipeline_marigold_depth.py

@staticmethod
def ensemble_depth(
    depth: ms.Tensor,
    scale_invariant: bool = True,
    shift_invariant: bool = True,
    output_uncertainty: bool = False,
    reduction: str = "median",
    regularizer_strength: float = 0.02,
    max_iter: int = 2,
    tol: float = 1e-3,
    max_res: int = 1024,
) -> Tuple[ms.Tensor, Optional[ms.Tensor]]:
    """
    Ensembles the depth maps represented by the `depth` tensor with expected shape `(B, 1, H, W)`, where B is the
    number of ensemble members for a given prediction of size `(H x W)`. Even though the function is designed for
    depth maps, it can also be used with disparity maps as long as the input tensor values are non-negative. The
    alignment happens when the predictions have one or more degrees of freedom, that is when they are either
    affine-invariant (`scale_invariant=True` and `shift_invariant=True`), or just scale-invariant (only
    `scale_invariant=True`). For absolute predictions (`scale_invariant=False` and `shift_invariant=False`)
    alignment is skipped and only ensembling is performed.

    Args:
        depth (`ms.Tensor`):
            Input ensemble depth maps.
        scale_invariant (`bool`, *optional*, defaults to `True`):
            Whether to treat predictions as scale-invariant.
        shift_invariant (`bool`, *optional*, defaults to `True`):
            Whether to treat predictions as shift-invariant.
        output_uncertainty (`bool`, *optional*, defaults to `False`):
            Whether to output uncertainty map.
        reduction (`str`, *optional*, defaults to `"median"`):
            Reduction method used to ensemble aligned predictions. The accepted values are: `"mean"` and
            `"median"`.
        regularizer_strength (`float`, *optional*, defaults to `0.02`):
            Strength of the regularizer that pulls the aligned predictions to the unit range from 0 to 1.
        max_iter (`int`, *optional*, defaults to `2`):
            Maximum number of the alignment solver steps. Refer to `scipy.optimize.minimize` function, `options`
            argument.
        tol (`float`, *optional*, defaults to `1e-3`):
            Alignment solver tolerance. The solver stops when the tolerance is reached.
        max_res (`int`, *optional*, defaults to `1024`):
            Resolution at which the alignment is performed; `None` matches the `processing_resolution`.
    Returns:
        A tensor of aligned and ensembled depth maps and optionally a tensor of uncertainties of the same shape:
        `(1, 1, H, W)`.
    """
    if depth.ndim != 4 or depth.shape[1] != 1:
        raise ValueError(f"Expecting 4D tensor of shape [B,1,H,W]; got {depth.shape}.")
    if reduction not in ("mean", "median"):
        raise ValueError(f"Unrecognized reduction method: {reduction}.")
    if not scale_invariant and shift_invariant:
        raise ValueError("Pure shift-invariant ensembling is not supported.")

    def init_param(depth: ms.Tensor):
        init_min = depth.reshape(ensemble_size, -1).min(dim=1)
        init_max = depth.reshape(ensemble_size, -1).max(dim=1)

        if scale_invariant and shift_invariant:
            init_s = 1.0 / (init_max - init_min).clamp(min=1e-6)
            init_t = -init_s * init_min
            param = mint.cat((init_s, init_t)).numpy()
        elif scale_invariant:
            init_s = 1.0 / init_max.clamp(min=1e-6)
            param = init_s.numpy()
        else:
            raise ValueError("Unrecognized alignment.")

        return param

    def align(depth: ms.Tensor, param: np.ndarray) -> ms.Tensor:
        if scale_invariant and shift_invariant:
            s, t = np.split(param, 2)
            s = ms.Tensor.from_numpy(s).to(depth.dtype).view(ensemble_size, 1, 1, 1)
            t = ms.Tensor.from_numpy(t).to(depth.dtype).view(ensemble_size, 1, 1, 1)
            out = depth * s + t
        elif scale_invariant:
            s = ms.Tensor.from_numpy(param).to(depth.dtype).view(ensemble_size, 1, 1, 1)
            out = depth * s
        else:
            raise ValueError("Unrecognized alignment.")
        return out

    def ensemble(
        depth_aligned: ms.Tensor, return_uncertainty: bool = False
    ) -> Tuple[ms.Tensor, Optional[ms.Tensor]]:
        uncertainty = None
        if reduction == "mean":
            prediction = mint.mean(depth_aligned, dim=0, keepdim=True)
            if return_uncertainty:
                uncertainty = mint.std(depth_aligned, dim=0, keepdim=True)
        elif reduction == "median":
            # ops.median has two return values and does not supported some data-type
            prediction = mint.median(depth_aligned.float(), dim=0, keepdim=True)[0]
            prediction = prediction.to(depth_aligned.dtype)
            if return_uncertainty:
                uncertainty = mint.median(mint.abs(depth_aligned - prediction).float(), dim=0, keepdim=True)[0]
                uncertainty = uncertainty.to(depth_aligned.dtype)
        else:
            raise ValueError(f"Unrecognized reduction method: {reduction}.")
        return prediction, uncertainty

    def cost_fn(param: np.ndarray, depth: ms.Tensor) -> float:
        cost = 0.0
        depth_aligned = align(depth, param)

        # todo: unavailable mint interface combinations
        for i, j in ops.combinations(mint.arange(ensemble_size)):
            diff = depth_aligned[i] - depth_aligned[j]
            cost += (diff**2).mean().sqrt().item()

        if regularizer_strength > 0:
            prediction, _ = ensemble(depth_aligned, return_uncertainty=False)
            err_near = (prediction.min()).abs().item()
            err_far = (1.0 - prediction.max()).abs().item()
            cost += (err_near + err_far) * regularizer_strength

        return cost

    def compute_param(depth: ms.Tensor):
        import scipy

        depth_to_align = depth.to(ms.float32)
        if max_res is not None and max(depth_to_align.shape[2:]) > max_res:
            depth_to_align = MarigoldImageProcessor.resize_to_max_edge(depth_to_align, max_res, "nearest-exact")

        param = init_param(depth_to_align)

        res = scipy.optimize.minimize(
            partial(cost_fn, depth=depth_to_align),
            param,
            method="BFGS",
            tol=tol,
            options={"maxiter": max_iter, "disp": False},
        )

        return res.x

    requires_aligning = scale_invariant or shift_invariant
    ensemble_size = depth.shape[0]

    if requires_aligning:
        param = compute_param(depth)
        depth = align(depth, param)

    depth, uncertainty = ensemble(depth, return_uncertainty=output_uncertainty)

    depth_max = depth.max()
    if scale_invariant and shift_invariant:
        depth_min = mint.min(depth)
    elif scale_invariant:
        depth_min = 0
    else:
        raise ValueError("Unrecognized alignment.")
    depth_range = (depth_max - depth_min).clamp(min=1e-6)
    depth = (depth - depth_min) / depth_range
    if output_uncertainty:
        uncertainty /= depth_range

    return depth, uncertainty  # [1,1,H,W], [1,1,H,W]

mindone.diffusers.pipelines.marigold.pipeline_marigold_depth.MarigoldDepthOutput dataclass

Bases: BaseOutput

Output class for Marigold monocular depth prediction pipeline.

PARAMETER	DESCRIPTION
prediction	Predicted depth maps with values in the range [0, 1]. The shape is \(numimages imes 1 imes height imes width\) for ms.Tensor or \(numimages imes height imes width imes 1\) for np.ndarray. TYPE: `np.ndarray`, `ms.Tensor`
uncertainty	Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is \(numimages imes 1 imes height imes width\) for ms.Tensor or \(numimages imes height imes width imes 1\) for np.ndarray. TYPE: `None`, `np.ndarray`, `ms.Tensor`
latent	Latent features corresponding to the predictions, compatible with the latents argument of the pipeline. The shape is \(numimages * numensemble imes 4 imes latentheight imes latentwidth\). TYPE: `None`, `ms.Tensor`

Source code in mindone/diffusers/pipelines/marigold/pipeline_marigold_depth.py

@dataclass
class MarigoldDepthOutput(BaseOutput):
    """
    Output class for Marigold monocular depth prediction pipeline.

    Args:
        prediction (`np.ndarray`, `ms.Tensor`):
            Predicted depth maps with values in the range [0, 1]. The shape is $numimages \times 1 \times height \times
            width$ for `ms.Tensor` or $numimages \times height \times width \times 1$ for `np.ndarray`.
        uncertainty (`None`, `np.ndarray`, `ms.Tensor`):
            Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $numimages
            \times 1 \times height \times width$ for `ms.Tensor` or $numimages \times height \times width \times 1$
            for `np.ndarray`.
        latent (`None`, `ms.Tensor`):
            Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline.
            The shape is $numimages * numensemble \times 4 \times latentheight \times latentwidth$.
    """

    prediction: Union[np.ndarray, ms.Tensor]
    uncertainty: Union[None, np.ndarray, ms.Tensor]
    latent: Union[None, ms.Tensor]

mindone.diffusers.pipelines.marigold.marigold_image_processing.MarigoldImageProcessor.visualize_depth(depth, val_min=0.0, val_max=1.0, color_map='Spectral') staticmethod

Visualizes depth maps, such as predictions of the MarigoldDepthPipeline.

PARAMETER	DESCRIPTION
val_min	Minimum value of the visualized depth range. TYPE: `float`, *optional*, defaults to `0.0` DEFAULT: 0.0
val_max	Maximum value of the visualized depth range. TYPE: `float`, *optional*, defaults to `1.0` DEFAULT: 1.0
color_map	Color map used to convert a single-channel depth prediction into colored representation. TYPE: `str`, *optional*, defaults to `"Spectral"` DEFAULT: 'Spectral'

Source code in mindone/diffusers/pipelines/marigold/marigold_image_processing.py

@staticmethod
def visualize_depth(
    depth: Union[
        PIL.Image.Image,
        np.ndarray,
        ms.Tensor,
        List[PIL.Image.Image],
        List[np.ndarray],
        List[ms.Tensor],
    ],
    val_min: float = 0.0,
    val_max: float = 1.0,
    color_map: str = "Spectral",
) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
    """
    Visualizes depth maps, such as predictions of the `MarigoldDepthPipeline`.

    Args:
        depth (`Union[PIL.Image.Image, np.ndarray, ms.Tensor, List[PIL.Image.Image], List[np.ndarray],
            List[ms.Tensor]]`): Depth maps.
        val_min (`float`, *optional*, defaults to `0.0`): Minimum value of the visualized depth range.
        val_max (`float`, *optional*, defaults to `1.0`): Maximum value of the visualized depth range.
        color_map (`str`, *optional*, defaults to `"Spectral"`): Color map used to convert a single-channel
                  depth prediction into colored representation.

    Returns: `PIL.Image.Image` or `List[PIL.Image.Image]` with depth maps visualization.
    """
    if val_max <= val_min:
        raise ValueError(f"Invalid values range: [{val_min}, {val_max}].")

    def visualize_depth_one(img, idx=None):
        prefix = "Depth" + (f"[{idx}]" if idx else "")
        if isinstance(img, PIL.Image.Image):
            if img.mode != "I;16":
                raise ValueError(f"{prefix}: invalid PIL mode={img.mode}.")
            img = np.array(img).astype(np.float32) / (2**16 - 1)
        if isinstance(img, np.ndarray) or ops.is_tensor(img):
            if img.ndim != 2:
                raise ValueError(f"{prefix}: unexpected shape={img.shape}.")
            if isinstance(img, np.ndarray):
                img = ms.Tensor.from_numpy(img)
            if not ops.is_floating_point(img):
                raise ValueError(f"{prefix}: unexected dtype={img.dtype}.")
        else:
            raise ValueError(f"{prefix}: unexpected type={type(img)}.")
        if val_min != 0.0 or val_max != 1.0:
            img = (img - val_min) / (val_max - val_min)
        img = MarigoldImageProcessor.colormap(img, cmap=color_map, bytes=True)  # [H,W,3]
        img = PIL.Image.fromarray(img.numpy())
        return img

    if depth is None or isinstance(depth, list) and any(o is None for o in depth):
        raise ValueError("Input depth is `None`")
    if isinstance(depth, (np.ndarray, ms.Tensor)):
        depth = MarigoldImageProcessor.expand_tensor_or_array(depth)
        if isinstance(depth, np.ndarray):
            depth = MarigoldImageProcessor.numpy_to_ms(depth)  # [N,H,W,1] -> [N,1,H,W]
        if not (depth.ndim == 4 and depth.shape[1] == 1):  # [N,1,H,W]
            raise ValueError(f"Unexpected input shape={depth.shape}, expecting [N,1,H,W].")
        return [visualize_depth_one(img[0], idx) for idx, img in enumerate(depth)]
    elif isinstance(depth, list):
        return [visualize_depth_one(img, idx) for idx, img in enumerate(depth)]
    else:
        raise ValueError(f"Unexpected input type: {type(depth)}")

Marigold Normals Estimation API

mindone.diffusers.MarigoldNormalsPipeline

Bases: DiffusionPipeline

Pipeline for monocular normals estimation using the Marigold method: https://marigoldmonodepth.github.io.

This model inherits from [DiffusionPipeline]. Check the superclass documentation for the generic methods the library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)

PARAMETER	DESCRIPTION
unet	Conditional U-Net to denoise the normals latent, conditioned on image latent. TYPE: `UNet2DConditionModel`
vae	Variational Auto-Encoder (VAE) Model to encode and decode images and predictions to and from latent representations. TYPE: `AutoencoderKL`
scheduler	A scheduler to be used in combination with unet to denoise the encoded image latents. TYPE: `DDIMScheduler` or `LCMScheduler`
text_encoder	Text-encoder, for empty text embedding. TYPE: `CLIPTextModel`
tokenizer	CLIP tokenizer. TYPE: `CLIPTokenizer`
prediction_type	Type of predictions made by the model. TYPE: `str`, *optional* DEFAULT: None
use_full_z_range	Whether the normals predicted by this model utilize the full range of the Z dimension, or only its positive half. TYPE: `bool`, *optional* DEFAULT: True
default_denoising_steps	The minimum number of denoising diffusion steps that are required to produce a prediction of reasonable quality with the given model. This value must be set in the model config. When the pipeline is called without explicitly setting num_inference_steps, the default value is used. This is required to ensure reasonable results with various model flavors compatible with the pipeline, such as those relying on very short denoising schedules (LCMScheduler) and those with full diffusion schedules (DDIMScheduler). TYPE: `int`, *optional* DEFAULT: None
default_processing_resolution	The recommended value of the processing_resolution parameter of the pipeline. This value must be set in the model config. When the pipeline is called without explicitly setting processing_resolution, the default value is used. This is required to ensure reasonable results with various model flavors trained with varying optimal processing resolution values. TYPE: `int`, *optional* DEFAULT: None

Source code in mindone/diffusers/pipelines/marigold/pipeline_marigold_normals.py

class MarigoldNormalsPipeline(DiffusionPipeline):
    """
    Pipeline for monocular normals estimation using the Marigold method: https://marigoldmonodepth.github.io.

    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)

    Args:
        unet (`UNet2DConditionModel`):
            Conditional U-Net to denoise the normals latent, conditioned on image latent.
        vae (`AutoencoderKL`):
            Variational Auto-Encoder (VAE) Model to encode and decode images and predictions to and from latent
            representations.
        scheduler (`DDIMScheduler` or `LCMScheduler`):
            A scheduler to be used in combination with `unet` to denoise the encoded image latents.
        text_encoder (`CLIPTextModel`):
            Text-encoder, for empty text embedding.
        tokenizer (`CLIPTokenizer`):
            CLIP tokenizer.
        prediction_type (`str`, *optional*):
            Type of predictions made by the model.
        use_full_z_range (`bool`, *optional*):
            Whether the normals predicted by this model utilize the full range of the Z dimension, or only its positive
            half.
        default_denoising_steps (`int`, *optional*):
            The minimum number of denoising diffusion steps that are required to produce a prediction of reasonable
            quality with the given model. This value must be set in the model config. When the pipeline is called
            without explicitly setting `num_inference_steps`, the default value is used. This is required to ensure
            reasonable results with various model flavors compatible with the pipeline, such as those relying on very
            short denoising schedules (`LCMScheduler`) and those with full diffusion schedules (`DDIMScheduler`).
        default_processing_resolution (`int`, *optional*):
            The recommended value of the `processing_resolution` parameter of the pipeline. This value must be set in
            the model config. When the pipeline is called without explicitly setting `processing_resolution`, the
            default value is used. This is required to ensure reasonable results with various model flavors trained
            with varying optimal processing resolution values.
    """

    model_cpu_offload_seq = "text_encoder->unet->vae"
    supported_prediction_types = ("normals",)

    def __init__(
        self,
        unet: UNet2DConditionModel,
        vae: AutoencoderKL,
        scheduler: Union[DDIMScheduler, LCMScheduler],
        text_encoder: CLIPTextModel,
        tokenizer: CLIPTokenizer,
        prediction_type: Optional[str] = None,
        use_full_z_range: Optional[bool] = True,
        default_denoising_steps: Optional[int] = None,
        default_processing_resolution: Optional[int] = None,
    ):
        super().__init__()

        if prediction_type not in self.supported_prediction_types:
            logger.warning(
                f"Potentially unsupported `prediction_type='{prediction_type}'`; values supported by the pipeline: "
                f"{self.supported_prediction_types}."
            )

        self.register_modules(
            unet=unet,
            vae=vae,
            scheduler=scheduler,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
        )
        self.register_to_config(
            use_full_z_range=use_full_z_range,
            default_denoising_steps=default_denoising_steps,
            default_processing_resolution=default_processing_resolution,
        )

        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)

        self.use_full_z_range = use_full_z_range
        self.default_denoising_steps = default_denoising_steps
        self.default_processing_resolution = default_processing_resolution

        self.empty_text_embedding = None

        self.image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor)

    def check_inputs(
        self,
        image: PipelineImageInput,
        num_inference_steps: int,
        ensemble_size: int,
        processing_resolution: int,
        resample_method_input: str,
        resample_method_output: str,
        batch_size: int,
        ensembling_kwargs: Optional[Dict[str, Any]],
        latents: Optional[ms.Tensor],
        generator: Optional[Union[np.random.Generator, List[np.random.Generator]]],
        output_type: str,
        output_uncertainty: bool,
    ) -> int:
        actual_vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        if actual_vae_scale_factor != self.vae_scale_factor:
            raise ValueError(
                f"`vae_scale_factor` computed at initialization ({self.vae_scale_factor}) differs from the actual one ({actual_vae_scale_factor})."
            )
        if num_inference_steps is None:
            raise ValueError("`num_inference_steps` is not specified and could not be resolved from the model config.")
        if num_inference_steps < 1:
            raise ValueError("`num_inference_steps` must be positive.")
        if ensemble_size < 1:
            raise ValueError("`ensemble_size` must be positive.")
        if ensemble_size == 2:
            logger.warning(
                "`ensemble_size` == 2 results are similar to no ensembling (1); "
                "consider increasing the value to at least 3."
            )
        if ensemble_size == 1 and output_uncertainty:
            raise ValueError(
                "Computing uncertainty by setting `output_uncertainty=True` also requires setting `ensemble_size` "
                "greater than 1."
            )
        if processing_resolution is None:
            raise ValueError(
                "`processing_resolution` is not specified and could not be resolved from the model config."
            )
        if processing_resolution < 0:
            raise ValueError(
                "`processing_resolution` must be non-negative: 0 for native resolution, or any positive value for "
                "downsampled processing."
            )
        if processing_resolution % self.vae_scale_factor != 0:
            raise ValueError(f"`processing_resolution` must be a multiple of {self.vae_scale_factor}.")
        if resample_method_input not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
            raise ValueError(
                "`resample_method_input` takes string values compatible with PIL library: "
                "nearest, nearest-exact, bilinear, bicubic, area."
            )
        if resample_method_output not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
            raise ValueError(
                "`resample_method_output` takes string values compatible with PIL library: "
                "nearest, nearest-exact, bilinear, bicubic, area."
            )
        if batch_size < 1:
            raise ValueError("`batch_size` must be positive.")
        if output_type not in ["pt", "np"]:
            raise ValueError("`output_type` must be one of `pt` or `np`.")
        if latents is not None and generator is not None:
            raise ValueError("`latents` and `generator` cannot be used together.")
        if ensembling_kwargs is not None:
            if not isinstance(ensembling_kwargs, dict):
                raise ValueError("`ensembling_kwargs` must be a dictionary.")
            if "reduction" in ensembling_kwargs and ensembling_kwargs["reduction"] not in ("closest", "mean"):
                raise ValueError("`ensembling_kwargs['reduction']` can be either `'closest'` or `'mean'`.")

        # image checks
        num_images = 0
        W, H = None, None
        if not isinstance(image, list):
            image = [image]
        for i, img in enumerate(image):
            if isinstance(img, np.ndarray) or ops.is_tensor(img):
                if img.ndim not in (2, 3, 4):
                    raise ValueError(f"`image[{i}]` has unsupported dimensions or shape: {img.shape}.")
                H_i, W_i = img.shape[-2:]
                N_i = 1
                if img.ndim == 4:
                    N_i = img.shape[0]
            elif isinstance(img, Image.Image):
                W_i, H_i = img.size
                N_i = 1
            else:
                raise ValueError(f"Unsupported `image[{i}]` type: {type(img)}.")
            if W is None:
                W, H = W_i, H_i
            elif (W, H) != (W_i, H_i):
                raise ValueError(
                    f"Input `image[{i}]` has incompatible dimensions {(W_i, H_i)} with the previous images {(W, H)}"
                )
            num_images += N_i

        # latents checks
        if latents is not None:
            if not ops.is_tensor(latents):
                raise ValueError("`latents` must be a ms.Tensor.")
            if latents.ndim != 4:
                raise ValueError(f"`latents` has unsupported dimensions or shape: {latents.shape}.")

            if processing_resolution > 0:
                max_orig = max(H, W)
                new_H = H * processing_resolution // max_orig
                new_W = W * processing_resolution // max_orig
                if new_H == 0 or new_W == 0:
                    raise ValueError(f"Extreme aspect ratio of the input image: [{W} x {H}]")
                W, H = new_W, new_H
            w = (W + self.vae_scale_factor - 1) // self.vae_scale_factor
            h = (H + self.vae_scale_factor - 1) // self.vae_scale_factor
            shape_expected = (num_images * ensemble_size, self.vae.config.latent_channels, h, w)

            if latents.shape != shape_expected:
                raise ValueError(f"`latents` has unexpected shape={latents.shape} expected={shape_expected}.")

        # generator checks
        if generator is not None:
            if isinstance(generator, list):
                if len(generator) != num_images * ensemble_size:
                    raise ValueError(
                        "The number of generators must match the total number of ensemble members for all input images."
                    )
            elif not isinstance(generator, np.random.Generator):
                raise ValueError(f"Unsupported generator type: {type(generator)}.")

        return num_images

    def progress_bar(self, iterable=None, total=None, desc=None, leave=True):
        if not hasattr(self, "_progress_bar_config"):
            self._progress_bar_config = {}
        elif not isinstance(self._progress_bar_config, dict):
            raise ValueError(
                f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}."
            )

        progress_bar_config = dict(**self._progress_bar_config)
        progress_bar_config["desc"] = progress_bar_config.get("desc", desc)
        progress_bar_config["leave"] = progress_bar_config.get("leave", leave)
        if iterable is not None:
            return tqdm(iterable, **progress_bar_config)
        elif total is not None:
            return tqdm(total=total, **progress_bar_config)
        else:
            raise ValueError("Either `total` or `iterable` has to be defined.")

    def __call__(
        self,
        image: PipelineImageInput,
        num_inference_steps: Optional[int] = None,
        ensemble_size: int = 1,
        processing_resolution: Optional[int] = None,
        match_input_resolution: bool = True,
        resample_method_input: str = "bilinear",
        resample_method_output: str = "bilinear",
        batch_size: int = 1,
        ensembling_kwargs: Optional[Dict[str, Any]] = None,
        latents: Optional[Union[ms.Tensor, List[ms.Tensor]]] = None,
        generator: Optional[Union[np.random.Generator, List[np.random.Generator]]] = None,
        output_type: str = "np",
        output_uncertainty: bool = False,
        output_latent: bool = False,
        return_dict: bool = False,
    ):
        """
        Function invoked when calling the pipeline.

        Args:
            image (`PIL.Image.Image`, `np.ndarray`, `ms.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`),
                `List[ms.Tensor]`: An input image or images used as an input for the normals estimation task. For
                arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible
                by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or
                three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the
                same width and height.
            num_inference_steps (`int`, *optional*, defaults to `None`):
                Number of denoising diffusion steps during inference. The default value `None` results in automatic
                selection.
            ensemble_size (`int`, defaults to `1`):
                Number of ensemble predictions. Higher values result in measurable improvements and visual degradation.
            processing_resolution (`int`, *optional*, defaults to `None`):
                Effective processing resolution. When set to `0`, matches the larger input image dimension. This
                produces crisper predictions, but may also lead to the overall loss of global context. The default
                value `None` resolves to the optimal value from the model config.
            match_input_resolution (`bool`, *optional*, defaults to `True`):
                When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer
                side of the output will equal to `processing_resolution`.
            resample_method_input (`str`, *optional*, defaults to `"bilinear"`):
                Resampling method used to resize input images to `processing_resolution`. The accepted values are:
                `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
            resample_method_output (`str`, *optional*, defaults to `"bilinear"`):
                Resampling method used to resize output predictions to match the input resolution. The accepted values
                are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
            batch_size (`int`, *optional*, defaults to `1`):
                Batch size; only matters when setting `ensemble_size` or passing a tensor of images.
            ensembling_kwargs (`dict`, *optional*, defaults to `None`)
                Extra dictionary with arguments for precise ensembling control. The following options are available:
                - reduction (`str`, *optional*, defaults to `"closest"`): Defines the ensembling function applied in
                  every pixel location, can be either `"closest"` or `"mean"`.
            latents (`ms.Tensor`, *optional*, defaults to `None`):
                Latent noise tensors to replace the random initialization. These can be taken from the previous
                function call's output.
            generator (`np.random.Generator`, or `List[np.random.Generator]`, *optional*, defaults to `None`):
                Random number generator object to ensure reproducibility.
            output_type (`str`, *optional*, defaults to `"np"`):
                Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted
                values are: `"np"` (numpy array) or `"pt"` (ms tensor).
            output_uncertainty (`bool`, *optional*, defaults to `False`):
                When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that
                the `ensemble_size` argument is set to a value above 2.
            output_latent (`bool`, *optional*, defaults to `False`):
                When enabled, the output's `latent` field contains the latent codes corresponding to the predictions
                within the ensemble. These codes can be saved, modified, and used for subsequent calls with the
                `latents` argument.
            return_dict (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~pipelines.marigold.MarigoldNormalsOutput`] instead of a plain tuple.

        Examples:

        Returns:
            [`~pipelines.marigold.MarigoldNormalsOutput`] or `tuple`:
                If `return_dict` is `True`, [`~pipelines.marigold.MarigoldNormalsOutput`] is returned, otherwise a
                `tuple` is returned where the first element is the prediction, the second element is the uncertainty
                (or `None`), and the third is the latent (or `None`).
        """

        # 0. Resolving variables.
        dtype = self.dtype

        # Model-specific optimal default values leading to fast and reasonable results.
        if num_inference_steps is None:
            num_inference_steps = self.default_denoising_steps
        if processing_resolution is None:
            processing_resolution = self.default_processing_resolution

        # 1. Check inputs.
        num_images = self.check_inputs(
            image,
            num_inference_steps,
            ensemble_size,
            processing_resolution,
            resample_method_input,
            resample_method_output,
            batch_size,
            ensembling_kwargs,
            latents,
            generator,
            output_type,
            output_uncertainty,
        )

        # 2. Prepare empty text conditioning.
        # Model invocation: self.tokenizer, self.text_encoder.
        if self.empty_text_embedding is None:
            prompt = ""
            text_inputs = self.tokenizer(
                prompt,
                padding="do_not_pad",
                max_length=self.tokenizer.model_max_length,
                truncation=True,
                return_tensors="np",
            )
            text_input_ids = ms.Tensor.from_numpy(text_inputs.input_ids)
            self.empty_text_embedding = self.text_encoder(text_input_ids)[0]  # [1,2,1024]

        # 3. Preprocess input images. This function loads input image or images of compatible dimensions `(H, W)`,
        # optionally downsamples them to the `processing_resolution` `(PH, PW)`, where
        # `max(PH, PW) == processing_resolution`, and pads the dimensions to `(PPH, PPW)` such that these values are
        # divisible by the latent space downscaling factor (typically 8 in Stable Diffusion). The default value `None`
        # of `processing_resolution` resolves to the optimal value from the model config. It is a recommended mode of
        # operation and leads to the most reasonable results. Using the native image resolution or any other processing
        # resolution can lead to loss of either fine details or global context in the output predictions.
        image, padding, original_resolution = self.image_processor.preprocess(
            image, processing_resolution, resample_method_input, dtype
        )  # [N,3,PPH,PPW]

        # 4. Encode input image into latent space. At this step, each of the `N` input images is represented with `E`
        # ensemble members. Each ensemble member is an independent diffused prediction, just initialized independently.
        # Latents of each such predictions across all input images and all ensemble members are represented in the
        # `pred_latent` variable. The variable `image_latent` is of the same shape: it contains each input image encoded
        # into latent space and replicated `E` times. The latents can be either generated (see `generator` to ensure
        # reproducibility), or passed explicitly via the `latents` argument. The latter can be set outside the pipeline
        # code. This behavior can be achieved by setting the `output_latent` argument to `True`. The latent space
        # dimensions are `(h, w)`. Encoding into latent space happens in batches of size `batch_size`.
        # Model invocation: self.vae.encoder.
        image_latent, pred_latent = self.prepare_latents(
            image, latents, generator, ensemble_size, batch_size
        )  # [N*E,4,h,w], [N*E,4,h,w]

        del image

        batch_empty_text_embedding = self.empty_text_embedding.to(dtype=dtype).tile((batch_size, 1, 1))  # [B,1024,2]

        # 5. Process the denoising loop. All `N * E` latents are processed sequentially in batches of size `batch_size`.
        # The unet model takes concatenated latent spaces of the input image and the predicted modality as an input, and
        # outputs noise for the predicted modality's latent space. The number of denoising diffusion steps is defined by
        # `num_inference_steps`. It is either set directly, or resolves to the optimal value specific to the loaded
        # model.
        # Model invocation: self.unet.
        pred_latents = []

        for i in self.progress_bar(
            range(0, num_images * ensemble_size, batch_size), leave=True, desc="Marigold predictions..."
        ):
            batch_image_latent = image_latent[i : i + batch_size]  # [B,4,h,w]
            batch_pred_latent = pred_latent[i : i + batch_size]  # [B,4,h,w]
            effective_batch_size = batch_image_latent.shape[0]
            text = batch_empty_text_embedding[:effective_batch_size]  # [B,2,1024]

            self.scheduler.set_timesteps(num_inference_steps)
            for t in self.progress_bar(self.scheduler.timesteps, leave=False, desc="Diffusion steps..."):
                batch_latent = mint.cat([batch_image_latent, batch_pred_latent], dim=1)  # [B,8,h,w]
                noise = self.unet(batch_latent, t, encoder_hidden_states=text, return_dict=False)[0]  # [B,4,h,w]
                batch_pred_latent = self.scheduler.step(noise, t, batch_pred_latent, generator=generator)[
                    0
                ]  # [B,4,h,w]

            pred_latents.append(batch_pred_latent)

        pred_latent = mint.cat(pred_latents, dim=0)  # [N*E,4,h,w]

        del (
            pred_latents,
            image_latent,
            batch_empty_text_embedding,
            batch_image_latent,
            batch_pred_latent,
            text,
            batch_latent,
            noise,
        )

        # 6. Decode predictions from latent into pixel space. The resulting `N * E` predictions have shape `(PPH, PPW)`,
        # which requires slight postprocessing. Decoding into pixel space happens in batches of size `batch_size`.
        # Model invocation: self.vae.decoder.
        prediction = mint.cat(
            [
                self.decode_prediction(pred_latent[i : i + batch_size])
                for i in range(0, pred_latent.shape[0], batch_size)
            ],
            dim=0,
        )  # [N*E,3,PPH,PPW]

        if not output_latent:
            pred_latent = None

        # 7. Remove padding. The output shape is (PH, PW).
        prediction = self.image_processor.unpad_image(prediction, padding)  # [N*E,3,PH,PW]

        # 8. Ensemble and compute uncertainty (when `output_uncertainty` is set). This code treats each of the `N`
        # groups of `E` ensemble predictions independently. For each group it computes an ensembled prediction of shape
        # `(PH, PW)` and an optional uncertainty map of the same dimensions. After computing this pair of outputs for
        # each group independently, it stacks them respectively into batches of `N` almost final predictions and
        # uncertainty maps.
        uncertainty = None
        if ensemble_size > 1:
            prediction = prediction.reshape(num_images, ensemble_size, *prediction.shape[1:])  # [N,E,3,PH,PW]
            prediction = [
                self.ensemble_normals(prediction[i], output_uncertainty, **(ensembling_kwargs or {}))
                for i in range(num_images)
            ]  # [ [[1,3,PH,PW], [1,1,PH,PW]], ... ]
            prediction, uncertainty = zip(*prediction)  # [[1,3,PH,PW], ... ], [[1,1,PH,PW], ... ]
            prediction = mint.cat(prediction, dim=0)  # [N,3,PH,PW]
            if output_uncertainty:
                uncertainty = mint.cat(uncertainty, dim=0)  # [N,1,PH,PW]
            else:
                uncertainty = None

        # 9. If `match_input_resolution` is set, the output prediction and the uncertainty are upsampled to match the
        # input resolution `(H, W)`. This step may introduce upsampling artifacts, and therefore can be disabled.
        # After upsampling, the native resolution normal maps are renormalized to unit length to reduce the artifacts.
        # Depending on the downstream use-case, upsampling can be also chosen based on the tolerated artifacts by
        # setting the `resample_method_output` parameter (e.g., to `"nearest"`).
        if match_input_resolution:
            prediction = self.image_processor.resize_antialias(
                prediction, original_resolution, resample_method_output, is_aa=False
            )  # [N,3,H,W]
            prediction = self.normalize_normals(prediction)  # [N,3,H,W]
            if uncertainty is not None and output_uncertainty:
                uncertainty = self.image_processor.resize_antialias(
                    uncertainty, original_resolution, resample_method_output, is_aa=False
                )  # [N,1,H,W]

        # 10. Prepare the final outputs.
        if output_type == "np":
            prediction = self.image_processor.ms_to_numpy(prediction)  # [N,H,W,3]
            if uncertainty is not None and output_uncertainty:
                uncertainty = self.image_processor.ms_to_numpy(uncertainty)  # [N,H,W,1]

        if not return_dict:
            return (prediction, uncertainty, pred_latent)

        return MarigoldNormalsOutput(
            prediction=prediction,
            uncertainty=uncertainty,
            latent=pred_latent,
        )

    # Copied from diffusers.pipelines.marigold.pipeline_marigold_depth.MarigoldDepthPipeline.prepare_latents
    def prepare_latents(
        self,
        image: ms.Tensor,
        latents: Optional[ms.Tensor],
        generator: Optional[np.random.Generator],
        ensemble_size: int,
        batch_size: int,
    ) -> Tuple[ms.Tensor, ms.Tensor]:
        def retrieve_latents(encoder_output):
            assert ops.is_tensor(
                encoder_output
            ), "Could not access latents of provided encoder_output which is not a tensor"
            if hasattr(self.vae, "diag_gauss_dist"):
                return self.vae.diag_gauss_dist.mode(encoder_output)
            else:
                return encoder_output

        image_latent = mint.cat(
            [
                retrieve_latents(self.vae.encode(image[i : i + batch_size])[0])
                for i in range(0, image.shape[0], batch_size)
            ],
            dim=0,
        )  # [N,4,h,w]
        image_latent = image_latent * self.vae.config.scaling_factor
        image_latent = image_latent.repeat_interleave(ensemble_size, dim=0)  # [N*E,4,h,w]

        pred_latent = latents
        if pred_latent is None:
            pred_latent = randn_tensor(
                image_latent.shape,
                generator=generator,
                dtype=image_latent.dtype,
            )  # [N*E,4,h,w]

        return image_latent, pred_latent

    def decode_prediction(self, pred_latent: ms.Tensor) -> ms.Tensor:
        if pred_latent.ndim != 4 or pred_latent.shape[1] != self.vae.config.latent_channels:
            raise ValueError(
                f"Expecting 4D tensor of shape [B,{self.vae.config.latent_channels},H,W]; got {pred_latent.shape}."
            )

        prediction = self.vae.decode(pred_latent / self.vae.config.scaling_factor, return_dict=False)[0]  # [B,3,H,W]

        prediction = mint.clip(prediction, -1.0, 1.0)

        if not self.use_full_z_range:
            prediction[:, 2, :, :] *= 0.5
            prediction[:, 2, :, :] += 0.5

        prediction = self.normalize_normals(prediction)  # [B,3,H,W]

        return prediction  # [B,3,H,W]

    @staticmethod
    def normalize_normals(normals: ms.Tensor, eps: float = 1e-6) -> ms.Tensor:
        if normals.ndim != 4 or normals.shape[1] != 3:
            raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.")

        norm = mint.norm(normals, dim=1, keepdim=True)
        normals /= norm.clamp(min=eps)

        return normals

    @staticmethod
    def ensemble_normals(
        normals: ms.Tensor, output_uncertainty: bool, reduction: str = "closest"
    ) -> Tuple[ms.Tensor, Optional[ms.Tensor]]:
        """
        Ensembles the normals maps represented by the `normals` tensor with expected shape `(B, 3, H, W)`, where B is
        the number of ensemble members for a given prediction of size `(H x W)`.

        Args:
            normals (`ms.Tensor`):
                Input ensemble normals maps.
            output_uncertainty (`bool`, *optional*, defaults to `False`):
                Whether to output uncertainty map.
            reduction (`str`, *optional*, defaults to `"closest"`):
                Reduction method used to ensemble aligned predictions. The accepted values are: `"closest"` and
                `"mean"`.

        Returns:
            A tensor of aligned and ensembled normals maps with shape `(1, 3, H, W)` and optionally a tensor of
            uncertainties of shape `(1, 1, H, W)`.
        """
        if normals.ndim != 4 or normals.shape[1] != 3:
            raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.")
        if reduction not in ("closest", "mean"):
            raise ValueError(f"Unrecognized reduction method: {reduction}.")

        mean_normals = normals.mean(dim=0, keep_dims=True)  # [1,3,H,W]
        mean_normals = MarigoldNormalsPipeline.normalize_normals(mean_normals)  # [1,3,H,W]

        sim_cos = (mean_normals * normals).sum(dim=1, keepdims=True)  # [E,1,H,W]
        sim_cos = sim_cos.clamp(-1.0, 1.0)  # required to avoid NaN in uncertainty with fp16

        uncertainty = None
        if output_uncertainty:
            uncertainty = sim_cos.arccos()  # [E,1,H,W]
            uncertainty = uncertainty.mean(dim=0, keep_dims=True) / ms.numpy.pi  # [1,1,H,W]

        if reduction == "mean":
            return mean_normals, uncertainty  # [1,3,H,W], [1,1,H,W]

        closest_indices = sim_cos.argmax(dim=0, keepdims=True)  # [1,1,H,W]
        closest_indices = closest_indices.tile((1, 3, 1, 1))  # [1,3,H,W]
        closest_normals = ops.gather_elements(normals, 0, closest_indices)  # [1,3,H,W]

        return closest_normals, uncertainty  # [1,3,H,W], [1,1,H,W]

mindone.diffusers.MarigoldNormalsPipeline.__call__(image, num_inference_steps=None, ensemble_size=1, processing_resolution=None, match_input_resolution=True, resample_method_input='bilinear', resample_method_output='bilinear', batch_size=1, ensembling_kwargs=None, latents=None, generator=None, output_type='np', output_uncertainty=False, output_latent=False, return_dict=False)

Function invoked when calling the pipeline.

PARAMETER	DESCRIPTION
num_inference_steps	Number of denoising diffusion steps during inference. The default value None results in automatic selection. TYPE: `int`, *optional*, defaults to `None` DEFAULT: None
ensemble_size	Number of ensemble predictions. Higher values result in measurable improvements and visual degradation. TYPE: `int`, defaults to `1` DEFAULT: 1
processing_resolution	Effective processing resolution. When set to 0, matches the larger input image dimension. This produces crisper predictions, but may also lead to the overall loss of global context. The default value None resolves to the optimal value from the model config. TYPE: `int`, *optional*, defaults to `None` DEFAULT: None
match_input_resolution	When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer side of the output will equal to processing_resolution. TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True
resample_method_input	Resampling method used to resize input images to processing_resolution. The accepted values are: "nearest", "nearest-exact", "bilinear", "bicubic", or "area". TYPE: `str`, *optional*, defaults to `"bilinear"` DEFAULT: 'bilinear'
resample_method_output	Resampling method used to resize output predictions to match the input resolution. The accepted values are "nearest", "nearest-exact", "bilinear", "bicubic", or "area". TYPE: `str`, *optional*, defaults to `"bilinear"` DEFAULT: 'bilinear'
batch_size	Batch size; only matters when setting ensemble_size or passing a tensor of images. TYPE: `int`, *optional*, defaults to `1` DEFAULT: 1
latents	Latent noise tensors to replace the random initialization. These can be taken from the previous function call's output. TYPE: `ms.Tensor`, *optional*, defaults to `None` DEFAULT: None
generator	Random number generator object to ensure reproducibility. TYPE: `np.random.Generator`, or `List[np.random.Generator]`, *optional*, defaults to `None` DEFAULT: None
output_type	Preferred format of the output's prediction and the optional uncertainty fields. The accepted values are: "np" (numpy array) or "pt" (ms tensor). TYPE: `str`, *optional*, defaults to `"np"` DEFAULT: 'np'
output_uncertainty	When enabled, the output's uncertainty field contains the predictive uncertainty map, provided that the ensemble_size argument is set to a value above 2. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
output_latent	When enabled, the output's latent field contains the latent codes corresponding to the predictions within the ensemble. These codes can be saved, modified, and used for subsequent calls with the latents argument. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
return_dict	Whether or not to return a [~pipelines.marigold.MarigoldNormalsOutput] instead of a plain tuple. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False

RETURNS	DESCRIPTION
	[~pipelines.marigold.MarigoldNormalsOutput] or tuple: If return_dict is True, [~pipelines.marigold.MarigoldNormalsOutput] is returned, otherwise a tuple is returned where the first element is the prediction, the second element is the uncertainty (or None), and the third is the latent (or None).

Source code in mindone/diffusers/pipelines/marigold/pipeline_marigold_normals.py

def __call__(
    self,
    image: PipelineImageInput,
    num_inference_steps: Optional[int] = None,
    ensemble_size: int = 1,
    processing_resolution: Optional[int] = None,
    match_input_resolution: bool = True,
    resample_method_input: str = "bilinear",
    resample_method_output: str = "bilinear",
    batch_size: int = 1,
    ensembling_kwargs: Optional[Dict[str, Any]] = None,
    latents: Optional[Union[ms.Tensor, List[ms.Tensor]]] = None,
    generator: Optional[Union[np.random.Generator, List[np.random.Generator]]] = None,
    output_type: str = "np",
    output_uncertainty: bool = False,
    output_latent: bool = False,
    return_dict: bool = False,
):
    """
    Function invoked when calling the pipeline.

    Args:
        image (`PIL.Image.Image`, `np.ndarray`, `ms.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`),
            `List[ms.Tensor]`: An input image or images used as an input for the normals estimation task. For
            arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is possible
            by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or
            three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the
            same width and height.
        num_inference_steps (`int`, *optional*, defaults to `None`):
            Number of denoising diffusion steps during inference. The default value `None` results in automatic
            selection.
        ensemble_size (`int`, defaults to `1`):
            Number of ensemble predictions. Higher values result in measurable improvements and visual degradation.
        processing_resolution (`int`, *optional*, defaults to `None`):
            Effective processing resolution. When set to `0`, matches the larger input image dimension. This
            produces crisper predictions, but may also lead to the overall loss of global context. The default
            value `None` resolves to the optimal value from the model config.
        match_input_resolution (`bool`, *optional*, defaults to `True`):
            When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer
            side of the output will equal to `processing_resolution`.
        resample_method_input (`str`, *optional*, defaults to `"bilinear"`):
            Resampling method used to resize input images to `processing_resolution`. The accepted values are:
            `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
        resample_method_output (`str`, *optional*, defaults to `"bilinear"`):
            Resampling method used to resize output predictions to match the input resolution. The accepted values
            are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
        batch_size (`int`, *optional*, defaults to `1`):
            Batch size; only matters when setting `ensemble_size` or passing a tensor of images.
        ensembling_kwargs (`dict`, *optional*, defaults to `None`)
            Extra dictionary with arguments for precise ensembling control. The following options are available:
            - reduction (`str`, *optional*, defaults to `"closest"`): Defines the ensembling function applied in
              every pixel location, can be either `"closest"` or `"mean"`.
        latents (`ms.Tensor`, *optional*, defaults to `None`):
            Latent noise tensors to replace the random initialization. These can be taken from the previous
            function call's output.
        generator (`np.random.Generator`, or `List[np.random.Generator]`, *optional*, defaults to `None`):
            Random number generator object to ensure reproducibility.
        output_type (`str`, *optional*, defaults to `"np"`):
            Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted
            values are: `"np"` (numpy array) or `"pt"` (ms tensor).
        output_uncertainty (`bool`, *optional*, defaults to `False`):
            When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that
            the `ensemble_size` argument is set to a value above 2.
        output_latent (`bool`, *optional*, defaults to `False`):
            When enabled, the output's `latent` field contains the latent codes corresponding to the predictions
            within the ensemble. These codes can be saved, modified, and used for subsequent calls with the
            `latents` argument.
        return_dict (`bool`, *optional*, defaults to `False`):
            Whether or not to return a [`~pipelines.marigold.MarigoldNormalsOutput`] instead of a plain tuple.

    Examples:

    Returns:
        [`~pipelines.marigold.MarigoldNormalsOutput`] or `tuple`:
            If `return_dict` is `True`, [`~pipelines.marigold.MarigoldNormalsOutput`] is returned, otherwise a
            `tuple` is returned where the first element is the prediction, the second element is the uncertainty
            (or `None`), and the third is the latent (or `None`).
    """

    # 0. Resolving variables.
    dtype = self.dtype

    # Model-specific optimal default values leading to fast and reasonable results.
    if num_inference_steps is None:
        num_inference_steps = self.default_denoising_steps
    if processing_resolution is None:
        processing_resolution = self.default_processing_resolution

    # 1. Check inputs.
    num_images = self.check_inputs(
        image,
        num_inference_steps,
        ensemble_size,
        processing_resolution,
        resample_method_input,
        resample_method_output,
        batch_size,
        ensembling_kwargs,
        latents,
        generator,
        output_type,
        output_uncertainty,
    )

    # 2. Prepare empty text conditioning.
    # Model invocation: self.tokenizer, self.text_encoder.
    if self.empty_text_embedding is None:
        prompt = ""
        text_inputs = self.tokenizer(
            prompt,
            padding="do_not_pad",
            max_length=self.tokenizer.model_max_length,
            truncation=True,
            return_tensors="np",
        )
        text_input_ids = ms.Tensor.from_numpy(text_inputs.input_ids)
        self.empty_text_embedding = self.text_encoder(text_input_ids)[0]  # [1,2,1024]

    # 3. Preprocess input images. This function loads input image or images of compatible dimensions `(H, W)`,
    # optionally downsamples them to the `processing_resolution` `(PH, PW)`, where
    # `max(PH, PW) == processing_resolution`, and pads the dimensions to `(PPH, PPW)` such that these values are
    # divisible by the latent space downscaling factor (typically 8 in Stable Diffusion). The default value `None`
    # of `processing_resolution` resolves to the optimal value from the model config. It is a recommended mode of
    # operation and leads to the most reasonable results. Using the native image resolution or any other processing
    # resolution can lead to loss of either fine details or global context in the output predictions.
    image, padding, original_resolution = self.image_processor.preprocess(
        image, processing_resolution, resample_method_input, dtype
    )  # [N,3,PPH,PPW]

    # 4. Encode input image into latent space. At this step, each of the `N` input images is represented with `E`
    # ensemble members. Each ensemble member is an independent diffused prediction, just initialized independently.
    # Latents of each such predictions across all input images and all ensemble members are represented in the
    # `pred_latent` variable. The variable `image_latent` is of the same shape: it contains each input image encoded
    # into latent space and replicated `E` times. The latents can be either generated (see `generator` to ensure
    # reproducibility), or passed explicitly via the `latents` argument. The latter can be set outside the pipeline
    # code. This behavior can be achieved by setting the `output_latent` argument to `True`. The latent space
    # dimensions are `(h, w)`. Encoding into latent space happens in batches of size `batch_size`.
    # Model invocation: self.vae.encoder.
    image_latent, pred_latent = self.prepare_latents(
        image, latents, generator, ensemble_size, batch_size
    )  # [N*E,4,h,w], [N*E,4,h,w]

    del image

    batch_empty_text_embedding = self.empty_text_embedding.to(dtype=dtype).tile((batch_size, 1, 1))  # [B,1024,2]

    # 5. Process the denoising loop. All `N * E` latents are processed sequentially in batches of size `batch_size`.
    # The unet model takes concatenated latent spaces of the input image and the predicted modality as an input, and
    # outputs noise for the predicted modality's latent space. The number of denoising diffusion steps is defined by
    # `num_inference_steps`. It is either set directly, or resolves to the optimal value specific to the loaded
    # model.
    # Model invocation: self.unet.
    pred_latents = []

    for i in self.progress_bar(
        range(0, num_images * ensemble_size, batch_size), leave=True, desc="Marigold predictions..."
    ):
        batch_image_latent = image_latent[i : i + batch_size]  # [B,4,h,w]
        batch_pred_latent = pred_latent[i : i + batch_size]  # [B,4,h,w]
        effective_batch_size = batch_image_latent.shape[0]
        text = batch_empty_text_embedding[:effective_batch_size]  # [B,2,1024]

        self.scheduler.set_timesteps(num_inference_steps)
        for t in self.progress_bar(self.scheduler.timesteps, leave=False, desc="Diffusion steps..."):
            batch_latent = mint.cat([batch_image_latent, batch_pred_latent], dim=1)  # [B,8,h,w]
            noise = self.unet(batch_latent, t, encoder_hidden_states=text, return_dict=False)[0]  # [B,4,h,w]
            batch_pred_latent = self.scheduler.step(noise, t, batch_pred_latent, generator=generator)[
                0
            ]  # [B,4,h,w]

        pred_latents.append(batch_pred_latent)

    pred_latent = mint.cat(pred_latents, dim=0)  # [N*E,4,h,w]

    del (
        pred_latents,
        image_latent,
        batch_empty_text_embedding,
        batch_image_latent,
        batch_pred_latent,
        text,
        batch_latent,
        noise,
    )

    # 6. Decode predictions from latent into pixel space. The resulting `N * E` predictions have shape `(PPH, PPW)`,
    # which requires slight postprocessing. Decoding into pixel space happens in batches of size `batch_size`.
    # Model invocation: self.vae.decoder.
    prediction = mint.cat(
        [
            self.decode_prediction(pred_latent[i : i + batch_size])
            for i in range(0, pred_latent.shape[0], batch_size)
        ],
        dim=0,
    )  # [N*E,3,PPH,PPW]

    if not output_latent:
        pred_latent = None

    # 7. Remove padding. The output shape is (PH, PW).
    prediction = self.image_processor.unpad_image(prediction, padding)  # [N*E,3,PH,PW]

    # 8. Ensemble and compute uncertainty (when `output_uncertainty` is set). This code treats each of the `N`
    # groups of `E` ensemble predictions independently. For each group it computes an ensembled prediction of shape
    # `(PH, PW)` and an optional uncertainty map of the same dimensions. After computing this pair of outputs for
    # each group independently, it stacks them respectively into batches of `N` almost final predictions and
    # uncertainty maps.
    uncertainty = None
    if ensemble_size > 1:
        prediction = prediction.reshape(num_images, ensemble_size, *prediction.shape[1:])  # [N,E,3,PH,PW]
        prediction = [
            self.ensemble_normals(prediction[i], output_uncertainty, **(ensembling_kwargs or {}))
            for i in range(num_images)
        ]  # [ [[1,3,PH,PW], [1,1,PH,PW]], ... ]
        prediction, uncertainty = zip(*prediction)  # [[1,3,PH,PW], ... ], [[1,1,PH,PW], ... ]
        prediction = mint.cat(prediction, dim=0)  # [N,3,PH,PW]
        if output_uncertainty:
            uncertainty = mint.cat(uncertainty, dim=0)  # [N,1,PH,PW]
        else:
            uncertainty = None

    # 9. If `match_input_resolution` is set, the output prediction and the uncertainty are upsampled to match the
    # input resolution `(H, W)`. This step may introduce upsampling artifacts, and therefore can be disabled.
    # After upsampling, the native resolution normal maps are renormalized to unit length to reduce the artifacts.
    # Depending on the downstream use-case, upsampling can be also chosen based on the tolerated artifacts by
    # setting the `resample_method_output` parameter (e.g., to `"nearest"`).
    if match_input_resolution:
        prediction = self.image_processor.resize_antialias(
            prediction, original_resolution, resample_method_output, is_aa=False
        )  # [N,3,H,W]
        prediction = self.normalize_normals(prediction)  # [N,3,H,W]
        if uncertainty is not None and output_uncertainty:
            uncertainty = self.image_processor.resize_antialias(
                uncertainty, original_resolution, resample_method_output, is_aa=False
            )  # [N,1,H,W]

    # 10. Prepare the final outputs.
    if output_type == "np":
        prediction = self.image_processor.ms_to_numpy(prediction)  # [N,H,W,3]
        if uncertainty is not None and output_uncertainty:
            uncertainty = self.image_processor.ms_to_numpy(uncertainty)  # [N,H,W,1]

    if not return_dict:
        return (prediction, uncertainty, pred_latent)

    return MarigoldNormalsOutput(
        prediction=prediction,
        uncertainty=uncertainty,
        latent=pred_latent,
    )

mindone.diffusers.MarigoldNormalsPipeline.ensemble_normals(normals, output_uncertainty, reduction='closest') staticmethod

Ensembles the normals maps represented by the normals tensor with expected shape (B, 3, H, W), where B is the number of ensemble members for a given prediction of size (H x W).

PARAMETER	DESCRIPTION
normals	Input ensemble normals maps. TYPE: `ms.Tensor`
output_uncertainty	Whether to output uncertainty map. TYPE: `bool`, *optional*, defaults to `False`
reduction	Reduction method used to ensemble aligned predictions. The accepted values are: "closest" and "mean". TYPE: `str`, *optional*, defaults to `"closest"` DEFAULT: 'closest'

RETURNS	DESCRIPTION
Tensor	A tensor of aligned and ensembled normals maps with shape (1, 3, H, W) and optionally a tensor of
Optional[Tensor]	uncertainties of shape (1, 1, H, W).

Source code in mindone/diffusers/pipelines/marigold/pipeline_marigold_normals.py

@staticmethod
def ensemble_normals(
    normals: ms.Tensor, output_uncertainty: bool, reduction: str = "closest"
) -> Tuple[ms.Tensor, Optional[ms.Tensor]]:
    """
    Ensembles the normals maps represented by the `normals` tensor with expected shape `(B, 3, H, W)`, where B is
    the number of ensemble members for a given prediction of size `(H x W)`.

    Args:
        normals (`ms.Tensor`):
            Input ensemble normals maps.
        output_uncertainty (`bool`, *optional*, defaults to `False`):
            Whether to output uncertainty map.
        reduction (`str`, *optional*, defaults to `"closest"`):
            Reduction method used to ensemble aligned predictions. The accepted values are: `"closest"` and
            `"mean"`.

    Returns:
        A tensor of aligned and ensembled normals maps with shape `(1, 3, H, W)` and optionally a tensor of
        uncertainties of shape `(1, 1, H, W)`.
    """
    if normals.ndim != 4 or normals.shape[1] != 3:
        raise ValueError(f"Expecting 4D tensor of shape [B,3,H,W]; got {normals.shape}.")
    if reduction not in ("closest", "mean"):
        raise ValueError(f"Unrecognized reduction method: {reduction}.")

    mean_normals = normals.mean(dim=0, keep_dims=True)  # [1,3,H,W]
    mean_normals = MarigoldNormalsPipeline.normalize_normals(mean_normals)  # [1,3,H,W]

    sim_cos = (mean_normals * normals).sum(dim=1, keepdims=True)  # [E,1,H,W]
    sim_cos = sim_cos.clamp(-1.0, 1.0)  # required to avoid NaN in uncertainty with fp16

    uncertainty = None
    if output_uncertainty:
        uncertainty = sim_cos.arccos()  # [E,1,H,W]
        uncertainty = uncertainty.mean(dim=0, keep_dims=True) / ms.numpy.pi  # [1,1,H,W]

    if reduction == "mean":
        return mean_normals, uncertainty  # [1,3,H,W], [1,1,H,W]

    closest_indices = sim_cos.argmax(dim=0, keepdims=True)  # [1,1,H,W]
    closest_indices = closest_indices.tile((1, 3, 1, 1))  # [1,3,H,W]
    closest_normals = ops.gather_elements(normals, 0, closest_indices)  # [1,3,H,W]

    return closest_normals, uncertainty  # [1,3,H,W], [1,1,H,W]

mindone.diffusers.pipelines.marigold.pipeline_marigold_normals.MarigoldNormalsOutput dataclass

Bases: BaseOutput

Output class for Marigold monocular normals prediction pipeline.

PARAMETER	DESCRIPTION
prediction	Predicted normals with values in the range [-1, 1]. The shape is \(numimages imes 3 imes height imes width\) for ms.Tensor or \(numimages imes height imes width imes 3\) for np.ndarray. TYPE: `np.ndarray`, `ms.Tensor`
uncertainty	Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is \(numimages imes 1 imes height imes width\) for ms.Tensor or \(numimages imes height imes width imes 1\) for np.ndarray. TYPE: `None`, `np.ndarray`, `ms.Tensor`
latent	Latent features corresponding to the predictions, compatible with the latents argument of the pipeline. The shape is \(numimages * numensemble imes 4 imes latentheight imes latentwidth\). TYPE: `None`, `ms.Tensor`

Source code in mindone/diffusers/pipelines/marigold/pipeline_marigold_normals.py

@dataclass
class MarigoldNormalsOutput(BaseOutput):
    """
    Output class for Marigold monocular normals prediction pipeline.

    Args:
        prediction (`np.ndarray`, `ms.Tensor`):
            Predicted normals with values in the range [-1, 1]. The shape is $numimages \times 3 \times height \times
            width$ for `ms.Tensor` or $numimages \times height \times width \times 3$ for `np.ndarray`.
        uncertainty (`None`, `np.ndarray`, `ms.Tensor`):
            Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $numimages
            \times 1 \times height \times width$ for `ms.Tensor` or $numimages \times height \times width \times 1$
            for `np.ndarray`.
        latent (`None`, `ms.Tensor`):
            Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline.
            The shape is $numimages * numensemble \times 4 \times latentheight \times latentwidth$.
    """

    prediction: Union[np.ndarray, ms.Tensor]
    uncertainty: Union[None, np.ndarray, ms.Tensor]
    latent: Union[None, ms.Tensor]

mindone.diffusers.pipelines.marigold.marigold_image_processing.MarigoldImageProcessor.visualize_normals(normals, flip_x=False, flip_y=False, flip_z=False) staticmethod

Visualizes surface normals, such as predictions of the MarigoldNormalsPipeline.

PARAMETER	DESCRIPTION
normals	Surface normals. TYPE: `Union[np.ndarray, ms.Tensor, List[np.ndarray], List[ms.Tensor]]`
flip_x	Flips the X axis of the normals frame of reference. Default direction is right. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
flip_y	Flips the Y axis of the normals frame of reference. Default direction is top. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
flip_z	Flips the Z axis of the normals frame of reference. Default direction is facing the observer. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False

Source code in mindone/diffusers/pipelines/marigold/marigold_image_processing.py

@staticmethod
def visualize_normals(
    normals: Union[
        np.ndarray,
        ms.Tensor,
        List[np.ndarray],
        List[ms.Tensor],
    ],
    flip_x: bool = False,
    flip_y: bool = False,
    flip_z: bool = False,
) -> Union[PIL.Image.Image, List[PIL.Image.Image]]:
    """
    Visualizes surface normals, such as predictions of the `MarigoldNormalsPipeline`.

    Args:
        normals (`Union[np.ndarray, ms.Tensor, List[np.ndarray], List[ms.Tensor]]`):
            Surface normals.
        flip_x (`bool`, *optional*, defaults to `False`): Flips the X axis of the normals frame of reference.
                  Default direction is right.
        flip_y (`bool`, *optional*, defaults to `False`):  Flips the Y axis of the normals frame of reference.
                  Default direction is top.
        flip_z (`bool`, *optional*, defaults to `False`): Flips the Z axis of the normals frame of reference.
                  Default direction is facing the observer.

    Returns: `PIL.Image.Image` or `List[PIL.Image.Image]` with surface normals visualization.
    """
    flip_vec = None
    if any((flip_x, flip_y, flip_z)):
        flip_vec = ms.Tensor(
            [
                (-1) ** flip_x,
                (-1) ** flip_y,
                (-1) ** flip_z,
            ],
            dtype=ms.float32,
        )

    def visualize_normals_one(img, idx=None):
        img = img.permute(1, 2, 0)
        if flip_vec is not None:
            img *= flip_vec
        img = (img + 1.0) * 0.5
        img = (img * 255).to(dtype=ms.uint8).numpy()
        img = PIL.Image.fromarray(img)
        return img

    if normals is None or isinstance(normals, list) and any(o is None for o in normals):
        raise ValueError("Input normals is `None`")
    if isinstance(normals, (np.ndarray, ms.Tensor)):
        normals = MarigoldImageProcessor.expand_tensor_or_array(normals)
        if isinstance(normals, np.ndarray):
            normals = MarigoldImageProcessor.numpy_to_ms(normals)  # [N,3,H,W]
        if not (normals.ndim == 4 and normals.shape[1] == 3):
            raise ValueError(f"Unexpected input shape={normals.shape}, expecting [N,3,H,W].")
        return [visualize_normals_one(img, idx) for idx, img in enumerate(normals)]
    elif isinstance(normals, list):
        return [visualize_normals_one(img, idx) for idx, img in enumerate(normals)]
    else:
        raise ValueError(f"Unexpected input type: {type(normals)}")

Marigold Intrinsic Image Decomposition API

mindone.diffusers.MarigoldIntrinsicsPipeline

Bases: DiffusionPipeline

Pipeline for Intrinsic Image Decomposition (IID) using the Marigold method: https://marigoldcomputervision.github.io.

This model inherits from [DiffusionPipeline]. Check the superclass documentation for the generic methods the library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)

PARAMETER	DESCRIPTION
unet	Conditional U-Net to denoise the targets latent, conditioned on image latent. TYPE: `UNet2DConditionModel`
vae	Variational Auto-Encoder (VAE) Model to encode and decode images and predictions to and from latent representations. TYPE: `AutoencoderKL`
scheduler	A scheduler to be used in combination with unet to denoise the encoded image latents. TYPE: `DDIMScheduler` or `LCMScheduler`
text_encoder	Text-encoder, for empty text embedding. TYPE: `CLIPTextModel`
tokenizer	CLIP tokenizer. TYPE: `CLIPTokenizer`
prediction_type	Type of predictions made by the model. TYPE: `str`, *optional* DEFAULT: None
target_properties	Properties of the predicted modalities, such as target_names, a List[str] used to define the number, order and names of the predicted modalities, and any other metadata that may be required to interpret the predictions. TYPE: `Dict[str, Any]`, *optional* DEFAULT: None
default_denoising_steps	The minimum number of denoising diffusion steps that are required to produce a prediction of reasonable quality with the given model. This value must be set in the model config. When the pipeline is called without explicitly setting num_inference_steps, the default value is used. This is required to ensure reasonable results with various model flavors compatible with the pipeline, such as those relying on very short denoising schedules (LCMScheduler) and those with full diffusion schedules (DDIMScheduler). TYPE: `int`, *optional* DEFAULT: None
default_processing_resolution	The recommended value of the processing_resolution parameter of the pipeline. This value must be set in the model config. When the pipeline is called without explicitly setting processing_resolution, the default value is used. This is required to ensure reasonable results with various model flavors trained with varying optimal processing resolution values. TYPE: `int`, *optional* DEFAULT: None

Source code in mindone/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py

class MarigoldIntrinsicsPipeline(DiffusionPipeline):
    """
    Pipeline for Intrinsic Image Decomposition (IID) using the Marigold method:
    https://marigoldcomputervision.github.io.

    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)

    Args:
        unet (`UNet2DConditionModel`):
            Conditional U-Net to denoise the targets latent, conditioned on image latent.
        vae (`AutoencoderKL`):
            Variational Auto-Encoder (VAE) Model to encode and decode images and predictions to and from latent
            representations.
        scheduler (`DDIMScheduler` or `LCMScheduler`):
            A scheduler to be used in combination with `unet` to denoise the encoded image latents.
        text_encoder (`CLIPTextModel`):
            Text-encoder, for empty text embedding.
        tokenizer (`CLIPTokenizer`):
            CLIP tokenizer.
        prediction_type (`str`, *optional*):
            Type of predictions made by the model.
        target_properties (`Dict[str, Any]`, *optional*):
            Properties of the predicted modalities, such as `target_names`, a `List[str]` used to define the number,
            order and names of the predicted modalities, and any other metadata that may be required to interpret the
            predictions.
        default_denoising_steps (`int`, *optional*):
            The minimum number of denoising diffusion steps that are required to produce a prediction of reasonable
            quality with the given model. This value must be set in the model config. When the pipeline is called
            without explicitly setting `num_inference_steps`, the default value is used. This is required to ensure
            reasonable results with various model flavors compatible with the pipeline, such as those relying on very
            short denoising schedules (`LCMScheduler`) and those with full diffusion schedules (`DDIMScheduler`).
        default_processing_resolution (`int`, *optional*):
            The recommended value of the `processing_resolution` parameter of the pipeline. This value must be set in
            the model config. When the pipeline is called without explicitly setting `processing_resolution`, the
            default value is used. This is required to ensure reasonable results with various model flavors trained
            with varying optimal processing resolution values.
    """

    model_cpu_offload_seq = "text_encoder->unet->vae"
    supported_prediction_types = ("intrinsics",)

    def __init__(
        self,
        unet: UNet2DConditionModel,
        vae: AutoencoderKL,
        scheduler: Union[DDIMScheduler, LCMScheduler],
        text_encoder: CLIPTextModel,
        tokenizer: CLIPTokenizer,
        prediction_type: Optional[str] = None,
        target_properties: Optional[Dict[str, Any]] = None,
        default_denoising_steps: Optional[int] = None,
        default_processing_resolution: Optional[int] = None,
    ):
        super().__init__()

        if prediction_type not in self.supported_prediction_types:
            logger.warning(
                f"Potentially unsupported `prediction_type='{prediction_type}'`; values supported by the pipeline: "
                f"{self.supported_prediction_types}."
            )

        self.register_modules(
            unet=unet,
            vae=vae,
            scheduler=scheduler,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
        )
        self.register_to_config(
            prediction_type=prediction_type,
            target_properties=target_properties,
            default_denoising_steps=default_denoising_steps,
            default_processing_resolution=default_processing_resolution,
        )

        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8

        self.target_properties = target_properties
        self.default_denoising_steps = default_denoising_steps
        self.default_processing_resolution = default_processing_resolution

        self.empty_text_embedding = None

        self.image_processor = MarigoldImageProcessor(vae_scale_factor=self.vae_scale_factor)

    @property
    def n_targets(self):
        return self.unet.config.out_channels // self.vae.config.latent_channels

    def check_inputs(
        self,
        image: PipelineImageInput,
        num_inference_steps: int,
        ensemble_size: int,
        processing_resolution: int,
        resample_method_input: str,
        resample_method_output: str,
        batch_size: int,
        ensembling_kwargs: Optional[Dict[str, Any]],
        latents: Optional[ms.Tensor],
        generator: Optional[Union[np.random.Generator, List[np.random.Generator]]],
        output_type: str,
        output_uncertainty: bool,
    ) -> int:
        actual_vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        if actual_vae_scale_factor != self.vae_scale_factor:
            raise ValueError(
                f"`vae_scale_factor` computed at initialization ({self.vae_scale_factor}) differs from the actual one ({actual_vae_scale_factor})."
            )
        if num_inference_steps is None:
            raise ValueError("`num_inference_steps` is not specified and could not be resolved from the model config.")
        if num_inference_steps < 1:
            raise ValueError("`num_inference_steps` must be positive.")
        if ensemble_size < 1:
            raise ValueError("`ensemble_size` must be positive.")
        if ensemble_size == 2:
            logger.warning(
                "`ensemble_size` == 2 results are similar to no ensembling (1); "
                "consider increasing the value to at least 3."
            )
        if ensemble_size == 1 and output_uncertainty:
            raise ValueError(
                "Computing uncertainty by setting `output_uncertainty=True` also requires setting `ensemble_size` "
                "greater than 1."
            )
        if processing_resolution is None:
            raise ValueError(
                "`processing_resolution` is not specified and could not be resolved from the model config."
            )
        if processing_resolution < 0:
            raise ValueError(
                "`processing_resolution` must be non-negative: 0 for native resolution, or any positive value for "
                "downsampled processing."
            )
        if processing_resolution % self.vae_scale_factor != 0:
            raise ValueError(f"`processing_resolution` must be a multiple of {self.vae_scale_factor}.")
        if resample_method_input not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
            raise ValueError(
                "`resample_method_input` takes string values compatible with PIL library: "
                "nearest, nearest-exact, bilinear, bicubic, area."
            )
        if resample_method_output not in ("nearest", "nearest-exact", "bilinear", "bicubic", "area"):
            raise ValueError(
                "`resample_method_output` takes string values compatible with PIL library: "
                "nearest, nearest-exact, bilinear, bicubic, area."
            )
        if batch_size < 1:
            raise ValueError("`batch_size` must be positive.")
        if output_type not in ["pt", "np"]:
            raise ValueError("`output_type` must be one of `pt` or `np`.")
        if latents is not None and generator is not None:
            raise ValueError("`latents` and `generator` cannot be used together.")
        if ensembling_kwargs is not None:
            if not isinstance(ensembling_kwargs, dict):
                raise ValueError("`ensembling_kwargs` must be a dictionary.")
            if "reduction" in ensembling_kwargs and ensembling_kwargs["reduction"] not in ("median", "mean"):
                raise ValueError("`ensembling_kwargs['reduction']` can be either `'median'` or `'mean'`.")

        # image checks
        num_images = 0
        W, H = None, None
        if not isinstance(image, list):
            image = [image]
        for i, img in enumerate(image):
            if isinstance(img, np.ndarray) or ms.is_tensor(img):
                if img.ndim not in (2, 3, 4):
                    raise ValueError(f"`image[{i}]` has unsupported dimensions or shape: {img.shape}.")
                H_i, W_i = img.shape[-2:]
                N_i = 1
                if img.ndim == 4:
                    N_i = img.shape[0]
            elif isinstance(img, Image.Image):
                W_i, H_i = img.size
                N_i = 1
            else:
                raise ValueError(f"Unsupported `image[{i}]` type: {type(img)}.")
            if W is None:
                W, H = W_i, H_i
            elif (W, H) != (W_i, H_i):
                raise ValueError(
                    f"Input `image[{i}]` has incompatible dimensions {(W_i, H_i)} with the previous images {(W, H)}"
                )
            num_images += N_i

        # latents checks
        if latents is not None:
            if not ms.is_tensor(latents):
                raise ValueError("`latents` must be a ms.Tensor.")
            if latents.ndim != 4:
                raise ValueError(f"`latents` has unsupported dimensions or shape: {latents.shape}.")

            if processing_resolution > 0:
                max_orig = max(H, W)
                new_H = H * processing_resolution // max_orig
                new_W = W * processing_resolution // max_orig
                if new_H == 0 or new_W == 0:
                    raise ValueError(f"Extreme aspect ratio of the input image: [{W} x {H}]")
                W, H = new_W, new_H
            w = (W + self.vae_scale_factor - 1) // self.vae_scale_factor
            h = (H + self.vae_scale_factor - 1) // self.vae_scale_factor
            shape_expected = (num_images * ensemble_size, self.unet.config.out_channels, h, w)

            if latents.shape != shape_expected:
                raise ValueError(f"`latents` has unexpected shape={latents.shape} expected={shape_expected}.")

        # generator checks
        if generator is not None:
            if isinstance(generator, list):
                if len(generator) != num_images * ensemble_size:
                    raise ValueError(
                        "The number of generators must match the total number of ensemble members for all input images."
                    )
            elif not isinstance(generator, np.random.Generator):
                raise ValueError(f"Unsupported generator type: {type(generator)}.")

        return num_images

    def progress_bar(self, iterable=None, total=None, desc=None, leave=True):
        if not hasattr(self, "_progress_bar_config"):
            self._progress_bar_config = {}
        elif not isinstance(self._progress_bar_config, dict):
            raise ValueError(
                f"`self._progress_bar_config` should be of type `dict`, but is {type(self._progress_bar_config)}."
            )

        progress_bar_config = dict(**self._progress_bar_config)
        progress_bar_config["desc"] = progress_bar_config.get("desc", desc)
        progress_bar_config["leave"] = progress_bar_config.get("leave", leave)
        if iterable is not None:
            return tqdm(iterable, **progress_bar_config)
        elif total is not None:
            return tqdm(total=total, **progress_bar_config)
        else:
            raise ValueError("Either `total` or `iterable` has to be defined.")

    def __call__(
        self,
        image: PipelineImageInput,
        num_inference_steps: Optional[int] = None,
        ensemble_size: int = 1,
        processing_resolution: Optional[int] = None,
        match_input_resolution: bool = True,
        resample_method_input: str = "bilinear",
        resample_method_output: str = "bilinear",
        batch_size: int = 1,
        ensembling_kwargs: Optional[Dict[str, Any]] = None,
        latents: Optional[Union[ms.Tensor, List[ms.Tensor]]] = None,
        generator: Optional[Union[np.random.Generator, List[np.random.Generator]]] = None,
        output_type: str = "np",
        output_uncertainty: bool = False,
        output_latent: bool = False,
        return_dict: bool = False,
    ):
        """
        Function invoked when calling the pipeline.

        Args:
            image (`PIL.Image.Image`, `np.ndarray`, `ms.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`),
                `List[ms.Tensor]`: An input image or images used as an input for the intrinsic decomposition task.
                For arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is
                possible by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or
                three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the
                same width and height.
            num_inference_steps (`int`, *optional*, defaults to `None`):
                Number of denoising diffusion steps during inference. The default value `None` results in automatic
                selection.
            ensemble_size (`int`, defaults to `1`):
                Number of ensemble predictions. Higher values result in measurable improvements and visual degradation.
            processing_resolution (`int`, *optional*, defaults to `None`):
                Effective processing resolution. When set to `0`, matches the larger input image dimension. This
                produces crisper predictions, but may also lead to the overall loss of global context. The default
                value `None` resolves to the optimal value from the model config.
            match_input_resolution (`bool`, *optional*, defaults to `True`):
                When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer
                side of the output will equal to `processing_resolution`.
            resample_method_input (`str`, *optional*, defaults to `"bilinear"`):
                Resampling method used to resize input images to `processing_resolution`. The accepted values are:
                `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
            resample_method_output (`str`, *optional*, defaults to `"bilinear"`):
                Resampling method used to resize output predictions to match the input resolution. The accepted values
                are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
            batch_size (`int`, *optional*, defaults to `1`):
                Batch size; only matters when setting `ensemble_size` or passing a tensor of images.
            ensembling_kwargs (`dict`, *optional*, defaults to `None`)
                Extra dictionary with arguments for precise ensembling control. The following options are available:
                - reduction (`str`, *optional*, defaults to `"median"`): Defines the ensembling function applied in
                  every pixel location, can be either `"median"` or `"mean"`.
            latents (`ms.Tensor`, *optional*, defaults to `None`):
                Latent noise tensors to replace the random initialization. These can be taken from the previous
                function call's output.
            generator (`np.Generator`, or `List[np.Generator]`, *optional*, defaults to `None`):
                Random number generator object to ensure reproducibility.
            output_type (`str`, *optional*, defaults to `"np"`):
                Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted
                values are: `"np"` (numpy array) or `"pt"` (mindspore tensor).
            output_uncertainty (`bool`, *optional*, defaults to `False`):
                When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that
                the `ensemble_size` argument is set to a value above 2.
            output_latent (`bool`, *optional*, defaults to `False`):
                When enabled, the output's `latent` field contains the latent codes corresponding to the predictions
                within the ensemble. These codes can be saved, modified, and used for subsequent calls with the
                `latents` argument.
            return_dict (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~pipelines.marigold.MarigoldIntrinsicsOutput`] instead of a plain tuple.

        Examples:

        Returns:
            [`~pipelines.marigold.MarigoldIntrinsicsOutput`] or `tuple`:
                If `return_dict` is `True`, [`~pipelines.marigold.MarigoldIntrinsicsOutput`] is returned, otherwise a
                `tuple` is returned where the first element is the prediction, the second element is the uncertainty
                (or `None`), and the third is the latent (or `None`).
        """

        # 0. Resolving variables.
        dtype = self.dtype

        # Model-specific optimal default values leading to fast and reasonable results.
        if num_inference_steps is None:
            num_inference_steps = self.default_denoising_steps
        if processing_resolution is None:
            processing_resolution = self.default_processing_resolution

        # 1. Check inputs.
        num_images = self.check_inputs(
            image,
            num_inference_steps,
            ensemble_size,
            processing_resolution,
            resample_method_input,
            resample_method_output,
            batch_size,
            ensembling_kwargs,
            latents,
            generator,
            output_type,
            output_uncertainty,
        )

        # 2. Prepare empty text conditioning.
        # Model invocation: self.tokenizer, self.text_encoder.
        if self.empty_text_embedding is None:
            prompt = ""
            text_inputs = self.tokenizer(
                prompt,
                padding="do_not_pad",
                max_length=self.tokenizer.model_max_length,
                truncation=True,
                return_tensors="np",
            )
            text_input_ids = ms.Tensor.from_numpy(text_inputs.input_ids)
            self.empty_text_embedding = self.text_encoder(text_input_ids)[0]  # [1,2,1024]

        # 3. Preprocess input images. This function loads input image or images of compatible dimensions `(H, W)`,
        # optionally downsamples them to the `processing_resolution` `(PH, PW)`, where
        # `max(PH, PW) == processing_resolution`, and pads the dimensions to `(PPH, PPW)` such that these values are
        # divisible by the latent space downscaling factor (typically 8 in Stable Diffusion). The default value `None`
        # of `processing_resolution` resolves to the optimal value from the model config. It is a recommended mode of
        # operation and leads to the most reasonable results. Using the native image resolution or any other processing
        # resolution can lead to loss of either fine details or global context in the output predictions.
        image, padding, original_resolution = self.image_processor.preprocess(
            image, processing_resolution, resample_method_input, dtype
        )  # [N,3,PPH,PPW]

        # 4. Encode input image into latent space. At this step, each of the `N` input images is represented with `E`
        # ensemble members. Each ensemble member is an independent diffused prediction, just initialized independently.
        # Latents of each such predictions across all input images and all ensemble members are represented in the
        # `pred_latent` variable. The variable `image_latent` contains each input image encoded into latent space and
        # replicated `E` times. The variable `pred_latent` contains latents initialization, where the latent space is
        # replicated `T` times relative to the single latent space of `image_latent`, where `T` is the number of the
        # predicted targets. The latents can be either generated (see `generator` to ensure reproducibility), or passed
        # explicitly via the `latents` argument. The latter can be set outside the pipeline code. This behavior can be
        # achieved by setting the `output_latent` argument to `True`. The latent space dimensions are `(h, w)`. Encoding
        # into latent space happens in batches of size `batch_size`.
        # Model invocation: self.vae.encoder.
        image_latent, pred_latent = self.prepare_latents(
            image, latents, generator, ensemble_size, batch_size
        )  # [N*E,4,h,w], [N*E,T*4,h,w]

        del image

        batch_empty_text_embedding = self.empty_text_embedding.to(dtype=dtype).tile((batch_size, 1, 1))  # [B,1024,2]

        # 5. Process the denoising loop. All `N * E` latents are processed sequentially in batches of size `batch_size`.
        # The unet model takes concatenated latent spaces of the input image and the predicted modality as an input, and
        # outputs noise for the predicted modality's latent space. The number of denoising diffusion steps is defined by
        # `num_inference_steps`. It is either set directly, or resolves to the optimal value specific to the loaded
        # model.
        # Model invocation: self.unet.
        pred_latents = []

        for i in self.progress_bar(
            range(0, num_images * ensemble_size, batch_size), leave=True, desc="Marigold predictions..."
        ):
            batch_image_latent = image_latent[i : i + batch_size]  # [B,4,h,w]
            batch_pred_latent = pred_latent[i : i + batch_size]  # [B,T*4,h,w]
            effective_batch_size = batch_image_latent.shape[0]
            text = batch_empty_text_embedding[:effective_batch_size]  # [B,2,1024]

            self.scheduler.set_timesteps(num_inference_steps)
            for t in self.progress_bar(self.scheduler.timesteps, leave=False, desc="Diffusion steps..."):
                batch_latent = mint.cat([batch_image_latent, batch_pred_latent], dim=1)  # [B,(1+T)*4,h,w]
                noise = self.unet(batch_latent, t, encoder_hidden_states=text, return_dict=False)[0]  # [B,T*4,h,w]
                batch_pred_latent = self.scheduler.step(noise, t, batch_pred_latent, generator=generator)[
                    0
                ]  # [B,T*4,h,w]

            pred_latents.append(batch_pred_latent)

        pred_latent = mint.cat(pred_latents, dim=0)  # [N*E,T*4,h,w]

        del (
            pred_latents,
            image_latent,
            batch_empty_text_embedding,
            batch_image_latent,
            batch_pred_latent,
            text,
            batch_latent,
            noise,
        )

        # 6. Decode predictions from latent into pixel space. The resulting `N * E` predictions have shape `(PPH, PPW)`,
        # which requires slight postprocessing. Decoding into pixel space happens in batches of size `batch_size`.
        # Model invocation: self.vae.decoder.
        pred_latent_for_decoding = pred_latent.reshape(
            num_images * ensemble_size * self.n_targets, self.vae.config.latent_channels, *pred_latent.shape[2:]
        )  # [N*E*T,4,PPH,PPW]
        prediction = mint.cat(
            [
                self.decode_prediction(pred_latent_for_decoding[i : i + batch_size])
                for i in range(0, pred_latent_for_decoding.shape[0], batch_size)
            ],
            dim=0,
        )  # [N*E*T,3,PPH,PPW]

        del pred_latent_for_decoding
        if not output_latent:
            pred_latent = None

        # 7. Remove padding. The output shape is (PH, PW).
        prediction = self.image_processor.unpad_image(prediction, padding)  # [N*E*T,3,PH,PW]

        # 8. Ensemble and compute uncertainty (when `output_uncertainty` is set). This code treats each of the `N*T`
        # groups of `E` ensemble predictions independently. For each group it computes an ensembled prediction of shape
        # `(PH, PW)` and an optional uncertainty map of the same dimensions. After computing this pair of outputs for
        # each group independently, it stacks them respectively into batches of `N*T` almost final predictions and
        # uncertainty maps.
        uncertainty = None
        if ensemble_size > 1:
            prediction = prediction.reshape(
                num_images, ensemble_size, self.n_targets, *prediction.shape[1:]
            )  # [N,E,T,3,PH,PW]
            prediction = [
                self.ensemble_intrinsics(prediction[i], output_uncertainty, **(ensembling_kwargs or {}))
                for i in range(num_images)
            ]  # [ [[T,3,PH,PW], [T,3,PH,PW]], ... ]
            prediction, uncertainty = zip(*prediction)  # [[T,3,PH,PW], ... ], [[T,3,PH,PW], ... ]
            prediction = mint.cat(prediction, dim=0)  # [N*T,3,PH,PW]
            if output_uncertainty:
                uncertainty = mint.cat(uncertainty, dim=0)  # [N*T,3,PH,PW]
            else:
                uncertainty = None

        # 9. If `match_input_resolution` is set, the output prediction and the uncertainty are upsampled to match the
        # input resolution `(H, W)`. This step may introduce upsampling artifacts, and therefore can be disabled.
        # Depending on the downstream use-case, upsampling can be also chosen based on the tolerated artifacts by
        # setting the `resample_method_output` parameter (e.g., to `"nearest"`).
        if match_input_resolution:
            prediction = self.image_processor.resize_antialias(
                prediction, original_resolution, resample_method_output, is_aa=False
            )  # [N*T,3,H,W]
            if uncertainty is not None and output_uncertainty:
                uncertainty = self.image_processor.resize_antialias(
                    uncertainty, original_resolution, resample_method_output, is_aa=False
                )  # [N*T,1,H,W]

        # 10. Prepare the final outputs.
        if output_type == "np":
            prediction = self.image_processor.ms_to_numpy(prediction)  # [N*T,H,W,3]
            if uncertainty is not None and output_uncertainty:
                uncertainty = self.image_processor.ms_to_numpy(uncertainty)  # [N*T,H,W,3]

        if not return_dict:
            return (prediction, uncertainty, pred_latent)

        return MarigoldIntrinsicsOutput(
            prediction=prediction,
            uncertainty=uncertainty,
            latent=pred_latent,
        )

    def prepare_latents(
        self,
        image: ms.Tensor,
        latents: Optional[ms.Tensor],
        generator: Optional[np.random.Generator],
        ensemble_size: int,
        batch_size: int,
    ) -> Tuple[ms.Tensor, ms.Tensor]:
        def retrieve_latents(encoder_output):
            assert ms.is_tensor(
                encoder_output
            ), "Could not access latents of provided encoder_output which is not a tensor"
            if hasattr(self.vae, "diag_gauss_dist"):
                return self.vae.diag_gauss_dist.mode(encoder_output)
            else:
                return encoder_output

        image_latent = mint.cat(
            [
                retrieve_latents(self.vae.encode(image[i : i + batch_size])[0])
                for i in range(0, image.shape[0], batch_size)
            ],
            dim=0,
        )  # [N,4,h,w]
        image_latent = image_latent * self.vae.config.scaling_factor
        image_latent = image_latent.repeat_interleave(ensemble_size, dim=0)  # [N*E,4,h,w]
        N_E, C, H, W = image_latent.shape

        pred_latent = latents
        if pred_latent is None:
            pred_latent = randn_tensor(
                (N_E, self.n_targets * C, H, W),
                generator=generator,
                dtype=image_latent.dtype,
            )  # [N*E,T*4,h,w]

        return image_latent, pred_latent

    def decode_prediction(self, pred_latent: ms.Tensor) -> ms.Tensor:
        if pred_latent.ndim != 4 or pred_latent.shape[1] != self.vae.config.latent_channels:
            raise ValueError(
                f"Expecting 4D tensor of shape [B,{self.vae.config.latent_channels},H,W]; got {pred_latent.shape}."
            )

        prediction = self.vae.decode(pred_latent / self.vae.config.scaling_factor, return_dict=False)[0]  # [B,3,H,W]

        prediction = mint.clip(prediction, -1.0, 1.0)  # [B,3,H,W]
        prediction = (prediction + 1.0) / 2.0

        return prediction  # [B,3,H,W]

    @staticmethod
    def ensemble_intrinsics(
        targets: ms.Tensor,
        output_uncertainty: bool = False,
        reduction: str = "median",
    ) -> Tuple[ms.Tensor, Optional[ms.Tensor]]:
        """
        Ensembles the intrinsic decomposition represented by the `targets` tensor with expected shape `(B, T, 3, H,
        W)`, where B is the number of ensemble members for a given prediction of size `(H x W)`, and T is the number of
        predicted targets.

        Args:
            targets (`ms.Tensor`):
                Input ensemble of intrinsic image decomposition maps.
            output_uncertainty (`bool`, *optional*, defaults to `False`):
                Whether to output uncertainty map.
            reduction (`str`, *optional*, defaults to `"mean"`):
                Reduction method used to ensemble aligned predictions. The accepted values are: `"median"` and
                `"mean"`.

        Returns:
            A tensor of aligned and ensembled intrinsic decomposition maps with shape `(T, 3, H, W)` and optionally a
            tensor of uncertainties of shape `(T, 3, H, W)`.
        """
        if targets.ndim != 5 or targets.shape[2] != 3:
            raise ValueError(f"Expecting 4D tensor of shape [B,T,3,H,W]; got {targets.shape}.")
        if reduction not in ("median", "mean"):
            raise ValueError(f"Unrecognized reduction method: {reduction}.")

        B, T, _, H, W = targets.shape
        uncertainty = None
        if reduction == "mean":
            prediction = mint.mean(targets, dim=0)  # [T,3,H,W]
            if output_uncertainty:
                uncertainty = mint.std(targets, dim=0)  # [T,3,H,W]
        elif reduction == "median":
            prediction = mint.median(targets, dim=0, keepdim=True)[0]  # [1,T,3,H,W]
            if output_uncertainty:
                uncertainty = mint.abs(targets - prediction)  # [B,T,3,H,W]
                uncertainty = mint.median(uncertainty, dim=0)[0]  # [T,3,H,W]
            prediction = prediction.squeeze(0)  # [T,3,H,W]
        else:
            raise ValueError(f"Unrecognized reduction method: {reduction}.")
        return prediction, uncertainty

mindone.diffusers.MarigoldIntrinsicsPipeline.__call__(image, num_inference_steps=None, ensemble_size=1, processing_resolution=None, match_input_resolution=True, resample_method_input='bilinear', resample_method_output='bilinear', batch_size=1, ensembling_kwargs=None, latents=None, generator=None, output_type='np', output_uncertainty=False, output_latent=False, return_dict=False)

Function invoked when calling the pipeline.

PARAMETER	DESCRIPTION
num_inference_steps	Number of denoising diffusion steps during inference. The default value None results in automatic selection. TYPE: `int`, *optional*, defaults to `None` DEFAULT: None
ensemble_size	Number of ensemble predictions. Higher values result in measurable improvements and visual degradation. TYPE: `int`, defaults to `1` DEFAULT: 1
processing_resolution	Effective processing resolution. When set to 0, matches the larger input image dimension. This produces crisper predictions, but may also lead to the overall loss of global context. The default value None resolves to the optimal value from the model config. TYPE: `int`, *optional*, defaults to `None` DEFAULT: None
match_input_resolution	When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer side of the output will equal to processing_resolution. TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True
resample_method_input	Resampling method used to resize input images to processing_resolution. The accepted values are: "nearest", "nearest-exact", "bilinear", "bicubic", or "area". TYPE: `str`, *optional*, defaults to `"bilinear"` DEFAULT: 'bilinear'
resample_method_output	Resampling method used to resize output predictions to match the input resolution. The accepted values are "nearest", "nearest-exact", "bilinear", "bicubic", or "area". TYPE: `str`, *optional*, defaults to `"bilinear"` DEFAULT: 'bilinear'
batch_size	Batch size; only matters when setting ensemble_size or passing a tensor of images. TYPE: `int`, *optional*, defaults to `1` DEFAULT: 1
latents	Latent noise tensors to replace the random initialization. These can be taken from the previous function call's output. TYPE: `ms.Tensor`, *optional*, defaults to `None` DEFAULT: None
generator	Random number generator object to ensure reproducibility. TYPE: `np.Generator`, or `List[np.Generator]`, *optional*, defaults to `None` DEFAULT: None
output_type	Preferred format of the output's prediction and the optional uncertainty fields. The accepted values are: "np" (numpy array) or "pt" (mindspore tensor). TYPE: `str`, *optional*, defaults to `"np"` DEFAULT: 'np'
output_uncertainty	When enabled, the output's uncertainty field contains the predictive uncertainty map, provided that the ensemble_size argument is set to a value above 2. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
output_latent	When enabled, the output's latent field contains the latent codes corresponding to the predictions within the ensemble. These codes can be saved, modified, and used for subsequent calls with the latents argument. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
return_dict	Whether or not to return a [~pipelines.marigold.MarigoldIntrinsicsOutput] instead of a plain tuple. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False

RETURNS	DESCRIPTION
	[~pipelines.marigold.MarigoldIntrinsicsOutput] or tuple: If return_dict is True, [~pipelines.marigold.MarigoldIntrinsicsOutput] is returned, otherwise a tuple is returned where the first element is the prediction, the second element is the uncertainty (or None), and the third is the latent (or None).

Source code in mindone/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py

def __call__(
    self,
    image: PipelineImageInput,
    num_inference_steps: Optional[int] = None,
    ensemble_size: int = 1,
    processing_resolution: Optional[int] = None,
    match_input_resolution: bool = True,
    resample_method_input: str = "bilinear",
    resample_method_output: str = "bilinear",
    batch_size: int = 1,
    ensembling_kwargs: Optional[Dict[str, Any]] = None,
    latents: Optional[Union[ms.Tensor, List[ms.Tensor]]] = None,
    generator: Optional[Union[np.random.Generator, List[np.random.Generator]]] = None,
    output_type: str = "np",
    output_uncertainty: bool = False,
    output_latent: bool = False,
    return_dict: bool = False,
):
    """
    Function invoked when calling the pipeline.

    Args:
        image (`PIL.Image.Image`, `np.ndarray`, `ms.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`),
            `List[ms.Tensor]`: An input image or images used as an input for the intrinsic decomposition task.
            For arrays and tensors, the expected value range is between `[0, 1]`. Passing a batch of images is
            possible by providing a four-dimensional array or a tensor. Additionally, a list of images of two- or
            three-dimensional arrays or tensors can be passed. In the latter case, all list elements must have the
            same width and height.
        num_inference_steps (`int`, *optional*, defaults to `None`):
            Number of denoising diffusion steps during inference. The default value `None` results in automatic
            selection.
        ensemble_size (`int`, defaults to `1`):
            Number of ensemble predictions. Higher values result in measurable improvements and visual degradation.
        processing_resolution (`int`, *optional*, defaults to `None`):
            Effective processing resolution. When set to `0`, matches the larger input image dimension. This
            produces crisper predictions, but may also lead to the overall loss of global context. The default
            value `None` resolves to the optimal value from the model config.
        match_input_resolution (`bool`, *optional*, defaults to `True`):
            When enabled, the output prediction is resized to match the input dimensions. When disabled, the longer
            side of the output will equal to `processing_resolution`.
        resample_method_input (`str`, *optional*, defaults to `"bilinear"`):
            Resampling method used to resize input images to `processing_resolution`. The accepted values are:
            `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
        resample_method_output (`str`, *optional*, defaults to `"bilinear"`):
            Resampling method used to resize output predictions to match the input resolution. The accepted values
            are `"nearest"`, `"nearest-exact"`, `"bilinear"`, `"bicubic"`, or `"area"`.
        batch_size (`int`, *optional*, defaults to `1`):
            Batch size; only matters when setting `ensemble_size` or passing a tensor of images.
        ensembling_kwargs (`dict`, *optional*, defaults to `None`)
            Extra dictionary with arguments for precise ensembling control. The following options are available:
            - reduction (`str`, *optional*, defaults to `"median"`): Defines the ensembling function applied in
              every pixel location, can be either `"median"` or `"mean"`.
        latents (`ms.Tensor`, *optional*, defaults to `None`):
            Latent noise tensors to replace the random initialization. These can be taken from the previous
            function call's output.
        generator (`np.Generator`, or `List[np.Generator]`, *optional*, defaults to `None`):
            Random number generator object to ensure reproducibility.
        output_type (`str`, *optional*, defaults to `"np"`):
            Preferred format of the output's `prediction` and the optional `uncertainty` fields. The accepted
            values are: `"np"` (numpy array) or `"pt"` (mindspore tensor).
        output_uncertainty (`bool`, *optional*, defaults to `False`):
            When enabled, the output's `uncertainty` field contains the predictive uncertainty map, provided that
            the `ensemble_size` argument is set to a value above 2.
        output_latent (`bool`, *optional*, defaults to `False`):
            When enabled, the output's `latent` field contains the latent codes corresponding to the predictions
            within the ensemble. These codes can be saved, modified, and used for subsequent calls with the
            `latents` argument.
        return_dict (`bool`, *optional*, defaults to `False`):
            Whether or not to return a [`~pipelines.marigold.MarigoldIntrinsicsOutput`] instead of a plain tuple.

    Examples:

    Returns:
        [`~pipelines.marigold.MarigoldIntrinsicsOutput`] or `tuple`:
            If `return_dict` is `True`, [`~pipelines.marigold.MarigoldIntrinsicsOutput`] is returned, otherwise a
            `tuple` is returned where the first element is the prediction, the second element is the uncertainty
            (or `None`), and the third is the latent (or `None`).
    """

    # 0. Resolving variables.
    dtype = self.dtype

    # Model-specific optimal default values leading to fast and reasonable results.
    if num_inference_steps is None:
        num_inference_steps = self.default_denoising_steps
    if processing_resolution is None:
        processing_resolution = self.default_processing_resolution

    # 1. Check inputs.
    num_images = self.check_inputs(
        image,
        num_inference_steps,
        ensemble_size,
        processing_resolution,
        resample_method_input,
        resample_method_output,
        batch_size,
        ensembling_kwargs,
        latents,
        generator,
        output_type,
        output_uncertainty,
    )

    # 2. Prepare empty text conditioning.
    # Model invocation: self.tokenizer, self.text_encoder.
    if self.empty_text_embedding is None:
        prompt = ""
        text_inputs = self.tokenizer(
            prompt,
            padding="do_not_pad",
            max_length=self.tokenizer.model_max_length,
            truncation=True,
            return_tensors="np",
        )
        text_input_ids = ms.Tensor.from_numpy(text_inputs.input_ids)
        self.empty_text_embedding = self.text_encoder(text_input_ids)[0]  # [1,2,1024]

    # 3. Preprocess input images. This function loads input image or images of compatible dimensions `(H, W)`,
    # optionally downsamples them to the `processing_resolution` `(PH, PW)`, where
    # `max(PH, PW) == processing_resolution`, and pads the dimensions to `(PPH, PPW)` such that these values are
    # divisible by the latent space downscaling factor (typically 8 in Stable Diffusion). The default value `None`
    # of `processing_resolution` resolves to the optimal value from the model config. It is a recommended mode of
    # operation and leads to the most reasonable results. Using the native image resolution or any other processing
    # resolution can lead to loss of either fine details or global context in the output predictions.
    image, padding, original_resolution = self.image_processor.preprocess(
        image, processing_resolution, resample_method_input, dtype
    )  # [N,3,PPH,PPW]

    # 4. Encode input image into latent space. At this step, each of the `N` input images is represented with `E`
    # ensemble members. Each ensemble member is an independent diffused prediction, just initialized independently.
    # Latents of each such predictions across all input images and all ensemble members are represented in the
    # `pred_latent` variable. The variable `image_latent` contains each input image encoded into latent space and
    # replicated `E` times. The variable `pred_latent` contains latents initialization, where the latent space is
    # replicated `T` times relative to the single latent space of `image_latent`, where `T` is the number of the
    # predicted targets. The latents can be either generated (see `generator` to ensure reproducibility), or passed
    # explicitly via the `latents` argument. The latter can be set outside the pipeline code. This behavior can be
    # achieved by setting the `output_latent` argument to `True`. The latent space dimensions are `(h, w)`. Encoding
    # into latent space happens in batches of size `batch_size`.
    # Model invocation: self.vae.encoder.
    image_latent, pred_latent = self.prepare_latents(
        image, latents, generator, ensemble_size, batch_size
    )  # [N*E,4,h,w], [N*E,T*4,h,w]

    del image

    batch_empty_text_embedding = self.empty_text_embedding.to(dtype=dtype).tile((batch_size, 1, 1))  # [B,1024,2]

    # 5. Process the denoising loop. All `N * E` latents are processed sequentially in batches of size `batch_size`.
    # The unet model takes concatenated latent spaces of the input image and the predicted modality as an input, and
    # outputs noise for the predicted modality's latent space. The number of denoising diffusion steps is defined by
    # `num_inference_steps`. It is either set directly, or resolves to the optimal value specific to the loaded
    # model.
    # Model invocation: self.unet.
    pred_latents = []

    for i in self.progress_bar(
        range(0, num_images * ensemble_size, batch_size), leave=True, desc="Marigold predictions..."
    ):
        batch_image_latent = image_latent[i : i + batch_size]  # [B,4,h,w]
        batch_pred_latent = pred_latent[i : i + batch_size]  # [B,T*4,h,w]
        effective_batch_size = batch_image_latent.shape[0]
        text = batch_empty_text_embedding[:effective_batch_size]  # [B,2,1024]

        self.scheduler.set_timesteps(num_inference_steps)
        for t in self.progress_bar(self.scheduler.timesteps, leave=False, desc="Diffusion steps..."):
            batch_latent = mint.cat([batch_image_latent, batch_pred_latent], dim=1)  # [B,(1+T)*4,h,w]
            noise = self.unet(batch_latent, t, encoder_hidden_states=text, return_dict=False)[0]  # [B,T*4,h,w]
            batch_pred_latent = self.scheduler.step(noise, t, batch_pred_latent, generator=generator)[
                0
            ]  # [B,T*4,h,w]

        pred_latents.append(batch_pred_latent)

    pred_latent = mint.cat(pred_latents, dim=0)  # [N*E,T*4,h,w]

    del (
        pred_latents,
        image_latent,
        batch_empty_text_embedding,
        batch_image_latent,
        batch_pred_latent,
        text,
        batch_latent,
        noise,
    )

    # 6. Decode predictions from latent into pixel space. The resulting `N * E` predictions have shape `(PPH, PPW)`,
    # which requires slight postprocessing. Decoding into pixel space happens in batches of size `batch_size`.
    # Model invocation: self.vae.decoder.
    pred_latent_for_decoding = pred_latent.reshape(
        num_images * ensemble_size * self.n_targets, self.vae.config.latent_channels, *pred_latent.shape[2:]
    )  # [N*E*T,4,PPH,PPW]
    prediction = mint.cat(
        [
            self.decode_prediction(pred_latent_for_decoding[i : i + batch_size])
            for i in range(0, pred_latent_for_decoding.shape[0], batch_size)
        ],
        dim=0,
    )  # [N*E*T,3,PPH,PPW]

    del pred_latent_for_decoding
    if not output_latent:
        pred_latent = None

    # 7. Remove padding. The output shape is (PH, PW).
    prediction = self.image_processor.unpad_image(prediction, padding)  # [N*E*T,3,PH,PW]

    # 8. Ensemble and compute uncertainty (when `output_uncertainty` is set). This code treats each of the `N*T`
    # groups of `E` ensemble predictions independently. For each group it computes an ensembled prediction of shape
    # `(PH, PW)` and an optional uncertainty map of the same dimensions. After computing this pair of outputs for
    # each group independently, it stacks them respectively into batches of `N*T` almost final predictions and
    # uncertainty maps.
    uncertainty = None
    if ensemble_size > 1:
        prediction = prediction.reshape(
            num_images, ensemble_size, self.n_targets, *prediction.shape[1:]
        )  # [N,E,T,3,PH,PW]
        prediction = [
            self.ensemble_intrinsics(prediction[i], output_uncertainty, **(ensembling_kwargs or {}))
            for i in range(num_images)
        ]  # [ [[T,3,PH,PW], [T,3,PH,PW]], ... ]
        prediction, uncertainty = zip(*prediction)  # [[T,3,PH,PW], ... ], [[T,3,PH,PW], ... ]
        prediction = mint.cat(prediction, dim=0)  # [N*T,3,PH,PW]
        if output_uncertainty:
            uncertainty = mint.cat(uncertainty, dim=0)  # [N*T,3,PH,PW]
        else:
            uncertainty = None

    # 9. If `match_input_resolution` is set, the output prediction and the uncertainty are upsampled to match the
    # input resolution `(H, W)`. This step may introduce upsampling artifacts, and therefore can be disabled.
    # Depending on the downstream use-case, upsampling can be also chosen based on the tolerated artifacts by
    # setting the `resample_method_output` parameter (e.g., to `"nearest"`).
    if match_input_resolution:
        prediction = self.image_processor.resize_antialias(
            prediction, original_resolution, resample_method_output, is_aa=False
        )  # [N*T,3,H,W]
        if uncertainty is not None and output_uncertainty:
            uncertainty = self.image_processor.resize_antialias(
                uncertainty, original_resolution, resample_method_output, is_aa=False
            )  # [N*T,1,H,W]

    # 10. Prepare the final outputs.
    if output_type == "np":
        prediction = self.image_processor.ms_to_numpy(prediction)  # [N*T,H,W,3]
        if uncertainty is not None and output_uncertainty:
            uncertainty = self.image_processor.ms_to_numpy(uncertainty)  # [N*T,H,W,3]

    if not return_dict:
        return (prediction, uncertainty, pred_latent)

    return MarigoldIntrinsicsOutput(
        prediction=prediction,
        uncertainty=uncertainty,
        latent=pred_latent,
    )

mindone.diffusers.MarigoldIntrinsicsPipeline.ensemble_intrinsics(targets, output_uncertainty=False, reduction='median') staticmethod

Ensembles the intrinsic decomposition represented by the targets tensor with expected shape (B, T, 3, H, W), where B is the number of ensemble members for a given prediction of size (H x W), and T is the number of predicted targets.

PARAMETER	DESCRIPTION
targets	Input ensemble of intrinsic image decomposition maps. TYPE: `ms.Tensor`
output_uncertainty	Whether to output uncertainty map. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
reduction	Reduction method used to ensemble aligned predictions. The accepted values are: "median" and "mean". TYPE: `str`, *optional*, defaults to `"mean"` DEFAULT: 'median'

RETURNS	DESCRIPTION
Tensor	A tensor of aligned and ensembled intrinsic decomposition maps with shape (T, 3, H, W) and optionally a
Optional[Tensor]	tensor of uncertainties of shape (T, 3, H, W).

Source code in mindone/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py

@staticmethod
def ensemble_intrinsics(
    targets: ms.Tensor,
    output_uncertainty: bool = False,
    reduction: str = "median",
) -> Tuple[ms.Tensor, Optional[ms.Tensor]]:
    """
    Ensembles the intrinsic decomposition represented by the `targets` tensor with expected shape `(B, T, 3, H,
    W)`, where B is the number of ensemble members for a given prediction of size `(H x W)`, and T is the number of
    predicted targets.

    Args:
        targets (`ms.Tensor`):
            Input ensemble of intrinsic image decomposition maps.
        output_uncertainty (`bool`, *optional*, defaults to `False`):
            Whether to output uncertainty map.
        reduction (`str`, *optional*, defaults to `"mean"`):
            Reduction method used to ensemble aligned predictions. The accepted values are: `"median"` and
            `"mean"`.

    Returns:
        A tensor of aligned and ensembled intrinsic decomposition maps with shape `(T, 3, H, W)` and optionally a
        tensor of uncertainties of shape `(T, 3, H, W)`.
    """
    if targets.ndim != 5 or targets.shape[2] != 3:
        raise ValueError(f"Expecting 4D tensor of shape [B,T,3,H,W]; got {targets.shape}.")
    if reduction not in ("median", "mean"):
        raise ValueError(f"Unrecognized reduction method: {reduction}.")

    B, T, _, H, W = targets.shape
    uncertainty = None
    if reduction == "mean":
        prediction = mint.mean(targets, dim=0)  # [T,3,H,W]
        if output_uncertainty:
            uncertainty = mint.std(targets, dim=0)  # [T,3,H,W]
    elif reduction == "median":
        prediction = mint.median(targets, dim=0, keepdim=True)[0]  # [1,T,3,H,W]
        if output_uncertainty:
            uncertainty = mint.abs(targets - prediction)  # [B,T,3,H,W]
            uncertainty = mint.median(uncertainty, dim=0)[0]  # [T,3,H,W]
        prediction = prediction.squeeze(0)  # [T,3,H,W]
    else:
        raise ValueError(f"Unrecognized reduction method: {reduction}.")
    return prediction, uncertainty

mindone.diffusers.pipelines.marigold.pipeline_marigold_intrinsics.MarigoldIntrinsicsOutput dataclass

Bases: BaseOutput

Output class for Marigold Intrinsic Image Decomposition pipeline.

PARAMETER	DESCRIPTION
prediction	Predicted image intrinsics with values in the range [0, 1]. The shape is \((numimages * numtargets) imes 3 imes height imes width\) for ms.Tensor or \((numimages * numtargets) imes height imes width imes 3\) for np.ndarray, where numtargets corresponds to the number of predicted target modalities of the intrinsic image decomposition. TYPE: `np.ndarray`, `ms.Tensor`
uncertainty	Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is \((numimages * numtargets) imes 3 imes height imes width\) for ms.Tensor or \((numimages * numtargets) imes height imes width imes 3\) for np.ndarray. TYPE: `None`, `np.ndarray`, `ms.Tensor`
latent	Latent features corresponding to the predictions, compatible with the latents argument of the pipeline. The shape is \((numimages * numensemble) imes (numtargets * 4) imes latentheight imes latentwidth\). TYPE: `None`, `ms.Tensor`

Source code in mindone/diffusers/pipelines/marigold/pipeline_marigold_intrinsics.py

@dataclass
class MarigoldIntrinsicsOutput(BaseOutput):
    """
    Output class for Marigold Intrinsic Image Decomposition pipeline.

    Args:
        prediction (`np.ndarray`, `ms.Tensor`):
            Predicted image intrinsics with values in the range [0, 1]. The shape is $(numimages * numtargets) \times 3
            \times height \times width$ for `ms.Tensor` or $(numimages * numtargets) \times height \times width
            \times 3$ for `np.ndarray`, where `numtargets` corresponds to the number of predicted target modalities of
            the intrinsic image decomposition.
        uncertainty (`None`, `np.ndarray`, `ms.Tensor`):
            Uncertainty maps computed from the ensemble, with values in the range [0, 1]. The shape is $(numimages *
            numtargets) \times 3 \times height \times width$ for `ms.Tensor` or $(numimages * numtargets) \times
            height \times width \times 3$ for `np.ndarray`.
        latent (`None`, `ms.Tensor`):
            Latent features corresponding to the predictions, compatible with the `latents` argument of the pipeline.
            The shape is $(numimages * numensemble) \times (numtargets * 4) \times latentheight \times latentwidth$.
    """

    prediction: Union[np.ndarray, ms.Tensor]
    uncertainty: Union[None, np.ndarray, ms.Tensor]
    latent: Union[None, ms.Tensor]

mindone.diffusers.pipelines.marigold.marigold_image_processing.MarigoldImageProcessor.visualize_intrinsics(prediction, target_properties, color_map='binary') staticmethod

Visualizes intrinsic image decomposition, such as predictions of the MarigoldIntrinsicsPipeline.

PARAMETER	DESCRIPTION
prediction	Intrinsic image decomposition. TYPE: `Union[np.ndarray, ms.Tensor, List[np.ndarray], List[ms.Tensor]]`
target_properties	Decomposition properties. Expected entries: target_names: List[str] and a dictionary with keys prediction_space: str, sub_target_names: List[Union[str, Null]] (must have 3 entries, null for missing modalities), up_to_scale: bool, one for each target and sub-target. TYPE: `Dict[str, Any]`
color_map	Color map used to convert a single-channel predictions into colored representations. When a dictionary is passed, each modality can be colored with its own color map. TYPE: `Union[str, Dict[str, str]]`, *optional*, defaults to `"Spectral"` DEFAULT: 'binary'

Source code in mindone/diffusers/pipelines/marigold/marigold_image_processing.py

@staticmethod
def visualize_intrinsics(
    prediction: Union[
        np.ndarray,
        ms.Tensor,
        List[np.ndarray],
        List[ms.Tensor],
    ],
    target_properties: Dict[str, Any],
    color_map: Union[str, Dict[str, str]] = "binary",
) -> List[Dict[str, PIL.Image.Image]]:
    """
    Visualizes intrinsic image decomposition, such as predictions of the `MarigoldIntrinsicsPipeline`.

    Args:
        prediction (`Union[np.ndarray, ms.Tensor, List[np.ndarray], List[ms.Tensor]]`):
            Intrinsic image decomposition.
        target_properties (`Dict[str, Any]`):
            Decomposition properties. Expected entries: `target_names: List[str]` and a dictionary with keys
            `prediction_space: str`, `sub_target_names: List[Union[str, Null]]` (must have 3 entries, null for
            missing modalities), `up_to_scale: bool`, one for each target and sub-target.
        color_map (`Union[str, Dict[str, str]]`, *optional*, defaults to `"Spectral"`):
            Color map used to convert a single-channel predictions into colored representations. When a dictionary
            is passed, each modality can be colored with its own color map.

    Returns: `List[Dict[str, PIL.Image.Image]]` with intrinsic image decomposition visualization.
    """
    if "target_names" not in target_properties:
        raise ValueError("Missing `target_names` in target_properties")
    if not isinstance(color_map, str) and not (
        isinstance(color_map, dict) and all(isinstance(k, str) and isinstance(v, str) for k, v in color_map.items())
    ):
        raise ValueError("`color_map` must be a string or a dictionary of strings")
    n_targets = len(target_properties["target_names"])

    def visualize_targets_one(images, idx=None):
        # img: [T, 3, H, W]
        out = {}
        for target_name, img in zip(target_properties["target_names"], images):
            img = img.permute(1, 2, 0)  # [H, W, 3]
            prediction_space = target_properties[target_name].get("prediction_space", "srgb")
            if prediction_space == "stack":
                sub_target_names = target_properties[target_name]["sub_target_names"]
                if len(sub_target_names) != 3 or any(
                    not (isinstance(s, str) or s is None) for s in sub_target_names
                ):
                    raise ValueError(f"Unexpected target sub-names {sub_target_names} in {target_name}")
                for i, sub_target_name in enumerate(sub_target_names):
                    if sub_target_name is None:
                        continue
                    sub_img = img[:, :, i]
                    sub_prediction_space = target_properties[sub_target_name].get("prediction_space", "srgb")
                    if sub_prediction_space == "linear":
                        sub_up_to_scale = target_properties[sub_target_name].get("up_to_scale", False)
                        if sub_up_to_scale:
                            sub_img = sub_img / max(sub_img.max().item(), 1e-6)
                        sub_img = sub_img ** (1 / 2.2)
                    cmap_name = (
                        color_map if isinstance(color_map, str) else color_map.get(sub_target_name, "binary")
                    )
                    sub_img = MarigoldImageProcessor.colormap(sub_img, cmap=cmap_name, bytes=True)
                    sub_img = PIL.Image.fromarray(sub_img.numpy())
                    out[sub_target_name] = sub_img
            elif prediction_space == "linear":
                up_to_scale = target_properties[target_name].get("up_to_scale", False)
                if up_to_scale:
                    img = img / max(img.max().item(), 1e-6)
                img = img ** (1 / 2.2)
            elif prediction_space == "srgb":
                pass
            img = (img * 255).to(dtype=ms.uint8).numpy()
            img = PIL.Image.fromarray(img)
            out[target_name] = img
        return out

    if prediction is None or isinstance(prediction, list) and any(o is None for o in prediction):
        raise ValueError("Input prediction is `None`")
    if isinstance(prediction, (np.ndarray, ms.Tensor)):
        prediction = MarigoldImageProcessor.expand_tensor_or_array(prediction)
        if isinstance(prediction, np.ndarray):
            prediction = MarigoldImageProcessor.numpy_to_ms(prediction)  # [N*T,3,H,W]
        if not (prediction.ndim == 4 and prediction.shape[1] == 3 and prediction.shape[0] % n_targets == 0):
            raise ValueError(f"Unexpected input shape={prediction.shape}, expecting [N*T,3,H,W].")
        N_T, _, H, W = prediction.shape
        N = N_T // n_targets
        prediction = prediction.reshape(N, n_targets, 3, H, W)
        return [visualize_targets_one(img, idx) for idx, img in enumerate(prediction)]
    elif isinstance(prediction, list):
        return [visualize_targets_one(img, idx) for idx, img in enumerate(prediction)]
    else:
        raise ValueError(f"Unexpected input type: {type(prediction)}")