DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 03:24:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add RelayVSR video upscaling with CUDA window attention (#2285)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-11-08 00:09:48 +08:00
1 parent b1f5a51a1e
commit 9234ad822a
21 files changed
+1764 -32
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/cuda/cute_attention.cu
+7 -4
View File
@@ -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"
src/core/cuda/registration.cu
+2
View File
@@ -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
}
}
src/core/cuda/window_attention.cuh
+284
View File
@@ -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, &params);
}
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, &params))
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
src/core/ggml_extensions.cpp
+59
View File
@@ -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, &params);
}
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, &params, &out))
return nullptr;
sd_ggml_register_extensions(backend);
auto result = ggml_ext_op_create(ctx, sd_ggml_window_attn_op(), inputs, 3, &params);
return ggml_backend_supports_op(backend, result) ? result : nullptr;
}
void sd_ggml_register_extensions(ggml_backend_t backend) {
if (!backend) {
return;
src/core/ggml_extensions.h
+12
View File
@@ -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].
src/core/tensor.hpp
+29 -13
View File
@@ -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;
}
src/core/util.cpp
+5 -7
View File
@@ -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,
src/model/common/vsr.hpp
+221
View File
@@ -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__
src/model/diffusion/relay_vsr.hpp
+293
View File
@@ -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__
src/model/upscaler/flash_vsr.hpp
+253
View File
@@ -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__
src/name_conversion.cpp
+33
View File
@@ -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);
src/name_conversion.h
+2 -1
View File
@@ -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__
src/pipeline/relay_vsr.cpp
+209
View File
@@ -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
src/pipeline/video_upscale.cpp
+143
View File
@@ -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;
}
src/pipeline/video_upscale.h
+46
View File
@@ -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__