diff --git a/Makefile b/Makefile index ae2351e5b..717c81de9 100644 --- a/Makefile +++ b/Makefile @@ -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 diff --git a/expose.h b/expose.h index f753a2ec2..202427586 100644 --- a/expose.h +++ b/expose.h @@ -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; diff --git a/koboldcpp.py b/koboldcpp.py index b6ba4a3a7..bfc16ee10 100644 --- a/koboldcpp.py +++ b/koboldcpp.py @@ -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
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 diff --git a/otherarch/sdcpp/examples/cli/main.cpp b/otherarch/sdcpp/examples/cli/main.cpp index 7892d5213..29cb391b6 100644 --- a/otherarch/sdcpp/examples/cli/main.cpp +++ b/otherarch/sdcpp/examples/cli/main.cpp @@ -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(), diff --git a/otherarch/sdcpp/examples/common/common.cpp b/otherarch/sdcpp/examples/common/common.cpp index ac1cf32f6..1dfb6aa70 100644 --- a/otherarch/sdcpp/examples/common/common.cpp +++ b/otherarch/sdcpp/examples/common/common.cpp @@ -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)',', ¶ms_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 'namedescription' 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 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(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(hires_custom_sigmas.size()); + params.circular_x = circular || circular_x; + params.circular_y = circular || circular_y; return params; } diff --git a/otherarch/sdcpp/examples/common/common.h b/otherarch/sdcpp/examples/common/common.h index 941fa3317..3a5b107bc 100644 --- a/otherarch/sdcpp/examples/common/common.h +++ b/otherarch/sdcpp/examples/common/common.h @@ -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; diff --git a/otherarch/sdcpp/include/stable-diffusion.h b/otherarch/sdcpp/include/stable-diffusion.h index d5bace723..3fc6f64b7 100644 --- a/otherarch/sdcpp/include/stable-diffusion.h +++ b/otherarch/sdcpp/include/stable-diffusion.h @@ -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 `namedescription` 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); diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp index 00c0a6dd7..4757a39b5 100644 --- a/otherarch/sdcpp/sdtype_adapter.cpp +++ b/otherarch/sdcpp/sdtype_adapter.cpp @@ -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(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) { diff --git a/otherarch/sdcpp/src/conditioning/conditioner.hpp b/otherarch/sdcpp/src/conditioning/conditioner.hpp index d63303a82..5d4ad7fbb 100644 --- a/otherarch/sdcpp/src/conditioning/conditioner.hpp +++ b/otherarch/sdcpp/src/conditioning/conditioner.hpp @@ -6,6 +6,7 @@ #include #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& 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& backends) {} + virtual void set_graph_cut_layer_split_enabled(bool enabled) {} + virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) {} + virtual void get_layer_split_param_tensors(std::map& tensors) {} virtual void set_flash_attention_enabled(bool enabled) = 0; virtual void set_weight_adapter(const std::shared_ptr& adapter) {} virtual void runner_done() {} @@ -178,6 +183,27 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { } } + void set_runtime_backends(const std::vector& 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& 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& 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& 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& 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& 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& 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& 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 weight_manager = nullptr) + std::shared_ptr 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& 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& limits) override { + if (t5) { + t5->set_graph_cut_layer_split_backend_vram_limits(limits); + } + } + + void get_layer_split_param_tensors(std::map& 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& 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& limits) override { + if (t5) { + t5->set_graph_cut_layer_split_backend_vram_limits(limits); + } + } + + void get_layer_split_param_tensors(std::map& 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& 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& limits) override { + llm->set_graph_cut_layer_split_backend_vram_limits(limits); + } + + void get_layer_split_param_tensors(std::map& 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& 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& limits) override { + if (llm) { + llm->set_graph_cut_layer_split_backend_vram_limits(limits); + } + } + + void get_layer_split_param_tensors(std::map& 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& 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& limits) override { + llm->set_graph_cut_layer_split_backend_vram_limits(limits); + } + + void get_layer_split_param_tensors(std::map& tensors) override { + llm->get_param_tensors(tensors, "text_encoders.llm"); + } + void set_weight_adapter(const std::shared_ptr& adapter) override { llm->set_weight_adapter(adapter); projector->set_weight_adapter(adapter); diff --git a/otherarch/sdcpp/src/convert.cpp b/otherarch/sdcpp/src/convert.cpp index 0b7fe2cfb..8e94a940c 100644 --- a/otherarch/sdcpp/src/convert.cpp +++ b/otherarch/sdcpp/src/convert.cpp @@ -1,14 +1,33 @@ +#include +#include +#include #include +#include +#include +#include #include #include +#include +#include #include +#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 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& 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& 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 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(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 tensor_write_plans_from_export_infos(const std::vector& tensors) { + std::vector 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(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(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& 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 workers; + workers.reserve(n_threads); + + auto reserve_memory = [&](size_t bytes) -> bool { + std::unique_lock 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 lock(work_mutex); + reserved_bytes -= std::min(reserved_bytes, bytes); + } + memory_cv.notify_all(); + }; + + auto fail = [&](const std::string& message) { + { + std::lock_guard 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 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 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(tensors_written), + static_cast(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& tensors, + StreamingModelWriter& writer, + int n_threads, + std::string* error) { + std::vector 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 tensors; - bool success = load_tensors_for_export(model_loader, ggml_ctx, type, type_rules, tensors); - ggml_backend_free(backend); - + std::vector tensors; + bool success = collect_tensors_for_export(model_loader, type, type_rules, tensors); std::string error; if (success) { + std::unique_ptr writer; if (output_is_safetensors) { - success = write_safetensors_file(output_path, tensors, &error); + writer = std::make_unique(); } else { - success = write_gguf_file(output_path, tensors, &error); + writer = std::make_unique(); } + 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); } diff --git a/otherarch/sdcpp/src/core/backend_fit.cpp b/otherarch/sdcpp/src/core/backend_fit.cpp new file mode 100644 index 000000000..460194740 --- /dev/null +++ b/otherarch/sdcpp/src/core/backend_fit.cpp @@ -0,0 +1,390 @@ +#include "backend_fit.h" + +#include +#include +#include +#include +#include + +#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 device_idxs; + }; + + struct Plan { + bool valid = false; + bool time_share = false; + std::vector 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 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 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 enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) { + std::vector out; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) { + continue; + } + Device d; + d.dev = dev; + d.name = ggml_backend_dev_name(dev); + d.description = ggml_backend_dev_description(dev); + size_t free_bytes = 0, total_bytes = 0; + ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); + d.free_bytes = (int64_t)free_bytes; + d.total_bytes = (int64_t)total_bytes; + + 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)(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(d.budget_bytes, 0); + out.push_back(d); + } + return out; + } + + Plan compute_plan(const std::vector& components, const std::vector& devices) { + Plan plan; + if (devices.empty()) { + return plan; + } + + std::vector 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 params_sum(devices.size(), 0); + std::vector max_reserve(devices.size(), 0); + std::vector 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(d.budget_bytes - comp.reserve_bytes, 0); + } + if (comp.params_bytes <= capacity) { + std::vector idxs(devices.size()); + for (size_t i = 0; i < idxs.size(); i++) { + idxs[i] = i; + } + std::sort(idxs.begin(), idxs.end(), [&](size_t a, size_t b) { + return devices[a].budget_bytes > devices[b].budget_bytes; + }); + 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& components, + const std::vector& 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& components, + const std::vector& 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 diff --git a/otherarch/sdcpp/src/core/backend_fit.h b/otherarch/sdcpp/src/core/backend_fit.h new file mode 100644 index 000000000..9ef298b3b --- /dev/null +++ b/otherarch/sdcpp/src/core/backend_fit.h @@ -0,0 +1,23 @@ +#ifndef __SD_BACKEND_FIT_H__ +#define __SD_BACKEND_FIT_H__ + +#include + +#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__ diff --git a/otherarch/sdcpp/src/core/ggml_extend.hpp b/otherarch/sdcpp/src/core/ggml_extend.hpp index 48a7ad740..3adae6a57 100644 --- a/otherarch/sdcpp/src/core/ggml_extend.hpp +++ b/otherarch/sdcpp/src/core/ggml_extend.hpp @@ -21,10 +21,12 @@ #include #include #include +#include #include #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 graph_cut_layer_split_backend_vram_limits_; + + std::vector 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 weight_adapter = nullptr; std::weak_ptr weight_manager; @@ -1771,6 +1780,9 @@ protected: sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_; std::unordered_set params_tensor_set_; + std::unordered_map graph_cut_layer_split_assignments_; + std::unordered_map graph_cut_layer_split_node_assignments_; + bool graph_cut_layer_split_primary_notice_logged_ = false; template static sd::Tensor take_or_empty(std::optional> 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 collect_used_param_tensors(ggml_cgraph* gf) { std::vector used_params; rebuild_params_tensor_set(); @@ -1881,12 +1907,8 @@ protected: seen_params.reserve(static_cast(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 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 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(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 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 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& 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& 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 { diff --git a/otherarch/sdcpp/src/core/ggml_extend_backend.cpp b/otherarch/sdcpp/src/core/ggml_extend_backend.cpp index f29bdb696..c66bb63f0 100644 --- a/otherarch/sdcpp/src/core/ggml_extend_backend.cpp +++ b/otherarch/sdcpp/src/core/ggml_extend_backend.cpp @@ -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 split_device_list(const std::string& value) { + std::vector 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 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 SDBackendManager::runtime_backends(SDBackendModule module) { + std::vector 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), ¶ms_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& 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 padded_split(std::max(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 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; } diff --git a/otherarch/sdcpp/src/core/ggml_extend_backend.h b/otherarch/sdcpp/src/core/ggml_extend_backend.h index 19b71d432..1f3bc8b3a 100644 --- a/otherarch/sdcpp/src/core/ggml_extend_backend.h +++ b/otherarch/sdcpp/src/core/ggml_extend_backend.h @@ -6,6 +6,7 @@ #include #include #include +#include #include "ggml-backend.h" #include "ggml.h" @@ -37,10 +38,16 @@ struct SDBackendHandleDeleter { using SDBackendHandle = std::unique_ptr; +enum class SDSplitMode { + LAYER, + ROW, +}; + class SDBackendManager { private: SDBackendAssignment runtime_assignment_; SDBackendAssignment params_assignment_; + SDBackendAssignment split_mode_assignment_; std::unordered_map 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 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& 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: diff --git a/otherarch/sdcpp/src/core/layer_split_partition.cpp b/otherarch/sdcpp/src/core/layer_split_partition.cpp new file mode 100644 index 000000000..654b3056b --- /dev/null +++ b/otherarch/sdcpp/src/core/layer_split_partition.cpp @@ -0,0 +1,257 @@ +#include "core/layer_split_partition.h" + +#include +#include +#include +#include +#include +#include + +#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& 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 graph_cut_layer_split_backend_capacities(const std::vector& backends, + const std::vector& backend_vram_limits, + size_t primary_backend_vram_limit) { + std::vector capacities(backends.size(), std::numeric_limits::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)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(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& split_backends, + const std::vector& backend_vram_limits, + size_t primary_backend_vram_limit, + std::unordered_map& param_assignments, + const std::function& 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> segment_params(plan.segments.size()); + std::vector segment_param_bytes(plan.segments.size(), 0); + std::unordered_set seen_params; + for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) { + std::vector 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 backend_capacities = graph_cut_layer_split_backend_capacities(split_backends, + backend_vram_limits, + primary_backend_vram_limit); + + std::vector 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& 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 diff --git a/otherarch/sdcpp/src/core/layer_split_partition.h b/otherarch/sdcpp/src/core/layer_split_partition.h new file mode 100644 index 000000000..9450b379e --- /dev/null +++ b/otherarch/sdcpp/src/core/layer_split_partition.h @@ -0,0 +1,44 @@ +#ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__ +#define __SD_CORE_LAYER_SPLIT_PARTITION_H__ + +#include +#include +#include +#include +#include + +#include "ggml-backend.h" +#include "ggml.h" + +#include "core/ggml_graph_cut.h" + +namespace sd { + + struct GraphCutLayerSplitAssignment { + std::vector> tensors_by_backend; + std::vector bytes_by_backend; + std::vector first_segment_by_backend; + std::vector last_segment_by_backend; + std::unordered_map 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& split_backends, + const std::vector& backend_vram_limits, + size_t primary_backend_vram_limit, + std::unordered_map& param_assignments, + const std::function& canonical_param_tensor, + GraphCutLayerSplitAssignment* assignment_out); + void log_graph_cut_layer_split_assignment(const char* desc, + const std::vector& split_backends, + const GraphCutLayerSplitAssignment& assignment); + +} // namespace sd + +#endif // __SD_CORE_LAYER_SPLIT_PARTITION_H__ diff --git a/otherarch/sdcpp/src/core/util.cpp b/otherarch/sdcpp/src/core/util.cpp index cf1460cb9..4672b139f 100644 --- a/otherarch/sdcpp/src/core/util.cpp +++ b/otherarch/sdcpp/src/core/util.cpp @@ -4,6 +4,8 @@ #include #include #include +#include +#include #include #include #include @@ -29,6 +31,7 @@ #include #endif +#include "ggml-backend.h" #include "ggml.h" #include "stable-diffusion.h" @@ -1042,3 +1045,26 @@ std::vector> 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(); +} diff --git a/otherarch/sdcpp/src/kcpp_sd_extensions.h b/otherarch/sdcpp/src/kcpp_sd_extensions.h index 0e0c6795a..516d4273b 100644 --- a/otherarch/sdcpp/src/kcpp_sd_extensions.h +++ b/otherarch/sdcpp/src/kcpp_sd_extensions.h @@ -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& lora_specs); diff --git a/otherarch/sdcpp/src/model/diffusion/control.hpp b/otherarch/sdcpp/src/model/diffusion/control.hpp index eeb8f5109..bf3c7e435 100644 --- a/otherarch/sdcpp/src/model/diffusion/control.hpp +++ b/otherarch/sdcpp/src/model/diffusion/control.hpp @@ -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 =================================== diff --git a/otherarch/sdcpp/src/model/diffusion/flux.hpp b/otherarch/sdcpp/src/model/diffusion/flux.hpp index e3cd6ba7c..e5d795f74 100644 --- a/otherarch/sdcpp/src/model/diffusion/flux.hpp +++ b/otherarch/sdcpp/src/model/diffusion/flux.hpp @@ -4,6 +4,7 @@ #include #include +#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 dct_vec; sd::Tensor 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 weight_manager = nullptr) + std::shared_ptr 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", diff --git a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp index 5dabbb176..52edc0d56 100644 --- a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp @@ -3,6 +3,7 @@ #include +#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 weight_manager = nullptr) + std::shared_ptr 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", diff --git a/otherarch/sdcpp/src/model/vae/vae.hpp b/otherarch/sdcpp/src/model/vae/vae.hpp index 8b8c46ded..b97aecee0 100644 --- a/otherarch/sdcpp/src/model/vae/vae.hpp +++ b/otherarch/sdcpp/src/model/vae/vae.hpp @@ -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, diff --git a/otherarch/sdcpp/src/model_io/gguf_io.cpp b/otherarch/sdcpp/src/model_io/gguf_io.cpp index c701d01f3..cd22312d5 100644 --- a/otherarch/sdcpp/src/model_io/gguf_io.cpp +++ b/otherarch/sdcpp/src/model_io/gguf_io.cpp @@ -1,7 +1,10 @@ #include "gguf_io.h" +#include #include +#include #include +#include #include #include @@ -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& 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(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(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(data), static_cast(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; + } +} diff --git a/otherarch/sdcpp/src/model_io/gguf_io.h b/otherarch/sdcpp/src/model_io/gguf_io.h index 81c981145..3a4ae8200 100644 --- a/otherarch/sdcpp/src/model_io/gguf_io.h +++ b/otherarch/sdcpp/src/model_io/gguf_io.h @@ -4,8 +4,12 @@ #include #include +#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& tensor_storages, @@ -14,4 +18,28 @@ bool write_gguf_file(const std::string& file_path, const std::vector& tensors, std::string* error = nullptr); +class GGUFStreamingWriter : public StreamingModelWriter { +public: + GGUFStreamingWriter() = default; + ~GGUFStreamingWriter(); + + bool write_metadata(const std::string& file_path, + const std::vector& 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 tensors_; + std::vector tensor_offsets_; + uint64_t file_size_ = 0; + ggml_context* meta_ctx_ = nullptr; + gguf_context* gguf_ctx_ = nullptr; +}; + #endif // __SD_MODEL_IO_GGUF_IO_H__ diff --git a/otherarch/sdcpp/src/model_io/safetensors_io.cpp b/otherarch/sdcpp/src/model_io/safetensors_io.cpp index 39131dbd8..197dc5f54 100644 --- a/otherarch/sdcpp/src/model_io/safetensors_io.cpp +++ b/otherarch/sdcpp/src/model_io/safetensors_io.cpp @@ -1,8 +1,10 @@ #include "safetensors_io.h" +#include #include #include #include +#include #include #include @@ -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 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 end = tensor_info["data_offsets"][1].get(); + 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& 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(ST_HEADER_SIZE_LEN); ++i) { + header_size[i] = static_cast((header_str.size() >> (8 * i)) & 0xFF); + } + file.write(reinterpret_cast(header_size), sizeof(header_size)); + file.write(header_str.data(), static_cast(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(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(data), static_cast(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_; +} diff --git a/otherarch/sdcpp/src/model_io/safetensors_io.h b/otherarch/sdcpp/src/model_io/safetensors_io.h index 08a1bc1f3..b4938ee18 100644 --- a/otherarch/sdcpp/src/model_io/safetensors_io.h +++ b/otherarch/sdcpp/src/model_io/safetensors_io.h @@ -4,6 +4,7 @@ #include #include +#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& tensors, std::string* error = nullptr); +class SafetensorsStreamingWriter : public StreamingModelWriter { +public: + SafetensorsStreamingWriter() = default; + + bool write_metadata(const std::string& file_path, + const std::vector& 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 tensors_; + std::vector tensor_offsets_; + uint64_t data_start_ = 0; + uint64_t file_size_ = 0; +}; + #endif // __SD_MODEL_IO_SAFETENSORS_IO_H__ diff --git a/otherarch/sdcpp/src/model_io/streaming_writer.h b/otherarch/sdcpp/src/model_io/streaming_writer.h new file mode 100644 index 000000000..06f65c59b --- /dev/null +++ b/otherarch/sdcpp/src/model_io/streaming_writer.h @@ -0,0 +1,26 @@ +#ifndef __SD_MODEL_IO_STREAMING_WRITER_H__ +#define __SD_MODEL_IO_STREAMING_WRITER_H__ + +#include +#include +#include +#include + +#include "tensor_storage.h" + +class StreamingModelWriter { +public: + virtual ~StreamingModelWriter() = default; + + virtual bool write_metadata(const std::string& file_path, + const std::vector& 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__ diff --git a/otherarch/sdcpp/src/model_io/tensor_storage.h b/otherarch/sdcpp/src/model_io/tensor_storage.h index c0cf079c5..307535a53 100644 --- a/otherarch/sdcpp/src/model_io/tensor_storage.h +++ b/otherarch/sdcpp/src/model_io/tensor_storage.h @@ -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 on_new_tensor_cb_t; #endif // __SD_TENSOR_STORAGE_H__ diff --git a/otherarch/sdcpp/src/model_loader.cpp b/otherarch/sdcpp/src/model_loader.cpp index fda1f3e82..6aa9afa8d 100644 --- a/otherarch/sdcpp/src/model_loader.cpp +++ b/otherarch/sdcpp/src/model_loader.cpp @@ -971,7 +971,8 @@ std::vector ModelLoader::mmap_tensors(std::map* target_tensor_names) { + const std::set* target_tensor_names, + bool log_progress) { process_model_files(enable_mmap, false); std::atomic 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(curr_num), static_cast(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(total_tensors_processed), static_cast(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 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& data, int n_threads, diff --git a/otherarch/sdcpp/src/model_loader.h b/otherarch/sdcpp/src/model_loader.h index 0d894a1a2..1973aaaec 100644 --- a/otherarch/sdcpp/src/model_loader.h +++ b/otherarch/sdcpp/src/model_loader.h @@ -27,6 +27,8 @@ struct MmapTensorStore { std::shared_ptr 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* target_tensor_names = nullptr); + const std::set* target_tensor_names = nullptr, + bool log_progress = true); bool load_tensors(std::map& tensors, std::set ignore_tensors = {}, bool use_mmap = false); @@ -76,6 +79,7 @@ public: std::vector& data, int n_threads = 0, bool use_mmap = false); + bool load_tensor(const TensorStorage& tensor_storage, ggml_tensor* dst_tensor); std::vector get_tensor_names() const { std::vector names; diff --git a/otherarch/sdcpp/src/model_manager.cpp b/otherarch/sdcpp/src/model_manager.cpp index 7095ec6a9..3a98bd545 100644 --- a/otherarch/sdcpp/src/model_manager.cpp +++ b/otherarch/sdcpp/src/model_manager.cpp @@ -100,12 +100,42 @@ size_t estimate_tensors_size(const std::map& 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 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(); - 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(); + 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& states) { - std::map> states_by_compute_backend; + std::map, std::vector> 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::vectorname.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& states = pair.second; if (states.empty()) { continue; @@ -285,7 +322,7 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector& 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& 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 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& tensors) { if (tensors.empty()) { return true; diff --git a/otherarch/sdcpp/src/model_manager.h b/otherarch/sdcpp/src/model_manager.h index 9225e3ea6..d80032614 100644 --- a/otherarch/sdcpp/src/model_manager.h +++ b/otherarch/sdcpp/src/model_manager.h @@ -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 tensor_states_by_name_; std::vector> params_storage_blocks_; std::vector> compute_staging_blocks_; + std::map split_buffer_types_; + bool warned_split_lora_skip_ = false; std::set common_ignore_tensors_; std::vector loras_; SDVersion lora_version_ = VERSION_COUNT; @@ -91,6 +95,7 @@ private: bool stage_tensors_to_compute_backend(const std::vector& 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* 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 ignore_tensors); void set_loras(std::vector 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 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 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& tensors, + ggml_backend_t compute_backend) override; bool prepare_params(const std::vector& tensors) override; void release_compute_backend_params(const std::vector& tensors) override; void release_params_backend_params(const std::vector& tensors) override; diff --git a/otherarch/sdcpp/src/name_conversion.cpp b/otherarch/sdcpp/src/name_conversion.cpp index 64b7c681d..2f9a4d184 100644 --- a/otherarch/sdcpp/src/name_conversion.cpp +++ b/otherarch/sdcpp/src/name_conversion.cpp @@ -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> 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 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) { diff --git a/otherarch/sdcpp/src/runtime/denoiser.hpp b/otherarch/sdcpp/src/runtime/denoiser.hpp index 812eebe61..115175b14 100644 --- a/otherarch/sdcpp/src/runtime/denoiser.hpp +++ b/otherarch/sdcpp/src/runtime/denoiser.hpp @@ -4,7 +4,10 @@ #include #include #include +#include #include +#include +#include #include #include @@ -1013,6 +1016,7 @@ struct Denoiser { const sd::Tensor& latent) = 0; virtual sd::Tensor inverse_noise_scaling(float sigma, const sd::Tensor& latent) = 0; + virtual float noise_level_to_sigma(float noise_level) = 0; virtual std::vector 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 inverse_noise_scaling(float sigma, const sd::Tensor& 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 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 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 sample_dpmpp_2m_sde(denoise_cb_t model, + sd::Tensor x, + const std::vector& sigmas, + std::shared_ptr rng, + float eta) { + sd::Tensor old_denoised; + bool have_old_denoised = false; + float h_last = 0.f; + + int steps = static_cast(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 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::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& 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(); + rng->manual_seed(mix64(seed, 0xBF58476D1CE4E5B9ULL)); + w_at_tmax_ = sd::Tensor::randn(shape_, rng) * std::sqrt(static_cast(t_max_ - t_min_)); + } + + sd::Tensor 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(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 w(double t) { + auto it = cache_.find(t); + if (it != cache_.end()) { + return it->second; + } + sd::Tensor zero = sd::Tensor::zeros(shape_); + sd::Tensor out = bridge(t_min_, t_max_, zero, w_at_tmax_, t, root_seed_, kMaxDepth); + cache_.emplace(t, out); + return out; + } + + sd::Tensor bridge(double a, + double c, + const sd::Tensor& w_a, + const sd::Tensor& w_c, + double t, + uint64_t node_seed, + int depth) { + if (depth <= 0 || c - a < 1e-9) { + float alpha = (c > a) ? static_cast((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(); + rng->manual_seed(node_seed); + auto z = sd::Tensor::randn(shape_, rng); + auto w_m = 0.5f * (w_a + w_c) + static_cast(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 shape_; + uint64_t root_seed_; + sd::Tensor w_at_tmax_; + std::map> 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 sample_dpmpp_2m_sde_bt(denoise_cb_t model, + sd::Tensor x, + const std::vector& sigmas, + std::shared_ptr rng, + float eta) { + double sigma_max = 0.0; + double sigma_min = std::numeric_limits::infinity(); + for (float s : sigmas) { + if (s > 0.0f) { + sigma_max = std::max(sigma_max, static_cast(s)); + sigma_min = std::min(sigma_min, static_cast(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 old_denoised; + bool have_old_denoised = false; + float h_last = 0.f; + + int steps = static_cast(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 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 sample_lcm(denoise_cb_t model, @@ -2490,6 +2705,10 @@ static sd::Tensor 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); diff --git a/otherarch/sdcpp/src/runtime/imatrix.cpp b/otherarch/sdcpp/src/runtime/imatrix.cpp index 313eadc68..a6f3bca95 100644 --- a/otherarch/sdcpp/src/runtime/imatrix.cpp +++ b/otherarch/sdcpp/src/runtime/imatrix.cpp @@ -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 tmp(nval); diff --git a/otherarch/sdcpp/src/stable-diffusion.cpp b/otherarch/sdcpp/src/stable-diffusion.cpp index a0ee643b3..f9f4e54e8 100644 --- a/otherarch/sdcpp/src/stable-diffusion.cpp +++ b/otherarch/sdcpp/src/stable-diffusion.cpp @@ -4,11 +4,14 @@ #include #include #include +#include #include +#include #include #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 + 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 - 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 +struct has_set_runtime_backends : std::false_type {}; +template +struct has_set_runtime_backends().set_runtime_backends( + std::declval&>()))>> : std::true_type {}; + static_assert(std::atomic::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 layer_split_vram_limits_for_backends(const std::vector& backends) { + std::vector 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 module_backends = backend_manager.runtime_backends(module); + if (module_backends.size() > 1) { + if constexpr (has_set_runtime_backends::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 + bool register_row_split_runner_params(const std::string& desc, + const std::shared_ptr& model, + SDBackendModule module, + const std::vector& module_backends, + std::map 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 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)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 split_tensors; + if constexpr (std::is_base_of_v) { + model->get_layer_split_param_tensors(split_tensors); + } else { + split_tensors = group_tensors; + } + + std::map row_split_map; + std::map 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 + bool register_layer_split_runner_params(const std::string& desc, + const std::shared_ptr& model, + SDBackendModule module, + const std::vector& module_backends, + std::map 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 =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 get_rng(rng_type_t rng_type) { if (rng_type == STD_DEFAULT_RNG) { return std::make_shared(); @@ -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 wtype_stat = model_loader.get_wtype_stat(); std::map conditioner_wtype_stat = model_loader.get_conditioner_wtype_stat(); std::map 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(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(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(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(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(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(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(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 normalize_ltx_video_latents(const sd::Tensor& 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 prepare_image_generation_latents(sd LOG_INFO("IMG2IMG"); if (request->strength < 1.f) { - size_t t_enc = static_cast(plan->sample_steps * request->strength); - if (t_enc == static_cast(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(plan->sample_steps * request->strength); + if (t_enc == static_cast(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 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(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; } diff --git a/otherarch/sdcpp/src/upscaler.cpp b/otherarch/sdcpp/src/upscaler.cpp index 88a8a6336..dbb99af36 100644 --- a/otherarch/sdcpp/src/upscaler.cpp +++ b/otherarch/sdcpp/src/upscaler.cpp @@ -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; diff --git a/otherarch/sdcpp/src/weight_manager.h b/otherarch/sdcpp/src/weight_manager.h index 28d6cf5c4..82f6d03e4 100644 --- a/otherarch/sdcpp/src/weight_manager.h +++ b/otherarch/sdcpp/src/weight_manager.h @@ -3,10 +3,14 @@ #include +#include "ggml-backend.h" + struct ggml_tensor; struct RunnerWeightManager { virtual ~RunnerWeightManager() = default; + virtual bool assign_compute_backend(const std::vector& tensors, + ggml_backend_t compute_backend) = 0; virtual bool prepare_params(const std::vector& tensors) = 0; virtual void release_compute_backend_params(const std::vector& tensors) = 0; virtual void release_params_backend_params(const std::vector& tensors) = 0;