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
  NEW: use_linear for more efficiency instead of the 1x1 convs
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","NEW",":","use_linear","for","more","efficiency","instead","of","the","1x1","convs"]
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=True,
  ):
      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],
                  disable_self_attn=disable_self_attn,
                  checkpoint=use_checkpoint,
              )
              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: from inspect import isfunction
  import math
  import torch
  import torch.nn.functional
  from torch import nn, einsum
  from einops import rearrange, repeat
  from typing import Optional, Any
  from ldm.modules.diffusionmodules.util import checkpoint
  import os
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)