diff --git a/README.md b/README.md index d99348cbf..a63dad9a8 100644 --- a/README.md +++ b/README.md @@ -61,12 +61,14 @@ API and command-line option may change frequently.*** - [SeFi-Image](./docs/sefi_image.md) - [HiDream-O1-Image](./docs/hidream_o1_image.md) - [Ideogram4](./docs/ideogram4.md) + - [LLaDA-Image](./docs/llada_image.md) - [Image Edit Models](./docs/edit.md) - [FLUX.1-Kontext-dev](./docs/kontext.md) - [Qwen Image Edit series](./docs/qwen_image_edit.md) - [LongCat Image Edit](./docs/longcat_image.md) - [Boogu Image Edit](./docs/boogu_image.md) - [Mage-Flow-Edit](./docs/mage_flow.md#image-editing) + - [LLaDA-Image Edit](./docs/llada_image.md#image-editing) - Video Models - [Wan2.1/Wan2.2](./docs/wan.md) - [MiniMax-H3](./docs/minimax_h3.md) diff --git a/assets/llada_image/edit_example.png b/assets/llada_image/edit_example.png new file mode 100644 index 000000000..aaff84f41 Binary files /dev/null and b/assets/llada_image/edit_example.png differ diff --git a/assets/llada_image/example.png b/assets/llada_image/example.png new file mode 100644 index 000000000..63cc05f50 Binary files /dev/null and b/assets/llada_image/example.png differ diff --git a/docs/edit.md b/docs/edit.md index 9791d046d..9f8a9c3f3 100644 --- a/docs/edit.md +++ b/docs/edit.md @@ -17,6 +17,7 @@ Depending on the architecture, different models handle reference images differen | [**Boogu Image Edit**](./boogu_image.md) | `z_image_omni` | | **Krea2 (Community Edit LoRAs)** | `krea2_ostris_edit` | | [**Mage-Flow-Edit**](./mage_flow.md#image-editing) | `mage_flow` | +| [**LLaDA-Image**](./llada_image.md#image-editing) | `llada_image` | | **Anima (Community Edit LoRAs)** | `cosmos_reference` | Stable-diffusion.spp also supports basic Unet-based editing models like instruct-pix2pix or CosXL-Edit. This document is not about those. diff --git a/docs/llada_image.md b/docs/llada_image.md new file mode 100644 index 000000000..af865e503 --- /dev/null +++ b/docs/llada_image.md @@ -0,0 +1,153 @@ +# How to Use + +LLaDA-Image is a 6B text-to-image and instruction-guided editing model. The denoiser is a +Lumina2/Z-Image-style NextDiT conditioned by a LLaDA2-MoE diffusion-LLM text encoder, and it +reuses the Flux.2 VAE. Two checkpoints are published: a 50-step base model and +LLaDA-Image-Turbo, a 4-step distilled model. + +## Download weights + +Four components are required: a transformer, a text encoder, a VAE, and a connectors file +holding the QueryFormer, the text projection and, for editing, the SigVQ image encoder. + +The two published checkpoints are **not** interchangeable. LLaDA-Image-Turbo and LLaDA-Image +ship different transformers, text encoders, QueryFormers and text projections; only the VAE, +the SigVQ encoder and the tokenizer are shared. Mixing the two produces degraded output rather +than a clean error, so keep each checkpoint's files together. + +Both need an external LLaDA2 `tokenizer.json`, which is not embedded in sd.cpp and is the same +file for either checkpoint. Take `tokenizer/tokenizer.json` from either repository and pass it +with `--tokenizer`. See [JSON tokenizers](tokenizers.md) for CLI and C API usage. + +### LLaDA-Image-Turbo (4 steps) + +Converted transformer, text encoder and pre-merged connectors are at +https://huggingface.co/fszontagh/LLaDA-Image-Turbo-GGUF: + +- `llada-image-turbo-f16.gguf` +- `llada-image-turbo-text_encoder-q8_0.gguf` +- `llada-image-turbo-connectors.safetensors` for text to image, or + `llada-image-turbo-connectors-edit.safetensors`, which also carries the SigVQ encoder that + editing needs. + +Other quantizations of the transformer and the text encoder are in the same repository. + +The VAE comes from the original repository, +https://huggingface.co/inclusionAI/LLaDA-Image-Turbo: `vae/diffusion_pytorch_model.safetensors`, +referred to below as `llada_vae.safetensors`. + +### LLaDA-Image (50 steps) + +Converted transformer, text encoder and pre-merged connectors are at +https://huggingface.co/fszontagh/LLaDA-Image-GGUF: + +- `llada-image-f16.gguf` +- `llada-image-text_encoder-q8_0.gguf` +- `llada-image-connectors.safetensors` for text to image, or + `llada-image-connectors-edit.safetensors`, which also carries the SigVQ encoder that editing + needs. + +Other quantizations of the transformer and the text encoder are in the same repository. + +The VAE comes from the original repository, +https://huggingface.co/inclusionAI/LLaDA-Image, and is the same file as the Turbo one. + +### Converting the weights yourself + +The transformer has to go in through `--diffusion-model` so that its tensor names keep the +prefix the loader expects, while the text encoder goes in through `-m`: + +```bash +./bin/sd-cli -M convert --diffusion-model transformer/diffusion_pytorch_model.safetensors.index.json \ + -o llada-image-f16.gguf --type f16 +./bin/sd-cli -M convert -m text_encoder/model.safetensors.index.json \ + -o llada-image-text_encoder-q8_0.gguf --type q8_0 +``` + +### Building the connector file yourself + +`--embeddings-connectors` takes one file, so the QueryFormer, the text projection and +(for editing) the SigVQ encoder have to be combined into a single Safetensors file, each +tensor name prefixed with its component name. Leaving `sigvq` out skips loading the 2.6 GB +encoder: + +```python +from safetensors.torch import load_file, save_file + +merged = {} +for prefix, path in [ + ("queryformer", "queryformer/diffusion_pytorch_model.safetensors"), + ("text_projection", "text_projection/diffusion_pytorch_model.safetensors"), + ("sigvq", "sigvq/diffusion_pytorch_model.safetensors"), +]: + for name, tensor in load_file(path).items(): + merged[f"{prefix}.{name}"] = tensor +save_file(merged, "llada_connectors.safetensors") +``` + +## Examples + +### Text to image + +```bash +./bin/sd-cli \ + --diffusion-model /path/to/llada-image-turbo-f16.gguf \ + --llm /path/to/llada-image-turbo-text_encoder-q8_0.gguf \ + --tokenizer /path/to/tokenizer.json \ + --vae /path/to/llada_vae.safetensors \ + --embeddings-connectors /path/to/llada-image-turbo-connectors.safetensors \ + --prompt "a lovely cat holding a sign says 'llada.cpp'" \ + --width 1024 \ + --height 1024 \ + --steps 4 \ + --cfg-scale 1.0 \ + --seed 42 \ + --output output.png +``` + +LLaDA-Image example + +### Image editing + +```bash +./bin/sd-cli \ + --diffusion-model /path/to/llada-image-turbo-f16.gguf \ + --llm /path/to/llada-image-turbo-text_encoder-q8_0.gguf \ + --tokenizer /path/to/tokenizer.json \ + --vae /path/to/llada_vae.safetensors \ + --embeddings-connectors /path/to/llada-image-turbo-connectors-edit.safetensors \ + --ref-image /path/to/input.png \ + --prompt "change the sign text to 'sd.cpp'" \ + --width 1024 \ + --height 1024 \ + --steps 4 \ + --cfg-scale 1.0 \ + --diffusion-fa \ + --output output.png +``` + +LLaDA-Image edit example + +See [edit.md](./edit.md) for the shared reference-image options. LLaDA-Image uses the +`llada_image` preset by default. + +## Notes + +- Use 4 steps and `--cfg-scale 1.0` for LLaDA-Image-Turbo; the guidance is distilled away, so + a higher CFG degrades output and doubles the text encoder cost. The 50-step base model uses + `--steps 50 --cfg-scale 5`. +- Width and height are rounded up to a multiple of 16. For editing the reference pipeline + requires them to be divisible by 32. +- Edit the 50-step base model at 1024x1024. At 512x512 it returns the reference image almost + unchanged instead of applying the instruction; LLaDA-Image-Turbo edits correctly at both. +- Editing runs the reference and the target in one sequence, so it needs roughly twice the + tokens of text to image at the same size. On 12 GB, editing at 1024x1024 needs + `--diffusion-fa`; without it the diffusion graph does not fit. +- The weights total about 16 GB, but segmented execution streams them, so a much smaller + budget works. At 512x512, `--max-vram 6` costs almost nothing over unconstrained execution, + and `--max-vram 3` still produces byte-identical output at roughly 2.5x the time. +- `--scheduler` defaults to `llada_image`, which reproduces the reference Kumaraswamy sigma + grid. `--extra-sample-args uniform=1` selects the uniform grid instead. +- Prompt templating is handled automatically; pass a plain description. +- VQ-conditioned generation (`generation_mode="vq"`, where the text encoder decodes image + tokens before diffusion) is not implemented. diff --git a/examples/common/common.cpp b/examples/common/common.cpp index e8ec40295..10884b4a5 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -1109,7 +1109,7 @@ ArgOptions SDGenerationParams::get_options() { &hires_upscaler}, {"", "--extra-sample-args", - "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng", + "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; llada_image supports uniform; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng", (int)',', &extra_sample_args}, {"", diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 9bbf8c757..28f33c1de 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -79,6 +79,7 @@ enum scheduler_t { FLUX2_SCHEDULER, FLUX_SCHEDULER, BETA_SCHEDULER, + LLADA_IMAGE_SCHEDULER, SCHEDULER_COUNT }; diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index 0566f93f9..a3bee137c 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -14,6 +14,7 @@ #include "core/util.h" #include "model/diffusion/model.hpp" #include "model/te/clip.hpp" +#include "model/te/llada_image_te.h" #include "model/te/llm.hpp" #include "model/te/t5.hpp" #include "model_loader.h" @@ -3159,6 +3160,203 @@ struct LTXAVTextProjectionRunner : public GGMLRunner { } }; +// LLaDA-Image's text path is a three-stage pipeline rather than a single encoder pass: +// the token embeddings feed a QueryFormer whose 256 queries are appended to the backbone +// input, and the backbone's final hidden states are projected to the denoiser's caption dim. +// Ref: LLaDAImagePipeline._encode_text. +struct LLaDAImageEmbedder : public Conditioner { + std::shared_ptr tokenizer; + std::shared_ptr llm; + std::shared_ptr query_former; + std::shared_ptr text_projection; + std::shared_ptr sigvq; + + std::string llm_prefix; + std::string query_former_prefix; + std::string text_projection_prefix; + std::string sigvq_prefix; + + LLaDAImageEmbedder(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string& llm_prefix = "text_encoders.llm", + const std::string& query_former_prefix = "queryformer", + const std::string& text_projection_prefix = "text_projection", + const std::string& sigvq_prefix = "sigvq", + std::shared_ptr weight_manager = nullptr, + const TokenizerConfig& tokenizers = {}) + : llm_prefix(llm_prefix), + query_former_prefix(query_former_prefix), + text_projection_prefix(text_projection_prefix), + sigvq_prefix(sigvq_prefix) { + if (!tokenizers.has(TokenizerConfig::MAIN)) { + throw std::runtime_error("LLaDA-Image requires an external LLaDA2 tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer"); + } + llm = std::make_shared(LLM::LLMArch::LLADA2_MOE, + backend, + tensor_storage_map, + llm_prefix, + false, + weight_manager); + // <|endoftext|> doubles as the pad token in LLaDA2's tokenizer.json. + tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 156892); + query_former = std::make_shared(backend, + tensor_storage_map, + query_former_prefix, + weight_manager); + text_projection = std::make_shared(backend, + tensor_storage_map, + text_projection_prefix, + weight_manager); + + // SigVQ is only present when the user supplies the editing weights. + for (const auto& [name, _] : tensor_storage_map) { + if (starts_with(name, sigvq_prefix + ".")) { + sigvq = std::make_shared(backend, + tensor_storage_map, + sigvq_prefix, + weight_manager); + break; + } + } + } + + void get_param_tensors(std::map& tensors) override { + llm->get_param_tensors(tensors, llm_prefix); + query_former->get_param_tensors(tensors, query_former_prefix); + text_projection->get_param_tensors(tensors, text_projection_prefix); + if (sigvq != nullptr) { + sigvq->get_param_tensors(tensors, sigvq_prefix); + } + } + + void get_param_tensor_ops(std::map& tensor_ops) override { + llm->get_param_tensor_ops(tensor_ops); + } + + void set_flash_attention_enabled(bool enabled) override { + llm->set_flash_attention_enabled(enabled); + query_former->set_flash_attention_enabled(enabled); + text_projection->set_flash_attention_enabled(enabled); + if (sigvq != nullptr) { + sigvq->set_flash_attention_enabled(enabled); + } + } + + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + llm->set_max_graph_vram_bytes(max_vram_bytes); + query_former->set_max_graph_vram_bytes(max_vram_bytes); + text_projection->set_max_graph_vram_bytes(max_vram_bytes); + if (sigvq != nullptr) { + sigvq->set_max_graph_vram_bytes(max_vram_bytes); + } + } + + void set_runtime_backends(const std::vector& backends) override { + llm->set_runtime_backends(backends); + } + + void set_graph_cut_layer_split_enabled(bool enabled) override { + llm->set_graph_cut_layer_split_enabled(enabled); + } + + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) override { + llm->set_graph_cut_layer_split_backend_vram_limits(limits); + } + + void get_layer_split_param_tensors(std::map& tensors) override { + llm->get_param_tensors(tensors, llm_prefix); + } + + void set_weight_adapter(const std::shared_ptr& adapter) override { + llm->set_weight_adapter(adapter); + query_former->set_weight_adapter(adapter); + text_projection->set_weight_adapter(adapter); + if (sigvq != nullptr) { + sigvq->set_weight_adapter(adapter); + } + } + + void runner_end() override { + llm->runner_end(); + query_former->runner_end(); + text_projection->runner_end(); + if (sigvq != nullptr) { + sigvq->runner_end(); + } + } + + SDCondition get_learned_condition(int n_threads, + const ConditionerParams& conditioner_params) override { + const int64_t num_queries = 256; + + std::string text = conditioner_params.text; + while (!text.empty() && std::isspace(static_cast(text.front()))) { + text.erase(text.begin()); + } + while (!text.empty() && std::isspace(static_cast(text.back()))) { + text.pop_back(); + } + std::string prompt = text.empty() + ? "HUMAN Generate an image.\nASSISTANT\n" + : "HUMAN Generate an image: " + text + "\nASSISTANT\n"; + + std::vector tokens; + if (!tokenizer->encode(prompt, tokens, nullptr)) { + return {}; + } + int64_t n_text = static_cast(tokens.size()); + GGML_ASSERT(n_text > 0); + + sd::Tensor text_ids({n_text}, std::vector(tokens.begin(), tokens.end())); + auto inputs_embeds = llm->compute_input_embeds(n_threads, text_ids); + auto query_embeds = query_former->compute(n_threads, inputs_embeds); + + // splice_image_embeds() replaces tokens in place, so the query slots have to exist in + // input_ids; their ids are irrelevant because the embeddings are overwritten. + std::vector padded(tokens.begin(), tokens.end()); + padded.resize(static_cast(n_text + num_queries), tokenizer->PAD_TOKEN_ID); + int64_t n_total = static_cast(padded.size()); + sd::Tensor input_ids({n_total}, padded); + + // Bidirectional everywhere except that the text tokens must not see the appended + // queries, matching backbone_attention_mask[:, :, :text_length, text_length:] = min. + const float mask_min = std::numeric_limits::lowest() / 4.0f; + sd::Tensor attention_mask({n_total, n_total}); + for (int64_t i1 = 0; i1 < n_total; ++i1) { + for (int64_t i0 = 0; i0 < n_total; ++i0) { + float value = (i1 < n_text && i0 >= n_text) ? mask_min : 0.0f; + attention_mask[i0 + n_total * i1] = value; + } + } + + LLM::ImageEmbeds image_embeds; + image_embeds.emplace_back(static_cast(n_text), query_embeds); + + std::set out_layers = {static_cast(llm->config.num_layers) + 1}; + auto hidden_states = llm->compute(n_threads, + input_ids, + attention_mask, + image_embeds, + out_layers); + + SDCondition result; + result.c_crossattn = text_projection->compute(n_threads, hidden_states); + + // Editing: SigVQ sees the reference at half the output resolution, as in + // LLaDAImagePipeline._encode_source_image. + if (sigvq != nullptr && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { + const auto& ref = conditioner_params.ref_images->front(); + auto resized = sd::ops::interpolate(ref, + {conditioner_params.width / 2, + conditioner_params.height / 2, + ref.shape()[2], + ref.shape()[3]}); + result.extra_c_crossattns.push_back(sigvq->compute(n_threads, resized)); + } + return result; + } +}; + struct LTXAVEmbedder : public Conditioner { static constexpr int64_t kHiddenSize = 3840; static constexpr int64_t kNumStates = 49; diff --git a/src/model.h b/src/model.h index 7a8bc757f..b419b0764 100644 --- a/src/model.h +++ b/src/model.h @@ -59,6 +59,7 @@ enum SDVersion { VERSION_KREA2, VERSION_MAGE_FLOW, VERSION_SENSENOVA_U1_5, + VERSION_LLADA_IMAGE, VERSION_ESRGAN, VERSION_COUNT, }; @@ -172,6 +173,13 @@ static inline bool sd_version_is_z_image(SDVersion version) { return false; } +static inline bool sd_version_is_llada_image(SDVersion version) { + if (version == VERSION_LLADA_IMAGE) { + return true; + } + return false; +} + static inline bool sd_version_is_boogu_image(SDVersion version) { if (version == VERSION_BOOGU_IMAGE) { return true; @@ -251,7 +259,7 @@ static inline bool sd_version_uses_flux_vae(SDVersion version) { } static inline bool sd_version_uses_flux2_vae(SDVersion version) { - if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version)) { + if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version) || sd_version_is_llada_image(version)) { return true; } return false; @@ -292,6 +300,7 @@ static inline bool sd_version_is_dit(SDVersion version) { version == VERSION_HIDREAM_O1 || sd_version_is_anima(version) || sd_version_is_z_image(version) || + sd_version_is_llada_image(version) || sd_version_is_boogu_image(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 2aca72f4a..408c44a6a 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -835,21 +835,30 @@ class RMSNorm : public UnaryBlock { protected: int64_t hidden_size; float eps; + bool elementwise_affine; std::string prefix; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override { - this->prefix = prefix; + this->prefix = prefix; + if (!elementwise_affine) { + return; + } enum ggml_type wtype = GGML_TYPE_F32; params["weight"] = ggml_new_tensor_1d(ctx, wtype, hidden_size); } public: RMSNorm(int64_t hidden_size, - float eps = 1e-06f) + float eps = 1e-06f, + bool elementwise_affine = true) : hidden_size(hidden_size), - eps(eps) {} + eps(eps), + elementwise_affine(elementwise_affine) {} ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + if (!elementwise_affine) { + return ggml_rms_norm(ctx->ggml_ctx, x, eps); + } ggml_tensor* w = params["weight"]; if (ctx->weight_adapter) { w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index 778bbf426..e44784936 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -929,6 +929,145 @@ namespace Rope { return ids; } + // LLaDA-Image shares Lumina2/z_image's axes layout, but assigns position (0,0,0) to the + // padding slots of the caption stream instead of continuing the caption ramp through them. + __STATIC_INLINE__ std::vector> gen_llada_image_ids(int h, + int w, + int patch_size, + int bs, + int context_len, + int seq_multi_of) { + int context_pad_len = bound_mod(context_len, seq_multi_of); + int padded_context_len = context_len + context_pad_len; + auto txt_ids = std::vector>(bs * padded_context_len, std::vector(3, 0.0f)); + for (int i = 0; i < bs * padded_context_len; i++) { + int pos = i % padded_context_len; + if (pos < context_len) { + txt_ids[i][0] = pos + 1.f; + } + } + + int axes_dim_num = 3; + int index = padded_context_len + 1; + auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index); + + int img_pad_len = bound_mod(static_cast(img_ids.size() / bs), seq_multi_of); + if (img_pad_len > 0) { + std::vector> img_pad_ids(bs * img_pad_len, std::vector(3, 0.f)); + img_ids = concat_ids(img_ids, img_pad_ids, bs); + } + + return concat_ids(txt_ids, img_ids, bs); + } + + // LLaDA-Image editing packs two caption copies (clean and noisy), the source and target + // latents anchored at their own caption's end position, and the SigVQ stream after both. + // Padding slots keep position (0,0,0), as in the text-only layout. + __STATIC_INLINE__ std::vector> gen_llada_image_edit_ids(int h, + int w, + int patch_size, + int context_len, + int sigvq_len, + int seq_multi_of) { + const int context_pad = bound_mod(context_len, seq_multi_of); + const int padded_context = context_len + context_pad; + const int h_len = (h + (patch_size / 2)) / patch_size; + const int w_len = (w + (patch_size / 2)) / patch_size; + const int image_len = h_len * w_len; + const int image_pad = bound_mod(image_len, seq_multi_of); + const int padded_image = image_len + image_pad; + const int sigvq_pad = bound_mod(sigvq_len, seq_multi_of); + + std::vector> cap_ids; + std::vector cap_end_positions; + int cursor = 1; + for (int copy = 0; copy < 2; ++copy) { + for (int i = 0; i < padded_context; ++i) { + std::vector id(3, 0.f); + if (i < context_len) { + id[0] = static_cast(cursor + i); + } + cap_ids.push_back(id); + } + cursor += context_len; + cap_end_positions.push_back(cursor); + cursor += 2; + } + + std::vector> img_ids; + for (int copy = 0; copy < 2; ++copy) { + auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy]); + img_ids.insert(img_ids.end(), ids.begin(), ids.end()); + img_ids.insert(img_ids.end(), image_pad, std::vector(3, 0.f)); + } + + const int sigvq_start = static_cast(cap_ids.size() + img_ids.size()) + 1; + std::vector> sigvq_ids; + for (int i = 0; i < sigvq_len + sigvq_pad; ++i) { + std::vector id(3, 0.f); + if (i < sigvq_len) { + id[0] = static_cast(sigvq_start + i); + } + sigvq_ids.push_back(id); + } + + std::vector> ids; + ids.reserve(cap_ids.size() + img_ids.size() + sigvq_ids.size()); + ids.insert(ids.end(), cap_ids.begin(), cap_ids.end()); + ids.insert(ids.end(), img_ids.begin(), img_ids.end()); + ids.insert(ids.end(), sigvq_ids.begin(), sigvq_ids.end()); + SD_UNUSED(padded_image); + return ids; + } + + __STATIC_INLINE__ std::vector gen_llada_image_edit_pe(int h, + int w, + int patch_size, + int context_len, + int sigvq_len, + int seq_multi_of, + int theta, + const std::vector& axes_dim) { + auto ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of); + return embed_nd(ids, 1, static_cast(theta), axes_dim, {}); + } + + __STATIC_INLINE__ std::vector gen_llada_image_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + int seq_multi_of, + int theta, + bool circular_h, + bool circular_w, + const std::vector& axes_dim) { + std::vector> ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of); + std::vector> wrap_dims; + if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) { + int pad_h = (patch_size - (h % patch_size)) % patch_size; + int pad_w = (patch_size - (w % patch_size)) % patch_size; + int h_len = (h + pad_h) / patch_size; + int w_len = (w + pad_w) / patch_size; + if (h_len > 0 && w_len > 0) { + size_t pos_len = ids.size() / bs; + wrap_dims.assign(axes_dim.size(), std::vector(pos_len, 0)); + size_t cursor = context_len + bound_mod(context_len, seq_multi_of); + size_t img_tokens = static_cast(h_len) * static_cast(w_len); + for (size_t token_i = 0; token_i < img_tokens; ++token_i) { + if (circular_h) { + wrap_dims[1][cursor + token_i] = h_len; + } + if (circular_w) { + wrap_dims[2][cursor + token_i] = w_len; + } + } + } + } + + return embed_nd(ids, bs, static_cast(theta), axes_dim, wrap_dims); + } + // Generate z_image positional embeddings __STATIC_INLINE__ std::vector gen_z_image_pe(int h, int w, diff --git a/src/model/diffusion/llada_image.h b/src/model/diffusion/llada_image.h new file mode 100644 index 000000000..3521d02f3 --- /dev/null +++ b/src/model/diffusion/llada_image.h @@ -0,0 +1,519 @@ +#ifndef __SD_MODEL_DIFFUSION_LLADA_IMAGE_H__ +#define __SD_MODEL_DIFFUSION_LLADA_IMAGE_H__ + +#include +#include + +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" +#include "model/diffusion/model.hpp" +#include "model/diffusion/z_image.hpp" +#include "model_loader.h" + +// Ref: https://github.com/inclusionAI/LLaDA-Image/blob/main/src/models/transformer_llada_image.py +// +// The denoiser is Lumina2/z_image's NextDiT with identical hyperparameters, so the blocks are +// reused from ZImage. Two things differ: every norm here is non-parametric (the checkpoint +// carries no norm weights at all), and latents arrive already patchified from the Flux2 VAE, +// so patch_size is 1 over 128 channels. + +namespace LLaDAImage { + constexpr int LLADA_IMAGE_GRAPH_SIZE = 20480; + + struct LLaDAImageConfig { + int patch_size = 1; + int64_t hidden_size = 3840; + int64_t in_channels = 128; + int64_t out_channels = 128; + int64_t num_layers = 30; + int64_t num_refiner_layers = 2; + int64_t head_dim = 128; + int64_t num_heads = 30; + int64_t num_kv_heads = 30; + int64_t multiple_of = 256; + float ffn_dim_multiplier = 8.0f / 3.0f; + float norm_eps = 1e-5f; + bool qk_norm = true; + int64_t cap_feat_dim = 2560; + int64_t semantic_feat_dim = 4096; + int theta = 256; + std::vector axes_dim = {32, 48, 48}; + int64_t axes_dim_sum = 128; + + static int64_t count_blocks(const String2TensorStorage& tensor_storage_map, + const std::string& prefix, + const std::string& block_prefix) { + int64_t count = 0; + for (const auto& [name, _] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + size_t pos = name.find(block_prefix); + if (pos == std::string::npos) { + continue; + } + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + count = std::max(count, atoi(items[1].c_str()) + 1); + } + } + return count; + } + + static LLaDAImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + LLaDAImageConfig config; + int64_t detected_q_dim = 0; + int64_t detected_kv_dim = 0; + + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) { + int64_t patch_area = config.patch_size * config.patch_size; + config.in_channels = tensor_storage.ne[0] / patch_area; + config.hidden_size = tensor_storage.ne[1]; + } else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) { + config.cap_feat_dim = tensor_storage.ne[0]; + config.hidden_size = tensor_storage.ne[1]; + } else if (ends_with(name, "sigvq_embedder.1.weight") && tensor_storage.n_dims == 2) { + config.semantic_feat_dim = tensor_storage.ne[0]; + } else if (ends_with(name, "layers.0.attention.to_q.weight") && tensor_storage.n_dims == 2) { + detected_q_dim = tensor_storage.ne[1]; + } else if (ends_with(name, "layers.0.attention.to_k.weight") && tensor_storage.n_dims == 2) { + detected_kv_dim = tensor_storage.ne[1]; + } else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) { + int64_t patch_area = config.patch_size * config.patch_size; + config.out_channels = tensor_storage.ne[1] / patch_area; + } + } + + int64_t detected_layers = count_blocks(tensor_storage_map, prefix, "layers."); + int64_t detected_refiner = std::max(count_blocks(tensor_storage_map, prefix, "noise_refiner."), + count_blocks(tensor_storage_map, prefix, "context_refiner.")); + if (detected_layers > 0) { + config.num_layers = detected_layers; + } + if (detected_refiner > 0) { + config.num_refiner_layers = detected_refiner; + } + if (detected_q_dim > 0) { + config.num_heads = detected_q_dim / config.head_dim; + } + if (detected_kv_dim > 0) { + config.num_kv_heads = detected_kv_dim / config.head_dim; + } else if (detected_q_dim > 0) { + config.num_kv_heads = config.num_heads; + } + + LOG_VERBOSE("llada_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64 ", cap_feat_dim = %" PRId64 ", semantic_feat_dim = %" PRId64, + config.num_layers, + config.num_refiner_layers, + config.hidden_size, + config.num_heads, + config.num_kv_heads, + config.in_channels, + config.out_channels, + config.cap_feat_dim, + config.semantic_feat_dim); + return config; + } + }; + + class LLaDAImageModel : public GGMLBlock { + protected: + LLaDAImageConfig config; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); + params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); + params["sigvq_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); + } + + std::shared_ptr make_block(bool modulation) { + return std::make_shared(0, + config.hidden_size, + config.head_dim, + config.num_heads, + config.num_kv_heads, + config.multiple_of, + config.ffn_dim_multiplier, + config.norm_eps, + config.qk_norm, + modulation, + false, + true); + } + + public: + LLaDAImageModel() = default; + LLaDAImageModel(LLaDAImageConfig config) + : config(config) { + blocks["x_embedder"] = std::make_shared(config.patch_size * config.patch_size * config.in_channels, config.hidden_size); + blocks["t_embedder"] = std::make_shared(MIN(config.hidden_size, 1024), 256, ZImage::ADALN_EMBED_DIM); + + blocks["cap_embedder.0"] = std::make_shared(config.cap_feat_dim, config.norm_eps, false); + blocks["cap_embedder.1"] = std::make_shared(config.cap_feat_dim, config.hidden_size); + + blocks["semantic_embedder.0"] = std::make_shared(config.semantic_feat_dim, config.norm_eps, false); + blocks["semantic_embedder.1"] = std::make_shared(config.semantic_feat_dim, config.hidden_size); + blocks["sigvq_embedder.0"] = std::make_shared(config.semantic_feat_dim, config.norm_eps, false); + blocks["sigvq_embedder.1"] = std::make_shared(config.semantic_feat_dim, config.hidden_size); + + for (int i = 0; i < config.num_refiner_layers; i++) { + blocks["noise_refiner." + std::to_string(i)] = make_block(true); + blocks["context_refiner." + std::to_string(i)] = make_block(false); + blocks["sigvq_refiner." + std::to_string(i)] = make_block(false); + } + for (int i = 0; i < config.num_layers; i++) { + blocks["layers." + std::to_string(i)] = make_block(true); + } + + blocks["final_layer"] = std::make_shared(config.hidden_size, config.patch_size, config.out_channels); + } + + ggml_tensor* forward_core(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timestep, + ggml_tensor* context, + ggml_tensor* pe) { + auto x_embedder = std::dynamic_pointer_cast(blocks["x_embedder"]); + auto t_embedder = std::dynamic_pointer_cast(blocks["t_embedder"]); + auto cap_embedder_0 = std::dynamic_pointer_cast(blocks["cap_embedder.0"]); + auto cap_embedder_1 = std::dynamic_pointer_cast(blocks["cap_embedder.1"]); + auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); + + auto txt_pad_token = params["cap_pad_token"]; + auto img_pad_token = params["x_pad_token"]; + + int64_t N = x->ne[2]; + int64_t n_img_token = x->ne[1]; + int64_t n_txt_token = context->ne[1]; + + // sdcpp's flow denoiser already hands over sigma * 1000, which is the range the + // reference reaches via its own t_scale, so no further scaling here. + auto t_emb = t_embedder->forward(ctx, timestep); + + auto txt = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context)); // [N, n_txt_token, hidden_size] + auto img = x_embedder->forward(ctx, x); // [N, n_img_token, hidden_size] + sd::ggml_graph_cut::mark_graph_cut(txt, "llada_image.prelude", "txt"); + sd::ggml_graph_cut::mark_graph_cut(img, "llada_image.prelude", "img"); + sd::ggml_graph_cut::mark_graph_cut(t_emb, "llada_image.prelude", "t_emb"); + + int64_t n_txt_pad_token = Rope::bound_mod(static_cast(n_txt_token), ZImage::SEQ_MULTI_OF); + if (n_txt_pad_token > 0) { + auto txt_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, txt_pad_token, txt_pad_token->ne[0], n_txt_pad_token, N, 1); + txt = ggml_concat(ctx->ggml_ctx, txt, txt_pad_tokens, 1); + } + + int64_t n_img_pad_token = Rope::bound_mod(static_cast(n_img_token), ZImage::SEQ_MULTI_OF); + if (n_img_pad_token > 0) { + auto img_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, img_pad_token, img_pad_token->ne[0], n_img_pad_token, N, 1); + img = ggml_concat(ctx->ggml_ctx, img, img_pad_tokens, 1); + } + + GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]); + + auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]); + auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]); + + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)]); + + txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr); + sd::ggml_graph_cut::mark_graph_cut(txt, "llada_image.context_refiner." + std::to_string(i), "txt"); + } + + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)]); + + img = block->forward(ctx, img, img_pe, nullptr, t_emb); + sd::ggml_graph_cut::mark_graph_cut(img, "llada_image.noise_refiner." + std::to_string(i), "img"); + } + + auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); + sd::ggml_graph_cut::mark_graph_cut(txt_img, "llada_image.prelude", "txt_img"); + + for (int i = 0; i < config.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); + + txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb); + sd::ggml_graph_cut::mark_graph_cut(txt_img, "llada_image.layers." + std::to_string(i), "txt_img"); + } + + txt_img = final_layer->forward(ctx, txt_img, t_emb); + + return ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token); + } + + ggml_tensor* pad_stream(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pad_token) { + int64_t n_pad = Rope::bound_mod(static_cast(x->ne[1]), ZImage::SEQ_MULTI_OF); + if (n_pad == 0) { + return x; + } + auto pads = ggml_repeat_4d(ctx->ggml_ctx, pad_token, pad_token->ne[0], n_pad, x->ne[2], 1); + return ggml_concat(ctx->ggml_ctx, x, pads, 1); + } + + // Editing runs one joint sequence carrying two timesteps: the caption and source latent + // are clean (t = 0) while the second caption copy and the target latent are noisy. adaLN + // is a linear map of the timestep embedding, so feeding a per-token embedding selects the + // right modulation exactly, without duplicating the modulation projections. + ggml_tensor* forward_editing(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timestep, + ggml_tensor* context, + ggml_tensor* semantic, + ggml_tensor* source_latent, + ggml_tensor* pe) { + ggml_context* gctx = ctx->ggml_ctx; + + auto x_embedder = std::dynamic_pointer_cast(blocks["x_embedder"]); + auto t_embedder = std::dynamic_pointer_cast(blocks["t_embedder"]); + auto cap_embedder_0 = std::dynamic_pointer_cast(blocks["cap_embedder.0"]); + auto cap_embedder_1 = std::dynamic_pointer_cast(blocks["cap_embedder.1"]); + auto sigvq_embed_0 = std::dynamic_pointer_cast(blocks["sigvq_embedder.0"]); + auto sigvq_embed_1 = std::dynamic_pointer_cast(blocks["sigvq_embedder.1"]); + auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); + + auto t_noisy = t_embedder->forward(ctx, timestep); + auto t_clean = t_embedder->forward(ctx, ggml_scale(gctx, timestep, 0.f)); + + auto per_token = [&](ggml_tensor* emb, int64_t n) { + return ggml_repeat_4d(gctx, emb, emb->ne[0], n, 1, 1); + }; + + auto cap = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context)); + cap = pad_stream(ctx, cap, params["cap_pad_token"]); + int64_t cap_len = cap->ne[1]; + cap = ggml_concat(gctx, cap, cap, 1); + + auto src = pad_stream(ctx, x_embedder->forward(ctx, source_latent), params["x_pad_token"]); + auto tgt_embed = x_embedder->forward(ctx, x); + int64_t n_img_token = tgt_embed->ne[1]; + auto tgt = pad_stream(ctx, tgt_embed, params["x_pad_token"]); + int64_t img_len = tgt->ne[1]; + auto img = ggml_concat(gctx, src, tgt, 1); + + auto sig = sigvq_embed_1->forward(ctx, sigvq_embed_0->forward(ctx, semantic)); + sig = pad_stream(ctx, sig, params["sigvq_pad_token"]); + int64_t sig_len = sig->ne[1]; + + GGML_ASSERT(cap_len * 2 + img_len * 2 + sig_len == pe->ne[3]); + + auto cap_pe = ggml_ext_slice(gctx, pe, 3, 0, cap_len * 2); + auto img_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2); + auto sig_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]); + + auto img_adaln = ggml_concat(gctx, per_token(t_clean, img_len), per_token(t_noisy, img_len), 1); + + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)]); + cap = block->forward(ctx, cap, cap_pe, nullptr, nullptr); + } + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)]); + img = block->forward(ctx, img, img_pe, nullptr, img_adaln); + } + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["sigvq_refiner." + std::to_string(i)]); + sig = block->forward(ctx, sig, sig_pe, nullptr, nullptr); + } + + auto seq = ggml_concat(gctx, ggml_concat(gctx, cap, img, 1), sig, 1); + + auto cap_adaln = ggml_concat(gctx, per_token(t_clean, cap_len), per_token(t_noisy, cap_len), 1); + auto seq_adaln = ggml_concat(gctx, + ggml_concat(gctx, cap_adaln, img_adaln, 1), + per_token(t_clean, sig_len), + 1); + + for (int i = 0; i < config.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); + seq = block->forward(ctx, seq, pe, nullptr, seq_adaln); + sd::ggml_graph_cut::mark_graph_cut(seq, "llada_image.layers." + std::to_string(i), "seq"); + } + + seq = final_layer->forward(ctx, seq, seq_adaln); + + // Only the target latent is denoised; the source half of the image stream is context. + // The stream is padded to SEQ_MULTI_OF, so drop the pad tokens: they are not part of + // the latent grid that unpatchify reconstructs. + int64_t target_start = cap_len * 2 + img_len; + return ggml_ext_slice(gctx, seq, 1, target_start, target_start + n_img_token); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timestep, + ggml_tensor* context, + ggml_tensor* pe) { + // x: [N, C, H, W] + // timestep: [N,] + // context: [N, L, cap_feat_dim] + // pe: [L, d_head/2, 2, 2] + // return: [N, C, H, W] + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + + int patch_size = config.patch_size; + + auto img = DiT::pad_and_patchify(ctx, x, patch_size, patch_size, false); + + auto out = forward_core(ctx, img, timestep, context, pe); + + out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false); + + // The reference pipeline negates the model output before the scheduler step. + return ggml_ext_scale(ctx->ggml_ctx, out, -1.f); + } + }; + + struct LLaDAImageRunner : public DiffusionModelRunner { + public: + LLaDAImageConfig config; + LLaDAImageModel llada_image; + std::vector pe_vec; + + LLaDAImageRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(LLaDAImageConfig::detect_from_weights(tensor_storage_map, prefix)) { + llada_image = LLaDAImageModel(config); + llada_image.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "llada_image"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + llada_image.get_param_tensors(tensors, prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor) { + ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_GRAPH_SIZE); + ggml_tensor* x = make_input(x_tensor); + ggml_tensor* timesteps = make_input(timesteps_tensor); + GGML_ASSERT(x->ne[3] == 1); + GGML_ASSERT(!context_tensor.empty()); + ggml_tensor* context = make_input(context_tensor); + + pe_vec = Rope::gen_llada_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ZImage::SEQ_MULTI_OF, + config.theta, + circular_y_enabled, + circular_x_enabled, + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); + set_backend_tensor_data(pe, pe_vec.data()); + auto runner_ctx = get_context(); + + ggml_tensor* out = llada_image.forward(&runner_ctx, x, timesteps, context, pe); + + ggml_build_forward_expand(gf, out); + + return gf; + } + + sd::Tensor compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context) { + // x: [N, in_channels, h, w] + // timesteps: [N, ] + // context: [N, max_position, cap_feat_dim] + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context); + }; + + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + } + + ggml_cgraph* build_edit_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor, + const sd::Tensor& semantic_tensor, + const sd::Tensor& source_tensor) { + ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_GRAPH_SIZE); + ggml_tensor* x = make_input(x_tensor); + ggml_tensor* timesteps = make_input(timesteps_tensor); + ggml_tensor* context = make_input(context_tensor); + ggml_tensor* semantic = make_input(semantic_tensor); + ggml_tensor* source = make_input(source_tensor); + GGML_ASSERT(x->ne[3] == 1); + + pe_vec = Rope::gen_llada_image_edit_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(context->ne[1]), + static_cast(semantic->ne[1]), + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); + set_backend_tensor_data(pe, pe_vec.data()); + auto runner_ctx = get_context(); + + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + auto target = DiT::pad_and_patchify(&runner_ctx, x, config.patch_size, config.patch_size, false); + auto src = DiT::pad_and_patchify(&runner_ctx, source, config.patch_size, config.patch_size, false); + + auto out = llada_image.forward_editing(&runner_ctx, target, timesteps, context, semantic, src, pe); + out = DiT::unpatchify_and_crop(runner_ctx.ggml_ctx, out, H, W, config.patch_size, config.patch_size, false); + out = ggml_ext_scale(runner_ctx.ggml_ctx, out, -1.f); + + ggml_build_forward_expand(gf, out); + return gf; + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + + const auto* extra = std::get_if(&diffusion_params.extra); + bool has_semantic = extra != nullptr && extra->semantic != nullptr && !extra->semantic->empty(); + bool has_ref_latent = diffusion_params.ref_latents != nullptr && !diffusion_params.ref_latents->empty(); + if (has_semantic != has_ref_latent) { + LOG_WARN( + "llada_image: editing needs both the SigVQ features and the reference latent " + "(have semantic: %d, reference latent: %d); falling back to text to image", + static_cast(has_semantic), + static_cast(has_ref_latent)); + } + if (has_semantic && has_ref_latent) { + auto get_graph = [&]() -> ggml_cgraph* { + return build_edit_graph(*diffusion_params.x, + *diffusion_params.timesteps, + tensor_or_empty(diffusion_params.context), + *extra->semantic, + diffusion_params.ref_latents->front()); + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), + diffusion_params.x->dim()); + } + + return compute(n_threads, + *diffusion_params.x, + *diffusion_params.timesteps, + tensor_or_empty(diffusion_params.context)); + } + }; + +} // namespace LLaDAImage + +#endif // __SD_MODEL_DIFFUSION_LLADA_IMAGE_H__ diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index a4b3c38fd..b6334e8e6 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -39,6 +39,9 @@ const std::unordered_map REF_IMAGE_PRESETS = { {"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}}, {"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}}, {"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}}, + // pass_to_vlm routes the reference image to the conditioner, which is where LLaDA-Image's + // SigVQ encoder lives; it does its own half-resolution resize. + {"llada_image", {true, true, Rope::RefIndexMode::FIXED, true, true, -1, RefImageResizeMode::NONE, -1, -1}}, {"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}}, }; @@ -127,6 +130,11 @@ struct HunyuanVideoDiffusionExtra { const sd::Tensor* timestep_r = nullptr; }; +struct LLaDAImageDiffusionExtra { + // SigVQ semantic features of the reference image; present only in editing mode. + const sd::Tensor* semantic = nullptr; +}; + using DiffusionExtraParams = std::variant; + HunyuanVideoDiffusionExtra, + LLaDAImageDiffusionExtra>; struct DiffusionParams { const sd::Tensor* x = nullptr; diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index 4ae47e268..a1012cad8 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -131,16 +131,30 @@ namespace ZImage { int64_t num_heads; int64_t num_kv_heads; bool qk_norm; + bool split_qkv; public: - JointAttention(int64_t hidden_size, int64_t head_dim, int64_t num_heads, int64_t num_kv_heads, bool qk_norm) - : head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm) { - blocks["qkv"] = std::make_shared(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false); - float scale = 1.f; - blocks["out"] = std::make_shared(num_heads * head_dim, hidden_size, false, false, false, scale); + JointAttention(int64_t hidden_size, + int64_t head_dim, + int64_t num_heads, + int64_t num_kv_heads, + bool qk_norm, + bool norm_elementwise_affine = true, + bool split_qkv = false) + : head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm), split_qkv(split_qkv) { + float scale = 1.f; + if (split_qkv) { + blocks["to_q"] = std::make_shared(hidden_size, num_heads * head_dim, false); + blocks["to_k"] = std::make_shared(hidden_size, num_kv_heads * head_dim, false); + blocks["to_v"] = std::make_shared(hidden_size, num_kv_heads * head_dim, false); + blocks["to_out.0"] = std::make_shared(num_heads * head_dim, hidden_size, false, false, false, scale); + } else { + blocks["qkv"] = std::make_shared(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false); + blocks["out"] = std::make_shared(num_heads * head_dim, hidden_size, false, false, false, scale); + } if (qk_norm) { - blocks["q_norm"] = std::make_shared(head_dim); - blocks["k_norm"] = std::make_shared(head_dim); + blocks["q_norm"] = std::make_shared(head_dim, 1e-06f, norm_elementwise_affine); + blocks["k_norm"] = std::make_shared(head_dim, 1e-06f, norm_elementwise_affine); } } @@ -151,8 +165,35 @@ namespace ZImage { // x: [N, n_token, hidden_size] int64_t n_token = x->ne[1]; int64_t N = x->ne[2]; - auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); - auto out_proj = std::dynamic_pointer_cast(blocks["out"]); + auto out_proj = std::dynamic_pointer_cast(blocks[split_qkv ? "to_out.0" : "out"]); + + if (split_qkv) { + auto q_proj = std::dynamic_pointer_cast(blocks["to_q"]); + auto k_proj = std::dynamic_pointer_cast(blocks["to_k"]); + auto v_proj = std::dynamic_pointer_cast(blocks["to_v"]); + + if (sd_backend_is(ctx->backend, "ROCm")) { + out_proj->set_scale(1.f / 16.f); + out_proj->set_force_prec_f32(true); + q_proj->set_force_prec_f32(true); + k_proj->set_force_prec_f32(true); + v_proj->set_force_prec_f32(true); + } + + auto q = ggml_reshape_4d(ctx->ggml_ctx, q_proj->forward(ctx, x), head_dim, num_heads, n_token, N); + auto k = ggml_reshape_4d(ctx->ggml_ctx, k_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N); + auto v = ggml_reshape_4d(ctx->ggml_ctx, v_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N); + + if (qk_norm) { + q = std::dynamic_pointer_cast(blocks["q_norm"])->forward(ctx, q); + k = std::dynamic_pointer_cast(blocks["k_norm"])->forward(ctx, k); + } + + auto out = Rope::attention(ctx, q, k, v, pe, mask, 1.f / 128.f); + return out_proj->forward(ctx, out); + } + + auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); if (sd_backend_is(ctx->backend, "ROCm")) { out_proj->set_scale(1.f / 16.f); @@ -252,9 +293,12 @@ namespace ZImage { ggml_tensor* x, ggml_tensor* scale) { // x: [N, L, C] - // scale: [N, C] - scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C] - x = ggml_add(ctx, x, ggml_mul(ctx, x, scale)); + // scale: [N, C], or [N, L, C] when the caller modulates per token (LLaDA-Image editing + // feeds a per-token timestep embedding so each segment carries its own modulation). + if (scale->ne[1] != x->ne[1]) { + scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C] + } + x = ggml_add(ctx, x, ggml_mul(ctx, x, scale)); return x; } @@ -272,14 +316,16 @@ namespace ZImage { float ffn_dim_multiplier, float norm_eps, bool qk_norm, - bool modulation = true) + bool modulation = true, + bool norm_elementwise_affine = true, + bool split_qkv = false) : modulation(modulation) { - blocks["attention"] = std::make_shared(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm); + blocks["attention"] = std::make_shared(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm, norm_elementwise_affine, split_qkv); blocks["feed_forward"] = std::make_shared(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier); - blocks["attention_norm1"] = std::make_shared(hidden_size, norm_eps); - blocks["ffn_norm1"] = std::make_shared(hidden_size, norm_eps); - blocks["attention_norm2"] = std::make_shared(hidden_size, norm_eps); - blocks["ffn_norm2"] = std::make_shared(hidden_size, norm_eps); + blocks["attention_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["ffn_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["attention_norm2"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["ffn_norm2"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); if (modulation) { blocks["adaLN_modulation.0"] = std::make_shared(MIN(hidden_size, ADALN_EMBED_DIM), 4 * hidden_size); } diff --git a/src/model/te/llada_image_te.h b/src/model/te/llada_image_te.h new file mode 100644 index 000000000..e7492832c --- /dev/null +++ b/src/model/te/llada_image_te.h @@ -0,0 +1,604 @@ +#ifndef __SD_MODEL_TE_LLADA_IMAGE_TE_H__ +#define __SD_MODEL_TE_LLADA_IMAGE_TE_H__ + +#include +#include +#include + +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "model/common/ggml_block.hpp" +#include "model_loader.h" + +// The conditioning components LLaDA-Image puts around its LLaDA2-MoE backbone. +// Ref: LLaDAImageQueryFormerModel / LLaDAImageTextProjectionModel in +// https://github.com/inclusionAI/LLaDA-Image/blob/main/src/models/transformer_llada_image.py +// +// QueryFormer turns the LLaDA token embeddings into 256 learned queries that the pipeline +// appends to the backbone input; TextProjection maps the backbone hidden states to the +// denoiser's caption dimension. Neither uses RoPE, and every norm is parameter-free. +// Both MLPs use the tanh GELU approximation, so ggml_gelu (not ggml_gelu_erf). +// +// SigVQ is the editing-only image encoder: a 40-layer ViT whose output is quantized against a +// 16384-entry codebook, with the resulting ids embedded and projected into the semantic features +// the denoiser consumes. Its MLP uses the exact erf GELU, unlike the two above. + +namespace LLaDAImageTE { + constexpr int LLADA_IMAGE_TE_GRAPH_SIZE = 16384; + + struct QueryFormerConfig { + int64_t num_queries = 256; + int64_t hidden_size = 2048; + int64_t num_layers = 1; + int64_t num_heads = 16; + int64_t intermediate_size = 8192; + float norm_eps = 1e-6f; + }; + + struct TextProjectionConfig { + int64_t hidden_size = 2048; + int64_t intermediate_size = 8960; + int64_t num_layers = 6; + int64_t num_heads = 32; + int64_t projection_dim = 2560; + float norm_eps = 1e-6f; + }; + + // Cross-attention with a single fused in_proj over q (from the queries) and k/v (from the + // token embeddings). The checkpoint stores in_proj as one [3*hidden, hidden] parameter. + struct QueryAttention : public GGMLBlock { + protected: + int64_t hidden_size; + int64_t num_heads; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + std::string prefix = "") override { + GGMLBlock::init_params(ctx, tensor_storage_map, prefix); + enum ggml_type wtype = get_type(prefix + "in_proj_weight", tensor_storage_map, GGML_TYPE_F32); + params["in_proj_weight"] = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size * 3); + params["in_proj_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size * 3); + } + + public: + QueryAttention(int64_t hidden_size, int64_t num_heads) + : hidden_size(hidden_size), num_heads(num_heads) { + blocks["out_proj"] = std::make_shared(hidden_size, hidden_size, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* query, + ggml_tensor* context, + ggml_tensor* mask = nullptr) { + // query: [N, num_queries, hidden_size], context: [N, n_token, hidden_size] + ggml_context* gctx = ctx->ggml_ctx; + auto out_proj = std::dynamic_pointer_cast(blocks["out_proj"]); + + auto w = params["in_proj_weight"]; + auto b = params["in_proj_bias"]; + + auto slice_w = [&](int64_t index) { + return ggml_ext_slice(gctx, w, 1, index * hidden_size, (index + 1) * hidden_size); + }; + auto slice_b = [&](int64_t index) { + return ggml_ext_slice(gctx, b, 0, index * hidden_size, (index + 1) * hidden_size); + }; + + auto q = ggml_ext_linear(gctx, query, slice_w(0), slice_b(0)); + auto k = ggml_ext_linear(gctx, context, slice_w(1), slice_b(1)); + auto v = ggml_ext_linear(gctx, context, slice_w(2), slice_b(2)); + + auto x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask); // [N, num_queries, hidden_size] + return out_proj->forward(ctx, x); + } + }; + + struct QueryFormerBlock : public GGMLBlock { + protected: + QueryFormerConfig config; + + public: + QueryFormerBlock(const QueryFormerConfig& config) + : config(config) { + blocks["norm_q"] = std::make_shared(config.hidden_size, config.norm_eps, false); + blocks["norm_k"] = std::make_shared(config.hidden_size, config.norm_eps, false); + blocks["cross_attn"] = std::make_shared(config.hidden_size, config.num_heads); + blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps, false); + blocks["mlp.fc1"] = std::make_shared(config.hidden_size, config.intermediate_size, true); + blocks["mlp.fc2"] = std::make_shared(config.intermediate_size, config.hidden_size, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* query, + ggml_tensor* context, + ggml_tensor* mask = nullptr) { + auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); + auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); + auto cross_attn = std::dynamic_pointer_cast(blocks["cross_attn"]); + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto fc1 = std::dynamic_pointer_cast(blocks["mlp.fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["mlp.fc2"]); + + // The reference overwrites query_embeds with its normalized value before the + // residual add, so both residuals here are on normalized activations. + query = norm_q->forward(ctx, query); + auto ctx_n = norm_k->forward(ctx, context); + query = ggml_add(ctx->ggml_ctx, query, cross_attn->forward(ctx, query, ctx_n, mask)); + query = norm1->forward(ctx, query); + + auto h = fc1->forward(ctx, query); + h = ggml_gelu(ctx->ggml_ctx, h); + h = fc2->forward(ctx, h); + return ggml_add(ctx->ggml_ctx, query, h); + } + }; + + struct QueryFormerModel : public GGMLBlock { + protected: + QueryFormerConfig config; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + params["meta_queries"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, config.num_queries); + } + + public: + QueryFormerModel() = default; + QueryFormerModel(const QueryFormerConfig& config) + : config(config) { + for (int i = 0; i < config.num_layers; i++) { + blocks["query_blocks." + std::to_string(i)] = std::make_shared(config); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* inputs_embeds, + ggml_tensor* mask = nullptr) { + // inputs_embeds: [N, n_token, hidden_size] -> [N, num_queries, hidden_size] + auto query = params["meta_queries"]; + query = ggml_reshape_3d(ctx->ggml_ctx, query, config.hidden_size, config.num_queries, 1); + + for (int i = 0; i < config.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["query_blocks." + std::to_string(i)]); + query = block->forward(ctx, query, inputs_embeds, mask); + } + return query; + } + }; + + struct TextProjectionAttention : public GGMLBlock { + protected: + int64_t num_heads; + int64_t head_dim; + + public: + TextProjectionAttention(const TextProjectionConfig& config) + : num_heads(config.num_heads), head_dim(config.hidden_size / config.num_heads) { + blocks["q_proj"] = std::make_shared(config.hidden_size, config.hidden_size, true); + blocks["k_proj"] = std::make_shared(config.hidden_size, config.hidden_size, true); + blocks["v_proj"] = std::make_shared(config.hidden_size, config.hidden_size, true); + blocks["out_proj"] = std::make_shared(config.hidden_size, config.hidden_size, true); + blocks["q_norm"] = std::make_shared(head_dim, config.norm_eps, false); + blocks["k_norm"] = std::make_shared(head_dim, config.norm_eps, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [N, n_token, hidden_size] + ggml_context* gctx = ctx->ggml_ctx; + int64_t n_token = x->ne[1]; + int64_t N = x->ne[2]; + + auto q_proj = std::dynamic_pointer_cast(blocks["q_proj"]); + auto k_proj = std::dynamic_pointer_cast(blocks["k_proj"]); + auto v_proj = std::dynamic_pointer_cast(blocks["v_proj"]); + auto out_proj = std::dynamic_pointer_cast(blocks["out_proj"]); + auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); + auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); + + auto q = q_proj->forward(ctx, x); + auto k = k_proj->forward(ctx, x); + auto v = v_proj->forward(ctx, x); + + q = ggml_reshape_4d(gctx, q, head_dim, num_heads, n_token, N); + k = ggml_reshape_4d(gctx, k, head_dim, num_heads, n_token, N); + q = q_norm->forward(ctx, q); + k = k_norm->forward(ctx, k); + q = ggml_reshape_3d(gctx, q, head_dim * num_heads, n_token, N); + k = ggml_reshape_3d(gctx, k, head_dim * num_heads, n_token, N); + + auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads); + return out_proj->forward(ctx, out); + } + }; + + struct TextProjectionBlock : public GGMLBlock { + public: + TextProjectionBlock(const TextProjectionConfig& config) { + blocks["self_attn"] = std::make_shared(config); + blocks["layer_norm1"] = std::make_shared(config.hidden_size, config.norm_eps, false); + blocks["layer_norm2"] = std::make_shared(config.hidden_size, config.norm_eps, false); + blocks["mlp.fc1"] = std::make_shared(config.hidden_size, config.intermediate_size, true); + blocks["mlp.fc2"] = std::make_shared(config.intermediate_size, config.hidden_size, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto self_attn = std::dynamic_pointer_cast(blocks["self_attn"]); + auto layer_norm1 = std::dynamic_pointer_cast(blocks["layer_norm1"]); + auto layer_norm2 = std::dynamic_pointer_cast(blocks["layer_norm2"]); + auto fc1 = std::dynamic_pointer_cast(blocks["mlp.fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["mlp.fc2"]); + + x = ggml_add(ctx->ggml_ctx, x, self_attn->forward(ctx, layer_norm1->forward(ctx, x))); + + auto h = fc1->forward(ctx, layer_norm2->forward(ctx, x)); + h = ggml_gelu(ctx->ggml_ctx, h); + h = fc2->forward(ctx, h); + return ggml_add(ctx->ggml_ctx, x, h); + } + }; + + struct TextProjectionModel : public GGMLBlock { + protected: + TextProjectionConfig config; + + public: + TextProjectionModel() = default; + TextProjectionModel(const TextProjectionConfig& config) + : config(config) { + for (int i = 0; i < config.num_layers; i++) { + blocks["layers." + std::to_string(i)] = std::make_shared(config); + } + blocks["projector"] = std::make_shared(config.hidden_size, config.projection_dim, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [N, n_token, hidden_size] -> [N, n_token, projection_dim] + for (int i = 0; i < config.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); + x = block->forward(ctx, x); + } + auto projector = std::dynamic_pointer_cast(blocks["projector"]); + return projector->forward(ctx, x); + } + }; + + struct SigVQConfig { + int64_t image_size = 2048; + int64_t patch_size = 16; + int64_t in_channels = 3; + int64_t hidden_size = 1536; + int64_t intermediate_size = 6144; + int64_t num_layers = 40; + int64_t num_heads = 16; + int64_t codebook_size = 16384; + int64_t codebook_embed_dim = 2048; + int64_t semantic_embed_dim = 4096; + float norm_eps = 1e-6f; + }; + + struct SigVQAttention : public GGMLBlock { + protected: + int64_t num_heads; + int64_t head_dim; + + public: + SigVQAttention(const SigVQConfig& config) + : num_heads(config.num_heads), head_dim(config.hidden_size / config.num_heads) { + blocks["qkv"] = std::make_shared(config.hidden_size, config.hidden_size * 3, true); + blocks["proj"] = std::make_shared(config.hidden_size, config.hidden_size, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [N, n_token, hidden_size] + ggml_context* gctx = ctx->ggml_ctx; + auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); + auto out_proj = std::dynamic_pointer_cast(blocks["proj"]); + + int64_t hidden_size = num_heads * head_dim; + auto qkv = qkv_proj->forward(ctx, x); + auto q = ggml_ext_slice(gctx, qkv, 0, 0, hidden_size); + auto k = ggml_ext_slice(gctx, qkv, 0, hidden_size, hidden_size * 2); + auto v = ggml_ext_slice(gctx, qkv, 0, hidden_size * 2, hidden_size * 3); + + auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads); + return out_proj->forward(ctx, out); + } + }; + + struct SigVQBlock : public GGMLBlock { + public: + SigVQBlock(const SigVQConfig& config) { + blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps); + blocks["norm2"] = std::make_shared(config.hidden_size, config.norm_eps); + blocks["attn"] = std::make_shared(config); + blocks["mlp.fc1"] = std::make_shared(config.hidden_size, config.intermediate_size, true); + blocks["mlp.fc2"] = std::make_shared(config.intermediate_size, config.hidden_size, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto fc1 = std::dynamic_pointer_cast(blocks["mlp.fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["mlp.fc2"]); + + x = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, norm1->forward(ctx, x))); + auto h = fc1->forward(ctx, norm2->forward(ctx, x)); + h = ggml_gelu_erf(ctx->ggml_ctx, h); + h = fc2->forward(ctx, h); + return ggml_add(ctx->ggml_ctx, x, h); + } + }; + + struct SigVQModel : public GGMLBlock { + protected: + SigVQConfig config; + + public: + SigVQModel() = default; + SigVQModel(const SigVQConfig& config) + : config(config) { + blocks["visual.patch_embed.proj"] = std::make_shared(config.in_channels, + config.hidden_size, + std::make_pair(config.patch_size, config.patch_size), + std::make_pair(config.patch_size, config.patch_size)); + for (int i = 0; i < config.num_layers; i++) { + blocks["visual.blocks." + std::to_string(i)] = std::make_shared(config); + } + blocks["vqmodel.quant_conv"] = std::make_shared(config.hidden_size, + config.codebook_embed_dim, + std::make_pair(1, 1)); + blocks["prior_projector.net.0.proj"] = std::make_shared(config.semantic_embed_dim, config.semantic_embed_dim, true); + blocks["prior_projector.net.2"] = std::make_shared(config.semantic_embed_dim, config.semantic_embed_dim, true); + } + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + params["visual.embeddings.position_embedding.weight"] = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, (config.image_size / config.patch_size) * (config.image_size / config.patch_size)); + params["vqmodel.quantize.embedding.weight"] = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.codebook_embed_dim, config.codebook_size); + params["prior_token_embedding.weight"] = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.semantic_embed_dim, config.codebook_size); + } + + // Bilinear-resamples the square position-embedding grid onto the image's patch grid. + // The reference uses grid_sample(align_corners=False, padding_mode="border"); the source + // coordinate for output index j is therefore (j + 0.5) * side / out - 0.5, clamped. + ggml_tensor* resample_pos_embed(GGMLRunnerContext* ctx, + ggml_tensor* pos_idx, + ggml_tensor* pos_weight) { + auto pos_embed = params["visual.embeddings.position_embedding.weight"]; + auto gathered = ggml_get_rows(ctx->ggml_ctx, pos_embed, pos_idx); + return ggml_mul(ctx->ggml_ctx, gathered, pos_weight); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* pixel_values, + const std::vector& pos_idx, + const std::vector& pos_weight) { + // pixel_values: [N, in_channels, H, W] -> [N, grid_h * grid_w, semantic_embed_dim] + ggml_context* gctx = ctx->ggml_ctx; + + auto patch_embed = std::dynamic_pointer_cast(blocks["visual.patch_embed.proj"]); + auto quant_conv = std::dynamic_pointer_cast(blocks["vqmodel.quant_conv"]); + auto proj_0 = std::dynamic_pointer_cast(blocks["prior_projector.net.0.proj"]); + auto proj_2 = std::dynamic_pointer_cast(blocks["prior_projector.net.2"]); + + auto x = patch_embed->forward(ctx, pixel_values); // [N, hidden_size, grid_h, grid_w] + int64_t grid_w = x->ne[0]; + int64_t grid_h = x->ne[1]; + int64_t n_token = grid_h * grid_w; + int64_t N = x->ne[3]; + + x = ggml_reshape_3d(gctx, x, n_token, config.hidden_size, N); + x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3)); // [N, n_token, hidden_size] + + ggml_tensor* pos = nullptr; + for (size_t i = 0; i < pos_idx.size(); i++) { + auto corner = resample_pos_embed(ctx, pos_idx[i], pos_weight[i]); + pos = pos == nullptr ? corner : ggml_add(gctx, pos, corner); + } + x = ggml_add(gctx, x, ggml_reshape_3d(gctx, pos, config.hidden_size, n_token, N)); + + for (int i = 0; i < config.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["visual.blocks." + std::to_string(i)]); + x = block->forward(ctx, x); + } + + // quant_conv is 1x1, so run it as a per-token projection rather than reshaping to 2-D. + x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3)); // [N, hidden_size, n_token] + x = ggml_reshape_4d(gctx, x, n_token, 1, config.hidden_size, N); + x = quant_conv->forward(ctx, x); // [N, codebook_embed_dim, 1, n_token] + x = ggml_reshape_3d(gctx, x, n_token, config.codebook_embed_dim, N); + x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3)); // [N, n_token, codebook_embed_dim] + + // Both sides are L2-normalized, so the nearest codebook entry by euclidean distance + // is the one with the largest dot product. + auto codebook = ggml_l2_norm(gctx, params["vqmodel.quantize.embedding.weight"], 1e-12f); + auto normed = ggml_l2_norm(gctx, x, 1e-12f); + auto logits = ggml_mul_mat(gctx, codebook, normed); // [N, n_token, codebook_size] + auto token_ids = ggml_argmax(gctx, ggml_reshape_2d(gctx, logits, config.codebook_size, n_token * N)); + + auto semantic = ggml_get_rows(gctx, params["prior_token_embedding.weight"], token_ids); + semantic = ggml_reshape_3d(gctx, semantic, config.semantic_embed_dim, n_token, N); + + auto h = proj_0->forward(ctx, semantic); + h = ggml_silu(gctx, h); + return proj_2->forward(ctx, h); + } + }; + + struct QueryFormerRunner : public GGMLRunner { + public: + QueryFormerConfig config; + QueryFormerModel query_former; + + QueryFormerRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager) { + query_former = QueryFormerModel(config); + query_former.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "llada_image_queryformer"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) { + query_former.get_param_tensors(tensors, prefix); + } + + sd::Tensor compute(int n_threads, const sd::Tensor& inputs_embeds) { + auto get_graph = [&]() -> ggml_cgraph* { + ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE); + ggml_tensor* x = make_input(inputs_embeds); + auto runner_ctx = get_context(); + ggml_tensor* out = query_former.forward(&runner_ctx, x); + ggml_build_forward_expand(gf, out); + return gf; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), + inputs_embeds.dim()); + } + }; + + struct TextProjectionRunner : public GGMLRunner { + public: + TextProjectionConfig config; + TextProjectionModel text_projection; + + TextProjectionRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager) { + text_projection = TextProjectionModel(config); + text_projection.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "llada_image_text_projection"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) { + text_projection.get_param_tensors(tensors, prefix); + } + + sd::Tensor compute(int n_threads, const sd::Tensor& hidden_states) { + auto get_graph = [&]() -> ggml_cgraph* { + ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE); + ggml_tensor* x = make_input(hidden_states); + auto runner_ctx = get_context(); + ggml_tensor* out = text_projection.forward(&runner_ctx, x); + ggml_build_forward_expand(gf, out); + return gf; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), + hidden_states.dim()); + } + }; + + struct SigVQRunner : public GGMLRunner { + public: + SigVQConfig config; + SigVQModel sigvq; + std::array, 4> pos_idx_data; + std::array, 4> pos_weight_data; + + SigVQRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager) { + sigvq = SigVQModel(config); + sigvq.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "llada_image_sigvq"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) { + sigvq.get_param_tensors(tensors, prefix); + } + + // Precomputes the four bilinear taps that resample the square position-embedding grid + // onto a grid_h x grid_w patch grid, matching grid_sample(align_corners=False, + // padding_mode="border"). + void build_pos_embed_taps(int64_t grid_h, int64_t grid_w) { + const int64_t side = config.image_size / config.patch_size; + for (auto& v : pos_idx_data) { + v.clear(); + } + for (auto& v : pos_weight_data) { + v.clear(); + } + + auto clamp_index = [side](int64_t v) { + return static_cast(std::min(std::max(v, 0), side - 1)); + }; + + for (int64_t i = 0; i < grid_h; ++i) { + double src_h = (static_cast(i) + 0.5) * side / static_cast(grid_h) - 0.5; + int64_t h_floor = static_cast(std::floor(src_h)); + double dh = src_h - static_cast(h_floor); + for (int64_t j = 0; j < grid_w; ++j) { + double src_w = (static_cast(j) + 0.5) * side / static_cast(grid_w) - 0.5; + int64_t w_floor = static_cast(std::floor(src_w)); + double dw = src_w - static_cast(w_floor); + + int32_t h0 = clamp_index(h_floor); + int32_t h1 = clamp_index(h_floor + 1); + int32_t w0 = clamp_index(w_floor); + int32_t w1 = clamp_index(w_floor + 1); + + pos_idx_data[0].push_back(h0 * static_cast(side) + w0); + pos_idx_data[1].push_back(h0 * static_cast(side) + w1); + pos_idx_data[2].push_back(h1 * static_cast(side) + w0); + pos_idx_data[3].push_back(h1 * static_cast(side) + w1); + + pos_weight_data[0].push_back(static_cast((1.0 - dh) * (1.0 - dw))); + pos_weight_data[1].push_back(static_cast((1.0 - dh) * dw)); + pos_weight_data[2].push_back(static_cast(dh * (1.0 - dw))); + pos_weight_data[3].push_back(static_cast(dh * dw)); + } + } + } + + sd::Tensor compute(int n_threads, const sd::Tensor& pixel_values) { + auto get_graph = [&]() -> ggml_cgraph* { + ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE); + ggml_tensor* x = make_input(pixel_values); + + int64_t grid_h = x->ne[1] / config.patch_size; + int64_t grid_w = x->ne[0] / config.patch_size; + build_pos_embed_taps(grid_h, grid_w); + + std::vector pos_idx; + std::vector pos_weight; + for (int i = 0; i < 4; i++) { + auto idx = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, static_cast(pos_idx_data[i].size())); + set_backend_tensor_data(idx, pos_idx_data[i].data()); + auto w = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, 1, static_cast(pos_weight_data[i].size())); + set_backend_tensor_data(w, pos_weight_data[i].data()); + pos_idx.push_back(idx); + pos_weight.push_back(w); + } + + auto runner_ctx = get_context(); + ggml_tensor* out = sigvq.forward(&runner_ctx, x, pos_idx, pos_weight); + ggml_build_forward_expand(gf, out); + return gf; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), 3); + } + }; + +} // namespace LLaDAImageTE + +#endif // __SD_MODEL_TE_LLADA_IMAGE_TE_H__ diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index 03deb610a..5c3e2802b 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -49,6 +49,7 @@ namespace LLM { GEMMA2_2B, GEMMA4_12B, GPT_OSS_20B, + LLADA2_MOE, ARCH_COUNT, }; @@ -62,6 +63,7 @@ namespace LLM { "gemma2_2b", "gemma4_12b", "gpt_oss_20b", + "llada2_moe", }; enum class MLPActivation { @@ -125,6 +127,17 @@ namespace LLM { std::vector sliding_attention; int64_t num_experts = 0; int64_t num_experts_per_tok = 0; + bool qkv_fused = false; + bool bidirectional = false; + float partial_rotary = 1.f; + + // DeepSeek-V3-style grouped-sigmoid MoE routing (LLaDA2) + int64_t moe_intermediate_size = 0; + int64_t num_shared_experts = 0; + int64_t first_k_dense_replace = 0; + int64_t n_group = 0; + int64_t topk_group = 0; + float routed_scaling_factor = 1.f; LLMVisionConfig vision; bool have_vision_weight = false; bool llama_cpp_style = false; @@ -212,6 +225,31 @@ namespace LLM { config.intermediate_size = 9216; config.num_layers = 26; config.vocab_size = 256000; + } else if (arch == LLMArch::LLADA2_MOE) { + config.head_dim = 128; + config.num_heads = 16; + config.num_kv_heads = 4; + config.qkv_bias = false; + config.attention_out_bias = false; + config.qk_norm = true; + config.rms_norm_eps = 1e-6f; + config.hidden_size = 2048; + config.intermediate_size = 5120; + config.num_layers = 20; + config.vocab_size = 173568; + config.max_position_embeddings = 16384; + config.rope_thetas = {600000.f}; + config.qkv_fused = true; + config.bidirectional = true; + config.partial_rotary = 0.5f; + config.num_experts = 256; + config.num_experts_per_tok = 8; + config.moe_intermediate_size = 512; + config.num_shared_experts = 1; + config.first_k_dense_replace = 1; + config.n_group = 8; + config.topk_group = 4; + config.routed_scaling_factor = 2.5f; } else if (arch == LLMArch::GPT_OSS_20B) { config.head_dim = 64; config.num_heads = 64; @@ -419,6 +457,195 @@ namespace LLM { } }; + // LLaDA2's MoE differs from GPT-OSS's in three ways that all change the result: + // routing scores are sigmoid (not softmax over the selected logits), expert selection is + // group-limited and uses a bias term that the returned weights do NOT include, and the + // experts carry no biases. Ref: LLaDA2MoeGate / LLaDA2MoeSparseMoeBlock in + // modeling_llada2uni_moe.py. + struct LLaDA2MoEMLP : public GGMLBlock { + protected: + int64_t hidden_size; + int64_t moe_intermediate_size; + int64_t num_experts; + int64_t num_experts_per_tok; + int64_t n_group; + int64_t topk_group; + float routed_scaling_factor; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + std::string prefix = "") override { + GGMLBlock::init_params(ctx, tensor_storage_map, prefix); + + auto supported_type = [](ggml_type wtype, int64_t in_features) { + if (in_features % ggml_blck_size(wtype) != 0) { + return GGML_TYPE_F32; + } + return wtype; + }; + + // The reference runs the router in fp32; keep the weight in fp32 so the sigmoid + // scores and the group sums match. + params["gate.weight"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_experts); + params["gate.expert_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts); + + ggml_type gate_type = supported_type(get_type(prefix + "experts.gate_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size); + ggml_type up_type = supported_type(get_type(prefix + "experts.up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size); + ggml_type down_type = supported_type(get_type(prefix + "experts.down_proj.weight", tensor_storage_map, GGML_TYPE_F32), moe_intermediate_size); + + // HF ships the stacked experts as 3-D nn.Parameters, while the ComfyUI GGUF repack + // flattens the expert axis into ne[1]. Declare whichever the file holds - the two are + // bit-identical, and forward() reshapes to 3-D for ggml_mul_mat_id either way. + auto declare_experts = [&](const std::string& name, ggml_type type, int64_t in_dim, int64_t out_dim) { + auto storage = tensor_storage_map.find(prefix + name); + if (storage != tensor_storage_map.end() && storage->second.n_dims == 2) { + GGML_ASSERT(storage->second.nelements() == in_dim * out_dim * num_experts); + params[name] = ggml_new_tensor_2d(ctx, type, in_dim, out_dim * num_experts); + } else { + params[name] = ggml_new_tensor_3d(ctx, type, in_dim, out_dim, num_experts); + } + }; + + declare_experts("experts.gate_proj.weight", gate_type, hidden_size, moe_intermediate_size); + declare_experts("experts.up_proj.weight", up_type, hidden_size, moe_intermediate_size); + declare_experts("experts.down_proj.weight", down_type, moe_intermediate_size, hidden_size); + } + + public: + LLaDA2MoEMLP(const LLMConfig& config) + : hidden_size(config.hidden_size), + moe_intermediate_size(config.moe_intermediate_size), + num_experts(config.num_experts), + num_experts_per_tok(config.num_experts_per_tok), + n_group(config.n_group), + topk_group(config.topk_group), + routed_scaling_factor(config.routed_scaling_factor) { + if (config.num_shared_experts > 0) { + blocks["shared_experts"] = std::make_shared(config.hidden_size, + config.moe_intermediate_size * config.num_shared_experts, + false, + config.mlp_activation); + } + } + + // Reproduces group_limited_topk(): keep the topk_group groups with the highest + // "sum of the two best scores in the group", then take the global top-k among them. + ggml_tensor* group_limited_mask(GGMLRunnerContext* ctx, + ggml_tensor* routing_scores, + int64_t n_token_total) { + ggml_context* gctx = ctx->ggml_ctx; + const int64_t per_group = num_experts / n_group; + + // [experts_per_group, n_group * tokens] so top-2 runs per (group, token) row. + auto grouped = ggml_reshape_2d(gctx, routing_scores, per_group, n_group * n_token_total); + auto best2_idx = ggml_argsort_top_k(gctx, grouped, 2); // [2, n_group * tokens] + auto grouped_val = ggml_reshape_3d(gctx, grouped, 1, per_group, n_group * n_token_total); + auto best2 = ggml_get_rows(gctx, grouped_val, best2_idx); // [1, 2, n_group * tokens] + best2 = ggml_reshape_2d(gctx, best2, 2, n_group * n_token_total); + auto group_score = ggml_reshape_2d(gctx, ggml_sum_rows(gctx, best2), n_group, n_token_total); // [n_group, tokens] + + // Threshold = the topk_group-th largest group score, taken from the sorted top-k. + auto top_groups = ggml_argsort_top_k(gctx, group_score, (int)topk_group); // [topk_group, tokens] + auto group_val = ggml_reshape_3d(gctx, group_score, 1, n_group, n_token_total); + auto top_scores = ggml_get_rows(gctx, group_val, top_groups); // [1, topk_group, tokens] + top_scores = ggml_reshape_2d(gctx, top_scores, topk_group, n_token_total); + auto threshold = ggml_view_2d(gctx, + top_scores, + 1, + n_token_total, + top_scores->nb[1], + (topk_group - 1) * top_scores->nb[0]); // [1, tokens] + threshold = ggml_cont(gctx, threshold); + + // keep = 1 - step(threshold - score). step(0) == 0, so the group sitting exactly on + // the threshold is kept without needing an epsilon. + auto diff = ggml_sub(gctx, ggml_repeat(gctx, threshold, group_score), group_score); + auto keep = ggml_scale_bias(gctx, ggml_step(gctx, diff), -1.f, 1.f); // [n_group, tokens] + + // 0 for kept groups, a large negative for dropped ones, broadcast over the group. + auto additive = ggml_scale_bias(gctx, keep, 1e30f, -1e30f); + additive = ggml_reshape_3d(gctx, additive, 1, n_group, n_token_total); + auto expanded = ggml_repeat_4d(gctx, additive, per_group, n_group, n_token_total, 1); + return ggml_reshape_2d(gctx, expanded, num_experts, n_token_total); + } + + ggml_tensor* expert_linear(GGMLRunnerContext* ctx, + const std::string& weight_name, + ggml_tensor* x, + ggml_tensor* selected_experts) { + ggml_tensor* w = params[weight_name]; + if (w->ne[2] != num_experts) { + // Flattened layout: split the expert axis back out. ne[0] is untouched, so this + // stays valid for quantized types. + w = ggml_reshape_3d(ctx->ggml_ctx, w, w->ne[0], w->ne[1] / num_experts, num_experts); + } + return ggml_mul_mat_id(ctx->ggml_ctx, w, x, selected_experts); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [N, n_token, hidden_size] + GGML_ASSERT(num_experts > 0 && num_experts_per_tok > 0); + GGML_ASSERT(n_group > 0 && topk_group > 0 && num_experts % n_group == 0); + + ggml_context* gctx = ctx->ggml_ctx; + const int64_t n_token = x->ne[1]; + const int64_t N = x->ne[2]; + const int64_t n_token_total = n_token * N; + + auto identity = x; + + auto logits = ggml_mul_mat(gctx, params["gate.weight"], x); + logits = ggml_reshape_2d(gctx, logits, num_experts, n_token_total); + auto scores = ggml_sigmoid(gctx, logits); // [num_experts, tokens] + + // The bias steers selection only; the combine weights come from the unbiased scores. + auto routing = ggml_add(gctx, scores, params["gate.expert_bias"]); + routing = ggml_add(gctx, routing, group_limited_mask(ctx, routing, n_token_total)); + + auto selected_experts = ggml_argsort_top_k(gctx, routing, (int)num_experts_per_tok); // [top_k, tokens] + auto score_rows = ggml_reshape_3d(gctx, scores, 1, num_experts, n_token_total); + auto weights = ggml_get_rows(gctx, score_rows, selected_experts); // [1, top_k, tokens] + weights = ggml_reshape_2d(gctx, weights, num_experts_per_tok, n_token_total); + + if (num_experts_per_tok > 1) { + auto denom = ggml_scale_bias(gctx, ggml_sum_rows(gctx, weights), 1.f, 1e-20f); // [1, tokens] + weights = ggml_div(gctx, weights, ggml_repeat(gctx, denom, weights)); + } + weights = ggml_scale(gctx, weights, routed_scaling_factor); + weights = ggml_reshape_3d(gctx, weights, 1, num_experts_per_tok, n_token_total); + + auto xf = ggml_reshape_3d(gctx, x, hidden_size, 1, n_token_total); + auto gate = expert_linear(ctx, "experts.gate_proj.weight", xf, selected_experts); + auto up = expert_linear(ctx, "experts.up_proj.weight", xf, selected_experts); + auto activated = ggml_swiglu_split(gctx, gate, up); + auto experts = expert_linear(ctx, "experts.down_proj.weight", activated, selected_experts); + experts = ggml_mul(gctx, experts, weights); + + ggml_tensor* out = nullptr; + for (int64_t i = 0; i < num_experts_per_tok; ++i) { + auto expert_out = ggml_view_2d(gctx, + experts, + hidden_size, + n_token_total, + experts->nb[2], + i * experts->nb[1]); + out = out == nullptr ? expert_out : ggml_add(gctx, out, expert_out); + } + if (num_experts_per_tok == 1) { + out = ggml_cont(gctx, out); + } + out = ggml_reshape_3d(gctx, out, hidden_size, n_token, N); + + auto shared_it = blocks.find("shared_experts"); + if (shared_it != blocks.end()) { + auto shared_experts = std::dynamic_pointer_cast(shared_it->second); + out = ggml_add(gctx, out, shared_experts->forward(ctx, identity)); + } + + return out; + } + }; + struct GPTOSSMLP : public GGMLBlock { protected: int64_t hidden_size; @@ -605,21 +832,31 @@ namespace LLM { } txt_token_end = image_embeds[i].first; - auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); + // An embed can sit flush against the previous one or at the very start/end of the + // sequence, leaving no text tokens to splice around it. + if (txt_token_end > txt_token_start) { + auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); + if (input_embed == nullptr) { + input_embed = txt_embed; + } else { + input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1); + } + } + if (input_embed == nullptr) { - input_embed = txt_embed; + input_embed = image_embeds[i].second; } else { - input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1); + input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1); } - - input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1); } txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1]; txt_token_end = raw_x->ne[1]; - auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); - input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1); + if (txt_token_end > txt_token_start) { + auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); + input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1); + } GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]); return input_embed; } @@ -1122,6 +1359,7 @@ namespace LLM { bool k_eq_v; bool v_norm; bool unscaled_attention; + bool qkv_fused; float rms_norm_eps; int rope_pairs; @@ -1147,12 +1385,20 @@ namespace LLM { k_eq_v(global_layer && config.global_k_eq_v), v_norm(config.v_norm), unscaled_attention(config.unscaled_attention), + qkv_fused(config.qkv_fused), rms_norm_eps(config.rms_norm_eps), rope_pairs(0) { - blocks["q_proj"] = std::make_shared(config.hidden_size, num_heads * head_dim, config.qkv_bias); - blocks["k_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); - if (!k_eq_v) { - blocks["v_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); + if (qkv_fused) { + // The checkpoint ships q, k and v as one tensor and the loader cannot split a + // source tensor, so keep it fused and slice it in forward(). + GGML_ASSERT(!k_eq_v); + blocks["query_key_value"] = std::make_shared(config.hidden_size, (num_heads + num_kv_heads * 2) * head_dim, config.qkv_bias); + } else { + blocks["q_proj"] = std::make_shared(config.hidden_size, num_heads * head_dim, config.qkv_bias); + blocks["k_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); + if (!k_eq_v) { + blocks["v_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); + } } blocks["o_proj"] = std::make_shared(num_heads * head_dim, config.hidden_size, config.attention_out_bias); if (config.qk_norm) { @@ -1161,7 +1407,7 @@ namespace LLM { } // Proportional RoPE rotates only the leading `rope_pairs` dimension pairs of the head; // the rest are left unrotated through freq_factors (see rope_freq_factors()). - float partial = global_layer ? config.global_partial_rotary : 1.f; + float partial = global_layer && config.global_partial_rotary != 1.f ? config.global_partial_rotary : config.partial_rotary; rope_pairs = static_cast(partial * head_dim / 2.f); } @@ -1186,14 +1432,28 @@ namespace LLM { // x: [N, n_token, hidden_size] int64_t n_token = x->ne[1]; int64_t N = x->ne[2]; - auto q_proj = std::dynamic_pointer_cast(blocks["q_proj"]); - auto k_proj = std::dynamic_pointer_cast(blocks["k_proj"]); - auto v_proj = k_eq_v ? nullptr : std::dynamic_pointer_cast(blocks["v_proj"]); auto out_proj = std::dynamic_pointer_cast(blocks["o_proj"]); - auto q = q_proj->forward(ctx, x); // [N, n_token, num_heads*head_dim] - auto k = k_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim] - auto v = k_eq_v ? k : v_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim] + ggml_tensor* q = nullptr; + ggml_tensor* k = nullptr; + ggml_tensor* v = nullptr; + if (qkv_fused) { + auto qkv_proj = std::dynamic_pointer_cast(blocks["query_key_value"]); + auto qkv = qkv_proj->forward(ctx, x); // [N, n_token, (num_heads + num_kv_heads*2)*head_dim] + int64_t q_len = num_heads * head_dim; + int64_t k_len = num_kv_heads * head_dim; + q = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, 0, q_len); + k = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len, q_len + k_len); + v = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len + k_len, q_len + k_len * 2); + } else { + auto q_proj = std::dynamic_pointer_cast(blocks["q_proj"]); + auto k_proj = std::dynamic_pointer_cast(blocks["k_proj"]); + auto v_proj = k_eq_v ? nullptr : std::dynamic_pointer_cast(blocks["v_proj"]); + + q = q_proj->forward(ctx, x); // [N, n_token, num_heads*head_dim] + k = k_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim] + v = k_eq_v ? k : v_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim] + } q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, n_token, N); // [N, n_token, num_heads, head_dim] k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_kv_heads, n_token, N); // [N, n_token, num_kv_heads, head_dim] @@ -1336,6 +1596,38 @@ namespace LLM { 1.f, 32.f, 1.f); + } else if (arch == LLMArch::LLADA2_MOE) { + // LLaDA2 slices the head (query[..., :rotary_dim]) instead of zero-padding + // inv_freq like gemma does, so rotate_half pairs i with i + rotary_dim/2 and the + // frequencies use rotary_dim as the exponent denominator. Passing n_dims = + // rotary_dim reproduces both; freq_factors would give the wrong pairing. + int rotary_dim = rope_pairs * 2; + q = ggml_rope_ext(ctx->ggml_ctx, + q, + input_pos, + nullptr, + rotary_dim, + GGML_ROPE_TYPE_NEOX, + static_cast(max_position_embeddings), + rope_thetas[0], + 1.f, + 0.f, + 1.f, + 32.f, + 1.f); + k = ggml_rope_ext(ctx->ggml_ctx, + k, + input_pos, + nullptr, + rotary_dim, + GGML_ROPE_TYPE_NEOX, + static_cast(max_position_embeddings), + rope_thetas[0], + 1.f, + 0.f, + 1.f, + 32.f, + 1.f); } else if (arch == LLMArch::QWEN3_VL) { int sections[4] = {24, 20, 20, 0}; q = ggml_rope_multi(ctx->ggml_ctx, q, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_IMROPE, 262144, 5000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); @@ -1432,6 +1724,8 @@ namespace LLM { blocks["self_attn"] = std::make_shared(config, sliding_attention == 0); if (config.arch == LLMArch::GPT_OSS_20B) { blocks["mlp"] = std::make_shared(config); + } else if (config.arch == LLMArch::LLADA2_MOE && layer_index >= config.first_k_dense_replace) { + blocks["mlp"] = std::make_shared(config); } else { blocks["mlp"] = std::make_shared(config.hidden_size, config.intermediate_size, @@ -1485,6 +1779,10 @@ namespace LLM { if (arch == LLMArch::GPT_OSS_20B) { auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); x = mlp->forward(ctx, x); + } else if (auto moe_mlp = std::dynamic_pointer_cast(blocks["mlp"])) { + // LLaDA2 is dense for the first first_k_dense_replace layers and MoE afterwards, + // so the block type varies per layer rather than per arch. + x = moe_mlp->forward(ctx, x); } else { auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); x = mlp->forward(ctx, x); @@ -1650,6 +1948,11 @@ namespace LLM { return x; } + ggml_tensor* embed(GGMLRunnerContext* ctx, ggml_tensor* input_ids) { + auto model = std::dynamic_pointer_cast(blocks["model"]); + return model->embed(ctx, input_ids); + } + std::shared_ptr vision_model() { GGML_ASSERT(enable_vision); return std::dynamic_pointer_cast(blocks["visual"]); @@ -1990,7 +2293,8 @@ namespace LLM { config.arch == LLMArch::GEMMA3_12B || config.arch == LLMArch::GEMMA4_12B || config.arch == LLMArch::GEMMA2_2B || - config.arch == LLMArch::GPT_OSS_20B) { + config.arch == LLMArch::GPT_OSS_20B || + config.arch == LLMArch::LLADA2_MOE) { input_pos_vec.resize(n_tokens); for (int i = 0; i < n_tokens; ++i) { input_pos_vec[i] = i; @@ -2042,8 +2346,9 @@ namespace LLM { attention_mask_vec.resize(n_tokens * n_tokens); for (int i0 = 0; i0 < n_tokens; i0++) { for (int i1 = 0; i1 < n_tokens; i1++) { + // Diffusion LLMs attend in both directions; only causal LMs get the triangle. float value = 0.f; - if (i0 > i1) { + if (!config.bidirectional && i0 > i1) { value = -INFINITY; } attention_mask_vec[i1 * n_tokens + i0] = value; @@ -2115,6 +2420,22 @@ namespace LLM { input_ids.dim() + 1); } + // LLaDA-Image's QueryFormer consumes the raw token embeddings before the backbone runs, + // so it needs the embedding lookup on its own. + sd::Tensor compute_input_embeds(const int n_threads, + const sd::Tensor& input_ids) { + auto get_graph = [&]() -> ggml_cgraph* { + ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE); + ggml_tensor* ids = make_input(input_ids); + auto runner_ctx = get_context(); + ggml_tensor* out = model.embed(&runner_ctx, ids); + ggml_build_forward_expand(gf, out); + return gf; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), + input_ids.dim() + 1); + } + int64_t get_num_image_tokens(int64_t t, int64_t h, int64_t w) { int64_t grid_t = 1; int64_t grid_h = h / config.vision.patch_size; @@ -2370,11 +2691,13 @@ namespace LLM { pad_id = 199999; } else if (arch == LLMArch::GEMMA2_2B) { pad_id = 0; + } else if (arch == LLMArch::LLADA2_MOE) { + pad_id = 156892; } tokenizer = tokenizers.create(TokenizerConfig::MAIN, model.config.vocab_size, pad_id); if (!tokenizer) { - if (arch == LLMArch::GPT_OSS_20B || arch == LLMArch::GEMMA2_2B) { - throw std::runtime_error("GPT-OSS and Gemma 2 require an external tokenizer.json in the main tokenizer slot"); + if (arch == LLMArch::GPT_OSS_20B || arch == LLMArch::GEMMA2_2B || arch == LLMArch::LLADA2_MOE) { + throw std::runtime_error("GPT-OSS, Gemma 2 and LLaDA2 require an external tokenizer.json in the main tokenizer slot"); } if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) { tokenizer = std::make_shared(); diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 9a38ceb19..319abb376 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -505,6 +505,9 @@ SDVersion ModelLoader::get_sd_version() const { if (tensor_storage.name.find("model.diffusion_model.double_blocks.0.img_mlp.gate_proj.weight") != std::string::npos) { return VERSION_OVIS_IMAGE; } + if (tensor_storage.name.find("model.diffusion_model.sigvq_embedder.1.weight") != std::string::npos) { + return VERSION_LLADA_IMAGE; + } if (tensor_storage.name.find("model.diffusion_model.cap_embedder.0.weight") != std::string::npos) { return VERSION_Z_IMAGE; } @@ -1536,6 +1539,9 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage // Pass, do not convert. For Unet } else if (contains(name, "embedding")) { // Pass, do not convert embedding + } else if (ends_with(name, "_pad_token")) { + // Pass, do not convert. LLaDA-Image stores its pad tokens far outside the f16 + // range, so any format with an f16 scale or payload turns them into inf. } else { return true; } diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index d0ff45c8e..4a7347ad6 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -103,6 +103,8 @@ std::string convert_open_clip_to_hf_clip_name(std::string name) { return name; } +std::string convert_llada2_moe_te_name(std::string name); + std::string convert_cond_stage_model_name(std::string name, std::string prefix) { static const std::vector> clip_name_map{ {"transformer.text_projection.weight", "transformer.text_model.text_projection"}, @@ -177,6 +179,7 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix) replace_with_name_map(name, llm_vision_name_map); } else { replace_with_name_map(name, llm_name_map); + name = convert_llada2_moe_te_name(name); } } else { name = convert_open_clip_to_hf_clip_name(name); @@ -749,6 +752,52 @@ std::string convert_hunyuan_video_to_original_flux(std::string name) { return name; } +// LLaDA-Image's LLaDA2-MoE text encoder. Both published layouts use these names; the ComfyUI +// GGUF repack differs only by appending ".weight" to the bare 3-D expert parameters. +// Called with the "text_encoders." prefix already stripped, so the name still carries "llm.". +std::string convert_llada2_moe_te_name(std::string name) { + static const std::vector> name_map = { + {"model.language_model.word_embeddings.", "model.embed_tokens."}, + {"model.language_model.norm.", "model.norm."}, + {"model.language_model.lm_head.", "lm_head."}, + {"model.language_model.layers.", "model.layers."}, + {"attention.query_key_value.", "self_attn.query_key_value."}, + {"attention.dense.", "self_attn.o_proj."}, + {"attention.query_layernorm.", "self_attn.q_norm."}, + {"attention.key_layernorm.", "self_attn.k_norm."}, + }; + replace_with_name_map(name, name_map); + + // The HF checkpoint stores the stacked experts as bare nn.Parameters with no ".weight". + static const std::vector bare_expert_params = { + "mlp.experts.gate_proj", + "mlp.experts.up_proj", + "mlp.experts.down_proj", + }; + for (const auto& suffix : bare_expert_params) { + if (ends_with(name, suffix)) { + name += ".weight"; + break; + } + } + + return name; +} + +// The attention projections keep their diffusers names (JointAttention's split_qkv mode), so +// only the patch-size-keyed dicts need flattening. Latents arrive already patchified from the +// Flux2 VAE, so the only patch key is 1-1. +std::string convert_diffusers_dit_to_original_llada_image(std::string name) { + static const std::vector> prefix_map = { + {"all_x_embedder.1-1.", "x_embedder."}, + {"all_final_layer.1-1.", "final_layer."}, + }; + + replace_with_prefix_map(name, prefix_map); + + return name; +} + std::string convert_diffusers_dit_to_original_lumina2(std::string name) { int num_layers = 30; int num_refiner_layers = 2; @@ -896,6 +945,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S name = convert_hunyuan_video_to_original_flux(name); } else if (sd_version_is_z_image(version)) { name = convert_diffusers_dit_to_original_lumina2(name); + } else if (sd_version_is_llada_image(version)) { + name = convert_diffusers_dit_to_original_llada_image(name); } else if (sd_version_is_anima(version)) { name = convert_other_dit_to_original_anima(name); } else if (sd_version_is_krea2(version)) { diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index ecad31a53..8b75f2fdd 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -100,6 +100,7 @@ const char* model_version_to_str[] = { "Krea2", "Mage Flow", "SenseNova U1.5", + "LLaDA-Image", "ESRGAN", }; @@ -1286,6 +1287,7 @@ bool StableDiffusionGGML::build_denoiser() { sd_version_is_anima(version) || sd_version_is_ernie_image(version) || sd_version_is_z_image(version) || + sd_version_is_llada_image(version) || sd_version_is_boogu_image(version) || sd_version_is_pid(version) || sd_version_is_ideogram4(version)) { @@ -1306,6 +1308,8 @@ bool StableDiffusionGGML::build_denoiser() { default_flow_shift = 3.16f; } else if (sd_version_is_mage_flow(version)) { default_flow_shift = 6.f; + } else if (sd_version_is_llada_image(version)) { + default_flow_shift = 1.0f; // unused: LLADA_IMAGE_SCHEDULER builds a fixed grid } else { default_flow_shift = 3.f; } @@ -2381,6 +2385,9 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + "text_encoders.llm", + "queryformer", + "text_projection", + "sigvq", + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); } else if (sd_version_is_boogu_image(version)) { result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), tensor_storage_map, diff --git a/src/pipeline/request.cpp b/src/pipeline/request.cpp index 2874bbae9..33530782e 100644 --- a/src/pipeline/request.cpp +++ b/src/pipeline/request.cpp @@ -73,6 +73,8 @@ namespace sd::pipeline { return LTX2_SCHEDULER; } else if (sd != nullptr && sd_version_is_ideogram4(sd->version)) { return LOGIT_NORMAL_SCHEDULER; + } else if (sd != nullptr && sd_version_is_llada_image(sd->version)) { + return LLADA_IMAGE_SCHEDULER; } return DISCRETE_SCHEDULER; } diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index 6e6f22d52..416859e6d 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -786,6 +786,54 @@ struct FluxScheduler : SigmaScheduler { }; // https://github.com/black-forest-labs/flux2/blob/main/src/flux2/sampling.py#L244 +// LLaDA-Image does not use a shift-based flow schedule. The reference pipeline builds a +// Kumaraswamy-shaped grid over t = linspace(0.001, 1, n + 1)[:-1]: +// schedule = (1 - (1 - t^1.17)^0.8)^1.1 +// sigma = 1 - schedule +// Its scheduler config can also set use_uniform_sigmas, which replaces the whole curve with a +// plain linspace(1, 0, n + 1)[:-1] pre-shift grid. +struct LLaDAImageScheduler : SigmaScheduler { + bool uniform_sigmas = false; + + explicit LLaDAImageScheduler(const char* extra_sample_args = nullptr) { + parse_extra_sample_args(extra_sample_args); + } + + void parse_extra_sample_args(const char* extra_sample_args) { + for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "llada_image scheduler arg")) { + if (key == "uniform") { + if (!parse_strict_bool(value, uniform_sigmas)) { + LOG_WARN("ignoring invalid llada_image scheduler arg '%s=%s'", key.c_str(), value.c_str()); + } + } + } + } + + std::vector get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override { + std::vector sigmas; + sigmas.reserve(n + 1); + + if (n == 0) { + sigmas.push_back(1.0f); + return sigmas; + } + + for (uint32_t i = 0; i < n; ++i) { + float progress = static_cast(i) / static_cast(n); + if (uniform_sigmas) { + sigmas.push_back(1.0f - progress); + } else { + float t = 0.001f + progress * (1.0f - 0.001f); + float schedule = powf(1.0f - powf(1.0f - powf(t, 1.17f), 0.8f), 1.1f); + sigmas.push_back(1.0f - schedule); + } + } + + sigmas.push_back(0.0f); + return sigmas; + } +}; + struct Flux2Scheduler : SigmaScheduler { int image_seq_len = 0; @@ -1123,6 +1171,11 @@ struct Denoiser { scheduler = std::make_shared(image_seq_len); break; } + case LLADA_IMAGE_SCHEDULER: { + LOG_INFO("get_sigmas with LLaDA-Image scheduler"); + scheduler = std::make_shared(extra_sample_args); + break; + } case FLUX_SCHEDULER: { LOG_INFO("get_sigmas with Flux scheduler"); scheduler = std::make_shared(image_seq_len, extra_sample_args); diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 5a56f4029..d12b57a8e 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -124,6 +124,7 @@ const char* scheduler_to_str[] = { "flux2", "flux", "beta", + "llada_image", }; static_assert(SCHEDULER_COUNT == sizeof(scheduler_to_str) / sizeof(scheduler_to_str[0]),