mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-09 18:36:33 +00:00
203 lines
9.1 KiB
C++
203 lines
9.1 KiB
C++
#include "pipeline/video_upscale.h"
|
|
|
|
#include <climits>
|
|
|
|
#include "core/ggml_runner.h"
|
|
#include "core/probe/conv_direct.h"
|
|
#include "core/probe/flash_attn.h"
|
|
#include "core/util.h"
|
|
|
|
namespace sd {
|
|
namespace pipeline {
|
|
|
|
bool VideoUpscaleRuntime::flash_attention_for(SDBackendModule module) {
|
|
const auto mode = flash_attn_assignment.get(module);
|
|
if (mode == "on" || mode == "off")
|
|
return mode == "on";
|
|
const auto devices = backends.runtime_backends(module);
|
|
bool enabled = !devices.empty(), any_enabled = false;
|
|
for (auto backend : devices) {
|
|
auto it = flash_attn_by_backend.find(backend);
|
|
if (it == flash_attn_by_backend.end())
|
|
it = flash_attn_by_backend.emplace(backend, sd_probe_flash_attention(backend, n_threads)).first;
|
|
enabled = enabled && it->second;
|
|
any_enabled = any_enabled || it->second;
|
|
}
|
|
if (any_enabled && !enabled)
|
|
LOG_INFO("flash attention auto: %s -> off (compute devices disagree)", sd_backend_module_name(module));
|
|
return enabled;
|
|
}
|
|
|
|
bool VideoUpscaleRuntime::conv_direct_for(SDBackendModule module, bool conv3d) {
|
|
const auto mode = conv_direct_assignment.get(module);
|
|
if (mode == "on" || mode == "off")
|
|
return mode == "on";
|
|
auto& cache = conv3d ? conv3d_direct_by_backend : conv2d_direct_by_backend;
|
|
const auto devices = backends.runtime_backends(module);
|
|
bool enabled = !devices.empty(), any_enabled = false;
|
|
for (auto backend : devices) {
|
|
auto it = cache.find(backend);
|
|
if (it == cache.end())
|
|
it = cache.emplace(backend, sd_probe_conv_direct(backend, conv3d)).first;
|
|
enabled = enabled && it->second;
|
|
any_enabled = any_enabled || it->second;
|
|
}
|
|
if (any_enabled && !enabled)
|
|
LOG_INFO("conv%dd direct auto: %s -> off (compute devices disagree)", conv3d ? 3 : 2, sd_backend_module_name(module));
|
|
return enabled;
|
|
}
|
|
|
|
void VideoUpscaleRuntime::configure_runner(GGMLRunner& runner, SDBackendModule module, bool conv2d, bool conv3d) {
|
|
const bool flash = flash_attention_for(module);
|
|
const bool direct2d = conv2d && conv_direct_for(module, false);
|
|
const bool direct3d = conv3d && conv_direct_for(module, true);
|
|
runner.set_flash_attention_enabled(flash);
|
|
runner.set_conv2d_direct_enabled(direct2d);
|
|
runner.set_conv3d_direct_enabled(direct3d);
|
|
LOG_INFO("video upscale %s: flash attention %s, conv2d direct %s, conv3d direct %s",
|
|
sd_backend_module_name(module), flash ? "on" : "off", direct2d ? "on" : "off", direct3d ? "on" : "off");
|
|
}
|
|
|
|
VideoUpscaler::~VideoUpscaler() {
|
|
if (model_)
|
|
model_->end();
|
|
// ModelManager retains parameter pointers owned by the model runners.
|
|
runtime_.manager.reset();
|
|
model_.reset();
|
|
}
|
|
|
|
bool VideoUpscaler::init(const sd_video_upscaler_ctx_params_t& params) {
|
|
if (!params.diffusion_model_path || !*params.diffusion_model_path ||
|
|
!params.vae_path || !*params.vae_path || params.upscale_factor <= 0) {
|
|
LOG_ERROR("video upscale requires diffusion and decoder model paths and a positive upscale factor");
|
|
return false;
|
|
}
|
|
upscale_factor_ = params.upscale_factor;
|
|
runtime_.n_threads = params.n_threads > 0 ? params.n_threads : sd_get_num_physical_cores();
|
|
executor_ = std::make_unique<sd::ParallelExecutor>(runtime_.n_threads);
|
|
sd::ParallelScope scope(executor_.get());
|
|
std::string error;
|
|
if (!sd_parse_flash_attention_assignment(params.flash_attn_mode ? params.flash_attn_mode : "", &runtime_.flash_attn_assignment, &error) ||
|
|
!sd_parse_conv_direct_assignment(params.conv_direct_mode ? params.conv_direct_mode : "", &runtime_.conv_direct_assignment, &error)) {
|
|
LOG_ERROR("video upscale operator configuration: %s", error.c_str());
|
|
return false;
|
|
}
|
|
if (!runtime_.backends.init(params.backend, params.params_backend, nullptr, &error)) {
|
|
LOG_ERROR("video upscaler backend: %s", error.c_str());
|
|
return false;
|
|
}
|
|
for (auto module : {SDBackendModule::DIFFUSION, SDBackendModule::VAE}) {
|
|
if (!runtime_.backends.runtime_backend(module) || !runtime_.backends.params_backend(module))
|
|
return false;
|
|
}
|
|
ModelLoader loader;
|
|
if (!loader.init_from_file(params.diffusion_model_path, "diffusion.") ||
|
|
!loader.init_from_file(params.vae_path, "vae."))
|
|
return false;
|
|
const auto& weights = loader.get_tensor_storage_map();
|
|
if (is_relay_vsr_model(weights))
|
|
model_ = create_relay_vsr_upscaler();
|
|
else if (is_seedvr2_model(weights))
|
|
model_ = create_seedvr2_upscaler();
|
|
else {
|
|
LOG_ERROR("unsupported video upscale model architecture");
|
|
return false;
|
|
}
|
|
runtime_.manager = std::make_shared<ModelManager>();
|
|
runtime_.manager->set_n_threads(runtime_.n_threads);
|
|
if (!runtime_.manager->set_loader(loader))
|
|
return false;
|
|
return model_->init(runtime_, weights, upscale_factor_);
|
|
}
|
|
|
|
bool VideoUpscaler::upscale(const sd_video_upscale_params_t& params, sd_video_generation_result_t** result_out) {
|
|
if (!result_out)
|
|
return false;
|
|
*result_out = nullptr;
|
|
if (!model_ || !params.frames || params.frame_count <= 0 || params.fps <= 0) {
|
|
LOG_ERROR("video upscale requires nonempty RGB frames and positive FPS");
|
|
return false;
|
|
}
|
|
const auto& first = params.frames[0];
|
|
if (!first.width || !first.height || first.width > uint32_t(INT_MAX / upscale_factor_) ||
|
|
first.height > uint32_t(INT_MAX / upscale_factor_))
|
|
return false;
|
|
for (int i = 0; i < params.frame_count; ++i) {
|
|
const auto& frame = params.frames[i];
|
|
if (!frame.data || frame.channel != 3 || frame.width != first.width || frame.height != first.height) {
|
|
LOG_ERROR("video upscale input frames must have equal dimensions and three RGB channels");
|
|
return false;
|
|
}
|
|
}
|
|
sd::ParallelScope scope(executor_.get());
|
|
auto result = std::unique_ptr<sd_video_generation_result_t, decltype(&free_sd_video_generation_result)>(
|
|
static_cast<sd_video_generation_result_t*>(calloc(1, sizeof(sd_video_generation_result_t))), free_sd_video_generation_result);
|
|
if (!result)
|
|
return false;
|
|
result->frames = static_cast<sd_image_t*>(calloc(params.frame_count, sizeof(sd_image_t)));
|
|
if (!result->frames)
|
|
return false;
|
|
result->frame_count = params.frame_count;
|
|
result->fps = params.fps;
|
|
if (!model_->upscale(params, result.get()))
|
|
return false;
|
|
*result_out = result.release();
|
|
return true;
|
|
}
|
|
|
|
} // namespace pipeline
|
|
} // namespace sd
|
|
|
|
struct video_upscaler_ctx_t {
|
|
sd::pipeline::VideoUpscaler upscaler;
|
|
};
|
|
|
|
void sd_video_upscaler_ctx_params_init(sd_video_upscaler_ctx_params_t* params) {
|
|
if (!params)
|
|
return;
|
|
*params = {};
|
|
params->n_threads = sd_get_num_physical_cores();
|
|
}
|
|
|
|
void sd_video_upscale_params_init(sd_video_upscale_params_t* params) {
|
|
if (!params)
|
|
return;
|
|
*params = {};
|
|
params->fps = 24;
|
|
params->seed = 42;
|
|
params->window_size = 16;
|
|
params->temporal_chunk_size = 5;
|
|
}
|
|
|
|
video_upscaler_ctx_t* new_video_upscaler_ctx(const sd_video_upscaler_ctx_params_t* params) {
|
|
if (!params)
|
|
return nullptr;
|
|
try {
|
|
auto ctx = std::make_unique<video_upscaler_ctx_t>();
|
|
if (!ctx->upscaler.init(*params))
|
|
return nullptr;
|
|
return ctx.release();
|
|
} catch (const std::exception& error) {
|
|
LOG_ERROR("video upscaler initialization failed: %s", error.what());
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
bool upscale_video(video_upscaler_ctx_t* ctx, const sd_video_upscale_params_t* params, sd_video_generation_result_t** result_out) {
|
|
if (!result_out)
|
|
return false;
|
|
*result_out = nullptr;
|
|
if (!ctx || !params)
|
|
return false;
|
|
try {
|
|
return ctx->upscaler.upscale(*params, result_out);
|
|
} catch (const std::exception& error) {
|
|
LOG_ERROR("video upscale failed: %s", error.what());
|
|
return false;
|
|
}
|
|
}
|
|
|
|
void free_video_upscaler_ctx(video_upscaler_ctx_t* ctx) {
|
|
delete ctx;
|
|
}
|