code: def ema_scope(self, context=None):
      if self.use_ema:
          self.model_ema.store(self.model.parameters())
          self.model_ema.copy_to(self.model)
          if context is not None:
              print(f"{context}: Switched to EMA weights")
      try:
          yield None
      finally:
          if self.use_ema:
              self.model_ema.restore(self.model.parameters())
              if context is not None:
                  print(f"{context}: Restored training weights")
  
imports: import math
  from contextlib import contextmanager
  from typing import Any, Dict, List, Optional, Tuple, Union
  import pytorch_lightning
  import torch
  from omegaconf import ListConfig, OmegaConf
  from safetensors.torch import load_file
  from torch.optim.lr_scheduler import LambdaLR
  from ..modules import UNCONDITIONAL_CONFIG
  from ..modules.autoencoding.temporal_ae import VideoDecoder
  from ..modules.diffusionmodules.wrappers import OPENAIUNETWRAPPER
  from ..modules.ema import LitEma
  from ..util import default, disabled_train, get_obj_from_str, instantiate_from_config, log_txt_as_img
label_desc: Use image node_id 12 for calling the DiffusionEngine obj's underlying member method code with example usage: obj.ema_scope(context) without return types