noise_schedule.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cstdint>
#include <vector>
 
#include "pulsatrix/device_backend.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
class NoiseSchedule {
public:
    explicit NoiseSchedule(int64_t num_timesteps, float beta_start = 1e-4f, float beta_end = 0.02f);
 
    [[nodiscard]] int64_t num_timesteps() const { return num_timesteps_; }
 
    [[nodiscard]] float beta(int64_t t) const;
 
    [[nodiscard]] float alpha(int64_t t) const;
 
    [[nodiscard]] float alpha_bar(int64_t t) const;
 
    [[nodiscard]] Tensor add_noise(const Tensor& x0, const Tensor& epsilon, int64_t t) const;
 
    [[nodiscard]] Tensor denoise_step(const Tensor& x_t, const Tensor& predicted_epsilon, const Tensor& z,
                                      int64_t t) const;
 
private:
    [[nodiscard]] size_t index_of(int64_t t, const char* context) const;
 
    int64_t num_timesteps_;
    std::vector<float> beta_;       
    std::vector<float> alpha_;      
    std::vector<float> alpha_bar_;  
};
 
}  // namespace pulsatrix