sd: sync with master-767-885f01a (#2310)

* sd: minor API path handling cleanup

* sd: sync with master-749-b11c95a

* sd: use original API parameters at the internal C++ API

* sd: split_mode and auto_fit backend support

* sd: sync with master-758-c674225

* sd: sync with master-765-bb84971

* sd: sync with master-767-885f01a
This commit is contained in:
Wagner Bruna 2026-07-08 06:08:19 -03:00 committed by GitHub
parent dad4ff5737
commit e11d3ddef0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
39 changed files with 3124 additions and 339 deletions

View file

@ -703,7 +703,7 @@ budget.o: common/reasoning-budget.cpp common/reasoning-budget.h
chat.o: common/chat.cpp common/chat.h
$(CXX) $(CXXFLAGS) -c $< -o $@
SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.hpp src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/ordered_map.hpp src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/kcpp_sd_extensions.h src/model/adapter/lora.hpp src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/common/block.hpp src/model/common/rope.hpp src/model/diffusion/anima.hpp src/model/diffusion/boogu.hpp src/model/diffusion/control.hpp src/model/diffusion/dit.hpp src/model/diffusion/ernie_image.hpp src/model/diffusion/flux.hpp src/model/diffusion/hidream_o1.hpp src/model/diffusion/ideogram4.hpp src/model/diffusion/krea2.hpp src/model/diffusion/lens.hpp src/model/diffusion/ltxv.hpp src/model/diffusion/minit2i.hpp src/model/diffusion/mmdit.hpp src/model/diffusion/model.hpp src/model/diffusion/pid.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/sefi_image.hpp src/model/diffusion/unet.hpp src/model/diffusion/wan.hpp src/model/diffusion/z_image.hpp src/model.h src/model_io/binary_io.h src/model_io/gguf_io.cpp src/model_io/gguf_io.h src/model_io/gguf_reader_ext.h src/model_io/pickle_io.cpp src/model_io/pickle_io.h src/model_io/safetensors_io.cpp src/model_io/safetensors_io.h src/model_io/tensor_storage.h src/model_io/torch_legacy_io.cpp src/model_io/torch_legacy_io.h src/model_io/torch_zip_io.cpp src/model_io/torch_zip_io.h src/model_loader.cpp src/model_loader.h src/model_manager.cpp src/model_manager.h src/model/te/clip.hpp src/model/te/llm.hpp src/model/te/t5.hpp src/model/upscaler/esrgan.hpp src/model/upscaler/ltx_latent_upscaler.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/runtime/cache_dit.hpp src/runtime/condition_cache_utils.hpp src/runtime/denoiser.hpp src/runtime/easycache.hpp src/runtime/gits_noise.h src/runtime/guidance.cpp src/runtime/guidance.h src/runtime/imatrix.cpp src/runtime/imatrix.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/ucache.hpp src/stable-diffusion.cpp src/tokenizers/bpe_tokenizer.cpp src/tokenizers/bpe_tokenizer.h src/tokenizers/clip_tokenizer.cpp src/tokenizers/clip_tokenizer.h src/tokenizers/gemma_tokenizer.cpp src/tokenizers/gemma_tokenizer.h src/tokenizers/gpt_oss_tokenizer.cpp src/tokenizers/gpt_oss_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer.cpp src/tokenizers/tokenizer.h src/tokenizers/tokenize_util.cpp src/tokenizers/tokenize_util.h src/tokenizers/vocab/vocab.h src/upscaler.cpp src/upscaler.h src/weight_manager.h
SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/core/backend_fit.cpp src/core/backend_fit.h src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.hpp src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/layer_split_partition.cpp src/core/layer_split_partition.h src/core/ordered_map.hpp src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/kcpp_sd_extensions.h src/model/adapter/lora.hpp src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/common/block.hpp src/model/common/rope.hpp src/model/diffusion/anima.hpp src/model/diffusion/boogu.hpp src/model/diffusion/control.hpp src/model/diffusion/dit.hpp src/model/diffusion/ernie_image.hpp src/model/diffusion/flux.hpp src/model/diffusion/hidream_o1.hpp src/model/diffusion/ideogram4.hpp src/model/diffusion/krea2.hpp src/model/diffusion/lens.hpp src/model/diffusion/ltxv.hpp src/model/diffusion/minit2i.hpp src/model/diffusion/mmdit.hpp src/model/diffusion/model.hpp src/model/diffusion/pid.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/sefi_image.hpp src/model/diffusion/unet.hpp src/model/diffusion/wan.hpp src/model/diffusion/z_image.hpp src/model.h src/model_io/binary_io.h src/model_io/gguf_io.cpp src/model_io/gguf_io.h src/model_io/gguf_reader_ext.h src/model_io/pickle_io.cpp src/model_io/pickle_io.h src/model_io/safetensors_io.cpp src/model_io/safetensors_io.h src/model_io/streaming_writer.h src/model_io/tensor_storage.h src/model_io/torch_legacy_io.cpp src/model_io/torch_legacy_io.h src/model_io/torch_zip_io.cpp src/model_io/torch_zip_io.h src/model_loader.cpp src/model_loader.h src/model_manager.cpp src/model_manager.h src/model/te/clip.hpp src/model/te/llm.hpp src/model/te/t5.hpp src/model/upscaler/esrgan.hpp src/model/upscaler/ltx_latent_upscaler.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/runtime/cache_dit.hpp src/runtime/condition_cache_utils.hpp src/runtime/denoiser.hpp src/runtime/easycache.hpp src/runtime/gits_noise.h src/runtime/guidance.cpp src/runtime/guidance.h src/runtime/imatrix.cpp src/runtime/imatrix.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/ucache.hpp src/stable-diffusion.cpp src/tokenizers/bpe_tokenizer.cpp src/tokenizers/bpe_tokenizer.h src/tokenizers/clip_tokenizer.cpp src/tokenizers/clip_tokenizer.h src/tokenizers/gemma_tokenizer.cpp src/tokenizers/gemma_tokenizer.h src/tokenizers/gpt_oss_tokenizer.cpp src/tokenizers/gpt_oss_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer.cpp src/tokenizers/tokenizer.h src/tokenizers/tokenize_util.cpp src/tokenizers/tokenize_util.h src/tokenizers/vocab/vocab.h src/upscaler.cpp src/upscaler.h src/weight_manager.h
SDCPP_MAIN_BASENAMES := examples/cli/image_metadata.cpp examples/cli/image_metadata.h examples/cli/main.cpp examples/cli/msf_gif.h examples/common/common.cpp examples/common/common.h examples/common/log.cpp examples/common/log.h examples/common/media_io.cpp examples/common/media_io.h examples/common/resource_owners.hpp src/tokenizers/vocab/clip_merges.hpp src/tokenizers/vocab/gemma2_merges.hpp src/tokenizers/vocab/gemma2_vocab.hpp src/tokenizers/vocab/gemma_merges.hpp src/tokenizers/vocab/gemma_vocab.hpp src/tokenizers/vocab/gpt_oss_merges.hpp src/tokenizers/vocab/gpt_oss_vocab.hpp src/tokenizers/vocab/mistral_merges.hpp src/tokenizers/vocab/mistral_vocab.hpp src/tokenizers/vocab/qwen_merges.hpp src/tokenizers/vocab/t5.hpp src/tokenizers/vocab/umt5.hpp src/tokenizers/vocab/vocab.cpp src/convert.cpp src/version.cpp

View file

@ -186,14 +186,12 @@ struct sd_load_model_inputs
{
const char * model_filename = nullptr;
const char * executable_path = nullptr;
const int kcpp_main_device = -1;
const char * backend = nullptr;
const int threads = 0;
const int quant = 0;
const bool flash_attention = false;
const bool offload_cpu = false;
const char * params_backend = nullptr;
const bool use_mmap = false;
const int kcpp_vae_device = -1;
const int kcpp_clip_device = -1;
const bool diffusion_conv_direct = false;
const bool vae_conv_direct = false;
const bool taesd = false;
@ -211,8 +209,10 @@ struct sd_load_model_inputs
const char * upscaler_filename = nullptr;
const int img_hard_limit = 0;
const int img_soft_limit = 0;
const float max_vram = 0.f;
const char * max_vram = nullptr;
const char * split_mode = nullptr;
const bool stream_layers = false;
const bool auto_fit = false;
const char * devices_override = nullptr;
const bool quiet = false;
const int debugmode = 0;

View file

@ -378,14 +378,12 @@ class generation_outputs(ctypes.Structure):
class sd_load_model_inputs(ctypes.Structure):
_fields_ = [("model_filename", ctypes.c_char_p),
("executable_path", ctypes.c_char_p),
("kcpp_main_device", ctypes.c_int),
("backend", ctypes.c_char_p),
("threads", ctypes.c_int),
("quant", ctypes.c_int),
("flash_attention", ctypes.c_bool),
("offload_cpu", ctypes.c_bool),
("params_backend", ctypes.c_char_p),
("use_mmap", ctypes.c_bool),
("kcpp_vae_device", ctypes.c_int),
("kcpp_clip_device", ctypes.c_int),
("diffusion_conv_direct", ctypes.c_bool),
("vae_conv_direct", ctypes.c_bool),
("taesd", ctypes.c_bool),
@ -403,8 +401,10 @@ class sd_load_model_inputs(ctypes.Structure):
("upscaler_filename", ctypes.c_char_p),
("img_hard_limit", ctypes.c_int),
("img_soft_limit", ctypes.c_int),
("max_vram", ctypes.c_float),
("max_vram", ctypes.c_char_p),
("split_mode", ctypes.c_char_p),
("stream_layers", ctypes.c_bool),
("auto_fit", ctypes.c_bool),
("devices_override", ctypes.c_char_p),
("quiet", ctypes.c_bool),
("debugmode", ctypes.c_int)]
@ -2487,6 +2487,21 @@ def sd_resolve_device(name, default_=-1):
name = str(max(name, -2))
return sd_get_device_number(name)
def sd_get_device_override(deviceid, module=''):
'''formats a device id and a module name in sd.cpp --backend syntax'''
global cached_sd_info
devices = cached_sd_info.get('devices', [])
device_name = ''
if deviceid <= -2:
device_name = "CPU"
elif deviceid >= 0 and deviceid < len(devices):
device_name = devices[deviceid]['name']
if device_name and module:
result = module + '=' + device_name
else:
result = device_name
return result;
def sd_load_model(model_filename,vae_filename,t5xxl_filename,clip1_filename,clip2_filename,photomaker_filename,upscaler_filename,audio_vae_filename):
global args
inputs = sd_load_model_inputs()
@ -2501,10 +2516,14 @@ def sd_load_model(model_filename,vae_filename,t5xxl_filename,clip1_filename,clip
inputs.threads = thds
inputs.quant = args.sdquant
inputs.flash_attention = args.sdflashattention
inputs.offload_cpu = args.sdoffloadcpu
inputs.params_backend = b'CPU' if args.sdoffloadcpu else b''
inputs.use_mmap = args.usemmap
inputs.kcpp_vae_device = sd_resolve_device(args.sdvaedevice, default_sdvaedevice)
inputs.kcpp_clip_device = sd_resolve_device(args.sdclipdevice, default_sdclipdevice)
backends = [
sd_get_device_override(sd_resolve_device(args.sdmaingpu, 'main')),
sd_get_device_override(sd_resolve_device(args.sdclipdevice, default_sdclipdevice), 'CLIP'),
sd_get_device_override(sd_resolve_device(args.sdvaedevice, default_sdvaedevice), 'VAE'),
]
inputs.backend = ','.join([b for b in backends if b]).encode("UTF-8")
sdconvdirect = sd_convdirect_option(args.sdconvdirect)
inputs.diffusion_conv_direct = sdconvdirect == 'full'
inputs.vae_conv_direct = sdconvdirect in ['vaeonly', 'full']
@ -2517,7 +2536,7 @@ def sd_load_model(model_filename,vae_filename,t5xxl_filename,clip1_filename,clip
inputs.clip2_filename = clip2_filename.encode("UTF-8")
inputs.photomaker_filename = photomaker_filename.encode("UTF-8")
inputs.upscaler_filename = upscaler_filename.encode("UTF-8")
inputs.max_vram = (args.sdvramlimit/1024.0) if args.sdvramlimit > 0 else 0
inputs.max_vram = str((args.sdvramlimit/1024.0) if args.sdvramlimit > 0 else '').encode('UTF-8')
inputs.stream_layers = False
lora_filenames, lora_multipliers = prepare_initial_lora_multipliers()
@ -2535,7 +2554,6 @@ def sd_load_model(model_filename,vae_filename,t5xxl_filename,clip1_filename,clip
inputs.img_hard_limit = args.sdclamped
inputs.img_soft_limit = args.sdclampedsoft
inputs = set_backend_props(inputs)
inputs.kcpp_main_device = sd_resolve_device(args.sdmaingpu, 'main')
ret = handle.sd_load_model(inputs)
return ret
@ -6293,8 +6311,6 @@ Change Mode<br>
response_body = (json.dumps([{"name":name,"label":name} for name in cached_sd_info.get('available_schedulers', [])]).encode())
elif clean_path.endswith('/sdapi/v1/latent-upscale-modes'):
response_body = (json.dumps([]).encode())
elif clean_path.endswith('/sdapi/v1/upscalers'):
response_body = (json.dumps([]).encode())
#vits compatible
elif clean_path=='/voice/check':
@ -11867,8 +11883,8 @@ def kcpp_main_process(launch_args, g_memory=None, gui_launcher=False):
friendlysdmodelname = os.path.basename(imgmodel)
friendlysdmodelname = os.path.splitext(friendlysdmodelname)[0]
friendlysdmodelname = sanitize_string(friendlysdmodelname)
loadok = sd_load_model(imgmodel,imgvae,imgt5xxl,imgclip1,imgclip2,imgphotomaker,imgupscaler,imgaudiovae)
cached_sd_info = sd_get_info()
loadok = sd_load_model(imgmodel,imgvae,imgt5xxl,imgclip1,imgclip2,imgphotomaker,imgupscaler,imgaudiovae)
print("Load Image Model OK: " + str(loadok))
if not loadok:
exitcounter = 999

View file

@ -653,7 +653,8 @@ int main(int argc, const char* argv[]) {
cli_params.output_path.c_str(),
ctx_params.wtype,
ctx_params.tensor_type_rules.c_str(),
cli_params.convert_name);
cli_params.convert_name,
ctx_params.n_threads);
if (!success) {
LOG_ERROR("convert '%s'/'%s' to '%s' failed",
ctx_params.model_path.c_str(),

View file

@ -443,6 +443,12 @@ ArgOptions SDContextParams::get_options() {
"weight type per tensor pattern (example: \"^vae\\.=f16,model\\.=q8_0\")",
(int)',',
&tensor_type_rules},
{"",
"--model-args",
"extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, "
"chroma_t5_mask_pad, qwen_image_zero_cond_t",
(int)',',
&model_args},
{"",
"--photo-maker",
"path to PHOTOMAKER model",
@ -468,6 +474,13 @@ ArgOptions SDContextParams::get_options() {
"parameter backend assignment, e.g. disk, cpu, or diffusion=disk,clip=cpu",
(int)',',
&params_backend},
{"",
"--split-mode",
"weight distribution for modules assigned multiple devices (--backend \"diffusion=cuda0&cuda1\"): "
"layer (whole transformer blocks per device, default) or row (matmul rows split across devices, CUDA only). "
"Accepts a single mode or per-module assignments, e.g. row or diffusion=row,te=layer",
(int)',',
&split_mode},
{"",
"--rpc-servers",
"comma-separated list of RPC servers to connect to for offloading, in the format host:port, e.g. localhost:50052,192.168.1.3:50052",
@ -486,10 +499,6 @@ ArgOptions SDContextParams::get_options() {
"number of threads to use during computation (default: -1). "
"If threads <= 0, then threads will be set to the number of CPU physical cores",
&n_threads},
{"",
"--chroma-t5-mask-pad",
"t5 mask pad size of chroma",
&chroma_t5_mask_pad},
};
options.bool_options = {
@ -501,6 +510,12 @@ ArgOptions SDContextParams::get_options() {
"--eager-load",
"load all params into the params backend at model-load time instead of lazily on first use (defaults to false)",
true, &eager_load},
{"",
"--auto-fit",
"pick the diffusion/te/vae device placements automatically from the model size and the per-device "
"memory budgets (--max-vram; defaults to free memory minus a small margin). Overrides --backend and "
"--params-backend; may split modules across GPUs (--split-mode still selects layer or row)",
true, &auto_fit},
{"",
"--force-sdxl-vae-conv-scale",
"force use of conv scale on sdxl vae",
@ -541,30 +556,6 @@ ArgOptions SDContextParams::get_options() {
"--vae-conv-direct",
"use ggml_conv2d_direct in the vae model",
true, &vae_conv_direct},
{"",
"--circular",
"enable circular padding for convolutions",
true, &circular},
{"",
"--circularx",
"enable circular RoPE wrapping on x-axis (width) only",
true, &circular_x},
{"",
"--circulary",
"enable circular RoPE wrapping on y-axis (height) only",
true, &circular_y},
{"",
"--chroma-disable-dit-mask",
"disable dit mask for chroma",
false, &chroma_use_dit_mask},
{"",
"--qwen-image-zero-cond-t",
"enable zero_cond_t for qwen image",
true, &qwen_image_zero_cond_t},
{"",
"--chroma-enable-t5-mask",
"enable t5 mask for chroma",
true, &chroma_use_t5_mask},
};
auto on_type_arg = [&](int argc, const char** argv, int index) {
@ -663,6 +654,18 @@ ArgOptions SDContextParams::get_options() {
"but it usually offers faster inference speed and, in some cases, lower memory usage. "
"The at_runtime mode, on the other hand, is exactly the opposite.",
on_lora_apply_mode_arg},
{"",
"--list-devices",
"list available ggml backend devices (one 'name<TAB>description' per line) and exit; "
"the names are the device names accepted by --backend and --params-backend",
[](int /*argc*/, const char** /*argv*/, int /*index*/) {
size_t device_list_size = sd_list_devices(nullptr, 0);
std::vector<char> devices(device_list_size + 1);
sd_list_devices(devices.data(), devices.size());
fputs(devices.data(), stdout);
std::exit(0);
return 0;
}},
};
return options;
@ -818,6 +821,9 @@ std::string SDContextParams::to_string() const {
<< " eager_load: " << (eager_load ? "true" : "false") << ",\n"
<< " backend: \"" << backend << "\",\n"
<< " params_backend: \"" << params_backend << "\",\n"
<< " split_mode: \"" << split_mode << "\",\n"
<< " model_args: \"" << model_args << "\",\n"
<< " auto_fit: " << (auto_fit ? "true" : "false") << ",\n"
<< " enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
<< " control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n"
<< " clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n"
@ -826,13 +832,6 @@ std::string SDContextParams::to_string() const {
<< " diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n"
<< " diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n"
<< " vae_conv_direct: " << (vae_conv_direct ? "true" : "false") << ",\n"
<< " circular: " << (circular ? "true" : "false") << ",\n"
<< " circular_x: " << (circular_x ? "true" : "false") << ",\n"
<< " circular_y: " << (circular_y ? "true" : "false") << ",\n"
<< " chroma_use_dit_mask: " << (chroma_use_dit_mask ? "true" : "false") << ",\n"
<< " qwen_image_zero_cond_t: " << (qwen_image_zero_cond_t ? "true" : "false") << ",\n"
<< " chroma_use_t5_mask: " << (chroma_use_t5_mask ? "true" : "false") << ",\n"
<< " chroma_t5_mask_pad: " << chroma_t5_mask_pad << ",\n"
<< " prediction: " << sd_prediction_name(prediction) << ",\n"
<< " lora_apply_mode: " << sd_lora_apply_mode_name(lora_apply_mode) << ",\n"
<< " force_sdxl_vae_conv_scale: " << (force_sdxl_vae_conv_scale ? "true" : "false") << "\n"
@ -885,20 +884,17 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
sd_ctx_params.tae_preview_only = taesd_preview;
sd_ctx_params.diffusion_conv_direct = diffusion_conv_direct;
sd_ctx_params.vae_conv_direct = vae_conv_direct;
sd_ctx_params.circular_x = circular || circular_x;
sd_ctx_params.circular_y = circular || circular_y;
sd_ctx_params.force_sdxl_vae_conv_scale = force_sdxl_vae_conv_scale;
sd_ctx_params.chroma_use_dit_mask = chroma_use_dit_mask;
sd_ctx_params.chroma_use_t5_mask = chroma_use_t5_mask;
sd_ctx_params.chroma_t5_mask_pad = chroma_t5_mask_pad;
sd_ctx_params.qwen_image_zero_cond_t = qwen_image_zero_cond_t;
sd_ctx_params.vae_format = str_to_vae_format(vae_format);
sd_ctx_params.max_vram = max_vram.c_str();
sd_ctx_params.stream_layers = stream_layers;
sd_ctx_params.eager_load = eager_load;
sd_ctx_params.backend = effective_backend.c_str();
sd_ctx_params.params_backend = effective_params_backend.c_str();
sd_ctx_params.split_mode = split_mode.c_str();
sd_ctx_params.auto_fit = auto_fit;
sd_ctx_params.rpc_servers = rpc_servers.c_str();
sd_ctx_params.model_args = model_args.empty() ? nullptr : model_args.c_str();
return sd_ctx_params;
}
@ -1077,7 +1073,7 @@ ArgOptions SDGenerationParams::get_options() {
&sample_params.guidance.slg.layer_end},
{"",
"--eta",
"noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)",
"noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde, dpm++2s_a, dpm++2m_sde and dpm++2m_sde_bt)",
&sample_params.eta},
{"",
"--flow-shift",
@ -1109,7 +1105,7 @@ ArgOptions SDGenerationParams::get_options() {
&high_noise_sample_params.guidance.slg.layer_end},
{"",
"--high-noise-eta",
"(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)",
"(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde, dpm++2s_a, dpm++2m_sde and dpm++2m_sde_bt)",
&high_noise_sample_params.eta},
{"",
"--strength",
@ -1160,6 +1156,18 @@ ArgOptions SDGenerationParams::get_options() {
"disable auto resize of ref images",
false,
&auto_resize_ref_image},
{"",
"--circular",
"enable circular padding on both axes for tileable output",
true, &circular},
{"",
"--circularx",
"enable circular padding on x-axis (width) only",
true, &circular_x},
{"",
"--circulary",
"enable circular padding on y-axis (height) only",
true, &circular_y},
{"",
"--disable-image-metadata",
"do not embed generation metadata on image files",
@ -1480,12 +1488,12 @@ ArgOptions SDGenerationParams::get_options() {
on_seed_arg},
{"",
"--sampling-method",
"sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
"sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
"(default: euler for Flux/SD3/Wan, euler_a otherwise)",
on_sample_method_arg},
{"",
"--high-noise-sampling-method",
"(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
"(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
" default: euler for Flux/SD3/Wan, euler_a otherwise",
on_high_noise_sample_method_arg},
{"",
@ -2446,6 +2454,8 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
params.hires.upscale_tile_size = hires_upscale_tile_size;
params.hires.custom_sigmas = hires_custom_sigmas.empty() ? nullptr : hires_custom_sigmas.data();
params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size());
params.circular_x = circular || circular_x;
params.circular_y = circular || circular_y;
return params;
}
@ -2511,6 +2521,8 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
params.hires.upscale_tile_size = hires_upscale_tile_size;
params.hires.custom_sigmas = hires_custom_sigmas.empty() ? nullptr : hires_custom_sigmas.data();
params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size());
params.circular_x = circular || circular_x;
params.circular_y = circular || circular_y;
return params;
}

View file

@ -151,6 +151,9 @@ struct SDContextParams {
bool eager_load = false;
std::string backend;
std::string params_backend;
std::string split_mode;
std::string model_args;
bool auto_fit = false;
std::string rpc_servers;
std::string effective_backend;
std::string effective_params_backend;
@ -163,16 +166,6 @@ struct SDContextParams {
bool diffusion_conv_direct = false;
bool vae_conv_direct = false;
bool circular = false;
bool circular_x = false;
bool circular_y = false;
bool chroma_use_dit_mask = true;
bool chroma_use_t5_mask = false;
int chroma_t5_mask_pad = 1;
bool qwen_image_zero_cond_t = false;
prediction_t prediction = PREDICTION_COUNT;
lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO;
@ -244,6 +237,10 @@ struct SDGenerationParams {
int upscale_repeats = 1;
int upscale_tile_size = 128;
bool circular = false;
bool circular_x = false;
bool circular_y = false;
bool hires_enabled = false;
std::string hires_upscaler = "Latent";
std::string hires_upscaler_model_path;

View file

@ -54,6 +54,8 @@ enum sample_method_t {
EULER_CFG_PP_SAMPLE_METHOD,
EULER_A_CFG_PP_SAMPLE_METHOD,
EULER_GE_SAMPLE_METHOD,
DPMPP2M_SDE_SAMPLE_METHOD,
DPMPP2M_SDE_BT_SAMPLE_METHOD,
SAMPLE_METHOD_COUNT
};
@ -214,20 +216,17 @@ typedef struct {
bool tae_preview_only;
bool diffusion_conv_direct;
bool vae_conv_direct;
bool circular_x;
bool circular_y;
bool force_sdxl_vae_conv_scale;
bool chroma_use_dit_mask;
bool chroma_use_t5_mask;
int chroma_t5_mask_pad;
bool qwen_image_zero_cond_t;
enum sd_vae_format_t vae_format;
const char* max_vram; // GiB budget or backend assignment spec for graph-cut segmented param offload (0 = disabled, -1 = auto)
bool stream_layers; // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram)
bool eager_load; // Load all params into the params backend at model-load time instead of lazily on first use
const char* backend;
const char* params_backend;
const char* split_mode; // weight distribution for multi-device modules: layer (default) or row, or per-module assignments e.g. "diffusion=row"
bool auto_fit;
const char* rpc_servers;
const char* model_args;
} sd_ctx_params_t;
typedef struct {
@ -381,6 +380,8 @@ typedef struct {
sd_cache_params_t cache;
sd_hires_params_t hires;
int qwen_image_layers;
bool circular_x;
bool circular_y;
} sd_img_gen_params_t;
typedef struct {
@ -407,6 +408,8 @@ typedef struct {
sd_tiling_params_t vae_tiling_params;
sd_cache_params_t cache;
sd_hires_params_t hires;
bool circular_x;
bool circular_y;
} sd_vid_gen_params_t;
typedef struct sd_ctx_t sd_ctx_t;
@ -518,7 +521,8 @@ SD_API bool convert_with_components(const char* model_path,
const char* output_path,
enum sd_type_t output_type,
const char* tensor_type_rules,
bool convert_name);
bool convert_name,
int n_threads);
SD_API bool preprocess_canny(sd_image_t image,
float high_threshold,
@ -535,6 +539,12 @@ SD_API void disable_imatrix_collection(void);
SD_API const char* sd_commit(void);
SD_API const char* sd_version(void);
// List available ggml backend devices, one `name<TAB>description` per line.
// The names are the device names accepted by the --backend / --params-backend
// assignment specs. Returns the number of bytes required, excluding the null
// terminator. Passing nullptr or buffer_size 0 only queries the required size.
SD_API size_t sd_list_devices(char* buffer, size_t buffer_size);
// for C API, caller needs to call free_sd_images to free the memory after use
// This helps avoid CRT problems on Windows when memory is allocated in the library but freed in the caller, which may use a different CRT.
SD_API void free_sd_images(sd_image_t* result_images, int num_images);

View file

@ -270,35 +270,6 @@ std::string load_gpt_oss_vocab_json()
return load_embd_file(cache, "embd_res/gpt_oss_vocab_json.embd");
}
static std::string get_device_override(int value, const char * module = nullptr)
{
std::string device_name;
if (value <= -2) {
device_name = "CPU";
} else if (value >= 0) {
size_t gpu_index = static_cast<size_t>(value);
if (gpu_index >= ggml_backend_dev_count()) {
printf("\nWARNING: device %zu doesn't exist, falling back to default for %s\n",
gpu_index,
module ? module : "the main device");
} else {
auto dev = ggml_backend_dev_get(gpu_index);
device_name = ggml_backend_dev_name(dev);
}
}
std::string result;
if (device_name == "") {
result = ""; // no override: sdcpp will use the main device
} else if (module) {
printf("Selecting %s as %s image generation device\n", device_name.c_str(), module);
result = std::string{","} + module + "=" + device_name;
} else {
printf("Selecting %s as the main image generation device\n", device_name.c_str());
result = device_name;
}
return result;
}
bool sdtype_load_model(const sd_load_model_inputs inputs) {
sd_is_quiet = inputs.quiet;
set_sd_quiet(sd_is_quiet);
@ -323,7 +294,9 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) {
cfg_square_limit = inputs.img_soft_limit;
printf("\nImageGen Init - Load Model: %s\n",inputs.model_filename);
std::string backends = get_device_override(inputs.kcpp_main_device);
std::string backend = inputs.backend ? inputs.backend : "";
std::string params_backend = inputs.params_backend ? inputs.params_backend : "";
std::string split_mode = inputs.split_mode ? inputs.split_mode : "";
int lora_apply_mode = LORA_APPLY_AT_RUNTIME;
bool lora_dynamic = false;
@ -402,20 +375,33 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) {
{
printf("Conv2D Direct for VAE model is enabled\n");
}
if (inputs.use_mmap && inputs.offload_cpu) {
if(backend != "")
{
printf("Backend assignment: \"%s\"\n", backend.c_str());
}
if (inputs.use_mmap && params_backend == "CPU") {
printf("Offloading weights to system RAM with mmap\n");
if (!lora_dynamic && inputs.lora_len > 0) {
printf("Note: static LoRAs can reduce mmap memory savings!\n");
}
} else if (inputs.offload_cpu) {
} else if (inputs.params_backend == "CPU") {
printf("Offloading weights to system RAM\n");
} else if (inputs.use_mmap) {
printf("Using mmap for I/O\n");
}
if(inputs.auto_fit) {
printf("Using auto-fit");
}
if(params_backend != "" && params_backend != "CPU") {
printf("Parameters backend assignment: \"%s\"\n", params_backend.c_str());
}
if(split_mode != "") {
printf("Using split mode: \"%s\"\n", split_mode.c_str());
}
std::string max_vram;
if(inputs.max_vram != 0.f) {
printf("Using max VRAM = %0.2f GB\n", inputs.max_vram);
max_vram = std::to_string(inputs.max_vram);
if(inputs.max_vram && *inputs.max_vram) {
max_vram = inputs.max_vram;
printf("Using max VRAM = %s GB\n", max_vram.c_str());
}
if(inputs.quant > 0)
{
@ -479,21 +465,15 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) {
params.diffusion_flash_attn = sd_params->diffusion_flash_attn;
params.diffusion_conv_direct = sd_params->diffusion_conv_direct;
params.vae_conv_direct = sd_params->vae_conv_direct;
params.chroma_use_dit_mask = true;
params.model_args = "chroma_use_dit_mask=true";
params.max_vram = max_vram.c_str();
params.stream_layers = inputs.stream_layers;
params.eager_load = true; //kcpp should preload everything
params.enable_mmap = inputs.use_mmap;
params.params_backend = inputs.offload_cpu ? "CPU" : "";
backends += get_device_override(inputs.kcpp_vae_device, "VAE");
backends += get_device_override(inputs.kcpp_clip_device, "CLIP");
if (backends.rfind(",", 0) == 0) {
backends = "auto" + backends;
}
params.backend = backends.c_str();
if (inputs.debugmode==1) {
printf("\nSetting sd backend list to \"%s\", params backend list to \"%s\"", params.backend, params.params_backend);
}
params.backend = backend.c_str();
params.params_backend = params_backend.c_str();
params.split_mode = split_mode.c_str();
params.auto_fit = inputs.auto_fit;
params.lora_apply_mode = (lora_apply_mode_t)lora_apply_mode;
// also switches flash attn for the vae and conditioner
@ -1030,8 +1010,6 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
sd_params->sample_method = sd_get_default_sample_method(sd_ctx);
}
SetCircularAxesAll(sd_ctx, inputs.circular_x, inputs.circular_y);
sd_params->cache_mode = inputs.cache_mode ? inputs.cache_mode : "";
sd_params->cache_options = inputs.cache_options ? inputs.cache_options : "";
@ -1289,6 +1267,8 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
params.vae_tiling_params.temporal_tiling = true;
}
parse_cache_options(params.cache, sd_params->cache_mode, sd_params->cache_options);
params.circular_x = inputs.circular_x;
params.circular_y = inputs.circular_y;
LoraMap lora_map = sd_params->lora_map;
if (sd_params->lora_dynamic) {

View file

@ -6,6 +6,7 @@
#include <optional>
#include "core/tensor_ggml.hpp"
#include "core/util.h"
#include "model/te/clip.hpp"
#include "model/te/llm.hpp"
#include "model/te/t5.hpp"
@ -116,6 +117,10 @@ public:
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
virtual void set_stream_layers_enabled(bool enabled) {}
virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
virtual void set_graph_cut_layer_split_enabled(bool enabled) {}
virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {}
virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
virtual void set_flash_attention_enabled(bool enabled) = 0;
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
virtual void runner_done() {}
@ -178,6 +183,27 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
text_model->set_runtime_backends(backends);
if (sd_version_is_sdxl(version)) {
text_model2->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
text_model->set_graph_cut_layer_split_enabled(enabled);
if (sd_version_is_sdxl(version)) {
text_model2->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
text_model->set_graph_cut_layer_split_backend_vram_limits(limits);
if (sd_version_is_sdxl(version)) {
text_model2->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void set_flash_attention_enabled(bool enabled) override {
text_model->set_flash_attention_enabled(enabled);
if (sd_version_is_sdxl(version)) {
@ -635,6 +661,48 @@ struct SD3CLIPEmbedder : public Conditioner {
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (clip_l) {
clip_l->set_runtime_backends(backends);
}
if (clip_g) {
clip_g->set_runtime_backends(backends);
}
if (t5) {
t5->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_graph_cut_layer_split_enabled(enabled);
}
if (clip_g) {
clip_g->set_graph_cut_layer_split_enabled(enabled);
}
if (t5) {
t5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (clip_l) {
clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
}
if (clip_g) {
clip_g->set_graph_cut_layer_split_backend_vram_limits(limits);
}
if (t5) {
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_flash_attention_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_flash_attention_enabled(enabled);
@ -994,6 +1062,39 @@ struct FluxCLIPEmbedder : public Conditioner {
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (clip_l) {
clip_l->set_runtime_backends(backends);
}
if (t5) {
t5->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_graph_cut_layer_split_enabled(enabled);
}
if (t5) {
t5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (clip_l) {
clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
}
if (t5) {
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_flash_attention_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_flash_attention_enabled(enabled);
@ -1191,8 +1292,27 @@ struct T5CLIPEmbedder : public Conditioner {
bool use_mask = false,
int mask_pad = 0,
bool is_umt5 = false,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const char* model_args = nullptr)
: use_mask(use_mask), mask_pad(mask_pad), t5_tokenizer(is_umt5) {
for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
if (key == "chroma_use_t5_mask") {
bool parsed = false;
if (parse_strict_bool(value, parsed)) {
this->use_mask = parsed;
} else {
LOG_WARN("ignoring invalid Chroma T5 model arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "chroma_t5_mask_pad") {
int parsed = 0;
if (parse_strict_int(value, parsed)) {
this->mask_pad = parsed;
} else {
LOG_WARN("ignoring invalid Chroma T5 model arg '%s=%s'", key.c_str(), value.c_str());
}
}
}
bool use_t5 = false;
for (auto pair : tensor_storage_map) {
if (pair.first.find("text_encoders.t5xxl") != std::string::npos) {
@ -1226,6 +1346,30 @@ struct T5CLIPEmbedder : public Conditioner {
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (t5) {
t5->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (t5) {
t5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (t5) {
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_flash_attention_enabled(bool enabled) override {
if (t5) {
t5->set_flash_attention_enabled(enabled);
@ -1418,6 +1562,30 @@ struct MiniT2IConditioner : public Conditioner {
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (t5) {
t5->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (t5) {
t5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (t5) {
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_flash_attention_enabled(bool enabled) override {
if (t5) {
t5->set_flash_attention_enabled(enabled);
@ -1502,6 +1670,22 @@ struct AnimaConditioner : public Conditioner {
llm->set_stream_layers_enabled(enabled);
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
llm->set_graph_cut_layer_split_enabled(enabled);
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
}
@ -1647,6 +1831,26 @@ struct LLMEmbedder : public Conditioner {
llm->set_stream_layers_enabled(enabled);
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (llm) {
llm->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (llm) {
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
}
@ -2316,6 +2520,22 @@ struct LTXAVEmbedder : public Conditioner {
projector->set_max_graph_vram_bytes(max_vram_bytes);
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
llm->set_graph_cut_layer_split_enabled(enabled);
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
llm->set_weight_adapter(adapter);
projector->set_weight_adapter(adapter);

View file

@ -1,14 +1,33 @@
#include <algorithm>
#include <condition_variable>
#include <cstdint>
#include <cstring>
#include <exception>
#include <fstream>
#include <memory>
#include <mutex>
#include <regex>
#include <string>
#include <thread>
#include <vector>
#include "core/util.h"
#include "model_io/gguf_io.h"
#include "model_io/safetensors_io.h"
#include "model_io/streaming_writer.h"
#include "model_loader.h"
#include "util.h"
#include "ggml_extend_backend.h"
struct TensorExportInfo {
TensorStorage storage;
ggml_type type;
};
struct TensorExportJob {
TensorExportInfo info;
std::vector<uint8_t> data;
std::string error;
bool success = false;
};
static ggml_type get_export_tensor_type(ModelLoader& model_loader,
const TensorStorage& tensor_storage,
@ -33,47 +52,262 @@ static ggml_type get_export_tensor_type(ModelLoader& model_loader,
return tensor_type;
}
static bool load_tensors_for_export(ModelLoader& model_loader,
ggml_context* ggml_ctx,
ggml_type type,
const TensorTypeRules& tensor_type_rules,
std::vector<TensorWriteInfo>& tensors) {
std::mutex tensor_mutex;
auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
const std::string& name = tensor_storage.name;
ggml_type tensor_type = get_export_tensor_type(model_loader, tensor_storage, type, tensor_type_rules);
static bool collect_tensors_for_export(ModelLoader& model_loader,
ggml_type type,
const TensorTypeRules& tensor_type_rules,
std::vector<TensorExportInfo>& tensors) {
tensors.clear();
tensors.reserve(model_loader.get_tensor_storage_map().size());
for (const auto& kv : model_loader.get_tensor_storage_map()) {
const TensorStorage& tensor_storage = kv.second;
TensorExportInfo info;
info.storage = tensor_storage;
info.type = get_export_tensor_type(model_loader, tensor_storage, type, tensor_type_rules);
tensors.push_back(std::move(info));
}
LOG_INFO("collected %zu tensors for export", tensors.size());
return true;
}
std::lock_guard<std::mutex> lock(tensor_mutex);
ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
if (tensor == nullptr) {
LOG_ERROR("ggml_new_tensor failed");
static size_t export_tensor_nbytes(const TensorExportInfo& info) {
TensorStorage output_storage = info.storage;
output_storage.type = info.type;
return static_cast<size_t>(output_storage.nbytes());
}
static TensorWritePlan tensor_write_plan_from_export_info(const TensorExportInfo& info) {
TensorWritePlan plan;
plan.name = info.storage.name;
plan.type = info.type;
plan.n_dims = info.storage.n_dims;
for (int i = 0; i < SD_MAX_DIMS; i++) {
plan.ne[i] = info.storage.ne[i];
}
return plan;
}
static std::vector<TensorWritePlan> tensor_write_plans_from_export_infos(const std::vector<TensorExportInfo>& tensors) {
std::vector<TensorWritePlan> plans;
plans.reserve(tensors.size());
for (const TensorExportInfo& info : tensors) {
plans.push_back(tensor_write_plan_from_export_info(info));
}
return plans;
}
static bool preallocate_output_file(const std::string& output_path, uint64_t file_size, std::string* error) {
if (file_size == 0) {
return true;
}
std::fstream file(output_path, std::ios::binary | std::ios::in | std::ios::out);
if (!file.is_open()) {
if (error != nullptr) {
*error = "failed to open output file '" + output_path + "' for preallocation";
}
return false;
}
// This portable fallback sets the final file size. A platform-specific
// posix_fallocate/ftruncate path can replace it later.
file.seekp(static_cast<std::streamoff>(file_size - 1), std::ios::beg);
file.put('\0');
file.flush();
if (!file) {
if (error != nullptr) {
*error = "failed to preallocate output file '" + output_path + "'";
}
return false;
}
return true;
}
static bool load_tensor_for_export(ModelLoader& model_loader, TensorExportJob& job) {
size_t mem_size = 1 * 1024 * 1024;
mem_size += ggml_tensor_overhead();
TensorStorage output_storage = job.info.storage;
output_storage.type = job.info.type;
mem_size += static_cast<size_t>(output_storage.nbytes());
ggml_context* ggml_ctx = ggml_init({mem_size, nullptr, false});
if (ggml_ctx == nullptr) {
job.error = "ggml_init failed for tensor '" + job.info.storage.name + "'";
return false;
}
ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, job.info.type, job.info.storage.n_dims, job.info.storage.ne);
if (tensor == nullptr) {
ggml_free(ggml_ctx);
job.error = "ggml_new_tensor failed for tensor '" + job.info.storage.name + "'";
return false;
}
ggml_set_name(tensor, job.info.storage.name.c_str());
const size_t tensor_nbytes = ggml_nbytes(tensor);
if (tensor_nbytes > 0 && !model_loader.load_tensor(job.info.storage, tensor)) {
ggml_free(ggml_ctx);
job.error = "failed to load tensor '" + job.info.storage.name + "'";
return false;
}
job.data.resize(tensor_nbytes);
if (tensor_nbytes > 0) {
memcpy(job.data.data(), tensor->data, tensor_nbytes);
}
ggml_free(ggml_ctx);
return true;
}
static bool stream_tensor_data(ModelLoader& model_loader,
const std::string& output_path,
const std::vector<TensorExportInfo>& tensors,
const StreamingModelWriter& writer,
int n_threads,
std::string* error) {
n_threads = n_threads > 0 ? n_threads : sd_get_num_physical_cores();
n_threads = std::max(1, n_threads);
LOG_INFO("streaming convert with %d threads", n_threads);
int64_t start_time = ggml_time_ms();
uint64_t bytes_written = 0;
size_t tensors_written = 0;
size_t next_tensor_index = 0;
bool failed = false;
std::string failure;
const size_t memory_budget = 1024ull * 1024ull * 1024ull;
size_t reserved_bytes = 0;
std::mutex work_mutex;
std::mutex progress_mutex;
std::condition_variable memory_cv;
std::vector<std::thread> workers;
workers.reserve(n_threads);
auto reserve_memory = [&](size_t bytes) -> bool {
std::unique_lock<std::mutex> lock(work_mutex);
memory_cv.wait(lock, [&]() {
return failed || reserved_bytes == 0 || reserved_bytes + bytes <= memory_budget;
});
if (failed) {
return false;
}
ggml_set_name(tensor, name.c_str());
if (!tensor->data) {
GGML_ASSERT(ggml_nelements(tensor) == 0);
// Avoid crashing writers by setting a dummy pointer for zero-sized tensors.
LOG_DEBUG("setting dummy pointer for zero-sized tensor %s", name.c_str());
tensor->data = ggml_get_mem_buffer(ggml_ctx);
}
TensorWriteInfo write_info;
write_info.tensor = tensor;
write_info.n_dims = tensor_storage.n_dims;
for (int i = 0; i < tensor_storage.n_dims; ++i) {
write_info.ne[i] = tensor_storage.ne[i];
}
*dst_tensor = tensor;
tensors.push_back(std::move(write_info));
reserved_bytes += bytes;
return true;
};
bool success = model_loader.load_tensors(on_new_tensor_cb);
LOG_INFO("load tensors done");
return success;
auto release_memory = [&](size_t bytes) {
{
std::lock_guard<std::mutex> lock(work_mutex);
reserved_bytes -= std::min(reserved_bytes, bytes);
}
memory_cv.notify_all();
};
auto fail = [&](const std::string& message) {
{
std::lock_guard<std::mutex> lock(work_mutex);
if (!failed) {
failed = true;
failure = message;
}
}
memory_cv.notify_all();
};
for (int worker = 0; worker < n_threads; worker++) {
workers.emplace_back([&]() {
std::fstream output_file(output_path, std::ios::binary | std::ios::in | std::ios::out);
if (!output_file.is_open()) {
fail("failed to open output file '" + output_path + "' for tensor writing");
return;
}
while (true) {
size_t tensor_index = 0;
{
std::lock_guard<std::mutex> lock(work_mutex);
if (failed || next_tensor_index >= tensors.size()) {
return;
}
tensor_index = next_tensor_index++;
}
const size_t tensor_bytes = export_tensor_nbytes(tensors[tensor_index]);
if (!reserve_memory(tensor_bytes)) {
return;
}
TensorExportJob job;
job.info = tensors[tensor_index];
try {
job.success = load_tensor_for_export(model_loader, job);
} catch (const std::exception& e) {
job.error = e.what();
job.success = false;
}
if (!job.success) {
release_memory(tensor_bytes);
fail(job.error.empty() ? "streaming conversion failed" : job.error);
return;
}
std::string write_error;
if (!writer.write_tensor(output_file,
tensor_index,
job.data.empty() ? nullptr : job.data.data(),
job.data.size(),
&write_error)) {
release_memory(tensor_bytes);
fail(write_error.empty() ? "streaming conversion write failed" : write_error);
return;
}
{
std::lock_guard<std::mutex> lock(progress_mutex);
bytes_written += job.data.size();
tensors_written++;
float elapsed_seconds = (ggml_time_ms() - start_time) / 1000.0f;
pretty_bytes_progress(static_cast<int>(tensors_written),
static_cast<int>(tensors.size()),
bytes_written,
elapsed_seconds);
}
release_memory(tensor_bytes);
}
});
}
for (auto& worker : workers) {
worker.join();
}
printf("\n");
if (failed) {
if (error != nullptr) {
*error = failure;
}
return false;
}
LOG_INFO("streaming conversion completed, taking %.2fs", (ggml_time_ms() - start_time) / 1000.f);
return true;
}
static bool write_model_file_streaming(ModelLoader& model_loader,
const std::string& output_path,
const std::vector<TensorExportInfo>& tensors,
StreamingModelWriter& writer,
int n_threads,
std::string* error) {
std::vector<TensorWritePlan> plans = tensor_write_plans_from_export_infos(tensors);
if (!writer.write_metadata(output_path, plans, error)) {
return false;
}
if (!preallocate_output_file(output_path, writer.file_size(), error)) {
return false;
}
model_loader.process_model_files(false, false);
return stream_tensor_data(model_loader, output_path, tensors, writer, n_threads, error);
}
static bool init_convert_path(ModelLoader& model_loader, const char* path, const char* prefix, bool& loaded_any) {
@ -91,42 +325,29 @@ static bool init_convert_path(ModelLoader& model_loader, const char* path, const
static bool export_loaded_model(ModelLoader& model_loader,
const char* output_path,
sd_type_t output_type,
const char* tensor_type_rules) {
const char* tensor_type_rules,
int n_threads) {
ggml_type type = sd_type_to_ggml_type(output_type);
bool output_is_safetensors = ends_with(output_path, ".safetensors");
TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
auto backend = sd_backend_cpu_init();
size_t mem_size = 1 * 1024 * 1024; // for padding
mem_size += model_loader.get_tensor_storage_map().size() * ggml_tensor_overhead();
mem_size += model_loader.get_params_mem_size(backend, type);
LOG_INFO("model tensors mem size: %.2fMB", mem_size / 1024.f / 1024.f);
ggml_context* ggml_ctx = ggml_init({mem_size, nullptr, false});
if (ggml_ctx == nullptr) {
LOG_ERROR("ggml_init failed for converter");
ggml_backend_free(backend);
return false;
}
std::vector<TensorWriteInfo> tensors;
bool success = load_tensors_for_export(model_loader, ggml_ctx, type, type_rules, tensors);
ggml_backend_free(backend);
std::vector<TensorExportInfo> tensors;
bool success = collect_tensors_for_export(model_loader, type, type_rules, tensors);
std::string error;
if (success) {
std::unique_ptr<StreamingModelWriter> writer;
if (output_is_safetensors) {
success = write_safetensors_file(output_path, tensors, &error);
writer = std::make_unique<SafetensorsStreamingWriter>();
} else {
success = write_gguf_file(output_path, tensors, &error);
writer = std::make_unique<GGUFStreamingWriter>();
}
success = write_model_file_streaming(model_loader, output_path, tensors, *writer, n_threads, &error);
}
if (!success && !error.empty()) {
LOG_ERROR("%s", error.c_str());
}
ggml_free(ggml_ctx);
return success;
}
@ -139,7 +360,8 @@ bool convert_with_components(const char* model_path,
const char* output_path,
sd_type_t output_type,
const char* tensor_type_rules,
bool convert_name) {
bool convert_name,
int n_threads) {
ModelLoader model_loader;
bool loaded_any = false;
@ -161,7 +383,7 @@ bool convert_with_components(const char* model_path,
model_loader.convert_tensors_name();
}
return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules);
return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules, n_threads);
}
bool convert(const char* input_path,
@ -179,5 +401,6 @@ bool convert(const char* input_path,
output_path,
output_type,
tensor_type_rules,
convert_name);
convert_name,
0);
}

View file

@ -0,0 +1,390 @@
#include "backend_fit.h"
#include <algorithm>
#include <cctype>
#include <cstdint>
#include <utility>
#include <vector>
#include "core/ggml_extend_backend.h"
#include "core/util.h"
#include "ggml-backend.h"
namespace sd::backend_fit {
namespace {
constexpr int64_t MiB = 1024ll * 1024;
enum class ComponentKind {
DIT = 0,
VAE = 1,
CONDITIONER = 2,
};
struct Component {
ComponentKind kind;
const char* name;
int64_t params_bytes = 0;
int64_t reserve_bytes = 0;
bool splittable = false;
};
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;
};
struct Decision {
ComponentKind kind;
bool on_cpu = false;
std::vector<size_t> device_idxs;
};
struct Plan {
bool valid = false;
bool time_share = false;
std::vector<Decision> decisions;
};
bool classify_tensor(const std::string& name, ComponentKind& out) {
auto contains = [&](const char* s) { return name.find(s) != std::string::npos; };
if (contains("model.diffusion_model.") || contains("unet.")) {
out = ComponentKind::DIT;
return true;
}
if (contains("first_stage_model.") ||
name.rfind("vae.", 0) == 0 ||
name.rfind("tae.", 0) == 0) {
out = ComponentKind::VAE;
return true;
}
if (contains("text_encoders") ||
contains("cond_stage_model") ||
contains("te.text_model.") ||
contains("conditioner") ||
name.rfind("text_encoder.", 0) == 0 ||
name.rfind("text_embedding_projection.", 0) == 0 ||
contains(".aggregate_embed.")) {
out = ComponentKind::CONDITIONER;
return true;
}
return false;
}
std::vector<Component> estimate_components(ModelLoader& loader, ggml_type override_wtype) {
const auto& storage = loader.get_tensor_storage_map();
int64_t bytes[3] = {0, 0, 0};
for (const auto& [name, ts_const] : storage) {
TensorStorage ts = ts_const;
if (is_unused_tensor(ts.name)) {
continue;
}
ComponentKind kind;
if (!classify_tensor(ts.name, kind)) {
continue;
}
if (override_wtype != GGML_TYPE_COUNT &&
loader.tensor_should_be_converted(ts, override_wtype)) {
ts.type = override_wtype;
} else if (ts.expected_type != GGML_TYPE_COUNT && ts.expected_type != ts.type) {
ts.type = ts.expected_type;
}
bytes[int(kind)] += (int64_t)ts.nbytes() + 64;
}
std::vector<Component> out;
out.push_back({ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, true});
out.push_back({ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, false});
out.push_back({ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, true});
return out;
}
std::vector<Device> enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) {
std::vector<Device> 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;
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); });
float gib = budgets.default_gib;
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>((int64_t)(gib * 1024.0 * 1024.0 * 1024.0), d.free_bytes);
} else if (gib < 0.f) {
d.budget_bytes = d.free_bytes + (int64_t)(gib * 1024.0 * 1024.0 * 1024.0);
} else {
d.budget_bytes = d.free_bytes - 512 * MiB;
}
d.budget_bytes = std::max<int64_t>(d.budget_bytes, 0);
out.push_back(d);
}
return out;
}
Plan compute_plan(const std::vector<Component>& components, const std::vector<Device>& devices) {
Plan plan;
if (devices.empty()) {
return plan;
}
std::vector<size_t> order(components.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 components[a].params_bytes > components[b].params_bytes;
});
{
std::vector<int64_t> params_sum(devices.size(), 0);
std::vector<int64_t> max_reserve(devices.size(), 0);
std::vector<Decision> decisions(components.size());
bool ok = true;
for (size_t ci : order) {
const Component& comp = components[ci];
decisions[ci].kind = comp.kind;
if (comp.params_bytes == 0) {
continue;
}
int best = -1;
for (size_t di = 0; di < devices.size(); di++) {
int64_t need = params_sum[di] + comp.params_bytes + std::max(max_reserve[di], comp.reserve_bytes);
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;
}
}
if (best < 0) {
ok = false;
break;
}
params_sum[best] += comp.params_bytes;
max_reserve[best] = std::max(max_reserve[best], comp.reserve_bytes);
decisions[ci].device_idxs.push_back((size_t)best);
}
if (ok) {
plan.valid = true;
plan.time_share = false;
plan.decisions = std::move(decisions);
return plan;
}
}
plan.decisions.assign(components.size(), {});
for (size_t ci : order) {
const Component& comp = components[ci];
Decision& decision = plan.decisions[ci];
decision.kind = comp.kind;
if (comp.params_bytes == 0) {
continue;
}
int best = -1;
for (size_t di = 0; di < devices.size(); di++) {
if (comp.params_bytes + comp.reserve_bytes <= devices[di].budget_bytes &&
(best < 0 || devices[di].budget_bytes > devices[best].budget_bytes)) {
best = (int)di;
}
}
if (best >= 0) {
decision.device_idxs.push_back((size_t)best);
continue;
}
if (comp.splittable && devices.size() > 1) {
int64_t capacity = 0;
for (const Device& d : devices) {
capacity += std::max<int64_t>(d.budget_bytes - comp.reserve_bytes, 0);
}
if (comp.params_bytes <= capacity) {
std::vector<size_t> 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;
});
decision.device_idxs = std::move(idxs);
continue;
}
}
decision.on_cpu = true;
}
plan.valid = true;
plan.time_share = true;
return plan;
}
void print_plan(const Plan& plan,
const std::vector<Component>& components,
const std::vector<Device>& devices) {
LOG_INFO("auto-fit plan%s:", plan.time_share ? " (time-share: params load per phase and free after)" : "");
LOG_INFO(" devices:");
for (const Device& d : devices) {
LOG_INFO(" %-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));
}
LOG_INFO(" components:");
for (size_t ci = 0; ci < components.size(); ci++) {
const Component& comp = components[ci];
const Decision& decision = plan.decisions[ci];
std::string target;
if (comp.params_bytes == 0) {
target = "(not present)";
} else 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)";
}
}
LOG_INFO(" %-12s params %6lld MiB, compute reserve %5lld MiB -> %s",
comp.name,
(long long)(comp.params_bytes / MiB),
(long long)(comp.reserve_bytes / MiB),
target.c_str());
}
}
void append_assignment(std::string& spec, const char* key, const std::string& value) {
if (!spec.empty()) {
spec += ",";
}
spec += key;
spec += "=";
spec += value;
}
void append_component_decision(const std::vector<Component>& components,
const std::vector<Device>& devices,
const Plan& plan,
ComponentKind kind,
const char* module_key,
std::string& runtime_spec,
std::string& params_spec) {
for (size_t ci = 0; ci < components.size(); ci++) {
if (components[ci].kind != kind || components[ci].params_bytes == 0) {
continue;
}
const Decision& decision = plan.decisions[ci];
if (decision.on_cpu) {
append_assignment(runtime_spec, module_key, "cpu");
return;
}
if (decision.device_idxs.empty()) {
return;
}
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(runtime_spec, module_key, device_list);
if (plan.time_share) {
append_assignment(params_spec, module_key, "disk");
}
return;
}
}
} // namespace
bool derive_backend_specs(ModelLoader& loader,
ggml_type override_wtype,
sd::ggml_graph_cut::MaxVramAssignment& budgets,
std::string& runtime_spec,
std::string& params_spec) {
if (!runtime_spec.empty() || !params_spec.empty()) {
LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend");
}
{
std::string error;
if (!budgets.canonicalize_backend_keys(&error)) {
LOG_ERROR("%s", error.c_str());
return false;
}
}
auto components = estimate_components(loader, override_wtype);
auto devices = enumerate_gpu_devices(budgets);
auto plan = compute_plan(components, devices);
if (!plan.valid) {
LOG_WARN("auto-fit: no usable GPU devices; using the default backend");
runtime_spec.clear();
params_spec.clear();
return true;
}
print_plan(plan, components, devices);
std::string derived_runtime_spec;
std::string derived_params_spec;
append_component_decision(components, devices, plan, ComponentKind::DIT, "diffusion", derived_runtime_spec, derived_params_spec);
append_component_decision(components, devices, plan, ComponentKind::CONDITIONER, "te", derived_runtime_spec, derived_params_spec);
append_component_decision(components, devices, plan, ComponentKind::VAE, "vae", derived_runtime_spec, derived_params_spec);
runtime_spec = std::move(derived_runtime_spec);
params_spec = std::move(derived_params_spec);
LOG_INFO("auto-fit: --backend \"%s\"%s%s%s",
runtime_spec.empty() ? "(default)" : runtime_spec.c_str(),
params_spec.empty() ? "" : " --params-backend \"",
params_spec.c_str(),
params_spec.empty() ? "" : "\"");
return true;
}
bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) {
if (prefer_temporal_tiling) {
if (tiling_params.temporal_tiling) {
return false;
}
tiling_params.temporal_tiling = true;
} else {
if (tiling_params.enabled) {
return false;
}
tiling_params.enabled = true;
if (tiling_params.tile_size_x <= 0) {
tiling_params.tile_size_x = 256;
}
if (tiling_params.tile_size_y <= 0) {
tiling_params.tile_size_y = 256;
}
}
LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling",
tiling_params.temporal_tiling ? "temporal" : "spatial");
return true;
}
} // namespace sd::backend_fit

View file

@ -0,0 +1,23 @@
#ifndef __SD_BACKEND_FIT_H__
#define __SD_BACKEND_FIT_H__
#include <string>
#include "core/ggml_graph_cut.h"
#include "model_loader.h"
#include "stable-diffusion.h"
namespace sd::backend_fit {
bool derive_backend_specs(ModelLoader& loader,
ggml_type override_wtype,
sd::ggml_graph_cut::MaxVramAssignment& budgets,
std::string& runtime_spec,
std::string& params_spec);
bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params,
bool prefer_temporal_tiling);
} // namespace sd::backend_fit
#endif // __SD_BACKEND_FIT_H__

View file

@ -21,10 +21,12 @@
#include <sstream>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "core/ggml_extend_backend.h"
#include "core/ggml_graph_cut.h"
#include "core/layer_split_partition.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml.h"
@ -1745,6 +1747,13 @@ protected:
size_t max_graph_vram_bytes = 0;
bool stream_layers_enabled = false;
size_t observed_max_effective_budget_ = 0;
bool graph_cut_layer_split_enabled = false;
std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
std::vector<ggml_backend_t> extra_runtime_backends; // borrowed (SDBackendManager-owned)
ggml_backend_sched_t sched = nullptr; // owned, multi-device only
ggml_backend_t cpu_fallback_backend = nullptr; // owned, sched requires a trailing CPU backend
bool multi_device_eval_callback_warned = false;
std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
std::weak_ptr<RunnerWeightManager> weight_manager;
@ -1771,6 +1780,9 @@ protected:
sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_;
std::unordered_set<const ggml_tensor*> params_tensor_set_;
std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_assignments_;
std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_node_assignments_;
bool graph_cut_layer_split_primary_notice_logged_ = false;
template <typename T>
static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
@ -1869,6 +1881,20 @@ protected:
params_tensor_set_dirty_ = false;
}
ggml_tensor* canonical_param_tensor(ggml_tensor* tensor) {
if (tensor == nullptr) {
return nullptr;
}
if (params_tensor_set_.find(tensor) != params_tensor_set_.end()) {
return tensor;
}
if (tensor->view_src != nullptr &&
params_tensor_set_.find(tensor->view_src) != params_tensor_set_.end()) {
return tensor->view_src;
}
return nullptr;
}
std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* gf) {
std::vector<ggml_tensor*> used_params;
rebuild_params_tensor_set();
@ -1881,12 +1907,8 @@ protected:
seen_params.reserve(static_cast<size_t>(n_leafs));
for (int i = 0; i < n_leafs; ++i) {
ggml_tensor* leaf = sd::ggml_graph_cut::leaf_tensor(gf, i);
ggml_tensor* param_leaf = leaf;
if (param_leaf != nullptr && params_tensor_set_.find(param_leaf) == params_tensor_set_.end()) {
param_leaf = param_leaf->view_src;
}
ggml_tensor* param_leaf = canonical_param_tensor(leaf);
if (param_leaf != nullptr &&
params_tensor_set_.find(param_leaf) != params_tensor_set_.end() &&
seen_params.insert(param_leaf).second) {
used_params.push_back(param_leaf);
}
@ -2013,7 +2035,127 @@ protected:
return true;
}
// Pass explicit buffer types: synthesized defaults can make CUDA devices
// report supporting each other's buffers and skip a required copy.
bool ensure_sched(ggml_cgraph* gf) {
if (sched != nullptr) {
return true;
}
std::vector<ggml_backend_t> backends;
backends.reserve(extra_runtime_backends.size() + 2);
backends.push_back(runtime_backend);
for (ggml_backend_t backend : extra_runtime_backends) {
backends.push_back(backend);
}
if (cpu_fallback_backend == nullptr && !sd_backend_is_cpu(runtime_backend)) {
cpu_fallback_backend = sd_backend_cpu_init();
}
if (cpu_fallback_backend != nullptr) {
backends.push_back(cpu_fallback_backend);
}
std::vector<ggml_backend_buffer_type_t> bufts;
bufts.reserve(backends.size());
ggml_backend_dev_t main_dev = ggml_backend_get_device(runtime_backend);
for (ggml_backend_t backend : backends) {
ggml_backend_buffer_type_t buft = nullptr;
if (backend == cpu_fallback_backend && main_dev != nullptr) {
buft = ggml_backend_dev_host_buffer_type(main_dev);
}
if (buft == nullptr) {
buft = ggml_backend_get_default_buffer_type(backend);
}
bufts.push_back(buft);
}
size_t graph_size = MAX_GRAPH_SIZE;
if (gf != nullptr) {
graph_size = std::max<size_t>(graph_size, (size_t)ggml_graph_n_nodes(gf));
}
sched = ggml_backend_sched_new(backends.data(),
bufts.data(),
(int)backends.size(),
graph_size,
/*parallel=*/false,
/*op_offload=*/false);
if (sched == nullptr) {
LOG_ERROR("%s: failed to create backend sched", get_desc().c_str());
return false;
}
return true;
}
ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const {
if (tensor == nullptr || tensor->buffer == nullptr) {
return nullptr;
}
if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS ||
ggml_backend_buffer_is_host(tensor->buffer)) {
return nullptr;
}
ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
if (dev == nullptr) {
return nullptr;
}
if (ggml_backend_get_device(runtime_backend) == dev) {
return runtime_backend;
}
for (ggml_backend_t backend : extra_runtime_backends) {
if (ggml_backend_get_device(backend) == dev) {
return backend;
}
}
return nullptr;
}
// Weightless ops have no scheduler anchor, so pin them to the most recent
// weight device. Views must stay unpinned or cross-device copies can be
// skipped for their consumers.
void pin_multi_device_nodes(ggml_cgraph* gf) {
if (sched == nullptr || gf == nullptr) {
return;
}
ggml_backend_t current = runtime_backend;
const int n_nodes = ggml_graph_n_nodes(gf);
for (int i = 0; i < n_nodes; i++) {
ggml_tensor* node = ggml_graph_node(gf, i);
auto node_assignment = graph_cut_layer_split_node_assignments_.find(node);
if (node_assignment != graph_cut_layer_split_node_assignments_.end()) {
current = node_assignment->second;
}
for (int s = 0; s < GGML_MAX_SRC; s++) {
ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
if (weight_backend != nullptr) {
if (node_assignment == graph_cut_layer_split_node_assignments_.end()) {
current = weight_backend;
}
}
}
if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
continue;
}
if (ggml_backend_supports_op(current, node)) {
ggml_backend_sched_set_tensor_backend(sched, node, current);
}
}
}
bool is_multi_device() const {
return !extra_runtime_backends.empty();
}
bool alloc_compute_buffer(ggml_cgraph* gf) {
if (is_multi_device()) {
// The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve
// the graph here: reserve runs split_graph, which rewires the
// graph's src pointers to sched-internal copy tensors, and the
// later ggml_backend_sched_alloc_graph would split the already
// rewired graph, silently corrupting every cross-backend input. A
// graph must be split at most once; the alloc in execute_graph
// performs the real allocation.
return ensure_sched(gf);
}
if (compute_allocr != nullptr) {
return true;
}
@ -2229,12 +2371,14 @@ protected:
plan.valid &&
max_graph_vram_bytes > 0 &&
plan.segments.size() > 1 &&
!sd_backend_is_cpu(runtime_backend);
!sd_backend_is_cpu(runtime_backend) &&
!is_multi_device();
}
bool can_attempt_graph_cut_segmented_compute() const {
return max_graph_vram_bytes > 0 &&
!sd_backend_is_cpu(runtime_backend);
!sd_backend_is_cpu(runtime_backend) &&
!is_multi_device();
}
bool resolve_graph_cut_plan(ggml_cgraph* gf,
@ -2314,6 +2458,123 @@ protected:
return true;
}
bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
GraphCutPlan* plan_out) {
GGML_ASSERT(plan_out != nullptr);
GGML_ASSERT(gf != nullptr);
*plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
gf,
&graph_cut_plan_cache_,
0,
params_tensor_set_,
get_desc().c_str());
return true;
}
bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
graph_cut_layer_split_node_assignments_.clear();
if (!graph_cut_layer_split_enabled) {
return true;
}
if (!is_multi_device()) {
LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str());
return false;
}
GraphCutPlan plan;
if (!resolve_graph_cut_layer_split_plan(gf, &plan)) {
return false;
}
if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
auto manager = weight_manager.lock();
if (manager == nullptr) {
LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
return false;
}
std::vector<ggml_tensor*> graph_params = collect_used_param_tensors(gf);
if (!graph_params.empty() &&
!manager->assign_compute_backend(graph_params, runtime_backend)) {
LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
get_desc().c_str(),
sd::layer_split_backend_device_display_name(runtime_backend).c_str());
return false;
}
for (ggml_tensor* param : graph_params) {
if (param != nullptr) {
graph_cut_layer_split_assignments_[param] = runtime_backend;
}
}
const int n_nodes = ggml_graph_n_nodes(gf);
for (int i = 0; i < n_nodes; i++) {
ggml_tensor* node = ggml_graph_node(gf, i);
if (node != nullptr) {
graph_cut_layer_split_node_assignments_[node] = runtime_backend;
}
}
if (!graph_cut_layer_split_primary_notice_logged_) {
LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
get_desc().c_str(),
sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
graph_params.size());
graph_cut_layer_split_primary_notice_logged_ = true;
} else {
LOG_DEBUG("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
get_desc().c_str(),
sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
graph_params.size());
}
return true;
}
std::vector<ggml_backend_t> split_backends;
split_backends.reserve(extra_runtime_backends.size() + 1);
split_backends.push_back(runtime_backend);
for (ggml_backend_t backend : extra_runtime_backends) {
if (backend != nullptr) {
split_backends.push_back(backend);
}
}
auto manager = weight_manager.lock();
if (manager == nullptr) {
LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
return false;
}
sd::GraphCutLayerSplitAssignment assignment;
auto canonicalize_param = [this](ggml_tensor* tensor) {
return canonical_param_tensor(tensor);
};
if (!sd::partition_graph_cut_layer_split(get_desc().c_str(),
gf,
plan,
split_backends,
graph_cut_layer_split_backend_vram_limits_,
max_graph_vram_bytes,
graph_cut_layer_split_assignments_,
canonicalize_param,
&assignment)) {
return false;
}
for (size_t i = 0; i < split_backends.size(); i++) {
if (assignment.tensors_by_backend[i].empty()) {
continue;
}
if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) {
LOG_ERROR("%s graph-cut layer split failed to assign params to %s",
get_desc().c_str(),
sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
return false;
}
}
graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments);
sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment);
return true;
}
struct PersistentExternalBinding {
ggml_backend_buffer_t buffer = nullptr;
void* data = nullptr;
@ -2490,7 +2751,14 @@ protected:
};
ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer);
if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) {
if (is_multi_device()) {
ggml_backend_sched_reset(sched);
pin_multi_device_nodes(gf); // reset clears the pins; re-apply before alloc
if (!ggml_backend_sched_alloc_graph(sched, gf)) {
LOG_ERROR("%s sched alloc compute graph failed", get_desc().c_str());
return std::nullopt;
}
} else if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) {
LOG_ERROR("%s alloc compute graph failed", get_desc().c_str());
return std::nullopt;
}
@ -2499,11 +2767,27 @@ protected:
if (sd_backend_is_cpu(runtime_backend)) {
sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
}
if (cpu_fallback_backend != nullptr) {
sd_backend_cpu_set_n_threads(cpu_fallback_backend, n_threads);
}
ggml_status status = sd_backend_graph_compute_with_eval_callback(runtime_backend,
gf,
sd_get_backend_eval_callback(),
sd_get_backend_eval_callback_data());
ggml_status status;
if (is_multi_device()) {
if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
LOG_WARN("%s: eval callback is not supported with multiple runtime backends; ignoring",
get_desc().c_str());
multi_device_eval_callback_warned = true;
}
status = ggml_backend_sched_graph_compute(sched, gf);
if (status == GGML_STATUS_SUCCESS) {
ggml_backend_sched_synchronize(sched);
}
} else {
status = sd_backend_graph_compute_with_eval_callback(runtime_backend,
gf,
sd_get_backend_eval_callback(),
sd_get_backend_eval_callback_data());
}
if (status != GGML_STATUS_SUCCESS) {
LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
return std::nullopt;
@ -2680,6 +2964,10 @@ public:
free_params_ctx();
free_compute_ctx();
free_cache_ctx_and_buffer();
if (cpu_fallback_backend != nullptr) {
ggml_backend_free(cpu_fallback_backend);
cpu_fallback_backend = nullptr;
}
}
virtual GGMLRunnerContext get_context() {
@ -2720,10 +3008,20 @@ public:
ggml_gallocr_free(compute_allocr);
compute_allocr = nullptr;
}
if (sched != nullptr) {
ggml_backend_sched_free(sched);
sched = nullptr;
}
}
// do copy after alloc graph
void set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
if (is_multi_device()) {
// The sched only assigns a backend (and thus a buffer) to tensors
// that participate in the graph; flag standalone data tensors as
// inputs so they get one.
ggml_set_input(tensor);
}
backend_tensor_data_map[tensor] = data;
}
@ -2814,6 +3112,11 @@ public:
GGML_ASSERT(gf != nullptr);
rebuild_params_tensor_set();
if (!assign_graph_cut_layer_split_backends(gf)) {
free_compute_ctx();
return std::nullopt;
}
if (can_attempt_graph_cut_segmented_compute()) {
GraphCutPlan plan;
if (!resolve_graph_cut_plan(gf, &plan)) {
@ -2859,8 +3162,50 @@ public:
}
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",
get_desc().c_str());
return;
}
stream_layers_enabled = enabled;
}
void set_graph_cut_layer_split_enabled(bool enabled) {
graph_cut_layer_split_enabled = enabled;
if (!enabled) {
graph_cut_layer_split_assignments_.clear();
graph_cut_layer_split_node_assignments_.clear();
graph_cut_layer_split_primary_notice_logged_ = false;
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {
graph_cut_layer_split_backend_vram_limits_ = limits;
graph_cut_layer_split_assignments_.clear();
graph_cut_layer_split_node_assignments_.clear();
graph_cut_layer_split_primary_notice_logged_ = false;
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
extra_runtime_backends.clear();
for (ggml_backend_t backend : backends) {
if (backend == nullptr || backend == runtime_backend) {
continue;
}
if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) ==
extra_runtime_backends.end()) {
extra_runtime_backends.push_back(backend);
}
}
graph_cut_layer_split_assignments_.clear();
graph_cut_layer_split_node_assignments_.clear();
graph_cut_layer_split_primary_notice_logged_ = false;
if (is_multi_device() && stream_layers_enabled) {
LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring",
get_desc().c_str());
stream_layers_enabled = false;
}
}
};
class GGMLBlock {

View file

@ -665,12 +665,52 @@ SDBackendManager::~SDBackendManager() {
void SDBackendManager::reset() {
backends_.clear();
runtime_assignment_ = {};
params_assignment_ = {};
runtime_assignment_ = {};
params_assignment_ = {};
split_mode_assignment_ = {};
}
static std::vector<std::string> split_device_list(const std::string& value) {
std::vector<std::string> names;
for (const std::string& raw : split_copy(value, '&')) {
const std::string name = trim_copy(raw);
if (!name.empty()) {
names.push_back(name);
}
}
return names;
}
static std::string primary_device_name(const std::string& value) {
std::vector<std::string> names = split_device_list(value);
return names.empty() ? std::string() : names.front();
}
ggml_backend_t SDBackendManager::runtime_backend(SDBackendModule module) {
return init_cached_backend(runtime_assignment_.get(module));
return init_cached_backend(primary_device_name(runtime_assignment_.get(module)));
}
std::vector<ggml_backend_t> SDBackendManager::runtime_backends(SDBackendModule module) {
std::vector<ggml_backend_t> backends;
for (const std::string& name : split_device_list(runtime_assignment_.get(module))) {
ggml_backend_t backend = init_cached_backend(name);
if (backend == nullptr) {
LOG_ERROR("failed to initialize backend '%s' for module %s",
name.c_str(),
sd_backend_module_name(module));
continue;
}
if (std::find(backends.begin(), backends.end(), backend) == backends.end()) {
backends.push_back(backend);
}
}
if (backends.empty()) {
ggml_backend_t backend = runtime_backend(module);
if (backend != nullptr) {
backends.push_back(backend);
}
}
return backends;
}
ggml_backend_t SDBackendManager::params_backend(SDBackendModule module) {
@ -696,6 +736,10 @@ bool SDBackendManager::params_backend_is_disk(SDBackendModule module) const {
return is_disk_backend_token(params_assignment_.get(module));
}
bool SDBackendManager::params_backend_follows_runtime(SDBackendModule module) const {
return params_assignment_.get(module).empty();
}
bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule module) {
ggml_backend_t backend = runtime_backend(module);
if (backend == nullptr) {
@ -715,6 +759,7 @@ bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule modu
bool SDBackendManager::init(const char* backend_spec,
const char* params_backend_spec,
const char* split_mode_spec,
std::string* error) {
reset();
@ -724,12 +769,53 @@ bool SDBackendManager::init(const char* backend_spec,
if (!sd_parse_backend_assignment(SAFE_STR(params_backend_spec), &params_assignment_, error)) {
return false;
}
if (!sd_parse_backend_assignment(SAFE_STR(split_mode_spec), &split_mode_assignment_, error)) {
return false;
}
return validate(error);
}
SDSplitMode SDBackendManager::split_mode(SDBackendModule module) const {
return lower_copy(trim_copy(split_mode_assignment_.get(module))) == "row" ? SDSplitMode::ROW
: SDSplitMode::LAYER;
}
ggml_backend_buffer_type_t SDBackendManager::split_buffer_type(ggml_backend_t backend,
const std::vector<float>& tensor_split) {
if (backend == nullptr) {
return nullptr;
}
ggml_backend_dev_t dev = ggml_backend_get_device(backend);
if (dev == nullptr) {
return nullptr;
}
ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
if (reg == nullptr) {
return nullptr;
}
auto fn = (ggml_backend_split_buffer_type_t)ggml_backend_reg_get_proc_address(reg, "ggml_backend_split_buffer_type");
if (fn == nullptr) {
return nullptr;
}
int main_device = -1;
const size_t dev_count = ggml_backend_reg_dev_count(reg);
for (size_t i = 0; i < dev_count; ++i) {
if (ggml_backend_reg_dev_get(reg, i) == dev) {
main_device = (int)i;
break;
}
}
if (main_device < 0) {
return nullptr;
}
std::vector<float> padded_split(std::max<size_t>(tensor_split.size(), 64), 0.0f);
std::copy(tensor_split.begin(), tensor_split.end(), padded_split.begin());
return fn(main_device, padded_split.data());
}
bool SDBackendManager::validate(std::string* error) const {
auto validate_runtime_name = [&](const std::string& name) -> bool {
auto validate_single_runtime_name = [&](const std::string& name) -> bool {
if (is_default_backend_token(name)) {
return true;
}
@ -747,15 +833,56 @@ bool SDBackendManager::validate(std::string* error) const {
}
return false;
};
auto validate_runtime_name = [&](const std::string& name) -> bool {
if (name.find('&') == std::string::npos) {
return validate_single_runtime_name(name);
}
std::vector<std::string> names = split_device_list(name);
if (names.empty()) {
if (error != nullptr) {
*error = "invalid backend device list '" + name + "'";
}
return false;
}
for (const std::string& entry : names) {
if (is_default_backend_token(entry)) {
if (error != nullptr) {
*error = "default backend token is not allowed in a device list '" + name + "'";
}
return false;
}
if (!validate_single_runtime_name(entry)) {
return false;
}
}
return true;
};
auto validate_params_name = [&](const std::string& name) -> bool {
if (is_disk_backend_token(name)) {
return true;
}
return validate_runtime_name(name);
if (name.find('&') != std::string::npos) {
if (error != nullptr) {
*error = "params_backend does not accept device lists ('" + name + "')";
}
return false;
}
return validate_single_runtime_name(name);
};
auto validate_split_mode_name = [&](const std::string& name) -> bool {
const std::string lower = lower_copy(trim_copy(name));
if (lower.empty() || lower == "layer" || lower == "row") {
return true;
}
if (error != nullptr) {
*error = "invalid split mode '" + name + "' (expected layer or row)";
}
return false;
};
if (!validate_runtime_name(runtime_assignment_.default_name) ||
!validate_params_name(params_assignment_.default_name)) {
!validate_params_name(params_assignment_.default_name) ||
!validate_split_mode_name(split_mode_assignment_.default_name)) {
return false;
}
for (const auto& kv : runtime_assignment_.module_names) {
@ -768,6 +895,11 @@ bool SDBackendManager::validate(std::string* error) const {
return false;
}
}
for (const auto& kv : split_mode_assignment_.module_names) {
if (!validate_split_mode_name(kv.second)) {
return false;
}
}
return true;
}

View file

@ -6,6 +6,7 @@
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "ggml-backend.h"
#include "ggml.h"
@ -37,10 +38,16 @@ struct SDBackendHandleDeleter {
using SDBackendHandle = std::unique_ptr<struct ggml_backend, SDBackendHandleDeleter>;
enum class SDSplitMode {
LAYER,
ROW,
};
class SDBackendManager {
private:
SDBackendAssignment runtime_assignment_;
SDBackendAssignment params_assignment_;
SDBackendAssignment split_mode_assignment_;
std::unordered_map<std::string, SDBackendHandle> backends_;
public:
@ -52,15 +59,23 @@ public:
bool init(const char* backend_spec,
const char* params_backend_spec,
const char* split_mode_spec,
std::string* error);
void reset();
ggml_backend_t runtime_backend(SDBackendModule module);
ggml_backend_t params_backend(SDBackendModule module);
std::vector<ggml_backend_t> runtime_backends(SDBackendModule module);
SDSplitMode split_mode(SDBackendModule module) const;
ggml_backend_buffer_type_t split_buffer_type(ggml_backend_t backend,
const std::vector<float>& tensor_split);
bool runtime_backend_is_cpu(SDBackendModule module);
bool params_backend_is_cpu(SDBackendModule module);
bool params_backend_is_disk(SDBackendModule module) const;
bool params_backend_follows_runtime(SDBackendModule module) const;
bool runtime_backend_supports_host_buffer(SDBackendModule module);
private:

View file

@ -0,0 +1,257 @@
#include "core/layer_split_partition.h"
#include <algorithm>
#include <cstdlib>
#include <cstring>
#include <limits>
#include <unordered_set>
#include <utility>
#include "core/util.h"
namespace sd {
static bool layer_split_path_segment_starts_at(const std::string& name, size_t pos) {
return pos == 0 || name[pos - 1] == '.';
}
static bool layer_split_has_path_segment(const std::string& name, const char* segment) {
size_t pos = name.find(segment);
while (pos != std::string::npos) {
if (layer_split_path_segment_starts_at(name, pos)) {
return true;
}
pos = name.find(segment, pos + 1);
}
return false;
}
int layer_split_tensor_block_index(const std::string& name) {
static const char* unet_block_segments[] = {"input_blocks.", "output_blocks.", "middle_block.",
"down_blocks.", "up_blocks.", "mid_block."};
for (const char* segment : unet_block_segments) {
if (layer_split_has_path_segment(name, segment)) {
return -1;
}
}
static const char* block_keywords[] = {"transformer_blocks.", "joint_blocks.", "double_blocks.",
"single_blocks.", "blocks.", "block.", "layers."};
for (const char* keyword : block_keywords) {
size_t pos = name.find(keyword);
while (pos != std::string::npos) {
if (!layer_split_path_segment_starts_at(name, pos)) {
pos = name.find(keyword, pos + 1);
continue;
}
pos += std::strlen(keyword);
size_t end = pos;
while (end < name.size() && name[end] >= '0' && name[end] <= '9') {
end++;
}
if (end > pos && (end == name.size() || name[end] == '.')) {
return std::atoi(name.substr(pos, end - pos).c_str());
}
break;
}
}
return -1;
}
std::string layer_split_backend_device_display_name(ggml_backend_t backend) {
ggml_backend_dev_t dev = ggml_backend_get_device(backend);
const char* name = dev != nullptr ? ggml_backend_dev_name(dev) : ggml_backend_name(backend);
return name != nullptr ? name : "unknown";
}
static size_t graph_cut_layer_split_backend_vram_limit(const std::vector<size_t>& backend_vram_limits,
size_t backend_index,
size_t primary_backend_vram_limit) {
if (backend_index < backend_vram_limits.size()) {
return backend_vram_limits[backend_index];
}
return backend_index == 0 ? primary_backend_vram_limit : 0;
}
static std::vector<int64_t> graph_cut_layer_split_backend_capacities(const std::vector<ggml_backend_t>& backends,
const std::vector<size_t>& backend_vram_limits,
size_t primary_backend_vram_limit) {
std::vector<int64_t> capacities(backends.size(), std::numeric_limits<int64_t>::max() / 4);
constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024;
for (size_t i = 0; i < backends.size(); i++) {
ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]);
size_t free_bytes = 0, total_bytes = 0;
if (dev != nullptr) {
ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
}
if (free_bytes > 0) {
capacities[i] = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes, 0);
}
size_t limit_bytes = graph_cut_layer_split_backend_vram_limit(backend_vram_limits,
i,
primary_backend_vram_limit);
if (limit_bytes > 0) {
capacities[i] = std::min<int64_t>(capacities[i], (int64_t)limit_bytes);
}
}
return capacities;
}
bool partition_graph_cut_layer_split(const char* desc,
ggml_cgraph* gf,
const sd::ggml_graph_cut::Plan& plan,
const std::vector<ggml_backend_t>& split_backends,
const std::vector<size_t>& backend_vram_limits,
size_t primary_backend_vram_limit,
std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
GraphCutLayerSplitAssignment* assignment_out) {
GGML_ASSERT(gf != nullptr);
GGML_ASSERT(assignment_out != nullptr);
GGML_ASSERT(canonical_param_tensor != nullptr);
GGML_ASSERT(!split_backends.empty());
GraphCutLayerSplitAssignment assignment;
assignment.segment_count = plan.segments.size();
assignment.tensors_by_backend.resize(split_backends.size());
assignment.bytes_by_backend.resize(split_backends.size(), 0);
assignment.first_segment_by_backend.resize(split_backends.size(), plan.segments.size());
assignment.last_segment_by_backend.resize(split_backends.size(), 0);
std::vector<std::vector<ggml_tensor*>> segment_params(plan.segments.size());
std::vector<int64_t> segment_param_bytes(plan.segments.size(), 0);
std::unordered_set<ggml_tensor*> seen_params;
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
std::vector<ggml_tensor*> params = sd::ggml_graph_cut::param_tensors(gf, plan.segments[seg_idx]);
for (ggml_tensor* raw_param : params) {
ggml_tensor* param = canonical_param_tensor(raw_param);
if (param == nullptr || !seen_params.insert(param).second) {
continue;
}
segment_params[seg_idx].push_back(param);
segment_param_bytes[seg_idx] += (int64_t)ggml_nbytes(param);
}
}
int64_t total_param_bytes = 0;
for (int64_t bytes : segment_param_bytes) {
total_param_bytes += bytes;
}
if (total_param_bytes <= 0) {
LOG_ERROR("%s graph-cut layer split found no graph params to assign", desc);
return false;
}
std::vector<int64_t> backend_capacities = graph_cut_layer_split_backend_capacities(split_backends,
backend_vram_limits,
primary_backend_vram_limit);
std::vector<ggml_backend_t> backend_by_segment(plan.segments.size(), split_backends[0]);
size_t current_backend = 0;
int64_t current_used = 0;
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
int64_t bytes = segment_param_bytes[seg_idx];
while (current_backend + 1 < split_backends.size() &&
bytes > 0 &&
current_used + bytes > backend_capacities[current_backend]) {
current_backend++;
current_used = 0;
}
if (bytes > 0 && current_used + bytes > backend_capacities[current_backend]) {
LOG_ERROR("%s graph-cut layer split: segment %zu needs %.1f MB on %s, but only %.1f MB is available under current VRAM limits",
desc,
seg_idx,
(current_used + bytes) / (1024.0 * 1024.0),
layer_split_backend_device_display_name(split_backends[current_backend]).c_str(),
backend_capacities[current_backend] / (1024.0 * 1024.0));
return false;
}
current_used += bytes;
backend_by_segment[seg_idx] = split_backends[current_backend];
for (ggml_tensor* param : segment_params[seg_idx]) {
ggml_backend_t target_backend = split_backends[current_backend];
auto assigned_it = param_assignments.find(param);
if (assigned_it == param_assignments.end()) {
param_assignments[param] = target_backend;
assignment.has_new_param_assignment = true;
} else {
target_backend = assigned_it->second;
}
auto backend_it = std::find(split_backends.begin(), split_backends.end(), target_backend);
if (backend_it == split_backends.end()) {
LOG_ERROR("%s graph-cut layer split tensor '%s' is assigned to an unavailable backend",
desc,
ggml_get_name(param));
return false;
}
size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it);
assignment.first_segment_by_backend[backend_idx] = std::min(assignment.first_segment_by_backend[backend_idx], seg_idx);
assignment.last_segment_by_backend[backend_idx] = std::max(assignment.last_segment_by_backend[backend_idx], seg_idx + 1);
assignment.tensors_by_backend[backend_idx].push_back(param);
assignment.bytes_by_backend[backend_idx] += (int64_t)ggml_nbytes(param);
}
}
const int n_nodes = ggml_graph_n_nodes(gf);
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
ggml_backend_t backend = backend_by_segment[seg_idx];
const auto& segment = plan.segments[seg_idx];
for (int node_index : segment.internal_node_indices) {
if (node_index < 0 || node_index >= n_nodes) {
continue;
}
ggml_tensor* node = ggml_graph_node(gf, node_index);
if (node != nullptr) {
assignment.node_assignments[node] = backend;
}
}
for (int node_index : segment.output_node_indices) {
if (node_index < 0 || node_index >= n_nodes) {
continue;
}
ggml_tensor* node = ggml_graph_node(gf, node_index);
if (node != nullptr) {
assignment.node_assignments[node] = backend;
}
}
}
*assignment_out = std::move(assignment);
return true;
}
void log_graph_cut_layer_split_assignment(const char* desc,
const std::vector<ggml_backend_t>& split_backends,
const GraphCutLayerSplitAssignment& assignment) {
for (size_t i = 0; i < split_backends.size(); i++) {
if (i >= assignment.tensors_by_backend.size() ||
assignment.tensors_by_backend[i].empty()) {
continue;
}
size_t first_segment = assignment.first_segment_by_backend[i] == assignment.segment_count
? 0
: assignment.first_segment_by_backend[i];
size_t last_segment = assignment.last_segment_by_backend[i];
if (assignment.has_new_param_assignment) {
LOG_INFO("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
desc,
layer_split_backend_device_display_name(split_backends[i]).c_str(),
first_segment,
last_segment,
assignment.tensors_by_backend[i].size(),
assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
} else {
LOG_DEBUG("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
desc,
layer_split_backend_device_display_name(split_backends[i]).c_str(),
first_segment,
last_segment,
assignment.tensors_by_backend[i].size(),
assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
}
}
}
} // namespace sd

View file

@ -0,0 +1,44 @@
#ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__
#define __SD_CORE_LAYER_SPLIT_PARTITION_H__
#include <cstdint>
#include <functional>
#include <string>
#include <unordered_map>
#include <vector>
#include "ggml-backend.h"
#include "ggml.h"
#include "core/ggml_graph_cut.h"
namespace sd {
struct GraphCutLayerSplitAssignment {
std::vector<std::vector<ggml_tensor*>> tensors_by_backend;
std::vector<int64_t> bytes_by_backend;
std::vector<size_t> first_segment_by_backend;
std::vector<size_t> last_segment_by_backend;
std::unordered_map<const ggml_tensor*, ggml_backend_t> node_assignments;
size_t segment_count = 0;
bool has_new_param_assignment = false;
};
std::string layer_split_backend_device_display_name(ggml_backend_t backend);
int layer_split_tensor_block_index(const std::string& name);
bool partition_graph_cut_layer_split(const char* desc,
ggml_cgraph* gf,
const sd::ggml_graph_cut::Plan& plan,
const std::vector<ggml_backend_t>& split_backends,
const std::vector<size_t>& backend_vram_limits,
size_t primary_backend_vram_limit,
std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
GraphCutLayerSplitAssignment* assignment_out);
void log_graph_cut_layer_split_assignment(const char* desc,
const std::vector<ggml_backend_t>& split_backends,
const GraphCutLayerSplitAssignment& assignment);
} // namespace sd
#endif // __SD_CORE_LAYER_SPLIT_PARTITION_H__

View file

@ -4,6 +4,8 @@
#include <cmath>
#include <codecvt>
#include <cstdarg>
#include <cstdlib>
#include <cstring>
#include <exception>
#include <filesystem>
#include <fstream>
@ -29,6 +31,7 @@
#include <unistd.h>
#endif
#include "ggml-backend.h"
#include "ggml.h"
#include "stable-diffusion.h"
@ -1042,3 +1045,26 @@ std::vector<std::pair<std::string, float>> split_quotation_attention(
}
return result;
}
size_t sd_list_devices(char* buffer, size_t buffer_size) {
if (ggml_backend_dev_count() == 0) {
// dynamic-backend builds discover their backend modules at runtime
ggml_backend_load_all();
}
std::ostringstream oss;
for (size_t i = 0; i < ggml_backend_dev_count(); i++) {
ggml_backend_dev_t dev = ggml_backend_dev_get(i);
const char* name = ggml_backend_dev_name(dev);
const char* desc = ggml_backend_dev_description(dev);
oss << (name ? name : "") << '\t' << (desc ? desc : "") << '\n';
}
std::string devices = oss.str();
if (buffer != nullptr && buffer_size > 0) {
size_t copy_size = std::min(devices.size(), buffer_size - 1);
memcpy(buffer, devices.data(), copy_size);
buffer[copy_size] = '\0';
}
return devices.size();
}

View file

@ -27,8 +27,6 @@ namespace kcpp_sd {
model_info get_model_info(sd_ctx_t* ctx);
void SetCircularAxesAll(sd_ctx_t* ctx, bool circular_x, bool circular_y);
void set_lora_cache(sd_ctx_t *ctx, bool enable);
void apply_loras(sd_ctx_t *ctx, const std::vector<sd_lora_t>& lora_specs);

View file

@ -5,7 +5,8 @@
#include "model_loader.h"
#include "model_manager.h"
#define CONTROL_NET_GRAPH_SIZE 1536
// Match main UNet's MAX_GRAPH_SIZE so SDXL ControlNet (transformer_depth={1,2,10}) fits.
#define CONTROL_NET_GRAPH_SIZE MAX_GRAPH_SIZE
/*
=================================== ControlNet ===================================

View file

@ -4,6 +4,7 @@
#include <memory>
#include <vector>
#include "core/util.h"
#include "model/adapter/pulid.hpp"
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
@ -1400,18 +1401,28 @@ namespace Flux {
std::vector<float> dct_vec;
sd::Tensor<float> guidance_tensor;
SDVersion version;
bool use_mask = false;
bool use_mask = true;
FluxRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
SDVersion version = VERSION_FLUX,
bool use_mask = false,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const char* model_args = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(FluxConfig::detect_from_weights(tensor_storage_map, prefix, version)),
version(version),
use_mask(use_mask) {
version(version) {
for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
if (key == "chroma_use_dit_mask") {
bool parsed = true;
if (parse_strict_bool(value, parsed)) {
use_mask = parsed;
} else {
LOG_WARN("ignoring invalid Chroma DiT model arg '%s=%s'", key.c_str(), value.c_str());
}
}
}
if (config.is_chroma) {
LOG_INFO("Using pruned modulation (Chroma)");
}
@ -1718,7 +1729,6 @@ namespace Flux {
tensor_storage_map,
"model.diffusion_model",
VERSION_FLUX2,
false,
model_manager);
if (!model_manager->register_runner_params("Flux test",

View file

@ -3,6 +3,7 @@
#include <memory>
#include "core/util.h"
#include "model/common/block.hpp"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/flux.hpp"
@ -566,12 +567,21 @@ namespace Qwen {
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
SDVersion version = VERSION_QWEN_IMAGE,
bool zero_cond_t = false,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const char* model_args = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)),
version(version) {
config.zero_cond_t = config.zero_cond_t || zero_cond_t;
for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
if (key == "qwen_image_zero_cond_t") {
bool parsed = false;
if (parse_strict_bool(value, parsed)) {
config.zero_cond_t = config.zero_cond_t || parsed;
} else {
LOG_WARN("ignoring invalid Qwen Image model arg '%s=%s'", key.c_str(), value.c_str());
}
}
}
if (version == VERSION_QWEN_IMAGE_LAYERED) {
config.use_additional_t_cond = true;
}
@ -775,7 +785,6 @@ namespace Qwen {
tensor_storage_map,
"model.diffusion_model",
VERSION_QWEN_IMAGE,
false,
model_manager);
if (!model_manager->register_runner_params("Qwen image test",

View file

@ -133,7 +133,11 @@ public:
int64_t H = input.shape()[1] / scale_factor;
float tile_overlap;
int tile_size_x, tile_size_y;
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, 1.30539f);
// Image VAE encode is more sensitive to tile boundary context than decode.
// Keep the smaller legacy factor for video VAEs, but default image encode
// tiles to 64 latent pixels so a 512px SD image is encoded as one tile.
const float encode_tile_factor = (sd_version_is_wan(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f;
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor);
LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
output = tiled_compute(input,
n_threads,

View file

@ -1,7 +1,10 @@
#include "gguf_io.h"
#include <algorithm>
#include <cstdint>
#include <cstdio>
#include <fstream>
#include <ostream>
#include <string>
#include <vector>
@ -121,3 +124,115 @@ bool write_gguf_file(const std::string& file_path,
gguf_free(gguf_ctx);
return success;
}
GGUFStreamingWriter::~GGUFStreamingWriter() {
close();
}
bool GGUFStreamingWriter::write_metadata(const std::string& file_path,
const std::vector<TensorWritePlan>& tensors,
std::string* error) {
close();
tensors_ = tensors;
file_size_ = 0;
size_t meta_mem = 1 * 1024 * 1024 + tensors.size() * ggml_tensor_overhead();
meta_ctx_ = ggml_init({meta_mem, nullptr, true});
if (meta_ctx_ == nullptr) {
set_error(error, "ggml_init failed for GGUF metadata");
return false;
}
gguf_ctx_ = gguf_init_empty();
if (gguf_ctx_ == nullptr) {
set_error(error, "gguf_init_empty failed");
close();
return false;
}
for (const TensorWritePlan& plan : tensors) {
ggml_tensor* tensor = ggml_new_tensor(meta_ctx_, plan.type, plan.n_dims, plan.ne);
if (tensor == nullptr) {
set_error(error, "ggml_new_tensor failed for tensor '" + plan.name + "'");
close();
return false;
}
ggml_set_name(tensor, plan.name.c_str());
gguf_add_tensor(gguf_ctx_, tensor);
}
LOG_INFO("trying to save tensors to %s", file_path.c_str());
FILE* file = fopen(file_path.c_str(), "wb+");
if (file == nullptr) {
set_error(error, "failed to open output file '" + file_path + "'");
close();
return false;
}
// ggml exposes GGUF metadata writing through FILE* only. Keep FILE usage
// isolated here; tensor data is written through std::fstream by the shared
// streaming pipeline.
if (!gguf_write_to_file_ptr(gguf_ctx_, file, true)) {
fclose(file);
set_error(error, "failed to write GGUF metadata to '" + file_path + "'");
close();
return false;
}
fclose(file);
const uint64_t data_start = gguf_get_meta_size(gguf_ctx_);
tensor_offsets_.resize(tensors.size());
file_size_ = data_start;
for (size_t i = 0; i < tensors.size(); i++) {
tensor_offsets_[i] = data_start + gguf_get_tensor_offset(gguf_ctx_, static_cast<int64_t>(i));
file_size_ = std::max(file_size_, tensor_offsets_[i] + tensors[i].nbytes());
}
return true;
}
bool GGUFStreamingWriter::write_tensor(std::ostream& output,
size_t tensor_index,
const uint8_t* data,
size_t size,
std::string* error) const {
if (tensor_index >= tensors_.size() || tensor_index >= tensor_offsets_.size()) {
set_error(error, "invalid GGUF tensor index");
return false;
}
const TensorWritePlan& plan = tensors_[tensor_index];
if (size != plan.nbytes()) {
set_error(error, "size mismatch while writing tensor '" + plan.name + "'");
return false;
}
output.seekp(static_cast<std::streamoff>(tensor_offsets_[tensor_index]), std::ios::beg);
if (!output) {
set_error(error, "failed to seek output for tensor '" + plan.name + "'");
return false;
}
if (size > 0) {
output.write(reinterpret_cast<const char*>(data), static_cast<std::streamsize>(size));
}
if (!output) {
set_error(error, "failed to write tensor '" + plan.name + "'");
return false;
}
return true;
}
uint64_t GGUFStreamingWriter::file_size() const {
return file_size_;
}
void GGUFStreamingWriter::close() {
tensor_offsets_.clear();
tensors_.clear();
file_size_ = 0;
if (gguf_ctx_ != nullptr) {
gguf_free(gguf_ctx_);
gguf_ctx_ = nullptr;
}
if (meta_ctx_ != nullptr) {
ggml_free(meta_ctx_);
meta_ctx_ = nullptr;
}
}

View file

@ -4,8 +4,12 @@
#include <string>
#include <vector>
#include "streaming_writer.h"
#include "tensor_storage.h"
struct ggml_context;
struct gguf_context;
bool is_gguf_file(const std::string& file_path);
bool read_gguf_file(const std::string& file_path,
std::vector<TensorStorage>& tensor_storages,
@ -14,4 +18,28 @@ bool write_gguf_file(const std::string& file_path,
const std::vector<TensorWriteInfo>& tensors,
std::string* error = nullptr);
class GGUFStreamingWriter : public StreamingModelWriter {
public:
GGUFStreamingWriter() = default;
~GGUFStreamingWriter();
bool write_metadata(const std::string& file_path,
const std::vector<TensorWritePlan>& tensors,
std::string* error = nullptr) override;
bool write_tensor(std::ostream& output,
size_t tensor_index,
const uint8_t* data,
size_t size,
std::string* error = nullptr) const override;
uint64_t file_size() const override;
void close();
private:
std::vector<TensorWritePlan> tensors_;
std::vector<uint64_t> tensor_offsets_;
uint64_t file_size_ = 0;
ggml_context* meta_ctx_ = nullptr;
gguf_context* gguf_ctx_ = nullptr;
};
#endif // __SD_MODEL_IO_GGUF_IO_H__

View file

@ -1,8 +1,10 @@
#include "safetensors_io.h"
#include <algorithm>
#include <cstdint>
#include <exception>
#include <fstream>
#include <ostream>
#include <string>
#include <vector>
@ -41,7 +43,7 @@ bool is_safetensors_file(const std::string& file_path) {
}
size_t header_size_ = model_io::read_u64(header_size_buf);
if (header_size_ >= file_size_ || header_size_ <= 2) {
if (header_size_ > file_size_ - ST_HEADER_SIZE_LEN || header_size_ <= 2) {
return false;
}
@ -112,10 +114,11 @@ bool read_safetensors_file(const std::string& file_path,
}
size_t header_size_ = model_io::read_u64(header_size_buf);
if (header_size_ >= file_size_) {
if (header_size_ > file_size_ - ST_HEADER_SIZE_LEN) {
set_error(error, "invalid safetensor file '" + file_path + "'");
return false;
}
const size_t data_start = ST_HEADER_SIZE_LEN + header_size_;
// read header
std::vector<char> header_buf;
@ -154,6 +157,10 @@ bool read_safetensors_file(const std::string& file_path,
size_t begin = tensor_info["data_offsets"][0].get<size_t>();
size_t end = tensor_info["data_offsets"][1].get<size_t>();
if (begin > end || end > file_size_ - data_start) {
set_error(error, "data offsets out of bounds for tensor '" + name + "'");
return false;
}
ggml_type type = safetensors_dtype_to_ggml_type(dtype);
if (type == GGML_TYPE_COUNT) {
@ -185,7 +192,7 @@ bool read_safetensors_file(const std::string& file_path,
n_dims = 1;
}
TensorStorage tensor_storage(name, type, ne, n_dims, 0, ST_HEADER_SIZE_LEN + header_size_ + begin);
TensorStorage tensor_storage(name, type, ne, n_dims, 0, data_start + begin);
tensor_storage.reverse_ne();
size_t tensor_data_size = end - begin;
@ -314,3 +321,102 @@ bool write_safetensors_file(const std::string& file_path,
return true;
}
bool SafetensorsStreamingWriter::write_metadata(const std::string& file_path,
const std::vector<TensorWritePlan>& tensors,
std::string* error) {
file_path_ = file_path;
tensors_ = tensors;
tensor_offsets_.clear();
data_start_ = 0;
file_size_ = 0;
nlohmann::ordered_json header = nlohmann::ordered_json::object();
uint64_t data_offset = 0;
tensor_offsets_.resize(tensors.size());
for (size_t i = 0; i < tensors.size(); i++) {
const TensorWritePlan& plan = tensors[i];
std::string dtype;
if (!ggml_type_to_safetensors_dtype(plan.type, &dtype)) {
set_error(error,
"unsupported safetensors dtype '" + std::string(ggml_type_name(plan.type)) +
"' for tensor '" + plan.name + "'");
return false;
}
nlohmann::ordered_json json_tensor_info = nlohmann::ordered_json::object();
json_tensor_info["dtype"] = dtype;
nlohmann::ordered_json shape = nlohmann::ordered_json::array();
for (int j = 0; j < plan.n_dims; ++j) {
shape.push_back(plan.ne[plan.n_dims - 1 - j]);
}
json_tensor_info["shape"] = shape;
nlohmann::ordered_json data_offsets = nlohmann::ordered_json::array();
data_offsets.push_back(data_offset);
data_offsets.push_back(data_offset + plan.nbytes());
json_tensor_info["data_offsets"] = data_offsets;
header[plan.name] = json_tensor_info;
tensor_offsets_[i] = data_offset;
data_offset += plan.nbytes();
}
const std::string header_str = header.dump();
data_start_ = ST_HEADER_SIZE_LEN + header_str.size();
LOG_INFO("trying to save tensors to %s", file_path.c_str());
std::ofstream file(file_path, std::ios::binary | std::ios::trunc);
if (!file.is_open()) {
set_error(error, "failed to open '" + file_path + "' for writing");
return false;
}
uint8_t header_size[ST_HEADER_SIZE_LEN];
for (int i = 0; i < static_cast<int>(ST_HEADER_SIZE_LEN); ++i) {
header_size[i] = static_cast<uint8_t>((header_str.size() >> (8 * i)) & 0xFF);
}
file.write(reinterpret_cast<const char*>(header_size), sizeof(header_size));
file.write(header_str.data(), static_cast<std::streamsize>(header_str.size()));
if (!file) {
set_error(error, "failed to write safetensors header to '" + file_path + "'");
return false;
}
file_size_ = data_start_ + data_offset;
return true;
}
bool SafetensorsStreamingWriter::write_tensor(std::ostream& output,
size_t tensor_index,
const uint8_t* data,
size_t size,
std::string* error) const {
if (tensor_index >= tensors_.size() || tensor_index >= tensor_offsets_.size()) {
set_error(error, "invalid safetensors tensor index");
return false;
}
const TensorWritePlan& plan = tensors_[tensor_index];
if (size != plan.nbytes()) {
set_error(error, "size mismatch while writing tensor '" + plan.name + "'");
return false;
}
output.seekp(static_cast<std::streamoff>(data_start_ + tensor_offsets_[tensor_index]), std::ios::beg);
if (!output) {
set_error(error, "failed to seek output for tensor '" + plan.name + "'");
return false;
}
if (size > 0) {
output.write(reinterpret_cast<const char*>(data), static_cast<std::streamsize>(size));
}
if (!output) {
set_error(error, "failed to write tensor '" + plan.name + "' to '" + file_path_ + "'");
return false;
}
return true;
}
uint64_t SafetensorsStreamingWriter::file_size() const {
return file_size_;
}

View file

@ -4,6 +4,7 @@
#include <string>
#include <vector>
#include "streaming_writer.h"
#include "tensor_storage.h"
bool is_safetensors_file(const std::string& file_path);
@ -14,4 +15,26 @@ bool write_safetensors_file(const std::string& file_path,
const std::vector<TensorWriteInfo>& tensors,
std::string* error = nullptr);
class SafetensorsStreamingWriter : public StreamingModelWriter {
public:
SafetensorsStreamingWriter() = default;
bool write_metadata(const std::string& file_path,
const std::vector<TensorWritePlan>& tensors,
std::string* error = nullptr) override;
bool write_tensor(std::ostream& output,
size_t tensor_index,
const uint8_t* data,
size_t size,
std::string* error = nullptr) const override;
uint64_t file_size() const override;
private:
std::string file_path_;
std::vector<TensorWritePlan> tensors_;
std::vector<uint64_t> tensor_offsets_;
uint64_t data_start_ = 0;
uint64_t file_size_ = 0;
};
#endif // __SD_MODEL_IO_SAFETENSORS_IO_H__

View file

@ -0,0 +1,26 @@
#ifndef __SD_MODEL_IO_STREAMING_WRITER_H__
#define __SD_MODEL_IO_STREAMING_WRITER_H__
#include <cstdint>
#include <iosfwd>
#include <string>
#include <vector>
#include "tensor_storage.h"
class StreamingModelWriter {
public:
virtual ~StreamingModelWriter() = default;
virtual bool write_metadata(const std::string& file_path,
const std::vector<TensorWritePlan>& tensors,
std::string* error = nullptr) = 0;
virtual bool write_tensor(std::ostream& output,
size_t tensor_index,
const uint8_t* data,
size_t size,
std::string* error = nullptr) const = 0;
virtual uint64_t file_size() const = 0;
};
#endif // __SD_MODEL_IO_STREAMING_WRITER_H__

View file

@ -127,6 +127,25 @@ struct TensorWriteInfo {
ggml_tensor* tensor = nullptr;
};
struct TensorWritePlan {
std::string name;
ggml_type type = GGML_TYPE_F32;
int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
int n_dims = 0;
int64_t nelements() const {
int64_t n = 1;
for (int i = 0; i < SD_MAX_DIMS; i++) {
n *= ne[i];
}
return n;
}
uint64_t nbytes() const {
return nelements() * ggml_type_size(type) / ggml_blck_size(type);
}
};
typedef std::function<bool(const TensorStorage&, ggml_tensor**)> on_new_tensor_cb_t;
#endif // __SD_TENSOR_STORAGE_H__

View file

@ -971,7 +971,8 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
bool enable_mmap,
const std::set<std::string>* target_tensor_names) {
const std::set<std::string>* target_tensor_names,
bool log_progress) {
process_model_files(enable_mmap, false);
std::atomic<int64_t> read_time_ms(0);
@ -1240,7 +1241,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
}
size_t curr_num = total_tensors_processed + current_idx;
float elapsed_seconds = (ggml_time_ms() - t_start) / 1000.0f;
if (total_tensors_to_process > 0) {
if (log_progress && total_tensors_to_process > 0) {
pretty_bytes_progress(static_cast<int>(curr_num),
static_cast<int>(total_tensors_to_process),
bytes_processed.load(),
@ -1258,27 +1259,81 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
break;
}
total_tensors_processed += tensors_to_process.size();
if (total_tensors_to_process > 0) {
if (log_progress && total_tensors_to_process > 0) {
pretty_bytes_progress(static_cast<int>(total_tensors_processed),
static_cast<int>(total_tensors_to_process),
bytes_processed.load(),
(ggml_time_ms() - t_start) / 1000.0f);
}
if (total_tensors_processed < total_tensors_to_process && total_tensors_to_process > 0) {
if (log_progress && total_tensors_processed < total_tensors_to_process && total_tensors_to_process > 0) {
printf("\n");
}
}
int64_t end_time = ggml_time_ms();
LOG_INFO("loading tensors completed, taking %.2fs (read: %.2fs, memcpy: %.2fs, convert: %.2fs, copy_to_backend: %.2fs)",
(end_time - start_time) / 1000.f,
(read_time_ms.load() / (float)last_n_threads) / 1000.f,
(memcpy_time_ms.load() / (float)last_n_threads) / 1000.f,
(convert_time_ms.load() / (float)last_n_threads) / 1000.f,
(copy_to_backend_time_ms.load() / (float)last_n_threads) / 1000.f);
if (log_progress) {
LOG_INFO("loading tensors completed, taking %.2fs (read: %.2fs, memcpy: %.2fs, convert: %.2fs, copy_to_backend: %.2fs)",
(end_time - start_time) / 1000.f,
(read_time_ms.load() / (float)last_n_threads) / 1000.f,
(memcpy_time_ms.load() / (float)last_n_threads) / 1000.f,
(convert_time_ms.load() / (float)last_n_threads) / 1000.f,
(copy_to_backend_time_ms.load() / (float)last_n_threads) / 1000.f);
}
return success;
}
bool ModelLoader::load_tensor(const TensorStorage& tensor_storage, ggml_tensor* dst_tensor) {
if (dst_tensor == nullptr || dst_tensor->data == nullptr) {
LOG_ERROR("load tensor failed: null destination for '%s'", tensor_storage.name.c_str());
return false;
}
bool loaded = false;
std::set<std::string> target_tensor_names{tensor_storage.name};
auto on_new_tensor_cb = [&](const TensorStorage& current_tensor_storage, ggml_tensor** out_tensor) -> bool {
*out_tensor = nullptr;
if (current_tensor_storage.name != tensor_storage.name) {
return true;
}
if (current_tensor_storage.file_index != tensor_storage.file_index ||
current_tensor_storage.offset != tensor_storage.offset ||
current_tensor_storage.index_in_zip != tensor_storage.index_in_zip) {
LOG_ERROR("load tensor failed: storage mismatch for '%s'", tensor_storage.name.c_str());
return false;
}
if (current_tensor_storage.n_dims != tensor_storage.n_dims ||
current_tensor_storage.nelements() != tensor_storage.nelements()) {
LOG_ERROR("load tensor failed: metadata changed for '%s'", tensor_storage.name.c_str());
return false;
}
for (int i = 0; i < current_tensor_storage.n_dims; i++) {
if (current_tensor_storage.ne[i] != dst_tensor->ne[i]) {
LOG_ERROR("load tensor failed: shape mismatch for '%s'", tensor_storage.name.c_str());
return false;
}
}
*out_tensor = dst_tensor;
loaded = true;
return true;
};
if (!load_tensors(on_new_tensor_cb, false, &target_tensor_names, false)) {
LOG_ERROR("load tensor failed: '%s'", tensor_storage.name.c_str());
return false;
}
if (!loaded) {
LOG_ERROR("load tensor failed: tensor '%s' not found", tensor_storage.name.c_str());
return false;
}
return true;
}
bool ModelLoader::load_float_tensor(const std::string& name,
std::vector<float>& data,
int n_threads,

View file

@ -27,6 +27,8 @@ struct MmapTensorStore {
std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
};
bool is_unused_tensor(const std::string& name);
class ModelLoader {
protected:
SDVersion version_ = VERSION_COUNT;
@ -68,7 +70,8 @@ public:
bool writable = true);
bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
bool use_mmap = false,
const std::set<std::string>* target_tensor_names = nullptr);
const std::set<std::string>* target_tensor_names = nullptr,
bool log_progress = true);
bool load_tensors(std::map<std::string, ggml_tensor*>& tensors,
std::set<std::string> ignore_tensors = {},
bool use_mmap = false);
@ -76,6 +79,7 @@ public:
std::vector<float>& data,
int n_threads = 0,
bool use_mmap = false);
bool load_tensor(const TensorStorage& tensor_storage, ggml_tensor* dst_tensor);
std::vector<std::string> get_tensor_names() const {
std::vector<std::string> names;

View file

@ -100,12 +100,42 @@ size_t estimate_tensors_size(const std::map<std::string, ggml_tensor*>& tensors)
return size;
}
void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft) {
if (compute_backend == nullptr) {
return;
}
if (split_buft == nullptr) {
split_buffer_types_.erase(compute_backend);
return;
}
split_buffer_types_[compute_backend] = split_buft;
}
bool ModelManager::tensor_shape_supports_split_buffer(const ggml_tensor* tensor) {
return tensor != nullptr &&
tensor->view_src == nullptr &&
ggml_is_contiguous(tensor) &&
ggml_n_dims(tensor) == 2 &&
tensor->ne[0] >= 256 &&
tensor->ne[1] >= 256;
}
ggml_backend_buffer_type_t ModelManager::split_buffer_type_for(const TensorState& state) const {
if (!state.allow_split_buffer || !tensor_shape_supports_split_buffer(state.tensor)) {
return nullptr;
}
auto it = split_buffer_types_.find(state.compute_backend);
return it != split_buffer_types_.end() ? it->second : nullptr;
}
bool ModelManager::register_param_tensors(const std::string& desc,
std::map<std::string, ggml_tensor*> tensors,
ResidencyMode residency_mode,
ggml_backend_t compute_backend,
ggml_backend_t params_backend,
size_t* registered_tensor_size) {
size_t* registered_tensor_size,
bool allow_split_buffer,
bool params_follow_compute_backend) {
if (desc.empty()) {
LOG_ERROR("model manager tensor desc is empty");
return false;
@ -129,13 +159,15 @@ bool ModelManager::register_param_tensors(const std::string& desc,
}
ggml_set_name(tensor, name.c_str());
auto state = std::make_unique<TensorState>();
state->name = name;
state->tensor = tensor;
state->desc = desc;
state->residency_mode = residency_mode;
state->compute_backend = compute_backend;
state->params_backend = params_backend;
auto state = std::make_unique<TensorState>();
state->name = name;
state->tensor = tensor;
state->desc = desc;
state->residency_mode = residency_mode;
state->compute_backend = compute_backend;
state->params_backend = params_backend;
state->allow_split_buffer = allow_split_buffer;
state->params_follow_compute_backend = params_follow_compute_backend;
new_states.push_back(std::move(state));
}
@ -237,7 +269,7 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector<TensorState*
}
bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorState*>& states) {
std::map<ggml_backend_t, std::vector<TensorState*>> states_by_compute_backend;
std::map<std::pair<ggml_backend_t, ggml_backend_buffer_type_t>, std::vector<TensorState*>> states_by_staging_target;
for (TensorState* state : states) {
if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
continue;
@ -257,11 +289,16 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
LOG_ERROR("model manager tensor '%s' is not loaded to params backend", state->name.c_str());
return false;
}
states_by_compute_backend[state->compute_backend].push_back(state);
ggml_backend_buffer_type_t staging_buft = split_buffer_type_for(*state);
if (staging_buft == nullptr) {
staging_buft = ggml_backend_get_default_buffer_type(state->compute_backend);
}
states_by_staging_target[{state->compute_backend, staging_buft}].push_back(state);
}
for (const auto& pair : states_by_compute_backend) {
ggml_backend_t compute_backend = pair.first;
for (const auto& pair : states_by_staging_target) {
ggml_backend_t compute_backend = pair.first.first;
ggml_backend_buffer_type_t staging_buft = pair.first.second;
const std::vector<TensorState*>& states = pair.second;
if (states.empty()) {
continue;
@ -285,7 +322,7 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
staged_tensors.push_back({state, staging_tensor});
}
ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors(staging_ctx, compute_backend);
ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, staging_buft);
if (compute_buffer == nullptr) {
LOG_ERROR("model manager alloc compute params backend buffer failed, num_tensors = %zu",
staged_tensors.size());
@ -350,6 +387,17 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
LOG_ERROR("model manager compute backend is null for lora target tensor '%s'", state->name.c_str());
return false;
}
if (state->tensor->buffer != nullptr &&
ggml_backend_buffer_get_type(state->tensor->buffer) == split_buffer_type_for(*state)) {
if (!warned_split_lora_skip_) {
LOG_WARN(
"model manager skipping direct lora application to row-split tensors "
"(use --lora-apply-mode at_runtime with row split)");
warned_split_lora_skip_ = true;
}
state->applied_lora_epoch = current_lora_epoch_;
continue;
}
if (state->tensor->data == nullptr) {
LOG_ERROR("model manager lora target tensor '%s' is not prepared", state->name.c_str());
return false;
@ -694,6 +742,8 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat
if (compute_dev != nullptr) {
params_buft = ggml_backend_dev_host_buffer_type(compute_dev);
}
} else if (state.params_backend == state.compute_backend) {
params_buft = split_buffer_type_for(state);
}
if (params_buft == nullptr) {
params_buft = ggml_backend_get_default_buffer_type(state.params_backend);
@ -871,6 +921,54 @@ bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*
return true;
}
bool ModelManager::assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
ggml_backend_t compute_backend) {
if (tensors.empty()) {
return true;
}
if (compute_backend == nullptr) {
LOG_ERROR("model manager cannot assign tensors to a null compute backend");
return false;
}
std::vector<TensorState*> required_states;
if (!resolve_required_tensor_states(tensors, required_states)) {
return false;
}
for (TensorState* state : required_states) {
if (state == nullptr || state->tensor == nullptr) {
continue;
}
const bool params_follow_compute = state->params_follow_compute_backend ||
state->residency_mode == ResidencyMode::Disk;
const bool compute_changes = state->compute_backend != compute_backend;
const bool params_changes = params_follow_compute && state->params_backend != compute_backend;
if (!compute_changes && !params_changes) {
continue;
}
if (state->active_prepare_count > 0 || state->staged_to_compute_backend) {
LOG_ERROR("model manager cannot move active tensor '%s' to another compute backend",
state->name.c_str());
return false;
}
if (params_changes && state->loaded_to_params_backend) {
LOG_ERROR("model manager cannot move loaded tensor '%s' to another params backend",
state->name.c_str());
return false;
}
state->compute_backend = compute_backend;
if (params_follow_compute) {
state->params_backend = compute_backend;
}
}
return true;
}
bool ModelManager::prepare_params(const std::vector<ggml_tensor*>& tensors) {
if (tensors.empty()) {
return true;

View file

@ -33,10 +33,12 @@ private:
ggml_tensor* tensor = nullptr;
std::string desc;
ResidencyMode residency_mode = ResidencyMode::ParamBackend;
ggml_backend_t compute_backend = nullptr;
ggml_backend_t params_backend = nullptr;
bool metadata_validated = false;
ResidencyMode residency_mode = ResidencyMode::ParamBackend;
ggml_backend_t compute_backend = nullptr;
ggml_backend_t params_backend = nullptr;
bool allow_split_buffer = false;
bool params_follow_compute_backend = false;
bool metadata_validated = false;
int active_prepare_count = 0;
@ -63,6 +65,8 @@ private:
std::map<std::string, TensorState*> tensor_states_by_name_;
std::vector<std::unique_ptr<ParamsStorageBlock>> params_storage_blocks_;
std::vector<std::unique_ptr<ComputeStagingBlock>> compute_staging_blocks_;
std::map<ggml_backend_t, ggml_backend_buffer_type_t> split_buffer_types_;
bool warned_split_lora_skip_ = false;
std::set<std::string> common_ignore_tensors_;
std::vector<LoraSpec> loras_;
SDVersion lora_version_ = VERSION_COUNT;
@ -91,6 +95,7 @@ private:
bool stage_tensors_to_compute_backend(const std::vector<TensorState*>& states);
ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const;
ggml_backend_buffer_type_t split_buffer_type_for(const TensorState& state) const;
void release_compute_staging_blocks(bool force = false,
const std::unordered_set<TensorState*>* target_states = nullptr);
void release_params_storage_blocks(bool force = false,
@ -114,6 +119,9 @@ public:
void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; }
void set_common_ignore_tensors(std::set<std::string> ignore_tensors);
void set_loras(std::vector<LoraSpec> loras, SDVersion version);
void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft);
static bool tensor_shape_supports_split_buffer(const ggml_tensor* tensor);
std::set<std::string> tensor_names() const;
@ -122,7 +130,9 @@ public:
ResidencyMode residency_mode,
ggml_backend_t compute_backend,
ggml_backend_t params_backend,
size_t* registered_tensor_size = nullptr);
size_t* registered_tensor_size = nullptr,
bool allow_split_buffer = false,
bool params_follow_compute_backend = false);
template <typename Runner>
bool register_runner_params(const std::string& desc,
@ -162,6 +172,8 @@ public:
bool validate_registered_tensors();
bool load_all_params_eagerly();
bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
ggml_backend_t compute_backend) override;
bool prepare_params(const std::vector<ggml_tensor*>& tensors) override;
void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) override;
void release_params_backend_params(const std::vector<ggml_tensor*>& tensors) override;

View file

@ -736,6 +736,62 @@ std::string convert_diffusers_dit_to_original_krea2(std::string name) {
return name;
}
// Convert a diffusers-format ControlNet tensor name to the original (LDM/lllyasviel) layout
// declared by ControlNetBlock. Reuses the UNet down/mid conversion for the shared encoder
// (down_blocks, mid_block, time_embedding, add_embedding, conv_in) and adds the ControlNet-only
// mappings: input_hint_block, zero_convs, middle_block_out.
std::string convert_diffusers_controlnet_to_original_sdxl(std::string name) {
name = convert_diffusers_unet_to_original_sdxl(std::move(name));
static const std::vector<std::pair<std::string, std::string>> prefix_map = {
{"controlnet_cond_embedding.conv_in.", "input_hint_block.0."},
{"controlnet_cond_embedding.blocks.0.", "input_hint_block.2."},
{"controlnet_cond_embedding.blocks.1.", "input_hint_block.4."},
{"controlnet_cond_embedding.blocks.2.", "input_hint_block.6."},
{"controlnet_cond_embedding.blocks.3.", "input_hint_block.8."},
{"controlnet_cond_embedding.blocks.4.", "input_hint_block.10."},
{"controlnet_cond_embedding.blocks.5.", "input_hint_block.12."},
{"controlnet_cond_embedding.conv_out.", "input_hint_block.14."},
{"controlnet_mid_block.", "middle_block_out.0."},
};
for (const auto& p : prefix_map) {
if (starts_with(name, p.first)) {
return p.second + name.substr(p.first.size());
}
}
static const std::string controlnet_down_prefix = "controlnet_down_blocks.";
if (starts_with(name, controlnet_down_prefix)) {
size_t rest_start = controlnet_down_prefix.size();
size_t dot = name.find('.', rest_start);
if (dot != std::string::npos) {
std::string idx = name.substr(rest_start, dot - rest_start);
return "zero_convs." + idx + ".0" + name.substr(dot);
}
}
return name;
}
static bool is_diffusers_controlnet_name(const std::string& name) {
static const std::vector<std::string> heads = {
"controlnet_cond_embedding.",
"controlnet_down_blocks.",
"controlnet_mid_block.",
"down_blocks.",
"mid_block.",
"time_embedding.",
"add_embedding.",
"conv_in.",
};
for (const auto& h : heads) {
if (starts_with(name, h)) {
return true;
}
}
return false;
}
std::string convert_diffusion_model_name(std::string name, std::string prefix, SDVersion version) {
if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
name = convert_diffusers_unet_to_original_sd1(name);
@ -1338,6 +1394,9 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
name = name.substr(pos + 1);
}
}
if (sd_version_is_sdxl(version) && is_diffusers_controlnet_name(name)) {
name = convert_diffusers_controlnet_to_original_sdxl(name);
}
}
if (is_lora) {

View file

@ -4,7 +4,10 @@
#include <algorithm>
#include <cctype>
#include <cmath>
#include <cstring>
#include <functional>
#include <limits>
#include <map>
#include <string>
#include <utility>
@ -1013,6 +1016,7 @@ struct Denoiser {
const sd::Tensor<float>& latent) = 0;
virtual sd::Tensor<float> inverse_noise_scaling(float sigma,
const sd::Tensor<float>& latent) = 0;
virtual float noise_level_to_sigma(float noise_level) = 0;
virtual std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) {
auto bound_t_to_sigma = std::bind(&Denoiser::t_to_sigma, this, std::placeholders::_1);
@ -1160,6 +1164,10 @@ struct CompVisDenoiser : public Denoiser {
SD_UNUSED(sigma);
return latent;
}
float noise_level_to_sigma(float noise_level) {
return noise_level / (1.0f - noise_level);
}
};
struct CompVisVDenoiser : public CompVisDenoiser {
@ -1247,6 +1255,10 @@ struct DiscreteFlowDenoiser : public Denoiser {
sd::Tensor<float> inverse_noise_scaling(float sigma, const sd::Tensor<float>& latent) override {
return latent * (1.0f / (1.0f - sigma));
}
float noise_level_to_sigma(float noise_level) {
return noise_level;
}
};
struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
@ -1384,6 +1396,11 @@ struct MiniT2IFlowDenoiser : public Denoiser {
return latent;
}
float noise_level_to_sigma(float noise_level) {
SD_UNUSED(noise_level);
return 1.0f;
}
std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) override {
SD_UNUSED(image_seq_len);
SD_UNUSED(scheduler_type);
@ -1814,6 +1831,204 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
return x;
}
// DPM-Solver++(2M) SDE, midpoint variant. Ref: Lu et al. arXiv:2211.01095;
// k-diffusion sample_dpmpp_2m_sde.
static sd::Tensor<float> sample_dpmpp_2m_sde(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
float eta) {
sd::Tensor<float> old_denoised;
bool have_old_denoised = false;
float h_last = 0.f;
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) {
auto denoised_opt = model(x, sigmas[i], i + 1);
if (denoised_opt.pred.empty()) {
return {};
}
sd::Tensor<float> denoised = std::move(denoised_opt.pred);
if (sigmas[i + 1] == 0.f) {
x = denoised;
} else {
float t = -std::log(sigmas[i]);
float s = -std::log(sigmas[i + 1]);
float h = s - t;
float eta_h = eta * h;
float a = sigmas[i + 1] / sigmas[i] * std::exp(-eta_h);
float b = -std::expm1(-h - eta_h);
x = a * x + b * denoised;
if (have_old_denoised) {
float r = h_last / h;
x += (0.5f * b / r) * (denoised - old_denoised);
}
if (eta > 0.f) {
x += sd::Tensor<float>::randn_like(x, rng) * (sigmas[i + 1] * std::sqrt(-std::expm1(-2.f * eta_h)));
}
h_last = h;
}
old_denoised = denoised;
have_old_denoised = true;
}
return x;
}
// Seeded Brownian tree providing deterministic, step-count-stable Gaussian
// increments for stochastic samplers. Constructed once per generation; each
// call returns unit-variance noise for interval [sigma_a, sigma_b].
// Reference: torchsde BrownianTree; k-diffusion BatchedBrownianTree.
class BrownianTreeNoiseSampler {
public:
BrownianTreeNoiseSampler(const sd::Tensor<float>& x_template,
double sigma_min,
double sigma_max,
uint64_t seed)
: t_min_(sigma_min),
t_max_(sigma_max),
shape_(x_template.shape()),
root_seed_(mix64(seed, 0x9E3779B97F4A7C15ULL)) {
auto rng = std::make_shared<STDDefaultRNG>();
rng->manual_seed(mix64(seed, 0xBF58476D1CE4E5B9ULL));
w_at_tmax_ = sd::Tensor<float>::randn(shape_, rng) * std::sqrt(static_cast<float>(t_max_ - t_min_));
}
sd::Tensor<float> operator()(double sigma_a, double sigma_b) {
double a = clamp(std::min(sigma_a, sigma_b));
double b = clamp(std::max(sigma_a, sigma_b));
auto dW = w(b) - w(a);
float span = static_cast<float>(std::max(std::abs(sigma_b - sigma_a), 1e-12));
return dW * (1.0f / std::sqrt(span));
}
private:
static constexpr int kMaxDepth = 24;
static uint64_t mix64(uint64_t v, uint64_t salt) {
uint64_t z = v + salt;
z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
return z ^ (z >> 31);
}
double clamp(double t) const {
return std::min(std::max(t, t_min_), t_max_);
}
sd::Tensor<float> w(double t) {
auto it = cache_.find(t);
if (it != cache_.end()) {
return it->second;
}
sd::Tensor<float> zero = sd::Tensor<float>::zeros(shape_);
sd::Tensor<float> out = bridge(t_min_, t_max_, zero, w_at_tmax_, t, root_seed_, kMaxDepth);
cache_.emplace(t, out);
return out;
}
sd::Tensor<float> bridge(double a,
double c,
const sd::Tensor<float>& w_a,
const sd::Tensor<float>& w_c,
double t,
uint64_t node_seed,
int depth) {
if (depth <= 0 || c - a < 1e-9) {
float alpha = (c > a) ? static_cast<float>((t - a) / (c - a)) : 0.5f;
return (1.0f - alpha) * w_a + alpha * w_c;
}
double m = 0.5 * (a + c);
double std_dev = std::sqrt((c - m) * (m - a) / (c - a));
auto rng = std::make_shared<STDDefaultRNG>();
rng->manual_seed(node_seed);
auto z = sd::Tensor<float>::randn(shape_, rng);
auto w_m = 0.5f * (w_a + w_c) + static_cast<float>(std_dev) * z;
if (t == m) {
return w_m;
}
if (t < m) {
return bridge(a, m, w_a, w_m, t, mix64(node_seed, 1), depth - 1);
}
return bridge(m, c, w_m, w_c, t, mix64(node_seed, 2), depth - 1);
}
double t_min_;
double t_max_;
std::vector<int64_t> shape_;
uint64_t root_seed_;
sd::Tensor<float> w_at_tmax_;
std::map<double, sd::Tensor<float>> cache_;
};
// DPM-Solver++(2M) SDE, midpoint variant, with step-count-stable Brownian-tree
// noise. Same trajectory shape at any step count for a given seed. Aliased in
// k-diffusion / ComfyUI as sample_dpmpp_2m_sde_gpu.
// Ref: Lu et al. arXiv:2211.01095; torchsde BrownianTree.
static sd::Tensor<float> sample_dpmpp_2m_sde_bt(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
std::shared_ptr<RNG> rng,
float eta) {
double sigma_max = 0.0;
double sigma_min = std::numeric_limits<double>::infinity();
for (float s : sigmas) {
if (s > 0.0f) {
sigma_max = std::max(sigma_max, static_cast<double>(s));
sigma_min = std::min(sigma_min, static_cast<double>(s));
}
}
if (sigma_max <= sigma_min) {
return x;
}
uint64_t tree_seed = 0;
{
auto draw = rng->randn(2);
std::memcpy(&tree_seed, draw.data(), sizeof(tree_seed));
}
BrownianTreeNoiseSampler noise_sampler(x, sigma_min, sigma_max, tree_seed);
sd::Tensor<float> old_denoised;
bool have_old_denoised = false;
float h_last = 0.f;
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) {
auto denoised_opt = model(x, sigmas[i], i + 1);
if (denoised_opt.pred.empty()) {
return {};
}
sd::Tensor<float> denoised = std::move(denoised_opt.pred);
if (sigmas[i + 1] == 0.f) {
x = denoised;
} else {
float t = -std::log(sigmas[i]);
float s = -std::log(sigmas[i + 1]);
float h = s - t;
float eta_h = eta * h;
float a = sigmas[i + 1] / sigmas[i] * std::exp(-eta_h);
float b = -std::expm1(-h - eta_h);
x = a * x + b * denoised;
if (have_old_denoised) {
float r = h_last / h;
x += (0.5f * b / r) * (denoised - old_denoised);
}
if (eta > 0.f) {
x += noise_sampler(sigmas[i], sigmas[i + 1]) * (sigmas[i + 1] * std::sqrt(-std::expm1(-2.f * eta_h)));
}
h_last = h;
}
old_denoised = denoised;
have_old_denoised = true;
}
return x;
}
using SamplerExtraArgs = KeyValueArgs;
static sd::Tensor<float> sample_lcm(denoise_cb_t model,
@ -2490,6 +2705,10 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
return sample_res_2s(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
case ER_SDE_SAMPLE_METHOD:
return sample_er_sde(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
case DPMPP2M_SDE_SAMPLE_METHOD:
return sample_dpmpp_2m_sde(model, std::move(x), sigmas, rng, eta);
case DPMPP2M_SDE_BT_SAMPLE_METHOD:
return sample_dpmpp_2m_sde_bt(model, std::move(x), sigmas, rng, eta);
case DDIM_TRAILING_SAMPLE_METHOD:
// DDIM is equivalent to Euler Ancestral with the Simple scheduler
return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);

View file

@ -287,6 +287,10 @@ bool IMatrixCollector::load_imatrix(const char* fname) {
if (e.values.empty()) {
e.values.resize(nval, 0);
e.counts.resize(nval, 0);
} else if (e.values.size() != (size_t)nval) {
LOG_ERROR("inconsistent size for a repeated entry (%d vs %d)\n", (int)e.values.size(), nval);
stats_ = {};
return false;
}
std::vector<float> tmp(nval);

View file

@ -4,11 +4,14 @@
#include <filesystem>
#include <limits>
#include <set>
#include <type_traits>
#include <unordered_set>
#include <utility>
#include <vector>
#include "core/ggml_extend.hpp"
#include "core/ggml_graph_cut.h"
#include "core/layer_split_partition.h"
#include "core/rng.hpp"
#include "core/rng_mt19937.hpp"
@ -19,6 +22,7 @@
#include "stable-diffusion.h"
#include "conditioning/conditioner.hpp"
#include "core/backend_fit.h"
#include "extensions/generation_extension.h"
#include "model/adapter/lora.hpp"
#include "model/diffusion/anima.hpp"
@ -55,11 +59,11 @@
#include "name_conversion.h"
#include "runtime/latent-preview.h"
#include <atomic>
const char* sd_vae_format_name(enum sd_vae_format_t format);
static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback);
#include <atomic>
const char* model_version_to_str[] = {
"SD 1.x",
"SD 1.x Inpaint",
@ -172,6 +176,13 @@ static float get_cache_reuse_threshold(const sd_cache_params_t& params) {
/*=============================================== StableDiffusionGGML ================================================*/
template <typename T, typename = void>
struct has_set_runtime_backends : std::false_type {};
template <typename T>
struct has_set_runtime_backends<T,
std::void_t<decltype(std::declval<T&>().set_runtime_backends(
std::declval<const std::vector<ggml_backend_t>&>()))>> : std::true_type {};
static_assert(std::atomic<sd_cancel_mode_t>::is_always_lock_free,
"sd_cancel_mode_t must be lock-free");
@ -212,6 +223,8 @@ public:
bool eager_load = false;
std::string backend_spec;
std::string params_backend_spec;
std::string split_mode_spec;
bool auto_fit_enabled = false;
bool is_using_v_parameterization = false;
bool is_using_edm_v_parameterization = false;
@ -259,6 +272,15 @@ public:
return max_vram_assignment.bytes_for_backend(backend_for(module));
}
std::vector<size_t> layer_split_vram_limits_for_backends(const std::vector<ggml_backend_t>& backends) {
std::vector<size_t> limits;
limits.reserve(backends.size());
for (ggml_backend_t backend : backends) {
limits.push_back(max_vram_assignment.bytes_for_backend(backend));
}
return limits;
}
bool ensure_backend_pair(SDBackendModule module) {
if (backend_for(module) == nullptr) {
return false;
@ -279,18 +301,208 @@ public:
if (model_manager == nullptr) {
return true;
}
ModelManager::ResidencyMode residency_mode =
backend_manager.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend;
std::vector<ggml_backend_t> module_backends = backend_manager.runtime_backends(module);
if (module_backends.size() > 1) {
if constexpr (has_set_runtime_backends<T>::value) {
if (module == SDBackendModule::DIFFUSION || module == SDBackendModule::TE) {
if (backend_manager.split_mode(module) == SDSplitMode::ROW) {
return register_row_split_runner_params(desc,
model,
module,
module_backends,
std::move(group_tensors),
residency_mode,
params_mem_size);
}
return register_layer_split_runner_params(desc,
model,
module,
module_backends,
std::move(group_tensors),
residency_mode,
params_mem_size);
}
}
LOG_WARN("%s module does not support multiple runtime backends; using %s",
sd_backend_module_name(module),
sd::layer_split_backend_device_display_name(module_backends[0]).c_str());
}
return model_manager->register_param_tensors(desc,
std::move(group_tensors),
backend_manager.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend,
residency_mode,
backend_for(module),
params_backend_for(module),
params_mem_size);
}
template <typename T>
bool register_row_split_runner_params(const std::string& desc,
const std::shared_ptr<T>& model,
SDBackendModule module,
const std::vector<ggml_backend_t>& module_backends,
std::map<std::string, ggml_tensor*> group_tensors,
ModelManager::ResidencyMode residency_mode,
size_t* params_mem_size) {
ggml_backend_t main_backend = module_backends[0];
auto fall_back_to_layer_split = [&](const char* reason) {
LOG_WARN("%s: row split unavailable (%s); falling back to layer split", desc.c_str(), reason);
return register_layer_split_runner_params(desc,
model,
module,
module_backends,
std::move(group_tensors),
residency_mode,
params_mem_size);
};
ggml_backend_dev_t main_dev = ggml_backend_get_device(main_backend);
ggml_backend_reg_t reg = main_dev != nullptr ? ggml_backend_dev_backend_reg(main_dev) : nullptr;
if (reg == nullptr) {
return fall_back_to_layer_split("no backend registry");
}
const size_t reg_dev_count = ggml_backend_reg_dev_count(reg);
std::vector<float> tensor_split(reg_dev_count, 0.0f);
constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024;
for (ggml_backend_t backend : module_backends) {
ggml_backend_dev_t dev = ggml_backend_get_device(backend);
int reg_index = -1;
for (size_t i = 0; i < reg_dev_count; i++) {
if (ggml_backend_reg_dev_get(reg, i) == dev) {
reg_index = (int)i;
break;
}
}
if (reg_index < 0) {
return fall_back_to_layer_split("devices span different backend registries");
}
size_t free_bytes = 0, total_bytes = 0;
ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
int64_t usable_bytes = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes,
(int64_t)free_bytes / 8);
tensor_split[reg_index] = usable_bytes > 0 ? (float)((double)usable_bytes / (1024.0 * 1024.0)) : 1.0f;
}
ggml_backend_buffer_type_t split_buft = backend_manager.split_buffer_type(main_backend, tensor_split);
if (split_buft == nullptr) {
return fall_back_to_layer_split("backend has no split buffer type");
}
model_manager->set_split_buffer_type(main_backend, split_buft);
std::map<std::string, ggml_tensor*> split_tensors;
if constexpr (std::is_base_of_v<Conditioner, T>) {
model->get_layer_split_param_tensors(split_tensors);
} else {
split_tensors = group_tensors;
}
std::map<std::string, ggml_tensor*> row_split_map;
std::map<std::string, ggml_tensor*> regular_map;
size_t row_split_bytes = 0;
for (const auto& kv : group_tensors) {
if (split_tensors.count(kv.first) != 0 &&
sd::layer_split_tensor_block_index(kv.first) >= 0 &&
ModelManager::tensor_shape_supports_split_buffer(kv.second)) {
row_split_map[kv.first] = kv.second;
row_split_bytes += ggml_nbytes(kv.second);
} else {
regular_map[kv.first] = kv.second;
}
}
if (row_split_map.empty()) {
return fall_back_to_layer_split("no row-splittable transformer block weights found");
}
LOG_INFO("%s row split: %zu tensors (%.1f MB) split across %zu devices (main %s)",
desc.c_str(),
row_split_map.size(),
row_split_bytes / (1024.f * 1024.f),
module_backends.size(),
sd::layer_split_backend_device_display_name(main_backend).c_str());
if (!model_manager->register_param_tensors(desc,
std::move(row_split_map),
residency_mode,
main_backend,
params_backend_for(module),
params_mem_size,
/*allow_split_buffer=*/true)) {
return false;
}
return model_manager->register_param_tensors(desc,
std::move(regular_map),
residency_mode,
main_backend,
params_backend_for(module),
params_mem_size);
}
// Register graph-cut layer-split tensors on the primary backend first.
// The first real graph assigns each param tensor to a runtime backend
// before weights are loaded or staged.
template <typename T>
bool register_layer_split_runner_params(const std::string& desc,
const std::shared_ptr<T>& model,
SDBackendModule module,
const std::vector<ggml_backend_t>& module_backends,
std::map<std::string, ggml_tensor*> group_tensors,
ModelManager::ResidencyMode residency_mode,
size_t* params_mem_size) {
bool has_cpu_device = false;
for (ggml_backend_t backend : module_backends) {
has_cpu_device = has_cpu_device || sd_backend_is_cpu(backend);
}
if (has_cpu_device) {
// The scheduler reserves the CPU slot for its fallback backend, and
// CPU weight participation is what --params-backend <module>=cpu is
// for; a CPU device in a split list is almost certainly a mistake.
LOG_WARN(
"%s: layer split across a CPU device is not supported; using %s "
"(use --params-backend %s=cpu to keep weights in RAM)",
desc.c_str(),
sd::layer_split_backend_device_display_name(module_backends[0]).c_str(),
sd_backend_module_name(module));
return model_manager->register_param_tensors(desc,
std::move(group_tensors),
residency_mode,
module_backends[0],
params_backend_for(module),
params_mem_size);
}
model->set_runtime_backends(module_backends);
model->set_graph_cut_layer_split_backend_vram_limits(layer_split_vram_limits_for_backends(module_backends));
model->set_graph_cut_layer_split_enabled(true);
const bool params_follow_runtime = backend_manager.params_backend_follows_runtime(module) ||
backend_manager.params_backend_is_disk(module);
ggml_backend_t initial_params_backend = params_follow_runtime ? module_backends[0] : params_backend_for(module);
if (initial_params_backend == nullptr) {
return false;
}
LOG_INFO("%s graph-cut layer split: deferring %zu tensors across %zu runtime backends until first graph",
desc.c_str(),
group_tensors.size(),
module_backends.size());
return model_manager->register_param_tensors(desc,
std::move(group_tensors),
residency_mode,
module_backends[0],
initial_params_backend,
params_mem_size,
false,
params_follow_runtime);
}
bool init_backend() {
std::string error;
if (!backend_manager.init(backend_spec.c_str(),
params_backend_spec.c_str(),
split_mode_spec.c_str(),
&error)) {
LOG_ERROR("backend config failed: %s", error.c_str());
return false;
@ -298,6 +510,26 @@ public:
return ensure_backend_pair(SDBackendModule::DIFFUSION);
}
bool row_split_active() {
for (SDBackendModule module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) {
if (backend_manager.split_mode(module) == SDSplitMode::ROW &&
backend_manager.runtime_backends(module).size() > 1) {
return true;
}
}
return false;
}
bool graph_cut_layer_split_active() {
for (SDBackendModule module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) {
if (backend_manager.split_mode(module) == SDSplitMode::LAYER &&
backend_manager.runtime_backends(module).size() > 1) {
return true;
}
}
return false;
}
std::shared_ptr<RNG> get_rng(rng_type_t rng_type) {
if (rng_type == STD_DEFAULT_RNG) {
return std::make_shared<STDDefaultRNG>();
@ -359,6 +591,8 @@ public:
eager_load = sd_ctx_params->eager_load;
backend_spec = SAFE_STR(sd_ctx_params->backend);
params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
split_mode_spec = SAFE_STR(sd_ctx_params->split_mode);
auto_fit_enabled = sd_ctx_params->auto_fit;
max_vram_assignment.reset(0.f);
{
std::string error;
@ -384,21 +618,6 @@ public:
ggml_log_set(ggml_log_callback_default, nullptr);
if (!init_backend()) {
return false;
}
{
std::string error;
if (!max_vram_assignment.canonicalize_backend_keys(&error)) {
LOG_ERROR("%s", error.c_str());
return false;
}
}
if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) {
LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring");
stream_layers = false;
}
std::string clip_vision_fixed = SAFE_STR(sd_ctx_params->clip_vision_path);
std::string clipg_path_fixed = SAFE_STR(sd_ctx_params->clip_g_path);
std::string clipl_path_fixed = SAFE_STR(sd_ctx_params->clip_l_path);
@ -788,6 +1007,35 @@ public:
model_loader.set_wtype_override(wtype, tensor_type_rules);
}
if (auto_fit_enabled) {
if (!sd::backend_fit::derive_backend_specs(model_loader,
wtype,
max_vram_assignment,
backend_spec,
params_backend_spec)) {
return false;
}
}
if (!init_backend()) {
return false;
}
{
std::string error;
if (!max_vram_assignment.canonicalize_backend_keys(&error)) {
LOG_ERROR("%s", error.c_str());
return false;
}
}
if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) {
LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring");
stream_layers = false;
}
if (eager_load && graph_cut_layer_split_active()) {
LOG_WARN("--eager-load is not supported with graph-cut layer split; weights will be prepared lazily");
eager_load = false;
}
std::map<ggml_type, uint32_t> wtype_stat = model_loader.get_wtype_stat();
std::map<ggml_type, uint32_t> conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat();
std::map<ggml_type, uint32_t> diffusion_model_wtype_stat = model_loader.get_diffusion_model_wtype_stat();
@ -829,12 +1077,17 @@ public:
// Avoid full-model LoRA merge buffers on constrained setups.
const bool params_offloaded = params_backend_for(SDBackendModule::DIFFUSION) != backend_for(SDBackendModule::DIFFUSION);
const bool streaming_constrained = stream_layers || params_offloaded;
if (have_quantized_weight || streaming_constrained) {
if (have_quantized_weight || streaming_constrained || row_split_active()) {
apply_lora_immediately = false;
} else {
apply_lora_immediately = true;
}
} else if (sd_ctx_params->lora_apply_mode == LORA_APPLY_IMMEDIATELY) {
if (row_split_active()) {
LOG_WARN(
"row-split tensors do not support the immediately LoRA apply mode; "
"LoRAs will not be applied to them (use --lora-apply-mode at_runtime)");
}
apply_lora_immediately = true;
} else {
apply_lora_immediately = false;
@ -860,10 +1113,6 @@ public:
use_tae = true;
}
if (sd_ctx_params->circular_x || sd_ctx_params->circular_y) {
LOG_INFO("Using circular padding for convolutions");
}
{
if (!ensure_backend_pair(SDBackendModule::TE) ||
!ensure_backend_pair(SDBackendModule::DIFFUSION)) {
@ -922,10 +1171,11 @@ public:
if (is_chroma) {
cond_stage_model = std::make_shared<T5CLIPEmbedder>(backend_for(SDBackendModule::TE),
tensor_storage_map,
sd_ctx_params->chroma_use_t5_mask,
sd_ctx_params->chroma_t5_mask_pad,
false,
model_manager);
1,
false,
model_manager,
sd_ctx_params->model_args);
} else if (version == VERSION_OVIS_IMAGE) {
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
tensor_storage_map,
@ -942,8 +1192,8 @@ public:
tensor_storage_map,
"model.diffusion_model",
version,
sd_ctx_params->chroma_use_dit_mask,
model_manager);
model_manager,
sd_ctx_params->model_args);
} else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
bool is_chroma = false;
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
@ -956,8 +1206,8 @@ public:
tensor_storage_map,
"model.diffusion_model",
version,
sd_ctx_params->chroma_use_dit_mask,
model_manager);
model_manager,
sd_ctx_params->model_args);
} else if (sd_version_is_ltxav(version)) {
cond_stage_model = std::make_shared<LTXAVEmbedder>(backend_for(SDBackendModule::TE),
tensor_storage_map,
@ -1015,8 +1265,8 @@ public:
tensor_storage_map,
"model.diffusion_model",
version,
sd_ctx_params->qwen_image_zero_cond_t,
model_manager);
model_manager,
sd_ctx_params->model_args);
} else if (sd_version_is_longcat(version)) {
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
tensor_storage_map,
@ -1028,8 +1278,8 @@ public:
tensor_storage_map,
"model.diffusion_model",
version,
sd_ctx_params->chroma_use_dit_mask,
model_manager);
model_manager,
sd_ctx_params->model_args);
} else if (version == VERSION_HIDREAM_O1) {
cond_stage_model = std::make_shared<HiDreamO1::HiDreamO1Conditioner>(backend_for(SDBackendModule::TE),
tensor_storage_map,
@ -1368,16 +1618,6 @@ public:
high_noise_diffusion_model->set_flash_attention_enabled(true);
}
}
diffusion_model->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y);
if (high_noise_diffusion_model) {
high_noise_diffusion_model->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y);
}
if (control_net) {
control_net->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y);
}
circular_x = sd_ctx_params->circular_x;
circular_y = sd_ctx_params->circular_y;
}
LOG_DEBUG("validating model metadata");
@ -2718,7 +2958,16 @@ public:
}
auto latents = first_stage_model->diffusion_to_vae_latents(x);
first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling);
return first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
if (decoded.empty() && auto_fit_enabled) {
bool prefer_temporal_tiling = decode_video && std::dynamic_pointer_cast<LTXVideoVAE>(first_stage_model) != nullptr;
if (sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) {
first_stage_model->free_compute_buffer();
first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling);
decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
}
}
return decoded;
}
sd::Tensor<float> normalize_ltx_video_latents(const sd::Tensor<float>& x) {
@ -2762,22 +3011,6 @@ public:
return !!flow_denoiser;
}
//added for kcpp
void SetCircularAxesAll(bool circular_x, bool circular_y)
{
diffusion_model->set_circular_axes(circular_x, circular_y);
if (high_noise_diffusion_model) {
high_noise_diffusion_model->set_circular_axes(circular_x, circular_y);
}
if (control_net) {
control_net->set_circular_axes(circular_x, circular_y);
}
if (first_stage_model) {
first_stage_model->set_circular_axes(circular_x, circular_y);
}
}
//end added for kcpp
};
/*================================================= SD API ==================================================*/
@ -2842,6 +3075,8 @@ const char* sample_method_to_str[] = {
"euler_cfg_pp",
"euler_a_cfg_pp",
"euler_ge",
"dpm++2m_sde",
"dpm++2m_sde_bt",
};
const char* sd_sample_method_name(enum sample_method_t sample_method) {
@ -3083,15 +3318,13 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
sd_ctx_params->eager_load = false;
sd_ctx_params->enable_mmap = false;
sd_ctx_params->diffusion_flash_attn = false;
sd_ctx_params->circular_x = false;
sd_ctx_params->circular_y = false;
sd_ctx_params->chroma_use_dit_mask = true;
sd_ctx_params->chroma_use_t5_mask = false;
sd_ctx_params->chroma_t5_mask_pad = 1;
sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO;
sd_ctx_params->backend = nullptr;
sd_ctx_params->params_backend = nullptr;
sd_ctx_params->split_mode = nullptr;
sd_ctx_params->auto_fit = false;
sd_ctx_params->rpc_servers = nullptr;
sd_ctx_params->model_args = nullptr;
sd_ctx_params->pulid_weights_path = nullptr;
}
@ -3130,13 +3363,11 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
"eager_load: %s\n"
"backend: %s\n"
"params_backend: %s\n"
"split_mode: %s\n"
"model_args: %s\n"
"auto_fit: %s\n"
"flash_attn: %s\n"
"diffusion_flash_attn: %s\n"
"circular_x: %s\n"
"circular_y: %s\n"
"chroma_use_dit_mask: %s\n"
"chroma_use_t5_mask: %s\n"
"chroma_t5_mask_pad: %d\n"
"vae_format: %s\n",
SAFE_STR(sd_ctx_params->model_path),
SAFE_STR(sd_ctx_params->clip_l_path),
@ -3166,13 +3397,11 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
BOOL_STR(sd_ctx_params->eager_load),
SAFE_STR(sd_ctx_params->backend),
SAFE_STR(sd_ctx_params->params_backend),
SAFE_STR(sd_ctx_params->split_mode),
SAFE_STR(sd_ctx_params->model_args),
BOOL_STR(sd_ctx_params->auto_fit),
BOOL_STR(sd_ctx_params->flash_attn),
BOOL_STR(sd_ctx_params->diffusion_flash_attn),
BOOL_STR(sd_ctx_params->circular_x),
BOOL_STR(sd_ctx_params->circular_y),
BOOL_STR(sd_ctx_params->chroma_use_dit_mask),
BOOL_STR(sd_ctx_params->chroma_use_t5_mask),
sd_ctx_params->chroma_t5_mask_pad,
sd_vae_format_name(sd_ctx_params->vae_format));
return buf;
@ -3250,6 +3479,8 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
sd_img_gen_params->batch_count = 1;
sd_img_gen_params->control_strength = 0.9f;
sd_img_gen_params->qwen_image_layers = 3;
sd_img_gen_params->circular_x = false;
sd_img_gen_params->circular_y = false;
sd_img_gen_params->pm_params = {nullptr, 0, nullptr, 20.f};
sd_img_gen_params->pulid_params = {nullptr, 1.0f};
sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
@ -3283,6 +3514,8 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
"control_strength: %.2f\n"
"photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n"
"VAE tiling: %s (temporal=%s, extra_tiling_args=%s)\n"
"circular_x: %s\n"
"circular_y: %s\n"
"hires: {enabled=%s, upscaler=%s, model_path=%s, scale=%.2f, target=%dx%d, steps=%d, denoising_strength=%.2f}\n",
SAFE_STR(sd_img_gen_params->prompt),
SAFE_STR(sd_img_gen_params->negative_prompt),
@ -3304,6 +3537,8 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
BOOL_STR(sd_img_gen_params->vae_tiling_params.enabled),
BOOL_STR(sd_img_gen_params->vae_tiling_params.temporal_tiling),
SAFE_STR(sd_img_gen_params->vae_tiling_params.extra_tiling_args),
BOOL_STR(sd_img_gen_params->circular_x),
BOOL_STR(sd_img_gen_params->circular_y),
BOOL_STR(sd_img_gen_params->hires.enabled),
sd_hires_upscaler_name(sd_img_gen_params->hires.upscaler),
SAFE_STR(sd_img_gen_params->hires.model_path),
@ -3353,6 +3588,8 @@ void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params) {
sd_vid_gen_params->hires.upscale_tile_size = 128;
sd_vid_gen_params->hires.custom_sigmas = nullptr;
sd_vid_gen_params->hires.custom_sigmas_count = 0;
sd_vid_gen_params->circular_x = false;
sd_vid_gen_params->circular_y = false;
sd_cache_params_init(&sd_vid_gen_params->cache);
}
@ -3609,6 +3846,8 @@ static float resolve_eta(sd_ctx_t* sd_ctx,
case DPMPP2S_A_SAMPLE_METHOD:
case ER_SDE_SAMPLE_METHOD:
case EULER_A_CFG_PP_SAMPLE_METHOD:
case DPMPP2M_SDE_SAMPLE_METHOD:
case DPMPP2M_SDE_BT_SAMPLE_METHOD:
return 1.0f;
default:;
}
@ -4298,6 +4537,25 @@ struct CircularAxesState {
bool circular_y = false;
};
static void apply_circular_axes_to_diffusion(sd_ctx_t* sd_ctx, bool circular_x, bool circular_y) {
sd_ctx->sd->circular_x = circular_x;
sd_ctx->sd->circular_y = circular_y;
if (sd_ctx->sd->diffusion_model) {
sd_ctx->sd->diffusion_model->set_circular_axes(circular_x, circular_y);
}
if (sd_ctx->sd->high_noise_diffusion_model) {
sd_ctx->sd->high_noise_diffusion_model->set_circular_axes(circular_x, circular_y);
}
if (sd_ctx->sd->control_net) {
sd_ctx->sd->control_net->set_circular_axes(circular_x, circular_y);
}
if (circular_x || circular_y) {
LOG_INFO("Using circular padding for convolutions (x=%s, y=%s)",
circular_x ? "true" : "false",
circular_y ? "true" : "false");
}
}
static CircularAxesState configure_image_vae_axes(sd_ctx_t* sd_ctx,
const sd_img_gen_params_t* sd_img_gen_params,
const GenerationRequest& request) {
@ -4408,13 +4666,56 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
LOG_INFO("IMG2IMG");
if (request->strength < 1.f) {
size_t t_enc = static_cast<size_t>(plan->sample_steps * request->strength);
if (t_enc == static_cast<size_t>(plan->sample_steps)) {
t_enc--;
bool strength_as_noise_level = false;
bool force_first_sigma = false;
for (const auto& [key, value] : parse_key_value_args(sd_img_gen_params->sample_params.extra_sample_args, "img2img arg")) {
if (key == "strength_as_noise_level") {
if (!parse_strict_bool(value, strength_as_noise_level)) {
LOG_WARN("ignoring invalid img2img sample arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "force_first_sigma") {
if (!parse_strict_bool(value, force_first_sigma)) {
LOG_WARN("ignoring invalid img2img sample arg '%s=%s'", key.c_str(), value.c_str());
}
}
}
size_t t_enc;
float target_sigma = -1;
if (!strength_as_noise_level) {
t_enc = static_cast<size_t>(plan->sample_steps * request->strength);
if (t_enc == static_cast<size_t>(plan->sample_steps)) {
t_enc--;
}
} else {
LOG_DEBUG("Interpreting denoise strength as relative noise level");
// assume x_noised = K * (x * (1-noise_level) + noise * noise_level) = K * lerp(x, noise, noise_level)
// K = 1, noise_level = sigma for flow models
// K = 1+sigma, noise_level=sigma/(1+sigma) for diffusion models
float target_noise_level = request->strength;
target_sigma = sd_ctx->sd->denoiser->noise_level_to_sigma(target_noise_level);
size_t start_index = 0;
for (size_t i = 0; i < plan->sigmas.size(); ++i) {
if (plan->sigmas[i] <= target_sigma) {
start_index = i;
break;
}
}
if (start_index >= plan->sigmas.size() - 1) {
start_index = plan->sigmas.size() - 2; // Leave at least 1 step
}
t_enc = plan->sample_steps - start_index - 1;
}
LOG_INFO("target t_enc is %zu steps", t_enc);
std::vector<float> sigma_sched;
sigma_sched.assign(plan->sigmas.begin() + plan->sample_steps - t_enc - 1, plan->sigmas.end());
if (target_sigma > 0 && force_first_sigma && strength_as_noise_level) {
LOG_DEBUG("force_first_sigma to %.4f (from %.4f)", target_sigma, sigma_sched[0]);
sigma_sched[0] = target_sigma;
}
plan->sigmas = std::move(sigma_sched);
plan->sample_steps = static_cast<int>(plan->sigmas.size() - 1);
}
@ -5011,6 +5312,7 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx,
sd_ctx->sd->sampler_rng->manual_seed(request.seed);
sd_ctx->sd->set_flow_shift(sd_img_gen_params->sample_params.flow_shift);
sd_ctx->sd->apply_loras(sd_img_gen_params->loras, sd_img_gen_params->lora_count);
apply_circular_axes_to_diffusion(sd_ctx, sd_img_gen_params->circular_x, sd_img_gen_params->circular_y);
ImageVaeAxesGuard axes_guard(sd_ctx, sd_img_gen_params, request);
@ -5900,6 +6202,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
}
int64_t t0 = ggml_time_ms();
sd_ctx->sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params;
apply_circular_axes_to_diffusion(sd_ctx, sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y);
GenerationRequest request(sd_ctx, sd_vid_gen_params);
bool has_input_audio = sd_vid_gen_params->input_audio != nullptr &&
sd_vid_gen_params->input_audio->data != nullptr &&
@ -6327,10 +6630,6 @@ namespace kcpp_sd {
return res;
}
void SetCircularAxesAll(sd_ctx_t* ctx, bool circular_x, bool circular_y) {
ctx->sd->SetCircularAxesAll(circular_x, circular_y);
}
void set_lora_cache(sd_ctx_t *ctx, bool enable) {
ctx->sd->kcpp_lora_cache_populate = enable;
}

View file

@ -46,6 +46,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
std::string error;
if (!backend_manager.init(backend_spec.c_str(),
params_backend_spec.c_str(),
/*split_mode_spec=*/nullptr,
&error)) {
LOG_ERROR("upscaler backend config failed: %s", error.c_str());
return false;

View file

@ -3,10 +3,14 @@
#include <vector>
#include "ggml-backend.h"
struct ggml_tensor;
struct RunnerWeightManager {
virtual ~RunnerWeightManager() = default;
virtual bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
ggml_backend_t compute_backend) = 0;
virtual bool prepare_params(const std::vector<ggml_tensor*>& tensors) = 0;
virtual void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
virtual void release_params_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;