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