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.
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_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,
      encoder_channels=None,
  ):
      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.num_classes = num_classes
      self.use_checkpoint = use_checkpoint
      self.dtype = th.float16 if use_fp16 else th.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),
          nn.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.ModuleList(
          [
              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,
                          encoder_channels=encoder_channels,
                      )
                  )
              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,
              encoder_channels=encoder_channels,
          ),
          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.ModuleList([])
      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,
                          encoder_channels=encoder_channels,
                      )
                  )
              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, swish=1.0),
          nn.Identity(),
          zero_module(
              conv_nd(dims, input_ch, out_channels, 3, padding=1)
          ),
      )
      self.use_fp16 = use_fp16
  
imports: import math
  from abc import abstractmethod
  import torch
  import torch.nn
  import torch.nn.functional
  from .fp16_util import convert_module_to_f16, convert_module_to_f32
  from .nn import avg_pool_nd, conv_nd, linear, normalization, timestep_embedding, zero_module
label_desc: Use image node_id 1 to create a new UNetModel object from inherited base classes: nn with example: obj = UNetModel(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, encoder_channels)