code: def model_wrapper(
      model,
      noise_schedule,
      model_type="noise",
      model_kwargs=None,
      guidance_type="uncond",
      # condition=None,
      # unconditional_condition=None,
      guidance_scale=1.0,
      classifier_fn=None,
      classifier_kwargs=None,
  ):
      """Create a wrapper function for the noise prediction model.
  
      DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
      firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
  
      We support four types of the diffusion model by setting `model_type`:
  
          1. "noise": noise prediction model. (Trained by predicting noise).
  
          2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
  
          3. "v": velocity prediction model. (Trained by predicting the velocity).
              The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
  
              [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
                  arXiv preprint arXiv:2202.00512 (2022).
              [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
                  arXiv preprint arXiv:2210.02303 (2022).
  
          4. "score": marginal score function. (Trained by denoising score matching).
              Note that the score function and the noise prediction model follows a simple relationship:
              ```
                  noise(x_t, t) = -sigma_t * score(x_t, t)
              ```
  
      We support three types of guided sampling by DPMs by setting `guidance_type`:
          1. "uncond": unconditional sampling by DPMs.
              The input `model` has the following format:
              ``
                  model(x, t_input, **model_kwargs) -> noise | x_start | v | score
              ``
  
          2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
              The input `model` has the following format:
              ``
                  model(x, t_input, **model_kwargs) -> noise | x_start | v | score
              ``
  
              The input `classifier_fn` has the following format:
              ``
                  classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
              ``
  
              [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
                  in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
  
          3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
              The input `model` has the following format:
              ``
                  model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
              ``
              And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
  
              [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
                  arXiv preprint arXiv:2207.12598 (2022).
  
  
      The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
      or continuous-time labels (i.e. epsilon to T).
  
      We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
      ``
          def model_fn(x, t_continuous) -> noise:
              t_input = get_model_input_time(t_continuous)
              return noise_pred(model, x, t_input, **model_kwargs)
      ``
      where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
  
      ===============================================================
  
      Args:
          model: A diffusion model with the corresponding format described above.
          noise_schedule: A noise schedule object, such as NoiseScheduleVP.
          model_type: A `str`. The parameterization type of the diffusion model.
                      "noise" or "x_start" or "v" or "score".
          model_kwargs: A `dict`. A dict for the other inputs of the model function.
          guidance_type: A `str`. The type of the guidance for sampling.
                      "uncond" or "classifier" or "classifier-free".
          condition: A pytorch tensor. The condition for the guided sampling.
                      Only used for "classifier" or "classifier-free" guidance type.
          unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
                      Only used for "classifier-free" guidance type.
          guidance_scale: A `float`. The scale for the guided sampling.
          classifier_fn: A classifier function. Only used for the classifier guidance.
          classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
      Returns:
          A noise prediction model that accepts the noised data and the continuous time as the inputs.
      """
  
      model_kwargs = model_kwargs or {}
      classifier_kwargs = classifier_kwargs or {}
  
      def get_model_input_time(t_continuous):
          """
          Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
          For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
          For continuous-time DPMs, we just use `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, None, **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
              )
          elif model_type == "score":
              sigma_t = noise_schedule.marginal_std(t_continuous)
              dims = x.dim()
              return -expand_dims(sigma_t, dims) * output
  
      def cond_grad_fn(x, t_input, condition):
          """
          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, condition, unconditional_condition):
          """
          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, condition)
              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)
                  if isinstance(condition, dict):
                      assert isinstance(unconditional_condition, dict)
                      c_in = {}
                      for k in condition:
                          if isinstance(condition[k], list):
                              c_in[k] = [
                                  torch.cat(
                                      [
                                          unconditional_condition[k][i],
                                          condition[k][i],
                                      ]
                                  )
                                  for i in range(len(condition[k]))
                              ]
                          else:
                              c_in[k] = torch.cat(
                                  [
                                      unconditional_condition[k],
                                      condition[k],
                                  ]
                              )
                  elif isinstance(condition, list):
                      c_in = []
                      assert isinstance(unconditional_condition, list)
                      for i in range(len(condition)):
                          c_in.append(
                              torch.cat(
                                  [
                                      unconditional_condition[i],
                                      condition[i],
                                  ]
                              )
                          )
                  else:
                      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 torch
  import math
  import tqdm
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, guidance_scale, classifier_fn, classifier_kwargs) and returns: model_fn, unknown, t_continuous, output, torch, noise_pred_fn, unknown, unknown, unknown, unknown, noise_pred_fn, unknown