Bound PyTorch ZIP storage entry lookups

CMakeLists.txt

@@ -87,6 +87,7 @@ function(sd_native_carry_patch patch)
 endfunction()
 
 sd_native_carry_patch(pickle_line_span.patch src/model_io/pickle_io.cpp)
+sd_native_carry_patch(torch_zip_entry_lookup.patch src/model_io/torch_zip_io.cpp)
 sd_native_carry_patch(fp8_metal_f16_load.patch src/model_manager.cpp src/core/backend_fit.cpp)
 sd_native_carry_patch(lora_immediate_fallback.patch src/pipeline/diffusion_engine.cpp)
 sd_native_carry_patch(intrinsic_tiny_decoder_policy.patch src/stable-diffusion.cpp src/pipeline/diffusion_engine.cpp)

patches/README.md

@@ -18,6 +18,7 @@ reports.
 
 | Patch | Sources | Fixes | Test |
 | --- | --- | --- | --- |
+| `torch_zip_entry_lookup.patch` | `src/model_io/torch_zip_io.cpp` | Cache archive-order storage lookups so PyTorch ZIP metadata opens at most twice the archive entry count; preserve first-match names, indices and tensor bounds ([owner #24](https://github.com/leehack/stable-diffusion-native/issues/24)). **Patch-specific maintainer exception:** the upstream report is retained locally at the maintainer's request; [owner #24](https://github.com/leehack/stable-diffusion-native/issues/24) tracks retirement until an upstream report is authorized. Retire when upstream provides equivalent bounded lookup behavior. | `tests/test_torch_zip_lookup.py` |
 | `pickle_line_span.patch` | `src/model_io/pickle_io.cpp` | A legacy PyTorch checkpoint above 2 GiB is rejected: the pickle scanner narrows the rest of the file to `int` before looking for a newline ([#23](https://github.com/leehack/stable-diffusion-native/issues/23)). | `tests/test_pickle_span.py` |
 | `fp8_metal_f16_load.patch` | `src/model_manager.cpp`, `src/core/backend_fit.cpp` | An FP8 (E4M3 or E5M2) checkpoint aborts the process on Metal: ggml-metal claims FP8 `MUL_MAT` but has no FP8 kernel. Weights a Metal backend computes with now load as F16, which holds every FP8 value, so they take twice their file size; other backends keep FP8 ([#29](https://github.com/leehack/stable-diffusion-native/issues/29)). Auto-fit plans them at that size, and FP8 LoRA tensors, which crashed ggml-cpu, load as F16 on every backend ([#31](https://github.com/leehack/stable-diffusion-native/issues/31)). | `tests/test_fp8_runtime.py` |
 | `lora_immediate_fallback.patch` | `src/pipeline/diffusion_engine.cpp` | The `immediately` LoRA apply mode, which `auto` selects for unquantized weights, crashes the CPU merging into FP8 weights (ggml-cpu can neither multiply nor write FP8) and aborts Metal merging into F16 weights (ggml-metal adds only equal types), which every model has: convolution weights load as F16. FP8 weights, and every model on a backend that cannot run that add outside host memory, now take the LoRA `at_runtime`, with a warning when `immediately` was requested ([#31](https://github.com/leehack/stable-diffusion-native/issues/31)). | `tests/test_fp8_runtime.py` |

patches/torch_zip_entry_lookup.patch

@@ -0,0 +1,120 @@
+--- a/src/model_io/torch_zip_io.cpp
++++ b/src/model_io/torch_zip_io.cpp
+@@ -1,5 +1,6 @@
+ #include "torch_zip_io.h"
+
++#include <climits>
+ #include <cstdint>
+ #include <cstdlib>
+ #include <string>
+@@ -25,18 +26,41 @@
+     return true;
+ }
+
+-static bool find_zip_entry(zip_t* zip, const std::string& entry_name, int* index, uint64_t* size) {
+-    size_t n = zip_entries_total(zip);
+-    for (size_t i = 0; i < n; ++i) {
+-        zip_entry_openbyindex(zip, i);
+-        std::string name = zip_entry_name(zip);
++struct ZipEntryLookup {
++    struct Entry { int index; uint64_t size; };
++    // Keep the first normalized name, as a fresh archive-order scan would.
++    std::unordered_map<std::string, Entry> entries;
++    size_t scanned = 0;
++};
++
++static bool find_zip_entry(zip_t* zip, ZipEntryLookup& lookup, const std::string& entry_name, int* index, uint64_t* size) {
++    auto found = lookup.entries.find(entry_name);
++    if (found != lookup.entries.end()) {
++        *index = found->second.index;
++        *size = found->second.size;
++        return true;
++    }
++    const size_t n = static_cast<size_t>(zip_entries_total(zip));
++    while (lookup.scanned < n) {
++        size_t i = lookup.scanned++;
++        if (zip_entry_openbyindex(zip, i) != 0) {
++            zip_entry_close(zip);
++            return false;
++        }
++        const char* raw_name = zip_entry_name(zip);
++        if (raw_name == nullptr) {
++            zip_entry_close(zip);
++            return false;
++        }
++        std::string name = raw_name;
++        uint64_t entry_size = zip_entry_size(zip);
++        lookup.entries.emplace(name, ZipEntryLookup::Entry{(int)i, entry_size});
++        zip_entry_close(zip);
+         if (name == entry_name) {
+             *index = (int)i;
+-            *size  = zip_entry_size(zip);
+-            zip_entry_close(zip);
++            *size = entry_size;
+             return true;
+         }
+-        zip_entry_close(zip);
+     }
+     return false;
+ }
+@@ -44,6 +68,7 @@
+ static bool parse_zip_data_pkl(const uint8_t* buffer,
+                                size_t buffer_size,
+                                zip_t* zip,
++                               ZipEntryLookup& lookup,
+                                const std::string& dir,
+                                std::vector<TensorStorage>& tensor_storages,
+                                std::string* error) {
+@@ -65,7 +90,7 @@
+         const std::string entry_name = dir + "data/" + tensor_storage.storage_key;
+         int zip_index                = -1;
+         uint64_t entry_size          = 0;
+-        if (!find_zip_entry(zip, entry_name, &zip_index, &entry_size)) {
++        if (!find_zip_entry(zip, lookup, entry_name, &zip_index, &entry_size)) {
+             set_error(error, "storage entry '" + entry_name + "' was not found");
+             return false;
+         }
+@@ -102,12 +127,31 @@
+     }
+
+     tensor_storages.clear();
++    ZipEntryLookup lookup;
+     bool success        = true;
+     bool found_data_pkl = false;
+-    int n               = (int)zip_entries_total(zip);
++    const ssize_t entry_count = zip_entries_total(zip);
++    if (entry_count < 0 || static_cast<uint64_t>(entry_count) > INT_MAX) {
++        set_error(error, "torch zip entry count exceeds supported index range");
++        zip_close(zip);
++        return false;
++    }
++    int n = static_cast<int>(entry_count);
+     for (int i = 0; i < n; ++i) {
+-        zip_entry_openbyindex(zip, i);
+-        std::string name = zip_entry_name(zip);
++        if (zip_entry_openbyindex(zip, i) != 0) {
++            set_error(error, "failed to open torch zip entry " + std::to_string(i));
++            zip_entry_close(zip);
++            success = false;
++            break;
++        }
++        const char* raw_name = zip_entry_name(zip);
++        if (raw_name == nullptr) {
++            set_error(error, "torch zip entry has no name");
++            zip_entry_close(zip);
++            success = false;
++            break;
++        }
++        std::string name = raw_name;
+         size_t pos       = name.find("data.pkl");
+         if (pos != std::string::npos) {
+             found_data_pkl  = true;
+@@ -119,7 +163,7 @@
+             if (pkl_data == nullptr || pkl_size == 0) {
+                 set_error(error, "failed to read '" + name + "' from '" + file_path + "'");
+                 success = false;
+-            } else if (!parse_zip_data_pkl((const uint8_t*)pkl_data, pkl_size, zip, dir, tensor_storages, error)) {
++            } else if (!parse_zip_data_pkl((const uint8_t*)pkl_data, pkl_size, zip, lookup, dir, tensor_storages, error)) {
+                 success = false;
+             }
+

tests/native/torch_zip_lookup_test.cpp

@@ -0,0 +1,128 @@
+// Real upstream ZIP/pickle readers with observed API calls; no tensor backend.
+#include <climits>
+#include <cstdarg>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <iostream>
+#include <stdexcept>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "model_io/pickle_io.h"
+#include "zip.h"
+#include "json.hpp"
+
+extern "C" size_t ggml_type_size(enum ggml_type type) {
+    switch (type) {
+        case GGML_TYPE_F32: case GGML_TYPE_I32: return 4;
+        case GGML_TYPE_F16: case GGML_TYPE_BF16: return 2;
+        default: throw std::runtime_error("unexpected metadata scalar type");
+    }
+}
+extern "C" int64_t ggml_blck_size(enum ggml_type type) {
+    (void)ggml_type_size(type);
+    return 1;
+}
+std::string sd_format(const char* format, ...) {
+    char buffer[64];
+    va_list args;
+    va_start(args, format);
+    const int count = std::vsnprintf(buffer, sizeof(buffer), format, args);
+    va_end(args);
+    if (count < 0 || static_cast<size_t>(count) >= sizeof(buffer))
+        throw std::runtime_error("error format overflow");
+    return buffer;
+}
+
+namespace {
+size_t opens = 0, names = 0, archives_opened = 0, archives_closed = 0;
+size_t metadata_reads = 0, payload_reads = 0;
+std::unordered_set<void*> pickle_allocations;
+std::string fault;
+zip_t* observed_zip_open(const char* path, int level, char mode) {
+    zip_t* zip = zip_open(path, level, mode);
+    if (zip) ++archives_opened;
+    return zip;
+}
+void observed_zip_close(zip_t* zip) {
+    if (zip) ++archives_closed;
+    zip_close(zip);
+}
+ssize_t observed_entry_count(zip_t* zip) {
+    if (fault == "negative-count") return -1;
+    if (fault == "large-count") return static_cast<ssize_t>(INT_MAX) + 1;
+    return zip_entries_total(zip);
+}
+int observed_entry_open(zip_t* zip, size_t index) {
+    ++opens;
+    // A failed open may already have allocated entry state. Exercise cleanup
+    // after the real call, rather than replacing the ZIP implementation.
+    const int result = zip_entry_openbyindex(zip, index);
+    if ((fault == "outer-open" && opens == 1) ||
+        (fault == "lookup-open" && opens == 2)) return ZIP_ENOENT;
+    return result;
+}
+const char* observed_entry_name(zip_t* zip) {
+    ++names;
+    if ((fault == "outer-name" && names == 1) ||
+        (fault == "lookup-name" && names == 2)) return nullptr;
+    return zip_entry_name(zip);
+}
+ssize_t observed_entry_read(zip_t* zip, void** buffer, size_t* size) {
+    const std::string name = zip_entry_name(zip);
+    if (name.find("data.pkl") == std::string::npos) {
+        ++payload_reads;
+        throw std::runtime_error("tensor payload extraction refused");
+    }
+    ++metadata_reads;
+    const ssize_t result = zip_entry_read(zip, buffer, size);
+    if (*buffer) pickle_allocations.insert(*buffer);
+    return result;
+}
+void observed_free(void* ptr) {
+    pickle_allocations.erase(ptr);
+    std::free(ptr);
+}
+}  // namespace
+
+#define zip_open observed_zip_open
+#define zip_close observed_zip_close
+#define zip_entries_total observed_entry_count
+#define zip_entry_openbyindex observed_entry_open
+#define zip_entry_name observed_entry_name
+#define zip_entry_read observed_entry_read
+#define free observed_free
+#include "model_io/torch_zip_io.cpp"
+#undef free
+#undef zip_entry_read
+#undef zip_entry_name
+#undef zip_entry_openbyindex
+#undef zip_entries_total
+#undef zip_close
+#undef zip_open
+
+int main(int argc, char** argv) {
+    if (argc < 3) return 2;
+    fault = argv[1];
+    for (int i = 2; i < argc; ++i) {
+        opens = names = archives_opened = archives_closed = metadata_reads = payload_reads = 0;
+        if (!pickle_allocations.empty()) return 4;
+        std::string error;
+        std::vector<TensorStorage> tensors;
+        const bool ok = read_torch_zip_file(argv[i], tensors, &error);
+        nlohmann::json inventory = nlohmann::json::array();
+        for (const auto& t : tensors) {
+            inventory.push_back({{"name", t.name}, {"type", t.type},
+                {"dims", std::vector<int64_t>(t.ne, t.ne + SD_MAX_DIMS)}, {"n_dims", t.n_dims},
+                {"index", t.index_in_zip}, {"offset", t.offset}, {"bytes", t.nbytes_to_read()},
+                {"storage_key", t.storage_key}, {"f64", t.is_f64}, {"i64", t.is_i64}});
+        }
+        std::cout << nlohmann::json{{"ok", ok}, {"error", error}, {"inventory", inventory},
+            {"entry_opens", opens}, {"entry_names", names}, {"archives_opened", archives_opened},
+            {"archives_closed", archives_closed}, {"live_pickle_buffers", pickle_allocations.size()},
+            {"metadata_reads", metadata_reads}, {"payload_reads", payload_reads}}.dump() << '\n';
+    }
+}

tests/native/torch_zip_metadata_shim/core/util.h

@@ -0,0 +1,5 @@
+#pragma once
+#include <string>
+
+// The metadata parser uses util only to format unsupported-opcode errors.
+std::string sd_format(const char* format, ...);