UNet2DModel

The UNet model was originally introduced by Ronneberger et al. for biomedical image segmentation, but it is also commonly used in 🤗 Diffusers because it outputs images that are the same size as the input. It is one of the most important components of a diffusion system because it facilitates the actual diffusion process. There are several variants of the UNet model in 🤗 Diffusers, depending on it's number of dimensions and whether it is a conditional model or not. This is a 2D UNet model.

The abstract from the paper is:

There is large consent that successful training of deep networks requires many thousand annotated training samples. In this paper, we present a network and training strategy that relies on the strong use of data augmentation to use the available annotated samples more efficiently. The architecture consists of a contracting path to capture context and a symmetric expanding path that enables precise localization. We show that such a network can be trained end-to-end from very few images and outperforms the prior best method (a sliding-window convolutional network) on the ISBI challenge for segmentation of neuronal structures in electron microscopic stacks. Using the same network trained on transmitted light microscopy images (phase contrast and DIC) we won the ISBI cell tracking challenge 2015 in these categories by a large margin. Moreover, the network is fast. Segmentation of a 512x512 image takes less than a second on a recent GPU. The full implementation (based on Caffe) and the trained networks are available at http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net.

mindone.diffusers.UNet2DModel

Bases: ModelMixin, ConfigMixin

A 2D UNet model that takes a noisy sample and a timestep and returns a sample shaped output.

This model inherits from [ModelMixin]. Check the superclass documentation for it's generic methods implemented for all models (such as downloading or saving).

PARAMETER	DESCRIPTION
sample_size	Height and width of input/output sample. Dimensions must be a multiple of 2 ** (len(block_out_channels) - 1). TYPE: `int` or `Tuple[int, int]`, *optional*, defaults to `None` DEFAULT: None
in_channels	Number of channels in the input sample. TYPE: `int`, *optional*, defaults to 3 DEFAULT: 3
out_channels	Number of channels in the output. TYPE: `int`, *optional*, defaults to 3 DEFAULT: 3
center_input_sample	Whether to center the input sample. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False
time_embedding_type	Type of time embedding to use. TYPE: `str`, *optional*, defaults to `"positional"` DEFAULT: 'positional'
freq_shift	Frequency shift for Fourier time embedding. TYPE: `int`, *optional*, defaults to 0 DEFAULT: 0
flip_sin_to_cos	Whether to flip sin to cos for Fourier time embedding. TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True
down_block_types	Tuple of downsample block types. TYPE: `Tuple[str]`, *optional*, defaults to `("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")` DEFAULT: ('DownBlock2D', 'AttnDownBlock2D', 'AttnDownBlock2D', 'AttnDownBlock2D')
mid_block_type	Block type for middle of UNet, it can be either UNetMidBlock2D or UnCLIPUNetMidBlock2D. TYPE: `str`, *optional*, defaults to `"UNetMidBlock2D"`
up_block_types	Tuple of upsample block types. TYPE: `Tuple[str]`, *optional*, defaults to `("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")` DEFAULT: ('AttnUpBlock2D', 'AttnUpBlock2D', 'AttnUpBlock2D', 'UpBlock2D')
block_out_channels	Tuple of block output channels. TYPE: `Tuple[int]`, *optional*, defaults to `(224, 448, 672, 896)` DEFAULT: (224, 448, 672, 896)
layers_per_block	The number of layers per block. TYPE: `int`, *optional*, defaults to `2` DEFAULT: 2
mid_block_scale_factor	The scale factor for the mid block. TYPE: `float`, *optional*, defaults to `1` DEFAULT: 1
downsample_padding	The padding for the downsample convolution. TYPE: `int`, *optional*, defaults to `1` DEFAULT: 1
downsample_type	The downsample type for downsampling layers. Choose between "conv" and "resnet" TYPE: `str`, *optional*, defaults to `conv` DEFAULT: 'conv'
upsample_type	The upsample type for upsampling layers. Choose between "conv" and "resnet" TYPE: `str`, *optional*, defaults to `conv` DEFAULT: 'conv'
dropout	The dropout probability to use. TYPE: `float`, *optional*, defaults to 0.0 DEFAULT: 0.0
act_fn	The activation function to use. TYPE: `str`, *optional*, defaults to `"silu"` DEFAULT: 'silu'
attention_head_dim	The attention head dimension. TYPE: `int`, *optional*, defaults to `8` DEFAULT: 8
norm_num_groups	The number of groups for normalization. TYPE: `int`, *optional*, defaults to `32` DEFAULT: 32
attn_norm_num_groups	If set to an integer, a group norm layer will be created in the mid block's [Attention] layer with the given number of groups. If left as None, the group norm layer will only be created if resnet_time_scale_shift is set to default, and if created will have norm_num_groups groups. TYPE: `int`, *optional*, defaults to `None` DEFAULT: None
norm_eps	The epsilon for normalization. TYPE: `float`, *optional*, defaults to `1e-5` DEFAULT: 1e-05
resnet_time_scale_shift	Time scale shift config for ResNet blocks (see [~models.resnet.ResnetBlock2D]). Choose from default or scale_shift. TYPE: `str`, *optional*, defaults to `"default"` DEFAULT: 'default'
class_embed_type	The type of class embedding to use which is ultimately summed with the time embeddings. Choose from None, "timestep", or "identity". TYPE: `str`, *optional*, defaults to `None` DEFAULT: None
num_class_embeds	Input dimension of the learnable embedding matrix to be projected to time_embed_dim when performing class conditioning with class_embed_type equal to None. TYPE: `int`, *optional*, defaults to `None` DEFAULT: None

