o  i@sLddlmZddlmZddlZddlmZddlmZGdddeZ dS)) annotations)SequenceN) ControlNet)get_timestep_embeddingcseZdZdZ            dCdDfd+d, Z - dEdFd7d8Zd9d:Zd;d<Zd=d>Zd?d@Z dAdBZ Z S)GControlNetMaisia Control network for diffusion models based on Zhang and Agrawala "Adding Conditional Control to Text-to-Image Diffusion Models" (https://arxiv.org/abs/2302.05543) Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. num_res_blocks: number of residual blocks (see ResnetBlock) per level. 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. 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. conditioning_embedding_in_channels: number of input channels for the conditioning embedding. conditioning_embedding_num_channels: number of channels for the blocks in the conditioning embedding. use_checkpointing: if True, use activation checkpointing to save memory. 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. rrr @r r FFTTr ư>FNr `T spatial_dimsint in_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"conditioning_embedding_in_channels#conditioning_embedding_num_channelsuse_checkpointing include_fcuse_combined_linearuse_flash_attentionreturnNonecs:t||||||||| | | | | ||||||||_dSN)super__init__r,)selfrrrrrrrr r"r$r%r&r(r)r*r+r,r-r.r/ __class__w/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/generation/maisi/networks/controlnet_maisi.pyr43s, zControlNetMaisi.__init__?x torch.Tensor timestepscontrolnet_condconditioning_scalecontexttorch.Tensor | None class_labels'tuple[list[torch.Tensor], torch.Tensor]c s||||}||}|jrtjjj|j|dd}n||}||7}||||\} }||||}| || \} } fdd| D} | 9} | | fS)NF) use_reentrantcsg|]}|qSr8r8).0hr?r8r9 wsz+ControlNetMaisi.forward..) !_prepare_time_and_class_embedding_apply_initial_convolutionr,torchutils checkpointcontrolnet_cond_embedding_apply_down_blocks_apply_mid_block_apply_controlnet_blocks) r5r;r=r>r?r@rBembrFdown_block_res_samplesmid_block_res_sampler8rGr9forwardas  zControlNetMaisi.forwardcCsft||jd}|j|jd}||}|jdur1|dur!td||}|j|jd}||}|S)Nr)dtypez9class_labels should be provided when num_class_embeds > 0)rblock_out_channelstorV time_embedr( ValueErrorclass_embedding)r5r;r=rBt_embrR class_embr8r8r9rI|s   z1ControlNetMaisi._prepare_time_and_class_embeddingcCs||}|Sr2)conv_in)r5r;rFr8r8r9rJs z*ControlNetMaisi._apply_initial_convolutioncCsZ|dur |jdur td|g}|jD]}||||d\}}|D]}||q q||fS)NFzAmodel should have with_conditioning = True if context is provided hidden_statestembr@)r$rZ down_blocksappend)r5rRr@rFrSdownsample_block res_samplesresidualr8r8r9rOs  z"ControlNetMaisi._apply_down_blockscCs|j|||d}|S)Nr_) middle_block)r5rRr@rFr8r8r9rPsz ControlNetMaisi._apply_mid_blockcCs>g}t||jD] \}}||}||q||}||fSr2)zipcontrolnet_down_blocksrccontrolnet_mid_block)r5rFrS!controlnet_down_block_res_samplesdown_block_res_samplecontrolnet_blockrTr8r8r9rQs   z(ControlNetMaisi._apply_controlnet_blocks)rr r r r FrFrNNFrrTFFF)*rrrrrrrrrrrrrrr r!r"r#r$r!r%rr&r'r(r'r)r!r*rr+rr,r!r-r!r.r!r/r!r0r1)r:NN)r;r<r=r<r>r<r?rr@rArBrAr0rC) __name__ __module__ __qualname____doc__r4rUrIrJrOrPrQ __classcell__r8r8r6r9rs< 3  r) __future__rcollections.abcrrKZmonai.networks.nets.controlnetr(monai.networks.nets.diffusion_model_unetrrr8r8r8r9s