o - iI@sddlmZddlmZddlZddlmZddlmZddlm Z m Z m Z ddl m Z Gdd d ejZd d ZGd d d ejZdS)) annotations)SequenceN)nn) Convolution)get_down_block get_mid_blockget_timestep_embedding)ensure_tuple_repcs*eZdZdZd fdd Zd d ZZS) ControlNetConditioningEmbeddingzA Network to encode the conditioning into a latent space. spatial_dimsint in_channels out_channelschannels Sequence[int]c stt|||ddddddd|_tg|_tt|dD],}||}||d}|j t|||dddddd|j t|||ddddddq"t t||d|dddd d |_ dS) NrAZSWISH)r r rstrides kernel_sizepadding adn_orderingactTr r rrrr conv_only) super__init__rconv_inr ModuleListblocksrangelenappend zero_moduleconv_out)selfr r rriZ channel_inZ channel_out __class__`/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/controlnet.pyr1sd      z(ControlNetConditioningEmbedding.__init__cCs,||}|jD]}||}q||}|SN)rr!r&)r'Z conditioning embeddingblockr+r+r,forwardjs    z'ControlNetConditioningEmbedding.forward)r r r r rr rr)__name__ __module__ __qualname____doc__rr0 __classcell__r+r+r)r,r ,s9r cCs|D]}tj|q|Sr-) parametersrinitzeros_)modulepr+r+r,r%us r%csdeZdZdZ           dd?d6d7Zd@dAd:d;ZZS)B ControlNeta 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. 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. 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). rrrr @r?r?FFTTr>ư>FrNr>`Tr r r num_res_blocksSequence[int] | intrrattention_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_channels include_fcuse_combined_linearuse_flash_attentionreturnNonecst| dur| durtd| dur| durtdtfdd|Dr2td|dt|t|krDtd |d |t| trPt| t|} t| t|krctd |d |d t|trot|t|}t|t|krtd |d|d||_||_ ||_ ||_ | |_ | |_ t|||dddddd|_|dd}tt|d|tt|||_| |_| durt| ||_t||||dd|_tg|_tg|_|d}t|||ddddd}t|j}|j|t t|D]}|}||}|t|dk}t!d*id|d|d|d|d||dd|d| d|d||o:| d ||oB| d!| |d"| d#| d$|d%|d&|d'|}|j|t ||D]}t|||ddddd}t|}|j|qh|st|||ddddd}t|}|j|q|d(}t"||||| | d(| | ||||d) |_#t|||ddddd}t|}||_$dS)+NTzControlNet expects dimension of the cross-attention conditioning (cross_attention_dim) to be specified when with_conditioning=True.FzRControlNet expects with_conditioning=True when specifying the cross_attention_dim.c3s|] }|dkVqdS)rNr+).0 out_channelrKr+r, sz&ControlNet.__init__..zVControlNet expects all channels to be a multiple of norm_num_groups, but got channels=z and norm_num_groups=zZControlNet expects channels to have the same length as attention_levels, but got channels=z and attention_levels=zTnum_head_channels should have the same length as attention_levels, but got channels=zq . For the i levels without attention, i.e. `attention_level[i]=False`, the num_head_channels[i] will be ignored.z`num_res_blocks` should be a single integer or a tuple of integers with the same length as `num_channels`, but got num_res_blocks=z and channels=.rrrr)r r rrr r r temb_channelsrGrKrLZadd_downsamplerNZ with_attnZwith_cross_attnrPrSrTrWrZr[r\r) r r rerKrLrRrPrSrTrWrZr[r\r+)%rr ValueErroranyr# isinstancer r r block_out_channelsrGrIrPrRrrr SequentialLinearSiLU time_embedrV Embeddingclass_embeddingr controlnet_cond_embeddingr down_blockscontrolnet_down_blocksr%convr$r"rr middle_blockcontrolnet_mid_block)r'r r rGrrIrKrLrNrPrRrSrTrVrWrXrYrZr[r\Ztime_embed_dimoutput_channelcontrolnet_blockr( input_channelis_final_blockZ down_block_Zmid_block_channelr)rar,rsf                zControlNet.__init__?x torch.Tensor timestepscontrolnet_condconditioning_scalecontexttorch.Tensor | None class_labels'tuple[list[torch.Tensor], torch.Tensor]cs<t||jd}|j|jd}||}|jdur1|dur!td||} | j|jd} || }||} | |}| |7} |durL|j durLtd| g} |j D]} | | ||d\} } | D]}| |q_qR|j | ||d} g}t| |jD] \}}||}| |qx|} || }fdd | D} |9}| |fS) a Args: x: input tensor (N, C, H, W, [D]). timesteps: timestep tensor (N,). controlnet_cond: controlnet conditioning tensor (N, C, H, W, [D]) conditioning_scale: conditioning scale. context: context tensor (N, 1, cross_attention_dim), where cross_attention_dim is specified in the model init. class_labels: context tensor (N, ). r)dtypeNz9class_labels should be provided when num_class_embeds > 0FzAmodel should have with_conditioning = True if context is provided)Z hidden_statesZtembrcsg|]}|qSr+r+)r_hrr+r, sz&ControlNet.forward..)rritorrmrVrfrorrprRrqr$rtziprrru)r'r|r~rrrrZt_embZembZ class_embrZdown_block_res_samplesZdownsample_blockZ res_samplesresidualZ!controlnet_down_block_res_samplesZdown_block_res_samplerwZmid_block_res_sampler+rr,r0bs<         zControlNet.forwardold_state_dictdictcs|tfdd|Drtd||dS|rBD]}||vr,td|dqtd|D]}|vrAtd|dq3D] }||vrQ|||<qDd d D}|D]}||d |d <||d |d<q[|rtd||dS)z Load a state dict from a ControlNet trained with [MONAI Generative](https://github.com/Project-MONAI/GenerativeModels). Args: old_state_dict: state dict from the old ControlNet model. c3s|]}|vVqdSr-r+r_knew_state_dictr+r,rbsz1ControlNet.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)zout_proj.weight.out_proj.weight)replacerr+r+r,rs z2ControlNet.load_old_state_dict..z.to_out.0.weightrz.to_out.0.biasz.out_proj.biasz!remaining keys in old_state_dict:) state_dictallprintload_state_dictpopkeys)r'rverboserattention_blocksr/r+rr,load_old_state_dicts4  zControlNet.load_old_state_dict)r<r=r@r>rAFrBFrNNFrrCTFF)(r r r r rGrHrrrIrJrKr rLrMrNrOrPrQrRrOrSr rTrUrVrUrWrOrXr rYrrZrOr[rOr\rOr]r^)r{NN)r|r}r~r}rr}rrMrrrrr]r)F)rrr]r^)r1r2r3r4rr0rr5r+r+r)r,r;{s4 P Gr;) __future__rcollections.abcrtorchrmonai.networks.blocksrZ(monai.networks.nets.diffusion_model_unetrrr monai.utilsr Moduler r%r;r+r+r+r,s     I