diff --git a/AGENTS.md b/AGENTS.md index ebebee7..eacf879 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -8,7 +8,8 @@ Guidance for coding agents working in `stable-diffusion-native`. consumed by `llamadart`. The Dart API and model handling live in `llamadart`. - Upstream is the `third_party/stable-diffusion.cpp` submodule. Update the submodule pin instead of patching vendored sources; report upstream bugs - there. + there. Never edit the submodule. A fix the pin lacks is carried only as a + `patches/` entry (`patches/README.md`) with a test that fails without it. - What upstream's C API lacks goes in `src/`, which the top-level `CMakeLists.txt` compiles into the same library. A wrapper export is an `sd_dart_*` function declared with `SD_API` in `src/sd_dart_wrapper.h`, with diff --git a/CMakeLists.txt b/CMakeLists.txt index 5823873..4182c9d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,4 +1,4 @@ -cmake_minimum_required(VERSION 3.13) +cmake_minimum_required(VERSION 3.18) project(stable_diffusion_native LANGUAGES C CXX) set(SD_NATIVE_UPSTREAM_DIR "${CMAKE_CURRENT_SOURCE_DIR}/third_party/stable-diffusion.cpp") @@ -17,6 +17,66 @@ define_property(DIRECTORY PROPERTY EP_PREFIX INHERITED set_property(DIRECTORY PROPERTY EP_PREFIX "${CMAKE_BINARY_DIR}/ep") add_subdirectory("${SD_NATIVE_UPSTREAM_DIR}" sd) +# Fixes the pinned upstream lacks (patches/README.md). Each patch is applied to +# build-tree copies of the upstream sources it changes, and the library +# compiles the copies, so the submodule stays pristine. Configuring fails when +# a patch no longer applies or changes nothing. +find_package(Git REQUIRED) +function(sd_native_carry_patch patch) + get_filename_component(name "${patch}" NAME_WE) + set(patched_root "${CMAKE_BINARY_DIR}/patched") + set(patched_dir "${patched_root}/${name}") + file(REMOVE_RECURSE "${patched_dir}") + set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/patches/${patch}") + get_target_property(sources stable-diffusion SOURCES) + foreach(path IN LISTS ARGN) + set(upstream_file "${SD_NATIVE_UPSTREAM_DIR}/${path}") + if(NOT "${upstream_file}" IN_LIST sources) + message(FATAL_ERROR "patches/${patch}: ${path} is not a stable-diffusion source") + endif() + file(READ "${upstream_file}" text) + # A Windows checkout can convert upstream's line endings; patches use LF. + string(REPLACE "\r\n" "\n" text "${text}") + file(WRITE "${patched_dir}/${path}" "${text}") + endforeach() + + # The build tree can sit inside this repository's checkout, where git apply + # would skip every path outside it and still succeed. + get_filename_component(ceiling "${patched_root}" REALPATH) + execute_process( + COMMAND "${CMAKE_COMMAND}" -E env "GIT_CEILING_DIRECTORIES=${ceiling}" + "${GIT_EXECUTABLE}" apply --verbose "${CMAKE_CURRENT_SOURCE_DIR}/patches/${patch}" + WORKING_DIRECTORY "${patched_dir}" + RESULT_VARIABLE result + OUTPUT_VARIABLE output + ERROR_VARIABLE output) + if(NOT result EQUAL 0) + message(FATAL_ERROR "patches/${patch} does not apply to the pinned upstream:\n${output}") + endif() + + foreach(path IN LISTS ARGN) + set(upstream_file "${SD_NATIVE_UPSTREAM_DIR}/${path}") + set(patched_file "${patched_dir}/${path}") + file(READ "${upstream_file}" original) + string(REPLACE "\r\n" "\n" original "${original}") + file(READ "${patched_file}" patched) + if(original STREQUAL patched) + message(FATAL_ERROR "patches/${patch} left ${path} unchanged:\n${output}") + endif() + list(REMOVE_ITEM sources "${upstream_file}") + list(APPEND sources "${patched_file}") + # Quoted includes in the copy must still find the upstream files next to it. + get_filename_component(upstream_file_dir "${upstream_file}" DIRECTORY) + set_source_files_properties("${patched_file}" TARGET_DIRECTORY stable-diffusion + PROPERTIES INCLUDE_DIRECTORIES "${upstream_file_dir}") + endforeach() + set_property(TARGET stable-diffusion PROPERTY SOURCES ${sources}) + message(STATUS "Carrying patches/${patch}") +endfunction() + +sd_native_carry_patch(pickle_line_span.patch src/model_io/pickle_io.cpp) + # Wrapper exports are compiled into the upstream library target, so they share # its SD_API export macro and reach upstream without patching it. target_sources(stable-diffusion PRIVATE diff --git a/README.md b/README.md index 869d545..79bba96 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,8 @@ consumed by [llamadart](https://github.com/leehack/llamadart) and other wrappers. It stays independent from `llamadart` so the artifacts can be reused elsewhere. -- Pin an upstream stable-diffusion.cpp commit as a submodule. +- Pin an upstream stable-diffusion.cpp commit as a submodule, plus the few + [carried fixes](patches/README.md) the pin lacks. - Build one shared library per platform with upstream's C API (`stable-diffusion.h`) as the FFI boundary. - Export only that API and the [wrapper API](#wrapper-api) this repository diff --git a/patches/README.md b/patches/README.md new file mode 100644 index 0000000..d91826f --- /dev/null +++ b/patches/README.md @@ -0,0 +1,15 @@ +# Carried upstream patches + +Fixes the pinned stable-diffusion.cpp lacks, carried until a pin includes them. +`CMakeLists.txt` applies each patch to build-tree copies of the upstream +sources it names and compiles the copies, so the submodule stays pristine. +Configuring fails when a patch no longer applies or changes nothing; then drop +or rebase it. Build info and the release manifest list the patches next to the +upstream commit. + +A patch here needs a regression test that fails against the unpatched pin and +an upstream issue or pull request to retire it. + +| Patch | Sources | Fixes | Test | +| --- | --- | --- | --- | +| `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` | diff --git a/patches/pickle_line_span.patch b/patches/pickle_line_span.patch new file mode 100644 index 0000000..b713cab --- /dev/null +++ b/patches/pickle_line_span.patch @@ -0,0 +1,134 @@ +diff --git a/src/model_io/pickle_io.cpp b/src/model_io/pickle_io.cpp +--- a/src/model_io/pickle_io.cpp ++++ b/src/model_io/pickle_io.cpp +@@ -1,5 +1,7 @@ + #include "pickle_io.h" + ++#include ++#include + #include + #include + #include +@@ -115,6 +117,13 @@ static void set_error(std::string* error, const std::string& message) { + } + } + ++// A legacy checkpoint passes the rest of the file, tensor storage included, so ++// the span can exceed INT_MAX. No pickle line is that long. ++static int find_line_end(const uint8_t* p, const uint8_t* end) { ++ const ptrdiff_t max_len = std::numeric_limits::max(); ++ return find_char(p, (int)std::min(end - p, max_len), '\n'); ++} ++ + bool skip_pickle_object(const uint8_t* buffer, size_t buffer_size, size_t* object_size) { + const uint8_t* p = buffer; + const uint8_t* end = buffer + buffer_size; +@@ -148,7 +157,7 @@ bool skip_pickle_object(const uint8_t* buffer, size_t buffer_size, size_t* objec + case 'F': // FLOAT = b'F' # push decimal float line + case 'S': // STRING = b'S' # push quoted string line + case 'V': { // UNICODE = b'V' # push raw-unicode string line +- int len = find_char(p, (int)(end - p), '\n'); ++ int len = find_line_end(p, end); + if (len < 0) { + return false; + } +@@ -206,7 +215,7 @@ bool skip_pickle_object(const uint8_t* buffer, size_t buffer_size, size_t* objec + p += 1 + p[0]; + break; + case 'P': { // PERSID = b'P' # persistent id, newline-terminated +- int len = find_char(p, (int)(end - p), '\n'); ++ int len = find_line_end(p, end); + if (len < 0) { + return false; + } +@@ -216,12 +225,12 @@ bool skip_pickle_object(const uint8_t* buffer, size_t buffer_size, size_t* objec + p += 8; + break; + case 'c': { // GLOBAL = b'c' # push module/name global reference +- int len = find_char(p, (int)(end - p), '\n'); ++ int len = find_line_end(p, end); + if (len < 0) { + return false; + } + p += len + 1; +- len = find_char(p, (int)(end - p), '\n'); ++ len = find_line_end(p, end); + if (len < 0) { + return false; + } +@@ -261,12 +270,12 @@ bool skip_pickle_object(const uint8_t* buffer, size_t buffer_size, size_t* objec + case 'o': // OBJ = b'o' # build class instance from mark + break; + case 'i': { // INST = b'i' # build class instance from module/name +- int len = find_char(p, (int)(end - p), '\n'); ++ int len = find_line_end(p, end); + if (len < 0) { + return false; + } + p += len + 1; +- len = find_char(p, (int)(end - p), '\n'); ++ len = find_line_end(p, end); + if (len < 0) { + return false; + } +@@ -680,7 +689,7 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer, + p += 4; + break; + case 'I': { // INT = b'I' # push decimal integer line +- int len = find_char(p, (int)(end - p), '\n'); ++ int len = find_line_end(p, end); + if (len < 0) { + return false; + } +@@ -695,7 +704,7 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer, + } + } break; + case 'L': { // LONG = b'L' # push decimal long integer line +- int len = find_char(p, (int)(end - p), '\n'); ++ int len = find_line_end(p, end); + if (len < 0) { + return false; + } +@@ -707,7 +716,7 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer, + stack.push_back(make_int_value(std::strtoll(s.c_str(), nullptr, 10))); + } break; + case 'F': { // FLOAT = b'F' # push decimal float line +- int len = find_char(p, (int)(end - p), '\n'); ++ int len = find_line_end(p, end); + if (len < 0) { + return false; + } +@@ -799,7 +808,7 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer, + p += len; + } break; + case 'S': { // STRING = b'S' # push quoted string line +- int len = find_char(p, (int)(end - p), '\n'); ++ int len = find_line_end(p, end); + if (len < 0) { + return false; + } +@@ -811,7 +820,7 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer, + stack.push_back(make_string_value(s)); + } break; + case 'V': { // UNICODE = b'V' # push raw-unicode string line +- int len = find_char(p, (int)(end - p), '\n'); ++ int len = find_line_end(p, end); + if (len < 0) { + return false; + } +@@ -819,13 +828,13 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer, + p += len + 1; + } break; + case 'c': { // GLOBAL = b'c' # push module/name global reference +- int len = find_char(p, (int)(end - p), '\n'); ++ int len = find_line_end(p, end); + if (len < 0) { + return false; + } + std::string module((const char*)p, len); + p += len + 1; +- len = find_char(p, (int)(end - p), '\n'); ++ len = find_line_end(p, end); + if (len < 0) { + return false; + } diff --git a/tests/native/pickle_span_test.cpp b/tests/native/pickle_span_test.cpp new file mode 100644 index 0000000..ee87a8f --- /dev/null +++ b/tests/native/pickle_span_test.cpp @@ -0,0 +1,161 @@ +// Runs upstream's pickle scanner, compiled as the library compiles it, on the +// spans read_torch_legacy_file gives it. That reader passes the rest of the +// whole checkpoint, tensor storage included, so a small pickle can sit in a +// span longer than INT_MAX. +// +// Modes: +// bounded each pickle in a buffer of exactly its size +// large malformed pickles at the start of spans above INT_MAX +// large-ok valid pickles at the start of spans above INT_MAX +// +// Large spans are mapped read-only, so they cost no memory; only the first +// page is writable and holds the pickle. + +#include "model_io/pickle_io.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#define CHECK(condition) \ + do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: CHECK(%s) failed\n", __FILE__, __LINE__, #condition); \ + std::_Exit(1); \ + } \ + } while (0) + +// Stands in for upstream's util.cpp; pickle_io.cpp calls it only to format errors. +std::string sd_format(const char* format, ...) { + char buffer[1024]; + va_list args; + va_start(args, format); + std::vsnprintf(buffer, sizeof(buffer), format, args); + va_end(args); + return buffer; +} + +namespace { + +struct Pickle { + const char* name; + std::string bytes; +}; + +// Keeps embedded NULs; drops the literal's terminator. +template +std::string bytes(const char (&data)[N]) { + return std::string(data, N - 1); +} + +// Every opcode whose operand skip_pickle_object scans to a newline, plus a +// length-prefixed string, which does not depend on the span. +std::vector valid_pickles() { + return { + {"GLOBAL", bytes("\x80\x02" "ccollections\nOrderedDict\n)R.")}, + {"INST", bytes("\x80\x02(icollections\nOrderedDict\n.")}, + {"PERSID", bytes("\x80\x02P0\n.")}, + {"UNICODE", bytes("\x80\x02Vweight\n.")}, + {"BINUNICODE", bytes("\x80\x02X\x06\x00\x00\x00weight.")}, + }; +} + +std::vector malformed_pickles() { + return { + {"GLOBAL without a name line", bytes("\x80\x02" "ccollections\nOrderedDict")}, + {"GLOBAL without a module line", bytes("\x80\x02" "ccollections")}, + {"PERSID without a newline", bytes("\x80\x02P0")}, + {"BINUNICODE past the end", bytes("\x80\x02X\xff\x00\x00\x00weight.")}, + {"unknown opcode", bytes("\x80\x02\x00.")}, + {"no STOP", bytes("\x80\x02)")}, + }; +} + +bool skip(const uint8_t* data, size_t size, size_t* consumed) { + *consumed = 0; + return skip_pickle_object(data, size, consumed); +} + +void check_bounded() { + for (const Pickle& pickle : valid_pickles()) { + // An exact-size heap copy lets a sanitizer catch reads past the end. + std::vector data(pickle.bytes.begin(), pickle.bytes.end()); + size_t consumed = 0; + std::printf("bounded %s\n", pickle.name); + CHECK(skip(data.data(), data.size(), &consumed)); + CHECK(consumed == data.size()); + } + for (const Pickle& pickle : malformed_pickles()) { + std::vector data(pickle.bytes.begin(), pickle.bytes.end()); + size_t consumed = 0; + std::printf("bounded %s\n", pickle.name); + CHECK(!skip(data.data(), data.size(), &consumed)); + } +} + +// Above INT_MAX, and above UINT32_MAX, where narrowing to int keeps only the +// low 32 bits. +const size_t kLargeSpans[] = {(size_t(1) << 31) + 4096, (size_t(1) << 32) + 16}; + +uint8_t* map_span(size_t span, const std::string& pickle) { + const size_t page = (size_t)sysconf(_SC_PAGESIZE); + void* mapped = mmap(nullptr, span, PROT_READ, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + CHECK(mapped != MAP_FAILED); + CHECK(mprotect(mapped, page, PROT_READ | PROT_WRITE) == 0); + CHECK(pickle.size() <= page); + std::memcpy(mapped, pickle.data(), pickle.size()); + return static_cast(mapped); +} + +void check_large(const std::vector& pickles, bool valid) { + for (size_t span : kLargeSpans) { + for (const Pickle& pickle : pickles) { + uint8_t* data = map_span(span, pickle.bytes); + size_t consumed = 0; + std::printf("span %zu %s\n", span, pickle.name); + std::fflush(stdout); + const bool skipped = skip(data, span, &consumed); + CHECK(munmap(data, span) == 0); + if (valid) { + CHECK(skipped); + CHECK(consumed == pickle.bytes.size()); + } else { + CHECK(!skipped); + } + } + } +} + +// The rest of a malformed pickle in a large span is zeros, which the scanner +// rejects as an unknown opcode; a scan for a missing newline would read up to +// INT_MAX bytes of them, so those cases stay bounded only. +std::vector malformed_large_pickles() { + return { + {"no STOP", bytes("\x80\x02)")}, + {"BINUNICODE past the pickle", bytes("\x80\x02X\x00\x00\x00\x80")}, + }; +} + +} // namespace + +int main(int argc, char** argv) { + const std::string mode = argc > 1 ? argv[1] : ""; + if (mode == "bounded") { + check_bounded(); + } else if (mode == "large") { + check_large(malformed_large_pickles(), false); + } else if (mode == "large-ok") { + check_large(valid_pickles(), true); + } else { + std::fprintf(stderr, "usage: %s bounded|large|large-ok\n", argv[0]); + return 2; + } + return 0; +} diff --git a/tests/test_pickle_span.py b/tests/test_pickle_span.py new file mode 100644 index 0000000..b5f8da6 --- /dev/null +++ b/tests/test_pickle_span.py @@ -0,0 +1,108 @@ +"""Builds and runs `tests/native/pickle_span_test.cpp`. + +The test compiles upstream's `pickle_io.cpp` from the pinned submodule with +`patches/pickle_line_span.patch` applied, as the library does, so it needs the +submodules' sources but no built runtime and no model. Spans above INT_MAX are +reserved read-only and cost no memory. +""" + +import os +import shutil +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +UPSTREAM = REPO_ROOT / "third_party" / "stable-diffusion.cpp" +SCANNER = Path("src") / "model_io" / "pickle_io.cpp" +PATCH = REPO_ROOT / "patches" / "pickle_line_span.patch" +TEST_SOURCE = REPO_ROOT / "tests" / "native" / "pickle_span_test.cpp" +INCLUDE_DIRS = [UPSTREAM / "src", UPSTREAM / "src" / "model_io", UPSTREAM / "include", + UPSTREAM / "ggml" / "include"] +COMPILER = os.environ.get("CXX") or shutil.which("c++") +GIT = shutil.which("git") +ADDRESS_SANITIZER = "-fsanitize=address" + + +@unittest.skipIf(sys.platform == "win32", "uses a GCC or Clang command line and mmap") +class PickleSpanTest(unittest.TestCase): + def require(self, condition: bool, reason: str) -> None: + """Skips locally, but fails in CI, where a skip would hide the test.""" + if condition: + return + if os.environ.get("CI"): + self.fail(reason) + self.skipTest(reason) + + def setUp(self) -> None: + self.require(COMPILER is not None, "no C++ compiler; set CXX") + self.require(GIT is not None, "no git") + self.require((INCLUDE_DIRS[3] / "ggml.h").is_file() and (UPSTREAM / SCANNER).is_file(), + "third_party/stable-diffusion.cpp is not checked out; run " + "`git submodule update --init --recursive`") + + def patched_scanner(self, directory: Path) -> Path: + """Applies the carried patch as CMakeLists.txt does.""" + root = directory / "patched" + scanner = root / SCANNER + scanner.parent.mkdir(parents=True) + scanner.write_bytes((UPSTREAM / SCANNER).read_bytes().replace(b"\r\n", b"\n")) + env = {**os.environ, "GIT_CEILING_DIRECTORIES": str(directory.resolve())} + subprocess.run([GIT, "apply", str(PATCH)], cwd=root, env=env, check=True) + self.assertNotEqual((UPSTREAM / SCANNER).read_bytes(), scanner.read_bytes()) + return scanner + + def run_mode(self, mode: str, *flags: str, + patched: bool = True) -> subprocess.CompletedProcess: + includes = [f"-I{directory}" for directory in INCLUDE_DIRS] + with tempfile.TemporaryDirectory() as directory: + directory = Path(directory) + scanner = self.patched_scanner(directory) if patched else UPSTREAM / SCANNER + binary = directory / "pickle_span_test" + subprocess.run([COMPILER, "-std=c++17", "-O1", "-g", *flags, *includes, + str(scanner), str(TEST_SOURCE), "-o", str(binary)], check=True) + return subprocess.run([str(binary), mode], capture_output=True, text=True, + timeout=600) + + def assert_passes(self, mode: str, *flags: str) -> None: + result = self.run_mode(mode, *flags) + self.assertEqual(0, result.returncode, result.stdout + result.stderr) + + def supports(self, flag: str) -> bool: + with tempfile.TemporaryDirectory() as directory: + source = Path(directory) / "empty.cpp" + source.write_text("int main() { return 0; }\n") + binary = Path(directory) / "empty" + built = subprocess.run([COMPILER, flag, str(source), "-o", str(binary)], + capture_output=True) + return built.returncode == 0 and subprocess.run( + [str(binary)], capture_output=True).returncode == 0 + + def test_skips_valid_and_rejects_malformed_bounded_pickles(self) -> None: + self.assert_passes("bounded") + + def test_bounded_pickles_stay_in_bounds_under_address_sanitizer(self) -> None: + self.require(self.supports(ADDRESS_SANITIZER), + f"{COMPILER} cannot build or run {ADDRESS_SANITIZER} binaries") + self.assert_passes("bounded", ADDRESS_SANITIZER) + + def test_rejects_malformed_pickles_in_spans_above_int_max(self) -> None: + self.assert_passes("large") + + def test_skips_valid_pickles_in_spans_above_int_max(self) -> None: + self.assert_passes("large-ok") + + # The pinned upstream narrows the remaining span to int before scanning for + # a newline (#23). Once a pin no longer does, the carried patch can go. + def test_pinned_upstream_still_needs_the_patch(self) -> None: + result = self.run_mode("large-ok", patched=False) + self.assertNotEqual( + 0, result.returncode, + "the pinned upstream now skips valid pickles in spans above INT_MAX; " + "remove patches/pickle_line_span.patch") + + +if __name__ == "__main__": + unittest.main() diff --git a/tools/build.py b/tools/build.py index 3009941..f88a9c4 100644 --- a/tools/build.py +++ b/tools/build.py @@ -208,6 +208,14 @@ def upstream_commit() -> str: ).strip() +def upstream_info() -> dict: + return { + "repository": "leejet/stable-diffusion.cpp", + "commit": upstream_commit(), + "patches": sorted(p.name for p in (REPO_ROOT / "patches").glob("*.patch")), + } + + def build(target: Target, jobs: int) -> Path: if not HEADER.is_file(): fail("third_party/stable-diffusion.cpp is missing; run " @@ -240,10 +248,7 @@ def build(target: Target, jobs: int) -> Path: "os": target.os, "library": target.library, "accelerators": list(target.accelerators), - "upstream": { - "repository": "leejet/stable-diffusion.cpp", - "commit": upstream_commit(), - }, + "upstream": upstream_info(), "cmakeArgs": [a for a in cmake_args if not a.startswith("-DCMAKE_SHARED_LINKER_FLAGS")], "exportedSymbols": api_symbols(*HEADERS), } diff --git a/tools/package_release.py b/tools/package_release.py index 1f788cd..88f1436 100644 --- a/tools/package_release.py +++ b/tools/package_release.py @@ -20,7 +20,7 @@ from pathlib import Path import apple_xcframework -from build import BIN_ROOT, HEADERS, REPO_ROOT, TARGETS, UPSTREAM_DIR, upstream_commit +from build import BIN_ROOT, HEADERS, REPO_ROOT, TARGETS, UPSTREAM_DIR, upstream_info PACKAGE = "stable-diffusion-native" TAG_PATTERN = re.compile(r"^v\d+\.\d+\.\d+(-[1-9]\d*)?$") @@ -120,10 +120,7 @@ def main() -> None: "schemaVersion": 1, "package": PACKAGE, "tag": args.tag, - "upstream": { - "repository": "leejet/stable-diffusion.cpp", - "commit": upstream_commit(), - }, + "upstream": upstream_info(), "artifacts": artifacts, } (args.dist / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")