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

src/model_loader.h

@@ -2,6 +2,7 @@
 #define __MODEL_LOADER_H__
 
 #include <cstdint>
+#include <filesystem>
 #include <map>
 #include <memory>
 #include <set>
@@ -30,6 +31,46 @@ struct MmapTensorStore {
 bool is_unused_tensor(const std::string& name);
 
 class ModelLoader {
+public:
+    using FileId       = uint64_t;
+    using FileVersions = std::map<FileId, uint64_t>;
+    enum class FileScope { Catalog,
+                           Isolated };
+
+private:
+    struct FileStamp {
+        std::string path;
+        uintmax_t size = 0;
+        std::filesystem::file_time_type modified;
+    };
+
+    struct FileRecord {
+        FileId id         = 0;
+        uint64_t revision = 0;
+        std::string path;
+        std::string prefix;
+        FileScope scope = FileScope::Catalog;
+        std::vector<FileStamp> dependencies;
+        String2TensorStorage tensors;
+        std::map<std::string, std::string> metadata;
+    };
+
+    std::vector<FileRecord> files_;
+    uint64_t revision_        = 0;
+    bool names_converted_     = false;
+    ggml_type wtype_override_ = GGML_TYPE_COUNT;
+    std::string tensor_type_rules_;
+    std::vector<FileStamp> parsed_dependencies_;
+    std::map<std::string, std::set<std::string>> parsed_tensor_names_;
+
+    static bool read_file_stamp(const std::string& path, FileStamp& stamp);
+    static bool file_unchanged(const FileStamp& stamp);
+    bool parse_file(const std::string& path, const std::string& prefix);
+    bool add_file_impl(const std::string& path, const std::string& prefix, FileId* id, bool force, FileScope scope);
+    ModelLoader file_reader(FileId id, SDVersion version) const;
+    void rebuild_catalog();
+    void invalidate_file_data();
+
 protected:
     SDVersion version_ = VERSION_COUNT;
     std::vector<std::string> file_paths_;
@@ -52,16 +93,27 @@ protected:
 public:
     ModelLoader();
 
+    bool add_file(const std::string& path, const std::string& prefix = "", FileId* id = nullptr, bool force = false, FileScope scope = FileScope::Catalog);
+    bool del_file(FileId id);
+    uint64_t file_revision(FileId id) const;
+    std::string file_path(FileId id) const;
+    String2TensorStorage file_tensors(FileId id, SDVersion version) const;
+    bool load_file_tensors(FileId id, SDVersion version, on_new_tensor_cb_t callback, const std::set<std::string>& names, bool use_mmap = false) const;
+    bool refresh_files(bool include_isolated = true);
+    bool files_changed(bool& changed, bool include_isolated = true) const;
+    bool validate_sources(const std::set<std::string>* tensor_names = nullptr) const;
+    uint64_t revision() const { return revision_; }
+    FileVersions file_versions(const std::vector<std::string>& prefixes = {}) const;
     bool init_from_file(const std::string& file_path, const std::string& prefix = "");
     void convert_tensors_name();
     bool init_from_file_and_convert_name(const std::string& file_path,
                                          const std::string& prefix = "",
                                          SDVersion version         = VERSION_COUNT);
-    SDVersion get_sd_version();
-    std::map<ggml_type, uint32_t> get_wtype_stat();
-    std::map<ggml_type, uint32_t> get_conditioner_wtype_stat();
-    std::map<ggml_type, uint32_t> get_diffusion_model_wtype_stat();
-    std::map<ggml_type, uint32_t> get_vae_wtype_stat();
+    SDVersion get_sd_version() const;
+    std::map<ggml_type, uint32_t> get_wtype_stat() const;
+    std::map<ggml_type, uint32_t> get_conditioner_wtype_stat() const;
+    std::map<ggml_type, uint32_t> get_diffusion_model_wtype_stat() const;
+    std::map<ggml_type, uint32_t> get_vae_wtype_stat() const;
     String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
     const String2TensorStorage& get_tensor_storage_map() const { return tensor_storage_map; }
     const std::map<std::string, std::string>& get_metadata() const { return metadata_; }
@@ -92,8 +144,8 @@ public:
         return names;
     }
 
-    bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type);
-    int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
+    bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) const;
+    int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT) const;
     ~ModelLoader() = default;
 };