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/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 facdf690cf40..a2302837a82d 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), @@ -2697,6 +2698,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 = {} @@ -2866,6 +2868,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) @@ -2937,6 +2940,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/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..0869e2f00c70 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")) { + 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) @@ -1241,7 +1256,7 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) 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 = 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; 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,