class_docstr: ## U-Net
class_docstr_tok: ["#","#","U-Net"]
docstr: * `image_channels` is the number of channels in the image. $3$ for RGB.
  * `n_channels` is number of channels in the initial feature map that we transform the image into
  * `ch_mults` is the list of channel numbers at each resolution. The number of channels is `ch_mults[i] * n_channels`
  * `is_attn` is a list of booleans that indicate whether to use attention at each resolution
  * `n_blocks` is the number of `UpDownBlocks` at each resolution
code: def __init__(
      self,
      image_channels: int = 3,
      n_channels: int = 64,
      ch_mults: Union[Tuple[int, ...], List[int]] = (1, 2, 2, 4),
      is_attn: Union[Tuple[bool, ...], List[bool]] = (
          False,
          False,
          True,
          True,
      ),
      n_blocks: int = 2,
  ):
      """
      * `image_channels` is the number of channels in the image. $3$ for RGB.
      * `n_channels` is number of channels in the initial feature map that we transform the image into
      * `ch_mults` is the list of channel numbers at each resolution. The number of channels is `ch_mults[i] * n_channels`
      * `is_attn` is a list of booleans that indicate whether to use attention at each resolution
      * `n_blocks` is the number of `UpDownBlocks` at each resolution
      """
      super().__init__()
  
      # Number of resolutions
      n_resolutions = len(ch_mults)
  
      # Project image into feature map
      self.image_proj = nn.Conv2d(
          image_channels, n_channels, kernel_size=(3, 3), padding=(1, 1)
      )
  
      # Time embedding layer. Time embedding has `n_channels * 4` channels
      self.time_emb = TimeEmbedding(n_channels * 4)
  
      # #### First half of U-Net - decreasing resolution
      down = []
      # Number of channels
      out_channels = in_channels = n_channels
      # For each resolution
      for i in range(n_resolutions):
          # Number of output channels at this resolution
          out_channels = in_channels * ch_mults[i]
          # Add `n_blocks`
          for _ in range(n_blocks):
              down.append(
                  DownBlock(
                      in_channels,
                      out_channels,
                      n_channels * 4,
                      is_attn[i],
                  )
              )
              in_channels = out_channels
          # Down sample at all resolutions except the last
          if i < n_resolutions - 1:
              down.append(Downsample(in_channels))
  
      # Combine the set of modules
      self.down = nn.ModuleList(down)
  
      # Middle block
      self.middle = MiddleBlock(
          out_channels,
          n_channels * 4,
      )
  
      # #### Second half of U-Net - increasing resolution
      up = []
      # Number of channels
      in_channels = out_channels
      # For each resolution
      for i in reversed(range(n_resolutions)):
          # `n_blocks` at the same resolution
          out_channels = in_channels
          for _ in range(n_blocks):
              up.append(
                  UpBlock(
                      in_channels,
                      out_channels,
                      n_channels * 4,
                      is_attn[i],
                  )
              )
          # Final block to reduce the number of channels
          out_channels = in_channels // ch_mults[i]
          up.append(
              UpBlock(
                  in_channels, out_channels, n_channels * 4, is_attn[i]
              )
          )
          in_channels = out_channels
          # Up sample at all resolutions except last
          if i > 0:
              up.append(Upsample(in_channels))
  
      # Combine the set of modules
      self.up = nn.ModuleList(up)
  
      # Final normalization and convolution layer
      self.norm = nn.GroupNorm(8, n_channels)
      self.act = Swish()
      self.final = nn.Conv2d(
          in_channels,
          image_channels,
          kernel_size=(3, 3),
          padding=(1, 1),
      )
  
imports: import math
  from typing import Optional, Tuple, Union, List
  import torch
  from torch import nn
  from labml_helpers.module import Module
label_desc: Use image node_id 1 to create a new UNet object from inherited base classes: Module with example: obj = UNet(image_channels, n_channels, ch_mults, is_attn, n_blocks)