scheduler.h

Go to the documentation of this file.

 
#ifndef _NNDEPLOY_MODEL_STABLE_DIFFUSION_SCHEDULER_H_
#define _NNDEPLOY_MODEL_STABLE_DIFFUSION_SCHEDULER_H_
 
#include <random>
 
#include "nndeploy/base/any.h"
#include "nndeploy/base/common.h"
#include "nndeploy/base/glic_stl_include.h"
#include "nndeploy/base/log.h"
#include "nndeploy/base/macro.h"
#include "nndeploy/base/object.h"
#include "nndeploy/base/opencv_include.h"
#include "nndeploy/base/param.h"
#include "nndeploy/base/status.h"
#include "nndeploy/base/string.h"
#include "nndeploy/dag/edge.h"
#include "nndeploy/dag/graph.h"
#include "nndeploy/dag/loop.h"
#include "nndeploy/dag/node.h"
#include "nndeploy/device/buffer.h"
#include "nndeploy/device/device.h"
#include "nndeploy/device/memory_pool.h"
#include "nndeploy/device/tensor.h"
#include "nndeploy/stable_diffusion/type.h"
 
namespace nndeploy {
namespace stable_diffusion {
 
class NNDEPLOY_CC_API SchedulerParam : public base::Param {
 public:
  SchedulerParam() : base::Param() {}
  virtual ~SchedulerParam() {}
 
  PARAM_COPY(SchedulerParam);
  PARAM_COPY_TO(SchedulerParam);
 
  SchedulerParam &operator=(const SchedulerParam &param) {
    if (this == &param) return *this;
    version_ = param.version_;
    num_train_timesteps_ = param.num_train_timesteps_;
    clip_sample_ = param.clip_sample_;
    num_inference_steps_ = param.num_inference_steps_;
    unet_channels_ = param.unet_channels_;
    image_height_ = param.image_height_;
    image_width_ = param.image_width_;
    guidance_scale_ = param.guidance_scale_;
    vae_scale_factor_ = param.vae_scale_factor_;
    init_noise_sigma_ = param.init_noise_sigma_;
    return *this;
  }
 
 public:
  std::string version_ = "v1.5";
  int num_train_timesteps_ = 1000;  // 训练时间步数
  bool clip_sample_ = false;        // 是否裁剪样本
  int num_inference_steps_ = 50;    // 推断步数
  int unet_channels_ = 4;           // channel
  int image_height_ = 512;          // height
  int image_width_ = 512;           // width
  float guidance_scale_ = 7.5;      // 指导比例
  float vae_scale_factor_ = 0.18215;
  float init_noise_sigma_ = 1.0f;
 
 public:
  virtual base::Status serialize(rapidjson::Value &json,
                                 rapidjson::Document::AllocatorType &allocator);
  virtual base::Status deserialize(rapidjson::Value &json);
};
 
class NNDEPLOY_CC_API Scheduler {
 public:
  Scheduler(SchedulerType type) : scheduler_type_(type) {}
  virtual ~Scheduler() {}
 
  virtual base::Status init(SchedulerParam *param) = 0;
  virtual base::Status deinit() = 0;
 
  virtual base::Status setTimesteps() = 0;
 
  virtual base::Status scaleModelInput(device::Tensor *sample, int index) = 0;
  virtual base::Status step(device::Tensor *sample, device::Tensor *timestep,
                            device::Tensor *latents,
                            device::Tensor *pre_sample) = 0;
 
  virtual std::vector<int> &getTimesteps() = 0;
 
 protected:
  SchedulerType scheduler_type_ = kSchedulerTypeNotSupport;
  SchedulerParam *scheduler_param_;
};
 
class SchedulerCreator {
 public:
  virtual ~SchedulerCreator() {};
  virtual Scheduler *createScheduler(SchedulerType type) = 0;
};
 
template <typename T>
class TypeSchedulerCreator : public SchedulerCreator {
  virtual Scheduler *createScheduler(SchedulerType type) { return new T(type); }
};
 
std::map<SchedulerType, std::shared_ptr<SchedulerCreator>> &
getGlobalSchedulerCreatorMap();
 
template <typename T>
class TypeSchedulerRegister {
 public:
  explicit TypeSchedulerRegister(SchedulerType type) {
    getGlobalSchedulerCreatorMap()[type] = std::shared_ptr<T>(new T());
  }
};
 
extern NNDEPLOY_CC_API Scheduler *createScheduler(SchedulerType type);
 
base::Status initializeLatents(std::mt19937 &generator, float init_noise_sigma,
                               device ::Tensor *latents);
 
}  // namespace stable_diffusion
}  // namespace nndeploy
 
#endif