Source code in mindone/diffusers/models/unets/unet_2d.py

class UNet2DModel(ModelMixin, ConfigMixin):
    r"""
    A 2D UNet model that takes a noisy sample and a timestep and returns a sample shaped output.

    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
    for all models (such as downloading or saving).

    Parameters:
        sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
            Height and width of input/output sample. Dimensions must be a multiple of `2 ** (len(block_out_channels) -
            1)`.
        in_channels (`int`, *optional*, defaults to 3): Number of channels in the input sample.
        out_channels (`int`, *optional*, defaults to 3): Number of channels in the output.
        center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.
        time_embedding_type (`str`, *optional*, defaults to `"positional"`): Type of time embedding to use.
        freq_shift (`int`, *optional*, defaults to 0): Frequency shift for Fourier time embedding.
        flip_sin_to_cos (`bool`, *optional*, defaults to `True`):
            Whether to flip sin to cos for Fourier time embedding.
        down_block_types (`Tuple[str]`, *optional*, defaults to `("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")`):
            Tuple of downsample block types.
        mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2D"`):
            Block type for middle of UNet, it can be either `UNetMidBlock2D` or `UnCLIPUNetMidBlock2D`.
        up_block_types (`Tuple[str]`, *optional*, defaults to `("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")`):
            Tuple of upsample block types.
        block_out_channels (`Tuple[int]`, *optional*, defaults to `(224, 448, 672, 896)`):
            Tuple of block output channels.
        layers_per_block (`int`, *optional*, defaults to `2`): The number of layers per block.
        mid_block_scale_factor (`float`, *optional*, defaults to `1`): The scale factor for the mid block.
        downsample_padding (`int`, *optional*, defaults to `1`): The padding for the downsample convolution.
        downsample_type (`str`, *optional*, defaults to `conv`):
            The downsample type for downsampling layers. Choose between "conv" and "resnet"
        upsample_type (`str`, *optional*, defaults to `conv`):
            The upsample type for upsampling layers. Choose between "conv" and "resnet"
        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
        attention_head_dim (`int`, *optional*, defaults to `8`): The attention head dimension.
        norm_num_groups (`int`, *optional*, defaults to `32`): The number of groups for normalization.
        attn_norm_num_groups (`int`, *optional*, defaults to `None`):
            If set to an integer, a group norm layer will be created in the mid block's [`Attention`] layer with the
            given number of groups. If left as `None`, the group norm layer will only be created if
            `resnet_time_scale_shift` is set to `default`, and if created will have `norm_num_groups` groups.
        norm_eps (`float`, *optional*, defaults to `1e-5`): The epsilon for normalization.
        resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
            for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`.
        class_embed_type (`str`, *optional*, defaults to `None`):
            The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,
            `"timestep"`, or `"identity"`.
        num_class_embeds (`int`, *optional*, defaults to `None`):
            Input dimension of the learnable embedding matrix to be projected to `time_embed_dim` when performing class
            conditioning with `class_embed_type` equal to `None`.
    """

    @register_to_config
    def __init__(
        self,
        sample_size: Optional[Union[int, Tuple[int, int]]] = None,
        in_channels: int = 3,
        out_channels: int = 3,
        center_input_sample: bool = False,
        time_embedding_type: str = "positional",
        freq_shift: int = 0,
        flip_sin_to_cos: bool = True,
        down_block_types: Tuple[str, ...] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"),
        up_block_types: Tuple[str, ...] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D"),
        block_out_channels: Tuple[int, ...] = (224, 448, 672, 896),
        layers_per_block: int = 2,
        mid_block_scale_factor: float = 1,
        downsample_padding: int = 1,
        downsample_type: str = "conv",
        upsample_type: str = "conv",
        dropout: float = 0.0,
        act_fn: str = "silu",
        attention_head_dim: Optional[int] = 8,
        norm_num_groups: int = 32,
        attn_norm_num_groups: Optional[int] = None,
        norm_eps: float = 1e-5,
        resnet_time_scale_shift: str = "default",
        add_attention: bool = True,
        class_embed_type: Optional[str] = None,
        num_class_embeds: Optional[int] = None,
        num_train_timesteps: Optional[int] = None,
    ):
        super().__init__()

        self.sample_size = sample_size
        time_embed_dim = block_out_channels[0] * 4

        # Check inputs
        if len(down_block_types) != len(up_block_types):
            raise ValueError(
                f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."  # noqa: E501
            )

        if len(block_out_channels) != len(down_block_types):
            raise ValueError(
                f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."  # noqa: E501
            )

        # input
        self.conv_in = nn.Conv2d(
            in_channels, block_out_channels[0], kernel_size=3, pad_mode="pad", padding=1, has_bias=True
        )

        # time
        if time_embedding_type == "fourier":
            self.time_proj = GaussianFourierProjection(embedding_size=block_out_channels[0], scale=16)
            timestep_input_dim = 2 * block_out_channels[0]
        elif time_embedding_type == "positional":
            self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
            timestep_input_dim = block_out_channels[0]
        elif time_embedding_type == "learned":
            self.time_proj = nn.Embedding(num_train_timesteps, block_out_channels[0])
            timestep_input_dim = block_out_channels[0]

        self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)

        # class embedding
        if class_embed_type is None and num_class_embeds is not None:
            self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
        elif class_embed_type == "timestep":
            self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
        elif class_embed_type == "identity":
            self.class_embedding = nn.Identity()
        else:
            self.class_embedding = None

        # down
        down_blocks = []
        output_channel = block_out_channels[0]
        for i, down_block_type in enumerate(down_block_types):
            input_channel = output_channel
            output_channel = block_out_channels[i]
            is_final_block = i == len(block_out_channels) - 1

            down_block = get_down_block(
                down_block_type,
                num_layers=layers_per_block,
                in_channels=input_channel,
                out_channels=output_channel,
                temb_channels=time_embed_dim,
                add_downsample=not is_final_block,
                resnet_eps=norm_eps,
                resnet_act_fn=act_fn,
                resnet_groups=norm_num_groups,
                attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel,
                downsample_padding=downsample_padding,
                resnet_time_scale_shift=resnet_time_scale_shift,
                downsample_type=downsample_type,
                dropout=dropout,
            )
            down_blocks.append(down_block)
        self.down_blocks = nn.CellList(down_blocks)

        # mid
        self.mid_block = UNetMidBlock2D(
            in_channels=block_out_channels[-1],
            temb_channels=time_embed_dim,
            dropout=dropout,
            resnet_eps=norm_eps,
            resnet_act_fn=act_fn,
            output_scale_factor=mid_block_scale_factor,
            resnet_time_scale_shift=resnet_time_scale_shift,
            attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1],
            resnet_groups=norm_num_groups,
            attn_groups=attn_norm_num_groups,
            add_attention=add_attention,
        )

        # up
        up_blocks = []
        layers_per_resnet_in_up_blocks = []
        reversed_block_out_channels = list(reversed(block_out_channels))
        output_channel = reversed_block_out_channels[0]
        for i, up_block_type in enumerate(up_block_types):
            prev_output_channel = output_channel
            output_channel = reversed_block_out_channels[i]
            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]

            is_final_block = i == len(block_out_channels) - 1

            up_block = get_up_block(
                up_block_type,
                num_layers=layers_per_block + 1,
                in_channels=input_channel,
                out_channels=output_channel,
                prev_output_channel=prev_output_channel,
                temb_channels=time_embed_dim,
                add_upsample=not is_final_block,
                resnet_eps=norm_eps,
                resnet_act_fn=act_fn,
                resnet_groups=norm_num_groups,
                attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel,
                resnet_time_scale_shift=resnet_time_scale_shift,
                upsample_type=upsample_type,
                dropout=dropout,
            )
            up_blocks.append(up_block)
            prev_output_channel = output_channel
            layers_per_resnet_in_up_blocks.append(len(up_block.resnets))
        self.up_blocks = nn.CellList(up_blocks)
        self.layers_per_resnet_in_up_blocks = layers_per_resnet_in_up_blocks

        # out
        num_groups_out = norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4, 32)
        self.conv_norm_out = GroupNorm(num_channels=block_out_channels[0], num_groups=num_groups_out, eps=norm_eps)
        self.conv_act = get_activation(act_fn)()
        self.conv_out = nn.Conv2d(
            block_out_channels[0], out_channels, kernel_size=3, pad_mode="pad", padding=1, has_bias=True
        )

        self.center_input_sample = self.config.center_input_sample
        self.class_embed_type = self.config.class_embed_type

    def construct(
        self,
        sample: ms.Tensor,
        timestep: Union[ms.Tensor, float, int],
        class_labels: Optional[ms.Tensor] = None,
        return_dict: bool = False,
    ) -> Union[UNet2DOutput, Tuple]:
        r"""
        The [`UNet2DModel`] forward method.

        Args:
            sample (`ms.Tensor`):
                The noisy input tensor with the following shape `(batch, channel, height, width)`.
            timestep (`ms.Tensor` or `float` or `int`): The number of timesteps to denoise an input.
            class_labels (`ms.Tensor`, *optional*, defaults to `None`):
                Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
            return_dict (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~models.unets.unet_2d.UNet2DOutput`] instead of a plain tuple.

        Returns:
            [`~models.unets.unet_2d.UNet2DOutput`] or `tuple`:
                If `return_dict` is True, an [`~models.unets.unet_2d.UNet2DOutput`] is returned, otherwise a `tuple` is
                returned where the first element is the sample tensor.
        """
        # 0. center input if necessary
        if self.center_input_sample:
            sample = 2 * sample - 1.0

        # 1. time
        timesteps = timestep
        if not ops.is_tensor(timesteps):
            timesteps = ms.Tensor([timesteps], dtype=ms.int64)
        elif ops.is_tensor(timesteps) and len(timesteps.shape) == 0:
            timesteps = timesteps[None]

        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
        if timesteps.shape[0] == 1:
            timesteps = timesteps.tile((sample.shape[0],))

        t_emb = self.time_proj(timesteps)

        # timesteps does not contain any weights and will always return f32 tensors
        # but time_embedding might actually be running in fp16. so we need to cast here.
        # there might be better ways to encapsulate this.
        t_emb = t_emb.to(dtype=self.dtype)
        emb = self.time_embedding(t_emb)

        if self.class_embedding is not None:
            if class_labels is None:
                raise ValueError("class_labels should be provided when doing class conditioning")

            if self.class_embed_type == "timestep":
                class_labels = self.time_proj(class_labels)

                # `Timesteps` does not contain any weights and will always return f32 tensors
                # there might be better ways to encapsulate this.
                class_labels = class_labels.to(dtype=sample.dtype)

            class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
            emb = emb + class_emb
        elif self.class_embedding is None and class_labels is not None:
            raise ValueError("class_embedding needs to be initialized in order to use class conditioning")

        # 2. pre-process
        skip_sample = sample
        sample = self.conv_in(sample)

        # 3. down
        down_block_res_samples = (sample,)
        for downsample_block in self.down_blocks:
            if hasattr(downsample_block, "skip_conv"):
                sample, res_samples, skip_sample = downsample_block(
                    hidden_states=sample, temb=emb, skip_sample=skip_sample
                )
            else:
                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)

            down_block_res_samples += res_samples

        # 4. mid
        sample = self.mid_block(sample, emb)

        # 5. up
        skip_sample = None
        for i, upsample_block in enumerate(self.up_blocks):
            res_samples = down_block_res_samples[-self.layers_per_resnet_in_up_blocks[i] :]
            down_block_res_samples = down_block_res_samples[: -self.layers_per_resnet_in_up_blocks[i]]

            if hasattr(upsample_block, "skip_conv"):
                sample, skip_sample = upsample_block(sample, res_samples, emb, skip_sample)
            else:
                sample = upsample_block(sample, res_samples, emb)

        # 6. post-process
        sample = self.conv_norm_out(sample)
        sample = self.conv_act(sample)
        sample = self.conv_out(sample)

        if skip_sample is not None:
            sample += skip_sample

        if self.config["time_embedding_type"] == "fourier":
            timesteps = timesteps.reshape((sample.shape[0],) + (1,) * len(sample.shape[1:]))
            sample = sample / timesteps

        if not return_dict:
            return (sample,)

        return UNet2DOutput(sample=sample)

