class_docstr: ## U-Net model
docstr: :param in_channels: is the number of channels in the input feature map
  :param out_channels: is the number of channels in the output feature map
  :param channels: is the base channel count for the model
  :param n_res_blocks: number of residual blocks at each level
  :param attention_levels: are the levels at which attention should be performed
  :param channel_multipliers: are the multiplicative factors for number of channels for each level
  :param n_heads: is the number of attention heads in the transformers
  :param tf_layers: is the number of transformer layers in the transformers
  :param d_cond: is the size of the conditional embedding in the transformers
docstr_tok: [":","param","in_channels",":","is","the","number","of","channels","in","the","input","feature","map",":","param","out_channels",":","is","the","number","of","channels","in","the","output","feature","map",":","param","channels",":","is","the","base","channel","count","for","the","model",":","param","n_res_blocks",":","number","of","residual","blocks","at","each","level",":","param","attention_levels",":","are","the","levels","at","which","attention","should","be","performed",":","param","channel_multipliers",":","are","the","multiplicative","factors","for","number","of","channels","for","each","level",":","param","n_heads",":","is","the","number","of","attention","heads","in","the","transformers",":","param","tf_layers",":","is","the","number","of","transformer","layers","in","the","transformers",":","param","d_cond",":","is","the","size","of","the","conditional","embedding","in","the","transformers"]
code: def __init__(
      self,
      *,
      in_channels: int,
      out_channels: int,
      channels: int,
      n_res_blocks: int,
      attention_levels: List[int],
      channel_multipliers: List[int],
      n_heads: int,
      tf_layers: int = 1,
      d_cond: int = 768,
  ):
      """
      :param in_channels: is the number of channels in the input feature map
      :param out_channels: is the number of channels in the output feature map
      :param channels: is the base channel count for the model
      :param n_res_blocks: number of residual blocks at each level
      :param attention_levels: are the levels at which attention should be performed
      :param channel_multipliers: are the multiplicative factors for number of channels for each level
      :param n_heads: is the number of attention heads in the transformers
      :param tf_layers: is the number of transformer layers in the transformers
      :param d_cond: is the size of the conditional embedding in the transformers
      """
      super().__init__()
      self.channels = channels
  
      # Number of levels
      levels = len(channel_multipliers)
      # Size time embeddings
      d_time_emb = channels * 4
      self.time_embed = nn.Sequential(
          nn.Linear(channels, d_time_emb),
          nn.SiLU(),
          nn.Linear(d_time_emb, d_time_emb),
      )
  
      # Input half of the U-Net
      self.input_blocks = nn.ModuleList()
      # Initial $3 \times 3$ convolution that maps the input to `channels`.
      # The blocks are wrapped in `TimestepEmbedSequential` module because
      # different modules have different forward function signatures;
      # for example, convolution only accepts the feature map and
      # residual blocks accept the feature map and time embedding.
      # `TimestepEmbedSequential` calls them accordingly.
      self.input_blocks.append(
          TimestepEmbedSequential(
              nn.Conv2d(in_channels, channels, 3, padding=1)
          )
      )
      # Number of channels at each block in the input half of U-Net
      input_block_channels = [channels]
      # Number of channels at each level
      channels_list = [channels * m for m in channel_multipliers]
      # Prepare levels
      for i in range(levels):
          # Add the residual blocks and attentions
          for _ in range(n_res_blocks):
              # Residual block maps from previous number of channels to the number of
              # channels in the current level
              layers = [
                  ResBlock(
                      channels,
                      d_time_emb,
                      out_channels=channels_list[i],
                  )
              ]
              channels = channels_list[i]
              # Add transformer
              if i in attention_levels:
                  layers.append(
                      SpatialTransformer(
                          channels, n_heads, tf_layers, d_cond
                      )
                  )
              # Add them to the input half of the U-Net and keep track of the number of channels of
              # its output
              self.input_blocks.append(TimestepEmbedSequential(*layers))
              input_block_channels.append(channels)
          # Down sample at all levels except last
          if i != levels - 1:
              self.input_blocks.append(
                  TimestepEmbedSequential(DownSample(channels))
              )
              input_block_channels.append(channels)
  
      # The middle of the U-Net
      self.middle_block = TimestepEmbedSequential(
          ResBlock(channels, d_time_emb),
          SpatialTransformer(channels, n_heads, tf_layers, d_cond),
          ResBlock(channels, d_time_emb),
      )
  
      # Second half of the U-Net
      self.output_blocks = nn.ModuleList([])
      # Prepare levels in reverse order
      for i in reversed(range(levels)):
          # Add the residual blocks and attentions
          for j in range(n_res_blocks + 1):
              # Residual block maps from previous number of channels plus the
              # skip connections from the input half of U-Net to the number of
              # channels in the current level.
              layers = [
                  ResBlock(
                      channels + input_block_channels.pop(),
                      d_time_emb,
                      out_channels=channels_list[i],
                  )
              ]
              channels = channels_list[i]
              # Add transformer
              if i in attention_levels:
                  layers.append(
                      SpatialTransformer(
                          channels, n_heads, tf_layers, d_cond
                      )
                  )
              # Up-sample at every level after last residual block
              # except the last one.
              # Note that we are iterating in reverse; i.e. `i == 0` is the last.
              if i != 0 and j == n_res_blocks:
                  layers.append(UpSample(channels))
              # Add to the output half of the U-Net
              self.output_blocks.append(
                  TimestepEmbedSequential(*layers)
              )
  
      # Final normalization and $3 \times 3$ convolution
      self.out = nn.Sequential(
          normalization(channels),
          nn.SiLU(),
          nn.Conv2d(channels, out_channels, 3, padding=1),
      )
  
imports: import math
  from typing import List
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from labml_nn.diffusion.stable_diffusion.model.unet_attention import SpatialTransformer
label_desc: Use image node_id 1 to create a new UNetModel object from inherited base classes: nn with example: obj = UNetModel()