o - io@sddlmZddlmZddlmZddlZddlmZddl mm Z ddl m Z mZmZddlmZmZeddd \ZZd gZGd d d ejZGd ddejZGdddejZGdddejZGdddejZGdd d ejZdS)) annotations)Sequence)ListN) ConvolutionSpatialAttentionBlockUpsample)ensure_tuple_repoptional_importzeinops.layers.torch Rearrange)name AutoencoderKLcs,eZdZdZd fdd Zd d d ZZS) AsymmetricPadz Pad the input tensor asymmetrically along every spatial dimension. Args: spatial_dims: number of spatial dimensions, could be 1, 2, or 3. spatial_dimsintreturnNonecstd||_dS)N)r)super__init__pad)selfr __class__c/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/autoencoderkl.pyr%s zAsymmetricPad.__init__x torch.TensorcCstjj||jddd}|S)Nconstantg)modevalue)nn functionalrrrrrrforward)szAsymmetricPad.forward)rrrrrrrr__name__ __module__ __qualname____doc__rr# __classcell__rrrrr sr cs,eZdZdZd fdd Zdd d ZZS)AEKLDownsamplez Convolution-based downsampling layer. Args: spatial_dims: number of spatial dimensions (1D, 2D, 3D). in_channels: number of input channels. rr in_channelsrrc s2tt|d|_t|||ddddd|_dS)N)rrTrr, out_channelsstrides kernel_sizepadding conv_only)rrr rrconv)rrr,rrrr7s   zAEKLDownsample.__init__rrcCs||}||}|SN)rr5r"rrrr#E  zAEKLDownsample.forward)rrr,rrrr$r%rrrrr+.sr+cs,eZdZdZdfd d ZdddZZS) AEKLResBlocka) Residual block consisting of a cascade of 2 convolutions + activation + normalisation block, and a residual connection between input and output. Args: spatial_dims: number of spatial dimensions, could be 1, 2, or 3. 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. rrr,norm_num_groupsnorm_epsfloatr0rrc st||_|dur|n||_tj|||dd|_t||j|jddddd|_tj|||dd|_ t||j|jddddd|_ ||j|jkrYt||j|jddddd|_ dSt |_ dS)NT num_groups num_channelsepsaffinerr.r/r) rrr,r0r GroupNormnorm1rconv1norm2conv2 nin_shortcutIdentity)rrr,r9r:r0rrrrYsF     zAEKLResBlock.__init__rrcCsR|}||}t|}||}||}t|}||}||}||Sr6)rBFsilurCrDrErF)rrhrrrr#s       zAEKLResBlock.forward) rrr,rr9rr:r;r0rrrr$r%rrrrr8Ks *r8cs6eZdZdZ    ddfdd ZdddZZS) Encodera Convolutional cascade that downsamples the image into a spatial latent space. Args: spatial_dims: number of spatial dimensions, could be 1, 2, or 3. in_channels: number of input channels. channels: sequence of block output channels. out_channels: number of channels in the bottom layer (latent space) of the autoencoder. 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. with_nonlocal_attn: if True use non-local attention block. 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). TFrrr,channels Sequence[int]r0num_res_blocksr9r:r;attention_levelsSequence[bool]with_nonlocal_attnbool include_fcuse_combined_linearuse_flash_attentionrrc  st||_||_||_||_||_||_||_||_ g} | t |||dddddd|d}t t |D]E}|}||}|t |dk}t |j|D]"}| t|||||d|}||rr| t||||| | | dqP|s~| t||dq9| dur| t||d |||d d| t||d ||| | | d| t||d |||d d| tj||d |dd | t |j|d |dddddt| |_dS) Nrrr.Tr/rr,r9r:r0rr>r9r:rSrTrU)rr,r<)rrrr,rLr0rNr9r:rOappendrrangelenr8rr+r rA ModuleListblocks)rrr,rLr0rNr9r:rOrQrSrTrUr]output_channeli input_channelis_final_block_rrrrs        zEncoder.__init__rrcC|jD]}||}q|Sr6r]rrblockrrrr#r7zEncoder.forward)TTFF)rrr,rrLrMr0rrNrMr9rr:r;rOrPrQrRrSrRrTrRrUrRrrr$r%rrrrrKsxrKcs8eZdZdZ     ddfdd Zd ddZZS)!Decodera Convolutional cascade upsampling from a spatial latent space into an image space. Args: spatial_dims: number of spatial dimensions, could be 1, 2, or 3. 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. with_nonlocal_attn: if True use non-local attention block. use_convtranspose: if True, use ConvTranspose to upsample feature maps in decoder. 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). TFrrrLrMr,r0rNr9r:r;rOrPrQrRuse_convtransposerSrTrUrrc s t||_||_||_||_||_||_||_||_ t t |}g}| t |||dddddd| duri| t||d|||dd| t||d||| | | d| t||d|||ddt t |}t t |}|d}tt|D]c}|}||}|t|dk}t||D]"}| t|||||d|}||r| t||||| | | dq|s| r| t|d||d qt |||ddddd}| t|d ||d d |dd q| tj|||dd| t |||dddddt||_dS)Nrrr.Tr/rVrWdeconv)rrr,r0 nontrainablenearestg@)rrr,r0 interp_mode scale_factor post_conv align_cornersr<)rrrrLr,r0rNr9r:rOlistreversedrYrr8rrZr[rr rAr\r])rrrLr,r0rNr9r:rOrQrhrSrTrUZreversed_block_out_channelsr]Zreversed_attention_levelsZreversed_num_res_blocksZ block_out_chr_Z block_in_chrarbrnrrrr9s           zDecoder.__init__rrcCrcr6rdrerrrr#r7zDecoder.forward)TFTFF)rrrLrMr,rr0rrNrMr9rr:r;rOrPrQrRrhrRrSrRrTrRrUrRrrr$r%rrrrrg$srgcseZdZdZ         d>d?fd#d$ Zd@d(d)ZdAd,d-ZdBd.d/ZdCd1d2ZdDd4d5Z dBd6d7Z dCd8d9Z dEdFdTFrrr,r0rNSequence[int] | intrLrMrOrPlatent_channelsr9r:r;with_encoder_nonlocal_attnrRwith_decoder_nonlocal_attnuse_checkpointrhrSrTrUrrcsttfdd|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)rNr).0 out_channelr9rr sz)AutoencoderKL.__init__..zDAutoencoderKL expects all channels being multiple of norm_num_groupszBAutoencoderKL 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`.) rr,rLr0rNr9r:rOrQrSrTrU) rrLr,r0rNr9r:rOrQrhrSrTrUrrTr/)rrany ValueErrorr[ isinstancerrrKencoderrgdecoderr quant_conv_muquant_conv_log_sigmapost_quant_convryr|)rrr,r0rNrLrOryr9r:rzr{r|rhrSrTrUrrrrs     zAutoencoderKL.__init__rr!tuple[torch.Tensor, torch.Tensor]cCs`|jrtjjj|j|dd}n||}||}||}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 F use_reentrantg>g4@r-) r|torchutils checkpointrrrclampexp)rrrJz_muZ z_log_varz_sigmarrrencodeRs   zAutoencoderKL.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)rrrr?Zz_vaerrrsamplingfs zAutoencoderKL.samplingcCs||\}}||}|S)z Encodes and decodes an input image. Args: x: BxCx[SPATIAL DIMENSIONS] tensor. Returns: reconstructed image, of the same shape as input )rdecode)rrrrbreconstructionrrr reconstructws zAutoencoderKL.reconstructzcCs8||}|jrtjjj|j|dd}|S||}|S)z Based on a latent space sample, forwards it through the Decoder. Args: z: Bx[Z_CHANNELS]x[LATENT SPACE SHAPE] Returns: decoded image tensor Fr)rr|rrrr)rrdecrrrrs  zAutoencoderKL.decode/tuple[torch.Tensor, torch.Tensor, torch.Tensor]cCs.||\}}|||}||}|||fSr6)rrr)rrrrrrrrrr#s   zAutoencoderKL.forwardcCs||\}}|||}|Sr6)rr)rrrrrrrrencode_stage_2_inputss z#AutoencoderKL.encode_stage_2_inputscCs||}|Sr6)r)rrimagerrrdecode_stage_2_outputss z$AutoencoderKL.decode_stage_2_outputsold_state_dictdictcs|tfdd|Drtd||dS|rBD]}||vr,td|dqtd|D]}|vrAtd|dq3D] }||vrQ|||<qDd d D}|D]r}||d |d <||d |d<||d|d<||d|d<||d|d<||d|d<t|djd|d<t|dj|d<q[D]}d|vr| dd}|||<q|rtd| |jdddS)z Load a state dict from an AutoencoderKL trained with [MONAI Generative](https://github.com/Project-MONAI/GenerativeModels). Args: old_state_dict: state dict from the old AutoencoderKL model. c3s|]}|vVqdSr6rr}knew_state_dictrrrsz4AutoencoderKL.load_old_state_dict..z#All keys match, loading state dict.Nzkey z not found in old state dictz.----------------------------------------------z not found in new state dictcSs g|] }d|vr|ddqS)zattn.to_q.weight.attn.to_q.weight)replacerrrr s z5AutoencoderKL.load_old_state_dict..z .to_q.weightrz .to_k.weightz.attn.to_k.weightz .to_v.weightz.attn.to_v.weightz .to_q.biasz.attn.to_q.biasz .to_k.biasz.attn.to_k.biasz .to_v.biasz.attn.to_v.biasz.attn.out_proj.weightrz.attn.out_proj.biaspostconvr5z!remaining keys in old_state_dict:T)strict) state_dictallprintload_state_dictpopreyeshapezerosrkeys)rrverboserZattention_blocksrfold_namerrrload_old_state_dictsR  z!AutoencoderKL.load_old_state_dict)rrrrrsrvr.rtrwTTFFTFF)"rrr,rr0rrNrxrLrMrOrPryrr9rr:r;rzrRr{rRr|rRrhrRrSrRrTrRrUrRrr)rrrr)rrrrrrr$)rrrr)rrrr)F)rrrr)r&r'r(r)rrrrrr#rrrr*rrrrr s4 `      ) __future__rcollections.abcrtypingrrtorch.nnr torch.nn.functionalr!rHZmonai.networks.blocksrrr monai.utilsrr r rb__all__Moduler r+r8rKrgr rrrrs$   G5