code: def add_special_tokens(
      hidden_states, attention_mask, sos_token, eos_token
  ):
      batch_size = hidden_states.shape[0]
  
      if attention_mask is not None:
          # Add two more steps to attn mask
          new_attn_mask_step = attention_mask.new_ones((batch_size, 1))
          attention_mask = torch.concat(
              [new_attn_mask_step, attention_mask, new_attn_mask_step],
              dim=-1,
          )
  
      # Add the SOS / EOS tokens at the start / end of the sequence respectively
      sos_token = sos_token.expand(batch_size, 1, -1)
      eos_token = eos_token.expand(batch_size, 1, -1)
      hidden_states = torch.concat(
          [sos_token, hidden_states, eos_token], dim=1
      )
      return hidden_states, attention_mask
  
imports: from dataclasses import dataclass
  from typing import Any, Dict, List, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.utils.checkpoint
  from ...configuration_utils import ConfigMixin, register_to_config
  from ...loaders import UNet2DConditionLoadersMixin
  from ...models.activations import get_activation
  from ...models.attention_processor import ADDED_KV_ATTENTION_PROCESSORS, CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnAddedKVProcessor, AttnProcessor
  from ...models.embeddings import TimestepEmbedding, Timesteps
  from ...models.modeling_utils import ModelMixin
  from ...models.resnet import Downsample2D, ResnetBlock2D, Upsample2D
  from ...models.transformer_2d import Transformer2DModel
  from ...models.unet_2d_blocks import DownBlock2D, UpBlock2D
  from ...models.unet_2d_condition import UNet2DConditionOutput
  from ...utils import BaseOutput, is_torch_version, logging
label_desc: Use image node_id 1 for calling a global function with example usage: add_special_tokens(hidden_states, attention_mask, sos_token, eos_token) and returns: hidden_states,  attention_mask