#ifndef __SD_UPSCALER_H__
#define __SD_UPSCALER_H__

#include "esrgan.hpp"
#include "ggml_extend_backend.h"
#include "stable-diffusion.h"
#include "tensor.hpp"

#include <memory>
#include <string>

struct UpscalerGGML {
    SDBackendManager backend_manager;
    ggml_type model_data_type = GGML_TYPE_F16;
    std::shared_ptr<ESRGAN> esrgan_upscaler;
    std::string esrgan_path;
    int n_threads;
    bool direct                 = false;
    int tile_size               = 128;
    size_t max_graph_vram_bytes = 0;
    bool stream_layers_enabled  = false;
    std::string backend_spec;
    std::string params_backend_spec;

    UpscalerGGML(int n_threads,
                 bool direct                     = false,
                 int tile_size                   = 128,
                 std::string backend_spec        = "",
                 std::string params_backend_spec = "");

    bool load_from_file(const std::string& esrgan_path,
                        bool offload_params_to_cpu,
                        int n_threads);
    void set_max_graph_vram_bytes(size_t max_vram_bytes);
    void set_stream_layers_enabled(bool enabled);
    sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
    sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
};

#endif  // __SD_UPSCALER_H__