o - i@@sddlmZddlZddlmZddlmZddlZddl Z ddl m Z ddl m Z m Z mZmZmZmZddlmZddlmZmZed d d \ZZd gZdYddZGddde jZGddde jZdZd[ddZGddde jZGd d!d!eZ Gd"d#d#e jZ!Gd$d%d%e jZ"Gd&d'd'e jZ#Gd(d)d)e jZ$Gd*d+d+e jZ%Gd,d-d-e jZ&Gd.d/d/e jZ'Gd0d1d1e jZ(Gd2d3d3e jZ) 4 5 6 4 4d\d]dMdNZ* 4 5 6 4 4d\d^dPdQZ+ 4 5 6 4 4d\d_dTdUZ,GdVd d e jZ-GdWdXdXe jZ.dS)`) annotationsN)Sequence)reduce)nn) ConvolutionCrossAttentionBlockMLPBlockSABlockSpatialAttentionBlockUpsample)Pool)ensure_tuple_repoptional_importzeinops.layers.torch Rearrange)nameDiffusionModelUNetmodule nn.ModulereturncCs|D]}|q|S)z< Zero out the parameters of a module and return it. ) parametersdetachzero_)rprj/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/diffusion_model_unet.py zero_module3s rcs<eZdZdZ      ddfdd Zdd ddZZS)!DiffusionUNetTransformerBlocka A Transformer block that allows for the input dimension to differ from the hidden dimension. Args: num_channels: number of channels in the input and output. num_attention_heads: number of heads to use for multi-head attention. num_head_channels: number of channels in each attention head. dropout: dropout probability to use. cross_attention_dim: size of the context vector for cross attention. upcast_attention: if True, upcast attention operations to full precision. 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). 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. NFT num_channelsintnum_attention_headsnum_head_channelsdropoutfloatcross_attention_dim int | Noneupcast_attentionbooluse_flash_attention include_fcuse_combined_linearrNonec stt|||||||rtjnd|| |d |_t||dd|d|_t||||||||r3tjnd|d|_ t ||_ t ||_ t ||_dS)N) hidden_sizehidden_input_size num_headsdim_head dropout_rateattention_dtyper)r*r(GEGLU)r,mlp_dimactr0)r,r.r-context_input_sizer/r0r1r()super__init__r torchr#attn1rffrattn2r LayerNormnorm1norm2norm3) selfrr r!r"r$r&r(r)r* __class__rrr8Ns4     z&DiffusionUNetTransformerBlock.__init__x torch.Tensorcontexttorch.Tensor | NonecCsD||||}|j|||d|}||||}|SNrF)r:r>r<r?r;r@)rArDrFrrrforwardusz%DiffusionUNetTransformerBlock.forward)rNFFTF)rrr rr!rr"r#r$r%r&r'r(r'r)r'r*r'rr+NrDrErFrGrrE__name__ __module__ __qualname____doc__r8rJ __classcell__rrrBrr<s'rcsBeZdZdZ         d$d%fdd Zd&d'd"d#ZZS)(SpatialTransformera Transformer block for image-like data. First, project the input (aka embedding) and reshape to b, t, d. Then apply standard transformer action. Finally, reshape to image. Args: spatial_dims: number of spatial dimensions. in_channels: number of channels in the input and output. num_attention_heads: number of heads to use for multi-head attention. num_head_channels: number of channels in each attention head. num_layers: number of layers of Transformer blocks to use. dropout: dropout probability to use. norm_num_groups: number of groups for the normalization. norm_eps: epsilon for the normalization. cross_attention_dim: number of context dimensions to use. upcast_attention: if True, upcast attention operations to full precision. 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). r ư>NFT spatial_dimsr in_channelsr r! num_layersr"r#norm_num_groupsnorm_epsr$r%r&r'r)r*r(rr+c st||_||_tj|||dd|_t||ddddd|_t f ddt |D|_ t t||ddddd|_ dS)NT num_groupsrepsaffinerTrrWrX out_channelsstrides kernel_sizepadding conv_onlyc s&g|]}td qS)) rr r!r"r$r&r)r*r()r).0_ r$r"r) inner_dimr r!r&r*r(rr s z/SpatialTransformer.__init__..)r7r8rWrXr GroupNormnormrproj_in ModuleListrangetransformer_blocksrproj_out)rArWrXr r!rYr"rZr[r$r&r)r*r(rBrhrr8s<    zSpatialTransformer.__init__rDrErFrGc Cs@d}}}}}|jdkr|j\}}}}|jdkr#|j\}}}}}|}||}||}|jd} |jdkrH|dddd|||| }|jdkr_|ddddd||||| }|jD]} | ||d}qb|jdkr||||| dddd}|jdkr|||||| ddddd}||}||S)NrTrr2rI) rWshaperlrmpermutereshaperp contiguousrq) rArDrFbatchchannelheightwidthdepthresidualriblockrrrrJs*       $    $ zSpatialTransformer.forward) rTrrUrVNFTFF)rWrrXrr rr!rrYrr"r#rZrr[r#r$r%r&r'r)r'r*r'r(r'rr+rKrLrMrrrBrrSs>rS' timestepsrE embedding_dimr max_periodcCs|jdkr td|d}t| tjd|tj|jd}t||}|dddf |dddf}tj t |t |gdd}|ddkrTtj j|d }|S) at Create sinusoidal timestep embeddings following the implementation in Ho et al. "Denoising Diffusion Probabilistic Models" https://arxiv.org/abs/2006.11239. Args: timesteps: a 1-D Tensor of N indices, one per batch element. embedding_dim: the dimension of the output. max_period: controls the minimum frequency of the embeddings. rTzTimesteps should be a 1d-arrayrsr)startenddtypedeviceNrrdim)rrTrr)ndim ValueErrormathlogr9arangefloat32rexpr#catcossinr functionalpad)rrrZhalf_dimexponentfreqsargs embeddingrrrget_timestep_embeddings "$ rcs2eZdZdZ ddfdd ZddddZZS)DiffusionUnetDownsamplea Downsampling layer. Args: spatial_dims: number of spatial dimensions. num_channels: number of input channels. use_conv: if True uses Convolution instead of Pool average to perform downsampling. In case that use_conv is False, the number of output channels must be the same as the number of input channels. out_channels: number of output channels. padding: controls the amount of implicit zero-paddings on both sides for padding number of points for each dimension. NrTrWrruse_convr'rar%rdrr+c stt||_|p ||_||_|r"t||j|jdd|dd|_dS|j|jkr,tdttj |fddd|_dS)NrsrtTr`z?num_channels and out_channels must be equal when use_conv=False)rcstride) r7r8rrarroprr AVG)rArWrrrardrBrrr8s"   z DiffusionUnetDownsample.__init__rDrEembrGcCs>~|jd|jkrtd|jdd|jd||}|S)NrTzInput number of channels (z/) is not equal to expected number of channels ())rurrr)rArDroutputrrrrJ4s zDiffusionUnetDownsample.forward)NrT) rWrrrrr'rar%rdrrr+rKrDrErrGrrErMrrrBrrs rcs$eZdZdZd d fdd ZZS) WrappedUpsamplezS Wraps MONAI upsample block to allow for calling with timestep embeddings. NrDrErrGrcs~t|}|SrK)r7rJ)rArDrZ upsampledrBrrrJDs zWrappedUpsample.forwardrKr)rNrOrPrQrJrRrrrBrr?srcs8eZdZdZ     ddfdd ZdddZZS)DiffusionUNetResnetBlocka Residual block with timestep conditioning. Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. temb_channels: number of timestep embedding channels. out_channels: number of output channels. up: if True, performs upsampling. down: if True, performs downsampling. norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. NFrUrVrWrrX temb_channelsrar%upr'downrZr[r#rr+c s*t||_||_||_|p||_||_||_tj |||dd|_ t |_ t |||jddddd|_d|_|_|jrKt|d||dddd |_n |rUt||d d |_t||j|_tj ||j|dd|_tt ||j|jddddd|_||j|krt|_dSt |||jddd dd|_dS) NTr\rTrtr` nontrainablenearest@)rWmoderXra interp_mode scale_factor align_cornersF)rr)r7r8rWchannelsZ emb_channelsrarrrrkr>SiLU nonlinearityrconv1upsample downsamplerrLinear time_emb_projr?rconv2Identityskip_connection) rArWrXrrarrrZr[rBrrr8Ysp        z!DiffusionUNetResnetBlock.__init__rDrErcCs|}||}||}|jdur||}||}n|jdur+||}||}||}|jdkrH|||ddddddf}n|||dddddddf}||}||}||}||}| ||}|S)Nrs) r>rrrrrWrr?rr)rArDrhtembrrrrrJs&          &&   z DiffusionUNetResnetBlock.forward)NFFrUrV)rWrrXrrrrar%rr'rr'rZrr[r#rr+)rDrErrErrErMrrrBrrJsIrcs>eZdZdZ      d!d"fdd Z d#d$dd ZZS)% DownBlocka Unet's down block containing resnet and downsamplers blocks. Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. temb_channels: number of timestep embedding channels. num_res_blocks: number of residual blocks. norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. add_downsample: if True add downsample block. resblock_updown: if True use residual blocks for downsampling. downsample_padding: padding used in the downsampling block. rTrUrVTFrWrrXrarnum_res_blocksrZr[r#add_downsampler'resblock_updowndownsample_paddingrr+c st| |_g} t|D]} | dkr|n|}| t||||||dqt| |_|rL|| r@t||||||dd|_ dSt ||d|| d|_ dSd|_ dS)NrrWrXrarrZr[TrWrXrarrZr[rrWrrrard) r7r8rroappendrrrnresnets downsamplerr) rArWrXrarrrZr[rrrrirBrrr8sH      zDownBlock.__init__N hidden_statesrErrFrG'tuple[torch.Tensor, list[torch.Tensor]]cCsN~g}|jD] }|||}||q|jdur#|||}||||fSrK)rrr)rArrrF output_statesresnetrrrrJs      zDownBlock.forward)rTrUrVTFrT)rWrrXrrarrrrrrZrr[r#rr'rr'rrrr+rKrrErrErFrGrrrMrrrBrrs9rcsFeZdZdZ          d%d&fdd Z d'd(d#d$ZZS)) AttnDownBlocka Unet's down block containing resnet, downsamplers and self-attention blocks. Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. temb_channels: number of timestep embedding channels. num_res_blocks: number of residual blocks. norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. add_downsample: if True add downsample block. resblock_updown: if True use residual blocks for downsampling. downsample_padding: padding used in the downsampling block. num_head_channels: number of channels in each attention head. 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). rTrUrVTFrWrrXrarrrZr[r#rr'rrr!r)r*r(rr+c st| |_g}g}t|D]&}|dkr|n|}|t||||||d|t||| ||| | |dqt||_ t||_ ||rc| rWt||||||dd|_ dSt ||d|| d|_ dSd|_ dS)NrrrWrr!rZr[r)r*r(Trr) r7r8rrorrr rrn attentionsrrr)rArWrXrarrrZr[rrrr!r)r*r(rrrrBrrr8-sd       zAttnDownBlock.__init__NrrErrFrGrcCsf~g}t|j|jD]\}}|||}||}||q |jdur/|||}||||fSrKziprrrxrrrArrrFrrattnrrrrJws      zAttnDownBlock.forward) rTrUrVTFrTrTTFF)rWrrXrrarrrrrrZrr[r#rr'rr'rrr!rr)r'r*r'r(r'rr+rKrrMrrrBrrsKrcsNeZdZdZ              d+d,fd!d" Z d-d.d)d*ZZS)/CrossAttnDownBlocka Unet's down block containing resnet, downsamplers and cross-attention blocks. Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. temb_channels: number of timestep embedding channels. num_res_blocks: number of residual blocks. norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. add_downsample: if True add downsample block. resblock_updown: if True use residual blocks for downsampling. downsample_padding: padding used in the downsampling block. num_head_channels: number of channels in each attention head. transformer_num_layers: number of layers of Transformer blocks to use. cross_attention_dim: number of context dimensions to use. upcast_attention: if True, upcast attention operations to full precision. dropout_cattn: if different from zero, this will be the dropout value for the cross-attention layers. 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). rTrUrVTFNrrWrrXrarrrZr[r#rr'rrr!transformer_num_layersr$r%r& dropout_cattnr)r*r(rr+cst| |_g}g}t|D]-}|dkr|n|}|t||||||d|t|||| | | ||| |||||d qt||_ t||_ ||rj| r^t||||||dd|_ dSt ||d|| d|_ dSd|_ dS)Nrr rWrXr r!rYrZr[r$r&r"r)r*r(Trr) r7r8rrorrrSrrnrrrr)rArWrXrarrrZr[rrrr!rr$r&rr)r*r(rrrrBrrr8sn        zCrossAttnDownBlock.__init__rrErrFrGrcCshg}t|j|jD]\}}|||}|||d}||q |jdur0|||}||||fSrHrrrrrrJs     zCrossAttnDownBlock.forward)rTrUrVTFrTrTrTNFrTFF)&rWrrXrrarrrrrrZrr[r#rr'rr'rrr!rrrr$r%r&r'rr#r)r'r*r'r(r'rr+rKrrMrrrBrrs&Urcs>eZdZdZ      dd fdd Z d!d"ddZZS)# AttnMidBlockaZ Unet's mid block containing resnet and self-attention blocks. Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. temb_channels: number of timestep embedding channels. norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. num_head_channels: number of channels in each attention head. 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). rUrVrTTFrWrrXrrZr[r#r!r)r'r*r(rr+c sTtt||||||d|_t|||||||| d|_t||||||d|_dS)Nrr)r7r8rresnet_1r attentionresnet_2) rArWrXrrZr[r!r)r*r(rBrrr8s6   zAttnMidBlock.__init__NrrErrFrGcCs,~|||}||}|||}|SrK)rrrxrrArrrFrrrrJCs   zAttnMidBlock.forward)rUrVrTTFF)rWrrXrrrrZrr[r#r!rr)r'r*r'r(r'rr+rKrrErrErFrGrrErMrrrBrrs+rcsFeZdZdZ          d%d&fdd Z d'd(d#d$ZZS))CrossAttnMidBlocka= Unet's mid block containing resnet and cross-attention blocks. Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. temb_channels: number of timestep embedding channels norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. num_head_channels: number of channels in each attention head. transformer_num_layers: number of layers of Transformer blocks to use. cross_attention_dim: number of context dimensions to use. upcast_attention: if True, upcast attention operations to full precision. 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). rUrVrTNFrTrWrrXrrZr[r#r!rr$r%r&r'rr)r*r(rr+csbtt||||||d|_t|||||||||| | | | | d |_t||||||d|_dS)Nrr)r7r8rrrSrr)rArWrXrrZr[r!rr$r&rr)r*r(rBrrr8bs@  zCrossAttnMidBlock.__init__rrErrFrGcCs*|||}|j||d}|||}|SrH)rrrrrrrrJs  zCrossAttnMidBlock.forward) rUrVrTrTNFrTFF)rWrrXrrrrZrr[r#r!rrrr$r%r&r'rr#r)r'r*r'r(r'rr+rKrrMrrrBrrNs3rcs<eZdZdZ     d"d#fdd Z d$d%d d!ZZS)&UpBlocka Unet's up block containing resnet and upsamplers blocks. Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. prev_output_channel: number of channels from residual connection. out_channels: number of output channels. temb_channels: number of timestep embedding channels. num_res_blocks: number of residual blocks. norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. add_upsample: if True add downsample block. resblock_updown: if True use residual blocks for upsampling. rTrUrVTFrWrrXprev_output_channelrarrrZr[r# add_upsampler'rrr+c  st| |_g} t|D]#} | |dkr|n|} | dkr |n|}| t||| ||||dqt| |_|| rf| rLt||||||dd|_ dSt |||ddddd}t |d||d d |dd |_ dSd|_ dS) NrTrrTrWrXrarrZr[rrtr`rrrrWrrXrarr post_convr) r7r8rrorrrrnr upsamplerrr)rArWrXrrarrrZr[rrrrres_skip_channelsresnet_in_channelsrrBrrr8sb       zUpBlock.__init__NrrEres_hidden_states_listlist[torch.Tensor]rrFrGcCsX~|jD]}|d}|dd}tj||gdd}|||}q|jdur*|||}|SNrrrTr)rr9rr)rArrrrFrres_hidden_statesrrrrJs     zUpBlock.forward)rTrUrVTF)rWrrXrrrrarrrrrrZrr[r#rr'rr'rr+rK rrErrrrErFrGrrErMrrrBrrsKrcsDeZdZdZ         d&d'fdd Z d(d)d$d%ZZS)* AttnUpBlocka Unet's up block containing resnet, upsamplers, and self-attention blocks. Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. prev_output_channel: number of channels from residual connection. out_channels: number of output channels. temb_channels: number of timestep embedding channels. num_res_blocks: number of residual blocks. norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. add_upsample: if True add downsample block. resblock_updown: if True use residual blocks for upsampling. num_head_channels: number of channels in each attention head. 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). rTrUrVTFrWrrXrrarrrZr[r#rr'rr!r)r*r(rr+c st| |_g}g}t|D]2}||dkr|n|}|dkr"|n|}|t|||||||d|t||| ||| | |dqt||_ t||_ || r}| rct||||||dd|_ dSt |||ddddd}t |d ||d d |dd |_ dSd|_ dS) NrTrrrTrrtr`rrrr)r7r8rrorrr rrnrrrrr)rArWrXrrarrrZr[rrr!r)r*r(rrrrrrrBrrr8!s~        zAttnUpBlock.__init__NrrErrrrFrGcCsp~t|j|jD]"\}}|d}|dd}tj||gdd}|||}||}q|jdur6|||}|Sr)rrrr9rrxrrArrrrFrrrrrrrJzs    zAttnUpBlock.forward) rTrUrVTFrTTFF)rWrrXrrrrarrrrrrZrr[r#rr'rr'r!rr)r'r*r'r(r'rr+rKrrMrrrBrr s^rcsLeZdZdZ             d,d-fd!d" Z d.d/d*d+ZZS)0CrossAttnUpBlocka Unet's up block containing resnet, upsamplers, and self-attention blocks. Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. prev_output_channel: number of channels from residual connection. out_channels: number of output channels. temb_channels: number of timestep embedding channels. num_res_blocks: number of residual blocks. norm_num_groups: number of groups for the group normalization. norm_eps: epsilon for the group normalization. add_upsample: if True add downsample block. resblock_updown: if True use residual blocks for upsampling. num_head_channels: number of channels in each attention head. transformer_num_layers: number of layers of Transformer blocks to use. cross_attention_dim: number of context dimensions to use. upcast_attention: if True, upcast attention operations to full precision. dropout_cattn: if different from zero, this will be the dropout value for the cross-attention layers. 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). rTrUrVTFNrrWrrXrrarrrZr[r#rr'rr!rr$r%r&rr)r*r(rr+cst| |_g}g}t|D]9}||dkr|n|}|dkr"|n|}|t|||||||d|t|||| | ||| | |||||d qt||_ t||_ || r| rjt||||||dd|_ dSt |||ddddd}t |d ||d d |dd |_ dSd|_ dS) NrTrr) rWrXr r!rZr[rYr$r&r"r)r*r(Trrtr`rrrr)r7r8rrorrrSrrnrrrrr)rArWrXrrarrrZr[rrr!rr$r&rr)r*r(rrrrrrrBrrr8s         zCrossAttnUpBlock.__init__rrErrrrFrGcCsnt|j|jD]"\}}|d}|dd}tj||gdd}|||}|||d}q|jdur5|||}|S)NrrrTrrI)rrrr9rrrrrrrJ s    zCrossAttnUpBlock.forward) rTrUrVTFrTrTNFrTFF)&rWrrXrrrrarrrrrrZrr[r#rr'rr'r!rrrr$r%r&r'rr#r)r'r*r'r(r'rr+rKrrMrrrBrrs$ grFrTrWrXrarrrZr[r#rr'r with_attnwith_cross_attnr!rr$r%r&rr)r*r(cCs| rt|||||||||| |||d S| rMtdid|d|d|d|d|d|d|d |d |d | d | d | d|d|d|d|d|St|||||||||d S)N) rWrXrarrrZr[rrr!r)r*r(rWrXrarrrZr[rrr!rr$r&rr)r*r() rWrXrarrrZr[rrr)rrr)rWrXrarrrZr[rrrrr!rr$r&rr)r*r(rrrget_down_block#s~     rwith_conditioningcCs@|rt||||||||| | | | | d St||||||| | | d S)N) rWrXrrZr[r!rr$r&rr)r*r() rWrXrrZr[r!r)r*r()rrrWrXrrZr[rr!rr$r&rr)r*r(rrr get_mid_blockjs6rrrcCs| rt|||||||||| | |||dS| rQtdid|d|d|d|d|d|d|d |d |d | d | d | d|d|d|d|d|d|St|||||||||| d S)N)rWrXrrarrrZr[rrr!r)r*r(rWrXrrarrrZr[rrr!rr$r&rr)r*r() rWrXrrarrrZr[rrr)rrr)rWrXrrarrrZr[rrrrr!rr$r&rr)r*r(rrr get_up_blocks     rcsdeZdZdZ           d;d<fd*d+ Z d=d>d5d6Zd?d@d9d:ZZS)Arap Unet network with timestep embedding and attention mechanisms for conditioning 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 Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. num_res_blocks: number of residual blocks (see _ResnetBlock) per level. channels: tuple of block output channels. attention_levels: list of levels to add attention. norm_num_groups: number of groups for the normalization. norm_eps: epsilon for the normalization. resblock_updown: if True use residual blocks for up/downsampling. num_head_channels: number of channels in each attention head. with_conditioning: if True add spatial transformers to perform conditioning. transformer_num_layers: number of layers of Transformer blocks to use. cross_attention_dim: number of context dimensions to use. num_class_embeds: if specified (as an int), then this model will be class-conditional with `num_class_embeds` classes. upcast_attention: if True, upcast attention operations to full precision. dropout_cattn: if different from zero, this will be the dropout value for the cross-attention layers. 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 True. 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). rsrsrsrsrU@rrFFTTrUrVFrTNrTrWrrXrarSequence[int] | intr Sequence[int]attention_levelsSequence[bool]rZr[r#rr'r!int | Sequence[int]rrr$r%num_class_embedsr&rr)r*r(rr+c st| dur| durtd| dur| durtd|dks%|dkr)tdtfdd |Dr8td t|t|krDtd t| trPt| t|} t| t|kr\td t|trht|t|}t|t|krttd ||_||_ ||_ ||_ ||_ | |_ | |_t|||dddddd|_|dd}tt|d|tt|||_||_|durt|||_tg|_|d}tt|D]d}|}||}|t|dk}td+id|d|d|d|d||dd|d| d| d||o| d||o| d| |d| d | d!|d"|d#|d$|d%|}|j|qt||d&||| | d&| | |||||d'|_tg|_ t!t"|}t!t"|}t!t"|}t!t"| }|d}tt|D]w}|}||}|t#|dt|d}|t|dk}t$d+id|d|d(|d|d|d||ddd|d)| d| d||o| d||o| d||d| d | d!|d"|d#|d$|d%|}|j |qsttj%|d|dd*tt&t||d|ddddd|_'dS),NTz|DiffusionModelUNet expects dimension of the cross-attention conditioning (cross_attention_dim) when using with_conditioning.FzZDiffusionModelUNet expects with_conditioning=True when specifying the cross_attention_dim.g?rz#Dropout cannot be negative or >1.0!c3|] }|dkVqdSrNrrf out_channelrZrr %z.DiffusionModelUNet.__init__..zMDiffusionModelUNet expects all num_channels being multiple of norm_num_groupszKDiffusionModelUNet expects num_channels being same size of attention_levelsnum_head_channels should have the same length as attention_levels. For the i levels without attention, i.e. `attention_level[i]=False`, the num_head_channels[i] will be ignored.zj`num_res_blocks` should be a single integer or a tuple of integers with the same length as `num_channels`.rrTrtr`r2rWrXrarrrZr[rrrrr!rr$r&rr)r*r(rrrrrr\r)(r7r8ranylen isinstancerr rXblock_out_channelsrarrr!rrconv_inr Sequentialrr time_embedr Embeddingclass_embeddingrn down_blocksrorrr middle_block up_blockslistreversedminrrkrout) rArWrXrarrrrZr[rr!rrr$rr&rr)r*r(time_embed_dimoutput_channelr input_channelis_final_block down_blockreversed_block_out_channelsreversed_num_res_blocksreversed_attention_levelsZreversed_num_head_channelsrZup_blockrBrrr8sp                      zDiffusionModelUNet.__init__rDrErrFrG class_labelsdown_block_additional_residualstuple[torch.Tensor] | Nonemid_block_additional_residualcCsbt||jd}|j|jd}||}|jdur1|dur!td||} | j|jd} || }||} |durC|j durCtd| g} |j D]} | | ||d\} } | D]}| |qVqI|durzg}t | |D] \}}||}||g7}qj|} |j | ||d} |dur| |} |jD]}t|j }| |d} | d|} || | ||d} q|| }|S) a Args: x: input tensor (N, C, SpatialDims). timesteps: timestep tensor (N,). context: context tensor (N, 1, ContextDim). class_labels: context tensor (N, ). down_block_additional_residuals: additional residual tensors for down blocks (N, C, FeatureMapsDims). mid_block_additional_residual: additional residual tensor for mid block (N, C, FeatureMapsDims). rrN9class_labels should be provided when num_class_embeds > 0FAmodel should have with_conditioning = True if context is providedrrrF)rrrrF)rrtorrrrrrrrrrrrr rr)rArDrrFr#r$r&t_embr class_embrdown_block_res_samplesdownsample_block res_samplesr~Znew_down_block_res_samplesdown_block_res_sampleZdown_block_additional_residualZupsample_blockidxrrrrrJsH             zDiffusionModelUNet.forwardold_state_dictdictcsD|tfdd|Drtd||dS|rBD]}||vr,td|dqtd|D]}|vrAtd|dq3D] }||vrQ|||<qDd d D}|D]x}||d |d <||d |d<||d|d<||d|d<||d|d<||d|d<|dvr|dvr||d|d<||d|d<q[dd D}|D]}||d|d<||d|d<q݈D]}d |vr|d d!}|||<q|rtd"||dS)#z Load a state dict from a DiffusionModelUNet trained with [MONAI Generative](https://github.com/Project-MONAI/GenerativeModels). Args: old_state_dict: state dict from the old DecoderOnlyTransformer model. c3s|]}|vVqdSrKrrfknew_state_dictrrrsz9DiffusionModelUNet.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_k.weight.attn.to_k.weightreplacer5rrrrj3s z:DiffusionModelUNet.load_old_state_dict..z .to_q.weightz.attn.to_q.weightz .to_k.weightr9z .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.weightz.attn.out_proj.biasz.proj_attn.weightz.proj_attn.biascSs(g|]}d|vrd|vr|ddqS)zout_proj.weightrp.out_proj.weightr:r;r5rrrrjAs  z.to_out.0.weightr=z.to_out.0.biasz.out_proj.biaspostconvconvz!remaining keys in old_state_dict:) state_dictallprintload_state_dictpopr<keys)rAr3verboser6attention_blocksrZcross_attention_blocksold_namerr7rload_old_state_dictsZ    z&DiffusionModelUNet.load_old_state_dict)rrrrUrVFrFrTNNFrTFF)(rWrrXrrarrrrrrrrZrr[r#rr'r!rrr'rrr$r%rr%r&r'rr#r)r'r*r'r(r'rr+)NNNN)rDrErrErFrGr#rGr$r%r&rGrrE)F)r3r4rr+)rNrOrPrQr8rJrIrRrrrBrrs4"I McsTeZdZdZ           d4d5fd*d+ Z d6d7d2d3ZZS)8DiffusionModelEncoderaU Classification Network based on the Encoder of the Diffusion Model, followed by fully connected layers. This network is based on Wolleb et al. "Diffusion Models for Medical Anomaly Detection" (https://arxiv.org/abs/2203.04306). Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. input_shape: spatial shape of the input (without batch and channel dims). num_res_blocks: number of residual blocks (see _ResnetBlock) per level. channels: tuple of block output channels. attention_levels: list of levels to add attention. norm_num_groups: number of groups for the normalization. norm_eps: epsilon for the normalization. resblock_updown: if True use residual blocks for downsampling. num_head_channels: number of channels in each attention head. with_conditioning: if True add spatial transformers to perform conditioning. transformer_num_layers: number of layers of Transformer blocks to use. cross_attention_dim: number of context dimensions to use. num_class_embeds: if specified (as an int), then this model will be class-conditional with `num_class_embeds` classes. upcast_attention: if True, upcast attention operations to full precision. rrrrrrUrVFrrTNTrWrrXra input_shaperrrrrrrZr[r#rr'r!rrrr$r%rr&r)r*r(rr+c st| dur|durtd|dur| durtdtfdd|Dr,tdt|t|kr8tdt| trDt| t|} t|trPt|t|}t| t|kr\td ||_||_ ||_ ||_ ||_ | |_ | |_t|||d d d d dd |_|d d}tt|d |tt|||_||_|durt|||_tg|_|d }tt|D]]}|}||}|t|k}td(id|d|d|d|d||dd| d| d| d||o| d||o| d| |d| d|d|d|d|d |}|j|q|D] }d!d"|D}qttd#d$||d%}tt|d&tt d'td&|j |_!dS))NTzDiffusionModelEncoder expects dimension of the cross-attention conditioning (cross_attention_dim) when using with_conditioning.Fz]DiffusionModelEncoder expects with_conditioning=True when specifying the cross_attention_dim.c3rrrrrrrrr z1DiffusionModelEncoder.__init__..zPDiffusionModelEncoder expects all num_channels being multiple of norm_num_groupszNDiffusionModelEncoder expects num_channels being same size of attention_levelsr rrTrtr`r2rWrXrarrrZr[rrrrr!rr$r&r)r*r(cSsg|] }tt|dqS)rs)rnpceil)rfi_rrrrjsz2DiffusionModelEncoder.__init__..cSs||SrKr)rDyrrrsz0DiffusionModelEncoder.__init__..rrig?r)"r7r8rr r r rr rXrrarrr!rrrrrrrrrrrrnrrorrrReLUDropoutr)rArWrXrarLrrrrZr[rr!rrr$rr&r)r*r(rrrrrrrgZlast_dim_flattenedrBrrr8ms           $ zDiffusionModelEncoder.__init__rDrErrFrGr#c Cst||jd}|j|jd}||}|jdur1|dur!td||}|j|jd}||}||}|durC|j durCtd|j D] } | |||d\}} qF| |j dd}| |} | S) z Args: x: input tensor (N, C, SpatialDims). timesteps: timestep tensor (N,). context: context tensor (N, 1, ContextDim). class_labels: context tensor (N, ). rr'Nr(Fr)r*rr)rrr+rrrrrrrrrwrur) rArDrrFr#r,rr-rr/rgrrrrrJs"      zDiffusionModelEncoder.forward)rKrrrrUrVFrFrTNNFTFF)(rWrrXrrarrLrrrrrrrrZrr[r#rr'r!rrr'rrr$r%rr%r&r'r)r'r*r'r(r'rr+)NN) rDrErrErFrGr#rGrrErMrrrBrrJUs,}rJ)rrrr)r)rrErrrrrrE)FrTFF)(rWrrXrrarrrrrrZrr[r#rr'rr'rr'rr'r!rrrr$r%r&r'rr#r)r'r*r'r(r'rr)rWrrXrrrrZrr[r#rr'r!rrrr$r%r&r'rr#r)r'r*r'r(r'rr)*rWrrXrrrrarrrrrrZrr[r#rr'rr'rr'rr'r!rrrr$r%r&r'rr#r)r'r*r'r(r'rr)/ __future__rrcollections.abcr functoolsrnumpyrMr9rmonai.networks.blocksrrrr r r monai.networks.layers.factoriesr monai.utilsr rrrg__all__rModulerrSrrrrrrrrrrrrrrrrrJrrrrsf        E t/ sZrFPm" Q > Kt