Commit 95b3090 for stable-diffusion.cpp
commit 95b3090664d019fcf8fa032a44d1d4793ba244b8
Author: leejet <leejet714@gmail.com>
Date: Sun Oct 18 16:02:28 2026 +0800
feat: support MOSS split QKV weights and batch host loading (#2160)
diff --git a/src/model_manager.cpp b/src/model_manager.cpp
index 5f7dbec..4562ebd 100644
--- a/src/model_manager.cpp
+++ b/src/model_manager.cpp
@@ -1022,28 +1022,60 @@ bool ModelManager::load_tensor_parts(TensorState& state) {
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*>>;
+ using ReadBatch = std::map<std::string, std::vector<ggml_tensor*>>;
std::map<ReadGroup, std::vector<ReadBatch>> groups;
+ size_t part_count = 0;
+ for (const auto* state : states) {
+ if (state != nullptr && state->sources.size() > 1 && ggml_backend_buffer_is_host(state->tensor->buffer)) {
+ part_count += state->sources.size();
+ }
+ }
+ auto parts_ctx = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
+ part_count > 0 ? ggml_init({part_count * ggml_tensor_overhead(), nullptr, true}) : nullptr, ggml_free);
+ if (part_count > 0 && !parts_ctx) {
+ return false;
+ }
+ auto add_destination = [&](const TensorState& state, const std::string& name, ggml_tensor* tensor) {
+ auto& batches = groups[{state.source_file, state.source_version}];
+ // The loader supplies one destination per name; incompatible destinations need another batch.
+ auto batch = std::find_if(batches.begin(), batches.end(), [&](const ReadBatch& candidate) {
+ auto found = candidate.find(name);
+ return found == candidate.end() ||
+ (found->second.front()->type == tensor->type && ggml_are_same_shape(found->second.front(), tensor));
+ });
+ if (batch == batches.end()) {
+ batches.emplace_back();
+ batch = std::prev(batches.end());
+ }
+ (*batch)[name].push_back(tensor);
+ };
for (auto* state : states) {
if (state == nullptr)
continue;
if (state->sources.size() > 1) {
- if (!load_tensor_parts(*state)) {
+ if (ggml_backend_buffer_is_host(state->tensor->buffer)) {
+ const size_t size = ggml_nbytes(state->tensor);
+ size_t offset = 0;
+ for (const auto& source : state->sources) {
+ auto* part = ggml_new_tensor(parts_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*>(state->tensor->data) + offset;
+ part->buffer = state->tensor->buffer;
+ add_destination(*state, source.name, part);
+ offset += part_size;
+ }
+ if (offset != size) {
+ return false;
+ }
+ } else 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) {
- auto found = candidate.find(state->name);
- return found == candidate.end() || found->second.front()->tensor->type == state->tensor->type;
- });
- if (batch == batches.end()) {
- batches.emplace_back();
- batch = std::prev(batches.end());
- }
- (*batch)[state->name].push_back(state);
+ add_destination(*state, state->name, state->tensor);
}
for (auto& group : groups) {
for (auto& batch : group.second) {
@@ -1057,7 +1089,7 @@ bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
auto found = batch.find(source.name);
if (found == batch.end())
return true;
- *dst = found->second.front()->tensor;
+ *dst = found->second.front();
std::lock_guard<std::mutex> lock(mutex);
loaded.insert(source.name);
return true;
@@ -1068,15 +1100,18 @@ bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
if (!success || loaded != names)
return false;
for (auto& entry : batch) {
- auto* first = entry.second.front()->tensor;
- for (auto* state : entry.second) {
- if (state->tensor != first)
- ggml_backend_tensor_copy(first, state->tensor);
- state->loaded_to_params_backend = true;
+ auto* first = entry.second.front();
+ for (auto* tensor : entry.second) {
+ if (tensor != first)
+ ggml_backend_tensor_copy(first, tensor);
}
}
}
}
+ for (auto* state : states) {
+ if (state != nullptr)
+ state->loaded_to_params_backend = true;
+ }
return true;
}
diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp
index 8f9f1fc..6a519ba 100644
--- a/src/name_conversion.cpp
+++ b/src/name_conversion.cpp
@@ -1013,6 +1013,16 @@ static std::string convert_diffusers_dit_to_original_pixart(std::string name) {
return prefix + name;
}
+static std::string convert_moss_image_dit_name(std::string name) {
+ static const std::vector<std::pair<std::string, std::string>> name_map = {
+ {".proj.q.weight", ".proj.qkv.weight"},
+ {".proj.k.weight", ".proj.qkv.weight.1"},
+ {".proj.v.weight", ".proj.qkv.weight.2"},
+ };
+ replace_with_name_map(name, name_map);
+ return name;
+}
+
static std::string convert_ming_image_dit_name(std::string name) {
static const std::vector<std::pair<std::string, std::string>> name_map = {
{"all_x_embedder.2-1.", "x_embedder."},
@@ -1053,6 +1063,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
name = convert_diffusers_dit_to_original_flux(name);
} else if (sd_version_is_hunyuan_video(version)) {
name = convert_hunyuan_video_to_original_flux(name);
+ } else if (version == VERSION_MOSS_IMAGE) {
+ name = convert_moss_image_dit_name(name);
} else if (version == VERSION_MING_IMAGE) {
name = convert_ming_image_dit_name(name);
} else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version)) {