o %ݫi@sdZddlZddlmZddlZddlmZddlmmZ ddl m Z ddl m Z d1dd Zd d Zd d Zd2ddZGdddejZGdddejZGdddejeZGdddejZGdddejZGdddeZGdddejZGdd d ejZd3d!d"ZGd#d$d$ejZGd%d&d&ejZGd'd(d(ejZGd)d*d*ejZ d+d,gZ!Gd-d.d.ejZ"Gd/d0d0e Z#dS)4aA UNet model implementation for use with diffusion models Adapted from OpenAI guided diffusion, with slight modifications and additional features https://github.com/openai/guided-diffusion MIT License Copyright (c) 2021 OpenAI Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. Authors * Artem Ploujnikov 2022 N)abstractmethod) pad_divisible)NormalizingAutoencoderTcCs"|r|D]}|q|S)a* Zero out the parameters of a module and return it. Arguments --------- module: torch.nn.Module a module use_fixup_init: bool whether to zero out the parameters. If set to false, the function is a no-op Returns ------- The fixed module ) parametersdetachzero_)moduleuse_fixup_initpr I/home/ubuntu/.local/lib/python3.10/site-packages/speechbrain/nnet/unet.pyfixup-s rcOV|dkr tj|i|S|dkrtj|i|S|dkr$tj|i|Std|)a Create a 1D, 2D, or 3D convolution module. Arguments --------- dims: int The number of dimensions *args: tuple **kwargs: dict Any remaining arguments are passed to the constructor Returns ------- The constructed Conv layer runsupported dimensions: )nnConv1dConv2dConv3d ValueErrordimsargskwargsr r r conv_ndCsrcOr)z8 Create a 1D, 2D, or 3D average pooling module. rrrr)r AvgPool1d AvgPool2d AvgPool3drrr r r avg_pool_nd\sr 'cCs|d}tt| tjd|tjd|j|jd}|dddf|d}tj t |t |gdd}|drRtj |t |ddddfgdd}|S) a Create sinusoidal timestep embeddings. Arguments --------- timesteps: torch.Tensor a 1-D Tensor of N indices, one per batch element. These may be fractional. dim: int the dimension of the output. max_period: int controls the minimum frequency of the embeddings. Returns ------- result: torch.Tensor an [N x dim] Tensor of positional embeddings. rr)startenddtype)deviceNdimr) torchexpmathlogarangefloat32tor%floatcatcossin zeros_like) timestepsr( max_periodhalffreqsr embeddingr r r timestep_embeddingis" r:c s>eZdZdZ d dedededeffdd Zd d ZZS) AttentionPool2daTwo-dimensional attentional pooling Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py Arguments --------- spatial_dim: int the size of the spatial dimension embed_dim: int the embedding dimension num_heads_channels: int the number of attention heads output_dim: int the output dimension Example ------- >>> attn_pool = AttentionPool2d( ... spatial_dim=64, ... embed_dim=16, ... num_heads_channels=2, ... output_dim=4 ... ) >>> x = torch.randn(4, 1, 64, 64) >>> x_pool = attn_pool(x) >>> x_pool.shape torch.Size([4, 4]) N spatial_dim embed_dimnum_heads_channels output_dimcspttt||dd|d|_td|d|d|_td||p'|d|_ |||_ t |j |_ dS)Nrrg?r) super__init__r Parameterr)randnpositional_embeddingrqkv_projc_proj num_heads QKVAttention attention)selfr<r=r>r? __class__r r rAs  zAttentionPool2d.__init__cCs|j^}}}|||d}tj|jddd|gdd}||jdddddf|j}||}| |}| |}|dddddfS)zComputes the attention forward pass Arguments --------- x: torch.Tensor the tensor to be attended to Returns ------- result: torch.Tensor the attention output r&T)r(keepdimr'Nr) shapereshaper)r1meanrDr/r$rErIrF)rJxbc_spatialr r r forwards $   zAttentionPool2d.forwardN)__name__ __module__ __qualname____doc__intrArU __classcell__r r rKr r;s"r;c@seZdZdZedddZdS) TimestepBlockzT Any module where forward() takes timestep embeddings as a second argument. NcCsdS)z Apply the module to `x` given `emb` timestep embeddings. Arguments --------- x: torch.Tensor the data tensor emb: torch.Tensor the embedding tensor Nr )rJrQembr r r rUszTimestepBlock.forwardrV)rWrXrYrZrrUr r r r r]sr]c@seZdZdZdddZdS)TimestepEmbedSequentialaA sequential module that passes timestep embeddings to the children that support it as an extra input. Example ------- >>> from speechbrain.nnet.linear import Linear >>> class MyBlock(TimestepBlock): ... def __init__(self, input_size, output_size, emb_size): ... super().__init__() ... self.lin = Linear( ... n_neurons=output_size, ... input_size=input_size ... ) ... self.emb_proj = Linear( ... n_neurons=output_size, ... input_size=emb_size, ... ) ... def forward(self, x, emb): ... return self.lin(x) + self.emb_proj(emb) >>> tes = TimestepEmbedSequential( ... MyBlock(128, 64, 16), ... Linear( ... n_neurons=32, ... input_size=64 ... ) ... ) >>> x = torch.randn(4, 10, 128) >>> emb = torch.randn(4, 10, 16) >>> out = tes(x, emb) >>> out.shape torch.Size([4, 10, 32]) NcCs,|D]}t|tr|||}q||}q|S)aComputes a sequential pass with sequential embeddings where applicable Arguments --------- x: torch.Tensor the data tensor emb: torch.Tensor timestep embeddings Returns ------- The processed input ) isinstancer])rJrQr^layerr r r rUs    zTimestepEmbedSequential.forwardrV)rWrXrYrZrUr r r r r_s!r_c*eZdZdZdfdd ZddZZS) Upsamplea An upsampling layer with an optional convolution. Arguments --------- channels: torch.Tensor channels in the inputs and outputs. use_conv: bool a bool determining if a convolution is applied. dims: int determines if the signal is 1D, 2D, or 3D. If 3D, then upsampling occurs in the inner-two dimensions. out_channels: int Number of output channels. If None, same as input channels. Example ------- >>> ups = Upsample(channels=4, use_conv=True, dims=2, out_channels=8) >>> x = torch.randn(8, 4, 32, 32) >>> x_up = ups(x) >>> x_up.shape torch.Size([8, 8, 64, 64]) rNcsJt||_|p ||_||_||_|r#t||j|jddd|_dSdS)Nrrpadding)r@rAchannels out_channelsuse_convrrconv)rJrfrhrrgrKr r rA3s   zUpsample.__init__cCst|jd|jks J|jdkr(tj||jd|jdd|jddfdd}ntj|ddd}|jr8||}|S)zComputes the upsampling pass Arguments --------- x: torch.Tensor layer inputs Returns ------- result: torch.Tensor upsampled outputsrrrnearest)mode) scale_factorrl)rNrfrF interpolaterhrirJrQr r r rU>s & zUpsample.forwardrNrWrXrYrZrArUr\r r rKr rcs rccrb) Downsamplea A downsampling layer with an optional convolution. Arguments --------- channels: int channels in the inputs and outputs. use_conv: bool a bool determining if a convolution is applied. dims: int determines if the signal is 1D, 2D, or 3D. If 3D, then downsampling occurs in the inner-two dimensions. out_channels: int Number of output channels. If None, same as input channels. Example ------- >>> ups = Downsample(channels=4, use_conv=True, dims=2, out_channels=8) >>> x = torch.randn(8, 4, 32, 32) >>> x_up = ups(x) >>> x_up.shape torch.Size([8, 8, 16, 16]) rNcs|t||_|p ||_||_||_|dkrdnd}|r,t||j|jd|dd|_dS|j|jks4Jt|||d|_dS)Nrr)rrrr)stridere) kernel_sizert) r@rArfrgrhrropr )rJrfrhrrgrtrKr r rAos"    zDownsample.__init__cCs|jd|jks J||S)zComputes the downsampling pass Arguments --------- x: torch.Tensor layer inputs Returns ------- result: torch.Tensor downsampled outputs r)rNrfrvrpr r r rUs zDownsample.forwardrqrrr r rKr rsVsrscs:eZdZdZ       d fdd Zd d d ZZS) ResBlocka A residual block that can optionally change the number of channels. Arguments --------- channels: int the number of input channels. emb_channels: int the number of timestep embedding channels. dropout: float the rate of dropout. out_channels: int if specified, the number of out channels. use_conv: bool if True and out_channels is specified, use a spatial convolution instead of a smaller 1x1 convolution to change the channels in the skip connection. dims: int determines if the signal is 1D, 2D, or 3D. up: bool if True, use this block for upsampling. down: bool if True, use this block for downsampling. norm_num_groups: int the number of groups for group normalization use_fixup_init: bool whether to use FixUp initialization Example ------- >>> res = ResBlock( ... channels=4, ... emb_channels=8, ... dropout=0.1, ... norm_num_groups=2, ... use_conv=True, ... ) >>> x = torch.randn(2, 4, 32, 32) >>> emb = torch.randn(2, 8) >>> res_out = res(x, emb) >>> res_out.shape torch.Size([2, 4, 32, 32]) NFr Tc spt||_||_||_|p||_||_tt | |t t |||jddd|_ |p/||_ |rBt|d||_t|d||_n|rSt|d||_t|d||_nt|_|_|durntt t||j|_nd|_tt | |jt tj|dtt ||j|jddd| d|_|j|krt|_dS|rt |||jddd|_dSt |||jd|_dS)NrrrdF)r r )r@rArf emb_channelsdropoutrgrhr Sequential GroupNormSiLUr in_layersupdownrch_updx_updrsIdentityLinear emb_layersDropoutr out_layersskip_connection) rJrfrzr{rgrhrupdownnorm_num_groupsr rKr r rAs\         zResBlock.__init__cCs|jr#|jdd|jd}}||}||}||}||}n||}|durN|||j}t|jt|jkrM|d}t|jt|jks?nt |}||}| |}| ||S)a Apply the block to a torch.Tensor, conditioned on a timestep embedding. Arguments --------- x: torch.Tensor an [N x C x ...] Tensor of features. emb: torch.Tensor an [N x emb_channels] Tensor of timestep embeddings. Returns ------- result: torch.Tensor an [N x C x ...] Tensor of outputs. Nr&).N) rrrrrtyper$lenrNr)r4rr)rJrQr^in_restin_convhemb_outr r r rUs"      zResBlock.forward)NFrFFrxTrVrrr r rKr rws1Drwcs2eZdZdZ    d fdd Zdd ZZS) AttentionBlocka An attention block that allows spatial positions to attend to each other. Originally ported from here, but adapted to the N-d case. https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66. Arguments --------- channels: int the number of channels num_heads: int the number of attention heads num_head_channels: int the number of channels in each attention head norm_num_groups: int the number of groups used for group normalization use_fixup_init: bool whether to use FixUp initialization Example ------- >>> attn = AttentionBlock( ... channels=8, ... num_heads=4, ... num_head_channels=4, ... norm_num_groups=2 ... ) >>> x = torch.randn(4, 8, 16, 16) >>> out = attn(x) >>> out.shape torch.Size([4, 8, 16, 16]) rr&rxTcst||_|dkr||_n||dks Jd|d||||_t|||_td||dd|_t |j|_ t td||d||_ dS)Nr&rzq,k,v channels z' is not divisible by num_head_channels rr) r@rArfrGrr}normrqkvrHrIrproj_out)rJrfrGnum_head_channelsrr rKr r rAKs   zAttentionBlock.__init__cCsV|j^}}}|||d}|||}||}||}||j||g|RS)zCompletes the forward pass Arguments --------- x: torch.Tensor the data to be attended to Returns ------- result: torch.Tensor The data, with attention applied r&)rNrOrrrIr)rJrQrRrSspatialrrr r r rUbs   zAttentionBlock.forward)rr&rxTrrr r rKr r)s$rc(eZdZdZfddZddZZS)rHa A module which performs QKV attention and splits in a different order. Arguments --------- n_heads : int Number of attention heads. Example ------- >>> attn = QKVAttention(4) >>> n = 4 >>> c = 8 >>> h = 64 >>> w = 16 >>> qkv = torch.randn(4, (3 * h * c), w) >>> out = attn(qkv) >>> out.shape torch.Size([4, 512, 16]) cst||_dSrV)r@rAn_heads)rJrrKr r rAs  zQKVAttention.__init__c Cs|j\}}}|d|jdksJ|d|j}|jddd\}}}dtt|} td|| ||j|||| ||j||} tj| dd | j } td| | ||j||} | |d|S)aApply QKV attention. Arguments --------- qkv: torch.Tensor an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs. Returns ------- result: torch.Tensor an [N x (H * C) x T] tensor after attention. rrrr'z bct,bcs->btsr&z bts,bcs->bct) rNrchunkr+sqrtr)einsumviewsoftmaxr0rr$rO) rJrbswidthlengthchqkvscaleweightar r r rUs zQKVAttention.forwardrrr r rKr rHws rHcCs`i}|dur+|D] \}}|dus||r*d|vr |||<q t|d|d||<q t|S)aBuilds a dictionary of embedding modules for embedding projections Arguments --------- emb_config: dict a configuration dictionary proj_dim: int the target projection dimension use_emb: dict an optional dictionary of "switches" to turn embeddings on and off Returns ------- result: torch.nn.ModuleDict a ModuleDict with a module for each embedding Nemb_projemb_dim)rproj_dim)itemsgetEmbeddingProjectionr ModuleDict) emb_configruse_embrkey item_configr r r build_emb_projs   rcsXeZdZdZ           dfd d Zdd dZ    dddZZS) UNetModela The full UNet model with attention and timestep embedding. Arguments --------- in_channels: int channels in the input torch.Tensor. model_channels: int base channel count for the model. out_channels: int channels in the output torch.Tensor. num_res_blocks: int number of residual blocks per downsample. attention_resolutions: int a collection of downsample rates at which attention will take place. May be a set, list, or tuple. For example, if this contains 4, then at 4x downsampling, attention will be used. dropout: float the dropout probability. channel_mult: int channel multiplier for each level of the UNet. conv_resample: bool if True, use learned convolutions for upsampling and downsampling dims: int determines if the signal is 1D, 2D, or 3D. emb_dim: int time embedding dimension (defaults to model_channels * 4) cond_emb: dict embeddings on which the model will be conditioned Example: { "speaker": { "emb_dim": 256 }, "label": { "emb_dim": 12 } } use_cond_emb: dict a dictionary with keys corresponding to keys in cond_emb and values corresponding to Booleans that turn embeddings on and off. This is useful in combination with hparams files to turn embeddings on and off with simple switches Example: {"speaker": False, "label": True} num_heads: int the number of attention heads in each attention layer. num_head_channels: int if specified, ignore num_heads and instead use a fixed channel width per attention head. num_heads_upsample: int works with num_heads to set a different number of heads for upsampling. Deprecated. norm_num_groups: int Number of groups in the norm, default 32 resblock_updown: bool use residual blocks for up/downsampling. use_fixup_init: bool whether to use FixUp initialization Example ------- >>> model = UNetModel( ... in_channels=3, ... model_channels=32, ... out_channels=1, ... num_res_blocks=1, ... attention_resolutions=[1] ... ) >>> x = torch.randn(4, 3, 16, 32) >>> ts = torch.tensor([10, 100, 50, 25]) >>> out = model(x, ts) >>> out.shape torch.Size([4, 1, 16, 32]) rrrrjTrNrr&rxFcsjt|dkr | }||_||_||_||_||_||_||_||_ t j |_ | |_ ||_||_| |_| |_| dur>|d} t|| |_t| | | d|_t|d|}}ttt| ||dddg|_||_|g}d}t|D]}\}}t|D]<}t|| |t||| ||dg}t||}||vr| t!|| |||d |j t||j|7_| |qz|t"|dkr|}|j t|rt|| ||| d ||d nt#||| |d |}| ||d 9}|j|7_qrtt|| || ||dt!|| |||d t|| || ||d|_$|j|7_tg|_%t&t|dddD]s\}}t|dD]g}|'}t||| |t||| ||dg}t||}||vr_| t!|||||d |r||kr|}| |rzt|| ||| d ||dnt(||| |d |d }|j% t||j|7_q0q&t)t*||t+t,t| ||ddd|d|_-dS)Nr&rj)rrrrrrrdrgrrr rGrrr Trgrrrr rrgrrrr rgrrrr ry).r@rA in_channelsmodel_channelsrgnum_res_blocksattention_resolutionsr{ channel_mult conv_resampler)r.r$rGrnum_heads_upsamplecond_emb use_cond_embr time_embedr cond_emb_projr[r ModuleListr_r input_blocks _feature_size enumeraterangerwappendrrrs middle_block output_blockslistpoprcr|r}r~rout)rJrrrgrrr{rrrrrrrGrrrresblock_updownr rinput_chinput_block_chansdslevelmult_layersout_chiichrKr r rA sJ                .  zUNetModel.__init__c Csg}|t||j}|dur#|D]\}}|j||}||7}q||j} |jD] } | | |} || q,| | |} |j D]} t j | | gdd} | | |} qB| |j} || S)aApply the model to an input batch. Arguments --------- x: torch.Tensor an [N x C x ...] Tensor of inputs. timesteps: torch.Tensor a 1-D batch of timesteps. cond_emb: dict a string -> tensor dictionary of conditional embeddings (multiple embeddings are supported) Returns ------- result: torch.Tensor an [N x C x ...] Tensor of outputs. Nrr')rr:rrrrr$rrrrr)r1rr) rJrQr5rhsr^rvaluerrr r r r rUs$           zUNetModel.forwardcCs||||dS)zForward function suitable for wrapping by diffusion. For this model, `length`/`out_mask_value`/`latent_mask_value` are unused and discarded. See :meth:`~UNetModel.forward` for details.)rr )rJrQr5rrout_mask_valuelatent_mask_valuer r r diffusion_forward szUNetModel.diffusion_forward) rrTrNNNrr&r&rxFTrV)NNNN)rWrXrYrZrArUrr\r r rKr rs.W E,rcsFeZdZdZ         dfd d ZdddZZS)EncoderUNetModelam The half UNet model with attention and timestep embedding. For usage, see UNetModel. Arguments --------- in_channels: int channels in the input torch.Tensor. model_channels: int base channel count for the model. out_channels: int channels in the output torch.Tensor. num_res_blocks: int number of residual blocks per downsample. attention_resolutions: int a collection of downsample rates at which attention will take place. May be a set, list, or tuple. For example, if this contains 4, then at 4x downsampling, attention will be used. dropout: float the dropout probability. channel_mult: int channel multiplier for each level of the UNet. conv_resample: bool if True, use learned convolutions for upsampling and downsampling dims: int determines if the signal is 1D, 2D, or 3D. num_heads: int the number of attention heads in each attention layer. num_head_channels: int if specified, ignore num_heads and instead use a fixed channel width per attention head. num_heads_upsample: int works with num_heads to set a different number of heads for upsampling. Deprecated. norm_num_groups: int Number of groups in the norm, default 32. resblock_updown: bool use residual blocks for up/downsampling. pool: str Type of pooling to use, one of: ["adaptive", "attention", "spatial", "spatial_v2"]. attention_pool_dim: int The dimension on which to apply attention pooling. out_kernel_size: int the kernel size of the output convolution use_fixup_init: bool whether to use FixUp initialization Example ------- >>> model = EncoderUNetModel( ... in_channels=3, ... model_channels=32, ... out_channels=1, ... num_res_blocks=1, ... attention_resolutions=[1] ... ) >>> x = torch.randn(4, 3, 16, 32) >>> ts = torch.tensor([10, 100, 50, 25]) >>> out = model(x, ts) >>> out.shape torch.Size([4, 1, 2, 4]) rrTrrr&rxFNrcst| dkr | } ||_||_||_||_||_||_||_||_ t j |_ | |_ | |_| |_||_|d}tt||tt|||_t|d|}ttt| ||dddg|_||_|g}d}t|D]}\}}t|D]<}t|||t||| | |dg}t||}||vr|t || | | |d|jt||j|7_||qu|t!|dkr|}|jt|rt||||| d | |d nt"||| |d |}|||d 9}|j|7_qmtt|||| |d t || | | |dt|||| |d |_#|j|7_||_$d|_%|dur2ttj&|| ddtt| |||dd|_'dS|dkrWtt&| |tt(dt)t| ||d|dt*|_'dS|dkrz| dkscJtt&| |tt+|||| ||_'dS|dkrtt|jdt,td|j|_'d |_%dS|dkrtt|jdt&| dttd|j|_'d |_%dSt-d|d)Nr&rjrrrrdrrTrrr)rr Fư> num_channels num_groupsepssamerureadaptive)rrryrIri spatial_v2z Unexpected z pooling).r@rArrrgrrr{rrr)r.r$rGrrout_kernel_sizerr|rr~rr[rr_rrrrrrwrrrrsrpoolspatial_poolingr}rAdaptiveAvgPool2drFlattenr;ReLUNotImplementedError)rJrrrgrrr{rrrrGrrrrrattention_pool_dimrr rrrrrrrrrrKr r rAbsD                          zEncoderUNetModel.__init__cCsd}|dur|t||j}g}||j}|jD]}|||}|jr1|||jjddq| ||}|jrT|||jjddt j |dd}| |S||j}| |S)O Apply the model to an input batch. Arguments --------- x: torch.Tensor an [N x C x ...] Tensor of inputs. timesteps: torch.Tensor a 1-D batch of timesteps. Returns ------- result: torch.Tensor an [N x K] Tensor of outputs. N)rrr'r&)axis) rr:rrr$rrrrPrr)r1r)rJrQr5r^resultsrr r r r rU!s&        zEncoderUNetModel.forward) rrTrrr&r&rxFNNrTrVrrr r rKr rs$K@rcr)raA simple module that computes the projection of an embedding vector onto the specified number of dimensions Arguments --------- emb_dim: int the original embedding dimensionality proj_dim: int the dimensionality of the target projection space Example ------- >>> mod_emb_proj = EmbeddingProjection(emb_dim=16, proj_dim=64) >>> emb = torch.randn(4, 16) >>> emb_proj = mod_emb_proj(emb) >>> emb_proj.shape torch.Size([4, 64]) cs@t||_||_t|||_t|_t|||_ dSrV) r@rArrrrinputr~actoutput)rJrrrKr r rA]s  zEmbeddingProjection.__init__cCs"||}||}||}|S)zComputes the forward pass Arguments --------- emb: torch.Tensor the original embedding tensor Returns ------- result: torch.Tensor the target embedding space )rrr)rJr^rQr r r rUes  zEmbeddingProjection.forwardrrr r rKr rGs rcsBeZdZdZ         dfd d ZdddZZS)DecoderUNetModela The half UNet model with attention and timestep embedding. For usage, see UNet. Arguments --------- in_channels: int channels in the input torch.Tensor. model_channels: int base channel count for the model. out_channels: int channels in the output torch.Tensor. num_res_blocks: int number of residual blocks per downsample. attention_resolutions: int a collection of downsample rates at which attention will take place. May be a set, list, or tuple. For example, if this contains 4, then at 4x downsampling, attention will be used. dropout: float the dropout probability. channel_mult: int channel multiplier for each level of the UNet. conv_resample: bool if True, use learned convolutions for upsampling and downsampling dims: int determines if the signal is 1D, 2D, or 3D. num_heads: int the number of attention heads in each attention layer. num_head_channels: int if specified, ignore num_heads and instead use a fixed channel width per attention head. num_heads_upsample: int works with num_heads to set a different number of heads for upsampling. Deprecated. resblock_updown: bool use residual blocks for up/downsampling. norm_num_groups: int Number of groups to use in norm, default 32 out_kernel_size: int Output kernel size, default 3 use_fixup_init: bool whether to use FixUp initialization Example ------- >>> model = DecoderUNetModel( ... in_channels=1, ... model_channels=32, ... out_channels=3, ... num_res_blocks=1, ... attention_resolutions=[1] ... ) >>> x = torch.randn(4, 1, 2, 4) >>> ts = torch.tensor([10, 100, 50, 25]) >>> out = model(x, ts) >>> out.shape torch.Size([4, 3, 16, 32]) rrTrrr&Frxrcs8t| dkr | } ||_||_||_||_||_||_||_||_ t j |_ | |_ | |_| |_|d}tt||tt|||_t|d|}tt| ||ddd|_tt|||| ||dt|| | ||dt|||| ||d|_t|_||_d}tt |D]s\}}t!|D]7}t|||t||| ||d g}t||}||vr|"t|| | ||d|j"t||j|7_q|t#|dkr|}|j"t| rt||||| d ||d nt$||| |d |}|d 9}|j|7_qttj%||ddtt| |||dd|_&|j|7_dS)Nr&rjrrrrdrrrTrrrrrrr)'r@rArrrgrrr{rrr)r.r$rGrrrr|rr~rr[r_r input_blockrwrrrupsample_blocksrrreversedrrrrcr}r)rJrrrgrrr{rrrrGrrrrrr rrrrrrrrrKr r rAs           zDecoderUNetModel.__init__NcCsfd}|dur|t||j}||j}|||}|||}|jD]}|||}q$||}|S)rN) rr:rrr$rrrr)rJrQr5r^rr r r r rUEs       zDecoderUNetModel.forward) rrTrrr&r&FrxrTrVrrr r rKr rxs Drrrcs,eZdZdZdfdd Zd ddZZS) DownsamplingPaddingaA wrapper module that applies the necessary padding for the downsampling factor Arguments --------- factor: int the downsampling / divisibility factor len_dim: int the index of the dimension in which the length will vary dims: list the list of dimensions to be included in padding Example ------- >>> padding = DownsamplingPadding(factor=4, dims=[1, 2], len_dim=1) >>> x = torch.randn(4, 7, 14) >>> length = torch.tensor([1., 0.8, 1., 0.7]) >>> x, length_new = padding(x, length) >>> x.shape torch.Size([4, 8, 16]) >>> length_new tensor([0.8750, 0.7000, 0.8750, 0.6125]) rNcs,t||_||_|durt}||_dSrV)r@rAfactorlen_dimDEFAULT_PADDING_DIMSr)rJrrrrKr r rAs  zDownsamplingPadding.__init__cCs<|}|jD]}t|||j|d\}}||jkr|}q||fS)aVApplies the padding Arguments --------- x: torch.Tensor the sample length: torch.Tensor the length tensor Returns ------- x_pad: torch.Tensor the padded tensor lens: torch.Tensor the new, adjusted lengths, if applicable )r)rrrr)rJrQrupdated_lengthr( length_padr r r rUs  zDownsamplingPadding.forwardrqrVrrr r rKr rgsrcs>eZdZdZ           dfd d ZZS)UNetNormalizingAutoencodera A convenience class for a UNet-based Variational Autoencoder (VAE) - useful in constructing Latent Diffusion models Arguments --------- in_channels: int the number of input channels model_channels: int the number of channels in the convolutional layers of the UNet encoder and decoder encoder_out_channels: int the number of channels the encoder will output latent_channels: int the number of channels in the latent space encoder_num_res_blocks: int the number of residual blocks in the encoder encoder_attention_resolutions: list the resolutions at which to apply attention layers in the encoder decoder_num_res_blocks: int the number of residual blocks in the decoder decoder_attention_resolutions: list the resolutions at which to apply attention layers in the encoder dropout: float the dropout probability channel_mult: tuple channel multipliers for each layer dims: int the convolution dimension to use (1, 2 or 3) num_heads: int the number of attention heads num_head_channels: int the number of channels in attention heads num_heads_upsample: int the number of upsampling heads norm_num_groups: int Number of norm groups, default 32 resblock_updown: bool whether to use residual blocks for upsampling and downsampling out_kernel_size: int the kernel size for output convolution layers (if applicable) len_dim: int Size of the output. out_mask_value: float Value to fill when masking the output. latent_mask_value: float Value to fill when masking the latent variable. use_fixup_norm: bool whether to use FixUp normalization downsampling_padding: int Amount of padding to apply in downsampling, default 2 ** len(channel_mult) Example ------- >>> unet_ae = UNetNormalizingAutoencoder( ... in_channels=1, ... model_channels=4, ... encoder_out_channels=16, ... latent_channels=3, ... encoder_num_res_blocks=1, ... encoder_attention_resolutions=[], ... decoder_num_res_blocks=1, ... decoder_attention_resolutions=[], ... norm_num_groups=2, ... ) >>> x = torch.randn(4, 1, 32, 32) >>> x_enc = unet_ae.encode(x) >>> x_enc.shape torch.Size([4, 3, 4, 4]) >>> x_dec = unet_ae.decode(x_enc) >>> x_dec.shape torch.Size([4, 1, 32, 32]) rrrrr&rxFrNcst|||||| | | | | |||||d}t|t| ||dd}|dur)dt| }t|}t|||||| t| | | | |||||d}tj ||||||ddS)N)rrrgrrr{rrrGrrrrrr r)rrrgrur)rrgrrrr{rrrGrrrrrr )encoderlatent_paddingdecoderrrr) rrr|rrrrrr@rA)rJrrencoder_out_channelslatent_channelsencoder_num_res_blocksencoder_attention_resolutionsdecoder_num_res_blocksdecoder_attention_resolutionsr{rrrGrrrrrrrruse_fixup_normdownsampling_padding encoder_unetr encoder_padr rKr r rAsl  z#UNetNormalizingAutoencoder.__init__)rrrrr&r&rxFrrrrFN)rWrXrYrZrAr\r r rKr rs"Sr)T)r!)NN)$rZr+abcrr)torch.nnrtorch.nn.functional functionalrnspeechbrain.utils.data_utilsr autoencodersrrrr r:Moduler;r]r|r_rcrsrwrrHrrrrrrrrr r r r sD      !E8<>N 8 P,1m;