From 1f9dd9c398d9201213a058a9c1aeb79ffa5c029a Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Mon, 13 Jul 2026 09:42:22 -0300 Subject: [PATCH] sd: sync with master-775-b5d8120 (#2321) * sd: expose extra_sample_args parameter * sd: sync with master-773-1b04283 * sd: sync with master-775-b5d8120 --- Makefile | 2 +- expose.h | 1 + koboldcpp.py | 4 + otherarch/sdcpp/include/stable-diffusion.h | 5 + otherarch/sdcpp/sdtype_adapter.cpp | 3 + .../sdcpp/src/conditioning/conditioner.hpp | 116 ++- otherarch/sdcpp/src/model.h | 11 +- otherarch/sdcpp/src/model/common/block.hpp | 2 +- otherarch/sdcpp/src/model/common/rope.hpp | 34 + otherarch/sdcpp/src/model/diffusion/boogu.hpp | 6 +- otherarch/sdcpp/src/model/diffusion/dit.hpp | 53 +- otherarch/sdcpp/src/model/diffusion/krea2.hpp | 153 +++- .../src/model/diffusion/lingbot_video.hpp | 692 ++++++++++++++++++ .../sdcpp/src/model/diffusion/qwen_image.hpp | 4 +- .../sdcpp/src/model/diffusion/z_image.hpp | 4 +- .../sdcpp/src/model_io/safetensors_io.cpp | 58 ++ otherarch/sdcpp/src/model_io/safetensors_io.h | 3 + otherarch/sdcpp/src/model_loader.cpp | 25 + otherarch/sdcpp/src/model_loader.h | 1 + otherarch/sdcpp/src/model_manager.cpp | 96 +++ otherarch/sdcpp/src/model_manager.h | 3 + otherarch/sdcpp/src/stable-diffusion.cpp | 155 +++- 22 files changed, 1366 insertions(+), 65 deletions(-) create mode 100644 otherarch/sdcpp/src/model/diffusion/lingbot_video.hpp diff --git a/Makefile b/Makefile index 717c81de9..d142d066c 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/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/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_MAIN_BASENAMES := examples/cli/image_metadata.cpp examples/cli/image_metadata.h examples/cli/main.cpp examples/cli/msf_gif.h examples/common/common.cpp examples/common/common.h examples/common/log.cpp examples/common/log.h examples/common/media_io.cpp examples/common/media_io.h examples/common/resource_owners.hpp src/tokenizers/vocab/clip_merges.hpp src/tokenizers/vocab/gemma2_merges.hpp src/tokenizers/vocab/gemma2_vocab.hpp src/tokenizers/vocab/gemma_merges.hpp src/tokenizers/vocab/gemma_vocab.hpp src/tokenizers/vocab/gpt_oss_merges.hpp src/tokenizers/vocab/gpt_oss_vocab.hpp src/tokenizers/vocab/mistral_merges.hpp src/tokenizers/vocab/mistral_vocab.hpp src/tokenizers/vocab/qwen_merges.hpp src/tokenizers/vocab/t5.hpp src/tokenizers/vocab/umt5.hpp src/tokenizers/vocab/vocab.cpp src/convert.cpp src/version.cpp diff --git a/expose.h b/expose.h index 202427586..56d7c0317 100644 --- a/expose.h +++ b/expose.h @@ -240,6 +240,7 @@ struct sd_generation_inputs const char * sample_method = nullptr; const char * scheduler = nullptr; const float eta = -1.0f; + const char * extra_sample_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 8f3716ecd..86d0279d9 100644 --- a/koboldcpp.py +++ b/koboldcpp.py @@ -431,6 +431,7 @@ class sd_generation_inputs(ctypes.Structure): ("sample_method", ctypes.c_char_p), ("scheduler", ctypes.c_char_p), ("eta", ctypes.c_float), + ("extra_sample_args", ctypes.c_char_p), ("clip_skip", ctypes.c_int), ("vid_req_frames", ctypes.c_int), ("vid_fps", ctypes.c_int), @@ -2677,6 +2678,7 @@ def gendefaults_parse_meta_field(value): # match sd.cpp flag 'cache-option': 'cache_options', 'cache_option': 'cache_options', + 'extra-sample-args': 'extra_sample_args', } parsed = parse_json_object(value, 'gendefaults') or {} result = {} @@ -2845,6 +2847,7 @@ def sd_generate(genparams): seed = random.randint(100000, 999999) 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 "") 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) @@ -2910,6 +2913,7 @@ def sd_generate(genparams): inputs.sample_method = sd_sampler_canonical_name(sample_method).encode("UTF-8") inputs.scheduler = scheduler.encode("UTF-8") inputs.eta = -1.0 if eta is None else eta + inputs.extra_sample_args = extra_sample_args.encode("UTF-8") inputs.clip_skip = clip_skip inputs.vid_req_frames = vid_req_frames inputs.vid_fps = vid_fps diff --git a/otherarch/sdcpp/include/stable-diffusion.h b/otherarch/sdcpp/include/stable-diffusion.h index 3fc6f64b7..95d507c77 100644 --- a/otherarch/sdcpp/include/stable-diffusion.h +++ b/otherarch/sdcpp/include/stable-diffusion.h @@ -429,6 +429,11 @@ SD_API const char* sd_get_system_info(); SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx); SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx); +// ControlNet hot-swap APIs are not safe to call while generation is in flight. +SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path); +SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx); +SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx); + SD_API const char* sd_type_name(enum sd_type_t type); SD_API enum sd_type_t str_to_sd_type(const char* str); SD_API const char* sd_rng_type_name(enum rng_type_t rng_type); diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp index 407fb1ca6..9e6a5a1bd 100644 --- a/otherarch/sdcpp/sdtype_adapter.cpp +++ b/otherarch/sdcpp/sdtype_adapter.cpp @@ -117,6 +117,7 @@ struct SDParams { float distilled_guidance = -1.0f; float shifted_timestep = 0; float flow_shift = -1.0f; + std::string extra_sample_args = ""; float eta = -1.0f; float strength = 0.75f; int64_t seed = 42; @@ -1001,6 +1002,7 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) sd_params->sample_steps = inputs.sample_steps; 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 : ""; sd_params->eta = inputs.eta; sd_params->seed = inputs.seed; sd_params->width = inputs.width; @@ -1259,6 +1261,7 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) if (sd_params->flow_shift > 0.f && sd_params->flow_shift != INFINITY) { params.sample_params.flow_shift = sd_params->flow_shift; } + params.sample_params.extra_sample_args = sd_params->extra_sample_args.c_str(); params.seed = sd_params->seed; params.strength = sd_params->strength; params.vae_tiling_params.enabled = dotile; diff --git a/otherarch/sdcpp/src/conditioning/conditioner.hpp b/otherarch/sdcpp/src/conditioning/conditioner.hpp index 5d4ad7fbb..61b5791c9 100644 --- a/otherarch/sdcpp/src/conditioning/conditioner.hpp +++ b/otherarch/sdcpp/src/conditioning/conditioner.hpp @@ -1797,7 +1797,7 @@ struct LLMEmbedder : public Conditioner { arch = LLM::LLMArch::GPT_OSS_20B; } else if (sd_version_is_pid(version)) { arch = LLM::LLMArch::GEMMA2_2B; - } else if (sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) { + } else if (sd_version_is_lingbot_video(version) || sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) { arch = LLM::LLMArch::QWEN3_VL; } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) { arch = LLM::LLMArch::QWEN3; @@ -2009,7 +2009,84 @@ struct LLMEmbedder : public Conditioner { int64_t t0 = ggml_time_ms(); - if (sd_version_is_qwen_image(version)) { + if (sd_version_is_lingbot_video(version)) { + const int pad_token = 151643; + const std::string prompt_prefix = + "<|im_start|>system\nGiven a user input that may include a text prompt alone, " + "a text prompt with an image reference, or a text prompt with a video reference " + "or a video reference alone, generate an \"Enhanced prompt\" that provides detailed " + "visual descriptions suitable for video generation. Evaluate the level of detail " + "in the user's input: if it is simple, enrich it by adding specifics about colors, " + "shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal " + "progression, and spatial relationships to create vivid, concrete, and temporally " + "coherent scenes to create vivid and concrete scenes. Please generate only the " + "enhanced description for the prompt below and avoid including any additional " + "commentary or evaluations:<|im_end|>\n<|im_start|>user\n"; + + auto prefix_tokens = tokenizer->encode(prompt_prefix, nullptr); + prompt_template_encode_start_idx = 0; + for (int token : prefix_tokens) { + if (token != pad_token) { + prompt_template_encode_start_idx++; + } + } + LOG_DEBUG("prompt_template_encode_start_idx %d", prompt_template_encode_start_idx); + + prompt = prompt_prefix; + if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { + LOG_INFO("LingBotVideoI2VPipeline"); + const 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; + 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)); + + 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); + } + + 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); + auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); + GGML_ASSERT(!image_embed.empty()); + + std::string image_prefix = prompt + img_prompt + "<|vision_start|>"; + int image_embed_idx = static_cast(tokenizer->encode(image_prefix, nullptr).size()); + image_embeds.emplace_back(image_embed_idx, image_embed); + + img_prompt += "<|vision_start|>"; + int64_t num_image_tokens = image_embed.shape()[1]; + img_prompt.reserve(img_prompt.size() + static_cast(num_image_tokens) * placeholder.size() + 32); + for (int j = 0; j < num_image_tokens; j++) { + img_prompt += placeholder; + } + img_prompt += "<|vision_end|>"; + } + prompt += img_prompt; + } + + prompt += conditioner_params.text; + prompt_attn_range = {0, 0}; + prompt += "<|im_end|>\n<|im_start|>assistant\n"; + } else if (sd_version_is_qwen_image(version)) { if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { LOG_INFO("QwenImageEditPlusPipeline"); prompt_template_encode_start_idx = 64; @@ -2141,6 +2218,41 @@ struct LLMEmbedder : public Conditioner { out_layers = {2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35}; prompt = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n"; + if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { + std::string img_prompt = ""; + const std::string placeholder = "<|image_pad|>"; + + 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; + 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)); + + LOG_DEBUG("resize conditioner 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); + auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); + GGML_ASSERT(!image_embed.empty()); + + std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>"; + int image_embed_idx = static_cast(tokenizer->encode(image_prefix, nullptr).size()); + image_embeds.emplace_back(image_embed_idx, image_embed); + + img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>"; + int64_t num_image_tokens = image_embed.shape()[1]; + img_prompt.reserve(img_prompt.size() + static_cast(num_image_tokens) * placeholder.size() + 32); + for (int j = 0; j < num_image_tokens; j++) { + img_prompt += placeholder; + } + img_prompt += "<|vision_end|>"; + } + prompt += img_prompt; + } prompt_attn_range.first = static_cast(prompt.size()); prompt += conditioner_params.text; diff --git a/otherarch/sdcpp/src/model.h b/otherarch/sdcpp/src/model.h index 2de82006a..cb2b2039e 100644 --- a/otherarch/sdcpp/src/model.h +++ b/otherarch/sdcpp/src/model.h @@ -35,6 +35,7 @@ enum SDVersion { VERSION_WAN2, VERSION_WAN2_2_I2V, VERSION_WAN2_2_TI2V, + VERSION_LINGBOT_VIDEO, VERSION_QWEN_IMAGE, VERSION_QWEN_IMAGE_LAYERED, VERSION_ANIMA, @@ -127,6 +128,13 @@ static inline bool sd_version_is_wan(SDVersion version) { return false; } +static inline bool sd_version_is_lingbot_video(SDVersion version) { + if (version == VERSION_LINGBOT_VIDEO) { + return true; + } + return false; +} + static inline bool sd_version_is_qwen_image(SDVersion version) { if (version == VERSION_QWEN_IMAGE || version == VERSION_QWEN_IMAGE_LAYERED) { return true; @@ -226,7 +234,7 @@ static inline bool sd_version_uses_flux2_vae(SDVersion version) { } static inline bool sd_version_uses_wan_vae(SDVersion version) { - if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) { + if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) { return true; } return false; @@ -249,6 +257,7 @@ static inline bool sd_version_is_dit(SDVersion version) { sd_version_is_ltxav(version) || sd_version_is_sd3(version) || sd_version_is_wan(version) || + sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || version == VERSION_HIDREAM_O1 || sd_version_is_anima(version) || diff --git a/otherarch/sdcpp/src/model/common/block.hpp b/otherarch/sdcpp/src/model/common/block.hpp index 15bfa3767..ec8f40d0b 100644 --- a/otherarch/sdcpp/src/model/common/block.hpp +++ b/otherarch/sdcpp/src/model/common/block.hpp @@ -294,7 +294,7 @@ public: auto net_0 = std::dynamic_pointer_cast(blocks["net.0"]); auto net_2 = std::dynamic_pointer_cast(blocks["net.2"]); - if (sd_backend_is(ctx->backend, "Vulkan")) { + if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { net_2->set_force_prec_f32(true); } diff --git a/otherarch/sdcpp/src/model/common/rope.hpp b/otherarch/sdcpp/src/model/common/rope.hpp index d6cbb2e6a..04ecde3c5 100644 --- a/otherarch/sdcpp/src/model/common/rope.hpp +++ b/otherarch/sdcpp/src/model/common/rope.hpp @@ -759,6 +759,40 @@ namespace Rope { return embed_nd(ids, bs, static_cast(theta), axes_dim); } + __STATIC_INLINE__ std::vector> gen_lingbot_video_ids(int t, + int h, + int w, + int pt, + int ph, + int pw, + int bs, + int context_len) { + auto vid_ids_repeated = gen_vid_ids(t, h, w, pt, ph, pw, bs, context_len + 1); + + std::vector> txt_ids(bs * context_len, std::vector(3, 0.0f)); + for (int i = 0; i < bs; ++i) { + for (int j = 0; j < context_len; ++j) { + txt_ids[i * context_len + j][0] = static_cast(j + 1); + } + } + + return concat_ids(vid_ids_repeated, txt_ids, bs); + } + + __STATIC_INLINE__ std::vector gen_lingbot_video_pe(int t, + int h, + int w, + int pt, + int ph, + int pw, + int bs, + int context_len, + int theta, + const std::vector& axes_dim) { + std::vector> ids = gen_lingbot_video_ids(t, h, w, pt, ph, pw, bs, context_len); + return embed_nd(ids, bs, static_cast(theta), axes_dim); + } + __STATIC_INLINE__ std::vector> gen_qwen2vl_ids(int grid_h, int grid_w, int merge_size, diff --git a/otherarch/sdcpp/src/model/diffusion/boogu.hpp b/otherarch/sdcpp/src/model/diffusion/boogu.hpp index 27e13aebd..9ab2dccbc 100644 --- a/otherarch/sdcpp/src/model/diffusion/boogu.hpp +++ b/otherarch/sdcpp/src/model/diffusion/boogu.hpp @@ -199,7 +199,7 @@ namespace Boogu { auto linear_2 = std::dynamic_pointer_cast(blocks["linear_2"]); auto linear_3 = std::dynamic_pointer_cast(blocks["linear_3"]); - if (sd_backend_is(ctx->backend, "Vulkan")) { + if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { linear_2->set_force_prec_f32(true); } @@ -259,7 +259,7 @@ namespace Boogu { auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); auto to_out_0 = std::dynamic_pointer_cast(blocks["to_out.0"]); - if (sd_backend_is(ctx->backend, "Vulkan")) { + if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { to_out_0->set_force_prec_f32(true); } @@ -383,7 +383,7 @@ namespace Boogu { auto instruct_out = std::dynamic_pointer_cast(blocks["processor.instruct_out"]); auto img_out = std::dynamic_pointer_cast(blocks["processor.img_out"]); - if (sd_backend_is(ctx->backend, "Vulkan")) { + if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { to_out_0->set_force_prec_f32(true); } diff --git a/otherarch/sdcpp/src/model/diffusion/dit.hpp b/otherarch/sdcpp/src/model/diffusion/dit.hpp index a88e7546d..e7d7b67f8 100644 --- a/otherarch/sdcpp/src/model/diffusion/dit.hpp +++ b/otherarch/sdcpp/src/model/diffusion/dit.hpp @@ -104,14 +104,15 @@ namespace DiT { return x; } - inline ggml_tensor* patchify(ggml_context* ctx, - ggml_tensor* x, - int pt, - int ph, - int pw, - int64_t N = 1) { + inline ggml_tensor* patchify_3d(ggml_context* ctx, + ggml_tensor* x, + int pt, + int ph, + int pw, + int64_t N = 1, + bool patch_last = true) { // x: [N*C, T, H, W] - // return: [N, h*w, C*pt*ph*pw] + // return: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C] int64_t C = x->ne[3] / N; int64_t T = x->ne[2]; int64_t H = x->ne[1]; @@ -123,15 +124,20 @@ namespace DiT { GGML_ASSERT(C * N == x->ne[3]); GGML_ASSERT(t_len * pt == T && h_len * ph == H && w_len * pw == W); - x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N); // [N*C*t_len, pt, h_len*ph, w_len*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, h_len*ph, pt, w_len*pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N); // [N*C*t_len*h_len, ph, pt, w_len*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt, ph, w_len*pw] - x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt*ph, w_len, pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, w_len, pt*ph, pw] - x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N); // [N, C, t_len*h_len*w_len, pt*ph*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, t_len*h_len*w_len, C, pt*ph*pw] - x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1); // [N, t_len*h_len*w_len, C*pt*ph*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N); // [N*C*t_len, pt, h_len*ph, w_len*pw] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, h_len*ph, pt, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N); // [N*C*t_len*h_len, ph, pt, w_len*pw] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt, ph, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt*ph, w_len, pw] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, w_len, pt*ph, pw] + x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N); // [N, C, t_len*h_len*w_len, pt*ph*pw] + if (patch_last) { + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, t_len*h_len*w_len, C, pt*ph*pw] + } else { + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [N, t_len*h_len*w_len, pt*ph*pw, C] + } + // [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C] + x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1); return x; } @@ -142,16 +148,23 @@ namespace DiT { int64_t w_len, int pt, int ph, - int pw) { - // x: [N, t_len*h_len*w_len, C*pt*ph*pw] + int pw, + bool patch_last = true) { + // x: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, pt*ph*pw*C] // return: [N*C, t_len*pt, h_len*ph, w_len*pw] int64_t N = x->ne[2]; int64_t C = x->ne[0] / pt / ph / pw; GGML_ASSERT(C * pt * ph * pw == x->ne[0]); - x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] + if (patch_last) { + x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*pw] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] + } else { + x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] + } + x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw] x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw] x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw] diff --git a/otherarch/sdcpp/src/model/diffusion/krea2.hpp b/otherarch/sdcpp/src/model/diffusion/krea2.hpp index 02e655590..8ca0bb51f 100644 --- a/otherarch/sdcpp/src/model/diffusion/krea2.hpp +++ b/otherarch/sdcpp/src/model/diffusion/krea2.hpp @@ -267,7 +267,7 @@ namespace Krea2 { auto knorm = std::dynamic_pointer_cast(blocks["qknorm.knorm"]); auto wo = std::dynamic_pointer_cast(blocks["wo"]); - if (sd_backend_is(ctx->backend, "Vulkan")) { + if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { wo->set_force_prec_f32(true); } @@ -421,29 +421,88 @@ namespace Krea2 { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* vec, - ggml_tensor* pe) { + ggml_tensor* pe, + ggml_tensor* vec_refs = nullptr, + int64_t ref_start = -1) { auto mod = std::dynamic_pointer_cast(blocks["mod"]); auto prenorm = std::dynamic_pointer_cast(blocks["prenorm"]); auto postnorm = std::dynamic_pointer_cast(blocks["postnorm"]); auto attn = std::dynamic_pointer_cast(blocks["attn"]); auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); - auto mods = mod->forward(ctx, vec); - auto attn_input = Flux::modulate(ctx->ggml_ctx, - prenorm->forward(ctx, x), - mods[1], - mods[0], - true); - auto attn_out = attn->forward(ctx, attn_input, pe); - x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2])); + if (ref_start >= 0 && vec_refs) { + // same as normal, but since vec is different for refs and the rest, needs a lot of views and concats + auto mods_main = mod->forward(ctx, vec); + auto mods_refs = mod->forward(ctx, vec_refs); - auto mlp_input = Flux::modulate(ctx->ggml_ctx, - postnorm->forward(ctx, x), - mods[4], - mods[3], - true); - auto mlp_out = mlp->forward(ctx, mlp_input); - x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mods[5])); + int64_t D = x->ne[0]; + int64_t N = x->ne[1]; + int64_t B = x->ne[2]; + size_t nb1 = x->nb[1]; + size_t nb2 = x->nb[2]; + + int64_t len_main = ref_start; + int64_t len_refs = N - ref_start; + + auto pre_x = prenorm->forward(ctx, x); + + auto pre_x_main = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_main, B, nb1, nb2, 0); + auto pre_x_refs = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_refs, B, nb1, nb2, len_main * nb1); + + auto attn_in_main = Flux::modulate(ctx->ggml_ctx, pre_x_main, mods_main[1], mods_main[0], true); + auto attn_in_refs = Flux::modulate(ctx->ggml_ctx, pre_x_refs, mods_refs[1], mods_refs[0], true); + + auto attn_input = ggml_concat(ctx->ggml_ctx, attn_in_main, attn_in_refs, 1); + + auto attn_out = attn->forward(ctx, attn_input, pe); + + auto attn_out_main = ggml_view_3d(ctx->ggml_ctx, attn_out, D, len_main, B, attn_out->nb[1], attn_out->nb[2], 0); + auto attn_out_refs = ggml_view_3d(ctx->ggml_ctx, attn_out, D, len_refs, B, attn_out->nb[1], attn_out->nb[2], len_main * attn_out->nb[1]); + + auto res_main = ggml_mul(ctx->ggml_ctx, attn_out_main, mods_main[2]); + auto res_refs = ggml_mul(ctx->ggml_ctx, attn_out_refs, mods_refs[2]); + + auto attn_res = ggml_concat(ctx->ggml_ctx, res_main, res_refs, 1); + + x = ggml_add(ctx->ggml_ctx, x, attn_res); + + auto post_x = postnorm->forward(ctx, x); + + auto post_x_main = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_main, B, post_x->nb[1], post_x->nb[2], 0); + auto post_x_refs = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_refs, B, post_x->nb[1], post_x->nb[2], len_main * post_x->nb[1]); + + auto mlp_in_main = Flux::modulate(ctx->ggml_ctx, post_x_main, mods_main[4], mods_main[3], true); + auto mlp_in_refs = Flux::modulate(ctx->ggml_ctx, post_x_refs, mods_refs[4], mods_refs[3], true); + + auto mlp_input = ggml_concat(ctx->ggml_ctx, mlp_in_main, mlp_in_refs, 1); + auto mlp_out = mlp->forward(ctx, mlp_input); + + auto mlp_out_main = ggml_view_3d(ctx->ggml_ctx, mlp_out, D, len_main, B, mlp_out->nb[1], mlp_out->nb[2], 0); + auto mlp_out_refs = ggml_view_3d(ctx->ggml_ctx, mlp_out, D, len_refs, B, mlp_out->nb[1], mlp_out->nb[2], len_main * mlp_out->nb[1]); + + auto mlp_res_main = ggml_mul(ctx->ggml_ctx, mlp_out_main, mods_main[5]); + auto mlp_res_refs = ggml_mul(ctx->ggml_ctx, mlp_out_refs, mods_refs[5]); + + auto mlp_res = ggml_concat(ctx->ggml_ctx, mlp_res_main, mlp_res_refs, 1); + x = ggml_add(ctx->ggml_ctx, x, mlp_res); + } else { + auto mods = mod->forward(ctx, vec); + auto attn_input = Flux::modulate(ctx->ggml_ctx, + prenorm->forward(ctx, x), + mods[1], + mods[0], + true); + auto attn_out = attn->forward(ctx, attn_input, pe); + x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2])); + + auto mlp_input = Flux::modulate(ctx->ggml_ctx, + postnorm->forward(ctx, x), + mods[4], + mods[3], + true); + auto mlp_out = mlp->forward(ctx, mlp_input); + x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mods[5])); + } return x; } @@ -555,7 +614,8 @@ namespace Krea2 { ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, - ggml_tensor* pe) { + ggml_tensor* pe, + std::vector ref_latents = {}) { int64_t W = x->ne[0]; int64_t H = x->ne[1]; int64_t N = x->ne[3]; @@ -570,6 +630,13 @@ namespace Krea2 { auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size, true); int64_t img_len = img->ne[1]; + if (ref_latents.size() > 0) { + for (ggml_tensor* ref : ref_latents) { + ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size, true); + img = ggml_concat(ctx->ggml_ctx, img, ref, 1); + } + } + int64_t ref_len = img->ne[1] - img_len; img = first->forward(ctx, img); auto t = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast(config.timestep_dim), 10000, 1000.f); @@ -577,19 +644,30 @@ namespace Krea2 { t = ggml_reshape_3d(ctx->ggml_ctx, t, t->ne[0], 1, t->ne[1]); auto tvec = tproj->forward(ctx, t); + ggml_tensor* tvec_0 = nullptr; + if (ref_latents.size() > 0) { + // "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); + t_0 = tmlp->forward(ctx, t_0); + t_0 = ggml_reshape_3d(ctx->ggml_ctx, t_0, t_0->ne[0], 1, t_0->ne[1]); + tvec_0 = tproj->forward(ctx, t_0); + } + auto txt = txtfusion->forward(ctx, context); txt = txtmlp->forward(ctx, txt); int64_t txt_len = txt->ne[1]; auto hidden_states = ggml_concat(ctx->ggml_ctx, txt, img, 1); + int64_t ref_start = hidden_states->ne[1] - ref_len; for (int i = 0; i < config.layers; ++i) { auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); - hidden_states = block->forward(ctx, hidden_states, tvec, pe); + hidden_states = block->forward(ctx, hidden_states, tvec, pe, tvec_0, ref_start); sd::ggml_graph_cut::mark_graph_cut(hidden_states, "krea2.blocks." + std::to_string(i), "hidden_states"); } - hidden_states = last->forward(ctx, hidden_states, t); hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, txt_len, txt_len + img_len); + hidden_states = last->forward(ctx, hidden_states, t); hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, true); return hidden_states; } @@ -601,10 +679,16 @@ namespace Krea2 { int bs, int context_len, float theta, - const std::vector& axes_dim) { + const std::vector& axes_dim, + const std::vector& ref_latents, + Rope::RefIndexMode ref_index_mode) { auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {}); auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false); auto ids = Rope::concat_ids(txt_ids, img_ids, bs); + if (ref_latents.size() > 0) { + auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0); + ids = Rope::concat_ids(ids, refs_ids, bs); + } return Rope::embed_nd(ids, bs, theta, axes_dim); } @@ -633,7 +717,9 @@ namespace Krea2 { ggml_cgraph* build_graph(const sd::Tensor& x_tensor, const sd::Tensor& timesteps_tensor, - const sd::Tensor& context_tensor) { + const sd::Tensor& context_tensor, + const std::vector>& ref_latents_tensor = {}, + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) { ggml_cgraph* gf = new_graph_custom(KREA2_GRAPH_SIZE); ggml_tensor* x = make_input(x_tensor); ggml_tensor* timesteps = make_input(timesteps_tensor); @@ -641,19 +727,27 @@ namespace Krea2 { GGML_ASSERT(!context_tensor.empty()); ggml_tensor* context = make_input(context_tensor); + std::vector ref_latents; + ref_latents.reserve(ref_latents_tensor.size()); + for (const auto& ref_latent_tensor : ref_latents_tensor) { + ref_latents.push_back(make_input(ref_latent_tensor)); + } + pe_vec = gen_krea2_pe(static_cast(x->ne[1]), static_cast(x->ne[0]), config.patch_size, static_cast(x->ne[3]), static_cast(context->ne[1]), config.theta, - config.axes_dim); + config.axes_dim, + ref_latents, + 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); + ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents); ggml_build_forward_expand(gf, out); return gf; } @@ -661,9 +755,11 @@ namespace Krea2 { sd::Tensor compute(int n_threads, const sd::Tensor& x, const sd::Tensor& timesteps, - const sd::Tensor& context) { + const sd::Tensor& context, + const std::vector>& ref_latents = {}, + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) { auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context); + return build_graph(x, timesteps, context, ref_latents, ref_index_mode); }; return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); } @@ -672,10 +768,13 @@ namespace Krea2 { const DiffusionParams& diffusion_params) override { GGML_ASSERT(diffusion_params.x != nullptr); GGML_ASSERT(diffusion_params.timesteps != nullptr); + static const std::vector> empty_ref_latents; return compute(n_threads, *diffusion_params.x, *diffusion_params.timesteps, - tensor_or_empty(diffusion_params.context)); + tensor_or_empty(diffusion_params.context), + diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents, + diffusion_params.ref_index_mode); } }; } // namespace Krea2 diff --git a/otherarch/sdcpp/src/model/diffusion/lingbot_video.hpp b/otherarch/sdcpp/src/model/diffusion/lingbot_video.hpp new file mode 100644 index 000000000..1868daa05 --- /dev/null +++ b/otherarch/sdcpp/src/model/diffusion/lingbot_video.hpp @@ -0,0 +1,692 @@ +#ifndef __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__ +#define __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/util.h" +#include "model/common/rope.hpp" +#include "model/diffusion/dit.hpp" +#include "model/diffusion/flux.hpp" +#include "model/diffusion/model.hpp" +#include "model/diffusion/qwen_image.hpp" + +namespace LingBotVideo { + constexpr int LINGBOT_VIDEO_GRAPH_SIZE = 65536; + + struct LingBotVideoConfig { + int patch_t = 1; + int patch_h = 2; + int patch_w = 2; + int64_t in_channels = 16; + int64_t out_channels = 16; + int64_t hidden_size = 2048; + int64_t num_attention_heads = 16; + int64_t depth = 24; + int64_t intermediate_size = 6144; + int64_t text_dim = 2560; + int64_t freq_dim = 256; + float norm_eps = 1e-6f; + int rope_theta = 256; + std::vector axes_dim = {32, 48, 48}; + int axes_dim_sum = 128; + bool qkv_bias = false; + bool out_bias = true; + bool patch_embed_bias = true; + bool timestep_mlp_bias = true; + int64_t num_experts = 0; + int64_t num_experts_per_tok = 8; + int64_t moe_intermediate_size = 512; + int64_t decoder_sparse_step = 1; + int64_t n_shared_experts = 0; + bool norm_topk_prob = true; + float routed_scaling_factor = 1.0f; + int64_t n_group = 0; + int64_t topk_group = 0; + std::set sparse_layers; + + static LingBotVideoConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix) { + LingBotVideoConfig config; + config.depth = 0; + + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + + if (ends_with(name, "patch_embedder.weight") && tensor_storage.n_dims == 2) { + int64_t patch_dim = tensor_storage.ne[0]; + config.hidden_size = tensor_storage.ne[1]; + int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w; + if (patch_dim % patch_volume == 0) { + config.in_channels = patch_dim / patch_volume; + } + } else if (ends_with(name, "text_embedder.linear_1.weight") && tensor_storage.n_dims == 2) { + config.text_dim = tensor_storage.ne[0]; + } else if (ends_with(name, "time_embedder.linear_1.weight") && tensor_storage.n_dims == 2) { + config.freq_dim = tensor_storage.ne[0]; + } else if (ends_with(name, "blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) { + int64_t head_dim = tensor_storage.ne[0]; + if (head_dim > 0) { + config.num_attention_heads = config.hidden_size / head_dim; + } + } else if (name.find(".attn.to_q.bias") != std::string::npos) { + config.qkv_bias = true; + } else if (name.find(".ffn.gate_proj.weight") != std::string::npos && tensor_storage.n_dims == 2) { + config.intermediate_size = tensor_storage.ne[1]; + } else if (name.find(".ffn.experts.w1") != std::string::npos && tensor_storage.n_dims == 3) { + config.num_experts = tensor_storage.ne[2]; + config.moe_intermediate_size = tensor_storage.ne[1]; + } else if (name.find(".ffn.shared_experts.gate_proj.weight") != std::string::npos && tensor_storage.n_dims == 2) { + if (config.moe_intermediate_size > 0) { + config.n_shared_experts = tensor_storage.ne[1] / config.moe_intermediate_size; + } + } else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) { + int64_t out_dim = tensor_storage.ne[1]; + int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w; + config.out_channels = patch_volume > 0 ? out_dim / patch_volume : config.out_channels; + } + + size_t block_pos = name.find("blocks."); + if (block_pos != std::string::npos) { + auto items = split_string(name.substr(block_pos), '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + if (block_index + 1 > config.depth) { + config.depth = block_index + 1; + } + if (name.find("blocks." + std::to_string(block_index) + ".ffn.experts.w1") != std::string::npos) { + config.sparse_layers.insert(block_index); + } + } + } + } + + if (config.depth == 0) { + config.depth = 24; + } + config.axes_dim_sum = 0; + for (int axis_dim : config.axes_dim) { + config.axes_dim_sum += axis_dim; + } + if (!config.sparse_layers.empty()) { + config.num_experts = 128; + config.num_experts_per_tok = 8; + config.moe_intermediate_size = 768; + config.decoder_sparse_step = 1; + config.n_shared_experts = 1; + config.norm_topk_prob = true; + config.n_group = 4; + config.topk_group = 2; + config.routed_scaling_factor = 2.5f; + } + LOG_DEBUG("lingbot_video: depth = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64 ", text_dim = %" PRId64 ", experts = %" PRId64 ", experts_per_tok = %" PRId64 ", n_group = %" PRId64 ", topk_group = %" PRId64 ", route_scale = %.2f, sparse_layers = %zu", + config.depth, + config.hidden_size, + config.num_attention_heads, + config.text_dim, + config.num_experts, + config.num_experts_per_tok, + config.n_group, + config.topk_group, + config.routed_scaling_factor, + config.sparse_layers.size()); + return config; + } + }; + + struct LingBotVideoTextEmbedder : public GGMLBlock { + LingBotVideoTextEmbedder(int64_t text_dim, + int64_t hidden_size, + float eps = 1e-6f) { + blocks["norm"] = std::make_shared(text_dim, eps); + blocks["linear_1"] = std::make_shared(text_dim, hidden_size, true); + blocks["linear_2"] = std::make_shared(hidden_size, hidden_size, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto linear_1 = std::dynamic_pointer_cast(blocks["linear_1"]); + auto linear_2 = std::dynamic_pointer_cast(blocks["linear_2"]); + + x = norm->forward(ctx, x); + x = linear_1->forward(ctx, x); + x = ggml_silu_inplace(ctx->ggml_ctx, x); + x = linear_2->forward(ctx, x); + return x; + } + }; + + struct LingBotVideoAttention : public GGMLBlock { + int64_t num_heads; + int64_t head_dim; + + LingBotVideoAttention(int64_t hidden_size, + int64_t num_heads, + bool qkv_bias = false, + bool out_bias = true, + float eps = 1e-6f) + : num_heads(num_heads), + head_dim(hidden_size / num_heads) { + int64_t inner_dim = num_heads * head_dim; + blocks["to_q"] = std::make_shared(hidden_size, inner_dim, qkv_bias); + blocks["to_k"] = std::make_shared(hidden_size, inner_dim, qkv_bias); + blocks["to_v"] = std::make_shared(hidden_size, inner_dim, qkv_bias); + blocks["norm_q"] = std::make_shared(head_dim, eps); + blocks["norm_k"] = std::make_shared(head_dim, eps); + blocks["to_out"] = std::make_shared(inner_dim, hidden_size, out_bias); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* pe, + ggml_tensor* attention_mask = nullptr) { + // x: [N, video_tokens + text_tokens, hidden_size] + 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 norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); + auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); + auto to_out = std::dynamic_pointer_cast(blocks["to_out"]); + + int64_t S = x->ne[1]; + int64_t N = x->ne[2]; + + auto q = to_q->forward(ctx, x); + auto k = to_k->forward(ctx, x); + auto v = to_v->forward(ctx, x); + + q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, S, N); + k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, S, N); + v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, S, N); + + q = norm_q->forward(ctx, q); + k = norm_k->forward(ctx, k); + + x = Rope::attention(ctx, q, k, v, pe, attention_mask); + x = to_out->forward(ctx, x); + return x; + } + }; + + struct LingBotVideoMLP : public UnaryBlock { + LingBotVideoMLP(int64_t hidden_size, + int64_t intermediate_size) { + blocks["gate_proj"] = std::make_shared(hidden_size, intermediate_size, false); + blocks["up_proj"] = std::make_shared(hidden_size, intermediate_size, false); + blocks["down_proj"] = std::make_shared(intermediate_size, hidden_size, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto gate_proj = std::dynamic_pointer_cast(blocks["gate_proj"]); + auto up_proj = std::dynamic_pointer_cast(blocks["up_proj"]); + auto down_proj = std::dynamic_pointer_cast(blocks["down_proj"]); + + auto gate = gate_proj->forward(ctx, x); + gate = ggml_silu_inplace(ctx->ggml_ctx, gate); + auto up = up_proj->forward(ctx, x); + x = ggml_mul(ctx->ggml_ctx, gate, up); + x = down_proj->forward(ctx, x); + return x; + } + }; + + struct LingBotVideoSparseMoeBlock : public UnaryBlock { + int64_t hidden_size; + int64_t intermediate_size; + int64_t num_experts; + int64_t num_experts_per_tok; + bool has_shared_experts; + bool norm_topk_prob; + float routed_scaling_factor; + int64_t n_group; + int64_t topk_group; + std::vector group_expert_mask_vec; + bool has_correction_bias = false; + + LingBotVideoSparseMoeBlock(const LingBotVideoConfig& config) + : hidden_size(config.hidden_size), + intermediate_size(config.moe_intermediate_size), + num_experts(config.num_experts), + num_experts_per_tok(config.num_experts_per_tok), + has_shared_experts(config.n_shared_experts > 0), + norm_topk_prob(config.norm_topk_prob), + routed_scaling_factor(config.routed_scaling_factor), + n_group(config.n_group), + topk_group(config.topk_group) { + if (n_group > 1) { + GGML_ASSERT(num_experts % n_group == 0); + int64_t experts_per_group = num_experts / n_group; + group_expert_mask_vec.assign(static_cast(num_experts * n_group), 0.f); + for (int64_t group = 0; group < n_group; ++group) { + int64_t expert_begin = group * experts_per_group; + int64_t expert_end = expert_begin + experts_per_group; + for (int64_t expert = expert_begin; expert < expert_end; ++expert) { + group_expert_mask_vec[static_cast(group * num_experts + expert)] = 1.f; + } + } + } + if (has_shared_experts) { + blocks["shared_experts"] = std::make_shared(hidden_size, + intermediate_size * config.n_shared_experts); + } + } + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + ggml_type router_type = get_type(prefix + "router.weight", tensor_storage_map, GGML_TYPE_F32); + ggml_type w1_type = get_type(prefix + "experts.w1", tensor_storage_map, GGML_TYPE_F32); + ggml_type w2_type = get_type(prefix + "experts.w2", tensor_storage_map, GGML_TYPE_F32); + ggml_type w3_type = get_type(prefix + "experts.w3", tensor_storage_map, GGML_TYPE_F32); + params["router.weight"] = ggml_new_tensor_2d(ctx, router_type, hidden_size, num_experts); + if (tensor_storage_map.find(prefix + "router.e_score_correction_bias") != tensor_storage_map.end()) { + params["router.e_score_correction_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts); + has_correction_bias = true; + } + params["experts.w1"] = ggml_new_tensor_3d(ctx, w1_type, hidden_size, intermediate_size, num_experts); + params["experts.w2"] = ggml_new_tensor_3d(ctx, w2_type, intermediate_size, hidden_size, num_experts); + params["experts.w3"] = ggml_new_tensor_3d(ctx, w3_type, hidden_size, intermediate_size, num_experts); + } + + ggml_tensor* expert_linear(GGMLRunnerContext* ctx, + const std::string& weight_name, + ggml_tensor* x, + ggml_tensor* selected_experts) { + return ggml_mul_mat_id(ctx->ggml_ctx, params[weight_name], x, selected_experts); + } + + ggml_tensor* select_experts(GGMLRunnerContext* ctx, ggml_tensor* choice_scores) { + ggml_context* gctx = ctx->ggml_ctx; + if (n_group <= 1 || topk_group <= 0) { + return ggml_argsort_top_k(gctx, choice_scores, static_cast(num_experts_per_tok)); + } + + GGML_ASSERT(choice_scores->ne[0] == num_experts); + GGML_ASSERT(num_experts % n_group == 0); + GGML_ASSERT(topk_group > 0 && topk_group <= n_group); + GGML_ASSERT(!group_expert_mask_vec.empty()); + + const int64_t n_token_total = choice_scores->ne[1]; + const int64_t experts_per_group = num_experts / n_group; + const int group_score_k = 2; + GGML_ASSERT(experts_per_group >= group_score_k); + + ggml_tensor* grouped_scores = ggml_reshape_3d(gctx, choice_scores, experts_per_group, n_group, n_token_total); + ggml_tensor* group_top_ids = ggml_argsort_top_k(gctx, grouped_scores, group_score_k); + grouped_scores = ggml_reshape_3d(gctx, grouped_scores, 1, experts_per_group, n_group * n_token_total); + group_top_ids = ggml_cont(gctx, group_top_ids); + group_top_ids = ggml_reshape_2d(gctx, group_top_ids, group_score_k, n_group * n_token_total); + + ggml_tensor* group_top_values = ggml_get_rows(gctx, grouped_scores, group_top_ids); + group_top_values = ggml_reshape_3d(gctx, group_top_values, group_score_k, n_group, n_token_total); + + ggml_tensor* group_scores = nullptr; + for (int rank = 0; rank < group_score_k; ++rank) { + ggml_tensor* value = ggml_view_3d(gctx, + group_top_values, + 1, + n_group, + n_token_total, + group_top_values->nb[1], + group_top_values->nb[2], + rank * group_top_values->nb[0]); + group_scores = group_scores == nullptr ? value : ggml_add(gctx, group_scores, value); + } + group_scores = ggml_reshape_2d(gctx, group_scores, n_group, n_token_total); + + ggml_tensor* selected_groups = ggml_argsort_top_k(gctx, group_scores, static_cast(topk_group)); + selected_groups = ggml_cont(gctx, selected_groups); + + ggml_tensor* group_expert_mask = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, num_experts, n_group, 1); + ctx->bind_backend_tensor_data(group_expert_mask, group_expert_mask_vec.data()); + ggml_tensor* group_expert_mask_template = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, num_experts, n_group, n_token_total); + group_expert_mask = ggml_repeat(gctx, group_expert_mask, group_expert_mask_template); + + ggml_tensor* selected_group_masks = ggml_get_rows(gctx, group_expert_mask, selected_groups); + ggml_tensor* selected_mask = nullptr; + for (int64_t rank = 0; rank < topk_group; ++rank) { + ggml_tensor* mask = ggml_view_3d(gctx, + selected_group_masks, + num_experts, + 1, + n_token_total, + selected_group_masks->nb[1], + selected_group_masks->nb[2], + rank * selected_group_masks->nb[1]); + selected_mask = selected_mask == nullptr ? mask : ggml_add(gctx, selected_mask, mask); + } + selected_mask = ggml_reshape_2d(gctx, selected_mask, num_experts, n_token_total); + + ggml_tensor* excluded_group_mask = ggml_sub(gctx, selected_mask, ggml_ext_ones_like(gctx, selected_mask)); + ggml_tensor* masked_scores = ggml_add(gctx, choice_scores, ggml_scale(gctx, excluded_group_mask, 1.0e9f)); + return ggml_argsort_top_k(gctx, masked_scores, static_cast(num_experts_per_tok)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [N, tokens, hidden_size] + GGML_ASSERT(num_experts > 0); + GGML_ASSERT(num_experts_per_tok > 0 && num_experts_per_tok <= num_experts); + + ggml_tensor* residual = x; + const int64_t n_token = x->ne[1]; + const int64_t N = x->ne[2]; + const int64_t n_token_total = n_token * N; + + ggml_tensor* router_logits = ggml_mul_mat(ctx->ggml_ctx, params["router.weight"], x); + router_logits = ggml_reshape_2d(ctx->ggml_ctx, router_logits, num_experts, n_token_total); + ggml_tensor* probs = ggml_sigmoid(ctx->ggml_ctx, router_logits); + ggml_tensor* choice_scores = probs; + if (has_correction_bias) { + choice_scores = ggml_add(ctx->ggml_ctx, choice_scores, params["router.e_score_correction_bias"]); + } + + ggml_tensor* selected_experts = select_experts(ctx, choice_scores); + ggml_tensor* weights = ggml_get_rows(ctx->ggml_ctx, + ggml_reshape_3d(ctx->ggml_ctx, probs, 1, num_experts, n_token_total), + selected_experts); + weights = ggml_reshape_2d(ctx->ggml_ctx, weights, num_experts_per_tok, n_token_total); + if (norm_topk_prob && num_experts_per_tok > 1) { + auto weights_sum = ggml_sum_rows(ctx->ggml_ctx, weights); + weights_sum = ggml_clamp(ctx->ggml_ctx, weights_sum, 6.103515625e-5f, INFINITY); + weights = ggml_div(ctx->ggml_ctx, weights, weights_sum); + } + if (routed_scaling_factor != 1.0f) { + weights = ggml_scale(ctx->ggml_ctx, weights, routed_scaling_factor); + } + weights = ggml_reshape_3d(ctx->ggml_ctx, weights, 1, num_experts_per_tok, n_token_total); + + x = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size, 1, n_token_total); + auto gate = expert_linear(ctx, "experts.w1", x, selected_experts); + gate = ggml_silu_inplace(ctx->ggml_ctx, gate); + auto up = expert_linear(ctx, "experts.w3", x, selected_experts); + auto act = ggml_mul(ctx->ggml_ctx, gate, up); + auto out = expert_linear(ctx, "experts.w2", act, selected_experts); + out = ggml_mul(ctx->ggml_ctx, out, weights); + + ggml_tensor* summed = nullptr; + for (int64_t i = 0; i < num_experts_per_tok; ++i) { + auto expert_out = ggml_view_2d(ctx->ggml_ctx, + out, + hidden_size, + n_token_total, + out->nb[2], + i * out->nb[1]); + summed = summed == nullptr ? expert_out : ggml_add(ctx->ggml_ctx, summed, expert_out); + } + if (num_experts_per_tok == 1) { + summed = ggml_cont(ctx->ggml_ctx, summed); + } + summed = ggml_reshape_3d(ctx->ggml_ctx, summed, hidden_size, n_token, N); + + if (has_shared_experts) { + auto shared_experts = std::dynamic_pointer_cast(blocks["shared_experts"]); + summed = ggml_add(ctx->ggml_ctx, summed, shared_experts->forward(ctx, residual)); + } + return summed; + } + }; + + struct LingBotVideoBlock : public GGMLBlock { + int64_t hidden_size; + + LingBotVideoBlock(const LingBotVideoConfig& config, + bool sparse) + : hidden_size(config.hidden_size) { + blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps); + blocks["attn"] = std::make_shared(config.hidden_size, + config.num_attention_heads, + config.qkv_bias, + config.out_bias, + config.norm_eps); + blocks["norm_post_attn"] = std::make_shared(config.hidden_size, config.norm_eps); + blocks["norm2"] = std::make_shared(config.hidden_size, config.norm_eps); + if (sparse) { + blocks["ffn"] = std::make_shared(config); + } else { + blocks["ffn"] = std::make_shared(config.hidden_size, config.intermediate_size); + } + blocks["norm_post_ffn"] = std::make_shared(config.hidden_size, config.norm_eps); + } + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + params["scale_shift_table"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size * 6, 1); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* temb6, + ggml_tensor* pe, + ggml_tensor* attention_mask = nullptr) { + // x: [N, tokens, hidden_size], temb6: [N, tokens, 6 * hidden_size] + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto norm_post_attn = std::dynamic_pointer_cast(blocks["norm_post_attn"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto ffn = std::dynamic_pointer_cast(blocks["ffn"]); + auto norm_post_ffn = std::dynamic_pointer_cast(blocks["norm_post_ffn"]); + + auto table = ggml_reshape_3d(ctx->ggml_ctx, params["scale_shift_table"], hidden_size * 6, 1, 1); + auto mod = ggml_add(ctx->ggml_ctx, temb6, table); + auto mods = ggml_ext_chunk(ctx->ggml_ctx, mod, 6, 0); + + auto shift_msa = mods[0]; + auto scale_msa = mods[1]; + auto gate_msa = ggml_tanh(ctx->ggml_ctx, mods[2]); + auto shift_mlp = mods[3]; + auto scale_mlp = mods[4]; + auto gate_mlp = ggml_tanh(ctx->ggml_ctx, mods[5]); + + auto attn_in = Flux::modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa, true); + auto attn_out = attn->forward(ctx, attn_in, pe, attention_mask); + attn_out = norm_post_attn->forward(ctx, attn_out); + x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_msa, attn_out)); + + auto ffn_in = Flux::modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp, true); + auto ffn_out = ffn->forward(ctx, ffn_in); + ffn_out = norm_post_ffn->forward(ctx, ffn_out); + x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_mlp, ffn_out)); + return x; + } + }; + + class LingBotVideoModel : public GGMLBlock { + public: + LingBotVideoConfig config; + + LingBotVideoModel() = default; + LingBotVideoModel(LingBotVideoConfig config) + : config(config) { + int64_t patch_dim = config.patch_t * config.patch_h * config.patch_w * config.in_channels; + blocks["patch_embedder"] = std::make_shared(patch_dim, config.hidden_size, config.patch_embed_bias); + blocks["time_embedder"] = std::make_shared(config.freq_dim, + config.hidden_size, + config.hidden_size, + 0, + config.timestep_mlp_bias); + blocks["time_modulation.1"] = std::make_shared(config.hidden_size, 6 * config.hidden_size, true); + blocks["text_embedder"] = std::make_shared(config.text_dim, + config.hidden_size, + config.norm_eps); + for (int i = 0; i < config.depth; i++) { + bool sparse = config.sparse_layers.find(i) != config.sparse_layers.end(); + blocks["blocks." + std::to_string(i)] = std::make_shared(config, sparse); + } + blocks["norm_out"] = std::make_shared(config.hidden_size, config.norm_eps, false); + blocks["norm_out_modulation.1"] = std::make_shared(config.hidden_size, 2 * config.hidden_size, true); + blocks["proj_out"] = std::make_shared(config.hidden_size, + config.patch_t * config.patch_h * config.patch_w * config.out_channels, + true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timestep, + ggml_tensor* context, + ggml_tensor* pe) { + // x: [N*C, T, H, W], context: [N, text_tokens, text_dim] + GGML_ASSERT(context != nullptr); + GGML_ASSERT(x->ne[3] == config.in_channels); + GGML_ASSERT(x->ne[2] % config.patch_t == 0); + GGML_ASSERT(x->ne[1] % config.patch_h == 0); + GGML_ASSERT(x->ne[0] % config.patch_w == 0); + + auto patch_embedder = std::dynamic_pointer_cast(blocks["patch_embedder"]); + auto time_embedder = std::dynamic_pointer_cast(blocks["time_embedder"]); + auto time_modulation = std::dynamic_pointer_cast(blocks["time_modulation.1"]); + auto text_embedder = std::dynamic_pointer_cast(blocks["text_embedder"]); + auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); + auto norm_out_modulation = std::dynamic_pointer_cast(blocks["norm_out_modulation.1"]); + auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); + + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t T = x->ne[2]; + int64_t N = 1; + int64_t t_len = T / config.patch_t; + int64_t h_len = H / config.patch_h; + int64_t w_len = W / config.patch_w; + int64_t n_img = t_len * h_len * w_len; + + auto img = DiT::patchify_3d(ctx->ggml_ctx, x, config.patch_t, config.patch_h, config.patch_w, N, false); + img = patch_embedder->forward(ctx, img); + + auto txt = text_embedder->forward(ctx, context); + auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1); + + auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, + timestep, + static_cast(config.freq_dim), + 10000, + 1.0f); + auto t_emb = time_embedder->forward(ctx, timestep_proj); + auto temb_template = ggml_new_tensor_3d(ctx->ggml_ctx, t_emb->type, t_emb->ne[0], hidden_states->ne[1], t_emb->ne[1]); + auto temb_tokens = ggml_repeat(ctx->ggml_ctx, + ggml_reshape_3d(ctx->ggml_ctx, t_emb, t_emb->ne[0], 1, t_emb->ne[1]), + temb_template); + auto temb6 = time_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens)); + + sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.prelude", "hidden_states"); + + for (int i = 0; i < config.depth; i++) { + auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); + hidden_states = block->forward(ctx, hidden_states, temb6, pe); + sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.blocks." + std::to_string(i), "hidden_states"); + } + + auto final_mods = ggml_ext_chunk(ctx->ggml_ctx, + norm_out_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens)), + 2, + 0); + hidden_states = norm_out->forward(ctx, hidden_states); + hidden_states = Flux::modulate(ctx->ggml_ctx, hidden_states, final_mods[0], final_mods[1], true); + hidden_states = proj_out->forward(ctx, hidden_states); + hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img); + + auto out = DiT::unpatchify_3d(ctx->ggml_ctx, + hidden_states, + t_len, + h_len, + w_len, + config.patch_t, + config.patch_h, + config.patch_w, + false); + return out; + } + }; + + struct LingBotVideoRunner : public DiffusionModelRunner { + LingBotVideoConfig config; + LingBotVideoModel lingbot_video; + std::vector pe_vec; + + LingBotVideoRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr, + const char* model_args = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(LingBotVideoConfig::detect_from_weights(tensor_storage_map, prefix)) { + SD_UNUSED(model_args); + + lingbot_video = LingBotVideoModel(config); + lingbot_video.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "lingbot_video"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + lingbot_video.get_param_tensors(tensors, prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor) { + ggml_cgraph* gf = new_graph_custom(LINGBOT_VIDEO_GRAPH_SIZE); + ggml_tensor* x = make_input(x_tensor); + ggml_tensor* timesteps = make_input(timesteps_tensor); + GGML_ASSERT(x_tensor.dim() == 5); + GGML_ASSERT(x->ne[3] == config.in_channels); + GGML_ASSERT(!context_tensor.empty()); + ggml_tensor* context = make_input(context_tensor); + + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t T = x->ne[2]; + int64_t N = x_tensor.shape()[4]; + GGML_ASSERT(N == 1); + pe_vec = Rope::gen_lingbot_video_pe(static_cast(T), + static_cast(H), + static_cast(W), + config.patch_t, + config.patch_h, + config.patch_w, + static_cast(N), + static_cast(context->ne[1]), + config.rope_theta, + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); + set_backend_tensor_data(pe, pe_vec.data()); + + auto runner_ctx = get_context(); + ggml_tensor* out = lingbot_video.forward(&runner_ctx, x, timesteps, context, pe); + ggml_build_forward_expand(gf, out); + return gf; + } + + sd::Tensor compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context) { + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context); + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + return compute(n_threads, + *diffusion_params.x, + *diffusion_params.timesteps, + tensor_or_empty(diffusion_params.context)); + } + }; +} // namespace LingBotVideo + +#endif // __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__ diff --git a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp index 52edc0d56..a7c946fd9 100644 --- a/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/qwen_image.hpp @@ -183,7 +183,7 @@ namespace Qwen { auto to_v = std::dynamic_pointer_cast(blocks["to_v"]); auto to_out_0 = std::dynamic_pointer_cast(blocks["to_out.0"]); - if (sd_backend_is(ctx->backend, "Vulkan")) { + if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { to_out_0->set_force_prec_f32(true); } @@ -517,7 +517,7 @@ namespace Qwen { if (input->ne[3] == 1) { input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]); } - return DiT::patchify(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N); + return DiT::patchify_3d(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N); }; auto img = patchify_input(x); diff --git a/otherarch/sdcpp/src/model/diffusion/z_image.hpp b/otherarch/sdcpp/src/model/diffusion/z_image.hpp index 362192f19..9fb333cbb 100644 --- a/otherarch/sdcpp/src/model/diffusion/z_image.hpp +++ b/otherarch/sdcpp/src/model/diffusion/z_image.hpp @@ -150,6 +150,8 @@ namespace ZImage { if (sd_backend_is(ctx->backend, "ROCm")) { out_proj->set_scale(1.f / 16.f); + out_proj->set_force_prec_f32(true); + qkv_proj->set_force_prec_f32(true); } auto qkv = qkv_proj->forward(ctx, x); // [N, n_token, (num_heads + num_kv_heads*2)*head_dim] @@ -227,7 +229,7 @@ namespace ZImage { auto w2 = std::dynamic_pointer_cast(blocks["w2"]); auto w3 = std::dynamic_pointer_cast(blocks["w3"]); - if (sd_backend_is(ctx->backend, "Vulkan")) { + if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { w2->set_force_prec_f32(true); } diff --git a/otherarch/sdcpp/src/model_io/safetensors_io.cpp b/otherarch/sdcpp/src/model_io/safetensors_io.cpp index 9c3c2869b..e38a38378 100644 --- a/otherarch/sdcpp/src/model_io/safetensors_io.cpp +++ b/otherarch/sdcpp/src/model_io/safetensors_io.cpp @@ -3,9 +3,11 @@ #include #include #include +#include #include #include #include +#include #include #include "binary_io.h" @@ -13,6 +15,8 @@ #include "json.hpp" #include "model_io/kcpp_sdcpp_quantized_safetensors.hpp" +namespace fs = std::filesystem; + static constexpr size_t ST_HEADER_SIZE_LEN = 8; static void set_error(std::string* error, const std::string& message) { @@ -21,6 +25,14 @@ static void set_error(std::string* error, const std::string& message) { } } +static std::string resolve_index_shard_path(const std::string& index_path, const std::string& shard_path) { + fs::path shard_fs_path(shard_path); + if (shard_fs_path.is_absolute()) { + return shard_fs_path.lexically_normal().string(); + } + return (fs::path(index_path).parent_path() / shard_fs_path).lexically_normal().string(); +} + bool is_safetensors_file(const std::string& file_path) { std::ifstream file(file_path, std::ios::binary); if (!file.is_open()) { @@ -251,6 +263,52 @@ bool read_safetensors_file(const std::string& file_path, return true; } +bool read_safetensors_index_file(const std::string& file_path, + std::vector& shard_paths, + std::string* error) { + shard_paths.clear(); + + std::ifstream file(file_path); + if (!file.is_open()) { + set_error(error, "failed to open '" + file_path + "'"); + return false; + } + + nlohmann::json index; + try { + index = nlohmann::json::parse(file); + } catch (const std::exception&) { + set_error(error, "parsing safetensors index failed: '" + file_path + "'"); + return false; + } + + if (!index.is_object() || !index.contains("weight_map") || !index["weight_map"].is_object()) { + set_error(error, "invalid safetensors index '" + file_path + "'"); + return false; + } + + std::unordered_set seen_shard_paths; + for (const auto& item : index["weight_map"].items()) { + if (!item.value().is_string()) { + set_error(error, "invalid shard path for tensor '" + item.key() + "'"); + return false; + } + + std::string shard_path = resolve_index_shard_path(file_path, + item.value().get()); + if (seen_shard_paths.insert(shard_path).second) { + shard_paths.push_back(std::move(shard_path)); + } + } + + if (shard_paths.empty()) { + set_error(error, "safetensors index has no tensors: '" + file_path + "'"); + return false; + } + + return true; +} + static bool ggml_type_to_safetensors_dtype(ggml_type type, std::string* dtype) { switch (type) { case GGML_TYPE_F16: diff --git a/otherarch/sdcpp/src/model_io/safetensors_io.h b/otherarch/sdcpp/src/model_io/safetensors_io.h index b4938ee18..b18f0ae75 100644 --- a/otherarch/sdcpp/src/model_io/safetensors_io.h +++ b/otherarch/sdcpp/src/model_io/safetensors_io.h @@ -11,6 +11,9 @@ 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); +bool read_safetensors_index_file(const std::string& file_path, + std::vector& shard_paths, + std::string* error = nullptr); bool write_safetensors_file(const std::string& file_path, const std::vector& tensors, std::string* error = nullptr); diff --git a/otherarch/sdcpp/src/model_loader.cpp b/otherarch/sdcpp/src/model_loader.cpp index ac490d21c..ce4b33f62 100644 --- a/otherarch/sdcpp/src/model_loader.cpp +++ b/otherarch/sdcpp/src/model_loader.cpp @@ -251,6 +251,9 @@ bool ModelLoader::init_from_file(const std::string& file_path, const std::string } else if (is_gguf_file(file_path)) { LOG_INFO("load %s using gguf format", file_path.c_str()); return init_from_gguf_file(file_path, prefix); + } else if (ends_with(file_path, ".json")) { + LOG_INFO("load %s using safetensors index format", file_path.c_str()); + return init_from_safetensors_index_file(file_path, prefix); } else if (is_safetensors_file(file_path)) { LOG_INFO("load %s using safetensors format", file_path.c_str()); return init_from_safetensors_file(file_path, prefix); @@ -358,6 +361,25 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const return true; } +bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, const std::string& prefix) { + LOG_DEBUG("init from safetensors index '%s', prefix = '%s'", file_path.c_str(), prefix.c_str()); + + std::vector shard_paths; + std::string error; + if (!read_safetensors_index_file(file_path, shard_paths, &error)) { + LOG_ERROR("%s", error.c_str()); + return false; + } + + for (const std::string& shard_path : shard_paths) { + if (!init_from_file(shard_path, prefix)) { + return false; + } + } + + return true; +} + /*================================================= TorchLegacyModelLoader ==================================================*/ bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix) { @@ -540,6 +562,9 @@ SDVersion ModelLoader::get_sd_version() { if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) { is_wan = true; } + if (tensor_storage.name.find("model.diffusion_model.patch_embedder.weight") != std::string::npos) { + return VERSION_LINGBOT_VIDEO; + } if (tensor_storage.name.find("model.diffusion_model.patch_embedding.weight") != std::string::npos) { patch_embedding_channels = tensor_storage.ne[3]; } diff --git a/otherarch/sdcpp/src/model_loader.h b/otherarch/sdcpp/src/model_loader.h index 1973aaaec..f66f27b0f 100644 --- a/otherarch/sdcpp/src/model_loader.h +++ b/otherarch/sdcpp/src/model_loader.h @@ -43,6 +43,7 @@ protected: bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = ""); bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = ""); + bool init_from_safetensors_index_file(const std::string& file_path, const std::string& prefix = ""); bool init_from_torch_zip_file(const std::string& file_path, const std::string& prefix = ""); bool init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix = ""); bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = ""); diff --git a/otherarch/sdcpp/src/model_manager.cpp b/otherarch/sdcpp/src/model_manager.cpp index 3a98bd545..33bc18373 100644 --- a/otherarch/sdcpp/src/model_manager.cpp +++ b/otherarch/sdcpp/src/model_manager.cpp @@ -179,6 +179,102 @@ bool ModelManager::register_param_tensors(const std::string& desc, return true; } +bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* registered_tensor_size) { + if (desc.empty()) { + return true; + } + + std::unordered_set target_states; + size_t released_size = 0; + for (auto& state : tensor_states_) { + if (state == nullptr || state->desc != desc) { + continue; + } + if (state->active_prepare_count > 0) { + LOG_ERROR("model manager cannot unregister active %s tensor '%s'", + desc.c_str(), + state->name.c_str()); + return false; + } + target_states.insert(state.get()); + if (state->tensor != nullptr) { + released_size += ggml_nbytes(state->tensor); + } + } + + if (target_states.empty()) { + return true; + } + + release_compute_staging_blocks(false); + + std::vector storage_blocks_to_release; + std::unordered_set affected_storage_states; + for (const auto& block : params_storage_blocks_) { + if (block == nullptr) { + continue; + } + bool has_target_state = false; + for (TensorState* state : block->states) { + if (state != nullptr && target_states.count(state) > 0) { + has_target_state = true; + break; + } + } + if (!has_target_state) { + continue; + } + storage_blocks_to_release.push_back(block.get()); + for (TensorState* state : block->states) { + if (state != nullptr) { + affected_storage_states.insert(state); + } + } + } + + for (TensorState* state : affected_storage_states) { + if (state == nullptr) { + continue; + } + if (state->active_prepare_count > 0 || state->staged_to_compute_backend) { + LOG_ERROR("model manager cannot unregister %s while tensor '%s' is active", + desc.c_str(), + state->name.c_str()); + return false; + } + } + + for (ParamsStorageBlock* block : storage_blocks_to_release) { + if (block != nullptr) { + free_params_storage_block(*block); + erase_params_storage_block(block); + } + } + + for (auto it = tensor_states_by_name_.begin(); it != tensor_states_by_name_.end();) { + if (target_states.count(it->second) > 0) { + it = tensor_states_by_name_.erase(it); + } else { + ++it; + } + } + tensor_states_.erase(std::remove_if(tensor_states_.begin(), + tensor_states_.end(), + [&](const std::unique_ptr& s) { + return s == nullptr || target_states.count(s.get()) > 0; + }), + tensor_states_.end()); + + if (registered_tensor_size != nullptr) { + if (released_size > *registered_tensor_size) { + *registered_tensor_size = 0; + } else { + *registered_tensor_size -= released_size; + } + } + return true; +} + bool ModelManager::load_all_params_eagerly() { std::vector all_states; all_states.reserve(tensor_states_.size()); diff --git a/otherarch/sdcpp/src/model_manager.h b/otherarch/sdcpp/src/model_manager.h index d80032614..85e982de2 100644 --- a/otherarch/sdcpp/src/model_manager.h +++ b/otherarch/sdcpp/src/model_manager.h @@ -134,6 +134,9 @@ public: bool allow_split_buffer = false, bool params_follow_compute_backend = false); + bool unregister_param_tensors(const std::string& desc, + size_t* registered_tensor_size = nullptr); + template bool register_runner_params(const std::string& desc, Runner& runner, diff --git a/otherarch/sdcpp/src/stable-diffusion.cpp b/otherarch/sdcpp/src/stable-diffusion.cpp index 5e9c7ef66..e5b331ba7 100644 --- a/otherarch/sdcpp/src/stable-diffusion.cpp +++ b/otherarch/sdcpp/src/stable-diffusion.cpp @@ -34,6 +34,7 @@ #include "model/diffusion/ideogram4.hpp" #include "model/diffusion/krea2.hpp" #include "model/diffusion/lens.hpp" +#include "model/diffusion/lingbot_video.hpp" #include "model/diffusion/ltxv.hpp" #include "model/diffusion/minit2i.hpp" #include "model/diffusion/mmdit.hpp" @@ -89,6 +90,7 @@ const char* model_version_to_str[] = { "Wan 2.x", "Wan 2.2 I2V", "Wan 2.2 TI2V", + "LingBot Video", "Qwen Image", "Qwen Image Layered", "Anima", @@ -226,9 +228,13 @@ public: std::string split_mode_spec; bool auto_fit_enabled = false; + bool diffusion_conv_direct = false; + bool is_using_v_parameterization = false; bool is_using_edm_v_parameterization = false; + size_t control_net_params_mem_size = 0; + std::shared_ptr model_manager; std::shared_ptr denoiser = std::make_shared(); @@ -498,6 +504,76 @@ public: params_follow_runtime); } + bool unload_control_net() { + if (control_net == nullptr) { + return true; + } + if (model_manager != nullptr) { + if (!model_manager->unregister_param_tensors("ControlNet", &control_net_params_mem_size)) { + return false; + } + } + control_net.reset(); + control_net_params_mem_size = 0; + return true; + } + + bool load_control_net_from_file(const std::string& path) { + if (path.empty()) { + LOG_ERROR("sd_ctx_load_control_net: empty path"); + return false; + } + if (model_manager == nullptr) { + LOG_ERROR("sd_ctx_load_control_net: model_manager not initialized"); + return false; + } + + if (!unload_control_net()) { + return false; + } + + ModelLoader& shared_loader = model_manager->loader(); + if (!shared_loader.init_from_file(path)) { + LOG_ERROR("sd_ctx_load_control_net: failed to load '%s'", path.c_str()); + return false; + } + shared_loader.convert_tensors_name(); + + if (!ensure_backend_pair(SDBackendModule::CONTROL_NET)) { + LOG_ERROR("sd_ctx_load_control_net: control_net backend unavailable"); + return false; + } + + control_net = std::make_shared(backend_for(SDBackendModule::CONTROL_NET), + params_backend_for(SDBackendModule::CONTROL_NET), + shared_loader.get_tensor_storage_map(), + version, + "", + model_manager); + if (diffusion_conv_direct) { + LOG_INFO("Using Conv2d direct in the control net"); + control_net->set_conv2d_direct_enabled(true); + } + if (!register_runner_params("ControlNet", + control_net, + SDBackendModule::CONTROL_NET, + &control_net_params_mem_size)) { + LOG_ERROR("sd_ctx_load_control_net: register_runner_params failed"); + control_net.reset(); + control_net_params_mem_size = 0; + return false; + } + if (!model_manager->validate_registered_tensors()) { + LOG_ERROR("sd_ctx_load_control_net: registered tensors validation failed"); + unload_control_net(); + return false; + } + LOG_INFO("sd_ctx_load_control_net: loaded '%s' (%.2f MB)", + path.c_str(), + control_net_params_mem_size / 1024.0 / 1024.0); + return true; + } + bool init_backend() { std::string error; if (!backend_manager.init(backend_spec.c_str(), @@ -1102,10 +1178,12 @@ public: model_loader.process_model_files(enable_mmap, needs_writable_mmap); load_alphas_cumprod(model_loader); + diffusion_conv_direct = sd_ctx_params->diffusion_conv_direct; + size_t text_encoder_params_mem_size = 0; size_t unet_params_mem_size = 0; size_t vae_params_mem_size = 0; - size_t control_net_params_mem_size = 0; + control_net_params_mem_size = 0; size_t extension_params_mem_size = 0; bool tae_preview_only = sd_ctx_params->tae_preview_only; @@ -1155,7 +1233,7 @@ public: tensor_storage_map, version, "", - false, + true, model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, @@ -1254,6 +1332,25 @@ public: return false; } } + } else if (sd_version_is_lingbot_video(version)) { + bool enable_vision = false; + for (const auto& [name, _] : tensor_storage_map) { + if (starts_with(name, "text_encoders.llm.visual.")) { + enable_vision = true; + break; + } + } + cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), + tensor_storage_map, + version, + "", + enable_vision, + model_manager); + diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + model_manager, + sd_ctx_params->model_args); } else if (sd_version_is_qwen_image(version)) { bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED; cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), @@ -1753,6 +1850,7 @@ public: } } else if (sd_version_is_sd3(version) || sd_version_is_wan(version) || + sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || version == VERSION_HIDREAM_O1 || sd_version_is_anima(version) || @@ -2601,7 +2699,7 @@ public: sd::Tensor timesteps_tensor({static_cast(timesteps_vec.size())}, timesteps_vec); sd::Tensor guidance_tensor({1}, std::vector{guidance.distilled_guidance}); sd::Tensor noised_input = x * c_in; - if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version))) { + if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) { noised_input = noised_input * denoise_mask + init_latent * (1.0f - denoise_mask); } @@ -2847,7 +2945,7 @@ public: int get_diffusion_model_down_factor() { int down_factor = 8; // unet if (sd_version_is_dit(version)) { - if (sd_version_is_wan(version)) { + if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) { down_factor = 2; } else { down_factor = 1; @@ -2910,7 +3008,7 @@ public: int latent_frames = frames; if (sd_version_is_ltxav(version)) { latent_frames = ((frames - 1) / 8) + 1; - } else if (sd_version_is_wan(version)) { + } else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) { latent_frames = ((frames - 1) / 4) + 1; } return latent_frames; @@ -2923,7 +3021,7 @@ public: if (sd_version_is_ltxav(version)) { return (latent_frames - 1) * 8 + 1; } - if (sd_version_is_wan(version)) { + if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) { return (latent_frames - 1) * 4 + 1; } return latent_frames; @@ -3599,7 +3697,7 @@ struct sd_ctx_t { }; static bool sd_version_supports_video_generation(SDVersion version) { - return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_ltxav(version); + return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version); } static bool sd_version_supports_image_generation(SDVersion version) { @@ -3771,6 +3869,27 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) { return sd_version_supports_video_generation(sd_ctx->sd->version); } +SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr || path == nullptr) { + return false; + } + return sd_ctx->sd->load_control_net_from_file(path); +} + +SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { + return false; + } + return sd_ctx->sd->unload_control_net(); +} + +SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { + return false; + } + return sd_ctx->sd->control_net != nullptr; +} + enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) { if (sd_ctx != nullptr && sd_ctx->sd != nullptr) { if (sd_version_is_pid(sd_ctx->sd->version)) { @@ -5800,6 +5919,27 @@ static std::optional prepare_video_generation_latents(sd sd::ops::fill_slice(&latents.denoise_mask, 2, latents.init_latent.shape()[2] - 1, latents.init_latent.shape()[2], 0.0f); } + int64_t t2 = ggml_time_ms(); + LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); + } else if (sd_version_is_lingbot_video(sd_ctx->sd->version) && !start_image.empty()) { + LOG_INFO("LingBot Video IMG2VID"); + + int64_t t1 = ggml_time_ms(); + auto init_img = start_image.reshape({start_image.shape()[0], start_image.shape()[1], 1, start_image.shape()[2], 1}); + auto init_image_latent = sd_ctx->sd->encode_first_stage(init_img); + if (init_image_latent.empty()) { + LOG_ERROR("failed to encode init video frame"); + return std::nullopt; + } + + latents.init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->frames, true); + sd::ops::slice_assign(&latents.init_latent, 2, 0, init_image_latent.shape()[2], init_image_latent); + + latents.denoise_mask = sd::full({latents.init_latent.shape()[0], latents.init_latent.shape()[1], latents.init_latent.shape()[2], 1, 1}, 1.f); + sd::ops::fill_slice(&latents.denoise_mask, 2, 0, init_image_latent.shape()[2], 0.0f); + + latents.ref_images.push_back(start_image); + int64_t t2 = ggml_time_ms(); LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); } else if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-VACE-1.3B" || @@ -5912,6 +6052,7 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx, condition_params.clip_skip = request.clip_skip; condition_params.text = request.prompt; condition_params.zero_out_masked = true; + condition_params.ref_images = &latents.ref_images; int64_t prepare_start_ms = ggml_time_ms(); embeds.cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,