o  iJ@sddlmZddlmZddlZddlmZddlmZddlm Z m Z m Z m Z m Z ddlmZddlmZd gZGd d d ejZdS) ) annotations)SequenceN)nn) Convolution)get_down_block get_mid_blockget_timestep_embedding get_up_block zero_module)ensure_tuple_rep)convert_to_tensorDiffusionModelUNetMaisicseZdZdZ               dFdGfd,d- Zd.d/Zd0d1Zd2d3Zd4d5Zd6d7Z dHdIdDdEZ Z S)Jr a, U-Net 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. Can be a single integer or a sequence of integers. num_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. Can be a single integer or a sequence of integers. 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. include_fc: whether to include the final linear layer. Default to False. use_combined_linear: whether to use a single linear layer for qkv projection, default to False. use_flash_attention: If True, use flash attention for a memory efficient attention mechanism. dropout_cattn: If different from zero, this will be the dropout value for the cross-attention layers. include_top_region_index_input: If True, use top region index input. include_bottom_region_index_input: If True, use bottom region index input. include_spacing_input: If True, use spacing input. rrr @rrFFTTrư>FN spatial_dimsint in_channels out_channelsnum_res_blocksSequence[int] | int num_channels Sequence[int]attention_levelsSequence[bool]norm_num_groupsnorm_epsfloatresblock_updownboolnum_head_channelsint | Sequence[int]with_conditioningtransformer_num_layerscross_attention_dim int | Nonenum_class_embedsupcast_attention include_fcuse_combined_linearuse_flash_attention dropout_cattninclude_top_region_index_input!include_bottom_region_index_inputinclude_spacing_inputreturnNonec$st| dur| durtd| dur| durtd|dks%|dkr)tdtfdd |Dr>td |d t|t|krTtd t|d t|t| tr`t| t|} t| t|krltdt|trxt|t|}t|t|krtd|durtj std||_ ||_ ||_ ||_||_| |_| |_t|||dddddd|_|dd}||d||_||_|durt|||_||_||_||_|}|jr|d||_||7}|jr|d||_||7}|jr|d||_||7}t g|_!|d}t"t|D]e}|}||}|t|dk}t#d.id|d|d|d|d||dd|d| d| d||oP| d ||oX| d!| |d"| d#| d$|d%|d&|d'|d(|}|j!$|qt%||d)||| | d)| | |||||d*|_&t g|_'t(t)|}t(t)|}t(t)|} t(t)| }!|d}t"t|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'$|#qt,tj-|d|dd-t.t/t||d|ddddd|_0dS)/NTzDiffusionModelUNetMaisi expects dimension of the cross-attention conditioning (cross_attention_dim) when using with_conditioning.Fz_DiffusionModelUNetMaisi expects with_conditioning=True when specifying the cross_attention_dim.g?rz#Dropout cannot be negative or >1.0!c3s|] }|dkVqdS)rN).0 out_channelr"r8/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/generation/maisi/networks/diffusion_model_unet_maisi.py {sz3DiffusionModelUNetMaisi.__init__..zjDiffusionModelUNetMaisi expects all num_channels being multiple of norm_num_groups, but get num_channels: z and norm_num_groups: zhDiffusionModelUNetMaisi expects num_channels being same size of attention_levels, but get num_channels: z and attention_levels: znum_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`.zatorch.cuda.is_available() should be True but is False. Flash attention is only available for GPU.rr)rrrstrides kernel_sizepadding conv_onlyrrr temb_channelsrr"r#add_downsampler% with_attnwith_cross_attnr'r*r+r.r/r0r1r2)rrrDr"r#r)r'r*r+r.r/r0r1r2prev_output_channel add_upsample) num_groupsrepsaffiner8)1super__init__ ValueErroranylen isinstancerr torchcuda is_availablerblock_out_channelsrrr r'r)rconv_in_create_embedding_module time_embedr-r Embeddingclass_embeddingr3r4r5top_region_index_layerbottom_region_index_layer spacing_layer ModuleList down_blocksrangerappendr middle_block up_blockslistreversedminr Sequential GroupNormSiLUr out)$selfrrrrrr r"r#r%r'r)r*r+r-r.r/r0r1r2r3r4r5time_embed_dimZnew_time_embed_dimoutput_channeli input_channelis_final_block down_blockreversed_block_out_channelsreversed_num_res_blocksreversed_attention_levelsreversed_num_head_channelsrIup_block __class__r;r<rOTs                      z DiffusionModelUNetMaisi.__init__cCs&tt||tt||}|S)N)rriLinearrk)rm input_dim embed_dimmodelr8r8r<rY2s"z0DiffusionModelUNetMaisi._create_embedding_modulecCsft||jd}|j|jd}||}|jdur1|dur!td||}|j|jd}||7}|S)Nr)dtypez9class_labels should be provided when num_class_embeds > 0)rrWtorrZr-rPr\)rmx timesteps class_labelst_embemb class_embr8r8r<_get_time_and_class_embedding6s   z5DiffusionModelUNetMaisi._get_time_and_class_embeddingcCsj|jr||}tj||fdd}|jr"||}tj||fdd}|jr3||}tj||fdd}|S)Nr)dim)r3r]rTcatr4r^r5r_)rmrZ top_index bottom_indexspacingZ_embr8r8r<_get_input_embeddingsGs   z-DiffusionModelUNetMaisi._get_input_embeddingsc Cs|dur |jdur td|g}|jD]}||||d\}}||q|dur?g}t||D] \} } | | 7} || q/|}||fS)NFzAmodel should have with_conditioning = True if context is provided) hidden_statestembcontext)r)rPraextendziprc) rmhrrdown_block_additional_residualsdown_block_res_samplesdownsample_block res_samplesnew_down_block_res_samplesdown_block_res_sampledown_block_additional_residualr8r8r<_apply_down_blocksSs    z*DiffusionModelUNetMaisi._apply_down_blockscCsD|jD]}t|j }||d}|d|}|||||d}q|S)N)rres_hidden_states_listrr)rerRresnets)rmrrrrupsample_blockidxrr8r8r<_apply_up_blocksgs    z(DiffusionModelUNetMaisi._apply_up_blocksr torch.Tensorrrtorch.Tensor | Nonerrtuple[torch.Tensor] | Nonemid_block_additional_residualtop_region_index_tensorbottom_region_index_tensorspacing_tensorc Cs||||} || ||| } ||} || | ||\} } || | |} |dur-| |7} || | || } || } t| } | S)a{ Forward pass through the UNet model. Args: x: Input tensor of shape (N, C, SpatialDims). timesteps: Timestep tensor of shape (N,). context: Context tensor of shape (N, 1, ContextDim). class_labels: Class labels tensor of shape (N,). down_block_additional_residuals: Additional residual tensors for down blocks of shape (N, C, FeatureMapsDims). mid_block_additional_residual: Additional residual tensor for mid block of shape (N, C, FeatureMapsDims). top_region_index_tensor: Tensor representing top region index of shape (N, 4). bottom_region_index_tensor: Tensor representing bottom region index of shape (N, 4). spacing_tensor: Tensor representing spacing of shape (N, 3). Returns: A tensor representing the output of the UNet model. N)rrrXrrdrrlr )rmrrrrrrrrrrrZ_updated_down_block_res_samplesZh_tensorr8r8r<forwardps  zDiffusionModelUNetMaisi.forward)rrrrrFrFrNNFFFFrFFF).rrrrrrrrrrr r!r"rr#r$r%r&r'r(r)r&r*rr+r,r-r,r.r&r/r&r0r&r1r&r2r$r3r&r4r&r5r&r6r7)NNNNNNN)rrrrrrrrrrrrrrrrrrr6r) __name__ __module__ __qualname____doc__rOrYrrrrr __classcell__r8r8ryr<r 5sH#_  ) __future__rcollections.abcrrTrmonai.networks.blocksr(monai.networks.nets.diffusion_model_unetrrrr r monai.utilsr monai.utils.type_conversionr __all__Moduler r8r8r8r<s