Update unet/conditional_unet_model.py

unet/conditional_unet_model.py

@@ -25,12 +25,13 @@ class UNet2DOutput(BaseOutput):
 
     sample: torch.FloatTensor
 
-
 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) -
@@ -105,7 +106,7 @@ class UNet2DModel(ModelMixin, ConfigMixin):
         class_embed_type: Optional[str] = None,
         num_class_embeds: Optional[int] = None,
         num_train_timesteps: Optional[int] = None,
-        set_W_to_weight: Optional[bool] = True,
+        set_W_to_weight: Optional[bool] = True
     ):
         super().__init__()
 
@@ -237,6 +238,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)`.
@@ -245,6 +247,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