o - io@sddlmZddlmZddlZddlmZddlmZmZddl m Z ddl m Z m Z mZmZmZmZmZmZddlmZd gZGd d d ejZGd d d ejZGdddejZGdddejZ     d3d4d0d1ZGd2d d ejZdS)5) annotations)SequenceN)nn) ConvolutionSpatialAttentionBlock)SPADE)DiffusionUnetDownsampleDiffusionUNetResnetBlockSpatialTransformerWrappedUpsampleget_down_block get_mid_blockget_timestep_embedding zero_module)ensure_tuple_repSPADEDiffusionModelUNetcs:eZdZdZ      dd fdd Zd!ddZZS)"SPADEDiffResBlocka  Residual block with timestep conditioning and SPADE norm. Enables SPADE normalisation for semantic conditioning (Park et. al (2019): https://github.com/NVlabs/SPADE) Args: spatial_dims: The number of spatial dimensions. in_channels: number of input channels. temb_channels: number of timestep embedding channels. label_nc: number of semantic channels for SPADE normalisation. 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. spade_intermediate_channels: number of intermediate channels for SPADE block layer NF ư> spatial_dimsint in_channels temb_channelslabel_nc out_channels int | Noneupbooldownnorm_num_groupsnorm_epsfloatspade_intermediate_channelsreturnNonec sBt||_||_||_|p||_||_||_t||d|| dd| d|d|_ t |_ t |||jddddd|_d|_|_|jrQt|d||d d dd |_n |r[t||d d |_t ||j|_t||jd|| dd| d|d|_tt ||j|jddddd|_||j|krt |_dSt |||jddddd|_dS)NGROUPT) num_groupsepsaffine)rnorm_ncnorm norm_paramshidden_channels kernel_sizerrrrstridesr/padding conv_only nontrainablenearest@)rmoderr interp_mode scale_factor align_cornersF)use_convr)super__init__rchannels emb_channelsrrrrnorm1rSiLU nonlinearityrconv1upsample downsampler rLinear time_emb_projnorm2rconv2Identityskip_connection) selfrrrrrrrr r!r# __class__p/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/spade_diffusion_model_unet.pyr>Js           zSPADEDiffResBlock.__init__x torch.TensorembsegcCs|}|||}||}|jdur||}||}n|jdur,||}||}||}|jdkrI|||ddddddf}n|||dddddddf}||}|||}||}||}| ||}|S)N) rArCrErFrDrrHrIrJrL)rMrRrTrUhtemboutputrPrPrQforwards&          &&   zSPADEDiffResBlock.forward)NFFrrr)rrrrrrrrrrrrrrr rr!r"r#rr$r%)rRrSrTrSrUrSr$rS__name__ __module__ __qualname____doc__r>rZ __classcell__rPrPrNrQr8s]rcs>eZdZdZ      d&d'fdd Z d(d)d$d%ZZS)* SPADEUpBlocka Unet's up block containing resnet and upsamplers blocks. Enables SPADE normalisation for semantic conditioning (Park et. al (2019): https://github.com/NVlabs/SPADE) 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. label_nc: number of semantic channels for SPADE normalisation. 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. spade_intermediate_channels: number of intermediate channels for SPADE block layer. r0rrTFrrrrprev_output_channelrrrnum_res_blocksr r!r" add_upsamplerresblock_updownr#r$r%c  st| |_g} t|D]%}||dkr|n|}|dkr |n|}| t|||||||| | dqt| |_|| rh| rNt |||||| dd|_ dSt |||ddddd}t |d||d d |dd |_ dSd|_ dS) Nr0rrrrrrr r!r#Trrrrr r!rr*r1r5r6r7rr8rrr9r: post_convr;) r=r>rerangeappendrr ModuleListresnetsr upsamplerrr )rMrrrbrrrrcr r!rdrer#rmires_skip_channelsresnet_in_channelsrirNrPrQr>sf       zSPADEUpBlock.__init__N hidden_statesrSres_hidden_states_listlist[torch.Tensor]rXrUcontexttorch.Tensor | NonecCsZ~|jD]}|d}|dd}tj||gdd}||||}q|jdur+|||}|SNr0dim)rmtorchcatrn)rMrrrsrXrUruresnetres_hidden_statesrPrPrQrZs    zSPADEUpBlock.forward)r0rrTFr)rrrrrbrrrrrrrrcrr rr!r"rdrrerr#rr$r%N rrrSrsrtrXrSrUrSrurvr$rSr[rPrPrNrQrasOracsFeZdZdZ          d*d+fdd Z d,d-d(d)ZZS).SPADEAttnUpBlocka Unet's up block containing resnet, upsamplers, and self-attention blocks. Enables SPADE normalisation for semantic conditioning (Park et. al (2019): https://github.com/NVlabs/SPADE) 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. label_nc: number of semantic channels for SPADE normalisation 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. 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). r0rrTFrrrrrbrrrrcr r!r"rdrrenum_head_channelsr# include_fcuse_combined_linearuse_flash_attentionr$r%c st| |_g}g}t|D]4}||dkr|n|}|dkr"|n|}|t|||||||| | d|t||| || |||dqt||_ t||_ || r| ret |||||| dd|_ dSt |||ddddd}t|d ||d d |dd |_ dSd|_ dS) Nr0rrf)r num_channelsrr r!rrrTrgr*r1r5r6r7rh)r=r>rerjrkrrrrlrm attentionsr rnrr )rMrrrbrrrrcr r!rdrerr#rrrrmrrorprqrirNrPrQr>Ns        zSPADEAttnUpBlock.__init__NrrrSrsrtrXrUrurvc Csr~t|j|jD]#\}}|d}|dd}tj||gdd}||||}||}q|jdur7|||}|Srwziprmrr{r| contiguousrn rMrrrsrXrUrur}attnr~rPrPrQrZs    zSPADEAttnUpBlock.forward) r0rrTFr0rTFF)"rrrrrbrrrrrrrrcrr rr!r"rdrrerrrr#rrrrrrrr$r%rrr[rPrPrNrQr5s arcsNeZdZdZ             d.d/fd"d# Z  d0d1d,d-ZZS)2SPADECrossAttnUpBlocka Unet's up block containing resnet, upsamplers, and self-attention blocks. Enables SPADE normalisation for semantic conditioning (Park et. al (2019): https://github.com/NVlabs/SPADE) 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. label_nc: number of semantic channels for SPADE normalisation. 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. spade_intermediate_channels: number of intermediate channels for SPADE block layer. 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). r0rrTFNrrrrrbrrrrcr r!r"rdrrertransformer_num_layerscross_attention_dimrupcast_attentionr#rrrr$r%cst| |_g}g}t|D]:}||dkr|n|}|dkr"|n|}|t||||||| ||d|t|||| | || | |||||d qt||_ t||_ || r| rkt |||||| dd|_ dSt |||ddddd}t|d ||d d |dd |_ dSd|_ dS) Nr0r)rrrrr r!rr#) rrnum_attention_headsrr r! num_layersrrrrrTrgr*r1r5r6r7rh)r=r>rerjrkrr rrlrrmr rnrr )rMrrrbrrrrcr r!rdrerrrrr#rrrrmrrorprqrirNrPrQr>s         zSPADECrossAttnUpBlock.__init__rrrSrsrtrXrUrvruc Cstt|j|jD]%\}}|d}|dd}tj||gdd}||||}|||d}q|jdur8|||}|S)Nrxr0ry)rurrrPrPrQrZ<s    zSPADECrossAttnUpBlock.forward) r0rrTFr0r0NFrTFF)(rrrrrbrrrrrrrrcrr rr!r"rdrrerrrrrrrrrr#rrrrrrrr$r%)NN) rrrSrsrtrXrSrUrvrurvr$rSr[rPrPrNrQrs&!grFrTrrrrbrrrcr r!r"rdrre with_attnwith_cross_attnrrrrrrr#rrrr$ nn.ModulecCs| r7tdid|d|d|d|d|d|d|d|d |d |d | d | d |d|d|d|S| rqtdid|d|d|d|d|d|d|d|d |d |d | d | d| d|d|d |d|St||||||||||| |d S)Nrrrbrrrrcr r!rdrerr#rrrrrr) rrrbrrrrcr r!rdrer#rP)rrra)rrrbrrrcr r!rdrerrrrrrrr#rrrrPrPrQget_spade_up_blockRs          rcsXeZdZdZ           d9d:fd+d, Z d;dsz3SPADEDiffusionModelUNet.__init__..zRSPADEDiffusionModelUNet expects all num_channels being multiple of norm_num_groupszPSPADEDiffusionModelUNet expects num_channels being same size of attention_levelsznum_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`.rr0r*r1rrrrrcr r!add_downsamplererrrrrrrrrrx) rrrr r!rrrrrrrrrbrdrr#)r'rr(r)rP))r=r> ValueErroranylen isinstancerrrblock_out_channelsrrcrrrrrconv_inr SequentialrGrB time_embedr Embeddingclass_embeddingrl down_blocksrjr rkr middle_block up_blockslistreversedminr GroupNormrout)!rMrrrrrcr?rr r!rerrrrrrr#rrrtime_embed_dimoutput_channelro input_channelis_final_block down_blockreversed_block_out_channelsreversed_num_res_blocksreversed_attention_levelsreversed_num_head_channelsrbup_blockrNrrQr>sd                      z SPADEDiffusionModelUNet.__init__rRrS timestepsrUrurv class_labelsdown_block_additional_residualstuple[torch.Tensor] | Nonemid_block_additional_residualcCsft||jd}|j|jd}||} |jdur1|dur!td||} | j|jd} | | } ||} |durC|j durCtd| g} |j D]} | | | |d\} }|D]}| |qVqI|dur{| g}t | |D] \}}||}| |qk|} |j | | |d} |dur| |} |jD]}t|j }| |d}| d|} || ||| |d} q|| }|S) a Args: x: input tensor (N, C, SpatialDims). timesteps: timestep tensor (N,). seg: Bx[LABEL_NC]x[SPATIAL DIMENSIONS] tensor of segmentations for SPADE norm. 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). r)dtypeNz9class_labels should be provided when num_class_embeds > 0FzAmodel should have with_conditioning = True if context is provided)rrrXru)rrrsrUrXru)rrtorrrrrrrrrkrrrrrmr)rMrRrrUrurrrt_embrT class_embrWdown_block_res_samplesdownsample_block res_samplesresidualnew_down_block_res_samplesdown_block_res_sampledown_block_additional_residualupsample_blockidxrYrPrPrQrZsH             zSPADEDiffusionModelUNet.forward)rrrrrFrFr0NNFrTFF)*rrrrrrrrrcrr?rrrr rr!r"rerrrrrrrrrrrrrr#rrrrrrrr$r%)NNNN)rRrSrrSrUrSrurvrrvrrrrvr$rSr[rPrPrNrQrs2"G)FrTFF),rrrrrbrrrrrrcrr rr!r"rdrrerrrrrrrrrrrrrrrr#rrrrrrrr$r) __future__rcollections.abcrr{rmonai.networks.blocksrrZ monai.networks.blocks.spade_normr(monai.networks.nets.diffusion_model_unetrr r r r r rr monai.utilsr__all__ModulerrarrrrrPrPrPrQs.    (  s $ O