QVAC-25501 fix: support header-only safetensors fitting

examples/fit-params/README.md

@@ -4,6 +4,9 @@
 device memory, using measured metadata-only dry runs: the real generation
 pipeline is executed with graph building and memory measurement only, so
 no weight data is read and no ggml weight or compute buffers are allocated.
+Safetensors inputs may contain only their complete header, including auxiliary
+models and LoRAs. A successful fit does not verify that weight data is present;
+generation still requires complete model files.
 Shaped host tensors are still materialized to carry state between graph builds;
 allocation failures are reported as fit errors. The measured per-module memory
 includes projected persistent cache buffers and is packed against the free

src/core/util.cpp

@@ -342,6 +342,7 @@ int32_t sd_get_num_physical_cores() {
 static sd_progress_cb_t sd_progress_cb = nullptr;
 void* sd_progress_cb_data              = nullptr;
 static thread_local bool sd_progress_suppressed = false;
+static thread_local bool sd_metadata_only_read  = false;
 
 static sd_abort_cb_t sd_abort_cb = nullptr;
 static void* sd_abort_cb_data   = nullptr;
@@ -702,6 +703,19 @@ void sd_set_progress_suppressed(bool suppressed) {
     sd_progress_suppressed = suppressed;
 }
 
+bool sd_get_metadata_only_read() {
+    return sd_metadata_only_read;
+}
+
+SDMetadataOnlyReadScope::SDMetadataOnlyReadScope()
+    : previous_(sd_metadata_only_read) {
+    sd_metadata_only_read = true;
+}
+
+SDMetadataOnlyReadScope::~SDMetadataOnlyReadScope() {
+    sd_metadata_only_read = previous_;
+}
+
 sd_image_t tensor_to_sd_image(const sd::Tensor<float>& tensor, int frame_index) {
     const auto& shape = tensor.shape();
     GGML_ASSERT(shape.size() == 4 || shape.size() == 5);

src/core/util.h

@@ -94,6 +94,18 @@ void* sd_get_progress_callback_data();
 bool sd_get_progress_suppressed();
 void sd_set_progress_suppressed(bool suppressed);
 
+bool sd_get_metadata_only_read();
+
+class SDMetadataOnlyReadScope {
+    bool previous_;
+
+public:
+    SDMetadataOnlyReadScope();
+    ~SDMetadataOnlyReadScope();
+    SDMetadataOnlyReadScope(const SDMetadataOnlyReadScope&)            = delete;
+    SDMetadataOnlyReadScope& operator=(const SDMetadataOnlyReadScope&) = delete;
+};
+
 sd_preview_cb_t sd_get_preview_callback();
 void* sd_get_preview_callback_data();
 preview_t sd_get_preview_mode();

src/model/adapter/lora.hpp

@@ -131,7 +131,7 @@ struct LoraModel : public GGMLRunner {
         for (const auto& pair : lora_tensors) {
             lora_params.push_back(pair.second);
         }
-        if (!model_manager->prepare_params(lora_params)) {
+        if (!GGMLRunner::measure_mode_enabled() && !model_manager->prepare_params(lora_params)) {
             LOG_ERROR("lora model manager prepare params failed");
             return false;
         }
@@ -409,6 +409,11 @@ struct LoraModel : public GGMLRunner {
         return compatible;
     }
 
+    static float read_scale(ggml_tensor* tensor) {
+        // Scalar values do not affect graph shapes during memory measurement.
+        return GGMLRunner::measure_mode_enabled() ? 1.0f : ggml_ext_backend_tensor_get_f32(tensor);
+    }
+
     ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
         ggml_tensor* updown = nullptr;
         int index           = 0;
@@ -461,12 +466,12 @@ struct LoraModel : public GGMLRunner {
             int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
             iter         = lora_tensors.find(scale_name);
             if (iter != lora_tensors.end()) {
-                scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
+                scale_value = read_scale(iter->second);
                 applied_lora_tensors.insert(scale_name);
             } else {
                 iter = lora_tensors.find(alpha_name);
                 if (iter != lora_tensors.end()) {
-                    float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
+                    float alpha = read_scale(iter->second);
                     scale_value = alpha / rank;
                     // LOG_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
                     applied_lora_tensors.insert(alpha_name);
@@ -615,7 +620,7 @@ struct LoraModel : public GGMLRunner {
             int64_t rank = hada_1_down->ne[ggml_n_dims(hada_1_down) - 1];
             iter         = lora_tensors.find(alpha_name);
             if (iter != lora_tensors.end()) {
-                float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
+                float alpha = read_scale(iter->second);
                 scale_value = alpha / rank;
                 applied_lora_tensors.insert(alpha_name);
             }
@@ -728,7 +733,7 @@ struct LoraModel : public GGMLRunner {
             float scale_value = 1.0f;
             iter              = lora_tensors.find(alpha_name);
             if (iter != lora_tensors.end()) {
-                float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
+                float alpha = read_scale(iter->second);
                 scale_value = alpha / rank;
                 applied_lora_tensors.insert(alpha_name);
             }
@@ -889,7 +894,7 @@ struct LoraModel : public GGMLRunner {
                 float scale_value = 1.0f;
                 iter              = lora_tensors.find(alpha_name);
                 if (iter != lora_tensors.end()) {
-                    float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
+                    float alpha = read_scale(iter->second);
                     scale_value = alpha / rank;
                 }
 
@@ -1016,12 +1021,12 @@ struct LoraModel : public GGMLRunner {
             int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
             iter         = lora_tensors.find(scale_name);
             if (iter != lora_tensors.end()) {
-                scale_value       = ggml_ext_backend_tensor_get_f32(iter->second);
+                scale_value       = read_scale(iter->second);
                 scale_tensor_name = scale_name;
             } else {
                 iter = lora_tensors.find(alpha_name);
                 if (iter != lora_tensors.end()) {
-                    float alpha       = ggml_ext_backend_tensor_get_f32(iter->second);
+                    float alpha       = read_scale(iter->second);
                     scale_value       = alpha / rank;
                     scale_tensor_name = alpha_name;
                     // LOG_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);

src/model_io/safetensors_io.cpp

@@ -5,6 +5,7 @@
 #include <exception>
 #include <filesystem>
 #include <fstream>
+#include <limits>
 #include <ostream>
 #include <string>
 #include <unordered_set>
@@ -180,9 +181,19 @@ bool read_safetensors_file(const std::string& file_path,
             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) {
+        const auto& offsets     = tensor_info["data_offsets"];
+        const auto valid_offset = [](const nlohmann::json& value) {
+            return value.is_number_unsigned() ? value.get<uint64_t>() <= std::numeric_limits<size_t>::max()
+                                              : value.is_number_integer() && value.get<int64_t>() >= 0;
+        };
+        if (!offsets.is_array() || offsets.size() != 2 || !valid_offset(offsets[0]) || !valid_offset(offsets[1])) {
+            set_error(error, "invalid data offsets for tensor '" + name + "'");
+            return false;
+        }
+        size_t begin = offsets[0].get<size_t>();
+        size_t end   = offsets[1].get<size_t>();
+        if (begin > end || end > std::numeric_limits<size_t>::max() - data_start ||
+            (!sd_get_metadata_only_read() && end > file_size_ - data_start)) {
             set_error(error, "data offsets out of bounds for tensor '" + name + "'");
             return false;
         }
@@ -193,15 +204,30 @@ bool read_safetensors_file(const std::string& file_path,
             return false;
         }
 
-        if (shape.size() > SD_MAX_DIMS) {
+        if (!shape.is_array() || shape.size() > SD_MAX_DIMS) {
             set_error(error, "invalid tensor '" + name + "'");
             return false;
         }
 
         int n_dims              = (int)shape.size();
         int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
+        // Bound intermediate products too, including F64/I64 conversion and zero-sized tensors.
+        const uint64_t max_nelements =
+            std::min<uint64_t>(std::numeric_limits<int64_t>::max(), std::numeric_limits<size_t>::max()) /
+            (2 * ggml_type_size(type));
+        uint64_t nelements = 1;
         for (int i = 0; i < n_dims; i++) {
-            ne[i] = shape[i].get<int64_t>();
+            if (!shape[i].is_number_unsigned()) {
+                set_error(error, "invalid shape for tensor '" + name + "'");
+                return false;
+            }
+            const uint64_t dim = shape[i].get<uint64_t>();
+            if (dim > max_nelements || (dim != 0 && nelements > max_nelements / dim)) {
+                set_error(error, "invalid shape for tensor '" + name + "'");
+                return false;
+            }
+            nelements *= std::max<uint64_t>(dim, 1);
+            ne[i] = static_cast<int64_t>(dim);
         }
 
         if (n_dims == 5) {
@@ -226,11 +252,11 @@ bool read_safetensors_file(const std::string& file_path,
         if (dtype == "F8_E4M3") {
             tensor_storage.is_f8_e4m3 = true;
             // f8 -> f16
-            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
+            tensor_size_ok = (tensor_storage.nbytes() / 2 == tensor_data_size);
         } else if (dtype == "F8_E5M2") {
             tensor_storage.is_f8_e5m2 = true;
             // f8 -> f16
-            tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
+            tensor_size_ok = (tensor_storage.nbytes() / 2 == tensor_data_size);
         } else if (dtype == "F64") {
             tensor_storage.is_f64 = true;
             // f64 -> f32

src/stable-diffusion.cpp

@@ -4054,6 +4054,7 @@ enum sd_fit_status_t sd_fit_params(const sd_ctx_params_t* sd_ctx_params,
         return SD_FIT_FAILURE;
     }
 
+    ggml_time_init();
     int64_t t0 = ggml_time_ms();
 
     sd_ctx_params_t dry_params = *sd_ctx_params;
@@ -4064,7 +4065,14 @@ enum sd_fit_status_t sd_fit_params(const sd_ctx_params_t* sd_ctx_params,
     sd_ctx_t* sd_ctx        = &sd_ctx_storage;
     sd_ctx->sd              = new StableDiffusionGGML();
     sd_ctx->sd->fit_dry_run = true;
-    if (!sd_ctx->sd->init(&dry_params)) {
+    bool init_ok = false;
+    try {
+        SDMetadataOnlyReadScope metadata_only;
+        init_ok = sd_ctx->sd->init(&dry_params);
+    } catch (const std::exception& error) {
+        LOG_ERROR("fit-params: dry-run model init failed: %s", error.what());
+    }
+    if (!init_ok) {
         LOG_ERROR("fit-params: dry-run model init failed");
         delete sd_ctx->sd;
         sd_ctx->sd = nullptr;
@@ -4128,6 +4136,7 @@ enum sd_fit_status_t sd_fit_params(const sd_ctx_params_t* sd_ctx_params,
 
     auto measure = [&](const sd_tiling_params_t& tiling,
                        std::vector<GGMLRunner::graph_memory_measurement>& records) -> bool {
+        SDMetadataOnlyReadScope metadata_only;
         records.clear();
         struct MeasureModeGuard {
             explicit MeasureModeGuard(std::vector<GGMLRunner::graph_memory_measurement>* records) {

tests/CMakeLists.txt

@@ -11,6 +11,11 @@ target_include_directories(test-fit-params PRIVATE "${PROJECT_SOURCE_DIR}/src")
 target_link_libraries(test-fit-params PRIVATE stable-diffusion ${CMAKE_THREAD_LIBS_INIT})
 add_test(NAME test-fit-params COMMAND test-fit-params)
 
+add_executable(test-safetensors-metadata test_safetensors_metadata.cpp)
+target_include_directories(test-safetensors-metadata PRIVATE "${PROJECT_SOURCE_DIR}/src")
+target_link_libraries(test-safetensors-metadata PRIVATE stable-diffusion ${CMAKE_THREAD_LIBS_INIT})
+add_test(NAME test-safetensors-metadata COMMAND test-safetensors-metadata)
+
 add_executable(test-ltx-vae-temporal test-ltx-vae-temporal.cpp)
 target_link_libraries(test-ltx-vae-temporal PRIVATE stable-diffusion)
 add_test(NAME test-ltx-vae-temporal COMMAND test-ltx-vae-temporal)