code: def __init__(self, plan_dir, checkpoint):
      super().__init__()
      with open(os.path.join(plan_dir, "conf.yaml"), "rb") as fp:
          config = OmegaConf.load(fp.name)
      max_batch_size = config.batch_size
  
      self.tokenizer = CLIPTokenizer.from_pretrained(
          "openai/clip-vit-large-patch14"
      )
      self.max_length = config.clip.max_length
      self.rng = torch.Generator(
          device=torch.cuda.current_device(),
      )
  
      self.set_beta_schedule()
  
      stream = cuda.Stream()
  
      self.image_encoder = self.load_vae_from_checkpoint(checkpoint)
  
      self.text_encoder = Engine(os.path.join(plan_dir, "clip.plan"))
      shape_dict = {
          "tokens": config.clip.tokens,
          "logits": config.clip.logits,
      }
      self.text_encoder.set_engine(stream, shape_dict)
  
      self.unet = Engine(os.path.join(plan_dir, "unet.plan"))
      shape_dict = {
          "x": config.unet.x,
          "t": (max_batch_size * 2,),
          "context": config.unet.context,
          "logits": config.unet.logits,
      }
      self.unet.set_engine(stream, shape_dict)
  
imports: import logging
  import os
  import tempfile
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from omegaconf import OmegaConf
  from polygraphy import cuda
  from transformers import CLIPTokenizer
  from nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.ddpm import LatentDiffusion
  from nemo.collections.multimodal.modules.nerf.guidance.txt2img_guidance_base import Txt2ImgGuidanceBase
  from nemo.collections.multimodal.modules.nerf.utils.trt_engine import Engine, device_view
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.util import extract_into_tensor, make_beta_schedule
  from nemo.collections.multimodal.parts.stable_diffusion.utils import default
  from nemo.collections.multimodal.parts.utils import randn_like
  from nemo.collections.nlp.parts.nlp_overrides import NLPSaveRestoreConnector
label_desc: Use image node_id 1 to create a new LatentDiffusionWrapper object from inherited base classes: Txt2ImgGuidanceBase with example: obj = LatentDiffusionWrapper(plan_dir, checkpoint)