o , inL@sddlmZddlmZddlZddlmZddlmmZ ddl m Z m Z m Z ddlmZddlmZddlmZdgZGd d d ejZGd d d ejZGd ddejZdS)) annotations)SequenceN) ConvolutionSpatialAttentionBlockUpsample)SPADE)Encoder)ensure_tuple_repSPADEAutoencoderKLcs,eZdZdZdfd d ZdddZZS) SPADEResBlocka2 Residual block consisting of a cascade of 2 convolutions + activation + normalisation block, and a residual connection between input and output. Enables SPADE normalisation for semantic conditioning (Park et. al (2019): https://github.com/NVlabs/SPADE) Args: spatial_dims: number of spatial dimensions (1D, 2D, 3D). in_channels: input channels to the layer. norm_num_groups: number of groups involved for the group normalisation layer. Ensure that your number of channels is divisible by this number. norm_eps: epsilon for the normalisation. out_channels: number of output channels. label_nc: number of semantic channels for SPADE normalisation spade_intermediate_channels: number of intermediate channels for SPADE block layer spatial_dimsint in_channelsnorm_num_groupsnorm_epsfloat out_channelslabel_ncspade_intermediate_channelsreturnNonec st||_|dur|n||_t||d|d|d|d|d|_t||j|jddddd|_t||d|d|d|d|d|_t||j|jddddd|_ ||j|jkret||j|jddd dd|_ dSt |_ dS) NGROUPF) num_groupsaffineeps)rnorm_ncnorm norm_paramshidden_channels kernel_sizer Tr rrstridesr padding conv_onlyr) super__init__rrrnorm1rconv1norm2conv2 nin_shortcutnnIdentity)selfr rrrrrr __class__i/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/spade_autoencoderkl.pyr'-sf         zSPADEResBlock.__init__x torch.TensorsegcCsV|}|||}t|}||}|||}t|}||}||}||SN)r(Fsilur)r*r+r,)r/r4r6hr2r2r3forwardms       zSPADEResBlock.forward)r r rr rr rrrr rr rr rrr4r5r6r5rr5__name__ __module__ __qualname____doc__r'r; __classcell__r2r2r0r3r s@r cs8eZdZdZ     d!d"fdd Zd#dd ZZS)$ SPADEDecodera Convolutional cascade upsampling from a spatial latent space into an image space. Enables SPADE normalisation for semantic conditioning (Park et. al (2019): https://github.com/NVlabs/SPADE) Args: spatial_dims: number of spatial dimensions (1D, 2D, 3D). channels: sequence of block output channels. in_channels: number of channels in the bottom layer (latent space) of the autoencoder. out_channels: number of output channels. num_res_blocks: number of residual blocks (see ResBlock) per level. norm_num_groups: number of groups for the GroupNorm layers, channels must be divisible by this number. norm_eps: epsilon for the normalization. attention_levels: indicate which level from channels contain an attention block. label_nc: number of semantic channels for SPADE normalisation. with_nonlocal_attn: if True use non-local attention block. spade_intermediate_channels: number of intermediate channels for SPADE block layer. include_fc: whether to include the final linear layer. Default to True. use_combined_linear: whether to use a single linear layer for qkv projection, default to False. use_flash_attention: if True, use Pytorch's inbuilt flash attention for a memory efficient attention mechanism (see https://pytorch.org/docs/2.2/generated/torch.nn.functional.scaled_dot_product_attention.html). TFr r channels Sequence[int]rrnum_res_blocksrrrattention_levelsSequence[bool]rwith_nonlocal_attnboolr include_fcuse_combined_linearuse_flash_attentionrrc st||_||_||_||_||_||_||_||_ | |_ t t |}g}| t|||dddddd| durp| t||d|||d| | d| t||d||| | |d| t||d|||d| | dt t |}t t |}|d}tt|D]W}|}||}|t|dk}t||D]$}| t|||||| | d|}||r| t||||| | |dq|st|||ddddd}| t|d||d d |dd q| tj|||dd | t|||dddddt||_dS) Nrr!rTr")r rrrrrr)r num_channelsrrrLrMrN nontrainablenearestg@)r moderr interp_mode scale_factor post_conv align_corners)rrOrr)r&r'r rErrrGrrrHrlistreversedappendrr rrangelenrr- GroupNorm ModuleListblocks)r/r rErrrGrrrHrrJrrLrMrNreversed_block_out_channelsr^reversed_attention_levelsreversed_num_res_blocks block_out_chi block_in_chis_final_block_rUr0r2r3r's           zSPADEDecoder.__init__r4r5r6cCs.|jD]}t|tr|||}q||}q|Sr7)r^ isinstancer )r/r4r6blockr2r2r3r;,s    zSPADEDecoder.forward)TrDTFF)r r rErFrr rr rGrFrr rrrHrIrr rJrKrr rLrKrMrKrNrKrrr<r=r2r2r0r3rC{s!rCcseZdZdZ         dd)d*Zd?d-d.Zd@d0d1ZdAd3d4ZdBd6d7Z dCd8d9Z dAd:d;Z Z S)Dr a/ Autoencoder model with KL-regularized latent space based on Rombach et al. "High-Resolution Image Synthesis with Latent Diffusion Models" https://arxiv.org/abs/2112.10752 and Pinaya et al. "Brain Imaging Generation with Latent Diffusion Models" https://arxiv.org/abs/2209.07162 Enables SPADE normalisation for semantic conditioning (Park et. al (2019): https://github.com/NVlabs/SPADE) Args: spatial_dims: number of spatial dimensions (1D, 2D, 3D). label_nc: number of semantic channels for SPADE normalisation. in_channels: number of input channels. out_channels: number of output channels. num_res_blocks: number of residual blocks (see ResBlock) per level. channels: sequence of block output channels. attention_levels: sequence of levels to add attention. latent_channels: latent embedding dimension. norm_num_groups: number of groups for the GroupNorm layers, channels must be divisible by this number. norm_eps: epsilon for the normalization. with_encoder_nonlocal_attn: if True use non-local attention block in the encoder. with_decoder_nonlocal_attn: if True use non-local attention block in the decoder. spade_intermediate_channels: number of intermediate channels for SPADE block layer. r!rjrjrj @rmrmFFTTrrlư>TrDFr r rrrrGSequence[int] | intrErFrHrIlatent_channelsrrrwith_encoder_nonlocal_attnrKwith_decoder_nonlocal_attnrrLrMrNrrcs ttfdd|Drtdt|t|kr tdt|tr,t|t|}t|t|kr8tdt|||||| || |||d |_ t |||||| ||| | |||d|_ t |||ddd d d |_ t |||ddd d d |_t |||ddd d d |_||_dS) Nc3s|] }|dkVqdS)rNr2).0 out_channelrr2r3 bsz.SPADEAutoencoderKL.__init__..zISPADEAutoencoderKL expects all channels being multiple of norm_num_groupszGSPADEAutoencoderKL expects channels being same size of attention_levelszf`num_res_blocks` should be a single integer or a tuple of integers with the same length as `channels`.) r rrErrGrrrHrJrLrMrN)r rErrrGrrrHrrJrrLrMrNr!rTr")r&r'any ValueErrorr[rgr r rencoderrCdecoderr quant_conv_muquant_conv_log_sigmapost_quant_convrq)r/r rrrrGrErHrqrrrrrsrrLrMrNr0rvr3r'Ls     zSPADEAutoencoderKL.__init__r4r5!tuple[torch.Tensor, torch.Tensor]cCsB||}||}||}t|dd}t|d}||fS)z Forwards an image through the spatial encoder, obtaining the latent mean and sigma representations. Args: x: BxCx[SPATIAL DIMS] tensor g>g4@rj)rzr|r}torchclampexp)r/r4r:z_mu z_log_varz_sigmar2r2r3encodes   zSPADEAutoencoderKL.encoderrcCst|}|||}|S)aE From the mean and sigma representations resulting of encoding an image through the latent space, obtains a noise sample resulting from sampling gaussian noise, multiplying by the variance (sigma) and adding the mean. Args: z_mu: Bx[Z_CHANNELS]x[LATENT SPACE SIZE] mean vector obtained by the encoder when you encode an image z_sigma: Bx[Z_CHANNELS]x[LATENT SPACE SIZE] variance vector obtained by the encoder when you encode an image Returns: sample of shape Bx[Z_CHANNELS]x[LATENT SPACE SIZE] )r randn_like)r/rrrz_vaer2r2r3samplings zSPADEAutoencoderKL.samplingr6cCs||\}}|||}|S)a Encodes and decodes an input image. Args: x: BxCx[SPATIAL DIMENSIONS] tensor. seg: Bx[LABEL_NC]x[SPATIAL DIMENSIONS] tensor of segmentations for SPADE norm. Returns: reconstructed image, of the same shape as input )rdecode)r/r4r6rrfreconstructionr2r2r3 reconstructs zSPADEAutoencoderKL.reconstructzcCs||}|||}|S)a! Based on a latent space sample, forwards it through the Decoder. Args: z: Bx[Z_CHANNELS]x[LATENT SPACE SHAPE] seg: Bx[LABEL_NC]x[SPATIAL DIMENSIONS] tensor of segmentations for SPADE norm. Returns: decoded image tensor )r~r{)r/rr6decr2r2r3rs zSPADEAutoencoderKL.decode/tuple[torch.Tensor, torch.Tensor, torch.Tensor]cCs0||\}}|||}|||}|||fSr7)rrr)r/r4r6rrrrr2r2r3r;s   zSPADEAutoencoderKL.forwardcCs||\}}|||}|Sr7)rr)r/r4rrrr2r2r3encode_stage_2_inputss z(SPADEAutoencoderKL.encode_stage_2_inputscCs|||}|Sr7)r)r/rr6imager2r2r3decode_stage_2_outputss z)SPADEAutoencoderKL.decode_stage_2_outputs)r!r!rirkrnrrlroTTrDTFF)"r r rr rr rr rGrprErFrHrIrqr rr rrrrrKrsrKrr rLrKrMrKrNrKrr)r4r5rr)rr5rr5rr5r<)rr5r6r5rr5)r4r5r6r5rr)r4r5rr5) r>r?r@rAr'rrrrr;rrrBr2r2r0r3r 5s0 `     ) __future__rcollections.abcrrtorch.nnr-torch.nn.functional functionalr8monai.networks.blocksrrrZ monai.networks.blocks.spade_normrZ!monai.networks.nets.autoencoderklr monai.utilsr __all__Moduler rCr r2r2r2r3s     _;