class_docstr: An attention block that allows spatial positions to attend to each other.
  Originally ported from here, but adapted to the N-d case.
  https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
class_docstr_tok: ["An","attention","block","that","allows","spatial","positions","to","attend","to","each","other",".","Originally","ported","from","here",",","but","adapted","to","the","N-d","case",".","https",":","\/\/github.com\/hojonathanho\/diffusion\/blob\/1e0dceb3b3495bbe19116a5e1b3596cd0706c543\/diffusion_tf\/models\/unet.py","#","L66","."]
code: def __init__(
      self,
      channels,
      num_heads=1,
      num_head_channels=-1,
      use_checkpoint=False,
      use_new_attention_order=False,
  ):
      super().__init__()
      self.channels = channels
      if num_head_channels == -1:
          self.num_heads = num_heads
      else:
          assert (
              channels % num_head_channels == 0
          ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
          self.num_heads = channels // num_head_channels
      self.use_checkpoint = use_checkpoint
      self.norm = normalization(channels)
      self.qkv = conv_nd(1, channels, channels * 3, 1)
      if use_new_attention_order:
          # split qkv before split heads
          self.attention = QKVAttention(self.num_heads)
      else:
          # split heads before split qkv
          self.attention = QKVAttentionLegacy(self.num_heads)
  
      self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
  
imports: from abc import abstractmethod
  from functools import partial
  import math
  from typing import Iterable
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from ldm.modules.diffusionmodules.util import checkpoint, conv_nd, linear, avg_pool_nd, zero_module, normalization, timestep_embedding
  from ldm.modules.attention import SpatialTransformer
label_desc: Use image node_id 1 to create a new AttentionBlock object from inherited base classes: nn with example: obj = AttentionBlock(channels, num_heads, num_head_channels, use_checkpoint, use_new_attention_order)