class_docstr: The full UNet model with attention and timestep embedding.
  :param in_channels: channels in the input Tensor.
  :param model_channels: base channel count for the model.
  :param out_channels: channels in the output Tensor.
  :param num_res_blocks: number of residual blocks per downsample.
  :param attention_resolutions: a collection of downsample rates at which
      attention will take place. May be a set, list, or tuple.
      For example, if this contains 4, then at 4x downsampling, attention
      will be used.
  :param dropout: the dropout probability.
  :param channel_mult: channel multiplier for each level of the UNet.
  :param conv_resample: if True, use learned convolutions for upsampling and
      downsampling.
  :param dims: determines if the signal is 1D, 2D, or 3D.
  :param num_classes: if specified (as an int), then this model will be
      class-conditional with `num_classes` classes.
  :param use_checkpoint: use gradient checkpointing to reduce memory usage.
  :param num_heads: the number of attention heads in each attention layer.
  :param num_heads_channels: if specified, ignore num_heads and instead use
                             a fixed channel width per attention head.
  :param num_heads_upsample: works with num_heads to set a different number
                             of heads for upsampling. Deprecated.
  :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
  :param resblock_updown: use residual blocks for up/downsampling.
  :param use_new_attention_order: use a different attention pattern for potentially
                                  increased efficiency.
class_docstr_tok: ["The","full","UNet","model","with","attention","and","timestep","embedding",".",":","param","in_channels",":","channels","in","the","input","Tensor",".",":","param","model_channels",":","base","channel","count","for","the","model",".",":","param","out_channels",":","channels","in","the","output","Tensor",".",":","param","num_res_blocks",":","number","of","residual","blocks","per","downsample",".",":","param","attention_resolutions",":","a","collection","of","downsample","rates","at","which","attention","will","take","place",".","May","be","a","set",",","list",",","or","tuple",".","For","example",",","if","this","contains","4",",","then","at","4x","downsampling",",","attention","will","be","used",".",":","param","dropout",":","the","dropout","probability",".",":","param","channel_mult",":","channel","multiplier","for","each","level","of","the","UNet",".",":","param","conv_resample",":","if","True",",","use","learned","convolutions","for","upsampling","and","downsampling",".",":","param","dims",":","determines","if","the","signal","is","1D",",","2D",",","or","3D",".",":","param","num_classes",":","if","specified","(","as","an","int",")",",","then","this","model","will","be","class-conditional","with","`","num_classes","`","classes",".",":","param","use_checkpoint",":","use","gradient","checkpointing","to","reduce","memory","usage",".",":","param","num_heads",":","the","number","of","attention","heads","in","each","attention","layer",".",":","param","num_heads_channels",":","if","specified",",","ignore","num_heads","and","instead","use","a","fixed","channel","width","per","attention","head",".",":","param","num_heads_upsample",":","works","with","num_heads","to","set","a","different","number","of","heads","for","upsampling",".","Deprecated",".",":","param","use_scale_shift_norm",":","use","a","FiLM-like","conditioning","mechanism",".",":","param","resblock_updown",":","use","residual","blocks","for","up\/downsampling",".",":","param","use_new_attention_order",":","use","a","different","attention","pattern","for","potentially","increased","efficiency","."]
docstr: Apply the model to an input batch.
  :param x: an [N x C x ...] Tensor of inputs.
  :param timesteps: a 1-D batch of timesteps.
  :param context: conditioning plugged in via crossattn
  :param y: an [N] Tensor of labels, if class-conditional.
  :return: an [N x C x ...] Tensor of outputs.
docstr_tok: ["Apply","the","model","to","an","input","batch",".",":","param","x",":","an","[","N","x","C","x","...","]","Tensor","of","inputs",".",":","param","timesteps",":","a","1-D","batch","of","timesteps",".",":","param","context",":","conditioning","plugged","in","via","crossattn",":","param","y",":","an","[","N","]","Tensor","of","labels",",","if","class-conditional",".",":","return",":","an","[","N","x","C","x","...","]","Tensor","of","outputs","."]
code: def forward(
      self,
      x: th.Tensor,
      timesteps: Optional[th.Tensor] = None,
      context: Optional[th.Tensor] = None,
      y: Optional[th.Tensor] = None,
      **kwargs,
  ) -> th.Tensor:
      """
      Apply the model to an input batch.
      :param x: an [N x C x ...] Tensor of inputs.
      :param timesteps: a 1-D batch of timesteps.
      :param context: conditioning plugged in via crossattn
      :param y: an [N] Tensor of labels, if class-conditional.
      :return: an [N x C x ...] Tensor of outputs.
      """
      assert (y is not None) == (
          self.num_classes is not None
      ), "must specify y if and only if the model is class-conditional"
      hs = []
      t_emb = timestep_embedding(
          timesteps, self.model_channels, repeat_only=False
      )
      emb = self.time_embed(t_emb)
  
      if self.num_classes is not None:
          assert y.shape[0] == x.shape[0]
          emb = emb + self.label_emb(y)
  
      h = x
      for module in self.input_blocks:
          h = module(h, emb, context)
          hs.append(h)
      h = self.middle_block(h, emb, context)
      for module in self.output_blocks:
          h = th.cat([h, hs.pop()], dim=1)
          h = module(h, emb, context)
      h = h.type(x.dtype)
  
      return self.out(h)
  
imports: import logging
  import math
  from abc import abstractmethod
  from typing import Iterable, List, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange
  from torch.utils.checkpoint import checkpoint
  from ...modules.attention import SpatialTransformer
  from ...modules.diffusionmodules.util import avg_pool_nd, conv_nd, linear, normalization, timestep_embedding, zero_module
  from ...modules.video_attention import SpatialVideoTransformer
  from ...util import exists
label_desc: Use image node_id 2 for calling the UNetModel obj's underlying member method code with example usage: obj.forward(x, timesteps, context, y) and returns: self