class_docstr: The half UNet model with attention and timestep embedding.
  For usage, see UNet.
  
  Arguments
  ---------
  in_channels: int
      channels in the input Tensor.
  model_channels: int
      base channel count for the model.
  out_channels: int
      channels in the output 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
  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}
  dims: int
      determines if the signal is 1D, 2D, or 3D.
  num_heads: int
      the number of attention heads in each attention layer.
  num_heads_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.
  
  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])
class_docstr_tok: ["The","half","UNet","model","with","attention","and","timestep","embedding",".","For","usage",",","see","UNet",".","Arguments","--","--","--","--","-","in_channels",":","int","channels","in","the","input","Tensor",".","model_channels",":","int","base","channel","count","for","the","model",".","out_channels",":","int","channels","in","the","output","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","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","}","dims",":","int","determines","if","the","signal","is","1D",",","2D",",","or","3D",".","num_heads",":","int","the","number","of","attention","heads","in","each","attention","layer",".","num_heads_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",".","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","]",")"]
code: def __init__(
      self,
      in_channels,
      model_channels,
      out_channels,
      num_res_blocks,
      attention_resolutions,
      dropout=0,
      channel_mult=(1, 2, 4, 8),
      conv_resample=True,
      dims=2,
      num_heads=1,
      num_head_channels=-1,
      num_heads_upsample=-1,
      resblock_updown=False,
      norm_num_groups=32,
      out_kernel_size=3,
      use_fixup_init=True,
  ):
      super().__init__()
  
      if num_heads_upsample == -1:
          num_heads_upsample = num_heads
  
      self.in_channels = in_channels
      self.model_channels = model_channels
      self.out_channels = out_channels
      self.num_res_blocks = num_res_blocks
      self.attention_resolutions = attention_resolutions
      self.dropout = dropout
      self.channel_mult = channel_mult
      self.conv_resample = conv_resample
      self.dtype = torch.float32
      self.num_heads = num_heads
      self.num_head_channels = num_head_channels
      self.num_heads_upsample = num_heads_upsample
  
      emb_dim = model_channels * 4
      self.time_embed = nn.Sequential(
          nn.Linear(model_channels, emb_dim),
          nn.SiLU(),
          nn.Linear(emb_dim, emb_dim),
      )
  
      ch = int(channel_mult[0] * model_channels)
  
      self.input_block = TimestepEmbedSequential(
          conv_nd(dims, in_channels, ch, 3, padding=1)
      )
  
      self.middle_block = TimestepEmbedSequential(
          ResBlock(
              ch,
              emb_dim,
              dropout,
              dims=dims,
              norm_num_groups=norm_num_groups,
              use_fixup_init=use_fixup_init,
          ),
          AttentionBlock(
              ch,
              num_heads=num_heads,
              num_head_channels=num_head_channels,
              norm_num_groups=norm_num_groups,
              use_fixup_init=use_fixup_init,
          ),
          ResBlock(
              ch,
              emb_dim,
              dropout,
              dims=dims,
              norm_num_groups=norm_num_groups,
              use_fixup_init=use_fixup_init,
          ),
      )
  
      self.upsample_blocks = nn.ModuleList()
      self._feature_size = ch
      ds = 1
  
      for level, mult in enumerate(reversed(channel_mult)):
          for _ in range(num_res_blocks):
              layers = [
                  ResBlock(
                      ch,
                      emb_dim,
                      dropout,
                      out_channels=int(mult * model_channels),
                      dims=dims,
                      norm_num_groups=norm_num_groups,
                      use_fixup_init=use_fixup_init,
                  )
              ]
              ch = int(mult * model_channels)
              if ds in attention_resolutions:
                  layers.append(
                      AttentionBlock(
                          ch,
                          num_heads=num_heads,
                          num_head_channels=num_head_channels,
                          norm_num_groups=norm_num_groups,
                          use_fixup_init=use_fixup_init,
                      )
                  )
              self.upsample_blocks.append(
                  TimestepEmbedSequential(*layers)
              )
              self._feature_size += ch
          if level != len(channel_mult) - 1:
              out_ch = ch
              self.upsample_blocks.append(
                  TimestepEmbedSequential(
                      ResBlock(
                          ch,
                          emb_dim,
                          dropout,
                          out_channels=out_ch,
                          dims=dims,
                          up=True,
                          norm_num_groups=norm_num_groups,
                          use_fixup_init=use_fixup_init,
                      )
                      if resblock_updown
                      else Upsample(
                          ch,
                          conv_resample,
                          dims=dims,
                          out_channels=out_ch,
                      )
                  )
              )
              ch = out_ch
              ds *= 2
              self._feature_size += ch
  
      self.out = nn.Sequential(
          nn.GroupNorm(
              num_channels=ch, num_groups=norm_num_groups, eps=1e-6
          ),
          nn.SiLU(),
          conv_nd(
              dims,
              ch,
              out_channels,
              kernel_size=out_kernel_size,
              padding="same",
          ),
      )
      self._feature_size += ch
  
imports: from abc import abstractmethod
  from speechbrain.utils.data_utils import pad_divisible
  from .autoencoders import NormalizingAutoencoder
  import math
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 1 to create a new DecoderUNetModel object from inherited base classes: nn with example: obj = DecoderUNetModel(in_channels, model_channels, out_channels, num_res_blocks, attention_resolutions, dropout, channel_mult, conv_resample, dims, num_heads, num_head_channels, num_heads_upsample, resblock_updown, norm_num_groups, out_kernel_size, use_fixup_init)