mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-12 03:49:09 +00:00
feat: add RelayVSR video upscaling with CUDA window attention (#2285)
This commit is contained in:
leejet authored and GitHub committed
1 parent
b1f5a51a1e
commit
9234ad822a
21 files changed
+1764
-32
No files matched your search
@@ -0,0 +1,209 @@
|
||||
#include "pipeline/video_upscale.h"
|
||||
|
||||
#include <climits>
|
||||
|
||||
#include "core/rng_philox.hpp"
|
||||
#include "generation_info.h"
|
||||
#include "model/diffusion/relay_vsr.hpp"
|
||||
#include "model/upscaler/flash_vsr.hpp"
|
||||
#include "name_conversion.h"
|
||||
|
||||
namespace sd {
|
||||
namespace pipeline {
|
||||
|
||||
static sd::Tensor<float> prepare_vsr_frame(const sd_image_t& frame, int width, int height) {
|
||||
auto result = sd_image_to_tensor(frame, -1, -1, false);
|
||||
result /= 127.5f;
|
||||
result -= 1.f;
|
||||
if (width > int(frame.width)) {
|
||||
auto edge = sd::ops::slice(result, 0, frame.width - 1, frame.width);
|
||||
auto padding = sd::zeros<float>({width - frame.width, frame.height, 3, 1}) + edge;
|
||||
result = sd::ops::concat(result, padding, 0);
|
||||
}
|
||||
if (height > int(frame.height)) {
|
||||
auto edge = sd::ops::slice(result, 1, frame.height - 1, frame.height);
|
||||
auto padding = sd::zeros<float>({width, height - frame.height, 3, 1}) + edge;
|
||||
result = sd::ops::concat(result, padding, 1);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
static sd::Tensor<float> prepare_vsr_lq(const sd::Tensor<float>& frame, int scale) {
|
||||
auto input = frame;
|
||||
if (scale == 4)
|
||||
input = sd::ops::interpolate(frame, {frame.shape()[0] * 4, frame.shape()[1] * 4, 3, 1}, sd::ops::InterpolateMode::Bilinear);
|
||||
const int patch = 4 * scale;
|
||||
const int width = int(input.shape()[0]), height = int(input.shape()[1]);
|
||||
input.reshape_({patch, width / patch, patch, height / patch, 3, 1});
|
||||
input.permute_({1, 3, 0, 2, 4, 5});
|
||||
input.reshape_({width / patch, height / patch, 3 * patch * patch, 1});
|
||||
return input;
|
||||
}
|
||||
|
||||
class RelayVSRUpscaler : public VideoUpscaleModel {
|
||||
std::unique_ptr<RelayVSR::WanVSRRunner> wan_;
|
||||
std::unique_ptr<RelayVSR::FlashVSRRunner> flash_;
|
||||
int scale_ = 0;
|
||||
int n_threads_ = 1;
|
||||
|
||||
public:
|
||||
bool init(VideoUpscaleRuntime& runtime, const String2TensorStorage& source_weights, int upscale_factor) override {
|
||||
if (upscale_factor != 2 && upscale_factor != 4) {
|
||||
LOG_ERROR("RelayVSR requires an explicit upscale factor of 2 or 4 matching both checkpoints");
|
||||
return false;
|
||||
}
|
||||
scale_ = upscale_factor;
|
||||
n_threads_ = runtime.n_threads;
|
||||
auto wan_backend = runtime.backends.runtime_backend(SDBackendModule::DIFFUSION);
|
||||
auto flash_backend = runtime.backends.runtime_backend(SDBackendModule::VAE);
|
||||
auto wan_params = runtime.backends.params_backend(SDBackendModule::DIFFUSION);
|
||||
auto flash_params = runtime.backends.params_backend(SDBackendModule::VAE);
|
||||
String2TensorStorage weights;
|
||||
std::map<std::string, std::string> source_names;
|
||||
for (const auto& [name, storage] : source_weights) {
|
||||
auto canonical = name;
|
||||
if (starts_with(canonical, "diffusion."))
|
||||
canonical = "wan." + canonical.substr(10);
|
||||
else if (starts_with(canonical, "vae."))
|
||||
canonical = "flash." + canonical.substr(4);
|
||||
canonical = convert_relayvsr_tensor_name(canonical);
|
||||
if (!source_names.emplace(canonical, name).second)
|
||||
throw std::runtime_error("Duplicate RelayVSR weight name: " + canonical);
|
||||
weights[canonical] = storage;
|
||||
weights[canonical].name = canonical;
|
||||
}
|
||||
wan_ = std::make_unique<RelayVSR::WanVSRRunner>(wan_backend, weights, upscale_factor, runtime.manager);
|
||||
flash_ = std::make_unique<RelayVSR::FlashVSRRunner>(flash_backend, weights, upscale_factor, runtime.manager);
|
||||
auto register_model = [&](auto& runner, ModelComponent component, SDBackendModule module, ggml_backend_t compute, ggml_backend_t params) {
|
||||
std::map<std::string, ggml_tensor*> tensors, sources;
|
||||
runner->get_param_tensors(tensors);
|
||||
for (const auto& [name, tensor] : tensors) {
|
||||
auto source = source_names.find(name);
|
||||
if (source == source_names.end())
|
||||
throw std::runtime_error("Missing RelayVSR tensor: " + name);
|
||||
sources.emplace(source->second, tensor);
|
||||
ggml_set_name(tensor, source->second.c_str());
|
||||
}
|
||||
runner->set_flash_attention_enabled(true);
|
||||
return runtime.manager->register_param_tensors(component, std::move(sources),
|
||||
runtime.backends.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend,
|
||||
compute, params);
|
||||
};
|
||||
if (!register_model(wan_, ModelComponent::Diffusion, SDBackendModule::DIFFUSION, wan_backend, wan_params) ||
|
||||
!register_model(flash_, ModelComponent::VAE, SDBackendModule::VAE, flash_backend, flash_params) ||
|
||||
!runtime.manager->validate_registered_tensors())
|
||||
return false;
|
||||
return true;
|
||||
}
|
||||
|
||||
void end() override {
|
||||
if (wan_)
|
||||
wan_->runner_end();
|
||||
if (flash_)
|
||||
flash_->runner_end();
|
||||
}
|
||||
|
||||
bool upscale(const sd_video_upscale_params_t& params, sd_video_generation_result_t* result) override {
|
||||
const auto* frames = params.frames;
|
||||
const int frame_count = params.frame_count, fps = params.fps, window_size = params.window_size;
|
||||
const int64_t seed = params.seed;
|
||||
if (window_size < 2) {
|
||||
LOG_ERROR("RelayVSR requires a reference window of at least two frames");
|
||||
return false;
|
||||
}
|
||||
if (frames[0].width > uint32_t(INT_MAX / scale_ - 32) || frames[0].height > uint32_t(INT_MAX / scale_ - 32))
|
||||
return false;
|
||||
struct Reset {
|
||||
RelayVSRUpscaler* model;
|
||||
~Reset() {
|
||||
model->wan_->reset();
|
||||
model->flash_->reset();
|
||||
}
|
||||
} reset{this};
|
||||
wan_->reset();
|
||||
flash_->reset();
|
||||
const int stride = 16 / scale_;
|
||||
const int alignment = 2 * stride;
|
||||
const int width = (int(frames[0].width) + alignment - 1) / alignment * alignment;
|
||||
const int height = (int(frames[0].height) + alignment - 1) / alignment * alignment;
|
||||
PhiloxRNG rng;
|
||||
rng.manual_seed(uint64_t(seed));
|
||||
auto keyframe = [&](int index, int slot) {
|
||||
LOG_INFO("RelayVSR reference frame %d/%d", index + 1, frame_count);
|
||||
auto input = prepare_vsr_frame(frames[index], width, height);
|
||||
sd::Tensor<float> noise({width / stride, height / stride, 1, 48});
|
||||
auto values = rng.randn(uint32_t(noise.numel()));
|
||||
std::copy(values.begin(), values.end(), noise.data());
|
||||
auto latent = wan_->compute(n_threads_, noise, prepare_vsr_lq(input, scale_));
|
||||
if (latent.empty())
|
||||
return false;
|
||||
latent.reshape_({width / stride, height / stride, 48, 1});
|
||||
return !flash_->compute(n_threads_, latent, slot).empty();
|
||||
};
|
||||
int start_slot = 0;
|
||||
if (!keyframe(0, start_slot))
|
||||
return false;
|
||||
for (int first = 0; first < frame_count;) {
|
||||
const int last = first + std::min(window_size - 1, frame_count - 1 - first);
|
||||
int end_slot = last == first ? start_slot : 1 - start_slot;
|
||||
if (last != first && !keyframe(last, end_slot))
|
||||
return false;
|
||||
flash_->history = false;
|
||||
for (int i = first; i <= last; ++i) {
|
||||
auto input = prepare_vsr_frame(frames[i], width, height);
|
||||
auto residual = flash_->compute(n_threads_, input, -1, start_slot, end_slot);
|
||||
if (residual.empty())
|
||||
return false;
|
||||
if (first != 0 && i == first)
|
||||
continue;
|
||||
auto output = sd::ops::interpolate(input, {width * scale_, height * scale_, 3, 1}, sd::ops::InterpolateMode::Bicubic);
|
||||
auto& image = result->frames[i];
|
||||
image.width = frames[i].width * scale_;
|
||||
image.height = frames[i].height * scale_;
|
||||
image.channel = 3;
|
||||
output += residual;
|
||||
output += 1.f;
|
||||
output *= 127.5f;
|
||||
output = sd::ops::slice(sd::ops::slice(output, 0, 0, image.width), 1, 0, image.height);
|
||||
if (!std::all_of(output.values().begin(), output.values().end(), [](float value) { return std::isfinite(value); }))
|
||||
throw std::runtime_error("RelayVSR produced non-finite pixels");
|
||||
output = sd::ops::clamp(output, 0.f, 255.f);
|
||||
output.permute_({2, 0, 1, 3});
|
||||
image.data = static_cast<uint8_t*>(malloc(size_t(image.width) * image.height * 3));
|
||||
if (!image.data)
|
||||
throw std::bad_alloc();
|
||||
std::transform(output.values().begin(), output.values().end(), image.data, [](float value) { return uint8_t(value); });
|
||||
LOG_INFO("RelayVSR output frame %d/%d", i + 1, frame_count);
|
||||
if (auto progress = sd_get_progress_callback())
|
||||
progress(i + 1, frame_count, 0.f, sd_get_progress_callback_data());
|
||||
}
|
||||
if (last == frame_count - 1)
|
||||
break;
|
||||
first = last;
|
||||
start_slot = end_slot;
|
||||
}
|
||||
auto info = sd::upscale_generation_info(scale_);
|
||||
info["kind"] = "video";
|
||||
info["upscale"]["model"] = "relayvsr";
|
||||
info["upscale"]["window_size"] = window_size;
|
||||
info["seed"] = seed;
|
||||
info["fps"] = fps;
|
||||
info["video_frames"] = frame_count;
|
||||
result->generation_info = sd::copy_generation_info(info);
|
||||
if (!result->generation_info)
|
||||
throw std::bad_alloc();
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
bool is_relay_vsr_model(const String2TensorStorage& weights) {
|
||||
return (weights.count("diffusion.positive_conditioning") || weights.count("diffusion.prompt_embeds")) &&
|
||||
(weights.count("diffusion.lq_proj_in.conv1.weight") || weights.count("diffusion.lq_proj.conv1.weight"));
|
||||
}
|
||||
|
||||
std::unique_ptr<VideoUpscaleModel> create_relay_vsr_upscaler() {
|
||||
return std::make_unique<RelayVSRUpscaler>();
|
||||
}
|
||||
|
||||
} // namespace pipeline
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,143 @@
|
||||
#include "pipeline/video_upscale.h"
|
||||
|
||||
#include <climits>
|
||||
|
||||
#include "core/util.h"
|
||||
|
||||
namespace sd {
|
||||
namespace pipeline {
|
||||
|
||||
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 (!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 {
|
||||
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;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
#ifndef __SD_PIPELINE_VIDEO_UPSCALE_H__
|
||||
#define __SD_PIPELINE_VIDEO_UPSCALE_H__
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "core/ggml_extend_backend.h"
|
||||
#include "core/tensor.hpp"
|
||||
#include "model_manager.h"
|
||||
#include "stable-diffusion.h"
|
||||
|
||||
namespace sd {
|
||||
namespace pipeline {
|
||||
|
||||
struct VideoUpscaleRuntime {
|
||||
SDBackendManager backends;
|
||||
std::shared_ptr<ModelManager> manager;
|
||||
int n_threads = 1;
|
||||
};
|
||||
|
||||
class VideoUpscaleModel {
|
||||
public:
|
||||
virtual ~VideoUpscaleModel() = default;
|
||||
virtual bool init(VideoUpscaleRuntime& runtime, const String2TensorStorage& weights, int upscale_factor) = 0;
|
||||
virtual void end() = 0;
|
||||
virtual bool upscale(const sd_video_upscale_params_t& params, sd_video_generation_result_t* result) = 0;
|
||||
};
|
||||
|
||||
bool is_relay_vsr_model(const String2TensorStorage& weights);
|
||||
std::unique_ptr<VideoUpscaleModel> create_relay_vsr_upscaler();
|
||||
|
||||
class VideoUpscaler {
|
||||
VideoUpscaleRuntime runtime_;
|
||||
int upscale_factor_ = 0;
|
||||
std::unique_ptr<sd::ParallelExecutor> executor_;
|
||||
std::unique_ptr<VideoUpscaleModel> model_;
|
||||
|
||||
public:
|
||||
~VideoUpscaler();
|
||||
bool init(const sd_video_upscaler_ctx_params_t& params);
|
||||
bool upscale(const sd_video_upscale_params_t& params, sd_video_generation_result_t** result_out);
|
||||
};
|
||||
|
||||
} // namespace pipeline
|
||||
} // namespace sd
|
||||
|
||||
#endif // __SD_PIPELINE_VIDEO_UPSCALE_H__
|
||||
Reference in new issue
Block a user