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

refactor: unify ggml operation helpers around runner context (#2314)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-11-10 01:58:04 +08:00
1 parent 73fef08224
commit a3c49b13be
76 files changed
+1983 -1831
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/vae/tae.hpp
+29 -33
View File
@@ -179,9 +179,9 @@ public:
// z: [n, z_channels, h, w]
// return: [n, out_channels, h*8, w*8]
auto h = ggml_ext_scale(ctx->ggml_ctx, z, 1.0f / 3.0f);
auto h = ggml_ext_scale(ctx, z, 1.0f / 3.0f);
h = ggml_tanh_inplace(ctx->ggml_ctx, h);
h = ggml_ext_scale(ctx->ggml_ctx, h, 3.0f);
h = ggml_ext_scale(ctx, h, 3.0f);
for (int i = 0; i < num_blocks * 3 + 10; i++) {
if (blocks.find(std::to_string(i)) == blocks.end()) {
@@ -322,10 +322,8 @@ public:
};
ggml_tensor*
patchify(ggml_context* ctx,
ggml_tensor* x,
int64_t patch_size,
int64_t b = 1) {
patchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t patch_size, int64_t b = 1) {
auto ctx = runner_ctx->ggml_ctx;
// x: [f, b*c, h*q, w*r]
// return: [f, b*c*r*q, h, w]
if (patch_size == 1) {
@@ -342,19 +340,17 @@ patchify(ggml_context* ctx,
int64_t w = W / r;
int64_t h = H / q;
x = ggml_reshape_4d(ctx, x, W, q, h, C * f); // [W, q, h, C*f]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [W, h, q, C*f]
x = ggml_reshape_4d(ctx, x, r, w, h, q * C * f); // [r, w, h, q*C*f]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [w, h, r, q*C*f]
x = ggml_reshape_4d(ctx, x, w, h, r * q * C, f); // [f, b*c*r*q, h, w]
x = ggml_reshape_4d(ctx, x, W, q, h, C * f); // [W, q, h, C*f]
x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [W, h, q, C*f]
x = ggml_reshape_4d(ctx, x, r, w, h, q * C * f); // [r, w, h, q*C*f]
x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3)); // [w, h, r, q*C*f]
x = ggml_reshape_4d(ctx, x, w, h, r * q * C, f); // [f, b*c*r*q, h, w]
return x;
}
ggml_tensor* unpatchify(ggml_context* ctx,
ggml_tensor* x,
int64_t patch_size,
int64_t b = 1) {
ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t patch_size, int64_t b = 1) {
auto ctx = runner_ctx->ggml_ctx;
// x: [f, b*c*r*q, h, w]
// return: [f, b*c, h*q, w*r]
if (patch_size == 1) {
@@ -367,10 +363,10 @@ ggml_tensor* unpatchify(ggml_context* ctx,
int64_t h = x->ne[1];
int64_t w = x->ne[0];
x = ggml_reshape_4d(ctx, x, w, h, r, q * c * b * f); // [q*c*b*f, r, h, w]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [r, w, h, q*c*b*f]
x = ggml_reshape_4d(ctx, x, r * w, h, q, c * b * f); // [c*b*f, q, h, r*w]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [r*w, q, h, c*b*f]
x = ggml_reshape_4d(ctx, x, w, h, r, q * c * b * f); // [q*c*b*f, r, h, w]
x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3)); // [r, w, h, q*c*b*f]
x = ggml_reshape_4d(ctx, x, r * w, h, q, c * b * f); // [c*b*f, q, h, r*w]
x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [r*w, q, h, c*b*f]
x = ggml_reshape_4d(ctx, x, r * w, q * h, c * b, f);
return x;
@@ -412,7 +408,7 @@ public:
auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["0"]);
if (patch_size > 1) {
z = patchify(ctx->ggml_ctx, z, patch_size, 1);
z = patchify(ctx, z, patch_size, 1);
}
auto h = first_conv->forward(ctx, z);
@@ -427,7 +423,7 @@ public:
h = conv->forward(ctx, h);
for (int j = 0; j < num_blocks; j++) {
auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0);
auto mem = ggml_ext_pad_ext(ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
h = block->forward(ctx, h, mem);
}
@@ -487,9 +483,9 @@ public:
auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["1"]);
// Clamp()
auto h = ggml_ext_scale(ctx->ggml_ctx,
auto h = ggml_ext_scale(ctx,
ggml_tanh_inplace(ctx->ggml_ctx,
ggml_ext_scale(ctx->ggml_ctx, z, 1.0f / 3.0f)),
ggml_ext_scale(ctx, z, 1.0f / 3.0f)),
3.0f,
true);
@@ -498,7 +494,7 @@ public:
int index = 3;
for (int i = 0; i < num_layers; i++) {
for (int j = 0; j < num_blocks; j++) {
auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0);
auto mem = ggml_ext_pad_ext(ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
if (is_wide) {
auto block = std::dynamic_pointer_cast<WideMemBlock>(blocks[std::to_string(index++)]);
@@ -521,7 +517,7 @@ public:
auto last_conv = std::dynamic_pointer_cast<Conv2d>(blocks[std::to_string(++index)]);
h = last_conv->forward(ctx, h);
if (patch_size > 1) {
h = unpatchify(ctx->ggml_ctx, h, patch_size, 1);
h = unpatchify(ctx, h, patch_size, 1);
}
// shape(W, H, 3, (t_upscale - 1) + T) => shape(W, H, 3, T)
h = ggml_view_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2], h->ne[3] - (t_upscale - 1), h->nb[1], h->nb[2], h->nb[3], (t_upscale - 1) * h->nb[3]);
@@ -579,7 +575,7 @@ public:
int64_t chunk_frames = 5 * decoder->t_upscale;
int64_t pad = (chunk_frames - (num_frames % chunk_frames)) % chunk_frames;
result = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, result, 0, 0, 0, 0, 0, 0, 0, static_cast<int>(pad), false, false);
result = ggml_ext_pad_ext(ctx, result, 0, 0, 0, 0, 0, 0, 0, static_cast<int>(pad), false, false);
int64_t num_chunks = (num_frames + pad) / chunk_frames;
auto to_trim = decoder->t_upscale - 1;
@@ -591,7 +587,7 @@ public:
i * chunk_frames * result->nb[3]);
to_concat.push_back(chunk);
}
result = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3);
result = ggml_ext_vec_concat(ctx, to_concat, 3);
result = ggml_view_4d(ctx->ggml_ctx, result,
result->ne[0], result->ne[1], result->ne[2],
result->ne[3] - decoder->t_upscale * 3,
@@ -645,7 +641,7 @@ public:
to_concat.push_back(chunk_padded);
}
ggml_tensor* x_in = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3);
ggml_tensor* x_in = ggml_ext_vec_concat(ctx, to_concat, 3);
out = encoder->forward(ctx, x_in);
} else {
std::vector<ggml_tensor*> to_concat = {};
@@ -662,7 +658,7 @@ public:
to_concat.push_back(chunk_out);
}
out = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3);
out = ggml_ext_vec_concat(ctx, to_concat, 3);
}
// Return x[:, :-3] - drop the last 3 elements in the T dimension
@@ -738,21 +734,21 @@ public:
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
auto decoder = std::dynamic_pointer_cast<TinyDecoder>(blocks["decoder.layers"]);
if (taef2) {
z = unpatchify(ctx->ggml_ctx, z, 2);
z = unpatchify(ctx, z, 2);
}
auto x = decoder->forward(ctx, z);
return f16 ? unpatchify(ctx->ggml_ctx, x, 2) : x;
return f16 ? unpatchify(ctx, x, 2) : x;
}
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
if (f16) {
GGML_ASSERT(x->ne[0] % 2 == 0 && x->ne[1] % 2 == 0);
x = patchify(ctx->ggml_ctx, x, 2);
x = patchify(ctx, x, 2);
}
auto encoder = std::dynamic_pointer_cast<TinyEncoder>(blocks["encoder.layers"]);
auto z = encoder->forward(ctx, x);
if (taef2) {
z = patchify(ctx->ggml_ctx, z, 2);
z = patchify(ctx, z, 2);
}
return z;
}