From 19b181e6280fa3fcea49280deeec1012e5bbff29 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sun, 26 Feb 2023 16:17:34 +0530 Subject: [PATCH] adding: print statements to all ff_attn_fn affected blocks. --- src/diffusers/models/attention.py | 1 + src/diffusers/models/transformer_2d.py | 1 + src/diffusers/models/unet_2d_blocks.py | 1 + src/diffusers/models/unet_2d_condition.py | 1 + 4 files changed, 4 insertions(+) diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py index 3b86f9e2b6..00460a0b39 100644 --- a/src/diffusers/models/attention.py +++ b/src/diffusers/models/attention.py @@ -211,6 +211,7 @@ class BasicTransformerBlock(nn.Module): final_dropout: bool = False, ): super().__init__() + print(f"Using {activation_fn} as activation_fn in BasicTransformerBlock.") self.only_cross_attention = only_cross_attention self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" diff --git a/src/diffusers/models/transformer_2d.py b/src/diffusers/models/transformer_2d.py index 2fdf70ce76..7fb7689dc0 100644 --- a/src/diffusers/models/transformer_2d.py +++ b/src/diffusers/models/transformer_2d.py @@ -104,6 +104,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin): self.num_attention_heads = num_attention_heads self.attention_head_dim = attention_head_dim inner_dim = num_attention_heads * attention_head_dim + print(f"Using {activation_fn} as activation_fn in Transformer2DModel.") # 1. Transformer2DModel can process both standard continous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)` # Define whether input is continuous or discrete depending on configuration diff --git a/src/diffusers/models/unet_2d_blocks.py b/src/diffusers/models/unet_2d_blocks.py index 230d24274a..5388f644f2 100644 --- a/src/diffusers/models/unet_2d_blocks.py +++ b/src/diffusers/models/unet_2d_blocks.py @@ -468,6 +468,7 @@ class UNetMidBlock2DCrossAttn(nn.Module): ff_activation_fn="geglu", ): super().__init__() + print(f"Using {ff_activation_fn} as ff_activation_fn in UNetMidBlock2DCrossAttn") self.has_cross_attention = True self.attn_num_head_channels = attn_num_head_channels diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py index 76a3a64530..27fb5d75d0 100644 --- a/src/diffusers/models/unet_2d_condition.py +++ b/src/diffusers/models/unet_2d_condition.py @@ -151,6 +151,7 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin) ff_activation_fn="geglu", ): super().__init__() + print(f"Using {ff_activation_fn} as ff_activation_fn in UNet2DConditionModel") self.sample_size = sample_size