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,
      attn_type="softmax",
      use_checkpoint=True,
      # sdp_backend=SDPBackend.FLASH_ATTENTION
      sdp_backend=None,
  ):
      super().__init__()
      logpy.debug(
          f"constructing {self.__class__.__name__} of depth {depth} w/ "
          f"{in_channels} channels and {n_heads} heads."
      )
  
      if exists(context_dim) and not isinstance(context_dim, list):
          context_dim = [context_dim]
      if exists(context_dim) and isinstance(context_dim, list):
          if depth != len(context_dim):
              logpy.warn(
                  f"{self.__class__.__name__}: Found context dims "
                  f"{context_dim} of depth {len(context_dim)}, which does not "
                  f"match the specified 'depth' of {depth}. Setting context_dim "
                  f"to {depth * [context_dim[0]]} now."
              )
              # depth does not match context dims.
              assert all(
                  map(lambda x: x == context_dim[0], context_dim)
              ), "need homogenous context_dim to match depth automatically"
              context_dim = depth * [context_dim[0]]
      elif context_dim is None:
          context_dim = [None] * depth
      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,
                  attn_mode=attn_type,
                  checkpoint=use_checkpoint,
                  sdp_backend=sdp_backend,
              )
              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.proj_out = zero_module(nn.Linear(inner_dim, in_channels))
      self.use_linear = use_linear
  
imports: import logging
  import math
  from inspect import isfunction
  from typing import Any, Optional
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from packaging import version
  from torch import nn
  from torch.utils.checkpoint 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, attn_type, use_checkpoint, sdp_backend)