|
15 | 15 | #include <atomic> |
16 | 16 | #include <chrono> |
17 | 17 | #include <queue> |
| 18 | +#include <map> |
18 | 19 |
|
19 | 20 | #ifdef _WIN32 |
20 | 21 | #include <winsock2.h> |
@@ -86,49 +87,61 @@ static std::vector<local_model> list_local_models(const std::string & dir) { |
86 | 87 | } |
87 | 88 |
|
88 | 89 | std::vector<local_model> models; |
89 | | - auto scan_subdir = [&models](const std::string & subdir_path, const std::string & name) { |
90 | | - auto files = fs_list(subdir_path, false); |
| 90 | + |
| 91 | + struct dir_model_files { |
91 | 92 | common_file_info model_file; |
92 | 93 | common_file_info first_shard_file; |
93 | 94 | common_file_info mmproj_file; |
94 | | - for (const auto & file : files) { |
95 | | - if (string_ends_with(file.name, ".gguf")) { |
96 | | - if (file.name.find("mmproj") != std::string::npos) { |
97 | | - mmproj_file = file; |
98 | | - } else if (file.name.find("-00001-of-") != std::string::npos) { |
99 | | - first_shard_file = file; |
100 | | - } else { |
101 | | - model_file = file; |
102 | | - } |
103 | | - } |
| 95 | + }; |
| 96 | + |
| 97 | + std::map<std::filesystem::path, dir_model_files> model_directories; |
| 98 | + |
| 99 | + for (const auto & entry : std::filesystem::recursive_directory_iterator( |
| 100 | + dir, std::filesystem::directory_options::skip_permission_denied)) { |
| 101 | + if (!entry.is_regular_file()) { |
| 102 | + continue; |
104 | 103 | } |
105 | | - // single file model |
106 | | - local_model model{ |
107 | | - /* name */ name, |
108 | | - /* path */ first_shard_file.path.empty() ? model_file.path : first_shard_file.path, |
109 | | - /* path_mmproj */ mmproj_file.path // can be empty |
110 | | - }; |
111 | | - if (!model.path.empty()) { |
112 | | - models.push_back(model); |
| 104 | + |
| 105 | + const auto & path = entry.path(); |
| 106 | + if (!string_ends_with(path.filename().string(), ".gguf")) { |
| 107 | + continue; |
113 | 108 | } |
114 | | - }; |
115 | 109 |
|
116 | | - auto files = fs_list(dir, true); |
117 | | - for (const auto & file : files) { |
118 | | - if (file.is_dir) { |
119 | | - scan_subdir(file.path, file.name); |
120 | | - } else if (string_ends_with(file.name, ".gguf")) { |
121 | | - // single file model |
122 | | - std::string name = file.name; |
123 | | - string_replace_all(name, ".gguf", ""); |
124 | | - local_model model{ |
125 | | - /* name */ name, |
126 | | - /* path */ file.path, |
127 | | - /* path_mmproj */ "" |
128 | | - }; |
129 | | - models.push_back(model); |
| 110 | + auto & files = model_directories[path.parent_path()]; |
| 111 | + const auto filename = path.filename().string(); |
| 112 | + if (filename.find("mmproj") != std::string::npos) { |
| 113 | + files.mmproj_file = {path.string(), filename, 0, false}; |
| 114 | + } else if (filename.find("-00001-of-") != std::string::npos) { |
| 115 | + files.first_shard_file = {path.string(), filename, 0, false}; |
| 116 | + } else { |
| 117 | + files.model_file = {path.string(), filename, 0, false}; |
| 118 | + } |
| 119 | + } |
| 120 | + |
| 121 | + for (const auto & [parent_path, files] : model_directories) { |
| 122 | + std::string model_path = files.first_shard_file.path.empty() ? files.model_file.path : files.first_shard_file.path; |
| 123 | + if (model_path.empty()) { |
| 124 | + continue; |
| 125 | + } |
| 126 | + |
| 127 | + std::string name; |
| 128 | + std::error_code ec; |
| 129 | + auto rel_parent = std::filesystem::relative(parent_path, dir, ec); |
| 130 | + if (!ec && !rel_parent.empty() && rel_parent.string() != ".") { |
| 131 | + name = rel_parent.generic_string(); |
| 132 | + } else { |
| 133 | + std::filesystem::path model_file_path(model_path); |
| 134 | + name = model_file_path.stem().string(); |
130 | 135 | } |
| 136 | + |
| 137 | + local_model model{ |
| 138 | + /* name */ name, |
| 139 | + /* path */ model_path, |
| 140 | + /* path_mmproj */ files.mmproj_file.path |
| 141 | + }; |
| 142 | + models.push_back(model); |
131 | 143 | } |
| 144 | + |
132 | 145 | return models; |
133 | 146 | } |
134 | 147 |
|
|
0 commit comments