Error loading Trained diffusion model while implementing ControlNet

[Aman0phy]

I am implementing control net for the image generation.
I get an error while loading a diffusion model (trained in other script)

Below are relevant parts of error and code
For the error If you scroll to the bottom, there is error related to activation function GEGLU, but this seems to be coming from the MONAI code as I have not provided any activation function.

Please let me know if anyone has ideas about the error
Relevant part of Code:
diffusion_model = DiffusionModelUNet( spatial_dims=2, in_channels=1, out_channels=1, num_res_blocks=2, num_channels=[128, 256, 512], attention_levels=[False, True, True], with_conditioning=True, cross_attention_dim=256, num_head_channels=[0,256,512], )
diffusion_model.load_state_dict(torch.load('name.pt'))

Error:

ValueError Traceback (most recent call last)
/tmp/ipykernel_2518302/380197457.py in
9 with_conditioning=True,
10 cross_attention_dim=256,
---> 11 num_head_channels=[0,256,512],
12 )
13

~/GenerativeModels/generative/networks/nets/diffusion_model_unet.py in init(self, spatial_dims, in_channels, out_channels, num_res_blocks, num_channels, attention_levels, norm_num_groups, norm_eps, resblock_updown, num_head_channels, with_conditioning, transformer_num_layers, cross_attention_dim, num_class_embeds, upcast_attention, use_flash_attention)
1773 cross_attention_dim=cross_attention_dim,
1774 upcast_attention=upcast_attention,
-> 1775 use_flash_attention=use_flash_attention,
1776 )
1777

~/GenerativeModels/generative/networks/nets/diffusion_model_unet.py in get_down_block(spatial_dims, in_channels, out_channels, temb_channels, num_res_blocks, norm_num_groups, norm_eps, add_downsample, resblock_updown, with_attn, with_cross_attn, num_head_channels, transformer_num_layers, cross_attention_dim, upcast_attention, use_flash_attention)
1509 cross_attention_dim=cross_attention_dim,
1510 upcast_attention=upcast_attention,
-> 1511 use_flash_attention=use_flash_attention,
1512 )
1513 else:

~/GenerativeModels/generative/networks/nets/diffusion_model_unet.py in init(self, spatial_dims, in_channels, out_channels, temb_channels, num_res_blocks, norm_num_groups, norm_eps, add_downsample, resblock_updown, downsample_padding, num_head_channels, transformer_num_layers, cross_attention_dim, upcast_attention, use_flash_attention)
962 cross_attention_dim=cross_attention_dim,
963 upcast_attention=upcast_attention,
--> 964 use_flash_attention=use_flash_attention,
965 )
966 )

~/GenerativeModels/generative/networks/nets/diffusion_model_unet.py in init(self, spatial_dims, in_channels, num_attention_heads, num_head_channels, num_layers, dropout, norm_num_groups, norm_eps, cross_attention_dim, upcast_attention, use_flash_attention)
296 use_flash_attention=use_flash_attention,
297 )
--> 298 for _ in range(num_layers)
299 ]
300 )

~/GenerativeModels/generative/networks/nets/diffusion_model_unet.py in (.0)
296 use_flash_attention=use_flash_attention,
297 )
--> 298 for _ in range(num_layers)
299 ]
300 )

~/GenerativeModels/generative/networks/nets/diffusion_model_unet.py in init(self, num_channels, num_attention_heads, num_head_channels, dropout, cross_attention_dim, upcast_attention, use_flash_attention)
209 use_flash_attention=use_flash_attention,
210 ) # is a self-attention
--> 211 self.ff = MLPBlock(hidden_size=num_channels, mlp_dim=num_channels * 4, act="GEGLU", dropout_rate=dropout)
212 self.attn2 = CrossAttention(
213 query_dim=num_channels,

~/anaconda3/envs/monai/lib/python3.7/site-packages/monai/networks/blocks/mlp.py in init(self, hidden_size, mlp_dim, dropout_rate, act, dropout_mode)
56 self.linear1 = nn.Linear(hidden_size, mlp_dim)
57 self.linear2 = nn.Linear(mlp_dim, hidden_size)
---> 58 self.fn = get_act_layer(act)
59 self.drop1 = nn.Dropout(dropout_rate)
60 dropout_opt = look_up_option(dropout_mode, SUPPORTED_DROPOUT_MODE)

~/anaconda3/envs/monai/lib/python3.7/site-packages/monai/networks/layers/utils.py in get_act_layer(name)
70 return torch.nn.Identity()
71 act_name, act_args = split_args(name)
---> 72 act_type = Act[act_name]
73 return act_type(**act_args)
74

~/anaconda3/envs/monai/lib/python3.7/site-packages/monai/networks/layers/factories.py in getitem(self, args)
145 name_obj, *args = args
146
--> 147 return self.get_constructor(name_obj, *args)
148
149 def getattr(self, key):

~/anaconda3/envs/monai/lib/python3.7/site-packages/monai/networks/layers/factories.py in get_constructor(self, factory_name, *args)
126 raise TypeError(f"factory_name must a str but is {type(factory_name).name}.")
127
--> 128 func = look_up_option(factory_name.upper(), self.factories)
129 return func(*args)
130

~/anaconda3/envs/monai/lib/python3.7/site-packages/monai/utils/module.py in look_up_option(opt_str, supported, default, print_all_options)
129 f"By '{opt_str}', did you mean '{guess_at_spelling}'?\n"
130 + f"'{opt_str}' is not a valid value.\n"
--> 131 + supported_msg
132 )
133 raise ValueError(f"Unsupported option '{opt_str}', " + supported_msg)

ValueError: By 'GEGLU', did you mean 'GELU'?
'GEGLU' is not a valid value.
Available options are {'TANH', 'LOGSOFTMAX', 'RELU6', 'SWISH', 'MEMSWISH', 'GELU', 'SOFTMAX', 'SIGMOID', 'RELU', 'ELU', 'SELU', 'MISH', 'PRELU', 'CELU', 'LEAKYRELU'}.

[KumoLiu]

Hi @Aman0phy, what MONAI version did you use? I think we have GEGLU.

from monai.networks.layers.factories import Act

Act.components.keys()

#dict_keys(['ELU', 'RELU', 'LEAKYRELU', 'PRELU', 'RELU6', 'SELU', 'CELU', 'GELU', 'SIGMOID', 'TANH', 'SOFTMAX', 'LOGSOFTMAX', 'SWISH', 'MEMSWISH', 'MISH', 'GEGLU'])

[Aman0phy]

I am using MONAI version: 1.0.1

[KumoLiu]

So that makes sense, we introduced the GEGLU in 1.2. Could you please upgrade the MONAI version?

MONAI/monai/networks/layers/factories.py

Line 312 in c33f1ba

@Act.factory_function("geglu")

[Aman0phy]

Thanks It worked with the update