diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 5f7dbec9..4562ebdf 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& states) { using ReadGroup = std::pair; - using ReadBatch = std::map>; + using ReadBatch = std::map>; std::map> groups; - for (auto* state : states) { - if (state == nullptr) - continue; - if (state->sources.size() > 1) { - if (!load_tensor_parts(*state)) { - return false; - } - continue; + 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& batches = groups[{state->source_file, state->source_version}]; - // The loader supplies one destination per name; only conflicting types need another batch. + } + auto parts_ctx = std::unique_ptr( + 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(state->name); - return found == candidate.end() || found->second.front()->tensor->type == state->tensor->type; + 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)[state->name].push_back(state); + (*batch)[name].push_back(tensor); + }; + for (auto* state : states) { + if (state == nullptr) + continue; + if (state->sources.size() > 1) { + 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(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; + } + 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& 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 lock(mutex); loaded.insert(source.name); return true; @@ -1068,15 +1100,18 @@ bool ModelManager::load_tensors(const std::vector& 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 8f9f1fc4..6a519ba2 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> 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> 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)) {