diff --git a/docs/backend.md b/docs/backend.md index c936b2232..eff576d1e 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -142,6 +142,47 @@ GPUs with the layer/row split mechanism (`--split-mode` selects which, layer by default). Components that fit nowhere fall back to the CPU. If a VAE decode still runs out of memory, tiling is enabled and the decode retried once. +## Measured fitting (`sd-fit-params`) + +`sd-fit-params` is a standalone tool that derives the same kind of placement, +but from *measured* memory instead of auto-fit's fixed compute reserves. It +runs the real generation pipeline in a metadata-only dry run: every module's +compute graphs are built for the requested generation request and their +compute buffer sizes are measured without allocating ggml weight/compute +buffers or reading weight data. Shaped host tensors are materialized between +graph builds, and allocation failures are returned as fit errors. Because +compute memory depends on the generation parameters, they are inputs to the +tool, and the printed arguments are valid for workloads up to that size. + +Logs go to stderr, the fitted arguments go to stdout: + +```shell +sd-fit-params -m model.gguf -W 1024 -H 1024 | tee args.txt +cat args.txt | xargs sd-cli -m model.gguf -p "a cat" -W 1024 -H 1024 +``` + +`--fit-print` prints the measured per-device / per-module memory table instead +of arguments. Budgets reuse `--max-vram`, with measured fitting retaining a +512 MiB safety margin after applying automatic or explicit per-device limits. +If the default placement already fits, nothing needs to change and the tool +prints an empty line. Explicit `--backend` / `--params-backend` assignments are +rejected because measured fitting derives a new placement rather than validating +an existing one. + +The same measurement is available to library users through `sd_fit_params()` +in `stable-diffusion.h`, which takes the context params plus an +`sd_fit_workload_t` and returns the derived specs and a report. The workload can +be the simple scalar fields (`prompt`, `width`, `height`, `video_frames`) or one +complete representative request through `image_gen_params` or +`video_gen_params`. Use the full request form when conditioning, LoRAs, hires, +cache settings, image/video/audio inputs, or VAE tiling settings materially +affect the graph being measured. Result strings are owned by `sd_fit_result_t` +and must be released with `sd_fit_result_free()`. Persistent cache buffers and +the concurrent diffusion/ControlNet compute phase are included in the measured +peak. CPU fallbacks are accepted only when the projected CPU parameters and +compute buffers fit currently available host memory. Callers must leave +`backend` and `params_backend` unset while requesting a fit. + ## Modules | Module | Purpose | Accepted names | diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index b347f1785..e0dbb75cf 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1,6 +1,7 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR}) add_subdirectory(cli) +add_subdirectory(fit-params) add_subdirectory(server) if(NOT SD_BUILD_SHARED_LIBS) # ABot validation harnesses link engine internals not exported by the shared library. diff --git a/examples/fit-params/CMakeLists.txt b/examples/fit-params/CMakeLists.txt new file mode 100644 index 000000000..2fea96f59 --- /dev/null +++ b/examples/fit-params/CMakeLists.txt @@ -0,0 +1,26 @@ +set(TARGET sd-fit-params) + +add_executable(${TARGET} + ../common/common.cpp + ../common/log.cpp + ../common/media_io.cpp + main.cpp +) +if(APPLE) + sd_set_macos_rpaths(${TARGET}) +endif() +target_include_directories(${TARGET} PRIVATE + "${CMAKE_CURRENT_SOURCE_DIR}/.." + "${PROJECT_SOURCE_DIR}/src" +) +install(TARGETS ${TARGET} RUNTIME) +target_link_libraries(${TARGET} PRIVATE stable-diffusion zip ${CMAKE_THREAD_LIBS_INIT}) +if(SD_WEBP) + target_compile_definitions(${TARGET} PRIVATE SD_USE_WEBP) + target_link_libraries(${TARGET} PRIVATE webp libwebpmux) +endif() +if(SD_WEBM) + target_compile_definitions(${TARGET} PRIVATE SD_USE_WEBM) + target_link_libraries(${TARGET} PRIVATE webm) +endif() +target_compile_features(${TARGET} PUBLIC c_std_11 cxx_std_17) diff --git a/examples/fit-params/README.md b/examples/fit-params/README.md new file mode 100644 index 000000000..2a16be97d --- /dev/null +++ b/examples/fit-params/README.md @@ -0,0 +1,190 @@ +# sd-fit-params + +`sd-fit-params` computes the CLI arguments that make a model fit into free +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. +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 +memory of every GPU device. A 512 MiB safety margin is retained after applying +either detected or explicit `--max-vram` limits, and the resulting placement is +printed to stdout as `--backend` / `--params-backend` / `--vae-tiling` / +`--stream-layers` arguments. + +Because compute memory depends on the generation parameters, pass the same +width/height (and video frames) you intend to generate with. Example usage: + +``` bash +# First, run sd-fit-params and store the results in a file: +> ./build/bin/sd-fit-params -m sd_v1-5.gguf -W 1024 -H 1024 --max-vram 4 | tee args.txt +[INFO ] fit_params.cpp:93 - fit-params: measured memory plan +[INFO ] fit_params.cpp:93 - devices: +[INFO ] fit_params.cpp:93 - MTL0 Apple M4 free 12123 MiB, budget 3584 MiB +[INFO ] fit_params.cpp:93 - modules (measured for this workload): +[INFO ] fit_params.cpp:93 - diffusion params 1398 MiB, compute 8360 MiB +[INFO ] fit_params.cpp:93 - te params 125 MiB, compute 1 MiB +[INFO ] fit_params.cpp:93 - vae params 159 MiB, compute 6656 MiB (tiled 416 MiB) +[INFO ] fit_params.cpp:93 - placement (time-share: params load per phase and free after): +[INFO ] fit_params.cpp:93 - diffusion -> cpu +[INFO ] fit_params.cpp:93 - te -> MTL0, params on disk +[INFO ] fit_params.cpp:93 - vae -> MTL0, params on disk, vae tiling +[INFO ] stable-diffusion.cpp - fit-params: fitting params to free memory took 0.26s +printing fitted CLI arguments to stdout... +--backend "diffusion=cpu,te=MTL0,vae=MTL0" --params-backend "te=disk,vae=disk" --vae-tiling + +# Next, use those results for sd-cli: +> cat args.txt | xargs ./build/bin/sd-cli -m sd_v1-5.gguf -p "a cat" -W 1024 -H 1024 +``` + +Useful flags: + +- `-W` / `-H` / `--video-frames`: the workload the fit must accommodate +- `--max-vram ` or `--max-vram cuda0=8,cuda1=14`: per-device limits. + The planner retains 512 MiB of headroom after applying the limit. Positive + values cap the graph-splitting budget, negative values use auto budget + detection, and `0` disables graph splitting. +- `--fit-print`: print the measured memory table to stdout instead of arguments +- `-p`: representative prompt (token count affects text encoder memory) +- generation inputs including init/control/reference images, LoRAs, and hires + options flow into the measurement as they would into a real run; explicit + `--backend` / `--params-backend` placement is rejected + +## Planner order + +The planner tries the fastest and most resident placements first, then falls +back to progressively lower-VRAM choices. Device budgets retain a 512 MiB +margin after both automatic and explicit `--max-vram` limits. If no GPU device +is available, the tool verifies that the workload fits available host memory +before keeping the default CPU backend. + +The checks run in this order: + +1. Default placement: put every module on the first GPU. This succeeds when + the sum of all module parameters plus the peak measured compute phase fits + that device budget. Diffusion and ControlNet buffers are added because both + remain live during denoising; other sequential module buffers use their + maximum. If it succeeds, the tool prints an empty line + because no extra CLI arguments are needed. +2. Resident multi-device placement: sort modules by parameter size, largest + first, and place each module on one GPU while keeping all parameters resident. + For each GPU, resident parameters accumulate and sequential compute buffers + use their maximum. Diffusion and ControlNet compute buffers are summed when + assigned to the same GPU. +3. Resident VAE tiling: while trying the resident plan, if a module has a + measured tiled compute size and full-resolution compute does not fit, retry + that module with tiled compute. This currently applies to VAE measurements + and emits `--vae-tiling`. +4. Time-share single-device placement: if resident placement fails, plan each + module as a separate phase. A module can run on a GPU with + `--params-backend =disk` when its parameters plus its compute buffer + fit one device budget. +5. Time-share VAE tiling: if the non-tiled time-share check fails and the module + has a tiled compute measurement, retry with the tiled compute size and emit + `--vae-tiling` if it fits. +6. Multi-GPU split: if the module is splittable and more than one GPU exists, + split its parameters across all GPUs when the sum of each device budget minus + that module's compute buffer can hold the module parameters. The emitted + backend uses `&`, for example `diffusion=CUDA0&CUDA1`, and parameters are + loaded per phase from disk. +7. Diffusion CPU params plus layer streaming: if split placement still does not + fit, and the module is a splittable diffusion module, choose the GPU with the + largest graph-splitting budget and keep diffusion parameters in CPU RAM while + streaming layers to the runtime GPU. This emits + `--params-backend diffusion=cpu`, preserves the original `--max-vram`, and + adds `--stream-layers`. This fallback is only considered when graph splitting + is enabled by a positive or negative `--max-vram`; `--max-vram 0` disables it. +8. CPU runtime fallback: if none of the GPU options above fit, put the module on + the CPU runtime backend. The planner returns `SD_FIT_FAILURE` if the CPU + parameters and compute phases exceed currently available host memory. + +If the default placement already fits, the tool prints an empty line and +reports that no changes are needed. Explicit `--backend` / `--params-backend` +assignments are rejected because the tool derives placement rather than +validating an existing assignment. + +See `docs/backend.md` for the placement spec syntax and the heuristic +`--auto-fit` alternative built into `sd-cli`. + +## Library API + +Library callers can use the same measured fitting through `sd_fit_params()` in +`stable-diffusion.h`. Start from initialized context params and workload params, +then free the result when done: + +```c +sd_ctx_params_t ctx; +sd_ctx_params_init(&ctx); +ctx.diffusion_model_path = "/models/model.gguf"; +ctx.max_vram = "8"; + +sd_fit_workload_t workload; +sd_fit_workload_init(&workload); +workload.prompt = "a cat"; +workload.width = 1024; +workload.height = 1024; +workload.video_frames = 1; + +sd_fit_result_t result; +enum sd_fit_status_t status = sd_fit_params(&ctx, &workload, &result); +if (status == SD_FIT_SUCCESS && result.changed) { + printf("backend=%s\n", result.backend ? result.backend : ""); + printf("params_backend=%s\n", result.params_backend ? result.params_backend : ""); + printf("vae_tiling=%d\n", result.vae_tiling); + printf("stream_layers=%d\n", result.stream_layers); +} +sd_fit_result_free(&result); +``` + +For simple text-to-image or text-to-video fitting, the scalar workload fields +are enough. `sd_fit_workload_init()` defaults to a 512x512 image workload +(`video_frames = 1`) and default VAE tiling params. + +For the most accurate plan, pass a full representative request: + +```c +sd_img_gen_params_t image_request; +sd_img_gen_params_init(&image_request); +image_request.prompt = "a cat"; +image_request.width = 1024; +image_request.height = 1024; +image_request.batch_count = 1; +image_request.vae_tiling_params = workload.vae_tiling_params; + +workload.image_gen_params = &image_request; +``` + +Use `workload.video_gen_params` with `sd_vid_gen_params_t` for video. Set at +most one of `image_gen_params` and `video_gen_params`; setting both returns +`SD_FIT_ERROR`. When a full request is present, it supplies conditioning, +LoRAs, hires/cache options, image/video/audio inputs, VAE tiling settings, and +other generation fields. The scalar workload fields remain as a fallback for +callers that only need a basic request. A video-only model selects the video +measurement pipeline even if an image request was supplied; shared request +fields are promoted to a video request. Supplying a video request for an +image-only model returns `SD_FIT_ERROR`. + +`sd_fit_params()` returns: + +- `SD_FIT_SUCCESS`: a placement was found, or the current/default placement + already fits. +- `SD_FIT_FAILURE`: no placement was projected to fit, or `ctx.backend` / + `ctx.params_backend` was already set by the caller. +- `SD_FIT_ERROR`: invalid inputs or a hard measurement error such as an + unreadable model. + +`sd_fit_result_t` owns `backend`, `params_backend`, and `report`; always call +`sd_fit_result_free()`. If `result.changed` is false, the current/default +placement already fits and the placement strings are null. If +`result.stream_layers` is true, preserve the caller's nonzero `ctx.max_vram` +when applying the result and also enable `--stream-layers`; the max-VRAM value +is not duplicated in the result. + +## Debugging the planner + +Set `SD_FIT_DEBUG_DEVICES` to plan against simulated devices instead of the +real ones, e.g. `SD_FIT_DEBUG_DEVICES="CUDA0:24,CUDA1:16"` (`name:free_gib`). +Measurement still runs on the real machine; only device enumeration is +replaced. Useful to preview placements for other hardware and to exercise +multi-device planning paths. Debug only: the emitted specs reference the +simulated device names. diff --git a/examples/fit-params/main.cpp b/examples/fit-params/main.cpp new file mode 100644 index 000000000..11df52533 --- /dev/null +++ b/examples/fit-params/main.cpp @@ -0,0 +1,347 @@ +#include + +#include +#include +#include +#include +#include + +#include "stable-diffusion.h" + +#include "common/common.h" +#include "common/media_io.h" + +namespace fs = std::filesystem; + +struct SDFitCliParams { + bool verbose = false; + bool color = false; + bool fit_print = false; + + ArgOptions get_options() { + ArgOptions options; + options.bool_options = { + {"-v", "--verbose", "print extra info", true, &verbose}, + {"", "--color", "colors the logging tags", true, &color}, + {"", "--fit-print", "print the measured memory report to stdout instead of fitted arguments", true, &fit_print}, + }; + return options; + } +}; + +static void print_usage(int argc, const char* argv[], const std::vector& options_list) { + fprintf(stderr, "usage: %s [arguments]\n", argv[0]); + fprintf(stderr, "\n"); + fprintf(stderr, "Measures the memory the model needs for the requested generation parameters using\n"); + fprintf(stderr, "metadata-only dry runs (no weight data is read), then prints the CLI arguments that\n"); + fprintf(stderr, "make it fit into free device memory. Logs go to stderr, arguments to stdout:\n"); + fprintf(stderr, "\n"); + fprintf(stderr, " %s -m model.gguf -W 1024 -H 1024 | tee args.txt\n", argv[0]); + fprintf(stderr, " cat args.txt | xargs sd-cli -m model.gguf -p \"a cat\" -W 1024 -H 1024\n"); + fprintf(stderr, "\n"); + fprintf(stderr, "arguments:\n"); + for (const auto& options : options_list) { + options.print(); + } +} + +// keep stdout clean for the fitted arguments +static void fit_log_cb(enum sd_log_level_t level, const char* log, void* data) { + SDFitCliParams* params = (SDFitCliParams*)data; + if (!params->verbose && level == SD_LOG_DEBUG) { + return; + } + const char* level_str[] = {"DEBUG", "INFO", "WARN", "ERROR"}; + fprintf(stderr, "[%-5s] %s", level >= 0 && level <= 3 ? level_str[level] : "?", SAFE_STR(log)); + fflush(stderr); +} + +static bool load_images_from_dir(const std::string& dir, + std::vector& images, + int expected_width, + int expected_height, + int max_image_num, + bool verbose) { + if (!fs::exists(dir) || !fs::is_directory(dir)) { + fprintf(stderr, "'%s' is not a valid directory\n", dir.c_str()); + return false; + } + + std::vector entries; + for (const auto& entry : fs::directory_iterator(dir)) { + if (entry.is_regular_file()) { + entries.push_back(entry); + } + } + std::sort(entries.begin(), entries.end(), [](const auto& a, const auto& b) { + return a.path().filename().string() < b.path().filename().string(); + }); + + for (const auto& entry : entries) { + std::string path = entry.path().string(); + std::string ext = entry.path().extension().string(); + std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + if (ext != ".jpg" && ext != ".jpeg" && ext != ".png" && ext != ".bmp" && ext != ".webp") { + continue; + } + if (verbose) { + fprintf(stderr, "load image %zu from '%s'\n", images.size(), path.c_str()); + } + int width = 0; + int height = 0; + uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, expected_width, expected_height); + if (image_buffer == nullptr) { + fprintf(stderr, "load image from '%s' failed\n", path.c_str()); + return false; + } + images.emplace_back(sd_image_t{static_cast(width), + static_cast(height), + 3, + image_buffer}); + if (max_image_num > 0 && static_cast(images.size()) >= max_image_num) { + break; + } + } + return true; +} + +static bool load_image_file(SDGenerationParams& params, + const std::string& path, + SDImageOwner& image, + bool resize_image = true, + int channels = 3) { + if (path.empty()) { + return true; + } + + const bool use_size = resize_image && params.width_and_height_are_set(); + if (!load_sd_image_from_file(image.put(), + path.c_str(), + use_size ? params.width : 0, + use_size ? params.height : 0, + channels)) { + fprintf(stderr, "failed to load image from '%s'\n", path.c_str()); + return false; + } + params.set_width_and_height_if_unset(image.get().width, image.get().height); + return true; +} + +static bool load_audio_file(const std::string& path, SDAudioOwner& audio) { + std::vector samples; + uint32_t sample_rate = 0; + uint32_t channels = 0; + if (!load_wav_from_file(path, samples, sample_rate, channels)) { + fprintf(stderr, "failed to load WAV audio from '%s'\n", path.c_str()); + return false; + } + audio.reset(std::move(samples), sample_rate, channels); + return true; +} + +using LoadInput = bool (*)(SDGenerationParams&, SDMode, bool); + +struct LoadInputHandler { + const char* name; + LoadInput load; +}; + +static const LoadInputHandler LOAD_INPUT_HANDLERS[] = { + {"init_image", [](SDGenerationParams& params, SDMode, bool) { + return load_image_file(params, params.init_image_path, params.init_image); + }}, + {"end_image", [](SDGenerationParams& params, SDMode, bool) { + return load_image_file(params, params.end_image_path, params.end_image); + }}, + {"ref_images", [](SDGenerationParams& params, SDMode, bool) { + params.ref_images.clear(); + for (const auto& path : params.ref_image_paths) { + SDImageOwner image({0, 0, 3, nullptr}); + if (!load_image_file(params, path, image, false)) { + return false; + } + params.ref_images.push_back(std::move(image)); + } + return true; + }}, + {"validate", [](SDGenerationParams& params, SDMode mode, bool) { + return params.validate(mode); + }}, + {"ref_videos", [](SDGenerationParams& params, SDMode, bool verbose) { + params.ref_videos.clear(); + for (const auto& path : params.ref_video_paths) { + std::vector frames; + if (!load_images_from_dir(path, frames, 0, 0, 0, verbose) || frames.empty()) { + fprintf(stderr, "failed to load reference video frames from '%s'\n", path.c_str()); + return false; + } + params.ref_videos.push_back(std::move(frames)); + } + return true; + }}, + {"ref_video_audios", [](SDGenerationParams& params, SDMode, bool) { + params.ref_video_audios.clear(); + params.ref_video_audios.resize(params.ref_videos.size()); + for (size_t i = 0; i < params.ref_video_audio_paths.size(); ++i) { + if (!load_audio_file(params.ref_video_audio_paths[i], params.ref_video_audios[i])) { + return false; + } + } + return true; + }}, + {"ref_audios", [](SDGenerationParams& params, SDMode, bool) { + params.ref_audios.clear(); + params.ref_audios.resize(params.ref_audio_paths.size()); + for (size_t i = 0; i < params.ref_audio_paths.size(); ++i) { + if (!load_audio_file(params.ref_audio_paths[i], params.ref_audios[i])) { + return false; + } + } + return true; + }}, + {"mask_image", [](SDGenerationParams& params, SDMode, bool) { + return load_image_file(params, params.mask_image_path, params.mask_image, true, 1); + }}, + {"control_image", [](SDGenerationParams& params, SDMode, bool) { + return load_image_file(params, params.control_image_path, params.control_image); + }}, + {"ip_adapter_image", [](SDGenerationParams& params, SDMode, bool) { + return load_image_file(params, params.ip_adapter_image_path, params.ip_adapter_image, false); + }}, + {"control_video", [](SDGenerationParams& params, SDMode, bool verbose) { + if (params.control_video_path.empty()) { + return true; + } + params.control_frames.clear(); + return load_images_from_dir(params.control_video_path, + params.control_frames, + params.get_resolved_width(), + params.get_resolved_height(), + params.video_frames, + verbose); + }}, + {"pm_id_images", [](SDGenerationParams& params, SDMode, bool verbose) { + if (params.pm_id_images_dir.empty()) { + return true; + } + params.pm_id_images.clear(); + return load_images_from_dir(params.pm_id_images_dir, + params.pm_id_images, + 0, + 0, + 0, + verbose); + }}, +}; + +static bool load_generation_inputs(SDGenerationParams& params, SDMode mode, bool verbose) { + for (const auto& handler : LOAD_INPUT_HANDLERS) { + if (!handler.load(params, mode, verbose)) { + return false; + } + } + return true; +} + +int main(int argc, const char* argv[]) { + if (argc > 1 && std::string(argv[1]) == "--version") { + printf("%s\n", version_string().c_str()); + return 0; + } + + SDFitCliParams fit_params; + SDContextParams ctx_params; + SDGenerationParams gen_params; + + std::vector options_vec = {fit_params.get_options(), ctx_params.get_options(), gen_params.get_options()}; + if (!parse_options(argc, argv, options_vec)) { + print_usage(argc, argv, options_vec); + return 1; + } + + sd_set_log_callback(fit_log_cb, (void*)&fit_params); + + SDMode mode = gen_params.video_frames > 1 ? VID_GEN : IMG_GEN; + if (!ctx_params.resolve_and_validate(mode) || + !gen_params.resolve_and_validate(mode, ctx_params.lora_model_dir, ctx_params.hires_upscalers_dir)) { + print_usage(argc, argv, options_vec); + return 1; + } + if (!load_generation_inputs(gen_params, mode, fit_params.verbose)) { + return 1; + } + + sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(false); + + sd_fit_workload_t workload; + sd_fit_workload_init(&workload); + workload.prompt = gen_params.prompt.c_str(); + workload.width = gen_params.get_resolved_width(); + workload.height = gen_params.get_resolved_height(); + workload.video_frames = gen_params.video_frames; + workload.vae_tiling_params = gen_params.vae_tiling_params; + + sd_img_gen_params_t image_request; + sd_vid_gen_params_t video_request; + if (mode == VID_GEN) { + video_request = gen_params.to_sd_vid_gen_params_t(); + workload.video_gen_params = &video_request; + } else { + image_request = gen_params.to_sd_img_gen_params_t(); + workload.image_gen_params = &image_request; + } + + sd_fit_result_t result; + enum sd_fit_status_t status = sd_fit_params(&sd_ctx_params, &workload, &result); + if (status != SD_FIT_SUCCESS) { + if (fit_params.fit_print && result.report != nullptr) { + printf("%s", result.report); + } + fprintf(stderr, "failed to fit CLI arguments to free memory, exiting...\n"); + sd_fit_result_free(&result); + return 1; + } + + if (fit_params.fit_print) { + printf("%s", SAFE_STR(result.report)); + } else if (result.changed) { + std::string args; + if (result.backend != nullptr) { + args += std::string("--backend \"") + result.backend + "\""; + } + if (result.params_backend != nullptr) { + if (!args.empty()) { + args += " "; + } + args += std::string("--params-backend \"") + result.params_backend + "\""; + } + if (result.vae_tiling && !workload.vae_tiling_params.enabled) { + if (!args.empty()) { + args += " "; + } + args += "--vae-tiling"; + } + if (result.stream_layers && strlen(SAFE_STR(sd_ctx_params.max_vram)) > 0) { + if (!args.empty()) { + args += " "; + } + args += std::string("--max-vram \"") + sd_ctx_params.max_vram + "\""; + } + if (result.stream_layers && !sd_ctx_params.stream_layers) { + if (!args.empty()) { + args += " "; + } + args += "--stream-layers"; + } + fprintf(stderr, "printing fitted CLI arguments to stdout...\n"); + printf("%s\n", args.c_str()); + } else { + fprintf(stderr, "current parameters already fit into free device memory, no changes needed\n"); + printf("\n"); + } + + sd_fit_result_free(&result); + return 0; +} diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 258465976..856641410 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -500,6 +500,45 @@ SD_API void sd_hires_params_init(sd_hires_params_t* hires_params); SD_API void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params); SD_API char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params); +// the workload the fitted parameters must accommodate: compute memory depends on it +typedef struct { + const char* prompt; // representative prompt, token count affects text encoder memory + int width; + int height; + int video_frames; // <= 1 for image generation + sd_tiling_params_t vae_tiling_params; + // Optional complete representative request. Set at most one. When present, + // it supplies conditioning, LoRAs, hires, cache, and other generation inputs; + // the scalar fields above remain the fallback for callers that only need a + // basic text-to-image/video workload. + const sd_img_gen_params_t* image_gen_params; + const sd_vid_gen_params_t* video_gen_params; +} sd_fit_workload_t; + +enum sd_fit_status_t { + SD_FIT_SUCCESS = 0, // found a placement projected to fit (or no changes needed) + SD_FIT_FAILURE = 1, // could not find a placement projected to fit + SD_FIT_ERROR = 2, // a hard error occurred, e.g. the model could not be read +}; + +typedef struct { + bool changed; // false = current/default placement already fits, specs below are NULL + char* backend; // fitted runtime placement spec for --backend, owned by the result + char* params_backend; // fitted params placement spec for --params-backend, owned by the result + bool vae_tiling; // recommend enabling VAE tiling + bool stream_layers; // recommend enabling graph-cut streaming layers + char* report; // human readable per-device / per-module memory table +} sd_fit_result_t; + +SD_API void sd_fit_workload_init(sd_fit_workload_t* workload); + +// fit model/context placement to free device memory using measured no-alloc dry runs; +// reads only model metadata, never weight data; backend and params_backend must be unset +SD_API enum sd_fit_status_t sd_fit_params(const sd_ctx_params_t* sd_ctx_params, + const sd_fit_workload_t* workload, + sd_fit_result_t* result); +SD_API void sd_fit_result_free(sd_fit_result_t* result); + SD_API sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params); SD_API void free_sd_ctx(sd_ctx_t* sd_ctx); SD_API void free_sd_audio(sd_audio_t* audio); diff --git a/src/core/fit_params.cpp b/src/core/fit_params.cpp new file mode 100644 index 000000000..abfa1d0cd --- /dev/null +++ b/src/core/fit_params.cpp @@ -0,0 +1,630 @@ +#include "fit_params.h" + +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#elif defined(__APPLE__) +#include +#elif defined(__linux__) +#include +#endif + +#include "core/util.h" +#include "ggml-backend.h" + +namespace sd::fit_params { + namespace { + + constexpr int64_t MiB = 1024ll * 1024; + constexpr int64_t GiB = 1024ll * MiB; + constexpr int64_t MEMORY_RESERVE = 512 * MiB; + + struct Device { + ggml_backend_dev_t dev = nullptr; + std::string name; + std::string description; + int64_t free_bytes = 0; + int64_t total_bytes = 0; + int64_t budget_bytes = 0; + bool graph_budget_enabled = false; + }; + + struct Decision { + bool placed = false; + bool on_cpu = false; + bool disk_params = false; + bool cpu_params = false; + bool tiled = false; + bool stream_layers = false; + std::vector device_idxs; + }; + + struct ComputePhases { + int64_t serial = 0; + int64_t diffusion = 0; + int64_t control_net = 0; + + void add(SDBackendModule module, int64_t bytes) { + if (module == SDBackendModule::DIFFUSION) { + diffusion = std::max(diffusion, bytes); + } else if (module == SDBackendModule::CONTROL_NET) { + control_net = std::max(control_net, bytes); + } else { + serial = std::max(serial, bytes); + } + } + + int64_t peak() const { + return std::max(serial, diffusion + control_net); + } + }; + + int64_t available_host_memory() { + const char* debug_gib = getenv("SD_FIT_DEBUG_HOST_MEMORY_GIB"); + if (debug_gib != nullptr && debug_gib[0] != '\0') { + return std::max((int64_t)(std::strtod(debug_gib, nullptr) * GiB), 0); + } +#if defined(_WIN32) + MEMORYSTATUSEX status{}; + status.dwLength = sizeof(status); + return GlobalMemoryStatusEx(&status) ? (int64_t)status.ullAvailPhys : -1; +#elif defined(__APPLE__) + vm_statistics64_data_t stats{}; + mach_msg_type_number_t count = HOST_VM_INFO64_COUNT; + vm_size_t page_size = 0; + if (host_page_size(mach_host_self(), &page_size) != KERN_SUCCESS || + host_statistics64(mach_host_self(), HOST_VM_INFO64, + reinterpret_cast(&stats), &count) != KERN_SUCCESS) { + return -1; + } + return (int64_t)(stats.free_count + stats.inactive_count + stats.speculative_count) * + (int64_t)page_size; +#elif defined(__linux__) + std::ifstream meminfo("/proc/meminfo"); + std::string key; + int64_t value = 0; + std::string unit; + while (meminfo >> key >> value >> unit) { + if (key == "MemAvailable:") { + return value * 1024; + } + } + struct sysinfo info {}; + if (sysinfo(&info) == 0) { + return (int64_t)(info.freeram + info.bufferram) * (int64_t)info.mem_unit; + } + return -1; +#else + ggml_backend_dev_t cpu = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); + if (cpu == nullptr) { + return -1; + } + size_t free_bytes = 0; + size_t total_bytes = 0; + ggml_backend_dev_memory(cpu, &free_bytes, &total_bytes); + return free_bytes > 0 ? (int64_t)free_bytes : -1; +#endif + } + + void apply_device_budget(Device& d, sd::ggml_graph_cut::MaxVramAssignment& budgets) { + float gib = budgets.default_gib; + { + std::string budget_key = d.name; + std::transform(budget_key.begin(), budget_key.end(), budget_key.begin(), [](unsigned char c) { return (char)std::tolower(c); }); + auto it = budgets.backend_gib.find(budget_key); + if (it != budgets.backend_gib.end()) { + gib = it->second; + } + } + if (gib > 0.f) { + d.budget_bytes = std::min((int64_t)(gib * GiB), d.free_bytes) - MEMORY_RESERVE; + d.graph_budget_enabled = true; + } else if (gib < 0.f) { + d.budget_bytes = d.free_bytes + (int64_t)(gib * GiB) - MEMORY_RESERVE; + d.graph_budget_enabled = true; + } else { + d.budget_bytes = d.free_bytes - MEMORY_RESERVE; + d.graph_budget_enabled = false; + } + d.budget_bytes = std::max(d.budget_bytes, 0); + } + + // debug override to exercise multi-device planning on any machine, + // e.g. SD_FIT_DEBUG_DEVICES="CUDA0:24,CUDA1:16" (name:free_gib) + std::vector simulated_devices(const char* spec, sd::ggml_graph_cut::MaxVramAssignment& budgets) { + std::vector out; + std::string s = spec; + size_t pos = 0; + while (pos < s.size()) { + size_t comma = s.find(',', pos); + std::string entry = s.substr(pos, comma == std::string::npos ? std::string::npos : comma - pos); + pos = comma == std::string::npos ? s.size() : comma + 1; + size_t colon = entry.find(':'); + if (colon == std::string::npos) { + continue; + } + Device d; + d.name = entry.substr(0, colon); + d.description = "simulated device"; + d.free_bytes = (int64_t)(std::stof(entry.substr(colon + 1)) * 1024.0 * 1024.0 * 1024.0); + d.total_bytes = d.free_bytes; + apply_device_budget(d, budgets); + out.push_back(d); + } + return out; + } + + std::vector enumerate_gpu_devices(sd::ggml_graph_cut::MaxVramAssignment& budgets) { + const char* debug_devices = getenv("SD_FIT_DEBUG_DEVICES"); + if (debug_devices != nullptr && debug_devices[0] != '\0') { + LOG_WARN("fit-params: planning against simulated devices (SD_FIT_DEBUG_DEVICES)"); + return simulated_devices(debug_devices, budgets); + } + + std::vector out; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) { + continue; + } + Device d; + d.dev = dev; + d.name = ggml_backend_dev_name(dev); + d.description = ggml_backend_dev_description(dev); + size_t free_bytes = 0, total_bytes = 0; + ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); + d.free_bytes = (int64_t)free_bytes; + d.total_bytes = (int64_t)total_bytes; + apply_device_budget(d, budgets); + out.push_back(d); + } + return out; + } + + std::string module_spec_key(SDBackendModule module) { + // sd_backend_module_name returns tokens sd_parse_backend_assignment accepts + return sd_backend_module_name(module); + } + + void append_assignment(std::string& spec, const std::string& key, const std::string& value) { + if (!spec.empty()) { + spec += ","; + } + spec += key; + spec += "="; + spec += value; + } + + void report_line(std::string& report, const char* fmt, ...) { + char line[512]; + va_list args; + va_start(args, fmt); + vsnprintf(line, sizeof(line), fmt, args); + va_end(args); + LOG_INFO("%s", line); + report += line; + report += "\n"; + } + + int64_t host_memory_requirement(const std::vector& modules, + const std::vector* decisions) { + int64_t params = 0; + ComputePhases compute; + for (size_t i = 0; i < modules.size(); ++i) { + const bool on_cpu = decisions == nullptr || + (*decisions)[i].on_cpu || + (*decisions)[i].cpu_params; + if (!on_cpu) { + continue; + } + params += (int64_t)modules[i].params_bytes; + if (decisions == nullptr || (*decisions)[i].on_cpu) { + compute.add(modules[i].module, (int64_t)modules[i].compute_bytes); + } + } + return params + compute.peak(); + } + + bool host_memory_fits(const std::vector& modules, + const std::vector* decisions, + std::string& report) { + const int64_t required = host_memory_requirement(modules, decisions); + if (required == 0) { + return true; + } + const int64_t available = available_host_memory(); + const int64_t budget = available < 0 ? -1 : std::max(available - MEMORY_RESERVE, 0); + if (budget < 0) { + report_line(report, " host memory availability could not be determined; refusing an unverified CPU placement"); + return false; + } + report_line(report, " host memory: available %lld MiB, budget %lld MiB, projected use %lld MiB", + (long long)(available / MiB), + (long long)(budget / MiB), + (long long)(required / MiB)); + return required <= budget; + } + + } // namespace + + bool plan_placement(const std::vector& modules, + sd::ggml_graph_cut::MaxVramAssignment& budgets, + FitPlan* plan) { + if (plan == nullptr) { + return false; + } + *plan = {}; + + { + std::string error; + if (!budgets.canonicalize_backend_keys(&error)) { + LOG_ERROR("%s", error.c_str()); + return false; + } + } + + std::vector devices = enumerate_gpu_devices(budgets); + + report_line(plan->report, "fit-params: measured memory plan"); + report_line(plan->report, " devices:"); + for (const Device& d : devices) { + report_line(plan->report, " %-12s %-32s free %6lld MiB, budget %6lld MiB", + d.name.c_str(), d.description.c_str(), + (long long)(d.free_bytes / MiB), (long long)(d.budget_bytes / MiB)); + } + report_line(plan->report, " modules (measured for this workload):"); + for (const ModuleMemory& m : modules) { + if (m.params_bytes == 0 && m.compute_bytes == 0) { + continue; + } + if (m.compute_bytes_tiled > 0) { + report_line(plan->report, " %-12s params %6lld MiB, compute %6lld MiB (tiled %6lld MiB)", + module_spec_key(m.module).c_str(), + (long long)(m.params_bytes / MiB), + (long long)(m.compute_bytes / MiB), + (long long)(m.compute_bytes_tiled / MiB)); + } else { + report_line(plan->report, " %-12s params %6lld MiB, compute %6lld MiB", + module_spec_key(m.module).c_str(), + (long long)(m.params_bytes / MiB), + (long long)(m.compute_bytes / MiB)); + } + } + + if (devices.empty()) { + report_line(plan->report, " no usable GPU devices; checking the default CPU backend"); + plan->valid = host_memory_fits(modules, nullptr, plan->report); + plan->changed = false; + return true; + } + + // check-first: the default placement puts every module on the default (first GPU) device + { + int64_t params_sum = 0; + ComputePhases compute; + for (const ModuleMemory& m : modules) { + params_sum += (int64_t)m.params_bytes; + compute.add(m.module, (int64_t)m.compute_bytes); + } + if (params_sum + compute.peak() <= devices[0].budget_bytes) { + report_line(plan->report, " projected use %lld MiB <= budget %lld MiB on %s, no changes needed", + (long long)((params_sum + compute.peak()) / MiB), + (long long)(devices[0].budget_bytes / MiB), + devices[0].name.c_str()); + plan->valid = true; + plan->changed = false; + return true; + } + } + + std::vector order(modules.size()); + for (size_t i = 0; i < order.size(); i++) { + order[i] = i; + } + std::sort(order.begin(), order.end(), [&](size_t a, size_t b) { + return modules[a].params_bytes > modules[b].params_bytes; + }); + + std::vector decisions(modules.size()); + bool time_share = false; + + // resident plan: every module keeps its params loaded, compute buffers coexist per device + { + std::vector params_sum(devices.size(), 0); + std::vector compute_phases(devices.size()); + bool ok = true; + bool vae_tiling = false; + std::vector resident(modules.size()); + auto find_device = [&](const ModuleMemory& m, int64_t compute) -> int { + int best = -1; + for (size_t di = 0; di < devices.size(); di++) { + ComputePhases candidate = compute_phases[di]; + candidate.add(m.module, compute); + int64_t need = params_sum[di] + (int64_t)m.params_bytes + + candidate.peak(); + if (need <= devices[di].budget_bytes && + (best < 0 || devices[di].budget_bytes - params_sum[di] > devices[best].budget_bytes - params_sum[best])) { + best = (int)di; + } + } + return best; + }; + for (size_t mi : order) { + const ModuleMemory& m = modules[mi]; + if (m.params_bytes == 0 && m.compute_bytes == 0) { + resident[mi].placed = true; + continue; + } + int64_t compute = (int64_t)m.compute_bytes; + int best = find_device(m, compute); + if (best < 0 && m.compute_bytes_tiled > 0 && m.compute_bytes_tiled < m.compute_bytes) { + // full-resolution decode does not fit anywhere, tiling may keep the module resident + compute = (int64_t)m.compute_bytes_tiled; + best = find_device(m, compute); + if (best >= 0) { + resident[mi].tiled = true; + vae_tiling = true; + } + } + if (best < 0) { + ok = false; + break; + } + params_sum[best] += (int64_t)m.params_bytes; + compute_phases[best].add(m.module, compute); + resident[mi].placed = true; + resident[mi].device_idxs.push_back((size_t)best); + } + if (ok) { + decisions = std::move(resident); + plan->vae_tiling = vae_tiling; + } else { + time_share = true; + } + } + + // time-share plan: phases run sequentially, heavy modules load per phase and free after + if (time_share) { + std::vector compute_phases(devices.size()); + auto compute_with_concurrent_phase = [&](size_t device_idx, + SDBackendModule module, + int64_t compute) { + if (module == SDBackendModule::DIFFUSION) { + return compute + compute_phases[device_idx].control_net; + } + if (module == SDBackendModule::CONTROL_NET) { + return compute + compute_phases[device_idx].diffusion; + } + return compute; + }; + auto concurrent_compute_extra = [&](size_t device_idx, + SDBackendModule module) { + if (module == SDBackendModule::DIFFUSION) { + return compute_phases[device_idx].control_net; + } + if (module == SDBackendModule::CONTROL_NET) { + return compute_phases[device_idx].diffusion; + } + return int64_t{0}; + }; + auto split_graphs_fit = [&](const ModuleMemory& m, + const std::vector& device_idxs, + int64_t compute) { + if (m.split_graph_segment_params.empty()) { + return false; + } + std::vector capacities; + capacities.reserve(device_idxs.size()); + for (size_t device_idx : device_idxs) { + capacities.push_back(std::max(devices[device_idx].budget_bytes - compute, 0)); + } + for (const auto& graph_segments : m.split_graph_segment_params) { + size_t device_pos = 0; + int64_t used = 0; + for (size_t segment_bytes : graph_segments) { + while (device_pos + 1 < capacities.size() && + used + (int64_t)segment_bytes > capacities[device_pos]) { + ++device_pos; + used = 0; + } + if (used + (int64_t)segment_bytes > capacities[device_pos]) { + return false; + } + used += (int64_t)segment_bytes; + } + } + return true; + }; + auto streamed_graphs_fit = [&](const ModuleMemory& m, int64_t budget) { + if (m.split_graph_segment_params.empty() || + m.split_graph_segment_params.size() != m.split_graph_segment_compute.size()) { + return false; + } + for (size_t graph_idx = 0; graph_idx < m.split_graph_segment_params.size(); ++graph_idx) { + const auto& params = m.split_graph_segment_params[graph_idx]; + const auto& compute = m.split_graph_segment_compute[graph_idx]; + if (params.size() != compute.size()) { + return false; + } + for (size_t segment_idx = 0; segment_idx < params.size(); ++segment_idx) { + if ((int64_t)params[segment_idx] + (int64_t)compute[segment_idx] > budget) { + return false; + } + } + } + return true; + }; + for (size_t mi : order) { + const ModuleMemory& m = modules[mi]; + Decision& decision = decisions[mi]; + decision = {}; + if (m.params_bytes == 0 && m.compute_bytes == 0) { + decision.placed = true; + continue; + } + int best = -1; + for (size_t di = 0; di < devices.size(); di++) { + const int64_t compute = compute_with_concurrent_phase(di, m.module, (int64_t)m.compute_bytes); + if ((int64_t)m.params_bytes + compute <= devices[di].budget_bytes && + (best < 0 || devices[di].budget_bytes > devices[best].budget_bytes)) { + best = (int)di; + } + } + if (best >= 0) { + decision.placed = true; + decision.disk_params = true; + decision.device_idxs.push_back((size_t)best); + compute_phases[best].add(m.module, (int64_t)m.compute_bytes); + continue; + } + if (m.compute_bytes_tiled > 0) { + for (size_t di = 0; di < devices.size(); di++) { + const int64_t compute = compute_with_concurrent_phase(di, m.module, (int64_t)m.compute_bytes_tiled); + if ((int64_t)m.params_bytes + compute <= devices[di].budget_bytes && + (best < 0 || devices[di].budget_bytes > devices[best].budget_bytes)) { + best = (int)di; + } + } + if (best >= 0) { + decision.placed = true; + decision.disk_params = true; + decision.tiled = true; + plan->vae_tiling = true; + decision.device_idxs.push_back((size_t)best); + compute_phases[best].add(m.module, (int64_t)m.compute_bytes_tiled); + continue; + } + } + if (m.splittable && devices.size() > 1) { + int64_t capacity = 0; + std::vector idxs(devices.size()); + for (size_t i = 0; i < idxs.size(); i++) { + idxs[i] = i; + } + std::sort(idxs.begin(), idxs.end(), [&](size_t a, size_t b) { + return devices[a].budget_bytes > devices[b].budget_bytes; + }); + for (size_t di = 0; di < devices.size(); ++di) { + const int64_t compute = compute_with_concurrent_phase(di, m.module, (int64_t)m.compute_bytes); + capacity += std::max(devices[di].budget_bytes - compute, 0); + } + int64_t split_compute = 0; + for (size_t di : idxs) { + split_compute = std::max(split_compute, + compute_with_concurrent_phase(di, m.module, (int64_t)m.compute_bytes)); + } + if ((int64_t)m.params_bytes <= capacity && split_graphs_fit(m, idxs, split_compute)) { + decision.placed = true; + decision.disk_params = true; + decision.device_idxs = std::move(idxs); + for (size_t di : decision.device_idxs) { + compute_phases[di].add(m.module, (int64_t)m.compute_bytes); + } + continue; + } + } + if (m.module == SDBackendModule::DIFFUSION && m.splittable) { + for (size_t di = 0; di < devices.size(); di++) { + if (devices[di].graph_budget_enabled && devices[di].budget_bytes > 0 && + streamed_graphs_fit(m, + devices[di].budget_bytes - + concurrent_compute_extra(di, m.module)) && + (best < 0 || devices[di].budget_bytes > devices[best].budget_bytes)) { + best = (int)di; + } + } + if (best >= 0) { + decision.placed = true; + decision.cpu_params = true; + decision.stream_layers = true; + plan->stream_layers = true; + decision.device_idxs.push_back((size_t)best); + compute_phases[best].add(m.module, (int64_t)m.compute_bytes); + continue; + } + } + decision.placed = true; + decision.on_cpu = true; + } + } + + if (!host_memory_fits(modules, &decisions, plan->report)) { + report_line(plan->report, " no placement fits available host memory"); + plan->valid = false; + return true; + } + + report_line(plan->report, " placement%s:", time_share ? " (time-share: params load per phase and free after)" : ""); + for (size_t mi = 0; mi < modules.size(); mi++) { + const ModuleMemory& m = modules[mi]; + const Decision& decision = decisions[mi]; + if (m.params_bytes == 0 && m.compute_bytes == 0) { + continue; + } + std::string target; + if (decision.on_cpu) { + target = "cpu"; + } else { + for (size_t k = 0; k < decision.device_idxs.size(); k++) { + if (k > 0) { + target += " & "; + } + target += devices[decision.device_idxs[k]].name; + } + if (decision.device_idxs.size() > 1) { + target += " (split)"; + } + } + report_line(plan->report, " %-12s -> %s%s%s%s", + module_spec_key(m.module).c_str(), + target.c_str(), + decision.disk_params ? ", params on disk" : "", + decision.cpu_params ? ", params on cpu, stream layers" : "", + decision.tiled ? ", vae tiling" : ""); + } + + for (size_t mi = 0; mi < modules.size(); mi++) { + const ModuleMemory& m = modules[mi]; + const Decision& decision = decisions[mi]; + if (m.params_bytes == 0 && m.compute_bytes == 0) { + continue; + } + const std::string key = module_spec_key(m.module); + if (decision.on_cpu) { + append_assignment(plan->runtime_spec, key, "cpu"); + continue; + } + if (decision.device_idxs.empty()) { + continue; + } + std::string device_list; + for (size_t k = 0; k < decision.device_idxs.size(); k++) { + if (k > 0) { + device_list += "&"; + } + device_list += devices[decision.device_idxs[k]].name; + } + append_assignment(plan->runtime_spec, key, device_list); + if (decision.disk_params) { + append_assignment(plan->params_spec, key, "disk"); + } else if (decision.cpu_params) { + append_assignment(plan->params_spec, key, "cpu"); + } + } + + plan->valid = true; + plan->changed = true; + plan->time_share = time_share; + return true; + } + +} // namespace sd::fit_params diff --git a/src/core/fit_params.h b/src/core/fit_params.h new file mode 100644 index 000000000..bf56ba25a --- /dev/null +++ b/src/core/fit_params.h @@ -0,0 +1,44 @@ +#ifndef __SD_FIT_PARAMS_H__ +#define __SD_FIT_PARAMS_H__ + +#include +#include + +#include "core/ggml_extend_backend.h" +#include "core/ggml_graph_cut.h" + +namespace sd::fit_params { + + // measured memory requirements for one module at the requested workload + struct ModuleMemory { + SDBackendModule module; + size_t params_bytes = 0; // weights registered for the module + size_t compute_bytes = 0; // largest measured compute buffer among the module's graphs + size_t compute_bytes_tiled = 0; // VAE only: compute buffer with tiling enabled, 0 if not measured + bool splittable = false; + // Ordered, de-duplicated parameter bytes per graph-cut segment for each + // measured graph. Multi-device placement must fit these indivisible + // segments, not just the aggregate module size. + std::vector> split_graph_segment_params; + std::vector> split_graph_segment_compute; + }; + + struct FitPlan { + bool valid = false; + bool changed = false; // false = current/default placement already fits + bool time_share = false; + bool vae_tiling = false; + bool stream_layers = false; + std::string runtime_spec; + std::string params_spec; + std::string report; // human readable per-device / per-module table + }; + + // derive placement specs from measured module memory and per-device budgets + bool plan_placement(const std::vector& modules, + sd::ggml_graph_cut::MaxVramAssignment& budgets, + FitPlan* plan); + +} // namespace sd::fit_params + +#endif // __SD_FIT_PARAMS_H__ diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp index 017d3d8f3..67f182ad5 100644 --- a/src/core/ggml_extend.hpp +++ b/src/core/ggml_extend.hpp @@ -1747,6 +1747,19 @@ struct GGMLRunnerContext { }; struct GGMLRunner { + struct graph_memory_measurement { + std::string desc; + const GGMLRunner* runner = nullptr; + ggml_backend_t backend = nullptr; + SDBackendModule module = SDBackendModule::UNSET; + size_t compute_bytes = 0; + size_t cache_bytes = 0; + size_t params_bytes = 0; + std::vector split_segment_params_bytes; + std::vector split_segment_compute_bytes; + bool valid = false; + }; + protected: typedef std::function get_graph_cb_t; using GraphCutSegment = sd::ggml_graph_cut::Segment; @@ -1783,6 +1796,22 @@ struct GGMLRunner { std::unordered_set runner_param_tensor_set; bool params_tensor_set_dirty_ = true; + // Thread-local so nested runners in one fitting call are intercepted without + // affecting generation or fitting calls running on other threads. + static inline thread_local bool measure_mode_ = false; + static inline thread_local std::vector* measure_collector_ = nullptr; + static inline thread_local size_t measure_generation_ = 0; + graph_memory_measurement last_measurement_; + SDBackendModule fit_module_ = SDBackendModule::UNSET; + size_t measure_generation_seen_ = 0; + + struct measured_cache_tensor { + ggml_type type = GGML_TYPE_F32; + std::vector shape; + size_t alloc_bytes = 0; + }; + std::map measured_cache_tensors_; + std::vector one_vec = {1.f}; ggml_tensor* one_tensor = nullptr; @@ -2270,6 +2299,112 @@ struct GGMLRunner { return true; } + // measure the compute buffer size and used param bytes of a built graph without + // allocating anything: params are temporarily marked as externally owned so + // gallocr does not reserve them (same trick as measure_segment_compute_buffer) + void measure_graph_memory(ggml_cgraph* gf) { + last_measurement_ = {}; + + struct TensorRuntimeBinding { + ggml_backend_buffer_t buffer = nullptr; + void* data = nullptr; + void* extra = nullptr; + }; + std::unordered_map saved_bindings; + auto mark_external = [&](ggml_tensor* t) { + if (t == nullptr || saved_bindings.find(t) != saved_bindings.end()) { + return; + } + saved_bindings[t] = {t->buffer, t->data, t->extra}; + t->data = reinterpret_cast(static_cast(1)); + }; + + std::vector used_params = collect_used_param_tensors(gf); + for (ggml_tensor* param : used_params) { + last_measurement_.params_bytes += ggml_nbytes(param); + } + const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); + for (int i = 0; i < n_leafs; ++i) { + ggml_tensor* leaf = sd::ggml_graph_cut::leaf_tensor(gf, i); + if (canonical_param_tensor(leaf) == nullptr) { + continue; + } + mark_external(leaf); + mark_external(leaf->view_src); + } + + ggml_gallocr_t allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_backend)); + size_t sizes[1] = {0}; + ggml_gallocr_reserve_n_size(allocr, gf, nullptr, nullptr, sizes); + last_measurement_.compute_bytes = sizes[0]; + ggml_gallocr_free(allocr); + + ggml_backend_buffer_type_t cache_buft = + ggml_backend_get_default_buffer_type(runtime_backend); + const size_t cache_alignment = ggml_backend_buft_get_alignment(cache_buft); + size_t previous_cache_bytes = 0; + for (const auto& entry : measured_cache_tensors_) { + previous_cache_bytes += entry.second.alloc_bytes; + } + const bool replaces_cache_buffer = !cache_tensor_map.empty(); + for (const auto& entry : cache_tensor_map) { + ggml_tensor* tensor = sd::ggml_graph_cut::cache_source_tensor(entry.second); + if (tensor == nullptr) { + continue; + } + measured_cache_tensor measured; + measured.type = tensor->type; + measured.shape.assign(tensor->ne, tensor->ne + ggml_n_dims(tensor)); + measured.alloc_bytes = ggml_backend_buft_get_alloc_size(cache_buft, tensor); + if (cache_alignment > 0) { + measured.alloc_bytes = GGML_PAD(measured.alloc_bytes, cache_alignment); + } + measured_cache_tensors_[entry.first] = std::move(measured); + } + for (const auto& entry : measured_cache_tensors_) { + last_measurement_.cache_bytes += entry.second.alloc_bytes; + } + const size_t live_cache_bytes = last_measurement_.cache_bytes + + (replaces_cache_buffer ? previous_cache_bytes : 0); + last_measurement_.compute_bytes += live_cache_bytes; + + for (const auto& kv : saved_bindings) { + kv.first->buffer = kv.second.buffer; + kv.first->data = kv.second.data; + kv.first->extra = kv.second.extra; + } + last_measurement_.valid = true; + last_measurement_.desc = get_desc(); + last_measurement_.runner = this; + last_measurement_.backend = runtime_backend; + last_measurement_.module = fit_module_; + + const auto split_plan = sd::ggml_graph_cut::build_plan(runtime_backend, + gf, + params_tensor_set_, + get_desc().c_str()); + if (split_plan.valid && split_plan.has_cuts && split_plan.segments.size() > 1) { + std::unordered_set seen_split_params; + last_measurement_.split_segment_params_bytes.reserve(split_plan.segments.size()); + last_measurement_.split_segment_compute_bytes.reserve(split_plan.segments.size()); + for (const auto& segment : split_plan.segments) { + size_t segment_bytes = 0; + for (ggml_tensor* raw_param : sd::ggml_graph_cut::param_tensors(gf, segment)) { + ggml_tensor* param = canonical_param_tensor(raw_param); + if (param != nullptr && seen_split_params.insert(param).second) { + segment_bytes += ggml_nbytes(param); + } + } + last_measurement_.split_segment_params_bytes.push_back(segment_bytes); + last_measurement_.split_segment_compute_bytes.push_back( + segment.compute_buffer_size + live_cache_bytes); + } + } + if (measure_collector_ != nullptr) { + measure_collector_->push_back(last_measurement_); + } + } + void free_cache_buffer() { if (cache_buffer != nullptr) { ggml_backend_buffer_free(cache_buffer); @@ -3111,6 +3246,9 @@ struct GGMLRunner { } void reset_compute_ctx() { + if (measure_mode_) { + cache_tensor_map.clear(); + } free_compute_ctx(); alloc_compute_ctx(); } @@ -3188,10 +3326,23 @@ struct GGMLRunner { } ggml_tensor* get_cache_tensor_by_name(const std::string& name) { - if (cache_ctx == nullptr) { + if (cache_ctx != nullptr) { + return ggml_get_tensor(cache_ctx, name.c_str()); + } + if (!measure_mode_) { return nullptr; } - return ggml_get_tensor(cache_ctx, name.c_str()); + auto it = measured_cache_tensors_.find(name); + if (it == measured_cache_tensors_.end() || compute_ctx == nullptr || it->second.shape.empty()) { + return nullptr; + } + ggml_tensor* tensor = ggml_new_tensor(compute_ctx, + it->second.type, + static_cast(it->second.shape.size()), + it->second.shape.data()); + ggml_set_name(tensor, name.c_str()); + tensor->data = reinterpret_cast(static_cast(1)); + return tensor; } template @@ -3220,6 +3371,12 @@ struct GGMLRunner { }; RunnerDoneGuard runner_done_guard(this, auto_free); + if (measure_mode_ && measure_generation_seen_ != measure_generation_) { + cache_tensor_map.clear(); + measured_cache_tensors_.clear(); + measure_generation_seen_ = measure_generation_; + } + ggml_cgraph* gf = nullptr; if (!prepare_compute_graph(get_graph, &gf)) { return std::nullopt; @@ -3227,6 +3384,21 @@ struct GGMLRunner { GGML_ASSERT(gf != nullptr); rebuild_params_tensor_set(); + if (measure_mode_) { + measure_graph_memory(gf); + // return a correctly shaped zero tensor so downstream host-side code + // (condition assembly, samplers) keeps working without weight data + std::optional> result = sd::Tensor(); + if (!no_return && ggml_graph_n_nodes(gf) > 0) { + ggml_tensor* out = ggml_get_tensor(compute_ctx, final_result_name.c_str()); + if (out == nullptr) { + return std::nullopt; + } + result = sd::zeros(sd::shape_from_ggml(out)); + } + return result; + } + if (!assign_graph_cut_layer_split_backends(gf)) { free_compute_ctx(); return std::nullopt; @@ -3398,6 +3570,28 @@ struct GGMLRunner { max_graph_vram_bytes = max_vram_bytes; } + // in measure mode compute() builds the graph, records memory requirements and + // returns a shaped zero tensor without loading weights or allocating buffers + static void set_measure_mode(bool enabled, std::vector* collector = nullptr) { + if (enabled && !measure_mode_) { + ++measure_generation_; + } + measure_mode_ = enabled; + measure_collector_ = enabled ? collector : nullptr; + } + + static bool measure_mode_enabled() { + return measure_mode_; + } + + graph_memory_measurement get_last_measurement() const { + return last_measurement_; + } + + void set_fit_module(SDBackendModule module) { + fit_module_ = module; + } + void set_stream_layers_enabled(bool enabled) { if (enabled && is_multi_device()) { LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring", diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index a11506ffc..74d11d5ec 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -993,6 +993,8 @@ const char* sd_backend_module_name(SDBackendModule module) { return "upscaler"; case SDBackendModule::DETECTOR: return "detector"; + case SDBackendModule::UNSET: + return "unset"; } return "unknown"; } diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index b8085c7b6..f81c3b98e 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -21,6 +21,7 @@ enum class SDBackendModule { PHOTOMAKER, UPSCALER, DETECTOR, + UNSET, }; struct SDBackendAssignment { diff --git a/src/core/util.cpp b/src/core/util.cpp index 23af71e73..4106cff52 100644 --- a/src/core/util.cpp +++ b/src/core/util.cpp @@ -341,6 +341,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 sd_abort_cb_t sd_abort_cb = nullptr; static void* sd_abort_cb_data = nullptr; @@ -537,6 +538,9 @@ static void print_progress_line(int step, int steps, const std::string& speed_te } void pretty_progress(int step, int steps, float time) { + if (sd_progress_suppressed) { + return; + } if (sd_progress_cb) { sd_progress_cb(step, steps, time, sd_progress_cb_data); return; @@ -554,6 +558,9 @@ void pretty_progress(int step, int steps, float time) { } void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds) { + if (sd_progress_suppressed) { + return; + } if (sd_progress_cb) { float time = elapsed_seconds / (step + 1e-6f); sd_progress_cb(step, steps, time, sd_progress_cb_data); @@ -687,6 +694,14 @@ void* sd_get_progress_callback_data() { return sd_progress_cb_data; } +bool sd_get_progress_suppressed() { + return sd_progress_suppressed; +} + +void sd_set_progress_suppressed(bool suppressed) { + sd_progress_suppressed = suppressed; +} + sd_image_t tensor_to_sd_image(const sd::Tensor& tensor, int frame_index) { const auto& shape = tensor.shape(); GGML_ASSERT(shape.size() == 4 || shape.size() == 5); diff --git a/src/core/util.h b/src/core/util.h index e882ece23..e0308d330 100644 --- a/src/core/util.h +++ b/src/core/util.h @@ -91,6 +91,8 @@ std::vector> split_quotation_attention( sd_progress_cb_t sd_get_progress_callback(); void* sd_get_progress_callback_data(); +bool sd_get_progress_suppressed(); +void sd_set_progress_suppressed(bool suppressed); sd_preview_cb_t sd_get_preview_callback(); void* sd_get_preview_callback_data(); diff --git a/src/model/adapter/lora.hpp b/src/model/adapter/lora.hpp index 26edbc07c..6c0f41a0a 100644 --- a/src/model/adapter/lora.hpp +++ b/src/model/adapter/lora.hpp @@ -1179,11 +1179,13 @@ struct LoraModel : public GGMLRunner { }; GGMLRunner::compute(get_graph, n_threads, false, false, false, true); stat(!warn_unused); - for (auto item : original_tensor_to_final_tensor) { - ggml_tensor* original_tensor = item.first; - ggml_tensor* final_tensor = item.second; + if (!GGMLRunner::measure_mode_enabled()) { + for (auto item : original_tensor_to_final_tensor) { + ggml_tensor* original_tensor = item.first; + ggml_tensor* final_tensor = item.second; - ggml_backend_tensor_copy(final_tensor, original_tensor); + ggml_backend_tensor_copy(final_tensor, original_tensor); + } } original_tensor_to_final_tensor.clear(); GGMLRunner::free_compute_buffer(); diff --git a/src/model/diffusion/control.hpp b/src/model/diffusion/control.hpp index bf3c7e435..9cbeed581 100644 --- a/src/model/diffusion/control.hpp +++ b/src/model/diffusion/control.hpp @@ -432,7 +432,10 @@ struct ControlNet : public GGMLRunner { controls.clear(); controls.reserve(control_outputs_ggml.size()); for (ggml_tensor* control : control_outputs_ggml) { - auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml(control), 4); + auto control_host = GGMLRunner::measure_mode_enabled() + ? sd::zeros(sd::shape_from_ggml(control)) + : sd::make_sd_tensor_from_ggml(control); + control_host = restore_trailing_singleton_dims(std::move(control_host), 4); GGML_ASSERT(!control_host.empty()); controls.push_back(std::move(control_host)); } diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 9cf04c874..a89cb3f7f 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -602,13 +602,12 @@ namespace MiniMaxH3VAE { sd::Tensor encode(int n_threads, const sd::Tensor& x, sd_tiling_params_t tiling_params, - bool encode_video = false, - bool circular_x = false, - bool circular_y = false) override { + bool circular_x = false, + bool circular_y = false) override { auto input = ensure_video_shape(x); auto tiling = h3_tiling(tiling_params); if (input.shape()[2] == 1) { - auto encoded = VAE::encode(n_threads, input, tiling, encode_video, circular_x, circular_y); + auto encoded = VAE::encode(n_threads, input, tiling, circular_x, circular_y); if (!encoded.empty() && encoded.shape()[2] > 1) { encoded = sd::ops::slice(encoded, 2, @@ -628,7 +627,7 @@ namespace MiniMaxH3VAE { sd::Tensor result; for (int64_t start = 0; start < input.shape()[2]; start += 17) { auto chunk = sd::ops::slice(input, 2, start, start + 17); - auto encoded = VAE::encode(n_threads, chunk, tiling, encode_video, circular_x, circular_y); + auto encoded = VAE::encode(n_threads, chunk, tiling, circular_x, circular_y); if (encoded.empty()) { return {}; } diff --git a/src/model/vae/vae.hpp b/src/model/vae/vae.hpp index e6cb82978..5a06a509c 100644 --- a/src/model/vae/vae.hpp +++ b/src/model/vae/vae.hpp @@ -118,9 +118,8 @@ struct VAE : public GGMLRunner { virtual sd::Tensor encode(int n_threads, const sd::Tensor& x, sd_tiling_params_t tiling_params, - bool encode_video = false, - bool circular_x = false, - bool circular_y = false) { + bool circular_x = false, + bool circular_y = false) { int64_t t0 = ggml_time_ms(); sd::Tensor input = x; sd::Tensor output; @@ -129,9 +128,7 @@ struct VAE : public GGMLRunner { scale_tensor_to_minus1_1(&input); } - // Video VAEs produce 4D/5D tensors that the 2D spatial tiling path cannot - // handle correctly; bypass spatial tiling for video encode. - if (tiling_params.enabled && !encode_video) { + if (tiling_params.enabled) { const int scale_factor = get_scale_factor(); int64_t W = input.shape()[0] / scale_factor; int64_t H = input.shape()[1] / scale_factor; @@ -183,9 +180,7 @@ struct VAE : public GGMLRunner { sd::Tensor output; set_tiling_params(tiling_params); - // Video VAEs produce 4D/5D tensors that the 2D spatial tiling path cannot - // handle correctly; bypass spatial tiling for video decode. - if (tiling_params.enabled && !decode_video) { + if (tiling_params.enabled) { const int scale_factor = get_scale_factor(); int64_t W = input.shape()[0] * scale_factor; int64_t H = input.shape()[1] * scale_factor; diff --git a/src/model_io/gguf_io.cpp b/src/model_io/gguf_io.cpp index cd22312d5..decea8996 100644 --- a/src/model_io/gguf_io.cpp +++ b/src/model_io/gguf_io.cpp @@ -1,5 +1,7 @@ #include "gguf_io.h" +#include + #include #include #include @@ -46,11 +48,30 @@ bool read_gguf_file(const std::string& file_path, gguf_context* ctx_gguf_ = nullptr; ggml_context* ctx_meta_ = nullptr; + static std::mutex gguf_probe_log_mutex; + + // ggml's reader rejects tensors with more than GGML_MAX_DIMS dimensions (e.g. + // the 5-D Wan patch_embedding.weight) that the GGUFReader fallback handles, + // so silence its log during the probe to avoid spurious error messages + { + std::lock_guard lock(gguf_probe_log_mutex); + ggml_log_callback saved_log_cb = nullptr; + void* saved_log_data = nullptr; + ggml_log_get(&saved_log_cb, &saved_log_data); + ggml_log_set([](ggml_log_level, const char*, void*) {}, nullptr); + ctx_gguf_ = gguf_init_from_file(file_path.c_str(), {true, &ctx_meta_}); + ggml_log_set(saved_log_cb, saved_log_data); + } - ctx_gguf_ = gguf_init_from_file(file_path.c_str(), {true, &ctx_meta_}); if (!ctx_gguf_) { GGUFReader gguf_reader; if (!gguf_reader.load(file_path)) { + // re-run the ggml reader with logging restored so its diagnostics are printed + ctx_gguf_ = gguf_init_from_file(file_path.c_str(), {true, &ctx_meta_}); + if (ctx_gguf_) { + gguf_free(ctx_gguf_); + ggml_free(ctx_meta_); + } set_error(error, "failed to open '" + file_path + "' with GGUFReader"); return false; } diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index e89ecf492..b85b35e06 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -21,6 +21,7 @@ #include "conditioning/conditioner.hpp" #include "core/backend_fit.h" +#include "core/fit_params.h" #include "extensions/generation_extension.h" #include "model/adapter/ip_adapter.hpp" #include "model/adapter/lora.hpp" @@ -267,6 +268,8 @@ class StableDiffusionGGML { std::string split_mode_spec; bool auto_fit_enabled = false; bool vae_auto_cpu_fallback_enabled = false; + bool fit_dry_run = false; // metadata-only init for memory measurement, never reads weight data + std::map fit_module_params_bytes; bool diffusion_conv_direct = false; @@ -342,12 +345,20 @@ class StableDiffusionGGML { if (model == nullptr) { return true; } + if constexpr (std::is_base_of_v) { + model->set_fit_module(module); + } std::map group_tensors; std::map tensor_ops; model->get_param_tensors(group_tensors); if constexpr (std::is_base_of_v) { model->get_param_tensor_ops(tensor_ops); } + for (const auto& kv : group_tensors) { + if (kv.second != nullptr) { + fit_module_params_bytes[module] += ggml_nbytes(kv.second); + } + } if (model_manager == nullptr) { return true; } @@ -901,6 +912,10 @@ class StableDiffusionGGML { } } } + if (fit_dry_run) { + enable_mmap = false; + eager_load = false; + } max_vram_assignment.reset(0.f); { std::string error; @@ -1031,7 +1046,9 @@ class StableDiffusionGGML { LOG_WARN("in mode 'immediately', LoRAs will cause extra memory usage with mmap"); } model_loader.process_model_files(enable_mmap, needs_writable_mmap); - load_alphas_cumprod(model_loader); + if (!fit_dry_run) { + load_alphas_cumprod(model_loader); + } diffusion_conv_direct = sd_ctx_params->diffusion_conv_direct; @@ -1775,8 +1792,9 @@ class StableDiffusionGGML { size_t total_params_size = total_params_ram_size + total_params_vram_size; LOG_INFO( - "total params memory size = %.2fMB (VRAM %.2fMB, RAM %.2fMB): " + "%stotal params memory size = %.2fMB (VRAM %.2fMB, RAM %.2fMB): " "text_encoders %.2fMB(%s), diffusion_model %.2fMB(%s), vae %.2fMB(%s), controlnet %.2fMB(%s), extensions %.2fMB(%s)", + fit_dry_run ? "projected " : "", // in a fit dry run nothing is allocated total_params_size / 1024.0 / 1024.0, total_params_vram_size / 1024.0 / 1024.0, total_params_ram_size / 1024.0 / 1024.0, @@ -1798,7 +1816,8 @@ class StableDiffusionGGML { if (pred_type == PREDICTION_COUNT) { if (sd_version_is_sd2(version)) { - pred_type = is_using_v_parameterization_for_sd2(sd_version_is_inpaint(version)) ? V_PRED : EPS_PRED; + // the v-pred probe runs a real compute with loaded weights; irrelevant for memory measurement + pred_type = !fit_dry_run && is_using_v_parameterization_for_sd2(sd_version_is_inpaint(version)) ? V_PRED : EPS_PRED; } else if (sd_version_is_sdxl(version)) { if (tensor_storage_map.find("edm_vpred.sigma_max") != tensor_storage_map.end()) { // CosXL models @@ -1983,6 +2002,7 @@ class StableDiffusionGGML { lora_spec.path, lora_spec.is_high_noise ? "model.high_noise_" : "", version); + lora->set_fit_module(module); LoraModel::filter_t lora_tensor_filter = module_filter; if (!lora_spec.tensor_name_prefix_filter.empty()) { lora_tensor_filter = [module_filter, prefix = lora_spec.tensor_name_prefix_filter](const std::string& tensor_name) { @@ -3120,8 +3140,8 @@ class StableDiffusionGGML { return latent_frames_to_video_frames(video_frames_to_latent_frames(frames)); } - sd::Tensor encode_to_vae_latents(const sd::Tensor& x, bool encode_video = false) { - auto latents = first_stage_model->encode(n_threads, x, vae_tiling_params, encode_video, circular_x, circular_y); + sd::Tensor encode_to_vae_latents(const sd::Tensor& x) { + auto latents = first_stage_model->encode(n_threads, x, vae_tiling_params, circular_x, circular_y); if (latents.empty()) { return {}; } @@ -3129,8 +3149,8 @@ class StableDiffusionGGML { return latents; } - sd::Tensor encode_first_stage(const sd::Tensor& x, bool encode_video = false) { - auto latents = encode_to_vae_latents(x, encode_video); + sd::Tensor encode_first_stage(const sd::Tensor& x) { + auto latents = encode_to_vae_latents(x); if (latents.empty()) { return {}; } @@ -3932,6 +3952,35 @@ static bool sd_version_supports_image_generation(SDVersion version) { return version != VERSION_ABOT_WORLD && !sd_version_supports_video_generation(version); } +static void sd_vid_gen_params_from_image_request(sd_vid_gen_params_t* video, + const sd_img_gen_params_t& image, + int video_frames) { + sd_vid_gen_params_init(video); + video->loras = image.loras; + video->lora_count = image.lora_count; + video->prompt = image.prompt; + video->negative_prompt = image.negative_prompt; + video->clip_skip = image.clip_skip; + video->init_image = image.init_image; + video->ref_images = image.ref_images; + video->ref_images_count = image.ref_images_count; + video->width = image.width; + video->height = image.height; + video->sample_params = image.sample_params; + video->strength = image.strength; + video->seed = image.seed; + video->video_frames = std::max(video_frames, 1); + video->vae_tiling_params = image.vae_tiling_params; + video->cache = image.cache; + video->hires = image.hires; + video->circular_x = image.circular_x; + video->circular_y = image.circular_y; + if (image.control_image.data != nullptr) { + video->control_frames = const_cast(&image.control_image); + video->control_frames_size = 1; + } +} + sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) { sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t)); if (sd_ctx == nullptr) { @@ -3970,6 +4019,287 @@ SD_API void sd_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode) { } } +void sd_fit_workload_init(sd_fit_workload_t* workload) { + if (workload == nullptr) { + return; + } + *workload = {}; + workload->width = 512; + workload->height = 512; + workload->video_frames = 1; + workload->vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr}; +} + +enum sd_fit_status_t sd_fit_params(const sd_ctx_params_t* sd_ctx_params, + const sd_fit_workload_t* workload, + sd_fit_result_t* result) { + if (result == nullptr) { + return SD_FIT_ERROR; + } + *result = {}; + if (sd_ctx_params == nullptr || workload == nullptr) { + return SD_FIT_ERROR; + } + if (workload->image_gen_params != nullptr && workload->video_gen_params != nullptr) { + LOG_ERROR("fit-params: set at most one complete generation request"); + return SD_FIT_ERROR; + } + if (strlen(SAFE_STR(sd_ctx_params->backend)) > 0 || + strlen(SAFE_STR(sd_ctx_params->params_backend)) > 0) { + LOG_WARN("fit-params: explicit backend placement cannot be validated; clear backend and params_backend before fitting"); + return SD_FIT_FAILURE; + } + + int64_t t0 = ggml_time_ms(); + + sd_ctx_params_t dry_params = *sd_ctx_params; + dry_params.auto_fit = false; + dry_params.eager_load = false; + + sd_ctx_t sd_ctx_storage{}; + 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)) { + LOG_ERROR("fit-params: dry-run model init failed"); + delete sd_ctx->sd; + sd_ctx->sd = nullptr; + return SD_FIT_ERROR; + } + + const int workload_width = workload->image_gen_params != nullptr ? workload->image_gen_params->width + : workload->video_gen_params != nullptr ? workload->video_gen_params->width + : workload->width; + const int workload_height = workload->image_gen_params != nullptr ? workload->image_gen_params->height + : workload->video_gen_params != nullptr ? workload->video_gen_params->height + : workload->height; + const int workload_frames = workload->video_gen_params != nullptr ? workload->video_gen_params->video_frames + : workload->video_frames; + const char* prompt = workload->prompt != nullptr && workload->prompt[0] != '\0' + ? workload->prompt + : "a photo of an astronaut riding a horse on the moon"; + const bool model_video_only = sd_version_supports_video_generation(sd_ctx->sd->version); + const bool animatediff_capable = sd_ctx->sd->animatediff_loaded && + sd_version_supports_animatediff(sd_ctx->sd->version); + const bool animatediff_video = animatediff_capable && workload_frames > 1; + if ((workload->video_gen_params != nullptr || workload_frames > 1) && + !model_video_only && !animatediff_capable) { + LOG_ERROR("fit-params: a video workload was supplied for an image-only model"); + delete sd_ctx->sd; + sd_ctx->sd = nullptr; + return SD_FIT_ERROR; + } + const bool video = workload->video_gen_params != nullptr || model_video_only || animatediff_video; + + sd_tiling_params_t requested_tiling = workload->vae_tiling_params; + if (workload->image_gen_params != nullptr) { + requested_tiling = workload->image_gen_params->vae_tiling_params; + } else if (workload->video_gen_params != nullptr) { + requested_tiling = workload->video_gen_params->vae_tiling_params; + } + + // Silence progress from this dry run without replacing the process-wide + // callback used by concurrent generation calls. + struct ProgressSuppressionGuard { + ProgressSuppressionGuard() + : previous(sd_get_progress_suppressed()) { + sd_set_progress_suppressed(true); + } + ~ProgressSuppressionGuard() { + sd_set_progress_suppressed(previous); + } + bool previous; + } progress_suppression_guard; + + // run the real generation pipeline in measure mode: every runner builds its graphs, + // records memory requirements and returns shaped zero tensors, no weights are read + // models with a vision tower condition on an input image; feed a dummy one so + // the clip_vision and VAE encode graphs are built and measured too + std::vector dummy_image_data; + sd_image_t dummy_init_image = {0, 0, 3, nullptr}; + if (sd_ctx->sd->clip_vision != nullptr && workload_width > 0 && workload_height > 0) { + dummy_image_data.assign((size_t)workload_width * workload_height * 3, 128); + dummy_init_image = {(uint32_t)workload_width, (uint32_t)workload_height, 3, dummy_image_data.data()}; + } + + auto measure = [&](const sd_tiling_params_t& tiling, + std::vector& records) -> bool { + records.clear(); + struct MeasureModeGuard { + explicit MeasureModeGuard(std::vector* records) { + GGMLRunner::set_measure_mode(true, records); + } + ~MeasureModeGuard() { + GGMLRunner::set_measure_mode(false); + } + } measure_mode_guard(&records); + bool ok = false; + sd_image_t* images = nullptr; + int num_images = 0; + sd_audio_t* audio = nullptr; + try { + if (video) { + sd_vid_gen_params_t gen; + if (workload->video_gen_params != nullptr) { + gen = *workload->video_gen_params; + } else if (workload->image_gen_params != nullptr) { + LOG_WARN("fit-params: promoting the image request to the video pipeline required by this model"); + sd_vid_gen_params_from_image_request(&gen, + *workload->image_gen_params, + workload_frames); + } else { + sd_vid_gen_params_init(&gen); + gen.prompt = prompt; + gen.width = workload_width; + gen.height = workload_height; + gen.video_frames = std::max(workload_frames, 1); + } + gen.sample_params.sample_steps = 2; + gen.sample_params.custom_sigmas = nullptr; + gen.sample_params.custom_sigmas_count = 0; + if (gen.high_noise_sample_params.sample_steps > 0) { + gen.high_noise_sample_params.sample_steps = 2; + gen.high_noise_sample_params.custom_sigmas = nullptr; + gen.high_noise_sample_params.custom_sigmas_count = 0; + } + gen.vae_tiling_params = tiling; + if (gen.init_image.data == nullptr && dummy_init_image.data != nullptr) { + gen.init_image = dummy_init_image; + } + ok = generate_video(sd_ctx, &gen, &images, &num_images, &audio); + } else { + sd_img_gen_params_t gen; + if (workload->image_gen_params != nullptr) { + gen = *workload->image_gen_params; + } else { + sd_img_gen_params_init(&gen); + gen.prompt = prompt; + gen.width = workload_width; + gen.height = workload_height; + } + gen.sample_params.sample_steps = 2; + gen.sample_params.custom_sigmas = nullptr; + gen.sample_params.custom_sigmas_count = 0; + gen.batch_count = 1; + gen.vae_tiling_params = tiling; + if (gen.ip_adapter_image.data == nullptr && dummy_init_image.data != nullptr) { + gen.ip_adapter_image = dummy_init_image; // image models have clip_vision only for ip-adapter + } + ok = generate_image(sd_ctx, &gen, &images, &num_images); + } + } catch (const std::bad_alloc&) { + LOG_ERROR("fit-params: host memory exhausted while materializing dry-run pipeline tensors"); + ok = false; + } catch (const std::exception& error) { + LOG_ERROR("fit-params: dry-run pipeline failed: %s", error.what()); + ok = false; + } + free_sd_audio(audio); + if (images != nullptr) { + for (int i = 0; i < num_images; i++) { + free(images[i].data); + } + free(images); + } + return ok; + }; + + std::vector records; + if (!measure(requested_tiling, records) || records.empty()) { + LOG_ERROR("fit-params: measurement dry run failed"); + delete sd_ctx->sd; + sd_ctx->sd = nullptr; + return SD_FIT_ERROR; + } + + std::map module_map; + for (const auto& kv : sd_ctx->sd->fit_module_params_bytes) { + auto& m = module_map[kv.first]; + m.module = kv.first; + m.params_bytes = kv.second; + } + for (const auto& record : records) { + SDBackendModule module = record.module; + if (module == SDBackendModule::UNSET) { + LOG_WARN("fit-params: ignoring unattributed graph measurement from %s", + record.desc.c_str()); + continue; + } + auto& m = module_map[module]; + m.module = module; + m.params_bytes = std::max(m.params_bytes, record.params_bytes); + m.compute_bytes = std::max(m.compute_bytes, record.compute_bytes); + if (!record.split_segment_params_bytes.empty()) { + m.split_graph_segment_params.push_back(record.split_segment_params_bytes); + m.split_graph_segment_compute.push_back(record.split_segment_compute_bytes); + } + } + for (auto module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) { + auto it = module_map.find(module); + if (it != module_map.end()) { + it->second.splittable = true; + } + } + + // price VAE tiling so the planner can fall back to it when full-resolution decode does not fit + { + auto it = module_map.find(SDBackendModule::VAE); + if (it != module_map.end() && !requested_tiling.enabled) { + sd_tiling_params_t tiled = requested_tiling; + tiled.enabled = true; + std::vector tiled_records; + if (measure(tiled, tiled_records)) { + for (const auto& record : tiled_records) { + if (record.module == SDBackendModule::VAE) { + it->second.compute_bytes_tiled = std::max(it->second.compute_bytes_tiled, record.compute_bytes); + } + } + } + } + } + + std::vector modules; + for (const auto& kv : module_map) { + modules.push_back(kv.second); + } + sd::fit_params::FitPlan plan; + bool planned = sd::fit_params::plan_placement(modules, sd_ctx->sd->max_vram_assignment, &plan); + + delete sd_ctx->sd; + sd_ctx->sd = nullptr; + + result->report = strdup(plan.report.c_str()); + result->vae_tiling = plan.vae_tiling; + result->stream_layers = plan.stream_layers; + if (!planned || !plan.valid) { + return SD_FIT_FAILURE; + } + + if (plan.changed) { + result->changed = true; + if (!plan.runtime_spec.empty()) { + result->backend = strdup(plan.runtime_spec.c_str()); + } + if (!plan.params_spec.empty()) { + result->params_backend = strdup(plan.params_spec.c_str()); + } + } + + int64_t t1 = ggml_time_ms(); + LOG_INFO("fit-params: fitting params to free memory took %.2fs", (t1 - t0) * 1.0f / 1000); + return SD_FIT_SUCCESS; +} + +void sd_fit_result_free(sd_fit_result_t* result) { + if (result == nullptr) { + return; + } + free(result->backend); + free(result->params_backend); + free(result->report); + *result = {}; +} + static sd_audio_t* waveform_to_sd_audio(const StableDiffusionGGML* sd, const sd::Tensor& waveform) { if (sd == nullptr || waveform.empty()) { @@ -6508,7 +6838,7 @@ static std::optional prepare_video_generation_latents(sd sd::ops::slice_assign(&image, 2, request->frames - 1, request->frames, end_image.unsqueeze(2)); } - auto concat_latent = sd_ctx->sd->encode_first_stage(image, /*encode_video=*/true); // [b, c, t, h/vae_scale_factor, w/vae_scale_factor]; encode_video bypasses spatial tiling + auto concat_latent = sd_ctx->sd->encode_first_stage(image); // [b, c, t, h/vae_scale_factor, w/vae_scale_factor] if (concat_latent.empty()) { LOG_ERROR("failed to encode video conditioning frames"); return std::nullopt; diff --git a/src/upscaler.cpp b/src/upscaler.cpp index 007396245..7cc69f22c 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -96,6 +96,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, LOG_ERROR("init esrgan model from metadata failed: '%s'", esrgan_path.c_str()); return false; } + esrgan_upscaler->set_fit_module(SDBackendModule::UPSCALER); esrgan_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); esrgan_upscaler->set_stream_layers_enabled(stream_layers_enabled); if (direct) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index fb62a7840..74c60e8d0 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -6,6 +6,11 @@ add_executable(test-vae-routing test-vae-routing.cpp) target_link_libraries(test-vae-routing PRIVATE stable-diffusion) add_test(NAME test-vae-routing COMMAND test-vae-routing) +add_executable(test-fit-params test-fit-params.cpp) +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-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) diff --git a/tests/test-fit-params.cpp b/tests/test-fit-params.cpp new file mode 100644 index 000000000..0206cf98b --- /dev/null +++ b/tests/test-fit-params.cpp @@ -0,0 +1,341 @@ +#include +#include +#include +#include +#include +#include + +#include "core/fit_params.h" +#include "core/ggml_extend.hpp" + +namespace { + +constexpr size_t GiB = 1024ull * 1024ull * 1024ull; + +bool expect(bool condition, const char* message) { + if (!condition) { + std::cerr << message << std::endl; + return false; + } + return true; +} + +bool set_test_env(const char* name, const std::string& value) { +#if defined(_WIN32) + return _putenv_s(name, value.c_str()) == 0; +#else + return setenv(name, value.c_str(), 1) == 0; +#endif +} + +void unset_test_env(const char* name) { +#if defined(_WIN32) + _putenv_s(name, ""); +#else + unsetenv(name); +#endif +} + +struct MeasureRunner : public GGMLRunner { + bool warm_cache_seen = false; + ggml_tensor* output = nullptr; + + explicit MeasureRunner(ggml_backend_t backend) + : GGMLRunner(backend) { + set_fit_module(SDBackendModule::VAE); + } + + std::string get_desc() override { + return "fit measurement test"; + } + + std::optional> run(bool no_return = false) { + const sd::Tensor input = sd::zeros({2, 3}); + const sd::Tensor initial_cache = sd::zeros({5}); + auto get_graph = [&]() { + ggml_cgraph* graph = new_graph_custom(32); + ggml_tensor* x = make_input(input); + output = ggml_scale(compute_ctx, x, 2.f); + ggml_build_forward_expand(graph, output); + + ggml_tensor* previous_cache = get_cache_tensor_by_name("state"); + warm_cache_seen = previous_cache != nullptr; + ggml_tensor* cache_input = previous_cache != nullptr + ? previous_cache + : make_input(initial_cache); + cache("state", ggml_scale(compute_ctx, cache_input, 2.f)); + return graph; + }; + return compute(get_graph, 1, false, true, true, no_return); + } +}; + +sd::fit_params::ModuleMemory module(SDBackendModule module, + size_t params_gib, + size_t compute_gib, + bool splittable = false, + size_t tiled_compute_mib = 0) { + sd::fit_params::ModuleMemory memory; + memory.module = module; + memory.params_bytes = params_gib * GiB; + memory.compute_bytes = compute_gib * GiB; + memory.splittable = splittable; + memory.compute_bytes_tiled = tiled_compute_mib * 1024ull * 1024ull; + if (splittable && params_gib > 0) { + std::vector segment_params(params_gib, GiB); + std::vector segment_compute(params_gib, std::min(compute_gib, 1) * GiB); + memory.split_graph_segment_params.push_back(std::move(segment_params)); + memory.split_graph_segment_compute.push_back(std::move(segment_compute)); + } + return memory; +} + +bool plan_with_devices(const char* devices, + float max_vram_gib, + const std::vector& modules, + sd::fit_params::FitPlan* plan, + float host_memory_gib = 64.f) { + if (!set_test_env("SD_FIT_DEBUG_DEVICES", devices) || + !set_test_env("SD_FIT_DEBUG_HOST_MEMORY_GIB", std::to_string(host_memory_gib))) { + return false; + } + sd::ggml_graph_cut::MaxVramAssignment budgets; + budgets.reset(max_vram_gib); + return sd::fit_params::plan_placement(modules, budgets, plan); +} + +bool test_default_fits() { + sd::fit_params::FitPlan plan; + bool ok = plan_with_devices("GPU0:8", 8.f, + {module(SDBackendModule::DIFFUSION, 1, 1, true), + module(SDBackendModule::TE, 1, 1, true), + module(SDBackendModule::VAE, 1, 1)}, + &plan); + return expect(ok && plan.valid, "default fit plan should be valid") && + expect(!plan.changed, "default fit should not emit placement changes") && + expect(plan.runtime_spec.empty(), "default fit runtime spec should be empty") && + expect(plan.params_spec.empty(), "default fit params spec should be empty"); +} + +bool test_resident_spread() { + sd::fit_params::FitPlan plan; + bool ok = plan_with_devices("GPU0:9,GPU1:9", 9.f, + {module(SDBackendModule::DIFFUSION, 5, 2, true), + module(SDBackendModule::TE, 3, 1, true), + module(SDBackendModule::VAE, 2, 3)}, + &plan); + return expect(ok && plan.valid, "resident spread plan should be valid") && + expect(plan.changed, "resident spread should emit placement changes") && + expect(!plan.time_share, "resident spread should not time-share params") && + expect(plan.runtime_spec == "diffusion=GPU0,te=GPU1,vae=GPU1", + ("unexpected resident runtime spec: " + plan.runtime_spec).c_str()) && + expect(plan.params_spec.empty(), "resident spread should not place params on disk"); +} + +bool test_time_share_cpu_fallback() { + sd::fit_params::FitPlan plan; + bool ok = plan_with_devices("GPU0:6", 0.f, + {module(SDBackendModule::DIFFUSION, 5, 2, true), + module(SDBackendModule::TE, 1, 1, true), + module(SDBackendModule::VAE, 1, 2)}, + &plan); + return expect(ok && plan.valid, "time-share plan should be valid") && + expect(plan.time_share, "oversized resident plan should time-share") && + expect(plan.runtime_spec == "diffusion=cpu,te=GPU0,vae=GPU0", + ("unexpected time-share runtime spec: " + plan.runtime_spec).c_str()) && + expect(plan.params_spec == "te=disk,vae=disk", + ("unexpected time-share params spec: " + plan.params_spec).c_str()); +} + +bool test_stream_layers_after_split_fails() { + sd::fit_params::FitPlan plan; + bool ok = plan_with_devices("GPU0:6,GPU1:6", 6.f, + {module(SDBackendModule::DIFFUSION, 20, 2, true)}, + &plan); + return expect(ok && plan.valid, "streaming plan should be valid") && + expect(plan.time_share, "streaming plan should be a time-share fallback") && + expect(plan.stream_layers, "streaming plan should request stream layers") && + expect(plan.runtime_spec == "diffusion=GPU0", + ("unexpected streaming runtime spec: " + plan.runtime_spec).c_str()) && + expect(plan.params_spec == "diffusion=cpu", + ("unexpected streaming params spec: " + plan.params_spec).c_str()); +} + +bool test_split_and_tiling() { + sd::fit_params::FitPlan split_plan; + bool split_ok = plan_with_devices("GPU0:7,GPU1:7", 7.f, + {module(SDBackendModule::DIFFUSION, 8, 2, true)}, + &split_plan); + if (!expect(split_ok && split_plan.valid, "split plan should be valid") || + !expect(split_plan.runtime_spec == "diffusion=GPU0&GPU1", + ("unexpected split runtime spec: " + split_plan.runtime_spec).c_str()) || + !expect(split_plan.params_spec == "diffusion=disk", + ("unexpected split params spec: " + split_plan.params_spec).c_str())) { + return false; + } + + sd::fit_params::FitPlan tiling_plan; + bool tiling_ok = plan_with_devices("GPU0:4", 4.f, + {module(SDBackendModule::VAE, 1, 5, false, 512)}, + &tiling_plan); + return expect(tiling_ok && tiling_plan.valid, "tiling plan should be valid") && + expect(tiling_plan.vae_tiling, "tiling plan should request VAE tiling") && + expect(tiling_plan.runtime_spec == "vae=GPU0", + ("unexpected tiling runtime spec: " + tiling_plan.runtime_spec).c_str()) && + expect(tiling_plan.params_spec.empty(), + ("unexpected tiling params spec: " + tiling_plan.params_spec).c_str()); +} + +bool test_split_rejects_indivisible_segment() { + auto memory = module(SDBackendModule::DIFFUSION, 8, 2, true); + memory.split_graph_segment_params = {{5 * GiB, 3 * GiB}}; + memory.split_graph_segment_compute = {{1 * GiB, 1 * GiB}}; + + sd::fit_params::FitPlan plan; + bool ok = plan_with_devices("GPU0:7,GPU1:7", 7.f, {memory}, &plan); + return expect(ok && plan.valid, "indivisible split fallback should remain valid") && + expect(plan.stream_layers, "indivisible split should fall back to streaming") && + expect(plan.runtime_spec == "diffusion=GPU0", + ("unexpected indivisible fallback runtime spec: " + plan.runtime_spec).c_str()) && + expect(plan.params_spec == "diffusion=cpu", + ("unexpected indivisible fallback params spec: " + plan.params_spec).c_str()); +} + +bool test_explicit_budget_keeps_headroom() { + sd::fit_params::FitPlan plan; + bool ok = plan_with_devices("GPU0:8", 8.f, + {module(SDBackendModule::VAE, 7, 1)}, + &plan); + return expect(ok && plan.valid, "headroom fallback plan should be valid") && + expect(plan.changed, "explicit max-vram must retain safety headroom") && + expect(plan.runtime_spec == "vae=cpu", + ("unexpected headroom fallback runtime spec: " + plan.runtime_spec).c_str()); +} + +bool test_controlnet_compute_is_concurrent() { + sd::fit_params::FitPlan plan; + bool ok = plan_with_devices("GPU0:8", 8.f, + {module(SDBackendModule::DIFFUSION, 2, 2, true), + module(SDBackendModule::CONTROL_NET, 2, 2)}, + &plan); + return expect(ok && plan.valid, "ControlNet plan should remain valid") && + expect(plan.changed, "concurrent ControlNet and diffusion buffers must not use the default placement") && + expect(plan.time_share, "concurrent ControlNet pressure should require the time-share tier"); +} + +bool test_cpu_fallback_checks_host_memory() { + sd::fit_params::FitPlan plan; + bool ok = plan_with_devices("GPU0:4", 4.f, + {module(SDBackendModule::DIFFUSION, 8, 2, true)}, + &plan, + 8.f); + return expect(ok, "host-capacity failure should be a completed planning attempt") && + expect(!plan.valid, "CPU fallback must fail when projected use exceeds host memory") && + expect(plan.report.find("no placement fits available host memory") != std::string::npos, + "host-capacity failure should be explained in the report"); +} + +bool test_measure_mode_preserves_outputs_and_projects_cache() { + ggml_backend_t backend = sd_backend_cpu_init(); + if (!expect(backend != nullptr, "CPU backend should initialize for measurement test")) { + return false; + } + + bool passed = true; + { + MeasureRunner runner(backend); + std::vector records; + GGMLRunner::set_measure_mode(true, &records); + + auto first = runner.run(); + passed &= expect(first.has_value() && first->dim() == 2, + "measure output should preserve the named result rank"); + passed &= expect(first.has_value() && first->shape()[0] == 2 && first->shape()[1] == 3, + "measure output should preserve the named result shape"); + passed &= expect(!records.empty() && records.back().cache_bytes > 0, + "measurements should include projected persistent cache bytes"); + + auto second = runner.run(true); + passed &= expect(second.has_value() && runner.warm_cache_seen, + "the next measured graph should observe projected warm cache state"); + passed &= expect(runner.output != nullptr && ggml_n_dims(runner.output) == 2, + "no-return graph tensors should remain alive for the caller"); + + GGMLRunner::set_measure_mode(false); + } + ggml_backend_free(backend); + return passed; +} + +bool test_measure_mode_is_thread_local() { + std::vector records; + GGMLRunner::set_measure_mode(true, &records); + const bool enabled_on_calling_thread = GGMLRunner::measure_mode_enabled(); + bool enabled_on_other_thread = true; + std::thread other_thread([&]() { + enabled_on_other_thread = GGMLRunner::measure_mode_enabled(); + }); + other_thread.join(); + GGMLRunner::set_measure_mode(false); + + return expect(enabled_on_calling_thread, "measurement should be enabled on the fitting thread") && + expect(!enabled_on_other_thread, "measurement must not affect another thread"); +} + +bool test_public_result_is_initialized_on_error() { + sd_fit_workload_t workload; + sd_fit_workload_init(&workload); + if (!expect(workload.image_gen_params == nullptr && workload.video_gen_params == nullptr, + "fit workload request pointers should default to null")) { + return false; + } + + sd_fit_result_t result{}; + result.changed = true; + result.backend = reinterpret_cast(1); + result.params_backend = reinterpret_cast(1); + result.report = reinterpret_cast(1); + const auto status = sd_fit_params(nullptr, &workload, &result); + return expect(status == SD_FIT_ERROR, "invalid fit arguments should return SD_FIT_ERROR") && + expect(!result.changed && result.backend == nullptr && result.params_backend == nullptr && result.report == nullptr, + "fit result should be initialized before argument validation"); +} + +bool test_public_rejects_explicit_placement() { + sd_ctx_params_t params; + sd_ctx_params_init(¶ms); + params.backend = "cpu"; + sd_fit_workload_t workload; + sd_fit_workload_init(&workload); + sd_fit_result_t result; + + const enum sd_fit_status_t status = sd_fit_params(¶ms, &workload, &result); + const bool passed = expect(status == SD_FIT_FAILURE, "explicit placement should be rejected") && + expect(!result.changed && result.backend == nullptr && result.params_backend == nullptr && result.report == nullptr, + "placement rejection should leave an initialized result"); + sd_fit_result_free(&result); + return passed; +} + +} // namespace + +int main() { + if (!test_default_fits() || + !test_resident_spread() || + !test_time_share_cpu_fallback() || + !test_stream_layers_after_split_fails() || + !test_split_and_tiling() || + !test_split_rejects_indivisible_segment() || + !test_explicit_budget_keeps_headroom() || + !test_controlnet_compute_is_concurrent() || + !test_cpu_fallback_checks_host_memory() || + !test_measure_mode_preserves_outputs_and_projects_cache() || + !test_measure_mode_is_thread_local() || + !test_public_rejects_explicit_placement() || + !test_public_result_is_initialized_on_error()) { + return 1; + } + unset_test_env("SD_FIT_DEBUG_DEVICES"); + unset_test_env("SD_FIT_DEBUG_HOST_MEMORY_GIB"); + return 0; +}