mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-11 19:36:33 +00:00
refactor: reorganize src model layout (#1615)
This commit is contained in:
leejet authored and GitHub committed
1 parent
dfb2390dd4
commit
f3fd359b58
81 files changed
+407
-385
No files matched your search
@@ -0,0 +1,375 @@
|
||||
#ifndef __SD_MODEL_UPSCALER_ESRGAN_HPP__
|
||||
#define __SD_MODEL_UPSCALER_ESRGAN_HPP__
|
||||
|
||||
#include "core/ggml_extend.hpp"
|
||||
#include "model.h"
|
||||
|
||||
/*
|
||||
=================================== ESRGAN ===================================
|
||||
References:
|
||||
https://github.com/xinntao/Real-ESRGAN/blob/master/inference_realesrgan.py
|
||||
https://github.com/XPixelGroup/BasicSR/blob/v1.4.2/basicsr/archs/rrdbnet_arch.py
|
||||
|
||||
*/
|
||||
|
||||
class ResidualDenseBlock : public GGMLBlock {
|
||||
protected:
|
||||
int num_feat;
|
||||
int num_grow_ch;
|
||||
|
||||
public:
|
||||
ResidualDenseBlock(int num_feat = 64, int num_grow_ch = 32)
|
||||
: num_feat(num_feat), num_grow_ch(num_grow_ch) {
|
||||
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
|
||||
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
|
||||
blocks["conv3"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
|
||||
blocks["conv4"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
|
||||
blocks["conv5"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 4 * num_grow_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
|
||||
}
|
||||
|
||||
ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [n, num_feat, h, w]
|
||||
// return: [n, num_feat, h, w]
|
||||
|
||||
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
|
||||
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv2"]);
|
||||
auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
|
||||
auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
|
||||
auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
|
||||
|
||||
auto x1 = lrelu(ctx, conv1->forward(ctx, x));
|
||||
auto x_cat = ggml_concat(ctx->ggml_ctx, x, x1, 2);
|
||||
auto x2 = lrelu(ctx, conv2->forward(ctx, x_cat));
|
||||
x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x2, 2);
|
||||
auto x3 = lrelu(ctx, conv3->forward(ctx, x_cat));
|
||||
x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x3, 2);
|
||||
auto x4 = lrelu(ctx, conv4->forward(ctx, x_cat));
|
||||
x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x4, 2);
|
||||
auto x5 = conv5->forward(ctx, x_cat);
|
||||
|
||||
x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, x5, 0.2f), x);
|
||||
return x5;
|
||||
}
|
||||
};
|
||||
|
||||
class RRDB : public GGMLBlock {
|
||||
public:
|
||||
RRDB(int num_feat, int num_grow_ch = 32) {
|
||||
blocks["rdb1"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
|
||||
blocks["rdb2"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
|
||||
blocks["rdb3"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [n, num_feat, h, w]
|
||||
// return: [n, num_feat, h, w]
|
||||
|
||||
auto rdb1 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb1"]);
|
||||
auto rdb2 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb2"]);
|
||||
auto rdb3 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb3"]);
|
||||
|
||||
auto out = rdb1->forward(ctx, x);
|
||||
out = rdb2->forward(ctx, out);
|
||||
out = rdb3->forward(ctx, out);
|
||||
|
||||
out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, out, 0.2f), x);
|
||||
return out;
|
||||
}
|
||||
};
|
||||
|
||||
class RRDBNet : public GGMLBlock {
|
||||
protected:
|
||||
int scale = 4;
|
||||
int num_block = 23;
|
||||
int num_in_ch = 3;
|
||||
int num_out_ch = 3;
|
||||
int num_feat = 64;
|
||||
int num_grow_ch = 32;
|
||||
|
||||
public:
|
||||
RRDBNet(int scale, int num_block, int num_in_ch, int num_out_ch, int num_feat, int num_grow_ch)
|
||||
: scale(scale), num_block(num_block), num_in_ch(num_in_ch), num_out_ch(num_out_ch), num_feat(num_feat), num_grow_ch(num_grow_ch) {
|
||||
blocks["conv_first"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_in_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
|
||||
for (int i = 0; i < num_block; i++) {
|
||||
std::string name = "body." + std::to_string(i);
|
||||
blocks[name] = std::shared_ptr<GGMLBlock>(new RRDB(num_feat, num_grow_ch));
|
||||
}
|
||||
blocks["conv_body"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
|
||||
if (scale >= 2) {
|
||||
blocks["conv_up1"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
|
||||
}
|
||||
if (scale == 4) {
|
||||
blocks["conv_up2"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
|
||||
}
|
||||
blocks["conv_hr"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
|
||||
blocks["conv_last"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_out_ch, {3, 3}, {1, 1}, {1, 1}));
|
||||
}
|
||||
|
||||
int get_scale() { return scale; }
|
||||
int get_num_block() { return num_block; }
|
||||
|
||||
ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [n, num_in_ch, h, w]
|
||||
// return: [n, num_out_ch, h*scale, w*scale]
|
||||
auto conv_first = std::dynamic_pointer_cast<Conv2d>(blocks["conv_first"]);
|
||||
auto conv_body = std::dynamic_pointer_cast<Conv2d>(blocks["conv_body"]);
|
||||
auto conv_hr = std::dynamic_pointer_cast<Conv2d>(blocks["conv_hr"]);
|
||||
auto conv_last = std::dynamic_pointer_cast<Conv2d>(blocks["conv_last"]);
|
||||
|
||||
auto feat = conv_first->forward(ctx, x);
|
||||
sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.prelude", "feat");
|
||||
auto body_feat = feat;
|
||||
for (int i = 0; i < num_block; i++) {
|
||||
std::string name = "body." + std::to_string(i);
|
||||
auto block = std::dynamic_pointer_cast<RRDB>(blocks[name]);
|
||||
|
||||
body_feat = block->forward(ctx, body_feat);
|
||||
sd::ggml_graph_cut::mark_graph_cut(body_feat, "esrgan.body." + std::to_string(i), "feat");
|
||||
}
|
||||
body_feat = conv_body->forward(ctx, body_feat);
|
||||
feat = ggml_add(ctx->ggml_ctx, feat, body_feat);
|
||||
sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.body.out", "feat");
|
||||
// upsample
|
||||
if (scale >= 2) {
|
||||
auto conv_up1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up1"]);
|
||||
feat = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
|
||||
sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up1", "feat");
|
||||
if (scale == 4) {
|
||||
auto conv_up2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up2"]);
|
||||
feat = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
|
||||
sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up2", "feat");
|
||||
}
|
||||
}
|
||||
// for all scales
|
||||
auto out = conv_last->forward(ctx, lrelu(ctx, conv_hr->forward(ctx, feat)));
|
||||
sd::ggml_graph_cut::mark_graph_cut(out, "esrgan.final", "out");
|
||||
return out;
|
||||
}
|
||||
};
|
||||
|
||||
struct ESRGAN : public GGMLRunner {
|
||||
std::unique_ptr<RRDBNet> rrdb_net;
|
||||
int scale = 4;
|
||||
int tile_size = 128; // avoid cuda OOM for 4gb VRAM
|
||||
|
||||
ESRGAN(ggml_backend_t backend,
|
||||
ggml_backend_t params_backend,
|
||||
int tile_size = 128,
|
||||
const String2TensorStorage& tensor_storage_map = {})
|
||||
: GGMLRunner(backend, params_backend) {
|
||||
this->tile_size = tile_size;
|
||||
}
|
||||
|
||||
std::string get_desc() override {
|
||||
return "esrgan";
|
||||
}
|
||||
|
||||
bool load_from_file(const std::string& file_path, int n_threads) {
|
||||
LOG_INFO("loading esrgan from '%s'", file_path.c_str());
|
||||
|
||||
ModelLoader model_loader;
|
||||
if (!model_loader.init_from_file_and_convert_name(file_path)) {
|
||||
LOG_ERROR("init esrgan model loader from file failed: '%s'", file_path.c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
// Get tensor names
|
||||
auto tensor_names = model_loader.get_tensor_names();
|
||||
|
||||
// Detect if it's ESRGAN format
|
||||
bool is_ESRGAN = std::find(tensor_names.begin(), tensor_names.end(), "model.0.weight") != tensor_names.end();
|
||||
|
||||
// Detect parameters from tensor names
|
||||
int detected_num_block = 0;
|
||||
if (is_ESRGAN) {
|
||||
for (const auto& name : tensor_names) {
|
||||
if (name.find("model.1.sub.") == 0) {
|
||||
size_t first_dot = name.find('.', 12);
|
||||
if (first_dot != std::string::npos) {
|
||||
size_t second_dot = name.find('.', first_dot + 1);
|
||||
if (second_dot != std::string::npos && name.substr(first_dot + 1, 3) == "RDB") {
|
||||
try {
|
||||
int idx = std::stoi(name.substr(12, first_dot - 12));
|
||||
detected_num_block = std::max(detected_num_block, idx + 1);
|
||||
} catch (...) {
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Original format
|
||||
for (const auto& name : tensor_names) {
|
||||
if (name.find("body.") == 0) {
|
||||
size_t pos = name.find('.', 5);
|
||||
if (pos != std::string::npos) {
|
||||
try {
|
||||
int idx = std::stoi(name.substr(5, pos - 5));
|
||||
detected_num_block = std::max(detected_num_block, idx + 1);
|
||||
} catch (...) {
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int detected_scale = 4; // default
|
||||
if (is_ESRGAN) {
|
||||
// For ESRGAN format, detect scale by highest model number
|
||||
int max_model_num = 0;
|
||||
for (const auto& name : tensor_names) {
|
||||
if (name.find("model.") == 0) {
|
||||
size_t dot_pos = name.find('.', 6);
|
||||
if (dot_pos != std::string::npos) {
|
||||
try {
|
||||
int num = std::stoi(name.substr(6, dot_pos - 6));
|
||||
max_model_num = std::max(max_model_num, num);
|
||||
} catch (...) {
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (max_model_num <= 4) {
|
||||
detected_scale = 1;
|
||||
} else if (max_model_num <= 7) {
|
||||
detected_scale = 2;
|
||||
} else {
|
||||
detected_scale = 4;
|
||||
}
|
||||
} else {
|
||||
// Original format
|
||||
bool has_conv_up2 = std::any_of(tensor_names.begin(), tensor_names.end(), [](const std::string& name) {
|
||||
return name == "conv_up2.weight";
|
||||
});
|
||||
bool has_conv_up1 = std::any_of(tensor_names.begin(), tensor_names.end(), [](const std::string& name) {
|
||||
return name == "conv_up1.weight";
|
||||
});
|
||||
if (has_conv_up2) {
|
||||
detected_scale = 4;
|
||||
} else if (has_conv_up1) {
|
||||
detected_scale = 2;
|
||||
} else {
|
||||
detected_scale = 1;
|
||||
}
|
||||
}
|
||||
|
||||
int detected_num_in_ch = 3;
|
||||
int detected_num_out_ch = 3;
|
||||
int detected_num_feat = 64;
|
||||
int detected_num_grow_ch = 32;
|
||||
|
||||
// Create RRDBNet with detected parameters
|
||||
rrdb_net = std::make_unique<RRDBNet>(detected_scale, detected_num_block, detected_num_in_ch, detected_num_out_ch, detected_num_feat, detected_num_grow_ch);
|
||||
rrdb_net->init(params_ctx, {}, "");
|
||||
|
||||
if (!alloc_params_buffer()) {
|
||||
LOG_ERROR("esrgan model buffer allocation failed");
|
||||
return false;
|
||||
}
|
||||
|
||||
std::map<std::string, ggml_tensor*> esrgan_tensors;
|
||||
rrdb_net->get_param_tensors(esrgan_tensors);
|
||||
|
||||
bool success;
|
||||
if (is_ESRGAN) {
|
||||
// Build name mapping for ESRGAN format
|
||||
std::map<std::string, std::string> expected_to_model;
|
||||
expected_to_model["conv_first.weight"] = "model.0.weight";
|
||||
expected_to_model["conv_first.bias"] = "model.0.bias";
|
||||
|
||||
for (int i = 0; i < detected_num_block; i++) {
|
||||
for (int j = 1; j <= 3; j++) {
|
||||
for (int k = 1; k <= 5; k++) {
|
||||
std::string expected_weight = "body." + std::to_string(i) + ".rdb" + std::to_string(j) + ".conv" + std::to_string(k) + ".weight";
|
||||
std::string model_weight = "model.1.sub." + std::to_string(i) + ".RDB" + std::to_string(j) + ".conv" + std::to_string(k) + ".0.weight";
|
||||
expected_to_model[expected_weight] = model_weight;
|
||||
|
||||
std::string expected_bias = "body." + std::to_string(i) + ".rdb" + std::to_string(j) + ".conv" + std::to_string(k) + ".bias";
|
||||
std::string model_bias = "model.1.sub." + std::to_string(i) + ".RDB" + std::to_string(j) + ".conv" + std::to_string(k) + ".0.bias";
|
||||
expected_to_model[expected_bias] = model_bias;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (detected_scale == 1) {
|
||||
expected_to_model["conv_body.weight"] = "model.1.sub." + std::to_string(detected_num_block) + ".weight";
|
||||
expected_to_model["conv_body.bias"] = "model.1.sub." + std::to_string(detected_num_block) + ".bias";
|
||||
expected_to_model["conv_hr.weight"] = "model.2.weight";
|
||||
expected_to_model["conv_hr.bias"] = "model.2.bias";
|
||||
expected_to_model["conv_last.weight"] = "model.4.weight";
|
||||
expected_to_model["conv_last.bias"] = "model.4.bias";
|
||||
} else {
|
||||
expected_to_model["conv_body.weight"] = "model.1.sub." + std::to_string(detected_num_block) + ".weight";
|
||||
expected_to_model["conv_body.bias"] = "model.1.sub." + std::to_string(detected_num_block) + ".bias";
|
||||
if (detected_scale >= 2) {
|
||||
expected_to_model["conv_up1.weight"] = "model.3.weight";
|
||||
expected_to_model["conv_up1.bias"] = "model.3.bias";
|
||||
}
|
||||
if (detected_scale == 4) {
|
||||
expected_to_model["conv_up2.weight"] = "model.6.weight";
|
||||
expected_to_model["conv_up2.bias"] = "model.6.bias";
|
||||
expected_to_model["conv_hr.weight"] = "model.8.weight";
|
||||
expected_to_model["conv_hr.bias"] = "model.8.bias";
|
||||
expected_to_model["conv_last.weight"] = "model.10.weight";
|
||||
expected_to_model["conv_last.bias"] = "model.10.bias";
|
||||
} else if (detected_scale == 2) {
|
||||
expected_to_model["conv_hr.weight"] = "model.5.weight";
|
||||
expected_to_model["conv_hr.bias"] = "model.5.bias";
|
||||
expected_to_model["conv_last.weight"] = "model.7.weight";
|
||||
expected_to_model["conv_last.bias"] = "model.7.bias";
|
||||
}
|
||||
}
|
||||
|
||||
std::map<std::string, ggml_tensor*> model_tensors;
|
||||
for (auto& p : esrgan_tensors) {
|
||||
auto it = expected_to_model.find(p.first);
|
||||
if (it != expected_to_model.end()) {
|
||||
model_tensors[it->second] = p.second;
|
||||
}
|
||||
}
|
||||
|
||||
success = model_loader.load_tensors(model_tensors, {}, n_threads);
|
||||
} else {
|
||||
success = model_loader.load_tensors(esrgan_tensors, {}, n_threads);
|
||||
}
|
||||
|
||||
if (!success) {
|
||||
LOG_ERROR("load esrgan tensors from model loader failed");
|
||||
return false;
|
||||
}
|
||||
|
||||
scale = rrdb_net->get_scale();
|
||||
LOG_INFO("esrgan model loaded with scale=%d, num_block=%d", scale, detected_num_block);
|
||||
return success;
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
|
||||
if (!rrdb_net)
|
||||
return nullptr;
|
||||
constexpr int kGraphNodes = 1 << 16; // 65k
|
||||
ggml_cgraph* gf = new_graph_custom(kGraphNodes);
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* out = rrdb_net->forward(&runner_ctx, x);
|
||||
ggml_build_forward_expand(gf, out);
|
||||
return gf;
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(const int n_threads,
|
||||
const sd::Tensor<float>& x) {
|
||||
auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
|
||||
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
|
||||
return result;
|
||||
}
|
||||
};
|
||||
|
||||
#endif // __SD_MODEL_UPSCALER_ESRGAN_HPP__
|
||||
@@ -0,0 +1,551 @@
|
||||
#ifndef __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
|
||||
#define __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
|
||||
|
||||
#include <cinttypes>
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "core/ggml_extend.hpp"
|
||||
#include "core/ggml_graph_cut.h"
|
||||
#include "core/util.h"
|
||||
#include "model.h"
|
||||
#include "model/diffusion/dit.hpp"
|
||||
|
||||
namespace LTXVUpsampler {
|
||||
constexpr int LTX_UPSAMPLER_GRAPH_SIZE = 10240;
|
||||
|
||||
struct LatentUpsamplerConfig {
|
||||
int64_t in_channels = 128;
|
||||
int64_t mid_channels = 1024;
|
||||
int num_blocks_per_stage = 4;
|
||||
int dims = 3;
|
||||
bool spatial_upsample = true;
|
||||
bool temporal_upsample = false;
|
||||
bool rational_resampler = false;
|
||||
float spatial_scale = 2.f;
|
||||
int spatial_up_num = 2;
|
||||
int spatial_down_den = 1;
|
||||
int temporal_up_factor = 1;
|
||||
};
|
||||
|
||||
static inline bool has_tensor(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& name) {
|
||||
return tensor_storage_map.find(name) != tensor_storage_map.end();
|
||||
}
|
||||
|
||||
static inline int64_t get_tensor_ne(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& name,
|
||||
int axis,
|
||||
int64_t fallback) {
|
||||
auto it = tensor_storage_map.find(name);
|
||||
if (it == tensor_storage_map.end() || axis < 0 || axis >= GGML_MAX_DIMS) {
|
||||
return fallback;
|
||||
}
|
||||
return it->second.ne[axis];
|
||||
}
|
||||
|
||||
static inline int64_t get_tensor_ne0(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& name,
|
||||
int64_t fallback) {
|
||||
return get_tensor_ne(tensor_storage_map, name, 0, fallback);
|
||||
}
|
||||
|
||||
static inline int count_module_blocks(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& module_name) {
|
||||
int max_block = -1;
|
||||
const std::string prefix = module_name + ".";
|
||||
for (const auto& pair : tensor_storage_map) {
|
||||
const std::string& name = pair.first;
|
||||
if (name.find(prefix) != 0) {
|
||||
continue;
|
||||
}
|
||||
size_t begin = prefix.size();
|
||||
size_t end = name.find('.', begin);
|
||||
if (end == std::string::npos) {
|
||||
continue;
|
||||
}
|
||||
int index = atoi(name.substr(begin, end - begin).c_str());
|
||||
max_block = std::max(max_block, index);
|
||||
}
|
||||
return max_block + 1;
|
||||
}
|
||||
|
||||
static inline LatentUpsamplerConfig detect_config_from_weights(const String2TensorStorage& tensor_storage_map) {
|
||||
LatentUpsamplerConfig config;
|
||||
config.mid_channels = get_tensor_ne0(tensor_storage_map, "initial_norm.weight", config.mid_channels);
|
||||
config.in_channels = get_tensor_ne0(tensor_storage_map, "final_conv.bias", config.in_channels);
|
||||
int detected_blocks = count_module_blocks(tensor_storage_map, "res_blocks");
|
||||
if (detected_blocks > 0) {
|
||||
config.num_blocks_per_stage = detected_blocks;
|
||||
}
|
||||
config.rational_resampler = has_tensor(tensor_storage_map, "upsampler.conv.weight");
|
||||
int64_t upsampler_out_channels = get_tensor_ne0(tensor_storage_map, "upsampler.0.bias", 0);
|
||||
config.spatial_upsample = config.rational_resampler || upsampler_out_channels == 4 * config.mid_channels;
|
||||
config.temporal_upsample = upsampler_out_channels == 2 * config.mid_channels;
|
||||
if (config.temporal_upsample) {
|
||||
config.temporal_up_factor = 2;
|
||||
}
|
||||
if (config.rational_resampler) {
|
||||
int64_t out_channels = get_tensor_ne(tensor_storage_map,
|
||||
"upsampler.conv.weight",
|
||||
3,
|
||||
config.mid_channels * 9);
|
||||
if (config.mid_channels > 0 && out_channels % config.mid_channels == 0) {
|
||||
int64_t ratio = out_channels / config.mid_channels;
|
||||
int num = static_cast<int>(std::round(std::sqrt(static_cast<double>(ratio))));
|
||||
if (num > 0 && static_cast<int64_t>(num) * num == ratio) {
|
||||
config.spatial_up_num = num;
|
||||
}
|
||||
}
|
||||
if (config.spatial_up_num == 3) {
|
||||
config.spatial_down_den = 2;
|
||||
config.spatial_scale = 1.5f;
|
||||
} else if (config.spatial_up_num == 4) {
|
||||
config.spatial_down_den = 1;
|
||||
config.spatial_scale = 4.f;
|
||||
} else {
|
||||
config.spatial_down_den = 1;
|
||||
config.spatial_scale = static_cast<float>(config.spatial_up_num);
|
||||
}
|
||||
}
|
||||
return config;
|
||||
}
|
||||
|
||||
class VideoGroupNorm : public GGMLBlock {
|
||||
protected:
|
||||
int num_groups;
|
||||
int64_t num_channels;
|
||||
float eps;
|
||||
std::string prefix;
|
||||
|
||||
void init_params(ggml_context* ctx,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "") override {
|
||||
SD_UNUSED(tensor_storage_map);
|
||||
this->prefix = prefix;
|
||||
params["weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_channels);
|
||||
params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_channels);
|
||||
}
|
||||
|
||||
public:
|
||||
VideoGroupNorm(int num_groups, int64_t num_channels, float eps = 1e-05f)
|
||||
: num_groups(num_groups),
|
||||
num_channels(num_channels),
|
||||
eps(eps) {}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// LTX video latent layout is [W, H, T, C]. ggml_group_norm treats ne[2]
|
||||
// as channels, so fold only H/T internally and restore the same layout.
|
||||
GGML_ASSERT(x->ne[3] == num_channels);
|
||||
const int64_t W = x->ne[0];
|
||||
const int64_t H = x->ne[1];
|
||||
const int64_t T = x->ne[2];
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, x);
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, num_channels, 1);
|
||||
x = ggml_group_norm(ctx->ggml_ctx, x, num_groups, eps);
|
||||
|
||||
ggml_tensor* weight = params["weight"];
|
||||
ggml_tensor* bias = params["bias"];
|
||||
if (ctx->weight_adapter) {
|
||||
weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
|
||||
bias = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias");
|
||||
}
|
||||
weight = ggml_reshape_4d(ctx->ggml_ctx, weight, 1, 1, num_channels, 1);
|
||||
bias = ggml_reshape_4d(ctx->ggml_ctx, bias, 1, 1, num_channels, 1);
|
||||
x = ggml_mul_inplace(ctx->ggml_ctx, x, weight);
|
||||
x = ggml_add_inplace(ctx->ggml_ctx, x, bias);
|
||||
return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, num_channels);
|
||||
}
|
||||
};
|
||||
|
||||
class ResBlock : public GGMLBlock {
|
||||
public:
|
||||
ResBlock(int64_t channels, int dims = 3) {
|
||||
GGML_ASSERT(dims == 3);
|
||||
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv3d(channels, channels, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
|
||||
blocks["norm1"] = std::shared_ptr<GGMLBlock>(new VideoGroupNorm(32, channels));
|
||||
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv3d(channels, channels, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
|
||||
blocks["norm2"] = std::shared_ptr<GGMLBlock>(new VideoGroupNorm(32, channels));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto conv1 = std::dynamic_pointer_cast<Conv3d>(blocks["conv1"]);
|
||||
auto norm1 = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["norm1"]);
|
||||
auto conv2 = std::dynamic_pointer_cast<Conv3d>(blocks["conv2"]);
|
||||
auto norm2 = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["norm2"]);
|
||||
|
||||
ggml_tensor* residual = x;
|
||||
|
||||
x = conv1->forward(ctx, x);
|
||||
x = norm1->forward(ctx, x);
|
||||
x = ggml_silu_inplace(ctx->ggml_ctx, x);
|
||||
x = conv2->forward(ctx, x);
|
||||
x = norm2->forward(ctx, x);
|
||||
x = ggml_add(ctx->ggml_ctx, x, residual);
|
||||
return ggml_silu(ctx->ggml_ctx, x);
|
||||
}
|
||||
};
|
||||
|
||||
class PixelShuffleND : public UnaryBlock {
|
||||
protected:
|
||||
int upscale_factor;
|
||||
|
||||
public:
|
||||
explicit PixelShuffleND(int upscale_factor)
|
||||
: upscale_factor(upscale_factor) {}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
GGML_ASSERT(upscale_factor > 0);
|
||||
int64_t h = x->ne[1];
|
||||
int64_t w = x->ne[0];
|
||||
GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0);
|
||||
// x: [b*f, c*p1*p2, h, w] -> [b*f, c, h*p1, w*p2]
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); // [b*f, h, w, c*p1*p2]
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]); // [b*f, h*w, c*p1*p2]
|
||||
return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true);
|
||||
}
|
||||
};
|
||||
|
||||
class TemporalPixelShuffleND : public UnaryBlock {
|
||||
protected:
|
||||
int upscale_factor;
|
||||
|
||||
public:
|
||||
explicit TemporalPixelShuffleND(int upscale_factor)
|
||||
: upscale_factor(upscale_factor) {}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
GGML_ASSERT(upscale_factor > 0);
|
||||
GGML_ASSERT(x->ne[3] % upscale_factor == 0);
|
||||
const int64_t W = x->ne[0];
|
||||
const int64_t H = x->ne[1];
|
||||
const int64_t F = x->ne[2];
|
||||
const int64_t C = x->ne[3] / upscale_factor;
|
||||
|
||||
// x: [b, c*p, f, h, w] -> [b, c, f*p, h, w]
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, x);
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, F, upscale_factor, C);
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
|
||||
return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, F * upscale_factor, C);
|
||||
}
|
||||
};
|
||||
|
||||
class BlurDownsample : public GGMLBlock {
|
||||
protected:
|
||||
int64_t channels;
|
||||
int stride;
|
||||
ggml_tensor* kernel = nullptr;
|
||||
std::vector<float> kernel_data;
|
||||
|
||||
void init_params(ggml_context* ctx,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "") override {
|
||||
SD_UNUSED(tensor_storage_map);
|
||||
if (stride == 1) {
|
||||
return;
|
||||
}
|
||||
kernel = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 5, 5, 1, channels);
|
||||
std::string name = prefix + "kernel";
|
||||
ggml_set_name(kernel, name.c_str());
|
||||
|
||||
static const float binomial[5] = {1.f, 4.f, 6.f, 4.f, 1.f};
|
||||
kernel_data.resize(static_cast<size_t>(5 * 5 * channels));
|
||||
for (int64_t c = 0; c < channels; ++c) {
|
||||
for (int y = 0; y < 5; ++y) {
|
||||
for (int x = 0; x < 5; ++x) {
|
||||
kernel_data[static_cast<size_t>(x + 5 * (y + 5 * c))] =
|
||||
binomial[y] * binomial[x] / 256.f;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
BlurDownsample(int64_t channels, int stride)
|
||||
: channels(channels),
|
||||
stride(stride) {
|
||||
GGML_ASSERT(stride >= 1);
|
||||
}
|
||||
|
||||
void load_fixed_tensors() {
|
||||
if (kernel == nullptr || kernel_data.empty()) {
|
||||
return;
|
||||
}
|
||||
ggml_backend_tensor_set(kernel, kernel_data.data(), 0, kernel_data.size() * sizeof(float));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
if (stride == 1) {
|
||||
return x;
|
||||
}
|
||||
GGML_ASSERT(kernel != nullptr);
|
||||
GGML_ASSERT(x->ne[2] == channels);
|
||||
if (ctx->conv2d_direct_enabled) {
|
||||
return ggml_conv_2d_dw_direct(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1);
|
||||
}
|
||||
return ggml_conv_2d_dw(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1);
|
||||
}
|
||||
};
|
||||
|
||||
class SpatialRationalResampler : public GGMLBlock {
|
||||
protected:
|
||||
int64_t mid_channels;
|
||||
int num;
|
||||
int den;
|
||||
|
||||
public:
|
||||
SpatialRationalResampler(int64_t mid_channels, int num, int den)
|
||||
: mid_channels(mid_channels),
|
||||
num(num),
|
||||
den(den) {
|
||||
GGML_ASSERT(num >= 1);
|
||||
GGML_ASSERT(den >= 1);
|
||||
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(mid_channels, num * num * mid_channels, {3, 3}, {1, 1}, {1, 1}));
|
||||
blocks["pixel_shuffle"] = std::shared_ptr<GGMLBlock>(new PixelShuffleND(num));
|
||||
blocks["blur_down"] = std::shared_ptr<GGMLBlock>(new BlurDownsample(mid_channels, den));
|
||||
}
|
||||
|
||||
void load_fixed_tensors() {
|
||||
auto blur_down = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
|
||||
blur_down->load_fixed_tensors();
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
|
||||
auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["pixel_shuffle"]);
|
||||
auto blur_down = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
|
||||
|
||||
// rearrange(x, "b c f h w -> (b f) c h w")
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
|
||||
x = conv->forward(ctx, x);
|
||||
x = pixel_shuffle->forward(ctx, x);
|
||||
x = blur_down->forward(ctx, x);
|
||||
return ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
|
||||
}
|
||||
};
|
||||
|
||||
class LatentUpsampler : public GGMLBlock {
|
||||
public:
|
||||
LatentUpsamplerConfig config;
|
||||
|
||||
explicit LatentUpsampler(LatentUpsamplerConfig config)
|
||||
: config(std::move(config)) {
|
||||
GGML_ASSERT(this->config.dims == 3);
|
||||
GGML_ASSERT(this->config.spatial_upsample || this->config.temporal_upsample);
|
||||
|
||||
blocks["initial_conv"] = std::shared_ptr<GGMLBlock>(new Conv3d(this->config.in_channels,
|
||||
this->config.mid_channels,
|
||||
{3, 3, 3},
|
||||
{1, 1, 1},
|
||||
{1, 1, 1}));
|
||||
blocks["initial_norm"] = std::shared_ptr<GGMLBlock>(new VideoGroupNorm(32, this->config.mid_channels));
|
||||
for (int i = 0; i < this->config.num_blocks_per_stage; ++i) {
|
||||
blocks["res_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResBlock(this->config.mid_channels, this->config.dims));
|
||||
}
|
||||
if (this->config.rational_resampler) {
|
||||
blocks["upsampler"] = std::shared_ptr<GGMLBlock>(new SpatialRationalResampler(this->config.mid_channels,
|
||||
this->config.spatial_up_num,
|
||||
this->config.spatial_down_den));
|
||||
} else if (this->config.temporal_upsample) {
|
||||
blocks["upsampler.0"] = std::shared_ptr<GGMLBlock>(new Conv3d(this->config.mid_channels,
|
||||
this->config.temporal_up_factor * this->config.mid_channels,
|
||||
{3, 3, 3},
|
||||
{1, 1, 1},
|
||||
{1, 1, 1}));
|
||||
blocks["upsampler.1"] = std::shared_ptr<GGMLBlock>(new TemporalPixelShuffleND(this->config.temporal_up_factor));
|
||||
} else {
|
||||
blocks["upsampler.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.mid_channels,
|
||||
4 * this->config.mid_channels,
|
||||
{3, 3},
|
||||
{1, 1},
|
||||
{1, 1}));
|
||||
blocks["upsampler.1"] = std::shared_ptr<GGMLBlock>(new PixelShuffleND(2));
|
||||
}
|
||||
for (int i = 0; i < this->config.num_blocks_per_stage; ++i) {
|
||||
blocks["post_upsample_res_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResBlock(this->config.mid_channels, this->config.dims));
|
||||
}
|
||||
blocks["final_conv"] = std::shared_ptr<GGMLBlock>(new Conv3d(this->config.mid_channels,
|
||||
this->config.in_channels,
|
||||
{3, 3, 3},
|
||||
{1, 1, 1},
|
||||
{1, 1, 1}));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [b, c, f, h, w]
|
||||
// return: [b, c, scaled_f, scaled_h, scaled_w]
|
||||
auto initial_conv = std::dynamic_pointer_cast<Conv3d>(blocks["initial_conv"]);
|
||||
auto initial_norm = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["initial_norm"]);
|
||||
auto final_conv = std::dynamic_pointer_cast<Conv3d>(blocks["final_conv"]);
|
||||
|
||||
x = initial_conv->forward(ctx, x);
|
||||
x = initial_norm->forward(ctx, x);
|
||||
x = ggml_silu(ctx->ggml_ctx, x);
|
||||
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.initial", "x");
|
||||
|
||||
for (int i = 0; i < config.num_blocks_per_stage; ++i) {
|
||||
auto block = std::dynamic_pointer_cast<ResBlock>(blocks["res_blocks." + std::to_string(i)]);
|
||||
x = block->forward(ctx, x);
|
||||
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.res_blocks." + std::to_string(i), "x");
|
||||
}
|
||||
|
||||
if (config.rational_resampler) {
|
||||
auto upsampler = std::dynamic_pointer_cast<SpatialRationalResampler>(blocks["upsampler"]);
|
||||
x = upsampler->forward(ctx, x);
|
||||
} else if (config.temporal_upsample) {
|
||||
auto upsample_conv = std::dynamic_pointer_cast<Conv3d>(blocks["upsampler.0"]);
|
||||
auto pixel_shuffle = std::dynamic_pointer_cast<TemporalPixelShuffleND>(blocks["upsampler.1"]);
|
||||
x = upsample_conv->forward(ctx, x); // [b, c*2, f, h, w]
|
||||
x = pixel_shuffle->forward(ctx, x); // [b, c, f*2, h, w]
|
||||
x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]); // x[:, :, 1:, :, :]
|
||||
} else {
|
||||
auto upsample_conv = std::dynamic_pointer_cast<Conv2d>(blocks["upsampler.0"]);
|
||||
auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["upsampler.1"]);
|
||||
|
||||
// rearrange(x, "b c f h w -> (b f) c h w"),
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // [b*f, c, h, w]
|
||||
x = upsample_conv->forward(ctx, x); // [b*f, c*4, h, w]
|
||||
x = pixel_shuffle->forward(ctx, x); // [b*f, c, h*2, w*2]
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // [b*c, f, h, w]
|
||||
}
|
||||
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.spatial_up", "x");
|
||||
|
||||
for (int i = 0; i < config.num_blocks_per_stage; ++i) {
|
||||
auto block = std::dynamic_pointer_cast<ResBlock>(blocks["post_upsample_res_blocks." + std::to_string(i)]);
|
||||
x = block->forward(ctx, x);
|
||||
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.post_blocks." + std::to_string(i), "x");
|
||||
}
|
||||
|
||||
x = final_conv->forward(ctx, x);
|
||||
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.final", "x");
|
||||
return x;
|
||||
}
|
||||
|
||||
void load_fixed_tensors() {
|
||||
if (!config.rational_resampler) {
|
||||
return;
|
||||
}
|
||||
auto upsampler = std::dynamic_pointer_cast<SpatialRationalResampler>(blocks["upsampler"]);
|
||||
upsampler->load_fixed_tensors();
|
||||
}
|
||||
};
|
||||
|
||||
struct LatentUpsamplerRunner : public GGMLRunner {
|
||||
std::unique_ptr<LatentUpsampler> model;
|
||||
|
||||
LatentUpsamplerRunner(ggml_backend_t backend,
|
||||
ggml_backend_t params_backend)
|
||||
: GGMLRunner(backend, params_backend) {}
|
||||
|
||||
std::string get_desc() override {
|
||||
return "ltx_latent_upsampler";
|
||||
}
|
||||
|
||||
bool load_from_file(const std::string& file_path, int n_threads) {
|
||||
LOG_INFO("loading LTX latent upsampler from '%s'", file_path.c_str());
|
||||
ModelLoader model_loader;
|
||||
if (!model_loader.init_from_file(file_path)) {
|
||||
LOG_ERROR("init LTX latent upsampler model loader from file failed: '%s'", file_path.c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
const auto& tensor_storage_map = model_loader.get_tensor_storage_map();
|
||||
bool has_regular_upsampler = has_tensor(tensor_storage_map, "upsampler.0.weight");
|
||||
bool has_rational_spatial = has_tensor(tensor_storage_map, "upsampler.conv.weight");
|
||||
if (!has_tensor(tensor_storage_map, "post_upsample_res_blocks.0.conv2.bias") ||
|
||||
(!has_regular_upsampler && !has_rational_spatial)) {
|
||||
LOG_ERROR("unsupported LTX latent upsampler weights: expected upsampler tensors");
|
||||
return false;
|
||||
}
|
||||
|
||||
LatentUpsamplerConfig config = detect_config_from_weights(tensor_storage_map);
|
||||
if (config.dims != 3 || (!config.spatial_upsample && !config.temporal_upsample) ||
|
||||
config.spatial_up_num < 1 || config.spatial_down_den < 1 || config.temporal_up_factor < 1) {
|
||||
LOG_ERROR("unsupported LTX latent upsampler config: dims=%d spatial=%d temporal=%d rational=%d scale=%.3f temporal_factor=%d",
|
||||
config.dims,
|
||||
config.spatial_upsample,
|
||||
config.temporal_upsample,
|
||||
config.rational_resampler,
|
||||
config.spatial_scale,
|
||||
config.temporal_up_factor);
|
||||
return false;
|
||||
}
|
||||
|
||||
model = std::make_unique<LatentUpsampler>(config);
|
||||
model->init(params_ctx, tensor_storage_map, "");
|
||||
if (!alloc_params_buffer()) {
|
||||
LOG_ERROR("LTX latent upsampler params buffer allocation failed");
|
||||
return false;
|
||||
}
|
||||
|
||||
std::map<std::string, ggml_tensor*> tensors;
|
||||
model->get_param_tensors(tensors);
|
||||
std::set<std::string> ignore_tensors;
|
||||
if (config.rational_resampler) {
|
||||
ignore_tensors.insert("upsampler.blur_down.kernel");
|
||||
}
|
||||
if (!model_loader.load_tensors(tensors, ignore_tensors, n_threads)) {
|
||||
LOG_ERROR("load LTX latent upsampler tensors failed");
|
||||
return false;
|
||||
}
|
||||
model->load_fixed_tensors();
|
||||
|
||||
LOG_INFO("LTX latent upsampler loaded: in_channels=%" PRId64 ", mid_channels=%" PRId64 ", blocks=%d, scale=%.3f, temporal_factor=%d, rational=%d",
|
||||
config.in_channels,
|
||||
config.mid_channels,
|
||||
config.num_blocks_per_stage,
|
||||
config.spatial_scale,
|
||||
config.temporal_up_factor,
|
||||
config.rational_resampler);
|
||||
return true;
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
|
||||
if (!model) {
|
||||
return nullptr;
|
||||
}
|
||||
ggml_cgraph* gf = new_graph_custom(LTX_UPSAMPLER_GRAPH_SIZE);
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* out = model->forward(&runner_ctx, x);
|
||||
ggml_build_forward_expand(gf, out);
|
||||
return gf;
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(const int n_threads,
|
||||
const sd::Tensor<float>& x) {
|
||||
if (!model) {
|
||||
LOG_ERROR("LTX latent upsampler is not loaded");
|
||||
return {};
|
||||
}
|
||||
if (x.dim() != 4 && x.dim() != 5) {
|
||||
LOG_ERROR("LTX latent upsampler expects 4D or 5D video latent, got dim=%lld",
|
||||
(long long)x.dim());
|
||||
return {};
|
||||
}
|
||||
if (x.dim() == 5 && x.shape()[4] != 1) {
|
||||
LOG_ERROR("LTX latent upsampler currently supports batch size 1, got batch=%lld",
|
||||
(long long)x.shape()[4]);
|
||||
return {};
|
||||
}
|
||||
if (x.shape()[3] != model->config.in_channels) {
|
||||
LOG_ERROR("LTX latent upsampler expected %" PRId64 " channels, got %lld",
|
||||
model->config.in_channels,
|
||||
(long long)x.shape()[3]);
|
||||
return {};
|
||||
}
|
||||
size_t expected_dim = static_cast<size_t>(x.dim());
|
||||
auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace LTXVUpsampler
|
||||
|
||||
#endif // __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
|
||||
Reference in new issue
Block a user