From 67c3fef2da2f570d4be557e455717bf22701fc77 Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Thu, 16 Jul 2026 23:41:14 -0300 Subject: [PATCH 1/5] sd: sync with master-782-b290693 --- Makefile | 4 +- otherarch/sdcpp/examples/cli/main.cpp | 89 +- otherarch/sdcpp/examples/common/common.cpp | 153 ++- otherarch/sdcpp/examples/common/common.h | 10 +- .../sdcpp/examples/common/resource_owners.hpp | 15 +- otherarch/sdcpp/include/stable-diffusion.h | 26 +- otherarch/sdcpp/sdtype_adapter.cpp | 5 +- .../sdcpp/src/conditioning/conditioner.hpp | 223 +++- .../sdcpp/src/core/ggml_extend_backend.cpp | 6 + .../sdcpp/src/core/ggml_extend_backend.h | 1 + otherarch/sdcpp/src/detailer.cpp | 1020 +++++++++++++++++ otherarch/sdcpp/src/detailer.h | 75 ++ otherarch/sdcpp/src/model/common/rope.hpp | 4 - otherarch/sdcpp/src/model/detector/yolov8.h | 362 ++++++ otherarch/sdcpp/src/model/diffusion/anima.hpp | 61 +- .../sdcpp/src/model/diffusion/animatediff.hpp | 182 +++ otherarch/sdcpp/src/model/diffusion/boogu.hpp | 2 +- otherarch/sdcpp/src/model/diffusion/flux.hpp | 4 +- otherarch/sdcpp/src/model/diffusion/krea2.hpp | 19 +- otherarch/sdcpp/src/model/diffusion/model.hpp | 32 +- otherarch/sdcpp/src/model/diffusion/pid.hpp | 207 +++- .../sdcpp/src/model/diffusion/qwen_image.hpp | 4 +- otherarch/sdcpp/src/model/diffusion/unet.hpp | 56 + .../sdcpp/src/model/diffusion/z_image.hpp | 4 +- .../sdcpp/src/model_io/safetensors_io.cpp | 15 +- otherarch/sdcpp/src/model_io/safetensors_io.h | 4 +- otherarch/sdcpp/src/model_loader.cpp | 2 +- otherarch/sdcpp/src/model_loader.h | 2 + otherarch/sdcpp/src/name_conversion.cpp | 24 +- otherarch/sdcpp/src/stable-diffusion.cpp | 314 ++++- 30 files changed, 2690 insertions(+), 235 deletions(-) create mode 100644 otherarch/sdcpp/src/detailer.cpp create mode 100644 otherarch/sdcpp/src/detailer.h create mode 100644 otherarch/sdcpp/src/model/detector/yolov8.h create mode 100644 otherarch/sdcpp/src/model/diffusion/animatediff.hpp diff --git a/Makefile b/Makefile index 3402953c2b4b..bba30467c57d 100644 --- a/Makefile +++ b/Makefile @@ -703,7 +703,7 @@ budget.o: common/reasoning-budget.cpp common/reasoning-budget.h chat.o: common/chat.cpp common/chat.h $(CXX) $(CXXFLAGS) -c $< -o $@ -SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/core/backend_fit.cpp src/core/backend_fit.h src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.hpp src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/layer_split_partition.cpp src/core/layer_split_partition.h src/core/ordered_map.hpp src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/kcpp_sd_extensions.h src/model/adapter/lora.hpp src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/common/block.hpp src/model/common/rope.hpp src/model/diffusion/anima.hpp src/model/diffusion/boogu.hpp src/model/diffusion/control.hpp src/model/diffusion/dit.hpp src/model/diffusion/ernie_image.hpp src/model/diffusion/flux.hpp src/model/diffusion/hidream_o1.hpp src/model/diffusion/ideogram4.hpp src/model/diffusion/krea2.hpp src/model/diffusion/lens.hpp src/model/diffusion/lingbot_video.hpp src/model/diffusion/ltxv.hpp src/model/diffusion/minit2i.hpp src/model/diffusion/mmdit.hpp src/model/diffusion/model.hpp src/model/diffusion/pid.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/sefi_image.hpp src/model/diffusion/unet.hpp src/model/diffusion/wan.hpp src/model/diffusion/z_image.hpp src/model.h src/model_io/binary_io.h src/model_io/gguf_io.cpp src/model_io/gguf_io.h src/model_io/gguf_reader_ext.h src/model_io/pickle_io.cpp src/model_io/pickle_io.h src/model_io/safetensors_io.cpp src/model_io/safetensors_io.h src/model_io/streaming_writer.h src/model_io/tensor_storage.h src/model_io/torch_legacy_io.cpp src/model_io/torch_legacy_io.h src/model_io/torch_zip_io.cpp src/model_io/torch_zip_io.h src/model_loader.cpp src/model_loader.h src/model_manager.cpp src/model_manager.h src/model/te/clip.hpp src/model/te/llm.hpp src/model/te/t5.hpp src/model/upscaler/esrgan.hpp src/model/upscaler/ltx_latent_upscaler.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/runtime/cache_dit.hpp src/runtime/condition_cache_utils.hpp src/runtime/denoiser.hpp src/runtime/easycache.hpp src/runtime/gits_noise.h src/runtime/guidance.cpp src/runtime/guidance.h src/runtime/imatrix.cpp src/runtime/imatrix.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/ucache.hpp src/stable-diffusion.cpp src/tokenizers/bpe_tokenizer.cpp src/tokenizers/bpe_tokenizer.h src/tokenizers/clip_tokenizer.cpp src/tokenizers/clip_tokenizer.h src/tokenizers/gemma_tokenizer.cpp src/tokenizers/gemma_tokenizer.h src/tokenizers/gpt_oss_tokenizer.cpp src/tokenizers/gpt_oss_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer.cpp src/tokenizers/tokenizer.h src/tokenizers/tokenize_util.cpp src/tokenizers/tokenize_util.h src/tokenizers/vocab/vocab.h src/upscaler.cpp src/upscaler.h src/weight_manager.h +SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/core/backend_fit.cpp src/core/backend_fit.h src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.hpp src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/layer_split_partition.cpp src/core/layer_split_partition.h src/core/ordered_map.hpp src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/detailer.cpp src/detailer.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/kcpp_sd_extensions.h src/model/adapter/lora.hpp src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/common/block.hpp src/model/common/rope.hpp src/model/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/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/minit2i.hpp src/model/diffusion/mmdit.hpp src/model/diffusion/model.hpp src/model/diffusion/pid.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/sefi_image.hpp src/model/diffusion/unet.hpp src/model/diffusion/wan.hpp src/model/diffusion/z_image.hpp src/model.h src/model_io/binary_io.h src/model_io/gguf_io.cpp src/model_io/gguf_io.h src/model_io/gguf_reader_ext.h src/model_io/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.h src/model_manager.cpp src/model_manager.h src/model/te/clip.hpp src/model/te/llm.hpp src/model/te/t5.hpp src/model/upscaler/esrgan.hpp src/model/upscaler/ltx_latent_upscaler.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/runtime/cache_dit.hpp src/runtime/condition_cache_utils.hpp src/runtime/denoiser.hpp src/runtime/easycache.hpp src/runtime/gits_noise.h src/runtime/guidance.cpp src/runtime/guidance.h src/runtime/imatrix.cpp src/runtime/imatrix.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/ucache.hpp src/stable-diffusion.cpp src/tokenizers/bpe_tokenizer.cpp src/tokenizers/bpe_tokenizer.h src/tokenizers/clip_tokenizer.cpp src/tokenizers/clip_tokenizer.h src/tokenizers/gemma_tokenizer.cpp src/tokenizers/gemma_tokenizer.h src/tokenizers/gpt_oss_tokenizer.cpp src/tokenizers/gpt_oss_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer.cpp src/tokenizers/tokenizer.h src/tokenizers/tokenize_util.cpp src/tokenizers/tokenize_util.h src/tokenizers/vocab/vocab.h src/upscaler.cpp src/upscaler.h src/weight_manager.h SDCPP_MAIN_BASENAMES := examples/cli/image_metadata.cpp examples/cli/image_metadata.h examples/cli/main.cpp examples/cli/msf_gif.h examples/common/common.cpp examples/common/common.h examples/common/log.cpp examples/common/log.h examples/common/media_io.cpp examples/common/media_io.h examples/common/resource_owners.hpp src/tokenizers/vocab/clip_merges.hpp src/tokenizers/vocab/gemma2_merges.hpp src/tokenizers/vocab/gemma2_vocab.hpp src/tokenizers/vocab/gemma_merges.hpp src/tokenizers/vocab/gemma_vocab.hpp src/tokenizers/vocab/gpt_oss_merges.hpp src/tokenizers/vocab/gpt_oss_vocab.hpp src/tokenizers/vocab/mistral_merges.hpp src/tokenizers/vocab/mistral_vocab.hpp src/tokenizers/vocab/qwen_merges.hpp src/tokenizers/vocab/t5.hpp src/tokenizers/vocab/umt5.hpp src/tokenizers/vocab/vocab.cpp src/convert.cpp src/version.cpp @@ -726,7 +726,7 @@ otherarch/sdcpp/src/%.o: otherarch/sdcpp/src/%.cpp otherarch/sdcpp/examples/%.o: otherarch/sdcpp/examples/%.cpp $(CXX) -I./otherarch/sdcpp/include -I./otherarch/sdcpp/examples -I./vendor/nlohmann $(CXXFLAGS) -c $< -o $@ -otherarch/sdcpp/sdtype_adapter.o: otherarch/sdcpp/sdtype_adapter.cpp +otherarch/sdcpp/sdtype_adapter.o: otherarch/sdcpp/sdtype_adapter.cpp otherarch/sdcpp/include/stable-diffusion.h otherarch/sdcpp/src/kcpp_sd_extensions.h model_adapter.h otherarch/utils.h $(CXX) -I./otherarch/sdcpp/include -I./otherarch/sdcpp/src $(CXXFLAGS) -c $< -o $@ otherarch/sdcpp/thirdparty/zip.o: otherarch/sdcpp/thirdparty/zip.c diff --git a/otherarch/sdcpp/examples/cli/main.cpp b/otherarch/sdcpp/examples/cli/main.cpp index 29cb391b66b5..68a3148cc605 100644 --- a/otherarch/sdcpp/examples/cli/main.cpp +++ b/otherarch/sdcpp/examples/cli/main.cpp @@ -199,7 +199,7 @@ struct SDCliParams { options.manual_options = { {"-M", "--mode", - "run mode, one of [img_gen, vid_gen, upscale, convert, metadata], default: img_gen", + "run mode, one of [img_gen, adetailer, vid_gen, upscale, convert, metadata], default: img_gen", on_mode_arg}, {"", "--preview", @@ -566,6 +566,65 @@ bool save_results(const SDCliParams& cli_params, return sucessful_reults != 0; } +static bool apply_adetailer(sd_ctx_t* sd_ctx, + const sd_ctx_params_t& sd_ctx_params, + const SDContextParams& ctx_params, + const SDGenerationParams& gen_params, + const sd_img_gen_params_t& img_gen_params, + SDMode mode, + SDImageVec& results, + int num_results) { + if (gen_params.ad_model_path.empty()) { + return true; + } + + sd_adetailer_params_t ad_params{}; + ad_params.prompt = gen_params.ad_prompt.empty() ? nullptr : gen_params.ad_prompt.c_str(); + ad_params.negative_prompt = gen_params.ad_negative_prompt.empty() ? nullptr : gen_params.ad_negative_prompt.c_str(); + ad_params.extra_ad_args = gen_params.extra_ad_args.c_str(); + + ADetailerCtxPtr ad_ctx(new_adetailer_ctx(gen_params.ad_model_path.c_str(), + ctx_params.n_threads, + sd_ctx_params.backend, + sd_ctx_params.params_backend)); + if (ad_ctx == nullptr) { + LOG_ERROR("new_adetailer_ctx failed"); + return false; + } + + for (int i = 0; i < num_results; ++i) { + if (results[i].data == nullptr) { + continue; + } + sd_img_gen_params_t ad_generation_params = img_gen_params; + ad_generation_params.seed = img_gen_params.seed + i; + if (mode == IMG_GEN) { + ad_generation_params.width = 512; + ad_generation_params.height = 512; + ad_generation_params.strength = 0.4f; + } + sd_image_t* detailed_images = nullptr; + int detailed_count = 0; + if (!adetail_image(ad_ctx.get(), + sd_ctx, + results[i], + &ad_params, + &ad_generation_params, + &detailed_images, + &detailed_count) || + detailed_count <= 0 || detailed_images == nullptr || detailed_images[0].data == nullptr) { + free_sd_images(detailed_images, detailed_count); + LOG_ERROR("ADetailer failed for image %d", i + 1); + return false; + } + free(results[i].data); + results[i] = detailed_images[0]; + detailed_images[0] = {0, 0, 0, nullptr}; + free_sd_images(detailed_images, detailed_count); + } + return true; +} + int main(int argc, const char* argv[]) { if (argc > 1 && std::string(argv[1]) == "--version") { std::cout << version_string() << "\n"; @@ -598,6 +657,11 @@ int main(int argc, const char* argv[]) { return 0; } + if (!gen_params.ad_model_path.empty() && cli_params.mode != IMG_GEN && cli_params.mode != ADETAILER) { + LOG_ERROR("--ad-model is only supported in image generation and adetailer modes"); + return 1; + } + if (gen_params.video_frames > 4) { size_t last_dot_pos = cli_params.preview_path.find_last_of("."); std::string base_path = cli_params.preview_path; @@ -806,15 +870,22 @@ int main(int argc, const char* argv[]) { gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx.get(), gen_params.sample_params.sample_method); } - if (cli_params.mode == IMG_GEN) { - sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t(); + sd_img_gen_params_t img_gen_params{}; + const bool use_img_gen_params = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER; + if (use_img_gen_params) { + img_gen_params = gen_params.to_sd_img_gen_params_t(); + } + if (cli_params.mode == IMG_GEN) { sd_image_t* generated_images = nullptr; if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) { generated_images = nullptr; num_results = 0; } results.adopt(generated_images, num_results); + } else if (cli_params.mode == ADETAILER) { + num_results = 1; + results.push_back(gen_params.init_image.release()); } else if (cli_params.mode == VID_GEN) { sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t(); sd_image_t* generated_video = nullptr; @@ -828,6 +899,18 @@ int main(int argc, const char* argv[]) { LOG_ERROR("generate failed"); return 1; } + + if (use_img_gen_params && + !apply_adetailer(sd_ctx.get(), + sd_ctx_params, + ctx_params, + gen_params, + img_gen_params, + cli_params.mode, + results, + num_results)) { + return 1; + } } int upscale_factor = 4; // unused for RealESRGAN_x4plus_anime_6B.pth diff --git a/otherarch/sdcpp/examples/common/common.cpp b/otherarch/sdcpp/examples/common/common.cpp index 1dfb6aa70cc7..a5cb48369a28 100644 --- a/otherarch/sdcpp/examples/common/common.cpp +++ b/otherarch/sdcpp/examples/common/common.cpp @@ -30,6 +30,7 @@ namespace fs = std::filesystem; const char* const modes_str[] = { "img_gen", + "adetailer", "vid_gen", "convert", "upscale", @@ -49,6 +50,9 @@ static sd_vae_format_t str_to_vae_format(const std::string& value) { if (value == "flux2") { return SD_VAE_FORMAT_FLUX2; } + if (value == "wan") { + return SD_VAE_FORMAT_WAN; + } return SD_VAE_FORMAT_COUNT; } @@ -400,7 +404,7 @@ ArgOptions SDContextParams::get_options() { &vae_path}, {"", "--vae-format", - "VAE latent format override: auto, flux, sd3, or flux2 (default: auto)", + "VAE latent format override: auto, flux, sd3, flux2, or wan (default: auto)", 0, &vae_format}, {"", @@ -423,6 +427,11 @@ ArgOptions SDContextParams::get_options() { "path to control net model", 0, &control_net_path}, + {"", + "--motion-module", + "path to AnimateDiff motion module (SD 1.5); enables video generation on --video-frames > 1", + 0, + &motion_module_path}, {"", "--embd-dir", "embeddings directory", @@ -672,7 +681,7 @@ ArgOptions SDContextParams::get_options() { } void SDContextParams::build_embedding_map() { - static const std::vector valid_ext = {".gguf", ".safetensors", ".pt"}; + static const std::vector valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"}; if (!fs::exists(embedding_dir) || !fs::is_directory(embedding_dir)) { return; @@ -737,7 +746,7 @@ bool SDContextParams::validate(SDMode mode) { } if (str_to_vae_format(vae_format) == SD_VAE_FORMAT_COUNT) { - LOG_ERROR("error: vae_format must be 'auto', 'flux', 'sd3', or 'flux2'"); + LOG_ERROR("error: vae_format must be 'auto', 'flux', 'sd3', 'flux2', or 'wan'"); return false; } @@ -867,6 +876,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.audio_vae_path = audio_vae_path.c_str(); sd_ctx_params.taesd_path = taesd_path.c_str(); sd_ctx_params.control_net_path = control_net_path.c_str(); + sd_ctx_params.motion_module_path = motion_module_path.c_str(); sd_ctx_params.embeddings = embedding_vec.data(); sd_ctx_params.embedding_count = static_cast(embedding_vec.size()); sd_ctx_params.photo_maker_path = photo_maker_path.c_str(); @@ -916,6 +926,26 @@ ArgOptions SDGenerationParams::get_options() { "the negative prompt (default: \"\")", 0, &negative_prompt}, + {"", + "--ad-model", + "path to a converted YOLOv8 detection model for ADetailer", + 0, + &ad_model_path}, + {"", + "--ad-prompt", + "ADetailer prompt; empty inherits the main prompt, supports [PROMPT], [SEP], and [SKIP]", + 0, + &ad_prompt}, + {"", + "--ad-negative-prompt", + "ADetailer negative prompt; empty inherits the main negative prompt, supports [PROMPT] and [SEP]", + 0, + &ad_negative_prompt}, + {"", + "--extra-ad-args", + "extra ADetailer args, key=value list. Supports input_size, confidence, nms, max_detections, mask_k_largest, mask_min_ratio, mask_max_ratio, dilate_erode, x_offset, y_offset, mask_mode, merge_masks, invert_mask, mask_blur, inpaint_padding, inpaint_width, inpaint_height, denoising_strength, steps, cfg_scale, sample_method, scheduler, sort_by", + (int)',', + &extra_ad_args}, {"-i", "--init-img", "path to the init image", @@ -975,6 +1005,11 @@ ArgOptions SDGenerationParams::get_options() { "extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)", (int)',', &extra_tiling_args}, + {"", + "--ref-image-args", + "Key-value list to set up the way the reference images are processed (empty = auto-detect from model weigths)", + (int)',', + &ref_image_args}, }; options.int_options = { @@ -1831,6 +1866,10 @@ bool SDGenerationParams::from_json_str( load_if_exists("prompt", prompt); load_if_exists("negative_prompt", negative_prompt); + load_if_exists("ad_model", ad_model_path); + load_if_exists("ad_prompt", ad_prompt); + load_if_exists("ad_negative_prompt", ad_negative_prompt); + load_if_exists("extra_ad_args", extra_ad_args); load_if_exists("cache_mode", cache_mode); load_if_exists("cache_option", cache_option); load_if_exists("scm_mask", scm_mask); @@ -2027,7 +2066,7 @@ void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_d return; } static const std::regex re(R"(]+):([^>]+)>)"); - static const std::vector valid_ext = {".gguf", ".safetensors", ".pt"}; + static const std::vector valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"}; std::smatch m; std::string tmp = prompt; @@ -2347,13 +2386,19 @@ bool SDGenerationParams::validate(SDMode mode) { } } - if (mode == UPSCALE) { + if (mode == UPSCALE || mode == ADETAILER) { if (init_image_path.length() == 0) { - LOG_ERROR("error: upscale mode needs an init image (--init-img)\n"); + LOG_ERROR("error: %s mode needs an init image (--init-img)\n", + mode == UPSCALE ? "upscale" : "adetailer"); return false; } } + if (mode == ADETAILER && ad_model_path.empty()) { + LOG_ERROR("error: adetailer mode needs a detector model (--ad-model)\n"); + return false; + } + return true; } @@ -2418,30 +2463,45 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() { pulid_id_weight, }; - params.loras = lora_vec.empty() ? nullptr : lora_vec.data(); - params.lora_count = static_cast(lora_vec.size()); - params.prompt = prompt.c_str(); - params.negative_prompt = negative_prompt.c_str(); - params.clip_skip = clip_skip; - params.init_image = init_image.get(); - params.ref_images = ref_image_views.empty() ? nullptr : ref_image_views.data(); - params.ref_images_count = static_cast(ref_image_views.size()); - params.auto_resize_ref_image = auto_resize_ref_image; - params.increase_ref_index = increase_ref_index; - params.mask_image = mask_image.get(); - params.width = get_resolved_width(); - params.height = get_resolved_height(); - params.sample_params = sample_params; - params.strength = strength; - params.seed = seed; - params.batch_count = batch_count; - params.qwen_image_layers = qwen_image_layers; - params.control_image = control_image.get(); - params.control_strength = control_strength; - params.pm_params = pm_params; - params.pulid_params = pulid_params; - params.vae_tiling_params = vae_tiling_params; - params.cache = cache_params; + 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 += ","; + } + ref_image_args += "ref_index_mode=increase"; + LOG_WARN("Notice: --increase-ref-index is deprecated. Use --ref-image-args \"ref_index_mode=increase\" instead."); + } + + params.loras = lora_vec.empty() ? nullptr : lora_vec.data(); + params.lora_count = static_cast(lora_vec.size()); + params.prompt = prompt.c_str(); + params.negative_prompt = negative_prompt.c_str(); + params.clip_skip = clip_skip; + params.init_image = init_image.get(); + params.ref_images = ref_image_views.empty() ? nullptr : ref_image_views.data(); + params.ref_images_count = static_cast(ref_image_views.size()); + params.ref_image_args = ref_image_args.c_str(); + params.mask_image = mask_image.get(); + params.width = get_resolved_width(); + params.height = get_resolved_height(); + params.sample_params = sample_params; + params.strength = strength; + params.seed = seed; + params.batch_count = batch_count; + params.qwen_image_layers = qwen_image_layers; + params.control_image = control_image.get(); + params.control_strength = control_strength; + params.pm_params = pm_params; + params.pulid_params = pulid_params; + params.vae_tiling_params = vae_tiling_params; + params.cache = cache_params; params.hires.enabled = hires_enabled; params.hires.upscaler = resolved_hires_upscaler; @@ -2561,6 +2621,10 @@ std::string SDGenerationParams::to_string() const { << " high_noise_loras: \"" << high_noise_loras_str << "\",\n" << " prompt: \"" << prompt << "\",\n" << " negative_prompt: \"" << negative_prompt << "\",\n" + << " ad_model_path: \"" << ad_model_path << "\",\n" + << " ad_prompt: \"" << ad_prompt << "\",\n" + << " ad_negative_prompt: \"" << ad_negative_prompt << "\",\n" + << " extra_ad_args: \"" << extra_ad_args << "\",\n" << " clip_skip: " << clip_skip << ",\n" << " width: " << width << ",\n" << " height: " << height << ",\n" @@ -2675,8 +2739,13 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params, int64_t seed, SDMode mode) { json root; - root["schema"] = "sdcpp.image.params/v1"; - root["mode"] = mode == VID_GEN ? "vid_gen" : "img_gen"; + root["schema"] = "sdcpp.image.params/v1"; + root["mode"] = "img_gen"; + if (mode == VID_GEN) { + root["mode"] = "vid_gen"; + } else if (mode == ADETAILER) { + root["mode"] = "adetailer"; + } root["generator"] = { {"name", "stable-diffusion.cpp"}, {"version", safe_json_string(sd_version())}, @@ -2690,6 +2759,14 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params, {"positive", gen_params.prompt}, {"negative", gen_params.negative_prompt}, }; + if (!gen_params.ad_model_path.empty()) { + root["adetailer"] = { + {"model", sd_basename(gen_params.ad_model_path)}, + {"prompt", gen_params.ad_prompt}, + {"negative_prompt", gen_params.ad_negative_prompt}, + {"extra_args", gen_params.extra_ad_args}, + }; + } root["sampling"] = build_sampling_metadata_json(gen_params.sample_params, gen_params.skip_layers, &gen_params.custom_sigmas); @@ -2844,6 +2921,18 @@ std::string get_image_params(const SDContextParams& ctx_params, if (!gen_params.extra_sample_args.empty()) { parameter_string += "Extra sample args: " + gen_params.extra_sample_args + ", "; } + if (!gen_params.ad_model_path.empty()) { + parameter_string += "ADetailer model: " + sd_basename(gen_params.ad_model_path) + ", "; + if (!gen_params.ad_prompt.empty()) { + parameter_string += "ADetailer prompt: " + gen_params.ad_prompt + ", "; + } + if (!gen_params.ad_negative_prompt.empty()) { + parameter_string += "ADetailer negative prompt: " + gen_params.ad_negative_prompt + ", "; + } + if (!gen_params.extra_ad_args.empty()) { + parameter_string += "ADetailer args: " + gen_params.extra_ad_args + ", "; + } + } parameter_string += "Seed: " + std::to_string(seed) + ", "; parameter_string += "Size: " + std::to_string(gen_params.get_resolved_width()) + "x" + std::to_string(gen_params.get_resolved_height()) + ", "; parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", "; diff --git a/otherarch/sdcpp/examples/common/common.h b/otherarch/sdcpp/examples/common/common.h index 3a5b107bcb85..769651987e97 100644 --- a/otherarch/sdcpp/examples/common/common.h +++ b/otherarch/sdcpp/examples/common/common.h @@ -16,10 +16,11 @@ #define BOOL_STR(b) ((b) ? "true" : "false") extern const char* const modes_str[]; -#define SD_ALL_MODES_STR "img_gen, vid_gen, convert, upscale, metadata" +#define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata" enum SDMode { IMG_GEN, + ADETAILER, VID_GEN, CONVERT, UPSCALE, @@ -132,6 +133,7 @@ struct SDContextParams { std::string taesd_path; std::string esrgan_path; std::string control_net_path; + std::string motion_module_path; std::string embedding_dir; std::string photo_maker_path; std::string pulid_weights_path; @@ -186,6 +188,10 @@ struct SDGenerationParams { // User-facing input fields. std::string prompt; std::string negative_prompt; + std::string ad_model_path; + 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; @@ -227,6 +233,8 @@ struct SDGenerationParams { sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr}; std::string extra_tiling_args; + std::string ref_image_args; + std::string pm_id_images_dir; std::string pm_id_embed_path; float pm_style_strength = 20.f; diff --git a/otherarch/sdcpp/examples/common/resource_owners.hpp b/otherarch/sdcpp/examples/common/resource_owners.hpp index d47134abe505..d7525a5faaf8 100644 --- a/otherarch/sdcpp/examples/common/resource_owners.hpp +++ b/otherarch/sdcpp/examples/common/resource_owners.hpp @@ -40,12 +40,21 @@ struct UpscalerCtxDeleter { } }; +struct ADetailerCtxDeleter { + void operator()(adetailer_ctx_t* ctx) const { + if (ctx != nullptr) { + free_adetailer_ctx(ctx); + } + } +}; + template using FreeUniquePtr = std::unique_ptr; -using FilePtr = std::unique_ptr; -using SDCtxPtr = std::unique_ptr; -using UpscalerCtxPtr = std::unique_ptr; +using FilePtr = std::unique_ptr; +using SDCtxPtr = std::unique_ptr; +using UpscalerCtxPtr = std::unique_ptr; +using ADetailerCtxPtr = std::unique_ptr; class SDImageOwner { private: diff --git a/otherarch/sdcpp/include/stable-diffusion.h b/otherarch/sdcpp/include/stable-diffusion.h index 95d507c77da0..536bce73f063 100644 --- a/otherarch/sdcpp/include/stable-diffusion.h +++ b/otherarch/sdcpp/include/stable-diffusion.h @@ -180,6 +180,7 @@ enum sd_vae_format_t { SD_VAE_FORMAT_FLUX, SD_VAE_FORMAT_SD3, SD_VAE_FORMAT_FLUX2, + SD_VAE_FORMAT_WAN, SD_VAE_FORMAT_COUNT, }; @@ -199,6 +200,7 @@ typedef struct { const char* audio_vae_path; const char* taesd_path; const char* control_net_path; + const char* motion_module_path; const sd_embedding_t* embeddings; uint32_t embedding_count; const char* photo_maker_path; @@ -363,8 +365,7 @@ typedef struct { sd_image_t init_image; sd_image_t* ref_images; int ref_images_count; - bool auto_resize_ref_image; - bool increase_ref_index; + const char* ref_image_args; sd_image_t mask_image; int width; int height; @@ -510,6 +511,27 @@ SD_API bool upscale(upscaler_ctx_t* upscaler_ctx, SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx); +typedef struct adetailer_ctx_t adetailer_ctx_t; + +typedef struct { + const char* prompt; + const char* negative_prompt; + const char* extra_ad_args; +} sd_adetailer_params_t; + +SD_API adetailer_ctx_t* new_adetailer_ctx(const char* detector_path, + int n_threads, + const char* backend, + const char* params_backend); +SD_API void free_adetailer_ctx(adetailer_ctx_t* adetailer_ctx); +SD_API bool adetail_image(adetailer_ctx_t* adetailer_ctx, + sd_ctx_t* sd_ctx, + sd_image_t input_image, + const sd_adetailer_params_t* adetailer_params, + const sd_img_gen_params_t* inpaint_params, + sd_image_t** images_out, + int* num_images_out); + SD_API bool convert(const char* input_path, const char* vae_path, const char* output_path, diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp index 6ca87a14c8d2..89c277e8d1d4 100644 --- a/otherarch/sdcpp/sdtype_adapter.cpp +++ b/otherarch/sdcpp/sdtype_adapter.cpp @@ -1238,10 +1238,13 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) } } + std::string ref_image_args; + ref_image_args += "resize_before_vae=on"; // auto_resize_ref_image = true + sd_img_gen_params_t params = {}; sd_img_gen_params_init (¶ms); params.batch_count = 1; - params.auto_resize_ref_image = true; + params.ref_image_args = ref_image_args.c_str(); params.prompt = sd_params->prompt.c_str(); params.negative_prompt = sd_params->negative_prompt.c_str(); params.clip_skip = sd_params->clip_skip; diff --git a/otherarch/sdcpp/src/conditioning/conditioner.hpp b/otherarch/sdcpp/src/conditioning/conditioner.hpp index 61b5791c9637..4dc4e0cd45b5 100644 --- a/otherarch/sdcpp/src/conditioning/conditioner.hpp +++ b/otherarch/sdcpp/src/conditioning/conditioner.hpp @@ -7,6 +7,7 @@ #include "core/tensor_ggml.hpp" #include "core/util.h" +#include "model/diffusion/model.hpp" #include "model/te/clip.hpp" #include "model/te/llm.hpp" #include "model/te/t5.hpp" @@ -106,6 +107,7 @@ struct ConditionerParams { int height = -1; bool zero_out_masked = false; const std::vector>* ref_images = nullptr; // for qwen image edit + RefImageParams ref_image_params; }; struct Conditioner { @@ -1993,6 +1995,54 @@ struct LLMEmbedder : public Conditioner { return new_hidden_states; } + void resize_image_dims(int height, int width, int& h_bar, int& w_bar, int factor, int min_size, int max_size, RefImageResizeMode mode) { + if (min_size > 0 && min_size == max_size) { + if (mode == RefImageResizeMode::AREA) { + double beta = std::sqrt(static_cast(min_size) / (static_cast(height) * width)); + h_bar = std::max(static_cast(factor), + static_cast(std::round(height * beta / factor)) * static_cast(factor)); + w_bar = std::max(static_cast(factor), + static_cast(std::round(width * beta / factor)) * static_cast(factor)); + } else if (mode == RefImageResizeMode::LONGEST_SIDE) { + int current_max_side = std::max(height, width); + double beta = static_cast(min_size) / current_max_side; + h_bar = std::max(static_cast(factor), + static_cast(std::round(height * beta / factor)) * static_cast(factor)); + w_bar = std::max(static_cast(factor), + static_cast(std::round(width * beta / factor)) * static_cast(factor)); + } + return; + } + + if (mode == RefImageResizeMode::AREA) { + double current_area = static_cast(h_bar) * w_bar; + if (max_size > 0 && current_area > max_size) { + double beta = std::sqrt((static_cast(height) * width) / static_cast(max_size)); + h_bar = std::max(static_cast(factor), + static_cast(std::floor(height / beta / factor)) * static_cast(factor)); + w_bar = std::max(static_cast(factor), + static_cast(std::floor(width / beta / factor)) * static_cast(factor)); + } else if (min_size > 0 && current_area < min_size) { + double beta = std::sqrt(static_cast(min_size) / (static_cast(height) * width)); + h_bar = static_cast(std::ceil(height * beta / factor)) * static_cast(factor); + w_bar = static_cast(std::ceil(width * beta / factor)) * static_cast(factor); + } + } else if (mode == RefImageResizeMode::LONGEST_SIDE) { + int current_max_side = std::max(height, width); + if (max_size > 0 && current_max_side > max_size) { + double beta = static_cast(max_size) / current_max_side; + h_bar = std::max(static_cast(factor), + static_cast(std::floor(height * beta / factor)) * static_cast(factor)); + w_bar = std::max(static_cast(factor), + static_cast(std::floor(width * beta / factor)) * static_cast(factor)); + } else if (min_size > 0 && current_max_side < min_size) { + double beta = static_cast(min_size) / current_max_side; + h_bar = static_cast(std::ceil(height * beta / factor)) * static_cast(factor); + w_bar = static_cast(std::ceil(width * beta / factor)) * static_cast(factor); + } + } + } + SDCondition get_learned_condition(int n_threads, const ConditionerParams& conditioner_params) override { std::string prompt; @@ -2007,7 +2057,8 @@ struct LLMEmbedder : public Conditioner { bool spell_quotes = false; std::set out_layers; - int64_t t0 = ggml_time_ms(); + int64_t t0 = ggml_time_ms(); + RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode; if (sd_version_is_lingbot_video(version)) { const int pad_token = 151643; @@ -2040,28 +2091,35 @@ struct LLMEmbedder : public Conditioner { for (int i = 0; i < conditioner_params.ref_images->size(); i++) { const auto& image = (*conditioner_params.ref_images)[i]; - double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; + const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; int height = static_cast(image.shape()[1]); int width = static_cast(image.shape()[0]); - int min_pixels = static_cast(4 * factor * factor); - int max_pixels = static_cast(16384 * factor * factor); - int h_bar = std::max(static_cast(factor), static_cast(std::round(height / factor) * factor)); - int w_bar = std::max(static_cast(factor), static_cast(std::round(width / factor) * factor)); + + int min_pixels = conditioner_params.ref_image_params.vlm_min_size; + if (min_pixels <= 0) { + if (resize_mode == RefImageResizeMode::AREA) { + min_pixels = static_cast(4 * factor * factor); + } else { + min_pixels = static_cast(2 * factor); + } + } + int max_pixels = conditioner_params.ref_image_params.vlm_max_size; + if (max_pixels <= 0) { + if (resize_mode == RefImageResizeMode::AREA) { + max_pixels = static_cast(16384 * factor * factor); + } else { + max_pixels = static_cast(128 * factor); + } + } + + int h_bar = std::max(factor, static_cast(std::round(static_cast(height) / factor) * factor)); + int w_bar = std::max(factor, static_cast(std::round(static_cast(width) / factor) * factor)); if (std::max(height, width) > 200 * std::min(height, width)) { LOG_WARN("LingBotVideo image aspect ratio is very large: %dx%d", width, height); } - if (h_bar * w_bar > max_pixels) { - double beta = std::sqrt((height * width) / static_cast(max_pixels)); - h_bar = std::max(static_cast(factor), - static_cast(std::floor(height / beta / factor)) * static_cast(factor)); - w_bar = std::max(static_cast(factor), - static_cast(std::floor(width / beta / factor)) * static_cast(factor)); - } else if (h_bar * w_bar < min_pixels) { - double beta = std::sqrt(static_cast(min_pixels) / (height * width)); - h_bar = static_cast(std::ceil(height * beta / factor)) * static_cast(factor); - w_bar = static_cast(std::ceil(width * beta / factor)) * static_cast(factor); - } + + resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode); LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); auto resized_image = clip_preprocess(image, w_bar, h_bar); @@ -2092,30 +2150,33 @@ struct LLMEmbedder : public Conditioner { prompt_template_encode_start_idx = 64; int image_embed_idx = 64 + 6; - int min_pixels = 384 * 384; - int max_pixels = 560 * 560; + int min_pixels = conditioner_params.ref_image_params.vlm_min_size; + if (min_pixels <= 0) { + min_pixels = 384; + if (resize_mode == RefImageResizeMode::AREA) { + min_pixels *= min_pixels; + } + } + int max_pixels = conditioner_params.ref_image_params.vlm_max_size; + if (max_pixels <= 0) { + max_pixels = 560; + if (resize_mode == RefImageResizeMode::AREA) { + max_pixels *= max_pixels; + } + } + std::string placeholder = "<|image_pad|>"; std::string img_prompt; for (int i = 0; i < conditioner_params.ref_images->size(); i++) { const auto& image = (*conditioner_params.ref_images)[i]; - double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; + const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; int height = static_cast(image.shape()[1]); int width = static_cast(image.shape()[0]); - int h_bar = static_cast(std::round(height / factor) * factor); - int w_bar = static_cast(std::round(width / factor) * factor); - - if (static_cast(h_bar) * w_bar > max_pixels) { - double beta = std::sqrt((height * width) / static_cast(max_pixels)); - h_bar = std::max(static_cast(factor), - static_cast(std::floor(height / beta / factor)) * static_cast(factor)); - w_bar = std::max(static_cast(factor), - static_cast(std::floor(width / beta / factor)) * static_cast(factor)); - } else if (static_cast(h_bar) * w_bar < min_pixels) { - double beta = std::sqrt(static_cast(min_pixels) / (height * width)); - h_bar = static_cast(std::ceil(height * beta / factor)) * static_cast(factor); - w_bar = static_cast(std::ceil(width * beta / factor)) * static_cast(factor); - } + int h_bar = static_cast(std::round(static_cast(height) / factor) * factor); + int w_bar = static_cast(std::round(static_cast(width) / factor) * factor); + + resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode); LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); @@ -2170,16 +2231,33 @@ struct LLMEmbedder : public Conditioner { std::string img_prompt; const std::string placeholder = "<|image_pad|>"; + int min_pixels = conditioner_params.ref_image_params.vlm_min_size; + if (min_pixels <= 0) { + min_pixels = 384; + if (resize_mode == RefImageResizeMode::AREA) { + min_pixels *= min_pixels; + } + } + int max_pixels = conditioner_params.ref_image_params.vlm_max_size; + if (max_pixels <= 0) { + max_pixels = 384; + if (resize_mode == RefImageResizeMode::AREA) { + max_pixels *= max_pixels; + } + } + for (int i = 0; i < conditioner_params.ref_images->size(); i++) { const auto& image = (*conditioner_params.ref_images)[i]; - double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; + const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; int height = static_cast(image.shape()[1]); int width = static_cast(image.shape()[0]); - double beta = std::sqrt((384.0 * 384.0) / (static_cast(height) * static_cast(width))); - int h_bar = std::max(static_cast(factor), - static_cast(std::round(height * beta / factor)) * static_cast(factor)); - int w_bar = std::max(static_cast(factor), - static_cast(std::round(width * beta / factor)) * static_cast(factor)); + + int h_bar = std::max(factor, + static_cast(std::round(static_cast(height) / factor)) * factor); + int w_bar = std::max(factor, + static_cast(std::round(static_cast(width) / factor)) * factor); + + resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode); LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); @@ -2221,17 +2299,33 @@ struct LLMEmbedder : public Conditioner { if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { std::string img_prompt = ""; const std::string placeholder = "<|image_pad|>"; + int min_pixels = conditioner_params.ref_image_params.vlm_min_size; + if (min_pixels <= 0) { + min_pixels = 384; + if (resize_mode == RefImageResizeMode::AREA) { + min_pixels *= min_pixels; + } + } + int max_pixels = conditioner_params.ref_image_params.vlm_max_size; + if (max_pixels <= 0) { + max_pixels = 1024; + if (resize_mode == RefImageResizeMode::AREA) { + max_pixels *= max_pixels; + } + } for (int i = 0; i < conditioner_params.ref_images->size(); i++) { const auto& image = (*conditioner_params.ref_images)[i]; - double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; + const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; int height = static_cast(image.shape()[1]); int width = static_cast(image.shape()[0]); - double beta = std::sqrt((384.0 * 384.0) / (static_cast(height) * static_cast(width))); - int h_bar = std::max(static_cast(factor), - static_cast(std::round(height * beta / factor)) * static_cast(factor)); - int w_bar = std::max(static_cast(factor), - static_cast(std::round(width * beta / factor)) * static_cast(factor)); + + int h_bar = std::max(factor, + static_cast(std::round(static_cast(height) / factor)) * factor); + int w_bar = std::max(factor, + static_cast(std::round(static_cast(width) / factor)) * factor); + + resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode); LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); @@ -2268,30 +2362,33 @@ struct LLMEmbedder : public Conditioner { min_length = 512 + prompt_template_encode_start_idx; int image_embed_idx = 36 + 6; - int min_pixels = 384 * 384; - int max_pixels = 560 * 560; + int min_pixels = conditioner_params.ref_image_params.vlm_min_size; + if (min_pixels <= 0) { + min_pixels = 384; + if (resize_mode == RefImageResizeMode::AREA) { + min_pixels *= min_pixels; + } + } + int max_pixels = conditioner_params.ref_image_params.vlm_max_size; + if (max_pixels <= 0) { + max_pixels = 560; + if (resize_mode == RefImageResizeMode::AREA) { + max_pixels *= max_pixels; + } + } + std::string placeholder = "<|image_pad|>"; std::string img_prompt; for (int i = 0; i < conditioner_params.ref_images->size(); i++) { const auto& image = (*conditioner_params.ref_images)[i]; - double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; + const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; int height = static_cast(image.shape()[1]); int width = static_cast(image.shape()[0]); - int h_bar = static_cast(std::round(height / factor) * factor); - int w_bar = static_cast(std::round(width / factor) * factor); - - if (static_cast(h_bar) * w_bar > max_pixels) { - double beta = std::sqrt((height * width) / static_cast(max_pixels)); - h_bar = std::max(static_cast(factor), - static_cast(std::floor(height / beta / factor)) * static_cast(factor)); - w_bar = std::max(static_cast(factor), - static_cast(std::floor(width / beta / factor)) * static_cast(factor)); - } else if (static_cast(h_bar) * w_bar < min_pixels) { - double beta = std::sqrt(static_cast(min_pixels) / (height * width)); - h_bar = static_cast(std::ceil(height * beta / factor)) * static_cast(factor); - w_bar = static_cast(std::ceil(width * beta / factor)) * static_cast(factor); - } + int h_bar = static_cast(std::round(static_cast(height) / factor) * factor); + int w_bar = static_cast(std::round(static_cast(width) / factor) * factor); + + resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode); LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); diff --git a/otherarch/sdcpp/src/core/ggml_extend_backend.cpp b/otherarch/sdcpp/src/core/ggml_extend_backend.cpp index c66bb63f01f1..a8316643872b 100644 --- a/otherarch/sdcpp/src/core/ggml_extend_backend.cpp +++ b/otherarch/sdcpp/src/core/ggml_extend_backend.cpp @@ -83,6 +83,10 @@ static bool parse_backend_module(const std::string& raw_name, SDBackendModule* m *module = SDBackendModule::UPSCALER; return true; } + if (name == "detector" || name == "adetailer" || name == "yolo") { + *module = SDBackendModule::DETECTOR; + return true; + } return false; } @@ -956,6 +960,8 @@ const char* sd_backend_module_name(SDBackendModule module) { return "photomaker"; case SDBackendModule::UPSCALER: return "upscaler"; + case SDBackendModule::DETECTOR: + return "detector"; } return "unknown"; } diff --git a/otherarch/sdcpp/src/core/ggml_extend_backend.h b/otherarch/sdcpp/src/core/ggml_extend_backend.h index 1f3bc8b3ae84..d5498e8f46e2 100644 --- a/otherarch/sdcpp/src/core/ggml_extend_backend.h +++ b/otherarch/sdcpp/src/core/ggml_extend_backend.h @@ -20,6 +20,7 @@ enum class SDBackendModule { CONTROL_NET, PHOTOMAKER, UPSCALER, + DETECTOR, }; struct SDBackendAssignment { diff --git a/otherarch/sdcpp/src/detailer.cpp b/otherarch/sdcpp/src/detailer.cpp new file mode 100644 index 000000000000..fa9884be33c0 --- /dev/null +++ b/otherarch/sdcpp/src/detailer.cpp @@ -0,0 +1,1020 @@ +#include "detailer.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/util.h" +#include "json.hpp" +#include "model_loader.h" + +struct adetailer_ctx_t { + ADetailerGGML* detailer = nullptr; +}; + +struct LetterboxInput { + sd::Tensor tensor; + float scale = 1.f; + int pad_x = 0; + int pad_y = 0; +}; + +struct Mask { + int width = 0; + int height = 0; + std::vector data; +}; + +struct CropRegion { + int x1 = 0; + int y1 = 0; + int x2 = 0; + int y2 = 0; +}; + +static std::string trim_copy(const std::string& value) { + size_t first = value.find_first_not_of(" \t\r\n"); + if (first == std::string::npos) { + return ""; + } + size_t last = value.find_last_not_of(" \t\r\n"); + return value.substr(first, last - first + 1); +} + +static std::string lower_copy(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return value; +} + +static bool parse_bool(const std::string& value, bool* result) { + const std::string lower = lower_copy(trim_copy(value)); + if (lower == "1" || lower == "true" || lower == "yes" || lower == "on") { + *result = true; + return true; + } + if (lower == "0" || lower == "false" || lower == "no" || lower == "off") { + *result = false; + return true; + } + return false; +} + +static bool parse_int(const std::string& value, int* result) { + try { + size_t used = 0; + int parsed = std::stoi(trim_copy(value), &used); + if (used != trim_copy(value).size()) { + return false; + } + *result = parsed; + return true; + } catch (...) { + return false; + } +} + +static bool parse_float(const std::string& value, float* result) { + try { + size_t used = 0; + float parsed = std::stof(trim_copy(value), &used); + if (used != trim_copy(value).size() || !std::isfinite(parsed)) { + return false; + } + *result = parsed; + return true; + } catch (...) { + return false; + } +} + +static float image_channel(const sd_image_t& image, int x, int y, int channel) { + if (image.data == nullptr || image.width == 0 || image.height == 0 || image.channel == 0) { + return 0.f; + } + x = std::clamp(x, 0, static_cast(image.width) - 1); + y = std::clamp(y, 0, static_cast(image.height) - 1); + int source_channel = image.channel == 1 ? 0 : std::min(channel, static_cast(image.channel) - 1); + size_t index = (static_cast(y) * image.width + x) * image.channel + source_channel; + return image.data[index] / 255.f; +} + +static float bilinear_channel(const sd_image_t& image, float x, float y, int channel) { + int x0 = static_cast(std::floor(x)); + int y0 = static_cast(std::floor(y)); + int x1 = x0 + 1; + int y1 = y0 + 1; + float wx = x - x0; + float wy = y - y0; + float a = image_channel(image, x0, y0, channel) * (1.f - wx) + image_channel(image, x1, y0, channel) * wx; + float b = image_channel(image, x0, y1, channel) * (1.f - wx) + image_channel(image, x1, y1, channel) * wx; + return a * (1.f - wy) + b * wy; +} + +static LetterboxInput make_letterbox_input(sd_image_t image, int input_size) { + LetterboxInput result; + result.tensor = sd::full({input_size, input_size, 3, 1}, 114.f / 255.f); + result.scale = std::min(static_cast(input_size) / image.width, + static_cast(input_size) / image.height); + int resized_width = std::max(1, static_cast(std::round(image.width * result.scale))); + int resized_height = std::max(1, static_cast(std::round(image.height * result.scale))); + result.pad_x = (input_size - resized_width) / 2; + result.pad_y = (input_size - resized_height) / 2; + + for (int y = 0; y < resized_height; ++y) { + float source_y = (y + 0.5f) / result.scale - 0.5f; + for (int x = 0; x < resized_width; ++x) { + float source_x = (x + 0.5f) / result.scale - 0.5f; + for (int c = 0; c < 3; ++c) { + result.tensor.index(x + result.pad_x, y + result.pad_y, c, 0) = bilinear_channel(image, source_x, source_y, c); + } + } + } + return result; +} + +static float logistic(float x) { + if (x >= 0.f) { + float z = std::exp(-x); + return 1.f / (1.f + z); + } + float z = std::exp(x); + return z / (1.f + z); +} + +static float dfl_expectation(const float* values, int64_t anchor_count, int channel_offset, int reg_max, int64_t anchor) { + float maximum = -std::numeric_limits::infinity(); + for (int i = 0; i < reg_max; ++i) { + maximum = std::max(maximum, values[anchor + anchor_count * (channel_offset + i)]); + } + float denominator = 0.f; + float numerator = 0.f; + for (int i = 0; i < reg_max; ++i) { + float probability = std::exp(values[anchor + anchor_count * (channel_offset + i)] - maximum); + denominator += probability; + numerator += probability * i; + } + return denominator > 0.f ? numerator / denominator : 0.f; +} + +static float box_iou(const ADetailerDetection& a, const ADetailerDetection& b) { + float x1 = std::max(a.x1, b.x1); + float y1 = std::max(a.y1, b.y1); + float x2 = std::min(a.x2, b.x2); + float y2 = std::min(a.y2, b.y2); + float intersection = std::max(0.f, x2 - x1) * std::max(0.f, y2 - y1); + float area_a = std::max(0.f, a.x2 - a.x1) * std::max(0.f, a.y2 - a.y1); + float area_b = std::max(0.f, b.x2 - b.x1) * std::max(0.f, b.y2 - b.y1); + float union_area = area_a + area_b - intersection; + return union_area > 0.f ? intersection / union_area : 0.f; +} + +static std::vector decode_detections(const sd::Tensor& raw, + const YOLOv8Config& config, + const LetterboxInput& letterbox, + sd_image_t image, + const ADetailerParams& params) { + std::vector candidates; + if (raw.empty() || raw.dim() < 2) { + return candidates; + } + int64_t anchor_count = raw.shape()[0]; + int64_t channels = raw.shape()[1]; + if (channels != config.reg_max * 4 + config.num_classes) { + LOG_ERROR("unexpected YOLOv8 output channels: %lld", static_cast(channels)); + return candidates; + } + + const std::array strides = {8, 16, 32}; + std::array offsets = {0, 0, 0, 0}; + for (int i = 0; i < 3; ++i) { + int grid = params.input_size / strides[i]; + offsets[i + 1] = offsets[i] + static_cast(grid) * grid; + } + if (offsets[3] != anchor_count) { + LOG_ERROR("unexpected YOLOv8 anchor count: %lld (expected %lld)", + static_cast(anchor_count), + static_cast(offsets[3])); + return candidates; + } + + const float* values = raw.data(); + for (int scale_index = 0; scale_index < 3; ++scale_index) { + int stride = strides[scale_index]; + int grid = params.input_size / stride; + for (int64_t anchor = offsets[scale_index]; anchor < offsets[scale_index + 1]; ++anchor) { + int64_t local = anchor - offsets[scale_index]; + int grid_x = static_cast(local % grid); + int grid_y = static_cast(local / grid); + + float confidence = 0.f; + int class_id = 0; + for (int c = 0; c < config.num_classes; ++c) { + float score = logistic(values[anchor + anchor_count * (config.reg_max * 4 + c)]); + if (score > confidence) { + confidence = score; + class_id = c; + } + } + if (confidence < params.confidence) { + continue; + } + + float left = dfl_expectation(values, anchor_count, 0 * config.reg_max, config.reg_max, anchor); + float top = dfl_expectation(values, anchor_count, 1 * config.reg_max, config.reg_max, anchor); + float right = dfl_expectation(values, anchor_count, 2 * config.reg_max, config.reg_max, anchor); + float bottom = dfl_expectation(values, anchor_count, 3 * config.reg_max, config.reg_max, anchor); + float center_x = (grid_x + 0.5f) * stride; + float center_y = (grid_y + 0.5f) * stride; + + ADetailerDetection detection; + detection.x1 = (center_x - left * stride - letterbox.pad_x) / letterbox.scale; + detection.y1 = (center_y - top * stride - letterbox.pad_y) / letterbox.scale; + detection.x2 = (center_x + right * stride - letterbox.pad_x) / letterbox.scale; + detection.y2 = (center_y + bottom * stride - letterbox.pad_y) / letterbox.scale; + detection.x1 = std::clamp(detection.x1, 0.f, static_cast(image.width)); + detection.y1 = std::clamp(detection.y1, 0.f, static_cast(image.height)); + detection.x2 = std::clamp(detection.x2, 0.f, static_cast(image.width)); + detection.y2 = std::clamp(detection.y2, 0.f, static_cast(image.height)); + detection.confidence = confidence; + detection.class_id = class_id; + if (detection.x2 > detection.x1 && detection.y2 > detection.y1) { + candidates.push_back(detection); + } + } + } + + std::sort(candidates.begin(), candidates.end(), [](const auto& a, const auto& b) { + return a.confidence > b.confidence; + }); + std::vector selected; + for (const auto& candidate : candidates) { + bool suppressed = false; + for (const auto& kept : selected) { + if (candidate.class_id == kept.class_id && box_iou(candidate, kept) > params.nms_threshold) { + suppressed = true; + break; + } + } + if (!suppressed) { + selected.push_back(candidate); + if (static_cast(selected.size()) >= params.max_detections) { + break; + } + } + } + return selected; +} + +static float detection_area(const ADetailerDetection& detection) { + return std::max(0.f, detection.x2 - detection.x1) * std::max(0.f, detection.y2 - detection.y1); +} + +static void filter_and_sort_detections(std::vector* detections, + int width, + int height, + const ADetailerParams& params) { + const float image_area = static_cast(width) * height; + detections->erase(std::remove_if(detections->begin(), detections->end(), [&](const auto& detection) { + float ratio = detection_area(detection) / image_area; + return ratio < params.mask_min_ratio || ratio > params.mask_max_ratio; + }), + detections->end()); + + if (params.mask_k_largest > 0 && static_cast(detections->size()) > params.mask_k_largest) { + std::partial_sort(detections->begin(), + detections->begin() + params.mask_k_largest, + detections->end(), + [](const auto& a, const auto& b) { return detection_area(a) > detection_area(b); }); + detections->resize(params.mask_k_largest); + } + + if (params.sort_by == ADETAILER_SORT_LEFT_TO_RIGHT) { + std::sort(detections->begin(), detections->end(), [](const auto& a, const auto& b) { return a.x1 < b.x1; }); + } else if (params.sort_by == ADETAILER_SORT_CENTER_TO_EDGE) { + const float cx = width * 0.5f; + const float cy = height * 0.5f; + std::sort(detections->begin(), detections->end(), [&](const auto& a, const auto& b) { + float ax = (a.x1 + a.x2) * 0.5f - cx; + float ay = (a.y1 + a.y2) * 0.5f - cy; + float bx = (b.x1 + b.x2) * 0.5f - cx; + float by = (b.y1 + b.y2) * 0.5f - cy; + return ax * ax + ay * ay < bx * bx + by * by; + }); + } else if (params.sort_by == ADETAILER_SORT_AREA) { + std::sort(detections->begin(), detections->end(), [](const auto& a, const auto& b) { + return detection_area(a) > detection_area(b); + }); + } +} + +static Mask bbox_mask(const ADetailerDetection& detection, int width, int height) { + Mask mask; + mask.width = width; + mask.height = height; + mask.data.assign(static_cast(width) * height, 0); + int x1 = std::clamp(static_cast(std::floor(detection.x1)), 0, width); + int y1 = std::clamp(static_cast(std::floor(detection.y1)), 0, height); + int x2 = std::clamp(static_cast(std::ceil(detection.x2)), 0, width); + int y2 = std::clamp(static_cast(std::ceil(detection.y2)), 0, height); + for (int y = y1; y < y2; ++y) { + std::fill(mask.data.begin() + static_cast(y) * width + x1, + mask.data.begin() + static_cast(y) * width + x2, + 255); + } + return mask; +} + +static Mask offset_mask(const Mask& source, int x_offset, int y_offset) { + Mask output{source.width, source.height, std::vector(source.data.size(), 0)}; + for (int y = 0; y < source.height; ++y) { + int target_y = y - y_offset; + if (target_y < 0 || target_y >= source.height) { + continue; + } + for (int x = 0; x < source.width; ++x) { + int target_x = x + x_offset; + if (target_x >= 0 && target_x < source.width) { + output.data[static_cast(target_y) * source.width + target_x] = source.data[static_cast(y) * source.width + x]; + } + } + } + return output; +} + +static Mask morphology_mask(const Mask& source, int amount) { + if (amount == 0) { + return source; + } + int kernel = std::abs(amount); + int before = kernel / 2; + int after = kernel - before - 1; + Mask output{source.width, source.height, std::vector(source.data.size(), amount > 0 ? 0 : 255)}; + for (int y = 0; y < source.height; ++y) { + for (int x = 0; x < source.width; ++x) { + uint8_t value = amount > 0 ? 0 : 255; + for (int ky = -before; ky <= after; ++ky) { + for (int kx = -before; kx <= after; ++kx) { + int sx = x + kx; + int sy = y + ky; + uint8_t sample = (sx >= 0 && sx < source.width && sy >= 0 && sy < source.height) + ? source.data[static_cast(sy) * source.width + sx] + : 0; + value = amount > 0 ? std::max(value, sample) : std::min(value, sample); + } + } + output.data[static_cast(y) * source.width + x] = value; + } + } + return output; +} + +static Mask gaussian_blur_mask(const Mask& source, int radius) { + if (radius <= 0) { + return source; + } + float sigma = static_cast(radius); + int kernel_radius = std::max(1, static_cast(std::ceil(2.5f * sigma))); + std::vector kernel(kernel_radius * 2 + 1); + float sum = 0.f; + for (int i = -kernel_radius; i <= kernel_radius; ++i) { + float value = std::exp(-(i * i) / (2.f * sigma * sigma)); + kernel[i + kernel_radius] = value; + sum += value; + } + for (float& value : kernel) { + value /= sum; + } + + std::vector horizontal(source.data.size()); + for (int y = 0; y < source.height; ++y) { + for (int x = 0; x < source.width; ++x) { + float value = 0.f; + for (int k = -kernel_radius; k <= kernel_radius; ++k) { + int sx = std::clamp(x + k, 0, source.width - 1); + value += source.data[static_cast(y) * source.width + sx] * kernel[k + kernel_radius]; + } + horizontal[static_cast(y) * source.width + x] = value; + } + } + + Mask output{source.width, source.height, std::vector(source.data.size())}; + for (int y = 0; y < source.height; ++y) { + for (int x = 0; x < source.width; ++x) { + float value = 0.f; + for (int k = -kernel_radius; k <= kernel_radius; ++k) { + int sy = std::clamp(y + k, 0, source.height - 1); + value += horizontal[static_cast(sy) * source.width + x] * kernel[k + kernel_radius]; + } + output.data[static_cast(y) * source.width + x] = static_cast(std::clamp(value, 0.f, 255.f) + 0.5f); + } + } + return output; +} + +static std::vector make_masks(const std::vector& detections, + int width, + int height, + const ADetailerParams& params) { + std::vector masks; + for (const auto& detection : detections) { + Mask mask = bbox_mask(detection, width, height); + if (params.x_offset != 0 || params.y_offset != 0) { + mask = offset_mask(mask, params.x_offset, params.y_offset); + } + mask = morphology_mask(mask, params.dilate_erode); + if (std::any_of(mask.data.begin(), mask.data.end(), [](uint8_t value) { return value != 0; })) { + masks.push_back(std::move(mask)); + } + } + if (params.merge_masks && !masks.empty()) { + Mask merged{width, height, std::vector(static_cast(width) * height, 0)}; + for (const Mask& mask : masks) { + for (size_t i = 0; i < merged.data.size(); ++i) { + merged.data[i] = std::max(merged.data[i], mask.data[i]); + } + } + masks = {std::move(merged)}; + } + if (params.invert_mask) { + for (Mask& mask : masks) { + for (uint8_t& value : mask.data) { + value = 255 - value; + } + } + } + return masks; +} + +static bool mask_bbox(const Mask& mask, CropRegion* region) { + int x1 = mask.width; + int y1 = mask.height; + int x2 = 0; + int y2 = 0; + for (int y = 0; y < mask.height; ++y) { + for (int x = 0; x < mask.width; ++x) { + if (mask.data[static_cast(y) * mask.width + x] != 0) { + x1 = std::min(x1, x); + y1 = std::min(y1, y); + x2 = std::max(x2, x + 1); + y2 = std::max(y2, y + 1); + } + } + } + if (x2 <= x1 || y2 <= y1) { + return false; + } + *region = {x1, y1, x2, y2}; + return true; +} + +static CropRegion expand_crop(CropRegion crop, + int image_width, + int image_height, + int padding, + int target_width, + int target_height) { + crop.x1 = std::max(0, crop.x1 - padding); + crop.y1 = std::max(0, crop.y1 - padding); + crop.x2 = std::min(image_width, crop.x2 + padding); + crop.y2 = std::min(image_height, crop.y2 + padding); + + float target_aspect = static_cast(target_width) / target_height; + int width = crop.x2 - crop.x1; + int height = crop.y2 - crop.y1; + int desired_width = width; + int desired_height = height; + if (static_cast(width) / height < target_aspect) { + desired_width = static_cast(std::ceil(height * target_aspect)); + } else { + desired_height = static_cast(std::ceil(width / target_aspect)); + } + desired_width = std::min(desired_width, image_width); + desired_height = std::min(desired_height, image_height); + int center_x = (crop.x1 + crop.x2) / 2; + int center_y = (crop.y1 + crop.y2) / 2; + crop.x1 = std::clamp(center_x - desired_width / 2, 0, image_width - desired_width); + crop.y1 = std::clamp(center_y - desired_height / 2, 0, image_height - desired_height); + crop.x2 = crop.x1 + desired_width; + crop.y2 = crop.y1 + desired_height; + return crop; +} + +static sd_image_t resize_crop_image(const sd_image_t& source, const CropRegion& crop, int width, int height) { + sd_image_t output = {static_cast(width), static_cast(height), 3, nullptr}; + output.data = static_cast(malloc(static_cast(width) * height * 3)); + if (output.data == nullptr) { + return {0, 0, 0, nullptr}; + } + float scale_x = static_cast(crop.x2 - crop.x1) / width; + float scale_y = static_cast(crop.y2 - crop.y1) / height; + for (int y = 0; y < height; ++y) { + float source_y = crop.y1 + (y + 0.5f) * scale_y - 0.5f; + for (int x = 0; x < width; ++x) { + float source_x = crop.x1 + (x + 0.5f) * scale_x - 0.5f; + for (int c = 0; c < 3; ++c) { + float value = bilinear_channel(source, source_x, source_y, c); + output.data[(static_cast(y) * width + x) * 3 + c] = static_cast(std::clamp(value * 255.f, 0.f, 255.f) + 0.5f); + } + } + } + return output; +} + +static sd_image_t resize_crop_mask(const Mask& source, const CropRegion& crop, int width, int height) { + sd_image_t output = {static_cast(width), static_cast(height), 1, nullptr}; + output.data = static_cast(malloc(static_cast(width) * height)); + if (output.data == nullptr) { + return {0, 0, 0, nullptr}; + } + float scale_x = static_cast(crop.x2 - crop.x1) / width; + float scale_y = static_cast(crop.y2 - crop.y1) / height; + for (int y = 0; y < height; ++y) { + int source_y = std::clamp(crop.y1 + static_cast((y + 0.5f) * scale_y), 0, source.height - 1); + for (int x = 0; x < width; ++x) { + int source_x = std::clamp(crop.x1 + static_cast((x + 0.5f) * scale_x), 0, source.width - 1); + output.data[static_cast(y) * width + x] = source.data[static_cast(source_y) * source.width + source_x]; + } + } + return output; +} + +static sd_image_t copy_as_rgb(const sd_image_t& source) { + CropRegion full{0, 0, static_cast(source.width), static_cast(source.height)}; + return resize_crop_image(source, full, source.width, source.height); +} + +static void composite_crop(sd_image_t* destination, + const sd_image_t& generated, + const Mask& feather_mask, + const CropRegion& crop) { + int crop_width = crop.x2 - crop.x1; + int crop_height = crop.y2 - crop.y1; + float scale_x = static_cast(generated.width) / crop_width; + float scale_y = static_cast(generated.height) / crop_height; + for (int y = crop.y1; y < crop.y2; ++y) { + float source_y = (y - crop.y1 + 0.5f) * scale_y - 0.5f; + for (int x = crop.x1; x < crop.x2; ++x) { + float alpha = feather_mask.data[static_cast(y) * feather_mask.width + x] / 255.f; + if (alpha <= 0.f) { + continue; + } + float source_x = (x - crop.x1 + 0.5f) * scale_x - 0.5f; + size_t destination_index = (static_cast(y) * destination->width + x) * 3; + for (int c = 0; c < 3; ++c) { + float generated_value = bilinear_channel(generated, source_x, source_y, c) * 255.f; + float original_value = destination->data[destination_index + c]; + float blended = original_value * (1.f - alpha) + generated_value * alpha; + destination->data[destination_index + c] = static_cast(std::clamp(blended, 0.f, 255.f) + 0.5f); + } + } + } +} + +static std::vector split_prompts(const std::string& prompt) { + std::vector result; + size_t start = 0; + while (true) { + size_t separator = prompt.find("[SEP]", start); + result.push_back(trim_copy(prompt.substr(start, separator == std::string::npos ? std::string::npos : separator - start))); + if (separator == std::string::npos) { + break; + } + start = separator + 5; + } + return result; +} + +static std::string resolve_prompt(const char* prompt_template, const char* base_prompt, int index) { + std::string base = base_prompt == nullptr ? "" : base_prompt; + std::string templ = prompt_template == nullptr ? "" : prompt_template; + auto prompts = split_prompts(templ); + std::string selected = prompts[std::min(index, static_cast(prompts.size()) - 1)]; + if (selected.empty()) { + return base; + } + size_t position = 0; + while ((position = selected.find("[PROMPT]", position)) != std::string::npos) { + selected.replace(position, 8, base); + position += base.size(); + } + return selected; +} + +static bool prompt_is_skip(const std::string& prompt) { + return trim_copy(prompt) == "[SKIP]"; +} + +static bool validate_params(const ADetailerParams& params) { + return params.input_size >= 32 && params.input_size % 32 == 0 && + params.confidence >= 0.f && params.confidence <= 1.f && + params.nms_threshold >= 0.f && params.nms_threshold <= 1.f && + params.max_detections > 0 && params.mask_k_largest >= 0 && + params.mask_min_ratio >= 0.f && params.mask_max_ratio <= 1.f && + params.mask_min_ratio <= params.mask_max_ratio && params.mask_blur >= 0 && + params.inpaint_padding >= 0 && params.inpaint_width > 0 && params.inpaint_height > 0 && + params.denoising_strength >= 0.f && params.denoising_strength <= 1.f && params.steps >= 0 && + (params.cfg_scale < 0.f || std::isfinite(params.cfg_scale)); +} + +static std::vector parse_class_names(const std::string& value, int class_count) { + std::vector names(static_cast(class_count)); + try { + nlohmann::json parsed = nlohmann::json::parse(value); + if (parsed.is_array()) { + for (size_t i = 0; i < names.size() && i < parsed.size(); ++i) { + if (parsed[i].is_string()) { + names[i] = parsed[i].get(); + } + } + } else if (parsed.is_object()) { + for (const auto& item : parsed.items()) { + int class_id = 0; + if (parse_int(item.key(), &class_id) && class_id >= 0 && class_id < class_count && + item.value().is_string()) { + names[static_cast(class_id)] = item.value().get(); + } + } + } + } catch (const std::exception&) { + } + return names; +} + +ADetailerGGML::ADetailerGGML(int n_threads, + std::string backend_spec, + std::string params_backend_spec) + : n_threads(n_threads > 0 ? n_threads : 1), + backend_spec(std::move(backend_spec)), + params_backend_spec(std::move(params_backend_spec)) { +} + +ADetailerGGML::~ADetailerGGML() { + model_manager.reset(); + detector.reset(); +} + +bool ADetailerGGML::load_from_file(const std::string& detector_path) { + std::string error; + if (!backend_manager.init(backend_spec.c_str(), params_backend_spec.c_str(), nullptr, &error)) { + LOG_ERROR("ADetailer backend config failed: %s", error.c_str()); + return false; + } + ggml_backend_t backend = backend_manager.runtime_backend(SDBackendModule::DETECTOR); + ggml_backend_t params_backend = backend_manager.params_backend(SDBackendModule::DETECTOR); + if (backend == nullptr || params_backend == nullptr) { + LOG_ERROR("failed to initialize detector backend"); + return false; + } + + model_manager = std::make_shared(); + model_manager->set_n_threads(n_threads); + model_manager->set_enable_mmap(false); + ModelLoader& loader = model_manager->loader(); + if (!loader.init_from_file(detector_path)) { + LOG_ERROR("failed to load ADetailer detector: '%s'", detector_path.c_str()); + return false; + } + + detector = std::make_shared(backend, loader.get_tensor_storage_map(), model_manager); + if (!detector || !detector->model || !detector->config.valid) { + LOG_ERROR("unsupported YOLOv8 detection weights: '%s'", detector_path.c_str()); + return false; + } + + class_names.assign(static_cast(detector->config.num_classes), ""); + auto names_item = loader.get_metadata().find("yolov8.names"); + if (names_item != loader.get_metadata().end()) { + class_names = parse_class_names(names_item->second, detector->config.num_classes); + } + + std::map tensors; + detector->get_param_tensors(tensors); + if (!model_manager->register_param_tensors("YOLOv8", + std::move(tensors), + backend_manager.params_backend_is_disk(SDBackendModule::DETECTOR) + ? ModelManager::ResidencyMode::Disk + : ModelManager::ResidencyMode::ParamBackend, + backend, + params_backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("failed to register YOLOv8 detector tensors"); + return false; + } + return true; +} + +std::vector ADetailerGGML::predict(sd_image_t image, + const ADetailerParams& params) { + LetterboxInput input = make_letterbox_input(image, params.input_size); + int64_t start = ggml_time_ms(); + sd::Tensor raw = detector->compute(n_threads, input.tensor); + detector->free_compute_buffer(); + if (raw.empty()) { + LOG_ERROR("YOLOv8 detector inference failed"); + return {}; + } + auto detections = decode_detections(raw, detector->config, input, image, params); + LOG_INFO("ADetailer detected %zu object(s), taking %.2fs", + detections.size(), + (ggml_time_ms() - start) / 1000.f); + for (size_t i = 0; i < detections.size(); ++i) { + const auto& detection = detections[i]; + std::string object = "class_" + std::to_string(detection.class_id); + if (detection.class_id >= 0 && static_cast(detection.class_id) < class_names.size() && + !class_names[static_cast(detection.class_id)].empty()) { + object = class_names[static_cast(detection.class_id)]; + } + LOG_INFO("ADetailer detection %zu: object=%s, class_id=%d, confidence=%.3f, bbox=[x1=%.1f, y1=%.1f, x2=%.1f, y2=%.1f]", + i + 1, + object.c_str(), + detection.class_id, + detection.confidence, + detection.x1, + detection.y1, + detection.x2, + detection.y2); + } + return detections; +} + +static bool parse_adetailer_params(const sd_adetailer_params_t& public_params, + const sd_img_gen_params_t& inpaint_params, + ADetailerParams* params) { + if (params == nullptr) { + return false; + } + *params = {}; + params->prompt = public_params.prompt; + params->negative_prompt = public_params.negative_prompt; + params->inpaint_width = inpaint_params.width; + params->inpaint_height = inpaint_params.height; + params->denoising_strength = inpaint_params.strength; + const char* extra_ad_args = public_params.extra_ad_args; + if (extra_ad_args == nullptr || extra_ad_args[0] == '\0') { + return validate_params(*params); + } + + std::stringstream stream(extra_ad_args); + std::string item; + while (std::getline(stream, item, ',')) { + item = trim_copy(item); + if (item.empty()) { + continue; + } + size_t equals = item.find('='); + if (equals == std::string::npos) { + LOG_ERROR("invalid --extra-ad-args item '%s'; expected key=value", item.c_str()); + return false; + } + std::string key = lower_copy(trim_copy(item.substr(0, equals))); + std::string value = trim_copy(item.substr(equals + 1)); + bool ok = true; + if (key == "input_size") { + ok = parse_int(value, ¶ms->input_size); + } else if (key == "confidence") { + ok = parse_float(value, ¶ms->confidence); + } else if (key == "nms" || key == "nms_threshold") { + ok = parse_float(value, ¶ms->nms_threshold); + } else if (key == "max_detections") { + ok = parse_int(value, ¶ms->max_detections); + } else if (key == "mask_k_largest") { + ok = parse_int(value, ¶ms->mask_k_largest); + } else if (key == "mask_min_ratio") { + ok = parse_float(value, ¶ms->mask_min_ratio); + } else if (key == "mask_max_ratio") { + ok = parse_float(value, ¶ms->mask_max_ratio); + } else if (key == "dilate_erode") { + ok = parse_int(value, ¶ms->dilate_erode); + } else if (key == "x_offset") { + ok = parse_int(value, ¶ms->x_offset); + } else if (key == "y_offset") { + ok = parse_int(value, ¶ms->y_offset); + } else if (key == "merge_masks") { + ok = parse_bool(value, ¶ms->merge_masks); + } else if (key == "invert_mask") { + ok = parse_bool(value, ¶ms->invert_mask); + } else if (key == "mask_blur") { + ok = parse_int(value, ¶ms->mask_blur); + } else if (key == "inpaint_padding") { + ok = parse_int(value, ¶ms->inpaint_padding); + } else if (key == "inpaint_width") { + ok = parse_int(value, ¶ms->inpaint_width); + } else if (key == "inpaint_height") { + ok = parse_int(value, ¶ms->inpaint_height); + } else if (key == "denoising_strength") { + ok = parse_float(value, ¶ms->denoising_strength); + } else if (key == "steps") { + ok = parse_int(value, ¶ms->steps); + } else if (key == "cfg_scale") { + ok = parse_float(value, ¶ms->cfg_scale); + } else if (key == "sample_method") { + params->sample_method = str_to_sample_method(value.c_str()); + ok = params->sample_method != SAMPLE_METHOD_COUNT; + } else if (key == "scheduler") { + params->scheduler = str_to_scheduler(value.c_str()); + ok = params->scheduler != SCHEDULER_COUNT; + } else if (key == "sort_by") { + std::string sort = lower_copy(value); + if (sort == "none") { + params->sort_by = ADETAILER_SORT_NONE; + } else if (sort == "left_to_right" || sort == "left-to-right") { + params->sort_by = ADETAILER_SORT_LEFT_TO_RIGHT; + } else if (sort == "center_to_edge" || sort == "center-to-edge") { + params->sort_by = ADETAILER_SORT_CENTER_TO_EDGE; + } else if (sort == "area") { + params->sort_by = ADETAILER_SORT_AREA; + } else { + ok = false; + } + } else if (key == "mask_mode") { + std::string mode = lower_copy(value); + if (mode == "none") { + params->merge_masks = false; + params->invert_mask = false; + } else if (mode == "merge") { + params->merge_masks = true; + params->invert_mask = false; + } else if (mode == "merge_invert" || mode == "merge-invert") { + params->merge_masks = true; + params->invert_mask = true; + } else { + ok = false; + } + } else { + LOG_ERROR("unknown --extra-ad-args key '%s'", key.c_str()); + return false; + } + if (!ok) { + LOG_ERROR("invalid --extra-ad-args value '%s=%s'", key.c_str(), value.c_str()); + return false; + } + } + if (!validate_params(*params)) { + LOG_ERROR("invalid --extra-ad-args parameter range"); + return false; + } + return true; +} + +adetailer_ctx_t* new_adetailer_ctx(const char* detector_path, + int n_threads, + const char* backend, + const char* params_backend) { + if (detector_path == nullptr || detector_path[0] == '\0') { + return nullptr; + } + auto* context = static_cast(calloc(1, sizeof(adetailer_ctx_t))); + if (context == nullptr) { + return nullptr; + } + context->detailer = new ADetailerGGML(n_threads, SAFE_STR(backend), SAFE_STR(params_backend)); + if (context->detailer == nullptr || !context->detailer->load_from_file(detector_path)) { + delete context->detailer; + free(context); + return nullptr; + } + return context; +} + +void free_adetailer_ctx(adetailer_ctx_t* context) { + if (context == nullptr) { + return; + } + delete context->detailer; + context->detailer = nullptr; + free(context); +} + +bool adetail_image(adetailer_ctx_t* context, + sd_ctx_t* sd_ctx, + sd_image_t input_image, + const sd_adetailer_params_t* public_params, + const sd_img_gen_params_t* inpaint_params, + sd_image_t** images_out, + int* num_images_out) { + if (images_out != nullptr) { + *images_out = nullptr; + } + if (num_images_out != nullptr) { + *num_images_out = 0; + } + if (context == nullptr || context->detailer == nullptr || sd_ctx == nullptr || public_params == nullptr || + inpaint_params == nullptr || input_image.data == nullptr) { + return false; + } + + ADetailerParams params; + if (!parse_adetailer_params(*public_params, *inpaint_params, ¶ms)) { + return false; + } + + auto detections = context->detailer->predict(input_image, params); + filter_and_sort_detections(&detections, input_image.width, input_image.height, params); + auto masks = make_masks(detections, input_image.width, input_image.height, params); + + sd_image_t current = copy_as_rgb(input_image); + if (current.data == nullptr) { + return false; + } + + int applied = 0; + for (size_t i = 0; i < masks.size(); ++i) { + std::string prompt = resolve_prompt(params.prompt, inpaint_params->prompt, static_cast(i)); + std::string negative_prompt = resolve_prompt(params.negative_prompt, + inpaint_params->negative_prompt, + static_cast(i)); + if (prompt_is_skip(prompt)) { + continue; + } + + CropRegion crop; + if (!mask_bbox(masks[i], &crop)) { + continue; + } + crop = expand_crop(crop, + current.width, + current.height, + params.inpaint_padding, + params.inpaint_width, + params.inpaint_height); + sd_image_t local_image = resize_crop_image(current, crop, params.inpaint_width, params.inpaint_height); + sd_image_t local_mask = resize_crop_mask(masks[i], crop, params.inpaint_width, params.inpaint_height); + if (local_image.data == nullptr || local_mask.data == nullptr) { + free(local_image.data); + free(local_mask.data); + free(current.data); + return false; + } + + sd_img_gen_params_t generation = *inpaint_params; + generation.prompt = prompt.c_str(); + generation.negative_prompt = negative_prompt.c_str(); + generation.init_image = local_image; + generation.mask_image = local_mask; + generation.width = params.inpaint_width; + generation.height = params.inpaint_height; + generation.strength = params.denoising_strength; + generation.seed = inpaint_params->seed + static_cast(i); + generation.batch_count = 1; + generation.ref_images = nullptr; + generation.ref_images_count = 0; + generation.control_image = {}; + generation.pm_params = {}; + generation.pulid_params = {}; + generation.hires.enabled = false; + if (params.steps > 0) { + generation.sample_params.sample_steps = params.steps; + generation.sample_params.custom_sigmas = nullptr; + generation.sample_params.custom_sigmas_count = 0; + } + if (params.cfg_scale >= 0.f) { + generation.sample_params.guidance.txt_cfg = params.cfg_scale; + } + if (params.sample_method != SAMPLE_METHOD_COUNT) { + generation.sample_params.sample_method = params.sample_method; + } + if (params.scheduler != SCHEDULER_COUNT) { + generation.sample_params.scheduler = params.scheduler; + } + + sd_image_t* generated = nullptr; + int generated_count = 0; + bool success = generate_image(sd_ctx, &generation, &generated, &generated_count); + free(local_image.data); + free(local_mask.data); + if (!success || generated_count <= 0 || generated == nullptr || generated[0].data == nullptr) { + free_sd_images(generated, generated_count); + free(current.data); + return false; + } + + Mask feather = gaussian_blur_mask(masks[i], params.mask_blur); + composite_crop(¤t, generated[0], feather, crop); + free_sd_images(generated, generated_count); + ++applied; + } + + auto* outputs = static_cast(calloc(1, sizeof(sd_image_t))); + if (outputs == nullptr) { + free(current.data); + return false; + } + outputs[0] = current; + if (images_out != nullptr) { + *images_out = outputs; + } else { + free_sd_images(outputs, 1); + } + if (num_images_out != nullptr) { + *num_images_out = 1; + } + LOG_INFO("ADetailer applied %d mask(s)", applied); + return true; +} diff --git a/otherarch/sdcpp/src/detailer.h b/otherarch/sdcpp/src/detailer.h new file mode 100644 index 000000000000..17a81a1ddea9 --- /dev/null +++ b/otherarch/sdcpp/src/detailer.h @@ -0,0 +1,75 @@ +#ifndef __SD_DETAILER_H__ +#define __SD_DETAILER_H__ + +#include +#include +#include + +#include "core/ggml_extend_backend.h" +#include "model/detector/yolov8.h" +#include "model_manager.h" +#include "stable-diffusion.h" + +struct ADetailerDetection { + float x1 = 0.f; + float y1 = 0.f; + float x2 = 0.f; + float y2 = 0.f; + float confidence = 0.f; + int class_id = 0; +}; + +enum ADetailerSort { + ADETAILER_SORT_NONE, + ADETAILER_SORT_LEFT_TO_RIGHT, + ADETAILER_SORT_CENTER_TO_EDGE, + ADETAILER_SORT_AREA, +}; + +struct ADetailerParams { + const char* prompt = nullptr; + const char* negative_prompt = nullptr; + int input_size = 640; + float confidence = 0.3f; + float nms_threshold = 0.45f; + int max_detections = 100; + int mask_k_largest = 0; + float mask_min_ratio = 0.f; + float mask_max_ratio = 1.f; + int dilate_erode = 4; + int x_offset = 0; + int y_offset = 0; + bool merge_masks = false; + bool invert_mask = false; + int mask_blur = 4; + int inpaint_padding = 32; + int inpaint_width = 512; + int inpaint_height = 512; + float denoising_strength = 0.4f; + int steps = 0; + float cfg_scale = -1.f; + sample_method_t sample_method = SAMPLE_METHOD_COUNT; + scheduler_t scheduler = SCHEDULER_COUNT; + ADetailerSort sort_by = ADETAILER_SORT_NONE; +}; + +struct ADetailerGGML { + SDBackendManager backend_manager; + std::shared_ptr model_manager; + std::shared_ptr detector; + std::vector class_names; + int n_threads = 1; + std::string backend_spec; + std::string params_backend_spec; + + ADetailerGGML(int n_threads, + std::string backend_spec, + std::string params_backend_spec); + ~ADetailerGGML(); + + bool load_from_file(const std::string& detector_path); + std::vector predict(sd_image_t image, + const ADetailerParams& params); +}; + +#endif // __SD_DETAILER_H__ diff --git a/otherarch/sdcpp/src/model/common/rope.hpp b/otherarch/sdcpp/src/model/common/rope.hpp index 04ecde3c59c1..d4142c181f91 100644 --- a/otherarch/sdcpp/src/model/common/rope.hpp +++ b/otherarch/sdcpp/src/model/common/rope.hpp @@ -18,10 +18,6 @@ namespace Rope { DECREASE, }; - __STATIC_INLINE__ RefIndexMode ref_index_mode_from_bool(bool increase_ref_index) { - return increase_ref_index ? RefIndexMode::INCREASE : RefIndexMode::FIXED; - } - template __STATIC_INLINE__ std::vector linspace(T start, T end, int num) { std::vector result(num); diff --git a/otherarch/sdcpp/src/model/detector/yolov8.h b/otherarch/sdcpp/src/model/detector/yolov8.h new file mode 100644 index 000000000000..a90fdf9a6d1e --- /dev/null +++ b/otherarch/sdcpp/src/model/detector/yolov8.h @@ -0,0 +1,362 @@ +#ifndef __SD_MODEL_DETECTOR_YOLOV8_H__ +#define __SD_MODEL_DETECTOR_YOLOV8_H__ + +#include +#include +#include +#include +#include +#include + +#include "core/ggml_extend.hpp" +#include "core/util.h" + +struct YOLOv8Config { + std::array out_channels{}; + std::map hidden_channels; + std::map repeats; + int detect_box_channels = 0; + int detect_cls_channels = 0; + int reg_max = 0; + int num_classes = 0; + bool valid = false; + + static YOLOv8Config detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix = "") { + YOLOv8Config config; + auto full_name = [&](const std::string& name) { + return prefix.empty() ? name : prefix + "." + name; + }; + auto find_weight = [&](const std::string& name) -> const TensorStorage* { + auto iter = tensor_storage_map.find(full_name(name)); + return iter == tensor_storage_map.end() ? nullptr : &iter->second; + }; + auto conv_out = [&](const std::string& name) -> int { + const TensorStorage* weight = find_weight(name); + return weight != nullptr && weight->n_dims == 4 ? static_cast(weight->ne[3]) : 0; + }; + + for (int layer : {0, 1, 3, 5, 7, 16, 19}) { + config.out_channels[layer] = conv_out("model." + std::to_string(layer) + ".conv.weight"); + } + for (int layer : {2, 4, 6, 8, 12, 15, 18, 21}) { + const std::string base = "model." + std::to_string(layer); + config.out_channels[layer] = conv_out(base + ".cv2.conv.weight"); + config.hidden_channels[layer] = conv_out(base + ".cv1.conv.weight") / 2; + + int repeat_count = 0; + while (find_weight(base + ".m." + std::to_string(repeat_count) + ".cv1.conv.weight") != nullptr) { + ++repeat_count; + } + config.repeats[layer] = repeat_count; + } + config.out_channels[9] = conv_out("model.9.cv2.conv.weight"); + + config.detect_box_channels = conv_out("model.22.cv2.0.0.conv.weight"); + config.detect_cls_channels = conv_out("model.22.cv3.0.0.conv.weight"); + const int box_outputs = conv_out("model.22.cv2.0.2.weight"); + config.num_classes = conv_out("model.22.cv3.0.2.weight"); + config.reg_max = box_outputs / 4; + + config.valid = config.out_channels[0] > 0 && config.out_channels[9] > 0 && + config.out_channels[15] > 0 && config.out_channels[18] > 0 && + config.out_channels[21] > 0 && config.detect_box_channels > 0 && + config.detect_cls_channels > 0 && box_outputs > 0 && box_outputs % 4 == 0 && + config.num_classes > 0; + for (int layer : {2, 4, 6, 8, 12, 15, 18, 21}) { + config.valid = config.valid && config.hidden_channels[layer] > 0 && config.repeats[layer] > 0; + } + + if (config.valid) { + LOG_DEBUG("yolov8: classes=%d, reg_max=%d, p3=%d, p4=%d, p5=%d", + config.num_classes, + config.reg_max, + config.out_channels[15], + config.out_channels[18], + config.out_channels[21]); + } + return config; + } +}; + +class YOLOConv : public UnaryBlock { + int out_channels_ = 0; + +public: + YOLOConv(int in_channels, int out_channels, int kernel, int stride = 1) + : out_channels_(out_channels) { + blocks["conv"] = std::shared_ptr(new Conv2d(in_channels, + out_channels, + {kernel, kernel}, + {stride, stride}, + {kernel / 2, kernel / 2}, + {1, 1}, + true)); + } + + int out_channels() const { + return out_channels_; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto conv = std::dynamic_pointer_cast(blocks["conv"]); + return ggml_silu_inplace(ctx->ggml_ctx, conv->forward(ctx, x)); + } +}; + +class YOLOBottleneck : public UnaryBlock { + bool shortcut_ = false; + +public: + YOLOBottleneck(int channels, bool shortcut) + : shortcut_(shortcut) { + blocks["cv1"] = std::shared_ptr(new YOLOConv(channels, channels, 3)); + blocks["cv2"] = std::shared_ptr(new YOLOConv(channels, channels, 3)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto cv1 = std::dynamic_pointer_cast(blocks["cv1"]); + auto cv2 = std::dynamic_pointer_cast(blocks["cv2"]); + auto out = cv2->forward(ctx, cv1->forward(ctx, x)); + return shortcut_ ? ggml_add(ctx->ggml_ctx, x, out) : out; + } +}; + +class YOLOC2f : public UnaryBlock { + int hidden_channels_ = 0; + int repeats_ = 0; + +public: + YOLOC2f(int in_channels, + int out_channels, + int hidden_channels, + int repeats, + bool shortcut) + : hidden_channels_(hidden_channels), repeats_(repeats) { + blocks["cv1"] = std::shared_ptr(new YOLOConv(in_channels, hidden_channels * 2, 1)); + blocks["cv2"] = std::shared_ptr(new YOLOConv(hidden_channels * (2 + repeats), out_channels, 1)); + for (int i = 0; i < repeats; ++i) { + blocks["m." + std::to_string(i)] = std::shared_ptr(new YOLOBottleneck(hidden_channels, shortcut)); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto cv1 = std::dynamic_pointer_cast(blocks["cv1"]); + auto cv2 = std::dynamic_pointer_cast(blocks["cv2"]); + auto split = cv1->forward(ctx, x); + + // split: [N, 2*C, H, W], ggml layout [W, H, 2*C, N]. + auto y0 = ggml_view_4d(ctx->ggml_ctx, + split, + split->ne[0], + split->ne[1], + hidden_channels_, + split->ne[3], + split->nb[1], + split->nb[2], + split->nb[3], + 0); + auto y1 = ggml_view_4d(ctx->ggml_ctx, + split, + split->ne[0], + split->ne[1], + hidden_channels_, + split->ne[3], + split->nb[1], + split->nb[2], + split->nb[3], + static_cast(hidden_channels_) * split->nb[2]); + auto joined = ggml_concat(ctx->ggml_ctx, y0, y1, 2); + auto last = y1; + for (int i = 0; i < repeats_; ++i) { + auto block = std::dynamic_pointer_cast(blocks["m." + std::to_string(i)]); + last = block->forward(ctx, last); + joined = ggml_concat(ctx->ggml_ctx, joined, last, 2); + } + return cv2->forward(ctx, joined); + } +}; + +class YOLOSPPF : public UnaryBlock { +public: + YOLOSPPF(int in_channels, int out_channels) { + blocks["cv1"] = std::shared_ptr(new YOLOConv(in_channels, in_channels / 2, 1)); + blocks["cv2"] = std::shared_ptr(new YOLOConv(in_channels * 2, out_channels, 1)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto cv1 = std::dynamic_pointer_cast(blocks["cv1"]); + auto cv2 = std::dynamic_pointer_cast(blocks["cv2"]); + x = cv1->forward(ctx, x); + auto y1 = ggml_pool_2d(ctx->ggml_ctx, x, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2); + auto y2 = ggml_pool_2d(ctx->ggml_ctx, y1, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2); + auto y3 = ggml_pool_2d(ctx->ggml_ctx, y2, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2); + auto out = ggml_concat(ctx->ggml_ctx, x, y1, 2); + out = ggml_concat(ctx->ggml_ctx, out, y2, 2); + out = ggml_concat(ctx->ggml_ctx, out, y3, 2); + return cv2->forward(ctx, out); + } +}; + +class YOLODetect : public GGMLBlock { + int num_classes_ = 0; + int reg_max_ = 0; + +public: + YOLODetect(const std::array& in_channels, + int box_channels, + int cls_channels, + int reg_max, + int num_classes) + : num_classes_(num_classes), reg_max_(reg_max) { + for (int i = 0; i < 3; ++i) { + const std::string box = "cv2." + std::to_string(i); + blocks[box + ".0"] = std::shared_ptr(new YOLOConv(in_channels[i], box_channels, 3)); + blocks[box + ".1"] = std::shared_ptr(new YOLOConv(box_channels, box_channels, 3)); + blocks[box + ".2"] = std::shared_ptr(new Conv2d(box_channels, reg_max * 4, {1, 1}, {1, 1}, {0, 0}, {1, 1}, true)); + + const std::string cls = "cv3." + std::to_string(i); + blocks[cls + ".0"] = std::shared_ptr(new YOLOConv(in_channels[i], cls_channels, 3)); + blocks[cls + ".1"] = std::shared_ptr(new YOLOConv(cls_channels, cls_channels, 3)); + blocks[cls + ".2"] = std::shared_ptr(new Conv2d(cls_channels, num_classes, {1, 1}, {1, 1}, {0, 0}, {1, 1}, true)); + } + } + + ggml_tensor* forward_scale(GGMLRunnerContext* ctx, ggml_tensor* x, int index) { + const std::string box = "cv2." + std::to_string(index); + auto box0 = std::dynamic_pointer_cast(blocks[box + ".0"]); + auto box1 = std::dynamic_pointer_cast(blocks[box + ".1"]); + auto box2 = std::dynamic_pointer_cast(blocks[box + ".2"]); + + const std::string cls = "cv3." + std::to_string(index); + auto cls0 = std::dynamic_pointer_cast(blocks[cls + ".0"]); + auto cls1 = std::dynamic_pointer_cast(blocks[cls + ".1"]); + auto cls2 = std::dynamic_pointer_cast(blocks[cls + ".2"]); + + auto boxes = box2->forward(ctx, box1->forward(ctx, box0->forward(ctx, x))); + auto classes = cls2->forward(ctx, cls1->forward(ctx, cls0->forward(ctx, x))); + return ggml_concat(ctx->ggml_ctx, boxes, classes, 2); + } + + int output_channels() const { + return reg_max_ * 4 + num_classes_; + } +}; + +class YOLOv8Model : public GGMLBlock { + YOLOv8Config config_; + + std::shared_ptr make_c2f(int layer, int in_channels, bool shortcut) { + return std::make_shared(in_channels, + config_.out_channels[layer], + config_.hidden_channels.at(layer), + config_.repeats.at(layer), + shortcut); + } + +public: + explicit YOLOv8Model(YOLOv8Config config) + : config_(std::move(config)) { + blocks["model.0"] = std::make_shared(3, config_.out_channels[0], 3, 2); + blocks["model.1"] = std::make_shared(config_.out_channels[0], config_.out_channels[1], 3, 2); + blocks["model.2"] = make_c2f(2, config_.out_channels[1], true); + blocks["model.3"] = std::make_shared(config_.out_channels[2], config_.out_channels[3], 3, 2); + blocks["model.4"] = make_c2f(4, config_.out_channels[3], true); + blocks["model.5"] = std::make_shared(config_.out_channels[4], config_.out_channels[5], 3, 2); + blocks["model.6"] = make_c2f(6, config_.out_channels[5], true); + blocks["model.7"] = std::make_shared(config_.out_channels[6], config_.out_channels[7], 3, 2); + blocks["model.8"] = make_c2f(8, config_.out_channels[7], true); + blocks["model.9"] = std::make_shared(config_.out_channels[8], config_.out_channels[9]); + + blocks["model.12"] = make_c2f(12, config_.out_channels[9] + config_.out_channels[6], false); + blocks["model.15"] = make_c2f(15, config_.out_channels[12] + config_.out_channels[4], false); + blocks["model.16"] = std::make_shared(config_.out_channels[15], config_.out_channels[16], 3, 2); + blocks["model.18"] = make_c2f(18, config_.out_channels[16] + config_.out_channels[12], false); + blocks["model.19"] = std::make_shared(config_.out_channels[18], config_.out_channels[19], 3, 2); + blocks["model.21"] = make_c2f(21, config_.out_channels[19] + config_.out_channels[9], false); + blocks["model.22"] = std::make_shared( + std::array{config_.out_channels[15], config_.out_channels[18], config_.out_channels[21]}, + config_.detect_box_channels, + config_.detect_cls_channels, + config_.reg_max, + config_.num_classes); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto run = [&](int layer, ggml_tensor* input) { + return std::dynamic_pointer_cast(blocks["model." + std::to_string(layer)])->forward(ctx, input); + }; + + auto x0 = run(0, x); + auto x1 = run(1, x0); + auto x2 = run(2, x1); + auto x3 = run(3, x2); + auto x4 = run(4, x3); + auto x5 = run(5, x4); + auto x6 = run(6, x5); + auto x7 = run(7, x6); + auto x8 = run(8, x7); + auto x9 = run(9, x8); + + auto x12 = run(12, ggml_concat(ctx->ggml_ctx, ggml_upscale(ctx->ggml_ctx, x9, 2, GGML_SCALE_MODE_NEAREST), x6, 2)); + auto x15 = run(15, ggml_concat(ctx->ggml_ctx, ggml_upscale(ctx->ggml_ctx, x12, 2, GGML_SCALE_MODE_NEAREST), x4, 2)); + auto x16 = run(16, x15); + auto x18 = run(18, ggml_concat(ctx->ggml_ctx, x16, x12, 2)); + auto x19 = run(19, x18); + auto x21 = run(21, ggml_concat(ctx->ggml_ctx, x19, x9, 2)); + + auto detect = std::dynamic_pointer_cast(blocks["model.22"]); + auto p3 = detect->forward_scale(ctx, x15, 0); + auto p4 = detect->forward_scale(ctx, x18, 1); + auto p5 = detect->forward_scale(ctx, x21, 2); + p3 = ggml_reshape_2d(ctx->ggml_ctx, p3, p3->ne[0] * p3->ne[1], detect->output_channels()); + p4 = ggml_reshape_2d(ctx->ggml_ctx, p4, p4->ne[0] * p4->ne[1], detect->output_channels()); + p5 = ggml_reshape_2d(ctx->ggml_ctx, p5, p5->ne[0] * p5->ne[1], detect->output_channels()); + return ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, p3, p4, 0), p5, 0); + } +}; + +struct YOLOv8Runner : public GGMLRunner { + YOLOv8Config config; + std::unique_ptr model; + + YOLOv8Runner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), + config(YOLOv8Config::detect_from_weights(tensor_storage_map)) { + if (config.valid) { + model = std::make_unique(config); + model->init(params_ctx, tensor_storage_map, ""); + } + } + + std::string get_desc() override { + return "yolov8"; + } + + void get_param_tensors(std::map& tensors) { + if (model) { + model->get_param_tensors(tensors); + } + } + + ggml_cgraph* build_graph(const sd::Tensor& input) { + if (!model) { + return nullptr; + } + ggml_cgraph* graph = new_graph_custom(1 << 16); + auto x = make_input(input); + auto runner_ctx = get_context(); + auto output = model->forward(&runner_ctx, x); + ggml_build_forward_expand(graph, output); + return graph; + } + + sd::Tensor compute(int n_threads, const sd::Tensor& input) { + auto get_graph = [&]() { return build_graph(input); }; + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false, false, false)); + } +}; + +#endif // __SD_MODEL_DETECTOR_YOLOV8_H__ diff --git a/otherarch/sdcpp/src/model/diffusion/anima.hpp b/otherarch/sdcpp/src/model/diffusion/anima.hpp index 1c1ec2ed23b1..4fe7e4650e1e 100644 --- a/otherarch/sdcpp/src/model/diffusion/anima.hpp +++ b/otherarch/sdcpp/src/model/diffusion/anima.hpp @@ -484,10 +484,11 @@ namespace Anima { ggml_tensor* timestep, ggml_tensor* encoder_hidden_states, ggml_tensor* image_pe, - ggml_tensor* t5_ids = nullptr, - ggml_tensor* t5_weights = nullptr, - ggml_tensor* adapter_q_pe = nullptr, - ggml_tensor* adapter_k_pe = nullptr) { + ggml_tensor* t5_ids = nullptr, + ggml_tensor* t5_weights = nullptr, + ggml_tensor* adapter_q_pe = nullptr, + ggml_tensor* adapter_k_pe = nullptr, + std::vector ref_latents = {}) { GGML_ASSERT(x->ne[3] == 1); auto x_embedder = std::dynamic_pointer_cast(blocks["x_embedder"]); @@ -502,8 +503,16 @@ namespace Anima { auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]); x = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2); // [N, C + 1, H, W] - x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw] - + x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw] + int64_t img_len = x->ne[1]; + if (ref_latents.size() > 0) { + for (ggml_tensor* ref : ref_latents) { + auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, ref->ne[0], ref->ne[1], 1, ref->ne[3]); + ref = ggml_concat(ctx->ggml_ctx, ref, padding_mask, 2); // [N, C + 1, H, W] + ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size); + x = ggml_concat(ctx->ggml_ctx, x, ref, 1); + } + } x = x_embedder->forward(ctx, x); auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast(config.hidden_size)); @@ -543,6 +552,7 @@ namespace Anima { x = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe); sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x"); } + x = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, img_len); x = final_layer->forward(ctx, x, embedded_timestep, temb); // [N, h*w, ph*pw*C] @@ -602,8 +612,8 @@ namespace Anima { const std::vector& axes_dim, float h_extrapolation_ratio, float w_extrapolation_ratio, - float t_extrapolation_ratio) { - static const std::vector empty_ref_latents; + float t_extrapolation_ratio, + const std::vector& ref_latents) { auto ids = Rope::gen_flux_ids(h, w, patch_size, @@ -611,7 +621,7 @@ namespace Anima { static_cast(axes_dim.size()), 0, {}, - empty_ref_latents, + ref_latents, Rope::RefIndexMode::FIXED, 1.0f, false); @@ -626,14 +636,20 @@ namespace Anima { ggml_cgraph* build_graph(const sd::Tensor& x_tensor, const sd::Tensor& timesteps_tensor, - const sd::Tensor& context_tensor = {}, - const sd::Tensor& t5_ids_tensor = {}, - const sd::Tensor& t5_weights_tensor = {}) { + const sd::Tensor& context_tensor = {}, + const sd::Tensor& t5_ids_tensor = {}, + const sd::Tensor& t5_weights_tensor = {}, + const std::vector>& ref_latents_tensor = {}) { ggml_tensor* x = make_input(x_tensor); ggml_tensor* timesteps = make_input(timesteps_tensor); ggml_tensor* context = make_optional_input(context_tensor); ggml_tensor* t5_ids = make_optional_input(t5_ids_tensor); ggml_tensor* t5_weights = make_optional_input(t5_weights_tensor); + std::vector ref_latents; + ref_latents.reserve(ref_latents_tensor.size()); + for (const auto& ref_latent_tensor : ref_latents_tensor) { + ref_latents.push_back(make_input(ref_latent_tensor)); + } GGML_ASSERT(x->ne[3] == 1); ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE); @@ -650,7 +666,8 @@ namespace Anima { config.axes_dim, 4.0f, 4.0f, - 1.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()); @@ -682,7 +699,8 @@ namespace Anima { t5_ids, t5_weights, adapter_q_pe, - adapter_k_pe); + adapter_k_pe, + ref_latents); ggml_build_forward_expand(gf, out); return gf; @@ -691,11 +709,13 @@ namespace Anima { sd::Tensor compute(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, - const sd::Tensor& context = {}, - const sd::Tensor& t5_ids = {}, - const sd::Tensor& t5_weights = {}) { + const sd::Tensor& context = {}, + const sd::Tensor& t5_ids = {}, + const sd::Tensor& t5_weights = {}, + const std::vector>& ref_latents = {}, + const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("cosmos_reference")) { auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, t5_ids, t5_weights); + return build_graph(x, timesteps, context, t5_ids, t5_weights, ref_latents); }; return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); } @@ -705,12 +725,15 @@ namespace Anima { GGML_ASSERT(diffusion_params.x != nullptr); GGML_ASSERT(diffusion_params.timesteps != nullptr); const auto* extra = diffusion_extra_as(diffusion_params); + static const std::vector> empty_ref_latents; return compute(n_threads, *diffusion_params.x, *diffusion_params.timesteps, tensor_or_empty(diffusion_params.context), tensor_or_empty(extra->t5_ids), - tensor_or_empty(extra->t5_weights)); + tensor_or_empty(extra->t5_weights), + diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents, + diffusion_params.ref_image_params); } }; } // namespace Anima diff --git a/otherarch/sdcpp/src/model/diffusion/animatediff.hpp b/otherarch/sdcpp/src/model/diffusion/animatediff.hpp new file mode 100644 index 000000000000..d470544fe8ea --- /dev/null +++ b/otherarch/sdcpp/src/model/diffusion/animatediff.hpp @@ -0,0 +1,182 @@ +#ifndef __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__ +#define __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__ + +#include "core/ggml_extend.hpp" +#include "model/common/block.hpp" + +// AnimateDiff (https://arxiv.org/abs/2307.04725) SD 1.5 motion modules. +namespace AnimateDiff { + + struct MotionModuleConfig { + int max_frames = 32; + int64_t num_heads = 8; + int norm_num_groups = 32; + std::vector down_channels = {320, 640, 1280, 1280}; + std::vector up_channels = {1280, 1280, 640, 320}; + int num_down_motion_per_block = 2; + int num_up_motion_per_block = 3; + bool enable_mid_block = false; + int64_t mid_channels = 1280; + }; + + class TemporalAttention : public GGMLBlock { + protected: + int64_t channels; + int64_t num_heads; + int max_frames; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + params["pos_encoder.pe"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, channels, max_frames, 1); + } + + public: + TemporalAttention(int64_t channels, int64_t num_heads, int max_frames) + : channels(channels), num_heads(num_heads), max_frames(max_frames) { + blocks["to_q"] = std::shared_ptr(new Linear(channels, channels, false)); + blocks["to_k"] = std::shared_ptr(new Linear(channels, channels, false)); + blocks["to_v"] = std::shared_ptr(new Linear(channels, channels, false)); + blocks["to_out.0"] = std::shared_ptr(new Linear(channels, channels, true)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto to_q = std::dynamic_pointer_cast(blocks["to_q"]); + auto to_k = std::dynamic_pointer_cast(blocks["to_k"]); + auto to_v = std::dynamic_pointer_cast(blocks["to_v"]); + auto to_out = std::dynamic_pointer_cast(blocks["to_out.0"]); + + int64_t C = x->ne[0]; + int64_t F = x->ne[1]; + + auto pe = params["pos_encoder.pe"]; + auto pe_f = (F == pe->ne[1]) + ? pe + : ggml_view_3d(ctx->ggml_ctx, pe, C, F, 1, pe->nb[1], pe->nb[2], 0); + auto x_pe = ggml_add(ctx->ggml_ctx, x, ggml_repeat(ctx->ggml_ctx, pe_f, x)); + + auto q = to_q->forward(ctx, x_pe); + auto k = to_k->forward(ctx, x_pe); + auto v = to_v->forward(ctx, x_pe); + + auto a = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, (int)num_heads, nullptr, false); + return to_out->forward(ctx, a); + } + }; + + class TemporalTransformerBlock : public GGMLBlock { + public: + TemporalTransformerBlock(int64_t channels, int64_t num_heads, int max_frames) { + blocks["attention_blocks.0"] = std::make_shared(channels, num_heads, max_frames); + blocks["attention_blocks.1"] = std::make_shared(channels, num_heads, max_frames); + blocks["norms.0"] = std::shared_ptr(new LayerNorm(channels)); + blocks["norms.1"] = std::shared_ptr(new LayerNorm(channels)); + blocks["ff"] = std::make_shared(channels, channels, 4, FeedForward::Activation::GEGLU); + blocks["ff_norm"] = std::shared_ptr(new LayerNorm(channels)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto attn0 = std::dynamic_pointer_cast(blocks["attention_blocks.0"]); + auto attn1 = std::dynamic_pointer_cast(blocks["attention_blocks.1"]); + auto norm0 = std::dynamic_pointer_cast(blocks["norms.0"]); + auto norm1 = std::dynamic_pointer_cast(blocks["norms.1"]); + auto ff = std::dynamic_pointer_cast(blocks["ff"]); + auto ff_norm = std::dynamic_pointer_cast(blocks["ff_norm"]); + + auto r = x; + x = ggml_add(ctx->ggml_ctx, attn0->forward(ctx, norm0->forward(ctx, x)), r); + + r = x; + x = ggml_add(ctx->ggml_ctx, attn1->forward(ctx, norm1->forward(ctx, x)), r); + + r = x; + x = ggml_add(ctx->ggml_ctx, ff->forward(ctx, ff_norm->forward(ctx, x)), r); + + return x; + } + }; + + class TemporalTransformer : public GGMLBlock { + public: + TemporalTransformer(int64_t channels, int64_t num_heads, int norm_num_groups, int max_frames) { + blocks["norm"] = std::shared_ptr(new GroupNorm(norm_num_groups, channels)); + blocks["proj_in"] = std::shared_ptr(new Linear(channels, channels, true)); + blocks["transformer_blocks.0"] = std::make_shared(channels, num_heads, max_frames); + blocks["proj_out"] = std::shared_ptr(new Linear(channels, channels, true)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t num_frames) { + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto proj_in = std::dynamic_pointer_cast(blocks["proj_in"]); + auto tb0 = std::dynamic_pointer_cast(blocks["transformer_blocks.0"]); + auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); + + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t C = x->ne[2]; + GGML_ASSERT(x->ne[3] == num_frames); + + auto residual = x; + auto h = norm->forward(ctx, x); + + h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1)); + h = ggml_reshape_3d(ctx->ggml_ctx, h, C, num_frames, W * H); + h = proj_in->forward(ctx, h); + h = tb0->forward(ctx, h); + h = proj_out->forward(ctx, h); + h = ggml_reshape_4d(ctx->ggml_ctx, h, C, num_frames, W, H); + h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1)); + + return ggml_add(ctx->ggml_ctx, h, residual); + } + }; + + class MotionModule : public GGMLBlock { + public: + MotionModule(int64_t channels, int64_t num_heads, int norm_num_groups, int max_frames) { + blocks["temporal_transformer"] = std::make_shared(channels, num_heads, norm_num_groups, max_frames); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t num_frames) { + auto tt = std::dynamic_pointer_cast(blocks["temporal_transformer"]); + return tt->forward(ctx, x, num_frames); + } + }; + + class AnimateDiffModel : public GGMLBlock { + public: + MotionModuleConfig config; + + AnimateDiffModel(const MotionModuleConfig& cfg) + : config(cfg) { + for (int i = 0; i < static_cast(cfg.down_channels.size()); ++i) { + int64_t ch = cfg.down_channels[i]; + for (int j = 0; j < cfg.num_down_motion_per_block; ++j) { + blocks["down_blocks." + std::to_string(i) + ".motion_modules." + std::to_string(j)] = + std::make_shared(ch, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames); + } + } + for (int i = 0; i < static_cast(cfg.up_channels.size()); ++i) { + int64_t ch = cfg.up_channels[i]; + for (int j = 0; j < cfg.num_up_motion_per_block; ++j) { + blocks["up_blocks." + std::to_string(i) + ".motion_modules." + std::to_string(j)] = + std::make_shared(ch, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames); + } + } + if (cfg.enable_mid_block) { + blocks["mid_block.motion_modules.0"] = + std::make_shared(cfg.mid_channels, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames); + } + } + + std::shared_ptr motion(const std::string& key) { + auto it = blocks.find(key); + if (it == blocks.end()) + return nullptr; + return std::dynamic_pointer_cast(it->second); + } + }; + +} // namespace AnimateDiff + +#endif // __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__ diff --git a/otherarch/sdcpp/src/model/diffusion/boogu.hpp b/otherarch/sdcpp/src/model/diffusion/boogu.hpp index 9ab2dccbc9b1..7bdbcd814204 100644 --- a/otherarch/sdcpp/src/model/diffusion/boogu.hpp +++ b/otherarch/sdcpp/src/model/diffusion/boogu.hpp @@ -827,7 +827,7 @@ namespace Boogu { *diffusion_params.x, *diffusion_params.timesteps, tensor_or_empty(diffusion_params.context), - diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents); + diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents); } }; } // namespace Boogu diff --git a/otherarch/sdcpp/src/model/diffusion/flux.hpp b/otherarch/sdcpp/src/model/diffusion/flux.hpp index e5d795f74406..edd7a89b89b0 100644 --- a/otherarch/sdcpp/src/model/diffusion/flux.hpp +++ b/otherarch/sdcpp/src/model/diffusion/flux.hpp @@ -1642,8 +1642,8 @@ namespace Flux { tensor_or_empty(diffusion_params.c_concat), tensor_or_empty(diffusion_params.y), tensor_or_empty(extra->guidance), - diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents, - diffusion_params.ref_index_mode, + diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents, + diffusion_params.ref_image_params.ref_index_mode, extra->skip_layers ? *extra->skip_layers : empty_skip_layers, tensor_or_empty(extra->pulid_id), extra->pulid_id_weight); diff --git a/otherarch/sdcpp/src/model/diffusion/krea2.hpp b/otherarch/sdcpp/src/model/diffusion/krea2.hpp index 8ca0bb51ff95..9121069ad507 100644 --- a/otherarch/sdcpp/src/model/diffusion/krea2.hpp +++ b/otherarch/sdcpp/src/model/diffusion/krea2.hpp @@ -615,7 +615,8 @@ namespace Krea2 { ggml_tensor* timestep, ggml_tensor* context, ggml_tensor* pe, - std::vector ref_latents = {}) { + std::vector ref_latents = {}, + bool zero_timestep_refs = false) { int64_t W = x->ne[0]; int64_t H = x->ne[1]; int64_t N = x->ne[3]; @@ -645,7 +646,7 @@ namespace Krea2 { auto tvec = tproj->forward(ctx, t); ggml_tensor* tvec_0 = nullptr; - if (ref_latents.size() > 0) { + if (ref_latents.size() > 0 && zero_timestep_refs) { // "index_timestep_zero" mode: use timestep = 0 for ref latents auto timestep_0 = ggml_scale(ctx->ggml_ctx, timestep, 0.0f); auto t_0 = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_0, static_cast(config.timestep_dim), 10000, 1000.f); @@ -719,7 +720,7 @@ namespace Krea2 { const sd::Tensor& timesteps_tensor, const sd::Tensor& context_tensor, const std::vector>& ref_latents_tensor = {}, - Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) { + const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) { ggml_cgraph* gf = new_graph_custom(KREA2_GRAPH_SIZE); ggml_tensor* x = make_input(x_tensor); ggml_tensor* timesteps = make_input(timesteps_tensor); @@ -741,13 +742,13 @@ namespace Krea2 { config.theta, config.axes_dim, ref_latents, - ref_index_mode); + 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()); auto runner_ctx = get_context(); - ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents); + ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents, ref_image_params.force_ref_timestep_zero); ggml_build_forward_expand(gf, out); return gf; } @@ -757,9 +758,9 @@ namespace Krea2 { const sd::Tensor& timesteps, const sd::Tensor& context, const std::vector>& ref_latents = {}, - Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) { + const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) { auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, ref_latents, ref_index_mode); + return build_graph(x, timesteps, context, ref_latents, ref_image_params); }; return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); } @@ -773,8 +774,8 @@ namespace Krea2 { *diffusion_params.x, *diffusion_params.timesteps, tensor_or_empty(diffusion_params.context), - diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents, - diffusion_params.ref_index_mode); + diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents, + diffusion_params.ref_image_params); } }; } // namespace Krea2 diff --git a/otherarch/sdcpp/src/model/diffusion/model.hpp b/otherarch/sdcpp/src/model/diffusion/model.hpp index 8ef00023eccf..6e04f9f1f5be 100644 --- a/otherarch/sdcpp/src/model/diffusion/model.hpp +++ b/otherarch/sdcpp/src/model/diffusion/model.hpp @@ -10,6 +10,36 @@ #include "model/common/rope.hpp" #include "model_manager.h" +enum class RefImageResizeMode { + NONE, + LONGEST_SIDE, + AREA, +}; + +struct RefImageParams { + bool pass_to_vlm = false; + bool pass_to_dit = true; + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED; + bool force_ref_timestep_zero = false; + bool resize_before_vae = true; + int vae_input_max_pixels = -1; + RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA; + int vlm_min_size = -1; + int vlm_max_size = -1; +}; + +const std::unordered_map REF_IMAGE_PRESETS = { + {"flux_kontext", {false, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::NONE, -1, -1}}, + {"longcat", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}}, + {"flux2", {false, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::NONE, -1, -1}}, + {"qwen", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}}, + {"qwen_layered", {true, true, Rope::RefIndexMode::DECREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}}, + {"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}}, + {"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}}, +}; + struct UNetDiffusionExtra { int num_video_frames = -1; const std::vector>* controls = nullptr; @@ -74,7 +104,7 @@ struct DiffusionParams { const sd::Tensor* c_concat = nullptr; const sd::Tensor* y = nullptr; const std::vector>* ref_latents = nullptr; - Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED; + RefImageParams ref_image_params = {false, false, Rope::RefIndexMode::FIXED, false}; DiffusionExtraParams extra = std::monostate{}; }; diff --git a/otherarch/sdcpp/src/model/diffusion/pid.hpp b/otherarch/sdcpp/src/model/diffusion/pid.hpp index 68dca00f815d..2a698e0a68ff 100644 --- a/otherarch/sdcpp/src/model/diffusion/pid.hpp +++ b/otherarch/sdcpp/src/model/diffusion/pid.hpp @@ -17,30 +17,38 @@ namespace Pid { constexpr float PID_PI = 3.14159265358979323846f; struct PixelDiTConfig { - int64_t in_channels = 3; - int64_t hidden_size = 1536; - int64_t num_groups = 24; - int64_t patch_mlp_hidden_dim = 4096; - int64_t pixel_hidden_size = 16; - int64_t pixel_attn_hidden_size = 1152; - int64_t pixel_num_groups = 16; - int64_t patch_depth = 14; - int64_t pixel_depth = 2; - int64_t patch_size = 16; - int64_t txt_embed_dim = 2304; - int64_t txt_max_length = 300; - float text_rope_theta = 10000.f; - int64_t lq_latent_channels = 16; - int64_t lq_hidden_dim = 512; - int64_t lq_num_res_blocks = 4; - int64_t lq_interval = 2; - int64_t lq_sr_scale = 4; - int64_t lq_latent_down_factor = 8; - int64_t rope_ref_grid_h = 64; - int64_t rope_ref_grid_w = 64; + int64_t in_channels = 3; + int64_t hidden_size = 1536; + int64_t num_groups = 24; + int64_t patch_mlp_hidden_dim = 4096; + int64_t pixel_hidden_size = 16; + int64_t pixel_attn_hidden_size = 1152; + int64_t pixel_num_groups = 16; + int64_t patch_depth = 14; + int64_t pixel_depth = 2; + int64_t patch_size = 16; + int64_t txt_embed_dim = 2304; + int64_t txt_max_length = 300; + float text_rope_theta = 10000.f; + int64_t lq_latent_channels = 16; + int64_t lq_hidden_dim = 512; + int64_t lq_num_res_blocks = 4; + int64_t lq_interval = 2; + int64_t lq_sr_scale = 4; + int64_t lq_latent_down_factor = 8; + int64_t lq_latent_unpatchify_factor = 1; + bool lq_replicate_padding = false; + bool lq_gate_per_token = false; + bool pit_lq_inject = false; + int64_t rope_ref_grid_h = 64; + int64_t rope_ref_grid_w = 64; static PixelDiTConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { PixelDiTConfig config; + int64_t latent_proj_in_channels = config.lq_latent_channels; + int64_t num_lq_gates = 0; + const std::string lq_prefix = prefix + ".lq_proj."; + config.pit_lq_inject = tensor_storage_map.find(lq_prefix + "pit_head.weight") != tensor_storage_map.end(); for (const auto& [name, tensor_storage] : tensor_storage_map) { if (!starts_with(name, prefix)) { continue; @@ -61,20 +69,56 @@ namespace Pid { config.pixel_depth = std::max(config.pixel_depth, block_index + 1); } } - if (name.find("lq_proj.latent_proj.0.weight") != std::string::npos) { - config.lq_latent_channels = tensor_storage.ne[2]; - config.lq_latent_down_factor = config.lq_latent_channels >= 64 ? 16 : 8; + if (name == lq_prefix + "latent_proj.0.weight") { + latent_proj_in_channels = tensor_storage.ne[2]; + config.lq_hidden_dim = tensor_storage.ne[3]; + } + if (starts_with(name, lq_prefix + "gate_modules.")) { + auto items = split_string(name.substr(lq_prefix.size()), '.'); + if (items.size() > 1) { + int gate_index = atoi(items[1].c_str()); + num_lq_gates = std::max(num_lq_gates, gate_index + 1); + } } if (name.find("patch_blocks.0.mlp_x.w1.weight") != std::string::npos) { config.patch_mlp_hidden_dim = tensor_storage.ne[1]; } } - LOG_DEBUG("pid: patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_latent_down_factor = %" PRId64, + if (num_lq_gates > 0) { + config.lq_interval = (config.patch_depth + num_lq_gates - 1) / num_lq_gates; + } + if (config.pit_lq_inject) { + if (latent_proj_in_channels == 16) { + config.lq_latent_channels = 16; + config.lq_latent_down_factor = 8; + config.lq_latent_unpatchify_factor = 1; + } else { + GGML_ASSERT(latent_proj_in_channels == 32); + config.lq_latent_channels = 128; + config.lq_latent_down_factor = 16; + config.lq_latent_unpatchify_factor = 2; + } + auto gate_weight = tensor_storage_map.find(lq_prefix + "gate_modules.0.content_proj.weight"); + if (gate_weight != tensor_storage_map.end()) { + config.lq_gate_per_token = gate_weight->second.ne[1] == 1; + } + config.lq_replicate_padding = true; + config.rope_ref_grid_h = 128; + config.rope_ref_grid_w = 128; + } else { + config.lq_latent_channels = latent_proj_in_channels; + config.lq_latent_down_factor = latent_proj_in_channels >= 64 ? 16 : 8; + } + LOG_DEBUG("pid: version = %s, patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_hidden_dim = %" PRId64 ", lq_latent_down_factor = %" PRId64 ", lq_latent_unpatchify_factor = %" PRId64 ", lq_interval = %" PRId64, + config.pit_lq_inject ? "1.5" : "1", config.patch_depth, config.pixel_depth, config.patch_mlp_hidden_dim, config.lq_latent_channels, - config.lq_latent_down_factor); + config.lq_hidden_dim, + config.lq_latent_down_factor, + config.lq_latent_unpatchify_factor, + config.lq_interval); return config; } }; @@ -135,6 +179,18 @@ namespace Pid { return ggml_add(ctx, ggml_add(ctx, x, ggml_mul(ctx, x, scale)), shift); } + inline ggml_tensor* replicate_pad_2d(ggml_context* ctx, ggml_tensor* x) { + auto left = ggml_ext_slice(ctx, x, 0, 0, 1); + auto right = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]); + x = ggml_concat(ctx, left, x, 0); + x = ggml_concat(ctx, x, right, 0); + + auto top = ggml_ext_slice(ctx, x, 1, 0, 1); + auto bottom = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]); + x = ggml_concat(ctx, top, x, 1); + return ggml_concat(ctx, x, bottom, 1); + } + struct PatchTokenEmbedder : public GGMLBlock { bool use_rms_norm; @@ -457,9 +513,9 @@ namespace Pid { struct SigmaAwareGate : public GGMLBlock { int64_t dim; - SigmaAwareGate(int64_t dim) + SigmaAwareGate(int64_t dim, bool per_token = false) : dim(dim) { - blocks["content_proj"] = std::make_shared(dim * 2, dim, true); + blocks["content_proj"] = std::make_shared(dim * 2, per_token ? 1 : dim, true); } void init_params(ggml_context* ctx, @@ -479,16 +535,20 @@ namespace Pid { auto alpha = ggml_exp(ctx->ggml_ctx, params["log_alpha"]); auto offset = ggml_neg(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, alpha, sigma)); auto gate = ggml_sigmoid(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, content_logit, offset)); - return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate, lq)); + return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, lq, gate)); } }; struct PiDResBlock : public GGMLBlock { - PiDResBlock(int64_t channels) { - blocks["block.0"] = std::make_shared(4, channels, 1e-5f); - blocks["block.2"] = std::make_shared(channels, channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); - blocks["block.3"] = std::make_shared(4, channels, 1e-5f); - blocks["block.5"] = std::make_shared(channels, channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + bool replicate_padding; + + PiDResBlock(int64_t channels, bool replicate_padding = false) + : replicate_padding(replicate_padding) { + std::pair padding = replicate_padding ? std::pair{0, 0} : std::pair{1, 1}; + blocks["block.0"] = std::make_shared(4, channels, 1e-5f); + blocks["block.2"] = std::make_shared(channels, channels, std::pair{3, 3}, std::pair{1, 1}, padding); + blocks["block.3"] = std::make_shared(4, channels, 1e-5f); + blocks["block.5"] = std::make_shared(channels, channels, std::pair{3, 3}, std::pair{1, 1}, padding); } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { @@ -497,9 +557,15 @@ namespace Pid { auto norm2 = std::dynamic_pointer_cast(blocks["block.3"]); auto conv2 = std::dynamic_pointer_cast(blocks["block.5"]); auto h = ggml_silu_inplace(ctx->ggml_ctx, norm1->forward(ctx, x)); - h = conv1->forward(ctx, h); - h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h)); - h = conv2->forward(ctx, h); + if (replicate_padding) { + h = replicate_pad_2d(ctx->ggml_ctx, h); + } + h = conv1->forward(ctx, h); + h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h)); + if (replicate_padding) { + h = replicate_pad_2d(ctx->ggml_ctx, h); + } + h = conv2->forward(ctx, h); return ggml_add(ctx->ggml_ctx, x, h); } }; @@ -509,16 +575,23 @@ namespace Pid { LQProjection2D(const PixelDiTConfig& config) : config(config) { - blocks["latent_proj.0"] = std::make_shared(config.lq_latent_channels, config.lq_hidden_dim, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); - blocks["latent_proj.2"] = std::make_shared(config.lq_hidden_dim, config.lq_hidden_dim, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + int64_t unpatchify_area = config.lq_latent_unpatchify_factor * config.lq_latent_unpatchify_factor; + GGML_ASSERT(config.lq_latent_channels % unpatchify_area == 0); + int64_t latent_proj_in_channels = config.lq_latent_channels / unpatchify_area; + std::pair padding = config.lq_replicate_padding ? std::pair{0, 0} : std::pair{1, 1}; + blocks["latent_proj.0"] = std::make_shared(latent_proj_in_channels, config.lq_hidden_dim, std::pair{3, 3}, std::pair{1, 1}, padding); + blocks["latent_proj.2"] = std::make_shared(config.lq_hidden_dim, config.lq_hidden_dim, std::pair{3, 3}, std::pair{1, 1}, padding); for (int i = 0; i < config.lq_num_res_blocks; ++i) { - blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared(config.lq_hidden_dim); + blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared(config.lq_hidden_dim, config.lq_replicate_padding); } int num_outputs = static_cast((config.patch_depth + config.lq_interval - 1) / config.lq_interval); for (int i = 0; i < num_outputs; ++i) { blocks["output_heads." + std::to_string(i)] = std::make_shared(config.lq_hidden_dim, config.hidden_size, true); - blocks["gate_modules." + std::to_string(i)] = std::make_shared(config.hidden_size); + blocks["gate_modules." + std::to_string(i)] = std::make_shared(config.hidden_size, config.lq_gate_per_token); + } + if (config.pit_lq_inject) { + blocks["pit_head"] = std::make_shared(config.lq_hidden_dim, config.hidden_size, true); } } @@ -543,9 +616,29 @@ namespace Pid { ggml_tensor* lq_latent, int64_t target_pH, int64_t target_pW) { - auto conv0 = std::dynamic_pointer_cast(blocks["latent_proj.0"]); - auto conv2 = std::dynamic_pointer_cast(blocks["latent_proj.2"]); - float z_to_patch_ratio = static_cast(config.lq_sr_scale * config.lq_latent_down_factor) / + auto conv0 = std::dynamic_pointer_cast(blocks["latent_proj.0"]); + auto conv2 = std::dynamic_pointer_cast(blocks["latent_proj.2"]); + int64_t unpatchify_factor = config.lq_latent_unpatchify_factor; + if (unpatchify_factor > 1) { + int64_t latent_h = lq_latent->ne[1]; + int64_t latent_w = lq_latent->ne[0]; + lq_latent = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, lq_latent, 2, 0, 1, 3)); + lq_latent = ggml_reshape_3d(ctx->ggml_ctx, + lq_latent, + lq_latent->ne[0], + lq_latent->ne[1] * lq_latent->ne[2], + lq_latent->ne[3]); + lq_latent = DiT::unpatchify(ctx->ggml_ctx, + lq_latent, + latent_h, + latent_w, + static_cast(unpatchify_factor), + static_cast(unpatchify_factor), + true); + } + + int64_t effective_down_factor = config.lq_latent_down_factor / unpatchify_factor; + float z_to_patch_ratio = static_cast(config.lq_sr_scale * effective_down_factor) / static_cast(config.patch_size); GGML_ASSERT(z_to_patch_ratio >= 1.0f); if (lq_latent->ne[0] != target_pW || lq_latent->ne[1] != target_pH) { @@ -558,9 +651,15 @@ namespace Pid { GGML_SCALE_MODE_NEAREST); } + if (config.lq_replicate_padding) { + lq_latent = replicate_pad_2d(ctx->ggml_ctx, lq_latent); + } auto feat = conv0->forward(ctx, lq_latent); feat = ggml_silu_inplace(ctx->ggml_ctx, feat); - feat = conv2->forward(ctx, feat); + if (config.lq_replicate_padding) { + feat = replicate_pad_2d(ctx->ggml_ctx, feat); + } + feat = conv2->forward(ctx, feat); for (int i = 0; i < config.lq_num_res_blocks; ++i) { auto block = std::dynamic_pointer_cast(blocks["latent_proj." + std::to_string(3 + i)]); feat = block->forward(ctx, feat); @@ -574,11 +673,15 @@ namespace Pid { int num_outputs = static_cast((config.patch_depth + config.lq_interval - 1) / config.lq_interval); std::vector outputs; - outputs.reserve(num_outputs); + outputs.reserve(num_outputs + (config.pit_lq_inject ? 1 : 0)); for (int i = 0; i < num_outputs; ++i) { auto head = std::dynamic_pointer_cast(blocks["output_heads." + std::to_string(i)]); outputs.push_back(head->forward(ctx, tokens)); } + if (config.pit_lq_inject) { + auto pit_head = std::dynamic_pointer_cast(blocks["pit_head"]); + outputs.push_back(pit_head->forward(ctx, tokens)); + } return outputs; } }; @@ -606,6 +709,9 @@ namespace Pid { } blocks["final_layer"] = std::make_shared(config.pixel_hidden_size, config.in_channels); blocks["lq_proj"] = std::make_shared(config); + if (config.pit_lq_inject) { + blocks["pit_lq_gate"] = std::make_shared(config.hidden_size, config.lq_gate_per_token); + } } void init_params(ggml_context* ctx, @@ -654,6 +760,11 @@ namespace Pid { y_emb = ggml_add(ctx->ggml_ctx, y_emb, y_pos); std::vector lq_features = lq_proj->forward(ctx, lq_latent, Hs, Ws); + ggml_tensor* pit_lq_feature = nullptr; + if (config.pit_lq_inject) { + pit_lq_feature = lq_features.back(); + lq_features.pop_back(); + } auto s = s_embedder->forward(ctx, x_patches); @@ -677,6 +788,10 @@ namespace Pid { sd::ggml_graph_cut::mark_graph_cut(y_emb, "pid.patch_blocks." + std::to_string(i), "y"); } s = ggml_silu(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, s, t_emb)); + if (pit_lq_feature != nullptr) { + auto pit_lq_gate = std::dynamic_pointer_cast(blocks["pit_lq_gate"]); + s = pit_lq_gate->forward(ctx, s, pit_lq_feature, degrade_sigma); + } auto s_cond = ggml_reshape_2d(ctx->ggml_ctx, s, config.hidden_size, L * B); auto pixels = pixel_embedder->forward(ctx, x, config.patch_size, pixel_pos_full); diff --git a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp index a7c946fd918e..c6c3622f9650 100644 --- a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp @@ -715,8 +715,8 @@ namespace Qwen { *diffusion_params.x, *diffusion_params.timesteps, tensor_or_empty(diffusion_params.context), - diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents, - diffusion_params.ref_index_mode); + diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents, + diffusion_params.ref_image_params.ref_index_mode); } void test() { diff --git a/otherarch/sdcpp/src/model/diffusion/unet.hpp b/otherarch/sdcpp/src/model/diffusion/unet.hpp index 253b3b4b298c..118ed6327663 100644 --- a/otherarch/sdcpp/src/model/diffusion/unet.hpp +++ b/otherarch/sdcpp/src/model/diffusion/unet.hpp @@ -6,6 +6,7 @@ #include "model.h" #include "model/common/block.hpp" +#include "model/diffusion/animatediff.hpp" #include "model/diffusion/model.hpp" /*==================================================== UnetModel =====================================================*/ @@ -29,6 +30,8 @@ struct UNetConfig { bool tiny_unet = false; int model_channels = 320; int adm_in_channels = 2816; // only for VERSION_SDXL/SVD + bool enable_animatediff = false; + bool animatediff_has_mid_block = false; static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix, @@ -84,6 +87,13 @@ struct UNetConfig { return &it->second; }; + if (find_weight("motion_module.down_blocks.0.motion_modules.0.temporal_transformer.proj_in.weight") != nullptr) { + config.enable_animatediff = true; + if (find_weight("motion_module.mid_block.motion_modules.0.temporal_transformer.proj_in.weight") != nullptr) { + config.animatediff_has_mid_block = true; + } + } + if (const TensorStorage* input = find_weight("input_blocks.0.0.weight")) { if (input->n_dims == 4) { config.in_channels = static_cast(input->ne[2]); @@ -473,6 +483,12 @@ class UnetModelBlock : public GGMLBlock { blocks["out.0"] = std::shared_ptr(new GroupNorm32(ch)); // ch == model_channels // out_1 is nn.SiLU() blocks["out.2"] = std::shared_ptr(new Conv2d(model_channels, out_channels, {3, 3}, {1, 1}, {1, 1})); + + if (this->config.enable_animatediff) { + AnimateDiff::MotionModuleConfig mm_cfg; + mm_cfg.enable_mid_block = this->config.animatediff_has_mid_block; + blocks["motion_module"] = std::make_shared(mm_cfg); + } } ggml_tensor* resblock_forward(std::string name, @@ -583,6 +599,42 @@ class UnetModelBlock : public GGMLBlock { ggml_set_name(h, "bench-start"); hs.push_back(h); + + auto motion_root = config.enable_animatediff && num_video_frames > 1 + ? std::dynamic_pointer_cast(blocks["motion_module"]) + : nullptr; + auto apply_motion_input = [&](int input_block_idx, ggml_tensor* h_in) -> ggml_tensor* { + if (!motion_root) + return h_in; + int di = (input_block_idx - 1) / 3; + int mj = (input_block_idx - 1) % 3; + if (di < 0 || di >= (int)channel_mult.size() || mj < 0 || mj >= num_res_blocks) + return h_in; + auto mm = motion_root->motion("down_blocks." + std::to_string(di) + ".motion_modules." + std::to_string(mj)); + if (!mm) + return h_in; + return mm->forward(ctx, h_in, num_video_frames); + }; + auto apply_motion_output = [&](int output_block_idx, ggml_tensor* h_in) -> ggml_tensor* { + if (!motion_root) + return h_in; + int ui = output_block_idx / 3; + int mj = output_block_idx % 3; + if (ui < 0 || ui >= (int)channel_mult.size() || mj < 0 || mj > num_res_blocks) + return h_in; + auto mm = motion_root->motion("up_blocks." + std::to_string(ui) + ".motion_modules." + std::to_string(mj)); + if (!mm) + return h_in; + return mm->forward(ctx, h_in, num_video_frames); + }; + auto apply_motion_mid = [&](ggml_tensor* h_in) -> ggml_tensor* { + if (!motion_root) + return h_in; + auto mm = motion_root->motion("mid_block.motion_modules.0"); + if (!mm) + return h_in; + return mm->forward(ctx, h_in, num_video_frames); + }; // input block 1-11 size_t len_mults = channel_mult.size(); int input_block_idx = 0; @@ -597,6 +649,7 @@ class UnetModelBlock : public GGMLBlock { std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1"; h = attention_layer_forward(name, ctx, h, context, num_video_frames); // [N, mult*model_channels, h, w] } + h = apply_motion_input(input_block_idx, h); sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h"); hs.push_back(h); } @@ -624,6 +677,7 @@ class UnetModelBlock : public GGMLBlock { h = attention_layer_forward("middle_block.1", ctx, h, context, num_video_frames); // [N, 4*model_channels, h/8, w/8] h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames); // [N, 4*model_channels, h/8, w/8] } + h = apply_motion_mid(h); } sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h"); if (controls.size() > 0) { @@ -660,6 +714,8 @@ class UnetModelBlock : public GGMLBlock { up_sample_idx++; } + h = apply_motion_output(output_block_idx, h); + if (i > 0 && j == num_res_blocks) { if (tiny_unet) { output_block_idx++; diff --git a/otherarch/sdcpp/src/model/diffusion/z_image.hpp b/otherarch/sdcpp/src/model/diffusion/z_image.hpp index 9fb333cbb5da..176adfc18e4c 100644 --- a/otherarch/sdcpp/src/model/diffusion/z_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/z_image.hpp @@ -648,8 +648,8 @@ namespace ZImage { *diffusion_params.x, *diffusion_params.timesteps, tensor_or_empty(diffusion_params.context), - diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents, - diffusion_params.ref_index_mode); + diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents, + diffusion_params.ref_image_params.ref_index_mode); } void test() { diff --git a/otherarch/sdcpp/src/model_io/safetensors_io.cpp b/otherarch/sdcpp/src/model_io/safetensors_io.cpp index e38a383787fd..4157961a4899 100644 --- a/otherarch/sdcpp/src/model_io/safetensors_io.cpp +++ b/otherarch/sdcpp/src/model_io/safetensors_io.cpp @@ -103,7 +103,8 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) { // https://huggingface.co/docs/safetensors/index bool read_safetensors_file(const std::string& file_path, std::vector& tensor_storages, - std::string* error) { + std::string* error, + std::map* metadata) { std::ifstream file(file_path, std::ios::binary); if (!file.is_open()) { set_error(error, "failed to open '" + file_path + "'"); @@ -153,6 +154,18 @@ bool read_safetensors_file(const std::string& file_path, return false; } + if (metadata != nullptr) { + metadata->clear(); + auto metadata_item = header_.find("__metadata__"); + if (metadata_item != header_.end() && metadata_item->is_object()) { + for (const auto& item : metadata_item->items()) { + if (item.value().is_string()) { + metadata->emplace(item.key(), item.value().get()); + } + } + } + } + tensor_storages.clear(); auto quant_layers = kcpp_safetensors_quant::read_quantization_metadata(header_); for (auto& item : header_.items()) { diff --git a/otherarch/sdcpp/src/model_io/safetensors_io.h b/otherarch/sdcpp/src/model_io/safetensors_io.h index b18f0ae75801..4291b54f6345 100644 --- a/otherarch/sdcpp/src/model_io/safetensors_io.h +++ b/otherarch/sdcpp/src/model_io/safetensors_io.h @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_IO_SAFETENSORS_IO_H__ #define __SD_MODEL_IO_SAFETENSORS_IO_H__ +#include #include #include @@ -10,7 +11,8 @@ bool is_safetensors_file(const std::string& file_path); bool read_safetensors_file(const std::string& file_path, std::vector& tensor_storages, - std::string* error = nullptr); + std::string* error = nullptr, + std::map* metadata = nullptr); bool read_safetensors_index_file(const std::string& file_path, std::vector& shard_paths, std::string* error = nullptr); diff --git a/otherarch/sdcpp/src/model_loader.cpp b/otherarch/sdcpp/src/model_loader.cpp index ce4b33f62f80..4c03788a1cd5 100644 --- a/otherarch/sdcpp/src/model_loader.cpp +++ b/otherarch/sdcpp/src/model_loader.cpp @@ -334,7 +334,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const std::vector tensor_storages; std::string error; - if (!read_safetensors_file(file_path, tensor_storages, &error)) { + if (!read_safetensors_file(file_path, tensor_storages, &error, &metadata_)) { LOG_ERROR("%s", error.c_str()); return false; } diff --git a/otherarch/sdcpp/src/model_loader.h b/otherarch/sdcpp/src/model_loader.h index f66f27b0ff86..9580aec205f8 100644 --- a/otherarch/sdcpp/src/model_loader.h +++ b/otherarch/sdcpp/src/model_loader.h @@ -36,6 +36,7 @@ class ModelLoader { std::vector file_data; bool model_files_processed = false; String2TensorStorage tensor_storage_map; + std::map metadata_; int n_threads_; size_t add_file_path(const std::string& file_path); @@ -63,6 +64,7 @@ class ModelLoader { std::map get_vae_wtype_stat(); String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; } const String2TensorStorage& get_tensor_storage_map() const { return tensor_storage_map; } + const std::map& get_metadata() const { return metadata_; } void set_n_threads(int n_threads); void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = ""); void process_model_files(bool enable_mmap = false, bool writable_mmap = true); diff --git a/otherarch/sdcpp/src/name_conversion.cpp b/otherarch/sdcpp/src/name_conversion.cpp index 2f9a4d184c20..3aa7880e2878 100644 --- a/otherarch/sdcpp/src/name_conversion.cpp +++ b/otherarch/sdcpp/src/name_conversion.cpp @@ -304,6 +304,12 @@ std::string convert_diffusers_unet_to_original_sd1(std::string name) { } } + static const std::vector> name_map{ + {"to_out.weight", "to_out.0.weight"}, + {"to_out.bias", "to_out.0.bias"}, + }; + replace_with_name_map(result, name_map); + return result; } @@ -1102,6 +1108,8 @@ std::string convert_sep_to_dot(std::string name) { "norm1_context", "ff_context", "x_embedder", + "cross_attn", + "output_proj", }; // record the positions of underscores that should NOT be replaced @@ -1341,13 +1349,25 @@ std::string convert_tensor_name(std::string name, SDVersion version) { // diffusion model { + bool matched = false; for (const auto& prefix : diffuison_model_prefix_vec) { if (starts_with(name, prefix)) { - name = convert_diffusion_model_name(name.substr(prefix.size()), prefix, version); - name = prefix + name; + name = convert_diffusion_model_name(name.substr(prefix.size()), prefix, version); + name = prefix + name; + matched = true; break; } } + if (is_lora && !matched && !diffuison_model_prefix_vec.empty()) { + if (starts_with(name, "down_blocks.") || starts_with(name, "up_blocks.") || + starts_with(name, "mid_block.") || starts_with(name, "conv_in.") || + starts_with(name, "conv_out.") || starts_with(name, "time_embedding.") || + starts_with(name, "conv_norm_out.")) { + const std::string& canonical_prefix = diffuison_model_prefix_vec.front(); + name = convert_diffusion_model_name(name, canonical_prefix, version); + name = canonical_prefix + name; + } + } } // cond_stage_model diff --git a/otherarch/sdcpp/src/stable-diffusion.cpp b/otherarch/sdcpp/src/stable-diffusion.cpp index e5b331ba7ca8..0ee9a4d6b6e4 100644 --- a/otherarch/sdcpp/src/stable-diffusion.cpp +++ b/otherarch/sdcpp/src/stable-diffusion.cpp @@ -26,6 +26,7 @@ #include "extensions/generation_extension.h" #include "model/adapter/lora.hpp" #include "model/diffusion/anima.hpp" +#include "model/diffusion/animatediff.hpp" #include "model/diffusion/boogu.hpp" #include "model/diffusion/control.hpp" #include "model/diffusion/ernie_image.hpp" @@ -65,6 +66,10 @@ const char* sd_vae_format_name(enum sd_vae_format_t format); static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback); +static bool sd_version_supports_animatediff(SDVersion version) { + return version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX; +} + const char* model_version_to_str[] = { "SD 1.x", "SD 1.x Inpaint", @@ -214,6 +219,8 @@ class StableDiffusionGGML { std::vector> generation_extensions; std::vector> runtime_lora_models; bool apply_lora_immediately = false; + bool animatediff_loaded = false; + int animatediff_num_frames = 0; std::map> kcpp_lora_cache; bool kcpp_lora_cache_populate = false; @@ -1053,6 +1060,16 @@ class StableDiffusionGGML { } } + if (strlen(SAFE_STR(sd_ctx_params->motion_module_path)) > 0) { + LOG_INFO("loading motion module (AnimateDiff) from '%s'", sd_ctx_params->motion_module_path); + if (!model_loader.init_from_file(sd_ctx_params->motion_module_path, + "model.diffusion_model.motion_module.")) { + LOG_WARN("loading motion module from '%s' failed", sd_ctx_params->motion_module_path); + } else { + animatediff_loaded = true; + } + } + if (strlen(SAFE_STR(sd_ctx_params->control_net_path)) > 0) { if (!model_loader.init_from_file(sd_ctx_params->control_net_path)) { LOG_ERROR("init control net model loader from file failed: '%s'", sd_ctx_params->control_net_path); @@ -1539,12 +1556,12 @@ class StableDiffusionGGML { false, version, model_manager); - } else if (sd_version_uses_wan_vae(version)) { + } else if (sd_version_uses_wan_vae(vae_version)) { return std::make_shared(backend_for(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", false, - version, + vae_version, model_manager); } else { auto model = std::make_shared(backend_for(SDBackendModule::VAE), @@ -2566,7 +2583,7 @@ class StableDiffusionGGML { const char* extra_sample_args, const std::vector& sigmas, const std::vector>& ref_latents, - bool increase_ref_index, + const RefImageParams& ref_image_params, const sd::Tensor& denoise_mask, const sd::Tensor& vace_context, float vace_strength, @@ -2727,9 +2744,9 @@ class StableDiffusionGGML { sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma); std::vector> controls; DiffusionParams diffusion_params; - diffusion_params.x = &noised_input; - diffusion_params.timesteps = ×teps_tensor; - diffusion_params.ref_index_mode = Rope::ref_index_mode_from_bool(increase_ref_index); + diffusion_params.x = &noised_input; + diffusion_params.timesteps = ×teps_tensor; + diffusion_params.ref_image_params = ref_image_params; sd::guidance::GuidanceInput step_guidance_input; step_guidance_input.step = step; step_guidance_input.schedule_size = sigmas.size(); @@ -2756,7 +2773,11 @@ class StableDiffusionGGML { diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images); if (sd_version_is_unet(version)) { - diffusion_params.extra = UNetDiffusionExtra{-1, &controls, control_strength}; + int nvf = -1; + if (animatediff_loaded && noised_input.dim() >= 4 && noised_input.shape()[3] > 1) { + nvf = static_cast(noised_input.shape()[3]); + } + diffusion_params.extra = UNetDiffusionExtra{nvf, &controls, control_strength}; } else if (sd_version_is_sd3(version)) { diffusion_params.extra = SkipLayerDiffusionExtra{local_skip_layers}; } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) { @@ -3110,6 +3131,125 @@ class StableDiffusionGGML { return !!flow_denoiser; } + std::string get_default_ref_image_preset(SDVersion version) const { + if (sd_version_is_longcat(version)) { + return "longcat"; + } else if (sd_version_is_flux(version)) { + return "flux_kontext"; + } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) { + return "flux2"; + } else if (version == VERSION_QWEN_IMAGE_LAYERED) { + return "qwen_layered"; + } else if (sd_version_is_qwen_image(version)) { + return "qwen"; + } else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) { + return "z_image_omni"; + } 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) + // since krea2 ostris edit support predates, it should probably be default + return "krea2_ostris_edit"; + } else if (sd_version_is_anima(version)) { + return "cosmos_reference"; + } + return "default"; + } + + RefImageParams resolve_ref_image_params(const char* ref_image_args) const { + RefImageParams params; + std::string preset_name = get_default_ref_image_preset(version); + + for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) { + if (key == "preset") { + std::string requested_preset_name = value; + if (REF_IMAGE_PRESETS.count(requested_preset_name)) { + preset_name = requested_preset_name; + } else if (value != "default") { + std::string valid_list; + for (auto const& [name, _] : REF_IMAGE_PRESETS) { + valid_list += (valid_list.empty() ? "" : ", ") + name; + } + LOG_WARN("ignoring invalid reference image preset '%s'. Valid options: [%s]", value.c_str(), valid_list.c_str()); + } + break; + } + } + if (preset_name != "default") { + LOG_INFO("Using '%s' preset for reference images", preset_name.c_str()); + params = REF_IMAGE_PRESETS.at(preset_name); + } + + for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) { + if (key == "pass_to_vlm") { + if (!parse_strict_bool(value, params.pass_to_vlm)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "pass_to_dit") { + if (!parse_strict_bool(value, params.pass_to_dit)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "ref_index_mode") { + if (value == "fixed") { + params.ref_index_mode = Rope::RefIndexMode::FIXED; + } else if (value == "increase") { + params.ref_index_mode = Rope::RefIndexMode::INCREASE; + } else if (value == "decrease") { + params.ref_index_mode = Rope::RefIndexMode::DECREASE; + } else { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "force_ref_timestep_zero") { + if (!parse_strict_bool(value, params.force_ref_timestep_zero)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "resize_before_vae") { + if (!parse_strict_bool(value, params.resize_before_vae)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "vae_input_max_pixels") { + if (!parse_strict_int(value, params.vae_input_max_pixels)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "vlm_resize_mode") { + if (value == "longest_side") { + params.vlm_resize_mode = RefImageResizeMode::LONGEST_SIDE; + } else if (value == "area") { + params.vlm_resize_mode = RefImageResizeMode::AREA; + } else if (value == "none") { + params.vlm_resize_mode = RefImageResizeMode::NONE; + } else { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "vlm_max_size") { + if (!parse_strict_int(value, params.vlm_max_size)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "vlm_min_size") { + if (!parse_strict_int(value, params.vlm_min_size)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key != "preset" && key != "vlm_size") { + LOG_WARN("ignoring unknown reference image arg '%s'", key.c_str()); + } + } + for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) { + if (key == "vlm_size") { + int vlm_size; + if (!parse_strict_int(value, vlm_size)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } else { + LOG_INFO("vlm_size override: setting both min and max size to %ld", (long)vlm_size); + params.vlm_min_size = vlm_size; + params.vlm_max_size = vlm_size; + } + break; + } + } + if (params.force_ref_timestep_zero && !sd_version_is_krea2(version)) { + LOG_WARN("force_ref_timestep_zero is only supported by Krea2 architecture for now"); + } + return params; + } }; /*================================================= SD API ==================================================*/ @@ -3342,6 +3482,8 @@ const char* sd_vae_format_name(enum sd_vae_format_t format) { return "sd3"; case SD_VAE_FORMAT_FLUX2: return "flux2"; + case SD_VAE_FORMAT_WAN: + return "wan"; default: return NONE_STR; } @@ -3355,6 +3497,8 @@ static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion return VERSION_SD3; case SD_VAE_FORMAT_FLUX2: return VERSION_FLUX2; + case SD_VAE_FORMAT_WAN: + return VERSION_WAN2; case SD_VAE_FORMAT_AUTO: default: return fallback; @@ -3571,6 +3715,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) { sd_sample_params_init(&sd_img_gen_params->sample_params); sd_img_gen_params->clip_skip = -1; sd_img_gen_params->ref_images_count = 0; + sd_img_gen_params->ref_image_args = ""; sd_img_gen_params->width = 512; sd_img_gen_params->height = 512; sd_img_gen_params->strength = 0.75f; @@ -3608,8 +3753,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) { "batch_count: %d\n" "qwen_image_layers: %d\n" "ref_images_count: %d\n" - "auto_resize_ref_image: %s\n" - "increase_ref_index: %s\n" + "ref_image_args: %s\n" "control_strength: %.2f\n" "photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n" "VAE tiling: %s (temporal=%s, extra_tiling_args=%s)\n" @@ -3627,8 +3771,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) { sd_img_gen_params->batch_count, sd_img_gen_params->qwen_image_layers, sd_img_gen_params->ref_images_count, - BOOL_STR(sd_img_gen_params->auto_resize_ref_image), - BOOL_STR(sd_img_gen_params->increase_ref_index), + SAFE_STR(sd_img_gen_params->ref_image_args), sd_img_gen_params->control_strength, sd_img_gen_params->pm_params.style_strength, sd_img_gen_params->pm_params.id_images_count, @@ -3866,6 +4009,9 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) { if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { return false; } + if (sd_ctx->sd->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->sd->version)) { + return true; + } return sd_version_supports_video_generation(sd_ctx->sd->version); } @@ -3997,8 +4143,6 @@ struct GenerationRequest { float strength = 1.f; float control_strength = 0.f; float eta = 0.f; - bool increase_ref_index = false; - bool auto_resize_ref_image = false; sd_guidance_params_t guidance = {}; sd_guidance_params_t high_noise_guidance = {}; sd_pm_params_t pm_params = {}; @@ -4024,8 +4168,6 @@ struct GenerationRequest { strength = sd_img_gen_params->strength; control_strength = sd_img_gen_params->control_strength; eta = sd_img_gen_params->sample_params.eta; - increase_ref_index = sd_img_gen_params->increase_ref_index; - auto_resize_ref_image = sd_img_gen_params->auto_resize_ref_image; has_ref_images = sd_img_gen_params->ref_images_count > 0; guidance = sd_img_gen_params->sample_params.guidance; pm_params = sd_img_gen_params->pm_params; @@ -4774,7 +4916,8 @@ static sd::Tensor ensure_image_tensor_channels(sd::Tensor image, i static std::optional prepare_image_generation_latents(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params, GenerationRequest* request, - SamplePlan* plan) { + SamplePlan* plan, + const RefImageParams& ref_image_params) { int64_t prepare_start_ms = ggml_time_ms(); sd::Tensor init_image_tensor; @@ -4886,6 +5029,22 @@ static std::optional prepare_image_generation_latents(sd } } + if (sd_ctx->sd->animatediff_num_frames > 1 && + init_latent.dim() >= 4 && init_latent.shape()[3] == 1) { + int n_frames = sd_ctx->sd->animatediff_num_frames; + std::vector shape(init_latent.shape().begin(), init_latent.shape().end()); + shape[3] = n_frames; + if (!init_image_tensor.empty()) { + sd::Tensor replicated(shape); + for (int f = 0; f < n_frames; ++f) { + sd::ops::slice_assign(&replicated, 3, f, f + 1, init_latent); + } + init_latent = std::move(replicated); + } else { + init_latent = sd::Tensor(std::move(shape)); + } + } + if (!control_image_tensor.empty()) { control_latent = sd_ctx->sd->encode_first_stage(control_image_tensor); if (control_latent.empty()) { @@ -4917,9 +5076,10 @@ static std::optional prepare_image_generation_latents(sd continue; } sd::Tensor ref_latent; - if (request->auto_resize_ref_image && !sd_version_is_pid(sd_ctx->sd->version)) { + if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd_ctx->sd->version)) { LOG_DEBUG("auto resize ref images"); - int vae_image_size = std::min(1024 * 1024, request->width * request->height); + int target_pixels = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : 1024 * 1024; + int vae_image_size = std::min(target_pixels, request->width * request->height); double vae_width = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]); double vae_height = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0]; @@ -5046,15 +5206,20 @@ static std::optional prepare_image_generation_embeds(sd_c const sd_img_gen_params_t* sd_img_gen_params, GenerationRequest* request, SamplePlan* plan, - ImageGenerationLatents* latents) { + ImageGenerationLatents* latents, + const RefImageParams& ref_image_params) { ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()}; ConditionerParams condition_params; - condition_params.text = request->prompt; - condition_params.clip_skip = request->clip_skip; - condition_params.width = request->width; - condition_params.height = request->height; - condition_params.ref_images = &latents->ref_images; + condition_params.text = request->prompt; + condition_params.clip_skip = request->clip_skip; + condition_params.width = request->width; + condition_params.height = request->height; + if (ref_image_params.pass_to_vlm) { + condition_params.ref_images = &latents->ref_images; + } + + condition_params.ref_image_params = ref_image_params; sd_ctx->sd->prepare_generation_extensions(request->pm_params, request->pulid_params, @@ -5064,7 +5229,7 @@ static std::optional prepare_image_generation_embeds(sd_c condition_params.zero_out_masked = false; auto cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, condition_params); - if (cond.c_concat.empty()) { + if (cond.c_concat.empty() && ref_image_params.pass_to_dit) { cond.c_concat = latents->concat_latent; // TODO: optimize } @@ -5093,7 +5258,7 @@ static std::optional prepare_image_generation_embeds(sd_c uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, condition_params); } - if (uncond.c_concat.empty()) { + if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) { uncond.c_concat = latents->concat_latent; // TODO: optimize } } @@ -5117,7 +5282,7 @@ static std::optional prepare_image_generation_embeds(sd_c } img_uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, condition_params); - if (img_uncond.c_concat.empty()) { + if (img_uncond.c_concat.empty() && ref_image_params.pass_to_dit) { img_uncond.c_concat = latents->img_uncond_concat_latent; // TODO: optimize } } @@ -5188,6 +5353,24 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx, if (cancelled) { break; } + } else if (sd_ctx->sd->animatediff_num_frames > 1 && + final_latents[i].dim() >= 4 && + final_latents[i].shape()[3] == sd_ctx->sd->animatediff_num_frames) { + int n_frames = sd_ctx->sd->animatediff_num_frames; + for (int f = 0; f < n_frames; ++f) { + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling latent decodings"); + cancelled = true; + break; + } + sd::Tensor frame_latent = sd::ops::slice(final_latents[i], 3, f, f + 1); + sd::Tensor image = sd_ctx->sd->decode_first_stage(frame_latent); + if (image.empty()) { + LOG_ERROR("decode_first_stage failed for AnimateDiff frame %d/%d", f + 1, n_frames); + return nullptr; + } + decoded_images.push_back(std::move(image)); + } } else { sd::Tensor image = sd_ctx->sd->decode_first_stage(final_latents[i]); if (image.empty()) { @@ -5434,13 +5617,16 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx, sd_ctx->sd->apply_loras(sd_img_gen_params->loras, sd_img_gen_params->lora_count); apply_circular_axes_to_diffusion(sd_ctx, sd_img_gen_params->circular_x, sd_img_gen_params->circular_y); + const RefImageParams ref_image_params = sd_ctx->sd->resolve_ref_image_params(sd_img_gen_params->ref_image_args); + ImageVaeAxesGuard axes_guard(sd_ctx, sd_img_gen_params, request); SamplePlan plan(sd_ctx, sd_img_gen_params, request); auto latents_opt = prepare_image_generation_latents(sd_ctx, sd_img_gen_params, &request, - &plan); + &plan, + ref_image_params); if (!latents_opt.has_value()) { return false; } @@ -5450,7 +5636,8 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx, sd_img_gen_params, &request, &plan, - &latents); + &latents, + ref_image_params); if (!embeds_opt.has_value()) { return false; } @@ -5496,7 +5683,7 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx, plan.extra_sample_args, plan.sigmas, latents.ref_latents, - request.increase_ref_index, + ref_image_params, latents.denoise_mask, sd::Tensor(), 1.f, @@ -5617,7 +5804,7 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx, plan.extra_sample_args, hires_sigma_sched, latents.ref_latents, - request.increase_ref_index, + ref_image_params, hires_denoise_mask, sd::Tensor(), 1.f, @@ -6053,6 +6240,9 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx, condition_params.text = request.prompt; condition_params.zero_out_masked = true; condition_params.ref_images = &latents.ref_images; + if (sd_version_is_lingbot_video(sd_ctx->sd->version)) { + condition_params.ref_image_params.vlm_resize_mode = RefImageResizeMode::AREA; + } int64_t prepare_start_ms = ggml_time_ms(); embeds.cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, @@ -6322,6 +6512,48 @@ static bool apply_ltxv_refine_image_conditioning(sd_ctx_t* sd_ctx, return true; } +static bool generate_animatediff_video(sd_ctx_t* sd_ctx, + const sd_vid_gen_params_t* sd_vid_gen_params, + sd_image_t** frames_out, + int* num_frames_out) { + int n_frames = sd_vid_gen_params->video_frames; + if (n_frames < 1) { + LOG_ERROR("AnimateDiff: --video-frames must be >= 1"); + return false; + } + if (n_frames > 32) { + LOG_WARN("AnimateDiff motion modules have a 32-frame positional-encoding context; capping to 32"); + n_frames = 32; + } + + sd_img_gen_params_t img_gen_params; + sd_img_gen_params_init(&img_gen_params); + img_gen_params.loras = sd_vid_gen_params->loras; + img_gen_params.lora_count = sd_vid_gen_params->lora_count; + img_gen_params.prompt = sd_vid_gen_params->prompt; + img_gen_params.negative_prompt = sd_vid_gen_params->negative_prompt; + img_gen_params.clip_skip = sd_vid_gen_params->clip_skip; + img_gen_params.width = sd_vid_gen_params->width; + img_gen_params.height = sd_vid_gen_params->height; + img_gen_params.sample_params = sd_vid_gen_params->sample_params; + img_gen_params.strength = sd_vid_gen_params->strength; + img_gen_params.init_image = sd_vid_gen_params->init_image; + img_gen_params.seed = sd_vid_gen_params->seed; + img_gen_params.batch_count = 1; + img_gen_params.control_strength = 1.0f; + img_gen_params.vae_tiling_params = sd_vid_gen_params->vae_tiling_params; + img_gen_params.cache = sd_vid_gen_params->cache; + img_gen_params.hires = sd_vid_gen_params->hires; + 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; + + sd_ctx->sd->animatediff_num_frames = n_frames; + bool ok = generate_image(sd_ctx, &img_gen_params, frames_out, num_frames_out); + sd_ctx->sd->animatediff_num_frames = 0; + return ok; +} + SD_API bool generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params, sd_image_t** frames_out, @@ -6336,12 +6568,20 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, if (audio_out != nullptr) { *audio_out = nullptr; } - - sd_ctx->sd->reset_cancel_flag(); - if (num_frames_out != nullptr) { *num_frames_out = 0; } + + if (sd_ctx->sd->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->sd->version)) { + LOG_INFO("AnimateDiff dispatch: %d frames, %dx%d", + sd_vid_gen_params->video_frames, sd_vid_gen_params->width, sd_vid_gen_params->height); + return generate_animatediff_video(sd_ctx, sd_vid_gen_params, frames_out, num_frames_out); + } + + sd_ctx->sd->reset_cancel_flag(); + + const RefImageParams ref_image_params; + int64_t t0 = ggml_time_ms(); sd_ctx->sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params; apply_circular_axes_to_diffusion(sd_ctx, sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y); @@ -6431,7 +6671,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, plan.high_noise_extra_sample_args, high_noise_sigmas, std::vector>{}, - false, + ref_image_params, latents.denoise_mask, latents.vace_context, request.vace_strength, @@ -6473,7 +6713,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, plan.extra_sample_args, plan.sigmas, std::vector>{}, - false, + ref_image_params, latents.denoise_mask, latents.vace_context, request.vace_strength, @@ -6622,7 +6862,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, plan.extra_sample_args, hires_sigma_sched, std::vector>{}, - false, + ref_image_params, hires_denoise_mask, sd::Tensor(), hires_request.vace_strength, From 2b81a174031680be9053c3e6a69e49963771cf23 Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Wed, 22 Jul 2026 23:14:52 -0300 Subject: [PATCH 2/5] sd: sync with master-788-8a51eb9 --- .../sdcpp/src/conditioning/conditioner.hpp | 125 ++- otherarch/sdcpp/src/core/ggml_extend.hpp | 27 +- otherarch/sdcpp/src/model.h | 21 +- otherarch/sdcpp/src/model/common/block.hpp | 2 +- otherarch/sdcpp/src/model/common/rope.hpp | 64 ++ otherarch/sdcpp/src/model/diffusion/flux.hpp | 8 +- .../sdcpp/src/model/diffusion/hunyuan.hpp | 681 ++++++++++++++ otherarch/sdcpp/src/model/diffusion/ltxv.hpp | 32 +- .../sdcpp/src/model/diffusion/mage_flow.hpp | 162 ++++ otherarch/sdcpp/src/model/diffusion/mmdit.hpp | 94 +- otherarch/sdcpp/src/model/diffusion/model.hpp | 12 +- .../sdcpp/src/model/diffusion/qwen_image.hpp | 16 +- otherarch/sdcpp/src/model/te/llm.hpp | 2 +- otherarch/sdcpp/src/model/te/t5.hpp | 4 +- otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp | 834 ++++++++++++++++++ otherarch/sdcpp/src/model/vae/mage_vae.hpp | 521 +++++++++++ otherarch/sdcpp/src/model/vae/tae.hpp | 11 +- otherarch/sdcpp/src/model/vae/vae.hpp | 4 +- otherarch/sdcpp/src/model_io/gguf_io.h | 2 +- otherarch/sdcpp/src/model_loader.cpp | 7 + otherarch/sdcpp/src/name_conversion.cpp | 81 +- otherarch/sdcpp/src/runtime/denoiser.hpp | 12 +- otherarch/sdcpp/src/stable-diffusion.cpp | 150 +++- 23 files changed, 2776 insertions(+), 96 deletions(-) create mode 100644 otherarch/sdcpp/src/model/diffusion/hunyuan.hpp create mode 100644 otherarch/sdcpp/src/model/diffusion/mage_flow.hpp create mode 100644 otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp create mode 100644 otherarch/sdcpp/src/model/vae/mage_vae.hpp diff --git a/otherarch/sdcpp/src/conditioning/conditioner.hpp b/otherarch/sdcpp/src/conditioning/conditioner.hpp index 4dc4e0cd45b5..e84e86297605 100644 --- a/otherarch/sdcpp/src/conditioning/conditioner.hpp +++ b/otherarch/sdcpp/src/conditioning/conditioner.hpp @@ -1782,6 +1782,7 @@ struct LLMEmbedder : public Conditioner { SDVersion version; std::shared_ptr tokenizer; std::shared_ptr llm; + std::shared_ptr byt5; LLMEmbedder(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -1799,7 +1800,12 @@ struct LLMEmbedder : public Conditioner { arch = LLM::LLMArch::GPT_OSS_20B; } else if (sd_version_is_pid(version)) { arch = LLM::LLMArch::GEMMA2_2B; - } else if (sd_version_is_lingbot_video(version) || sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) { + } else if (sd_version_is_lingbot_video(version) || + sd_version_is_ideogram4(version) || + sd_version_is_boogu_image(version) || + sd_version_is_sefi_image(version) || + sd_version_is_krea2(version) || + sd_version_is_mage_flow(version)) { arch = LLM::LLMArch::QWEN3_VL; } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) { arch = LLM::LLMArch::QWEN3; @@ -1819,54 +1825,97 @@ struct LLMEmbedder : public Conditioner { "text_encoders.llm", enable_vision, weight_manager); + if (sd_version_is_hunyuan_video(version)) { + const std::string byt5_prefix = "text_encoders.t5xxl.transformer"; + for (const auto& [name, _] : tensor_storage_map) { + if (starts_with(name, byt5_prefix + ".")) { + byt5 = std::make_shared(backend, + tensor_storage_map, + byt5_prefix, + false, + weight_manager); + break; + } + } + } } void get_param_tensors(std::map& tensors) override { llm->get_param_tensors(tensors, "text_encoders.llm"); + if (byt5) { + byt5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } } void set_max_graph_vram_bytes(size_t max_vram_bytes) override { llm->set_max_graph_vram_bytes(max_vram_bytes); + if (byt5) { + byt5->set_max_graph_vram_bytes(max_vram_bytes); + } } void set_stream_layers_enabled(bool enabled) override { llm->set_stream_layers_enabled(enabled); + if (byt5) { + byt5->set_stream_layers_enabled(enabled); + } } void set_runtime_backends(const std::vector& backends) override { llm->set_runtime_backends(backends); + if (byt5) { + byt5->set_runtime_backends(backends); + } } void set_graph_cut_layer_split_enabled(bool enabled) override { if (llm) { llm->set_graph_cut_layer_split_enabled(enabled); } + if (byt5) { + byt5->set_graph_cut_layer_split_enabled(enabled); + } } void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) override { if (llm) { llm->set_graph_cut_layer_split_backend_vram_limits(limits); } + if (byt5) { + byt5->set_graph_cut_layer_split_backend_vram_limits(limits); + } } void get_layer_split_param_tensors(std::map& tensors) override { llm->get_param_tensors(tensors, "text_encoders.llm"); + if (byt5) { + byt5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } } void set_flash_attention_enabled(bool enabled) override { llm->set_flash_attention_enabled(enabled); + if (byt5) { + byt5->set_flash_attention_enabled(enabled); + } } void set_weight_adapter(const std::shared_ptr& adapter) override { if (llm) { llm->set_weight_adapter(adapter); } + if (byt5) { + byt5->set_weight_adapter(adapter); + } } void runner_done() override { if (llm) { llm->runner_done(); } + if (byt5) { + byt5->runner_done(); + } } std::tuple, std::vector, std::vector> tokenize(std::string text, @@ -2060,7 +2109,24 @@ struct LLMEmbedder : public Conditioner { int64_t t0 = ggml_time_ms(); RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode; - if (sd_version_is_lingbot_video(version)) { + if (sd_version_is_hunyuan_video(version)) { + prompt_template_encode_start_idx = 98; + out_layers = {26}; + + prompt = + "<|im_start|>system\nYou are a helpful assistant. Describe the video by detailing the following aspects:\n" + "1. The main content and theme of the video.\n" + "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects.\n" + "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects.\n" + "4. background environment, light, style and atmosphere.\n" + "5. camera angles, movements, and transitions used in the video.<|im_end|>\n" + "<|im_start|>user\n"; + + prompt_attn_range.first = static_cast(prompt.size()); + prompt += conditioner_params.text; + prompt_attn_range.second = static_cast(prompt.size()); + prompt += "<|im_end|>\n<|im_start|>assistant\n"; + } else if (sd_version_is_lingbot_video(version)) { const int pad_token = 151643; const std::string prompt_prefix = "<|im_start|>system\nGiven a user input that may include a text prompt alone, " @@ -2144,22 +2210,22 @@ struct LLMEmbedder : public Conditioner { prompt += conditioner_params.text; prompt_attn_range = {0, 0}; prompt += "<|im_end|>\n<|im_start|>assistant\n"; - } else if (sd_version_is_qwen_image(version)) { + } else if (sd_version_is_qwen_image(version) || sd_version_is_mage_flow(version)) { if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { - LOG_INFO("QwenImageEditPlusPipeline"); + LOG_INFO("%s", sd_version_is_mage_flow(version) ? "MageFlowEditPipeline" : "QwenImageEditPlusPipeline"); prompt_template_encode_start_idx = 64; int image_embed_idx = 64 + 6; int min_pixels = conditioner_params.ref_image_params.vlm_min_size; if (min_pixels <= 0) { - min_pixels = 384; - if (resize_mode == RefImageResizeMode::AREA) { + min_pixels = sd_version_is_mage_flow(version) ? -1 : 384; + if (min_pixels > 0 && resize_mode == RefImageResizeMode::AREA) { min_pixels *= min_pixels; } } int max_pixels = conditioner_params.ref_image_params.vlm_max_size; if (max_pixels <= 0) { - max_pixels = 560; + max_pixels = sd_version_is_mage_flow(version) ? 384 : 560; if (resize_mode == RefImageResizeMode::AREA) { max_pixels *= max_pixels; } @@ -2187,7 +2253,7 @@ struct LLMEmbedder : public Conditioner { image_embeds.emplace_back(image_embed_idx, image_embed); image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; - img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>"; // [24669, 220, index, 25, 220, 151652] + img_prompt += (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>"; int64_t num_image_tokens = image_embed.shape()[1]; img_prompt.reserve(num_image_tokens * placeholder.size()); for (int j = 0; j < num_image_tokens; j++) { @@ -2215,6 +2281,9 @@ struct LLMEmbedder : public Conditioner { prompt += "<|im_end|>\n<|im_start|>assistant\n"; } + if (sd_version_is_mage_flow(version)) { + max_length = 2048 + prompt_template_encode_start_idx; + } } else if (sd_version_is_boogu_image(version)) { prompt_template_encode_start_idx = 0; @@ -2590,6 +2659,46 @@ struct LLMEmbedder : public Conditioner { spell_quotes, max_length); std::vector> extra_hidden_states_vec; + if (sd_version_is_hunyuan_video(version) && byt5) { + std::vector quoted_texts; + auto collect_quoted = [&](const std::string& open, const std::string& close) { + size_t begin = 0; + while ((begin = conditioner_params.text.find(open, begin)) != std::string::npos) { + size_t content_begin = begin + open.size(); + size_t end = conditioner_params.text.find(close, content_begin); + if (end == std::string::npos) { + break; + } + quoted_texts.push_back(conditioner_params.text.substr(content_begin, end - content_begin)); + begin = end + close.size(); + } + }; + collect_quoted("\"", "\""); + collect_quoted("\xE2\x80\x98", "\xE2\x80\x99"); + collect_quoted("\xE2\x80\x9C", "\xE2\x80\x9D"); + + if (!quoted_texts.empty()) { + std::string byt5_text; + for (const auto& text : quoted_texts) { + byt5_text += "Text \"" + text + "\". "; + } + std::vector tokens; + tokens.reserve(byt5_text.size() + 1); + for (unsigned char byte : byt5_text) { + tokens.push_back(static_cast(byte) + 3); + } + tokens.push_back(1); + sd::Tensor input_ids({static_cast(tokens.size())}, tokens); + auto byt5_hidden_states = byt5->compute(n_threads, + input_ids, + sd::Tensor(), + false, + true, + true); + GGML_ASSERT(!byt5_hidden_states.empty()); + extra_hidden_states_vec.push_back(std::move(byt5_hidden_states)); + } + } for (int i = 0; i < extra_prompts.size(); i++) { auto extra_hidden_states = encode_prompt(n_threads, extra_prompts[i], diff --git a/otherarch/sdcpp/src/core/ggml_extend.hpp b/otherarch/sdcpp/src/core/ggml_extend.hpp index 3adae6a57e2c..2280a535f58c 100644 --- a/otherarch/sdcpp/src/core/ggml_extend.hpp +++ b/otherarch/sdcpp/src/core/ggml_extend.hpp @@ -3309,7 +3309,7 @@ class UnaryBlock : public GGMLBlock { class Identity : public UnaryBlock { public: - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { return x; } }; @@ -3368,7 +3368,7 @@ class Linear : public UnaryBlock { force_prec_f32 = force_prec_f32_; } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* w = params["weight"]; ggml_tensor* b = nullptr; if (bias) { @@ -3422,7 +3422,7 @@ class Embedding : public UnaryBlock { } ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* input_ids) { + ggml_tensor* input_ids) override { // input_ids: [N, n_token] auto weight = params["weight"]; @@ -3482,11 +3482,11 @@ class Conv2d : public UnaryBlock { scale = scale_value; } - std::string get_desc() { + std::string get_desc() override { return "Conv2d"; } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* w = params["weight"]; ggml_tensor* b = nullptr; if (bias) { @@ -3569,11 +3569,11 @@ class Conv2d_grouped : public UnaryBlock { scale = scale_value; } - std::string get_desc() { + std::string get_desc() override { return "Conv2d_grouped"; } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* w = params["weight"]; ggml_tensor* b = nullptr; if (bias) { @@ -3609,18 +3609,19 @@ class Conv2d_grouped : public UnaryBlock { if (groups == in_channels && groups == out_channels) { ggml_tensor* res; if (ctx->conv2d_direct_enabled) { - res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, x, w, + res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, w, x, stride.second, stride.first, padding.second, padding.first, dilation.second, dilation.first); } else { - res = ggml_conv_2d_dw(ctx->ggml_ctx, x, w, + res = ggml_conv_2d_dw(ctx->ggml_ctx, w, x, stride.second, stride.first, padding.second, padding.first, dilation.second, dilation.first); } if (b) { - res = ggml_add(ctx->ggml_ctx, res, b); + b = ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, b->ne[0], 1); + res = ggml_add_inplace(ctx->ggml_ctx, res, b); } return res; } @@ -3725,7 +3726,7 @@ class Conv3d : public UnaryBlock { bias(bias), force_prec_f32(force_prec_f32) {} - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* w = params["weight"]; ggml_tensor* b = nullptr; if (ctx->weight_adapter) { @@ -3778,7 +3779,7 @@ class LayerNorm : public UnaryBlock { elementwise_affine(elementwise_affine), bias(bias) {} - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* w = nullptr; ggml_tensor* b = nullptr; @@ -3865,7 +3866,7 @@ class RMSNorm : public UnaryBlock { : hidden_size(hidden_size), eps(eps) {} - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { 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.h b/otherarch/sdcpp/src/model.h index cb2b2039e72c..b0cee896b446 100644 --- a/otherarch/sdcpp/src/model.h +++ b/otherarch/sdcpp/src/model.h @@ -38,6 +38,7 @@ enum SDVersion { VERSION_LINGBOT_VIDEO, VERSION_QWEN_IMAGE, VERSION_QWEN_IMAGE_LAYERED, + VERSION_HUNYUAN_VIDEO, VERSION_ANIMA, VERSION_FLUX2, VERSION_FLUX2_KLEIN, @@ -54,6 +55,7 @@ enum SDVersion { VERSION_IDEOGRAM4, VERSION_SEFI_IMAGE, VERSION_KREA2, + VERSION_MAGE_FLOW, VERSION_ESRGAN, VERSION_COUNT, }; @@ -142,6 +144,13 @@ static inline bool sd_version_is_qwen_image(SDVersion version) { return false; } +static inline bool sd_version_is_hunyuan_video(SDVersion version) { + if (version == VERSION_HUNYUAN_VIDEO) { + return true; + } + return false; +} + static inline bool sd_version_is_anima(SDVersion version) { if (version == VERSION_ANIMA) { return true; @@ -219,6 +228,10 @@ static inline bool sd_version_is_krea2(SDVersion version) { return false; } +static inline bool sd_version_is_mage_flow(SDVersion version) { + return version == VERSION_MAGE_FLOW; +} + 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; @@ -240,6 +253,10 @@ static inline bool sd_version_uses_wan_vae(SDVersion version) { return false; } +static inline bool sd_version_uses_hunyuan_video_vae(SDVersion version) { + return sd_version_is_hunyuan_video(version); +} + static inline bool sd_version_is_inpaint(SDVersion version) { if (version == VERSION_SD1_INPAINT || version == VERSION_SD2_INPAINT || @@ -259,6 +276,7 @@ static inline bool sd_version_is_dit(SDVersion version) { sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || + sd_version_is_hunyuan_video(version) || version == VERSION_HIDREAM_O1 || sd_version_is_anima(version) || sd_version_is_z_image(version) || @@ -270,7 +288,8 @@ static inline bool sd_version_is_dit(SDVersion version) { sd_version_is_pid(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version) || - sd_version_is_krea2(version)) { + sd_version_is_krea2(version) || + sd_version_is_mage_flow(version)) { return true; } return false; diff --git a/otherarch/sdcpp/src/model/common/block.hpp b/otherarch/sdcpp/src/model/common/block.hpp index ec8f40d0b94a..6fb40ec0f61e 100644 --- a/otherarch/sdcpp/src/model/common/block.hpp +++ b/otherarch/sdcpp/src/model/common/block.hpp @@ -450,7 +450,7 @@ class SpatialTransformer : public GGMLBlock { int64_t context_dim = 768; // hidden_size, 1024 for VERSION_SD2 bool use_linear = false; - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") { + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { auto iter = tensor_storage_map.find(prefix + "proj_out.weight"); if (iter != tensor_storage_map.end()) { int64_t inner_dim = n_head * d_head; diff --git a/otherarch/sdcpp/src/model/common/rope.hpp b/otherarch/sdcpp/src/model/common/rope.hpp index d4142c181f91..c36cc7e9c0f0 100644 --- a/otherarch/sdcpp/src/model/common/rope.hpp +++ b/otherarch/sdcpp/src/model/common/rope.hpp @@ -535,6 +535,33 @@ namespace Rope { return vid_ids_repeated; } + __STATIC_INLINE__ std::vector> gen_hunyuan_video_ids(int t, + int h, + int w, + int patch_t, + int patch_h, + int patch_w, + int bs, + int context_len) { + std::vector> txt_ids(bs * context_len, std::vector(3, 0.0f)); + auto img_ids = gen_vid_ids(t, h, w, patch_t, patch_h, patch_w, bs); + return concat_ids(txt_ids, img_ids, bs); + } + + __STATIC_INLINE__ std::vector gen_hunyuan_video_pe(int t, + int h, + int w, + int patch_t, + int patch_h, + int patch_w, + int bs, + int context_len, + float theta, + const std::vector& axes_dim) { + auto ids = gen_hunyuan_video_ids(t, h, w, patch_t, patch_h, patch_w, bs, context_len); + return embed_nd(ids, bs, theta, axes_dim); + } + __STATIC_INLINE__ std::vector> gen_qwen_image_ids(int t, int h, int w, @@ -627,6 +654,43 @@ namespace Rope { return embed_nd(ids, bs, static_cast(theta), axes_dim, wrap_dims); } + __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) { + const int axes_dim_num = static_cast(axes_dim.size()); + auto make_image_ids = [=](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)); + int h_start = -(image_h - image_h / 2); + int w_start = -(image_w - image_w / 2); + for (int b = 0; b < bs; ++b) { + for (int y = 0; y < image_h; ++y) { + for (int x = 0; x < image_w; ++x) { + auto& id = image_ids[static_cast(b) * image_h * image_w + y * image_w + x]; + id[0] = static_cast(image_index); + id[1] = static_cast(h_start + y); + id[2] = static_cast(w_start + x); + } + } + } + return image_ids; + }; + auto ids = gen_flux_txt_ids(bs, context_len, axes_dim_num, {}); + auto img_ids = make_image_ids(h, w, 0); + ids = concat_ids(ids, img_ids, bs); + for (size_t i = 0; i < ref_latents.size(); ++i) { + auto ref_ids = make_image_ids(static_cast(ref_latents[i]->ne[1]), + static_cast(ref_latents[i]->ne[0]), + static_cast(i + 1)); + ids = concat_ids(ids, ref_ids, bs); + } + return embed_nd(ids, bs, static_cast(theta), axes_dim); + } + __STATIC_INLINE__ std::vector> gen_lens_ids(int h, int w, int bs, diff --git a/otherarch/sdcpp/src/model/diffusion/flux.hpp b/otherarch/sdcpp/src/model/diffusion/flux.hpp index edd7a89b89b0..375576c14f57 100644 --- a/otherarch/sdcpp/src/model/diffusion/flux.hpp +++ b/otherarch/sdcpp/src/model/diffusion/flux.hpp @@ -706,11 +706,13 @@ namespace Flux { LastLayer(int64_t hidden_size, int64_t patch_size, int64_t out_channels, - bool prune_mod = false, - bool bias = true) + bool prune_mod = false, + bool bias = true, + int64_t patch_volume = 0) : prune_mod(prune_mod) { blocks["norm_final"] = std::shared_ptr(new LayerNorm(hidden_size, 1e-06f, false)); - blocks["linear"] = std::shared_ptr(new Linear(hidden_size, patch_size * patch_size * out_channels, bias)); + int64_t out_dim = (patch_volume > 0 ? patch_volume : patch_size * patch_size) * out_channels; + blocks["linear"] = std::shared_ptr(new Linear(hidden_size, out_dim, bias)); if (!prune_mod) { blocks["adaLN_modulation.1"] = std::shared_ptr(new Linear(hidden_size, 2 * hidden_size, bias)); } diff --git a/otherarch/sdcpp/src/model/diffusion/hunyuan.hpp b/otherarch/sdcpp/src/model/diffusion/hunyuan.hpp new file mode 100644 index 000000000000..f810856394c4 --- /dev/null +++ b/otherarch/sdcpp/src/model/diffusion/hunyuan.hpp @@ -0,0 +1,681 @@ +#ifndef __SD_MODEL_DIFFUSION_HUNYUAN_HPP__ +#define __SD_MODEL_DIFFUSION_HUNYUAN_HPP__ + +#include + +#include "model/common/block.hpp" +#include "model/diffusion/flux.hpp" +#include "model/diffusion/mmdit.hpp" +#include "model/diffusion/wan.hpp" +#include "model_manager.h" + +namespace Hunyuan { + constexpr int HUNYUAN_VIDEO_GRAPH_SIZE = 65536; + + // Ref: https://github.com/huggingface/diffusers/pull/12696 + struct IndividualTokenRefinerBlock : public GGMLBlock { + protected: + int64_t num_heads; + + public: + IndividualTokenRefinerBlock(int64_t num_heads, + int64_t head_dim, + int64_t mlp_ratio = 4, + bool attn_bias = true) + : num_heads(num_heads) { + int64_t hidden_size = num_heads * head_dim; + blocks["self_attn.qkv"] = std::make_shared(hidden_size, hidden_size * 3, attn_bias); + blocks["self_attn.proj"] = std::make_shared(hidden_size, hidden_size, attn_bias); + + blocks["norm1"] = std::make_shared(hidden_size, 1e-6f, true); + blocks["norm2"] = std::make_shared(hidden_size, 1e-6f, true); + + blocks["mlp.0"] = std::make_shared(hidden_size, hidden_size * mlp_ratio); + blocks["mlp.2"] = std::make_shared(hidden_size * mlp_ratio, hidden_size); + + // adaLN_modulation.0 is nn.SiLU() + blocks["adaLN_modulation.1"] = std::make_shared(hidden_size, hidden_size * 2); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* t_emb, ggml_tensor* mask) { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto self_attn_qkv = std::dynamic_pointer_cast(blocks["self_attn.qkv"]); + auto self_attn_proj = std::dynamic_pointer_cast(blocks["self_attn.proj"]); + auto mlp_fc1 = std::dynamic_pointer_cast(blocks["mlp.0"]); + auto mlp_fc2 = std::dynamic_pointer_cast(blocks["mlp.2"]); + auto adaLN_modulation_1 = std::dynamic_pointer_cast(blocks["adaLN_modulation.1"]); + + // self attn + auto qkv = self_attn_qkv->forward(ctx, norm1->forward(ctx, txt)); + auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv); + auto q = qkv_vec[0]; + auto k = qkv_vec[1]; + auto v = qkv_vec[2]; + + auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled); + attn_out = self_attn_proj->forward(ctx, attn_out); + + // adaLN_modulation + auto emb = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, t_emb)); + auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0); + + txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, attn_out, mods[0])); + + // mlp + auto mlp_out = mlp_fc1->forward(ctx, norm2->forward(ctx, txt)); + mlp_out = ggml_silu_inplace(ctx->ggml_ctx, mlp_out); + mlp_out = mlp_fc2->forward(ctx, mlp_out); + txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, mlp_out, mods[1])); + + return txt; + } + }; + + struct IndividualTokenRefiner : public GGMLBlock { + protected: + int num_layers; + + public: + IndividualTokenRefiner(int64_t num_heads, + int64_t head_dim, + int num_layers, + int64_t mlp_ratio = 4, + bool attn_bias = true) + : num_layers(num_layers) { + for (int i = 0; i < num_layers; i++) { + blocks["blocks." + std::to_string(i)] = std::make_shared(num_heads, head_dim, mlp_ratio, attn_bias); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* t_emb, ggml_tensor* mask) { + for (int i = 0; i < num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); + + txt = block->forward(ctx, txt, t_emb, mask); + } + + return txt; + } + }; + + struct TokenRefiner : public GGMLBlock { + public: + TokenRefiner(int64_t in_channels, + int64_t num_heads, + int64_t head_dim, + int num_layers, + int64_t mlp_ratio = 4, + bool attn_bias = true) { + int64_t hidden_size = num_heads * head_dim; + blocks["input_embedder"] = std::make_shared(in_channels, hidden_size); + blocks["t_embedder"] = std::make_shared(256, hidden_size); + blocks["c_embedder"] = std::make_shared(in_channels, hidden_size); + blocks["individual_token_refiner"] = std::make_shared(num_heads, head_dim, num_layers, mlp_ratio, attn_bias); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* timestep, ggml_tensor* mask) { + auto input_embedder = std::dynamic_pointer_cast(blocks["input_embedder"]); + auto t_embedder = std::dynamic_pointer_cast(blocks["t_embedder"]); + auto c_embedder = std::dynamic_pointer_cast(blocks["c_embedder"]); + auto individual_token_refiner = std::dynamic_pointer_cast(blocks["individual_token_refiner"]); + + auto t_emb = t_embedder->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f)); + + auto h = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, txt, 1, 0, 2, 3)); + auto pooled_projections = ggml_scale(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, h), 1.f / txt->ne[1]); + pooled_projections = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, pooled_projections, 1, 0, 2, 3)); + auto c_emb = c_embedder->forward(ctx, pooled_projections); + + t_emb = ggml_add(ctx->ggml_ctx, t_emb, c_emb); + txt = input_embedder->forward(ctx, txt); + txt = individual_token_refiner->forward(ctx, txt, t_emb, mask); + return txt; + } + }; + + struct ByT5Mapper : public UnaryBlock { + ByT5Mapper(int64_t in_dim, int64_t hidden_size) { + blocks["layernorm"] = std::make_shared(in_dim); + blocks["fc1"] = std::make_shared(in_dim, 2048); + blocks["fc2"] = std::make_shared(2048, 2048); + blocks["fc3"] = std::make_shared(2048, hidden_size); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto layernorm = std::dynamic_pointer_cast(blocks["layernorm"]); + auto fc1 = std::dynamic_pointer_cast(blocks["fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); + auto fc3 = std::dynamic_pointer_cast(blocks["fc3"]); + + x = fc1->forward(ctx, layernorm->forward(ctx, x)); + x = ggml_ext_gelu(ctx->ggml_ctx, x); + x = fc2->forward(ctx, x); + x = ggml_ext_gelu(ctx->ggml_ctx, x); + return fc3->forward(ctx, x); + } + }; + + struct HunyuanVideoConfig { + std::tuple patch_size = {1, 2, 2}; + int64_t in_channels = 65; + int64_t out_channels = 32; + int64_t hidden_size = 2048; + int64_t vec_in_dim = 0; + int64_t context_in_dim = 3584; + int64_t vision_in_dim = 0; + float mlp_ratio = 4.0f; + int num_heads = 16; + int depth = 54; + int depth_single_blocks = 0; + bool qkv_bias = true; + bool guidance_embed = false; + bool use_byt5 = false; + bool use_cond_type_embedding = false; + bool use_meanflow = false; + bool use_meanflow_sum = false; + float theta = 256; + std::vector axes_dim = {16, 56, 56}; + int axes_dim_sum = 128; + + int64_t patch_volume() const { + return static_cast(std::get<0>(patch_size)) * std::get<1>(patch_size) * std::get<2>(patch_size); + } + + static HunyuanVideoConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix) { + HunyuanVideoConfig config; + config.depth = 0; + config.depth_single_blocks = 0; + bool inferred = false; + + int64_t img_embed_dim = 0; + for (const auto& [name, storage] : tensor_storage_map) { + if (starts_with(name, prefix) && ends_with(name, "img_in.proj.bias")) { + img_embed_dim = storage.ne[0]; + break; + } + } + + for (const auto& entry : tensor_storage_map) { + const auto& name = entry.first; + const auto& storage = entry.second; + if (!starts_with(name, prefix)) { + continue; + } + + auto update_depth = [&](const char* block_prefix, int* depth) { + size_t pos = name.find(block_prefix); + if (pos == std::string::npos) { + return; + } + pos += strlen(block_prefix); + size_t end = name.find('.', pos); + if (end != std::string::npos) { + *depth = std::max(*depth, atoi(name.substr(pos, end - pos).c_str()) + 1); + } + }; + update_depth("double_blocks.", &config.depth); + update_depth("single_blocks.", &config.depth_single_blocks); + + if (ends_with(name, "img_in.proj.weight") && storage.n_dims == 5) { + config.patch_size = {static_cast(storage.ne[2]), + static_cast(storage.ne[1]), + static_cast(storage.ne[0])}; + config.in_channels = storage.ne[3]; + config.hidden_size = storage.ne[4]; + inferred = true; + } else if (ends_with(name, "img_in.proj.weight") && storage.n_dims == 4) { + config.patch_size = {static_cast(storage.ne[2]), + static_cast(storage.ne[1]), + static_cast(storage.ne[0])}; + if (img_embed_dim > 0 && storage.ne[3] % img_embed_dim == 0) { + config.hidden_size = img_embed_dim; + config.in_channels = storage.ne[3] / img_embed_dim; + } + inferred = true; + } else if (ends_with(name, "txt_in.input_embedder.weight")) { + config.context_in_dim = storage.ne[0]; + inferred = true; + } else if (ends_with(name, "vector_in.in_layer.weight")) { + config.vec_in_dim = storage.ne[0]; + } else if (ends_with(name, "vision_in.proj.0.weight")) { + config.vision_in_dim = storage.ne[0]; + } else if (ends_with(name, "double_blocks.0.img_attn.norm.key_norm.scale") || + ends_with(name, "double_blocks.0.img_attn.norm.key_norm.weight")) { + config.num_heads = static_cast(config.hidden_size / storage.ne[0]); + } else if (ends_with(name, "double_blocks.0.img_mlp.0.weight")) { + config.mlp_ratio = static_cast(storage.ne[1]) / static_cast(storage.ne[0]); + } + + config.guidance_embed = config.guidance_embed || name.find("guidance_in.") != std::string::npos; + config.use_byt5 = config.use_byt5 || name.find("byt5_in.") != std::string::npos; + config.use_meanflow = config.use_meanflow || name.find("time_r_in.") != std::string::npos; + } + + config.use_cond_type_embedding = tensor_storage_map.find(prefix + ".cond_type_embedding.weight") != tensor_storage_map.end(); + config.use_meanflow_sum = config.vision_in_dim > 0; + + auto final_iter = tensor_storage_map.find(prefix + ".final_layer.linear.weight"); + if (final_iter != tensor_storage_map.end()) { + config.out_channels = final_iter->second.ne[1] / config.patch_volume(); + } + config.qkv_bias = tensor_storage_map.find(prefix + ".double_blocks.0.img_attn.qkv.bias") != tensor_storage_map.end(); + + GGML_ASSERT(config.hidden_size % config.num_heads == 0); + GGML_ASSERT(config.hidden_size / config.num_heads == config.axes_dim_sum); + + if (inferred) { + LOG_DEBUG("hunyuan video: depth = %d, single depth = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", hidden_size = %" PRId64 ", context_in_dim = %" PRId64 ", patch_size = %dx%dx%d", + config.depth, + config.depth_single_blocks, + config.in_channels, + config.out_channels, + config.hidden_size, + config.context_in_dim, + std::get<0>(config.patch_size), + std::get<1>(config.patch_size), + std::get<2>(config.patch_size)); + } + return config; + } + }; + + class HunyuanVideoModel : public GGMLBlock { + protected: + HunyuanVideoConfig config; + + void init_params(struct ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + if (config.use_cond_type_embedding) { + ggml_type type = get_type(prefix + "cond_type_embedding.weight", tensor_storage_map, GGML_TYPE_F16); + GGMLBlock::params["cond_type_embedding.weight"] = ggml_new_tensor_2d(ctx, type, config.hidden_size, 3); + } + } + + public: + HunyuanVideoModel() {} + explicit HunyuanVideoModel(HunyuanVideoConfig config) + : config(std::move(config)) { + int64_t head_dim = this->config.hidden_size / this->config.num_heads; + blocks["txt_in"] = std::make_shared(this->config.context_in_dim, this->config.num_heads, head_dim, 2); + blocks["img_in"] = std::make_shared(static_cast(224) /*Not used*/, + this->config.patch_size, + this->config.in_channels, + this->config.hidden_size); + blocks["time_in"] = std::make_shared(256, this->config.hidden_size); + if (this->config.vec_in_dim > 0) { + blocks["vector_in"] = std::make_shared(this->config.vec_in_dim, this->config.hidden_size); + } + if (this->config.vision_in_dim > 0) { + blocks["vision_in"] = std::make_shared(this->config.vision_in_dim, this->config.hidden_size); + } + if (this->config.guidance_embed) { + blocks["guidance_in"] = std::make_shared(256, this->config.hidden_size); + } + if (this->config.use_byt5) { + blocks["byt5_in"] = std::make_shared(1472, this->config.hidden_size); + } + if (this->config.use_meanflow) { + blocks["time_r_in"] = std::make_shared(256, this->config.hidden_size); + } + + for (int i = 0; i < this->config.depth; i++) { + blocks["double_blocks." + std::to_string(i)] = std::make_shared(this->config.hidden_size, + this->config.num_heads, + this->config.mlp_ratio, + i, + this->config.qkv_bias); + } + + for (int i = 0; i < this->config.depth_single_blocks; i++) { + blocks["single_blocks." + std::to_string(i)] = std::make_shared(this->config.hidden_size, + this->config.num_heads, + this->config.mlp_ratio, + i, + 0.f); + } + + blocks["final_layer"] = std::make_shared(this->config.hidden_size, + std::get<2>(this->config.patch_size), + this->config.out_channels, + false, + true, + this->config.patch_volume()); + } + + ggml_tensor* pad_to_patch_size(struct ggml_context* ctx, + ggml_tensor* x) { + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t T = x->ne[2]; + + int pt = std::get<0>(config.patch_size); + int ph = std::get<1>(config.patch_size); + int pw = std::get<2>(config.patch_size); + int pad_t = (pt - static_cast(T % pt)) % pt; + int pad_h = (ph - static_cast(H % ph)) % ph; + int pad_w = (pw - static_cast(W % pw)) % pw; + x = ggml_pad(ctx, x, pad_w, pad_h, pad_t, 0); // [N*C, T + pad_t, H + pad_h, W + pad_w] + + return x; + } + + ggml_tensor* unpatchify(struct ggml_context* ctx, + ggml_tensor* x, + int64_t t_len, + int64_t h_len, + int64_t w_len) { + // x: [N, t_len*h_len*w_len, C*pt*ph*pw] + // return: [N*C, t_len*pt, h_len*ph, w_len*pw] + int64_t N = x->ne[3]; + int64_t pt = std::get<0>(config.patch_size); + int64_t ph = std::get<1>(config.patch_size); + int64_t pw = std::get<2>(config.patch_size); + int64_t C = x->ne[0] / pt / ph / pw; + + GGML_ASSERT(C * pt * ph * pw == x->ne[0]); + + x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] + x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, ph, pt, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N); // [N*C*t_len, h_len*ph, pt, w_len*pw] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, pt, h_len*ph, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N); // [N*C, t_len*pt, h_len*ph, w_len*pw] + return x; + } + + ggml_tensor* add_condition_type(GGMLRunnerContext* ctx, ggml_tensor* x, int type) { + if (!config.use_cond_type_embedding) { + return x; + } + auto weight = GGMLBlock::params["cond_type_embedding.weight"]; + auto row = ggml_view_1d(ctx->ggml_ctx, + weight, + weight->ne[0], + static_cast(type) * weight->nb[1]); + auto target = ggml_new_tensor_3d(ctx->ggml_ctx, row->type, config.hidden_size, x->ne[1], x->ne[2]); + auto embed = ggml_repeat(ctx->ggml_ctx, row, target); + embed = ggml_cast(ctx->ggml_ctx, embed, x->type); + return ggml_add(ctx->ggml_ctx, x, embed); + } + + ggml_tensor* forward_orig(GGMLRunnerContext* ctx, + ggml_tensor* img, + ggml_tensor* txt, + ggml_tensor* timestep, + ggml_tensor* pe, + ggml_tensor* guidance = nullptr, + ggml_tensor* y = nullptr, + ggml_tensor* txt_byt5 = nullptr, + ggml_tensor* clip_fea = nullptr, + ggml_tensor* timestep_r = nullptr, + int64_t N = 1) { + // img: [N*C, T, H, W], C => in_dim + // txt: [N, L, text_dim] + // timestep: [N,] or [T] + // return: [N, t_len*h_len*w_len, out_dim*pt*ph*pw] + + GGML_ASSERT(N == 1); + + auto img_in = std::dynamic_pointer_cast(blocks["img_in"]); + auto txt_in = std::dynamic_pointer_cast(blocks["txt_in"]); + auto time_in = std::dynamic_pointer_cast(blocks["time_in"]); + auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); + + img = img_in->forward(ctx, img); // [N*C, t_len*h_len*w_len, hidden_size] + txt = txt_in->forward(ctx, txt, timestep, nullptr); // [N, n_txt_token, hidden_size] + auto vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f)); + if (config.use_meanflow && timestep_r != nullptr) { + auto time_r_in = std::dynamic_pointer_cast(blocks["time_r_in"]); + auto vec_r = time_r_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_r, 256, 10000, 1000.f)); + vec = ggml_add(ctx->ggml_ctx, vec, vec_r); + if (!config.use_meanflow_sum) { + vec = ggml_scale(ctx->ggml_ctx, vec, 0.5f); + } + } + if (config.vec_in_dim > 0 && y != nullptr) { + auto vector_in = std::dynamic_pointer_cast(blocks["vector_in"]); + vec = ggml_add(ctx->ggml_ctx, vec, vector_in->forward(ctx, y)); + } + if (config.guidance_embed && guidance != nullptr) { + auto guidance_in = std::dynamic_pointer_cast(blocks["guidance_in"]); + auto guidance_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 256, 10000, 1.f); + vec = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, guidance_emb)); + } + + txt = add_condition_type(ctx, txt, 0); + if (config.use_byt5 && txt_byt5 != nullptr) { + auto byt5_in = std::dynamic_pointer_cast(blocks["byt5_in"]); + txt_byt5 = add_condition_type(ctx, byt5_in->forward(ctx, txt_byt5), 1); + txt = config.use_cond_type_embedding ? ggml_concat(ctx->ggml_ctx, txt_byt5, txt, 1) + : ggml_concat(ctx->ggml_ctx, txt, txt_byt5, 1); + } + if (config.vision_in_dim > 0 && clip_fea != nullptr) { + auto vision_in = std::dynamic_pointer_cast(blocks["vision_in"]); + clip_fea = add_condition_type(ctx, vision_in->forward(ctx, clip_fea), 2); + txt = ggml_concat(ctx->ggml_ctx, clip_fea, txt, 1); + } + + for (int i = 0; i < config.depth; i++) { + auto block = std::dynamic_pointer_cast(blocks["double_blocks." + std::to_string(i)]); + + auto img_txt = block->forward(ctx, img, txt, vec, pe, nullptr); + img = img_txt.first; // [N, n_img_token, hidden_size] + txt = img_txt.second; // [N, n_txt_token, hidden_size] + } + + if (config.depth_single_blocks > 0) { + auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_img_token, hidden_size] + for (int i = 0; i < config.depth_single_blocks; i++) { + auto block = std::dynamic_pointer_cast(blocks["single_blocks." + std::to_string(i)]); + txt_img = block->forward(ctx, txt_img, vec, pe, nullptr); + } + + txt_img = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, txt_img, 0, 2, 1, 3)); + img = ggml_view_3d(ctx->ggml_ctx, + txt_img, + txt_img->ne[0], + txt_img->ne[1], + img->ne[1], + txt_img->nb[1], + txt_img->nb[2], + txt_img->nb[2] * txt->ne[1]); + img = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, img, 0, 2, 1, 3)); + } + + img = final_layer->forward(ctx, img, vec); // (N, t_len*h_len*w_len, out_channels * patch_size ** 3) + + return img; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timestep, + ggml_tensor* context, + ggml_tensor* pe, + ggml_tensor* guidance = nullptr, + ggml_tensor* y = nullptr, + ggml_tensor* txt_byt5 = nullptr, + ggml_tensor* clip_fea = nullptr, + ggml_tensor* timestep_r = nullptr, + int64_t N = 1) { + // Forward pass of DiT. + // x: [N*C, T, H, W] + // timestep: [N,] + // context: [N, L, D] + // pe: [L, d_head/2, 2, 2] + // return: [N*C, T, H, W] + + GGML_ASSERT(N == 1); + + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t T = x->ne[2]; + x = pad_to_patch_size(ctx->ggml_ctx, x); + + int64_t pt = std::get<0>(config.patch_size); + int64_t ph = std::get<1>(config.patch_size); + int64_t pw = std::get<2>(config.patch_size); + int64_t t_len = (T + pt - 1) / pt; + int64_t h_len = (H + ph - 1) / ph; + int64_t w_len = (W + pw - 1) / pw; + + auto out = forward_orig(ctx, x, context, timestep, pe, guidance, y, txt_byt5, clip_fea, timestep_r, N); + + out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w] + + // slice + out = ggml_ext_slice(ctx->ggml_ctx, out, 2, 0, T); // [N*C, T, H + pad_h, W + pad_w] + out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w] + out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W); // [N*C, T, H, W] + + return out; + } + }; + + struct HunyuanVideoRunner : public DiffusionModelRunner { + public: + HunyuanVideoConfig config; + HunyuanVideoModel hunyuan_video; + std::vector pe_vec; + SDVersion version; + + HunyuanVideoRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + SDVersion version = VERSION_HUNYUAN_VIDEO, + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(HunyuanVideoConfig::detect_from_weights(tensor_storage_map, prefix)), + version(version) { + LOG_INFO("HunyuanVideo blocks: %d double, %d single", config.depth, config.depth_single_blocks); + + hunyuan_video = HunyuanVideoModel(config); + hunyuan_video.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "hunyuan_video"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + hunyuan_video.get_param_tensors(tensors, prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor, + const sd::Tensor& c_concat_tensor = {}, + const sd::Tensor& y_tensor = {}, + const sd::Tensor& guidance_tensor = {}, + const sd::Tensor& byt5_tensor = {}, + const sd::Tensor& vision_tensor = {}, + const sd::Tensor& timestep_r_tensor = {}) { + ggml_cgraph* gf = new_graph_custom(HUNYUAN_VIDEO_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* c_concat = make_optional_input(c_concat_tensor); + ggml_tensor* y = make_optional_input(y_tensor); + ggml_tensor* guidance = make_optional_input(guidance_tensor); + ggml_tensor* byt5 = make_optional_input(byt5_tensor); + ggml_tensor* vision = make_optional_input(vision_tensor); + ggml_tensor* timestep_r = make_optional_input(timestep_r_tensor); + + GGML_ASSERT(x->ne[3] == config.out_channels); + if (c_concat != nullptr) { + x = ggml_concat(compute_ctx, x, c_concat, 3); + } + GGML_ASSERT(x->ne[3] <= config.in_channels); + if (x->ne[3] < config.in_channels) { + x = ggml_pad(compute_ctx, x, 0, 0, 0, static_cast(config.in_channels - x->ne[3])); + } + + int text_len = static_cast(context->ne[1]); + if (byt5 != nullptr) { + text_len += static_cast(byt5->ne[1]); + } + if (vision != nullptr) { + text_len += static_cast(vision->ne[1]); + } + pe_vec = Rope::gen_hunyuan_video_pe(static_cast(x->ne[2]), + static_cast(x->ne[1]), + static_cast(x->ne[0]), + std::get<0>(config.patch_size), + std::get<1>(config.patch_size), + std::get<2>(config.patch_size), + 1, + text_len, + config.theta, + config.axes_dim); + int64_t pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + // LOG_DEBUG("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); + // pe->data = pe_vec.data(); + // print_ggml_tensor(pe, true, "pe"); + // pe->data = nullptr; + set_backend_tensor_data(pe, pe_vec.data()); + + auto runner_ctx = get_context(); + + ggml_tensor* out = hunyuan_video.forward(&runner_ctx, + x, + timesteps, + context, + pe, + guidance, + y, + byt5, + vision, + timestep_r); + + 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, + const sd::Tensor& c_concat = {}, + const sd::Tensor& y = {}, + const sd::Tensor& guidance = {}, + const sd::Tensor& byt5 = {}, + const sd::Tensor& vision = {}, + const sd::Tensor& timestep_r = {}) { + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context, c_concat, y, guidance, byt5, vision, timestep_r); + }; + + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + GGML_ASSERT(diffusion_params.context != nullptr); + const auto* extra = diffusion_extra_as(diffusion_params); + return compute(n_threads, + *diffusion_params.x, + *diffusion_params.timesteps, + *diffusion_params.context, + tensor_or_empty(diffusion_params.c_concat), + tensor_or_empty(diffusion_params.y), + tensor_or_empty(extra->guidance), + tensor_or_empty(extra->byt5), + tensor_or_empty(extra->vision), + tensor_or_empty(extra->timestep_r)); + } + }; + +} // namespace Hunyuan + +#endif // __SD_MODEL_DIFFUSION_HUNYUAN_HPP__ diff --git a/otherarch/sdcpp/src/model/diffusion/ltxv.hpp b/otherarch/sdcpp/src/model/diffusion/ltxv.hpp index b89ff32c6a55..39c633747e5a 100644 --- a/otherarch/sdcpp/src/model/diffusion/ltxv.hpp +++ b/otherarch/sdcpp/src/model/diffusion/ltxv.hpp @@ -800,7 +800,7 @@ namespace LTXV { auto gate_mlp = mods[5]; auto x_norm = rms_norm(ctx->ggml_ctx, x); - x_norm = modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa); + x_norm = LTXV::modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa); auto msa = attn1->forward(ctx, x_norm, nullptr, self_attention_mask, pe); x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, msa, gate_msa)); @@ -810,12 +810,12 @@ namespace LTXV { auto gate_q = mods[8]; auto q = rms_norm(ctx->ggml_ctx, x); - q = modulate(ctx->ggml_ctx, q, shift_q, scale_q); + q = LTXV::modulate(ctx->ggml_ctx, q, shift_q, scale_q); auto context_mod = context; if (prompt_timestep != nullptr) { auto prompt_mods = get_prompt_scale_shift_values(ctx, prompt_timestep); - context_mod = modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]); + context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]); } auto mca = attn2->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr); @@ -826,7 +826,7 @@ namespace LTXV { } auto y = rms_norm(ctx->ggml_ctx, x); - y = modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp); + y = LTXV::modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp); auto mlp_out = ff->forward(ctx, y); x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, mlp_out, gate_mlp)); return x; @@ -1177,11 +1177,11 @@ namespace LTXV { if (cross_attention_adaln) { auto q_mods = get_ada_values(ctx, table, timestep, dim, 9, 6, 3); auto q = rms_norm(ctx->ggml_ctx, x); - q = modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]); + q = LTXV::modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]); auto context_mod = context; if (prompt_timestep != nullptr && prompt_table != nullptr) { auto p_mods = get_ada_values(ctx, prompt_table, prompt_timestep, dim, 2); - context_mod = modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]); + context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]); } auto out = attn->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr); return apply_gate(ctx->ggml_ctx, out, q_mods[2]); @@ -1228,7 +1228,7 @@ namespace LTXV { auto v_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6); auto v_norm = rms_norm(ctx->ggml_ctx, vx); - v_norm = modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]); + v_norm = LTXV::modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]); auto v_sa = attn1->forward(ctx, v_norm, nullptr, self_attention_mask, v_pe); vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_sa, v_mods[2])); auto v_txt = apply_text_cross_attention(ctx, @@ -1246,7 +1246,7 @@ namespace LTXV { if (run_ax) { auto a_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6); auto a_norm = rms_norm(ctx->ggml_ctx, ax); - a_norm = modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]); + a_norm = LTXV::modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]); auto a_sa = audio_attn1->forward(ctx, a_norm, nullptr, nullptr, a_pe); ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_sa, a_mods[2])); auto a_txt = apply_text_cross_attention(ctx, @@ -1269,8 +1269,8 @@ namespace LTXV { auto a2v_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4); auto a2v_audio = get_ada_values(ctx, a2v_audio_table, a_cross_scale_shift_timestep, a_dim, 4); auto a2v_video = get_ada_values(ctx, a2v_video_table, v_cross_scale_shift_timestep, v_dim, 4); - auto vx_scaled = modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]); - auto ax_scaled = modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]); + auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]); + auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]); auto a2v_out = audio_to_video_attn->forward(ctx, vx_scaled, ax_scaled, nullptr, v_cross_pe, a_cross_pe); auto a2v_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 4, 5); auto a2v_gate = get_ada_values(ctx, a2v_gate_table, v_cross_gate_timestep, v_dim, 1)[0]; @@ -1282,8 +1282,8 @@ namespace LTXV { auto v2a_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4); auto v2a_audio = get_ada_values(ctx, v2a_audio_table, a_cross_scale_shift_timestep, a_dim, 4); auto v2a_video = get_ada_values(ctx, v2a_video_table, v_cross_scale_shift_timestep, v_dim, 4); - auto ax_scaled = modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]); - auto vx_scaled = modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]); + auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]); + auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]); auto v2a_out = video_to_audio_attn->forward(ctx, ax_scaled, vx_scaled, nullptr, a_cross_pe, v_cross_pe); auto v2a_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 4, 5); auto v2a_gate = get_ada_values(ctx, v2a_gate_table, a_cross_gate_timestep, a_dim, 1)[0]; @@ -1291,14 +1291,14 @@ namespace LTXV { } auto a_ff_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6, 3, 3); auto ax_scaled = rms_norm(ctx->ggml_ctx, ax); - ax_scaled = modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]); + ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]); auto a_ff_out = audio_ff->forward(ctx, ax_scaled); ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_ff_out, a_ff_mods[2])); } auto v_ff_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6, 3, 3); auto vx_scaled = rms_norm(ctx->ggml_ctx, vx); - vx_scaled = modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]); + vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]); auto v_ff_out = ff->forward(ctx, vx_scaled); vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_ff_out, v_ff_mods[2])); @@ -1657,14 +1657,14 @@ namespace LTXV { auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, config.hidden_size); vx = norm_out->forward(ctx, vx); - vx = modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]); + vx = LTXV::modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]); vx = proj_out->forward(ctx, vx); vx = unpatchify_video(ctx, vx, width, height, frames); if (ax != nullptr && audio_time > 0) { auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, config.audio_hidden_size); ax = audio_norm_out->forward(ctx, ax); - ax = modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]); + ax = LTXV::modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]); ax = audio_proj_out->forward(ctx, ax); ax = unpatchify_audio(ctx, ax, audio_time); } diff --git a/otherarch/sdcpp/src/model/diffusion/mage_flow.hpp b/otherarch/sdcpp/src/model/diffusion/mage_flow.hpp new file mode 100644 index 000000000000..6ac2d2aa5b7c --- /dev/null +++ b/otherarch/sdcpp/src/model/diffusion/mage_flow.hpp @@ -0,0 +1,162 @@ +#ifndef __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__ +#define __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__ + +#include +#include + +#include "model/diffusion/qwen_image.hpp" + +namespace MageFlow { + constexpr int MAGE_FLOW_GRAPH_SIZE = 20480; + + // Mage-Flow was trained with BF16-rounded timestep frequencies; using Qwen's F32 projection degrades generation quality. + struct MageFlowTimestepProjEmbeddings : public Qwen::QwenTimestepProjEmbeddings { + static constexpr int TIMESTEP_DIM = 256; + static constexpr int HALF_DIM = TIMESTEP_DIM / 2; + + std::vector frequencies; + std::vector timesteps_proj; + + explicit MageFlowTimestepProjEmbeddings(int64_t embedding_dim) + : QwenTimestepProjEmbeddings(embedding_dim), frequencies(HALF_DIM) { + for (int i = 0; i < HALF_DIM; ++i) { + float frequency = std::exp(-std::log(10000.f) * static_cast(i) / HALF_DIM); + frequencies[i] = ggml_bf16_to_fp32(ggml_fp32_to_bf16(frequency)); + } + } + + void prepare(const sd::Tensor& timesteps) { + size_t num_timesteps = static_cast(timesteps.numel()); + timesteps_proj.resize(static_cast(TIMESTEP_DIM) * num_timesteps); + for (size_t b = 0; b < num_timesteps; ++b) { + float sigma = ggml_bf16_to_fp32(ggml_fp32_to_bf16(timesteps.values()[b] / 1000.f)); + for (int i = 0; i < HALF_DIM; ++i) { + float argument = sigma * frequencies[i] * 1000.f; + timesteps_proj[b * TIMESTEP_DIM + i] = + ggml_bf16_to_fp32(ggml_fp32_to_bf16(std::cos(argument))); + timesteps_proj[b * TIMESTEP_DIM + HALF_DIM + i] = + ggml_bf16_to_fp32(ggml_fp32_to_bf16(std::sin(argument))); + } + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* timesteps, + ggml_tensor* addition_t_cond = nullptr) override { + GGML_ASSERT(addition_t_cond == nullptr); + GGML_ASSERT(timesteps_proj.size() == + static_cast(TIMESTEP_DIM * ggml_nelements(timesteps))); + auto projection = ggml_new_tensor_2d(ctx->ggml_ctx, + GGML_TYPE_F32, + TIMESTEP_DIM, + ggml_nelements(timesteps)); + ctx->bind_backend_tensor_data(projection, timesteps_proj.data()); + auto timestep_embedder = std::dynamic_pointer_cast(blocks["timestep_embedder"]); + return timestep_embedder->forward(ctx, projection); + } + }; + + struct MageFlowRunner : public DiffusionModelRunner { + public: + Qwen::QwenImageConfig config; + Qwen::QwenImageModel mage_flow; + std::shared_ptr time_text_embed; + std::vector pe_vec; + + MageFlowRunner(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.patch_size = 1; + config.in_channels = 128; + config.out_channels = 128; + config.num_layers = 12; + config.attention_head_dim = 128; + config.num_attention_heads = 24; + config.joint_attention_dim = 2560; + config.theta = 10000; + config.axes_dim = {16, 56, 56}; + config.axes_dim_sum = 128; + time_text_embed = std::make_shared( + config.num_attention_heads * config.attention_head_dim); + mage_flow = Qwen::QwenImageModel(config, time_text_embed); + mage_flow.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "mage_flow"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + mage_flow.get_param_tensors(tensors, prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor, + const std::vector>& ref_latents_tensor = {}) { + ggml_cgraph* gf = new_graph_custom(MAGE_FLOW_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); + + std::vector ref_latents; + ref_latents.reserve(ref_latents_tensor.size()); + for (const auto& ref_latent_tensor : ref_latents_tensor) { + ref_latents.push_back(make_input(ref_latent_tensor)); + } + + 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); + 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()); + + time_text_embed->prepare(timesteps_tensor); + auto runner_ctx = get_context(); + auto out = mage_flow.forward(&runner_ctx, + x, + timesteps, + nullptr, + context, + pe, + ref_latents); + 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, + const std::vector>& ref_latents = {}) { + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context, ref_latents); + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + static const std::vector> empty_ref_latents; + return compute(n_threads, + *diffusion_params.x, + *diffusion_params.timesteps, + tensor_or_empty(diffusion_params.context), + diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents); + } + }; +} // namespace MageFlow + +#endif // __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__ diff --git a/otherarch/sdcpp/src/model/diffusion/mmdit.hpp b/otherarch/sdcpp/src/model/diffusion/mmdit.hpp index d8e76dfb9c61..6731b5fb5257 100644 --- a/otherarch/sdcpp/src/model/diffusion/mmdit.hpp +++ b/otherarch/sdcpp/src/model/diffusion/mmdit.hpp @@ -136,11 +136,15 @@ struct MMDiTConfig { }; struct PatchEmbed : public GGMLBlock { - // 2D Image to Patch Embedding + // 2D/3D Image to Patch Embedding protected: + bool is_3d; bool flatten; bool dynamic_img_pad; - int patch_size; + int patch_t; + int patch_h; + int patch_w; + int64_t embed_dim; public: PatchEmbed(int64_t img_size = 224, @@ -149,42 +153,90 @@ struct PatchEmbed : public GGMLBlock { int64_t embed_dim = 1536, bool bias = true, bool flatten = true, - bool dynamic_img_pad = true) - : patch_size(patch_size), + bool dynamic_img_pad = true, + bool is_3d = false) + : patch_t(is_3d ? patch_size : 1), + patch_h(patch_size), + patch_w(patch_size), + embed_dim(embed_dim), flatten(flatten), - dynamic_img_pad(dynamic_img_pad) { + dynamic_img_pad(dynamic_img_pad), + is_3d(is_3d) { // img_size is always None // patch_size is always 2 // in_chans is always 16 // norm_layer is always False // strict_img_size is always true, but not used - blocks["proj"] = std::shared_ptr(new Conv2d(in_chans, - embed_dim, - {patch_size, patch_size}, - {patch_size, patch_size}, - {0, 0}, - {1, 1}, - bias)); + if (is_3d) { + blocks["proj"] = std::make_shared(in_chans, + embed_dim, + std::tuple{patch_size, patch_size, patch_size}, + std::tuple{patch_size, patch_size, patch_size}, + std::tuple{0, 0, 0}, + std::tuple{1, 1, 1}, + bias); + } else { + blocks["proj"] = std::make_shared(in_chans, + embed_dim, + std::pair{patch_size, patch_size}, + std::pair{patch_size, patch_size}, + std::pair{0, 0}, + std::pair{1, 1}, + bias); + } + } + + PatchEmbed(int64_t img_size, + std::tuple patch_size, + int64_t in_chans, + int64_t embed_dim, + bool bias = true, + bool flatten = true, + bool dynamic_img_pad = true) + : patch_t(std::get<0>(patch_size)), + patch_h(std::get<1>(patch_size)), + patch_w(std::get<2>(patch_size)), + embed_dim(embed_dim), + flatten(flatten), + dynamic_img_pad(dynamic_img_pad), + is_3d(true) { + SD_UNUSED(img_size); + blocks["proj"] = std::make_shared(in_chans, + embed_dim, + patch_size, + patch_size, + std::tuple{0, 0, 0}, + std::tuple{1, 1, 1}, + bias); } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - // x: [N, C, H, W] - // return: [N, H*W, embed_dim] - auto proj = std::dynamic_pointer_cast(blocks["proj"]); + // x: [N, C, H, W] or [N*C, T, H, W] + // return: [N, h_len*w_len, embed_dim] or [N, t_len*h_len*w_len, embed_dim] + auto proj = std::dynamic_pointer_cast(blocks["proj"]); if (dynamic_img_pad) { int64_t W = x->ne[0]; int64_t H = x->ne[1]; - int pad_h = (patch_size - H % patch_size) % patch_size; - int pad_w = (patch_size - W % patch_size) % patch_size; - x = ggml_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0); // TODO: reflect pad mode + int pad_t = 0; + int pad_h = (patch_h - static_cast(H % patch_h)) % patch_h; + int pad_w = (patch_w - static_cast(W % patch_w)) % patch_w; + if (is_3d) { + int64_t T = x->ne[2]; + pad_t = (patch_t - static_cast(T % patch_t)) % patch_t; + } + x = ggml_pad(ctx->ggml_ctx, x, pad_w, pad_h, pad_t, 0); // TODO: reflect pad mode } - x = proj->forward(ctx, x); + x = proj->forward(ctx, x); // [N, C, h_len, w_len] or [N*C, t_len, h_len, w_len] if (flatten) { - x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]); - x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + if (is_3d) { + x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], embed_dim, x->ne[3] / embed_dim); // [N, C, t_len*h_len*w_len] + } else { + x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]); // [N, C, h_len*w_len] + } + x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [N, h_len*w_len, C] } return x; } diff --git a/otherarch/sdcpp/src/model/diffusion/model.hpp b/otherarch/sdcpp/src/model/diffusion/model.hpp index 6e04f9f1f5be..31c5048330df 100644 --- a/otherarch/sdcpp/src/model/diffusion/model.hpp +++ b/otherarch/sdcpp/src/model/diffusion/model.hpp @@ -26,6 +26,7 @@ struct RefImageParams { RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA; int vlm_min_size = -1; int vlm_max_size = -1; + bool resize_vae_to_target = false; }; const std::unordered_map REF_IMAGE_PRESETS = { @@ -34,6 +35,7 @@ const std::unordered_map REF_IMAGE_PRESETS = { {"flux2", {false, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::NONE, -1, -1}}, {"qwen", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}}, {"qwen_layered", {true, true, Rope::RefIndexMode::DECREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}}, + {"mage_flow", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, -1, 384, true}}, {"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}}, @@ -87,6 +89,13 @@ struct MiniT2IDiffusionExtra { const sd::Tensor* mask = nullptr; }; +struct HunyuanVideoDiffusionExtra { + const sd::Tensor* guidance = nullptr; + const sd::Tensor* byt5 = nullptr; + const sd::Tensor* vision = nullptr; + const sd::Tensor* timestep_r = nullptr; +}; + using DiffusionExtraParams = std::variant; + MiniT2IDiffusionExtra, + HunyuanVideoDiffusionExtra>; struct DiffusionParams { const sd::Tensor* x = nullptr; diff --git a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp index c6c3622f9650..6f08d6180828 100644 --- a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp @@ -103,9 +103,9 @@ namespace Qwen { } } - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* timesteps, - ggml_tensor* addition_t_cond = nullptr) { + virtual ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* timesteps, + ggml_tensor* addition_t_cond = nullptr) { // timesteps: [N,] // return: [N, embedding_dim] auto timestep_embedder = std::dynamic_pointer_cast(blocks["timestep_embedder"]); @@ -416,10 +416,14 @@ namespace Qwen { public: QwenImageModel() {} - QwenImageModel(QwenImageConfig config) + QwenImageModel(QwenImageConfig config, + std::shared_ptr time_text_embed = nullptr) : config(config) { - int64_t inner_dim = config.num_attention_heads * config.attention_head_dim; - blocks["time_text_embed"] = std::shared_ptr(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond)); + int64_t inner_dim = config.num_attention_heads * config.attention_head_dim; + if (time_text_embed == nullptr) { + time_text_embed = std::make_shared(inner_dim, config.use_additional_t_cond); + } + blocks["time_text_embed"] = std::move(time_text_embed); blocks["txt_norm"] = std::shared_ptr(new RMSNorm(config.joint_attention_dim, 1e-6f)); blocks["img_in"] = std::shared_ptr(new Linear(config.in_channels, inner_dim)); blocks["txt_in"] = std::shared_ptr(new Linear(config.joint_attention_dim, inner_dim)); diff --git a/otherarch/sdcpp/src/model/te/llm.hpp b/otherarch/sdcpp/src/model/te/llm.hpp index 3c5e9bb541bc..e5c521ef1066 100644 --- a/otherarch/sdcpp/src/model/te/llm.hpp +++ b/otherarch/sdcpp/src/model/te/llm.hpp @@ -285,7 +285,7 @@ namespace LLM { bool add_unit_offset = false) : hidden_size(hidden_size), eps(eps), add_unit_offset(add_unit_offset) {} - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { 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/te/t5.hpp b/otherarch/sdcpp/src/model/te/t5.hpp index 6d2326f94348..90c781e1da3b 100644 --- a/otherarch/sdcpp/src/model/te/t5.hpp +++ b/otherarch/sdcpp/src/model/te/t5.hpp @@ -18,6 +18,7 @@ struct T5Config { int64_t num_layers = 24; int64_t model_dim = 4096; + int64_t inner_dim = 4096; int64_t ff_dim = 10240; int64_t num_heads = 64; int64_t vocab_size = 32128; @@ -53,6 +54,7 @@ struct T5Config { if (q->n_dims == 2) { config.model_dim = q->ne[0]; int64_t inner_dim = q->ne[1]; + config.inner_dim = inner_dim; // Flan-T5/T5 uses d_kv=64 for common sizes. if (inner_dim % 64 == 0) { config.num_heads = inner_dim / 64; @@ -357,7 +359,7 @@ struct T5 : public GGMLBlock { : config(config) { blocks["encoder"] = std::shared_ptr(new T5Stack(config.num_layers, config.model_dim, - config.model_dim, + config.inner_dim, config.ff_dim, config.num_heads, config.relative_attention)); diff --git a/otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp b/otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp new file mode 100644 index 000000000000..b161261813b7 --- /dev/null +++ b/otherarch/sdcpp/src/model/vae/hunyuan_vae.hpp @@ -0,0 +1,834 @@ +#ifndef __SD_MODEL_VAE_HUNYUAN_VAE_HPP__ +#define __SD_MODEL_VAE_HUNYUAN_VAE_HPP__ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "model/vae/wan_vae.hpp" +#include "model_manager.h" + +namespace Hunyuan { + constexpr int HUNYUAN_VIDEO_VAE_GRAPH_SIZE = 65536; + constexpr int HUNYUAN_VIDEO_VAE_GRAPH_SIZE_PER_LATENT_FRAME = 8192; + constexpr int HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE = 1; + + struct TemporalConvCarry { + const std::vector* input = nullptr; + std::vector* output = nullptr; + size_t input_index = 0; + + bool is_continuation() const { + return input != nullptr; + } + + ggml_tensor* take() { + GGML_ASSERT(input != nullptr && input_index < input->size()); + return (*input)[input_index++]; + } + + void push(ggml_tensor* tensor) { + if (output != nullptr) { + output->push_back(tensor); + } + } + + void finish() const { + GGML_ASSERT(input == nullptr || input_index == input->size()); + } + }; + + static ggml_tensor* repeat_interleave_channels(GGMLRunnerContext* ctx, + ggml_tensor* x, + int64_t repeats, + int64_t width, + int64_t height, + int64_t frames) { + GGML_ASSERT(repeats > 0); + GGML_ASSERT(width * height * frames == x->ne[0] * x->ne[1] * x->ne[2]); + int64_t channels = x->ne[3]; + if (repeats == 1) { + return ggml_reshape_4d(ctx->ggml_ctx, x, width, height, frames, channels); + } + x = ggml_reshape_3d(ctx->ggml_ctx, x, width * height * frames, 1, channels); + auto target = ggml_new_tensor_3d(ctx->ggml_ctx, x->type, width * height * frames, repeats, channels); + x = ggml_repeat(ctx->ggml_ctx, x, target); + return ggml_reshape_4d(ctx->ggml_ctx, x, width, height, frames, channels * repeats); + } + + class CausalConv3d : public GGMLBlock { + protected: + std::tuple kernel_size; + + public: + CausalConv3d(int64_t in_channels, + int64_t out_channels, + std::tuple kernel_size, + std::tuple stride = {1, 1, 1}, + std::tuple padding = {0, 0, 0}, + std::tuple dilation = {1, 1, 1}, + bool bias = true) + : kernel_size(kernel_size) { + blocks["conv"] = std::make_shared(in_channels, out_channels, kernel_size, stride, padding, dilation, bias); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + TemporalConvCarry* carry = nullptr) { + // x: [N*IC, ID, IH, IW] + // result: x: [N*OC, OD, OH, OW] + // assert N == 1 + + auto conv = std::dynamic_pointer_cast(blocks["conv"]); + + int pad_w = std::get<2>(kernel_size) / 2; + int pad_h = std::get<1>(kernel_size) / 2; + int pad_t = std::get<0>(kernel_size) - 1; + std::vector temporal_frames; + temporal_frames.reserve(x->ne[2] + pad_t); + if (pad_t > 0) { + if (carry != nullptr && carry->is_continuation()) { + auto previous = carry->take(); + GGML_ASSERT(previous->ne[2] <= pad_t); + for (int64_t frame = 0; frame < previous->ne[2]; frame++) { + temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, previous, 2, frame, frame + 1)); + } + for (int64_t frame = previous->ne[2]; frame < pad_t; frame++) { + temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1)); + } + } else { + auto first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); + for (int frame = 0; frame < pad_t; frame++) { + temporal_frames.push_back(first); + } + } + } + for (int64_t frame = 0; frame < x->ne[2]; frame++) { + temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + 1)); + } + + if (pad_t > 0 && carry != nullptr && carry->output != nullptr) { + ggml_tensor* next = nullptr; + for (int frame = pad_t; frame > 0; frame--) { + auto item = temporal_frames[temporal_frames.size() - frame]; + next = next == nullptr ? item : ggml_concat(ctx->ggml_ctx, next, item, 2); + } + carry->push(ggml_cont(ctx->ggml_ctx, next)); + } + + ggml_tensor* padded = nullptr; + for (auto frame : temporal_frames) { + padded = padded == nullptr ? frame : ggml_concat(ctx->ggml_ctx, padded, frame, 2); + } + auto replicate_pad = [&](ggml_tensor* input, int dim, int left, int right) { + if (left > 0) { + auto first = ggml_ext_slice(ctx->ggml_ctx, input, dim, 0, 1); + for (int i = 0; i < left; i++) { + input = ggml_concat(ctx->ggml_ctx, first, input, dim); + } + } + if (right > 0) { + auto last = ggml_ext_slice(ctx->ggml_ctx, input, dim, input->ne[dim] - 1, input->ne[dim]); + for (int i = 0; i < right; i++) { + input = ggml_concat(ctx->ggml_ctx, input, last, dim); + } + } + return input; + }; + padded = replicate_pad(padded, 0, pad_w, pad_w); + padded = replicate_pad(padded, 1, pad_h, pad_h); + return conv->forward(ctx, padded); + } + }; + + class AttnBlock : public UnaryBlock { + protected: + int64_t in_channels; + + public: + AttnBlock(int64_t in_channels) + : in_channels(in_channels) { + blocks["norm"] = std::make_shared(in_channels); + blocks["q"] = std::make_shared(in_channels, in_channels, std::tuple{1, 1, 1}); + blocks["k"] = std::make_shared(in_channels, in_channels, std::tuple{1, 1, 1}); + blocks["v"] = std::make_shared(in_channels, in_channels, std::tuple{1, 1, 1}); + blocks["proj_out"] = std::make_shared(in_channels, in_channels, std::tuple{1, 1, 1}); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x) override { + // x: [b*c, t, h, w] + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto q_proj = std::dynamic_pointer_cast(blocks["q"]); + auto k_proj = std::dynamic_pointer_cast(blocks["k"]); + auto v_proj = std::dynamic_pointer_cast(blocks["v"]); + auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); + + const int64_t b = x->ne[3] / in_channels; + + auto identity = x; + + x = norm->forward(ctx, x); + + const int64_t c = x->ne[3] / b; + const int64_t t = x->ne[2]; + const int64_t h = x->ne[1]; + const int64_t w = x->ne[0]; + + auto q = q_proj->forward(ctx, x); // [b*c, t, h, w] + auto k = k_proj->forward(ctx, x); // [b*c, t, h, w] + auto v = v_proj->forward(ctx, x); // [b*c, t, h, w] + + q = ggml_reshape_3d(ctx->ggml_ctx, q, w * h * t, c, b); // [b, c, t*h*w] + q = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 1, 0, 2, 3)); // [b, t*h*w, c] + + k = ggml_reshape_3d(ctx->ggml_ctx, k, w * h * t, c, b); // [b, c, t*h*w] + k = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 1, 0, 2, 3)); // [b, t*h*w, c] + + v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b); // [b, c, t*h*w] + v = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [b, t*h*w, c] + + x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [b, t*h*w, c] + + x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [b, c, t*h*w] + x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b); // [b*c, t, h, w] + + x = proj_out->forward(ctx, x); + + x = ggml_add(ctx->ggml_ctx, x, identity); + return x; + } + }; + + class ResnetBlock : public UnaryBlock { + protected: + int64_t in_channels; + int64_t out_channels; + + public: + ResnetBlock(int64_t in_channels, + int64_t out_channels) + : in_channels(in_channels), + out_channels(out_channels) { + blocks["norm1"] = std::make_shared(in_channels); + blocks["conv1"] = std::make_shared(in_channels, out_channels, std::tuple{3, 3, 3}); + + blocks["norm2"] = std::make_shared(out_channels); + blocks["conv2"] = std::make_shared(out_channels, out_channels, std::tuple{3, 3, 3}); + + if (out_channels != in_channels) { + blocks["nin_shortcut"] = std::make_shared(in_channels, out_channels, std::tuple{1, 1, 1}); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + return forward(ctx, x, nullptr); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + TemporalConvCarry* carry) { + // x: [B*IC, IT, OH, OW] + // return: [B*OC, OT, OH, OW] + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); + + auto h = x; + h = norm1->forward(ctx, h); + h = ggml_silu_inplace(ctx->ggml_ctx, h); // swish + h = conv1->forward(ctx, h, carry); + + h = norm2->forward(ctx, h); + h = ggml_silu_inplace(ctx->ggml_ctx, h); // swish + // dropout, skip for inference + h = conv2->forward(ctx, h, carry); + + // skip connection + if (out_channels != in_channels) { + auto nin_shortcut = std::dynamic_pointer_cast(blocks["nin_shortcut"]); + + x = nin_shortcut->forward(ctx, x); // [B*OC, OT, OH, OW] + } + + h = ggml_add(ctx->ggml_ctx, h, x); + return h; // [B*OC, OT, OH, OW] + } + }; + + class Upsample : public GGMLBlock { + protected: + int64_t in_channels; + int64_t out_channels; + int64_t factor_t; + int64_t factor_s; + int64_t factor; + int64_t repeats; + + public: + Upsample(int64_t in_channels, int64_t out_channels, bool add_temporal_upsample) + : in_channels(in_channels), out_channels(out_channels) { + if (add_temporal_upsample) { + factor_t = 2; + } else { + factor_t = 1; + } + factor_s = 2; + factor = factor_t * factor_s * factor_s; + GGML_ASSERT(out_channels * factor % in_channels == 0); + repeats = out_channels * factor / in_channels; + blocks["conv"] = std::make_shared(in_channels, out_channels * factor, std::tuple{3, 3, 3}); + } + + static ggml_tensor* _pixel_shuffle_3d(GGMLRunnerContext* ctx, + ggml_tensor* x, + int64_t factor_t, + int64_t factor_s, + int64_t B = 1) { + // x: [B*factor*C, T, H, W] + // return: [B*C, T*factor_t, H*factor_s, W*factor_s] + GGML_ASSERT(B == 1); + int64_t factor = factor_t * factor_s * factor_s; + int64_t C = x->ne[3] / factor; + int64_t T = x->ne[2]; + int64_t H = x->ne[1]; + int64_t W = x->ne[0]; + + x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, C, factor); // [factor, C, T*H, W] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // [C, factor, T*H, W] + x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor_s, factor_s * factor_t * C); // [C*factor_t*factor_s, factor_s, T*H, W] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); // [C*factor_t*factor_s, T*H, W, factor_s] + x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, H * T, factor_s, factor_t * C); // [C*factor_t, factor_s, T*H, W*factor_s] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // [C*factor_t, T*H, factor_s, W*factor_s] + x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, T, factor_t, C); // [C, factor_t, T, H*factor_s*W*factor_s] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // [C, T, factor_t, H*factor_s*W*factor_s] + x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s * H, factor_t * T, C); // [C, T*factor_t, H*factor_s, W*factor_s] + + return x; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + TemporalConvCarry* carry = nullptr) { + // x: [B*IC, T, H, W] + // return: [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s] + const int64_t B = x->ne[3] / in_channels; + GGML_ASSERT(B == 1); + + auto conv = std::dynamic_pointer_cast(blocks["conv"]); + + const bool continuation = carry != nullptr && carry->is_continuation(); + auto h = conv->forward(ctx, x, carry); // [B*factor*OC, T, H, W] + + ggml_tensor* shortcut = nullptr; + if (factor_t == 2 && !continuation) { + auto h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1); // [B*factor*OC, 1, H, W] + h_first = _pixel_shuffle_3d(ctx, h_first, 1, factor_s, B); // [B*2*OC, 1, H*factor_s, W*factor_s] + h_first = ggml_ext_slice(ctx->ggml_ctx, h_first, 3, 0, out_channels); // [B*OC, 1, H*factor_s, W*factor_s] + + auto x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); + x_first = repeat_interleave_channels(ctx, x_first, repeats / 2, x->ne[0], x->ne[1], 1); + x_first = _pixel_shuffle_3d(ctx, x_first, 1, factor_s, B); + + if (x->ne[2] == 1) { + return ggml_add(ctx->ggml_ctx, h_first, x_first); + } + + auto h_next = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]); // [B*factor*OC, T - 1, H, W] + h_next = _pixel_shuffle_3d(ctx, h_next, factor_t, factor_s, B); // [B*OC, (T - 1)*factor_t, H*factor_s, W*factor_s] + + h = ggml_concat(ctx->ggml_ctx, h_first, h_next, 2); // [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s] + + auto x_next = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]); + x_next = repeat_interleave_channels(ctx, x_next, repeats, x->ne[0], x->ne[1], x->ne[2] - 1); + x_next = _pixel_shuffle_3d(ctx, x_next, factor_t, factor_s, B); + + shortcut = ggml_concat(ctx->ggml_ctx, x_first, x_next, 2); // [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s] + } else { + h = _pixel_shuffle_3d(ctx, h, factor_t, factor_s, B); + shortcut = repeat_interleave_channels(ctx, x, repeats, x->ne[0], x->ne[1], x->ne[2]); + shortcut = _pixel_shuffle_3d(ctx, shortcut, factor_t, factor_s, B); // [B*OC, T*factor_t, H*factor_s, W*factor_s] + } + + return ggml_add(ctx->ggml_ctx, h, shortcut); + } + }; + + static ggml_tensor* pixel_unshuffle_3d(GGMLRunnerContext* ctx, + ggml_tensor* x, + int64_t factor_t, + int64_t factor_s) { + GGML_ASSERT(x->ne[0] % factor_s == 0); + GGML_ASSERT(x->ne[1] % factor_s == 0); + GGML_ASSERT(x->ne[2] % factor_t == 0); + int64_t W = x->ne[0] / factor_s; + int64_t H = x->ne[1] / factor_s; + int64_t T = x->ne[2] / factor_t; + int64_t C = x->ne[3]; + int64_t factor = factor_t * factor_s * factor_s; + + x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, factor_t, T, C); + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); + x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s, H * T, factor_t * C); + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); + x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s, W, H * T, factor_s * factor_t * C); + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); + x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor, C); + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); + return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, C * factor); + } + + static ggml_tensor* mean_channel_groups(GGMLRunnerContext* ctx, + ggml_tensor* x, + int64_t group_size) { + GGML_ASSERT(group_size > 0); + GGML_ASSERT(x->ne[3] % group_size == 0); + if (group_size == 1) { + return x; + } + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t T = x->ne[2]; + int64_t spatial = W * H * T; + int64_t groups = x->ne[3] / group_size; + x = ggml_reshape_3d(ctx->ggml_ctx, x, spatial, group_size, groups); + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = ggml_sum_rows(ctx->ggml_ctx, x); + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, groups); + return ggml_scale(ctx->ggml_ctx, x, 1.f / static_cast(group_size)); + } + + class Downsample : public GGMLBlock { + protected: + int64_t in_channels; + int64_t out_channels; + int64_t factor_t; + int64_t factor_s = 2; + int64_t factor; + int64_t group_size; + + public: + Downsample(int64_t in_channels, int64_t out_channels, bool add_temporal_downsample) + : in_channels(in_channels), + out_channels(out_channels), + factor_t(add_temporal_downsample ? 2 : 1), + factor(factor_t * factor_s * factor_s), + group_size(factor * in_channels / out_channels) { + GGML_ASSERT(out_channels % factor == 0); + GGML_ASSERT(factor * in_channels % out_channels == 0); + blocks["conv"] = std::make_shared(in_channels, + out_channels / factor, + std::tuple{3, 3, 3}); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto conv = std::dynamic_pointer_cast(blocks["conv"]); + auto h = conv->forward(ctx, x); + + ggml_tensor* h_first = nullptr; + ggml_tensor* x_first = nullptr; + if (factor_t == 2) { + h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1); + h_first = pixel_unshuffle_3d(ctx, h_first, 1, factor_s); + h_first = ggml_concat(ctx->ggml_ctx, h_first, h_first, 3); + + x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); + x_first = pixel_unshuffle_3d(ctx, x_first, 1, factor_s); + x_first = mean_channel_groups(ctx, x_first, group_size / 2); + + if (x->ne[2] == 1) { + return ggml_add(ctx->ggml_ctx, h_first, x_first); + } + h = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]); + x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]); + } + + GGML_ASSERT(h->ne[2] % factor_t == 0); + h = pixel_unshuffle_3d(ctx, h, factor_t, factor_s); + x = pixel_unshuffle_3d(ctx, x, factor_t, factor_s); + x = mean_channel_groups(ctx, x, group_size); + + if (factor_t == 2) { + h = ggml_concat(ctx->ggml_ctx, h_first, h, 2); + x = ggml_concat(ctx->ggml_ctx, x_first, x, 2); + } + return ggml_add(ctx->ggml_ctx, h, x); + } + }; + + class MidBlock : public UnaryBlock { + protected: + int64_t in_channels; + int num_layers; + bool add_attention; + + public: + MidBlock(int64_t in_channels, + int num_layers = 1, + bool add_attention = true) + : in_channels(in_channels), + num_layers(num_layers), + add_attention(add_attention) { + blocks["block_1"] = std::make_shared(in_channels, in_channels); + for (int i = 0; i < num_layers; i++) { + if (add_attention) { + blocks["attn_" + std::to_string(i + 1)] = std::make_shared(in_channels); + } + blocks["block_" + std::to_string(i + 2)] = std::make_shared(in_channels, in_channels); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + // x: [B*C, T, H, W] + // return: [B*C, T, H, W] + auto block_1 = std::dynamic_pointer_cast(blocks["block_1"]); + + x = block_1->forward(ctx, x); + for (int i = 0; i < num_layers; i++) { + if (add_attention) { + auto block = std::dynamic_pointer_cast(blocks["attn_" + std::to_string(i + 1)]); + x = block->forward(ctx, x); + } + auto block = std::dynamic_pointer_cast(blocks["block_" + std::to_string(i + 2)]); + x = block->forward(ctx, x); + } + return x; + } + }; + + class UpBlock : public UnaryBlock { + protected: + int num_layers; + int64_t upsample_out_channels; + + public: + UpBlock(int64_t in_channels, + int64_t out_channels, + int num_layers = 1, + int64_t upsample_out_channels = 0, + bool add_temporal_upsample = true) + : num_layers(num_layers), + upsample_out_channels(upsample_out_channels) { + for (int i = 0; i < num_layers; i++) { + int64_t IC = i == 0 ? in_channels : out_channels; + blocks["block." + std::to_string(i)] = std::make_shared(IC, out_channels); + } + if (upsample_out_channels > 0) { + blocks["upsample"] = std::make_shared(out_channels, upsample_out_channels, add_temporal_upsample); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + return forward(ctx, x, nullptr); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + TemporalConvCarry* carry) { + // x: [B*IC, T, H, W] + // return: [B*OC, T, H, W] or [B*OC, T, H*2, W*2] or [B*OC, T*2, H*2, W*2] + for (int i = 0; i < num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["block." + std::to_string(i)]); + x = block->forward(ctx, x, carry); + } + if (upsample_out_channels > 0) { + auto upsample = std::dynamic_pointer_cast(blocks["upsample"]); + x = upsample->forward(ctx, x, carry); + } + return x; + } + }; + + class DownBlock : public UnaryBlock { + protected: + int num_layers; + int64_t downsample_out_channels; + + public: + DownBlock(int64_t in_channels, + int64_t out_channels, + int num_layers, + int64_t downsample_out_channels = 0, + bool add_temporal_downsample = false) + : num_layers(num_layers), + downsample_out_channels(downsample_out_channels) { + for (int i = 0; i < num_layers; i++) { + int64_t IC = i == 0 ? in_channels : out_channels; + blocks["block." + std::to_string(i)] = std::make_shared(IC, out_channels); + } + if (downsample_out_channels > 0) { + blocks["downsample"] = std::make_shared(out_channels, + downsample_out_channels, + add_temporal_downsample); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + for (int i = 0; i < num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["block." + std::to_string(i)]); + x = block->forward(ctx, x); + } + if (downsample_out_channels > 0) { + auto downsample = std::dynamic_pointer_cast(blocks["downsample"]); + x = downsample->forward(ctx, x); + } + return x; + } + }; + + class Encoder : public GGMLBlock { + protected: + int64_t z_channels; + std::vector block_out_channels; + + public: + Encoder(int64_t in_channels = 3, + int64_t z_channels = 32, + std::vector block_out_channels = {128, 256, 512, 1024, 1024}, + int layers_per_block = 2, + int spatial_compression_ratio = 16, + int temporal_compression_ratio = 4, + bool downsample_match_channel = true) + : z_channels(z_channels), + block_out_channels(std::move(block_out_channels)) { + blocks["conv_in"] = std::make_shared(in_channels, + this->block_out_channels[0], + std::tuple{3, 3, 3}); + + int spatial_depth = static_cast(std::log2(static_cast(spatial_compression_ratio))); + int temporal_start = static_cast(std::log2(static_cast(spatial_compression_ratio / temporal_compression_ratio))); + int64_t channels = this->block_out_channels[0]; + for (int i = 0; i < static_cast(this->block_out_channels.size()); i++) { + int64_t out_channels = this->block_out_channels[i]; + if (i < spatial_depth) { + int64_t next_channels = downsample_match_channel ? this->block_out_channels[i + 1] : out_channels; + blocks["down." + std::to_string(i)] = std::make_shared(channels, + out_channels, + layers_per_block, + next_channels, + i >= temporal_start); + channels = next_channels; + } else { + blocks["down." + std::to_string(i)] = std::make_shared(channels, + out_channels, + layers_per_block); + channels = out_channels; + } + } + + blocks["mid"] = std::make_shared(channels); + blocks["norm_out"] = std::make_shared(channels); + blocks["conv_out"] = std::make_shared(channels, + z_channels * 2, + std::tuple{3, 3, 3}); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto conv_in = std::dynamic_pointer_cast(blocks["conv_in"]); + auto mid = std::dynamic_pointer_cast(blocks["mid"]); + auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); + auto conv_out = std::dynamic_pointer_cast(blocks["conv_out"]); + + x = conv_in->forward(ctx, x); + for (int i = 0; i < static_cast(block_out_channels.size()); i++) { + auto down = std::dynamic_pointer_cast(blocks["down." + std::to_string(i)]); + x = down->forward(ctx, x); + } + x = mid->forward(ctx, x); + + auto shortcut = mean_channel_groups(ctx, x, x->ne[3] / (z_channels * 2)); + x = norm_out->forward(ctx, x); + x = ggml_silu_inplace(ctx->ggml_ctx, x); + x = conv_out->forward(ctx, x); + x = ggml_add(ctx->ggml_ctx, x, shortcut); + return ggml_ext_slice(ctx->ggml_ctx, x, 3, 0, z_channels); + } + }; + + class Decoder : public GGMLBlock { + protected: + int64_t repeats; + std::vector block_out_channels; + + public: + Decoder(int64_t in_channels = 32, + int64_t out_channels = 3, + std::vector block_out_channels = {1024, 1024, 512, 256, 128}, + int layers_per_block = 2, + int spatial_compression_ratio = 16, + int temporal_compression_ratio = 4, + bool upsample_match_channel = true) + : block_out_channels(std::move(block_out_channels)) { + repeats = this->block_out_channels[0] / in_channels; + blocks["conv_in"] = std::make_shared(in_channels, this->block_out_channels[0], std::tuple{3, 3, 3}); + blocks["mid"] = std::make_shared(this->block_out_channels[0]); + + int64_t IC = this->block_out_channels[0]; + for (int i = 0; i < this->block_out_channels.size(); i++) { + int64_t OC = this->block_out_channels[i]; + bool add_spatial_upsample = i < std::log2(static_cast(spatial_compression_ratio)); + bool add_temporal_upsample = i < std::log2(static_cast(temporal_compression_ratio)); + + if (add_spatial_upsample || add_temporal_upsample) { + int64_t upsample_out_channels = upsample_match_channel ? this->block_out_channels[i + 1] : OC; + blocks["up." + std::to_string(i)] = std::make_shared(IC, OC, layers_per_block + 1, upsample_out_channels, add_temporal_upsample); + IC = upsample_out_channels; + } else { + blocks["up." + std::to_string(i)] = std::make_shared(IC, OC, layers_per_block + 1, 0, false); + } + } + + blocks["norm_out"] = std::make_shared(this->block_out_channels.back()); + blocks["conv_out"] = std::make_shared(this->block_out_channels.back(), out_channels, std::tuple{3, 3, 3}); + } + + struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) { + auto conv_in = std::dynamic_pointer_cast(blocks["conv_in"]); + auto mid_block = std::dynamic_pointer_cast(blocks["mid"]); + auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); + auto conv_out = std::dynamic_pointer_cast(blocks["conv_out"]); + + auto h = conv_in->forward(ctx, z); + + auto shortcut = repeat_interleave_channels(ctx, z, repeats, z->ne[0], z->ne[1], z->ne[2]); + h = ggml_add(ctx->ggml_ctx, h, shortcut); + + h = mid_block->forward(ctx, h); + + ggml_tensor* output = nullptr; + std::vector carry_input; + const int64_t frames = h->ne[2]; + for (int64_t start = 0; start < frames; start += HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE) { + const int64_t end = std::min(start + HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE, frames); + auto chunk = ggml_ext_slice(ctx->ggml_ctx, h, 2, start, end); + + std::vector carry_output; + TemporalConvCarry carry{ + start == 0 ? nullptr : &carry_input, + end == frames ? nullptr : &carry_output, + }; + + for (int i = 0; i < block_out_channels.size(); i++) { + auto up_block = std::dynamic_pointer_cast(blocks["up." + std::to_string(i)]); + chunk = up_block->forward(ctx, chunk, &carry); + } + + chunk = norm_out->forward(ctx, chunk); + chunk = ggml_silu_inplace(ctx->ggml_ctx, chunk); // nonlinearity/swish + chunk = conv_out->forward(ctx, chunk, &carry); + carry.finish(); + + output = output == nullptr ? chunk : ggml_concat(ctx->ggml_ctx, output, chunk, 2); + carry_input = std::move(carry_output); + } + return output; + } + }; + + class HunyuanVideoVAERunner : public VAE { + protected: + bool decode_only; + Encoder encoder; + Decoder decoder; + + public: + HunyuanVideoVAERunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map, + const std::string& prefix, + bool decode_only, + SDVersion version, + std::shared_ptr weight_manager = nullptr) + : VAE(version, backend, prefix, weight_manager), + decode_only(decode_only || + tensor_storage_map.find(prefix + ".encoder.conv_in.conv.weight") == tensor_storage_map.end()) { + if (!this->decode_only) { + encoder.init(params_ctx, tensor_storage_map, prefix + ".encoder"); + } + decoder.init(params_ctx, tensor_storage_map, prefix + ".decoder"); + } + + std::string get_desc() override { + return "hunyuan_video_vae"; + } + + void get_param_tensors(std::map& tensors) override { + if (!decode_only) { + encoder.get_param_tensors(tensors, weight_prefix + ".encoder"); + } + decoder.get_param_tensors(tensors, weight_prefix + ".decoder"); + } + + int get_encoder_output_channels(int input_channels) override { + SD_UNUSED(input_channels); + return 32; + } + + sd::Tensor vae_output_to_latents(const sd::Tensor& vae_output, + std::shared_ptr rng) override { + SD_UNUSED(rng); + return vae_output; + } + + sd::Tensor diffusion_to_vae_latents(const sd::Tensor& latents) override { + return latents / 1.03682f; + } + + sd::Tensor vae_to_diffusion_latents(const sd::Tensor& latents) override { + return latents * 1.03682f; + } + + ggml_cgraph* build_graph(const sd::Tensor& input_tensor, bool decode_graph) { + size_t graph_size = HUNYUAN_VIDEO_VAE_GRAPH_SIZE; + if (decode_graph) { + graph_size = std::max(graph_size, + HUNYUAN_VIDEO_VAE_GRAPH_SIZE_PER_LATENT_FRAME * + static_cast(input_tensor.shape()[2])); + } + ggml_cgraph* gf = new_graph_custom(graph_size); + ggml_tensor* input = make_input(input_tensor); + auto runner_ctx = get_context(); + ggml_tensor* output = decode_graph ? decoder.forward(&runner_ctx, input) + : encoder.forward(&runner_ctx, input); + ggml_build_forward_expand(gf, output); + return gf; + } + + sd::Tensor _compute(const int n_threads, + const sd::Tensor& input, + bool decode_graph) override { + if (!decode_graph && decode_only) { + LOG_ERROR("Hunyuan Video VAE encoder weights are not available"); + return {}; + } + + sd::Tensor expanded; + if (input.dim() == 4) { + expanded = input.unsqueeze(2); + } + const auto& graph_input = expanded.empty() ? input : expanded; + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(graph_input, decode_graph); + }; + auto output = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, + n_threads, + true, + true, + true), + graph_input.dim()); + if (!output.empty() && input.dim() == 4) { + output.squeeze_(2); + } + return output; + } + }; + +} // namespace Hunyuan + +#endif // __SD_MODEL_VAE_HUNYUAN_VAE_HPP__ diff --git a/otherarch/sdcpp/src/model/vae/mage_vae.hpp b/otherarch/sdcpp/src/model/vae/mage_vae.hpp new file mode 100644 index 000000000000..39075a8365cb --- /dev/null +++ b/otherarch/sdcpp/src/model/vae/mage_vae.hpp @@ -0,0 +1,521 @@ +#ifndef __SD_MODEL_VAE_MAGE_VAE_HPP__ +#define __SD_MODEL_VAE_MAGE_VAE_HPP__ + +#include "model/diffusion/dit.hpp" +#include "model/vae/vae.hpp" + +namespace MageVAE { + constexpr int MAGE_VAE_GRAPH_SIZE = 327680; + constexpr int HIDDEN_SIZE = 384; + constexpr int LATENT_CHANNELS = 128; + constexpr int PATCH_SIZE = 16; + + struct LayerNorm2d : public UnaryBlock { + int64_t channels; + bool affine; + std::string prefix; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + this->prefix = prefix; + if (affine) { + params["weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels); + params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels); + } + } + + LayerNorm2d(int64_t channels, bool affine = true) + : channels(channels), affine(affine) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + ggml_tensor* weight = affine ? params["weight"] : nullptr; + ggml_tensor* bias = affine ? params["bias"] : nullptr; + if (affine && ctx->weight_adapter) { + weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight"); + bias = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias"); + } + // [N, C, H, W] -> [N, H, W, C] so layer norm reduces over channels. + x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); + x = ggml_ext_layer_norm(ctx->ggml_ctx, x, weight, bias, 1e-6f); + return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); + } + }; + + inline ggml_tensor* modulate_2d(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* shift, + ggml_tensor* scale) { + shift = ggml_reshape_4d(ctx, shift, 1, 1, shift->ne[0], shift->ne[1]); + scale = ggml_reshape_4d(ctx, scale, 1, 1, scale->ne[0], scale->ne[1]); + return ggml_add(ctx, ggml_mul(ctx, x, ggml_add(ctx, scale, ggml_ext_ones(ctx, 1, 1, 1, 1))), shift); + } + + inline ggml_tensor* channel_attention(GGMLRunnerContext* ctx, + ggml_tensor* x, + Conv2d* projection) { + auto pooled = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]); + pooled = ggml_mean(ctx->ggml_ctx, pooled); + pooled = ggml_reshape_4d(ctx->ggml_ctx, pooled, 1, 1, x->ne[2], x->ne[3]); + pooled = ggml_sigmoid(ctx->ggml_ctx, projection->forward(ctx, pooled)); + return ggml_mul(ctx->ggml_ctx, x, pooled); + } + + struct TimestepEmbedder : public GGMLBlock { + TimestepEmbedder() { + blocks["mlp.0"] = std::make_shared(256, HIDDEN_SIZE); + blocks["mlp.2"] = std::make_shared(HIDDEN_SIZE, HIDDEN_SIZE); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timestep) { + auto linear_0 = std::dynamic_pointer_cast(blocks["mlp.0"]); + auto linear_2 = std::dynamic_pointer_cast(blocks["mlp.2"]); + auto x = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f); + x = linear_0->forward(ctx, x); + x = ggml_silu_inplace(ctx->ggml_ctx, x); + return linear_2->forward(ctx, x); + } + }; + + struct EncoderDiCoBlock : public UnaryBlock { + explicit EncoderDiCoBlock(int64_t channels) { + blocks["conv1"] = std::make_shared(channels, channels, std::pair{1, 1}); + blocks["conv2"] = std::make_shared(channels, channels, static_cast(channels), std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + blocks["conv3"] = std::make_shared(channels, channels, std::pair{1, 1}); + blocks["ca.1"] = std::make_shared(channels, channels, std::pair{1, 1}); + blocks["conv4"] = std::make_shared(channels, channels * 4, std::pair{1, 1}); + blocks["conv5"] = std::make_shared(channels * 4, channels, std::pair{1, 1}); + blocks["norm1"] = std::make_shared(channels); + blocks["norm2"] = std::make_shared(channels); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override { + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); + auto conv3 = std::dynamic_pointer_cast(blocks["conv3"]); + auto ca = std::dynamic_pointer_cast(blocks["ca.1"]); + auto conv4 = std::dynamic_pointer_cast(blocks["conv4"]); + auto conv5 = std::dynamic_pointer_cast(blocks["conv5"]); + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + + auto x = norm1->forward(ctx, input); + x = conv1->forward(ctx, x); + x = conv2->forward(ctx, x); + x = ggml_gelu(ctx->ggml_ctx, x); + x = channel_attention(ctx, x, ca.get()); + x = conv3->forward(ctx, x); + x = ggml_add(ctx->ggml_ctx, input, x); + auto h = norm2->forward(ctx, x); + h = conv4->forward(ctx, h); + h = ggml_gelu(ctx->ggml_ctx, h); + h = conv5->forward(ctx, h); + return ggml_add(ctx->ggml_ctx, x, h); + } + }; + + struct DiCoBlock : public GGMLBlock { + explicit DiCoBlock(int64_t channels) { + blocks["conv1"] = std::make_shared(channels, channels, std::pair{1, 1}); + blocks["conv2"] = std::make_shared(channels, channels, static_cast(channels), std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + blocks["conv3"] = std::make_shared(channels, channels, std::pair{1, 1}); + blocks["ca.1"] = std::make_shared(channels, channels, std::pair{1, 1}); + blocks["conv4"] = std::make_shared(channels, channels * 4, std::pair{1, 1}); + blocks["conv5"] = std::make_shared(channels * 4, channels, std::pair{1, 1}); + blocks["norm1"] = std::make_shared(channels, false); + blocks["norm2"] = std::make_shared(channels, false); + blocks["adaLN_modulation.1"] = std::make_shared(channels, channels * 6); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input, ggml_tensor* condition) { + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); + auto conv3 = std::dynamic_pointer_cast(blocks["conv3"]); + auto ca = std::dynamic_pointer_cast(blocks["ca.1"]); + auto conv4 = std::dynamic_pointer_cast(blocks["conv4"]); + auto conv5 = std::dynamic_pointer_cast(blocks["conv5"]); + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto ada = std::dynamic_pointer_cast(blocks["adaLN_modulation.1"]); + + auto params = ada->forward(ctx, ggml_silu(ctx->ggml_ctx, condition)); + auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 6, 0); + auto x = norm1->forward(ctx, input); + x = modulate_2d(ctx->ggml_ctx, x, chunks[0], chunks[1]); + x = conv1->forward(ctx, x); + x = conv2->forward(ctx, x); + x = ggml_gelu(ctx->ggml_ctx, x); + x = channel_attention(ctx, x, ca.get()); + x = conv3->forward(ctx, x); + auto gate_1 = ggml_reshape_4d(ctx->ggml_ctx, chunks[2], 1, 1, chunks[2]->ne[0], chunks[2]->ne[1]); + x = ggml_add(ctx->ggml_ctx, input, ggml_mul(ctx->ggml_ctx, x, gate_1)); + + auto h = norm2->forward(ctx, x); + h = modulate_2d(ctx->ggml_ctx, h, chunks[3], chunks[4]); + h = conv4->forward(ctx, h); + h = ggml_gelu(ctx->ggml_ctx, h); + h = conv5->forward(ctx, h); + auto gate_2 = ggml_reshape_4d(ctx->ggml_ctx, chunks[5], 1, 1, chunks[5]->ne[0], chunks[5]->ne[1]); + return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, h, gate_2)); + } + }; + + struct MageResnetBlock : public UnaryBlock { + explicit MageResnetBlock(int64_t channels) { + blocks["norm1"] = std::make_shared(channels); + blocks["conv1"] = std::make_shared(channels, channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + blocks["norm2"] = std::make_shared(channels); + blocks["conv2"] = std::make_shared(channels, channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); + auto x = conv1->forward(ctx, ggml_silu(ctx->ggml_ctx, norm1->forward(ctx, input))); + x = conv2->forward(ctx, ggml_silu(ctx->ggml_ctx, norm2->forward(ctx, x))); + return ggml_add(ctx->ggml_ctx, input, x); + } + }; + + inline ggml_tensor* replicate_pad_right_bottom(ggml_context* ctx, + ggml_tensor* x, + int pad_w, + int pad_h) { + if (pad_w > 0) { + auto edge = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]); + edge = ggml_repeat_4d(ctx, edge, pad_w, x->ne[1], x->ne[2], x->ne[3]); + x = ggml_concat(ctx, x, edge, 0); + } + if (pad_h > 0) { + auto edge = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]); + edge = ggml_repeat_4d(ctx, edge, x->ne[0], pad_h, x->ne[2], x->ne[3]); + x = ggml_concat(ctx, x, edge, 1); + } + return x; + } + + struct MageAttnBlock : public UnaryBlock { + int64_t channels; + int patch_size; + + MageAttnBlock(int64_t channels, int patch_size = 32) + : channels(channels), patch_size(patch_size) { + blocks["norm"] = std::make_shared(channels); + blocks["q"] = std::make_shared(channels, channels, std::pair{1, 1}); + blocks["k"] = std::make_shared(channels, channels, std::pair{1, 1}); + blocks["v"] = std::make_shared(channels, channels, std::pair{1, 1}); + blocks["proj_out"] = std::make_shared(channels, channels, std::pair{1, 1}); + } + + ggml_tensor* to_patches(ggml_context* ctx, ggml_tensor* x) { + x = DiT::patchify(ctx, x, patch_size, patch_size); + x = ggml_reshape_4d(ctx, x, patch_size * patch_size, channels, x->ne[1], x->ne[2]); + // [N, np, C, P] -> [N, np, P, C] for attention over P pixels. + x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); + return ggml_reshape_3d(ctx, x, channels, patch_size * patch_size, x->ne[2] * x->ne[3]); + } + + ggml_tensor* from_patches(ggml_context* ctx, + ggml_tensor* x, + int64_t patch_count, + int64_t batch_size, + int64_t h_patches, + int64_t w_patches) { + x = ggml_reshape_4d(ctx, x, channels, patch_size * patch_size, patch_count, batch_size); + // [N, np, P, C] -> [N, np, C, P] before spatial unpatchify. + x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); + x = ggml_reshape_3d(ctx, x, patch_size * patch_size * channels, patch_count, batch_size); + return DiT::unpatchify(ctx, x, h_patches, w_patches, patch_size, patch_size); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override { + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto q_proj = std::dynamic_pointer_cast(blocks["q"]); + auto k_proj = std::dynamic_pointer_cast(blocks["k"]); + auto v_proj = std::dynamic_pointer_cast(blocks["v"]); + auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); + + int64_t width = input->ne[0]; + int64_t height = input->ne[1]; + int64_t batch = input->ne[3]; + int pad_w = (patch_size - static_cast(width % patch_size)) % patch_size; + int pad_h = (patch_size - static_cast(height % patch_size)) % patch_size; + int64_t wp = (width + pad_w) / patch_size; + int64_t hp = (height + pad_h) / patch_size; + int64_t np = wp * hp; + + auto h = norm->forward(ctx, input); + auto q = replicate_pad_right_bottom(ctx->ggml_ctx, q_proj->forward(ctx, h), pad_w, pad_h); + auto k = replicate_pad_right_bottom(ctx->ggml_ctx, k_proj->forward(ctx, h), pad_w, pad_h); + auto v = replicate_pad_right_bottom(ctx->ggml_ctx, v_proj->forward(ctx, h), pad_w, pad_h); + q = to_patches(ctx->ggml_ctx, q); + k = to_patches(ctx->ggml_ctx, k); + v = to_patches(ctx->ggml_ctx, v); + h = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); + h = from_patches(ctx->ggml_ctx, h, np, batch, hp, wp); + if (pad_h > 0) { + h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, height); + } + if (pad_w > 0) { + h = ggml_ext_slice(ctx->ggml_ctx, h, 0, 0, width); + } + return ggml_add(ctx->ggml_ctx, input, proj_out->forward(ctx, h)); + } + }; + + struct Decoder : public UnaryBlock { + Decoder() { + blocks["conv_in"] = std::make_shared(LATENT_CHANNELS, HIDDEN_SIZE, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + blocks["block.0"] = std::make_shared(HIDDEN_SIZE); + blocks["block.1"] = std::make_shared(HIDDEN_SIZE); + blocks["block.2"] = std::make_shared(HIDDEN_SIZE); + blocks["block.3"] = std::make_shared(HIDDEN_SIZE); + blocks["block.4"] = std::make_shared(HIDDEN_SIZE); + blocks["norm_out"] = std::make_shared(HIDDEN_SIZE); + blocks["conv_out"] = std::make_shared(HIDDEN_SIZE, HIDDEN_SIZE, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + x = std::dynamic_pointer_cast(blocks["conv_in"])->forward(ctx, x); + for (int i = 0; i < 5; ++i) { + x = std::dynamic_pointer_cast(blocks["block." + std::to_string(i)])->forward(ctx, x); + } + x = std::dynamic_pointer_cast(blocks["norm_out"])->forward(ctx, x); + x = ggml_silu(ctx->ggml_ctx, x); + return std::dynamic_pointer_cast(blocks["conv_out"])->forward(ctx, x); + } + }; + + struct DConvEncoder : public UnaryBlock { + DConvEncoder() { + blocks["patch_cond_embed"] = std::make_shared(3, 768, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{PATCH_SIZE, PATCH_SIZE}); + for (int i = 0; i < 2; ++i) { + blocks["head_blocks." + std::to_string(i)] = std::make_shared(768); + } + blocks["proj_down"] = std::make_shared(768, HIDDEN_SIZE, std::pair{1, 1}); + blocks["z_proj"] = std::make_shared(LATENT_CHANNELS, HIDDEN_SIZE, std::pair{1, 1}); + blocks["fuse_proj"] = std::make_shared(HIDDEN_SIZE * 2, HIDDEN_SIZE, std::pair{1, 1}); + blocks["t_embedder"] = std::make_shared(); + for (int i = 0; i < 21; ++i) { + blocks["blocks." + std::to_string(i)] = std::make_shared(HIDDEN_SIZE); + } + blocks["norm_out"] = std::make_shared(HIDDEN_SIZE); + blocks["proj_out"] = std::make_shared(HIDDEN_SIZE, LATENT_CHANNELS * 2, std::pair{1, 1}); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image) override { + auto cond = std::dynamic_pointer_cast(blocks["patch_cond_embed"])->forward(ctx, image); + for (int i = 0; i < 2; ++i) { + cond = std::dynamic_pointer_cast(blocks["head_blocks." + std::to_string(i)])->forward(ctx, cond); + } + cond = std::dynamic_pointer_cast(blocks["proj_down"])->forward(ctx, cond); + auto z = ggml_ext_zeros(ctx->ggml_ctx, cond->ne[0], cond->ne[1], LATENT_CHANNELS, cond->ne[3]); + z = std::dynamic_pointer_cast(blocks["z_proj"])->forward(ctx, z); + z = ggml_concat(ctx->ggml_ctx, cond, z, 2); + z = std::dynamic_pointer_cast(blocks["fuse_proj"])->forward(ctx, z); + auto t = ggml_ext_zeros(ctx->ggml_ctx, image->ne[3], 1, 1, 1); + t = ggml_reshape_1d(ctx->ggml_ctx, t, image->ne[3]); + auto c = std::dynamic_pointer_cast(blocks["t_embedder"])->forward(ctx, t); + for (int i = 0; i < 21; ++i) { + z = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)])->forward(ctx, z, c); + } + z = std::dynamic_pointer_cast(blocks["norm_out"])->forward(ctx, z); + return std::dynamic_pointer_cast(blocks["proj_out"])->forward(ctx, z); + } + }; + + struct MLPResBlock : public GGMLBlock { + MLPResBlock() { + blocks["in_ln"] = std::make_shared(32, 1e-6f); + blocks["mlp.0"] = std::make_shared(32, 32); + blocks["mlp.2"] = std::make_shared(32, 32); + blocks["adaLN_modulation.1"] = std::make_shared(32, 96); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* condition) { + auto params = std::dynamic_pointer_cast(blocks["adaLN_modulation.1"])->forward(ctx, ggml_silu(ctx->ggml_ctx, condition)); + auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 3, 0); + auto h = std::dynamic_pointer_cast(blocks["in_ln"])->forward(ctx, x); + h = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, h, ggml_add(ctx->ggml_ctx, chunks[1], ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1))), chunks[0]); + h = std::dynamic_pointer_cast(blocks["mlp.0"])->forward(ctx, h); + h = ggml_silu(ctx->ggml_ctx, h); + h = std::dynamic_pointer_cast(blocks["mlp.2"])->forward(ctx, h); + return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, chunks[2], h)); + } + }; + + struct DConvDenoiser : public GGMLBlock { + DConvDenoiser() { + blocks["t_embedder"] = std::make_shared(); + blocks["y_embedder_x"] = std::make_shared(HIDDEN_SIZE, 32 * PATCH_SIZE * PATCH_SIZE, std::pair{1, 1}); + blocks["x_embedder.embedder.0"] = std::make_shared(3 + 32 + 64, 32); + blocks["s_embedder.proj1"] = std::make_shared(3, LATENT_CHANNELS, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{0, 0}, std::pair{1, 1}, false); + blocks["s_embedder.proj2"] = std::make_shared(LATENT_CHANNELS + HIDDEN_SIZE, HIDDEN_SIZE, std::pair{1, 1}); + for (int i = 0; i < 21; ++i) { + blocks["blocks." + std::to_string(i)] = std::make_shared(HIDDEN_SIZE); + } + blocks["dec_net.cond_embed"] = std::make_shared(HIDDEN_SIZE, PATCH_SIZE * PATCH_SIZE * 32); + blocks["dec_net.input_proj"] = std::make_shared(32, 32); + for (int i = 0; i < 3; ++i) { + blocks["dec_net.res_blocks." + std::to_string(i)] = std::make_shared(); + } + blocks["final_layer.norm"] = std::make_shared(32); + blocks["final_layer.linear"] = std::make_shared(32, 3); + blocks["y_embedder.decoder"] = std::make_shared(); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* dct) { + auto cond = std::dynamic_pointer_cast(blocks["y_embedder.decoder"])->forward(ctx, latent); + int64_t w = cond->ne[0]; + int64_t h = cond->ne[1]; + int64_t n = cond->ne[3]; + int64_t length = w * h; + + auto image = ggml_ext_zeros(ctx->ggml_ctx, w * PATCH_SIZE, h * PATCH_SIZE, 3, n); + auto t = ggml_ext_zeros(ctx->ggml_ctx, n, 1, 1, 1); + t = ggml_reshape_1d(ctx->ggml_ctx, t, n); + auto c = std::dynamic_pointer_cast(blocks["t_embedder"])->forward(ctx, t); + + auto s0 = std::dynamic_pointer_cast(blocks["s_embedder.proj1"])->forward(ctx, image); + s0 = ggml_concat(ctx->ggml_ctx, s0, cond, 2); + auto s = std::dynamic_pointer_cast(blocks["s_embedder.proj2"])->forward(ctx, s0); + for (int i = 0; i < 21; ++i) { + s = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)])->forward(ctx, s, c); + } + // [N, C, H, W] -> [N*H*W, C]. + s = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 2, 0, 3)); + s = ggml_reshape_2d(ctx->ggml_ctx, s, HIDDEN_SIZE, length * n); + + auto y = std::dynamic_pointer_cast(blocks["y_embedder_x"])->forward(ctx, cond); + // Split 32*P channels as [32, P], then produce [N*L, P, 32]. + y = ggml_reshape_4d(ctx->ggml_ctx, y, length, PATCH_SIZE * PATCH_SIZE, 32, n); + y = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, y, 2, 1, 0, 3)); + y = ggml_reshape_3d(ctx->ggml_ctx, y, 32, PATCH_SIZE * PATCH_SIZE, length * n); + auto zeros = ggml_ext_zeros(ctx->ggml_ctx, 3, PATCH_SIZE * PATCH_SIZE, length * n, 1); + dct = ggml_repeat_4d(ctx->ggml_ctx, dct, 64, PATCH_SIZE * PATCH_SIZE, length * n, 1); + auto x = ggml_concat(ctx->ggml_ctx, zeros, y, 0); + x = ggml_concat(ctx->ggml_ctx, x, dct, 0); + x = std::dynamic_pointer_cast(blocks["x_embedder.embedder.0"])->forward(ctx, x); + x = std::dynamic_pointer_cast(blocks["dec_net.input_proj"])->forward(ctx, x); + + auto dec_cond = std::dynamic_pointer_cast(blocks["dec_net.cond_embed"])->forward(ctx, s); + dec_cond = ggml_reshape_3d(ctx->ggml_ctx, dec_cond, 32, PATCH_SIZE * PATCH_SIZE, length * n); + for (int i = 0; i < 3; ++i) { + x = std::dynamic_pointer_cast(blocks["dec_net.res_blocks." + std::to_string(i)])->forward(ctx, x, dec_cond); + } + x = std::dynamic_pointer_cast(blocks["final_layer.norm"])->forward(ctx, x); + x = std::dynamic_pointer_cast(blocks["final_layer.linear"])->forward(ctx, x); + // [N*L, P, 3] -> [N, L, 3*P] for fold/unpatchify. + x = ggml_reshape_4d(ctx->ggml_ctx, x, 3, PATCH_SIZE * PATCH_SIZE, length, n); + x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = ggml_reshape_3d(ctx->ggml_ctx, x, 3 * PATCH_SIZE * PATCH_SIZE, length, n); + return DiT::unpatchify(ctx->ggml_ctx, x, h, w, PATCH_SIZE, PATCH_SIZE); + } + }; + + struct MageVAEModel : public GGMLBlock { + MageVAEModel() { + blocks["student.dconv_encoder"] = std::make_shared(); + blocks["pipeline"] = std::make_shared(); + } + + ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* image) { + return std::dynamic_pointer_cast(blocks["student.dconv_encoder"])->forward(ctx, image); + } + + ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* dct) { + return std::dynamic_pointer_cast(blocks["pipeline"])->forward(ctx, latent, dct); + } + }; + + struct MageVAERunner : public VAE { + MageVAEModel model; + std::vector dct_vec; + + MageVAERunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map, + const std::string& prefix, + std::shared_ptr weight_manager = nullptr) + : VAE(VERSION_MAGE_FLOW, backend, prefix, weight_manager) { + model = MageVAEModel(); + model.init(params_ctx, tensor_storage_map, prefix); + dct_vec.resize(64 * PATCH_SIZE * PATCH_SIZE); + constexpr float pi = 3.14159265358979323846f; + for (int py = 0; py < PATCH_SIZE; ++py) { + float y = static_cast(py) / static_cast(PATCH_SIZE - 1); + for (int px = 0; px < PATCH_SIZE; ++px) { + float x = static_cast(px) / static_cast(PATCH_SIZE - 1); + int pos = py * PATCH_SIZE + px; + for (int fy = 0; fy < 8; ++fy) { + for (int fx = 0; fx < 8; ++fx) { + int freq = fx * 8 + fy; + float freq_x = static_cast(fx) * 8.f / 7.f; + float freq_y = static_cast(fy) * 8.f / 7.f; + float coeff = 1.f / (1.f + freq_x * freq_y); + dct_vec[freq + 64 * pos] = std::cos(x * freq_x * pi) * + std::cos(y * freq_y * pi) * coeff; + } + } + } + } + } + + std::string get_desc() override { + return "mage_vae"; + } + + void get_param_tensors(std::map& tensors) override { + model.get_param_tensors(tensors, weight_prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& input_tensor, bool decode_graph) { + ggml_cgraph* gf = new_graph_custom(MAGE_VAE_GRAPH_SIZE); + auto input = make_input(input_tensor); + auto runner_ctx = get_context(); + ggml_tensor* dct = nullptr; + if (decode_graph) { + dct = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, 64, PATCH_SIZE * PATCH_SIZE, 1); + set_backend_tensor_data(dct, dct_vec.data()); + } + auto out = decode_graph ? model.decode(&runner_ctx, input, dct) : model.encode(&runner_ctx, input); + ggml_build_forward_expand(gf, out); + return gf; + } + + sd::Tensor _compute(const int n_threads, + const sd::Tensor& input, + bool decode_graph) override { + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(input, decode_graph); + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), input.dim()); + } + + int get_encoder_output_channels(int input_channels) override { + SD_UNUSED(input_channels); + return LATENT_CHANNELS * 2; + } + + sd::Tensor vae_output_to_latents(const sd::Tensor& vae_output, std::shared_ptr rng) override { + const auto chunks = sd::ops::chunk(vae_output, 2, 2); + const auto& mean = chunks[0]; + const auto& logvar = chunks[1]; + sd::Tensor stddev = sd::ops::exp(0.5f * sd::ops::clamp(logvar, -20.0f, 10.0f)); + sd::Tensor noise = sd::Tensor::randn_like(mean, rng); + sd::Tensor latents = mean + stddev * noise; + return latents; + } + + sd::Tensor diffusion_to_vae_latents(const sd::Tensor& latents) override { + return latents; + } + + sd::Tensor vae_to_diffusion_latents(const sd::Tensor& latents) override { + return latents; + } + }; +} // namespace MageVAE + +#endif // __SD_MODEL_VAE_MAGE_VAE_HPP__ diff --git a/otherarch/sdcpp/src/model/vae/tae.hpp b/otherarch/sdcpp/src/model/vae/tae.hpp index 558c0544cf09..710729fd7271 100644 --- a/otherarch/sdcpp/src/model/vae/tae.hpp +++ b/otherarch/sdcpp/src/model/vae/tae.hpp @@ -528,6 +528,9 @@ class TAEHV : public GGMLBlock { if (version == VERSION_WAN2_2_TI2V) { z_channels = 48; patch = 2; + } else if (sd_version_is_hunyuan_video(version)) { + z_channels = 32; + patch = 2; } else if (sd_version_is_ltxav(version)) { z_channels = 128; patch = 4; @@ -542,12 +545,12 @@ class TAEHV : public GGMLBlock { ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) { auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); - if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) { + if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) { // (W, H, C, T) -> (W, H, T, C) z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 1, 3, 2)); } auto result = decoder->forward(ctx, z); - if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) { + if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) { // (W, H, T, C) -> (W, H, C, T) result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2)); } @@ -556,7 +559,7 @@ class TAEHV : public GGMLBlock { ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) { auto encoder = std::dynamic_pointer_cast(blocks["encoder"]); - if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) { + if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) { // (W, H, T, C) -> (W, H, C, T) x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); } @@ -569,7 +572,7 @@ class TAEHV : public GGMLBlock { } } x = encoder->forward(ctx, x); - if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) { + if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) { // (W, H, C, T) -> (W, H, T, C) x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); } diff --git a/otherarch/sdcpp/src/model/vae/vae.hpp b/otherarch/sdcpp/src/model/vae/vae.hpp index b97aecee065e..e8b4eed38b64 100644 --- a/otherarch/sdcpp/src/model/vae/vae.hpp +++ b/otherarch/sdcpp/src/model/vae/vae.hpp @@ -74,7 +74,7 @@ struct VAE : public GGMLRunner { int scale_factor = 8; if (version == VERSION_LTXAV) { scale_factor = 32; - } else if (version == VERSION_WAN2_2_TI2V) { + } else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version)) { scale_factor = 16; } else if (sd_version_uses_flux2_vae(version)) { scale_factor = 16; @@ -136,7 +136,7 @@ struct VAE : public GGMLRunner { // Image VAE encode is more sensitive to tile boundary context than decode. // Keep the smaller legacy factor for video VAEs, but default image encode // tiles to 64 latent pixels so a 512px SD image is encoded as one tile. - const float encode_tile_factor = (sd_version_is_wan(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f; + const float encode_tile_factor = (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f; get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor); LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y); output = tiled_compute(input, diff --git a/otherarch/sdcpp/src/model_io/gguf_io.h b/otherarch/sdcpp/src/model_io/gguf_io.h index 3a4ae8200e5f..c6ac1754af50 100644 --- a/otherarch/sdcpp/src/model_io/gguf_io.h +++ b/otherarch/sdcpp/src/model_io/gguf_io.h @@ -21,7 +21,7 @@ bool write_gguf_file(const std::string& file_path, class GGUFStreamingWriter : public StreamingModelWriter { public: GGUFStreamingWriter() = default; - ~GGUFStreamingWriter(); + ~GGUFStreamingWriter() override; bool write_metadata(const std::string& file_path, const std::vector& tensors, diff --git a/otherarch/sdcpp/src/model_loader.cpp b/otherarch/sdcpp/src/model_loader.cpp index 4c03788a1cd5..b76dd8e49162 100644 --- a/otherarch/sdcpp/src/model_loader.cpp +++ b/otherarch/sdcpp/src/model_loader.cpp @@ -527,11 +527,18 @@ SDVersion ModelLoader::get_sd_version() { return VERSION_MINIT2I; } if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) { + auto img_in = tensor_storage_map.find("model.diffusion_model.img_in.weight"); + if (img_in != tensor_storage_map.end() && img_in->second.ne[0] == 128) { + return VERSION_MAGE_FLOW; + } if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) { return VERSION_QWEN_IMAGE_LAYERED; } return VERSION_QWEN_IMAGE; } + if (tensor_storage.name.find("model.diffusion_model.txt_in.individual_token_refiner.blocks.0.adaLN_modulation.1.weight") != std::string::npos) { + return VERSION_HUNYUAN_VIDEO; + } if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) { return VERSION_ANIMA; } diff --git a/otherarch/sdcpp/src/name_conversion.cpp b/otherarch/sdcpp/src/name_conversion.cpp index 3aa7880e2878..2b41a55760c6 100644 --- a/otherarch/sdcpp/src/name_conversion.cpp +++ b/otherarch/sdcpp/src/name_conversion.cpp @@ -664,6 +664,72 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) { return name; } +std::string convert_hunyuan_video_to_original_flux(std::string name) { + int num_layers = 54; + int num_single_layers = 0; + static std::unordered_map hy_name_map; + + if (hy_name_map.empty()) { + // --- double transformer blocks --- + for (int i = 0; i < num_layers; ++i) { + std::string block_prefix = "double_blocks." + std::to_string(i) + "."; + std::string dst_prefix = "double_blocks." + std::to_string(i) + "."; + + hy_name_map[block_prefix + "img_mod.linear"] = dst_prefix + "img_mod.lin"; + hy_name_map[block_prefix + "txt_mod.linear"] = dst_prefix + "txt_mod.lin"; + + // attn + hy_name_map[block_prefix + "img_attn_qkv"] = dst_prefix + "img_attn.qkv"; + hy_name_map[block_prefix + "txt_attn_qkv"] = dst_prefix + "txt_attn.qkv"; + + // norm + hy_name_map[block_prefix + "img_attn_q_norm.weight"] = dst_prefix + "img_attn.norm.query_norm.scale"; + hy_name_map[block_prefix + "img_attn_k_norm.weight"] = dst_prefix + "img_attn.norm.key_norm.scale"; + hy_name_map[block_prefix + "txt_attn_q_norm.weight"] = dst_prefix + "txt_attn.norm.query_norm.scale"; + hy_name_map[block_prefix + "txt_attn_k_norm.weight"] = dst_prefix + "txt_attn.norm.key_norm.scale"; + + // ff + hy_name_map[block_prefix + "img_mlp.fc1"] = dst_prefix + "img_mlp.0"; + hy_name_map[block_prefix + "img_mlp.fc2"] = dst_prefix + "img_mlp.2"; + + hy_name_map[block_prefix + "txt_mlp.fc1"] = dst_prefix + "txt_mlp.0"; + hy_name_map[block_prefix + "txt_mlp.fc2"] = dst_prefix + "txt_mlp.2"; + + // output projections + hy_name_map[block_prefix + "img_attn_proj"] = dst_prefix + "img_attn.proj"; + hy_name_map[block_prefix + "txt_attn_proj"] = dst_prefix + "txt_attn.proj"; + } + } + + hy_name_map["time_in.mlp.0"] = "time_in.in_layer"; + hy_name_map["time_in.mlp.2"] = "time_in.out_layer"; + hy_name_map["time_r_in.mlp.0"] = "time_r_in.in_layer"; + hy_name_map["time_r_in.mlp.2"] = "time_r_in.out_layer"; + hy_name_map["vector_in.mlp.0"] = "vector_in.in_layer"; + hy_name_map["vector_in.mlp.2"] = "vector_in.out_layer"; + hy_name_map["guidance_in.mlp.0"] = "guidance_in.in_layer"; + hy_name_map["guidance_in.mlp.2"] = "guidance_in.out_layer"; + + hy_name_map["txt_in.c_embedder.linear_1"] = "txt_in.c_embedder.in_layer"; + hy_name_map["txt_in.c_embedder.linear_2"] = "txt_in.c_embedder.out_layer"; + + hy_name_map["txt_in.t_embedder.mlp.0"] = "txt_in.t_embedder.in_layer"; + hy_name_map["txt_in.t_embedder.mlp.2"] = "txt_in.t_embedder.out_layer"; + + replace_with_prefix_map(name, hy_name_map); + + static const std::vector> generic_name_map = { + {"_attn_qkv.", "_attn.qkv."}, + {"_attn_proj.", "_attn.proj."}, + {"mlp.fc1.", "mlp.0."}, + {"mlp.fc2.", "mlp.2."}, + {".modulation.linear.", ".modulation.lin."}, + }; + replace_with_name_map(name, generic_name_map); + + return name; +} + std::string convert_diffusers_dit_to_original_lumina2(std::string name) { int num_layers = 30; int num_refiner_layers = 2; @@ -807,6 +873,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S name = convert_diffusers_dit_to_original_sd3(name); } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) { name = convert_diffusers_dit_to_original_flux(name); + } else if (sd_version_is_hunyuan_video(version)) { + 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_anima(version)) { @@ -980,6 +1048,9 @@ std::string convert_diffusers_to_original_wan_vae(std::string name) { } std::string convert_first_stage_model_name(std::string name, std::string prefix, SDVersion version) { + if (sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version)) { + return name; + } if (sd_version_uses_wan_vae(version)) { return convert_diffusers_to_original_wan_vae(name); } @@ -1343,8 +1414,14 @@ std::string convert_tensor_name(std::string name, SDVersion version) { replace_with_prefix_map(name, prefix_map); - if ((sd_version_is_boogu_image(version) || sd_version_is_krea2(version)) && starts_with(name, "text_encoders.llm.visual.")) { - name = convert_qwen3_vl_vision_name(std::move(name)); + if (sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version)) { + const std::string hf_vision_prefix = "text_encoders.llm.model.visual."; + if (starts_with(name, hf_vision_prefix)) { + name = "text_encoders.llm.visual." + name.substr(hf_vision_prefix.size()); + } + if (starts_with(name, "text_encoders.llm.visual.")) { + name = convert_qwen3_vl_vision_name(std::move(name)); + } } // diffusion model diff --git a/otherarch/sdcpp/src/runtime/denoiser.hpp b/otherarch/sdcpp/src/runtime/denoiser.hpp index 115175b14d98..48dd3b51d703 100644 --- a/otherarch/sdcpp/src/runtime/denoiser.hpp +++ b/otherarch/sdcpp/src/runtime/denoiser.hpp @@ -1153,9 +1153,9 @@ struct CompVisDenoiser : public Denoiser { return {c_skip, c_out, c_in}; } - virtual sd::Tensor noise_scaling(float sigma, - const sd::Tensor& noise, - const sd::Tensor& latent) override { + sd::Tensor noise_scaling(float sigma, + const sd::Tensor& noise, + const sd::Tensor& latent) override { GGML_ASSERT(noise.numel() == latent.numel()); return latent + noise * sigma; } @@ -1165,7 +1165,7 @@ struct CompVisDenoiser : public Denoiser { return latent; } - float noise_level_to_sigma(float noise_level) { + float noise_level_to_sigma(float noise_level) override { return noise_level / (1.0f - noise_level); } }; @@ -1256,7 +1256,7 @@ struct DiscreteFlowDenoiser : public Denoiser { return latent * (1.0f / (1.0f - sigma)); } - float noise_level_to_sigma(float noise_level) { + float noise_level_to_sigma(float noise_level) override { return noise_level; } }; @@ -1396,7 +1396,7 @@ struct MiniT2IFlowDenoiser : public Denoiser { return latent; } - float noise_level_to_sigma(float noise_level) { + float noise_level_to_sigma(float noise_level) override { SD_UNUSED(noise_level); return 1.0f; } diff --git a/otherarch/sdcpp/src/stable-diffusion.cpp b/otherarch/sdcpp/src/stable-diffusion.cpp index 0ee9a4d6b6e4..ae5fd53e36b7 100644 --- a/otherarch/sdcpp/src/stable-diffusion.cpp +++ b/otherarch/sdcpp/src/stable-diffusion.cpp @@ -32,11 +32,13 @@ #include "model/diffusion/ernie_image.hpp" #include "model/diffusion/flux.hpp" #include "model/diffusion/hidream_o1.hpp" +#include "model/diffusion/hunyuan.hpp" #include "model/diffusion/ideogram4.hpp" #include "model/diffusion/krea2.hpp" #include "model/diffusion/lens.hpp" #include "model/diffusion/lingbot_video.hpp" #include "model/diffusion/ltxv.hpp" +#include "model/diffusion/mage_flow.hpp" #include "model/diffusion/minit2i.hpp" #include "model/diffusion/mmdit.hpp" #include "model/diffusion/model.hpp" @@ -48,8 +50,10 @@ #include "model/upscaler/esrgan.hpp" #include "model/upscaler/ltx_latent_upscaler.hpp" #include "model/vae/auto_encoder_kl.hpp" +#include "model/vae/hunyuan_vae.hpp" #include "model/vae/ltx_audio_vae.hpp" #include "model/vae/ltx_vae.hpp" +#include "model/vae/mage_vae.hpp" #include "model/vae/tae.hpp" #include "model/vae/vae.hpp" #include "model/vae/wan_vae.hpp" @@ -98,6 +102,7 @@ const char* model_version_to_str[] = { "LingBot Video", "Qwen Image", "Qwen Image Layered", + "Hunyuan Video", "Anima", "Flux.2", "Flux.2 klein", @@ -114,6 +119,7 @@ const char* model_version_to_str[] = { "Ideogram 4", "SeFi-Image", "Krea2", + "Mage Flow", "ESRGAN", }; @@ -144,6 +150,7 @@ static bool sd_version_supports_ref_latent_img_cfg(SDVersion version) { return version == VERSION_FLUX || sd_version_is_flux2(version) || sd_version_is_qwen_image(version) || + sd_version_is_mage_flow(version) || sd_version_is_longcat(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version); @@ -805,7 +812,8 @@ class StableDiffusionGGML { bool is_boogu = sd_version_is_boogu_image(tempver); bool is_krea2 = sd_version_is_krea2(tempver); bool is_sefi = sd_version_is_sefi_image(tempver); - bool conditioner_is_llm = (is_qwenimg || iszimg || isflux2 || is_ovis || is_anima || is_ernie || is_longcat || is_lens || is_ltx || is_ideogram || is_boogu || is_krea2 || is_sefi); + bool is_mageflow = sd_version_is_mage_flow(tempver); + bool conditioner_is_llm = (is_qwenimg || iszimg || isflux2 || is_ovis || is_anima || is_ernie || is_longcat || is_lens || is_ltx || is_ideogram || is_boogu || is_krea2 || is_sefi || is_mageflow); bool has_llm_vision = (is_qwenimg || is_longcat || is_boogu); //kcpp qol fallback: if a llm was loaded as t5 by mistake @@ -1314,6 +1322,18 @@ class StableDiffusionGGML { tensor_storage_map, "model.diffusion_model", model_manager); + } else if (sd_version_is_hunyuan_video(version)) { + cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + model_manager); + diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + model_manager); } else if (sd_version_is_wan(version)) { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), tensor_storage_map, @@ -1382,6 +1402,17 @@ class StableDiffusionGGML { version, model_manager, sd_ctx_params->model_args); + } else if (sd_version_is_mage_flow(version)) { + cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + model_manager); + diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + model_manager); } else if (sd_version_is_longcat(version)) { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), tensor_storage_map, @@ -1519,7 +1550,7 @@ class StableDiffusionGGML { } auto create_tae = [&](bool decode_only) -> std::shared_ptr { - if (sd_version_uses_wan_vae(version) || sd_version_is_ltxav(version)) { + if (sd_version_uses_wan_vae(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) { return std::make_shared(backend_for(SDBackendModule::VAE), tensor_storage_map, "decoder", @@ -1556,6 +1587,18 @@ class StableDiffusionGGML { false, version, model_manager); + } else if (sd_version_is_mage_flow(vae_version)) { + return std::make_shared(backend_for(SDBackendModule::VAE), + tensor_storage_map, + "first_stage_model", + model_manager); + } else if (sd_version_uses_hunyuan_video_vae(vae_version)) { + return std::make_shared(backend_for(SDBackendModule::VAE), + tensor_storage_map, + "first_stage_model", + false, + vae_version, + model_manager); } else if (sd_version_uses_wan_vae(vae_version)) { return std::make_shared(backend_for(SDBackendModule::VAE), tensor_storage_map, @@ -1867,8 +1910,10 @@ class StableDiffusionGGML { } } else if (sd_version_is_sd3(version) || sd_version_is_wan(version) || + sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || + sd_version_is_mage_flow(version) || version == VERSION_HIDREAM_O1 || sd_version_is_anima(version) || sd_version_is_ernie_image(version) || @@ -1879,6 +1924,8 @@ class StableDiffusionGGML { pred_type = FLOW_PRED; if (sd_version_is_wan(version)) { default_flow_shift = 5.f; + } else if (sd_version_is_hunyuan_video(version)) { + default_flow_shift = 7.f; } else if (sd_version_is_ernie_image(version)) { default_flow_shift = 4.f; } else if (sd_version_is_pid(version)) { @@ -1887,6 +1934,8 @@ class StableDiffusionGGML { default_flow_shift = 1.0f; } else if (sd_version_is_boogu_image(version)) { default_flow_shift = 3.16f; + } else if (sd_version_is_mage_flow(version)) { + default_flow_shift = 6.f; } else { default_flow_shift = 3.f; } @@ -2715,6 +2764,10 @@ class StableDiffusionGGML { sd::Tensor timesteps_tensor({static_cast(timesteps_vec.size())}, timesteps_vec); sd::Tensor guidance_tensor({1}, std::vector{guidance.distilled_guidance}); + sd::Tensor hunyuan_timestep_r_tensor; + if (sd_version_is_hunyuan_video(version) && step + 1 < sigmas.size()) { + hunyuan_timestep_r_tensor = sd::Tensor::from_vector({sigmas[step + 1]}); + } sd::Tensor noised_input = x * c_in; if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) { noised_input = noised_input * denoise_mask + init_latent * (1.0f - denoise_mask); @@ -2789,6 +2842,12 @@ class StableDiffusionGGML { } else if (sd_version_is_wan(version)) { diffusion_params.extra = WanDiffusionExtra{vace_context.empty() ? nullptr : &vace_context, vace_strength}; + } else if (sd_version_is_hunyuan_video(version)) { + diffusion_params.extra = HunyuanVideoDiffusionExtra{ + &guidance_tensor, + condition.extra_c_crossattns.empty() ? nullptr : &condition.extra_c_crossattns[0], + condition.c_vector.empty() ? nullptr : &condition.c_vector, + hunyuan_timestep_r_tensor.empty() ? nullptr : &hunyuan_timestep_r_tensor}; } else if (version == VERSION_HIDREAM_O1) { diffusion_params.extra = HiDreamO1DiffusionExtra{ condition.c_input_ids.empty() ? nullptr : &condition.c_input_ids, @@ -2982,6 +3041,8 @@ class StableDiffusionGGML { latent_channel = 128; } else if (version == VERSION_WAN2_2_TI2V) { latent_channel = 48; + } else if (sd_version_is_hunyuan_video(version)) { + latent_channel = 32; } else if (version == VERSION_HIDREAM_O1) { latent_channel = 3; } else if (version == VERSION_CHROMA_RADIANCE) { @@ -2994,6 +3055,8 @@ class StableDiffusionGGML { latent_channel = 144; } else if (sd_version_uses_flux2_vae(version)) { latent_channel = 128; + } else if (sd_version_is_mage_flow(version)) { + latent_channel = 128; } else { latent_channel = 16; } @@ -3029,7 +3092,7 @@ class StableDiffusionGGML { int latent_frames = frames; if (sd_version_is_ltxav(version)) { latent_frames = ((frames - 1) / 8) + 1; - } else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) { + } else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_hunyuan_video(version)) { latent_frames = ((frames - 1) / 4) + 1; } return latent_frames; @@ -3042,7 +3105,7 @@ class StableDiffusionGGML { if (sd_version_is_ltxav(version)) { return (latent_frames - 1) * 8 + 1; } - if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) { + if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_hunyuan_video(version)) { return (latent_frames - 1) * 4 + 1; } return latent_frames; @@ -3142,6 +3205,8 @@ class StableDiffusionGGML { return "qwen_layered"; } else if (sd_version_is_qwen_image(version)) { return "qwen"; + } else if (sd_version_is_mage_flow(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_krea2(version)) { @@ -3840,7 +3905,7 @@ struct sd_ctx_t { }; static bool sd_version_supports_video_generation(SDVersion version) { - return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(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); } static bool sd_version_supports_image_generation(SDVersion version) { @@ -5078,10 +5143,17 @@ static std::optional prepare_image_generation_latents(sd sd::Tensor ref_latent; if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd_ctx->sd->version)) { LOG_DEBUG("auto resize ref images"); - int target_pixels = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : 1024 * 1024; - int vae_image_size = std::min(target_pixels, request->width * request->height); - double vae_width = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]); - double vae_height = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0]; + double vae_width; + double vae_height; + if (ref_image_params.resize_vae_to_target) { + vae_width = request->width; + vae_height = request->height; + } else { + int target_pixels = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : 1024 * 1024; + int vae_image_size = std::min(target_pixels, request->width * request->height); + vae_width = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]); + vae_height = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0]; + } int factor = sd_version_is_qwen_image(sd_ctx->sd->version) ? 32 : 16; vae_height = round(vae_height / factor) * factor; @@ -6007,6 +6079,66 @@ static std::optional prepare_video_generation_latents(sd } } + if (sd_version_is_hunyuan_video(sd_ctx->sd->version) && + (!start_image.empty() || !end_image.empty())) { + LOG_INFO("Hunyuan Video IMG2VID"); + + int64_t t1 = ggml_time_ms(); + auto concat_latent = sd_ctx->sd->generate_init_latent(request->width, + request->height, + request->frames, + true); + auto encode_condition_frame = [&](const sd::Tensor& image, + int64_t latent_frame, + const char* name) -> bool { + auto encoded = sd_ctx->sd->encode_first_stage(image); + if (encoded.empty()) { + LOG_ERROR("failed to encode Hunyuan Video %s conditioning frame", name); + return false; + } + if (encoded.dim() == 4) { + encoded.unsqueeze_(2); + } + if (encoded.dim() != 5 || + encoded.shape()[0] != concat_latent.shape()[0] || + encoded.shape()[1] != concat_latent.shape()[1] || + encoded.shape()[3] != concat_latent.shape()[3]) { + LOG_ERROR("invalid Hunyuan Video %s conditioning latent shape", name); + return false; + } + sd::ops::slice_assign(&concat_latent, + 2, + latent_frame, + latent_frame + 1, + sd::ops::slice(encoded, 2, 0, 1)); + return true; + }; + + if (!start_image.empty() && !encode_condition_frame(start_image, 0, "start")) { + return std::nullopt; + } + if (!end_image.empty() && + !encode_condition_frame(end_image, concat_latent.shape()[2] - 1, "end")) { + return std::nullopt; + } + + sd::Tensor concat_mask = sd::zeros({concat_latent.shape()[0], + concat_latent.shape()[1], + concat_latent.shape()[2], + 1, + 1}); + if (!start_image.empty()) { + sd::ops::fill_slice(&concat_mask, 2, 0, 1, 1.0f); + } + if (!end_image.empty()) { + sd::ops::fill_slice(&concat_mask, 2, concat_mask.shape()[2] - 1, concat_mask.shape()[2], 1.0f); + } + latents.concat_latent = sd::ops::concat(concat_latent, concat_mask, 3); + + int64_t t2 = ggml_time_ms(); + LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); + } + if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-I2V-14B" || sd_ctx->sd->diffusion_model->get_desc() == "Wan2.2-I2V-14B" || sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-I2V-1.3B" || From 48c7f6b7e60b0f934ded79351fbc245987eec13b Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Thu, 23 Jul 2026 19:08:33 -0300 Subject: [PATCH 3/5] sd: sync with master-789-5114672 --- otherarch/sdcpp/src/model/te/llm.hpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/otherarch/sdcpp/src/model/te/llm.hpp b/otherarch/sdcpp/src/model/te/llm.hpp index e5c521ef1066..60e2d9c15c72 100644 --- a/otherarch/sdcpp/src/model/te/llm.hpp +++ b/otherarch/sdcpp/src/model/te/llm.hpp @@ -200,7 +200,11 @@ namespace LLM { config.vision.in_channels = tensor_storage.ne[2]; config.vision.hidden_size = tensor_storage.ne[3]; } - if (contains(name, "visual.patch_embed.bias")) { + // HF-format checkpoints keep the patch embed unsplit under a single name. + if (contains(name, "visual.patch_embed.proj.weight")) { + config.vision.patch_size = static_cast(tensor_storage.ne[0]); + } + if (contains(name, "visual.patch_embed.bias") || contains(name, "visual.patch_embed.proj.bias")) { config.vision.hidden_size = tensor_storage.ne[0]; } if (contains(name, "visual.pos_embed.weight")) { From e322986a37c802a123d61a84a5bb6e224e06a413 Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Fri, 24 Jul 2026 07:02:43 -0300 Subject: [PATCH 4/5] sd: sync with master-795-87a0177 --- Makefile | 2 +- otherarch/sdcpp/examples/cli/main.cpp | 10 ++ otherarch/sdcpp/examples/common/common.cpp | 63 +++++---- otherarch/sdcpp/examples/common/common.h | 4 + otherarch/sdcpp/include/stable-diffusion.h | 3 + otherarch/sdcpp/sdtype_adapter.cpp | 2 +- otherarch/sdcpp/src/core/ggml_extend.hpp | 2 + .../sdcpp/src/model/adapter/ip_adapter.hpp | 88 +++++++++++++ otherarch/sdcpp/src/model/common/block.hpp | 33 ++++- otherarch/sdcpp/src/model/diffusion/model.hpp | 2 + otherarch/sdcpp/src/model/diffusion/unet.hpp | 29 +++-- otherarch/sdcpp/src/name_conversion.cpp | 43 +++++++ otherarch/sdcpp/src/stable-diffusion.cpp | 121 +++++++++++++++--- 13 files changed, 342 insertions(+), 60 deletions(-) create mode 100644 otherarch/sdcpp/src/model/adapter/ip_adapter.hpp diff --git a/Makefile b/Makefile index bba30467c57d..5c5fc8a67e36 100644 --- a/Makefile +++ b/Makefile @@ -703,7 +703,7 @@ budget.o: common/reasoning-budget.cpp common/reasoning-budget.h chat.o: common/chat.cpp common/chat.h $(CXX) $(CXXFLAGS) -c $< -o $@ -SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/core/backend_fit.cpp src/core/backend_fit.h src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.hpp src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/layer_split_partition.cpp src/core/layer_split_partition.h src/core/ordered_map.hpp src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/detailer.cpp src/detailer.h src/extensions/generation_extension.h src/extensions/photomaker_extension.cpp src/extensions/pulid_extension.cpp src/kcpp_sd_extensions.h src/model/adapter/lora.hpp src/model/adapter/pmid.hpp src/model/adapter/pulid.hpp src/model/common/block.hpp src/model/common/rope.hpp src/model/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/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/minit2i.hpp src/model/diffusion/mmdit.hpp src/model/diffusion/model.hpp src/model/diffusion/pid.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/sefi_image.hpp src/model/diffusion/unet.hpp src/model/diffusion/wan.hpp src/model/diffusion/z_image.hpp src/model.h src/model_io/binary_io.h src/model_io/gguf_io.cpp src/model_io/gguf_io.h src/model_io/gguf_reader_ext.h src/model_io/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.h src/model_manager.cpp src/model_manager.h src/model/te/clip.hpp src/model/te/llm.hpp src/model/te/t5.hpp src/model/upscaler/esrgan.hpp src/model/upscaler/ltx_latent_upscaler.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/ltx_audio_vae.hpp src/model/vae/ltx_vae.hpp src/model/vae/tae.hpp src/model/vae/vae.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/runtime/cache_dit.hpp src/runtime/condition_cache_utils.hpp src/runtime/denoiser.hpp src/runtime/easycache.hpp src/runtime/gits_noise.h src/runtime/guidance.cpp src/runtime/guidance.h src/runtime/imatrix.cpp src/runtime/imatrix.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/ucache.hpp src/stable-diffusion.cpp src/tokenizers/bpe_tokenizer.cpp src/tokenizers/bpe_tokenizer.h src/tokenizers/clip_tokenizer.cpp src/tokenizers/clip_tokenizer.h src/tokenizers/gemma_tokenizer.cpp src/tokenizers/gemma_tokenizer.h src/tokenizers/gpt_oss_tokenizer.cpp src/tokenizers/gpt_oss_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer.cpp src/tokenizers/tokenizer.h src/tokenizers/tokenize_util.cpp src/tokenizers/tokenize_util.h src/tokenizers/vocab/vocab.h src/upscaler.cpp src/upscaler.h src/weight_manager.h +SDCPP_COMMON_BASENAMES := include/stable-diffusion.h src/conditioning/conditioner.hpp src/core/backend_fit.cpp src/core/backend_fit.h src/core/ggml_extend_backend.cpp src/core/ggml_extend_backend.h src/core/ggml_extend.hpp src/core/ggml_graph_cut.cpp src/core/ggml_graph_cut.h src/core/layer_split_partition.cpp src/core/layer_split_partition.h src/core/ordered_map.hpp src/core/rng.hpp src/core/rng_mt19937.hpp src/core/rng_philox.hpp src/core/tensor_ggml.hpp src/core/tensor.hpp src/core/util.cpp src/core/util.h src/detailer.cpp src/detailer.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/pmid.hpp src/model/adapter/pulid.hpp src/model/common/block.hpp src/model/common/rope.hpp 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/minit2i.hpp src/model/diffusion/mmdit.hpp src/model/diffusion/model.hpp src/model/diffusion/pid.hpp src/model/diffusion/qwen_image.hpp src/model/diffusion/sefi_image.hpp src/model/diffusion/unet.hpp src/model/diffusion/wan.hpp src/model/diffusion/z_image.hpp src/model.h src/model_io/binary_io.h src/model_io/gguf_io.cpp src/model_io/gguf_io.h src/model_io/gguf_reader_ext.h src/model_io/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.h src/model_manager.cpp src/model_manager.h src/model/te/clip.hpp src/model/te/llm.hpp src/model/te/t5.hpp src/model/upscaler/esrgan.hpp src/model/upscaler/ltx_latent_upscaler.hpp src/model/vae/auto_encoder_kl.hpp src/model/vae/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/tae.hpp src/model/vae/vae.hpp src/model/vae/wan_vae.hpp src/name_conversion.cpp src/name_conversion.h src/runtime/cache_dit.hpp src/runtime/condition_cache_utils.hpp src/runtime/denoiser.hpp src/runtime/easycache.hpp src/runtime/gits_noise.h src/runtime/guidance.cpp src/runtime/guidance.h src/runtime/imatrix.cpp src/runtime/imatrix.h src/runtime/latent-preview.h src/runtime/preprocessing.hpp src/runtime/sample-cache.cpp src/runtime/sample-cache.h src/runtime/spectrum.hpp src/runtime/ucache.hpp src/stable-diffusion.cpp src/tokenizers/bpe_tokenizer.cpp src/tokenizers/bpe_tokenizer.h src/tokenizers/clip_tokenizer.cpp src/tokenizers/clip_tokenizer.h src/tokenizers/gemma_tokenizer.cpp src/tokenizers/gemma_tokenizer.h src/tokenizers/gpt_oss_tokenizer.cpp src/tokenizers/gpt_oss_tokenizer.h src/tokenizers/mistral_tokenizer.cpp src/tokenizers/mistral_tokenizer.h src/tokenizers/qwen2_tokenizer.cpp src/tokenizers/qwen2_tokenizer.h src/tokenizers/t5_unigram_tokenizer.cpp src/tokenizers/t5_unigram_tokenizer.h src/tokenizers/tokenizer.cpp src/tokenizers/tokenizer.h src/tokenizers/tokenize_util.cpp src/tokenizers/tokenize_util.h src/tokenizers/vocab/vocab.h src/upscaler.cpp src/upscaler.h src/weight_manager.h SDCPP_MAIN_BASENAMES := examples/cli/image_metadata.cpp examples/cli/image_metadata.h examples/cli/main.cpp examples/cli/msf_gif.h examples/common/common.cpp examples/common/common.h examples/common/log.cpp examples/common/log.h examples/common/media_io.cpp examples/common/media_io.h examples/common/resource_owners.hpp src/tokenizers/vocab/clip_merges.hpp src/tokenizers/vocab/gemma2_merges.hpp src/tokenizers/vocab/gemma2_vocab.hpp src/tokenizers/vocab/gemma_merges.hpp src/tokenizers/vocab/gemma_vocab.hpp src/tokenizers/vocab/gpt_oss_merges.hpp src/tokenizers/vocab/gpt_oss_vocab.hpp src/tokenizers/vocab/mistral_merges.hpp src/tokenizers/vocab/mistral_vocab.hpp src/tokenizers/vocab/qwen_merges.hpp src/tokenizers/vocab/t5.hpp src/tokenizers/vocab/umt5.hpp src/tokenizers/vocab/vocab.cpp src/convert.cpp src/version.cpp diff --git a/otherarch/sdcpp/examples/cli/main.cpp b/otherarch/sdcpp/examples/cli/main.cpp index 68a3148cc605..af62a9a93890 100644 --- a/otherarch/sdcpp/examples/cli/main.cpp +++ b/otherarch/sdcpp/examples/cli/main.cpp @@ -817,6 +817,16 @@ int main(int argc, const char* argv[]) { } } + if (gen_params.ip_adapter_image_path.size() > 0) { + if (!load_sd_image_from_file(gen_params.ip_adapter_image.put(), + gen_params.ip_adapter_image_path.c_str(), + 0, + 0)) { + LOG_ERROR("load image from '%s' failed", gen_params.ip_adapter_image_path.c_str()); + return 1; + } + } + if (!gen_params.control_video_path.empty()) { gen_params.control_frames.clear(); if (!load_images_from_dir(gen_params.control_video_path, diff --git a/otherarch/sdcpp/examples/common/common.cpp b/otherarch/sdcpp/examples/common/common.cpp index a5cb48369a28..82558e7c24a1 100644 --- a/otherarch/sdcpp/examples/common/common.cpp +++ b/otherarch/sdcpp/examples/common/common.cpp @@ -427,6 +427,11 @@ ArgOptions SDContextParams::get_options() { "path to control net model", 0, &control_net_path}, + {"", + "--ip-adapter", + "path to IP-Adapter model (requires --clip_vision)", + 0, + &ip_adapter_path}, {"", "--motion-module", "path to AnimateDiff motion module (SD 1.5); enables video generation on --video-frames > 1", @@ -876,6 +881,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.audio_vae_path = audio_vae_path.c_str(); sd_ctx_params.taesd_path = taesd_path.c_str(); sd_ctx_params.control_net_path = control_net_path.c_str(); + sd_ctx_params.ip_adapter_path = ip_adapter_path.c_str(); sd_ctx_params.motion_module_path = motion_module_path.c_str(); sd_ctx_params.embeddings = embedding_vec.data(); sd_ctx_params.embedding_count = static_cast(embedding_vec.size()); @@ -966,6 +972,11 @@ ArgOptions SDGenerationParams::get_options() { "path to control image, control net", 0, &control_image_path}, + {"", + "--ip-adapter-image", + "path to the IP-Adapter reference image", + 0, + &ip_adapter_image_path}, {"", "--control-video", "path to control video frames, It must be a directory path. The video frames inside should be stored as images in " @@ -1158,6 +1169,10 @@ ArgOptions SDGenerationParams::get_options() { "--control-strength", "strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image", &control_strength}, + {"", + "--ip-adapter-strength", + "strength to apply IP-Adapter (default: 1.0)", + &ip_adapter_strength}, {"", "--moe-boundary", "timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if `--high-noise-steps` is set to -1", @@ -2479,29 +2494,31 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() { LOG_WARN("Notice: --increase-ref-index is deprecated. Use --ref-image-args \"ref_index_mode=increase\" instead."); } - params.loras = lora_vec.empty() ? nullptr : lora_vec.data(); - params.lora_count = static_cast(lora_vec.size()); - params.prompt = prompt.c_str(); - params.negative_prompt = negative_prompt.c_str(); - params.clip_skip = clip_skip; - params.init_image = init_image.get(); - params.ref_images = ref_image_views.empty() ? nullptr : ref_image_views.data(); - params.ref_images_count = static_cast(ref_image_views.size()); - params.ref_image_args = ref_image_args.c_str(); - params.mask_image = mask_image.get(); - params.width = get_resolved_width(); - params.height = get_resolved_height(); - params.sample_params = sample_params; - params.strength = strength; - params.seed = seed; - params.batch_count = batch_count; - params.qwen_image_layers = qwen_image_layers; - params.control_image = control_image.get(); - params.control_strength = control_strength; - params.pm_params = pm_params; - params.pulid_params = pulid_params; - params.vae_tiling_params = vae_tiling_params; - params.cache = cache_params; + params.loras = lora_vec.empty() ? nullptr : lora_vec.data(); + params.lora_count = static_cast(lora_vec.size()); + params.prompt = prompt.c_str(); + params.negative_prompt = negative_prompt.c_str(); + params.clip_skip = clip_skip; + params.init_image = init_image.get(); + params.ref_images = ref_image_views.empty() ? nullptr : ref_image_views.data(); + params.ref_images_count = static_cast(ref_image_views.size()); + params.ref_image_args = ref_image_args.c_str(); + params.mask_image = mask_image.get(); + params.width = get_resolved_width(); + params.height = get_resolved_height(); + params.sample_params = sample_params; + params.strength = strength; + params.seed = seed; + params.batch_count = batch_count; + params.qwen_image_layers = qwen_image_layers; + params.control_image = control_image.get(); + params.control_strength = control_strength; + params.ip_adapter_image = ip_adapter_image.get(); + params.ip_adapter_strength = ip_adapter_strength; + params.pm_params = pm_params; + params.pulid_params = pulid_params; + params.vae_tiling_params = vae_tiling_params; + params.cache = cache_params; params.hires.enabled = hires_enabled; params.hires.upscaler = resolved_hires_upscaler; diff --git a/otherarch/sdcpp/examples/common/common.h b/otherarch/sdcpp/examples/common/common.h index 769651987e97..ea90c8c1bdfb 100644 --- a/otherarch/sdcpp/examples/common/common.h +++ b/otherarch/sdcpp/examples/common/common.h @@ -133,6 +133,7 @@ struct SDContextParams { std::string taesd_path; std::string esrgan_path; std::string control_net_path; + std::string ip_adapter_path; std::string motion_module_path; std::string embedding_dir; std::string photo_maker_path; @@ -200,6 +201,7 @@ struct SDGenerationParams { 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; @@ -208,6 +210,7 @@ struct SDGenerationParams { std::string end_image_path; std::string mask_image_path; std::string control_image_path; + std::string ip_adapter_image_path; std::vector ref_image_paths; std::string control_video_path; @@ -274,6 +277,7 @@ struct SDGenerationParams { std::vector ref_images; SDImageOwner mask_image; SDImageOwner control_image; + SDImageOwner ip_adapter_image; std::vector pm_id_images; std::vector control_frames; diff --git a/otherarch/sdcpp/include/stable-diffusion.h b/otherarch/sdcpp/include/stable-diffusion.h index 536bce73f063..60c99992f5f9 100644 --- a/otherarch/sdcpp/include/stable-diffusion.h +++ b/otherarch/sdcpp/include/stable-diffusion.h @@ -200,6 +200,7 @@ typedef struct { const char* audio_vae_path; const char* taesd_path; const char* control_net_path; + const char* ip_adapter_path; const char* motion_module_path; const sd_embedding_t* embeddings; uint32_t embedding_count; @@ -375,6 +376,8 @@ typedef struct { int batch_count; sd_image_t control_image; float control_strength; + sd_image_t ip_adapter_image; + float ip_adapter_strength; sd_pm_params_t pm_params; sd_pulid_params_t pulid_params; sd_tiling_params_t vae_tiling_params; diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp index 89c277e8d1d4..c42ce7654124 100644 --- a/otherarch/sdcpp/sdtype_adapter.cpp +++ b/otherarch/sdcpp/sdtype_adapter.cpp @@ -1249,7 +1249,7 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) params.negative_prompt = sd_params->negative_prompt.c_str(); params.clip_skip = sd_params->clip_skip; params.sample_params.guidance.txt_cfg = sd_params->cfg_scale; - // params.sample_params.guidance.img_cfg = sd_params->cfg_scale; //removed, breaks qwen img edit and more + params.sample_params.guidance.img_cfg = sd_params->cfg_scale; //removed, breaks qwen img edit and more if (sd_params->distilled_guidance >= 0.f) { params.sample_params.guidance.distilled_guidance = sd_params->distilled_guidance; } diff --git a/otherarch/sdcpp/src/core/ggml_extend.hpp b/otherarch/sdcpp/src/core/ggml_extend.hpp index 2280a535f58c..cc189f5b0cd9 100644 --- a/otherarch/sdcpp/src/core/ggml_extend.hpp +++ b/otherarch/sdcpp/src/core/ggml_extend.hpp @@ -1689,6 +1689,8 @@ struct GGMLRunnerContext { bool conv2d_direct_enabled = false; bool circular_x_enabled = false; bool circular_y_enabled = false; + ggml_tensor* ip_context = nullptr; + float ip_scale = 1.0f; std::shared_ptr weight_adapter = nullptr; std::vector>* debug_tensors = nullptr; std::function get_cache_tensor; diff --git a/otherarch/sdcpp/src/model/adapter/ip_adapter.hpp b/otherarch/sdcpp/src/model/adapter/ip_adapter.hpp new file mode 100644 index 000000000000..35fd1bb88ed5 --- /dev/null +++ b/otherarch/sdcpp/src/model/adapter/ip_adapter.hpp @@ -0,0 +1,88 @@ +#ifndef __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__ +#define __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__ + +#include "core/ggml_extend.hpp" +#include "model/common/block.hpp" +#include "model_loader.h" + +namespace IPAdapter { + + struct ImageProjModel : public GGMLBlock { + int64_t num_tokens = 4; + int64_t ctx_dim = 768; + int64_t clip_dim = 1024; + + ImageProjModel() {} + ImageProjModel(int64_t num_tokens, int64_t ctx_dim, int64_t clip_dim) + : num_tokens(num_tokens), ctx_dim(ctx_dim), clip_dim(clip_dim) { + blocks["proj"] = std::shared_ptr(new Linear(clip_dim, num_tokens * ctx_dim, true)); + blocks["norm"] = std::shared_ptr(new LayerNorm(ctx_dim)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image_embeds) { + auto proj = std::dynamic_pointer_cast(blocks["proj"]); + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + + int64_t n = image_embeds->ne[1]; + auto x = proj->forward(ctx, image_embeds); + x = ggml_reshape_3d(ctx->ggml_ctx, x, ctx_dim, num_tokens, n); + x = norm->forward(ctx, x); + return x; + } + }; + + struct IPAdapterRunner : public GGMLRunner { + ImageProjModel image_proj; + int64_t num_tokens = 4; + std::string prefix; + + IPAdapterRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map, + const std::string prefix, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), prefix(prefix) { + int64_t ctx_dim = 768; + int64_t clip_dim = 1024; + int64_t out_dim = 3072; + auto norm_iter = tensor_storage_map.find(prefix + ".image_proj.norm.weight"); + if (norm_iter != tensor_storage_map.end()) { + ctx_dim = norm_iter->second.ne[0]; + } + auto proj_iter = tensor_storage_map.find(prefix + ".image_proj.proj.weight"); + if (proj_iter != tensor_storage_map.end()) { + clip_dim = proj_iter->second.ne[0]; + out_dim = proj_iter->second.ne[1]; + } + num_tokens = out_dim / ctx_dim; + image_proj = ImageProjModel(num_tokens, ctx_dim, clip_dim); + image_proj.init(params_ctx, tensor_storage_map, prefix + ".image_proj"); + } + + std::string get_desc() override { + return "ip_adapter"; + } + + void get_param_tensors(std::map& tensors, const std::string = "") { + image_proj.get_param_tensors(tensors, prefix + ".image_proj"); + } + + ggml_cgraph* build_graph(const sd::Tensor& image_embeds_tensor) { + ggml_cgraph* gf = new_graph_custom(1024); + ggml_tensor* embeds = make_input(image_embeds_tensor); + auto runner_ctx = get_context(); + ggml_tensor* out = image_proj.forward(&runner_ctx, embeds); + ggml_build_forward_expand(gf, out); + return gf; + } + + sd::Tensor compute(int n_threads, const sd::Tensor& image_embeds) { + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(image_embeds); + }; + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true, true, true)); + } + }; + +} // namespace IPAdapter + +#endif // __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__ diff --git a/otherarch/sdcpp/src/model/common/block.hpp b/otherarch/sdcpp/src/model/common/block.hpp index 6fb40ec0f61e..be76dd7130e1 100644 --- a/otherarch/sdcpp/src/model/common/block.hpp +++ b/otherarch/sdcpp/src/model/common/block.hpp @@ -310,17 +310,33 @@ class CrossAttention : public GGMLBlock { int64_t context_dim; int64_t n_head; int64_t d_head; - bool xtra_dim = false; + bool xtra_dim = false; + bool enable_ip = false; + bool has_ip = false; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + GGMLBlock::init_params(ctx, tensor_storage_map, prefix); + if (enable_ip && + tensor_storage_map.find(prefix + "to_k_ip.weight") != tensor_storage_map.end()) { + has_ip = true; + int64_t inner_dim = d_head * n_head; + int64_t ip_dim = tensor_storage_map.at(prefix + "to_k_ip.weight").ne[0]; + blocks["to_k_ip"] = std::shared_ptr(new Linear(ip_dim, inner_dim, false)); + blocks["to_v_ip"] = std::shared_ptr(new Linear(ip_dim, inner_dim, false)); + } + } public: CrossAttention(int64_t query_dim, int64_t context_dim, int64_t n_head, - int64_t d_head) + int64_t d_head, + bool enable_ip = false) : n_head(n_head), d_head(d_head), query_dim(query_dim), - context_dim(context_dim) { + context_dim(context_dim), + enable_ip(enable_ip) { int64_t inner_dim = d_head * n_head; if (context_dim == 320 && d_head == 320) { // LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09"); @@ -363,6 +379,15 @@ class CrossAttention : public GGMLBlock { } x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, inner_dim] + if (has_ip && ctx->ip_context != nullptr && ctx->ip_scale != 0.0f) { + auto to_k_ip = std::dynamic_pointer_cast(blocks["to_k_ip"]); + auto to_v_ip = std::dynamic_pointer_cast(blocks["to_v_ip"]); + auto k_ip = to_k_ip->forward(ctx, ctx->ip_context); + auto v_ip = to_v_ip->forward(ctx, ctx->ip_context); + auto x_ip = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_ip, v_ip, n_head, nullptr, false, ctx->flash_attn_enabled); + x = ggml_add(ctx->ggml_ctx, x, ggml_scale(ctx->ggml_ctx, x_ip, ctx->ip_scale)); + } + x = to_out_0->forward(ctx, x); // [N, n_token, query_dim] return x; } @@ -387,7 +412,7 @@ class BasicTransformerBlock : public GGMLBlock { // inner_dim is always None or equal to dim // gated_ff is always True blocks["attn1"] = std::shared_ptr(new CrossAttention(dim, dim, n_head, d_head)); - blocks["attn2"] = std::shared_ptr(new CrossAttention(dim, context_dim, n_head, d_head)); + blocks["attn2"] = std::shared_ptr(new CrossAttention(dim, context_dim, n_head, d_head, true)); blocks["ff"] = std::shared_ptr(new FeedForward(dim, dim)); blocks["norm1"] = std::shared_ptr(new LayerNorm(dim)); blocks["norm2"] = std::shared_ptr(new LayerNorm(dim)); diff --git a/otherarch/sdcpp/src/model/diffusion/model.hpp b/otherarch/sdcpp/src/model/diffusion/model.hpp index 31c5048330df..d7a5f423123b 100644 --- a/otherarch/sdcpp/src/model/diffusion/model.hpp +++ b/otherarch/sdcpp/src/model/diffusion/model.hpp @@ -46,6 +46,8 @@ struct UNetDiffusionExtra { int num_video_frames = -1; const std::vector>* controls = nullptr; float control_strength = 0.f; + const sd::Tensor* ip_context = nullptr; + float ip_scale = 1.f; }; struct SkipLayerDiffusionExtra { diff --git a/otherarch/sdcpp/src/model/diffusion/unet.hpp b/otherarch/sdcpp/src/model/diffusion/unet.hpp index 118ed6327663..567573049a04 100644 --- a/otherarch/sdcpp/src/model/diffusion/unet.hpp +++ b/otherarch/sdcpp/src/model/diffusion/unet.hpp @@ -775,14 +775,17 @@ struct UNetModelRunner : public DiffusionModelRunner { const sd::Tensor& y_tensor = {}, int num_video_frames = -1, const std::vector>& controls_tensor = {}, - float control_strength = 0.f) { + float control_strength = 0.f, + const sd::Tensor& ip_context_tensor = {}, + float ip_scale = 1.f) { ggml_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE); - ggml_tensor* x = make_input(x_tensor); - ggml_tensor* timesteps = make_input(timesteps_tensor); - ggml_tensor* context = make_optional_input(context_tensor); - ggml_tensor* c_concat = make_optional_input(c_concat_tensor); - ggml_tensor* y = make_optional_input(y_tensor); + ggml_tensor* x = make_input(x_tensor); + ggml_tensor* timesteps = make_input(timesteps_tensor); + ggml_tensor* context = make_optional_input(context_tensor); + ggml_tensor* c_concat = make_optional_input(c_concat_tensor); + ggml_tensor* y = make_optional_input(y_tensor); + ggml_tensor* ip_context = make_optional_input(ip_context_tensor); std::vector controls; controls.reserve(controls_tensor.size()); for (const auto& control_tensor : controls_tensor) { @@ -793,7 +796,9 @@ struct UNetModelRunner : public DiffusionModelRunner { num_video_frames = static_cast(x->ne[3]); } - auto runner_ctx = get_context(); + auto runner_ctx = get_context(); + runner_ctx.ip_context = ip_context; + runner_ctx.ip_scale = ip_scale; ggml_tensor* out = unet.forward(&runner_ctx, x, @@ -818,14 +823,16 @@ struct UNetModelRunner : public DiffusionModelRunner { const sd::Tensor& y = {}, int num_video_frames = -1, const std::vector>& controls = {}, - float control_strength = 0.f) { + float control_strength = 0.f, + const sd::Tensor& ip_context = {}, + float ip_scale = 1.f) { // x: [N, in_channels, h, w] // timesteps: [N, ] // context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size] // c_concat: [N, in_channels, h, w] or [1, in_channels, h, w] // y: [N, adm_in_channels] or [1, adm_in_channels] auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength); + return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, ip_context, ip_scale); }; return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); @@ -845,7 +852,9 @@ struct UNetModelRunner : public DiffusionModelRunner { tensor_or_empty(diffusion_params.y), extra->num_video_frames, extra->controls ? *extra->controls : empty_controls, - extra->control_strength); + extra->control_strength, + extra->ip_context ? *extra->ip_context : sd::Tensor{}, + extra->ip_scale); } void test() { diff --git a/otherarch/sdcpp/src/name_conversion.cpp b/otherarch/sdcpp/src/name_conversion.cpp index 2b41a55760c6..63c916edd85a 100644 --- a/otherarch/sdcpp/src/name_conversion.cpp +++ b/otherarch/sdcpp/src/name_conversion.cpp @@ -1285,11 +1285,54 @@ static std::string convert_esrgan_tensor_name(std::string name) { return name; } +static const std::map& ip_adapter_index_map(SDVersion version) { + static const std::map sd15_map = { + {1, "input_blocks.1.1.transformer_blocks.0"}, {3, "input_blocks.2.1.transformer_blocks.0"}, {5, "input_blocks.4.1.transformer_blocks.0"}, {7, "input_blocks.5.1.transformer_blocks.0"}, {9, "input_blocks.7.1.transformer_blocks.0"}, {11, "input_blocks.8.1.transformer_blocks.0"}, {13, "output_blocks.3.1.transformer_blocks.0"}, {15, "output_blocks.4.1.transformer_blocks.0"}, {17, "output_blocks.5.1.transformer_blocks.0"}, {19, "output_blocks.6.1.transformer_blocks.0"}, {21, "output_blocks.7.1.transformer_blocks.0"}, {23, "output_blocks.8.1.transformer_blocks.0"}, {25, "output_blocks.9.1.transformer_blocks.0"}, {27, "output_blocks.10.1.transformer_blocks.0"}, {29, "output_blocks.11.1.transformer_blocks.0"}, {31, "middle_block.1.transformer_blocks.0"}}; + + static std::map sdxl_map; + if (sdxl_map.empty()) { + std::vector> order = { + {"input_blocks.4.1", 2}, {"input_blocks.5.1", 2}, {"input_blocks.7.1", 10}, {"input_blocks.8.1", 10}, {"output_blocks.0.1", 10}, {"output_blocks.1.1", 10}, {"output_blocks.2.1", 10}, {"output_blocks.3.1", 2}, {"output_blocks.4.1", 2}, {"output_blocks.5.1", 2}, {"middle_block.1", 10}}; + int idx = 1; + for (const auto& [block, depth] : order) { + for (int m = 0; m < depth; m++) { + sdxl_map[idx] = block + ".transformer_blocks." + std::to_string(m); + idx += 2; + } + } + } + return sd_version_is_sdxl(version) ? sdxl_map : sd15_map; +} + +static std::string convert_ip_adapter_name(std::string name, SDVersion version) { + if (starts_with(name, "image_proj.")) { + return "ip_adapter." + name; + } + if (starts_with(name, "ip_adapter.")) { + auto items = split_string(name, '.'); + if (items.size() < 4) { + return name; + } + int idx = atoi(items[1].c_str()); + const auto& mp = ip_adapter_index_map(version); + auto blk = mp.find(idx); + if (blk == mp.end()) { + return name; + } + return "model.diffusion_model." + blk->second + ".attn2." + items[2] + "." + items[3]; + } + return name; +} + std::string convert_tensor_name(std::string name, SDVersion version) { if (version == VERSION_ESRGAN) { return convert_esrgan_tensor_name(std::move(name)); } + if (starts_with(name, "ip_adapter.") || starts_with(name, "image_proj.")) { + return convert_ip_adapter_name(std::move(name), version); + } + bool is_lora = false; bool is_lycoris_underline = false; bool is_underline = false; diff --git a/otherarch/sdcpp/src/stable-diffusion.cpp b/otherarch/sdcpp/src/stable-diffusion.cpp index ae5fd53e36b7..1f2b8b4d231d 100644 --- a/otherarch/sdcpp/src/stable-diffusion.cpp +++ b/otherarch/sdcpp/src/stable-diffusion.cpp @@ -24,6 +24,7 @@ #include "conditioning/conditioner.hpp" #include "core/backend_fit.h" #include "extensions/generation_extension.h" +#include "model/adapter/ip_adapter.hpp" #include "model/adapter/lora.hpp" #include "model/diffusion/anima.hpp" #include "model/diffusion/animatediff.hpp" @@ -223,6 +224,10 @@ class StableDiffusionGGML { std::shared_ptr preview_vae; std::shared_ptr audio_vae_model; std::shared_ptr control_net; + std::shared_ptr ip_adapter; + sd::Tensor ip_adapter_tokens; + sd::Tensor ip_adapter_uncond_tokens; + float ip_adapter_strength = 1.0f; std::vector> generation_extensions; std::vector> runtime_lora_models; bool apply_lora_immediately = false; @@ -1087,6 +1092,13 @@ class StableDiffusionGGML { } } + if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) { + if (!model_loader.init_from_file(sd_ctx_params->ip_adapter_path)) { + LOG_ERROR("init ip-adapter model loader from file failed: '%s'", sd_ctx_params->ip_adapter_path); + return false; + } + } + model_loader.convert_tensors_name(); version = model_loader.get_sd_version(); @@ -1545,6 +1557,33 @@ class StableDiffusionGGML { } } + if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0 && clip_vision == nullptr) { + if (!ensure_backend_pair(SDBackendModule::CLIP_VISION)) { + return false; + } + clip_vision = std::make_shared(backend_for(SDBackendModule::CLIP_VISION), + tensor_storage_map, + model_manager); + clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CLIP_VISION)); + if (!register_runner_params("CLIP vision", + clip_vision, + SDBackendModule::CLIP_VISION)) { + return false; + } + } + + if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) { + ip_adapter = std::make_shared(backend_for(SDBackendModule::DIFFUSION), + tensor_storage_map, + "ip_adapter", + model_manager); + if (!register_runner_params("IP-Adapter", + ip_adapter, + SDBackendModule::DIFFUSION)) { + return false; + } + } + if (!ensure_backend_pair(SDBackendModule::VAE)) { return false; } @@ -2331,6 +2370,34 @@ class StableDiffusionGGML { return output; } + void compute_ip_adapter_tokens(const sd_image_t& image, float strength) { + ip_adapter_tokens = {}; + ip_adapter_uncond_tokens = {}; + ip_adapter_strength = strength; + if (ip_adapter == nullptr || clip_vision == nullptr || image.data == nullptr) { + return; + } + auto image_tensor = sd_image_to_tensor(image); + auto embed = get_clip_vision_output(image_tensor, true, -1); + if (embed.empty()) { + return; + } + ip_adapter_tokens = ip_adapter->compute(n_threads, embed); + if (ip_adapter_tokens.empty()) { + LOG_ERROR("IP-Adapter conditional image projection failed"); + return; + } + auto uncond_embed = sd::Tensor::zeros_like(embed); + ip_adapter_uncond_tokens = ip_adapter->compute(n_threads, uncond_embed); + if (ip_adapter_uncond_tokens.empty()) { + LOG_ERROR("IP-Adapter unconditional image projection failed"); + ip_adapter_tokens = {}; + return; + } + LOG_INFO("IP-Adapter: %lld image tokens, strength %.2f", + (long long)ip_adapter_tokens.shape()[1], strength); + } + std::vector process_timesteps(const std::vector& timesteps, const sd::Tensor& init_latent, const sd::Tensor& denoise_mask, @@ -2819,7 +2886,8 @@ class StableDiffusionGGML { auto run_condition = [&](const SDCondition& condition, const sd::Tensor* c_concat_override = nullptr, const std::vector* local_skip_layers = nullptr, - const std::vector>* ref_latents_override = nullptr) -> sd::Tensor { + const std::vector>* ref_latents_override = nullptr, + bool use_uncond_ip = false) -> sd::Tensor { diffusion_params.context = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn; diffusion_params.c_concat = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat); diffusion_params.y = condition.c_vector.empty() ? nullptr : &condition.c_vector; @@ -2830,7 +2898,13 @@ class StableDiffusionGGML { if (animatediff_loaded && noised_input.dim() >= 4 && noised_input.shape()[3] > 1) { nvf = static_cast(noised_input.shape()[3]); } - diffusion_params.extra = UNetDiffusionExtra{nvf, &controls, control_strength}; + UNetDiffusionExtra unet_extra{nvf, &controls, control_strength}; + const auto& ip_tokens = use_uncond_ip ? ip_adapter_uncond_tokens : ip_adapter_tokens; + if (!ip_tokens.empty()) { + unet_extra.ip_context = &ip_tokens; + unet_extra.ip_scale = ip_adapter_strength; + } + diffusion_params.extra = unet_extra; } else if (sd_version_is_sd3(version)) { diffusion_params.extra = SkipLayerDiffusionExtra{local_skip_layers}; } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) { @@ -2921,7 +2995,9 @@ class StableDiffusionGGML { } uncond_out = run_condition(uncond, uncond.c_concat.empty() ? nullptr : &uncond.c_concat, - uncond_skip_layers); + uncond_skip_layers, + nullptr, + true); if (uncond_out.empty()) { return {}; } @@ -2930,7 +3006,8 @@ class StableDiffusionGGML { 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); + uncond_without_ref_latents ? &empty_ref_latents : nullptr, + true); if (img_uncond_out.empty()) { return {}; } @@ -3778,21 +3855,22 @@ char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) { void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) { *sd_img_gen_params = {}; sd_sample_params_init(&sd_img_gen_params->sample_params); - sd_img_gen_params->clip_skip = -1; - sd_img_gen_params->ref_images_count = 0; - sd_img_gen_params->ref_image_args = ""; - sd_img_gen_params->width = 512; - sd_img_gen_params->height = 512; - sd_img_gen_params->strength = 0.75f; - sd_img_gen_params->seed = -1; - sd_img_gen_params->batch_count = 1; - sd_img_gen_params->control_strength = 0.9f; - sd_img_gen_params->qwen_image_layers = 3; - sd_img_gen_params->circular_x = false; - sd_img_gen_params->circular_y = false; - sd_img_gen_params->pm_params = {nullptr, 0, nullptr, 20.f}; - sd_img_gen_params->pulid_params = {nullptr, 1.0f}; - sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr}; + sd_img_gen_params->clip_skip = -1; + sd_img_gen_params->ref_images_count = 0; + sd_img_gen_params->ref_image_args = ""; + sd_img_gen_params->width = 512; + sd_img_gen_params->height = 512; + sd_img_gen_params->strength = 0.75f; + sd_img_gen_params->seed = -1; + sd_img_gen_params->batch_count = 1; + sd_img_gen_params->control_strength = 0.9f; + sd_img_gen_params->ip_adapter_strength = 1.0f; + sd_img_gen_params->qwen_image_layers = 3; + sd_img_gen_params->circular_x = false; + sd_img_gen_params->circular_y = false; + sd_img_gen_params->pm_params = {nullptr, 0, nullptr, 20.f}; + sd_img_gen_params->pulid_params = {nullptr, 1.0f}; + sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr}; sd_cache_params_init(&sd_img_gen_params->cache); sd_hires_params_init(&sd_img_gen_params->hires); } @@ -5297,6 +5375,7 @@ static std::optional prepare_image_generation_embeds(sd_c request->pulid_params, condition_params, plan->total_steps); + sd_ctx->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_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, @@ -5348,8 +5427,8 @@ static std::optional prepare_image_generation_embeds(sd_c } condition_params.text = request->negative_prompt; condition_params.zero_out_masked = zero_out_masked; + std::vector> empty_ref_images; if (use_ref_latent_img_cfg) { - std::vector> empty_ref_images; condition_params.ref_images = &empty_ref_images; } img_uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, @@ -6091,7 +6170,7 @@ static std::optional prepare_video_generation_latents(sd auto encode_condition_frame = [&](const sd::Tensor& image, int64_t latent_frame, const char* name) -> bool { - auto encoded = sd_ctx->sd->encode_first_stage(image); + auto encoded = sd_ctx->sd->encode_first_stage(image.unsqueeze(2)); if (encoded.empty()) { LOG_ERROR("failed to encode Hunyuan Video %s conditioning frame", name); return false; From 5d097c9de60222216fb057a97d9aa49a8abd14f9 Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Sat, 18 Jul 2026 10:38:40 -0300 Subject: [PATCH 5/5] sd: expose ref_image_args and make it trigger edit mode --- expose.h | 1 + koboldcpp.py | 4 ++++ otherarch/sdcpp/sdtype_adapter.cpp | 28 ++++++++++++++++++++-------- 3 files changed, 25 insertions(+), 8 deletions(-) diff --git a/expose.h b/expose.h index 56d7c0317918..8370a7217fc5 100644 --- a/expose.h +++ b/expose.h @@ -241,6 +241,7 @@ struct sd_generation_inputs const char * scheduler = nullptr; const float eta = -1.0f; const char * extra_sample_args = nullptr; + const char * ref_image_args = nullptr; const int clip_skip = -1; const int vid_req_frames = 1; const int vid_fps = 16; diff --git a/koboldcpp.py b/koboldcpp.py index ae378e959065..32a5e1d1e2ad 100644 --- a/koboldcpp.py +++ b/koboldcpp.py @@ -433,6 +433,7 @@ class sd_generation_inputs(ctypes.Structure): ("scheduler", ctypes.c_char_p), ("eta", ctypes.c_float), ("extra_sample_args", ctypes.c_char_p), + ("ref_image_args", ctypes.c_char_p), ("clip_skip", ctypes.c_int), ("vid_req_frames", ctypes.c_int), ("vid_fps", ctypes.c_int), @@ -2698,6 +2699,7 @@ def gendefaults_parse_meta_field(value): 'cache-option': 'cache_options', 'cache_option': 'cache_options', 'extra-sample-args': 'extra_sample_args', + 'ref-image-args': 'ref_image_args', } parsed = parse_json_object(value, 'gendefaults') or {} result = {} @@ -2867,6 +2869,7 @@ def build_prompt(prompt, forced): sample_method = (genparams.get("sampler_name") or "default") scheduler = (genparams.get("scheduler") or "default").lower() extra_sample_args = str(genparams.get("extra_sample_args") or "") + ref_image_args = str(genparams.get("ref_image_args") or "").strip() clip_skip = tryparseint(genparams.get("clip_skip", -1),-1) eta = tryparsefloat(genparams.get("eta", None), None) vid_req_frames = tryparseint(genparams.get("frames", 1),1) @@ -2938,6 +2941,7 @@ def build_prompt(prompt, forced): inputs.vid_fps = vid_fps inputs.video_output_type = video_output_type inputs.remove_limits = allow_remove_limits + inputs.ref_image_args = ref_image_args.encode("UTF-8") inputs.circular_x = tryparseint(adapter_obj.get("circular_x", genparams.get("circular_x",0)),0) inputs.circular_y = tryparseint(adapter_obj.get("circular_y", genparams.get("circular_y",0)),0) inputs.cache_mode = cache_mode.encode("UTF-8") diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp index c42ce7654124..95926ace3f16 100644 --- a/otherarch/sdcpp/sdtype_adapter.cpp +++ b/otherarch/sdcpp/sdtype_adapter.cpp @@ -125,6 +125,7 @@ struct SDParams { bool diffusion_flash_attn = false; bool diffusion_conv_direct = false; bool vae_conv_direct = false; + std::string ref_image_args = ""; LoraMap lora_map; bool lora_dynamic = false; @@ -995,7 +996,6 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) { extra_image_data.push_back(std::string(inputs.extra_images[i])); } - sd_params->prompt = inputs.prompt; sd_params->negative_prompt = inputs.negative_prompt; sd_params->cfg_scale = inputs.cfg_scale; @@ -1004,6 +1004,18 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) sd_params->shifted_timestep = inputs.shifted_timestep; sd_params->flow_shift = inputs.flow_shift; sd_params->extra_sample_args = inputs.extra_sample_args ? inputs.extra_sample_args : ""; + bool force_image_edit = false; + sd_params->ref_image_args = "resize_before_vae=on"; // auto_resize_ref_image = true; + if (inputs.ref_image_args && *inputs.ref_image_args) { + sd_params->ref_image_args += ","; + sd_params->ref_image_args += inputs.ref_image_args; + if (sd_params->ref_image_args.find("preset") != std::string::npos) { + force_image_edit = true; + if(!sd_is_quiet && sddebugmode==1) { + printf("ref_image_args=\"%s\", forcing edit mode", inputs.ref_image_args); + } + } + } sd_params->eta = inputs.eta; sd_params->seed = inputs.seed; sd_params->width = inputs.width; @@ -1060,11 +1072,14 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) //if a single extra image is provided, mask is NOT provided, and img2img image is NOT provided //and it's a (SD1.5, SDXL) model that doesn't support extra images (see extra_image_data later) //swap extra image data into img2img instead (graceful fallback) - if(!supports_reference_images(info) && extra_image_data.size()==1 && !is_img2img && img2img_mask=="") + if(!supports_reference_images(info) && !force_image_edit && extra_image_data.size()==1 && !is_img2img && img2img_mask=="") { is_img2img = true; img2img_data = extra_image_data[0]; extra_image_data.clear(); + if (!sd_is_quiet && sddebugmode == 1) { + printf("Switching reference image to img2img\n"); + } } if ((info.is_wan || info.is_ltx) && extra_image_data.size() == 0 && is_img2img) @@ -1147,7 +1162,7 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) { int nx2, ny2, nc2; int desiredchannels = 3; - if(supports_reference_images(info)) + if(supports_reference_images(info)||force_image_edit) { if(info.is_wan || info.is_ltx) { @@ -1163,7 +1178,7 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) wan_imgs.push_back(extraimage_reference); } } - else if(info.supports_ref_image) + else if(info.supports_ref_image||force_image_edit) { uint8_t * loaded = load_image_from_b64(extra_image_data[i],nx2,ny2); if(loaded) @@ -1238,13 +1253,10 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) } } - std::string ref_image_args; - ref_image_args += "resize_before_vae=on"; // auto_resize_ref_image = true - sd_img_gen_params_t params = {}; sd_img_gen_params_init (¶ms); params.batch_count = 1; - params.ref_image_args = ref_image_args.c_str(); + params.ref_image_args = sd_params->ref_image_args.c_str(); params.prompt = sd_params->prompt.c_str(); params.negative_prompt = sd_params->negative_prompt.c_str(); params.clip_skip = sd_params->clip_skip;