Update unet/conditional_unet_model.py

unet/conditional_unet_model.py

@@ -17,6 +17,7 @@ from diffusers.models.unets.unet_2d_blocks import UNetMidBlock2D, get_down_block
 class UNet2DOutput(BaseOutput):
     """
     The output of [`UNet2DModel`].
+
     Args:
         sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
             The hidden states output from the last layer of the model.
@@ -28,8 +29,10 @@ class UNet2DOutput(BaseOutput):
 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) -
@@ -236,6 +239,7 @@ class UNet2DModel(ModelMixin, ConfigMixin):
     ) -> Union[UNet2DOutput, Tuple]:
         r"""
         The [`UNet2DModel`] forward method.
+
         Args:
             sample (`torch.FloatTensor`):
                 The noisy input tensor with the following shape `(batch, channel, height, width)`.
@@ -244,6 +248,7 @@ class UNet2DModel(ModelMixin, ConfigMixin):
                 Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`~models.unet_2d.UNet2DOutput`] instead of a plain tuple.
+
         Returns:
             [`~models.unet_2d.UNet2DOutput`] or `tuple`:
                 If `return_dict` is True, an [`~models.unet_2d.UNet2DOutput`] is returned, otherwise a `tuple` is
@@ -433,8 +438,10 @@ class ClassConditionedUnetForShapes3D(ModelMixin, ConfigMixin):
 class MultiLabelConditionalUNet2DModelForShapes3D(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) -
@@ -661,6 +668,7 @@ class MultiLabelConditionalUNet2DModelForShapes3D(ModelMixin, ConfigMixin):
     ) -> Union[UNet2DOutput, Tuple]:
         r"""
         The [`UNet2DModel`] forward method.
+
         Args:
             sample (`torch.FloatTensor`):
                 The noisy input tensor with the following shape `(batch, channel, height, width)`.
@@ -669,6 +677,7 @@ class MultiLabelConditionalUNet2DModelForShapes3D(ModelMixin, ConfigMixin):
                 Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`~models.unet_2d.UNet2DOutput`] instead of a plain tuple.
+
         Returns:
             [`~models.unet_2d.UNet2DOutput`] or `tuple`:
                 If `return_dict` is True, an [`~models.unet_2d.UNet2DOutput`] is returned, otherwise a `tuple` is