diff --git a/CMakeLists.txt b/CMakeLists.txt index 6a571bc..4768a8f 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/patches/README.md b/patches/README.md index c3736cf..0a6fab3 100644 --- a/patches/README.md +++ b/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` | diff --git a/patches/torch_zip_entry_lookup.patch b/patches/torch_zip_entry_lookup.patch new file mode 100644 index 0000000..6054d49 --- /dev/null +++ b/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 + #include + #include + #include +@@ -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 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(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& 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(entry_count) > INT_MAX) { ++ set_error(error, "torch zip entry count exceeds supported index range"); ++ zip_close(zip); ++ return false; ++ } ++ int n = static_cast(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; + } + diff --git a/tests/native/torch_zip_lookup_test.cpp b/tests/native/torch_zip_lookup_test.cpp new file mode 100644 index 0000000..d734ef9 --- /dev/null +++ b/tests/native/torch_zip_lookup_test.cpp @@ -0,0 +1,128 @@ +// Real upstream ZIP/pickle readers with observed API calls; no tensor backend. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#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(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 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(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 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(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'; + } +} diff --git a/tests/native/torch_zip_metadata_shim/core/util.h b/tests/native/torch_zip_metadata_shim/core/util.h new file mode 100644 index 0000000..86d7080 --- /dev/null +++ b/tests/native/torch_zip_metadata_shim/core/util.h @@ -0,0 +1,5 @@ +#pragma once +#include + +// The metadata parser uses util only to format unsupported-opcode errors. +std::string sd_format(const char* format, ...); diff --git a/tests/test_torch_zip_lookup.py b/tests/test_torch_zip_lookup.py new file mode 100644 index 0000000..9de95ed --- /dev/null +++ b/tests/test_torch_zip_lookup.py @@ -0,0 +1,287 @@ +"""Model-free causal controls for the actual carried torch ZIP reader. + +Compiles unchanged zip.c/pickle_io.cpp and the parent/carried torch reader. +Fixtures are tiny tensor-only pickle byte streams, never Python-unpickled. +""" +import json +import os +from pathlib import Path +import shlex +import shutil +import struct +import subprocess +import sys +import tempfile +import unittest +from unittest import mock +import warnings +import zipfile + +REPO = Path(__file__).resolve().parent.parent +UPSTREAM = REPO / "third_party/stable-diffusion.cpp" +READER = Path("src/model_io/torch_zip_io.cpp") +PATCH = REPO / "patches/torch_zip_entry_lookup.patch" +NATIVE = REPO / "tests/native/torch_zip_lookup_test.cpp" +CXX = os.environ.get("CXX") or shutil.which("c++") +CC = os.environ.get("CC") or shutil.which("cc") +GIT = shutil.which("git") + + +def compile_source(command): + result = subprocess.run(command, capture_output=True, text=True, timeout=120) + if result.returncode: + raise AssertionError( + f"Compiler exited {result.returncode}: {shlex.join(command)}\n" + f"stdout (last 8192 characters):\n{result.stdout[-8192:]}\n" + f"stderr (last 8192 characters):\n{result.stderr[-8192:]}") + + +def string(value): + data = value.encode("utf-8") + return b"X" + struct.pack(" 512) name.resize(512);"), + "missing-bounds": ("if (tensor_nbytes < 0 ||\n tensor_storage.offset > entry_size ||\n static_cast(tensor_nbytes) > entry_size - tensor_storage.offset)", "if (false)"), + "missing-free": (" free(pkl_data);", " (void)pkl_data;"), + } + cases = { + "last-match": [("archive/data.pkl", pickle_data([tensor("a", "0"), tensor("b", "1"), tensor("again", "0")])), + ("archive/data/0", bytes(4)), ("archive/data/0", b"x"), ("archive/data/1", bytes(4))], + "truncated-name": [("archive/data.pkl", pickle_data([tensor("a", "x" * 600)])), + ("archive/data/" + "x" * 600, bytes(4))], + "missing-bounds": [("archive/data.pkl", pickle_data([tensor("a", "0", elements=4, offset=4)])), + ("archive/data/0", bytes(16))], + "missing-free": [("archive/data.pkl", pickle_data([tensor("a", "missing")]))], + } + for name, (before, after) in mutants.items(): + with self.subTest(mutant=name): + self.assertIn(before, source) + mutated = self.root / (name + ".cpp") + mutated.write_text(source.replace(before, after)) + binary = self.build("mutant-" + name, mutated) + archive = self.archive(name, cases[name]) + baseline = self.probe(self.fixed, archive)[0] + result = subprocess.run([str(binary), "none", str(archive)], capture_output=True, + text=True, timeout=20) + self.assertEqual(result.returncode, 0, result.stderr) + observed = json.loads(result.stdout) + fields = ("ok", "error", "inventory", "live_pickle_buffers") + self.assertNotEqual({k: baseline[k] for k in fields}, {k: observed[k] for k in fields}) + + +class CompilerDiagnosticsTest(unittest.TestCase): + def test_failed_compile_exposes_bounded_stdout_and_stderr(self): + failed = subprocess.CompletedProcess(["cc", "source.c"], 1, + "x" * 20000 + "stdout-marker", + "y" * 20000 + "stderr-marker") + with mock.patch.object(subprocess, "run", return_value=failed): + with self.assertRaises(AssertionError) as caught: + compile_source(["cc", "source.c"]) + message = str(caught.exception) + self.assertIn("Compiler exited 1: cc source.c", message) + self.assertIn("stdout-marker", message) + self.assertIn("stderr-marker", message) + self.assertLess(len(message), 17000) + + +if __name__ == "__main__": + unittest.main()