class_docstr: The full UNet model with attention and timestep embedding.
  
  :param in_channels: channels in the input Tensor.
  :param model_channels: base channel count for the model.
  :param out_channels: channels in the output Tensor.
  :param num_res_blocks: number of residual blocks per downsample.
  :param attention_resolutions: 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.
  :param dropout: the dropout probability.
  :param channel_mult: channel multiplier for each level of the UNet.
  :param conv_resample: if True, use learned convolutions for upsampling and
      downsampling.
  :param dims: determines if the signal is 1D, 2D, or 3D.
  :param num_classes: if specified (as an int), then this model will be
      class-conditional with `num_classes` classes.
  :param use_checkpoint: use gradient checkpointing to reduce memory usage.
  :param num_heads: the number of attention heads in each attention layer.
  :param num_heads_channels: if specified, ignore num_heads and instead use
                             a fixed channel width per attention head.
  :param num_heads_upsample: works with num_heads to set a different number
                             of heads for upsampling. Deprecated.
  :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
  :param resblock_updown: use residual blocks for up/downsampling.
  :param use_new_attention_order: use a different attention pattern for potentially
                                  increased efficiency.
class_docstr_tok: ["The","full","UNet","model","with","attention","and","timestep","embedding",".",":","param","in_channels",":","channels","in","the","input","Tensor",".",":","param","model_channels",":","base","channel","count","for","the","model",".",":","param","out_channels",":","channels","in","the","output","Tensor",".",":","param","num_res_blocks",":","number","of","residual","blocks","per","downsample",".",":","param","attention_resolutions",":","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",".",":","param","dropout",":","the","dropout","probability",".",":","param","channel_mult",":","channel","multiplier","for","each","level","of","the","UNet",".",":","param","conv_resample",":","if","True",",","use","learned","convolutions","for","upsampling","and","downsampling",".",":","param","dims",":","determines","if","the","signal","is","1D",",","2D",",","or","3D",".",":","param","num_classes",":","if","specified","(","as","an","int",")",",","then","this","model","will","be","class-conditional","with","`","num_classes","`","classes",".",":","param","use_checkpoint",":","use","gradient","checkpointing","to","reduce","memory","usage",".",":","param","num_heads",":","the","number","of","attention","heads","in","each","attention","layer",".",":","param","num_heads_channels",":","if","specified",",","ignore","num_heads","and","instead","use","a","fixed","channel","width","per","attention","head",".",":","param","num_heads_upsample",":","works","with","num_heads","to","set","a","different","number","of","heads","for","upsampling",".","Deprecated",".",":","param","use_scale_shift_norm",":","use","a","FiLM-like","conditioning","mechanism",".",":","param","resblock_updown",":","use","residual","blocks","for","up\/downsampling",".",":","param","use_new_attention_order",":","use","a","different","attention","pattern","for","potentially","increased","efficiency","."]
code: def __init__(
      self,
      image_size,
      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_classes=None,
      use_checkpoint=False,
      use_fp16=False,
      num_heads=1,
      num_head_channels=-1,
      num_heads_upsample=-1,
      use_scale_shift_norm=False,
      resblock_updown=False,
      use_new_attention_order=False,
  ):
      super().__init__()
  
      if num_heads_upsample == -1:
          num_heads_upsample = num_heads
  
      self.image_size = image_size
      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.num_classes = num_classes
      self.use_checkpoint = use_checkpoint
      self.dtype = paddle.float16 if use_fp16 else paddle.float32
      self.num_heads = num_heads
      self.num_head_channels = num_head_channels
      self.num_heads_upsample = num_heads_upsample
  
      time_embed_dim = model_channels * 4
      self.time_embed = nn.Sequential(
          linear(model_channels, time_embed_dim),
          SiLU(),
          linear(time_embed_dim, time_embed_dim),
      )
  
      if self.num_classes is not None:
          self.label_emb = nn.Embedding(num_classes, time_embed_dim)
  
      ch = input_ch = int(channel_mult[0] * model_channels)
      self.input_blocks = nn.LayerList(
          [
              TimestepEmbedSequential(
                  conv_nd(dims, in_channels, ch, 3, padding=1)
              )
          ]
      )
      self._feature_size = ch
      input_block_chans = [ch]
      ds = 1
      for level, mult in enumerate(channel_mult):
          for _ in range(num_res_blocks):
              layers = [
                  ResBlock(
                      ch,
                      time_embed_dim,
                      dropout,
                      out_channels=int(mult * model_channels),
                      dims=dims,
                      use_checkpoint=use_checkpoint,
                      use_scale_shift_norm=use_scale_shift_norm,
                  )
              ]
              ch = int(mult * model_channels)
              if ds in attention_resolutions:
                  layers.append(
                      AttentionBlock(
                          ch,
                          use_checkpoint=use_checkpoint,
                          num_heads=num_heads,
                          num_head_channels=num_head_channels,
                          use_new_attention_order=use_new_attention_order,
                      )
                  )
              self.input_blocks.append(TimestepEmbedSequential(*layers))
              self._feature_size += ch
              input_block_chans.append(ch)
          if level != len(channel_mult) - 1:
              out_ch = ch
              self.input_blocks.append(
                  TimestepEmbedSequential(
                      ResBlock(
                          ch,
                          time_embed_dim,
                          dropout,
                          out_channels=out_ch,
                          dims=dims,
                          use_checkpoint=use_checkpoint,
                          use_scale_shift_norm=use_scale_shift_norm,
                          down=True,
                      )
                      if resblock_updown
                      else Downsample(
                          ch,
                          conv_resample,
                          dims=dims,
                          out_channels=out_ch,
                      )
                  )
              )
              ch = out_ch
              input_block_chans.append(ch)
              ds *= 2
              self._feature_size += ch
  
      self.middle_block = TimestepEmbedSequential(
          ResBlock(
              ch,
              time_embed_dim,
              dropout,
              dims=dims,
              use_checkpoint=use_checkpoint,
              use_scale_shift_norm=use_scale_shift_norm,
          ),
          AttentionBlock(
              ch,
              use_checkpoint=use_checkpoint,
              num_heads=num_heads,
              num_head_channels=num_head_channels,
              use_new_attention_order=use_new_attention_order,
          ),
          ResBlock(
              ch,
              time_embed_dim,
              dropout,
              dims=dims,
              use_checkpoint=use_checkpoint,
              use_scale_shift_norm=use_scale_shift_norm,
          ),
      )
      self._feature_size += ch
  
      self.output_blocks = nn.LayerList([])
      for level, mult in list(enumerate(channel_mult))[::-1]:
          for i in range(num_res_blocks + 1):
              ich = input_block_chans.pop()
              layers = [
                  ResBlock(
                      ch + ich,
                      time_embed_dim,
                      dropout,
                      out_channels=int(model_channels * mult),
                      dims=dims,
                      use_checkpoint=use_checkpoint,
                      use_scale_shift_norm=use_scale_shift_norm,
                  )
              ]
              ch = int(model_channels * mult)
              if ds in attention_resolutions:
                  layers.append(
                      AttentionBlock(
                          ch,
                          use_checkpoint=use_checkpoint,
                          num_heads=num_heads_upsample,
                          num_head_channels=num_head_channels,
                          use_new_attention_order=use_new_attention_order,
                      )
                  )
              if level and i == num_res_blocks:
                  out_ch = ch
                  layers.append(
                      ResBlock(
                          ch,
                          time_embed_dim,
                          dropout,
                          out_channels=out_ch,
                          dims=dims,
                          use_checkpoint=use_checkpoint,
                          use_scale_shift_norm=use_scale_shift_norm,
                          up=True,
                      )
                      if resblock_updown
                      else Upsample(
                          ch,
                          conv_resample,
                          dims=dims,
                          out_channels=out_ch,
                      )
                  )
                  ds //= 2
              self.output_blocks.append(
                  TimestepEmbedSequential(*layers)
              )
              self._feature_size += ch
  
      self.out = nn.Sequential(
          normalization(ch),
          SiLU(),
          zero_module(
              conv_nd(dims, input_ch, out_channels, 3, padding=1)
          ),
      )
  
imports: import math
  from abc import abstractmethod
  import numpy
  import paddle
  import paddle.nn
  import paddle.nn.functional
  from .nn import avg_pool_nd
  from .nn import checkpoint
  from .nn import conv_nd
  from .nn import linear
  from .nn import normalization
  from .nn import SiLU
  from .nn import timestep_embedding
  from .nn import zero_module
label_desc: Use image node_id 1 to create a new UNetModel object from inherited base classes: nn with example: obj = UNetModel(image_size, in_channels, model_channels, out_channels, num_res_blocks, attention_resolutions, dropout, channel_mult, conv_resample, dims, num_classes, use_checkpoint, use_fp16, num_heads, num_head_channels, num_heads_upsample, use_scale_shift_norm, resblock_updown, use_new_attention_order)