code: def model_wrapper(
      model,
      noise_schedule,
      model_type="noise",
      model_kwargs={},
      guidance_type="uncond",
      condition=None,
      unconditional_condition=None,
      guidance_scale=1.0,
      classifier_fn=None,
      classifier_kwargs={},
  ):
      """Create a wrapper function for the noise prediction model."""
  
      def get_model_input_time(t_continuous):
          if noise_schedule.schedule == "discrete":
              return (
                  t_continuous - 1.0 / noise_schedule.total_N
              ) * 1000.0
          else:
              return t_continuous
  
      def noise_pred_fn(x, t_continuous, cond=None):
          if t_continuous.reshape((-1,)).shape[0] == 1:
              t_continuous = t_continuous.expand((x.shape[0]))
          t_input = get_model_input_time(t_continuous)
          if cond is None:
              output = model(x, t_input, **model_kwargs)
          else:
              output = model(x, t_input, cond, **model_kwargs)
          if model_type == "noise":
              return output
          elif model_type == "x_start":
              alpha_t, sigma_t = (
                  noise_schedule.marginal_alpha(t_continuous),
                  noise_schedule.marginal_std(t_continuous),
              )
              dims = x.dim()
              return (
                  x - expand_dims(alpha_t, dims) * output
              ) / expand_dims(sigma_t, dims)
          elif model_type == "v":
              alpha_t, sigma_t = (
                  noise_schedule.marginal_alpha(t_continuous),
                  noise_schedule.marginal_std(t_continuous),
              )
              dims = x.dim()
              return (
                  expand_dims(alpha_t, dims) * output
                  + expand_dims(sigma_t, dims) * x
              )
  
      def cond_grad_fn(x, t_input):
          """
          Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
          """
          with torch.enable_grad():
              x_in = x.detach().requires_grad_(True)
              log_prob = classifier_fn(
                  x_in, t_input, condition, **classifier_kwargs
              )
              return torch.autograd.grad(log_prob.sum(), x_in)[0]
  
      def model_fn(x, t_continuous):
          """
          The noise predicition model function that is used for DPM-Solver.
          """
          if t_continuous.reshape((-1,)).shape[0] == 1:
              t_continuous = t_continuous.expand((x.shape[0]))
          if guidance_type == "uncond":
              return noise_pred_fn(x, t_continuous)
          elif guidance_type == "classifier":
              assert classifier_fn is not None
              t_input = get_model_input_time(t_continuous)
              cond_grad = cond_grad_fn(x, t_input)
              sigma_t = noise_schedule.marginal_std(t_continuous)
              noise = noise_pred_fn(x, t_continuous)
              return (
                  noise
                  - guidance_scale
                  * expand_dims(sigma_t, dims=cond_grad.dim())
                  * cond_grad
              )
          elif guidance_type == "classifier-free":
              if (
                  guidance_scale == 1.0
                  or unconditional_condition is None
              ):
                  return noise_pred_fn(x, t_continuous, cond=condition)
              else:
                  x_in = torch.cat([x] * 2)
                  t_in = torch.cat([t_continuous] * 2)
                  c_in = torch.cat([unconditional_condition, condition])
                  noise_uncond, noise = noise_pred_fn(
                      x_in, t_in, cond=c_in
                  ).chunk(2)
                  return noise_uncond + guidance_scale * (
                      noise - noise_uncond
                  )
  
      assert model_type in ["noise", "x_start", "v"]
      assert guidance_type in [
          "uncond",
          "classifier",
          "classifier-free",
      ]
      return model_fn
  
imports: import math
  import torch
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.util import expand_dims, interpolate_fn
label_desc: Use image node_id 1 for calling a global function with example usage: model_wrapper(model, noise_schedule, model_type, model_kwargs, guidance_type, condition, unconditional_condition, guidance_scale, classifier_fn, classifier_kwargs) and returns: model_fn, unknown, t_continuous, output, torch, noise_pred_fn, unknown, unknown, unknown, noise_pred_fn, unknown