class_docstr: Transformer block for image-like data.
  First, project the input (aka embedding)
  and reshape to b, t, d.
  Then apply standard transformer action.
  Finally, reshape to image
class_docstr_tok: ["Transformer","block","for","image-like","data",".","First",",","project","the","input","(","aka","embedding",")","and","reshape","to","b",",","t",",","d.","Then","apply","standard","transformer","action",".","Finally",",","reshape","to","image"]
code: def __init__(
      self,
      in_channels,
      n_heads,
      d_head,
      depth=1,
      dropout=0.0,
      context_dim=None,
      disable_self_attn=False,
      use_linear=False,
      use_checkpoint=False,
      use_flash_attention=False,
  ):
      super().__init__()
      if exists(context_dim) and not isinstance(context_dim, list):
          context_dim = [context_dim]
      self.in_channels = in_channels
      inner_dim = n_heads * d_head
      self.norm = Normalize(in_channels)
  
      if not use_linear:
          self.proj_in = nn.Conv2d(
              in_channels, inner_dim, kernel_size=1, stride=1, padding=0
          )
      else:
          self.proj_in = nn.Linear(in_channels, inner_dim)
  
      self.transformer_blocks = nn.ModuleList(
          [
              BasicTransformerBlock(
                  inner_dim,
                  n_heads,
                  d_head,
                  dropout=dropout,
                  context_dim=context_dim[d],
                  use_checkpoint=use_checkpoint,
                  use_flash_attention=use_flash_attention,
                  disable_self_attn=disable_self_attn,
              )
              for d in range(depth)
          ]
      )
  
      if not use_linear:
          self.proj_out = zero_module(
              nn.Conv2d(
                  inner_dim,
                  in_channels,
                  kernel_size=1,
                  stride=1,
                  padding=0,
              )
          )
      else:
          self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
      self.use_linear = use_linear
  
imports: import math
  from inspect import isfunction
  import torch
  import torch.nn.functional
  from apex.contrib.group_norm import GroupNorm
  from einops import rearrange
  from torch import einsum, nn
  from torch._dynamo import disable
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.util import checkpoint
label_desc: Use image node_id 1 to create a new SpatialTransformer object from inherited base classes: nn with example: obj = SpatialTransformer(in_channels, n_heads, d_head, depth, dropout, context_dim, disable__attn, use_linear, use_checkpoint, use_flash_attention)