class_docstr: Modified Unet class that combines the output of controlling copy and frozen copy during forward pass.
class_docstr_tok: ["Modified","Unet","class","that","combines","the","output","of","controlling","copy","and","frozen","copy","during","forward","pass","."]
docstr: :param x: latents of diffusion process
  :param timesteps: diffusion step
  :param context: text embedding guiding the denoising process
  :param control: output from controlling copy of each corresponding layer
  :param only_mid_control: whether to add the output of controlling copy from middle block only
code: def forward(
      self,
      x,
      timesteps=None,
      context=None,
      control=None,
      only_mid_control=False,
      **kwargs,
  ):
      """
      :param x: latents of diffusion process
      :param timesteps: diffusion step
      :param context: text embedding guiding the denoising process
      :param control: output from controlling copy of each corresponding layer
      :param only_mid_control: whether to add the output of controlling copy from middle block only
      """
      hs = []
      with torch.no_grad():
          t_emb = timestep_embedding(
              timesteps, self.model_channels, repeat_only=False
          )
          emb = self.time_embed(t_emb)
          h = x.type(emb.dtype)
          for module in self.input_blocks:
              h = module(h, emb, context)
              hs.append(h)
          h = self.middle_block(h, emb, context)
  
      if control is not None:
          h += control.pop()
  
      for i, module in enumerate(self.output_blocks):
          if only_mid_control or control is None:
              h = torch.cat([h, hs.pop()], dim=1)
          else:
              h = torch.cat([h, hs.pop() + control.pop()], dim=1)
          h = module(h, emb, context)
  
      h = h.type(x.dtype)
      return self.out(h)
  
imports: from typing import Any
  import einops
  import torch
  import torch.nn
  from einops import rearrange, repeat
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from pytorch_lightning.utilities.rank_zero import rank_zero_only
  from torch._inductor import config
  from nemo.collections.multimodal.data.controlnet.controlnet_dataset import build_train_valid_datasets
  from nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.ddpm import LatentDiffusion
  from nemo.collections.multimodal.models.text_to_image.stable_diffusion.samplers.ddim import DDIMSampler
  from nemo.collections.multimodal.modules.stable_diffusion.attention import SpatialTransformer
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.openaimodel import AttentionBlock, Downsample, ResBlock, TimestepEmbedSequential, UNetModel
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.util import conv_nd, linear, timestep_embedding, zero_module
  from nemo.collections.multimodal.parts.stable_diffusion.utils import exists, log_txt_as_img
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.utils import logging
label_desc: Use image node_id 1 for calling the ControlledUnetModel obj's underlying member method code with example usage: obj.forward(x, timesteps, context, control, only_mid_control) and returns: self