mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-11 19: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
@@ -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__
|
||||
Reference in new issue
Block a user