From 9234ad822a376a4724638467736118ed61d68ee9 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 8 Nov 2026 00:09:48 +0800 Subject: [PATCH] feat: add RelayVSR video upscaling with CUDA window attention (#2285) --- README.md | 1 + docs/relay_vsr.md | 70 +++++++ examples/cli/main.cpp | 67 ++++++- examples/common/common.cpp | 1 + examples/common/common.h | 3 +- include/stable-diffusion.h | 31 +++ src/core/cuda/cute_attention.cu | 11 +- src/core/cuda/registration.cu | 2 + src/core/cuda/window_attention.cuh | 284 ++++++++++++++++++++++++++++ src/core/ggml_extensions.cpp | 59 ++++++ src/core/ggml_extensions.h | 12 ++ src/core/tensor.hpp | 42 +++-- src/core/util.cpp | 12 +- src/model/common/vsr.hpp | 221 ++++++++++++++++++++++ src/model/diffusion/relay_vsr.hpp | 293 +++++++++++++++++++++++++++++ src/model/upscaler/flash_vsr.hpp | 253 +++++++++++++++++++++++++ src/name_conversion.cpp | 33 ++++ src/name_conversion.h | 3 +- src/pipeline/relay_vsr.cpp | 209 ++++++++++++++++++++ src/pipeline/video_upscale.cpp | 143 ++++++++++++++ src/pipeline/video_upscale.h | 46 +++++ 21 files changed, 1764 insertions(+), 32 deletions(-) create mode 100644 docs/relay_vsr.md create mode 100644 src/core/cuda/window_attention.cuh create mode 100644 src/model/common/vsr.hpp create mode 100644 src/model/diffusion/relay_vsr.hpp create mode 100644 src/model/upscaler/flash_vsr.hpp create mode 100644 src/pipeline/relay_vsr.cpp create mode 100644 src/pipeline/video_upscale.cpp create mode 100644 src/pipeline/video_upscale.h diff --git a/README.md b/README.md index 28af00a6..2442865b 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 00000000..1352f41f --- /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 591a1511..bcf73595 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -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 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 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 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 02c30880..79349241 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 4ecbef4f..d4df4fad 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 1c61e4b1..a1d82e78 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 cdafc998..1eaefa4b 100644 --- a/src/core/cuda/cute_attention.cu +++ b/src/core/cuda/cute_attention.cu @@ -150,7 +150,7 @@ namespace sd_fattn_cute { } } - template + template struct Softmax { decltype(make_tensor(Shape>{})) maximum, sum; template @@ -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 1f481743..d92be448 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 00000000..600cd91b --- /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 + __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 + __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; + 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(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>{})); + auto kv_id = thread_copy.partition_S(make_identity_tensor(Shape, Int>{})); + 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>{}); + 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>{})); + 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>{}); + 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(acc_s, acc_o); + else + softmax.template update(acc_s, acc_o); + auto rP = to_half(acc_s); + auto tOrP = make_tensor(rP.data(), convert_layout_acc_Aregs(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>{})); +#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 +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; + constexpr int shared_bytes = sd_fattn_cute::Traits<64, 64, D>::SharedBytes; + static std::mutex mutex; + static std::unordered_map initialized; + { + std::lock_guard lock(mutex); + if (!initialized[context.device]) { + auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes); + if (error != cudaSuccess) + return error; + initialized[context.device] = true; + } + } + kernel<<(context.stream)>>>( + reinterpret_cast(inputs[0]), reinterpret_cast(inputs[1]), + reinterpret_cast(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_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(context.stream); + auto* base = static_cast(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(base); + const int64_t heads = src->ne[0] / dim; + copy_to_half<<<(ggml_nelements(src) + 255) / 256, 256, 0, stream>>>(static_cast(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(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 25fa2226..dfe5bca4 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 bc9ee9a5..87a51a7e 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 a925dbff..715068c3 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 out_coord = tensor_unravel_index(flat, out_shape); - std::vector src_coord(static_cast(dim()), 0); - for (size_t i = 0; i < dims.size(); ++i) { - src_coord[dims[i]] = out_coord[i]; + std::vector 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(y)]; + parallel_for(0, plane_count * output_height, std::max(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(y)]; for (int64_t x = 0; x < output_width; ++x) { const auto& x_axis = x_contributors[static_cast(x)]; double value = 0.0; @@ -1055,7 +1071,7 @@ namespace sd { output.data()[output_plane_offset + y * output_width + x] = static_cast(value); } } - } + }); return output; } diff --git a/src/core/util.cpp b/src/core/util.cpp index a691b423..9de1ace2 100644 --- a/src/core/util.cpp +++ b/src/core/util.cpp @@ -807,13 +807,11 @@ sd::Tensor sd_image_to_tensor(sd_image_t image, static_cast(image.height), static_cast(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 00000000..c203dff1 --- /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 +#include +#include +#include +#include +#include "model/common/ggml_block.hpp" +#include "model/common/rope.hpp" + +namespace VSR { + + struct GraphInputs { + std::map, sd::Tensor> rotations; + mutable std::unordered_map bindings; + + void begin_graph() { + bindings.clear(); + } + + void clear() { + bindings.clear(); + rotations.clear(); + } + + template + ggml_tensor* input(GGMLRunnerContext* ctx, const sd::Tensor& 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& times, const std::vector& axes) { + std::vector 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({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 queries; + sd::Tensor keys; + sd::Tensor mask; + }; + + struct LocalAttention { + int width; + int height; + std::vector> windows; + bool preserve_boundary; + mutable std::vector tiles; + GraphInputs inputs; + + void begin_graph() { + inputs.begin_graph(); + } + + LocalAttention(int width, int height, const std::vector>& 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({tile.width * tile.height}); + std::vector 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{top, left, top + wh, left + ww}; + } + return std::array{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({int64_t(keys.size())}); + std::copy(keys.begin(), keys.end(), tile.keys.data()); + tile.mask = sd::Tensor({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::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 00000000..a390aa34 --- /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(3 * patch * patch, hidden_dim1_, std::pair{1, 1}); + blocks["norm1"] = std::make_shared(hidden_dim1_, 1e-6f); + blocks["conv2"] = std::make_shared(hidden_dim1_, hidden_dim2_, std::pair{3, 3}, std::pair{2, 2}, std::pair{1, 1}); + blocks["norm2"] = std::make_shared(hidden_dim2_, 1e-6f); + blocks["out_proj"] = std::make_shared(hidden_dim2_, config.dim); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* unshuffled) { + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto out_proj = std::dynamic_pointer_cast(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(dim_, dim_); + blocks["norm_q"] = std::make_shared(dim_, 1e-6f); + blocks["norm_k"] = std::make_shared(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(blocks["to_q"]); + auto to_k = std::dynamic_pointer_cast(blocks["to_k"]); + auto to_v = std::dynamic_pointer_cast(blocks["to_v"]); + auto to_out_0 = std::dynamic_pointer_cast(blocks["to_out.0"]); + auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); + auto norm_k = std::dynamic_pointer_cast(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(dim_, 1e-6f, false); + blocks["norm2"] = std::make_shared(dim_, 1e-6f, true); + blocks["norm3"] = std::make_shared(dim_, 1e-6f, false); + blocks["attn1"] = std::make_shared(config); + blocks["attn2"] = std::make_shared(config); + blocks["ffn.net.0.proj"] = std::make_shared(dim_, config.ffn_dim); + blocks["ffn.net.2"] = std::make_shared(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(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto norm3 = std::dynamic_pointer_cast(blocks["norm3"]); + auto ffn_net_0_proj = std::dynamic_pointer_cast(blocks["ffn.net.0.proj"]); + auto ffn_net_2 = std::dynamic_pointer_cast(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(blocks["attn1"]); + auto attn2 = std::dynamic_pointer_cast(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(config.in_channels, config.dim, std::tuple{1, 2, 2}, std::tuple{1, 2, 2}); + blocks["lq_proj"] = std::make_shared(config); + blocks["condition_embedder.time_embedder.linear_1"] = std::make_shared(256, config.dim); + blocks["condition_embedder.time_embedder.linear_2"] = std::make_shared(config.dim, config.dim); + blocks["condition_embedder.time_proj"] = std::make_shared(config.dim, config.dim * 6); + blocks["condition_embedder.text_embedder.linear_1"] = std::make_shared(config.text_dim, config.dim); + blocks["condition_embedder.text_embedder.linear_2"] = std::make_shared(config.dim, config.dim); + for (int i = 0; i < config.num_layers; ++i) + blocks["blocks." + std::to_string(i)] = std::make_shared(config); + blocks["norm_out"] = std::make_shared(config.dim, 1e-6f, false); + blocks["proj_out"] = std::make_shared(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(blocks["patch_embedding"]); + auto condition_embedder_time_embedder_linear_1 = std::dynamic_pointer_cast(blocks["condition_embedder.time_embedder.linear_1"]); + auto condition_embedder_time_embedder_linear_2 = std::dynamic_pointer_cast(blocks["condition_embedder.time_embedder.linear_2"]); + auto condition_embedder_time_proj = std::dynamic_pointer_cast(blocks["condition_embedder.time_proj"]); + auto condition_embedder_text_embedder_linear_1 = std::dynamic_pointer_cast(blocks["condition_embedder.text_embedder.linear_1"]); + auto condition_embedder_text_embedder_linear_2 = std::dynamic_pointer_cast(blocks["condition_embedder.text_embedder.linear_2"]); + auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); + auto proj_out = std::dynamic_pointer_cast(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(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(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> attention; + sd::Tensor timestep; + int frame = 0; + + WanVSRRunner(ggml_backend_t backend, const String2TensorStorage& weights, int scale, std::shared_ptr 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& tensors) { model.get_param_tensors(tensors, "wan"); } + void reset() { + runner_end(); + free_cache_ctx_and_buffer(); + attention.clear(); + inputs.clear(); + frame = 0; + } + + sd::Tensor compute(int n_threads, const sd::Tensor& noise, const sd::Tensor& 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(width, height, std::vector>(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 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 00000000..470bfc30 --- /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(channels[i], channels[i + 1], std::pair{3, 3}, std::pair{stride, stride}, std::pair{1, 1}); + } + blocks["layers.6"] = std::make_shared(128, config.dim, std::pair{1, 1}); + blocks["norm"] = std::make_shared(config.dim, 1e-5f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto layers_6 = std::dynamic_pointer_cast(blocks["layers.6"]); + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + for (int i = 0; i < 3; ++i) + x = ggml_silu(ctx->ggml_ctx, std::dynamic_pointer_cast(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(48, config.dim, std::pair{1, 1}); + blocks["norm"] = std::make_shared(config.dim, 1e-5f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto projection = std::dynamic_pointer_cast(blocks["projection"]); + auto norm = std::dynamic_pointer_cast(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(config.dim, config.dim); + blocks["output"] = std::make_shared(config.dim, config.dim); + for (const auto& name : {"key", "value", "condition_key", "condition_value"}) + blocks[name] = std::make_shared(config.dim, 64 * config.num_kv_groups); + blocks["key_norm"] = std::make_shared(64, 1e-5f); + blocks["value_norm"] = std::make_shared(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(blocks["query"]); + auto output = std::dynamic_pointer_cast(blocks["output"]); + auto key_norm = std::dynamic_pointer_cast(blocks["key_norm"]); + auto value_norm = std::dynamic_pointer_cast(blocks["value_norm"]); + auto c = ctx->ggml_ctx; + const bool condition = condition_slot >= 0; + auto key = std::dynamic_pointer_cast(blocks[condition ? "condition_key" : "key"]); + auto value = std::dynamic_pointer_cast(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(config.dim, 1e-5f); + blocks["attention"] = std::make_shared(config); + blocks["feed_forward_norm"] = std::make_shared(config.dim, 1e-5f); + blocks["feed_forward.gate_up_proj"] = std::make_shared(config.dim, config.hidden_dim * 2, false); + blocks["feed_forward.down"] = std::make_shared(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(blocks["attention_norm"]); + auto feed_forward_norm = std::dynamic_pointer_cast(blocks["feed_forward_norm"]); + auto feed_forward_gate_up_proj = std::dynamic_pointer_cast(blocks["feed_forward.gate_up_proj"]); + auto feed_forward_down = std::dynamic_pointer_cast(blocks["feed_forward.down"]); + auto c = ctx->ggml_ctx; + auto attention = std::dynamic_pointer_cast(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(config); + blocks["condition_encoder"] = std::make_shared(config); + for (int i = 0; i < config.depth; ++i) + blocks["backbone.blocks." + std::to_string(i)] = std::make_shared(config); + blocks["spatial_head.0"] = std::make_shared(config.dim, config.head_dim); + blocks["residual_projection"] = std::make_shared(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(blocks["spatial_head.0"]); + auto residual_projection = std::dynamic_pointer_cast(blocks["residual_projection"]); + auto c = ctx->ggml_ctx; + const bool condition = condition_slot >= 0; + ggml_tensor* x; + if (condition) + x = std::dynamic_pointer_cast(blocks["condition_encoder"])->forward(ctx, input); + else + x = std::dynamic_pointer_cast(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(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> attention; + bool history = false; + + FlashVSRRunner(ggml_backend_t backend, const String2TensorStorage& weights, int scale, std::shared_ptr 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& tensors) { model.get_param_tensors(tensors, "flash"); } + void reset() { + runner_end(); + free_cache_ctx_and_buffer(); + attention.clear(); + inputs.clear(); + history = false; + } + + sd::Tensor compute(int n_threads, const sd::Tensor& 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> windows(planes, {24, 24}); + windows[0] = {48, 48}; + if (!condition) + windows[1] = {48, 48}; + attention[planes] = std::make_unique(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 times = condition ? std::vector{0} : history ? std::vector{0, 0, -1, 0} + : std::vector{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 1f19b118..78a5e106 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> 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>& 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 3fefcf78..8d90c7e2 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 00000000..bc5cad71 --- /dev/null +++ b/src/pipeline/relay_vsr.cpp @@ -0,0 +1,209 @@ +#include "pipeline/video_upscale.h" + +#include + +#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 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({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({width, height - frame.height, 3, 1}) + edge; + result = sd::ops::concat(result, padding, 1); + } + return result; + } + + static sd::Tensor prepare_vsr_lq(const sd::Tensor& 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 wan_; + std::unique_ptr 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 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(wan_backend, weights, upscale_factor, runtime.manager); + flash_ = std::make_unique(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 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 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(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 create_relay_vsr_upscaler() { + return std::make_unique(); + } + + } // namespace pipeline +} // namespace sd diff --git a/src/pipeline/video_upscale.cpp b/src/pipeline/video_upscale.cpp new file mode 100644 index 00000000..5b34a928 --- /dev/null +++ b/src/pipeline/video_upscale.cpp @@ -0,0 +1,143 @@ +#include "pipeline/video_upscale.h" + +#include + +#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(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(); + 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( + static_cast(calloc(1, sizeof(sd_video_generation_result_t))), free_sd_video_generation_result); + if (!result) + return false; + result->frames = static_cast(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(); + 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 00000000..62d0af10 --- /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 + +#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 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 create_relay_vsr_upscaler(); + + class VideoUpscaler { + VideoUpscaleRuntime runtime_; + int upscale_factor_ = 0; + std::unique_ptr executor_; + std::unique_ptr 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__