From 40fbf48a16bc7c55312b7dff24c9596b967015ae Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Sat, 3 Oct 2026 12:52:58 -0300 Subject: [PATCH] sd: sync with master-917-4c3cf75 (#2496) * sd: sync with master-897-ac45422 * sd: sync to master-907-500ef5f * sd: sync with master-913-b167b94 * sd: sync with master-917-4c3cf75 * sd: fix Qwen Image 2.1 edit mode detection * sd: disable TAE for Qwen Image 2.1 * sd: warn about sdvramlimit without sdoffloadcpu --- Makefile | 2 +- otherarch/sdcpp/examples/cli/main.cpp | 69 +- otherarch/sdcpp/examples/common/common.cpp | 158 ++-- otherarch/sdcpp/examples/common/common.h | 29 +- otherarch/sdcpp/examples/common/media_io.cpp | 47 +- otherarch/sdcpp/examples/common/media_io.h | 4 + otherarch/sdcpp/include/stable-diffusion.h | 12 + otherarch/sdcpp/sdtype_adapter.cpp | 3 + .../sdcpp/src/conditioning/conditioner.hpp | 215 +++++- .../src/conditioning/conditioning_cache.h | 107 +++ otherarch/sdcpp/src/core/backend_fit.cpp | 8 +- otherarch/sdcpp/src/core/backend_fit.h | 3 +- otherarch/sdcpp/src/core/ggml_extend.cpp | 51 +- otherarch/sdcpp/src/core/ggml_extend.h | 13 +- .../sdcpp/src/core/ggml_extend_backend.cpp | 19 + .../sdcpp/src/core/ggml_extend_backend.h | 4 + otherarch/sdcpp/src/core/ggml_runner.cpp | 76 +- otherarch/sdcpp/src/core/ggml_runner.h | 36 +- otherarch/sdcpp/src/core/runner_cache.cpp | 30 +- otherarch/sdcpp/src/core/runner_cache.h | 8 +- otherarch/sdcpp/src/detailer.cpp | 1 + otherarch/sdcpp/src/model.h | 19 +- otherarch/sdcpp/src/model/adapter/lora.hpp | 21 + .../sdcpp/src/model/common/ggml_block.hpp | 15 +- otherarch/sdcpp/src/model/common/rope.hpp | 685 ++++++++++-------- .../sdcpp/src/model/common/rope_circular.hpp | 65 ++ otherarch/sdcpp/src/model/diffusion/anima.hpp | 67 +- otherarch/sdcpp/src/model/diffusion/boogu.hpp | 42 +- .../sdcpp/src/model/diffusion/control.hpp | 5 +- .../sdcpp/src/model/diffusion/ernie_image.hpp | 16 +- otherarch/sdcpp/src/model/diffusion/flux.hpp | 26 +- .../sdcpp/src/model/diffusion/ideogram4.hpp | 71 +- otherarch/sdcpp/src/model/diffusion/krea2.hpp | 47 +- otherarch/sdcpp/src/model/diffusion/lens.hpp | 14 +- .../sdcpp/src/model/diffusion/llada_image.hpp | 525 ++++++++++++++ .../sdcpp/src/model/diffusion/mage_flow.hpp | 14 +- .../sdcpp/src/model/diffusion/minimax_h3.hpp | 9 +- .../sdcpp/src/model/diffusion/minit2i.hpp | 21 +- otherarch/sdcpp/src/model/diffusion/model.hpp | 20 +- otherarch/sdcpp/src/model/diffusion/pid.hpp | 42 +- .../sdcpp/src/model/diffusion/qwen_image.hpp | 22 +- .../src/model/diffusion/qwen_image_2_1.hpp | 356 +++++++-- .../sdcpp/src/model/diffusion/sensenova_u1.h | 15 +- .../sdcpp/src/model/diffusion/z_image.hpp | 104 ++- .../sdcpp/src/model/te/llada_image_te.hpp | 604 +++++++++++++++ otherarch/sdcpp/src/model/te/llm.hpp | 369 +++++++++- otherarch/sdcpp/src/model/vae/ltx_vae.hpp | 5 +- otherarch/sdcpp/src/model/vae/wan_vae.hpp | 49 +- otherarch/sdcpp/src/model_loader.cpp | 78 +- otherarch/sdcpp/src/model_loader.h | 6 +- otherarch/sdcpp/src/model_loader_files.cpp | 2 +- otherarch/sdcpp/src/model_manager.cpp | 98 ++- otherarch/sdcpp/src/model_manager.h | 3 +- otherarch/sdcpp/src/name_conversion.cpp | 59 +- .../sdcpp/src/pipeline/diffusion_engine.cpp | 221 +++++- .../sdcpp/src/pipeline/diffusion_engine.h | 6 + otherarch/sdcpp/src/pipeline/image.cpp | 32 +- .../sdcpp/src/pipeline/model_builders.cpp | 17 +- otherarch/sdcpp/src/pipeline/request.cpp | 2 + otherarch/sdcpp/src/pipeline/video.cpp | 29 +- otherarch/sdcpp/src/runtime/denoiser.hpp | 89 ++- .../sdcpp/src/runtime/image_preprocess.cpp | 447 ++++++++++++ .../sdcpp/src/runtime/image_preprocess.h | 88 +++ otherarch/sdcpp/src/runtime/latent-preview.h | 124 +++- otherarch/sdcpp/src/runtime/preprocessing.hpp | 25 +- otherarch/sdcpp/src/runtime/sample-cache.cpp | 20 + otherarch/sdcpp/src/runtime/sample-cache.h | 1 + otherarch/sdcpp/src/stable-diffusion.cpp | 41 +- otherarch/sdcpp/src/upscaler.cpp | 48 +- 69 files changed, 4586 insertions(+), 993 deletions(-) create mode 100644 otherarch/sdcpp/src/conditioning/conditioning_cache.h create mode 100644 otherarch/sdcpp/src/model/common/rope_circular.hpp create mode 100644 otherarch/sdcpp/src/model/diffusion/llada_image.hpp create mode 100644 otherarch/sdcpp/src/model/te/llada_image_te.hpp create mode 100644 otherarch/sdcpp/src/runtime/image_preprocess.cpp create mode 100644 otherarch/sdcpp/src/runtime/image_preprocess.h diff --git a/Makefile b/Makefile index e0080c418..dd6770796 100644 --- a/Makefile +++ b/Makefile @@ -704,7 +704,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/conditioning/wan_audio.cpp src/conditioning/wan_audio.h src/core/backend_fit.cpp src/core/backend_fit.h src/core/compute_workspace.cpp src/core/compute_workspace.h src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.cpp src/core/ggml_extend.h src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/ggml_runner.cpp src/core/ggml_runner.h src/core/ggml_tensor_utils.cpp src/core/ggml_tensor_utils.h src/core/layer_split_partition.cpp src/core/layer_split_partition.h src/core/ordered_map.hpp src/core/parallel.cpp src/core/parallel.h src/core/regex.cpp src/core/regex.h src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/runner_cache.cpp src/core/runner_cache.h src/core/segment_graph_bindings.cpp src/core/segment_graph_bindings.h src/core/segment_weight_pipeline.cpp src/core/segment_weight_pipeline.h src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/detailer.cpp src/detailer.h src/device_residency_manager.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/kcpp_sd_extensions.h src/model/adapter/ip_adapter.hpp src/model/adapter/lora.hpp src/model/adapter/lora_ops.cpp src/model/adapter/lora_ops.h src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/audio/wav2vec2.hpp src/model/common/block.hpp src/model/common/ggml_block.hpp src/model/common/rope.hpp src/model_component.h src/model/detector/yolov8.h src/model/diffusion/anima.hpp src/model/diffusion/animatediff.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/hunyuan.hpp src/model/diffusion/ideogram4.hpp src/model/diffusion/krea2.hpp src/model/diffusion/lens.hpp src/model/diffusion/lingbot_video.hpp src/model/diffusion/ltxv.hpp src/model/diffusion/mage_flow.hpp src/model/diffusion/minimax_h3.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_2_1.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/sefi_image.hpp src/model/diffusion/sensenova_u1.h src/model/diffusion/unet.hpp src/model/diffusion/wan_audio.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/kcpp_sdcpp_quantized_safetensors.hpp 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_files.cpp src/model_loader.h src/model_manager.cpp src/model_manager_files.cpp src/model_manager.h src/model_manager_prefetch.cpp 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/audio_vae.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/hunyuan_vae.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/mage_vae.hpp src/model/vae/minimax_h3_audio_vae.hpp src/model/vae/minimax_h3_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/vae_tiling.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/pipeline/diffusion_engine.cpp src/pipeline/diffusion_engine.h src/pipeline/generation.h src/pipeline/image.cpp src/pipeline/model_builders.cpp src/pipeline/model_builders.h src/pipeline/request.cpp src/pipeline/request.h src/pipeline/video.cpp src/runtime/audio_processing.cpp src/runtime/audio_processing.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/preview_interval.h src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/tiling.cpp src/runtime/tiling.h 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/hf_tokenizer.cpp src/tokenizers/hf_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/sensenova_u1_tokenizer.cpp src/tokenizers/sensenova_u1_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer_config.cpp src/tokenizers/tokenizer_config.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 thirdparty/darts.h thirdparty/miniz.h thirdparty/oniguruma/ascii.c thirdparty/oniguruma/oniguruma.h thirdparty/oniguruma/regcomp.c thirdparty/oniguruma/regenc.c thirdparty/oniguruma/regenc.h thirdparty/oniguruma/regerror.c thirdparty/oniguruma/regexec.c thirdparty/oniguruma/regint.h thirdparty/oniguruma/regparse.c thirdparty/oniguruma/regparse.h thirdparty/oniguruma/st.c thirdparty/oniguruma/st.h thirdparty/oniguruma/unicode.c thirdparty/oniguruma/unicode_fold1_key.c thirdparty/oniguruma/unicode_fold2_key.c thirdparty/oniguruma/unicode_fold3_key.c thirdparty/oniguruma/unicode_unfold_key.c thirdparty/oniguruma/utf8.c thirdparty/stb_image_resize.h thirdparty/stb_image_write.h thirdparty/utf8proc/utf8proc.c thirdparty/utf8proc/utf8proc.h thirdparty/zip.c thirdparty/zip.h +SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/conditioning/conditioning_cache.h src/conditioning/wan_audio.cpp src/conditioning/wan_audio.h src/core/backend_fit.cpp src/core/backend_fit.h src/core/compute_workspace.cpp src/core/compute_workspace.h src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.cpp src/core/ggml_extend.h src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/ggml_runner.cpp src/core/ggml_runner.h src/core/ggml_tensor_utils.cpp src/core/ggml_tensor_utils.h src/core/layer_split_partition.cpp src/core/layer_split_partition.h src/core/ordered_map.hpp src/core/parallel.cpp src/core/parallel.h src/core/regex.cpp src/core/regex.h src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/runner_cache.cpp src/core/runner_cache.h src/core/segment_graph_bindings.cpp src/core/segment_graph_bindings.h src/core/segment_weight_pipeline.cpp src/core/segment_weight_pipeline.h src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/detailer.cpp src/detailer.h src/device_residency_manager.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/kcpp_sd_extensions.h src/model/adapter/ip_adapter.hpp src/model/adapter/lora.hpp src/model/adapter/lora_ops.cpp src/model/adapter/lora_ops.h src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/audio/wav2vec2.hpp src/model/common/block.hpp src/model/common/ggml_block.hpp src/model/common/rope_circular.hpp src/model/common/rope.hpp src/model_component.h src/model/detector/yolov8.h src/model/diffusion/anima.hpp src/model/diffusion/animatediff.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/hunyuan.hpp src/model/diffusion/ideogram4.hpp src/model/diffusion/krea2.hpp src/model/diffusion/lens.hpp src/model/diffusion/lingbot_video.hpp src/model/diffusion/llada_image.hpp src/model/diffusion/ltxv.hpp src/model/diffusion/mage_flow.hpp src/model/diffusion/minimax_h3.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_2_1.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/sefi_image.hpp src/model/diffusion/sensenova_u1.h src/model/diffusion/unet.hpp src/model/diffusion/wan_audio.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/kcpp_sdcpp_quantized_safetensors.hpp 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_files.cpp src/model_loader.h src/model_manager.cpp src/model_manager_files.cpp src/model_manager.h src/model_manager_prefetch.cpp src/model/te/clip.hpp src/model/te/llada_image_te.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/audio_vae.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/hunyuan_vae.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/mage_vae.hpp src/model/vae/minimax_h3_audio_vae.hpp src/model/vae/minimax_h3_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/vae_tiling.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/pipeline/diffusion_engine.cpp src/pipeline/diffusion_engine.h src/pipeline/generation.h src/pipeline/image.cpp src/pipeline/model_builders.cpp src/pipeline/model_builders.h src/pipeline/request.cpp src/pipeline/request.h src/pipeline/video.cpp src/runtime/audio_processing.cpp src/runtime/audio_processing.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/image_preprocess.cpp src/runtime/image_preprocess.h src/runtime/imatrix.cpp src/runtime/imatrix.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/preview_interval.h src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/tiling.cpp src/runtime/tiling.h 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/hf_tokenizer.cpp src/tokenizers/hf_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/sensenova_u1_tokenizer.cpp src/tokenizers/sensenova_u1_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer_config.cpp src/tokenizers/tokenizer_config.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 thirdparty/darts.h thirdparty/miniz.h thirdparty/oniguruma/ascii.c thirdparty/oniguruma/oniguruma.h thirdparty/oniguruma/regcomp.c thirdparty/oniguruma/regenc.c thirdparty/oniguruma/regenc.h thirdparty/oniguruma/regerror.c thirdparty/oniguruma/regexec.c thirdparty/oniguruma/regint.h thirdparty/oniguruma/regparse.c thirdparty/oniguruma/regparse.h thirdparty/oniguruma/st.c thirdparty/oniguruma/st.h thirdparty/oniguruma/unicode.c thirdparty/oniguruma/unicode_fold1_key.c thirdparty/oniguruma/unicode_fold2_key.c thirdparty/oniguruma/unicode_fold3_key.c thirdparty/oniguruma/unicode_unfold_key.c thirdparty/oniguruma/utf8.c thirdparty/stb_image_resize.h thirdparty/stb_image_write.h thirdparty/utf8proc/utf8proc.c thirdparty/utf8proc/utf8proc.h thirdparty/zip.c thirdparty/zip.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/gemma_merges.hpp src/tokenizers/vocab/gemma_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/otherarch/sdcpp/examples/cli/main.cpp b/otherarch/sdcpp/examples/cli/main.cpp index eb96d9b5f..dc79b8ce1 100644 --- a/otherarch/sdcpp/examples/cli/main.cpp +++ b/otherarch/sdcpp/examples/cli/main.cpp @@ -41,7 +41,6 @@ struct SDCliParams { std::string metadata_format = "text"; sd_log_level_t log_level = SD_LOG_INFO; - bool canny_preprocess = false; bool convert_name = false; preview_t preview_method = PREVIEW_NONE; @@ -107,10 +106,6 @@ struct SDCliParams { }; options.bool_options = { - {"", - "--canny", - "apply canny preprocessor (edge detection)", - true, &canny_preprocess}, {"", "--convert-name", "convert tensor name (for convert mode)", @@ -268,7 +263,6 @@ struct SDCliParams { << " metadata_format: \"" << metadata_format << "\",\n" << " log_level: " << log_level_name(log_level) << ",\n" << " color: " << (color ? "true" : "false") << ",\n" - << " canny_preprocess: " << (canny_preprocess ? "true" : "false") << ",\n" << " convert_name: " << (convert_name ? "true" : "false") << ",\n" << " preview_method: " << previews_str[preview_method] << ",\n" << " preview_interval: " << preview_interval << ",\n" @@ -328,9 +322,7 @@ void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) { bool load_images_from_dir(const std::string dir, std::vector& images, - int expected_width = 0, - int expected_height = 0, - int max_image_num = 0) { + int max_image_num = 0) { if (!fs::exists(dir) || !fs::is_directory(dir)) { LOG_ERROR("'%s' is not a valid directory\n", dir.c_str()); return false; @@ -357,7 +349,8 @@ bool load_images_from_dir(const std::string dir, LOG_VERBOSE("load image %zu from '%s'", images.size(), path.c_str()); int width = 0; int height = 0; - uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, expected_width, expected_height); + int loaded_channel = 0; + uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, loaded_channel, 0, 0); if (image_buffer == nullptr) { LOG_ERROR("load image from '%s' failed", path.c_str()); return false; @@ -365,7 +358,7 @@ bool load_images_from_dir(const std::string dir, images.emplace_back(sd_image_t{(uint32_t)width, (uint32_t)height, - 3, + (uint32_t)loaded_channel, image_buffer}); if (max_image_num > 0 && static_cast(images.size()) >= max_image_num) { @@ -651,10 +644,11 @@ int main(int argc, const char* argv[]) { SDCliParams cli_params; SDContextParams ctx_params; + ctx_params.conditioning_cache_size = 0; SDGenerationParams gen_params; - parse_args(argc, argv, cli_params, ctx_params, gen_params); sd_set_log_callback(sd_log_cb, (void*)&cli_params); + parse_args(argc, argv, cli_params, ctx_params, gen_params); if (cli_params.mode == METADATA) { MetadataReadOptions options; @@ -750,16 +744,8 @@ int main(int argc, const char* argv[]) { auto load_image_and_update_size = [&](const std::string& path, SDImageOwner& image, - bool resize_image = true, int expected_channel = 3) -> bool { - int expected_width = 0; - int expected_height = 0; - if (resize_image && gen_params.width_and_height_are_set()) { - expected_width = gen_params.width; - expected_height = gen_params.height; - } - - if (!load_sd_image_from_file(image.put(), path.c_str(), expected_width, expected_height, expected_channel)) { + if (!load_sd_image_from_file(image.put(), path.c_str(), 0, 0, expected_channel)) { LOG_ERROR("load image from '%s' failed", path.c_str()); return false; } @@ -781,7 +767,8 @@ int main(int argc, const char* argv[]) { }; if (gen_params.init_image_path.size() > 0) { - if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image)) { + const bool native_init = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER; + if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image, native_init ? 0 : 3)) { return 1; } } @@ -795,8 +782,8 @@ int main(int argc, const char* argv[]) { if (gen_params.ref_image_paths.size() > 0) { gen_params.ref_images.clear(); for (auto& path : gen_params.ref_image_paths) { - SDImageOwner ref_image({0, 0, 3, nullptr}); - if (!load_image_and_update_size(path, ref_image, false)) { + SDImageOwner ref_image({0, 0, 0, nullptr}); + if (!load_image_and_update_size(path, ref_image, 0)) { return 1; } gen_params.ref_images.push_back(std::move(ref_image)); @@ -837,41 +824,22 @@ int main(int argc, const char* argv[]) { if (gen_params.mask_image_path.size() > 0) { if (!load_sd_image_from_file(gen_params.mask_image.put(), gen_params.mask_image_path.c_str(), - gen_params.get_resolved_width(), - gen_params.get_resolved_height(), + 0, + 0, 1)) { LOG_ERROR("load image from '%s' failed", gen_params.mask_image_path.c_str()); return 1; } - } else { - sd_image_t generated_mask = {0, 0, 1, nullptr}; - generated_mask.data = (uint8_t*)malloc(gen_params.get_resolved_width() * gen_params.get_resolved_height()); - if (generated_mask.data == nullptr) { - LOG_ERROR("malloc mask image failed"); - return 1; - } - generated_mask.width = gen_params.get_resolved_width(); - generated_mask.height = gen_params.get_resolved_height(); - memset(generated_mask.data, 255, gen_params.get_resolved_width() * gen_params.get_resolved_height()); - gen_params.mask_image.reset(generated_mask); } if (gen_params.control_image_path.size() > 0) { if (!load_sd_image_from_file(gen_params.control_image.put(), gen_params.control_image_path.c_str(), - gen_params.get_resolved_width(), - gen_params.get_resolved_height())) { + 0, + 0)) { LOG_ERROR("load image from '%s' failed", gen_params.control_image_path.c_str()); return 1; } - if (cli_params.canny_preprocess) { // apply preprocessor - preprocess_canny(gen_params.control_image.get(), - 0.08f, - 0.08f, - 0.8f, - 1.0f, - false); - } } if (gen_params.ip_adapter_image_path.size() > 0) { @@ -888,8 +856,6 @@ int main(int argc, const char* argv[]) { gen_params.control_frames.clear(); if (!load_images_from_dir(gen_params.control_video_path, gen_params.control_frames, - gen_params.get_resolved_width(), - gen_params.get_resolved_height(), gen_params.video_frames)) { return 1; } @@ -898,10 +864,7 @@ int main(int argc, const char* argv[]) { if (!gen_params.pm_id_images_dir.empty()) { gen_params.pm_id_images.clear(); if (!load_images_from_dir(gen_params.pm_id_images_dir, - gen_params.pm_id_images, - 0, - 0, - 0)) { + gen_params.pm_id_images)) { return 1; } } diff --git a/otherarch/sdcpp/examples/common/common.cpp b/otherarch/sdcpp/examples/common/common.cpp index 933c8622e..fa7354bc8 100644 --- a/otherarch/sdcpp/examples/common/common.cpp +++ b/otherarch/sdcpp/examples/common/common.cpp @@ -518,7 +518,8 @@ ArgOptions SDContextParams::get_options() { {"", "--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", + "chroma_t5_mask_pad, qwen_image_zero_cond_t, qwen_image_2_1_prefix_cache, " + "qwen_image_2_1_prefix_cache_type (auto or a type name from --type)", (int)',', &model_args}, {"", @@ -571,6 +572,10 @@ 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}, + {"", + "--conditioning-cache-size", + "maximum number of conditioning results cached per model context (default: " + std::to_string(conditioning_cache_size) + ", 0 disables caching)", + &conditioning_cache_size}, }; options.bool_options = { @@ -618,6 +623,10 @@ ArgOptions SDContextParams::get_options() { "--diffusion-fa", "use flash attention in the diffusion model only", true, &diffusion_flash_attn}, + {"", + "--sage-attn", + "use native CUDA SageAttention in the diffusion model, with flash/default attention fallback", + true, &sage_attn}, {"", "--diffusion-conv-direct", "use ggml_conv2d_direct in the diffusion model", @@ -818,6 +827,10 @@ bool SDContextParams::resolve(SDMode mode) { } bool SDContextParams::validate(SDMode mode) { + if (conditioning_cache_size < 0) { + LOG_ERROR("error: conditioning-cache-size must be non-negative"); + return false; + } if (mode == CONVERT) { const bool has_convert_input = model_path.length() != 0 || clip_l_path.length() != 0 || @@ -894,6 +907,7 @@ std::string SDContextParams::to_string() const { std::ostringstream oss; oss << "SDContextParams {\n" << " n_threads: " << n_threads << ",\n" + << " conditioning_cache_size: " << conditioning_cache_size << ",\n" << " model_path: \"" << model_path << "\",\n" << " clip_l_path: \"" << clip_l_path << "\",\n" << " clip_g_path: \"" << clip_g_path << "\",\n" @@ -938,6 +952,7 @@ std::string SDContextParams::to_string() const { << " vae_on_cpu: " << (vae_on_cpu ? "true" : "false") << ",\n" << " flash_attn: " << (flash_attn ? "true" : "false") << ",\n" << " diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n" + << " sage_attn: " << (sage_attn ? "true" : "false") << ",\n" << " linear_scale: " << linear_scale << ",\n" << " attn_scale: " << attn_scale << ",\n" << " diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n" @@ -987,6 +1002,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.pulid_weights_path = pulid_weights_path.c_str(); sd_ctx_params.tensor_type_rules = tensor_type_rules.c_str(); sd_ctx_params.n_threads = n_threads; + sd_ctx_params.conditioning_cache_size = conditioning_cache_size; sd_ctx_params.wtype = wtype; sd_ctx_params.rng_type = rng_type; sd_ctx_params.sampler_rng_type = sampler_rng_type; @@ -995,6 +1011,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.enable_mmap = enable_mmap; sd_ctx_params.flash_attn = flash_attn; sd_ctx_params.diffusion_flash_attn = diffusion_flash_attn; + sd_ctx_params.sage_attn = sage_attn; sd_ctx_params.linear_scale = linear_scale; sd_ctx_params.attn_scale = attn_scale; sd_ctx_params.tae_preview_only = taesd_preview; @@ -1109,7 +1126,7 @@ ArgOptions SDGenerationParams::get_options() { &hires_upscaler}, {"", "--extra-sample-args", - "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng", + "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; llada_image supports uniform; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng", (int)',', &extra_sample_args}, {"", @@ -1122,6 +1139,9 @@ ArgOptions SDGenerationParams::get_options() { "Key-value list to set up the way the reference images are processed (empty = auto-detect from model weigths)", (int)',', &ref_image_args}, + {"", "--image-preprocess", + "Image preprocessing rule: target=init|end|mask|control|ref|ip-adapter|id|control-frame,index=N,mode=auto|none|stretch|crop|crop-resize|fit-pad,filter=auto|nearest|nearest-exact|bilinear|bicubic|lanczos,antialias=auto|true|false,width=W,height=H,anchor=center|top|bottom|left|right,pad_color=#RRGGBB[AA],canny=true|false. Repeat for multiple rules.", + (int)';', &image_preprocess}, }; options.int_options = { @@ -1302,11 +1322,6 @@ ArgOptions SDGenerationParams::get_options() { "automatically increase the indices of references images based on the order they are listed (starting with 1).", true, &increase_ref_index}, - {"", - "--disable-auto-resize-ref-image", - "disable auto resize of ref images", - false, - &auto_resize_ref_image}, {"", "--circular", "enable circular padding on both axes for tileable output", @@ -1748,7 +1763,7 @@ ArgOptions SDGenerationParams::get_options() { on_scm_policy_arg}, {"", "--vae-tile-size", - "tile size for vae tiling, format [X]x[Y] (default: 32x32)", + "tile size for vae tiling in latent units, not image pixels, format [X]x[Y] (default: 32x32)", on_tile_size_arg}, {"", "--vae-relative-tile-size", @@ -1842,28 +1857,28 @@ bool decode_base64_image(const std::string& encoded_input, return false; } - int decoded_width = 0; - int decoded_height = 0; - uint8_t* raw_data = load_image_from_memory(reinterpret_cast(image_bytes.data()), - static_cast(image_bytes.size()), - decoded_width, - decoded_height, - expected_width, - expected_height, - target_channels); + int decoded_width = 0; + int decoded_height = 0; + int resolved_channel = target_channels; + uint8_t* raw_data = load_image_from_memory(reinterpret_cast(image_bytes.data()), + static_cast(image_bytes.size()), + decoded_width, + decoded_height, + resolved_channel, + expected_width, + expected_height, + target_channels); if (raw_data == nullptr) { return false; } - out_image.reset({(uint32_t)decoded_width, (uint32_t)decoded_height, (uint32_t)target_channels, raw_data}); + out_image.reset({(uint32_t)decoded_width, (uint32_t)decoded_height, (uint32_t)resolved_channel, raw_data}); return true; } static bool parse_image_json_field(const json& parent, const char* key, int channels, - int expected_width, - int expected_height, SDImageOwner& out_image) { if (!parent.contains(key)) { return true; @@ -1875,14 +1890,12 @@ static bool parse_image_json_field(const json& parent, if (!parent.at(key).is_string()) { return false; } - return decode_base64_image(parent.at(key).get(), channels, expected_width, expected_height, out_image); + return decode_base64_image(parent.at(key).get(), channels, 0, 0, out_image); } static bool parse_image_array_json_field(const json& parent, const char* key, int channels, - int expected_width, - int expected_height, std::vector& out_images) { if (!parent.contains(key)) { return true; @@ -1901,7 +1914,7 @@ static bool parse_image_array_json_field(const json& parent, return false; } SDImageOwner image; - if (!decode_base64_image(item.get(), channels, expected_width, expected_height, image)) { + if (!decode_base64_image(item.get(), channels, 0, 0, image)) { return false; } out_images.push_back(std::move(image)); @@ -2000,6 +2013,29 @@ static bool resolve_model_file_from_dir(const std::string& model_name, return false; } +bool SDGenerationParams::parse_image_preprocess_json(const std::string& json_str) { + const auto value = json::parse(json_str, nullptr, false); + std::string rules; + if (value.is_string()) { + rules = value.get(); + } else if (value.is_array()) { + for (const auto& item : value) { + if (!item.is_string()) { + LOG_ERROR("image_preprocess must contain rule strings"); + return false; + } + if (!rules.empty()) + rules += ";"; + rules += item.get(); + } + } else { + LOG_ERROR("image_preprocess must be a string or array of strings"); + return false; + } + image_preprocess = std::move(rules); + return true; +} + bool SDGenerationParams::from_json_str( const std::string& json_str, const std::function& lora_path_resolver) { @@ -2011,6 +2047,9 @@ bool SDGenerationParams::from_json_str( return false; } + if (j.contains("image_preprocess") && !parse_image_preprocess_json(j["image_preprocess"].dump())) + return false; + auto load_if_exists = [&](const char* key, auto& out) { if (j.contains(key)) { using T = std::decay_t; @@ -2048,6 +2087,7 @@ bool SDGenerationParams::from_json_str( load_if_exists("cache_mode", cache_mode); load_if_exists("cache_option", cache_option); load_if_exists("scm_mask", scm_mask); + load_if_exists("ref_image_args", ref_image_args); load_if_exists("clip_skip", clip_skip); load_if_exists("width", width); @@ -2065,7 +2105,6 @@ bool SDGenerationParams::from_json_str( load_if_exists("moe_boundary", moe_boundary); load_if_exists("vace_strength", vace_strength); - load_if_exists("auto_resize_ref_image", auto_resize_ref_image); load_if_exists("increase_ref_index", increase_ref_index); load_if_exists("embed_image_metadata", embed_image_metadata); @@ -2209,32 +2248,23 @@ bool SDGenerationParams::from_json_str( LOG_ERROR("invalid lora"); return false; } - if (!parse_image_json_field(j, "init_image", 3, width, height, init_image)) { - LOG_ERROR("invalid init_image"); + auto load_image = [&](const char* key, int channels, SDImageOwner& image) { + if (!parse_image_json_field(j, key, channels, image)) { + LOG_ERROR("invalid %s", key); + return false; + } + return true; + }; + if (!load_image("init_image", 0, init_image) || + !load_image("end_image", 3, end_image) || + !load_image("mask_image", 1, mask_image) || + !load_image("control_image", 3, control_image) || + !load_image("ip_adapter_image", 3, ip_adapter_image)) { return false; } - if (!parse_image_json_field(j, "end_image", 3, width, height, end_image)) { - LOG_ERROR("invalid end_image"); - return false; - } - if (!parse_image_array_json_field(j, "ref_images", 3, width, height, ref_images)) { - LOG_ERROR("invalid ref_images"); - return false; - } - if (!parse_image_array_json_field(j, "control_frames", 3, width, height, control_frames)) { - LOG_ERROR("invalid control_frames"); - return false; - } - if (!parse_image_json_field(j, "mask_image", 1, width, height, mask_image)) { - LOG_ERROR("invalid mask_image"); - return false; - } - if (!parse_image_json_field(j, "control_image", 3, width, height, control_image)) { - LOG_ERROR("invalid control_image"); - return false; - } - if (!parse_image_json_field(j, "ip_adapter_image", 3, width, height, ip_adapter_image)) { - LOG_ERROR("invalid ip_adapter_image"); + if (!parse_image_array_json_field(j, "ref_images", 0, ref_images) || + !parse_image_array_json_field(j, "control_frames", 3, control_frames)) { + LOG_ERROR("invalid input image array"); return false; } @@ -2478,6 +2508,10 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s } bool SDGenerationParams::validate(SDMode mode) { + if (!image_preprocess.empty() && mode != IMG_GEN && mode != VID_GEN) { + LOG_ERROR("--image-preprocess requires img_gen or vid_gen mode"); + return false; + } if (batch_count <= 0) { LOG_ERROR("error: batch_count must be greater than 0"); return false; @@ -2653,14 +2687,6 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() { pulid_id_weight, }; - if (!auto_resize_ref_image) { - if (!ref_image_args.empty()) { - ref_image_args += ","; - } - ref_image_args += "resize_before_vae=0"; - LOG_WARN("Notice: --disable-auto-resize-ref-image is deprecated. Use --ref-image-args \"resize_before_vae=off\" instead."); - } - if (increase_ref_index) { if (!ref_image_args.empty()) { ref_image_args += ","; @@ -2708,6 +2734,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() { params.hires.custom_sigmas_count = static_cast(hires_custom_sigmas.size()); params.circular_x = circular || circular_x; params.circular_y = circular || circular_y; + params.image_preprocess = {image_preprocess.c_str()}; return params; } @@ -2810,6 +2837,7 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() { params.hires.custom_sigmas_count = static_cast(hires_custom_sigmas.size()); params.circular_x = circular || circular_x; params.circular_y = circular || circular_y; + params.image_preprocess = {image_preprocess.c_str()}; return params; } @@ -2866,7 +2894,8 @@ std::string SDGenerationParams::to_string() const { << " ref_video_audio_paths: " << vec_str_to_string(ref_video_audio_paths) << ",\n" << " ref_audio_paths: " << vec_str_to_string(ref_audio_paths) << ",\n" << " control_video_path: \"" << control_video_path << "\",\n" - << " auto_resize_ref_image: " << (auto_resize_ref_image ? "true" : "false") << ",\n" + << " image_preprocess: " << image_preprocess << ",\n" + << " ref_image_args: " << ref_image_args << ",\n" << " increase_ref_index: " << (increase_ref_index ? "true" : "false") << ",\n" << " pm_id_images_dir: \"" << pm_id_images_dir << "\",\n" << " pm_id_embed_path: \"" << pm_id_embed_path << "\",\n" @@ -3017,12 +3046,13 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params, set_json_basename_if_not_empty(models, "control_net", ctx_params.control_net_path); root["models"] = std::move(models); - root["clip_skip"] = gen_params.clip_skip; - root["strength"] = gen_params.strength; - root["control_strength"] = gen_params.control_strength; - root["ip_adapter_strength"] = gen_params.ip_adapter_strength; - root["auto_resize_ref_image"] = gen_params.auto_resize_ref_image; - root["increase_ref_index"] = gen_params.increase_ref_index; + root["clip_skip"] = gen_params.clip_skip; + root["strength"] = gen_params.strength; + root["control_strength"] = gen_params.control_strength; + root["ip_adapter_strength"] = gen_params.ip_adapter_strength; + root["ref_image_args"] = gen_params.ref_image_args; + root["image_preprocess"] = gen_params.image_preprocess; + root["increase_ref_index"] = gen_params.increase_ref_index; if (mode == VID_GEN) { root["video"] = { {"frame_count", gen_params.video_frames}, diff --git a/otherarch/sdcpp/examples/common/common.h b/otherarch/sdcpp/examples/common/common.h index 3b66ae4ab..d0d39b738 100644 --- a/otherarch/sdcpp/examples/common/common.h +++ b/otherarch/sdcpp/examples/common/common.h @@ -116,7 +116,8 @@ bool decode_base64_image(const std::string& encoded_input, SDImageOwner& out_image); struct SDContextParams { - int n_threads = -1; + int n_threads = -1; + int conditioning_cache_size = 4; std::string model_path; std::string clip_l_path; std::string clip_g_path; @@ -170,6 +171,7 @@ struct SDContextParams { bool vae_on_cpu = false; bool flash_attn = false; bool diffusion_flash_attn = false; + bool sage_attn = false; bool diffusion_conv_direct = false; bool vae_conv_direct = false; @@ -199,18 +201,17 @@ struct SDGenerationParams { std::string ad_prompt; std::string ad_negative_prompt; std::string extra_ad_args; - int clip_skip = -1; // <= 0 represents unspecified - int width = -1; - int height = -1; - int batch_count = 1; - int qwen_image_layers = 3; - int64_t seed = 42; - float strength = 0.75f; - float control_strength = 0.9f; - float ip_adapter_strength = 1.0f; - bool auto_resize_ref_image = true; - bool increase_ref_index = false; - bool embed_image_metadata = true; + int clip_skip = -1; // <= 0 represents unspecified + int width = -1; + int height = -1; + int batch_count = 1; + int qwen_image_layers = 3; + int64_t seed = 42; + float strength = 0.75f; + float control_strength = 0.9f; + float ip_adapter_strength = 1.0f; + bool increase_ref_index = false; + bool embed_image_metadata = true; std::string init_image_path; std::string end_image_path; @@ -246,6 +247,7 @@ struct SDGenerationParams { std::string extra_tiling_args; std::string ref_image_args; + std::string image_preprocess; std::string pm_id_images_dir; std::string pm_id_embed_path; @@ -309,6 +311,7 @@ struct SDGenerationParams { ArgOptions get_options(); bool from_json_str(const std::string& json_str, const std::function& lora_path_resolver = {}); + bool parse_image_preprocess_json(const std::string& json_str); bool initialize_cache_params(); void extract_and_remove_lora(const std::string& lora_model_dir); bool width_and_height_are_set() const; diff --git a/otherarch/sdcpp/examples/common/media_io.cpp b/otherarch/sdcpp/examples/common/media_io.cpp index 834584991..498117313 100644 --- a/otherarch/sdcpp/examples/common/media_io.cpp +++ b/otherarch/sdcpp/examples/common/media_io.cpp @@ -275,6 +275,10 @@ uint8_t* decode_webp_image_to_buffer(const uint8_t* data, height = features.height; source_channel_count = features.has_alpha ? 4 : 3; + if (expected_channel == 0) { + expected_channel = source_channel_count; + } + const size_t pixel_count = static_cast(width) * static_cast(height); if (expected_channel == 1) { @@ -495,7 +499,8 @@ uint8_t* load_image_common(bool from_memory, int& height, int expected_width, int expected_height, - int expected_channel) { + int expected_channel, + int& out_channel) { const char* image_path; FreeUniquePtr image_buffer; int source_channel_count = 0; @@ -552,6 +557,32 @@ uint8_t* load_image_common(bool from_memory, LOG_ERROR("load image from '%s' failed", image_path); return nullptr; } + if (expected_channel == 0) { + expected_channel = source_channel_count == 2 ? 4 : (source_channel_count == 1 ? 3 : source_channel_count); + if (expected_channel != source_channel_count) { + FreeUniquePtr promoted((uint8_t*)malloc((size_t)width * height * expected_channel)); + if (promoted == nullptr) { + LOG_ERROR("error: allocate memory for channel promotion, image_path = %s", image_path); + return nullptr; + } + const size_t pixel_count = (size_t)width * (size_t)height; + for (size_t i = 0; i < pixel_count; ++i) { + if (source_channel_count == 1) { + promoted.get()[i * 3 + 0] = image_buffer.get()[i]; + promoted.get()[i * 3 + 1] = image_buffer.get()[i]; + promoted.get()[i * 3 + 2] = image_buffer.get()[i]; + } else { + promoted.get()[i * 4 + 0] = image_buffer.get()[i * 2]; + promoted.get()[i * 4 + 1] = image_buffer.get()[i * 2]; + promoted.get()[i * 4 + 2] = image_buffer.get()[i * 2]; + promoted.get()[i * 4 + 3] = image_buffer.get()[i * 2 + 1]; + } + } + image_buffer = std::move(promoted); + source_channel_count = expected_channel; + } + } + // stb reports the source channel count even when it converts the output. if (source_channel_count < expected_channel) { fprintf(stderr, "the number of channels for the input image must be >= %d," @@ -611,7 +642,7 @@ uint8_t* load_image_common(bool from_memory, } stbir_resize(image_buffer.get(), width, height, 0, resized_image_buffer.get(), expected_width, expected_height, 0, STBIR_TYPE_UINT8, - expected_channel, STBIR_ALPHA_CHANNEL_NONE, 0, + expected_channel, expected_channel == 4 ? 3 : STBIR_ALPHA_CHANNEL_NONE, 0, STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP, STBIR_FILTER_BOX, STBIR_FILTER_BOX, STBIR_COLORSPACE_SRGB, nullptr); @@ -619,6 +650,7 @@ uint8_t* load_image_common(bool from_memory, height = expected_height; image_buffer = std::move(resized_image_buffer); } + out_channel = expected_channel; return image_buffer.release(); } @@ -791,10 +823,11 @@ bool write_image_to_file(const std::string& path, uint8_t* load_image_from_file(const char* image_path, int& width, int& height, + int& out_channel, int expected_width, int expected_height, int expected_channel) { - return load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel); + return load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel, out_channel); } bool load_sd_image_from_file(sd_image_t* image, @@ -804,13 +837,14 @@ bool load_sd_image_from_file(sd_image_t* image, int expected_channel) { int width; int height; - image->data = load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel); + int resolved_channel = expected_channel; + image->data = load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel, resolved_channel); if (image->data == nullptr) { return false; } image->width = width; image->height = height; - image->channel = expected_channel; + image->channel = resolved_channel; return true; } @@ -818,10 +852,11 @@ uint8_t* load_image_from_memory(const char* image_bytes, int len, int& width, int& height, + int& out_channel, int expected_width, int expected_height, int expected_channel) { - return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel); + return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel, out_channel); } static void append_avi_metadata(std::vector& data, const std::string& parameters) { diff --git a/otherarch/sdcpp/examples/common/media_io.h b/otherarch/sdcpp/examples/common/media_io.h index f0f3d63aa..aeef9f405 100644 --- a/otherarch/sdcpp/examples/common/media_io.h +++ b/otherarch/sdcpp/examples/common/media_io.h @@ -32,9 +32,12 @@ bool write_image_to_file(const std::string& path, const std::string& parameters = "", int quality = 90); +// expected_channel == 0 preserves native channels (grayscale -> RGB, gray+alpha -> RGBA). +// out_channel receives the output channel count. uint8_t* load_image_from_file(const char* image_path, int& width, int& height, + int& out_channel, int expected_width = 0, int expected_height = 0, int expected_channel = 3); @@ -49,6 +52,7 @@ uint8_t* load_image_from_memory(const char* image_bytes, int len, int& width, int& height, + int& out_channel, int expected_width = 0, int expected_height = 0, int expected_channel = 3); diff --git a/otherarch/sdcpp/include/stable-diffusion.h b/otherarch/sdcpp/include/stable-diffusion.h index 9dbddf9bd..a716eb0ed 100644 --- a/otherarch/sdcpp/include/stable-diffusion.h +++ b/otherarch/sdcpp/include/stable-diffusion.h @@ -79,6 +79,7 @@ enum scheduler_t { FLUX2_SCHEDULER, FLUX_SCHEDULER, BETA_SCHEDULER, + LLADA_IMAGE_SCHEDULER, SCHEDULER_COUNT }; @@ -245,6 +246,8 @@ typedef struct { float linear_scale; // Override linear input scaling; 0 keeps the model default float attn_scale; // Override flash-attention K/V scaling; 0 keeps the model default const char* tokenizer; // tokenizer.json path or main=FILE,clip-l=FILE,clip-g=FILE assignments; required for PiD and Lens + bool sage_attn; + int conditioning_cache_size; // Maximum cached conditioning entries per context; 0 disables caching (default: 4) } sd_ctx_params_t; typedef struct { @@ -261,6 +264,11 @@ typedef struct { uint8_t* data; } sd_image_t; +typedef struct { + // Semicolon-separated target=...,key=value rules. NULL preserves defaults. + const char* rules; +} sd_image_preprocess_params_t; + typedef struct { sd_image_t* frames; int frame_count; @@ -408,6 +416,7 @@ typedef struct { int qwen_image_layers; bool circular_x; bool circular_y; + sd_image_preprocess_params_t image_preprocess; } sd_img_gen_params_t; typedef struct { @@ -442,6 +451,7 @@ typedef struct { sd_hires_params_t hires; bool circular_x; bool circular_y; + sd_image_preprocess_params_t image_preprocess; } sd_vid_gen_params_t; typedef struct sd_ctx_t sd_ctx_t; @@ -549,6 +559,8 @@ SD_API bool upscale(upscaler_ctx_t* upscaler_ctx, int* num_images_out); SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx); +// Reads model metadata only; returns 0 if the file is not a recognized RGB ESRGAN model. +SD_API int get_upscaler_model_scale(const char* model_path); typedef struct adetailer_ctx_t adetailer_ctx_t; diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp index 7b1e74328..218faf138 100644 --- a/otherarch/sdcpp/sdtype_adapter.cpp +++ b/otherarch/sdcpp/sdtype_adapter.cpp @@ -488,6 +488,9 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) { if(inputs.max_vram && *inputs.max_vram) { max_vram = inputs.max_vram; printf("Using max VRAM = %s GB\n", max_vram.c_str()); + if (params_backend == "") { + printf("Note: a VRAM limit may not be effective without offloading!\n"); + } } if(inputs.quant > 0) { diff --git a/otherarch/sdcpp/src/conditioning/conditioner.hpp b/otherarch/sdcpp/src/conditioning/conditioner.hpp index a0544eec5..6b750f939 100644 --- a/otherarch/sdcpp/src/conditioning/conditioner.hpp +++ b/otherarch/sdcpp/src/conditioning/conditioner.hpp @@ -14,6 +14,7 @@ #include "core/util.h" #include "model/diffusion/model.hpp" #include "model/te/clip.hpp" +#include "model/te/llada_image_te.hpp" #include "model/te/llm.hpp" #include "model/te/t5.hpp" #include "model_loader.h" @@ -2218,7 +2219,10 @@ struct LLMEmbedder : public Conditioner { false, deepstack_image_embeds, image_grids); - GGML_ASSERT(!hidden_states.empty()); + if (hidden_states.empty()) { + LOG_ERROR("LLM prompt encoding failed"); + return {}; + } hidden_states = apply_token_weights(std::move(hidden_states), weights); GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx); @@ -3162,6 +3166,7 @@ struct LLMEmbedder : public Conditioner { int64_t tag_count = static_cast(tags.size()); result.c_token_types = sd::Tensor({tag_count}, std::move(tags)); } + return result; } }; @@ -3236,6 +3241,214 @@ struct LTXAVTextProjectionRunner : public GGMLRunner { } }; +// LLaDA-Image's text path is a three-stage pipeline rather than a single encoder pass: +// the token embeddings feed a QueryFormer whose 256 queries are appended to the backbone +// input, and the backbone's final hidden states are projected to the denoiser's caption dim. +// Ref: LLaDAImagePipeline._encode_text. +struct LLaDAImageEmbedder : public Conditioner { + std::shared_ptr tokenizer; + std::shared_ptr llm; + std::shared_ptr query_former; + std::shared_ptr text_projection; + std::shared_ptr sigvq; + + std::string llm_prefix; + std::string query_former_prefix; + std::string text_projection_prefix; + std::string sigvq_prefix; + + LLaDAImageEmbedder(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string& llm_prefix = "text_encoders.llm", + const std::string& query_former_prefix = "queryformer", + const std::string& text_projection_prefix = "text_projection", + const std::string& sigvq_prefix = "sigvq", + std::shared_ptr weight_manager = nullptr, + const TokenizerConfig& tokenizers = {}) + : llm_prefix(llm_prefix), + query_former_prefix(query_former_prefix), + text_projection_prefix(text_projection_prefix), + sigvq_prefix(sigvq_prefix) { + if (!tokenizers.has(TokenizerConfig::MAIN)) { + throw std::runtime_error("LLaDA-Image requires an external LLaDA2 tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer"); + } + llm = std::make_shared(LLM::LLMArch::LLADA2_MOE, + backend, + tensor_storage_map, + llm_prefix, + false, + weight_manager); + // <|endoftext|> doubles as the pad token in LLaDA2's tokenizer.json. + tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 156892); + query_former = std::make_shared(backend, + tensor_storage_map, + query_former_prefix, + weight_manager); + text_projection = std::make_shared(backend, + tensor_storage_map, + text_projection_prefix, + weight_manager); + + // SigVQ is only present when the user supplies the editing weights. + for (const auto& [name, _] : tensor_storage_map) { + if (starts_with(name, sigvq_prefix + ".")) { + sigvq = std::make_shared(backend, + tensor_storage_map, + sigvq_prefix, + weight_manager); + break; + } + } + } + + void get_param_tensors(std::map& tensors) override { + llm->get_param_tensors(tensors, llm_prefix); + query_former->get_param_tensors(tensors, query_former_prefix); + text_projection->get_param_tensors(tensors, text_projection_prefix); + if (sigvq != nullptr) { + sigvq->get_param_tensors(tensors, sigvq_prefix); + } + } + + void get_param_tensor_ops(std::map& tensor_ops) override { + llm->get_param_tensor_ops(tensor_ops); + } + + void set_flash_attention_enabled(bool enabled) override { + llm->set_flash_attention_enabled(enabled); + query_former->set_flash_attention_enabled(enabled); + text_projection->set_flash_attention_enabled(enabled); + if (sigvq != nullptr) { + sigvq->set_flash_attention_enabled(enabled); + } + } + + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + llm->set_max_graph_vram_bytes(max_vram_bytes); + query_former->set_max_graph_vram_bytes(max_vram_bytes); + text_projection->set_max_graph_vram_bytes(max_vram_bytes); + if (sigvq != nullptr) { + sigvq->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, llm_prefix); + } + + void set_weight_adapter(const std::shared_ptr& adapter) override { + llm->set_weight_adapter(adapter); + query_former->set_weight_adapter(adapter); + text_projection->set_weight_adapter(adapter); + if (sigvq != nullptr) { + sigvq->set_weight_adapter(adapter); + } + } + + void runner_end() override { + llm->runner_end(); + query_former->runner_end(); + text_projection->runner_end(); + if (sigvq != nullptr) { + sigvq->runner_end(); + } + } + + SDCondition get_learned_condition(int n_threads, + const ConditionerParams& conditioner_params) override { + const int64_t num_queries = 256; + const bool has_ref_images = conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty(); + if (has_ref_images && sigvq == nullptr) { + LOG_ERROR("LLaDA-Image editing requires connectors with SigVQ weights"); + return {}; + } + + std::string text = conditioner_params.text; + while (!text.empty() && std::isspace(static_cast(text.front()))) { + text.erase(text.begin()); + } + while (!text.empty() && std::isspace(static_cast(text.back()))) { + text.pop_back(); + } + std::string prompt = text.empty() + ? "HUMAN Generate an image.\nASSISTANT\n" + : "HUMAN Generate an image: " + text + "\nASSISTANT\n"; + + std::vector tokens; + if (!tokenizer->encode(prompt, tokens, nullptr)) { + return {}; + } + int64_t n_text = static_cast(tokens.size()); + GGML_ASSERT(n_text > 0); + + sd::Tensor text_ids({n_text}, std::vector(tokens.begin(), tokens.end())); + auto inputs_embeds = llm->compute_input_embeds(n_threads, text_ids); + auto query_embeds = query_former->compute(n_threads, inputs_embeds); + + // splice_image_embeds() replaces tokens in place, so the query slots have to exist in + // input_ids; their ids are irrelevant because the embeddings are overwritten. + std::vector padded(tokens.begin(), tokens.end()); + padded.resize(static_cast(n_text + num_queries), tokenizer->PAD_TOKEN_ID); + int64_t n_total = static_cast(padded.size()); + sd::Tensor input_ids({n_total}, padded); + + // Bidirectional everywhere except that the text tokens must not see the appended + // queries, matching backbone_attention_mask[:, :, :text_length, text_length:] = min. + const float mask_min = std::numeric_limits::lowest() / 4.0f; + sd::Tensor attention_mask({n_total, n_total}); + for (int64_t i1 = 0; i1 < n_total; ++i1) { + for (int64_t i0 = 0; i0 < n_total; ++i0) { + float value = (i1 < n_text && i0 >= n_text) ? mask_min : 0.0f; + attention_mask[i0 + n_total * i1] = value; + } + } + + LLM::ImageEmbeds image_embeds; + image_embeds.emplace_back(static_cast(n_text), query_embeds); + + std::set out_layers = {static_cast(llm->config.num_layers) + 1}; + auto hidden_states = llm->compute(n_threads, + input_ids, + attention_mask, + image_embeds, + out_layers); + + SDCondition result; + result.c_crossattn = text_projection->compute(n_threads, hidden_states); + + // Editing: SigVQ sees the reference at half the output resolution, as in + // LLaDAImagePipeline._encode_source_image. + if (has_ref_images) { + const auto& ref = conditioner_params.ref_images->front(); + auto resized = sd::ops::interpolate(ref, + {conditioner_params.width / 2, + conditioner_params.height / 2, + ref.shape()[2], + ref.shape()[3]}, + sd::ops::InterpolateMode::Bilinear); + resized = resized * 2.f - 1.f; + auto semantic = sigvq->compute(n_threads, resized); + if (semantic.empty()) { + return {}; + } + result.extra_c_crossattns.push_back(std::move(semantic)); + } + return result; + } +}; + struct LTXAVEmbedder : public Conditioner { static constexpr int64_t kHiddenSize = 3840; static constexpr int64_t kNumStates = 49; diff --git a/otherarch/sdcpp/src/conditioning/conditioning_cache.h b/otherarch/sdcpp/src/conditioning/conditioning_cache.h new file mode 100644 index 000000000..6e6cbcaf7 --- /dev/null +++ b/otherarch/sdcpp/src/conditioning/conditioning_cache.h @@ -0,0 +1,107 @@ +#ifndef __SD_CONDITIONING_CONDITIONING_CACHE_H__ +#define __SD_CONDITIONING_CONDITIONING_CACHE_H__ + +#include +#include +#include + +#include "conditioning/conditioner.hpp" + +class ConditioningCache { + struct Entry { + ConditionerParams params; + std::vector> ref_images; + std::vector references; + SDCondition condition; + + Entry(const ConditionerParams& input, const SDCondition& output) + : params(input), condition(output) { + // Request-owned reference pointers must not outlive the request. + if (input.ref_images != nullptr) { + ref_images = *input.ref_images; + params.ref_images = &ref_images; + } + if (input.minimax_h3_references != nullptr) { + references = *input.minimax_h3_references; + params.minimax_h3_references = &references; + } + } + + Entry(const Entry&) = delete; + Entry& operator=(const Entry&) = delete; + }; + + size_t capacity_ = 4; + std::list entries_; + + static bool same_images(const std::vector>& a, + const std::vector>& b) { + return std::equal(a.begin(), a.end(), b.begin(), b.end(), + [](const sd::Tensor& x, const sd::Tensor& y) { + return x.shape() == y.shape() && x.values() == y.values(); + }); + } + + static bool same_params(const ConditionerParams& a, const ConditionerParams& b) { + const auto fields = [](const ConditionerParams& p) { + const auto& r = p.ref_image_params; + return std::tie(p.text, p.clip_skip, p.width, p.height, p.zero_out_masked, + r.pass_to_vlm, r.pass_to_dit, r.ref_index_mode, + r.force_ref_timestep_zero, r.resize_before_vae, r.vae_input_max_pixels, + r.vlm_resize_mode, r.vlm_min_size, r.vlm_max_size, r.resize_vae_to_target); + }; + if (fields(a) != fields(b) || + (a.ref_images == nullptr) != (b.ref_images == nullptr) || + (a.minimax_h3_references == nullptr) != (b.minimax_h3_references == nullptr)) { + return false; + } + if (a.ref_images != nullptr && !same_images(*a.ref_images, *b.ref_images)) { + return false; + } + if (a.minimax_h3_references != nullptr && + !std::equal(a.minimax_h3_references->begin(), a.minimax_h3_references->end(), + b.minimax_h3_references->begin(), b.minimax_h3_references->end(), + [](const MiniMaxH3PresentationItem& x, const MiniMaxH3PresentationItem& y) { + return x.kind == y.kind && x.timestamps == y.timestamps && same_images(x.frames, y.frames); + })) { + return false; + } + return true; + } + +public: + void set_capacity(size_t capacity) { + capacity_ = capacity; + while (entries_.size() > capacity_) { + entries_.pop_back(); + } + } + + void clear() { + entries_.clear(); + } + + SDCondition get(Conditioner& conditioner, int n_threads, const ConditionerParams& params) { + if (capacity_ == 0) { + return conditioner.get_learned_condition(n_threads, params); + } + for (auto it = entries_.begin(); it != entries_.end(); ++it) { + if (same_params(it->params, params)) { + entries_.splice(entries_.begin(), entries_, it); + LOG_INFO("conditioning cache hit"); + return entries_.front().condition; + } + } + auto condition = conditioner.get_learned_condition(n_threads, params); + if (!condition.empty()) { + if (entries_.size() == capacity_) { + entries_.pop_back(); + } + entries_.emplace_front(params, condition); + LOG_VERBOSE("conditioning cache stored (%zu/%zu)", entries_.size(), capacity_); + } + return condition; + } +}; + +#endif // __SD_CONDITIONING_CONDITIONING_CACHE_H__ diff --git a/otherarch/sdcpp/src/core/backend_fit.cpp b/otherarch/sdcpp/src/core/backend_fit.cpp index 85cfc727c..2ffc1cbb7 100644 --- a/otherarch/sdcpp/src/core/backend_fit.cpp +++ b/otherarch/sdcpp/src/core/backend_fit.cpp @@ -478,7 +478,11 @@ namespace sd::backend_fit { return true; } - bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) { + bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling, ggml_status status) { + // Execution failures can leave the device unusable; tiling only helps with allocation failures. + if (status != GGML_STATUS_ALLOC_FAILED) { + return false; + } const char* retry_mode = nullptr; if (prefer_temporal_tiling && !tiling_params.temporal_tiling) { tiling_params.temporal_tiling = true; @@ -498,7 +502,7 @@ namespace sd::backend_fit { return false; } - LOG_WARN("VAE decode failed (likely out of memory); retrying with %s tiling", + LOG_WARN("VAE decode ran out of memory; retrying with %s tiling", retry_mode); return true; } diff --git a/otherarch/sdcpp/src/core/backend_fit.h b/otherarch/sdcpp/src/core/backend_fit.h index 9ef298b3b..0c682a406 100644 --- a/otherarch/sdcpp/src/core/backend_fit.h +++ b/otherarch/sdcpp/src/core/backend_fit.h @@ -16,7 +16,8 @@ namespace sd::backend_fit { std::string& params_spec); bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, - bool prefer_temporal_tiling); + bool prefer_temporal_tiling, + ggml_status status); } // namespace sd::backend_fit diff --git a/otherarch/sdcpp/src/core/ggml_extend.cpp b/otherarch/sdcpp/src/core/ggml_extend.cpp index fd65f41c2..c329a3427 100644 --- a/otherarch/sdcpp/src/core/ggml_extend.cpp +++ b/otherarch/sdcpp/src/core/ggml_extend.cpp @@ -465,7 +465,11 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, int d1, int d2, bool force_prec_f32, - bool direct) { + bool direct, + float scale) { + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, scale); + } if (direct) { int64_t OC = w->ne[3] / IC; int64_t N = x->ne[3] / IC; @@ -502,6 +506,9 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, } } + if (scale != 1.f) { + x = ggml_ext_scale(ctx, x, 1.f / scale); + } if (b != nullptr) { b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]); // [OC, 1, 1, 1] x = ggml_add_inplace(ctx, x, b); @@ -622,7 +629,12 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, ggml_tensor* mask, bool skip_reshape, bool flash_attn, - float kv_scale) { // avoid overflow + float kv_scale, + bool sage_attn, + bool* used_flash_attn) { // avoid overflow + if (used_flash_attn != nullptr) { + *used_flash_attn = false; + } int64_t L_q; int64_t L_k; int64_t C; @@ -713,7 +725,37 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, return out; }; - if (flash_attn) { +#ifndef SD_USE_UPSTREAM_GGML + if (sage_attn && mask == nullptr && d_head > 0 && d_head <= 128) { + auto q_in = ggml_reshape_4d(ctx, ggml_ext_cont(ctx, q->type == GGML_TYPE_F32 ? q : ggml_cast(ctx, q, GGML_TYPE_F32)), d_head, L_q, n_head, N); + auto k_in = ggml_reshape_4d(ctx, ggml_ext_cont(ctx, k->type == GGML_TYPE_F32 ? k : ggml_cast(ctx, k, GGML_TYPE_F32)), d_head, L_k, n_kv_head, N); + auto v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3)); + const int64_t padded_head = d_head <= 64 ? 64 : 128; + if ((padded_head != d_head || kv_scale != 1.0f) && v_in->type != GGML_TYPE_F32) { + v_in = ggml_cast(ctx, v_in, GGML_TYPE_F32); + } + if (padded_head != d_head) { + // Keep the original head's softmax scale when padding for the CUDA kernel. + q_in = ggml_pad(ctx, q_in, padded_head - d_head, 0, 0, 0); + k_in = ggml_pad(ctx, k_in, padded_head - d_head, 0, 0, 0); + v_in = ggml_pad(ctx, v_in, padded_head - d_head, 0, 0, 0); + } + if (kv_scale != 1.0f) { + k_in = ggml_ext_scale(ctx, k_in, kv_scale); + v_in = ggml_ext_scale(ctx, v_in, kv_scale); + } + v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16); + auto out = ggml_sage_attn(ctx, q_in, k_in, v_in, scale / kv_scale, GGML_SAGE_ATTN_AUTO); + if (ggml_backend_supports_op(backend, out)) { + kqv = kv_scale != 1.0f ? ggml_ext_scale(ctx, out, 1.0f / kv_scale) : out; + if (padded_head != d_head) { + kqv = ggml_ext_slice(ctx, kqv, 0, 0, d_head); + } + } + } +#endif + + if (kqv == nullptr && (flash_attn || sage_attn)) { // LOG_VERBOSE("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N); bool can_use_flash_attn = true; if (mask != nullptr) { @@ -724,6 +766,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, if (can_use_flash_attn) { kqv = build_kqv(q, k, v, mask); if (kqv != nullptr) { + if (used_flash_attn != nullptr) { + *used_flash_attn = true; + } kqv = ggml_view_4d(ctx, kqv, d_head, diff --git a/otherarch/sdcpp/src/core/ggml_extend.h b/otherarch/sdcpp/src/core/ggml_extend.h index e843cbee0..5747d625a 100644 --- a/otherarch/sdcpp/src/core/ggml_extend.h +++ b/otherarch/sdcpp/src/core/ggml_extend.h @@ -154,7 +154,8 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, int d1 = 1, int d2 = 1, bool force_prec_f32 = false, - bool direct = false); + bool direct = false, + float scale = 1.f); // w: [OC，IC, KD, 1 * 1] // x: [N, IC, ID, IH*IW] @@ -217,10 +218,12 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, ggml_tensor* k, ggml_tensor* v, int64_t n_head, - ggml_tensor* mask = nullptr, - bool skip_reshape = false, - bool flash_attn = false, - float kv_scale = 1.0f); + ggml_tensor* mask = nullptr, + bool skip_reshape = false, + bool flash_attn = false, + float kv_scale = 1.0f, + bool sage_attn = false, + bool* used_flash_attn = nullptr); ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, ggml_tensor* x, diff --git a/otherarch/sdcpp/src/core/ggml_extend_backend.cpp b/otherarch/sdcpp/src/core/ggml_extend_backend.cpp index f3049dc23..b1bce9847 100644 --- a/otherarch/sdcpp/src/core/ggml_extend_backend.cpp +++ b/otherarch/sdcpp/src/core/ggml_extend_backend.cpp @@ -13,6 +13,7 @@ #endif #include "core/util.h" +#include "ggml-backend-impl.h" #include "ggml-impl.h" #include "stable-diffusion.h" @@ -433,6 +434,24 @@ bool sd_backend_is_cpu(ggml_backend_t backend) { return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU; } +ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, + void* ptr, + size_t size, + size_t max_tensor_size) { + ggml_backend_buffer_t buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size); + if (buffer != nullptr && buffer->context == nullptr) { + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device); + if (reg != nullptr && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) { + // Metal can wrap a failed mapping in a non-null buffer. Its free callback also + // dereferences the missing context, so only release the outer buffer. + buffer->iface.free_buffer = nullptr; + ggml_backend_buffer_free(buffer); + return nullptr; + } + } + return buffer; +} + bool sd_backend_supports_cuda_mma(ggml_backend_t backend) { #ifdef SD_USE_CUDA if (!sd_backend_is(backend, "CUDA")) { diff --git a/otherarch/sdcpp/src/core/ggml_extend_backend.h b/otherarch/sdcpp/src/core/ggml_extend_backend.h index b2ebd7eb2..df1d69240 100644 --- a/otherarch/sdcpp/src/core/ggml_extend_backend.h +++ b/otherarch/sdcpp/src/core/ggml_extend_backend.h @@ -88,6 +88,10 @@ private: bool sd_backend_is(ggml_backend_t backend, const std::string& name); bool sd_backend_is_cpu(ggml_backend_t backend); bool sd_backend_supports_cuda_mma(ggml_backend_t backend); +ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, + void* ptr, + size_t size, + size_t max_tensor_size); ggml_backend_t sd_backend_cpu_init(); bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads); ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend, diff --git a/otherarch/sdcpp/src/core/ggml_runner.cpp b/otherarch/sdcpp/src/core/ggml_runner.cpp index 50173aab7..4c156ed43 100644 --- a/otherarch/sdcpp/src/core/ggml_runner.cpp +++ b/otherarch/sdcpp/src/core/ggml_runner.cpp @@ -21,11 +21,12 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx, ggml_tensor* mask, bool skip_reshape, bool flash_attn, - float kv_scale) { + float kv_scale, + bool* used_flash_attn) { if (ctx->attn_scale > 0.f) { kv_scale = ctx->attn_scale; } - return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale); + return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled, used_flash_attn); } void GGMLRunner::alloc_params_ctx() { @@ -515,11 +516,13 @@ GGMLRunner::~GGMLRunner() { free_params_ctx(); } -GGMLRunnerContext GGMLRunner::get_context() { +GGMLRunnerContext GGMLRunner::get_context(ggml_cgraph* graph) { GGMLRunnerContext runner_ctx; runner_ctx.ggml_ctx = compute_ctx; + runner_ctx.graph = graph; runner_ctx.backend = runtime_backend; runner_ctx.flash_attn_enabled = flash_attn_enabled; + runner_ctx.sage_attn_enabled = sage_attn_enabled; runner_ctx.linear_scale = linear_scale; runner_ctx.attn_scale = attn_scale; runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled; @@ -531,8 +534,8 @@ GGMLRunnerContext GGMLRunner::get_context() { runner_ctx.get_cache_tensor = [this](const std::string& name) { return this->get_cache_tensor_by_name(name); }; - runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) { - this->cache(name, tensor); + runner_ctx.cache_tensor = [this, graph](const std::string& name, ggml_tensor* tensor) { + this->cache(name, tensor, graph); }; runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) { this->set_backend_tensor_data(tensor, data); @@ -574,7 +577,7 @@ ggml_tensor* GGMLRunner::to_backend(ggml_tensor* tensor) { } } -void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) { +void GGMLRunner::cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph) { if (tensor != nullptr && tensor->view_src != nullptr) { tensor = ggml_cont(compute_ctx, tensor); } @@ -582,6 +585,10 @@ void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) { ggml_set_output(tensor); } cache_.stage(name, tensor); + if (graph != nullptr && tensor != nullptr) { + // Schedule the cache output here so its source can be reused before graph end. + ggml_build_forward_expand(graph, tensor); + } } std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, @@ -589,6 +596,7 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, bool auto_runner_end, bool no_return, const std::function& read_outputs) { + last_compute_status_ = GGML_STATUS_FAILED; if (graph_active_) { LOG_ERROR("%s does not support reentrant graph execution", get_desc().c_str()); return std::nullopt; @@ -612,7 +620,9 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, GGMLRunner& runner; const bool& success; ~GraphEndGuard() { - runner.workspace_.segment_end(); + if (!runner.workspace_.segment_end()) { + runner.last_compute_status_ = GGML_STATUS_FAILED; + } runner.cache_.graph_end(false); runner.cut_cache_.clear(); runner.free_compute_ctx(); @@ -640,7 +650,12 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, std::optional> output; try { output = execute_graph(graph, n_threads, no_return, read_outputs); + } catch (const std::bad_alloc&) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + LOG_ERROR("%s graph allocation failed", get_desc().c_str()); + return std::nullopt; } catch (const std::exception& error) { + last_compute_status_ = GGML_STATUS_FAILED; LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(), ggml_backend_name(runtime_backend), error.what()); return std::nullopt; @@ -648,6 +663,7 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, success = output.has_value(); if (success) { cache_.graph_end(true); + last_compute_status_ = GGML_STATUS_SUCCESS; } return output; } @@ -765,6 +781,7 @@ bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) { } workspace_.synchronize(); if (status != GGML_STATUS_SUCCESS) { + last_compute_status_ = status; LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); return false; } @@ -817,13 +834,23 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n const auto& cached_plan = resolve_graph_cut_plan(graph); const auto full_measurement = measure(graph, cached_plan.compute_buffer_size); if (full_measurement.buffers.empty()) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; return std::nullopt; } + auto fits_monolithic = [&]() { + // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits. + constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL; + auto requests = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)); + for (auto& request : requests) { + request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom); + } + return fits(requests, params); + }; auto manager = residency_manager.lock(); const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) && manager != nullptr && manager->segmented_compute_enabled() && cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1 && - !fits(memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)), params); + !fits_monolithic(); ggml_graph_cut::Plan monolithic_plan; if (!segmented) { monolithic_plan.segments.emplace_back(); @@ -887,7 +914,9 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n SegmentGraphBindings& bindings; ggml_context* context; ~SegmentCleanup() { - runner.workspace_.segment_end(); + if (!runner.workspace_.segment_end()) { + runner.last_compute_status_ = GGML_STATUS_FAILED; + } bindings.restore(); weights.segment_end(); ggml_free(context); @@ -897,6 +926,7 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement; if (!workspace_.prepare(measurement)) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; return fail_segment("workspace preparation"); } const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment); @@ -911,7 +941,11 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n sync_runtime_residency(); requests = memory_requests(measurement.buffers, new_cache_bytes); } - return weights.ensure_segment_capacity(index, requests); + const bool ready = weights.ensure_segment_capacity(index, requests); + if (!ready && manager != nullptr) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + } + return ready; }; if (!weights.segment_start(index, ensure_capacity)) { return fail_segment("weight preparation"); @@ -920,12 +954,17 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n if (!workspace_.measurement_matches(segment_graph, measurement)) { measurement = measure(segment_graph, segment.compute_buffer_size); } - if (!workspace_.prepare(measurement) || !ensure_capacity()) { + if (!workspace_.prepare(measurement)) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return fail_segment("workspace preparation"); + } + if (!ensure_capacity()) { return fail_segment("workspace capacity check"); } if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) { pin_multi_device_nodes(scheduler, current); })) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; return fail_segment("workspace allocation"); } for (const auto& size : measurement.buffers) { @@ -944,10 +983,16 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n } LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(), index + 1, plan.segments.size(), segment.group_name.c_str()); - if (!execute_segment(segment_graph, n_threads) || - !cache_.capture(segment_graph) || - !cut_cache_.capture(graph, segment, get_desc().c_str())) { - return fail_segment("execution or output caching"); + if (!execute_segment(segment_graph, n_threads)) { + return fail_segment("execution"); + } + auto cache_status = cache_.capture(segment_graph); + if (cache_status == GGML_STATUS_SUCCESS) { + cache_status = cut_cache_.capture(graph, segment, get_desc().c_str()); + } + if (cache_status != GGML_STATUS_SUCCESS) { + last_compute_status_ = cache_status; + return fail_segment("output caching"); } sync_runtime_residency(); if (last) { @@ -963,6 +1008,7 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n } } if (!workspace_.segment_end()) { + last_compute_status_ = GGML_STATUS_FAILED; return fail_segment("workspace synchronization"); } // Final outputs and their callbacks may still be views of consumed cuts. diff --git a/otherarch/sdcpp/src/core/ggml_runner.h b/otherarch/sdcpp/src/core/ggml_runner.h index 2152ec8a2..02630c0d6 100644 --- a/otherarch/sdcpp/src/core/ggml_runner.h +++ b/otherarch/sdcpp/src/core/ggml_runner.h @@ -67,7 +67,9 @@ struct WeightAdapter { struct GGMLRunnerContext { ggml_backend_t backend = nullptr; ggml_context* ggml_ctx = nullptr; + ggml_cgraph* graph = nullptr; bool flash_attn_enabled = false; + bool sage_attn_enabled = false; float linear_scale = 0.f; float attn_scale = 0.f; bool conv2d_direct_enabled = false; @@ -101,6 +103,12 @@ struct GGMLRunnerContext { return get_cache_tensor(name); } + void expand_graph(ggml_tensor* tensor) const { + if (graph != nullptr && tensor != nullptr) { + ggml_build_forward_expand(graph, tensor); + } + } + void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const { if (!cache_tensor || tensor == nullptr) { return; @@ -121,15 +129,17 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx, ggml_tensor* k, ggml_tensor* v, int64_t n_head, - ggml_tensor* mask = nullptr, - bool skip_reshape = false, - bool flash_attn = false, - float kv_scale = 1.f); + ggml_tensor* mask = nullptr, + bool skip_reshape = false, + bool flash_attn = false, + float kv_scale = 1.f, + bool* used_flash_attn = nullptr); struct GGMLRunner { private: std::map logged_compute_bytes_; - size_t logged_segment_count_ = 0; + size_t logged_segment_count_ = 0; + ggml_status last_compute_status_ = GGML_STATUS_SUCCESS; sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes); std::vector memory_requests(const std::vector& sizes, @@ -176,6 +186,7 @@ protected: const std::string final_result_name = "ggml_runner_final_result_tensor"; bool flash_attn_enabled = false; + bool sage_attn_enabled = false; float linear_scale = 0.f; float attn_scale = 0.f; bool conv2d_direct_enabled = false; @@ -286,7 +297,8 @@ public: virtual ~GGMLRunner(); - virtual GGMLRunnerContext get_context(); + // Binding a graph schedules cache outputs at registration instead of graph end. + virtual GGMLRunnerContext get_context(ggml_cgraph* graph = nullptr); void reset_compute_ctx(); @@ -321,7 +333,7 @@ public: ggml_tensor* to_backend(ggml_tensor* tensor); - void cache(const std::string name, ggml_tensor* tensor); + void cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph = nullptr); ggml_tensor* get_cache_tensor_by_name(const std::string& name) { return cache_.get(name); @@ -333,10 +345,20 @@ public: bool no_return = false, const std::function& read_outputs = {}); + ggml_status last_compute_status() const { return last_compute_status_; } + void set_flash_attention_enabled(bool enabled) { flash_attn_enabled = enabled; } + void set_sage_attention_enabled(bool enabled) { + if (sage_attn_enabled != enabled) { + free_cache_ctx_and_buffer(); + graph_cut_plan_cache_.graph_cut_plans.clear(); + sage_attn_enabled = enabled; + } + } + void set_scale_overrides(float linear_scale, float attn_scale) { this->linear_scale = linear_scale; this->attn_scale = attn_scale; diff --git a/otherarch/sdcpp/src/core/runner_cache.cpp b/otherarch/sdcpp/src/core/runner_cache.cpp index 095621068..b640bb45a 100644 --- a/otherarch/sdcpp/src/core/runner_cache.cpp +++ b/otherarch/sdcpp/src/core/runner_cache.cpp @@ -26,10 +26,13 @@ namespace sd { std::unique_ptr CachedTensor::copy(ggml_backend_t backend, const std::string& name, - ggml_tensor* source) { + ggml_tensor* source, + ggml_status& status) { + status = GGML_STATUS_FAILED; if (ggml_graph_cut::tensor_buffer(source) == nullptr) { return nullptr; } + status = GGML_STATUS_ALLOC_FAILED; auto entry = std::make_unique(); entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true}); if (entry->context == nullptr) { @@ -50,6 +53,7 @@ namespace sd { } else { ggml_backend_tensor_copy(source, entry->tensor); } + status = GGML_STATUS_SUCCESS; return entry; } @@ -106,9 +110,9 @@ namespace sd { return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending; } - bool RunnerCache::capture(ggml_cgraph* graph) { + ggml_status RunnerCache::capture(ggml_cgraph* graph) { if (outputs_.empty()) { - return true; + return GGML_STATUS_SUCCESS; } const auto tensors = cache_graph_tensors(graph); for (const auto& output : outputs_) { @@ -116,14 +120,15 @@ namespace sd { continue; } GGML_ASSERT(ggml_is_contiguous(output.second)); - auto entry = CachedTensor::copy(backend_, output.first, output.second); + ggml_status status; + auto entry = CachedTensor::copy(backend_, output.first, output.second, status); if (entry == nullptr) { - return false; + return status; } pending_[output.first] = std::move(entry); } ggml_backend_synchronize(backend_); - return true; + return GGML_STATUS_SUCCESS; } void RunnerCache::graph_end(bool success) { @@ -180,9 +185,9 @@ namespace sd { } } - bool GraphCutTensorCache::capture(ggml_cgraph* graph, - const ggml_graph_cut::Segment& segment, - const char* log_desc) { + ggml_status GraphCutTensorCache::capture(ggml_cgraph* graph, + const ggml_graph_cut::Segment& segment, + const char* log_desc) { size_t copied_bytes = 0; size_t copied_count = 0; for (int index : segment.output_node_indices) { @@ -191,10 +196,11 @@ namespace sd { !segment.future_cut_names.count(output->name)) { continue; } - auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output)); + ggml_status status; + auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output), status); if (entry == nullptr) { LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name); - return false; + return status; } const size_t size = ggml_backend_buffer_get_size(entry->buffer); copied_bytes = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size; @@ -206,6 +212,6 @@ namespace sd { LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)", log_desc, copied_bytes / (1024.f * 1024.f), copied_count); } - return true; + return GGML_STATUS_SUCCESS; } } diff --git a/otherarch/sdcpp/src/core/runner_cache.h b/otherarch/sdcpp/src/core/runner_cache.h index bd331a4bb..f1b5cf91c 100644 --- a/otherarch/sdcpp/src/core/runner_cache.h +++ b/otherarch/sdcpp/src/core/runner_cache.h @@ -20,7 +20,8 @@ namespace sd { ~CachedTensor(); static std::unique_ptr copy(ggml_backend_t backend, const std::string& name, - ggml_tensor* source); + ggml_tensor* source, + ggml_status& status); }; using CachedTensors = std::map>; @@ -41,7 +42,8 @@ namespace sd { const std::map& outputs() const { return outputs_; } size_t pending_bytes(ggml_cgraph* graph) const; size_t resident_bytes(ggml_backend_dev_t device) const; - bool capture(ggml_cgraph* graph); + bool empty() const { return committed_.empty(); } + ggml_status capture(ggml_cgraph* graph); void graph_end(bool success); void clear(); }; @@ -57,7 +59,7 @@ namespace sd { size_t resident_bytes(ggml_backend_dev_t device) const; size_t estimate_output_bytes(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment) const; - bool capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc); + ggml_status capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc); void prune(const std::unordered_set& keep_names); void clear() { tensors_.clear(); } }; diff --git a/otherarch/sdcpp/src/detailer.cpp b/otherarch/sdcpp/src/detailer.cpp index bbf9c6880..72d01dc6f 100644 --- a/otherarch/sdcpp/src/detailer.cpp +++ b/otherarch/sdcpp/src/detailer.cpp @@ -970,6 +970,7 @@ bool adetail_image(adetailer_ctx_t* context, generation.pm_params = {}; generation.pulid_params = {}; generation.hires.enabled = false; + generation.image_preprocess = {}; if (params.steps > 0) { generation.sample_params.sample_steps = params.steps; generation.sample_params.custom_sigmas = nullptr; diff --git a/otherarch/sdcpp/src/model.h b/otherarch/sdcpp/src/model.h index 90044ecb5..fcba8620b 100644 --- a/otherarch/sdcpp/src/model.h +++ b/otherarch/sdcpp/src/model.h @@ -60,6 +60,7 @@ enum SDVersion { VERSION_KREA2, VERSION_MAGE_FLOW, VERSION_SENSENOVA_U1_5, + VERSION_LLADA_IMAGE, VERSION_ESRGAN, VERSION_COUNT, }; @@ -173,6 +174,13 @@ static inline bool sd_version_is_z_image(SDVersion version) { return false; } +static inline bool sd_version_is_llada_image(SDVersion version) { + if (version == VERSION_LLADA_IMAGE) { + return true; + } + return false; +} + static inline bool sd_version_is_boogu_image(SDVersion version) { if (version == VERSION_BOOGU_IMAGE) { return true; @@ -244,6 +252,14 @@ static inline bool sd_version_is_sensenova_u1(SDVersion version) { return version == VERSION_SENSENOVA_U1_5; } +static inline bool sd_version_supports_video_generation(SDVersion version) { + return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version); +} + +static inline bool sd_version_supports_image_generation(SDVersion version) { + return !sd_version_supports_video_generation(version); +} + static inline bool sd_version_uses_flux_vae(SDVersion version) { if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) { return true; @@ -252,7 +268,7 @@ static inline bool sd_version_uses_flux_vae(SDVersion version) { } static inline bool sd_version_uses_flux2_vae(SDVersion version) { - if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version)) { + if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version) || sd_version_is_llada_image(version)) { return true; } return false; @@ -293,6 +309,7 @@ static inline bool sd_version_is_dit(SDVersion version) { version == VERSION_HIDREAM_O1 || sd_version_is_anima(version) || sd_version_is_z_image(version) || + sd_version_is_llada_image(version) || sd_version_is_boogu_image(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || diff --git a/otherarch/sdcpp/src/model/adapter/lora.hpp b/otherarch/sdcpp/src/model/adapter/lora.hpp index 636d05189..1a5558b1d 100644 --- a/otherarch/sdcpp/src/model/adapter/lora.hpp +++ b/otherarch/sdcpp/src/model/adapter/lora.hpp @@ -163,6 +163,27 @@ struct LoraModel : public GGMLRunner { lora_tensors = std::move(new_lora_tensors); } + + std::unordered_map new_lora_tensors; + for (const auto& [old_name, tensor] : lora_tensors) { + std::string new_name = old_name; + if (starts_with(old_name, "lora.model.diffusion_model.transformer_blocks.")) { + // Qwen Image 2.1 stores the gate before the projection in fused MLP weights. + for (const auto& suffix : {std::string(".img_mlp.gate_layer.weight."), std::string(".img_mlp.proj.weight.")}) { + size_t pos = old_name.find(suffix); + if (pos == std::string::npos) { + continue; + } + std::string fused_name = old_name.substr(5, pos - 5) + ".img_mlp.gate_up.weight"; + if (model_tensor_names.find(fused_name) != model_tensor_names.end()) { + new_name = "lora." + fused_name + (suffix == ".img_mlp.proj.weight." ? ".1." : ".") + old_name.substr(pos + suffix.size()); + } + break; + } + } + new_lora_tensors[new_name] = tensor; + } + lora_tensors = std::move(new_lora_tensors); } ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) { diff --git a/otherarch/sdcpp/src/model/common/ggml_block.hpp b/otherarch/sdcpp/src/model/common/ggml_block.hpp index a06fe7638..991b88f28 100644 --- a/otherarch/sdcpp/src/model/common/ggml_block.hpp +++ b/otherarch/sdcpp/src/model/common/ggml_block.hpp @@ -839,21 +839,30 @@ class RMSNorm : public UnaryBlock { protected: int64_t hidden_size; float eps; + bool elementwise_affine; std::string prefix; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override { - this->prefix = prefix; + this->prefix = prefix; + if (!elementwise_affine) { + return; + } enum ggml_type wtype = GGML_TYPE_F32; params["weight"] = ggml_new_tensor_1d(ctx, wtype, hidden_size); } public: RMSNorm(int64_t hidden_size, - float eps = 1e-06f) + float eps = 1e-06f, + bool elementwise_affine = true) : hidden_size(hidden_size), - eps(eps) {} + eps(eps), + elementwise_affine(elementwise_affine) {} ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + if (!elementwise_affine) { + return ggml_rms_norm(ctx->ggml_ctx, x, eps); + } ggml_tensor* w = params["weight"]; if (ctx->weight_adapter) { w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); diff --git a/otherarch/sdcpp/src/model/common/rope.hpp b/otherarch/sdcpp/src/model/common/rope.hpp index 778bbf426..ec8224559 100644 --- a/otherarch/sdcpp/src/model/common/rope.hpp +++ b/otherarch/sdcpp/src/model/common/rope.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include "core/ggml_extend.h" #include "core/ggml_runner.h" @@ -16,6 +17,45 @@ namespace Rope { ErnieImage, }; + struct SpatialRegion { + size_t begin; + size_t count; + float height_period; + float width_period; + int height_axis = 1; + int width_axis = 2; + }; + + struct PositionLayout { + // Token ranges are relative to one batch item. + std::vector images; + size_t token_count = 0; + + void append_tokens(size_t count) { + token_count += count; + } + + void append_image(int height, int width, int frames = 1, float height_step = 1.f, float width_step = 1.f) { + size_t count = static_cast(height) * width * frames; + images.push_back({token_count, count, height * height_step, width * width_step}); + append_tokens(count); + } + }; + + struct Frequency { + size_t axis; + float omega; + }; + + struct Embedding { + std::vector values; + std::vector> ids; + PositionLayout positions; + std::vector frequencies; + EmbedNDLayout layout = EmbedNDLayout::Matrix; + int batch_size = 1; + }; + enum class RefIndexMode { FIXED, INCREASE, @@ -56,40 +96,25 @@ namespace Rope { return flat_vec; } - __STATIC_INLINE__ std::vector> rope(const std::vector& pos, - int dim, - float theta, - const std::vector& axis_wrap_dims = {}) { + __STATIC_INLINE__ std::vector rope_frequencies(int dim, float theta) { assert(dim % 2 == 0); - int half_dim = dim / 2; - + int half_dim = dim / 2; std::vector scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim); - std::vector omega(half_dim); for (int i = 0; i < half_dim; ++i) { omega[i] = 1.0f / ::powf(1.f * theta, scale[i]); } + return omega; + } + __STATIC_INLINE__ std::vector> rope(const std::vector& pos, + const std::vector& omega) { + int half_dim = static_cast(omega.size()); size_t pos_size = pos.size(); std::vector> out(pos_size, std::vector(half_dim)); for (size_t i = 0; i < pos_size; ++i) { for (size_t j = 0; j < half_dim; ++j) { float angle = pos[i] * omega[j]; - if (!axis_wrap_dims.empty()) { - size_t wrap_size = axis_wrap_dims.size(); - // mod batch size since we only store this for one item in the batch - size_t wrap_idx = wrap_size > 0 ? (i % wrap_size) : 0; - int wrap_dim = axis_wrap_dims[wrap_idx]; - if (wrap_dim > 0) { - constexpr float TWO_PI = 6.28318530717958647692f; - float cycles = omega[j] * wrap_dim / TWO_PI; - // closest periodic harmonic, necessary to ensure things neatly tile - // without this round, things don't tile at the boundaries and you end up - // with the model knowing what is "center" - float rounded = std::round(cycles); - angle = pos[i] * TWO_PI * rounded / wrap_dim; - } - } out[i][j] = angle; } @@ -108,6 +133,12 @@ namespace Rope { return result; } + __STATIC_INLINE__ std::vector> rope(const std::vector& pos, + int dim, + float theta) { + return rope(pos, rope_frequencies(dim, theta)); + } + // Generate IDs for image patches and text __STATIC_INLINE__ std::vector> gen_flux_txt_ids(int bs, int context_len, int axes_dim_num, std::set arange_dims) { auto txt_ids = std::vector>(bs * context_len, std::vector(axes_dim_num, 0.0f)); @@ -136,12 +167,16 @@ namespace Rope { int patch_size, int bs, int axes_dim_num, - int index = 0, - int h_offset = 0, - int w_offset = 0, - bool scale_rope = false) { + int index = 0, + int h_offset = 0, + int w_offset = 0, + bool scale_rope = false, + PositionLayout* layout = nullptr) { int h_len = (h + (patch_size / 2)) / patch_size; int w_len = (w + (patch_size / 2)) / patch_size; + if (layout) { + layout->append_image(h_len, w_len); + } std::vector> img_ids(h_len * w_len, std::vector(axes_dim_num, 0.0)); int h_start = h_offset; @@ -192,8 +227,8 @@ namespace Rope { int bs, const std::vector& axis_thetas, const std::vector& axes_dim, - const std::vector>& wrap_dims = {}, - EmbedNDLayout layout = EmbedNDLayout::Matrix) { + EmbedNDLayout layout = EmbedNDLayout::Matrix, + std::vector* frequencies = nullptr) { std::vector> trans_ids = transpose(ids); size_t pos_len = ids.size() / bs; size_t num_axes = axes_dim.size(); @@ -205,19 +240,25 @@ namespace Rope { for (int d : axes_dim) emb_dim += d / 2; + if (frequencies) { + frequencies->clear(); + frequencies->reserve(emb_dim); + } std::vector> emb(bs * pos_len, std::vector(emb_dim * 2 * 2, 0.0)); size_t offset = 0; for (size_t i = 0; i < num_axes; ++i) { - std::vector axis_wrap_dims; - if (!wrap_dims.empty() && i < (int)wrap_dims.size()) { - axis_wrap_dims = wrap_dims[i]; - } float axis_theta = 10000.0f; if (!axis_thetas.empty()) { axis_theta = axis_thetas[std::min(i, axis_thetas.size() - 1)]; } + auto omega = rope_frequencies(axes_dim[i], axis_theta); + if (frequencies) { + for (float frequency : omega) { + frequencies->push_back({i, frequency}); + } + } std::vector> rope_emb = - rope(trans_ids[i], axes_dim[i], axis_theta, axis_wrap_dims); // [bs*pos_len, axes_dim[i]/2 * 2 * 2] + rope(trans_ids[i], omega); // [bs*pos_len, axes_dim[i]/2 * 2 * 2] for (int b = 0; b < bs; ++b) { for (int j = 0; j < pos_len; ++j) { for (int k = 0; k < rope_emb[0].size(); ++k) { @@ -253,10 +294,10 @@ namespace Rope { int bs, float theta, const std::vector& axes_dim, - const std::vector>& wrap_dims = {}, - EmbedNDLayout layout = EmbedNDLayout::Matrix) { + EmbedNDLayout layout = EmbedNDLayout::Matrix, + std::vector* frequencies = nullptr) { std::vector axis_thetas(axes_dim.size(), theta); - return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout); + return embed_nd(ids, bs, axis_thetas, axes_dim, layout, frequencies); } __STATIC_INLINE__ std::vector embed_interleaved_mrope(const std::vector>& ids, @@ -264,7 +305,7 @@ namespace Rope { float theta, int head_dim, const std::vector& mrope_section, - const std::vector>& axis_wrap_dims = {}) { + std::vector* frequencies = nullptr) { GGML_ASSERT(bs > 0); GGML_ASSERT(head_dim % 2 == 0); GGML_ASSERT(mrope_section.size() >= 3); @@ -273,20 +314,26 @@ namespace Rope { size_t pos_len = ids.size() / bs; int half_dim = head_dim / 2; + auto omega = rope_frequencies(head_dim, theta); + if (frequencies) { + frequencies->clear(); + for (float frequency : omega) { + frequencies->push_back({0, frequency}); + } + } std::vector>> axis_embs; axis_embs.reserve(3); for (int axis = 0; axis < 3; ++axis) { - std::vector axis_wrap; - if (axis < static_cast(axis_wrap_dims.size())) { - axis_wrap = axis_wrap_dims[axis]; - } - axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap)); + axis_embs.push_back(rope(trans_ids[axis], omega)); } std::vector> emb = axis_embs[0]; for (int axis = 1; axis < 3; ++axis) { int length = std::min(mrope_section[axis] * 3, half_dim); for (int freq_idx = axis; freq_idx < length; freq_idx += 3) { + if (frequencies) { + (*frequencies)[freq_idx].axis = axis; + } for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) { for (int k = 0; k < 4; ++k) { emb[pos_idx][4 * freq_idx + k] = axis_embs[axis][pos_idx][4 * freq_idx + k]; @@ -298,13 +345,13 @@ namespace Rope { return flatten(emb); } - __STATIC_INLINE__ std::vector embed_2d_interleaved(int height, - int width, - int dim, - float theta = 10000.f, - float scale = 16.f, - int ref_grid_h = 0, - int ref_grid_w = 0) { + __STATIC_INLINE__ Embedding embed_2d_interleaved(int height, + int width, + int dim, + float theta = 10000.f, + float scale = 16.f, + int ref_grid_h = 0, + int ref_grid_w = 0) { assert(dim % 4 == 0); int half_dim = dim / 2; int dim_axis = dim / 2; @@ -318,6 +365,10 @@ namespace Rope { w_ntk = std::pow(static_cast(width) / static_cast(ref_grid_w), power); } + Embedding result; + result.positions.append_image(height, width, 1, + height > 1 ? scale / (height - 1) : 1.f, + width > 1 ? scale / (width - 1) : 1.f); std::vector x_pos; std::vector y_pos; x_pos.reserve(static_cast(height) * width); @@ -326,13 +377,20 @@ namespace Rope { float y = height == 1 ? 0.f : scale * static_cast(iy) / static_cast(height - 1); for (int ix = 0; ix < width; ++ix) { float x = width == 1 ? 0.f : scale * static_cast(ix) / static_cast(width - 1); + result.ids.push_back({0.f, y, x}); x_pos.push_back(x); y_pos.push_back(y); } } - auto x_emb = rope(x_pos, dim_axis, theta * w_ntk); - auto y_emb = rope(y_pos, dim_axis, theta * h_ntk); + auto x_freq = rope_frequencies(dim_axis, theta * w_ntk); + auto y_freq = rope_frequencies(dim_axis, theta * h_ntk); + auto x_emb = rope(x_pos, x_freq); + auto y_emb = rope(y_pos, y_freq); + for (int i = 0; i < axis_half_dim; ++i) { + result.frequencies.push_back({2, x_freq[i]}); + result.frequencies.push_back({1, y_freq[i]}); + } std::vector out(static_cast(height) * width * half_dim * 4); for (int pos = 0; pos < height * width; ++pos) { @@ -348,7 +406,8 @@ namespace Rope { } } } - return out; + result.values = std::move(out); + return result; } __STATIC_INLINE__ std::vector> gen_refs_ids(int patch_size, @@ -359,7 +418,8 @@ namespace Rope { RefIndexMode ref_index_mode, float ref_index_scale, bool scale_rope, - int base_offset = 0) { + int base_offset = 0, + PositionLayout* layout = nullptr) { std::vector> ids; int curr_h_offset = 0; int curr_w_offset = 0; @@ -386,7 +446,8 @@ namespace Rope { static_cast(index * ref_index_scale), h_offset + base_offset, w_offset + base_offset, - scale_rope); + scale_rope, + layout); ids = concat_ids(ids, ref_ids, bs); if (ref_index_mode == RefIndexMode::INCREASE) { @@ -409,88 +470,53 @@ namespace Rope { const std::vector& ref_latents, RefIndexMode ref_index_mode, float ref_index_scale, - bool is_longcat) { + bool is_longcat, + PositionLayout* layout = nullptr) { + if (layout) { + layout->append_tokens(context_len); + } int x_index = is_longcat ? 1 : 0; auto txt_ids = is_longcat ? gen_longcat_txt_ids(bs, context_len, axes_dim_num) : gen_flux_txt_ids(bs, context_len, axes_dim_num, txt_arange_dims); int offset = is_longcat ? context_len : 0; - auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset); + auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset, false, layout); auto ids = concat_ids(txt_ids, img_ids, bs); if (ref_latents.size() > 0) { - auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset); + auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset, layout); ids = concat_ids(ids, refs_ids, bs); } return ids; } // Generate flux positional embeddings - __STATIC_INLINE__ std::vector gen_flux_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - std::set txt_arange_dims, - const std::vector& ref_latents, - RefIndexMode ref_index_mode, - float ref_index_scale, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim, - bool is_longcat) { - std::vector> ids = gen_flux_ids(h, - w, - patch_size, - bs, - static_cast(axes_dim.size()), - context_len, - txt_arange_dims, - ref_latents, - ref_index_mode, - ref_index_scale, - is_longcat); - std::vector> wrap_dims; - if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { - int h_len = (h + (patch_size / 2)) / patch_size; - int w_len = (w + (patch_size / 2)) / patch_size; - if (h_len > 0 && w_len > 0) { - size_t pos_len = ids.size() / bs; - wrap_dims.assign(axes_dim.size(), std::vector(pos_len, 0)); - size_t cursor = context_len; // text first - const size_t img_tokens = static_cast(h_len) * static_cast(w_len); - for (size_t token_i = 0; token_i < img_tokens; ++token_i) { - if (circular_h) { - wrap_dims[1][cursor + token_i] = h_len; - } - if (circular_w) { - wrap_dims[2][cursor + token_i] = w_len; - } - } - cursor += img_tokens; - // reference latents - for (ggml_tensor* ref : ref_latents) { - if (ref == nullptr) { - continue; - } - int ref_h = static_cast(ref->ne[1]); - int ref_w = static_cast(ref->ne[0]); - int ref_h_l = (ref_h + (patch_size / 2)) / patch_size; - int ref_w_l = (ref_w + (patch_size / 2)) / patch_size; - size_t ref_tokens = static_cast(ref_h_l) * static_cast(ref_w_l); - for (size_t token_i = 0; token_i < ref_tokens; ++token_i) { - if (circular_h) { - wrap_dims[1][cursor + token_i] = ref_h_l; - } - if (circular_w) { - wrap_dims[2][cursor + token_i] = ref_w_l; - } - } - cursor += ref_tokens; - } - } - } - return embed_nd(ids, bs, static_cast(theta), axes_dim, wrap_dims); + __STATIC_INLINE__ Embedding gen_flux_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + std::set txt_arange_dims, + const std::vector& ref_latents, + RefIndexMode ref_index_mode, + float ref_index_scale, + int theta, + const std::vector& axes_dim, + bool is_longcat) { + Embedding result; + result.batch_size = bs; + result.ids = gen_flux_ids(h, + w, + patch_size, + bs, + static_cast(axes_dim.size()), + context_len, + txt_arange_dims, + ref_latents, + ref_index_mode, + ref_index_scale, + is_longcat, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } __STATIC_INLINE__ std::vector> gen_vid_ids(int t, @@ -500,14 +526,18 @@ namespace Rope { int ph, int pw, int bs, - int t_offset = 0, - int h_offset = 0, - int w_offset = 0, - bool scale_rope = false) { + int t_offset = 0, + int h_offset = 0, + int w_offset = 0, + bool scale_rope = false, + PositionLayout* layout = nullptr) { int t_len = (t + (pt / 2)) / pt; int h_len = (h + (ph / 2)) / ph; int w_len = (w + (pw / 2)) / pw; + if (layout) { + layout->append_image(h_len, w_len, t_len); + } std::vector> vid_ids(t_len * h_len * w_len, std::vector(3, 0.0)); if (scale_rope) { @@ -573,7 +603,11 @@ namespace Rope { int bs, int context_len, const std::vector& ref_latents, - RefIndexMode ref_index_mode) { + RefIndexMode ref_index_mode, + PositionLayout* layout = nullptr) { + if (layout) { + layout->append_tokens(context_len); + } int h_len = (h + (patch_size / 2)) / patch_size; int w_len = (w + (patch_size / 2)) / patch_size; int txt_id_start = std::max(h_len, w_len) / 2; @@ -585,90 +619,49 @@ namespace Rope { } } int axes_dim_num = 3; - auto img_ids = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true); + auto img_ids = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true, layout); auto ids = concat_ids(txt_ids_repeated, img_ids, bs); if (ref_latents.size() > 0) { int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1; - auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true); + auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true, 0, layout); ids = concat_ids(ids, refs_ids, bs); } return ids; } // Generate qwen_image positional embeddings - __STATIC_INLINE__ std::vector gen_qwen_image_pe(int t, - int h, - int w, - int patch_size, - int bs, - int context_len, - const std::vector& ref_latents, - RefIndexMode ref_index_mode, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim) { - std::vector> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode); - std::vector> wrap_dims; - // This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles - if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { - int pad_h = (patch_size - (h % patch_size)) % patch_size; - int pad_w = (patch_size - (w % patch_size)) % patch_size; - int h_len = (h + pad_h) / patch_size; - int w_len = (w + pad_w) / patch_size; - if (h_len > 0 && w_len > 0) { - const size_t total_tokens = ids.size(); - // Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic. - wrap_dims.assign(axes_dim.size(), std::vector(total_tokens / bs, 0)); - size_t cursor = context_len; // ignore text tokens - const size_t img_tokens = static_cast(t) * static_cast(h_len) * static_cast(w_len); - for (size_t token_i = 0; token_i < img_tokens; ++token_i) { - if (circular_h) { - wrap_dims[1][cursor + token_i] = h_len; - } - if (circular_w) { - wrap_dims[2][cursor + token_i] = w_len; - } - } - cursor += img_tokens; - // For each reference image, store wrap sizes as well - for (ggml_tensor* ref : ref_latents) { - if (ref == nullptr) { - continue; - } - int ref_h = static_cast(ref->ne[1]); - int ref_w = static_cast(ref->ne[0]); - int ref_pad_h = (patch_size - (ref_h % patch_size)) % patch_size; - int ref_pad_w = (patch_size - (ref_w % patch_size)) % patch_size; - int ref_h_len = (ref_h + ref_pad_h) / patch_size; - int ref_w_len = (ref_w + ref_pad_w) / patch_size; - size_t ref_n_tokens = static_cast(ref_h_len) * static_cast(ref_w_len); - for (size_t token_i = 0; token_i < ref_n_tokens; ++token_i) { - if (circular_h) { - wrap_dims[1][cursor + token_i] = ref_h_len; - } - if (circular_w) { - wrap_dims[2][cursor + token_i] = ref_w_len; - } - } - cursor += ref_n_tokens; - } - } - } - return embed_nd(ids, bs, static_cast(theta), axes_dim, wrap_dims); + __STATIC_INLINE__ Embedding gen_qwen_image_pe(int t, + int h, + int w, + int patch_size, + int bs, + int context_len, + const std::vector& ref_latents, + RefIndexMode ref_index_mode, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } - __STATIC_INLINE__ std::vector gen_mage_flow_pe(int h, - int w, - int bs, - int context_len, - const std::vector& ref_latents, - int theta, - const std::vector& axes_dim) { + __STATIC_INLINE__ Embedding gen_mage_flow_pe(int h, + int w, + int bs, + int context_len, + const std::vector& ref_latents, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.positions.append_tokens(context_len); const int axes_dim_num = static_cast(axes_dim.size()); - auto make_image_ids = [=](int image_h, int image_w, int image_index) { + auto make_image_ids = [=, &result](int image_h, int image_w, int image_index) { std::vector> image_ids(static_cast(bs) * image_h * image_w, std::vector(axes_dim_num, 0.f)); + result.positions.append_image(image_h, image_w); int h_start = -(image_h - image_h / 2); int w_start = -(image_w - image_w / 2); for (int b = 0; b < bs; ++b) { @@ -692,15 +685,18 @@ namespace Rope { static_cast(i + 1)); ids = concat_ids(ids, ref_ids, bs); } - return embed_nd(ids, bs, static_cast(theta), axes_dim); + result.ids = std::move(ids); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } __STATIC_INLINE__ std::vector> gen_lens_ids(int h, int w, int bs, int context_len, - bool scale_rope = true) { - auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope); + bool scale_rope = true, + PositionLayout* layout = nullptr) { + auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope, layout); int txt_id_start = scale_rope ? std::max(h / 2, w / 2) : 0; auto txt_ids = linspace(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len); @@ -711,44 +707,37 @@ namespace Rope { } } + if (layout) { + layout->append_tokens(context_len); + } return concat_ids(img_ids_repeated, txt_ids_repeated, bs); } - __STATIC_INLINE__ std::vector gen_lens_pe(int h, - int w, - int bs, - int context_len, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim) { - std::vector> ids = gen_lens_ids(h, w, bs, context_len, true); - std::vector> wrap_dims; - if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { - size_t pos_len = ids.size() / bs; - wrap_dims.assign(axes_dim.size(), std::vector(pos_len, 0)); - const size_t img_tokens = static_cast(h) * static_cast(w); - for (size_t token_i = 0; token_i < img_tokens; ++token_i) { - if (circular_h) { - wrap_dims[1][token_i] = h; - } - if (circular_w) { - wrap_dims[2][token_i] = w; - } - } - } - - return embed_nd(ids, bs, static_cast(theta), axes_dim, wrap_dims); + __STATIC_INLINE__ Embedding gen_lens_pe(int h, + int w, + int bs, + int context_len, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.ids = gen_lens_ids(h, w, bs, context_len, true, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } __STATIC_INLINE__ std::vector> gen_ernie_image_ids(int h, int w, int patch_size, int bs, - int context_len) { + int context_len, + PositionLayout* layout = nullptr) { int h_len = h / patch_size; int w_len = w / patch_size; + if (layout) { + layout->append_image(h_len, w_len); + } std::vector> img_ids(h_len * w_len, std::vector(3, 0.0f)); std::vector h_ids = linspace(0.f, static_cast(h_len - 1), h_len); std::vector w_ids = linspace(0.f, static_cast(w_len - 1), w_len); @@ -774,39 +763,25 @@ namespace Rope { } } + if (layout) { + layout->append_tokens(context_len); + } return concat_ids(img_ids_repeated, txt_ids, bs); } - __STATIC_INLINE__ std::vector gen_ernie_image_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim) { - std::vector> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len); - std::vector> wrap_dims; - if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { - int h_len = h / patch_size; - int w_len = w / patch_size; - if (h_len > 0 && w_len > 0) { - size_t pos_len = ids.size() / bs; - wrap_dims.assign(axes_dim.size(), std::vector(pos_len, 0)); - const size_t img_tokens = static_cast(h_len) * static_cast(w_len); - for (size_t token_i = 0; token_i < img_tokens; ++token_i) { - if (circular_h) { - wrap_dims[1][token_i] = h_len; - } - if (circular_w) { - wrap_dims[2][token_i] = w_len; - } - } - } - } - - return embed_nd(ids, bs, static_cast(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage); + __STATIC_INLINE__ Embedding gen_ernie_image_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.layout = EmbedNDLayout::ErnieImage; + result.ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } // Generate wan positional embeddings @@ -905,7 +880,8 @@ namespace Rope { int context_len, int seq_multi_of, const std::vector& ref_latents, - RefIndexMode ref_index_mode) { + RefIndexMode ref_index_mode, + PositionLayout* layout = nullptr) { SD_UNUSED(ref_index_mode); int padded_context_len = context_len + bound_mod(context_len, seq_multi_of); auto txt_ids = std::vector>(bs * padded_context_len, std::vector(3, 0.0f)); @@ -913,11 +889,17 @@ namespace Rope { txt_ids[i][0] = (i % padded_context_len) + 1.f; } + if (layout) { + layout->append_tokens(padded_context_len); + } int axes_dim_num = 3; int index = padded_context_len + 1; - auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index); + auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout); int img_pad_len = bound_mod(static_cast(img_ids.size() / bs), seq_multi_of); + if (layout) { + layout->append_tokens(img_pad_len); + } if (img_pad_len > 0) { std::vector> img_pad_ids(bs * img_pad_len, std::vector(3, 0.f)); img_ids = concat_ids(img_ids, img_pad_ids, bs); @@ -929,43 +911,160 @@ namespace Rope { return ids; } - // Generate z_image positional embeddings - __STATIC_INLINE__ std::vector gen_z_image_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - int seq_multi_of, - const std::vector& ref_latents, - RefIndexMode ref_index_mode, - int theta, - bool circular_h, - bool circular_w, - const std::vector& axes_dim) { - std::vector> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode); - std::vector> wrap_dims; - if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { - int pad_h = (patch_size - (h % patch_size)) % patch_size; - int pad_w = (patch_size - (w % patch_size)) % patch_size; - int h_len = (h + pad_h) / patch_size; - int w_len = (w + pad_w) / patch_size; - if (h_len > 0 && w_len > 0) { - size_t pos_len = ids.size() / bs; - wrap_dims.assign(axes_dim.size(), std::vector(pos_len, 0)); - size_t cursor = context_len + bound_mod(context_len, seq_multi_of); // skip text (and its padding) - size_t img_tokens = static_cast(h_len) * static_cast(w_len); - for (size_t token_i = 0; token_i < img_tokens; ++token_i) { - if (circular_h) { - wrap_dims[1][cursor + token_i] = h_len; - } - if (circular_w) { - wrap_dims[2][cursor + token_i] = w_len; - } - } + // LLaDA-Image shares Lumina2/z_image's axes layout, but assigns position (0,0,0) to the + // padding slots of the caption stream instead of continuing the caption ramp through them. + __STATIC_INLINE__ std::vector> gen_llada_image_ids(int h, + int w, + int patch_size, + int bs, + int context_len, + int seq_multi_of, + PositionLayout* layout = nullptr) { + int context_pad_len = bound_mod(context_len, seq_multi_of); + int padded_context_len = context_len + context_pad_len; + auto txt_ids = std::vector>(bs * padded_context_len, std::vector(3, 0.0f)); + for (int i = 0; i < bs * padded_context_len; i++) { + int pos = i % padded_context_len; + if (pos < context_len) { + txt_ids[i][0] = pos + 1.f; } } - return embed_nd(ids, bs, static_cast(theta), axes_dim, wrap_dims); + if (layout) { + layout->append_tokens(padded_context_len); + } + int axes_dim_num = 3; + int index = padded_context_len + 1; + auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout); + + int img_pad_len = bound_mod(static_cast(img_ids.size() / bs), seq_multi_of); + if (layout) { + layout->append_tokens(img_pad_len); + } + if (img_pad_len > 0) { + std::vector> img_pad_ids(bs * img_pad_len, std::vector(3, 0.f)); + img_ids = concat_ids(img_ids, img_pad_ids, bs); + } + + return concat_ids(txt_ids, img_ids, bs); + } + + // LLaDA-Image editing packs two caption copies (clean and noisy), the source and target + // latents anchored at their own caption's end position, and the SigVQ stream after both. + // Padding slots keep position (0,0,0), as in the text-only layout. + __STATIC_INLINE__ std::vector> gen_llada_image_edit_ids(int h, + int w, + int patch_size, + int context_len, + int sigvq_len, + int seq_multi_of, + PositionLayout* layout = nullptr) { + const int context_pad = bound_mod(context_len, seq_multi_of); + const int padded_context = context_len + context_pad; + const int h_len = (h + (patch_size / 2)) / patch_size; + const int w_len = (w + (patch_size / 2)) / patch_size; + const int image_len = h_len * w_len; + const int image_pad = bound_mod(image_len, seq_multi_of); + const int padded_image = image_len + image_pad; + const int sigvq_pad = bound_mod(sigvq_len, seq_multi_of); + + std::vector> cap_ids; + std::vector cap_end_positions; + int cursor = 1; + for (int copy = 0; copy < 2; ++copy) { + for (int i = 0; i < padded_context; ++i) { + std::vector id(3, 0.f); + if (i < context_len) { + id[0] = static_cast(cursor + i); + } + cap_ids.push_back(id); + } + cursor += context_len; + cap_end_positions.push_back(cursor); + cursor += 2; + } + + if (layout) { + layout->append_tokens(cap_ids.size()); + } + std::vector> img_ids; + for (int copy = 0; copy < 2; ++copy) { + auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy], 0, 0, false, layout); + img_ids.insert(img_ids.end(), ids.begin(), ids.end()); + img_ids.insert(img_ids.end(), image_pad, std::vector(3, 0.f)); + if (layout) { + layout->append_tokens(image_pad); + } + } + + const int sigvq_start = static_cast(cap_ids.size() + img_ids.size()) + 1; + std::vector> sigvq_ids; + for (int i = 0; i < sigvq_len + sigvq_pad; ++i) { + std::vector id(3, 0.f); + if (i < sigvq_len) { + id[0] = static_cast(sigvq_start + i); + } + sigvq_ids.push_back(id); + } + + std::vector> ids; + ids.reserve(cap_ids.size() + img_ids.size() + sigvq_ids.size()); + ids.insert(ids.end(), cap_ids.begin(), cap_ids.end()); + ids.insert(ids.end(), img_ids.begin(), img_ids.end()); + ids.insert(ids.end(), sigvq_ids.begin(), sigvq_ids.end()); + if (layout) { + layout->append_tokens(sigvq_ids.size()); + } + SD_UNUSED(padded_image); + return ids; + } + + __STATIC_INLINE__ Embedding gen_llada_image_edit_pe(int h, + int w, + int patch_size, + int context_len, + int sigvq_len, + int seq_multi_of, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = 1; + result.ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of, &result.positions); + result.values = embed_nd(result.ids, 1, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; + } + + __STATIC_INLINE__ Embedding gen_llada_image_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + int seq_multi_of, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; + } + + // Generate z_image positional embeddings + __STATIC_INLINE__ Embedding gen_z_image_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + int seq_multi_of, + const std::vector& ref_latents, + RefIndexMode ref_index_mode, + int theta, + const std::vector& axes_dim) { + Embedding result; + result.batch_size = bs; + result.ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode, &result.positions); + result.values = embed_nd(result.ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } __STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx, diff --git a/otherarch/sdcpp/src/model/common/rope_circular.hpp b/otherarch/sdcpp/src/model/common/rope_circular.hpp new file mode 100644 index 000000000..19dccb959 --- /dev/null +++ b/otherarch/sdcpp/src/model/common/rope_circular.hpp @@ -0,0 +1,65 @@ +#ifndef __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__ +#define __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__ + +#include "model/common/rope.hpp" + +namespace Rope { + __STATIC_INLINE__ void apply_circular(Embedding& embedding, bool circular_x, bool circular_y) { + if (!circular_x && !circular_y) { + return; + } + + GGML_ASSERT(embedding.batch_size > 0); + GGML_ASSERT(embedding.ids.size() % embedding.batch_size == 0); + size_t pos_len = embedding.ids.size() / embedding.batch_size; + size_t half_dim = embedding.frequencies.size(); + GGML_ASSERT(embedding.positions.token_count == pos_len); + GGML_ASSERT(embedding.values.size() == embedding.ids.size() * half_dim * 4); + + constexpr float TWO_PI = 6.28318530717958647692f; + for (const auto& region : embedding.positions.images) { + GGML_ASSERT(region.begin <= pos_len && region.count <= pos_len - region.begin); + for (size_t j = 0; j < half_dim; ++j) { + const auto& frequency = embedding.frequencies[j]; + float period = 0.f; + if (circular_y && frequency.axis == static_cast(region.height_axis)) { + period = region.height_period; + } else if (circular_x && frequency.axis == static_cast(region.width_axis)) { + period = region.width_period; + } + if (period <= 0) { + continue; + } + + // Quantize to periodic harmonics while preserving the original coordinate offsets. + float rounded = std::round(frequency.omega * period / TWO_PI); + for (int b = 0; b < embedding.batch_size; ++b) { + size_t begin = b * pos_len + region.begin; + for (size_t i = begin; i < begin + region.count; ++i) { + GGML_ASSERT(frequency.axis < embedding.ids[i].size()); + float angle = embedding.ids[i][frequency.axis] * TWO_PI * rounded / period; + float cos_val = std::cos(angle); + float sin_val = std::sin(angle); + if (embedding.layout == EmbedNDLayout::ErnieImage) { + size_t cos_offset = (i * half_dim + j) * 2; + size_t sin_offset = embedding.ids.size() * half_dim * 2 + cos_offset; + embedding.values[cos_offset] = cos_val; + embedding.values[cos_offset + 1] = cos_val; + embedding.values[sin_offset] = sin_val; + embedding.values[sin_offset + 1] = sin_val; + } else { + size_t offset = (i * half_dim + j) * 4; + embedding.values[offset] = cos_val; + embedding.values[offset + 1] = -sin_val; + embedding.values[offset + 2] = sin_val; + embedding.values[offset + 3] = cos_val; + } + } + } + } + } + } + +} // namespace Rope + +#endif // __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__ diff --git a/otherarch/sdcpp/src/model/diffusion/anima.hpp b/otherarch/sdcpp/src/model/diffusion/anima.hpp index a8977c966..9ace46d54 100644 --- a/otherarch/sdcpp/src/model/diffusion/anima.hpp +++ b/otherarch/sdcpp/src/model/diffusion/anima.hpp @@ -603,34 +603,37 @@ namespace Anima { return std::pow(extrapolation_ratio, static_cast(axis_dim) / static_cast(axis_dim - 2)); } - static std::vector gen_anima_image_pe_vec(int bs, - int h, - int w, - int patch_size, - int theta, - const std::vector& axes_dim, - float h_extrapolation_ratio, - float w_extrapolation_ratio, - float t_extrapolation_ratio, - const std::vector& ref_latents) { - auto ids = Rope::gen_flux_ids(h, - w, - patch_size, - bs, - static_cast(axes_dim.size()), - 0, - {}, - ref_latents, - Rope::RefIndexMode::FIXED, - 1.0f, - false); + static Rope::Embedding gen_anima_image_pe_vec(int bs, + int h, + int w, + int patch_size, + int theta, + const std::vector& axes_dim, + float h_extrapolation_ratio, + float w_extrapolation_ratio, + float t_extrapolation_ratio, + const std::vector& ref_latents) { + Rope::Embedding result; + result.batch_size = bs; + result.ids = Rope::gen_flux_ids(h, + w, + patch_size, + bs, + static_cast(axes_dim.size()), + 0, + {}, + ref_latents, + Rope::RefIndexMode::FIXED, + 1.0f, + false, &result.positions); std::vector axis_thetas = { static_cast(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]), static_cast(theta) * calc_ntk_factor(h_extrapolation_ratio, axes_dim[1]), static_cast(theta) * calc_ntk_factor(w_extrapolation_ratio, axes_dim[2]), }; - return Rope::embed_nd(ids, bs, axis_thetas, axes_dim); + result.values = Rope::embed_nd(result.ids, bs, axis_thetas, axes_dim, result.layout, &result.frequencies); + return result; } ggml_cgraph* build_graph(const sd::Tensor& x_tensor, @@ -657,16 +660,16 @@ namespace Anima { int64_t h_pad = x->ne[1] + pad_h; int64_t w_pad = x->ne[0] + pad_w; - image_pe_vec = gen_anima_image_pe_vec(1, - static_cast(h_pad), - static_cast(w_pad), - static_cast(config.patch_size), - config.theta, - config.axes_dim, - 4.0f, - 4.0f, - 1.0f, - ref_latents); + image_pe_vec = finish_rope_pe(gen_anima_image_pe_vec(1, + static_cast(h_pad), + static_cast(w_pad), + static_cast(config.patch_size), + config.theta, + config.axes_dim, + 4.0f, + 4.0f, + 1.0f, + ref_latents)); int64_t image_pos_len = static_cast(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2)); auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len); set_backend_tensor_data(image_pe, image_pe_vec.data()); diff --git a/otherarch/sdcpp/src/model/diffusion/boogu.hpp b/otherarch/sdcpp/src/model/diffusion/boogu.hpp index 8c91bb55d..1ca137098 100644 --- a/otherarch/sdcpp/src/model/diffusion/boogu.hpp +++ b/otherarch/sdcpp/src/model/diffusion/boogu.hpp @@ -720,15 +720,18 @@ namespace Boogu { } } - __STATIC_INLINE__ std::vector gen_boogu_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - const std::vector& ref_latents, - int theta, - const std::vector& axes_dim) { - std::vector> ids; + __STATIC_INLINE__ Rope::Embedding gen_boogu_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + const std::vector& ref_latents, + int theta, + const std::vector& axes_dim) { + Rope::Embedding result; + result.batch_size = bs; + result.positions.append_tokens(context_len); + auto& ids = result.ids; ids.reserve(static_cast(bs) * context_len); for (int b = 0; b < bs; b++) { for (int i = 0; i < context_len; i++) { @@ -741,15 +744,18 @@ namespace Boogu { for (ggml_tensor* ref : ref_latents) { int ref_h_tokens = patched_token_count(ref->ne[1], patch_size); int ref_w_tokens = patched_token_count(ref->ne[0], patch_size); + result.positions.append_image(ref_h_tokens, ref_w_tokens); append_spatial_ids(ids, bs, pe_shift, ref_h_tokens, ref_w_tokens); pe_shift += std::max(ref_h_tokens, ref_w_tokens); } int h_tokens = patched_token_count(h, patch_size); int w_tokens = patched_token_count(w, patch_size); + result.positions.append_image(h_tokens, w_tokens); append_spatial_ids(ids, bs, pe_shift, h_tokens, w_tokens); - return Rope::embed_nd(ids, bs, static_cast(theta), axes_dim); + result.values = Rope::embed_nd(ids, bs, static_cast(theta), axes_dim, result.layout, &result.frequencies); + return result; } struct BooguImageRunner : public DiffusionModelRunner { @@ -793,14 +799,14 @@ namespace Boogu { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = gen_boogu_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ref_latents, - config.theta, - config.axes_dim); + pe_vec = finish_rope_pe(gen_boogu_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim)); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/otherarch/sdcpp/src/model/diffusion/control.hpp b/otherarch/sdcpp/src/model/diffusion/control.hpp index 34387d1b7..3b08ed150 100644 --- a/otherarch/sdcpp/src/model/diffusion/control.hpp +++ b/otherarch/sdcpp/src/model/diffusion/control.hpp @@ -376,7 +376,7 @@ struct ControlNet : public GGMLRunner { hint = make_input(hint_tensor); } - auto runner_ctx = get_context(); + auto runner_ctx = get_context(gf); auto outs = control_net.forward(&runner_ctx, x, @@ -389,8 +389,7 @@ struct ControlNet : public GGMLRunner { if (guided_hint_input == nullptr && !outs.empty()) { guided_hint_output_ggml = outs[0]; ggml_set_output(guided_hint_output_ggml); - cache(guided_hint_cache_name(), guided_hint_output_ggml); - ggml_build_forward_expand(gf, guided_hint_output_ggml); + runner_ctx.persist_cache_tensor(guided_hint_cache_name(), guided_hint_output_ggml); } control_outputs_ggml.reserve(outs.size() > 0 ? outs.size() - 1 : 0); diff --git a/otherarch/sdcpp/src/model/diffusion/ernie_image.hpp b/otherarch/sdcpp/src/model/diffusion/ernie_image.hpp index 81bf142a0..b284cb83f 100644 --- a/otherarch/sdcpp/src/model/diffusion/ernie_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/ernie_image.hpp @@ -415,15 +415,13 @@ namespace ErnieImage { GGML_ASSERT(!context_tensor.empty()); ggml_tensor* context = make_input(context_tensor); - pe_vec = Rope::gen_ernie_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim)); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/otherarch/sdcpp/src/model/diffusion/flux.hpp b/otherarch/sdcpp/src/model/diffusion/flux.hpp index 67e79d3d9..825dabe13 100644 --- a/otherarch/sdcpp/src/model/diffusion/flux.hpp +++ b/otherarch/sdcpp/src/model/diffusion/flux.hpp @@ -1548,20 +1548,18 @@ namespace Flux { } else if (version == VERSION_OVIS_IMAGE) { txt_arange_dims = {1, 2}; } - pe_vec = Rope::gen_flux_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - txt_arange_dims, - ref_latents, - ref_index_mode, - config.ref_index_scale, - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim, - sd_version_is_longcat(version)); + pe_vec = finish_rope_pe(Rope::gen_flux_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + txt_arange_dims, + ref_latents, + ref_index_mode, + config.ref_index_scale, + config.theta, + config.axes_dim, + sd_version_is_longcat(version))); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); // LOG_VERBOSE("pos_len %d", pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); diff --git a/otherarch/sdcpp/src/model/diffusion/ideogram4.hpp b/otherarch/sdcpp/src/model/diffusion/ideogram4.hpp index 9e2aa3612..3254571db 100644 --- a/otherarch/sdcpp/src/model/diffusion/ideogram4.hpp +++ b/otherarch/sdcpp/src/model/diffusion/ideogram4.hpp @@ -149,18 +149,21 @@ namespace Ideogram4 { return std::make_shared(in_features, out_features, bias); } - __STATIC_INLINE__ std::vector gen_ideogram4_pe(int grid_h, - int grid_w, - int bs, - int context_len, - int head_dim, - int rope_theta, - const std::vector& mrope_section, - bool circular_x = false, - bool circular_y = false) { + __STATIC_INLINE__ Rope::Embedding gen_ideogram4_pe(int grid_h, + int grid_w, + int bs, + int context_len, + int head_dim, + int rope_theta, + const std::vector& mrope_section) { GGML_ASSERT(bs == 1); - std::vector> ids(static_cast(bs) * (context_len + grid_h * grid_w), - std::vector(3, 0.f)); + Rope::Embedding result; + result.batch_size = bs; + result.positions.append_tokens(context_len); + result.positions.append_image(grid_h, grid_w); + result.ids.assign(static_cast(bs) * (context_len + grid_h * grid_w), + std::vector(3, 0.f)); + auto& ids = result.ids; for (int i = 0; i < context_len; ++i) { ids[i] = {static_cast(i), static_cast(i), static_cast(i)}; @@ -175,29 +178,13 @@ namespace Ideogram4 { } } - std::vector> axis_wrap_dims(3); - if (circular_y || circular_x) { - size_t total_len = static_cast(bs) * (context_len + grid_h * grid_w); - axis_wrap_dims[1].assign(total_len, 0); - axis_wrap_dims[2].assign(total_len, 0); - if (circular_y) { - for (size_t idx = static_cast(context_len); idx < total_len; ++idx) { - axis_wrap_dims[1][idx] = grid_h; - } - } - if (circular_x) { - for (size_t idx = static_cast(context_len); idx < total_len; ++idx) { - axis_wrap_dims[2][idx] = grid_w; - } - } - } - - return Rope::embed_interleaved_mrope(ids, - bs, - static_cast(rope_theta), - head_dim, - mrope_section, - axis_wrap_dims); + result.values = Rope::embed_interleaved_mrope(ids, + bs, + static_cast(rope_theta), + head_dim, + mrope_section, + &result.frequencies); + return result; } class Ideogram4Attention : public GGMLBlock { @@ -509,15 +496,13 @@ namespace Ideogram4 { int64_t head_dim = config.emb_dim / config.num_heads; auto runner_ctx = get_context(); - pe_vec = gen_ideogram4_pe(static_cast(grid_h), - static_cast(grid_w), - static_cast(x->ne[3]), - static_cast(context_len), - static_cast(head_dim), - static_cast(config.rope_theta), - config.mrope_section, - runner_ctx.circular_x_enabled, - runner_ctx.circular_y_enabled); + pe_vec = finish_rope_pe(gen_ideogram4_pe(static_cast(grid_h), + static_cast(grid_w), + static_cast(x->ne[3]), + static_cast(context_len), + static_cast(head_dim), + static_cast(config.rope_theta), + config.mrope_section)); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/otherarch/sdcpp/src/model/diffusion/krea2.hpp b/otherarch/sdcpp/src/model/diffusion/krea2.hpp index 84ecfe0b1..e58df4108 100644 --- a/otherarch/sdcpp/src/model/diffusion/krea2.hpp +++ b/otherarch/sdcpp/src/model/diffusion/krea2.hpp @@ -689,23 +689,28 @@ namespace Krea2 { } }; - __STATIC_INLINE__ std::vector gen_krea2_pe(int h, - int w, - int patch_size, - int bs, - int context_len, - float theta, - const std::vector& axes_dim, - const std::vector& ref_latents, - Rope::RefIndexMode ref_index_mode) { + __STATIC_INLINE__ Rope::Embedding gen_krea2_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + float theta, + const std::vector& axes_dim, + const std::vector& ref_latents, + Rope::RefIndexMode ref_index_mode) { + Rope::Embedding result; + result.batch_size = bs; + result.positions.append_tokens(context_len); auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {}); - auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false); + auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false, &result.positions); auto ids = Rope::concat_ids(txt_ids, img_ids, bs); if (ref_latents.size() > 0) { - auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0); + auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0, &result.positions); ids = Rope::concat_ids(ids, refs_ids, bs); } - return Rope::embed_nd(ids, bs, theta, axes_dim); + result.ids = std::move(ids); + result.values = Rope::embed_nd(result.ids, bs, theta, axes_dim, result.layout, &result.frequencies); + return result; } struct Krea2Runner : public DiffusionModelRunner { @@ -749,15 +754,15 @@ namespace Krea2 { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = gen_krea2_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - config.axes_dim, - ref_latents, - ref_image_params.ref_index_mode); + pe_vec = finish_rope_pe(gen_krea2_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim, + ref_latents, + ref_image_params.ref_index_mode)); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/otherarch/sdcpp/src/model/diffusion/lens.hpp b/otherarch/sdcpp/src/model/diffusion/lens.hpp index 35b4b517d..8fb1512b4 100644 --- a/otherarch/sdcpp/src/model/diffusion/lens.hpp +++ b/otherarch/sdcpp/src/model/diffusion/lens.hpp @@ -384,14 +384,12 @@ namespace Lens { GGML_ASSERT(!context_tensor.empty()); ggml_tensor* context = make_input(context_tensor); - pe_vec = Rope::gen_lens_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim)); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/otherarch/sdcpp/src/model/diffusion/llada_image.hpp b/otherarch/sdcpp/src/model/diffusion/llada_image.hpp new file mode 100644 index 000000000..e69c3ecd3 --- /dev/null +++ b/otherarch/sdcpp/src/model/diffusion/llada_image.hpp @@ -0,0 +1,525 @@ +#ifndef __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__ +#define __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__ + +#include +#include + +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" +#include "model/diffusion/model.hpp" +#include "model/diffusion/z_image.hpp" +#include "model_loader.h" + +// Ref: https://github.com/inclusionAI/LLaDA-Image/blob/main/src/models/transformer_llada_image.py +// +// The denoiser is Lumina2/z_image's NextDiT with identical hyperparameters, so the blocks are +// reused from ZImage. Two things differ: every norm here is non-parametric (the checkpoint +// carries no norm weights at all), and latents arrive already patchified from the Flux2 VAE, +// so patch_size is 1 over 128 channels. + +namespace LLaDAImage { + constexpr int LLADA_IMAGE_GRAPH_SIZE = 20480; + + struct LLaDAImageConfig { + int patch_size = 1; + int64_t hidden_size = 3840; + int64_t in_channels = 128; + int64_t out_channels = 128; + int64_t num_layers = 30; + int64_t num_refiner_layers = 2; + int64_t head_dim = 128; + int64_t num_heads = 30; + int64_t num_kv_heads = 30; + int64_t multiple_of = 256; + float ffn_dim_multiplier = 8.0f / 3.0f; + float norm_eps = 1e-5f; + bool qk_norm = true; + int64_t cap_feat_dim = 2560; + int64_t semantic_feat_dim = 4096; + int theta = 256; + std::vector axes_dim = {32, 48, 48}; + int64_t axes_dim_sum = 128; + + static int64_t count_blocks(const String2TensorStorage& tensor_storage_map, + const std::string& prefix, + const std::string& block_prefix) { + int64_t count = 0; + for (const auto& [name, _] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + size_t pos = name.find(block_prefix); + if (pos == std::string::npos) { + continue; + } + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + count = std::max(count, atoi(items[1].c_str()) + 1); + } + } + return count; + } + + static LLaDAImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + LLaDAImageConfig config; + int64_t detected_q_dim = 0; + int64_t detected_kv_dim = 0; + + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) { + int64_t patch_area = config.patch_size * config.patch_size; + config.in_channels = tensor_storage.ne[0] / patch_area; + config.hidden_size = tensor_storage.ne[1]; + } else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) { + config.cap_feat_dim = tensor_storage.ne[0]; + config.hidden_size = tensor_storage.ne[1]; + } else if (ends_with(name, "sigvq_embedder.1.weight") && tensor_storage.n_dims == 2) { + config.semantic_feat_dim = tensor_storage.ne[0]; + } else if (ends_with(name, "layers.0.attention.to_q.weight") && tensor_storage.n_dims == 2) { + detected_q_dim = tensor_storage.ne[1]; + } else if (ends_with(name, "layers.0.attention.to_k.weight") && tensor_storage.n_dims == 2) { + detected_kv_dim = tensor_storage.ne[1]; + } else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) { + int64_t patch_area = config.patch_size * config.patch_size; + config.out_channels = tensor_storage.ne[1] / patch_area; + } + } + + int64_t detected_layers = count_blocks(tensor_storage_map, prefix, "layers."); + int64_t detected_refiner = std::max(count_blocks(tensor_storage_map, prefix, "noise_refiner."), + count_blocks(tensor_storage_map, prefix, "context_refiner.")); + if (detected_layers > 0) { + config.num_layers = detected_layers; + } + if (detected_refiner > 0) { + config.num_refiner_layers = detected_refiner; + } + if (detected_q_dim > 0) { + config.num_heads = detected_q_dim / config.head_dim; + } + if (detected_kv_dim > 0) { + config.num_kv_heads = detected_kv_dim / config.head_dim; + } else if (detected_q_dim > 0) { + config.num_kv_heads = config.num_heads; + } + + LOG_VERBOSE("llada_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64 ", cap_feat_dim = %" PRId64 ", semantic_feat_dim = %" PRId64, + config.num_layers, + config.num_refiner_layers, + config.hidden_size, + config.num_heads, + config.num_kv_heads, + config.in_channels, + config.out_channels, + config.cap_feat_dim, + config.semantic_feat_dim); + return config; + } + }; + + class LLaDAImageModel : public GGMLBlock { + protected: + LLaDAImageConfig config; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); + params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); + params["sigvq_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); + } + + std::shared_ptr make_block(bool modulation) { + return std::make_shared(0, + config.hidden_size, + config.head_dim, + config.num_heads, + config.num_kv_heads, + config.multiple_of, + config.ffn_dim_multiplier, + config.norm_eps, + config.qk_norm, + modulation, + false, + true); + } + + public: + LLaDAImageModel() = default; + LLaDAImageModel(LLaDAImageConfig config) + : config(config) { + blocks["x_embedder"] = std::make_shared(config.patch_size * config.patch_size * config.in_channels, config.hidden_size); + blocks["t_embedder"] = std::make_shared(MIN(config.hidden_size, 1024), 256, ZImage::ADALN_EMBED_DIM); + + blocks["cap_embedder.0"] = std::make_shared(config.cap_feat_dim, config.norm_eps, false); + blocks["cap_embedder.1"] = std::make_shared(config.cap_feat_dim, config.hidden_size); + + blocks["semantic_embedder.0"] = std::make_shared(config.semantic_feat_dim, config.norm_eps, false); + blocks["semantic_embedder.1"] = std::make_shared(config.semantic_feat_dim, config.hidden_size); + blocks["sigvq_embedder.0"] = std::make_shared(config.semantic_feat_dim, config.norm_eps, false); + blocks["sigvq_embedder.1"] = std::make_shared(config.semantic_feat_dim, config.hidden_size); + + for (int i = 0; i < config.num_refiner_layers; i++) { + blocks["noise_refiner." + std::to_string(i)] = make_block(true); + blocks["context_refiner." + std::to_string(i)] = make_block(false); + blocks["sigvq_refiner." + std::to_string(i)] = make_block(false); + } + for (int i = 0; i < config.num_layers; i++) { + blocks["layers." + std::to_string(i)] = make_block(true); + } + + blocks["final_layer"] = std::make_shared(config.hidden_size, config.patch_size, config.out_channels); + } + + ggml_tensor* forward_core(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timestep, + ggml_tensor* context, + ggml_tensor* pe) { + auto x_embedder = std::dynamic_pointer_cast(blocks["x_embedder"]); + auto t_embedder = std::dynamic_pointer_cast(blocks["t_embedder"]); + auto cap_embedder_0 = std::dynamic_pointer_cast(blocks["cap_embedder.0"]); + auto cap_embedder_1 = std::dynamic_pointer_cast(blocks["cap_embedder.1"]); + auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); + + auto txt_pad_token = params["cap_pad_token"]; + auto img_pad_token = params["x_pad_token"]; + + int64_t N = x->ne[2]; + int64_t n_img_token = x->ne[1]; + int64_t n_txt_token = context->ne[1]; + + // sdcpp's flow denoiser already hands over sigma * 1000, which is the range the + // reference reaches via its own t_scale, so no further scaling here. + auto t_emb = t_embedder->forward(ctx, timestep); + + auto txt = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context)); // [N, n_txt_token, hidden_size] + auto img = x_embedder->forward(ctx, x); // [N, n_img_token, hidden_size] + sd::ggml_graph_cut::mark_graph_cut(txt, "llada_image.prelude", "txt"); + sd::ggml_graph_cut::mark_graph_cut(img, "llada_image.prelude", "img"); + sd::ggml_graph_cut::mark_graph_cut(t_emb, "llada_image.prelude", "t_emb"); + + int64_t n_txt_pad_token = Rope::bound_mod(static_cast(n_txt_token), ZImage::SEQ_MULTI_OF); + if (n_txt_pad_token > 0) { + auto txt_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, txt_pad_token, txt_pad_token->ne[0], n_txt_pad_token, N, 1); + txt = ggml_concat(ctx->ggml_ctx, txt, txt_pad_tokens, 1); + } + + int64_t n_img_pad_token = Rope::bound_mod(static_cast(n_img_token), ZImage::SEQ_MULTI_OF); + if (n_img_pad_token > 0) { + auto img_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, img_pad_token, img_pad_token->ne[0], n_img_pad_token, N, 1); + img = ggml_concat(ctx->ggml_ctx, img, img_pad_tokens, 1); + } + + GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]); + + auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]); + auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]); + + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)]); + + txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr); + sd::ggml_graph_cut::mark_graph_cut(txt, "llada_image.context_refiner." + std::to_string(i), "txt"); + } + + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)]); + + img = block->forward(ctx, img, img_pe, nullptr, t_emb); + sd::ggml_graph_cut::mark_graph_cut(img, "llada_image.noise_refiner." + std::to_string(i), "img"); + } + + auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); + sd::ggml_graph_cut::mark_graph_cut(txt_img, "llada_image.prelude", "txt_img"); + + for (int i = 0; i < config.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); + + txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb); + sd::ggml_graph_cut::mark_graph_cut(txt_img, "llada_image.layers." + std::to_string(i), "txt_img"); + } + + txt_img = final_layer->forward(ctx, txt_img, t_emb); + + return ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token); + } + + ggml_tensor* pad_stream(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pad_token) { + int64_t n_pad = Rope::bound_mod(static_cast(x->ne[1]), ZImage::SEQ_MULTI_OF); + if (n_pad == 0) { + return x; + } + auto pads = ggml_repeat_4d(ctx->ggml_ctx, pad_token, pad_token->ne[0], n_pad, x->ne[2], 1); + return ggml_concat(ctx->ggml_ctx, x, pads, 1); + } + + // Editing runs one joint sequence carrying two timesteps: the caption and source latent + // are clean (t = 0) while the second caption copy and the target latent are noisy. adaLN + // is a linear map of the timestep embedding, so feeding a per-token embedding selects the + // right modulation exactly, without duplicating the modulation projections. + ggml_tensor* forward_editing(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timestep, + ggml_tensor* context, + ggml_tensor* semantic, + ggml_tensor* source_latent, + ggml_tensor* pe) { + ggml_context* gctx = ctx->ggml_ctx; + + auto x_embedder = std::dynamic_pointer_cast(blocks["x_embedder"]); + auto t_embedder = std::dynamic_pointer_cast(blocks["t_embedder"]); + auto cap_embedder_0 = std::dynamic_pointer_cast(blocks["cap_embedder.0"]); + auto cap_embedder_1 = std::dynamic_pointer_cast(blocks["cap_embedder.1"]); + auto sigvq_embed_0 = std::dynamic_pointer_cast(blocks["sigvq_embedder.0"]); + auto sigvq_embed_1 = std::dynamic_pointer_cast(blocks["sigvq_embedder.1"]); + auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); + + auto t_noisy = t_embedder->forward(ctx, timestep); + auto t_clean = t_embedder->forward(ctx, ggml_scale(gctx, timestep, 0.f)); + + auto per_token = [&](ggml_tensor* emb, int64_t n) { + return ggml_repeat_4d(gctx, emb, emb->ne[0], n, 1, 1); + }; + + auto cap = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context)); + cap = pad_stream(ctx, cap, params["cap_pad_token"]); + int64_t cap_len = cap->ne[1]; + cap = ggml_concat(gctx, cap, cap, 1); + + auto src = pad_stream(ctx, x_embedder->forward(ctx, source_latent), params["x_pad_token"]); + auto tgt_embed = x_embedder->forward(ctx, x); + int64_t n_img_token = tgt_embed->ne[1]; + auto tgt = pad_stream(ctx, tgt_embed, params["x_pad_token"]); + int64_t img_len = tgt->ne[1]; + auto img = ggml_concat(gctx, src, tgt, 1); + + ggml_tensor* sig = nullptr; + int64_t sig_len = 0; + if (semantic != nullptr) { + sig = sigvq_embed_1->forward(ctx, sigvq_embed_0->forward(ctx, semantic)); + sig = pad_stream(ctx, sig, params["sigvq_pad_token"]); + sig_len = sig->ne[1]; + } + + GGML_ASSERT(cap_len * 2 + img_len * 2 + sig_len == pe->ne[3]); + + auto cap_pe = ggml_ext_slice(gctx, pe, 3, 0, cap_len * 2); + auto img_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2); + + auto img_adaln = ggml_concat(gctx, per_token(t_clean, img_len), per_token(t_noisy, img_len), 1); + + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)]); + cap = block->forward(ctx, cap, cap_pe, nullptr, nullptr); + } + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)]); + img = block->forward(ctx, img, img_pe, nullptr, img_adaln); + } + if (sig != nullptr) { + auto sig_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]); + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["sigvq_refiner." + std::to_string(i)]); + sig = block->forward(ctx, sig, sig_pe, nullptr, nullptr); + } + } + + auto seq = ggml_concat(gctx, cap, img, 1); + + auto cap_adaln = ggml_concat(gctx, per_token(t_clean, cap_len), per_token(t_noisy, cap_len), 1); + auto seq_adaln = ggml_concat(gctx, cap_adaln, img_adaln, 1); + if (sig != nullptr) { + seq = ggml_concat(gctx, seq, sig, 1); + seq_adaln = ggml_concat(gctx, seq_adaln, per_token(t_clean, sig_len), 1); + } + + for (int i = 0; i < config.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); + seq = block->forward(ctx, seq, pe, nullptr, seq_adaln); + sd::ggml_graph_cut::mark_graph_cut(seq, "llada_image.layers." + std::to_string(i), "seq"); + } + + seq = final_layer->forward(ctx, seq, seq_adaln); + + // Only the target latent is denoised; the source half of the image stream is context. + // The stream is padded to SEQ_MULTI_OF, so drop the pad tokens: they are not part of + // the latent grid that unpatchify reconstructs. + int64_t target_start = cap_len * 2 + img_len; + return ggml_ext_slice(gctx, seq, 1, target_start, target_start + n_img_token); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timestep, + ggml_tensor* context, + ggml_tensor* pe) { + // x: [N, C, H, W] + // timestep: [N,] + // context: [N, L, cap_feat_dim] + // pe: [L, d_head/2, 2, 2] + // return: [N, C, H, W] + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + + int patch_size = config.patch_size; + + auto img = DiT::pad_and_patchify(ctx, x, patch_size, patch_size, false); + + auto out = forward_core(ctx, img, timestep, context, pe); + + out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false); + + // The reference pipeline negates the model output before the scheduler step. + return ggml_ext_scale(ctx->ggml_ctx, out, -1.f); + } + }; + + struct LLaDAImageRunner : public DiffusionModelRunner { + public: + LLaDAImageConfig config; + LLaDAImageModel llada_image; + std::vector pe_vec; + + LLaDAImageRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(LLaDAImageConfig::detect_from_weights(tensor_storage_map, prefix)) { + llada_image = LLaDAImageModel(config); + llada_image.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "llada_image"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + llada_image.get_param_tensors(tensors, prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor) { + ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_GRAPH_SIZE); + ggml_tensor* x = make_input(x_tensor); + ggml_tensor* timesteps = make_input(timesteps_tensor); + GGML_ASSERT(x->ne[3] == 1); + GGML_ASSERT(!context_tensor.empty()); + ggml_tensor* context = make_input(context_tensor); + + pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); + set_backend_tensor_data(pe, pe_vec.data()); + auto runner_ctx = get_context(); + + ggml_tensor* out = llada_image.forward(&runner_ctx, x, timesteps, context, pe); + + ggml_build_forward_expand(gf, out); + + return gf; + } + + sd::Tensor compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context) { + // x: [N, in_channels, h, w] + // timesteps: [N, ] + // context: [N, max_position, cap_feat_dim] + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context); + }; + + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + } + + ggml_cgraph* build_edit_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor, + const sd::Tensor& semantic_tensor, + const sd::Tensor& source_tensor) { + ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_GRAPH_SIZE); + ggml_tensor* x = make_input(x_tensor); + ggml_tensor* timesteps = make_input(timesteps_tensor); + ggml_tensor* context = make_input(context_tensor); + ggml_tensor* semantic = make_optional_input(semantic_tensor); + ggml_tensor* source = make_input(source_tensor); + GGML_ASSERT(x->ne[3] == 1); + + pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(context->ne[1]), + semantic != nullptr ? static_cast(semantic->ne[1]) : 0, + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); + set_backend_tensor_data(pe, pe_vec.data()); + auto runner_ctx = get_context(); + + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + auto target = DiT::pad_and_patchify(&runner_ctx, x, config.patch_size, config.patch_size, false); + auto src = DiT::pad_and_patchify(&runner_ctx, source, config.patch_size, config.patch_size, false); + + auto out = llada_image.forward_editing(&runner_ctx, target, timesteps, context, semantic, src, pe); + out = DiT::unpatchify_and_crop(runner_ctx.ggml_ctx, out, H, W, config.patch_size, config.patch_size, false); + out = ggml_ext_scale(runner_ctx.ggml_ctx, out, -1.f); + + ggml_build_forward_expand(gf, out); + return gf; + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + + const auto* extra = std::get_if(&diffusion_params.extra); + bool has_semantic = extra != nullptr && extra->semantic != nullptr && !extra->semantic->empty(); + bool has_ref_latent = diffusion_params.ref_latents != nullptr && !diffusion_params.ref_latents->empty(); + if (has_semantic && !has_ref_latent) { + LOG_WARN("llada_image: SigVQ features without a reference latent are not supported; falling back to text to image"); + } + if (has_ref_latent) { + const auto& source = diffusion_params.ref_latents->front(); + if (source.shape() != diffusion_params.x->shape()) { + LOG_ERROR("llada_image: reference latent must match the target shape; use resize_vae_to_target=1"); + return {}; + } + auto get_graph = [&]() -> ggml_cgraph* { + return build_edit_graph(*diffusion_params.x, + *diffusion_params.timesteps, + tensor_or_empty(diffusion_params.context), + tensor_or_empty(extra != nullptr ? extra->semantic : nullptr), + source); + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), + diffusion_params.x->dim()); + } + + return compute(n_threads, + *diffusion_params.x, + *diffusion_params.timesteps, + tensor_or_empty(diffusion_params.context)); + } + }; + +} // namespace LLaDAImage + +#endif // __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__ diff --git a/otherarch/sdcpp/src/model/diffusion/mage_flow.hpp b/otherarch/sdcpp/src/model/diffusion/mage_flow.hpp index 89df0f482..1b9b9f9be 100644 --- a/otherarch/sdcpp/src/model/diffusion/mage_flow.hpp +++ b/otherarch/sdcpp/src/model/diffusion/mage_flow.hpp @@ -110,13 +110,13 @@ namespace MageFlow { } int batch_size = static_cast(x->ne[3]); - pe_vec = Rope::gen_mage_flow_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - batch_size, - static_cast(context->ne[1]), - ref_latents, - config.theta, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + batch_size, + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim)); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); set_backend_tensor_data(pe, pe_vec.data()); diff --git a/otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp b/otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp index 80c29ab52..ed7c5ac77 100644 --- a/otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp +++ b/otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp @@ -264,6 +264,9 @@ namespace MiniMaxH3 { for (int64_t i = 0; i < num_layers; ++i) { auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); x = block->forward(ctx, x); + sd::ggml_graph_cut::mark_graph_cut(x, + "minimax_h3.token_refiner.blocks." + std::to_string(i), + "hidden_states"); } return std::dynamic_pointer_cast(blocks["final_norm"])->forward(ctx, x); } @@ -527,7 +530,11 @@ namespace MiniMaxH3 { GGML_ASSERT(context->ne[0] == config.text_dim); auto condition_proj = std::dynamic_pointer_cast(blocks["condition_proj"]); auto token_refiner = std::dynamic_pointer_cast(blocks["token_refiner"]); - return token_refiner->forward(ctx, condition_proj->forward(ctx, context)); + auto projected = condition_proj->forward(ctx, context); + sd::ggml_graph_cut::mark_graph_cut(projected, + "minimax_h3.condition_proj", + "hidden_states"); + return token_refiner->forward(ctx, projected); } ggml_tensor* time_embedding(GGMLRunnerContext* ctx, diff --git a/otherarch/sdcpp/src/model/diffusion/minit2i.hpp b/otherarch/sdcpp/src/model/diffusion/minit2i.hpp index 06c4db5a5..0ddacf02e 100644 --- a/otherarch/sdcpp/src/model/diffusion/minit2i.hpp +++ b/otherarch/sdcpp/src/model/diffusion/minit2i.hpp @@ -154,18 +154,26 @@ namespace MiniT2I { return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast(length - 1), length), head_dim, 10000.f)); } - inline std::vector make_vision_rope(int side, int head_dim) { + inline Rope::Embedding make_vision_rope(int side, int head_dim) { GGML_ASSERT(head_dim % 4 == 0); int dim = head_dim / 2; int quarter = dim / 2; int length = side * side; + Rope::Embedding result; + result.positions.append_image(side, side); std::vector out(static_cast(length) * (head_dim / 2) * 4); std::vector freqs(quarter); for (int i = 0; i < quarter; ++i) { freqs[i] = 1.0f / std::pow(10000.0f, static_cast(2 * i) / static_cast(dim)); } + for (int axis : {1, 2}) { + for (float frequency : freqs) { + result.frequencies.push_back({static_cast(axis), frequency}); + } + } for (int y = 0; y < side; ++y) { for (int x = 0; x < side; ++x) { + result.ids.push_back({0.f, static_cast(y), static_cast(x)}); int pos = y * side + x; size_t base = static_cast(pos) * (head_dim / 2) * 4; for (int i = 0; i < quarter; ++i) { @@ -182,7 +190,8 @@ namespace MiniT2I { } } } - return out; + result.values = std::move(out); + return result; } struct SwiGLUMlp : public GGMLBlock { @@ -475,6 +484,8 @@ namespace MiniT2I { int64_t cached_txt_len = -1; int64_t cached_hidden_size = -1; int64_t cached_head_dim = -1; + bool cached_circular_x = false; + bool cached_circular_y = false; MiniT2IRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -521,6 +532,8 @@ namespace MiniT2I { cached_txt_len == txt_len && cached_hidden_size == config.hidden_size && cached_head_dim == config.head_dim && + cached_circular_x == circular_x_enabled && + cached_circular_y == circular_y_enabled && cached_pos_embed != nullptr && cached_txt_pe != nullptr && cached_joint_pe != nullptr) { @@ -531,7 +544,7 @@ namespace MiniT2I { auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast(img_side), static_cast(config.hidden_size)); auto txt_pe_vec = make_text_rope(static_cast(txt_len), static_cast(config.head_dim)); - auto img_pe_vec = make_vision_rope(static_cast(img_side), static_cast(config.head_dim)); + auto img_pe_vec = finish_rope_pe(make_vision_rope(static_cast(img_side), static_cast(config.head_dim))); auto joint_pe_vec = txt_pe_vec; joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end()); @@ -561,6 +574,8 @@ namespace MiniT2I { cached_txt_len = txt_len; cached_hidden_size = config.hidden_size; cached_head_dim = config.head_dim; + cached_circular_x = circular_x_enabled; + cached_circular_y = circular_y_enabled; } ggml_cgraph* build_graph(const sd::Tensor& x_tensor, diff --git a/otherarch/sdcpp/src/model/diffusion/model.hpp b/otherarch/sdcpp/src/model/diffusion/model.hpp index cfbd5c40a..c82b9d718 100644 --- a/otherarch/sdcpp/src/model/diffusion/model.hpp +++ b/otherarch/sdcpp/src/model/diffusion/model.hpp @@ -7,7 +7,7 @@ #include "core/ggml_runner.h" #include "core/tensor_ggml.hpp" -#include "model/common/rope.hpp" +#include "model/common/rope_circular.hpp" #include "model_manager.h" enum class RefImageResizeMode { @@ -39,6 +39,9 @@ const std::unordered_map REF_IMAGE_PRESETS = { {"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}}, {"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}}, {"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}}, + // pass_to_vlm routes the reference image to the conditioner, which is where LLaDA-Image's + // SigVQ encoder lives; it does its own half-resolution resize. + {"llada_image", {true, true, Rope::RefIndexMode::FIXED, true, true, -1, RefImageResizeMode::NONE, -1, -1, true}}, {"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}}, }; @@ -68,6 +71,8 @@ struct AnimaDiffusionExtra { struct QwenImage21DiffusionExtra { const sd::Tensor* image_slots = nullptr; + // Nonzero IDs identify immutable prefix inputs within one sampling run. + uint64_t prefix_id = 0; }; struct WanDiffusionExtra { @@ -131,6 +136,11 @@ struct HunyuanVideoDiffusionExtra { const sd::Tensor* timestep_r = nullptr; }; +struct LLaDAImageDiffusionExtra { + // SigVQ semantic features of the reference image; present only in editing mode. + const sd::Tensor* semantic = nullptr; +}; + using DiffusionExtraParams = std::variant; + HunyuanVideoDiffusionExtra, + LLaDAImageDiffusionExtra>; struct DiffusionParams { const sd::Tensor* x = nullptr; @@ -173,6 +184,11 @@ struct DiffusionModelRunner : public GGMLRunner { protected: std::string prefix; + std::vector finish_rope_pe(Rope::Embedding embedding) { + Rope::apply_circular(embedding, circular_x_enabled, circular_y_enabled); + return std::move(embedding.values); + } + public: DiffusionModelRunner(ggml_backend_t backend, const std::string& prefix, diff --git a/otherarch/sdcpp/src/model/diffusion/pid.hpp b/otherarch/sdcpp/src/model/diffusion/pid.hpp index 59eac2950..6edc83422 100644 --- a/otherarch/sdcpp/src/model/diffusion/pid.hpp +++ b/otherarch/sdcpp/src/model/diffusion/pid.hpp @@ -135,13 +135,13 @@ namespace Pid { return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast(length - 1), length), dim, theta)); } - inline std::vector make_rope_2d(int height, - int width, - int dim, - float theta = 10000.f, - float scale = 16.f, - int ref_grid_h = 0, - int ref_grid_w = 0) { + inline Rope::Embedding make_rope_2d(int height, + int width, + int dim, + float theta = 10000.f, + float scale = 16.f, + int ref_grid_h = 0, + int ref_grid_w = 0) { GGML_ASSERT(dim % 4 == 0); return Rope::embed_2d_interleaved(height, width, dim, theta, scale, ref_grid_h, ref_grid_w); } @@ -867,13 +867,13 @@ namespace Pid { int64_t Hs = Hp / config.patch_size; int64_t Ws = Wp / config.patch_size; - pos_img_vec = make_rope_2d(static_cast(Hs), - static_cast(Ws), - static_cast(config.hidden_size / config.num_groups), - 10000.f, - 16.f, - static_cast(config.rope_ref_grid_h), - static_cast(config.rope_ref_grid_w)); + pos_img_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), + static_cast(Ws), + static_cast(config.hidden_size / config.num_groups), + 10000.f, + 16.f, + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w))); auto pos_img = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, @@ -904,13 +904,13 @@ namespace Pid { 1); set_backend_tensor_data(pixel_pos, pixel_pos_vec.data()); - pixel_pos_comp_vec = make_rope_2d(static_cast(Hs), - static_cast(Ws), - static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), - 10000.f, - 16.f, - static_cast(config.rope_ref_grid_h), - static_cast(config.rope_ref_grid_w)); + pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), + static_cast(Ws), + static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), + 10000.f, + 16.f, + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w))); auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, diff --git a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp index 07364958f..82d173005 100644 --- a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp @@ -635,18 +635,16 @@ namespace Qwen { ref_index_mode = Rope::RefIndexMode::DECREASE; } - pe_vec = Rope::gen_qwen_image_pe(time_len, - static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - batch_size, - static_cast(context->ne[1]), - ref_latents, - ref_index_mode, - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_qwen_image_pe(time_len, + static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + batch_size, + static_cast(context->ne[1]), + ref_latents, + ref_index_mode, + config.theta, + config.axes_dim)); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); // LOG_VERBOSE("pos_len %d", pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); diff --git a/otherarch/sdcpp/src/model/diffusion/qwen_image_2_1.hpp b/otherarch/sdcpp/src/model/diffusion/qwen_image_2_1.hpp index 5b5872a49..09311c290 100644 --- a/otherarch/sdcpp/src/model/diffusion/qwen_image_2_1.hpp +++ b/otherarch/sdcpp/src/model/diffusion/qwen_image_2_1.hpp @@ -68,6 +68,7 @@ namespace Qwen { std::vector segments; std::vector> positions; int64_t prefix_length = 0; + Rope::PositionLayout rope_layout; static QwenImage21Layout build(int64_t text_length, const sd::Tensor& image_slots, @@ -82,6 +83,7 @@ namespace Qwen { auto [height, width] = image_shapes[index]; int64_t start = static_cast(layout.positions.size()); layout.segments.push_back({start, start + height * width, context_start, index}); + layout.rope_layout.append_image(static_cast(height), static_cast(width)); for (int64_t h = 0; h < height; ++h) { for (int64_t w = 0; w < width; ++w) { layout.positions.push_back({static_cast(position), @@ -106,6 +108,7 @@ namespace Qwen { } else { int64_t start = static_cast(layout.positions.size()); layout.segments.push_back({start, start + i - begin, begin, -1}); + layout.rope_layout.append_tokens(i - begin); for (int64_t j = begin; j < i; ++j, ++position) { float p = static_cast(position); layout.positions.push_back({p, p, p}); @@ -121,6 +124,20 @@ namespace Qwen { } }; + struct QwenImage21PrefixCache { + enum class Mode { + NONE, + STORE, + REUSE + }; + Mode mode = Mode::NONE; + std::string name; + std::string cut_group; + int64_t prefix_length = 0; + ggml_type type = GGML_TYPE_F32; + bool* flash_attn_used = nullptr; + }; + class QwenImage21ZeroCenterRMSNorm : public RMSNorm { public: using RMSNorm::RMSNorm; @@ -160,27 +177,70 @@ namespace Qwen { } } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector& segments, const std::vector& masks) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector& segments, const std::vector& masks, const QwenImage21PrefixCache& cache) { int64_t heads = x->ne[0] / dim_head; auto project = [&](const char* name) { auto h = std::dynamic_pointer_cast(blocks[name])->forward(ctx, x); return ggml_reshape_4d(ctx->ggml_ctx, h, dim_head, heads, x->ne[1], x->ne[2]); }; - auto q = project("to_q"); - auto k = project("to_k"); - auto v = project("to_v"); - q = std::dynamic_pointer_cast(blocks["norm_q"])->forward(ctx, q); - k = std::dynamic_pointer_cast(blocks["norm_k"])->forward(ctx, k); - q = Rope::apply_rope(ctx->ggml_ctx, q, pe); - k = Rope::apply_rope(ctx->ggml_ctx, k, pe); + auto q = project("to_q"); + auto k = project("to_k"); + auto v = project("to_v"); + q = std::dynamic_pointer_cast(blocks["norm_q"])->forward(ctx, q); + k = std::dynamic_pointer_cast(blocks["norm_k"])->forward(ctx, k); + q = Rope::apply_rope(ctx->ggml_ctx, q, pe); + k = Rope::apply_rope(ctx->ggml_ctx, k, pe); + if (cache.mode == QwenImage21PrefixCache::Mode::STORE) { + // Preserve query-first attention evaluation while writing each layer's + // prefix before its full-sequence K/V can accumulate across layers. + ctx->expand_graph(q); + auto persist = [&](ggml_tensor* tensor, int axis, const char* name) { + auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length); + // Pack the contiguous data into wider rows so quantization blocks + // can exceed head_dim without padding or changing element order. + part = ggml_reshape_2d(ctx->ggml_ctx, part, x->ne[0], cache.prefix_length); + auto copy = ggml_cast(ctx->ggml_ctx, part, cache.type); + // Keep the copy in this layer's segment so graph cuts do not + // retain or recompute the full-sequence K/V in the final segment. + sd::ggml_graph_cut::mark_graph_cut(copy, cache.cut_group, name); + ctx->persist_cache_tensor(cache.name + "." + name, copy); + }; + persist(k, 1, "k"); + persist(v, 2, "v"); + } + auto attend = [&](ggml_tensor* aq, ggml_tensor* ak, ggml_tensor* av, ggml_tensor* mask) { + bool used_flash_attn = false; + auto out = ggml_ext_attention_ext(ctx, aq, ak, av, heads, mask, true, ctx->flash_attn_enabled, 1.f, &used_flash_attn); + if (cache.flash_attn_used != nullptr) { + *cache.flash_attn_used &= used_flash_attn; + } + return out; + }; ggml_tensor* result = nullptr; - for (size_t i = 0; i < segments.size(); ++i) { - const auto& segment = segments[i]; - auto sq = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end); - auto sk = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end); - auto sv = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end); - auto out = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled); - result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1); + if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) { + auto prefix_k = ctx->load_cache_tensor(cache.name + ".k"); + auto prefix_v = ctx->load_cache_tensor(cache.name + ".v"); + GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr); + if (prefix_k->type != k->type) { + prefix_k = ggml_cast(ctx->ggml_ctx, prefix_k, k->type); + } + if (prefix_v->type != v->type) { + prefix_v = ggml_cast(ctx->ggml_ctx, prefix_v, v->type); + } + prefix_k = ggml_reshape_4d(ctx->ggml_ctx, prefix_k, dim_head, cache.prefix_length, heads, k->ne[3]); + prefix_v = ggml_reshape_4d(ctx->ggml_ctx, prefix_v, dim_head, heads, cache.prefix_length, v->ne[3]); + k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1); + v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2); + result = attend(q, k, v, nullptr); + } else { + for (size_t i = 0; i < segments.size(); ++i) { + const auto& segment = segments[i]; + auto sq = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end); + auto sk = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end); + auto sv = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end); + auto out = attend(sq, sk, sv, masks[i]); + result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1); + } } auto to_out = std::dynamic_pointer_cast(blocks["to_out.0"]); if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { @@ -193,16 +253,16 @@ namespace Qwen { class QwenImage21TransformerBlock : public GGMLBlock { public: QwenImage21TransformerBlock(const QwenImage21Config& config) { - blocks["img_norm1"] = std::make_shared(config.hidden_size, 1e-6f, false); - blocks["img_norm2"] = std::make_shared(config.hidden_size, 1e-6f, false); - blocks["attn"] = std::make_shared(config); + blocks["img_norm1"] = std::make_shared(config.hidden_size, 1e-6f, false); + blocks["img_norm2"] = std::make_shared(config.hidden_size, 1e-6f, false); + blocks["attn"] = std::make_shared(config); if (config.fused_mlp) { blocks["img_mlp.gate_up"] = std::make_shared(config.hidden_size, 2 * config.intermediate_size, false); } else { blocks["img_mlp.proj"] = std::make_shared(config.hidden_size, config.intermediate_size, false); blocks["img_mlp.gate_layer"] = std::make_shared(config.hidden_size, config.intermediate_size, false); } - blocks["img_mlp.out"] = std::make_shared(config.intermediate_size, config.hidden_size, false); + blocks["img_mlp.out"] = std::make_shared(config.intermediate_size, config.hidden_size, false); } static ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* params, int64_t prefix_length, bool gate = false) { @@ -219,13 +279,14 @@ namespace Qwen { return ggml_concat(ctx, prefix, target, 1); } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks) { - auto h = std::dynamic_pointer_cast(blocks["img_norm1"])->forward(ctx, x); - h = modulate(ctx->ggml_ctx, h, modulation[0], layout.prefix_length); - h = std::dynamic_pointer_cast(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks); - x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], layout.prefix_length, true)); - h = std::dynamic_pointer_cast(blocks["img_norm2"])->forward(ctx, x); - h = modulate(ctx->ggml_ctx, h, modulation[2], layout.prefix_length); + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks, const QwenImage21PrefixCache& cache) { + const int64_t prefix_length = cache.mode == QwenImage21PrefixCache::Mode::REUSE ? 0 : layout.prefix_length; + auto h = std::dynamic_pointer_cast(blocks["img_norm1"])->forward(ctx, x); + h = modulate(ctx->ggml_ctx, h, modulation[0], prefix_length); + h = std::dynamic_pointer_cast(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks, cache); + x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], prefix_length, true)); + h = std::dynamic_pointer_cast(blocks["img_norm2"])->forward(ctx, x); + h = modulate(ctx->ggml_ctx, h, modulation[2], prefix_length); ggml_tensor* gate; auto fused = blocks.find("img_mlp.gate_up"); if (fused != blocks.end()) { @@ -237,9 +298,9 @@ namespace Qwen { gate = std::dynamic_pointer_cast(blocks["img_mlp.gate_layer"])->forward(ctx, h); h = std::dynamic_pointer_cast(blocks["img_mlp.proj"])->forward(ctx, h); } - h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate)); - h = std::dynamic_pointer_cast(blocks["img_mlp.out"])->forward(ctx, h); - return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], layout.prefix_length, true)); + h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate)); + h = std::dynamic_pointer_cast(blocks["img_mlp.out"])->forward(ctx, h); + return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], prefix_length, true)); } }; @@ -261,7 +322,7 @@ namespace Qwen { } } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks, const QwenImage21PrefixCache& cache) { auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), 0); // Runtime flow timesteps already use the [0, 1000] scale. time = ggml_ext_timestep_embedding(ctx->ggml_ctx, time, 256, 10000, 1.f); @@ -269,27 +330,37 @@ namespace Qwen { time = ggml_silu(ctx->ggml_ctx, time); auto modulation = std::dynamic_pointer_cast(blocks["modulation.1"])->forward(ctx, time); auto mod = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0); - auto text = std::dynamic_pointer_cast(blocks["txt_in"])->forward(ctx, context); auto img_in = std::dynamic_pointer_cast(blocks["img_in"]); ggml_tensor* joint = nullptr; - for (const auto& segment : layout.segments) { - ggml_tensor* h; - if (segment.image_index < 0) { - h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start, - segment.context_start + segment.end - segment.start); - } else { - auto image = segment.image_index == static_cast(refs.size()) ? x : refs[segment.image_index]; - h = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1)); + if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) { + joint = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, x, 1, 1)); + } else { + auto text = std::dynamic_pointer_cast(blocks["txt_in"])->forward(ctx, context); + for (const auto& segment : layout.segments) { + ggml_tensor* h; + if (segment.image_index < 0) { + h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start, + segment.context_start + segment.end - segment.start); + } else { + auto image = segment.image_index == static_cast(refs.size()) ? x : refs[segment.image_index]; + h = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1)); + } + joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1); } - joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1); } sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.prelude", "joint"); for (int i = 0; i < config.num_layers; ++i) { - auto block = std::dynamic_pointer_cast(blocks["transformer_blocks." + std::to_string(i)]); - joint = block->forward(ctx, joint, mod, pe, layout, masks); - sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.transformer_blocks." + std::to_string(i), "joint"); + const std::string layer = "transformer_blocks." + std::to_string(i); + auto layer_cache = cache; + layer_cache.name = cache.name + "." + std::to_string(i); + layer_cache.cut_group = "qwen_image_2_1." + layer; + auto block = std::dynamic_pointer_cast(blocks[layer]); + joint = block->forward(ctx, joint, mod, pe, layout, masks, layer_cache); + sd::ggml_graph_cut::mark_graph_cut(joint, layer_cache.cut_group, "joint"); + } + if (cache.mode != QwenImage21PrefixCache::Mode::REUSE) { + joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]); } - joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]); auto scale = std::dynamic_pointer_cast(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]); joint = std::dynamic_pointer_cast(blocks["norm_out.norm"])->forward(ctx, joint); joint = ggml_mul(ctx->ggml_ctx, joint, ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f)); @@ -303,11 +374,61 @@ namespace Qwen { QwenImage21Model model; std::vector pe_data; std::vector> mask_data; + ggml_type prefix_cache_type = GGML_TYPE_COUNT; + bool prefix_cache_enabled = true; + bool prefix_cache_disabled = false; + bool prefix_cache_auto_f32 = false; - QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr weight_manager = nullptr) + static bool supports_prefix_cache_type(ggml_type type) { + if (type == GGML_TYPE_F32) { + return true; + } + const auto* traits = ggml_get_type_traits(type); + if (traits->from_float_ref == nullptr || traits->to_float == nullptr) { + return false; + } + auto cpu = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); + if (cpu == nullptr) { + return false; + } + auto ctx = std::unique_ptr( + ggml_init({3 * ggml_tensor_overhead(), nullptr, true}), ggml_free); + if (ctx == nullptr) { + return false; + } + // Some reference quantizers have no runtime copy support. Query the + // device through the registry so dynamically loaded CPU backends work. + auto source = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, ggml_blck_size(type)); + auto encoded = ggml_cast(ctx.get(), source, type); + auto decoded = ggml_cast(ctx.get(), encoded, GGML_TYPE_F32); + return ggml_backend_dev_supports_op(cpu, encoded) && ggml_backend_dev_supports_op(cpu, decoded); + } + + QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr weight_manager = nullptr, const char* model_args = nullptr) : DiffusionModelRunner(backend, prefix, weight_manager), config(QwenImage21Config::detect_from_weights(weights, prefix)), model(config) { + for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) { + if (key == "qwen_image_2_1_prefix_cache" && !parse_strict_bool(value, prefix_cache_enabled)) { + LOG_WARN("ignoring invalid Qwen Image 2.1 model arg '%s=%s'", key.c_str(), value.c_str()); + } else if (key == "qwen_image_2_1_prefix_cache_type") { + if (value == "auto") { + prefix_cache_type = GGML_TYPE_COUNT; + continue; + } + const auto type = sd_type_to_ggml_type(str_to_sd_type(value.c_str())); + if (type == GGML_TYPE_COUNT) { + LOG_WARN("ignoring unknown Qwen Image 2.1 cache type '%s'", value.c_str()); + } else if (!supports_prefix_cache_type(type)) { + LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': runtime conversion to and from F32 is unavailable", value.c_str()); + } else if (config.hidden_size % ggml_blck_size(type) != 0) { + LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': block size %" PRId64 " does not divide hidden size %" PRId64, + value.c_str(), ggml_blck_size(type), config.hidden_size); + } else { + prefix_cache_type = type; + } + } + } model.init(params_ctx, weights, prefix); } @@ -317,6 +438,20 @@ namespace Qwen { model.get_param_tensors(tensors, prefix); } + bool has_prefix_cache(const QwenImage21PrefixCache& cache) { + for (int i = 0; i < config.num_layers; ++i) { + const auto name = cache.name + "." + std::to_string(i); + auto k = get_cache_tensor_by_name(name + ".k"); + auto v = get_cache_tensor_by_name(name + ".v"); + if (k == nullptr || v == nullptr || k->type != cache.type || v->type != cache.type || + k->ne[0] != config.hidden_size || k->ne[1] != cache.prefix_length || k->ne[2] != 1 || k->ne[3] != 1 || + v->ne[0] != config.hidden_size || v->ne[1] != cache.prefix_length || v->ne[2] != 1 || v->ne[3] != 1) { + return false; + } + } + return true; + } + sd::Tensor compute(int n_threads, const DiffusionParams& inputs) override { const auto& x = tensor_or_empty(inputs.x); const auto& context = tensor_or_empty(inputs.context); @@ -345,38 +480,113 @@ namespace Qwen { LOG_ERROR("%s", error.what()); return {}; } - pe_data = Rope::embed_nd(layout.positions, 1, 10000.f, config.axes_dim); - mask_data.clear(); - for (const auto& segment : layout.segments) { - sd::Tensor mask; - if (segment.image_index < 0) { - mask = sd::Tensor::zeros({segment.end, segment.end - segment.start}); - for (int64_t q = segment.start; q < segment.end; ++q) { - for (int64_t k = q + 1; k < segment.end; ++k) { - mask[k + segment.end * (q - segment.start)] = -INFINITY; - } + if (!runner_started()) { + prefix_cache_disabled = false; + prefix_cache_auto_f32 = false; + } + QwenImage21PrefixCache cache; + if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) { + cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) + + ".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled); + cache.prefix_length = layout.prefix_length; + if (prefix_cache_type != GGML_TYPE_COUNT) { + cache.type = prefix_cache_type; + } else if (!prefix_cache_auto_f32 && flash_attn_enabled && !sage_attn_enabled && + (attn_scale <= 0.f || attn_scale == 1.f)) { + cache.type = GGML_TYPE_F16; + } + cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE; + } + bool flash_attn_used = true; + auto run = [&](const QwenImage21PrefixCache& active_cache) { + flash_attn_used = true; + auto checked_cache = active_cache; + if (prefix_cache_type == GGML_TYPE_COUNT && active_cache.type == GGML_TYPE_F16) { + checked_cache.flash_attn_used = &flash_attn_used; + } + const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE; + const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0); + Rope::Embedding embedding; + embedding.ids.assign(first_position, layout.positions.end()); + const size_t offset = cached ? static_cast(layout.prefix_length) : 0; + embedding.positions.token_count = embedding.ids.size(); + for (auto region : layout.rope_layout.images) { + if (region.begin >= offset) { + region.begin -= offset; + embedding.positions.images.push_back(region); } } - mask_data.push_back(std::move(mask)); - } - auto build = [&]() { - auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, layout.positions.size()); - set_backend_tensor_data(pe, pe_data.data()); - std::vector masks, ref_inputs; - for (const auto& mask : mask_data) { - masks.push_back(mask.empty() ? nullptr : make_input(mask)); + embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies); + pe_data = finish_rope_pe(std::move(embedding)); + mask_data.clear(); + if (!cached) { + for (const auto& segment : layout.segments) { + sd::Tensor mask; + if (segment.image_index < 0) { + mask = sd::Tensor::zeros({segment.end, segment.end - segment.start}); + for (int64_t q = segment.start; q < segment.end; ++q) { + for (int64_t k = q + 1; k < segment.end; ++k) { + mask[k + segment.end * (q - segment.start)] = -INFINITY; + } + } + } + mask_data.push_back(std::move(mask)); + } } - for (const auto& ref : refs) { - ref_inputs.push_back(make_input(ref)); - } - auto ctx = get_context(); - auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), make_input(context), - ref_inputs, pe, layout, masks); - ggml_build_forward_expand(graph, out); - return graph; + auto build = [&]() { + auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, + layout.positions.size() - (cached ? layout.prefix_length : 0)); + set_backend_tensor_data(pe, pe_data.data()); + std::vector masks, ref_inputs; + for (const auto& mask : mask_data) { + masks.push_back(mask.empty() ? nullptr : make_input(mask)); + } + if (!cached) { + for (const auto& ref : refs) { + ref_inputs.push_back(make_input(ref)); + } + } + auto ctx = get_context(graph); + auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), cached ? nullptr : make_input(context), + ref_inputs, pe, layout, masks, checked_cache); + if (!flash_attn_used) { + return static_cast(nullptr); + } + ggml_build_forward_expand(graph, out); + return graph; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim()); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim()); + auto result = run(cache); + if (result.empty() && !flash_attn_used) { + // Casting an F16 cache back to F32 cannot recover its original values. + // Recompute the prefix before executing a graph that falls back from FA. + free_cache_ctx_and_buffer(); + prefix_cache_auto_f32 = true; + cache.type = GGML_TYPE_F32; + cache.mode = QwenImage21PrefixCache::Mode::STORE; + LOG_DEBUG("Qwen Image 2.1: Flash Attention unavailable; using F32 prefix caching for this sampling run"); + result = run(cache); + } + if (result.empty() && last_compute_status() == GGML_STATUS_ALLOC_FAILED && + (cache.mode != QwenImage21PrefixCache::Mode::NONE || !cache_.empty())) { + // The failed graph has ended before persistent inputs are released. + free_cache_ctx_and_buffer(); + prefix_cache_disabled = true; + LOG_WARN("Qwen Image 2.1: insufficient memory for prefix caching; retrying without it for this sampling run"); + return run(QwenImage21PrefixCache{}); + } + if (!result.empty() && cache.mode == QwenImage21PrefixCache::Mode::STORE) { + if (!has_prefix_cache(cache)) { + free_cache_ctx_and_buffer(); + prefix_cache_disabled = true; + LOG_WARN("Qwen Image 2.1: incomplete prefix cache; disabling it for this sampling run"); + } else { + LOG_DEBUG("Qwen Image 2.1: cached prefix %" PRIu64 " (%" PRId64 " tokens, %s)", extra->prefix_id, layout.prefix_length, ggml_type_name(cache.type)); + } + } + return result; } }; } diff --git a/otherarch/sdcpp/src/model/diffusion/sensenova_u1.h b/otherarch/sdcpp/src/model/diffusion/sensenova_u1.h index 22b789fac..1702ed854 100644 --- a/otherarch/sdcpp/src/model/diffusion/sensenova_u1.h +++ b/otherarch/sdcpp/src/model/diffusion/sensenova_u1.h @@ -442,16 +442,9 @@ namespace SenseNovaU1 { k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 2); v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2); } else { - // Keep dedicated graph outputs alive until the runner copies them - // into its persistent cache buffer after graph execution. - auto cache_k = ggml_dup_tensor(ctx->ggml_ctx, k); - cache_k = ggml_cpy(ctx->ggml_ctx, k, cache_k); - ggml_set_output(cache_k); - auto cache_v = ggml_dup_tensor(ctx->ggml_ctx, v); - cache_v = ggml_cpy(ctx->ggml_ctx, v, cache_v); - ggml_set_output(cache_v); - ctx->persist_cache_tensor(layer_cache + ".k", cache_k); - ctx->persist_cache_tensor(layer_cache + ".v", cache_v); + ctx->expand_graph(q); + ctx->persist_cache_tensor(layer_cache + ".k", k); + ctx->persist_cache_tensor(layer_cache + ".v", v); } q = ggml_cont(ctx->ggml_ctx, @@ -687,7 +680,7 @@ namespace SenseNovaU1 { ggml_set_name(attention_mask, "snu15.prefix.attention_mask"); set_backend_tensor_data(attention_mask, attention_mask_vec.data()); - auto runner_ctx = get_context(); + auto runner_ctx = get_context(graph); auto text_model = model.text_model(); auto hidden = text_model->embed(&runner_ctx, ids); hidden = text_model->forward(&runner_ctx, diff --git a/otherarch/sdcpp/src/model/diffusion/z_image.hpp b/otherarch/sdcpp/src/model/diffusion/z_image.hpp index 4ae47e268..d654bd388 100644 --- a/otherarch/sdcpp/src/model/diffusion/z_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/z_image.hpp @@ -131,16 +131,30 @@ namespace ZImage { int64_t num_heads; int64_t num_kv_heads; bool qk_norm; + bool split_qkv; public: - JointAttention(int64_t hidden_size, int64_t head_dim, int64_t num_heads, int64_t num_kv_heads, bool qk_norm) - : head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm) { - blocks["qkv"] = std::make_shared(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false); - float scale = 1.f; - blocks["out"] = std::make_shared(num_heads * head_dim, hidden_size, false, false, false, scale); + JointAttention(int64_t hidden_size, + int64_t head_dim, + int64_t num_heads, + int64_t num_kv_heads, + bool qk_norm, + bool norm_elementwise_affine = true, + bool split_qkv = false) + : head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm), split_qkv(split_qkv) { + float scale = 1.f; + if (split_qkv) { + blocks["to_q"] = std::make_shared(hidden_size, num_heads * head_dim, false); + blocks["to_k"] = std::make_shared(hidden_size, num_kv_heads * head_dim, false); + blocks["to_v"] = std::make_shared(hidden_size, num_kv_heads * head_dim, false); + blocks["to_out.0"] = std::make_shared(num_heads * head_dim, hidden_size, false, false, false, scale); + } else { + blocks["qkv"] = std::make_shared(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false); + blocks["out"] = std::make_shared(num_heads * head_dim, hidden_size, false, false, false, scale); + } if (qk_norm) { - blocks["q_norm"] = std::make_shared(head_dim); - blocks["k_norm"] = std::make_shared(head_dim); + blocks["q_norm"] = std::make_shared(head_dim, 1e-06f, norm_elementwise_affine); + blocks["k_norm"] = std::make_shared(head_dim, 1e-06f, norm_elementwise_affine); } } @@ -151,8 +165,35 @@ namespace ZImage { // x: [N, n_token, hidden_size] int64_t n_token = x->ne[1]; int64_t N = x->ne[2]; - auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); - auto out_proj = std::dynamic_pointer_cast(blocks["out"]); + auto out_proj = std::dynamic_pointer_cast(blocks[split_qkv ? "to_out.0" : "out"]); + + if (split_qkv) { + auto q_proj = std::dynamic_pointer_cast(blocks["to_q"]); + auto k_proj = std::dynamic_pointer_cast(blocks["to_k"]); + auto v_proj = std::dynamic_pointer_cast(blocks["to_v"]); + + if (sd_backend_is(ctx->backend, "ROCm")) { + out_proj->set_scale(1.f / 16.f); + out_proj->set_force_prec_f32(true); + q_proj->set_force_prec_f32(true); + k_proj->set_force_prec_f32(true); + v_proj->set_force_prec_f32(true); + } + + auto q = ggml_reshape_4d(ctx->ggml_ctx, q_proj->forward(ctx, x), head_dim, num_heads, n_token, N); + auto k = ggml_reshape_4d(ctx->ggml_ctx, k_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N); + auto v = ggml_reshape_4d(ctx->ggml_ctx, v_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N); + + if (qk_norm) { + q = std::dynamic_pointer_cast(blocks["q_norm"])->forward(ctx, q); + k = std::dynamic_pointer_cast(blocks["k_norm"])->forward(ctx, k); + } + + auto out = Rope::attention(ctx, q, k, v, pe, mask, 1.f / 128.f); + return out_proj->forward(ctx, out); + } + + auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); if (sd_backend_is(ctx->backend, "ROCm")) { out_proj->set_scale(1.f / 16.f); @@ -252,9 +293,12 @@ namespace ZImage { ggml_tensor* x, ggml_tensor* scale) { // x: [N, L, C] - // scale: [N, C] - scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C] - x = ggml_add(ctx, x, ggml_mul(ctx, x, scale)); + // scale: [N, C], or [N, L, C] when the caller modulates per token (LLaDA-Image editing + // feeds a per-token timestep embedding so each segment carries its own modulation). + if (scale->ne[1] != x->ne[1]) { + scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C] + } + x = ggml_add(ctx, x, ggml_mul(ctx, x, scale)); return x; } @@ -272,14 +316,16 @@ namespace ZImage { float ffn_dim_multiplier, float norm_eps, bool qk_norm, - bool modulation = true) + bool modulation = true, + bool norm_elementwise_affine = true, + bool split_qkv = false) : modulation(modulation) { - blocks["attention"] = std::make_shared(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm); + blocks["attention"] = std::make_shared(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm, norm_elementwise_affine, split_qkv); blocks["feed_forward"] = std::make_shared(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier); - blocks["attention_norm1"] = std::make_shared(hidden_size, norm_eps); - blocks["ffn_norm1"] = std::make_shared(hidden_size, norm_eps); - blocks["attention_norm2"] = std::make_shared(hidden_size, norm_eps); - blocks["ffn_norm2"] = std::make_shared(hidden_size, norm_eps); + blocks["attention_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["ffn_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["attention_norm2"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["ffn_norm2"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); if (modulation) { blocks["adaLN_modulation.0"] = std::make_shared(MIN(hidden_size, ADALN_EMBED_DIM), 4 * hidden_size); } @@ -596,18 +642,16 @@ namespace ZImage { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = Rope::gen_z_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - SEQ_MULTI_OF, - ref_latents, - ref_index_mode, - config.theta, - circular_y_enabled, - circular_x_enabled, - config.axes_dim); + pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + SEQ_MULTI_OF, + ref_latents, + ref_index_mode, + config.theta, + config.axes_dim)); int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); // LOG_VERBOSE("pos_len %d", pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); diff --git a/otherarch/sdcpp/src/model/te/llada_image_te.hpp b/otherarch/sdcpp/src/model/te/llada_image_te.hpp new file mode 100644 index 000000000..1c97a4529 --- /dev/null +++ b/otherarch/sdcpp/src/model/te/llada_image_te.hpp @@ -0,0 +1,604 @@ +#ifndef __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__ +#define __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__ + +#include +#include +#include + +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "model/common/ggml_block.hpp" +#include "model_loader.h" + +// The conditioning components LLaDA-Image puts around its LLaDA2-MoE backbone. +// Ref: LLaDAImageQueryFormerModel / LLaDAImageTextProjectionModel in +// https://github.com/inclusionAI/LLaDA-Image/blob/main/src/models/transformer_llada_image.py +// +// QueryFormer turns the LLaDA token embeddings into 256 learned queries that the pipeline +// appends to the backbone input; TextProjection maps the backbone hidden states to the +// denoiser's caption dimension. Neither uses RoPE, and every norm is parameter-free. +// Both MLPs use the tanh GELU approximation, so ggml_gelu (not ggml_gelu_erf). +// +// SigVQ is the editing-only image encoder: a 40-layer ViT whose output is quantized against a +// 16384-entry codebook, with the resulting ids embedded and projected into the semantic features +// the denoiser consumes. Its MLP uses the exact erf GELU, unlike the two above. + +namespace LLaDAImageTE { + constexpr int LLADA_IMAGE_TE_GRAPH_SIZE = 16384; + + struct QueryFormerConfig { + int64_t num_queries = 256; + int64_t hidden_size = 2048; + int64_t num_layers = 1; + int64_t num_heads = 16; + int64_t intermediate_size = 8192; + float norm_eps = 1e-6f; + }; + + struct TextProjectionConfig { + int64_t hidden_size = 2048; + int64_t intermediate_size = 8960; + int64_t num_layers = 6; + int64_t num_heads = 32; + int64_t projection_dim = 2560; + float norm_eps = 1e-6f; + }; + + // Cross-attention with a single fused in_proj over q (from the queries) and k/v (from the + // token embeddings). The checkpoint stores in_proj as one [3*hidden, hidden] parameter. + struct QueryAttention : public GGMLBlock { + protected: + int64_t hidden_size; + int64_t num_heads; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + std::string prefix = "") override { + GGMLBlock::init_params(ctx, tensor_storage_map, prefix); + enum ggml_type wtype = get_type(prefix + "in_proj_weight", tensor_storage_map, GGML_TYPE_F32); + params["in_proj_weight"] = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size * 3); + params["in_proj_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size * 3); + } + + public: + QueryAttention(int64_t hidden_size, int64_t num_heads) + : hidden_size(hidden_size), num_heads(num_heads) { + blocks["out_proj"] = std::make_shared(hidden_size, hidden_size, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* query, + ggml_tensor* context, + ggml_tensor* mask = nullptr) { + // query: [N, num_queries, hidden_size], context: [N, n_token, hidden_size] + ggml_context* gctx = ctx->ggml_ctx; + auto out_proj = std::dynamic_pointer_cast(blocks["out_proj"]); + + auto w = params["in_proj_weight"]; + auto b = params["in_proj_bias"]; + + auto slice_w = [&](int64_t index) { + return ggml_ext_slice(gctx, w, 1, index * hidden_size, (index + 1) * hidden_size); + }; + auto slice_b = [&](int64_t index) { + return ggml_ext_slice(gctx, b, 0, index * hidden_size, (index + 1) * hidden_size); + }; + + auto q = ggml_ext_linear(gctx, query, slice_w(0), slice_b(0)); + auto k = ggml_ext_linear(gctx, context, slice_w(1), slice_b(1)); + auto v = ggml_ext_linear(gctx, context, slice_w(2), slice_b(2)); + + auto x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask); // [N, num_queries, hidden_size] + return out_proj->forward(ctx, x); + } + }; + + struct QueryFormerBlock : public GGMLBlock { + protected: + QueryFormerConfig config; + + public: + QueryFormerBlock(const QueryFormerConfig& config) + : config(config) { + blocks["norm_q"] = std::make_shared(config.hidden_size, config.norm_eps, false); + blocks["norm_k"] = std::make_shared(config.hidden_size, config.norm_eps, false); + blocks["cross_attn"] = std::make_shared(config.hidden_size, config.num_heads); + blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps, false); + blocks["mlp.fc1"] = std::make_shared(config.hidden_size, config.intermediate_size, true); + blocks["mlp.fc2"] = std::make_shared(config.intermediate_size, config.hidden_size, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* query, + ggml_tensor* context, + ggml_tensor* mask = nullptr) { + auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); + auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); + auto cross_attn = std::dynamic_pointer_cast(blocks["cross_attn"]); + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto fc1 = std::dynamic_pointer_cast(blocks["mlp.fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["mlp.fc2"]); + + // The reference overwrites query_embeds with its normalized value before the + // residual add, so both residuals here are on normalized activations. + query = norm_q->forward(ctx, query); + auto ctx_n = norm_k->forward(ctx, context); + query = ggml_add(ctx->ggml_ctx, query, cross_attn->forward(ctx, query, ctx_n, mask)); + query = norm1->forward(ctx, query); + + auto h = fc1->forward(ctx, query); + h = ggml_gelu(ctx->ggml_ctx, h); + h = fc2->forward(ctx, h); + return ggml_add(ctx->ggml_ctx, query, h); + } + }; + + struct QueryFormerModel : public GGMLBlock { + protected: + QueryFormerConfig config; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + params["meta_queries"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, config.num_queries); + } + + public: + QueryFormerModel() = default; + QueryFormerModel(const QueryFormerConfig& config) + : config(config) { + for (int i = 0; i < config.num_layers; i++) { + blocks["query_blocks." + std::to_string(i)] = std::make_shared(config); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* inputs_embeds, + ggml_tensor* mask = nullptr) { + // inputs_embeds: [N, n_token, hidden_size] -> [N, num_queries, hidden_size] + auto query = params["meta_queries"]; + query = ggml_reshape_3d(ctx->ggml_ctx, query, config.hidden_size, config.num_queries, 1); + + for (int i = 0; i < config.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["query_blocks." + std::to_string(i)]); + query = block->forward(ctx, query, inputs_embeds, mask); + } + return query; + } + }; + + struct TextProjectionAttention : public GGMLBlock { + protected: + int64_t num_heads; + int64_t head_dim; + + public: + TextProjectionAttention(const TextProjectionConfig& config) + : num_heads(config.num_heads), head_dim(config.hidden_size / config.num_heads) { + blocks["q_proj"] = std::make_shared(config.hidden_size, config.hidden_size, true); + blocks["k_proj"] = std::make_shared(config.hidden_size, config.hidden_size, true); + blocks["v_proj"] = std::make_shared(config.hidden_size, config.hidden_size, true); + blocks["out_proj"] = std::make_shared(config.hidden_size, config.hidden_size, true); + blocks["q_norm"] = std::make_shared(head_dim, config.norm_eps, false); + blocks["k_norm"] = std::make_shared(head_dim, config.norm_eps, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [N, n_token, hidden_size] + ggml_context* gctx = ctx->ggml_ctx; + int64_t n_token = x->ne[1]; + int64_t N = x->ne[2]; + + auto q_proj = std::dynamic_pointer_cast(blocks["q_proj"]); + auto k_proj = std::dynamic_pointer_cast(blocks["k_proj"]); + auto v_proj = std::dynamic_pointer_cast(blocks["v_proj"]); + auto out_proj = std::dynamic_pointer_cast(blocks["out_proj"]); + auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); + auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); + + auto q = q_proj->forward(ctx, x); + auto k = k_proj->forward(ctx, x); + auto v = v_proj->forward(ctx, x); + + q = ggml_reshape_4d(gctx, q, head_dim, num_heads, n_token, N); + k = ggml_reshape_4d(gctx, k, head_dim, num_heads, n_token, N); + q = q_norm->forward(ctx, q); + k = k_norm->forward(ctx, k); + q = ggml_reshape_3d(gctx, q, head_dim * num_heads, n_token, N); + k = ggml_reshape_3d(gctx, k, head_dim * num_heads, n_token, N); + + auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads); + return out_proj->forward(ctx, out); + } + }; + + struct TextProjectionBlock : public GGMLBlock { + public: + TextProjectionBlock(const TextProjectionConfig& config) { + blocks["self_attn"] = std::make_shared(config); + blocks["layer_norm1"] = std::make_shared(config.hidden_size, config.norm_eps, false); + blocks["layer_norm2"] = std::make_shared(config.hidden_size, config.norm_eps, false); + blocks["mlp.fc1"] = std::make_shared(config.hidden_size, config.intermediate_size, true); + blocks["mlp.fc2"] = std::make_shared(config.intermediate_size, config.hidden_size, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto self_attn = std::dynamic_pointer_cast(blocks["self_attn"]); + auto layer_norm1 = std::dynamic_pointer_cast(blocks["layer_norm1"]); + auto layer_norm2 = std::dynamic_pointer_cast(blocks["layer_norm2"]); + auto fc1 = std::dynamic_pointer_cast(blocks["mlp.fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["mlp.fc2"]); + + x = ggml_add(ctx->ggml_ctx, x, self_attn->forward(ctx, layer_norm1->forward(ctx, x))); + + auto h = fc1->forward(ctx, layer_norm2->forward(ctx, x)); + h = ggml_gelu(ctx->ggml_ctx, h); + h = fc2->forward(ctx, h); + return ggml_add(ctx->ggml_ctx, x, h); + } + }; + + struct TextProjectionModel : public GGMLBlock { + protected: + TextProjectionConfig config; + + public: + TextProjectionModel() = default; + TextProjectionModel(const TextProjectionConfig& config) + : config(config) { + for (int i = 0; i < config.num_layers; i++) { + blocks["layers." + std::to_string(i)] = std::make_shared(config); + } + blocks["projector"] = std::make_shared(config.hidden_size, config.projection_dim, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [N, n_token, hidden_size] -> [N, n_token, projection_dim] + for (int i = 0; i < config.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); + x = block->forward(ctx, x); + } + auto projector = std::dynamic_pointer_cast(blocks["projector"]); + return projector->forward(ctx, x); + } + }; + + struct SigVQConfig { + int64_t image_size = 2048; + int patch_size = 16; + int64_t in_channels = 3; + int64_t hidden_size = 1536; + int64_t intermediate_size = 6144; + int64_t num_layers = 40; + int64_t num_heads = 16; + int64_t codebook_size = 16384; + int64_t codebook_embed_dim = 2048; + int64_t semantic_embed_dim = 4096; + float norm_eps = 1e-6f; + }; + + struct SigVQAttention : public GGMLBlock { + protected: + int64_t num_heads; + int64_t head_dim; + + public: + SigVQAttention(const SigVQConfig& config) + : num_heads(config.num_heads), head_dim(config.hidden_size / config.num_heads) { + blocks["qkv"] = std::make_shared(config.hidden_size, config.hidden_size * 3, true); + blocks["proj"] = std::make_shared(config.hidden_size, config.hidden_size, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [N, n_token, hidden_size] + ggml_context* gctx = ctx->ggml_ctx; + auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); + auto out_proj = std::dynamic_pointer_cast(blocks["proj"]); + + int64_t hidden_size = num_heads * head_dim; + auto qkv = qkv_proj->forward(ctx, x); + auto q = ggml_ext_slice(gctx, qkv, 0, 0, hidden_size); + auto k = ggml_ext_slice(gctx, qkv, 0, hidden_size, hidden_size * 2); + auto v = ggml_ext_slice(gctx, qkv, 0, hidden_size * 2, hidden_size * 3); + + auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads); + return out_proj->forward(ctx, out); + } + }; + + struct SigVQBlock : public GGMLBlock { + public: + SigVQBlock(const SigVQConfig& config) { + blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps); + blocks["norm2"] = std::make_shared(config.hidden_size, config.norm_eps); + blocks["attn"] = std::make_shared(config); + blocks["mlp.fc1"] = std::make_shared(config.hidden_size, config.intermediate_size, true); + blocks["mlp.fc2"] = std::make_shared(config.intermediate_size, config.hidden_size, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto fc1 = std::dynamic_pointer_cast(blocks["mlp.fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["mlp.fc2"]); + + x = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, norm1->forward(ctx, x))); + auto h = fc1->forward(ctx, norm2->forward(ctx, x)); + h = ggml_gelu_erf(ctx->ggml_ctx, h); + h = fc2->forward(ctx, h); + return ggml_add(ctx->ggml_ctx, x, h); + } + }; + + struct SigVQModel : public GGMLBlock { + protected: + SigVQConfig config; + + public: + SigVQModel() = default; + SigVQModel(const SigVQConfig& config) + : config(config) { + blocks["visual.patch_embed.proj"] = std::make_shared(config.in_channels, + config.hidden_size, + std::make_pair(config.patch_size, config.patch_size), + std::make_pair(config.patch_size, config.patch_size)); + for (int i = 0; i < config.num_layers; i++) { + blocks["visual.blocks." + std::to_string(i)] = std::make_shared(config); + } + blocks["vqmodel.quant_conv"] = std::make_shared(config.hidden_size, + config.codebook_embed_dim, + std::make_pair(1, 1)); + blocks["prior_projector.net.0.proj"] = std::make_shared(config.semantic_embed_dim, config.semantic_embed_dim, true); + blocks["prior_projector.net.2"] = std::make_shared(config.semantic_embed_dim, config.semantic_embed_dim, true); + } + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + params["visual.embeddings.position_embedding.weight"] = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, (config.image_size / config.patch_size) * (config.image_size / config.patch_size)); + params["vqmodel.quantize.embedding.weight"] = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.codebook_embed_dim, config.codebook_size); + params["prior_token_embedding.weight"] = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.semantic_embed_dim, config.codebook_size); + } + + // Bilinear-resamples the square position-embedding grid onto the image's patch grid. + // The reference uses grid_sample(align_corners=False, padding_mode="border"); the source + // coordinate for output index j is therefore (j + 0.5) * side / out - 0.5, clamped. + ggml_tensor* resample_pos_embed(GGMLRunnerContext* ctx, + ggml_tensor* pos_idx, + ggml_tensor* pos_weight) { + auto pos_embed = params["visual.embeddings.position_embedding.weight"]; + auto gathered = ggml_get_rows(ctx->ggml_ctx, pos_embed, pos_idx); + return ggml_mul(ctx->ggml_ctx, gathered, pos_weight); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* pixel_values, + const std::vector& pos_idx, + const std::vector& pos_weight) { + // pixel_values: [N, in_channels, H, W] -> [N, grid_h * grid_w, semantic_embed_dim] + ggml_context* gctx = ctx->ggml_ctx; + + auto patch_embed = std::dynamic_pointer_cast(blocks["visual.patch_embed.proj"]); + auto quant_conv = std::dynamic_pointer_cast(blocks["vqmodel.quant_conv"]); + auto proj_0 = std::dynamic_pointer_cast(blocks["prior_projector.net.0.proj"]); + auto proj_2 = std::dynamic_pointer_cast(blocks["prior_projector.net.2"]); + + auto x = patch_embed->forward(ctx, pixel_values); // [N, hidden_size, grid_h, grid_w] + int64_t grid_w = x->ne[0]; + int64_t grid_h = x->ne[1]; + int64_t n_token = grid_h * grid_w; + int64_t N = x->ne[3]; + + x = ggml_reshape_3d(gctx, x, n_token, config.hidden_size, N); + x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3)); // [N, n_token, hidden_size] + + ggml_tensor* pos = nullptr; + for (size_t i = 0; i < pos_idx.size(); i++) { + auto corner = resample_pos_embed(ctx, pos_idx[i], pos_weight[i]); + pos = pos == nullptr ? corner : ggml_add(gctx, pos, corner); + } + x = ggml_add(gctx, x, ggml_reshape_3d(gctx, pos, config.hidden_size, n_token, N)); + + for (int i = 0; i < config.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["visual.blocks." + std::to_string(i)]); + x = block->forward(ctx, x); + } + + // quant_conv is 1x1, so run it as a per-token projection rather than reshaping to 2-D. + x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3)); // [N, hidden_size, n_token] + x = ggml_reshape_4d(gctx, x, n_token, 1, config.hidden_size, N); + x = quant_conv->forward(ctx, x); // [N, codebook_embed_dim, 1, n_token] + x = ggml_reshape_3d(gctx, x, n_token, config.codebook_embed_dim, N); + x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3)); // [N, n_token, codebook_embed_dim] + + // Both sides are L2-normalized, so the nearest codebook entry by euclidean distance + // is the one with the largest dot product. + auto codebook = ggml_l2_norm(gctx, params["vqmodel.quantize.embedding.weight"], 1e-12f); + auto normed = ggml_l2_norm(gctx, x, 1e-12f); + auto logits = ggml_mul_mat(gctx, codebook, normed); // [N, n_token, codebook_size] + auto token_ids = ggml_argmax(gctx, ggml_reshape_2d(gctx, logits, config.codebook_size, n_token * N)); + + auto semantic = ggml_get_rows(gctx, params["prior_token_embedding.weight"], token_ids); + semantic = ggml_reshape_3d(gctx, semantic, config.semantic_embed_dim, n_token, N); + + auto h = proj_0->forward(ctx, semantic); + h = ggml_silu(gctx, h); + return proj_2->forward(ctx, h); + } + }; + + struct QueryFormerRunner : public GGMLRunner { + public: + QueryFormerConfig config; + QueryFormerModel query_former; + + QueryFormerRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager) { + query_former = QueryFormerModel(config); + query_former.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "llada_image_queryformer"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) { + query_former.get_param_tensors(tensors, prefix); + } + + sd::Tensor compute(int n_threads, const sd::Tensor& inputs_embeds) { + auto get_graph = [&]() -> ggml_cgraph* { + ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE); + ggml_tensor* x = make_input(inputs_embeds); + auto runner_ctx = get_context(); + ggml_tensor* out = query_former.forward(&runner_ctx, x); + ggml_build_forward_expand(gf, out); + return gf; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), + inputs_embeds.dim()); + } + }; + + struct TextProjectionRunner : public GGMLRunner { + public: + TextProjectionConfig config; + TextProjectionModel text_projection; + + TextProjectionRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager) { + text_projection = TextProjectionModel(config); + text_projection.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "llada_image_text_projection"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) { + text_projection.get_param_tensors(tensors, prefix); + } + + sd::Tensor compute(int n_threads, const sd::Tensor& hidden_states) { + auto get_graph = [&]() -> ggml_cgraph* { + ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE); + ggml_tensor* x = make_input(hidden_states); + auto runner_ctx = get_context(); + ggml_tensor* out = text_projection.forward(&runner_ctx, x); + ggml_build_forward_expand(gf, out); + return gf; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), + hidden_states.dim()); + } + }; + + struct SigVQRunner : public GGMLRunner { + public: + SigVQConfig config; + SigVQModel sigvq; + std::array, 4> pos_idx_data; + std::array, 4> pos_weight_data; + + SigVQRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager) { + sigvq = SigVQModel(config); + sigvq.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "llada_image_sigvq"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) { + sigvq.get_param_tensors(tensors, prefix); + } + + // Precomputes the four bilinear taps that resample the square position-embedding grid + // onto a grid_h x grid_w patch grid, matching grid_sample(align_corners=False, + // padding_mode="border"). + void build_pos_embed_taps(int64_t grid_h, int64_t grid_w) { + const int64_t side = config.image_size / config.patch_size; + for (auto& v : pos_idx_data) { + v.clear(); + } + for (auto& v : pos_weight_data) { + v.clear(); + } + + auto clamp_index = [side](int64_t v) { + return static_cast(std::min(std::max(v, 0), side - 1)); + }; + + for (int64_t i = 0; i < grid_h; ++i) { + double src_h = (static_cast(i) + 0.5) * side / static_cast(grid_h) - 0.5; + int64_t h_floor = static_cast(std::floor(src_h)); + double dh = src_h - static_cast(h_floor); + for (int64_t j = 0; j < grid_w; ++j) { + double src_w = (static_cast(j) + 0.5) * side / static_cast(grid_w) - 0.5; + int64_t w_floor = static_cast(std::floor(src_w)); + double dw = src_w - static_cast(w_floor); + + int32_t h0 = clamp_index(h_floor); + int32_t h1 = clamp_index(h_floor + 1); + int32_t w0 = clamp_index(w_floor); + int32_t w1 = clamp_index(w_floor + 1); + + pos_idx_data[0].push_back(h0 * static_cast(side) + w0); + pos_idx_data[1].push_back(h0 * static_cast(side) + w1); + pos_idx_data[2].push_back(h1 * static_cast(side) + w0); + pos_idx_data[3].push_back(h1 * static_cast(side) + w1); + + pos_weight_data[0].push_back(static_cast((1.0 - dh) * (1.0 - dw))); + pos_weight_data[1].push_back(static_cast((1.0 - dh) * dw)); + pos_weight_data[2].push_back(static_cast(dh * (1.0 - dw))); + pos_weight_data[3].push_back(static_cast(dh * dw)); + } + } + } + + sd::Tensor compute(int n_threads, const sd::Tensor& pixel_values) { + auto get_graph = [&]() -> ggml_cgraph* { + ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE); + ggml_tensor* x = make_input(pixel_values); + + int64_t grid_h = x->ne[1] / config.patch_size; + int64_t grid_w = x->ne[0] / config.patch_size; + build_pos_embed_taps(grid_h, grid_w); + + std::vector pos_idx; + std::vector pos_weight; + for (int i = 0; i < 4; i++) { + auto idx = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, static_cast(pos_idx_data[i].size())); + set_backend_tensor_data(idx, pos_idx_data[i].data()); + auto w = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, 1, static_cast(pos_weight_data[i].size())); + set_backend_tensor_data(w, pos_weight_data[i].data()); + pos_idx.push_back(idx); + pos_weight.push_back(w); + } + + auto runner_ctx = get_context(); + ggml_tensor* out = sigvq.forward(&runner_ctx, x, pos_idx, pos_weight); + ggml_build_forward_expand(gf, out); + return gf; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), 3); + } + }; + +} // namespace LLaDAImageTE + +#endif // __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__ diff --git a/otherarch/sdcpp/src/model/te/llm.hpp b/otherarch/sdcpp/src/model/te/llm.hpp index b289328ef..da59cfab6 100644 --- a/otherarch/sdcpp/src/model/te/llm.hpp +++ b/otherarch/sdcpp/src/model/te/llm.hpp @@ -49,6 +49,7 @@ namespace LLM { GEMMA2_2B, GEMMA4_12B, GPT_OSS_20B, + LLADA2_MOE, ARCH_COUNT, }; @@ -62,6 +63,7 @@ namespace LLM { "gemma2_2b", "gemma4_12b", "gpt_oss_20b", + "llada2_moe", }; enum class MLPActivation { @@ -125,6 +127,17 @@ namespace LLM { std::vector sliding_attention; int64_t num_experts = 0; int64_t num_experts_per_tok = 0; + bool qkv_fused = false; + bool bidirectional = false; + float partial_rotary = 1.f; + + // DeepSeek-V3-style grouped-sigmoid MoE routing (LLaDA2) + int64_t moe_intermediate_size = 0; + int64_t num_shared_experts = 0; + int64_t first_k_dense_replace = 0; + int64_t n_group = 0; + int64_t topk_group = 0; + float routed_scaling_factor = 1.f; LLMVisionConfig vision; bool have_vision_weight = false; bool llama_cpp_style = false; @@ -212,6 +225,31 @@ namespace LLM { config.intermediate_size = 9216; config.num_layers = 26; config.vocab_size = 256000; + } else if (arch == LLMArch::LLADA2_MOE) { + config.head_dim = 128; + config.num_heads = 16; + config.num_kv_heads = 4; + config.qkv_bias = false; + config.attention_out_bias = false; + config.qk_norm = true; + config.rms_norm_eps = 1e-6f; + config.hidden_size = 2048; + config.intermediate_size = 5120; + config.num_layers = 20; + config.vocab_size = 173568; + config.max_position_embeddings = 16384; + config.rope_thetas = {600000.f}; + config.qkv_fused = true; + config.bidirectional = true; + config.partial_rotary = 0.5f; + config.num_experts = 256; + config.num_experts_per_tok = 8; + config.moe_intermediate_size = 512; + config.num_shared_experts = 1; + config.first_k_dense_replace = 1; + config.n_group = 8; + config.topk_group = 4; + config.routed_scaling_factor = 2.5f; } else if (arch == LLMArch::GPT_OSS_20B) { config.head_dim = 64; config.num_heads = 64; @@ -419,6 +457,195 @@ namespace LLM { } }; + // LLaDA2's MoE differs from GPT-OSS's in three ways that all change the result: + // routing scores are sigmoid (not softmax over the selected logits), expert selection is + // group-limited and uses a bias term that the returned weights do NOT include, and the + // experts carry no biases. Ref: LLaDA2MoeGate / LLaDA2MoeSparseMoeBlock in + // modeling_llada2uni_moe.py. + struct LLaDA2MoEMLP : public GGMLBlock { + protected: + int64_t hidden_size; + int64_t moe_intermediate_size; + int64_t num_experts; + int64_t num_experts_per_tok; + int64_t n_group; + int64_t topk_group; + float routed_scaling_factor; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + std::string prefix = "") override { + GGMLBlock::init_params(ctx, tensor_storage_map, prefix); + + auto supported_type = [](ggml_type wtype, int64_t in_features) { + if (in_features % ggml_blck_size(wtype) != 0) { + return GGML_TYPE_F32; + } + return wtype; + }; + + // The reference runs the router in fp32; keep the weight in fp32 so the sigmoid + // scores and the group sums match. + params["gate.weight"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_experts); + params["gate.expert_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts); + + ggml_type gate_type = supported_type(get_type(prefix + "experts.gate_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size); + ggml_type up_type = supported_type(get_type(prefix + "experts.up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size); + ggml_type down_type = supported_type(get_type(prefix + "experts.down_proj.weight", tensor_storage_map, GGML_TYPE_F32), moe_intermediate_size); + + // HF ships the stacked experts as 3-D nn.Parameters, while the ComfyUI GGUF repack + // flattens the expert axis into ne[1]. Declare whichever the file holds - the two are + // bit-identical, and forward() reshapes to 3-D for ggml_mul_mat_id either way. + auto declare_experts = [&](const std::string& name, ggml_type type, int64_t in_dim, int64_t out_dim) { + auto storage = tensor_storage_map.find(prefix + name); + if (storage != tensor_storage_map.end() && storage->second.n_dims == 2) { + GGML_ASSERT(storage->second.nelements() == in_dim * out_dim * num_experts); + params[name] = ggml_new_tensor_2d(ctx, type, in_dim, out_dim * num_experts); + } else { + params[name] = ggml_new_tensor_3d(ctx, type, in_dim, out_dim, num_experts); + } + }; + + declare_experts("experts.gate_proj.weight", gate_type, hidden_size, moe_intermediate_size); + declare_experts("experts.up_proj.weight", up_type, hidden_size, moe_intermediate_size); + declare_experts("experts.down_proj.weight", down_type, moe_intermediate_size, hidden_size); + } + + public: + LLaDA2MoEMLP(const LLMConfig& config) + : hidden_size(config.hidden_size), + moe_intermediate_size(config.moe_intermediate_size), + num_experts(config.num_experts), + num_experts_per_tok(config.num_experts_per_tok), + n_group(config.n_group), + topk_group(config.topk_group), + routed_scaling_factor(config.routed_scaling_factor) { + if (config.num_shared_experts > 0) { + blocks["shared_experts"] = std::make_shared(config.hidden_size, + config.moe_intermediate_size * config.num_shared_experts, + false, + config.mlp_activation); + } + } + + // Reproduces group_limited_topk(): keep the topk_group groups with the highest + // "sum of the two best scores in the group", then take the global top-k among them. + ggml_tensor* group_limited_mask(GGMLRunnerContext* ctx, + ggml_tensor* routing_scores, + int64_t n_token_total) { + ggml_context* gctx = ctx->ggml_ctx; + const int64_t per_group = num_experts / n_group; + + // [experts_per_group, n_group * tokens] so top-2 runs per (group, token) row. + auto grouped = ggml_reshape_2d(gctx, routing_scores, per_group, n_group * n_token_total); + auto best2_idx = ggml_argsort_top_k(gctx, grouped, 2); // [2, n_group * tokens] + auto grouped_val = ggml_reshape_3d(gctx, grouped, 1, per_group, n_group * n_token_total); + auto best2 = ggml_get_rows(gctx, grouped_val, best2_idx); // [1, 2, n_group * tokens] + best2 = ggml_reshape_2d(gctx, best2, 2, n_group * n_token_total); + auto group_score = ggml_reshape_2d(gctx, ggml_sum_rows(gctx, best2), n_group, n_token_total); // [n_group, tokens] + + // Threshold = the topk_group-th largest group score, taken from the sorted top-k. + auto top_groups = ggml_argsort_top_k(gctx, group_score, (int)topk_group); // [topk_group, tokens] + auto group_val = ggml_reshape_3d(gctx, group_score, 1, n_group, n_token_total); + auto top_scores = ggml_get_rows(gctx, group_val, top_groups); // [1, topk_group, tokens] + top_scores = ggml_reshape_2d(gctx, top_scores, topk_group, n_token_total); + auto threshold = ggml_view_2d(gctx, + top_scores, + 1, + n_token_total, + top_scores->nb[1], + (topk_group - 1) * top_scores->nb[0]); // [1, tokens] + threshold = ggml_cont(gctx, threshold); + + // keep = 1 - step(threshold - score). step(0) == 0, so the group sitting exactly on + // the threshold is kept without needing an epsilon. + auto diff = ggml_sub(gctx, ggml_repeat(gctx, threshold, group_score), group_score); + auto keep = ggml_scale_bias(gctx, ggml_step(gctx, diff), -1.f, 1.f); // [n_group, tokens] + + // 0 for kept groups, a large negative for dropped ones, broadcast over the group. + auto additive = ggml_scale_bias(gctx, keep, 1e30f, -1e30f); + additive = ggml_reshape_3d(gctx, additive, 1, n_group, n_token_total); + auto expanded = ggml_repeat_4d(gctx, additive, per_group, n_group, n_token_total, 1); + return ggml_reshape_2d(gctx, expanded, num_experts, n_token_total); + } + + ggml_tensor* expert_linear(GGMLRunnerContext* ctx, + const std::string& weight_name, + ggml_tensor* x, + ggml_tensor* selected_experts) { + ggml_tensor* w = params[weight_name]; + if (w->ne[2] != num_experts) { + // Flattened layout: split the expert axis back out. ne[0] is untouched, so this + // stays valid for quantized types. + w = ggml_reshape_3d(ctx->ggml_ctx, w, w->ne[0], w->ne[1] / num_experts, num_experts); + } + return ggml_mul_mat_id(ctx->ggml_ctx, w, x, selected_experts); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [N, n_token, hidden_size] + GGML_ASSERT(num_experts > 0 && num_experts_per_tok > 0); + GGML_ASSERT(n_group > 0 && topk_group > 0 && num_experts % n_group == 0); + + ggml_context* gctx = ctx->ggml_ctx; + const int64_t n_token = x->ne[1]; + const int64_t N = x->ne[2]; + const int64_t n_token_total = n_token * N; + + auto identity = x; + + auto logits = ggml_mul_mat(gctx, params["gate.weight"], x); + logits = ggml_reshape_2d(gctx, logits, num_experts, n_token_total); + auto scores = ggml_sigmoid(gctx, logits); // [num_experts, tokens] + + // The bias steers selection only; the combine weights come from the unbiased scores. + auto routing = ggml_add(gctx, scores, params["gate.expert_bias"]); + routing = ggml_add(gctx, routing, group_limited_mask(ctx, routing, n_token_total)); + + auto selected_experts = ggml_argsort_top_k(gctx, routing, (int)num_experts_per_tok); // [top_k, tokens] + auto score_rows = ggml_reshape_3d(gctx, scores, 1, num_experts, n_token_total); + auto weights = ggml_get_rows(gctx, score_rows, selected_experts); // [1, top_k, tokens] + weights = ggml_reshape_2d(gctx, weights, num_experts_per_tok, n_token_total); + + if (num_experts_per_tok > 1) { + auto denom = ggml_scale_bias(gctx, ggml_sum_rows(gctx, weights), 1.f, 1e-20f); // [1, tokens] + weights = ggml_div(gctx, weights, ggml_repeat(gctx, denom, weights)); + } + weights = ggml_scale(gctx, weights, routed_scaling_factor); + weights = ggml_reshape_3d(gctx, weights, 1, num_experts_per_tok, n_token_total); + + auto xf = ggml_reshape_3d(gctx, x, hidden_size, 1, n_token_total); + auto gate = expert_linear(ctx, "experts.gate_proj.weight", xf, selected_experts); + auto up = expert_linear(ctx, "experts.up_proj.weight", xf, selected_experts); + auto activated = ggml_swiglu_split(gctx, gate, up); + auto experts = expert_linear(ctx, "experts.down_proj.weight", activated, selected_experts); + experts = ggml_mul(gctx, experts, weights); + + ggml_tensor* out = nullptr; + for (int64_t i = 0; i < num_experts_per_tok; ++i) { + auto expert_out = ggml_view_2d(gctx, + experts, + hidden_size, + n_token_total, + experts->nb[2], + i * experts->nb[1]); + out = out == nullptr ? expert_out : ggml_add(gctx, out, expert_out); + } + if (num_experts_per_tok == 1) { + out = ggml_cont(gctx, out); + } + out = ggml_reshape_3d(gctx, out, hidden_size, n_token, N); + + auto shared_it = blocks.find("shared_experts"); + if (shared_it != blocks.end()) { + auto shared_experts = std::dynamic_pointer_cast(shared_it->second); + out = ggml_add(gctx, out, shared_experts->forward(ctx, identity)); + } + + return out; + } + }; + struct GPTOSSMLP : public GGMLBlock { protected: int64_t hidden_size; @@ -605,21 +832,31 @@ namespace LLM { } txt_token_end = image_embeds[i].first; - auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); - if (input_embed == nullptr) { - input_embed = txt_embed; - } else { - input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1); + // An embed can sit flush against the previous one or at the very start/end of the + // sequence, leaving no text tokens to splice around it. + if (txt_token_end > txt_token_start) { + auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); + if (input_embed == nullptr) { + input_embed = txt_embed; + } else { + input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1); + } } - input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1); + if (input_embed == nullptr) { + input_embed = image_embeds[i].second; + } else { + input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1); + } } txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1]; txt_token_end = raw_x->ne[1]; - auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); - input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1); + if (txt_token_end > txt_token_start) { + auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); + input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1); + } GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]); return input_embed; } @@ -1122,6 +1359,7 @@ namespace LLM { bool k_eq_v; bool v_norm; bool unscaled_attention; + bool qkv_fused; float rms_norm_eps; int rope_pairs; @@ -1147,12 +1385,20 @@ namespace LLM { k_eq_v(global_layer && config.global_k_eq_v), v_norm(config.v_norm), unscaled_attention(config.unscaled_attention), + qkv_fused(config.qkv_fused), rms_norm_eps(config.rms_norm_eps), rope_pairs(0) { - blocks["q_proj"] = std::make_shared(config.hidden_size, num_heads * head_dim, config.qkv_bias); - blocks["k_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); - if (!k_eq_v) { - blocks["v_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); + if (qkv_fused) { + // The checkpoint ships q, k and v as one tensor and the loader cannot split a + // source tensor, so keep it fused and slice it in forward(). + GGML_ASSERT(!k_eq_v); + blocks["query_key_value"] = std::make_shared(config.hidden_size, (num_heads + num_kv_heads * 2) * head_dim, config.qkv_bias); + } else { + blocks["q_proj"] = std::make_shared(config.hidden_size, num_heads * head_dim, config.qkv_bias); + blocks["k_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); + if (!k_eq_v) { + blocks["v_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); + } } blocks["o_proj"] = std::make_shared(num_heads * head_dim, config.hidden_size, config.attention_out_bias); if (config.qk_norm) { @@ -1161,7 +1407,7 @@ namespace LLM { } // Proportional RoPE rotates only the leading `rope_pairs` dimension pairs of the head; // the rest are left unrotated through freq_factors (see rope_freq_factors()). - float partial = global_layer ? config.global_partial_rotary : 1.f; + float partial = global_layer && config.global_partial_rotary != 1.f ? config.global_partial_rotary : config.partial_rotary; rope_pairs = static_cast(partial * head_dim / 2.f); } @@ -1186,14 +1432,28 @@ namespace LLM { // x: [N, n_token, hidden_size] int64_t n_token = x->ne[1]; int64_t N = x->ne[2]; - auto q_proj = std::dynamic_pointer_cast(blocks["q_proj"]); - auto k_proj = std::dynamic_pointer_cast(blocks["k_proj"]); - auto v_proj = k_eq_v ? nullptr : std::dynamic_pointer_cast(blocks["v_proj"]); auto out_proj = std::dynamic_pointer_cast(blocks["o_proj"]); - auto q = q_proj->forward(ctx, x); // [N, n_token, num_heads*head_dim] - auto k = k_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim] - auto v = k_eq_v ? k : v_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim] + ggml_tensor* q = nullptr; + ggml_tensor* k = nullptr; + ggml_tensor* v = nullptr; + if (qkv_fused) { + auto qkv_proj = std::dynamic_pointer_cast(blocks["query_key_value"]); + auto qkv = qkv_proj->forward(ctx, x); // [N, n_token, (num_heads + num_kv_heads*2)*head_dim] + int64_t q_len = num_heads * head_dim; + int64_t k_len = num_kv_heads * head_dim; + q = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, 0, q_len); + k = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len, q_len + k_len); + v = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len + k_len, q_len + k_len * 2); + } else { + auto q_proj = std::dynamic_pointer_cast(blocks["q_proj"]); + auto k_proj = std::dynamic_pointer_cast(blocks["k_proj"]); + auto v_proj = k_eq_v ? nullptr : std::dynamic_pointer_cast(blocks["v_proj"]); + + q = q_proj->forward(ctx, x); // [N, n_token, num_heads*head_dim] + k = k_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim] + v = k_eq_v ? k : v_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim] + } q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, n_token, N); // [N, n_token, num_heads, head_dim] k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_kv_heads, n_token, N); // [N, n_token, num_kv_heads, head_dim] @@ -1336,6 +1596,38 @@ namespace LLM { 1.f, 32.f, 1.f); + } else if (arch == LLMArch::LLADA2_MOE) { + // LLaDA2 slices the head (query[..., :rotary_dim]) instead of zero-padding + // inv_freq like gemma does, so rotate_half pairs i with i + rotary_dim/2 and the + // frequencies use rotary_dim as the exponent denominator. Passing n_dims = + // rotary_dim reproduces both; freq_factors would give the wrong pairing. + int rotary_dim = rope_pairs * 2; + q = ggml_rope_ext(ctx->ggml_ctx, + q, + input_pos, + nullptr, + rotary_dim, + GGML_ROPE_TYPE_NEOX, + static_cast(max_position_embeddings), + rope_thetas[0], + 1.f, + 0.f, + 1.f, + 32.f, + 1.f); + k = ggml_rope_ext(ctx->ggml_ctx, + k, + input_pos, + nullptr, + rotary_dim, + GGML_ROPE_TYPE_NEOX, + static_cast(max_position_embeddings), + rope_thetas[0], + 1.f, + 0.f, + 1.f, + 32.f, + 1.f); } else if (arch == LLMArch::QWEN3_VL) { int sections[4] = {24, 20, 20, 0}; q = ggml_rope_multi(ctx->ggml_ctx, q, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_IMROPE, 262144, 5000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); @@ -1432,6 +1724,8 @@ namespace LLM { blocks["self_attn"] = std::make_shared(config, sliding_attention == 0); if (config.arch == LLMArch::GPT_OSS_20B) { blocks["mlp"] = std::make_shared(config); + } else if (config.arch == LLMArch::LLADA2_MOE && layer_index >= config.first_k_dense_replace) { + blocks["mlp"] = std::make_shared(config); } else { blocks["mlp"] = std::make_shared(config.hidden_size, config.intermediate_size, @@ -1485,6 +1779,10 @@ namespace LLM { if (arch == LLMArch::GPT_OSS_20B) { auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); x = mlp->forward(ctx, x); + } else if (auto moe_mlp = std::dynamic_pointer_cast(blocks["mlp"])) { + // LLaDA2 is dense for the first first_k_dense_replace layers and MoE afterwards, + // so the block type varies per layer rather than per arch. + x = moe_mlp->forward(ctx, x); } else { auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); x = mlp->forward(ctx, x); @@ -1650,6 +1948,11 @@ namespace LLM { return x; } + ggml_tensor* embed(GGMLRunnerContext* ctx, ggml_tensor* input_ids) { + auto model = std::dynamic_pointer_cast(blocks["model"]); + return model->embed(ctx, input_ids); + } + std::shared_ptr vision_model() { GGML_ASSERT(enable_vision); return std::dynamic_pointer_cast(blocks["visual"]); @@ -1990,7 +2293,8 @@ namespace LLM { config.arch == LLMArch::GEMMA3_12B || config.arch == LLMArch::GEMMA4_12B || config.arch == LLMArch::GEMMA2_2B || - config.arch == LLMArch::GPT_OSS_20B) { + config.arch == LLMArch::GPT_OSS_20B || + config.arch == LLMArch::LLADA2_MOE) { input_pos_vec.resize(n_tokens); for (int i = 0; i < n_tokens; ++i) { input_pos_vec[i] = i; @@ -2042,8 +2346,9 @@ namespace LLM { attention_mask_vec.resize(n_tokens * n_tokens); for (int i0 = 0; i0 < n_tokens; i0++) { for (int i1 = 0; i1 < n_tokens; i1++) { + // Diffusion LLMs attend in both directions; only causal LMs get the triangle. float value = 0.f; - if (i0 > i1) { + if (!config.bidirectional && i0 > i1) { value = -INFINITY; } attention_mask_vec[i1 * n_tokens + i0] = value; @@ -2115,6 +2420,22 @@ namespace LLM { input_ids.dim() + 1); } + // LLaDA-Image's QueryFormer consumes the raw token embeddings before the backbone runs, + // so it needs the embedding lookup on its own. + sd::Tensor compute_input_embeds(const int n_threads, + const sd::Tensor& input_ids) { + auto get_graph = [&]() -> ggml_cgraph* { + ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE); + ggml_tensor* ids = make_input(input_ids); + auto runner_ctx = get_context(); + ggml_tensor* out = model.embed(&runner_ctx, ids); + ggml_build_forward_expand(gf, out); + return gf; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), + input_ids.dim() + 1); + } + int64_t get_num_image_tokens(int64_t t, int64_t h, int64_t w) { int64_t grid_t = 1; int64_t grid_h = h / config.vision.patch_size; @@ -2370,11 +2691,13 @@ namespace LLM { pad_id = 199999; } else if (arch == LLMArch::GEMMA2_2B) { pad_id = 0; + } else if (arch == LLMArch::LLADA2_MOE) { + pad_id = 156892; } tokenizer = tokenizers.create(TokenizerConfig::MAIN, model.config.vocab_size, pad_id); if (!tokenizer) { - if (arch == LLMArch::GPT_OSS_20B || arch == LLMArch::GEMMA2_2B) { - throw std::runtime_error("GPT-OSS and Gemma 2 require an external tokenizer.json in the main tokenizer slot"); + if (arch == LLMArch::GPT_OSS_20B || arch == LLMArch::GEMMA2_2B || arch == LLMArch::LLADA2_MOE) { + throw std::runtime_error("GPT-OSS, Gemma 2 and LLaDA2 require an external tokenizer.json in the main tokenizer slot"); } if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) { tokenizer = std::make_shared(); diff --git a/otherarch/sdcpp/src/model/vae/ltx_vae.hpp b/otherarch/sdcpp/src/model/vae/ltx_vae.hpp index ac22c09b4..f680f2008 100644 --- a/otherarch/sdcpp/src/model/vae/ltx_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/ltx_vae.hpp @@ -1300,7 +1300,7 @@ struct LTXVideoVAE : public VAE { feat_map[feat_idx] = get_cache_tensor_by_name(temporal_feat_cache_name(feat_idx)); } - auto runner_ctx = get_context(); + auto runner_ctx = get_context(gf); int feat_count = 0; ggml_tensor* out = vae.decode_tiled_chunk(&runner_ctx, z, @@ -1313,8 +1313,7 @@ struct LTXVideoVAE : public VAE { for (int feat_idx = 0; feat_idx < feat_count && feat_idx < static_cast(feat_map.size()); ++feat_idx) { ggml_tensor* feat_cache = feat_map[static_cast(feat_idx)]; if (feat_cache != nullptr) { - cache(temporal_feat_cache_name(static_cast(feat_idx)), feat_cache); - ggml_build_forward_expand(gf, feat_cache); + runner_ctx.persist_cache_tensor(temporal_feat_cache_name(static_cast(feat_idx)), feat_cache); } } diff --git a/otherarch/sdcpp/src/model/vae/wan_vae.hpp b/otherarch/sdcpp/src/model/vae/wan_vae.hpp index 79efca4a0..58cff816b 100644 --- a/otherarch/sdcpp/src/model/vae/wan_vae.hpp +++ b/otherarch/sdcpp/src/model/vae/wan_vae.hpp @@ -24,6 +24,7 @@ namespace WAN { std::tuple padding; std::tuple dilation; bool bias; + float scale = 1.f; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { auto weight = tensor_storage_map.find(prefix + "weight"); @@ -60,6 +61,10 @@ namespace WAN { dilation(std::move(dilation)), bias(bias) {} + void set_scale(float scale_value) { + scale = scale_value; + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cache_x = nullptr) { // x: [N*IC, ID, IH, IW] // result: x: [N*OC, ID, IH, IW] @@ -82,11 +87,25 @@ namespace WAN { } x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); + if (w->ne[2] == 1 && x->ne[2] == 1 && x->ne[3] == in_channels) { + // One frame through a one-frame-deep kernel is a 2D conv; backends without + // im2col_3d (Metal) otherwise fall back to a much slower direct conv_3d. + if (!ggml_is_contiguous(x)) { + x = ggml_cont(ctx->ggml_ctx, x); + } + ggml_tensor* x2 = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], x->ne[1], in_channels, 1); + ggml_tensor* w2 = ggml_reshape_4d(ctx->ggml_ctx, w, w->ne[0], w->ne[1], in_channels, out_channels); + x2 = ggml_ext_conv_2d(ctx->ggml_ctx, x2, w2, b, + std::get<2>(stride), std::get<1>(stride), 0, 0, + std::get<2>(dilation), std::get<1>(dilation), + ctx->conv2d_direct_enabled, false, false, scale); + return ggml_reshape_4d(ctx->ggml_ctx, x2, x2->ne[0], x2->ne[1], 1, out_channels); + } return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels, std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), 0, 0, 0, std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation), - false, ctx->conv3d_direct_enabled); + false, ctx->conv3d_direct_enabled, scale); } }; @@ -1070,11 +1089,11 @@ namespace WAN { } if (version == VERSION_QWEN_IMAGE_2_1) { - wan2_2 = true; - dec_dim = 144; - z_dim = 64; - input_channels = 4; - dim_mult = {1, 2, 4, 8, 8}; + wan2_2 = true; + dec_dim = 144; + z_dim = 64; + input_channels = 4; + dim_mult = {1, 2, 4, 8, 8}; } if (is_2D) { @@ -1103,6 +1122,19 @@ namespace WAN { } else { blocks["conv2"] = std::shared_ptr(new CausalConv3d(z_dim, z_dim, {1, 1, 1})); } + if (version == VERSION_QWEN_IMAGE_2_1) { + // Keep large VAE activations within the FP16 convolution range. + const float conv_scale = 1.f / 128.f; + std::vector all_blocks; + get_all_blocks(all_blocks); + for (auto block : all_blocks) { + if (auto conv = dynamic_cast(block)) { + conv->set_scale(conv_scale); + } else if (auto conv = dynamic_cast(block)) { + conv->set_scale(conv_scale); + } + } + } } static ggml_tensor* patchify(ggml_context* ctx, @@ -1429,15 +1461,14 @@ namespace WAN { ggml_tensor* z = make_input(z_tensor); - auto runner_ctx = get_context(); + auto runner_ctx = get_context(gf); ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx); for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) { ggml_tensor* feat_cache = ae._feat_map[feat_idx]; if (feat_cache != nullptr) { - cache("feat_idx:" + std::to_string(feat_idx), feat_cache); - ggml_build_forward_expand(gf, feat_cache); + runner_ctx.persist_cache_tensor("feat_idx:" + std::to_string(feat_idx), feat_cache); } } diff --git a/otherarch/sdcpp/src/model_loader.cpp b/otherarch/sdcpp/src/model_loader.cpp index 9cad3fce7..792363cb6 100644 --- a/otherarch/sdcpp/src/model_loader.cpp +++ b/otherarch/sdcpp/src/model_loader.cpp @@ -339,7 +339,13 @@ bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, } for (const std::string& shard_path : shard_paths) { - if (!parse_file(shard_path, prefix)) { + FileStamp stamp; + if (!read_file_stamp(shard_path, stamp)) { + return false; + } + parsed_dependencies_.push_back(stamp); + LOG_INFO("load %s using safetensors format", shard_path.c_str()); + if (!init_from_safetensors_file(shard_path, prefix)) { return false; } } @@ -541,6 +547,9 @@ SDVersion ModelLoader::get_sd_version() const { if (tensor_storage.name.find("model.diffusion_model.double_blocks.0.img_mlp.gate_proj.weight") != std::string::npos) { return VERSION_OVIS_IMAGE; } + if (tensor_storage.name.find("model.diffusion_model.sigvq_embedder.1.weight") != std::string::npos) { + return VERSION_LLADA_IMAGE; + } if (tensor_storage.name.find("model.diffusion_model.cap_embedder.0.weight") != std::string::npos) { return VERSION_Z_IMAGE; } @@ -898,7 +907,8 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) { std::vector ModelLoader::mmap_tensors(std::map& tensors, std::set ignore_tensors, - bool writable_mmap) { + bool writable_mmap, + ggml_backend_dev_t device) { std::set names; for (const auto& entry : tensors) { names.insert(entry.first); @@ -920,6 +930,39 @@ std::vector ModelLoader::mmap_tensors(std::map file_buffer = device == nullptr ? fdata.mmbuffer : nullptr; + bool file_unmappable = false; + + auto buffer_for_file = [&]() -> ggml_backend_buffer_t { + if (file_buffer || file_unmappable) { + return file_buffer.get(); + } + auto cached = fdata.device_mmbuffers.find(device); + if (cached != fdata.device_mmbuffers.end()) { + file_buffer = cached->second; + return file_buffer.get(); + } + size_t max_tensor_size = 0; + for (const auto& ts : fdata.tensors) { + max_tensor_size = std::max(max_tensor_size, static_cast(ts.nbytes())); + } + ggml_backend_buffer_t buf = sd_backend_dev_buffer_from_host_ptr(device, + fdata.mmapped->writable_data(), + fdata.mmapped->size(), + max_tensor_size); + if (buf == nullptr) { + LOG_WARN("mmap: %s cannot map '%s', loading it instead", + ggml_backend_dev_name(device), fdata.path.c_str()); + file_unmappable = true; + return nullptr; + } + LOG_INFO("mmap: mapped '%s' for %s", fdata.path.c_str(), ggml_backend_dev_name(device)); + file_buffer = std::shared_ptr(buf, ggml_backend_buffer_free); + fdata.device_mmbuffers[device] = file_buffer; + return file_buffer.get(); + }; + const std::vector& file_tensors = fdata.tensors; size_t file_mapped_bytes = 0; @@ -969,10 +1012,13 @@ std::vector ModelLoader::mmap_tensors(std::map(ggml_backend_buffer_get_base(buf_mmap)); - dst_tensor->buffer = buf_mmap; - dst_tensor->data = mmap_data + tensor_offset; + ggml_backend_buffer_t buf_mmap = buffer_for_file(); + if (buf_mmap == nullptr) { + break; + } + uint8_t* mmap_data = static_cast(ggml_backend_buffer_get_base(buf_mmap)); + dst_tensor->buffer = buf_mmap; + dst_tensor->data = mmap_data + tensor_offset; file_mapped_bytes += tensor_size; file_mapped_tensors++; @@ -981,7 +1027,7 @@ std::vector ModelLoader::mmap_tensors(std::map 0) { mapped_tensors += file_mapped_tensors; mapped_bytes += file_mapped_bytes; - result.push_back({fdata.mmapped, fdata.mmbuffer}); + result.push_back({fdata.mmapped, file_buffer}); } } @@ -997,6 +1043,16 @@ std::vector ModelLoader::mmap_tensors(std::map ModelLoader::get_device_mmap_buffers() const { + std::vector buffers; + for (const auto& fdata : file_data) { + for (const auto& entry : fdata.device_mmbuffers) { + buffers.push_back(entry.second.get()); + } + } + return buffers; +} + bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, bool enable_mmap, const std::set* target_tensor_names, @@ -1142,6 +1198,11 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, if (dst_tensor->buffer != nullptr && dst_tensor->buffer == fdata.mmbuffer.get()) { continue; } + if (dst_tensor->buffer != nullptr && + std::any_of(fdata.device_mmbuffers.begin(), fdata.device_mmbuffers.end(), + [&](const auto& entry) { return entry.second.get() == dst_tensor->buffer; })) { + continue; + } size_t nbytes_to_read = tensor_storage.nbytes_to_read(); @@ -1599,6 +1660,9 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage // Pass, do not convert. For Unet } else if (contains(name, "embedding")) { // Pass, do not convert embedding + } else if (ends_with(name, "_pad_token")) { + // Pass, do not convert. LLaDA-Image stores its pad tokens far outside the f16 + // range, so any format with an f16 scale or payload turns them into inf. } else { return true; } diff --git a/otherarch/sdcpp/src/model_loader.h b/otherarch/sdcpp/src/model_loader.h index c410c1c97..0b94ad342 100644 --- a/otherarch/sdcpp/src/model_loader.h +++ b/otherarch/sdcpp/src/model_loader.h @@ -20,6 +20,8 @@ struct ModelFileData { std::vector tensors; std::shared_ptr mmapped; std::shared_ptr mmbuffer; + // mmapped wrapped by devices that can use host memory in place (buffer_from_host_ptr) + std::map> device_mmbuffers; bool is_zip; }; @@ -120,7 +122,9 @@ public: void process_model_files(bool enable_mmap = false, bool writable_mmap = true); std::vector mmap_tensors(std::map& tensors, std::set ignore_tensors = {}, - bool writable = true); + bool writable = true, + ggml_backend_dev_t device = nullptr); + std::vector get_device_mmap_buffers() const; bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, bool use_mmap = false, const std::set* target_tensor_names = nullptr, diff --git a/otherarch/sdcpp/src/model_loader_files.cpp b/otherarch/sdcpp/src/model_loader_files.cpp index f7f1cb37e..8e838b119 100644 --- a/otherarch/sdcpp/src/model_loader_files.cpp +++ b/otherarch/sdcpp/src/model_loader_files.cpp @@ -147,7 +147,7 @@ bool ModelLoader::add_file_impl(const std::string& path, const std::string& pref } if (tensor.index_in_zip < 0) { const auto& stamp = physical_files[tensor.file_index]; - if (tensor.offset > stamp.size || static_cast(tensor.nbytes_to_read()) > stamp.size - tensor.offset) { //kcpp int8 fp8 + if (tensor.offset > stamp.size || static_cast(tensor.nbytes_to_read()) > stamp.size - tensor.offset) { // kcpp int8 fp8 LOG_ERROR("tensor '%s' extends beyond its model file", tensor.name.c_str()); return false; } diff --git a/otherarch/sdcpp/src/model_manager.cpp b/otherarch/sdcpp/src/model_manager.cpp index 9a9fe0628..5bdf59cac 100644 --- a/otherarch/sdcpp/src/model_manager.cpp +++ b/otherarch/sdcpp/src/model_manager.cpp @@ -780,38 +780,52 @@ bool ModelManager::validate_tensor(const TensorState& state) const { bool ModelManager::mmap_params(const std::vector& states, std::vector& created_storage_blocks) { - std::map mmap_candidates; - std::map mmap_states; + // A GPU that computes on mmapped params in place cannot address a CPU buffer, and nothing + // stages them for it, so they are mapped through a buffer of that GPU's device. + struct MmapGroup { + std::map candidates; + std::map states; + }; + std::map groups; for (TensorState* state : states) { if (state == nullptr || !can_mmap_storage(*state) || state->tensor == nullptr || state->tensor->data != nullptr || state->tensor->view_src != nullptr) { continue; } - mmap_candidates[state->name] = state->tensor; - mmap_states[state->name] = state; - } - if (mmap_candidates.empty()) { - return true; - } - - auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, writable_mmap_); - if (mmap_store.empty()) { - return true; - } - - auto block = std::make_unique(); - block->mmap_tensor_stores = std::move(mmap_store); - ParamsStorageBlock* raw = block.get(); - for (const auto& pair : mmap_states) { - TensorState* state = pair.second; - if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) { - block->states.push_back(state); + ggml_backend_dev_t device = nullptr; + if (!sd_backend_is_cpu(state->compute_backend) && !sd_backend_is_cpu(state->params_backend)) { + device = ggml_backend_get_device(state->compute_backend); } + MmapGroup& group = groups[device]; + group.candidates[state->name] = state->tensor; + group.states[state->name] = state; } - if (!block->states.empty()) { - params_storage_blocks_.push_back(std::move(block)); - created_storage_blocks.push_back(raw); + for (auto& [device, group] : groups) { + // Device buffers wrap read-only mappings only; params that LoRAs are merged into in place + // are loaded instead. + if (device != nullptr && writable_mmap_) { + continue; + } + auto mmap_store = model_loader_.mmap_tensors(group.candidates, {}, writable_mmap_, device); + if (mmap_store.empty()) { + continue; + } + + auto block = std::make_unique(); + block->mmap_tensor_stores = std::move(mmap_store); + ParamsStorageBlock* raw = block.get(); + for (const auto& pair : group.states) { + TensorState* state = pair.second; + if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) { + block->states.push_back(state); + } + } + + if (!block->states.empty()) { + params_storage_blocks_.push_back(std::move(block)); + created_storage_blocks.push_back(raw); + } } return true; } @@ -1353,15 +1367,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe } size_t total_size = 0; - auto add_buffer = [&](ggml_backend_buffer_t buffer) { - if (buffer == nullptr || ggml_backend_buffer_is_host(buffer)) { + std::unordered_set seen; + auto add_buffer = [&](ggml_backend_buffer_t buffer) { + if (buffer == nullptr || ggml_backend_buffer_is_host(buffer) || !seen.insert(buffer).second) { return; } ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer); auto split_devices = split_buffer_devices_.find(buffer_type); const bool on_device = split_devices == split_buffer_devices_.end() - ? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device - : std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) { + ? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device + : std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) { return ggml_backend_get_device(entry.first) == compute_device; }); if (!on_device) { @@ -1371,9 +1386,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe total_size = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size; }; + // The loader may retain device mappings after their parameter blocks are released. + for (ggml_backend_buffer_t buffer : model_loader_.get_device_mmap_buffers()) { + add_buffer(buffer); + } for (const auto& block : params_storage_blocks_) { if (block != nullptr) { add_buffer(block->buffer); + for (const auto& store : block->mmap_tensor_stores) { + add_buffer(store.mmbuffer.get()); + } } } for (const auto& block : compute_staging_blocks_) { @@ -1613,7 +1635,8 @@ void ModelManager::remove_runtime_owner(uintptr_t owner_id) { ModelManager::CapacityCheck ModelManager::check_capacity( const DeviceMemoryRequest& request, - const std::vector& states) const { + const std::vector& states, + bool log_details) const { CapacityCheck result; if (request.compute_backend == nullptr || sd_backend_is_cpu(request.compute_backend)) { return result; @@ -1631,16 +1654,23 @@ ModelManager::CapacityCheck ModelManager::check_capacity( } size_t free_bytes = 0, total_bytes = 0; ggml_backend_dev_memory(device, &free_bytes, &total_bytes); + const size_t weights_resident = compute_backend_resident_bytes(backend); + const size_t other_runtime = other_runtime_resident_bytes(request.owner_id, backend); + const size_t resident = add(weights_resident, add(other_runtime, request.runtime_resident_bytes)); + if (log_details) { + LOG_WARN("model manager memory on %s: reported free %.2f MB / total %.2f MB, tracked weights %.2f MB / other runtime %.2f MB / current runtime %.2f MB", + ggml_backend_name(backend), + free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0), + weights_resident / (1024.0 * 1024.0), other_runtime / (1024.0 * 1024.0), + request.runtime_resident_bytes / (1024.0 * 1024.0)); + } if (free_bytes == 0 && total_bytes == 0) { return SIZE_MAX; } // Vulkan's heap budget subtraction can underflow when usage exceeds the budget. - if (total_bytes > 0 && free_bytes > total_bytes) { + if (total_bytes > 0 && free_bytes > total_bytes && sd_backend_is(backend, "Vulkan")) { return size_t{0}; } - const size_t resident = add(compute_backend_resident_bytes(backend), - add(other_runtime_resident_bytes(request.owner_id, backend), - request.runtime_resident_bytes)); if (total_bytes > 0) { free_bytes = std::min(free_bytes, resident < total_bytes ? total_bytes - resident : 0); } @@ -1786,7 +1816,7 @@ bool ModelManager::ensure_compute_backend_capacity( } } - const auto capacity = check_capacity(request, required_states); + const auto capacity = check_capacity(request, required_states, true); const std::string available_device = capacity.available_device_bytes == SIZE_MAX ? "unknown" : sd_format("%.2f MB", capacity.available_device_bytes / (1024.0 * 1024.0)); diff --git a/otherarch/sdcpp/src/model_manager.h b/otherarch/sdcpp/src/model_manager.h index 7f8df59c9..f3a6e3286 100644 --- a/otherarch/sdcpp/src/model_manager.h +++ b/otherarch/sdcpp/src/model_manager.h @@ -157,7 +157,8 @@ private: } }; CapacityCheck check_capacity(const DeviceMemoryRequest& request, - const std::vector& states) const; + const std::vector& states, + bool log_details = false) const; 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; diff --git a/otherarch/sdcpp/src/name_conversion.cpp b/otherarch/sdcpp/src/name_conversion.cpp index a3249ef54..325168235 100644 --- a/otherarch/sdcpp/src/name_conversion.cpp +++ b/otherarch/sdcpp/src/name_conversion.cpp @@ -103,6 +103,8 @@ std::string convert_open_clip_to_hf_clip_name(std::string name) { return name; } +std::string convert_llada2_moe_te_name(std::string name); + std::string convert_cond_stage_model_name(std::string name, std::string prefix) { static const std::vector> clip_name_map{ {"transformer.text_projection.weight", "transformer.text_model.text_projection"}, @@ -177,6 +179,7 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix) replace_with_name_map(name, llm_vision_name_map); } else { replace_with_name_map(name, llm_name_map); + name = convert_llada2_moe_te_name(name); } } else { name = convert_open_clip_to_hf_clip_name(name); @@ -749,6 +752,52 @@ std::string convert_hunyuan_video_to_original_flux(std::string name) { return name; } +// LLaDA-Image's LLaDA2-MoE text encoder. Both published layouts use these names; the ComfyUI +// GGUF repack differs only by appending ".weight" to the bare 3-D expert parameters. +// Called with the "text_encoders." prefix already stripped, so the name still carries "llm.". +std::string convert_llada2_moe_te_name(std::string name) { + static const std::vector> name_map = { + {"model.language_model.word_embeddings.", "model.embed_tokens."}, + {"model.language_model.norm.", "model.norm."}, + {"model.language_model.lm_head.", "lm_head."}, + {"model.language_model.layers.", "model.layers."}, + {"attention.query_key_value.", "self_attn.query_key_value."}, + {"attention.dense.", "self_attn.o_proj."}, + {"attention.query_layernorm.", "self_attn.q_norm."}, + {"attention.key_layernorm.", "self_attn.k_norm."}, + }; + replace_with_name_map(name, name_map); + + // The HF checkpoint stores the stacked experts as bare nn.Parameters with no ".weight". + static const std::vector bare_expert_params = { + "mlp.experts.gate_proj", + "mlp.experts.up_proj", + "mlp.experts.down_proj", + }; + for (const auto& suffix : bare_expert_params) { + if (ends_with(name, suffix)) { + name += ".weight"; + break; + } + } + + return name; +} + +// The attention projections keep their diffusers names (JointAttention's split_qkv mode), so +// only the patch-size-keyed dicts need flattening. Latents arrive already patchified from the +// Flux2 VAE, so the only patch key is 1-1. +std::string convert_diffusers_dit_to_original_llada_image(std::string name) { + static const std::vector> prefix_map = { + {"all_x_embedder.1-1.", "x_embedder."}, + {"all_final_layer.1-1.", "final_layer."}, + }; + + replace_with_prefix_map(name, prefix_map); + + return name; +} + std::string convert_diffusers_dit_to_original_lumina2(std::string name) { int num_layers = 30; int num_refiner_layers = 2; @@ -896,6 +945,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S name = convert_hunyuan_video_to_original_flux(name); } else if (sd_version_is_z_image(version)) { name = convert_diffusers_dit_to_original_lumina2(name); + } else if (sd_version_is_llada_image(version)) { + name = convert_diffusers_dit_to_original_llada_image(name); } else if (sd_version_is_anima(version)) { name = convert_other_dit_to_original_anima(name); } else if (sd_version_is_krea2(version)) { @@ -1298,9 +1349,8 @@ bool is_first_stage_model_name(const std::string& name) { } static std::string convert_esrgan_tensor_name(std::string name) { - static std::unordered_map esrgan_name_map; - - if (esrgan_name_map.empty()) { + static const auto esrgan_name_map = [] { + std::unordered_map esrgan_name_map; esrgan_name_map["model.0."] = "conv_first."; constexpr int max_num_blocks = 64; @@ -1326,7 +1376,8 @@ static std::string convert_esrgan_tensor_name(std::string name) { esrgan_name_map["model.7."] = "conv_last."; esrgan_name_map["model.8."] = "conv_hr."; esrgan_name_map["model.10."] = "conv_last."; - } + return esrgan_name_map; + }(); replace_with_prefix_map(name, esrgan_name_map); return name; diff --git a/otherarch/sdcpp/src/pipeline/diffusion_engine.cpp b/otherarch/sdcpp/src/pipeline/diffusion_engine.cpp index ffa155151..5fd6d8a58 100644 --- a/otherarch/sdcpp/src/pipeline/diffusion_engine.cpp +++ b/otherarch/sdcpp/src/pipeline/diffusion_engine.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -29,6 +30,7 @@ #include "stable-diffusion.h" #include "conditioning/conditioner.hpp" +#include "conditioning/conditioning_cache.h" #include "core/backend_fit.h" #include "extensions/generation_extension.h" #include "model/adapter/ip_adapter.hpp" @@ -101,6 +103,7 @@ const char* model_version_to_str[] = { "Krea2", "Mage Flow", "SenseNova U1.5", + "LLaDA-Image", "ESRGAN", }; @@ -134,6 +137,7 @@ static_assert(std::atomic::is_always_lock_free, StableDiffusionGGML::StableDiffusionGGML() : rng(std::make_shared()), + conditioning_cache_(std::make_unique()), denoiser(std::make_shared()) {} StableDiffusionGGML::~StableDiffusionGGML() = default; @@ -203,6 +207,8 @@ void StableDiffusionGGML::end_runners() { } bool StableDiffusionGGML::reset_runners(const RunnerGroups& groups) { + conditioning_cache_->clear(); + conditioning_loras_.clear(); end_runners(); clear_lora_adapters(); runtime_lora_models.clear(); @@ -883,7 +889,7 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf { to_replace = "taesd_f2.embd"; } - else if(sd_version_uses_wan_vae(tempver)) + else if(sd_version_uses_wan_vae(tempver) && tempver != VERSION_QWEN_IMAGE_2_1) // qwen 2.1 + tae crashing as of master-917 { to_replace = "taesd_w21.embd"; } @@ -1105,6 +1111,47 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf return true; } +bool StableDiffusionGGML::set_sage_attention_enabled(bool enabled) { + if (!diffusion_model) { + return false; + } + if (enabled) { +#ifndef SD_USE_UPSTREAM_GGML + auto* ctx = ggml_init({4 * ggml_tensor_overhead(), nullptr, true}); + if (ctx == nullptr) { + return false; + } + auto* q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1); + auto* k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1); + auto* v = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 128, 128, 1, 1); + auto* op = ggml_sage_attn(ctx, q, k, v, 1.f / sqrtf(128.f), GGML_SAGE_ATTN_AUTO); + bool supported = true; + for (auto backend : backend_manager.runtime_backends(SDBackendModule::DIFFUSION)) { + if (!ggml_backend_supports_op(backend, op)) { + LOG_ERROR("SageAttention is unavailable on %s; it requires patched GGML, CUDA Toolkit 12.0 or newer, and SM80 or newer kernels", + ggml_backend_name(backend)); + supported = false; + } + } + ggml_free(ctx); + if (!supported) { + return false; + } +#else + LOG_ERROR("SageAttention requires -DSD_USE_UPSTREAM_GGML=OFF and a CUDA backend"); + return false; +#endif + } + diffusion_model->set_sage_attention_enabled(enabled); + if (high_noise_diffusion_model) { + high_noise_diffusion_model->set_sage_attention_enabled(enabled); + } + if (enabled) { + LOG_INFO("Using SageAttention in the diffusion model; CUDA selects the supported kernel, unsupported layers use flash/default attention"); + } + return true; +} + bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) { #ifdef SD_USE_UPSTREAM_GGML LOG_WARN( @@ -1121,6 +1168,11 @@ bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) { return false; } } + if (sd_ctx_params->conditioning_cache_size < 0) { + LOG_ERROR("conditioning_cache_size must be non-negative"); + return false; + } + conditioning_cache_->set_capacity(static_cast(sd_ctx_params->conditioning_cache_size)); auto configuration = std::make_unique(*sd_ctx_params); n_threads = sd_ctx_params->n_threads; tensor_executor = std::make_unique(n_threads > 0 ? n_threads : sd_get_num_physical_cores()); @@ -1386,6 +1438,9 @@ bool StableDiffusionGGML::validate_and_load_runners() { high_noise_diffusion_model->set_flash_attention_enabled(true); } } + if (sd_ctx_params->sage_attn && !set_sage_attention_enabled(true)) { + return false; + } LOG_VERBOSE("validating model metadata"); std::set ignore_tensors; @@ -1549,6 +1604,7 @@ bool StableDiffusionGGML::build_denoiser() { sd_version_is_anima(version) || sd_version_is_ernie_image(version) || sd_version_is_z_image(version) || + sd_version_is_llada_image(version) || sd_version_is_boogu_image(version) || sd_version_is_pid(version) || sd_version_is_ideogram4(version)) { @@ -1569,6 +1625,8 @@ bool StableDiffusionGGML::build_denoiser() { default_flow_shift = 3.16f; } else if (sd_version_is_mage_flow(version)) { default_flow_shift = 6.f; + } else if (sd_version_is_llada_image(version)) { + default_flow_shift = 1.0f; // unused: LLADA_IMAGE_SCHEDULER builds a fixed grid } else { default_flow_shift = 3.f; } @@ -1986,8 +2044,19 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun int64_t t0 = ggml_time_ms(); end_runners(); clear_lora_adapters(); - if (!model_manager->prepare_lora_sources(all_loras)) + if (!model_manager->prepare_lora_sources(all_loras)) { + conditioning_cache_->clear(); return false; + } + if (!std::equal(all_loras.begin(), all_loras.end(), + conditioning_loras_.begin(), conditioning_loras_.end(), + [](const ModelManager::LoraSpec& a, const ModelManager::LoraSpec& b) { + return a.file_id == b.file_id && a.file_revision == b.file_revision && + a.multiplier == b.multiplier && a.is_high_noise == b.is_high_noise && + a.tensor_name_prefix_filter == b.tensor_name_prefix_filter; + })) { + conditioning_cache_->clear(); + } runtime_lora_models.erase(std::remove_if(runtime_lora_models.begin(), runtime_lora_models.end(), [&](const RuntimeLora& entry) { return std::none_of(all_loras.begin(), all_loras.end(), [&](const ModelManager::LoraSpec& spec) { return entry.matches(spec); @@ -1997,6 +2066,7 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun const bool success = apply_lora_immediately ? apply_loras_immediately(all_loras) : apply_loras_at_runtime(all_loras); if (!success) { + conditioning_cache_->clear(); clear_lora_adapters(); runtime_lora_models.clear(); return false; @@ -2006,9 +2076,14 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun if (!all_loras.empty()) { LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); } + conditioning_loras_ = std::move(all_loras); return true; } +SDCondition StableDiffusionGGML::get_learned_condition(const ConditionerParams& params) { + return conditioning_cache_->get(*cond_stage_model, n_threads, params); +} + void StableDiffusionGGML::reset_generation_extensions() { for (auto& extension : generation_extensions) { extension->reset_runtime_condition(); @@ -2193,6 +2268,8 @@ void StableDiffusionGGML::preview_image(int step, int patch_sz = 1; const float(*latent_rgb_proj)[3] = nullptr; float* latent_rgb_bias = nullptr; + const float* latent_alpha_proj = nullptr; + float latent_alpha_bias = 1.f; if (channels == 128) { if (sd_version_uses_flux2_vae(version)) { @@ -2206,6 +2283,16 @@ void StableDiffusionGGML::preview_image(int step, LOG_WARN("No latent to RGB projection known for this model"); return; } + } else if (channels == 64) { + if (version == VERSION_QWEN_IMAGE_2_1) { + latent_rgb_proj = qwen21_latent_rgb_proj; + latent_rgb_bias = qwen21_latent_rgb_bias; + latent_alpha_proj = qwen21_latent_alpha_proj; + latent_alpha_bias = qwen21_latent_alpha_bias; + } else { + LOG_WARN("No latent to RGB projection known for this model"); + return; + } } else if (channels == 48) { if (sd_version_is_wan(version)) { latent_rgb_proj = wan_22_latent_rgb_proj; @@ -2256,13 +2343,14 @@ void StableDiffusionGGML::preview_image(int step, uint32_t img_width = static_cast(_latents.shape()[0]) * patch_sz; uint32_t img_height = static_cast(_latents.shape()[1]) * patch_sz; - uint8_t* data = (uint8_t*)malloc(frames * img_width * img_height * 3 * sizeof(uint8_t)); + uint32_t img_channels = latent_alpha_proj != nullptr ? 4 : 3; + uint8_t* data = (uint8_t*)malloc(frames * img_width * img_height * img_channels * sizeof(uint8_t)); GGML_ASSERT(data != nullptr); - preview_latent_video(data, _latents, latent_rgb_proj, latent_rgb_bias, patch_sz); + preview_latent_video(data, _latents, latent_rgb_proj, latent_rgb_bias, patch_sz, latent_alpha_proj, latent_alpha_bias); sd_image_t* images = (sd_image_t*)malloc(frames * sizeof(sd_image_t)); GGML_ASSERT(images != nullptr); for (uint32_t i = 0; i < frames; i++) { - images[i] = {img_width, img_height, 3, data + i * img_width * img_height * 3}; + images[i] = {img_width, img_height, img_channels, data + i * img_width * img_height * img_channels}; } step_callback(step, frames, images, is_noisy, step_callback_data); free(data); @@ -2437,6 +2525,15 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptris_enabled(); }); + using QwenPrefixInputs = std::tuple*, const sd::Tensor*, + const std::vector>*>; + std::vector qwen_prefix_inputs; + RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; const bool apply_denoise_mask = !denoise_mask.empty() && @@ -2613,6 +2710,7 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr> empty_ref_latents; bool uncond_without_ref_latents = !img_uncond.empty() && @@ -2668,6 +2766,9 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr StableDiffusionGGML::sample(const std::shared_ptrbefore_diffusion(diffusion_params, step); } + if (cache_qwen_prefix) { + auto* extra = std::get_if(&diffusion_params.extra); + if (extra != nullptr) { + auto key = std::make_tuple(diffusion_params.context, extra->image_slots, + diffusion_params.ref_image_params.pass_to_dit ? diffusion_params.ref_latents : nullptr); + auto entry = std::find(qwen_prefix_inputs.begin(), qwen_prefix_inputs.end(), key); + extra->prefix_id = static_cast(entry - qwen_prefix_inputs.begin()) + 1; + if (entry == qwen_prefix_inputs.end()) { + qwen_prefix_inputs.push_back(key); + } + } + } auto output_opt = work_diffusion_model->compute(n_threads, diffusion_params); if (output_opt.empty()) { LOG_ERROR("diffusion model compute failed"); @@ -2726,41 +2850,69 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr* uncond_skip_layers = nullptr; - if (is_skiplayer_step && slg_uncond) { - LOG_VERBOSE("Skipping layers at uncond step %d\n", step); - uncond_skip_layers = &skip_layer_guidance.layers(); - } - uncond_out = run_condition(uncond, - uncond.c_concat.empty() ? nullptr : &uncond.c_concat, - uncond_skip_layers, - nullptr, - true); - if (uncond_out.empty()) { - return {}; + if (!skip_uncond) { + const std::vector* uncond_skip_layers = nullptr; + if (is_skiplayer_step && slg_uncond) { + LOG_VERBOSE("Skipping layers at uncond step %d\n", step); + uncond_skip_layers = &skip_layer_guidance.layers(); + } + uncond_out = run_condition(uncond, + uncond.c_concat.empty() ? nullptr : &uncond.c_concat, + uncond_skip_layers, + nullptr, + true); + if (uncond_out.empty()) { + return {}; + } + } else { + step_cache.invalidate_condition(&uncond); } } + if (!img_uncond.empty()) { - img_uncond_out = run_condition(img_uncond, - img_uncond.c_concat.empty() ? nullptr : &img_uncond.c_concat, - nullptr, - uncond_without_ref_latents ? &empty_ref_latents : nullptr, - true); - if (img_uncond_out.empty()) { - return {}; + if (!skip_img_uncond) { + img_uncond_out = run_condition(img_uncond, + img_uncond.c_concat.empty() ? nullptr : &img_uncond.c_concat, + nullptr, + uncond_without_ref_latents ? &empty_ref_latents : nullptr, + true); + if (img_uncond_out.empty()) { + return {}; + } + } else { + step_cache.invalidate_condition(&img_uncond); } } sd::guidance::GuidanceInput guidance_input; @@ -2770,7 +2922,7 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr StableDiffusionGGML::decode_first_stage(const sd::Tensordecode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); const bool prefer_temporal_tiling = decode_video && first_stage_model->can_temporal_tile_decode(); while (decoded.empty() && - sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { + sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling, + first_stage_model->last_compute_status())) { decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); } return decoded; @@ -3047,6 +3200,8 @@ std::string StableDiffusionGGML::get_default_ref_image_preset(SDVersion version) return "mage_flow"; } else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) { return "z_image_omni"; + } else if (sd_version_is_llada_image(version)) { + return "llada_image"; } else if (sd_version_is_krea2(version)) { // have to make a choice between "krea2_edit" mode (for lbouaraba/krea2edit) // and "krea2_ostris_edit" (for krea2 ostris edit) diff --git a/otherarch/sdcpp/src/pipeline/diffusion_engine.h b/otherarch/sdcpp/src/pipeline/diffusion_engine.h index dc89a5185..4abf01cf7 100644 --- a/otherarch/sdcpp/src/pipeline/diffusion_engine.h +++ b/otherarch/sdcpp/src/pipeline/diffusion_engine.h @@ -26,6 +26,7 @@ class RNG; struct Denoiser; struct LoraModel; struct ConditionerParams; +class ConditioningCache; struct SDCondition; struct RefImageParams; namespace Wav2Vec2 { @@ -184,6 +185,8 @@ public: std::recursive_mutex execution_mutex; std::unique_ptr config_; RunnerState runner_state_; + std::unique_ptr conditioning_cache_; + std::vector conditioning_loras_; bool executing_ = false; std::shared_ptr denoiser; @@ -318,6 +321,7 @@ public: bool init_model_loader(ModelLoader& model_loader, ModelConfig& configuration); bool init(const sd_ctx_params_t* sd_ctx_params); + bool set_sage_attention_enabled(bool enabled); bool uses_tae() const; @@ -366,6 +370,8 @@ public: bool apply_loras(const sd_lora_t* loras, uint32_t lora_count); + SDCondition get_learned_condition(const ConditionerParams& params); + void reset_generation_extensions(); void prepare_generation_extensions(const sd_pm_params_t& pm_params, diff --git a/otherarch/sdcpp/src/pipeline/image.cpp b/otherarch/sdcpp/src/pipeline/image.cpp index 0eff1d962..38b8f086a 100644 --- a/otherarch/sdcpp/src/pipeline/image.cpp +++ b/otherarch/sdcpp/src/pipeline/image.cpp @@ -10,6 +10,7 @@ #include "model/vae/vae.hpp" #include "request.h" #include "runtime/denoiser.hpp" +#include "runtime/image_preprocess.h" #include "upscaler.h" namespace sd::pipeline { @@ -440,8 +441,7 @@ namespace sd::pipeline { sd->compute_ip_adapter_tokens(sd_img_gen_params->ip_adapter_image, sd_img_gen_params->ip_adapter_strength); int64_t prepare_start_ms = ggml_time_ms(); condition_params.zero_out_masked = false; - auto cond = sd->cond_stage_model->get_learned_condition(sd->n_threads, - condition_params); + auto cond = sd->get_learned_condition(condition_params); if (cond.empty()) { LOG_ERROR("failed to encode prompt"); return std::nullopt; @@ -475,8 +475,11 @@ namespace sd::pipeline { } condition_params.text = request->negative_prompt; condition_params.zero_out_masked = zero_out_masked; - uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, - condition_params); + if (sd_version_is_llada_image(sd->version)) { + // LLaDA-Image CFG keeps the source latent but drops its SigVQ features. + condition_params.ref_images = nullptr; + } + uncond = sd->get_learned_condition(condition_params); if (uncond.empty()) { LOG_ERROR("failed to encode negative prompt"); return std::nullopt; @@ -504,8 +507,7 @@ namespace sd::pipeline { if (use_ref_latent_img_cfg) { condition_params.ref_images = &empty_ref_images; } - img_uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, - condition_params); + img_uncond = sd->get_learned_condition(condition_params); if (img_uncond.empty()) { LOG_ERROR("failed to encode image guidance prompt"); return std::nullopt; @@ -785,15 +787,9 @@ namespace sd::pipeline { return false; } - // MiniMax-H3 is video-only. Its denoiser always splits the packed latent into a video and an - // audio half, and only generate_video ever computes the audio length, so reaching this - // function with an H3 checkpoint is guaranteed to die on - // GGML_ASSERT(!audio_input_cache.empty()) with a core dump, after the several minutes it - // takes to load the weights, and with nothing in the output pointing at the missing --mode. - // (The AnimateDiff path below routes vid_gen back through here, but that is SD1.5 plus a - // motion module, never H3.) - if (sd_version_is_minimax_h3(sd->version)) { - LOG_ERROR("MiniMax-H3 is a video model and cannot be run in img_gen mode; use --mode vid_gen"); + if (!sd_version_supports_image_generation(sd->version)) { + LOG_ERROR("%s cannot be run with generate_image(); use generate_video() or --mode vid_gen in the CLI", + model_version_to_str[sd->version]); return false; } @@ -802,6 +798,12 @@ namespace sd::pipeline { int64_t t0 = ggml_time_ms(); sd->vae_tiling_params = sd_img_gen_params->vae_tiling_params; GenerationRequest request(sd, sd_img_gen_params); + sd::ImagePreprocessor preprocessing(sd_img_gen_params->image_preprocess.rules); + sd_img_gen_params_t processed_params = *sd_img_gen_params; + if (!preprocessing.prepare_inputs(processed_params, request.width, request.height)) + return false; + sd_img_gen_params = &processed_params; + request.pm_params = processed_params.pm_params; LOG_INFO("generate_image %dx%d", request.width, request.height); sd->rng->manual_seed(request.seed); diff --git a/otherarch/sdcpp/src/pipeline/model_builders.cpp b/otherarch/sdcpp/src/pipeline/model_builders.cpp index 31b0f2b4c..9b8819a05 100644 --- a/otherarch/sdcpp/src/pipeline/model_builders.cpp +++ b/otherarch/sdcpp/src/pipeline/model_builders.cpp @@ -20,6 +20,7 @@ #include "model/diffusion/krea2.hpp" #include "model/diffusion/lens.hpp" #include "model/diffusion/lingbot_video.hpp" +#include "model/diffusion/llada_image.hpp" #include "model/diffusion/ltxv.hpp" #include "model/diffusion/mage_flow.hpp" #include "model/diffusion/minimax_h3.hpp" @@ -290,7 +291,8 @@ namespace sd::model_builders { result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), tensor_storage_map, "model.diffusion_model", - weight_manager); + weight_manager, + sd_ctx_params->model_args); } else { result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), tensor_storage_map, @@ -370,6 +372,19 @@ namespace sd::model_builders { "model.diffusion_model", version, weight_manager); + } else if (sd_version_is_llada_image(version)) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + "text_encoders.llm", + "queryformer", + "text_projection", + "sigvq", + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); } else if (sd_version_is_boogu_image(version)) { result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), tensor_storage_map, diff --git a/otherarch/sdcpp/src/pipeline/request.cpp b/otherarch/sdcpp/src/pipeline/request.cpp index 8befcd7c5..075ae67f2 100644 --- a/otherarch/sdcpp/src/pipeline/request.cpp +++ b/otherarch/sdcpp/src/pipeline/request.cpp @@ -73,6 +73,8 @@ namespace sd::pipeline { return LTX2_SCHEDULER; } else if (sd != nullptr && sd_version_is_ideogram4(sd->version)) { return LOGIT_NORMAL_SCHEDULER; + } else if (sd != nullptr && sd_version_is_llada_image(sd->version)) { + return LLADA_IMAGE_SCHEDULER; } return DISCRETE_SCHEDULER; } diff --git a/otherarch/sdcpp/src/pipeline/video.cpp b/otherarch/sdcpp/src/pipeline/video.cpp index f5312174d..7c0f85c6b 100644 --- a/otherarch/sdcpp/src/pipeline/video.cpp +++ b/otherarch/sdcpp/src/pipeline/video.cpp @@ -15,6 +15,7 @@ #include "model/vae/vae.hpp" #include "request.h" #include "runtime/denoiser.hpp" +#include "runtime/image_preprocess.h" // kcpp #include "model/vae/ltx_audio_vae.hpp" @@ -555,11 +556,15 @@ namespace sd::pipeline { } runner_guard{sd->audio_vae_model.get()}; if (sd_vid_gen_params->init_image.data) { - start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, request->width, request->height); + start_image = ensure_image_tensor_channels( + sd_image_to_tensor(sd_vid_gen_params->init_image, request->width, request->height), + sd->get_image_channels()); } if (sd_vid_gen_params->end_image.data) { - end_image = sd_image_to_tensor(sd_vid_gen_params->end_image, request->width, request->height); + end_image = ensure_image_tensor_channels( + sd_image_to_tensor(sd_vid_gen_params->end_image, request->width, request->height), + sd->get_image_channels()); } if (sd_version_is_minimax_h3(sd->version)) { @@ -1289,8 +1294,7 @@ namespace sd::pipeline { } int64_t prepare_start_ms = ggml_time_ms(); - embeds.cond = sd->cond_stage_model->get_learned_condition(sd->n_threads, - condition_params); + embeds.cond = sd->get_learned_condition(condition_params); if (embeds.cond.empty()) { LOG_ERROR("failed to encode video prompt"); return std::nullopt; @@ -1315,8 +1319,7 @@ namespace sd::pipeline { } if (request.use_uncond) { condition_params.text = request.negative_prompt; - embeds.uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, - condition_params); + embeds.uncond = sd->get_learned_condition(condition_params); if (embeds.uncond.empty()) { LOG_ERROR("failed to encode negative video prompt"); return std::nullopt; @@ -1548,7 +1551,9 @@ namespace sd::pipeline { sd::Tensor video_mask = make_ltxav_video_denoise_mask(video_latent, 1.f); if (sd_vid_gen_params->init_image.data != nullptr) { - sd::Tensor start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, image_width, image_height); + sd::Tensor start_image = ensure_image_tensor_channels( + sd_image_to_tensor(sd_vid_gen_params->init_image, image_width, image_height), + sd->get_image_channels()); if (!apply_ltxav_condition_image_by_latent_index(sd, start_image, &video_latent, @@ -1561,7 +1566,9 @@ namespace sd::pipeline { } if (sd_vid_gen_params->end_image.data != nullptr) { - sd::Tensor end_image = sd_image_to_tensor(sd_vid_gen_params->end_image, image_width, image_height); + sd::Tensor end_image = ensure_image_tensor_channels( + sd_image_to_tensor(sd_vid_gen_params->end_image, image_width, image_height), + sd->get_image_channels()); sd::Tensor end_image_latent = encode_ltxav_condition_image(sd, end_image, "end"); if (end_image_latent.empty()) { return false; @@ -1654,6 +1661,7 @@ namespace sd::pipeline { img_gen_params.qwen_image_layers = 0; img_gen_params.circular_x = sd_vid_gen_params->circular_x; img_gen_params.circular_y = sd_vid_gen_params->circular_y; + img_gen_params.image_preprocess = sd_vid_gen_params->image_preprocess; sd->animatediff_num_frames = n_frames; bool ok = generate_image(sd, &img_gen_params, frames_out, num_frames_out); @@ -1684,6 +1692,11 @@ namespace sd::pipeline { sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params; sd->apply_circular_axes(sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y); GenerationRequest request(sd, sd_vid_gen_params); + sd::ImagePreprocessor preprocessing(sd_vid_gen_params->image_preprocess.rules); + sd_vid_gen_params_t processed_params = *sd_vid_gen_params; + if (!preprocessing.prepare_inputs(processed_params, request.width, request.height)) + return false; + sd_vid_gen_params = &processed_params; if (fps_out != nullptr) { *fps_out = request.fps; } diff --git a/otherarch/sdcpp/src/runtime/denoiser.hpp b/otherarch/sdcpp/src/runtime/denoiser.hpp index 6e6f22d52..15c126086 100644 --- a/otherarch/sdcpp/src/runtime/denoiser.hpp +++ b/otherarch/sdcpp/src/runtime/denoiser.hpp @@ -728,12 +728,20 @@ inline float flux_time_shift(float mu, float sigma, float t) { // https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py#L289 struct FluxScheduler : SigmaScheduler { - int image_seq_len = 0; - float base_shift = 0.5f; - float max_shift = 1.15f; + int image_seq_len = 0; + int base_image_seq_len = 256; + int max_image_seq_len = 4096; + float base_shift = 0.5f; + float max_shift = 1.15f; + float shift_terminal = 0.0f; - explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr) + FluxScheduler(int image_seq_len, SDVersion version, const char* extra_sample_args = nullptr) : image_seq_len(image_seq_len) { + if (version == VERSION_QWEN_IMAGE_2_1) { + max_image_seq_len = 8192; + max_shift = 0.9f; + shift_terminal = 0.02f; + } parse_extra_sample_args(extra_sample_args); } @@ -752,10 +760,8 @@ struct FluxScheduler : SigmaScheduler { } float compute_mu() const { - constexpr float base_shift_anchor = 256.0f; - constexpr float max_shift_anchor = 4096.0f; - float m = (max_shift - base_shift) / (max_shift_anchor - base_shift_anchor); - float b = base_shift - m * base_shift_anchor; + float m = (max_shift - base_shift) / static_cast(max_image_seq_len - base_image_seq_len); + float b = base_shift - m * static_cast(base_image_seq_len); return static_cast(image_seq_len) * m + b; } @@ -764,7 +770,7 @@ struct FluxScheduler : SigmaScheduler { sigmas.reserve(n + 1); float mu = compute_mu(); - LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu); + LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f, shift_terminal=%.3f", image_seq_len, n, mu, shift_terminal); if (n == 0) { sigmas.push_back(1.0f); @@ -780,12 +786,70 @@ struct FluxScheduler : SigmaScheduler { } } + if (shift_terminal > 0.0f && n > 1) { + // The terminal shift applies to the last model evaluation, not the final zero sigma. + float scale_factor = (1.0f - sigmas[n - 1]) / (1.0f - shift_terminal); + if (std::isfinite(scale_factor) && scale_factor > 0.0f) { + for (uint32_t i = 0; i < n; ++i) { + sigmas[i] = 1.0f - (1.0f - sigmas[i]) / scale_factor; + } + } + } + sigmas[n] = 0.0f; return sigmas; } }; // https://github.com/black-forest-labs/flux2/blob/main/src/flux2/sampling.py#L244 +// LLaDA-Image does not use a shift-based flow schedule. The reference pipeline builds a +// Kumaraswamy-shaped grid over t = linspace(0.001, 1, n + 1)[:-1]: +// schedule = (1 - (1 - t^1.17)^0.8)^1.1 +// sigma = 1 - schedule +// Its scheduler config can also set use_uniform_sigmas, which replaces the whole curve with a +// plain linspace(1, 0, n + 1)[:-1] pre-shift grid. +struct LLaDAImageScheduler : SigmaScheduler { + bool uniform_sigmas = false; + + explicit LLaDAImageScheduler(const char* extra_sample_args = nullptr) { + parse_extra_sample_args(extra_sample_args); + } + + void parse_extra_sample_args(const char* extra_sample_args) { + for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "llada_image scheduler arg")) { + if (key == "uniform") { + if (!parse_strict_bool(value, uniform_sigmas)) { + LOG_WARN("ignoring invalid llada_image scheduler arg '%s=%s'", key.c_str(), value.c_str()); + } + } + } + } + + std::vector get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override { + std::vector sigmas; + sigmas.reserve(n + 1); + + if (n == 0) { + sigmas.push_back(1.0f); + return sigmas; + } + + for (uint32_t i = 0; i < n; ++i) { + float progress = static_cast(i) / static_cast(n); + if (uniform_sigmas) { + sigmas.push_back(1.0f - progress); + } else { + float t = 0.001f + progress * (1.0f - 0.001f); + float schedule = powf(1.0f - powf(1.0f - powf(t, 1.17f), 0.8f), 1.1f); + sigmas.push_back(1.0f - schedule); + } + } + + sigmas.push_back(0.0f); + return sigmas; + } +}; + struct Flux2Scheduler : SigmaScheduler { int image_seq_len = 0; @@ -1123,9 +1187,14 @@ struct Denoiser { scheduler = std::make_shared(image_seq_len); break; } + case LLADA_IMAGE_SCHEDULER: { + LOG_INFO("get_sigmas with LLaDA-Image scheduler"); + scheduler = std::make_shared(extra_sample_args); + break; + } case FLUX_SCHEDULER: { LOG_INFO("get_sigmas with Flux scheduler"); - scheduler = std::make_shared(image_seq_len, extra_sample_args); + scheduler = std::make_shared(image_seq_len, version, extra_sample_args); break; } default: diff --git a/otherarch/sdcpp/src/runtime/image_preprocess.cpp b/otherarch/sdcpp/src/runtime/image_preprocess.cpp new file mode 100644 index 000000000..a8df7c3b5 --- /dev/null +++ b/otherarch/sdcpp/src/runtime/image_preprocess.cpp @@ -0,0 +1,447 @@ +#include "image_preprocess.h" +#include "core/util.h" + +#include +#include + +namespace sd { + + static constexpr std::pair image_targets[] = { + {"init", ImageTarget::Init}, + {"end", ImageTarget::End}, + {"mask", ImageTarget::Mask}, + {"control", ImageTarget::Control}, + {"ref", ImageTarget::Ref}, + {"ip-adapter", ImageTarget::IPAdapter}, + {"id", ImageTarget::ID}, + {"control-frame", ImageTarget::ControlFrame}, + }; + + static constexpr std::pair image_resize_modes[] = { + {"auto", ImageResizeMode::Auto}, + {"none", ImageResizeMode::None}, + {"stretch", ImageResizeMode::Stretch}, + {"crop", ImageResizeMode::Crop}, + {"crop-resize", ImageResizeMode::CropResize}, + {"fit-pad", ImageResizeMode::FitPad}, + }; + + template + static bool parse_enum(const std::string& text, const std::pair (&names)[N], T& value) { + for (const auto& entry : names) { + if (text == entry.first) { + value = entry.second; + return true; + } + } + return false; + } + + template + static const char* enum_name(T value, const std::pair (&names)[N]) { + for (const auto& entry : names) { + if (value == entry.second) + return entry.first; + } + return "unknown"; + } + + template + static bool one_of(T value, std::initializer_list choices) { + return std::find(choices.begin(), choices.end(), value) != choices.end(); + } + + static bool one_of(const std::string& value, std::initializer_list choices) { + for (const char* choice : choices) { + if (value == choice) + return true; + } + return false; + } + + static ImageResizeMode resolve_mode(const std::map& options, ImageResizeMode default_mode) { + auto it = options.find("mode"); + ImageResizeMode mode = ImageResizeMode::Auto; + if (it != options.end()) + parse_enum(it->second, image_resize_modes, mode); + if (mode != ImageResizeMode::Auto) + return mode; + return options.count("width") && default_mode == ImageResizeMode::None ? ImageResizeMode::Stretch : default_mode; + } + + bool ImagePreprocessor::fail(const std::string& message) const { + LOG_ERROR("image preprocessing: %s", message.c_str()); + valid_ = false; + return false; + } + + ImagePreprocessor::ImagePreprocessor(const char* text) { + if (text == nullptr || trim(text).empty()) + return; + for (const auto& part : split_string(text, ';')) { + ImagePreprocessRule rule; + std::set keys; + if (trim(part).empty()) { + fail("empty rule"); + return; + } + for (const auto& entry : split_string(part, ',')) { + size_t equal = entry.find('='); + if (equal == std::string::npos) { + fail("expected key=value: " + entry); + return; + } + std::string key = trim(entry.substr(0, equal)); + std::string value = trim(entry.substr(equal + 1)); + bool ok = !value.empty() && keys.insert(key).second; + int number = 0; + if (key == "target") { + ok &= parse_enum(value, image_targets, rule.target); + } else if (key == "index") { + ok &= parse_strict_int(value, rule.index) && rule.index >= 0; + } else { + if (key == "mode") { + ImageResizeMode mode; + ok &= parse_enum(value, image_resize_modes, mode); + } else if (key == "filter") { + ok &= one_of(value, {"auto", "nearest", "nearest-exact", "bilinear", "bicubic", "lanczos"}); + } else if (key == "antialias") { + ok &= one_of(value, {"auto", "true", "false"}); + } else if (key == "canny") { + ok &= one_of(value, {"true", "false"}); + } else if (key == "anchor") { + ok &= one_of(value, {"center", "top", "bottom", "left", "right"}); + } else if (key == "width" || key == "height") { + ok &= parse_strict_int(value, number) && number > 0; + } else if (key == "pad_color") { + ok &= value.size() == 7 || value.size() == 9; + ok &= !value.empty() && value[0] == '#'; + for (size_t i = 1; i < value.size(); ++i) + ok &= std::isxdigit(static_cast(value[i])) != 0; + } else { + ok = false; + } + rule.options[key] = value; + } + if (!ok) { + fail("invalid or duplicate option: " + entry); + return; + } + } + if (!keys.count("target") || rule.options.empty() || + rule.options.count("width") != rule.options.count("height") || + (rule.index >= 0 && !one_of(rule.target, {ImageTarget::Ref, ImageTarget::ID, ImageTarget::ControlFrame}))) { + fail("invalid target, index, or incomplete dimensions: " + part); + return; + } + rules_.push_back(std::move(rule)); + } + for (const auto& rule : rules_) { + const auto options = resolve_options(rule.target, std::max(0, rule.index)); + if (options.count("antialias") && options.at("antialias") == "true" && options.count("filter") && + one_of(options.at("filter"), {"nearest", "nearest-exact"})) { + fail("antialias requires bilinear, bicubic, or lanczos"); + return; + } + } + } + + std::map ImagePreprocessor::resolve_options(ImageTarget target, int index) const { + std::map options; + for (int specificity = 0; specificity < 2; ++specificity) { + for (const auto& rule : rules_) { + if (rule.target == target && + rule.index == (specificity == 0 ? -1 : index)) { + for (const auto& entry : rule.options) + options[entry.first] = entry.second; + } + } + } + return options; + } + + bool ImagePreprocessor::validate_inputs(const sd_img_gen_params_t& params) const { + const std::map counts = { + {ImageTarget::Init, params.init_image.data != nullptr}, + {ImageTarget::Mask, params.mask_image.data != nullptr}, + {ImageTarget::Control, params.control_image.data != nullptr}, + {ImageTarget::IPAdapter, params.ip_adapter_image.data != nullptr}, + {ImageTarget::Ref, params.ref_images != nullptr ? params.ref_images_count : 0}, + {ImageTarget::ID, params.pm_params.id_images != nullptr ? params.pm_params.id_images_count : 0}, + }; + for (const auto& rule : rules_) { + auto it = counts.find(rule.target); + int count = it == counts.end() ? 0 : it->second; + if (count <= 0 || rule.index >= count) { + return fail(std::string("rule targets an unavailable image: ") + enum_name(rule.target, image_targets)); + } + } + return valid_; + } + + bool ImagePreprocessor::validate_inputs(const sd_vid_gen_params_t& params) const { + const std::map counts = { + {ImageTarget::Init, params.init_image.data != nullptr}, + {ImageTarget::End, params.end_image.data != nullptr}, + {ImageTarget::Ref, params.ref_images != nullptr ? params.ref_images_count : 0}, + {ImageTarget::ControlFrame, params.control_frames != nullptr ? params.control_frames_size : 0}, + }; + for (const auto& rule : rules_) { + auto it = counts.find(rule.target); + int count = it == counts.end() ? 0 : it->second; + if (count <= 0 || rule.index >= count) + return fail(std::string("rule targets an unavailable video input: ") + enum_name(rule.target, image_targets)); + } + return valid_; + } + + static int anchor_offset(int remaining, const std::string& anchor, bool horizontal) { + if (anchor == (horizontal ? "left" : "top")) + return 0; + if (anchor == (horizontal ? "right" : "bottom")) + return remaining; + return remaining / 2; + } + + Tensor ImagePreprocessor::apply_transform(const Tensor& image, const std::map& options, ImageTransform p, const std::string& label, ops::InterpolateMode default_filter) const { + auto value = [&](const char* key, const char* fallback) { + auto it = options.find(key); + return it == options.end() ? std::string(fallback) : it->second; + }; + std::string filter = value("filter", "auto"); + ops::InterpolateMode mode = default_filter; + if (filter == "nearest") + mode = ops::InterpolateMode::Nearest; + if (filter == "nearest-exact") + mode = ops::InterpolateMode::NearestExact; + if (filter == "bilinear") + mode = ops::InterpolateMode::Bilinear; + if (filter == "bicubic") + mode = ops::InterpolateMode::Bicubic; + if (filter == "lanczos") + mode = ops::InterpolateMode::Lanczos; + bool filtered = ops::is_2d_filter_interpolate_mode(mode); + bool antialias = value("antialias", "auto") == "true" || + (value("antialias", "auto") == "auto" && filtered && + (p.resize_width < p.crop_width || p.resize_height < p.crop_height)); + if (antialias && !filtered) { + fail(label + ": antialias requires bilinear, bicubic, or lanczos"); + return {}; + } + auto cropped = ops::slice(ops::slice(image, 0, p.x, p.x + p.crop_width), 1, p.y, p.y + p.crop_height); + int channels = static_cast(image.shape()[2]); + bool resize = p.resize_width != p.crop_width || p.resize_height != p.crop_height; + if (resize && channels == 4 && filtered) { + for (int64_t i = 0, pixels = cropped.shape()[0] * cropped.shape()[1]; i < pixels; ++i) { + for (int c = 0; c < 3; ++c) + cropped[i + c * pixels] *= cropped[i + 3 * pixels]; + } + } + auto resized = ops::interpolate(cropped, {p.resize_width, p.resize_height, channels, 1}, mode, false, antialias); + if (resize && channels == 4 && filtered) { + for (int64_t i = 0, pixels = resized.shape()[0] * resized.shape()[1]; i < pixels; ++i) { + float alpha = std::clamp(resized[i + 3 * pixels], 0.f, 1.f); + for (int c = 0; c < 3; ++c) + resized[i + c * pixels] = alpha > 1e-6f ? resized[i + c * pixels] / alpha : 0.f; + } + } + resized = ops::clamp(resized, 0.f, 1.f); + Tensor output({p.width, p.height, channels, 1}); + std::string color = value("pad_color", "#000000ff"); + if (color.size() == 7) + color += "ff"; + uint8_t rgba[4]; + for (int c = 0; c < 4; ++c) + rgba[c] = static_cast(std::strtoul(color.substr(1 + c * 2, 2).c_str(), nullptr, 16)); + for (int c = 0; c < channels; ++c) { + float fill = rgba[channels == 1 ? 0 : c] / 255.f; + for (int y = 0; y < p.height; ++y) { + for (int x = 0; x < p.width; ++x) { + output.index(x, y, c, 0) = x >= p.pad_x && x < p.pad_x + p.resize_width && y >= p.pad_y && y < p.pad_y + p.resize_height + ? resized.index(x - p.pad_x, y - p.pad_y, c, 0) + : fill; + } + } + } + LOG_INFO("preprocess %s: %dx%d crop=(%d,%d,%d,%d) resize=%dx%d pad=(%d,%d) output=%dx%d filter=%s(%d) antialias=%s", + label.c_str(), p.source_width, p.source_height, p.x, p.y, p.crop_width, p.crop_height, + p.resize_width, p.resize_height, p.pad_x, p.pad_y, p.width, p.height, filter.c_str(), static_cast(mode), BOOL_STR(antialias)); + return output; + } + + Tensor ImagePreprocessor::apply_geometry(const Tensor& image, ImageTarget target, int index, int width, int height, ImageResizeMode default_mode, ops::InterpolateMode default_filter, ImageTransform* plan_out) const { + if (!valid_ || image.empty()) + return {}; + const std::string label = std::string(enum_name(target, image_targets)) + "[" + std::to_string(index) + "]"; + auto options = resolve_options(target, index); + if (image.dim() != 4 || image.shape()[3] != 1 || image.shape()[2] < 1 || image.shape()[2] > 4) { + fail(label + ": expected one image with 1 to 4 channels"); + return {}; + } + ImageTransform p; + p.source_width = p.crop_width = static_cast(image.shape()[0]); + p.source_height = p.crop_height = static_cast(image.shape()[1]); + int target_width = width > 0 ? width : p.source_width; + int target_height = height > 0 ? height : p.source_height; + if (options.count("width")) { + parse_strict_int(options.at("width"), target_width); + parse_strict_int(options.at("height"), target_height); + } + ImageResizeMode mode = resolve_mode(options, default_mode); + std::string anchor = options.count("anchor") ? options.at("anchor") : "center"; + p.width = p.resize_width = target_width; + p.height = p.resize_height = target_height; + if (mode == ImageResizeMode::None) { + if (options.count("width") && (target_width != p.source_width || target_height != p.source_height)) { + fail(label + ": mode=none conflicts with requested dimensions"); + return {}; + } + p.width = p.resize_width = p.source_width; + p.height = p.resize_height = p.source_height; + } else if (mode == ImageResizeMode::Crop || mode == ImageResizeMode::CropResize) { + if (mode == ImageResizeMode::Crop) { + p.crop_width = target_width; + p.crop_height = target_height; + } else if (int64_t(p.source_width) * target_height > int64_t(p.source_height) * target_width) { + p.crop_width = std::max(1, static_cast(int64_t(p.source_height) * target_width / target_height)); + } else { + p.crop_height = std::max(1, static_cast(int64_t(p.source_width) * target_height / target_width)); + } + if (p.crop_width > p.source_width || p.crop_height > p.source_height) { + fail(label + ": crop exceeds source dimensions"); + return {}; + } + p.x = anchor_offset(p.source_width - p.crop_width, anchor, true); + p.y = anchor_offset(p.source_height - p.crop_height, anchor, false); + } else if (mode == ImageResizeMode::FitPad) { + double scale = std::min(double(target_width) / p.source_width, double(target_height) / p.source_height); + p.resize_width = std::max(1, std::min(target_width, static_cast(std::round(p.source_width * scale)))); + p.resize_height = std::max(1, std::min(target_height, static_cast(std::round(p.source_height * scale)))); + p.pad_x = anchor_offset(target_width - p.resize_width, anchor, true); + p.pad_y = anchor_offset(target_height - p.resize_height, anchor, false); + } + if (p.width <= 0 || p.height <= 0) { + fail(label + ": invalid output dimensions"); + return {}; + } + uint64_t max_pixels = std::min(INT64_MAX, SIZE_MAX / sizeof(float)) / static_cast(image.shape()[2]); + if (uint64_t(p.width) * p.height > max_pixels || uint64_t(p.resize_width) * p.resize_height > max_pixels) { + fail(label + ": image allocation size overflows"); + return {}; + } + if (plan_out != nullptr) + *plan_out = p; + return apply_transform(image, options, p, label, default_filter); + } + + Tensor ImagePreprocessor::preprocess_input(sd_image_t image, ImageTarget target, int index, int width, int height) { + if (image.data == nullptr || image.width == 0 || image.height == 0 || image.width > INT_MAX || image.height > INT_MAX || image.channel < 1 || image.channel > 4) { + fail(std::string(enum_name(target, image_targets)) + ": invalid input image"); + return {}; + } + auto tensor = sd_image_to_tensor(image); + if (target == ImageTarget::Mask && has_init_transform_) { + auto options = resolve_options(target, index); + if (image.width != init_transform_.source_width || image.height != init_transform_.source_height) { + fail("mask and init source dimensions must match"); + return {}; + } + bool geometry_override = options.count("width") || options.count("anchor") || + (options.count("mode") && options.at("mode") != "auto"); + if (geometry_override) { + ImageTransform p; + auto init_options = resolve_options(ImageTarget::Init, 0); + ImageResizeMode default_mode = resolve_mode(init_options, ImageResizeMode::CropResize); + auto result = apply_geometry(tensor, target, index, init_transform_.width, init_transform_.height, default_mode, ops::InterpolateMode::NearestExact, &p); + if (result.empty()) + return {}; + const auto& q = init_transform_; + if (p.x != q.x || p.y != q.y || p.crop_width != q.crop_width || p.crop_height != q.crop_height || + p.resize_width != q.resize_width || p.resize_height != q.resize_height || p.pad_x != q.pad_x || p.pad_y != q.pad_y || p.width != q.width || p.height != q.height) { + fail("mask geometry conflicts with init; configure geometry on init and filter on mask"); + return {}; + } + return result; + } + return apply_transform(tensor, options, init_transform_, "mask[0]", ops::InterpolateMode::NearestExact); + } + auto result = apply_geometry(tensor, target, index, width, height, width > 0 ? ImageResizeMode::CropResize : ImageResizeMode::None, + target == ImageTarget::Mask ? ops::InterpolateMode::NearestExact : ops::InterpolateMode::Nearest, + target == ImageTarget::Init ? &init_transform_ : nullptr); + if (target == ImageTarget::Init) + has_init_transform_ = !result.empty(); + return result; + } + + ImagePreprocessor::~ImagePreprocessor() { + for (const auto& image : owned_images_) + std::free(image.data); + } + + bool ImagePreprocessor::prepare_image(sd_image_t& image, ImageTarget target, int index, int width, int height) { + if (image.data == nullptr) + return true; + auto options = resolve_options(target, index); + bool canny = options.count("canny") && options.at("canny") == "true"; + auto tensor = preprocess_input(image, target, index, width, height); + if (tensor.empty()) + return false; + auto output = tensor_to_sd_image(tensor); + if (output.data == nullptr) + return fail("could not allocate input preprocessing buffer"); + owned_images_.push_back(output); + if (canny && !preprocess_canny(output, 0.08f, 0.08f, 0.8f, 1.f, false)) + return fail("Canny preprocessing failed"); + image = output; + return true; + } + + bool ImagePreprocessor::prepare_array(sd_image_t*& images, int count, ImageTarget target, std::vector& storage, int width, int height) { + if (count < 0 || (count > 0 && images == nullptr)) + return fail(std::string("invalid image array: ") + enum_name(target, image_targets)); + if (count == 0) + return true; + storage.assign(images, images + count); + for (int i = 0; i < count; ++i) { + if (storage[i].data == nullptr) + return fail(std::string("empty image in array: ") + enum_name(target, image_targets)); + if (!prepare_image(storage[i], target, i, width, height)) + return false; + } + images = storage.data(); + return true; + } + + bool ImagePreprocessor::prepare_inputs(sd_img_gen_params_t& params, int width, int height) { + if (prepared_) + return fail("inputs have already been prepared"); + prepared_ = true; + if (!valid_ || !validate_inputs(params)) + return false; + if (!prepare_image(params.init_image, ImageTarget::Init, 0, width, height) || + !prepare_image(params.mask_image, ImageTarget::Mask, 0, width, height) || + !prepare_image(params.control_image, ImageTarget::Control, 0, width, height) || + !prepare_image(params.ip_adapter_image, ImageTarget::IPAdapter, 0, -1, -1) || + !prepare_array(params.ref_images, params.ref_images_count, ImageTarget::Ref, ref_images_) || + !prepare_array(params.pm_params.id_images, params.pm_params.id_images_count, ImageTarget::ID, id_images_)) + return false; + params.image_preprocess = {}; + return true; + } + + bool ImagePreprocessor::prepare_inputs(sd_vid_gen_params_t& params, int width, int height) { + if (prepared_) + return fail("inputs have already been prepared"); + prepared_ = true; + if (!valid_ || !validate_inputs(params)) + return false; + if (!prepare_image(params.init_image, ImageTarget::Init, 0, width, height) || + !prepare_image(params.end_image, ImageTarget::End, 0, width, height) || + !prepare_array(params.ref_images, params.ref_images_count, ImageTarget::Ref, ref_images_) || + !prepare_array(params.control_frames, params.control_frames_size, ImageTarget::ControlFrame, control_frames_, width, height)) + return false; + params.image_preprocess = {}; + return true; + } + +} // namespace sd diff --git a/otherarch/sdcpp/src/runtime/image_preprocess.h b/otherarch/sdcpp/src/runtime/image_preprocess.h new file mode 100644 index 000000000..b75c4e0e5 --- /dev/null +++ b/otherarch/sdcpp/src/runtime/image_preprocess.h @@ -0,0 +1,88 @@ +#ifndef __SD_RUNTIME_IMAGE_PREPROCESS_H__ +#define __SD_RUNTIME_IMAGE_PREPROCESS_H__ + +#include +#include +#include + +#include "core/tensor.hpp" +#include "stable-diffusion.h" + +namespace sd { + + enum class ImageTarget { + Init, + End, + Mask, + Control, + Ref, + IPAdapter, + ID, + ControlFrame, + }; + + enum class ImageResizeMode { + Auto, + None, + Stretch, + Crop, + CropResize, + FitPad, + }; + + struct ImageTransform { + int source_width = 0; + int source_height = 0; + int x = 0; + int y = 0; + int crop_width = 0; + int crop_height = 0; + int resize_width = 0; + int resize_height = 0; + int width = 0; + int height = 0; + int pad_x = 0; + int pad_y = 0; + }; + + struct ImagePreprocessRule { + ImageTarget target = ImageTarget::Init; + int index = -1; + std::map options; + }; + + class ImagePreprocessor { + std::vector rules_; + mutable bool valid_ = true; + ImageTransform init_transform_; + bool has_init_transform_ = false; + bool prepared_ = false; + std::vector owned_images_; + std::vector ref_images_; + std::vector id_images_; + std::vector control_frames_; + + bool fail(const std::string& message) const; + std::map resolve_options(ImageTarget target, int index) const; + Tensor apply_transform(const Tensor& image, const std::map& options, ImageTransform plan, const std::string& label, ops::InterpolateMode default_filter) const; + + bool prepare_image(sd_image_t& image, ImageTarget target, int index, int width, int height); + bool prepare_array(sd_image_t*& images, int count, ImageTarget target, std::vector& storage, int width = -1, int height = -1); + + public: + explicit ImagePreprocessor(const char* rules = nullptr); + ~ImagePreprocessor(); + ImagePreprocessor(const ImagePreprocessor&) = delete; + ImagePreprocessor& operator=(const ImagePreprocessor&) = delete; + bool prepare_inputs(sd_img_gen_params_t& params, int width, int height); + bool prepare_inputs(sd_vid_gen_params_t& params, int width, int height); + bool is_valid() const { return valid_; } + bool validate_inputs(const sd_img_gen_params_t& params) const; + bool validate_inputs(const sd_vid_gen_params_t& params) const; + Tensor apply_geometry(const Tensor& image, ImageTarget target, int index, int width, int height, ImageResizeMode default_mode = ImageResizeMode::Stretch, ops::InterpolateMode default_filter = ops::InterpolateMode::Nearest, ImageTransform* plan_out = nullptr) const; + Tensor preprocess_input(sd_image_t image, ImageTarget target, int index = 0, int width = -1, int height = -1); + }; + +} // namespace sd + +#endif // __SD_RUNTIME_IMAGE_PREPROCESS_H__ diff --git a/otherarch/sdcpp/src/runtime/latent-preview.h b/otherarch/sdcpp/src/runtime/latent-preview.h index c19a0eebc..5baf6f53c 100644 --- a/otherarch/sdcpp/src/runtime/latent-preview.h +++ b/otherarch/sdcpp/src/runtime/latent-preview.h @@ -4,6 +4,86 @@ #include "core/tensor.hpp" #include "ggml.h" +// RGB is projected to [-1, 1]; alpha is projected directly to [0, 1]. +const float qwen21_latent_rgb_proj[64][3] = { + {0.00860495522f, 0.01219501462f, -0.00321337196f}, + {0.01889233090f, 0.01246581216f, 0.01074959482f}, + {0.1255941446f, 0.1176879344f, -0.0332212352f}, + {0.0418238528f, 0.1043427754f, 0.0121666316f}, + {0.02025338f, 0.01453670296f, -0.000224336044f}, + {-0.01896720702f, -0.0206099030f, -0.0322728584f}, + {0.00438984796f, -0.01374969766f, 0.02849196f}, + {-0.0374495856f, -0.0286777126f, -0.0693192810f}, + {0.01511914734f, 0.0242979386f, 0.0553878870f}, + {-0.1138629518f, -0.020391466f, 0.001550520522f}, + {-0.0233650696f, -0.0417292018f, -0.0362361182f}, + {-0.0351603342f, -0.0243595924f, -0.00216261038f}, + {0.01093355288f, -0.0373466924f, 0.00241315350f}, + {0.01778704744f, -0.00401984678f, -0.0343259192f}, + {0.0486059334f, 0.0253144f, 0.0672564966f}, + {0.0309463558f, 0.0277963166f, 0.0520869622f}, + {0.0374485008f, 0.0551753676f, 0.0225853902f}, + {-0.0090809962f, -0.004756176f, 0.00636443612f}, + {-0.0270455652f, -0.0384966954f, -0.00905908082f}, + {-0.00553493756f, 0.01484553684f, -0.0211502468f}, + {0.01319502562f, 0.00948005666f, 0.0483789212f}, + {-0.00931847104f, -0.00276452734f, -0.01011985302f}, + {0.0180478258f, 0.01614954356f, -0.0209424690f}, + {-0.0214530434f, -0.00272961176f, 0.0217887476f}, + {-0.0636772304f, -0.0208893548f, 0.0479167742f}, + {-0.0250321236f, -0.0286715676f, 0.0530110146f}, + {-0.01853078078f, 0.01647272818f, -0.00207747588f}, + {0.0023101082f, 0.01228800748f, 0.01303505006f}, + {-0.01243671408f, -0.0258638728f, -0.0379116264f}, + {0.00598934710f, 0.00642563550f, -0.01234514304f}, + {-0.0296733996f, -0.0234698050f, 0.00060018212f}, + {-0.0322019498f, -0.0529200462f, -0.00344987414f}, + {-0.00205026458f, -0.00846599446f, 0.00455971038f}, + {-0.01082227064f, 0.0315661948f, -0.0677753362f}, + {0.0645553474f, 0.1109666998f, 0.0674744864f}, + {0.01036801108f, -0.00484841210f, -0.001529168474f}, + {0.01264353566f, 0.01548126338f, -0.00966374324f}, + {-0.0223892408f, -0.00871751526f, -0.000306421670f}, + {0.0271322742f, 0.03496524f, -0.0089692858f}, + {0.0512178672f, 0.0173080034f, 0.00804227746f}, + {0.01210987192f, 0.00758025926f, -0.00281712586f}, + {0.1897278390f, 0.1210261828f, 0.062603892f}, + {0.0208058822f, 0.00547548182f, 0.01262955638f}, + {0.00813332858f, 0.01015930914f, 0.01301771290f}, + {-0.000927236014f, -0.00152540594f, -0.00599213302f}, + {0.01663314616f, -0.00582789626f, 0.0163958132f}, + {-0.0252546342f, -0.0604193732f, -0.1606919922f}, + {-0.091722686f, -0.0409201224f, -0.0959576198f}, + {0.0282963112f, -0.01387223872f, -0.01648814464f}, + {0.0552316818f, 0.0967547788f, 0.0413586632f}, + {0.00922849292f, 0.00451467542f, -0.0529172378f}, + {0.0558600768f, 0.0122988308f, -0.01445942422f}, + {0.000210660902f, -0.01295958782f, -0.01804761764f}, + {0.0358136250f, -0.0472505970f, -0.1156405142f}, + {-0.0506390696f, -0.0471914842f, 0.0349791468f}, + {-0.0480143168f, 0.00628389868f, -0.0545163826f}, + {0.0315499582f, 0.0564846606f, -0.0430850488f}, + {-0.0362330316f, -0.01267788554f, 0.0061024772f}, + {0.0038627542f, 0.00911055916f, -0.00758526008f}, + {-0.0447103298f, -0.00835411408f, 0.01545872328f}, + {-0.015006738f, 0.00270612302f, -0.00784361356f}, + {-0.0221755048f, -0.0513344748f, -0.0475317424f}, + {-0.01036656294f, -0.00422146068f, -0.0213499052f}, + {0.01788952706f, 0.01191944190f, 0.0397205238f}, +}; +float qwen21_latent_rgb_bias[3] = {-0.043293118f, -0.02695978f, -0.11986706f}; + +const float qwen21_latent_alpha_proj[64] = { + -0.0416241114f, -0.00678954612f, -0.0169095515f, -0.0230551401f, 0.0100882595f, 0.00655586802f, 0.0401166874f, -0.0055510216f, + 0.0224234441f, -0.0389640963f, -0.0114492163f, -0.00721128977f, -0.0029064082f, 0.0150300547f, -0.00321615308f, -0.0498856338f, + -0.0215251401f, 0.0240220482f, 0.0117338008f, -0.0460420624f, 0.0387872889f, 0.0131517207f, 0.0147100836f, 0.0266985286f, + 0.0153097324f, -0.0418119757f, 0.0421013917f, 0.0401724499f, 0.00972515915f, 0.011718495f, 0.0117622291f, 0.0136505134f, + -0.0350017363f, -0.0100692606f, -0.0131582529f, -0.00660639315f, 0.00253611396f, -0.0195736368f, -0.04240184f, 0.0321299262f, + 0.0106089414f, -0.0179845306f, -0.00806212708f, 0.0135889057f, -0.0157393098f, -0.0267791344f, 0.0109068534f, 0.0283931966f, + -0.0435370078f, 0.00187883536f, -0.0108995378f, -0.0450757676f, -0.0699481501f, 0.0123562106f, -0.0222592249f, 0.0216155907f, + 0.0563755424f, -0.0073379912f, 0.0160012921f, 0.0411637742f, 0.0189607258f, -0.024025029f, -0.0161487905f, -0.016913203f}; +const float qwen21_latent_alpha_bias = 0.871268134f; + const float minimax_latent_rgb_proj[24][3] = { {0.19819857f, 0.11584999f, 0.07929777f}, {-0.16047224f, -0.10601170f, -0.15996324f}, @@ -324,7 +404,7 @@ const float sd_latent_rgb_proj[4][3] = { {-0.178022f, -0.200862f, -0.678514f}}; float sd_latent_rgb_bias[3] = {-0.017478f, -0.055834f, -0.105825f}; -void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size) { +void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size, const float* latent_alpha_proj = nullptr, float latent_alpha_bias = 1.f) { size_t buffer_head = 0; uint32_t latent_width = static_cast(latents->ne[0]); @@ -338,7 +418,8 @@ void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*l uint32_t rgb_width = latent_width * patch_size; uint32_t rgb_height = latent_height * patch_size; - uint32_t unpatched_dim = dim / (patch_size * patch_size); + uint32_t unpatched_dim = dim / (patch_size * patch_size); + const uint32_t output_channels = latent_alpha_proj != nullptr ? 4 : 3; for (uint32_t k = 0; k < frames; k++) { for (uint32_t rgb_x = 0; rgb_x < rgb_width; rgb_x++) { @@ -356,13 +437,16 @@ void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*l // should be incremented by 1 for each pixel size_t pixel_id = k * rgb_width * rgb_height + rgb_y * rgb_width + rgb_x; - float r = 0, g = 0, b = 0; + float r = 0, g = 0, b = 0, a = 0; if (latent_rgb_proj != nullptr) { for (uint32_t d = 0; d < unpatched_dim; d++) { float value = *(float*)((char*)latents->data + latent_id + (d * patch_size * patch_size + channel_offset) * latents->nb[ggml_n_dims(latents) - 1]); r += value * latent_rgb_proj[d][0]; g += value * latent_rgb_proj[d][1]; b += value * latent_rgb_proj[d][2]; + if (latent_alpha_proj != nullptr) { + a += value * latent_alpha_proj[d]; + } } } else { // interpret first 3 channels as RGB @@ -386,9 +470,13 @@ void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*l g = g >= 0 ? g <= 1 ? g : 1 : 0; b = b >= 0 ? b <= 1 ? b : 1 : 0; - buffer[pixel_id * 3 + 0] = (uint8_t)(r * 255); - buffer[pixel_id * 3 + 1] = (uint8_t)(g * 255); - buffer[pixel_id * 3 + 2] = (uint8_t)(b * 255); + buffer[pixel_id * output_channels + 0] = (uint8_t)(r * 255); + buffer[pixel_id * output_channels + 1] = (uint8_t)(g * 255); + buffer[pixel_id * output_channels + 2] = (uint8_t)(b * 255); + if (latent_alpha_proj != nullptr) { + a = std::min(1.0f, std::max(0.0f, a + latent_alpha_bias)); + buffer[pixel_id * output_channels + 3] = (uint8_t)(a * 255); + } } } } @@ -398,16 +486,17 @@ static inline bool preview_latent_tensor_is_video(const sd::Tensor& laten return latents.dim() == 5; } -void preview_latent_video(uint8_t* buffer, const sd::Tensor& latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size) { +void preview_latent_video(uint8_t* buffer, const sd::Tensor& latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size, const float* latent_alpha_proj = nullptr, float latent_alpha_bias = 1.f) { uint32_t latent_width = static_cast(latents.shape()[0]); uint32_t latent_height = static_cast(latents.shape()[1]); bool is_video = preview_latent_tensor_is_video(latents); uint32_t frames = is_video ? static_cast(latents.shape()[2]) : 1; uint32_t dim = is_video ? static_cast(latents.shape()[3]) : static_cast(latents.shape()[2]); - uint32_t rgb_width = latent_width * patch_size; - uint32_t rgb_height = latent_height * patch_size; - uint32_t unpatched_dim = dim / (patch_size * patch_size); + uint32_t rgb_width = latent_width * patch_size; + uint32_t rgb_height = latent_height * patch_size; + uint32_t unpatched_dim = dim / (patch_size * patch_size); + const uint32_t output_channels = latent_alpha_proj != nullptr ? 4 : 3; for (uint32_t k = 0; k < frames; k++) { for (uint32_t rgb_x = 0; rgb_x < rgb_width; rgb_x++) { @@ -427,7 +516,7 @@ void preview_latent_video(uint8_t* buffer, const sd::Tensor& latents, con : latents.values()[latent_x + latent_width * (latent_y + latent_height * latent_channel)]; }; - float r = 0.f, g = 0.f, b = 0.f; + float r = 0.f, g = 0.f, b = 0.f, a = 0.f; if (latent_rgb_proj != nullptr) { for (uint32_t d = 0; d < unpatched_dim; d++) { uint32_t latent_channel = d * patch_size * patch_size + channel_offset; @@ -435,6 +524,9 @@ void preview_latent_video(uint8_t* buffer, const sd::Tensor& latents, con r += value * latent_rgb_proj[d][0]; g += value * latent_rgb_proj[d][1]; b += value * latent_rgb_proj[d][2]; + if (latent_alpha_proj != nullptr) { + a += value * latent_alpha_proj[d]; + } } } else { r = latent_value(0); @@ -450,9 +542,13 @@ void preview_latent_video(uint8_t* buffer, const sd::Tensor& latents, con g = std::min(1.0f, std::max(0.0f, g * .5f + .5f)); b = std::min(1.0f, std::max(0.0f, b * .5f + .5f)); - buffer[pixel_id * 3 + 0] = (uint8_t)(r * 255); - buffer[pixel_id * 3 + 1] = (uint8_t)(g * 255); - buffer[pixel_id * 3 + 2] = (uint8_t)(b * 255); + buffer[pixel_id * output_channels + 0] = (uint8_t)(r * 255); + buffer[pixel_id * output_channels + 1] = (uint8_t)(g * 255); + buffer[pixel_id * output_channels + 2] = (uint8_t)(b * 255); + if (latent_alpha_proj != nullptr) { + a = std::min(1.0f, std::max(0.0f, a + latent_alpha_bias)); + buffer[pixel_id * output_channels + 3] = (uint8_t)(a * 255); + } } } } diff --git a/otherarch/sdcpp/src/runtime/preprocessing.hpp b/otherarch/sdcpp/src/runtime/preprocessing.hpp index 1990d15f4..ab4ca29b6 100644 --- a/otherarch/sdcpp/src/runtime/preprocessing.hpp +++ b/otherarch/sdcpp/src/runtime/preprocessing.hpp @@ -165,16 +165,18 @@ static inline sd::Tensor convolve_tensor(const sd::Tensor& input, return output; } -static inline sd::Tensor grayscale_tensor(const sd::Tensor& rgb_img) { - GGML_ASSERT(rgb_img.dim() == 4); - GGML_ASSERT(rgb_img.shape()[2] >= 3); - sd::Tensor grayscale({rgb_img.shape()[0], rgb_img.shape()[1], 1, rgb_img.shape()[3]}); - for (int64_t iy = 0; iy < rgb_img.shape()[1]; ++iy) { - for (int64_t ix = 0; ix < rgb_img.shape()[0]; ++ix) { - float r = preprocessing_get_4d(rgb_img, ix, iy, 0, 0); - float g = preprocessing_get_4d(rgb_img, ix, iy, 1, 0); - float b = preprocessing_get_4d(rgb_img, ix, iy, 2, 0); - float gray = 0.2989f * r + 0.5870f * g + 0.1140f * b; +static inline sd::Tensor grayscale_tensor(const sd::Tensor& image) { + GGML_ASSERT(image.dim() == 4); + GGML_ASSERT(image.shape()[2] >= 1); + sd::Tensor grayscale({image.shape()[0], image.shape()[1], 1, image.shape()[3]}); + for (int64_t iy = 0; iy < image.shape()[1]; ++iy) { + for (int64_t ix = 0; ix < image.shape()[0]; ++ix) { + float gray = preprocessing_get_4d(image, ix, iy, 0, 0); + if (image.shape()[2] >= 3) { + float g = preprocessing_get_4d(image, ix, iy, 1, 0); + float b = preprocessing_get_4d(image, ix, iy, 2, 0); + gray = 0.2989f * gray + 0.5870f * g + 0.1140f * b; + } preprocessing_set_4d(grayscale, gray, ix, iy, 0, 0); } } @@ -317,11 +319,12 @@ bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold, image_gray = non_max_supression(G, theta); threshold_hystersis(&image_gray, high_threshold, low_threshold, weak, strong); + const uint32_t color_channels = img.channel == 2 || img.channel == 4 ? img.channel - 1 : img.channel; for (uint32_t iy = 0; iy < img.height; ++iy) { for (uint32_t ix = 0; ix < img.width; ++ix) { float gray = preprocessing_get_4d(image_gray, ix, iy, 0, 0); gray = inverse ? 1.0f - gray : gray; - for (uint32_t c = 0; c < img.channel; ++c) { + for (uint32_t c = 0; c < color_channels; ++c) { preprocessing_set_4d(image, gray, ix, iy, c, 0); } } diff --git a/otherarch/sdcpp/src/runtime/sample-cache.cpp b/otherarch/sdcpp/src/runtime/sample-cache.cpp index bfa3d4487..0b92c7335 100644 --- a/otherarch/sdcpp/src/runtime/sample-cache.cpp +++ b/otherarch/sdcpp/src/runtime/sample-cache.cpp @@ -275,6 +275,26 @@ namespace sd_sample { } } + void SampleStepCacheDispatcher::invalidate_condition(const void* condition) { + if (condition == nullptr) { + return; + } + + switch (runtime.mode) { + case SampleCacheMode::EASYCACHE: + runtime.easycache.cache_diffs.erase(condition); + break; + case SampleCacheMode::UCACHE: + runtime.ucache.cache_diffs.erase(condition); + break; + case SampleCacheMode::CACHEDIT: + runtime.cachedit.cache_diffs.erase(condition); + break; + case SampleCacheMode::NONE: + break; + } + } + bool SampleStepCacheDispatcher::is_step_skipped() const { switch (runtime.mode) { case SampleCacheMode::EASYCACHE: diff --git a/otherarch/sdcpp/src/runtime/sample-cache.h b/otherarch/sdcpp/src/runtime/sample-cache.h index cb0fe7bc5..30ff66e5f 100644 --- a/otherarch/sdcpp/src/runtime/sample-cache.h +++ b/otherarch/sdcpp/src/runtime/sample-cache.h @@ -46,6 +46,7 @@ namespace sd_sample { bool before_condition(const void* condition, const sd::Tensor& input, sd::Tensor* output); void after_condition(const void* condition, const sd::Tensor& input, const sd::Tensor& output); + void invalidate_condition(const void* condition); bool is_step_skipped() const; }; diff --git a/otherarch/sdcpp/src/stable-diffusion.cpp b/otherarch/sdcpp/src/stable-diffusion.cpp index 7e3855d55..3870be256 100644 --- a/otherarch/sdcpp/src/stable-diffusion.cpp +++ b/otherarch/sdcpp/src/stable-diffusion.cpp @@ -137,6 +137,7 @@ const char* scheduler_to_str[] = { "flux2", "flux", "beta", + "llada_image", }; static_assert(SCHEDULER_COUNT == sizeof(scheduler_to_str) / sizeof(scheduler_to_str[0]), @@ -325,6 +326,7 @@ void sd_hires_params_init(sd_hires_params_t* hires_params) { void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { *sd_ctx_params = {}; sd_ctx_params->n_threads = sd_get_num_physical_cores(); + sd_ctx_params->conditioning_cache_size = 4; sd_ctx_params->wtype = SD_TYPE_COUNT; sd_ctx_params->rng_type = CUDA_RNG; sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT; @@ -336,6 +338,7 @@ 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->sage_attn = false; sd_ctx_params->linear_scale = 0.f; sd_ctx_params->attn_scale = 0.f; sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; @@ -376,6 +379,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "pulid_weights_path: %s\n" "tensor_type_rules: %s\n" "n_threads: %d\n" + "conditioning_cache_size: %d\n" "wtype: %s\n" "rng_type: %s\n" "sampler_rng_type: %s\n" @@ -391,6 +395,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "auto_fit: %s\n" "flash_attn: %s\n" "diffusion_flash_attn: %s\n" + "sage_attn: %s\n" "linear_scale: %g\n" "attn_scale: %g\n" "vae_format: %s\n", @@ -415,6 +420,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { SAFE_STR(sd_ctx_params->pulid_weights_path), SAFE_STR(sd_ctx_params->tensor_type_rules), sd_ctx_params->n_threads, + sd_ctx_params->conditioning_cache_size, sd_type_name(sd_ctx_params->wtype), sd_rng_type_name(sd_ctx_params->rng_type), sd_rng_type_name(sd_ctx_params->sampler_rng_type), @@ -430,6 +436,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { 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->sage_attn), sd_ctx_params->linear_scale, sd_ctx_params->attn_scale, sd_vae_format_name(sd_ctx_params->vae_format)); @@ -627,14 +634,6 @@ struct sd_ctx_t { StableDiffusionGGML* sd = nullptr; }; -static bool sd_version_supports_video_generation(SDVersion version) { - return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version); -} - -static bool sd_version_supports_image_generation(SDVersion version) { - return !sd_version_supports_video_generation(version); -} - sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) { sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t)); if (sd_ctx == nullptr) { @@ -756,28 +755,20 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, int* num_frames_out, sd_audio_t** audio_out, int* fps_out) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) { - if (fps_out != nullptr) { - *fps_out = 0; - } - return false; - } - - if (frames_out != nullptr) { + if (frames_out != nullptr) *frames_out = nullptr; - } - if (audio_out != nullptr) { + if (audio_out != nullptr) *audio_out = nullptr; - } - if (num_frames_out != nullptr) { + if (num_frames_out != nullptr) *num_frames_out = 0; + if (fps_out != nullptr) + *fps_out = 0; + if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) { + return false; } StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd); if (!execution.ready) { - if (fps_out != nullptr) { - *fps_out = 0; - } return false; } @@ -827,7 +818,7 @@ namespace kcpp_sd { model_info res = {}; SDVersion loadedsdver = (SDVersion)get_loaded_sd_version(ctx); res.is_wan = (loadedsdver == SDVersion::VERSION_WAN2 || loadedsdver == SDVersion::VERSION_WAN2_2_I2V || loadedsdver == SDVersion::VERSION_WAN2_2_TI2V); - res.is_qwenimg = (loadedsdver == SDVersion::VERSION_QWEN_IMAGE); + res.is_qwenimg = sd_version_is_qwen_image(loadedsdver); res.is_chroma = loaded_model_is_chroma(ctx); res.is_kontext = (loadedsdver==SDVersion::VERSION_FLUX && !res.is_chroma); res.is_flux2 = (loadedsdver == SDVersion::VERSION_FLUX2 || loadedsdver == SDVersion::VERSION_FLUX2_KLEIN); @@ -840,7 +831,7 @@ namespace kcpp_sd { res.is_ltx = sd_version_is_ltxav(loadedsdver); res.is_minimaxh3 = sd_version_is_minimax_h3(loadedsdver); res.is_boogu = sd_version_is_boogu_image(loadedsdver); - res.supports_ref_image = sd_version_supports_ref_latent_img_cfg(loadedsdver) || sd_version_is_pid(loadedsdver); + res.supports_ref_image = sd_version_supports_ref_latent_img_cfg(loadedsdver) || sd_version_is_pid(loadedsdver) || res.is_qwenimg; res.vae_scale_factor = ctx->sd->get_vae_scale_factor(); res.spatial_multiple = get_spatial_multiple(ctx); return res; diff --git a/otherarch/sdcpp/src/upscaler.cpp b/otherarch/sdcpp/src/upscaler.cpp index a69ff9e16..64dabaaf9 100644 --- a/otherarch/sdcpp/src/upscaler.cpp +++ b/otherarch/sdcpp/src/upscaler.cpp @@ -111,10 +111,22 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_tensor) { sd::ParallelScope tensor_scope(&tensor_executor); + if (input_tensor.empty() || input_tensor.dim() != 4 || + (input_tensor.shape()[2] != 3 && input_tensor.shape()[2] != 4)) { + LOG_ERROR("esrgan expects a 4D RGB or RGBA image tensor"); + return {}; + } + + const bool has_alpha = input_tensor.shape()[2] == 4; + sd::Tensor rgb; + if (has_alpha) { + rgb = sd::ops::slice(input_tensor, 2, 0, 3); + } + const sd::Tensor& model_input = has_alpha ? rgb : input_tensor; sd::Tensor upscaled; const int scale = esrgan_upscaler->config.scale; if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) { - upscaled = esrgan_upscaler->compute(n_threads, input_tensor); + upscaled = esrgan_upscaler->compute(n_threads, model_input); } else { auto on_processing = [&](const sd::Tensor& input_tile) -> sd::Tensor { auto output_tile = esrgan_upscaler->compute(n_threads, input_tile); @@ -125,7 +137,7 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te return output_tile; }; - upscaled = process_tiles_2d(input_tensor, + upscaled = process_tiles_2d(model_input, static_cast(input_tensor.shape()[0] * scale), static_cast(input_tensor.shape()[1] * scale), scale, @@ -141,6 +153,14 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te LOG_ERROR("esrgan compute failed"); return {}; } + if (has_alpha) { + auto alpha = sd::ops::slice(input_tensor, 2, 3, 4); + auto alpha_shape = alpha.shape(); + alpha_shape[0] = upscaled.shape()[0]; + alpha_shape[1] = upscaled.shape()[1]; + alpha = sd::ops::interpolate(alpha, alpha_shape, sd::ops::InterpolateMode::Bilinear); + upscaled = sd::ops::concat(upscaled, alpha, 2); + } return upscaled; } @@ -242,6 +262,30 @@ int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) { return upscaler_ctx->upscaler->esrgan_upscaler->config.scale; } +int get_upscaler_model_scale(const char* model_path) { + if (model_path == nullptr || model_path[0] == '\0') { + return 0; + } + try { + ModelLoader loader; + if (!loader.init_from_file_and_convert_name(model_path, "", VERSION_ESRGAN)) { + return 0; + } + const auto& tensors = loader.get_tensor_storage_map(); + auto first = tensors.find("conv_first.weight"); + auto last = tensors.find("conv_last.weight"); + if (first == tensors.end() || last == tensors.end() || + tensors.count("body.0.rdb1.conv1.weight") == 0 || + first->second.n_dims != 4 || last->second.n_dims != 4 || + first->second.ne[2] != 3 || last->second.ne[3] != 3) { + return 0; + } + return ESRGANConfig::detect_from_weights(tensors).scale; + } catch (const std::exception&) { + return 0; + } +} + void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx) { if (upscaler_ctx->upscaler != nullptr) { delete upscaler_ctx->upscaler;