mindone.diffusers.UNet2DModel.construct(sample, timestep, class_labels=None, return_dict=False)

The [UNet2DModel] forward method.

PARAMETER	DESCRIPTION
sample	The noisy input tensor with the following shape (batch, channel, height, width). TYPE: `ms.Tensor`
timestep	The number of timesteps to denoise an input. TYPE: `ms.Tensor` or `float` or `int`
class_labels	Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. TYPE: `ms.Tensor`, *optional*, defaults to `None` DEFAULT: None
return_dict	Whether or not to return a [~models.unets.unet_2d.UNet2DOutput] instead of a plain tuple. TYPE: `bool`, *optional*, defaults to `False` DEFAULT: False

RETURNS	DESCRIPTION
Union[UNet2DOutput, Tuple]	[~models.unets.unet_2d.UNet2DOutput] or tuple: If return_dict is True, an [~models.unets.unet_2d.UNet2DOutput] is returned, otherwise a tuple is returned where the first element is the sample tensor.

Source code in mindone/diffusers/models/unets/unet_2d.py

def construct(
    self,
    sample: ms.Tensor,
    timestep: Union[ms.Tensor, float, int],
    class_labels: Optional[ms.Tensor] = None,
    return_dict: bool = False,
) -> Union[UNet2DOutput, Tuple]:
    r"""
    The [`UNet2DModel`] forward method.

    Args:
        sample (`ms.Tensor`):
            The noisy input tensor with the following shape `(batch, channel, height, width)`.
        timestep (`ms.Tensor` or `float` or `int`): The number of timesteps to denoise an input.
        class_labels (`ms.Tensor`, *optional*, defaults to `None`):
            Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
        return_dict (`bool`, *optional*, defaults to `False`):
            Whether or not to return a [`~models.unets.unet_2d.UNet2DOutput`] instead of a plain tuple.

    Returns:
        [`~models.unets.unet_2d.UNet2DOutput`] or `tuple`:
            If `return_dict` is True, an [`~models.unets.unet_2d.UNet2DOutput`] is returned, otherwise a `tuple` is
            returned where the first element is the sample tensor.
    """
    # 0. center input if necessary
    if self.center_input_sample:
        sample = 2 * sample - 1.0

    # 1. time
    timesteps = timestep
    if not ops.is_tensor(timesteps):
        timesteps = ms.Tensor([timesteps], dtype=ms.int64)
    elif ops.is_tensor(timesteps) and len(timesteps.shape) == 0:
        timesteps = timesteps[None]

    # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
    if timesteps.shape[0] == 1:
        timesteps = timesteps.tile((sample.shape[0],))

    t_emb = self.time_proj(timesteps)

    # timesteps does not contain any weights and will always return f32 tensors
    # but time_embedding might actually be running in fp16. so we need to cast here.
    # there might be better ways to encapsulate this.
    t_emb = t_emb.to(dtype=self.dtype)
    emb = self.time_embedding(t_emb)

    if self.class_embedding is not None:
        if class_labels is None:
            raise ValueError("class_labels should be provided when doing class conditioning")

        if self.class_embed_type == "timestep":
            class_labels = self.time_proj(class_labels)

            # `Timesteps` does not contain any weights and will always return f32 tensors
            # there might be better ways to encapsulate this.
            class_labels = class_labels.to(dtype=sample.dtype)

        class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
        emb = emb + class_emb
    elif self.class_embedding is None and class_labels is not None:
        raise ValueError("class_embedding needs to be initialized in order to use class conditioning")

    # 2. pre-process
    skip_sample = sample
    sample = self.conv_in(sample)

    # 3. down
    down_block_res_samples = (sample,)
    for downsample_block in self.down_blocks:
        if hasattr(downsample_block, "skip_conv"):
            sample, res_samples, skip_sample = downsample_block(
                hidden_states=sample, temb=emb, skip_sample=skip_sample
            )
        else:
            sample, res_samples = downsample_block(hidden_states=sample, temb=emb)

        down_block_res_samples += res_samples

    # 4. mid
    sample = self.mid_block(sample, emb)

    # 5. up
    skip_sample = None
    for i, upsample_block in enumerate(self.up_blocks):
        res_samples = down_block_res_samples[-self.layers_per_resnet_in_up_blocks[i] :]
        down_block_res_samples = down_block_res_samples[: -self.layers_per_resnet_in_up_blocks[i]]

        if hasattr(upsample_block, "skip_conv"):
            sample, skip_sample = upsample_block(sample, res_samples, emb, skip_sample)
        else:
            sample = upsample_block(sample, res_samples, emb)

    # 6. post-process
    sample = self.conv_norm_out(sample)
    sample = self.conv_act(sample)
    sample = self.conv_out(sample)

    if skip_sample is not None:
        sample += skip_sample

    if self.config["time_embedding_type"] == "fourier":
        timesteps = timesteps.reshape((sample.shape[0],) + (1,) * len(sample.shape[1:]))
        sample = sample / timesteps

    if not return_dict:
        return (sample,)

    return UNet2DOutput(sample=sample)

mindone.diffusers.models.unets.unet_2d.UNet2DOutput dataclass

Bases: BaseOutput

The output of [UNet2DModel].

PARAMETER	DESCRIPTION
sample	The hidden states output from the last layer of the model. TYPE: `ms.Tensor` of shape `(batch_size, num_channels, height, width)`

Source code in mindone/diffusers/models/unets/unet_2d.py

@dataclass
class UNet2DOutput(BaseOutput):
    """
    The output of [`UNet2DModel`].

    Args:
        sample (`ms.Tensor` of shape `(batch_size, num_channels, height, width)`):
            The hidden states output from the last layer of the model.
    """

    sample: ms.Tensor