mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-09-25 20:50:39 +00:00
refactor: align PixArt weights with upstream layout (#2061)
This commit is contained in:
leejet
GitHub
parent
19bbbca1c7
commit
2f886889e6
@ -41,28 +41,28 @@ namespace PixArt {
|
||||
auto it = weights.find(prefix + "." + suffix);
|
||||
return it == weights.end() ? nullptr : &it->second;
|
||||
};
|
||||
if (auto w = find("pos_embed.proj.weight")) {
|
||||
if (auto w = find("x_embedder.proj.weight")) {
|
||||
config.hidden_size = w->ne[3];
|
||||
config.in_channels = w->ne[2];
|
||||
config.patch_size = w->ne[0];
|
||||
}
|
||||
if (auto w = find("proj_out.weight")) {
|
||||
if (auto w = find("final_layer.linear.weight")) {
|
||||
config.out_channels = w->ne[1] / (config.patch_size * config.patch_size);
|
||||
}
|
||||
if (auto w = find("caption_projection.linear_1.weight")) {
|
||||
if (auto w = find("y_embedder.y_proj.fc1.weight")) {
|
||||
config.caption_channels = w->ne[0];
|
||||
}
|
||||
if (auto w = find("transformer_blocks.0.attn2.to_k.weight")) {
|
||||
if (auto w = find("blocks.0.cross_attn.kv_linear.weight")) {
|
||||
config.cross_attention_dim = w->ne[0];
|
||||
}
|
||||
if (auto w = find("transformer_blocks.0.ff.net.0.proj.weight")) {
|
||||
if (auto w = find("blocks.0.mlp.fc1.weight")) {
|
||||
config.ffn_dim = w->ne[1];
|
||||
}
|
||||
if (find("adaln_single.emb.resolution_embedder.linear_1.weight") != nullptr) {
|
||||
if (find("csize_embedder.mlp.0.weight") != nullptr) {
|
||||
LOG_WARN("pixart: resolution/aspect-ratio micro conditions are not supported; output may differ from the reference");
|
||||
}
|
||||
int layers = 0;
|
||||
const std::string block_prefix = prefix + ".transformer_blocks.";
|
||||
const std::string block_prefix = prefix + ".blocks.";
|
||||
for (const auto& [name, _] : weights) {
|
||||
if (starts_with(name, block_prefix)) {
|
||||
layers = std::max(layers, atoi(name.substr(block_prefix.size()).c_str()) + 1);
|
||||
@ -108,37 +108,46 @@ namespace PixArt {
|
||||
class PixArtTimestepEmbedding : public GGMLBlock {
|
||||
public:
|
||||
PixArtTimestepEmbedding(int64_t in_channels, int64_t out_dim) {
|
||||
blocks["linear_1"] = std::make_shared<Linear>(in_channels, out_dim);
|
||||
blocks["linear_2"] = std::make_shared<Linear>(out_dim, out_dim);
|
||||
blocks["mlp.0"] = std::make_shared<Linear>(in_channels, out_dim);
|
||||
blocks["mlp.2"] = std::make_shared<Linear>(out_dim, out_dim);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
x = std::dynamic_pointer_cast<Linear>(blocks["linear_1"])->forward(ctx, x);
|
||||
x = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"])->forward(ctx, x);
|
||||
x = ggml_silu(ctx->ggml_ctx, x);
|
||||
return std::dynamic_pointer_cast<Linear>(blocks["linear_2"])->forward(ctx, x);
|
||||
return std::dynamic_pointer_cast<Linear>(blocks["mlp.2"])->forward(ctx, x);
|
||||
}
|
||||
};
|
||||
|
||||
class PixArtAttention : public GGMLBlock {
|
||||
int64_t num_heads;
|
||||
bool self_attention;
|
||||
|
||||
public:
|
||||
PixArtAttention(int64_t dim, int64_t num_heads, int64_t context_dim)
|
||||
: num_heads(num_heads) {
|
||||
blocks["to_q"] = std::make_shared<Linear>(dim, dim);
|
||||
blocks["to_k"] = std::make_shared<Linear>(context_dim, dim);
|
||||
blocks["to_v"] = std::make_shared<Linear>(context_dim, dim);
|
||||
blocks["to_out.0"] = std::make_shared<Linear>(dim, dim);
|
||||
PixArtAttention(int64_t dim, int64_t num_heads, int64_t context_dim, bool self_attention)
|
||||
: num_heads(num_heads), self_attention(self_attention) {
|
||||
if (self_attention) {
|
||||
blocks["qkv"] = std::make_shared<Linear>(dim, 3 * dim);
|
||||
} else {
|
||||
blocks["q_linear"] = std::make_shared<Linear>(dim, dim);
|
||||
blocks["kv_linear"] = std::make_shared<Linear>(context_dim, 2 * dim);
|
||||
}
|
||||
blocks["proj"] = std::make_shared<Linear>(dim, dim);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* context, ggml_tensor* mask = nullptr) {
|
||||
// x: [N, n_token, dim], context: [N, n_context, context_dim]
|
||||
auto q = std::dynamic_pointer_cast<Linear>(blocks["to_q"])->forward(ctx, x);
|
||||
auto k = std::dynamic_pointer_cast<Linear>(blocks["to_k"])->forward(ctx, context);
|
||||
auto v = std::dynamic_pointer_cast<Linear>(blocks["to_v"])->forward(ctx, context);
|
||||
|
||||
auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled);
|
||||
return std::dynamic_pointer_cast<Linear>(blocks["to_out.0"])->forward(ctx, out);
|
||||
std::vector<ggml_tensor*> qkv;
|
||||
if (self_attention) {
|
||||
auto projected = std::dynamic_pointer_cast<Linear>(blocks["qkv"])->forward(ctx, x);
|
||||
qkv = ggml_ext_chunk(ctx->ggml_ctx, projected, 3, 0);
|
||||
} else {
|
||||
auto q = std::dynamic_pointer_cast<Linear>(blocks["q_linear"])->forward(ctx, x);
|
||||
auto kv = std::dynamic_pointer_cast<Linear>(blocks["kv_linear"])->forward(ctx, context);
|
||||
auto parts = ggml_ext_chunk(ctx->ggml_ctx, kv, 2, 0);
|
||||
qkv = {q, parts[0], parts[1]};
|
||||
}
|
||||
auto out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, mask, false, ctx->flash_attn_enabled);
|
||||
return std::dynamic_pointer_cast<Linear>(blocks["proj"])->forward(ctx, out);
|
||||
}
|
||||
};
|
||||
|
||||
@ -155,10 +164,10 @@ namespace PixArt {
|
||||
public:
|
||||
PixArtBlock(int64_t dim, int64_t num_heads, int64_t context_dim, int64_t ffn_dim)
|
||||
: dim(dim) {
|
||||
blocks["attn1"] = std::make_shared<PixArtAttention>(dim, num_heads, dim);
|
||||
blocks["attn2"] = std::make_shared<PixArtAttention>(dim, num_heads, context_dim);
|
||||
blocks["ff.net.0.proj"] = std::make_shared<Linear>(dim, ffn_dim);
|
||||
blocks["ff.net.2"] = std::make_shared<Linear>(ffn_dim, dim);
|
||||
blocks["attn"] = std::make_shared<PixArtAttention>(dim, num_heads, dim, true);
|
||||
blocks["cross_attn"] = std::make_shared<PixArtAttention>(dim, num_heads, context_dim, false);
|
||||
blocks["mlp.fc1"] = std::make_shared<Linear>(dim, ffn_dim);
|
||||
blocks["mlp.fc2"] = std::make_shared<Linear>(ffn_dim, dim);
|
||||
}
|
||||
|
||||
static ggml_tensor* norm(ggml_context* ctx, ggml_tensor* x) {
|
||||
@ -178,10 +187,10 @@ namespace PixArt {
|
||||
auto m = ggml_add(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, mod, dim, 6, N), table);
|
||||
auto mv = ggml_ext_chunk(ctx->ggml_ctx, ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, m), dim * 6, N), 6, 0);
|
||||
|
||||
auto attn1 = std::dynamic_pointer_cast<PixArtAttention>(blocks["attn1"]);
|
||||
auto attn2 = std::dynamic_pointer_cast<PixArtAttention>(blocks["attn2"]);
|
||||
auto proj = std::dynamic_pointer_cast<Linear>(blocks["ff.net.0.proj"]);
|
||||
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["ff.net.2"]);
|
||||
auto attn1 = std::dynamic_pointer_cast<PixArtAttention>(blocks["attn"]);
|
||||
auto attn2 = std::dynamic_pointer_cast<PixArtAttention>(blocks["cross_attn"]);
|
||||
auto proj = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
|
||||
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"]);
|
||||
|
||||
auto gate = [&](ggml_tensor* y, ggml_tensor* g) {
|
||||
g = ggml_reshape_3d(ctx->ggml_ctx, g, dim, 1, N);
|
||||
@ -206,28 +215,28 @@ namespace PixArt {
|
||||
void init_params(ggml_context* ctx,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "") override {
|
||||
ggml_type wtype = get_type(prefix + "scale_shift_table", tensor_storage_map, GGML_TYPE_F32);
|
||||
params["scale_shift_table"] = ggml_new_tensor_2d(ctx, wtype, config.hidden_size, 2);
|
||||
ggml_type wtype = get_type(prefix + "final_layer.scale_shift_table", tensor_storage_map, GGML_TYPE_F32);
|
||||
params["final_layer.scale_shift_table"] = ggml_new_tensor_2d(ctx, wtype, config.hidden_size, 2);
|
||||
}
|
||||
|
||||
public:
|
||||
PixArtModel() = default;
|
||||
PixArtModel(const PixArtConfig& config)
|
||||
: config(config) {
|
||||
blocks["pos_embed.proj"] = std::make_shared<Conv2d>(config.in_channels,
|
||||
blocks["x_embedder.proj"] = std::make_shared<Conv2d>(config.in_channels,
|
||||
config.hidden_size,
|
||||
std::pair<int, int>{static_cast<int>(config.patch_size), static_cast<int>(config.patch_size)},
|
||||
std::pair<int, int>{static_cast<int>(config.patch_size), static_cast<int>(config.patch_size)});
|
||||
blocks["adaln_single.emb.timestep_embedder"] = std::make_shared<PixArtTimestepEmbedding>(ADALN_EMBED_DIM, config.hidden_size);
|
||||
blocks["adaln_single.linear"] = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size);
|
||||
blocks["caption_projection.linear_1"] = std::make_shared<Linear>(config.caption_channels, config.hidden_size);
|
||||
blocks["caption_projection.linear_2"] = std::make_shared<Linear>(config.hidden_size, config.cross_attention_dim);
|
||||
blocks["t_embedder"] = std::make_shared<PixArtTimestepEmbedding>(ADALN_EMBED_DIM, config.hidden_size);
|
||||
blocks["t_block.1"] = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size);
|
||||
blocks["y_embedder.y_proj.fc1"] = std::make_shared<Linear>(config.caption_channels, config.hidden_size);
|
||||
blocks["y_embedder.y_proj.fc2"] = std::make_shared<Linear>(config.hidden_size, config.cross_attention_dim);
|
||||
for (int i = 0; i < config.num_layers; ++i) {
|
||||
blocks["transformer_blocks." + std::to_string(i)] =
|
||||
blocks["blocks." + std::to_string(i)] =
|
||||
std::make_shared<PixArtBlock>(config.hidden_size, config.num_heads, config.cross_attention_dim, config.ffn_dim);
|
||||
}
|
||||
blocks["norm_out"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
|
||||
blocks["proj_out"] = std::make_shared<Linear>(config.hidden_size,
|
||||
blocks["final_layer.norm_final"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
|
||||
blocks["final_layer.linear"] = std::make_shared<Linear>(config.hidden_size,
|
||||
config.patch_size * config.patch_size * config.out_channels);
|
||||
}
|
||||
|
||||
@ -245,29 +254,29 @@ namespace PixArt {
|
||||
int64_t wp = W / p;
|
||||
int64_t hp = H / p;
|
||||
|
||||
auto h = std::dynamic_pointer_cast<Conv2d>(blocks["pos_embed.proj"])->forward(ctx, x); // [N, hidden, hp, wp]
|
||||
auto h = std::dynamic_pointer_cast<Conv2d>(blocks["x_embedder.proj"])->forward(ctx, x); // [N, hidden, hp, wp]
|
||||
h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 1, 2, 0, 3)); // [N, hp, wp, hidden] -> [N, hp*wp, hidden]
|
||||
h = ggml_reshape_3d(ctx->ggml_ctx, h, config.hidden_size, wp * hp, N); // [N, hp*wp, hidden]
|
||||
h = ggml_add(ctx->ggml_ctx, h, pos_embed);
|
||||
|
||||
auto t = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, ADALN_EMBED_DIM, 10000);
|
||||
auto emb = std::dynamic_pointer_cast<PixArtTimestepEmbedding>(blocks["adaln_single.emb.timestep_embedder"])->forward(ctx, t);
|
||||
auto emb = std::dynamic_pointer_cast<PixArtTimestepEmbedding>(blocks["t_embedder"])->forward(ctx, t);
|
||||
|
||||
auto mod = std::dynamic_pointer_cast<Linear>(blocks["adaln_single.linear"])
|
||||
auto mod = std::dynamic_pointer_cast<Linear>(blocks["t_block.1"])
|
||||
->forward(ctx, ggml_silu(ctx->ggml_ctx, emb)); // [N, 6 * hidden]
|
||||
|
||||
auto ctx_emb = std::dynamic_pointer_cast<Linear>(blocks["caption_projection.linear_1"])->forward(ctx, context);
|
||||
auto ctx_emb = std::dynamic_pointer_cast<Linear>(blocks["y_embedder.y_proj.fc1"])->forward(ctx, context);
|
||||
ctx_emb = ggml_ext_gelu(ctx->ggml_ctx, ctx_emb, true);
|
||||
ctx_emb = std::dynamic_pointer_cast<Linear>(blocks["caption_projection.linear_2"])->forward(ctx, ctx_emb);
|
||||
ctx_emb = std::dynamic_pointer_cast<Linear>(blocks["y_embedder.y_proj.fc2"])->forward(ctx, ctx_emb);
|
||||
|
||||
for (int i = 0; i < config.num_layers; ++i) {
|
||||
auto block = std::dynamic_pointer_cast<PixArtBlock>(blocks["transformer_blocks." + std::to_string(i)]);
|
||||
auto block = std::dynamic_pointer_cast<PixArtBlock>(blocks["blocks." + std::to_string(i)]);
|
||||
h = block->forward(ctx, h, mod, ctx_emb, context_mask);
|
||||
sd::ggml_graph_cut::mark_graph_cut(h, "pixart.transformer_blocks." + std::to_string(i), "h");
|
||||
sd::ggml_graph_cut::mark_graph_cut(h, "pixart.blocks." + std::to_string(i), "h");
|
||||
}
|
||||
|
||||
// scale_shift_table + emb -> (shift, scale) for the affine-free final norm
|
||||
auto tail_table = params["scale_shift_table"];
|
||||
auto tail_table = params["final_layer.scale_shift_table"];
|
||||
if (tail_table->type != GGML_TYPE_F32) {
|
||||
tail_table = ggml_cast(ctx->ggml_ctx, tail_table, GGML_TYPE_F32);
|
||||
}
|
||||
@ -277,9 +286,9 @@ namespace PixArt {
|
||||
auto parts = ggml_ext_chunk(ctx->ggml_ctx,
|
||||
ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, ss), config.hidden_size * 2, N),
|
||||
2, 0);
|
||||
h = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"])->forward(ctx, h);
|
||||
h = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer.norm_final"])->forward(ctx, h);
|
||||
h = modulate(ctx->ggml_ctx, h, parts[0], parts[1]);
|
||||
h = std::dynamic_pointer_cast<Linear>(blocks["proj_out"])->forward(ctx, h); // [N, hp*wp, p*p*out_ch]
|
||||
h = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, h); // [N, hp*wp, p*p*out_ch]
|
||||
h = DiT::unpatchify(ctx->ggml_ctx, h, hp, wp, static_cast<int>(p), static_cast<int>(p), false);
|
||||
return h; // [N, out_channels, H, W]
|
||||
}
|
||||
|
||||
@ -532,9 +532,16 @@ SDVersion ModelLoader::get_sd_version() const {
|
||||
if (tensor_storage.name.find("model.diffusion_model.layers.0.adaLN_sa_ln.weight") != std::string::npos) {
|
||||
return VERSION_ERNIE_IMAGE;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.t_block.1.weight") != std::string::npos &&
|
||||
tensor_storage_map.find("model.diffusion_model.x_embedder.proj.weight") != tensor_storage_map.end() &&
|
||||
tensor_storage_map.find("model.diffusion_model.audio_patchify_proj.weight") == tensor_storage_map.end()) {
|
||||
return VERSION_PIXART;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.adaln_single.emb.timestep_embedder.linear_1.bias") != std::string::npos) {
|
||||
// PixArt shares this signature with LTX-AV; pos_embed.proj is PixArt-only.
|
||||
if (tensor_storage_map.find("model.diffusion_model.pos_embed.proj.weight") != tensor_storage_map.end()) {
|
||||
// PixArt shares this timestep embedding with LTX-AV.
|
||||
if (tensor_storage_map.find("model.diffusion_model.pos_embed.proj.weight") != tensor_storage_map.end() &&
|
||||
tensor_storage_map.find("model.diffusion_model.adaln_single.linear.weight") != tensor_storage_map.end() &&
|
||||
tensor_storage_map.find("model.diffusion_model.audio_patchify_proj.weight") == tensor_storage_map.end()) {
|
||||
return VERSION_PIXART;
|
||||
}
|
||||
return VERSION_LTXAV;
|
||||
@ -1094,10 +1101,12 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
|
||||
if (tensors_to_process.empty()) {
|
||||
continue;
|
||||
}
|
||||
if (log_progress) {
|
||||
LOG_VERBOSE("loading %zu/%zu tensors from %s",
|
||||
tensors_to_process.size(),
|
||||
file_tensors.size(),
|
||||
file_path.c_str());
|
||||
}
|
||||
|
||||
bool is_zip = fdata.is_zip;
|
||||
|
||||
|
||||
@ -252,11 +252,7 @@ bool ModelManager::register_param_tensors(ModelComponent component,
|
||||
state->component = component;
|
||||
state->source_file = source_file;
|
||||
state->source_version = source_version;
|
||||
auto source = sources.find(name);
|
||||
if (source != sources.end()) {
|
||||
state->source = source->second;
|
||||
state->has_source = true;
|
||||
}
|
||||
state->sources = find_tensor_sources(*state, sources);
|
||||
state->residency_mode = residency_mode;
|
||||
state->compute_backend = compute_backend;
|
||||
state->params_backend = params_backend;
|
||||
@ -757,12 +753,33 @@ bool ModelManager::validate_tensor(const TensorState& state) const {
|
||||
return true;
|
||||
}
|
||||
|
||||
if (!state.has_source) {
|
||||
if (state.sources.empty()) {
|
||||
LOG_ERROR("%s tensor '%s' not in model metadata", model_component_name(state.component), state.name.c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
const TensorStorage& tensor_storage = state.source;
|
||||
TensorStorage tensor_storage = state.sources.front();
|
||||
if (state.sources.size() > 1) {
|
||||
const int dim = tensor_storage.n_dims - 1;
|
||||
if (dim < 0 || dim >= GGML_MAX_DIMS) {
|
||||
return false;
|
||||
}
|
||||
tensor_storage.ne[dim] = 0;
|
||||
for (const auto& part : state.sources) {
|
||||
if (part.n_dims != tensor_storage.n_dims || part.ne[dim] < 0 ||
|
||||
part.ne[dim] > state.tensor->ne[dim] - tensor_storage.ne[dim]) {
|
||||
LOG_ERROR("invalid tensor part '%s' for '%s'", part.name.c_str(), state.name.c_str());
|
||||
return false;
|
||||
}
|
||||
for (int i = 0; i < GGML_MAX_DIMS; ++i) {
|
||||
if (i != dim && part.ne[i] != tensor_storage.ne[i]) {
|
||||
LOG_ERROR("incompatible tensor part '%s' for '%s'", part.name.c_str(), state.name.c_str());
|
||||
return false;
|
||||
}
|
||||
}
|
||||
tensor_storage.ne[dim] += part.ne[dim];
|
||||
}
|
||||
}
|
||||
if (state.tensor->ne[0] != tensor_storage.ne[0] ||
|
||||
state.tensor->ne[1] != tensor_storage.ne[1] ||
|
||||
state.tensor->ne[2] != tensor_storage.ne[2] ||
|
||||
@ -831,7 +848,7 @@ bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
|
||||
}
|
||||
|
||||
bool ModelManager::can_mmap_storage(const TensorState& state) const {
|
||||
if (state.source_file != 0 || !enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
|
||||
if (state.sources.size() > 1 || state.source_file != 0 || !enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
|
||||
return false;
|
||||
}
|
||||
if (state.compute_backend == nullptr || state.params_backend == nullptr) {
|
||||
@ -941,6 +958,62 @@ bool ModelManager::alloc_params_buffers(const std::vector<TensorState*>& states,
|
||||
return true;
|
||||
}
|
||||
|
||||
bool ModelManager::load_tensor_parts(TensorState& state) {
|
||||
auto ctx = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
|
||||
ggml_init({state.sources.size() * ggml_tensor_overhead(), nullptr, true}), ggml_free);
|
||||
if (!ctx) {
|
||||
return false;
|
||||
}
|
||||
const size_t size = ggml_nbytes(state.tensor);
|
||||
std::vector<uint8_t> buffer;
|
||||
void* data = state.tensor->data;
|
||||
if (!ggml_backend_buffer_is_host(state.tensor->buffer)) {
|
||||
buffer.resize(size);
|
||||
data = buffer.data();
|
||||
}
|
||||
std::map<std::string, ggml_tensor*> parts;
|
||||
std::set<std::string> names;
|
||||
size_t offset = 0;
|
||||
for (const auto& source : state.sources) {
|
||||
auto part = ggml_new_tensor(ctx.get(), state.tensor->type, source.n_dims, source.ne);
|
||||
const size_t part_size = ggml_nbytes(part);
|
||||
if (part_size > size - offset) {
|
||||
return false;
|
||||
}
|
||||
part->data = static_cast<uint8_t*>(data) + offset;
|
||||
parts[source.name] = part;
|
||||
names.insert(source.name);
|
||||
offset += part_size;
|
||||
}
|
||||
if (offset != size) {
|
||||
return false;
|
||||
}
|
||||
std::set<std::string> loaded;
|
||||
std::mutex mutex;
|
||||
auto callback = [&](const TensorStorage& source, ggml_tensor** dst) {
|
||||
*dst = nullptr;
|
||||
auto part = parts.find(source.name);
|
||||
if (part != parts.end()) {
|
||||
*dst = part->second;
|
||||
std::lock_guard<std::mutex> lock(mutex);
|
||||
loaded.insert(source.name);
|
||||
}
|
||||
return true;
|
||||
};
|
||||
const bool success = state.source_file == 0
|
||||
? model_loader_.load_tensors(callback, enable_mmap_, &names, false)
|
||||
: model_loader_.load_file_tensors(state.source_file, state.source_version, callback, names, enable_mmap_);
|
||||
if (!success || loaded != names) {
|
||||
return false;
|
||||
}
|
||||
if (!buffer.empty()) {
|
||||
// Upload the assembled tensor once, including for row-split backend buffers.
|
||||
ggml_backend_tensor_set(state.tensor, buffer.data(), 0, size);
|
||||
}
|
||||
state.loaded_to_params_backend = true;
|
||||
return true;
|
||||
}
|
||||
|
||||
bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
|
||||
using ReadGroup = std::pair<ModelLoader::FileId, SDVersion>;
|
||||
using ReadBatch = std::map<std::string, std::vector<TensorState*>>;
|
||||
@ -948,6 +1021,12 @@ bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
|
||||
for (auto* state : states) {
|
||||
if (state == nullptr)
|
||||
continue;
|
||||
if (state->sources.size() > 1) {
|
||||
if (!load_tensor_parts(*state)) {
|
||||
return false;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
auto& batches = groups[{state->source_file, state->source_version}];
|
||||
// The loader supplies one destination per name; only conflicting types need another batch.
|
||||
auto batch = std::find_if(batches.begin(), batches.end(), [&](const ReadBatch& candidate) {
|
||||
|
||||
@ -38,8 +38,7 @@ private:
|
||||
std::string name;
|
||||
ggml_tensor* tensor = nullptr;
|
||||
ModelComponent component = ModelComponent::Count;
|
||||
TensorStorage source;
|
||||
bool has_source = false;
|
||||
std::vector<TensorStorage> sources;
|
||||
ModelLoader::FileId source_file = 0;
|
||||
SDVersion source_version = VERSION_COUNT;
|
||||
|
||||
@ -138,6 +137,8 @@ private:
|
||||
bool apply_loras_to_params(const std::vector<TensorState*>& states);
|
||||
bool mmap_params(const std::vector<TensorState*>& states,
|
||||
std::vector<ParamsStorageBlock*>& created_storage_blocks);
|
||||
static std::vector<TensorStorage> find_tensor_sources(const TensorState& state, const String2TensorStorage& sources);
|
||||
bool load_tensor_parts(TensorState& state);
|
||||
bool can_mmap_storage(const TensorState& state) const;
|
||||
bool alloc_params_buffers(const std::vector<TensorState*>& states,
|
||||
std::vector<ParamsStorageBlock*>& created_storage_blocks);
|
||||
|
||||
@ -15,6 +15,26 @@ static bool same_tensor_source(const TensorStorage& a, const TensorStorage& b) {
|
||||
a.int8_convrot_group_size == b.int8_convrot_group_size;
|
||||
}
|
||||
|
||||
std::vector<TensorStorage> ModelManager::find_tensor_sources(const TensorState& state, const String2TensorStorage& sources) {
|
||||
auto first = sources.find(state.name);
|
||||
if (first == sources.end()) {
|
||||
return {};
|
||||
}
|
||||
std::vector<TensorStorage> result{first->second};
|
||||
if (state.component == ModelComponent::LoRA ||
|
||||
std::equal(first->second.ne, first->second.ne + GGML_MAX_DIMS, state.tensor->ne)) {
|
||||
return result;
|
||||
}
|
||||
for (size_t i = 1;; ++i) {
|
||||
auto part = sources.find(state.name + "." + std::to_string(i));
|
||||
if (part == sources.end()) {
|
||||
break;
|
||||
}
|
||||
result.push_back(part->second);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
void ModelManager::invalidate_sources(const std::unordered_set<TensorState*>& states) {
|
||||
auto affected = states;
|
||||
for (const auto& block : params_storage_blocks_) {
|
||||
@ -76,20 +96,16 @@ bool ModelManager::set_loader(ModelLoader loader) {
|
||||
}
|
||||
std::unordered_set<TensorState*> changed;
|
||||
for (const auto& state : tensor_states_) {
|
||||
const auto& sources = sources_for(*state);
|
||||
auto source = sources.find(state->name);
|
||||
const bool found = source != sources.end();
|
||||
if (found != state->has_source || (found && !same_tensor_source(state->source, source->second)) ||
|
||||
const auto sources = find_tensor_sources(*state, sources_for(*state));
|
||||
if (sources.size() != state->sources.size() ||
|
||||
!std::equal(sources.begin(), sources.end(), state->sources.begin(), same_tensor_source) ||
|
||||
(lora_changed && state->component != ModelComponent::LoRA && state->applied_lora_epoch != UINT64_MAX)) {
|
||||
changed.insert(state.get());
|
||||
}
|
||||
}
|
||||
invalidate_sources(changed);
|
||||
for (auto* state : changed) {
|
||||
const auto& sources = sources_for(*state);
|
||||
auto source = sources.find(state->name);
|
||||
state->has_source = source != sources.end();
|
||||
state->source = state->has_source ? source->second : TensorStorage{};
|
||||
state->sources = find_tensor_sources(*state, sources_for(*state));
|
||||
}
|
||||
if (lora_changed) {
|
||||
++current_lora_epoch_;
|
||||
@ -134,9 +150,8 @@ ModelLoader::FileVersions ModelManager::source_versions(const std::set<ModelComp
|
||||
versions[state->source_file] = loader.file_revision(state->source_file);
|
||||
continue;
|
||||
}
|
||||
auto source = sources.find(state->name);
|
||||
if (source != sources.end()) {
|
||||
versions[source->second.file_id] = source->second.file_revision;
|
||||
for (const auto& source : find_tensor_sources(*state, sources)) {
|
||||
versions[source.file_id] = source.file_revision;
|
||||
}
|
||||
}
|
||||
return versions;
|
||||
|
||||
@ -932,6 +932,68 @@ static bool is_diffusers_controlnet_name(const std::string& name) {
|
||||
return false;
|
||||
}
|
||||
|
||||
static std::string convert_diffusers_dit_to_original_pixart(std::string name) {
|
||||
static const std::vector<std::pair<std::string, std::string>> prefix_map = {
|
||||
{"pos_embed.proj.", "x_embedder.proj."},
|
||||
{"adaln_single.emb.timestep_embedder.linear_1.", "t_embedder.mlp.0."},
|
||||
{"adaln_single.emb.timestep_embedder.linear_2.", "t_embedder.mlp.2."},
|
||||
{"adaln_single.emb.resolution_embedder.linear_1.", "csize_embedder.mlp.0."},
|
||||
{"adaln_single.emb.resolution_embedder.linear_2.", "csize_embedder.mlp.2."},
|
||||
{"adaln_single.emb.aspect_ratio_embedder.linear_1.", "ar_embedder.mlp.0."},
|
||||
{"adaln_single.emb.aspect_ratio_embedder.linear_2.", "ar_embedder.mlp.2."},
|
||||
{"adaln_single.linear.", "t_block.1."},
|
||||
{"caption_projection.linear_1.", "y_embedder.y_proj.fc1."},
|
||||
{"caption_projection.linear_2.", "y_embedder.y_proj.fc2."},
|
||||
{"proj_out.", "final_layer.linear."},
|
||||
};
|
||||
for (const auto& entry : prefix_map) {
|
||||
if (starts_with(name, entry.first)) {
|
||||
return entry.second + name.substr(entry.first.size());
|
||||
}
|
||||
}
|
||||
if (name == "scale_shift_table") {
|
||||
return "final_layer.scale_shift_table";
|
||||
}
|
||||
const std::string block_prefix = "transformer_blocks.";
|
||||
if (!starts_with(name, block_prefix)) {
|
||||
return name;
|
||||
}
|
||||
const size_t block_end = name.find('.', block_prefix.size());
|
||||
if (block_end == std::string::npos) {
|
||||
return name;
|
||||
}
|
||||
const std::string prefix = "blocks." + name.substr(block_prefix.size(), block_end - block_prefix.size()) + ".";
|
||||
name = name.substr(block_end + 1);
|
||||
static const std::vector<std::pair<std::string, std::string>> block_map = {
|
||||
{"attn1.to_q.", "attn.qkv."},
|
||||
{"attn1.to_out.0.", "attn.proj."},
|
||||
{"attn2.to_q.", "cross_attn.q_linear."},
|
||||
{"attn2.to_k.", "cross_attn.kv_linear."},
|
||||
{"attn2.to_out.0.", "cross_attn.proj."},
|
||||
{"ff.net.0.proj.", "mlp.fc1."},
|
||||
{"ff.net.2.", "mlp.fc2."},
|
||||
};
|
||||
for (const auto& entry : block_map) {
|
||||
if (starts_with(name, entry.first)) {
|
||||
return prefix + entry.second + name.substr(entry.first.size());
|
||||
}
|
||||
}
|
||||
static const std::vector<std::pair<std::string, std::string>> part_map = {
|
||||
{"attn1.to_k.weight", "attn.qkv.weight.1"},
|
||||
{"attn1.to_k.bias", "attn.qkv.bias.1"},
|
||||
{"attn1.to_v.weight", "attn.qkv.weight.2"},
|
||||
{"attn1.to_v.bias", "attn.qkv.bias.2"},
|
||||
{"attn2.to_v.weight", "cross_attn.kv_linear.weight.1"},
|
||||
{"attn2.to_v.bias", "cross_attn.kv_linear.bias.1"},
|
||||
};
|
||||
for (const auto& entry : part_map) {
|
||||
if (name == entry.first || starts_with(name, entry.first + ".")) {
|
||||
return prefix + entry.second + name.substr(entry.first.size());
|
||||
}
|
||||
}
|
||||
return prefix + name;
|
||||
}
|
||||
|
||||
std::string convert_diffusion_model_name(std::string name, std::string prefix, SDVersion version) {
|
||||
if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
|
||||
name = convert_diffusers_unet_to_original_sd1(name);
|
||||
@ -951,6 +1013,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
|
||||
name = convert_other_dit_to_original_anima(name);
|
||||
} else if (sd_version_is_krea2(version)) {
|
||||
name = convert_diffusers_dit_to_original_krea2(name);
|
||||
} else if (sd_version_is_pixart(version)) {
|
||||
name = convert_diffusers_dit_to_original_pixart(name);
|
||||
}
|
||||
return name;
|
||||
}
|
||||
|
||||
@ -579,7 +579,7 @@ namespace sd::model_builders {
|
||||
weight_manager);
|
||||
if (sd_version_is_pixart(version)) {
|
||||
// Alpha-512 and Sigma share tensor layouts; Alpha-512 needs an explicit scale override.
|
||||
if (tensor_storage_map.count("model.diffusion_model.adaln_single.emb.resolution_embedder.linear_1.weight") != 0) {
|
||||
if (tensor_storage_map.count("model.diffusion_model.csize_embedder.mlp.0.weight") != 0) {
|
||||
model->scale_factor = 0.18215f;
|
||||
}
|
||||
for (const auto& [key, value] : parse_key_value_args(sd_ctx_params->model_args, "model arg")) {
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user