code: def __init__(
      self,
      input_size=(3, 32, 32),
      hidden_channels=64,
      dim_mults=(1, 2, 4, 8),
      groups=32,
      heads=4,
      dim_head=32,
      dropout_rate=0.0,
      num_res_blocks=2,
      attn_resolutions=(16,),
      attention_cls=SelfAttention,
  ):
      super().__init__()
      in_channels, in_height, in_width = input_size
      dims = [
          hidden_channels,
          *map(lambda m: hidden_channels * m, dim_mults),
      ]
      in_out = list(zip(dims[:-1], dims[1:]))
  
      self.time_pos_emb = SinusoidalPosEmb(hidden_channels)
      self.mlp = nn.Sequential(
          nn.Linear(hidden_channels, hidden_channels * 4),
          Mish(),
          nn.Linear(hidden_channels * 4, hidden_channels),
      )
  
      self.first_conv = nn.Conv2d(
          in_channels, hidden_channels, kernel_size=3, padding=1
      )
  
      h, w = in_height, in_width
      self.down_res_blocks = nn.ModuleList([])
      self.down_attn_blocks = nn.ModuleList([])
      self.down_spatial_blocks = nn.ModuleList([])
      for ind, (dim_in, dim_out) in enumerate(in_out):
          res_blocks = nn.ModuleList(
              [
                  ResnetBlock(
                      dim=dim_in,
                      dim_out=dim_out,
                      time_emb_dim=hidden_channels,
                      groups=groups,
                      dropout_rate=dropout_rate,
                  )
              ]
          )
          res_blocks.extend(
              [
                  ResnetBlock(
                      dim=dim_out,
                      dim_out=dim_out,
                      time_emb_dim=hidden_channels,
                      groups=groups,
                      dropout_rate=dropout_rate,
                  )
                  for _ in range(num_res_blocks - 1)
              ]
          )
          self.down_res_blocks.append(res_blocks)
  
          attn_blocks = nn.ModuleList([])
          if h in attn_resolutions and w in attn_resolutions:
              attn_blocks.extend(
                  [
                      Residual(
                          attention_cls(
                              dim_out,
                              heads=heads,
                              dim_head=dim_head,
                              groups=groups,
                          )
                      )
                      for _ in range(num_res_blocks)
                  ]
              )
          self.down_attn_blocks.append(attn_blocks)
  
          if ind < (len(in_out) - 1):
              spatial_blocks = nn.ModuleList([Downsample(dim_out)])
              h, w = h // 2, w // 2
          else:
              spatial_blocks = nn.ModuleList()
          self.down_spatial_blocks.append(spatial_blocks)
  
      mid_dim = dims[-1]
      self.mid_block1 = ResnetBlock(
          dim=mid_dim,
          dim_out=mid_dim,
          time_emb_dim=hidden_channels,
          groups=groups,
          dropout_rate=dropout_rate,
      )
      self.mid_attn = Residual(
          attention_cls(
              mid_dim, heads=heads, dim_head=dim_head, groups=groups
          )
      )
      self.mid_block2 = ResnetBlock(
          dim=mid_dim,
          dim_out=mid_dim,
          time_emb_dim=hidden_channels,
          groups=groups,
          dropout_rate=dropout_rate,
      )
  
      self.ups_res_blocks = nn.ModuleList([])
      self.ups_attn_blocks = nn.ModuleList([])
      self.ups_spatial_blocks = nn.ModuleList([])
      for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):
          res_blocks = nn.ModuleList(
              [
                  ResnetBlock(
                      dim=dim_out * 2,
                      dim_out=dim_out,
                      time_emb_dim=hidden_channels,
                      groups=groups,
                      dropout_rate=dropout_rate,
                  )
                  for _ in range(num_res_blocks)
              ]
          )
          res_blocks.extend(
              [
                  ResnetBlock(
                      dim=dim_out + dim_in,
                      dim_out=dim_in,
                      time_emb_dim=hidden_channels,
                      groups=groups,
                      dropout_rate=dropout_rate,
                  )
              ]
          )
          self.ups_res_blocks.append(res_blocks)
  
          attn_blocks = nn.ModuleList([])
          if h in attn_resolutions and w in attn_resolutions:
              attn_blocks.extend(
                  [
                      Residual(
                          attention_cls(
                              dim_out,
                              heads=heads,
                              dim_head=dim_head,
                              groups=groups,
                          )
                      )
                      for _ in range(num_res_blocks)
                  ]
              )
              attn_blocks.append(
                  Residual(
                      attention_cls(
                          dim_in,
                          heads=heads,
                          dim_head=dim_head,
                          groups=groups,
                      )
                  )
              )
          self.ups_attn_blocks.append(attn_blocks)
  
          spatial_blocks = nn.ModuleList()
          if ind < (len(in_out) - 1):
              spatial_blocks.append(Upsample(dim_in))
              h, w = h * 2, w * 2
          self.ups_spatial_blocks.append(spatial_blocks)
  
      self.final_conv = nn.Sequential(
          nn.GroupNorm(groups, hidden_channels),
          Mish(),
          nn.Conv2d(hidden_channels, in_channels, 1),
      )
  
imports: import math
  import kornia
  import torch
  import torch.nn.functional
  from einops import rearrange
  from torch import nn
label_desc: Use image node_id 1 to create a new Unet object from inherited base classes: nn with example: obj = Unet(input_size, hidden_channels, dim_mults, groups, heads, dim_head, dropout_rate, num_res_blocks, attn_resolutions, attention_cls)