TIME: 2025-05-09 21:58:45.534412 DEVICE: cuda {'autoencoder_params': {'attn_down': [False, False], 'codebook_size': 20, 'down_channels': [32, 64, 128], 'down_sample': [True, True], 'mid_channels': [128, 128], 'norm_channels': 32, 'num_down_layers': 4, 'num_heads': 16, 'num_mid_layers': 4, 'num_up_layers': 4, 'z_channels': 3}, 'dataset_params': {'im_channels': 3, 'im_size': 128}, 'diffusion_params': {'beta_end': 0.0195, 'beta_start': 0.0015, 'num_timesteps': 1000}, 'ldm_params': {'attn_down': [True, True, True], 'conv_out_channels': 128, 'down_channels': [128, 256, 256, 256], 'down_sample': [False, False, False], 'mid_channels': [256, 256], 'norm_channels': 32, 'num_down_layers': 2, 'num_heads': 16, 'num_mid_layers': 2, 'num_up_layers': 2, 'time_emb_dim': 256}, 'paths': {'images_dir': '/scratch/IITB/ai-at-ieor/23m1521/datasets/Vaani/Images'}, 'train_params': {'autoencoder_acc_steps': 1, 'autoencoder_batch_size': 8, 'autoencoder_epochs': 30, 'autoencoder_img_save_steps': 8, 'autoencoder_lr': 0.0001, 'codebook_weight': 1, 'commitment_beta': 0.2, 'disc_start': 1000, 'disc_weight': 0.5, 'kl_weight': 5e-06, 'ldm_batch_size': 1, 'ldm_ckpt_name': 'ddpm_ckpt.pth', 'ldm_epochs': 10, 'ldm_lr': 1e-05, 'num_grid_rows': 3, 'num_samples': 9, 'perceptual_weight': 1, 'save_latents': True, 'seed': 4422, 'task_name': 'VaaniLDM', 'vqvae_ckpt_name': 'vqvaq_ckpt.pth', 'vqvae_latent_dir_name': 'vqvae_latents'}, 'training': {'_continue_': True}} Files found: 128807 IMAGE SHAPE: torch.Size([3, 128, 128]) BATCH SHAPE: torch.Size([8, 3, 128, 128]) ====================================================================================================================================================== Layer (type (var_name)) Input Shape Output Shape Param # Trainable Param % ====================================================================================================================================================== VQVAE (VQVAE) [8, 3, 128, 128] [8, 3, 128, 128] 60 True 0.00% ├─Conv2d (encoder_conv_in) [8, 3, 128, 128] [8, 32, 128, 128] 896 True 0.01% ├─ModuleList (encoder_layers) -- -- -- True -- │ └─DownBlock (0) [8, 32, 128, 128] [8, 64, 64, 64] -- True -- │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (0) [8, 32, 128, 128] [8, 64, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 32, 128, 128] [8, 32, 128, 128] 64 True 0.00% │ │ │ │ └─SiLU (1) [8, 32, 128, 128] [8, 32, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 32, 128, 128] [8, 64, 128, 128] 18,496 True 0.30% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (0) [8, 64, 128, 128] [8, 64, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 128, 128] [8, 64, 128, 128] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 128, 128] [8, 64, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 128, 128] [8, 64, 128, 128] 36,928 True 0.59% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (0) [8, 32, 128, 128] [8, 64, 128, 128] 2,112 True 0.03% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (1) [8, 64, 128, 128] [8, 64, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 128, 128] [8, 64, 128, 128] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 128, 128] [8, 64, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 128, 128] [8, 64, 128, 128] 36,928 True 0.59% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (1) [8, 64, 128, 128] [8, 64, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 128, 128] [8, 64, 128, 128] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 128, 128] [8, 64, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 128, 128] [8, 64, 128, 128] 36,928 True 0.59% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (1) [8, 64, 128, 128] [8, 64, 128, 128] 4,160 True 0.07% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (2) [8, 64, 128, 128] [8, 64, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 128, 128] [8, 64, 128, 128] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 128, 128] [8, 64, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 128, 128] [8, 64, 128, 128] 36,928 True 0.59% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (2) [8, 64, 128, 128] [8, 64, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 128, 128] [8, 64, 128, 128] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 128, 128] [8, 64, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 128, 128] [8, 64, 128, 128] 36,928 True 0.59% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (2) [8, 64, 128, 128] [8, 64, 128, 128] 4,160 True 0.07% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (3) [8, 64, 128, 128] [8, 64, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 128, 128] [8, 64, 128, 128] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 128, 128] [8, 64, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 128, 128] [8, 64, 128, 128] 36,928 True 0.59% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (3) [8, 64, 128, 128] [8, 64, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 128, 128] [8, 64, 128, 128] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 128, 128] [8, 64, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 128, 128] [8, 64, 128, 128] 36,928 True 0.59% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (3) [8, 64, 128, 128] [8, 64, 128, 128] 4,160 True 0.07% │ │ └─Conv2d (down_sample_conv) [8, 64, 128, 128] [8, 64, 64, 64] 65,600 True 1.05% │ └─DownBlock (1) [8, 64, 64, 64] [8, 128, 32, 32] -- True -- │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (0) [8, 64, 64, 64] [8, 128, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 64, 64] [8, 64, 64, 64] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 64, 64] [8, 64, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 64, 64] [8, 128, 64, 64] 73,856 True 1.19% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (0) [8, 128, 64, 64] [8, 128, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 64, 64] [8, 128, 64, 64] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 64, 64] [8, 128, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 64, 64] [8, 128, 64, 64] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (0) [8, 64, 64, 64] [8, 128, 64, 64] 8,320 True 0.13% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (1) [8, 128, 64, 64] [8, 128, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 64, 64] [8, 128, 64, 64] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 64, 64] [8, 128, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 64, 64] [8, 128, 64, 64] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (1) [8, 128, 64, 64] [8, 128, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 64, 64] [8, 128, 64, 64] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 64, 64] [8, 128, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 64, 64] [8, 128, 64, 64] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (1) [8, 128, 64, 64] [8, 128, 64, 64] 16,512 True 0.27% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (2) [8, 128, 64, 64] [8, 128, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 64, 64] [8, 128, 64, 64] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 64, 64] [8, 128, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 64, 64] [8, 128, 64, 64] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (2) [8, 128, 64, 64] [8, 128, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 64, 64] [8, 128, 64, 64] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 64, 64] [8, 128, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 64, 64] [8, 128, 64, 64] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (2) [8, 128, 64, 64] [8, 128, 64, 64] 16,512 True 0.27% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (3) [8, 128, 64, 64] [8, 128, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 64, 64] [8, 128, 64, 64] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 64, 64] [8, 128, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 64, 64] [8, 128, 64, 64] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (3) [8, 128, 64, 64] [8, 128, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 64, 64] [8, 128, 64, 64] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 64, 64] [8, 128, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 64, 64] [8, 128, 64, 64] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (3) [8, 128, 64, 64] [8, 128, 64, 64] 16,512 True 0.27% │ │ └─Conv2d (down_sample_conv) [8, 128, 64, 64] [8, 128, 32, 32] 262,272 True 4.22% ├─ModuleList (encoder_mids) -- -- -- True -- │ └─MidBlock (0) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (0) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (0) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (0) [8, 128, 32, 32] [8, 128, 32, 32] 16,512 True 0.27% │ │ └─ModuleList (attention_norms) -- -- (recursive) True (recursive) │ │ │ └─GroupNorm (0) [8, 128, 1024] [8, 128, 1024] 256 True 0.00% │ │ └─ModuleList (attentions) -- -- (recursive) True (recursive) │ │ │ └─MultiheadAttention (0) [8, 1024, 128] [8, 1024, 128] 66,048 True 1.06% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (1) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (1) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (1) [8, 128, 32, 32] [8, 128, 32, 32] 16,512 True 0.27% │ │ └─ModuleList (attention_norms) -- -- (recursive) True (recursive) │ │ │ └─GroupNorm (1) [8, 128, 1024] [8, 128, 1024] 256 True 0.00% │ │ └─ModuleList (attentions) -- -- (recursive) True (recursive) │ │ │ └─MultiheadAttention (1) [8, 1024, 128] [8, 1024, 128] 66,048 True 1.06% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (2) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (2) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 16,512 True 0.27% │ │ └─ModuleList (attention_norms) -- -- (recursive) True (recursive) │ │ │ └─GroupNorm (2) [8, 128, 1024] [8, 128, 1024] 256 True 0.00% │ │ └─ModuleList (attentions) -- -- (recursive) True (recursive) │ │ │ └─MultiheadAttention (2) [8, 1024, 128] [8, 1024, 128] 66,048 True 1.06% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (3) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (3) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (3) [8, 128, 32, 32] [8, 128, 32, 32] 16,512 True 0.27% │ │ └─ModuleList (attention_norms) -- -- (recursive) True (recursive) │ │ │ └─GroupNorm (3) [8, 128, 1024] [8, 128, 1024] 256 True 0.00% │ │ └─ModuleList (attentions) -- -- (recursive) True (recursive) │ │ │ └─MultiheadAttention (3) [8, 1024, 128] [8, 1024, 128] 66,048 True 1.06% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (4) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (4) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (4) [8, 128, 32, 32] [8, 128, 32, 32] 16,512 True 0.27% ├─GroupNorm (encoder_norm_out) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% ├─Conv2d (encoder_conv_out) [8, 128, 32, 32] [8, 3, 32, 32] 3,459 True 0.06% ├─Conv2d (pre_quant_conv) [8, 3, 32, 32] [8, 3, 32, 32] 12 True 0.00% ├─Conv2d (post_quant_conv) [8, 3, 32, 32] [8, 3, 32, 32] 12 True 0.00% ├─Conv2d (decoder_conv_in) [8, 3, 32, 32] [8, 128, 32, 32] 3,584 True 0.06% ├─ModuleList (decoder_mids) -- -- -- True -- │ └─MidBlock (0) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (0) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (0) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (0) [8, 128, 32, 32] [8, 128, 32, 32] 16,512 True 0.27% │ │ └─ModuleList (attention_norms) -- -- (recursive) True (recursive) │ │ │ └─GroupNorm (0) [8, 128, 1024] [8, 128, 1024] 256 True 0.00% │ │ └─ModuleList (attentions) -- -- (recursive) True (recursive) │ │ │ └─MultiheadAttention (0) [8, 1024, 128] [8, 1024, 128] 66,048 True 1.06% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (1) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (1) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (1) [8, 128, 32, 32] [8, 128, 32, 32] 16,512 True 0.27% │ │ └─ModuleList (attention_norms) -- -- (recursive) True (recursive) │ │ │ └─GroupNorm (1) [8, 128, 1024] [8, 128, 1024] 256 True 0.00% │ │ └─ModuleList (attentions) -- -- (recursive) True (recursive) │ │ │ └─MultiheadAttention (1) [8, 1024, 128] [8, 1024, 128] 66,048 True 1.06% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (2) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (2) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 16,512 True 0.27% │ │ └─ModuleList (attention_norms) -- -- (recursive) True (recursive) │ │ │ └─GroupNorm (2) [8, 128, 1024] [8, 128, 1024] 256 True 0.00% │ │ └─ModuleList (attentions) -- -- (recursive) True (recursive) │ │ │ └─MultiheadAttention (2) [8, 1024, 128] [8, 1024, 128] 66,048 True 1.06% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (3) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (3) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (3) [8, 128, 32, 32] [8, 128, 32, 32] 16,512 True 0.27% │ │ └─ModuleList (attention_norms) -- -- (recursive) True (recursive) │ │ │ └─GroupNorm (3) [8, 128, 1024] [8, 128, 1024] 256 True 0.00% │ │ └─ModuleList (attentions) -- -- (recursive) True (recursive) │ │ │ └─MultiheadAttention (3) [8, 1024, 128] [8, 1024, 128] 66,048 True 1.06% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (4) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (4) [8, 128, 32, 32] [8, 128, 32, 32] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 32, 32] [8, 128, 32, 32] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 32, 32] [8, 128, 32, 32] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 32, 32] [8, 128, 32, 32] 147,584 True 2.37% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (4) [8, 128, 32, 32] [8, 128, 32, 32] 16,512 True 0.27% ├─ModuleList (decoder_layers) -- -- -- True -- │ └─UpBlock (0) [8, 128, 32, 32] [8, 64, 64, 64] -- True -- │ │ └─ConvTranspose2d (up_sample_conv) [8, 128, 32, 32] [8, 128, 64, 64] 262,272 True 4.22% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (0) [8, 128, 64, 64] [8, 64, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 128, 64, 64] [8, 128, 64, 64] 256 True 0.00% │ │ │ │ └─SiLU (1) [8, 128, 64, 64] [8, 128, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 128, 64, 64] [8, 64, 64, 64] 73,792 True 1.19% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (0) [8, 64, 64, 64] [8, 64, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 64, 64] [8, 64, 64, 64] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 64, 64] [8, 64, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 64, 64] [8, 64, 64, 64] 36,928 True 0.59% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (0) [8, 128, 64, 64] [8, 64, 64, 64] 8,256 True 0.13% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (1) [8, 64, 64, 64] [8, 64, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 64, 64] [8, 64, 64, 64] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 64, 64] [8, 64, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 64, 64] [8, 64, 64, 64] 36,928 True 0.59% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (1) [8, 64, 64, 64] [8, 64, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 64, 64] [8, 64, 64, 64] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 64, 64] [8, 64, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 64, 64] [8, 64, 64, 64] 36,928 True 0.59% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (1) [8, 64, 64, 64] [8, 64, 64, 64] 4,160 True 0.07% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (2) [8, 64, 64, 64] [8, 64, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 64, 64] [8, 64, 64, 64] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 64, 64] [8, 64, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 64, 64] [8, 64, 64, 64] 36,928 True 0.59% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (2) [8, 64, 64, 64] [8, 64, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 64, 64] [8, 64, 64, 64] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 64, 64] [8, 64, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 64, 64] [8, 64, 64, 64] 36,928 True 0.59% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (2) [8, 64, 64, 64] [8, 64, 64, 64] 4,160 True 0.07% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (3) [8, 64, 64, 64] [8, 64, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 64, 64] [8, 64, 64, 64] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 64, 64] [8, 64, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 64, 64] [8, 64, 64, 64] 36,928 True 0.59% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (3) [8, 64, 64, 64] [8, 64, 64, 64] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 64, 64] [8, 64, 64, 64] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 64, 64] [8, 64, 64, 64] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 64, 64] [8, 64, 64, 64] 36,928 True 0.59% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (3) [8, 64, 64, 64] [8, 64, 64, 64] 4,160 True 0.07% │ └─UpBlock (1) [8, 64, 64, 64] [8, 32, 128, 128] -- True -- │ │ └─ConvTranspose2d (up_sample_conv) [8, 64, 64, 64] [8, 64, 128, 128] 65,600 True 1.05% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (0) [8, 64, 128, 128] [8, 32, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 64, 128, 128] [8, 64, 128, 128] 128 True 0.00% │ │ │ │ └─SiLU (1) [8, 64, 128, 128] [8, 64, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 64, 128, 128] [8, 32, 128, 128] 18,464 True 0.30% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (0) [8, 32, 128, 128] [8, 32, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 32, 128, 128] [8, 32, 128, 128] 64 True 0.00% │ │ │ │ └─SiLU (1) [8, 32, 128, 128] [8, 32, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 32, 128, 128] [8, 32, 128, 128] 9,248 True 0.15% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (0) [8, 64, 128, 128] [8, 32, 128, 128] 2,080 True 0.03% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (1) [8, 32, 128, 128] [8, 32, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 32, 128, 128] [8, 32, 128, 128] 64 True 0.00% │ │ │ │ └─SiLU (1) [8, 32, 128, 128] [8, 32, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 32, 128, 128] [8, 32, 128, 128] 9,248 True 0.15% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (1) [8, 32, 128, 128] [8, 32, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 32, 128, 128] [8, 32, 128, 128] 64 True 0.00% │ │ │ │ └─SiLU (1) [8, 32, 128, 128] [8, 32, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 32, 128, 128] [8, 32, 128, 128] 9,248 True 0.15% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (1) [8, 32, 128, 128] [8, 32, 128, 128] 1,056 True 0.02% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (2) [8, 32, 128, 128] [8, 32, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 32, 128, 128] [8, 32, 128, 128] 64 True 0.00% │ │ │ │ └─SiLU (1) [8, 32, 128, 128] [8, 32, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 32, 128, 128] [8, 32, 128, 128] 9,248 True 0.15% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (2) [8, 32, 128, 128] [8, 32, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 32, 128, 128] [8, 32, 128, 128] 64 True 0.00% │ │ │ │ └─SiLU (1) [8, 32, 128, 128] [8, 32, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 32, 128, 128] [8, 32, 128, 128] 9,248 True 0.15% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (2) [8, 32, 128, 128] [8, 32, 128, 128] 1,056 True 0.02% │ │ └─ModuleList (resnet_conv_first) -- -- (recursive) True (recursive) │ │ │ └─Sequential (3) [8, 32, 128, 128] [8, 32, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 32, 128, 128] [8, 32, 128, 128] 64 True 0.00% │ │ │ │ └─SiLU (1) [8, 32, 128, 128] [8, 32, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 32, 128, 128] [8, 32, 128, 128] 9,248 True 0.15% │ │ └─ModuleList (resnet_conv_second) -- -- (recursive) True (recursive) │ │ │ └─Sequential (3) [8, 32, 128, 128] [8, 32, 128, 128] -- True -- │ │ │ │ └─GroupNorm (0) [8, 32, 128, 128] [8, 32, 128, 128] 64 True 0.00% │ │ │ │ └─SiLU (1) [8, 32, 128, 128] [8, 32, 128, 128] -- -- -- │ │ │ │ └─Conv2d (2) [8, 32, 128, 128] [8, 32, 128, 128] 9,248 True 0.15% │ │ └─ModuleList (residual_input_conv) -- -- (recursive) True (recursive) │ │ │ └─Conv2d (3) [8, 32, 128, 128] [8, 32, 128, 128] 1,056 True 0.02% ├─GroupNorm (decoder_norm_out) [8, 32, 128, 128] [8, 32, 128, 128] 64 True 0.00% ├─Conv2d (decoder_conv_out) [8, 32, 128, 128] [8, 3, 128, 128] 867 True 0.01% ====================================================================================================================================================== Total params: 6,219,770 Trainable params: 6,219,770 Non-trainable params: 0 Total mult-adds (Units.GIGABYTES): 146.86 ====================================================================================================================================================== Input size (MB): 1.57 Forward/backward pass size (MB): 3719.89 Params size (MB): 22.77 Estimated Total Size (MB): 3744.23 ======================================================================================================================================================