From 8425cd48774e41c95fa0a2378c3033916aeccbbb Mon Sep 17 00:00:00 2001 From: hi-sushanta Date: Thu, 5 Oct 2023 06:34:18 +0530 Subject: [PATCH 1/7] I added a new doc string to the class. This is more flexible to understanding other developers what are doing and where it's using. --- src/diffusers/models/unet_2d_blocks.py | 36 ++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/src/diffusers/models/unet_2d_blocks.py b/src/diffusers/models/unet_2d_blocks.py index 8aebb3aad615..31a6a5241e09 100644 --- a/src/diffusers/models/unet_2d_blocks.py +++ b/src/diffusers/models/unet_2d_blocks.py @@ -453,6 +453,42 @@ def get_up_block( class AutoencoderTinyBlock(nn.Module): + """ + It's tiny autoencoder block. The block consists of three + convolutional layers with ReLU activation, followed by a + skip connection and a final ReLU activation. + + The skip connection helps to preserve the spatial information of the input signal, which + can be beneficial for autoencoders that are used for image compression or denoising. + This block is designed to be used in autoencoder architectures that are lightweight and + fast to train. + + Parameters: + * `in_channels:` The number of input channels. + * `out_channels:` The number of output channels. + * `act_fn:` The activation function to use. Supported values are `relu`, `tanh`, and `sigmoid`. + + Example: + ```python + >>> # Define an autoencoder architecture using the AutoencoderTinyBlock. + >>> model = nn.Sequential( + >>> AutoencoderTinyBlock(3, 64, "relu"), + >>> AutoencoderTinyBlock(64, 128, "relu"), + >>> AutoencoderTinyBlock(128, 256, "relu"), + >>> ) + + >>> # Train the model on a dataset of images. + >>> # ... + + >>> # Use the model to reconstruct an image from a noisy input. + >>> noisy_image = ... + >>> reconstructed_image = model(noisy_image) + ``` + + Output: + A tensor with the same shape as the input tensor, but with the number of channels equal to `out_channels`. + """ + def __init__(self, in_channels: int, out_channels: int, act_fn: str): super().__init__() act_fn = get_activation(act_fn) From 746a8e8e51efef6c73c34153c0585391ad000391 Mon Sep 17 00:00:00 2001 From: Chi Date: Sat, 7 Oct 2023 05:54:19 +0530 Subject: [PATCH 2/7] Update src/diffusers/models/unet_2d_blocks.py This changes suggest by maintener. Co-authored-by: Sayak Paul --- src/diffusers/models/unet_2d_blocks.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/src/diffusers/models/unet_2d_blocks.py b/src/diffusers/models/unet_2d_blocks.py index 6cd11c29ae11..5a109e338315 100644 --- a/src/diffusers/models/unet_2d_blocks.py +++ b/src/diffusers/models/unet_2d_blocks.py @@ -481,22 +481,6 @@ class AutoencoderTinyBlock(nn.Module): * `out_channels:` The number of output channels. * `act_fn:` The activation function to use. Supported values are `relu`, `tanh`, and `sigmoid`. - Example: - ```python - >>> # Define an autoencoder architecture using the AutoencoderTinyBlock. - >>> model = nn.Sequential( - >>> AutoencoderTinyBlock(3, 64, "relu"), - >>> AutoencoderTinyBlock(64, 128, "relu"), - >>> AutoencoderTinyBlock(128, 256, "relu"), - >>> ) - - >>> # Train the model on a dataset of images. - >>> # ... - - >>> # Use the model to reconstruct an image from a noisy input. - >>> noisy_image = ... - >>> reconstructed_image = model(noisy_image) - ``` Output: A tensor with the same shape as the input tensor, but with the number of channels equal to `out_channels`. From 6e5688647b0ba681f4b47e1a96fabe6ef733ff8e Mon Sep 17 00:00:00 2001 From: Chi Date: Sat, 7 Oct 2023 17:16:49 +0530 Subject: [PATCH 3/7] Update src/diffusers/models/unet_2d_blocks.py Add suggested text Co-authored-by: Sayak Paul --- src/diffusers/models/unet_2d_blocks.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/src/diffusers/models/unet_2d_blocks.py b/src/diffusers/models/unet_2d_blocks.py index 5a109e338315..f4ef776d1c8a 100644 --- a/src/diffusers/models/unet_2d_blocks.py +++ b/src/diffusers/models/unet_2d_blocks.py @@ -467,14 +467,7 @@ def get_up_block( class AutoencoderTinyBlock(nn.Module): """ - It's tiny autoencoder block. The block consists of three - convolutional layers with ReLU activation, followed by a - skip connection and a final ReLU activation. - - The skip connection helps to preserve the spatial information of the input signal, which - can be beneficial for autoencoders that are used for image compression or denoising. - This block is designed to be used in autoencoder architectures that are lightweight and - fast to train. + Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU blocks. Parameters: * `in_channels:` The number of input channels. From 627fd9fee2c0074c1ce2ca68de7afecc2f7a6b38 Mon Sep 17 00:00:00 2001 From: Chi Date: Sat, 7 Oct 2023 20:44:27 +0530 Subject: [PATCH 4/7] Update unet_2d_blocks.py I changed the Parameter to Args text. --- src/diffusers/models/unet_2d_blocks.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/diffusers/models/unet_2d_blocks.py b/src/diffusers/models/unet_2d_blocks.py index f4ef776d1c8a..b4fc95afd083 100644 --- a/src/diffusers/models/unet_2d_blocks.py +++ b/src/diffusers/models/unet_2d_blocks.py @@ -469,11 +469,10 @@ class AutoencoderTinyBlock(nn.Module): """ Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU blocks. - Parameters: - * `in_channels:` The number of input channels. - * `out_channels:` The number of output channels. - * `act_fn:` The activation function to use. Supported values are `relu`, `tanh`, and `sigmoid`. - + Args: + in_channels (`int`): The number of input channels. + out_channels (`int`): The number of output channels. + act_fn (`str`):` The activation function to use. Supported values are `relu`, `tanh`, and `sigmoid`. Output: A tensor with the same shape as the input tensor, but with the number of channels equal to `out_channels`. From ae4f7f272c902177ad736144dd03132cdc2f1c00 Mon Sep 17 00:00:00 2001 From: Chi Date: Sun, 8 Oct 2023 19:53:14 +0530 Subject: [PATCH 5/7] Update unet_2d_blocks.py proper indentation set in this file. --- src/diffusers/models/unet_2d_blocks.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/diffusers/models/unet_2d_blocks.py b/src/diffusers/models/unet_2d_blocks.py index b4fc95afd083..28adb7ec6d1e 100644 --- a/src/diffusers/models/unet_2d_blocks.py +++ b/src/diffusers/models/unet_2d_blocks.py @@ -467,15 +467,15 @@ def get_up_block( class AutoencoderTinyBlock(nn.Module): """ - Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU blocks. + Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU blocks. - Args: - in_channels (`int`): The number of input channels. - out_channels (`int`): The number of output channels. - act_fn (`str`):` The activation function to use. Supported values are `relu`, `tanh`, and `sigmoid`. + Args: + in_channels (`int`): The number of input channels. + out_channels (`int`): The number of output channels. + act_fn (`str`): The activation function to use. Supported values are `relu`, `tanh`, and `sigmoid`. - Output: - A tensor with the same shape as the input tensor, but with the number of channels equal to `out_channels`. + Returns: + `torch.FloatTensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to `out_channels`. """ def __init__(self, in_channels: int, out_channels: int, act_fn: str): From f0bea430b8104f3a21ec1c7b8849aafd8463b26b Mon Sep 17 00:00:00 2001 From: Chi Date: Sun, 8 Oct 2023 20:03:51 +0530 Subject: [PATCH 6/7] Update unet_2d_blocks.py a little bit of change in the act_fun argument line. --- src/diffusers/models/unet_2d_blocks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/models/unet_2d_blocks.py b/src/diffusers/models/unet_2d_blocks.py index 28adb7ec6d1e..d0342aac93af 100644 --- a/src/diffusers/models/unet_2d_blocks.py +++ b/src/diffusers/models/unet_2d_blocks.py @@ -472,7 +472,7 @@ class AutoencoderTinyBlock(nn.Module): Args: in_channels (`int`): The number of input channels. out_channels (`int`): The number of output channels. - act_fn (`str`): The activation function to use. Supported values are `relu`, `tanh`, and `sigmoid`. + act_fn (`str`):` The activation function to use. Supported values are `"swish"`, `"mish"`, `"gelu"`, and `"relu"`. Returns: `torch.FloatTensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to `out_channels`. From 3546f6d417565d16fe1692d961e9794254dff641 Mon Sep 17 00:00:00 2001 From: hi-sushanta Date: Mon, 9 Oct 2023 19:27:39 +0530 Subject: [PATCH 7/7] I run the black command to reformat style in the code --- src/diffusers/models/unet_2d_blocks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/unet_2d_blocks.py b/src/diffusers/models/unet_2d_blocks.py index d0342aac93af..18b0e41af738 100644 --- a/src/diffusers/models/unet_2d_blocks.py +++ b/src/diffusers/models/unet_2d_blocks.py @@ -467,13 +467,13 @@ def get_up_block( class AutoencoderTinyBlock(nn.Module): """ - Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU blocks. + Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU blocks. Args: in_channels (`int`): The number of input channels. out_channels (`int`): The number of output channels. act_fn (`str`):` The activation function to use. Supported values are `"swish"`, `"mish"`, `"gelu"`, and `"relu"`. - + Returns: `torch.FloatTensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to `out_channels`. """