Commit 9234ad8 for stable-diffusion.cpp
commit 9234ad822a376a4724638467736118ed61d68ee9
Author: leejet <leejet714@gmail.com>
Date: Sun Nov 8 00:09:48 2026 +0800
feat: add RelayVSR video upscaling with CUDA window attention (#2285)
diff --git a/README.md b/README.md
index 28af00a..2442865 100644
--- a/README.md
+++ b/README.md
@@ -96,6 +96,7 @@ API and command-line option may change frequently.***
- Latent Consistency Models support (LCM/LCM-LoRA)
- Faster and memory efficient latent decoding with [TAESD](./docs/taesd.md)
- Upscale images generated with [ESRGAN](./docs/esrgan.md)
+ - Upscale videos with [RelayVSR](./docs/relay_vsr.md) (experimental, 2x and 4x)
- Supported backends
- CPU (AVX, AVX2 and AVX512 support for x86 architectures)
- CUDA
diff --git a/docs/relay_vsr.md b/docs/relay_vsr.md
new file mode 100644
index 0000000..1352f41
--- /dev/null
+++ b/docs/relay_vsr.md
@@ -0,0 +1,70 @@
+# RelayVSR video super-resolution
+
+RelayVSR uses a Wan model to sample sparse reference latents, followed by the
+Flash conditional decoder to reconstruct every input frame. The model code
+follows [RelayVSR](https://github.com/kopperx/RelayVSR); the loader maps
+[ComfyUI repackaged weights](https://huggingface.co/Comfy-Org/RelayVSR) to the
+original module names.
+
+Use `-M vid_upscale` with an ordered directory of RGB images. Filenames are
+sorted lexicographically. Supply the source frame rate with `--fps`. Output
+keeps the input frame count. Use a video output supported by the CLI, such as
+`.avi` or `.webm`, or a numbered PNG output such as `output/%04d.png`.
+
+The scale is an explicit, required argument. It is not read from checkpoint
+metadata. Use matching Wan and Flash checkpoints:
+
+| Scale | `--diffusion-model` | `--vae` |
+| --- | --- | --- |
+| 2 | `relayvsr_x2_lightx2v_int8_convrot.safetensors` | `relayvsr_flash_vae_x2_fp32.safetensors` |
+| 4 | `relayvsr_x4_int8_convrot.safetensors` | `relayvsr_flash_vae_x4_fp32.safetensors` |
+
+```sh
+sd-cli -M vid_upscale \
+ --diffusion-model models/diffusion_models/relayvsr_x2_lightx2v_int8_convrot.safetensors \
+ --vae models/vae/relayvsr_flash_vae_x2_fp32.safetensors \
+ --input-frames input_frames --upscale-factor 2 --fps 24 \
+ --vsr-window-size 16 -s 42 -o output.avi
+```
+
+`--vsr-window-size` includes both reference endpoints. The default 16 samples
+references at frame indices 0, 15, 30, and so on, always including the last
+frame. Shared endpoints update Flash history in the following window but are
+emitted only once. Single-frame input is also supported.
+
+Wan performs one flow-matching Euler step at timestep 1000 with guidance 1,
+using the prompt embedding stored in the checkpoint. No text encoder or
+ordinary Wan VAE is needed. Generic generation sampling settings do not
+configure this fixed sampling path.
+
+The Flash scale cannot be verified from its tensor shapes: the x2 and x4 files
+have identical shapes. The caller must select the correct pair. Wan's input
+projection is checked against the requested scale. Inputs are replicate-padded
+internally and outputs are cropped to exactly the input size times the scale.
+
+INT8 convrot uses the existing [quantized-weight execution path](int8_convrot.md).
+Backend and parameter placement use `--backend` and `--params-backend`, with
+Wan assigned to `diffusion` and Flash to `vae`. CUDA builds with CUTLASS use
+fused window attention on compute capability 8.0 or newer. It loads spatial
+neighborhoods directly, with FP16 Q/K/V and FP32 accumulation, and preserves
+Wan's boundary-shifted windows and Flash's boundary-truncated windows.
+Other backends use bounded spatial tiles with ordinary attention operators.
+Set `GGML_CUDA_DISABLE_WINDOW_ATTENTION=1` to use that fallback on CUDA for
+comparison. Floating-point results can differ between these implementations.
+
+Library callers initialize `sd_video_upscaler_ctx_params_t` with
+`sd_video_upscaler_ctx_params_init()`, set the model paths, explicit scale and
+backend options, and create a `video_upscaler_ctx_t` with
+`new_video_upscaler_ctx()`. Model architecture is detected from
+weight names and shapes; the scale remains caller-supplied.
+
+Initialize `sd_video_upscale_params_t` with `sd_video_upscale_params_init()`,
+then set the RGB frames, frame count, FPS, seed and RelayVSR reference window.
+Call `upscale_video()` and free the result with
+`free_sd_video_generation_result()`. Defaults are 24 FPS, seed 42 and a
+16-frame reference window. No sampler or scheduler selection is required.
+
+Each call starts fresh
+temporal state. Model caches are bounded by the temporal windows; the public
+API currently holds the complete input and output frame arrays in host memory.
+Audio is not copied by this path.
diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
index 591a151..bcf7359 100644
--- a/examples/cli/main.cpp
+++ b/examples/cli/main.cpp
@@ -3,6 +3,7 @@
#include <time.h>
#include <algorithm>
#include <cctype>
+#include <climits>
#include <filesystem>
#include <functional>
#include <iostream>
@@ -33,10 +34,13 @@ const char* previews_str[] = {
std::regex format_specifier_regex("(?:[^%]|^)(?:%%)*(%\\d{0,3}d)");
struct SDCliParams {
- SDMode mode = IMG_GEN;
- std::string output_path = "output.png";
- int output_begin_idx = -1;
- int compression_quality = 90;
+ std::string input_frames;
+ int video_upscale_factor = 0;
+ int vsr_window_size = 16;
+ SDMode mode = IMG_GEN;
+ std::string output_path = "output.png";
+ int output_begin_idx = -1;
+ int compression_quality = 90;
std::string image_path;
std::string metadata_format = "text";
@@ -62,6 +66,7 @@ struct SDCliParams {
ArgOptions options;
options.string_options = {
+ {"", "--input-frames", "directory of ordered RGB frames for vid_upscale", 0, &input_frames},
{"-o",
"--output",
"path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png). Single-file video outputs support .avi, .webm, and animated .webp",
@@ -90,6 +95,8 @@ struct SDCliParams {
};
options.int_options = {
+ {"", "--upscale-factor", "required video upscale factor (RelayVSR: 2 or 4)", &video_upscale_factor},
+ {"", "--vsr-window-size", "RelayVSR window including both endpoints (default: 16)", &vsr_window_size},
{"",
"--preview-interval",
"preview interval: in each sampling pass, positive N updates every Nth denoiser step and -N previews only completed logical step N; 0 previews the final completed step of the first pass (base-resolution or high-noise). Default: 1",
@@ -190,7 +197,7 @@ struct SDCliParams {
options.manual_options = {
{"-M",
"--mode",
- "run mode, one of [img_gen, adetailer, vid_gen, upscale, convert, metadata], default: img_gen",
+ "run mode, one of [img_gen, adetailer, vid_gen, vid_upscale, upscale, convert, metadata], default: img_gen",
on_mode_arg},
{"",
"--preview",
@@ -220,6 +227,10 @@ struct SDCliParams {
}
bool validate() {
+ if (mode == VID_UPSCALE && (input_frames.empty() || video_upscale_factor <= 0 || vsr_window_size < 2)) {
+ LOG_ERROR("vid_upscale requires --input-frames, positive --upscale-factor, and --vsr-window-size >= 2");
+ return false;
+ }
if (mode != METADATA) {
if (output_path.length() == 0) {
LOG_ERROR("error: the following arguments are required: output_path");
@@ -529,7 +540,7 @@ bool save_results(const SDCliParams& cli_params,
return sucessful_reults != 0;
}
- if (cli_params.mode == VID_GEN && num_results > 1) {
+ if ((cli_params.mode == VID_GEN && num_results > 1) || cli_params.mode == VID_UPSCALE) {
if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm")
ext = ".avi";
std::string params = gen_params.embed_image_metadata
@@ -664,6 +675,50 @@ int main(int argc, const char* argv[]) {
return 0;
}
+ if (cli_params.mode == VID_UPSCALE) {
+ if (ctx_params.diffusion_model_path.empty() || ctx_params.vae_path.empty() || gen_params.fps <= 0) {
+ LOG_ERROR("vid_upscale requires --diffusion-model, --vae and positive --fps");
+ return 1;
+ }
+ std::vector<SDImageOwner> owned_frames;
+ if (!load_images_from_dir(cli_params.input_frames, owned_frames) || owned_frames.empty()) {
+ LOG_ERROR("vid_upscale input contains no readable frames");
+ return 1;
+ }
+ if (owned_frames.size() > size_t(INT_MAX))
+ return 1;
+ std::vector<sd_image_t> frames;
+ for (const auto& frame : owned_frames)
+ frames.push_back(frame.get());
+ sd_video_upscaler_ctx_params_t params;
+ sd_video_upscaler_ctx_params_init(¶ms);
+ params.diffusion_model_path = ctx_params.diffusion_model_path.c_str();
+ params.vae_path = ctx_params.vae_path.c_str();
+ params.upscale_factor = cli_params.video_upscale_factor;
+ params.n_threads = ctx_params.n_threads;
+ params.backend = ctx_params.effective_backend.c_str();
+ params.params_backend = ctx_params.effective_params_backend.c_str();
+ std::unique_ptr<video_upscaler_ctx_t, decltype(&free_video_upscaler_ctx)> upscaler(
+ new_video_upscaler_ctx(¶ms),
+ free_video_upscaler_ctx);
+ if (!upscaler)
+ return 1;
+ sd_video_generation_result_t* output = nullptr;
+ sd_video_upscale_params_t upscale_params;
+ sd_video_upscale_params_init(&upscale_params);
+ upscale_params.frames = frames.data();
+ upscale_params.frame_count = int(frames.size());
+ upscale_params.fps = gen_params.fps;
+ upscale_params.window_size = cli_params.vsr_window_size;
+ upscale_params.seed = gen_params.seed;
+ if (!upscale_video(upscaler.get(), &upscale_params, &output))
+ return 1;
+ std::string video_info = output->generation_info;
+ SDImageVec results;
+ results.adopt_video(output);
+ return save_results(cli_params, ctx_params, gen_params, results, int(results.size()), nullptr, video_info) ? 0 : 1;
+ }
+
if (!gen_params.ad_model_path.empty() && cli_params.mode != IMG_GEN && cli_params.mode != ADETAILER) {
LOG_ERROR("--ad-model is only supported in image generation and adetailer modes");
return 1;
diff --git a/examples/common/common.cpp b/examples/common/common.cpp
index 02c3088..7934924 100644
--- a/examples/common/common.cpp
+++ b/examples/common/common.cpp
@@ -35,6 +35,7 @@ const char* const modes_str[] = {
"convert",
"upscale",
"metadata",
+ "vid_upscale",
};
static sd_vae_format_t str_to_vae_format(const std::string& value) {
diff --git a/examples/common/common.h b/examples/common/common.h
index 4ecbef4..d4df4fa 100644
--- a/examples/common/common.h
+++ b/examples/common/common.h
@@ -22,7 +22,7 @@ std::string argv_to_utf8(int index, const char** argv);
#define BOOL_STR(b) ((b) ? "true" : "false")
extern const char* const modes_str[];
-#define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata"
+#define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata, vid_upscale"
enum SDMode {
IMG_GEN,
@@ -31,6 +31,7 @@ enum SDMode {
CONVERT,
UPSCALE,
METADATA,
+ VID_UPSCALE,
MODE_COUNT
};
diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h
index 1c61e4b..a1d82e7 100644
--- a/include/stable-diffusion.h
+++ b/include/stable-diffusion.h
@@ -615,6 +615,37 @@ SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
// Reads model metadata only; returns 0 if the file is not a recognized RGB ESRGAN model.
SD_API int get_upscaler_model_scale(const char* model_path);
+typedef struct video_upscaler_ctx_t video_upscaler_ctx_t;
+
+typedef struct {
+ const char* diffusion_model_path;
+ const char* vae_path;
+ // Explicit output scale; RelayVSR requires 2 or 4 matching both checkpoints.
+ int upscale_factor;
+ int n_threads;
+ const char* backend;
+ const char* params_backend;
+} sd_video_upscaler_ctx_params_t;
+
+typedef struct {
+ const sd_image_t* frames;
+ int frame_count;
+ int fps;
+ int64_t seed;
+ // RelayVSR reference window, including both endpoints; defaults to 16.
+ int window_size;
+} sd_video_upscale_params_t;
+
+SD_API void sd_video_upscaler_ctx_params_init(sd_video_upscaler_ctx_params_t* params);
+SD_API void sd_video_upscale_params_init(sd_video_upscale_params_t* params);
+SD_API video_upscaler_ctx_t* new_video_upscaler_ctx(const sd_video_upscaler_ctx_params_t* params);
+SD_API void free_video_upscaler_ctx(video_upscaler_ctx_t* ctx);
+// Frames are RGB, have equal dimensions, and remain owned by the caller.
+// Free the result with free_sd_video_generation_result().
+SD_API bool upscale_video(video_upscaler_ctx_t* ctx,
+ const sd_video_upscale_params_t* params,
+ sd_video_generation_result_t** result_out);
+
typedef struct adetailer_ctx_t adetailer_ctx_t;
typedef struct {
diff --git a/src/core/cuda/cute_attention.cu b/src/core/cuda/cute_attention.cu
index cdafc99..1eaefa4 100644
--- a/src/core/cuda/cute_attention.cu
+++ b/src/core/cuda/cute_attention.cu
@@ -150,7 +150,7 @@ namespace sd_fattn_cute {
}
}
- template <int Rows>
+ template <int Rows, bool Masked = false>
struct Softmax {
decltype(make_tensor<float>(Shape<Int<Rows>>{})) maximum, sum;
template <bool First, typename Scores, typename Output>
@@ -180,8 +180,9 @@ namespace sd_fattn_cute {
}
#pragma unroll
for (int n = 0; n < size<1>(scores); ++n) {
- const float value = exp2f(scores(m, n) * 1.4426950408889634f - next_max * 1.4426950408889634f);
- scores(m, n) = value;
+ const float safe_max = Masked && next_max == -INFINITY ? 0.f : next_max;
+ const float value = exp2f(scores(m, n) * 1.4426950408889634f - safe_max * 1.4426950408889634f);
+ scores(m, n) = value;
row_sum += value;
}
// Defer the four-lane sum reduction until the final output normalization.
@@ -196,7 +197,7 @@ namespace sd_fattn_cute {
float total = sum(m);
total += __shfl_xor_sync(0xffffffff, total, 2);
total += __shfl_xor_sync(0xffffffff, total, 1);
- const float inverse = 1.f / total;
+ const float inverse = Masked && total == 0.f ? 0.f : 1.f / total;
#pragma unroll
for (int n = 0; n < size<1>(output); ++n)
output(m, n) *= inverse;
@@ -790,3 +791,5 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) {
desc.priority = 400;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
+
+#include "window_attention.cuh"
diff --git a/src/core/cuda/registration.cu b/src/core/cuda/registration.cu
index 1f48174..d92be44 100644
--- a/src/core/cuda/registration.cu
+++ b/src/core/cuda/registration.cu
@@ -4,6 +4,7 @@
void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device);
void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
+void sd_ggml_register_cuda_window_attention(ggml_backend_dev_t device);
#endif
namespace SD_CUDA_NAMESPACE {
@@ -39,6 +40,7 @@ namespace SD_CUDA_NAMESPACE {
sd_ggml_register_cuda_i8_projection(device);
sd_ggml_register_cuda_i8_gated_residual(device);
sd_ggml_register_cuda_attention(device);
+ sd_ggml_register_cuda_window_attention(device);
#endif
}
}
diff --git a/src/core/cuda/window_attention.cuh b/src/core/cuda/window_attention.cuh
new file mode 100644
index 0000000..600cd91
--- /dev/null
+++ b/src/core/cuda/window_attention.cuh
@@ -0,0 +1,284 @@
+/*
+ * Adapted from FlashAttention v2.8.3:
+ * https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src
+ * Sources: kernel_traits.h, utils.h, softmax.h, flash_fwd_kernel.h.
+ * Copyright (c) 2023-2024, Tri Dao.
+ *
+ * BSD 3-Clause License
+ *
+ * Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * * Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * * Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * * Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+#ifndef __SD_CORE_CUDA_WINDOW_ATTENTION_CUH__
+#define __SD_CORE_CUDA_WINDOW_ATTENTION_CUH__
+
+namespace sd_fattn_cute {
+ struct WindowRegion {
+ int left, top, width, height;
+ };
+
+ static __device__ WindowRegion window_region(const SDWindowAttentionParams& p, int plane, int x, int y) {
+ int w = p.window_width[plane], h = p.window_height[plane];
+ if (p.preserve_boundary) {
+ w = min(w, p.width);
+ h = min(h, p.height);
+ return {min(max(x - w / 2, 0), p.width - w), min(max(y - h / 2, 0), p.height - h), w, h};
+ }
+ int left = max(x - w / 2, 0), top = max(y - h / 2, 0);
+ return {left, top, min(p.width, x + w / 2) - left, min(p.height, y + h / 2) - top};
+ }
+
+ template <typename Dst, typename Id, typename Index>
+ __device__ __forceinline__ void window_copy(const cutlass::half_t* source, Dst& dst, const Id& id, int dim, Index index) {
+#pragma unroll
+ for (int m = 0; m < size<1>(dst); ++m) {
+ const int row = index(get<0>(id(0, m, 0)));
+#pragma unroll
+ for (int k = 0; k < size<2>(dst); ++k) {
+ if (row >= 0) {
+ const auto* src = source + int64_t(row) * dim + get<1>(id(0, m, k));
+ const unsigned address = unsigned(__cvta_generic_to_shared(&dst(0, m, k)));
+ asm volatile("cp.async.cg.shared.global [%0], [%1], 16;"
+ :
+ : "r"(address), "l"(src));
+ } else {
+ cute::clear(dst(_, m, k));
+ }
+ }
+ }
+ }
+
+ template <int D>
+ __global__ __launch_bounds__(128) void window_kernel(const cutlass::half_t* q, const cutlass::half_t* k, const cutlass::half_t* v, float* output, SDWindowAttentionParams p, int kv_heads) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+ constexpr int BM = 64, BN = 64, TW = 8;
+ using T = Traits<BM, BN, D>;
+ using Element = typename T::Element;
+ extern __shared__ char storage[];
+ const int tid = threadIdx.x, head = blockIdx.z;
+ const int qx = blockIdx.x * TW, qy = blockIdx.y * (BM / TW);
+ const int tokens = p.width * p.height;
+ q += int64_t(head) * tokens * D;
+ const int64_t kv_offset = int64_t(head / (p.heads / kv_heads)) * tokens * p.planes * D;
+ k += kv_offset;
+ v += kv_offset;
+ auto sQ = make_tensor(make_smem_ptr(reinterpret_cast<Element*>(storage)), typename T::QLayout{});
+ auto sK = make_tensor(sQ.data() + size(sQ), typename T::KVLayout{});
+ auto sV = make_tensor(sK.data() + size(sK), typename T::KVLayout{});
+ auto sVt = make_tensor(sV.data(), typename T::VtLayout{});
+ auto sVplain = make_tensor(sV.data().get(), get_nonswizzle_portion(typename T::VtLayout{}));
+ typename T::GlobalCopy global_copy;
+ auto thread_copy = global_copy.get_thread_slice(tid);
+ auto tQsQ = thread_copy.partition_D(sQ);
+ auto tKsK = thread_copy.partition_D(sK);
+ auto tVsV = thread_copy.partition_D(sV);
+ auto q_id = thread_copy.partition_S(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
+ auto kv_id = thread_copy.partition_S(make_identity_tensor(Shape<Int<BN>, Int<D>>{}));
+ auto query_index = [&](int row) {
+ const int x = qx + row % TW, y = qy + row / TW;
+ return x < p.width && y < p.height ? y * p.width + x : -1;
+ };
+ window_copy(q, tQsQ, q_id, D, query_index);
+ cute::cp_async_fence();
+
+ typename T::Mma mma;
+ auto thread_mma = mma.get_thread_slice(tid);
+ auto rQ = thread_mma.partition_fragment_A(sQ);
+ auto rK = thread_mma.partition_fragment_B(sK);
+ auto rV = thread_mma.partition_fragment_B(sVplain);
+ auto acc_o = partition_fragment_C(mma, Shape<Int<BM>, Int<D>>{});
+ auto copyQ = make_tiled_copy_A(typename T::CopyAtom{}, mma);
+ auto copyK = make_tiled_copy_B(typename T::CopyAtom{}, mma);
+ auto copyV = make_tiled_copy_B(typename T::CopyAtomT{}, mma);
+ auto threadQ = copyQ.get_thread_slice(tid);
+ auto threadK = copyK.get_thread_slice(tid);
+ auto threadV = copyV.get_thread_slice(tid);
+ auto tSsQ = threadQ.partition_S(sQ);
+ auto tSsK = threadK.partition_S(sK);
+ auto tOsV = threadV.partition_S(sVt);
+ auto score_id = thread_mma.partition_C(make_identity_tensor(Shape<Int<BM>, Int<BN>>{}));
+ clear(acc_o);
+ Softmax<2 * size<1>(acc_o), true> softmax;
+ const float scale = rsqrtf(float(D));
+ for (int plane = 0; plane < p.planes; ++plane) {
+ const auto first = window_region(p, plane, qx, qy);
+ const auto last = window_region(p, plane, min(qx + TW - 1, p.width - 1), min(qy + BM / TW - 1, p.height - 1));
+ const int kw = last.left + last.width - first.left;
+ const int kh = last.top + last.height - first.top;
+ for (int start = 0; start < kw * kh; start += BN) {
+ auto key_index = [&](int row) {
+ row += start;
+ return row < kw * kh ? plane * tokens + (first.top + row / kw) * p.width + first.left + row % kw : -1;
+ };
+ window_copy(k, tKsK, kv_id, D, key_index);
+ window_copy(v, tVsV, kv_id, D, key_index);
+ cute::cp_async_fence();
+ cute::cp_async_wait<0>();
+ __syncthreads();
+ auto acc_s = partition_fragment_C(mma, Shape<Int<BM>, Int<BN>>{});
+ clear(acc_s);
+ gemm(acc_s, rQ, rK, tSsQ, tSsK, mma, copyQ, copyK, threadQ, threadK);
+#pragma unroll
+ for (int i = 0; i < size(acc_s); ++i) {
+ const int qr = get<0>(score_id(i)), kr = start + get<1>(score_id(i));
+ const int x = qx + qr % TW, y = qy + qr / TW;
+ const auto region = window_region(p, plane, x, y);
+ const int kx = first.left + kr % kw, ky = first.top + kr / kw;
+ const bool valid = x < p.width && y < p.height && kr < kw * kh &&
+ kx >= region.left && kx < region.left + region.width &&
+ ky >= region.top && ky < region.top + region.height;
+ acc_s(i) = valid ? acc_s(i) * scale : -INFINITY;
+ }
+ if (plane == 0 && start == 0)
+ softmax.template update<true>(acc_s, acc_o);
+ else
+ softmax.template update<false>(acc_s, acc_o);
+ auto rP = to_half(acc_s);
+ auto tOrP = make_tensor(rP.data(), convert_layout_acc_Aregs<typename T::Mma>(rP.layout()));
+ gemm_rs(acc_o, tOrP, rV, tOsV, mma, copyV, threadV);
+ __syncthreads();
+ }
+ }
+ softmax.normalize(acc_o);
+ auto output_id = thread_mma.partition_C(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
+#pragma unroll
+ for (int i = 0; i < size(acc_o); ++i) {
+ const int token = query_index(get<0>(output_id(i)));
+ if (token >= 0)
+ output[(int64_t(token) * p.heads + head) * D + get<1>(output_id(i))] = acc_o(i);
+ }
+#else
+ // This translation unit also targets older GPUs; runtime dispatch excludes this kernel there.
+ asm volatile("trap;");
+#endif
+ }
+}
+
+static bool window_attention_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+ static const bool disabled = [] {
+ const char* value = std::getenv("GGML_CUDA_DISABLE_WINDOW_ATTENTION");
+ return value && std::atoi(value) != 0;
+ }();
+ SDWindowAttentionParams params;
+ const auto* info = cuda_device(device);
+ return !disabled && info && info->properties.major >= 8 && sd_ggml_window_attn_describe(node, ¶ms);
+}
+
+static ggml_status window_attention_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
+ *workspace = 0;
+ for (int i = 0; i < 3; ++i)
+ *workspace += (size_t(ggml_nelements(nodes[0]->src[i])) * sizeof(half) + 255) / 256 * 256;
+ return GGML_STATUS_SUCCESS;
+}
+
+static void window_attention_destroy(void*, void*) {}
+
+template <int D>
+static cudaError_t window_attention_launch(const ggml_ext_cuda_context& context, const half* const* inputs, float* output, SDWindowAttentionParams params, int kv_heads) {
+ auto kernel = sd_fattn_cute::window_kernel<D>;
+ constexpr int shared_bytes = sd_fattn_cute::Traits<64, 64, D>::SharedBytes;
+ static std::mutex mutex;
+ static std::unordered_map<int, bool> initialized;
+ {
+ std::lock_guard<std::mutex> lock(mutex);
+ if (!initialized[context.device]) {
+ auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes);
+ if (error != cudaSuccess)
+ return error;
+ initialized[context.device] = true;
+ }
+ }
+ kernel<<<dim3((params.width + 7) / 8, (params.height + 7) / 8, params.heads), 128, shared_bytes, static_cast<cudaStream_t>(context.stream)>>>(
+ reinterpret_cast<const cutlass::half_t*>(inputs[0]), reinterpret_cast<const cutlass::half_t*>(inputs[1]),
+ reinterpret_cast<const cutlass::half_t*>(inputs[2]), output, params, kv_heads);
+ return cudaGetLastError();
+}
+
+static ggml_status window_attention_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) {
+ auto device = ggml_backend_get_device(execution->backend);
+ auto reg = ggml_backend_dev_backend_reg(device);
+ auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+ ggml_ext_cuda_context context;
+ if (!query || !query(execution->backend, &context))
+ return GGML_STATUS_FAILED;
+ SDWindowAttentionParams params;
+ auto* node = nodes[0];
+ if (!sd_ggml_window_attn_describe(node, ¶ms))
+ return GGML_STATUS_FAILED;
+ auto stream = static_cast<cudaStream_t>(context.stream);
+ auto* base = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+ const int dim = int(node->ne[0] / params.heads);
+ const half* inputs[3];
+ for (int i = 0; i < 3; ++i) {
+ const auto* src = node->src[i];
+ auto* packed = reinterpret_cast<half*>(base);
+ const int64_t heads = src->ne[0] / dim;
+ copy_to_half<float><<<(ggml_nelements(src) + 255) / 256, 256, 0, stream>>>(static_cast<const char*>(src->data), packed,
+ ggml_nelements(src), dim, src->ne[1], heads, src->nb[1], dim * sizeof(float), ggml_nbytes(src), 1.f);
+ if (cudaGetLastError() != cudaSuccess)
+ return GGML_STATUS_FAILED;
+ inputs[i] = packed;
+ base += (size_t(ggml_nelements(src)) * sizeof(half) + 255) / 256 * 256;
+ }
+ const int kv_heads = int(node->src[1]->ne[0] / dim);
+ auto* output = static_cast<float*>(node->data);
+ const auto error = dim == 64 ? window_attention_launch<64>(context, inputs, output, params, kv_heads)
+ : window_attention_launch<128>(context, inputs, output, params, kv_heads);
+ return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+}
+
+void sd_ggml_register_cuda_window_attention(ggml_backend_dev_t device) {
+ const auto* info = cuda_device(device);
+ if (!info || info->properties.major < 8)
+ return;
+ int previous;
+ if (cudaGetDevice(&previous) != cudaSuccess || cudaSetDevice(info->ordinal) != cudaSuccess)
+ return;
+ cudaFuncAttributes attributes64{}, attributes128{};
+ const auto error64 = cudaFuncGetAttributes(&attributes64, sd_fattn_cute::window_kernel<64>);
+ const auto error128 = cudaFuncGetAttributes(&attributes128, sd_fattn_cute::window_kernel<128>);
+ const auto restored = cudaSetDevice(previous);
+ // JIT can produce an SM80+ binary from older PTX containing only the trap.
+ if (restored != cudaSuccess || error64 != cudaSuccess || error128 != cudaSuccess ||
+ std::min(attributes64.binaryVersion, attributes64.ptxVersion) < 80 ||
+ std::min(attributes128.binaryVersion, attributes128.ptxVersion) < 80)
+ return;
+ ggml_ext_kernel_desc desc = {};
+ desc.abi_version = GGML_EXT_ABI_VERSION;
+ desc.name = "sd.window_attention.cuda.v1";
+ desc.op = GGML_OP_CUSTOM;
+ desc.custom_op = sd_ggml_window_attn_op();
+ desc.priority = 100;
+ desc.supports = window_attention_supports;
+ desc.prepare = window_attention_prepare;
+ desc.destroy = window_attention_destroy;
+ desc.execute = window_attention_execute;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+}
+
+#endif
diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp
index 25fa222..dfe5bca 100644
--- a/src/core/ggml_extensions.cpp
+++ b/src/core/ggml_extensions.cpp
@@ -501,6 +501,65 @@ ggml_tensor* ggml_sol_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, gg
return ggml_ext_op_create(ctx, sd_ggml_sol_attn_op(), inputs, 3, ¶ms);
}
+static bool window_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
+ if (count != 3 || !attributes)
+ return false;
+ SDWindowAttentionParams p;
+ std::memcpy(&p, attributes, sizeof(p));
+ if (p.width <= 0 || p.height <= 0 || int64_t(p.width) * p.height > INT32_MAX / 4 ||
+ p.heads <= 0 || p.planes < 1 || p.planes > 4 || (p.preserve_boundary != 0 && p.preserve_boundary != 1))
+ return false;
+ for (int i = 0; i < p.planes; ++i)
+ if (p.window_height[i] <= 0 || p.window_width[i] <= 0 ||
+ (!p.preserve_boundary && (p.window_height[i] % 2 || p.window_width[i] % 2)))
+ return false;
+ for (int i = 0; i < 3; ++i) {
+ const auto* t = inputs[i];
+ if (!t || t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || t->ne[2] != 1 || t->ne[3] != 1 ||
+ t->ne[1] != int64_t(p.width) * p.height * (i ? p.planes : 1))
+ return false;
+ }
+ const auto* q = inputs[0];
+ const auto* k = inputs[1];
+ if (q->ne[0] % p.heads || (q->ne[0] / p.heads != 64 && q->ne[0] / p.heads != 128) ||
+ k->ne[0] <= 0 || k->ne[0] > q->ne[0] || q->ne[0] % k->ne[0] ||
+ k->ne[0] % (q->ne[0] / p.heads) || !ggml_are_same_shape(k, inputs[2]))
+ return false;
+ *out = {GGML_TYPE_F32, {q->ne[0], q->ne[1], 1, 1}};
+ return true;
+}
+
+ggml_ext_op_t sd_ggml_window_attn_op() {
+ static const auto op = [] {
+ const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.window_attn.v1", sizeof(SDWindowAttentionParams), window_attn_infer};
+ auto result = ggml_ext_register_op(&desc);
+ GGML_ASSERT(result);
+ return result;
+ }();
+ return op;
+}
+
+bool sd_ggml_window_attn_describe(const ggml_tensor* node, SDWindowAttentionParams* params) {
+ if (!node || ggml_ext_op_get(node) != sd_ggml_window_attn_op() || node->src[3] || !ggml_is_contiguous(node))
+ return false;
+ const auto* attributes = ggml_ext_op_attributes(node);
+ ggml_ext_tensor_desc out;
+ if (!window_attn_infer(node->src, 3, attributes, &out) || node->type != out.type || std::memcmp(node->ne, out.ne, sizeof(out.ne)))
+ return false;
+ std::memcpy(params, attributes, sizeof(*params));
+ return true;
+}
+
+ggml_tensor* sd_ggml_window_attn(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, const SDWindowAttentionParams& params) {
+ ggml_tensor* inputs[] = {q, k, v};
+ ggml_ext_tensor_desc out;
+ if (!backend || !window_attn_infer(inputs, 3, ¶ms, &out))
+ return nullptr;
+ sd_ggml_register_extensions(backend);
+ auto result = ggml_ext_op_create(ctx, sd_ggml_window_attn_op(), inputs, 3, ¶ms);
+ return ggml_backend_supports_op(backend, result) ? result : nullptr;
+}
+
void sd_ggml_register_extensions(ggml_backend_t backend) {
if (!backend) {
return;
diff --git a/src/core/ggml_extensions.h b/src/core/ggml_extensions.h
index bc9ee9a..87a51a7 100644
--- a/src/core/ggml_extensions.h
+++ b/src/core/ggml_extensions.h
@@ -80,6 +80,18 @@ struct SDSolAttentionParams {
float tau;
};
+struct SDWindowAttentionParams {
+ int32_t width, height, heads, planes;
+ int16_t window_height[4], window_width[4];
+ int32_t preserve_boundary;
+};
+
+ggml_ext_op_t sd_ggml_window_attn_op();
+bool sd_ggml_window_attn_describe(const ggml_tensor* node, SDWindowAttentionParams* params);
+// Contiguous F32 Q [D*heads, W*H], K/V [D*kv_heads, W*H*planes].
+// Returns Q's layout, or nullptr when the backend has no compatible kernel.
+ggml_tensor* sd_ggml_window_attn(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, const SDWindowAttentionParams& params);
+
ggml_ext_op_t sd_ggml_sol_attn_op();
bool sd_ggml_sol_attn_describe(const ggml_tensor* node, SDSolAttentionParams* params);
// Noncausal Sol-Attn: contiguous F32 Q/K/V [128, tokens, heads, batch], output [128, heads, tokens, batch].
diff --git a/src/core/tensor.hpp b/src/core/tensor.hpp
index a925dbf..715068c 100644
--- a/src/core/tensor.hpp
+++ b/src/core/tensor.hpp
@@ -222,14 +222,28 @@ namespace sd {
return result;
}
- for (int64_t flat = 0; flat < result.numel(); ++flat) {
- std::vector<int64_t> out_coord = tensor_unravel_index(flat, out_shape);
- std::vector<int64_t> src_coord(static_cast<size_t>(dim()), 0);
- for (size_t i = 0; i < dims.size(); ++i) {
- src_coord[dims[i]] = out_coord[i];
+ std::vector<int64_t> source_strides(dims.size(), 1), strides(dims.size());
+ for (size_t i = 1; i < dims.size(); ++i)
+ source_strides[i] = source_strides[i - 1] * shape_[i - 1];
+ for (size_t i = 0; i < dims.size(); ++i)
+ strides[i] = source_strides[dims[i]];
+ parallel_for(0, result.numel(), 65536, [&](int64_t begin, int64_t end) {
+ auto coord = tensor_unravel_index(begin, out_shape);
+ int64_t source = 0;
+ for (size_t i = 0; i < dims.size(); ++i)
+ source += coord[i] * strides[i];
+ for (int64_t flat = begin; flat < end; ++flat) {
+ result.data()[flat] = data_[source];
+ for (size_t i = 0; i < dims.size(); ++i) {
+ ++coord[i];
+ source += strides[i];
+ if (coord[i] < out_shape[i])
+ break;
+ source -= coord[i] * strides[i];
+ coord[i] = 0;
+ }
}
- result[flat] = index(src_coord);
- }
+ });
return result;
}
@@ -1037,11 +1051,13 @@ namespace sd {
auto x_contributors = make_interpolate_contributors(input_width, output_width, mode, antialias);
auto y_contributors = make_interpolate_contributors(input_height, output_height, mode, antialias);
- for (int64_t plane = 0; plane < plane_count; ++plane) {
- const int64_t input_plane_offset = plane * input_plane;
- const int64_t output_plane_offset = plane * output_plane;
- for (int64_t y = 0; y < output_height; ++y) {
- const auto& y_axis = y_contributors[static_cast<size_t>(y)];
+ parallel_for(0, plane_count * output_height, std::max<int64_t>(1, 16384 / output_width), [&](int64_t begin, int64_t end) {
+ for (int64_t row = begin; row < end; ++row) {
+ const int64_t plane = row / output_height;
+ const int64_t y = row % output_height;
+ const int64_t input_plane_offset = plane * input_plane;
+ const int64_t output_plane_offset = plane * output_plane;
+ const auto& y_axis = y_contributors[static_cast<size_t>(y)];
for (int64_t x = 0; x < output_width; ++x) {
const auto& x_axis = x_contributors[static_cast<size_t>(x)];
double value = 0.0;
@@ -1055,7 +1071,7 @@ namespace sd {
output.data()[output_plane_offset + y * output_width + x] = static_cast<T>(value);
}
}
- }
+ });
return output;
}
diff --git a/src/core/util.cpp b/src/core/util.cpp
index a691b42..9de1ace 100644
--- a/src/core/util.cpp
+++ b/src/core/util.cpp
@@ -807,13 +807,11 @@ sd::Tensor<float> sd_image_to_tensor(sd_image_t image,
static_cast<int64_t>(image.height),
static_cast<int64_t>(image.channel),
1});
- for (uint32_t iw = 0; iw < image.width; ++iw) {
- for (uint32_t ih = 0; ih < image.height; ++ih) {
- for (uint32_t ic = 0; ic < image.channel; ++ic) {
- tensor.index(iw, ih, ic, 0) = sd_image_get_f32(image, iw, ih, ic, scale);
- }
- }
- }
+ const int64_t pixels = int64_t(image.width) * image.height;
+ sd::tensor_for_each(tensor.numel(), [&](int64_t i) {
+ float value = image.data[(i % pixels) * image.channel + i / pixels];
+ tensor.data()[i] = scale ? value / 255.f : value;
+ });
if (target_width >= 0 && target_height >= 0 &&
(tensor.shape()[0] != target_width || tensor.shape()[1] != target_height)) {
tensor = sd::ops::interpolate(tensor,
diff --git a/src/model/common/vsr.hpp b/src/model/common/vsr.hpp
new file mode 100644
index 0000000..c203dff
--- /dev/null
+++ b/src/model/common/vsr.hpp
@@ -0,0 +1,221 @@
+#ifndef __SD_MODEL_COMMON_VSR_HPP__
+#define __SD_MODEL_COMMON_VSR_HPP__
+
+#include <array>
+#include <limits>
+#include <map>
+#include <stdexcept>
+#include <unordered_map>
+#include "model/common/ggml_block.hpp"
+#include "model/common/rope.hpp"
+
+namespace VSR {
+
+ struct GraphInputs {
+ std::map<std::vector<int>, sd::Tensor<float>> rotations;
+ mutable std::unordered_map<const void*, ggml_tensor*> bindings;
+
+ void begin_graph() {
+ bindings.clear();
+ }
+
+ void clear() {
+ bindings.clear();
+ rotations.clear();
+ }
+
+ template <typename T>
+ ggml_tensor* input(GGMLRunnerContext* ctx, const sd::Tensor<T>& value) const {
+ auto existing = bindings.find(value.data());
+ if (existing != bindings.end())
+ return existing->second;
+ auto result = sd::make_ggml_tensor(ctx->ggml_ctx, value, false);
+ ctx->set_backend_tensor_data(result, value.data());
+ bindings.emplace(value.data(), result);
+ return result;
+ }
+
+ ggml_tensor* rope(GGMLRunnerContext* ctx, int width, int height, const std::vector<int>& times, const std::vector<int>& axes) {
+ std::vector<int> key = {width, height, axes[0], axes[1], axes[2]};
+ key.insert(key.end(), times.begin(), times.end());
+ auto existing = rotations.find(key);
+ if (existing != rotations.end())
+ return input(ctx, existing->second);
+ const int dim = axes[0] + axes[1] + axes[2];
+ auto& pe = rotations.emplace(std::move(key), sd::Tensor<float>({2, 2, dim / 2, int64_t(width) * height * int64_t(times.size())})).first->second;
+ float* out = pe.data();
+ for (int t : times) {
+ for (int y = 0; y < height; ++y) {
+ for (int x = 0; x < width; ++x) {
+ const int pos[] = {t, y, x};
+ for (int a = 0; a < 3; ++a) {
+ for (int i = 0; i < axes[a]; i += 2) {
+ const double angle = pos[a] * std::pow(10000.0, -double(i) / axes[a]);
+ *out++ = float(std::cos(angle));
+ *out++ = float(-std::sin(angle));
+ *out++ = float(std::sin(angle));
+ *out++ = float(std::cos(angle));
+ }
+ }
+ }
+ }
+ }
+ return input(ctx, pe);
+ }
+ };
+
+ static ggml_tensor* tokens(ggml_context* ctx, ggml_tensor* x) {
+ x = ggml_reshape_2d(ctx, x, x->ne[0] * x->ne[1], x->ne[2] * x->ne[3]);
+ return ggml_cont(ctx, ggml_transpose(ctx, x));
+ }
+
+ static ggml_tensor* image(ggml_context* ctx, ggml_tensor* x, int width, int height) {
+ const auto channels = x->ne[0];
+ x = ggml_cont(ctx, ggml_transpose(ctx, x));
+ return ggml_reshape_4d(ctx, x, width, height, channels, 1);
+ }
+
+ static ggml_tensor* pixel_shuffle(ggml_context* ctx, ggml_tensor* x, int width, int height, int scale, bool channel_first = true) {
+ const int64_t channels = x->ne[0] / (scale * scale);
+ if (!channel_first) {
+ x = ggml_reshape_4d(ctx, x, channels, scale, scale, width * height);
+ x = ggml_cont(ctx, ggml_permute(ctx, x, 2, 0, 1, 3));
+ }
+ x = ggml_reshape_4d(ctx, x, scale, scale, channels, width * height);
+ x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 3, 1));
+ x = ggml_reshape_4d(ctx, x, scale * width, height, scale, channels);
+ x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));
+ return ggml_reshape_4d(ctx, x, scale * width, scale * height, channels, 1);
+ }
+
+ struct AttentionTile {
+ int width;
+ int height;
+ sd::Tensor<int32_t> queries;
+ sd::Tensor<int32_t> keys;
+ sd::Tensor<float> mask;
+ };
+
+ struct LocalAttention {
+ int width;
+ int height;
+ std::vector<std::pair<int, int>> windows;
+ bool preserve_boundary;
+ mutable std::vector<AttentionTile> tiles;
+ GraphInputs inputs;
+
+ void begin_graph() {
+ inputs.begin_graph();
+ }
+
+ LocalAttention(int width, int height, const std::vector<std::pair<int, int>>& windows, bool preserve_boundary)
+ : width(width), height(height), windows(windows), preserve_boundary(preserve_boundary) {}
+
+ void prepare_tiles() const {
+ if (!tiles.empty())
+ return;
+ constexpr int tile_size = 16;
+ for (int y = 0; y < height; y += tile_size) {
+ for (int x = 0; x < width; x += tile_size) {
+ AttentionTile tile;
+ tile.width = std::min(tile_size, width - x);
+ tile.height = std::min(tile_size, height - y);
+ tile.queries = sd::Tensor<int32_t>({tile.width * tile.height});
+ std::vector<int32_t> keys;
+ auto bounds = [&](int qy, int qx, int plane) {
+ auto [wh, ww] = windows[plane];
+ if (preserve_boundary) {
+ wh = std::min(wh, height);
+ ww = std::min(ww, width);
+ const int top = std::clamp(qy - wh / 2, 0, height - wh);
+ const int left = std::clamp(qx - ww / 2, 0, width - ww);
+ return std::array<int, 4>{top, left, top + wh, left + ww};
+ }
+ return std::array<int, 4>{std::max(0, qy - wh / 2), std::max(0, qx - ww / 2),
+ std::min(height, qy + wh / 2), std::min(width, qx + ww / 2)};
+ };
+ for (int p = 0; p < int(windows.size()); ++p) {
+ auto begin = bounds(y, x, p);
+ auto end = bounds(y + tile.height - 1, x + tile.width - 1, p);
+ for (int ky = begin[0]; ky < end[2]; ++ky)
+ for (int kx = begin[1]; kx < end[3]; ++kx)
+ keys.push_back(p * height * width + ky * width + kx);
+ }
+ tile.keys = sd::Tensor<int32_t>({int64_t(keys.size())});
+ std::copy(keys.begin(), keys.end(), tile.keys.data());
+ tile.mask = sd::Tensor<float>({int64_t(keys.size()), tile.width * tile.height});
+ for (int q = 0; q < tile.width * tile.height; ++q) {
+ int qy = y + q / tile.width;
+ int qx = x + q % tile.width;
+ tile.queries.data()[q] = qy * width + qx;
+ for (size_t k = 0; k < keys.size(); ++k) {
+ int p = keys[k] / (width * height);
+ int ky = (keys[k] % (width * height)) / width;
+ int kx = keys[k] % width;
+ auto b = bounds(qy, qx, p);
+ tile.mask.data()[q * keys.size() + k] = ky >= b[0] && ky < b[2] && kx >= b[1] && kx < b[3]
+ ? 0.f
+ : -std::numeric_limits<float>::infinity();
+ }
+ }
+ tiles.push_back(std::move(tile));
+ }
+ }
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, int heads) const {
+ auto c = ctx->ggml_ctx;
+#ifndef SD_USE_UPSTREAM_GGML
+ if (ctx->flash_attn_enabled && windows.size() <= 4) {
+ SDWindowAttentionParams params = {};
+ params.width = width;
+ params.height = height;
+ params.heads = heads;
+ params.planes = int(windows.size());
+ params.preserve_boundary = preserve_boundary;
+ for (size_t i = 0; i < windows.size(); ++i) {
+ params.window_height[i] = windows[i].first;
+ params.window_width[i] = windows[i].second;
+ }
+ if (auto out = sd_ggml_window_attn(c, ctx->backend, q, k, v, params))
+ return out;
+ }
+#endif
+ prepare_tiles();
+ ggml_tensor* output = nullptr;
+ ggml_tensor* row = nullptr;
+ int row_width = 0;
+ for (const auto& tile : tiles) {
+ auto qi = inputs.input(ctx, tile.queries);
+ auto ki = inputs.input(ctx, tile.keys);
+ auto tq = ggml_get_rows(c, q, qi);
+ auto tk = ggml_get_rows(c, k, ki);
+ auto tv = ggml_get_rows(c, v, ki);
+ auto mask = inputs.input(ctx, tile.mask);
+ auto out = ggml_ext_attention_ext(c, ctx->backend, tq, tk, tv, heads, mask, false, ctx->flash_attn_enabled);
+ out = ggml_reshape_3d(c, out, out->ne[0], tile.width, tile.height);
+ row = row ? ggml_concat(c, row, out, 1) : out;
+ row_width += tile.width;
+ if (row_width == width) {
+ output = output ? ggml_concat(c, output, row, 2) : row;
+ row = nullptr;
+ row_width = 0;
+ }
+ }
+ return ggml_reshape_2d(c, output, output->ne[0], width * height);
+ }
+ };
+
+ static ggml_tensor* apply_rope(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, int heads) {
+ auto c = ctx->ggml_ctx;
+ const int64_t dim = x->ne[0];
+ const int64_t count = x->ne[1];
+ x = ggml_reshape_4d(c, x, dim / heads, heads, count, 1);
+ x = Rope::apply_rope(ctx, x, pe);
+ x = ggml_cont(c, ggml_permute(c, x, 0, 2, 1, 3));
+ return ggml_reshape_2d(c, x, dim, count);
+ }
+
+} // namespace VSR
+
+#endif // __SD_MODEL_COMMON_VSR_HPP__
diff --git a/src/model/diffusion/relay_vsr.hpp b/src/model/diffusion/relay_vsr.hpp
new file mode 100644
index 0000000..a390aa3
--- /dev/null
+++ b/src/model/diffusion/relay_vsr.hpp
@@ -0,0 +1,293 @@
+#ifndef __SD_MODEL_DIFFUSION_RELAY_VSR_HPP__
+#define __SD_MODEL_DIFFUSION_RELAY_VSR_HPP__
+
+#include "model/common/vsr.hpp"
+
+namespace RelayVSR {
+
+ struct WanVSRConfig {
+ int dim = 3072;
+ int ffn_dim = 14336;
+ int num_heads = 24;
+ int num_layers = 30;
+ int in_channels = 48;
+ int text_dim = 4096;
+ int text_len = 512;
+ int hidden_dim1 = 768;
+ int hidden_dim2 = 3072;
+ int upscale_factor = 4;
+
+ static WanVSRConfig detect_from_weights(const String2TensorStorage& weights, const std::string& prefix, int scale) {
+ WanVSRConfig config;
+ config.upscale_factor = scale;
+ auto weight = [&](const std::string& name) -> const TensorStorage& { return weights.at(prefix + "." + name); };
+ config.dim = int(weight("condition_embedder.time_embedder.linear_1.weight").ne[1]);
+ config.ffn_dim = int(weight("blocks.0.ffn.net.0.proj.weight").ne[1]);
+ config.num_heads = config.dim / 128;
+ config.text_dim = int(weight("prompt_embeds").ne[0]);
+ config.text_len = int(weight("prompt_embeds").ne[1]);
+ config.hidden_dim1 = int(weight("lq_proj.conv1.weight").ne[3]);
+ config.hidden_dim2 = int(weight("lq_proj.conv2.weight").ne[3]);
+ config.num_layers = 0;
+ while (weights.count(prefix + ".blocks." + std::to_string(config.num_layers) + ".attn1.to_q.weight"))
+ ++config.num_layers;
+ if ((scale != 2 && scale != 4) || weight("lq_proj.conv1.weight").ne[2] != 3 * (4 * scale) * (4 * scale) ||
+ config.dim != 3072 || config.num_layers != 30 || config.text_dim != 4096 ||
+ weight("proj_out.weight").ne[1] != 192) {
+ throw std::runtime_error("RelayVSR Wan weights do not match the requested scale or architecture");
+ }
+ LOG_VERBOSE("relayvsr: dim = %d, num_layers = %d, ffn_dim = %d, scale = %d", config.dim, config.num_layers, config.ffn_dim, scale);
+ return config;
+ }
+ };
+
+ class LQ4xProj : public GGMLBlock {
+ int hidden_dim1_;
+ int hidden_dim2_;
+
+ public:
+ explicit LQ4xProj(const WanVSRConfig& config)
+ : hidden_dim1_(config.hidden_dim1), hidden_dim2_(config.hidden_dim2) {
+ const int patch = 4 * config.upscale_factor;
+ blocks["conv1"] = std::make_shared<Conv2d>(3 * patch * patch, hidden_dim1_, std::pair<int, int>{1, 1});
+ blocks["norm1"] = std::make_shared<RMSNorm>(hidden_dim1_, 1e-6f);
+ blocks["conv2"] = std::make_shared<Conv2d>(hidden_dim1_, hidden_dim2_, std::pair<int, int>{3, 3}, std::pair<int, int>{2, 2}, std::pair<int, int>{1, 1});
+ blocks["norm2"] = std::make_shared<RMSNorm>(hidden_dim2_, 1e-6f);
+ blocks["out_proj"] = std::make_shared<Linear>(hidden_dim2_, config.dim);
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* unshuffled) {
+ auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
+ auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv2"]);
+ auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
+ auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
+ auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
+ auto c = ctx->ggml_ctx;
+ auto x = conv1->forward(ctx, unshuffled);
+ int width = int(x->ne[0]), height = int(x->ne[1]);
+ x = ggml_silu(c, norm1->forward(ctx, VSR::tokens(c, x)));
+ x = conv2->forward(ctx, VSR::image(c, x, width, height));
+ x = ggml_silu(c, norm2->forward(ctx, VSR::tokens(c, x)));
+ return out_proj->forward(ctx, x);
+ }
+ };
+
+ class WanAttention : public GGMLBlock {
+ int dim_;
+ int heads_;
+
+ public:
+ explicit WanAttention(const WanVSRConfig& config)
+ : dim_(config.dim), heads_(config.num_heads) {
+ for (const auto& name : {"to_q", "to_k", "to_v", "to_out.0"})
+ blocks[name] = std::make_shared<Linear>(dim_, dim_);
+ blocks["norm_q"] = std::make_shared<RMSNorm>(dim_, 1e-6f);
+ blocks["norm_k"] = std::make_shared<RMSNorm>(dim_, 1e-6f);
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* context, ggml_tensor* query_rope, ggml_tensor* key_rope, const VSR::LocalAttention* local, const std::string& cache_name) {
+ auto to_q = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
+ auto to_k = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
+ auto to_v = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
+ auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
+ auto norm_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
+ auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
+ auto c = ctx->ggml_ctx;
+ auto q = norm_q->forward(ctx, to_q->forward(ctx, x));
+ auto past_k = ctx->get_cache_tensor(cache_name + ".key");
+ auto past_v = ctx->get_cache_tensor(cache_name + ".value");
+ ggml_tensor* k = nullptr;
+ ggml_tensor* v = nullptr;
+ if (context && past_k) {
+ k = past_k;
+ v = past_v;
+ } else {
+ k = norm_k->forward(ctx, to_k->forward(ctx, context ? context : x));
+ v = to_v->forward(ctx, context ? context : x);
+ if (!context && past_k) {
+ k = ggml_concat(c, past_k, k, 1);
+ v = ggml_concat(c, past_v, v, 1);
+ }
+ auto retained = [&](ggml_tensor* value) {
+ if (context || value->ne[1] <= x->ne[1] * 3)
+ return value;
+ // The next attention window adds its current frame to these three previous frames.
+ const int64_t keep = x->ne[1] * 3;
+ return ggml_cont(c, ggml_view_2d(c, value, value->ne[0], keep, value->nb[1], (value->ne[1] - keep) * value->nb[1]));
+ };
+ ctx->cache_tensor(cache_name + ".key", retained(k));
+ ctx->cache_tensor(cache_name + ".value", retained(v));
+ }
+ ggml_tensor* out;
+ if (local) {
+ q = VSR::apply_rope(ctx, q, query_rope, heads_);
+ k = VSR::apply_rope(ctx, k, key_rope, heads_);
+ out = local->forward(ctx, q, k, v, heads_);
+ } else {
+ out = ggml_ext_attention_ext(c, ctx->backend, q, k, v, heads_, nullptr, false, ctx->flash_attn_enabled);
+ }
+ return to_out_0->forward(ctx, out);
+ }
+ };
+
+ class WanTransformerBlock : public GGMLBlock {
+ int dim_;
+
+ void init_params(ggml_context* ctx, const String2TensorStorage& weights, const std::string prefix) override {
+ params["scale_shift_table"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, dim_, 6, 1);
+ }
+
+ public:
+ explicit WanTransformerBlock(const WanVSRConfig& config)
+ : dim_(config.dim) {
+ blocks["norm1"] = std::make_shared<LayerNorm>(dim_, 1e-6f, false);
+ blocks["norm2"] = std::make_shared<LayerNorm>(dim_, 1e-6f, true);
+ blocks["norm3"] = std::make_shared<LayerNorm>(dim_, 1e-6f, false);
+ blocks["attn1"] = std::make_shared<WanAttention>(config);
+ blocks["attn2"] = std::make_shared<WanAttention>(config);
+ blocks["ffn.net.0.proj"] = std::make_shared<Linear>(dim_, config.ffn_dim);
+ blocks["ffn.net.2"] = std::make_shared<Linear>(config.ffn_dim, dim_);
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* text, ggml_tensor* timestep_proj, ggml_tensor* qpe, ggml_tensor* kpe, const VSR::LocalAttention& local, const std::string& name) {
+ auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
+ auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
+ auto norm3 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm3"]);
+ auto ffn_net_0_proj = std::dynamic_pointer_cast<Linear>(blocks["ffn.net.0.proj"]);
+ auto ffn_net_2 = std::dynamic_pointer_cast<Linear>(blocks["ffn.net.2"]);
+ auto c = ctx->ggml_ctx;
+ auto modulation = ggml_add(c, timestep_proj, params["scale_shift_table"]);
+ auto pieces = ggml_ext_chunk(c, modulation, 6, 1);
+ auto attn1 = std::dynamic_pointer_cast<WanAttention>(blocks["attn1"]);
+ auto attn2 = std::dynamic_pointer_cast<WanAttention>(blocks["attn2"]);
+ auto y = attn1->forward(ctx, ggml_ext_modulate(ctx, norm1->forward(ctx, x), pieces[0], pieces[1]), nullptr, qpe, kpe, &local, name + ".attn1");
+ x = ggml_add(c, x, ggml_mul(c, y, pieces[2]));
+ x = ggml_add(c, x, attn2->forward(ctx, norm2->forward(ctx, x), text, nullptr, nullptr, nullptr, name + ".attn2"));
+ y = ggml_ext_modulate(ctx, norm3->forward(ctx, x), pieces[3], pieces[4]);
+ y = ffn_net_0_proj->forward(ctx, y);
+ y = ffn_net_2->forward(ctx, ggml_ext_gelu(c, y));
+ return ggml_add(c, x, ggml_mul(c, y, pieces[5]));
+ }
+ };
+
+ class WanTransformer3DModel : public GGMLBlock {
+ WanVSRConfig config_;
+
+ void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override {
+ params["scale_shift_table"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, config_.dim, 2, 1);
+ params["prompt_embeds"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config_.text_dim, config_.text_len);
+ }
+
+ public:
+ explicit WanTransformer3DModel(const WanVSRConfig& config)
+ : config_(config) {
+ blocks["patch_embedding"] = std::make_shared<Conv3d>(config.in_channels, config.dim, std::tuple<int, int, int>{1, 2, 2}, std::tuple<int, int, int>{1, 2, 2});
+ blocks["lq_proj"] = std::make_shared<LQ4xProj>(config);
+ blocks["condition_embedder.time_embedder.linear_1"] = std::make_shared<Linear>(256, config.dim);
+ blocks["condition_embedder.time_embedder.linear_2"] = std::make_shared<Linear>(config.dim, config.dim);
+ blocks["condition_embedder.time_proj"] = std::make_shared<Linear>(config.dim, config.dim * 6);
+ blocks["condition_embedder.text_embedder.linear_1"] = std::make_shared<Linear>(config.text_dim, config.dim);
+ blocks["condition_embedder.text_embedder.linear_2"] = std::make_shared<Linear>(config.dim, config.dim);
+ for (int i = 0; i < config.num_layers; ++i)
+ blocks["blocks." + std::to_string(i)] = std::make_shared<WanTransformerBlock>(config);
+ blocks["norm_out"] = std::make_shared<LayerNorm>(config.dim, 1e-6f, false);
+ blocks["proj_out"] = std::make_shared<Linear>(config.dim, config.in_channels * 4);
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* lq, ggml_tensor* timestep, ggml_tensor* qpe, ggml_tensor* kpe, const VSR::LocalAttention& local) {
+ auto patch_embedding = std::dynamic_pointer_cast<Conv3d>(blocks["patch_embedding"]);
+ auto condition_embedder_time_embedder_linear_1 = std::dynamic_pointer_cast<Linear>(blocks["condition_embedder.time_embedder.linear_1"]);
+ auto condition_embedder_time_embedder_linear_2 = std::dynamic_pointer_cast<Linear>(blocks["condition_embedder.time_embedder.linear_2"]);
+ auto condition_embedder_time_proj = std::dynamic_pointer_cast<Linear>(blocks["condition_embedder.time_proj"]);
+ auto condition_embedder_text_embedder_linear_1 = std::dynamic_pointer_cast<Linear>(blocks["condition_embedder.text_embedder.linear_1"]);
+ auto condition_embedder_text_embedder_linear_2 = std::dynamic_pointer_cast<Linear>(blocks["condition_embedder.text_embedder.linear_2"]);
+ auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
+ auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
+ auto c = ctx->ggml_ctx;
+ const int width = int(latent->ne[0]) / 2, height = int(latent->ne[1]) / 2;
+ auto x = patch_embedding->forward(ctx, latent);
+ x = VSR::tokens(c, ggml_reshape_4d(c, x, width, height, config_.dim, 1));
+ x = ggml_add(c, x, std::dynamic_pointer_cast<LQ4xProj>(blocks["lq_proj"])->forward(ctx, lq));
+ auto e = condition_embedder_time_embedder_linear_1->forward(ctx, timestep);
+ e = condition_embedder_time_embedder_linear_2->forward(ctx, ggml_silu(c, e));
+ auto e0 = condition_embedder_time_proj->forward(ctx, ggml_silu(c, e));
+ e0 = ggml_reshape_3d(c, e0, config_.dim, 6, 1);
+ auto text = condition_embedder_text_embedder_linear_1->forward(ctx, params["prompt_embeds"]);
+ text = condition_embedder_text_embedder_linear_2->forward(ctx, ggml_ext_gelu(c, text));
+ for (int i = 0; i < config_.num_layers; ++i) {
+ auto name = "blocks." + std::to_string(i);
+ x = std::dynamic_pointer_cast<WanTransformerBlock>(blocks[name])->forward(ctx, x, text, e0, qpe, kpe, local, name);
+ sd::ggml_graph_cut::mark_graph_cut(x, "relayvsr." + name, "hidden");
+ }
+ e = ggml_reshape_3d(c, e, config_.dim, 1, 1);
+ auto modulation = ggml_add(c, params["scale_shift_table"], e);
+ auto pieces = ggml_ext_chunk(c, modulation, 2, 1);
+ x = norm_out->forward(ctx, x);
+ x = ggml_ext_modulate(ctx, x, pieces[0], pieces[1]);
+ x = proj_out->forward(ctx, x);
+ x = VSR::pixel_shuffle(c, x, width, height, 2, false);
+ return ggml_reshape_4d(c, x, width * 2, height * 2, 1, config_.in_channels);
+ }
+ };
+
+ struct WanVSRRunner : GGMLRunner {
+ WanVSRConfig config;
+ WanTransformer3DModel model;
+ VSR::GraphInputs inputs;
+ std::map<int, std::unique_ptr<VSR::LocalAttention>> attention;
+ sd::Tensor<float> timestep;
+ int frame = 0;
+
+ WanVSRRunner(ggml_backend_t backend, const String2TensorStorage& weights, int scale, std::shared_ptr<DeviceResidencyManager> manager)
+ : GGMLRunner(backend, manager), config(WanVSRConfig::detect_from_weights(weights, "wan", scale)), model(config), timestep({256}) {
+ model.init(params_ctx, weights, "wan");
+ for (int i = 0; i < 128; ++i) {
+ double angle = 1000.0 * std::pow(10000.0, -double(i) / 128);
+ timestep.data()[i] = float(std::cos(angle));
+ timestep.data()[i + 128] = float(std::sin(angle));
+ }
+ }
+
+ std::string get_desc() override { return "RelayVSR Wan"; }
+ void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) { model.get_param_tensors(tensors, "wan"); }
+ void reset() {
+ runner_end();
+ free_cache_ctx_and_buffer();
+ attention.clear();
+ inputs.clear();
+ frame = 0;
+ }
+
+ sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& noise, const sd::Tensor<float>& lq) {
+ int width = int(noise.shape()[0]) / 2, height = int(noise.shape()[1]) / 2;
+ int count = std::min(frame + 1, 4);
+ if (!attention.count(count))
+ attention[count] = std::make_unique<VSR::LocalAttention>(width, height, std::vector<std::pair<int, int>>(count, {22, 40}), true);
+ auto graph = [&]() {
+ inputs.begin_graph();
+ attention[count]->begin_graph();
+ auto gf = new_graph_custom(MAX_GRAPH_SIZE);
+ auto ctx = get_context(gf);
+ const int origin = std::max(0, frame - 7);
+ std::vector<int> times;
+ for (int i = frame - count + 1; i <= frame; ++i)
+ times.push_back(i - origin);
+ auto qpe = inputs.rope(&ctx, width, height, {frame - origin}, {44, 42, 42});
+ auto kpe = inputs.rope(&ctx, width, height, times, {44, 42, 42});
+ auto out = model.forward(&ctx, make_input(noise), make_input(lq), make_input(timestep), qpe, kpe, *attention[count]);
+ ggml_build_forward_expand(gf, out);
+ return gf;
+ };
+ auto prediction = restore_trailing_singleton_dims(GGMLRunner::compute(graph, n_threads, false), 4);
+ if (prediction.empty())
+ return {};
+ ++frame;
+ for (int64_t i = 0; i < prediction.numel(); ++i)
+ prediction.data()[i] = noise.data()[i] - prediction.data()[i];
+ return prediction;
+ }
+ };
+
+} // namespace RelayVSR
+
+#endif // __SD_MODEL_DIFFUSION_RELAY_VSR_HPP__
diff --git a/src/model/upscaler/flash_vsr.hpp b/src/model/upscaler/flash_vsr.hpp
new file mode 100644
index 0000000..470bfc3
--- /dev/null
+++ b/src/model/upscaler/flash_vsr.hpp
@@ -0,0 +1,253 @@
+#ifndef __SD_MODEL_UPSCALER_FLASH_VSR_HPP__
+#define __SD_MODEL_UPSCALER_FLASH_VSR_HPP__
+
+#include "model/common/vsr.hpp"
+
+namespace RelayVSR {
+
+ struct FlashVSRConfig {
+ int dim = 512;
+ int depth = 10;
+ int num_heads = 8;
+ int num_kv_groups = 2;
+ int hidden_dim = 1408;
+ int head_dim = 1024;
+ int upscale_factor = 4;
+
+ static FlashVSRConfig detect_from_weights(const String2TensorStorage& weights, const std::string& prefix, int scale) {
+ FlashVSRConfig config;
+ config.upscale_factor = scale;
+ auto weight = [&](const std::string& name) -> const TensorStorage& { return weights.at(prefix + "." + name); };
+ config.dim = int(weight("backbone.blocks.0.attention.query.weight").ne[1]);
+ config.num_heads = config.dim / 64;
+ config.num_kv_groups = int(weight("backbone.blocks.0.attention.key.weight").ne[1]) / 64;
+ config.hidden_dim = int(weight("backbone.blocks.0.feed_forward.down.weight").logical_ne0());
+ config.head_dim = int(weight("spatial_head.0.weight").ne[1]);
+ config.depth = 0;
+ while (weights.count(prefix + ".backbone.blocks." + std::to_string(config.depth) + ".attention.query.weight"))
+ ++config.depth;
+ if ((scale != 2 && scale != 4) || config.dim != 512 || config.depth != 10 || config.num_kv_groups != 2 ||
+ weight("condition_encoder.projection.weight").ne[2] != 48 || weight("residual_projection.weight").ne[1] != 768)
+ throw std::runtime_error("Unsupported RelayVSR Flash architecture");
+ LOG_VERBOSE("flashvsr: dim = %d, depth = %d, hidden_dim = %d, scale = %d", config.dim, config.depth, config.hidden_dim, scale);
+ return config;
+ }
+ };
+
+ class LRFrameEncoder : public GGMLBlock {
+ public:
+ explicit LRFrameEncoder(const FlashVSRConfig& config) {
+ int channels[] = {3, 64, 128, 128};
+ for (int i = 0; i < 3; ++i) {
+ int stride = i < 2 ? 2 : 4 / config.upscale_factor;
+ blocks["layers." + std::to_string(i * 2)] = std::make_shared<Conv2d>(channels[i], channels[i + 1], std::pair<int, int>{3, 3}, std::pair<int, int>{stride, stride}, std::pair<int, int>{1, 1});
+ }
+ blocks["layers.6"] = std::make_shared<Conv2d>(128, config.dim, std::pair<int, int>{1, 1});
+ blocks["norm"] = std::make_shared<RMSNorm>(config.dim, 1e-5f);
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+ auto layers_6 = std::dynamic_pointer_cast<Conv2d>(blocks["layers.6"]);
+ auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
+ for (int i = 0; i < 3; ++i)
+ x = ggml_silu(ctx->ggml_ctx, std::dynamic_pointer_cast<Conv2d>(blocks["layers." + std::to_string(i * 2)])->forward(ctx, x));
+ x = layers_6->forward(ctx, x);
+ return norm->forward(ctx, VSR::tokens(ctx->ggml_ctx, x));
+ }
+ };
+
+ class ConditionEncoder : public GGMLBlock {
+ public:
+ explicit ConditionEncoder(const FlashVSRConfig& config) {
+ blocks["projection"] = std::make_shared<Conv2d>(48, config.dim, std::pair<int, int>{1, 1});
+ blocks["norm"] = std::make_shared<RMSNorm>(config.dim, 1e-5f);
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+ auto projection = std::dynamic_pointer_cast<Conv2d>(blocks["projection"]);
+ auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
+ return norm->forward(ctx, VSR::tokens(ctx->ggml_ctx, projection->forward(ctx, x)));
+ }
+ };
+
+ class ConditionedGroupedQueryAttention : public GGMLBlock {
+ FlashVSRConfig config_;
+
+ public:
+ explicit ConditionedGroupedQueryAttention(const FlashVSRConfig& config)
+ : config_(config) {
+ blocks["query"] = std::make_shared<Linear>(config.dim, config.dim);
+ blocks["output"] = std::make_shared<Linear>(config.dim, config.dim);
+ for (const auto& name : {"key", "value", "condition_key", "condition_value"})
+ blocks[name] = std::make_shared<Linear>(config.dim, 64 * config.num_kv_groups);
+ blocks["key_norm"] = std::make_shared<RMSNorm>(64, 1e-5f);
+ blocks["value_norm"] = std::make_shared<RMSNorm>(64, 1e-5f);
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* qpe, ggml_tensor* kpe, const VSR::LocalAttention& local, const std::string& name, int condition_slot, int start_slot, int end_slot, bool history) {
+ auto query = std::dynamic_pointer_cast<Linear>(blocks["query"]);
+ auto output = std::dynamic_pointer_cast<Linear>(blocks["output"]);
+ auto key_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["key_norm"]);
+ auto value_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["value_norm"]);
+ auto c = ctx->ggml_ctx;
+ const bool condition = condition_slot >= 0;
+ auto key = std::dynamic_pointer_cast<Linear>(blocks[condition ? "condition_key" : "key"]);
+ auto value = std::dynamic_pointer_cast<Linear>(blocks[condition ? "condition_value" : "value"]);
+ auto q = query->forward(ctx, x);
+ auto k = key->forward(ctx, x);
+ auto v = value->forward(ctx, x);
+ const int64_t count = x->ne[1];
+ k = key_norm->forward(ctx, ggml_reshape_3d(c, k, 64, config_.num_kv_groups, count));
+ v = value_norm->forward(ctx, ggml_reshape_3d(c, v, 64, config_.num_kv_groups, count));
+ k = ggml_reshape_2d(c, k, 64 * config_.num_kv_groups, count);
+ v = ggml_reshape_2d(c, v, 64 * config_.num_kv_groups, count);
+ auto cache_prefix = name + (condition ? ".condition." + std::to_string(condition_slot) : ".video");
+ auto past_k = history ? ctx->get_cache_tensor(name + ".video.key") : nullptr;
+ auto past_v = history ? ctx->get_cache_tensor(name + ".video.value") : nullptr;
+ ctx->cache_tensor(cache_prefix + ".key", k);
+ ctx->cache_tensor(cache_prefix + ".value", v);
+ if (!condition) {
+ auto join = [&](ggml_tensor* current, ggml_tensor* past, const char* suffix) {
+ auto start = ctx->get_cache_tensor(name + ".condition." + std::to_string(start_slot) + suffix);
+ auto end = ctx->get_cache_tensor(name + ".condition." + std::to_string(end_slot) + suffix);
+ if (!start || !end)
+ throw std::runtime_error("RelayVSR Flash endpoint cache is missing");
+ auto value = ggml_concat(c, start, end, 1);
+ if (past)
+ value = ggml_concat(c, value, past, 1);
+ return ggml_concat(c, value, current, 1);
+ };
+ k = join(k, past_k, ".key");
+ v = join(v, past_v, ".value");
+ }
+ q = VSR::apply_rope(ctx, q, qpe, config_.num_heads);
+ k = VSR::apply_rope(ctx, k, kpe, config_.num_kv_groups);
+ return output->forward(ctx, local.forward(ctx, q, k, v, config_.num_heads));
+ }
+ };
+
+ class ConditionedTransformerBlock : public GGMLBlock {
+ public:
+ explicit ConditionedTransformerBlock(const FlashVSRConfig& config) {
+ blocks["attention_norm"] = std::make_shared<RMSNorm>(config.dim, 1e-5f);
+ blocks["attention"] = std::make_shared<ConditionedGroupedQueryAttention>(config);
+ blocks["feed_forward_norm"] = std::make_shared<RMSNorm>(config.dim, 1e-5f);
+ blocks["feed_forward.gate_up_proj"] = std::make_shared<Linear>(config.dim, config.hidden_dim * 2, false);
+ blocks["feed_forward.down"] = std::make_shared<Linear>(config.hidden_dim, config.dim, false);
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* qpe, ggml_tensor* kpe, const VSR::LocalAttention& local, const std::string& name, int condition_slot, int start_slot, int end_slot, bool history) {
+ auto attention_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm"]);
+ auto feed_forward_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["feed_forward_norm"]);
+ auto feed_forward_gate_up_proj = std::dynamic_pointer_cast<Linear>(blocks["feed_forward.gate_up_proj"]);
+ auto feed_forward_down = std::dynamic_pointer_cast<Linear>(blocks["feed_forward.down"]);
+ auto c = ctx->ggml_ctx;
+ auto attention = std::dynamic_pointer_cast<ConditionedGroupedQueryAttention>(blocks["attention"]);
+ auto y = attention->forward(ctx, attention_norm->forward(ctx, x), qpe, kpe, local, name, condition_slot, start_slot, end_slot, history);
+ x = ggml_add(c, x, y);
+ y = feed_forward_gate_up_proj->forward(ctx, feed_forward_norm->forward(ctx, x));
+ auto parts = ggml_ext_chunk(c, y, 2, 0);
+ y = ggml_mul(c, ggml_silu(c, parts[0]), parts[1]);
+ return ggml_add(c, x, feed_forward_down->forward(ctx, y));
+ }
+ };
+
+ class FlashVSRConditionNet : public GGMLBlock {
+ FlashVSRConfig config_;
+
+ void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override {
+ params["video_type_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, config_.dim, 1, 1);
+ params["condition_type_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, config_.dim, 1, 1);
+ }
+
+ public:
+ explicit FlashVSRConditionNet(const FlashVSRConfig& config)
+ : config_(config) {
+ blocks["lr_encoder"] = std::make_shared<LRFrameEncoder>(config);
+ blocks["condition_encoder"] = std::make_shared<ConditionEncoder>(config);
+ for (int i = 0; i < config.depth; ++i)
+ blocks["backbone.blocks." + std::to_string(i)] = std::make_shared<ConditionedTransformerBlock>(config);
+ blocks["spatial_head.0"] = std::make_shared<Linear>(config.dim, config.head_dim);
+ blocks["residual_projection"] = std::make_shared<Linear>(config.head_dim, 768);
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input, ggml_tensor* qpe, ggml_tensor* kpe, const VSR::LocalAttention& local, int condition_slot, int start_slot, int end_slot, bool history) {
+ auto spatial_head_0 = std::dynamic_pointer_cast<Linear>(blocks["spatial_head.0"]);
+ auto residual_projection = std::dynamic_pointer_cast<Linear>(blocks["residual_projection"]);
+ auto c = ctx->ggml_ctx;
+ const bool condition = condition_slot >= 0;
+ ggml_tensor* x;
+ if (condition)
+ x = std::dynamic_pointer_cast<ConditionEncoder>(blocks["condition_encoder"])->forward(ctx, input);
+ else
+ x = std::dynamic_pointer_cast<LRFrameEncoder>(blocks["lr_encoder"])->forward(ctx, input);
+ x = ggml_add(c, x, params[condition ? "condition_type_embedding" : "video_type_embedding"]);
+ for (int i = 0; i < config_.depth; ++i) {
+ auto name = "backbone.blocks." + std::to_string(i);
+ x = std::dynamic_pointer_cast<ConditionedTransformerBlock>(blocks[name])->forward(ctx, x, qpe, kpe, local, name, condition_slot, start_slot, end_slot, history);
+ sd::ggml_graph_cut::mark_graph_cut(x, "flashvsr." + name, "hidden");
+ }
+ if (condition)
+ return x;
+ x = residual_projection->forward(ctx, ggml_silu(c, spatial_head_0->forward(ctx, x)));
+ return VSR::pixel_shuffle(c, x, local.width, local.height, 16);
+ }
+ };
+
+ struct FlashVSRRunner : GGMLRunner {
+ FlashVSRConfig config;
+ FlashVSRConditionNet model;
+ VSR::GraphInputs inputs;
+ std::map<int, std::unique_ptr<VSR::LocalAttention>> attention;
+ bool history = false;
+
+ FlashVSRRunner(ggml_backend_t backend, const String2TensorStorage& weights, int scale, std::shared_ptr<DeviceResidencyManager> manager)
+ : GGMLRunner(backend, manager), config(FlashVSRConfig::detect_from_weights(weights, "flash", scale)), model(config) {
+ model.init(params_ctx, weights, "flash");
+ }
+ std::string get_desc() override { return "RelayVSR Flash"; }
+ void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) { model.get_param_tensors(tensors, "flash"); }
+ void reset() {
+ runner_end();
+ free_cache_ctx_and_buffer();
+ attention.clear();
+ inputs.clear();
+ history = false;
+ }
+
+ sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& input, int condition_slot, int start_slot = 0, int end_slot = 1) {
+ const bool condition = condition_slot >= 0;
+ const int stride = condition ? 1 : 16 / config.upscale_factor;
+ const int width = int(input.shape()[0]) / stride, height = int(input.shape()[1]) / stride;
+ const int planes = condition ? 1 : history ? 4
+ : 3;
+ if (!attention.count(planes)) {
+ std::vector<std::pair<int, int>> windows(planes, {24, 24});
+ windows[0] = {48, 48};
+ if (!condition)
+ windows[1] = {48, 48};
+ attention[planes] = std::make_unique<VSR::LocalAttention>(width, height, windows, false);
+ }
+ auto graph = [&]() {
+ inputs.begin_graph();
+ attention[planes]->begin_graph();
+ auto gf = new_graph_custom(MAX_GRAPH_SIZE);
+ auto ctx = get_context(gf);
+ auto qpe = inputs.rope(&ctx, width, height, {0}, {16, 24, 24});
+ std::vector<int> times = condition ? std::vector<int>{0} : history ? std::vector<int>{0, 0, -1, 0}
+ : std::vector<int>{0, 0, 0};
+ auto kpe = inputs.rope(&ctx, width, height, times, {16, 24, 24});
+ auto out = model.forward(&ctx, make_input(input), qpe, kpe, *attention[planes], condition_slot, start_slot, end_slot, history && !condition);
+ ggml_build_forward_expand(gf, out);
+ return gf;
+ };
+ auto result = restore_trailing_singleton_dims(GGMLRunner::compute(graph, n_threads, false), condition ? 2 : 4);
+ if (!condition && !result.empty())
+ history = true;
+ return result;
+ }
+ };
+
+} // namespace RelayVSR
+
+#endif // __SD_MODEL_UPSCALER_FLASH_VSR_HPP__
diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp
index 1f19b11..78a5e10 100644
--- a/src/name_conversion.cpp
+++ b/src/name_conversion.cpp
@@ -5,6 +5,39 @@
#include "core/util.h"
#include "name_conversion.h"
+std::string convert_relayvsr_tensor_name(std::string name) {
+ if (!starts_with(name, "wan."))
+ return name;
+ const std::vector<std::pair<std::string, std::string>> mappings = {
+ {"positive_conditioning", "prompt_embeds"},
+ {"lq_proj_in.linear.", "lq_proj.out_proj."},
+ {"lq_proj_in.", "lq_proj."},
+ {"text_embedding.0.", "condition_embedder.text_embedder.linear_1."},
+ {"text_embedding.2.", "condition_embedder.text_embedder.linear_2."},
+ {"time_embedding.0.", "condition_embedder.time_embedder.linear_1."},
+ {"time_embedding.2.", "condition_embedder.time_embedder.linear_2."},
+ {"time_projection.1.", "condition_embedder.time_proj."},
+ {"head.head.", "proj_out."},
+ {"head.modulation", "scale_shift_table"},
+ {".norm3.", ".norm2."},
+ {".modulation", ".scale_shift_table"},
+ {".ffn.0.", ".ffn.net.0.proj."},
+ {".ffn.2.", ".ffn.net.2."},
+ {".self_attn.", ".attn1."},
+ {".cross_attn.", ".attn2."},
+ {".q.", ".to_q."},
+ {".k.", ".to_k."},
+ {".v.", ".to_v."},
+ {".o.", ".to_out.0."},
+ };
+ for (const auto& mapping : mappings) {
+ auto pos = name.find(mapping.first);
+ if (pos != std::string::npos)
+ name.replace(pos, mapping.first.size(), mapping.second);
+ }
+ return name;
+}
+
void replace_with_name_map(std::string& name, const std::vector<std::pair<std::string, std::string>>& name_map) {
for (auto kv : name_map) {
size_t pos = name.find(kv.first);
diff --git a/src/name_conversion.h b/src/name_conversion.h
index 3fefcf7..8d90c7e 100644
--- a/src/name_conversion.h
+++ b/src/name_conversion.h
@@ -10,5 +10,6 @@ bool is_diffusion_model_name(const std::string& name);
bool is_first_stage_model_name(const std::string& name);
std::string convert_tensor_name(std::string name, SDVersion version);
+std::string convert_relayvsr_tensor_name(std::string name);
-#endif // __NAME_CONVERSTION_H__
\ No newline at end of file
+#endif // __NAME_CONVERSTION_H__
diff --git a/src/pipeline/relay_vsr.cpp b/src/pipeline/relay_vsr.cpp
new file mode 100644
index 0000000..bc5cad7
--- /dev/null
+++ b/src/pipeline/relay_vsr.cpp
@@ -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
diff --git a/src/pipeline/video_upscale.cpp b/src/pipeline/video_upscale.cpp
new file mode 100644
index 0000000..5b34a92
--- /dev/null
+++ b/src/pipeline/video_upscale.cpp
@@ -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;
+}
diff --git a/src/pipeline/video_upscale.h b/src/pipeline/video_upscale.h
new file mode 100644
index 0000000..62d0af1
--- /dev/null
+++ b/src/pipeline/video_upscale.h
@@ -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__