refactor: unify model source and weight lifecycle management (#1956)

src/model_io/safetensors_io.cpp

@@ -144,7 +144,8 @@ static bool read_comfy_quant_config(std::ifstream& file,
 bool read_safetensors_file(const std::string& file_path,
                            std::vector<TensorStorage>& tensor_storages,
                            std::string* error,
-                           std::map<std::string, std::string>* metadata) {
+                           std::map<std::string, std::string>* metadata,
+                           std::set<std::string>* tensor_names) {
     std::ifstream file(file_path, std::ios::binary);
     if (!file.is_open()) {
         set_error(error, "failed to open '" + file_path + "'");
@@ -246,10 +247,6 @@ bool read_safetensors_file(const std::string& file_path,
         std::string dtype    = tensor_info["dtype"];
         nlohmann::json shape = tensor_info["shape"];
 
-        if (dtype == "U8") {
-            continue;
-        }
-
         size_t begin = tensor_info["data_offsets"][0].get<size_t>();
         size_t end   = tensor_info["data_offsets"][1].get<size_t>();
         if (begin > end || end > file_size_ - data_start) {
@@ -257,6 +254,26 @@ bool read_safetensors_file(const std::string& file_path,
             return false;
         }
 
+        if (tensor_names != nullptr) {
+            tensor_names->insert(name);
+        }
+        if (dtype == "U8") {
+            uint64_t bytes = 1;
+            for (const auto& dimension : shape) {
+                const int64_t size = dimension.get<int64_t>();
+                if (size < 0 || (bytes != 0 && static_cast<uint64_t>(size) > UINT64_MAX / bytes)) {
+                    set_error(error, "invalid dimensions for tensor '" + name + "'");
+                    return false;
+                }
+                bytes *= size;
+            }
+            if (bytes != end - begin) {
+                set_error(error, "size mismatch for tensor '" + name + "'");
+                return false;
+            }
+            continue;
+        }
+
         ggml_type type = safetensors_dtype_to_ggml_type(dtype);
         if (type == GGML_TYPE_COUNT) {
             set_error(error, "unsupported dtype '" + dtype + "' (tensor '" + name + "')");
@@ -270,8 +287,20 @@ bool read_safetensors_file(const std::string& file_path,
 
         int n_dims              = (int)shape.size();
         int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+        uint64_t elements       = 1;
         for (int i = 0; i < n_dims; i++) {
             ne[i] = shape[i].get<int64_t>();
+            if (ne[i] < 0 || (elements != 0 && static_cast<uint64_t>(ne[i]) > INT64_MAX / elements)) {
+                set_error(error, "invalid dimensions for tensor '" + name + "'");
+                return false;
+            }
+            elements *= ne[i];
+        }
+        const uint64_t storage_size = ggml_type_size(type) * ((dtype == "F64" || dtype == "I64") ? 2 : 1);
+        if (elements % ggml_blck_size(type) != 0 ||
+            elements / ggml_blck_size(type) > INT64_MAX / storage_size) {
+            set_error(error, "invalid storage size for tensor '" + name + "'");
+            return false;
         }
 
         if (n_dims == 5) {

src/model_io/safetensors_io.h

@@ -2,6 +2,7 @@
 #define __SD_MODEL_IO_SAFETENSORS_IO_H__
 
 #include <map>
+#include <set>
 #include <string>
 #include <vector>
 
@@ -12,7 +13,8 @@ bool is_safetensors_file(const std::string& file_path);
 bool read_safetensors_file(const std::string& file_path,
                            std::vector<TensorStorage>& tensor_storages,
                            std::string* error                           = nullptr,
-                           std::map<std::string, std::string>* metadata = nullptr);
+                           std::map<std::string, std::string>* metadata = nullptr,
+                           std::set<std::string>* tensor_names          = nullptr);
 bool read_safetensors_index_file(const std::string& file_path,
                                  std::vector<std::string>& shard_paths,
                                  std::string* error = nullptr);

src/model_io/tensor_storage.h

@@ -28,9 +28,11 @@ struct TensorStorage {
     int n_dims                  = 0;
 
     std::string storage_key;
-    size_t file_index = 0;
-    int index_in_zip  = -1;  // >= means stored in a zip file
-    uint64_t offset   = 0;   // offset in file
+    size_t file_index      = 0;
+    uint64_t file_id       = 0;
+    uint64_t file_revision = 0;
+    int index_in_zip       = -1;  // >= means stored in a zip file
+    uint64_t offset        = 0;   // offset in file
 
     TensorStorage() = default;