mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-11 11:26:34 +00:00
refactor: reorganize src model layout (#1615)
This commit is contained in:
leejet authored and GitHub committed
1 parent
dfb2390dd4
commit
f3fd359b58
81 files changed
+407
-385
No files matched your search
@@ -0,0 +1,261 @@
|
||||
#ifndef __SD_MODEL_VAE_VAE_HPP__
|
||||
#define __SD_MODEL_VAE_VAE_HPP__
|
||||
|
||||
#include "core/tensor_ggml.hpp"
|
||||
#include "model/common/block.hpp"
|
||||
|
||||
struct VAE : public GGMLRunner {
|
||||
protected:
|
||||
SDVersion version;
|
||||
bool scale_input = true;
|
||||
virtual sd::Tensor<float> _compute(const int n_threads,
|
||||
const sd::Tensor<float>& z,
|
||||
bool decode_graph) = 0;
|
||||
|
||||
static inline void scale_tensor_to_minus1_1(sd::Tensor<float>* tensor) {
|
||||
GGML_ASSERT(tensor != nullptr);
|
||||
for (int64_t i = 0; i < tensor->numel(); ++i) {
|
||||
(*tensor)[i] = (*tensor)[i] * 2.0f - 1.0f;
|
||||
}
|
||||
}
|
||||
|
||||
static inline void scale_tensor_to_0_1(sd::Tensor<float>* tensor) {
|
||||
GGML_ASSERT(tensor != nullptr);
|
||||
for (int64_t i = 0; i < tensor->numel(); ++i) {
|
||||
float value = ((*tensor)[i] + 1.0f) * 0.5f;
|
||||
(*tensor)[i] = std::max(0.0f, std::min(1.0f, value));
|
||||
}
|
||||
}
|
||||
|
||||
sd::Tensor<float> tiled_compute(const sd::Tensor<float>& input,
|
||||
int n_threads,
|
||||
int output_width,
|
||||
int output_height,
|
||||
int scale,
|
||||
int p_tile_size_x,
|
||||
int p_tile_size_y,
|
||||
float tile_overlap_factor,
|
||||
bool circular_x,
|
||||
bool circular_y,
|
||||
bool decode_graph,
|
||||
const char* error_message,
|
||||
bool silent = false) {
|
||||
auto on_processing = [&](const sd::Tensor<float>& input_tile) {
|
||||
auto output_tile = _compute(n_threads, input_tile, decode_graph);
|
||||
if (output_tile.empty()) {
|
||||
LOG_ERROR("%s", error_message);
|
||||
return sd::Tensor<float>();
|
||||
}
|
||||
return output_tile;
|
||||
};
|
||||
return ::process_tiles_2d(input,
|
||||
output_width,
|
||||
output_height,
|
||||
scale,
|
||||
p_tile_size_x,
|
||||
p_tile_size_y,
|
||||
tile_overlap_factor,
|
||||
circular_x,
|
||||
circular_y,
|
||||
on_processing,
|
||||
silent);
|
||||
}
|
||||
|
||||
public:
|
||||
VAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend)
|
||||
: version(version), GGMLRunner(backend, params_backend) {}
|
||||
|
||||
int get_scale_factor() {
|
||||
int scale_factor = 8;
|
||||
if (version == VERSION_LTXAV) {
|
||||
scale_factor = 32;
|
||||
} else if (version == VERSION_WAN2_2_TI2V) {
|
||||
scale_factor = 16;
|
||||
} else if (sd_version_uses_flux2_vae(version)) {
|
||||
scale_factor = 16;
|
||||
} else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) {
|
||||
scale_factor = 1;
|
||||
}
|
||||
return scale_factor;
|
||||
}
|
||||
|
||||
virtual int get_encoder_output_channels(int input_channels) = 0;
|
||||
|
||||
void get_tile_sizes(int& tile_size_x,
|
||||
int& tile_size_y,
|
||||
float& tile_overlap,
|
||||
const sd_tiling_params_t& params,
|
||||
int64_t latent_x,
|
||||
int64_t latent_y,
|
||||
float encoding_factor = 1.0f) {
|
||||
tile_overlap = std::max(std::min(params.target_overlap, 0.5f), 0.0f);
|
||||
auto get_tile_size = [&](int requested_size, float factor, int64_t latent_size) {
|
||||
const int default_tile_size = 32;
|
||||
const int min_tile_dimension = 4;
|
||||
int tile_size = default_tile_size;
|
||||
// factor <= 1 means simple fraction of the latent dimension
|
||||
// factor > 1 means number of tiles across that dimension
|
||||
if (factor > 0.f) {
|
||||
if (factor > 1.0)
|
||||
factor = 1 / (factor - factor * tile_overlap + tile_overlap);
|
||||
tile_size = static_cast<int>(std::round(latent_size * factor));
|
||||
} else if (requested_size >= min_tile_dimension) {
|
||||
tile_size = requested_size;
|
||||
}
|
||||
tile_size = static_cast<int>(tile_size * encoding_factor);
|
||||
return std::max(std::min(tile_size, static_cast<int>(latent_size)), min_tile_dimension);
|
||||
};
|
||||
|
||||
tile_size_x = get_tile_size(params.tile_size_x, params.rel_size_x, latent_x);
|
||||
tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
|
||||
}
|
||||
|
||||
sd::Tensor<float> encode(int n_threads,
|
||||
const sd::Tensor<float>& x,
|
||||
sd_tiling_params_t tiling_params,
|
||||
bool circular_x = false,
|
||||
bool circular_y = false) {
|
||||
int64_t t0 = ggml_time_ms();
|
||||
sd::Tensor<float> input = x;
|
||||
sd::Tensor<float> output;
|
||||
if (scale_input) {
|
||||
scale_tensor_to_minus1_1(&input);
|
||||
}
|
||||
|
||||
if (tiling_params.enabled) {
|
||||
const int scale_factor = get_scale_factor();
|
||||
int64_t W = input.shape()[0] / scale_factor;
|
||||
int64_t H = input.shape()[1] / scale_factor;
|
||||
float tile_overlap;
|
||||
int tile_size_x, tile_size_y;
|
||||
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, 1.30539f);
|
||||
LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
|
||||
output = tiled_compute(input,
|
||||
n_threads,
|
||||
static_cast<int>(W),
|
||||
static_cast<int>(H),
|
||||
scale_factor,
|
||||
tile_size_x,
|
||||
tile_size_y,
|
||||
tile_overlap,
|
||||
circular_x,
|
||||
circular_y,
|
||||
false,
|
||||
"vae encode compute failed while processing a tile");
|
||||
} else {
|
||||
output = _compute(n_threads, input, false);
|
||||
}
|
||||
|
||||
free_compute_buffer();
|
||||
|
||||
if (output.empty()) {
|
||||
LOG_ERROR("vae encode compute failed");
|
||||
return {};
|
||||
}
|
||||
int64_t t1 = ggml_time_ms();
|
||||
LOG_DEBUG("computing vae encode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
|
||||
return std::move(output);
|
||||
}
|
||||
|
||||
sd::Tensor<float> decode(int n_threads,
|
||||
const sd::Tensor<float>& x,
|
||||
sd_tiling_params_t tiling_params,
|
||||
bool decode_video = false,
|
||||
bool circular_x = false,
|
||||
bool circular_y = false,
|
||||
bool silent = false) {
|
||||
int64_t t0 = ggml_time_ms();
|
||||
sd::Tensor<float> input = x;
|
||||
sd::Tensor<float> output;
|
||||
set_tiling_params(tiling_params);
|
||||
|
||||
if (tiling_params.enabled) {
|
||||
const int scale_factor = get_scale_factor();
|
||||
int64_t W = input.shape()[0] * scale_factor;
|
||||
int64_t H = input.shape()[1] * scale_factor;
|
||||
float tile_overlap;
|
||||
int tile_size_x, tile_size_y;
|
||||
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, input.shape()[0], input.shape()[1]);
|
||||
if (!silent) {
|
||||
LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
|
||||
}
|
||||
output = tiled_compute(
|
||||
input,
|
||||
n_threads,
|
||||
static_cast<int>(W),
|
||||
static_cast<int>(H),
|
||||
scale_factor,
|
||||
tile_size_x,
|
||||
tile_size_y,
|
||||
tile_overlap,
|
||||
circular_x,
|
||||
circular_y,
|
||||
true,
|
||||
"vae decode compute failed while processing a tile",
|
||||
silent);
|
||||
} else {
|
||||
output = _compute(n_threads, input, true);
|
||||
}
|
||||
|
||||
free_compute_buffer();
|
||||
|
||||
if (output.empty()) {
|
||||
LOG_ERROR("vae decode compute failed");
|
||||
return {};
|
||||
}
|
||||
if (scale_input) {
|
||||
scale_tensor_to_0_1(&output);
|
||||
}
|
||||
int64_t t1 = ggml_time_ms();
|
||||
LOG_DEBUG("computing vae decode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
|
||||
return std::move(output);
|
||||
}
|
||||
|
||||
virtual sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) = 0;
|
||||
virtual sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) = 0;
|
||||
virtual sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) = 0;
|
||||
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) = 0;
|
||||
virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
|
||||
virtual void set_temporal_tiling_enabled(bool enabled) { SD_UNUSED(enabled); };
|
||||
virtual void set_tiling_params(const sd_tiling_params_t& params) {
|
||||
set_temporal_tiling_enabled(params.temporal_tiling);
|
||||
};
|
||||
};
|
||||
|
||||
struct FakeVAE : public VAE {
|
||||
FakeVAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend)
|
||||
: VAE(version, backend, params_backend) {}
|
||||
|
||||
int get_encoder_output_channels(int input_channels) {
|
||||
return input_channels;
|
||||
}
|
||||
|
||||
sd::Tensor<float> _compute(const int n_threads,
|
||||
const sd::Tensor<float>& z,
|
||||
bool decode_graph) override {
|
||||
SD_UNUSED(n_threads);
|
||||
SD_UNUSED(decode_graph);
|
||||
return z;
|
||||
}
|
||||
|
||||
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
|
||||
SD_UNUSED(rng);
|
||||
return vae_output;
|
||||
}
|
||||
|
||||
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
|
||||
return latents;
|
||||
}
|
||||
|
||||
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
|
||||
return latents;
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {}
|
||||
|
||||
std::string get_desc() override {
|
||||
return "fake_vae";
|
||||
}
|
||||
};
|
||||
|
||||
#endif // __SD_MODEL_VAE_VAE_HPP__
|
||||
Reference in new issue
Block a user