feat: add LLaDA-Image support (#1968)

Co-authored-by: leejet <leejet714@gmail.com>

README.md

@@ -63,12 +63,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)

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.

docs/llada_image.md

@@ -0,0 +1,156 @@
+# 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
+```
+
+<img width="256" alt="LLaDA-Image example" src="../assets/llada_image/example.png" />
+
+### 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
+```
+
+<img width="256" alt="LLaDA-Image edit example" src="../assets/llada_image/edit_example.png" />
+
+See [edit.md](./edit.md) for the shared reference-image options. LLaDA-Image uses the
+`llada_image` preset by default, resizing the reference image to the output width and height
+before VAE encoding. SigVQ uses bilinear resizing to half the output resolution and inputs
+normalized to `[-1, 1]`. CFG keeps the source latent in both branches and uses SigVQ features
+only in the positive branch. Editing requires connectors that include the SigVQ weights.
+
+## 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.

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},
         {"",

include/stable-diffusion.h

@@ -79,6 +79,7 @@ enum scheduler_t {
     FLUX2_SCHEDULER,
     FLUX_SCHEDULER,
     BETA_SCHEDULER,
+    LLADA_IMAGE_SCHEDULER,
     SCHEDULER_COUNT
 };
 

src/conditioning/conditioner.hpp

@@ -14,6 +14,7 @@
 #include "core/util.h"
 #include "model/diffusion/model.hpp"
 #include "model/te/clip.hpp"
+#include "model/te/llada_image_te.hpp"
 #include "model/te/llm.hpp"
 #include "model/te/t5.hpp"
 #include "model_loader.h"
@@ -3236,6 +3237,214 @@ struct LTXAVTextProjectionRunner : public GGMLRunner {
     }
 };
 
+// LLaDA-Image's text path is a three-stage pipeline rather than a single encoder pass:
+// the token embeddings feed a QueryFormer whose 256 queries are appended to the backbone
+// input, and the backbone's final hidden states are projected to the denoiser's caption dim.
+// Ref: LLaDAImagePipeline._encode_text.
+struct LLaDAImageEmbedder : public Conditioner {
+    std::shared_ptr<Tokenizer> tokenizer;
+    std::shared_ptr<LLM::LLMRunner> llm;
+    std::shared_ptr<LLaDAImageTE::QueryFormerRunner> query_former;
+    std::shared_ptr<LLaDAImageTE::TextProjectionRunner> text_projection;
+    std::shared_ptr<LLaDAImageTE::SigVQRunner> 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<RunnerWeightManager> 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::LLMRunner>(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<LLaDAImageTE::QueryFormerRunner>(backend,
+                                                                         tensor_storage_map,
+                                                                         query_former_prefix,
+                                                                         weight_manager);
+        text_projection = std::make_shared<LLaDAImageTE::TextProjectionRunner>(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<LLaDAImageTE::SigVQRunner>(backend,
+                                                                    tensor_storage_map,
+                                                                    sigvq_prefix,
+                                                                    weight_manager);
+                break;
+            }
+        }
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& 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<ggml_tensor*, enum ggml_op>& 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<ggml_backend_t>& 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<size_t>& limits) override {
+        llm->set_graph_cut_layer_split_backend_vram_limits(limits);
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        llm->get_param_tensors(tensors, llm_prefix);
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        llm->set_weight_adapter(adapter);
+        query_former->set_weight_adapter(adapter);
+        text_projection->set_weight_adapter(adapter);
+        if (sigvq != nullptr) {
+            sigvq->set_weight_adapter(adapter);
+        }
+    }
+
+    void runner_end() override {
+        llm->runner_end();
+        query_former->runner_end();
+        text_projection->runner_end();
+        if (sigvq != nullptr) {
+            sigvq->runner_end();
+        }
+    }
+
+    SDCondition get_learned_condition(int n_threads,
+                                      const ConditionerParams& conditioner_params) override {
+        const int64_t num_queries = 256;
+        const bool has_ref_images = conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty();
+        if (has_ref_images && sigvq == nullptr) {
+            LOG_ERROR("LLaDA-Image editing requires connectors with SigVQ weights");
+            return {};
+        }
+
+        std::string text = conditioner_params.text;
+        while (!text.empty() && std::isspace(static_cast<unsigned char>(text.front()))) {
+            text.erase(text.begin());
+        }
+        while (!text.empty() && std::isspace(static_cast<unsigned char>(text.back()))) {
+            text.pop_back();
+        }
+        std::string prompt = text.empty()
+                                 ? "<role>HUMAN</role> Generate an image.\n<role>ASSISTANT</role>\n<IMAGE1>"
+                                 : "<role>HUMAN</role> Generate an image: " + text + "\n<role>ASSISTANT</role>\n<IMAGE1>";
+
+        std::vector<int> tokens;
+        if (!tokenizer->encode(prompt, tokens, nullptr)) {
+            return {};
+        }
+        int64_t n_text = static_cast<int64_t>(tokens.size());
+        GGML_ASSERT(n_text > 0);
+
+        sd::Tensor<int32_t> text_ids({n_text}, std::vector<int32_t>(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<int32_t> padded(tokens.begin(), tokens.end());
+        padded.resize(static_cast<size_t>(n_text + num_queries), tokenizer->PAD_TOKEN_ID);
+        int64_t n_total = static_cast<int64_t>(padded.size());
+        sd::Tensor<int32_t> 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<float>::lowest() / 4.0f;
+        sd::Tensor<float> 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<int>(n_text), query_embeds);
+
+        std::set<int> out_layers = {static_cast<int>(llm->config.num_layers) + 1};
+        auto hidden_states       = llm->compute(n_threads,
+                                                input_ids,
+                                                attention_mask,
+                                                image_embeds,
+                                                out_layers);
+
+        SDCondition result;
+        result.c_crossattn = text_projection->compute(n_threads, hidden_states);
+
+        // Editing: SigVQ sees the reference at half the output resolution, as in
+        // LLaDAImagePipeline._encode_source_image.
+        if (has_ref_images) {
+            const auto& ref = conditioner_params.ref_images->front();
+            auto resized    = sd::ops::interpolate(ref,
+                                                   {conditioner_params.width / 2,
+                                                    conditioner_params.height / 2,
+                                                    ref.shape()[2],
+                                                    ref.shape()[3]},
+                                                   sd::ops::InterpolateMode::Bilinear);
+            resized         = resized * 2.f - 1.f;
+            auto semantic   = sigvq->compute(n_threads, resized);
+            if (semantic.empty()) {
+                return {};
+            }
+            result.extra_c_crossattns.push_back(std::move(semantic));
+        }
+        return result;
+    }
+};
+
 struct LTXAVEmbedder : public Conditioner {
     static constexpr int64_t kHiddenSize = 3840;
     static constexpr int64_t kNumStates  = 49;

src/model.h

@@ -60,6 +60,7 @@ enum SDVersion {
     VERSION_KREA2,
     VERSION_MAGE_FLOW,
     VERSION_SENSENOVA_U1_5,
+    VERSION_LLADA_IMAGE,
     VERSION_ESRGAN,
     VERSION_COUNT,
 };
@@ -173,6 +174,13 @@ static inline bool sd_version_is_z_image(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_llada_image(SDVersion version) {
+    if (version == VERSION_LLADA_IMAGE) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_boogu_image(SDVersion version) {
     if (version == VERSION_BOOGU_IMAGE) {
         return true;
@@ -252,7 +260,7 @@ static inline bool sd_version_uses_flux_vae(SDVersion version) {
 }
 
 static inline bool sd_version_uses_flux2_vae(SDVersion version) {
-    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version)) {
+    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version) || sd_version_is_llada_image(version)) {
         return true;
     }
     return false;
@@ -293,6 +301,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) ||

src/model/common/ggml_block.hpp

@@ -839,21 +839,30 @@ class RMSNorm : public UnaryBlock {
 protected:
     int64_t hidden_size;
     float eps;
+    bool elementwise_affine;
     std::string prefix;
 
     void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
-        this->prefix         = prefix;
+        this->prefix = prefix;
+        if (!elementwise_affine) {
+            return;
+        }
         enum ggml_type wtype = GGML_TYPE_F32;
         params["weight"]     = ggml_new_tensor_1d(ctx, wtype, hidden_size);
     }
 
 public:
     RMSNorm(int64_t hidden_size,
-            float eps = 1e-06f)
+            float eps               = 1e-06f,
+            bool elementwise_affine = true)
         : hidden_size(hidden_size),
-          eps(eps) {}
+          eps(eps),
+          elementwise_affine(elementwise_affine) {}
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+        if (!elementwise_affine) {
+            return ggml_rms_norm(ctx->ggml_ctx, x, eps);
+        }
         ggml_tensor* w = params["weight"];
         if (ctx->weight_adapter) {
             w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");

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<std::vector<float>> 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<std::vector<float>>(bs * padded_context_len, std::vector<float>(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<int>(img_ids.size() / bs), seq_multi_of);
+        if (img_pad_len > 0) {
+            std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(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<std::vector<float>> 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<std::vector<float>> cap_ids;
+        std::vector<int> cap_end_positions;
+        int cursor = 1;
+        for (int copy = 0; copy < 2; ++copy) {
+            for (int i = 0; i < padded_context; ++i) {
+                std::vector<float> id(3, 0.f);
+                if (i < context_len) {
+                    id[0] = static_cast<float>(cursor + i);
+                }
+                cap_ids.push_back(id);
+            }
+            cursor += context_len;
+            cap_end_positions.push_back(cursor);
+            cursor += 2;
+        }
+
+        std::vector<std::vector<float>> 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<float>(3, 0.f));
+        }
+
+        const int sigvq_start = static_cast<int>(cap_ids.size() + img_ids.size()) + 1;
+        std::vector<std::vector<float>> sigvq_ids;
+        for (int i = 0; i < sigvq_len + sigvq_pad; ++i) {
+            std::vector<float> id(3, 0.f);
+            if (i < sigvq_len) {
+                id[0] = static_cast<float>(sigvq_start + i);
+            }
+            sigvq_ids.push_back(id);
+        }
+
+        std::vector<std::vector<float>> 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<float> 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<int>& 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<float>(theta), axes_dim, {});
+    }
+
+    __STATIC_INLINE__ std::vector<float> 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<int>& axes_dim) {
+        std::vector<std::vector<float>> ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of);
+        std::vector<std::vector<int>> 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<int>(pos_len, 0));
+                size_t cursor     = context_len + bound_mod(context_len, seq_multi_of);
+                size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(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<float>(theta), axes_dim, wrap_dims);
+    }
+
     // Generate z_image positional embeddings
     __STATIC_INLINE__ std::vector<float> gen_z_image_pe(int h,
                                                         int w,

src/model/diffusion/llada_image.hpp

@@ -0,0 +1,527 @@
+#ifndef __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__
+#define __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__
+
+#include <algorithm>
+#include <cinttypes>
+
+#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<int> 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<int64_t>(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<ZImage::JointTransformerBlock> make_block(bool modulation) {
+            return std::make_shared<ZImage::JointTransformerBlock>(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<Linear>(config.patch_size * config.patch_size * config.in_channels, config.hidden_size);
+            blocks["t_embedder"] = std::make_shared<TimestepEmbedder>(MIN(config.hidden_size, 1024), 256, ZImage::ADALN_EMBED_DIM);
+
+            blocks["cap_embedder.0"] = std::make_shared<RMSNorm>(config.cap_feat_dim, config.norm_eps, false);
+            blocks["cap_embedder.1"] = std::make_shared<Linear>(config.cap_feat_dim, config.hidden_size);
+
+            blocks["semantic_embedder.0"] = std::make_shared<RMSNorm>(config.semantic_feat_dim, config.norm_eps, false);
+            blocks["semantic_embedder.1"] = std::make_shared<Linear>(config.semantic_feat_dim, config.hidden_size);
+            blocks["sigvq_embedder.0"]    = std::make_shared<RMSNorm>(config.semantic_feat_dim, config.norm_eps, false);
+            blocks["sigvq_embedder.1"]    = std::make_shared<Linear>(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<ZImage::FinalLayer>(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<Linear>(blocks["x_embedder"]);
+            auto t_embedder     = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
+            auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
+            auto cap_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"]);
+            auto final_layer    = std::dynamic_pointer_cast<ZImage::FinalLayer>(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<int>(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<int>(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<ZImage::JointTransformerBlock>(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<ZImage::JointTransformerBlock>(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<ZImage::JointTransformerBlock>(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<int>(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<Linear>(blocks["x_embedder"]);
+            auto t_embedder     = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
+            auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
+            auto cap_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"]);
+            auto sigvq_embed_0  = std::dynamic_pointer_cast<RMSNorm>(blocks["sigvq_embedder.0"]);
+            auto sigvq_embed_1  = std::dynamic_pointer_cast<Linear>(blocks["sigvq_embedder.1"]);
+            auto final_layer    = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"]);
+
+            auto t_noisy = t_embedder->forward(ctx, timestep);
+            auto t_clean = t_embedder->forward(ctx, ggml_scale(gctx, timestep, 0.f));
+
+            auto per_token = [&](ggml_tensor* emb, int64_t n) {
+                return ggml_repeat_4d(gctx, emb, emb->ne[0], n, 1, 1);
+            };
+
+            auto cap        = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context));
+            cap             = pad_stream(ctx, cap, params["cap_pad_token"]);
+            int64_t cap_len = cap->ne[1];
+            cap             = ggml_concat(gctx, cap, cap, 1);
+
+            auto src            = pad_stream(ctx, x_embedder->forward(ctx, source_latent), params["x_pad_token"]);
+            auto tgt_embed      = x_embedder->forward(ctx, x);
+            int64_t n_img_token = tgt_embed->ne[1];
+            auto tgt            = pad_stream(ctx, tgt_embed, params["x_pad_token"]);
+            int64_t img_len     = tgt->ne[1];
+            auto img            = ggml_concat(gctx, src, tgt, 1);
+
+            ggml_tensor* sig = nullptr;
+            int64_t sig_len  = 0;
+            if (semantic != nullptr) {
+                sig     = sigvq_embed_1->forward(ctx, sigvq_embed_0->forward(ctx, semantic));
+                sig     = pad_stream(ctx, sig, params["sigvq_pad_token"]);
+                sig_len = sig->ne[1];
+            }
+
+            GGML_ASSERT(cap_len * 2 + img_len * 2 + sig_len == pe->ne[3]);
+
+            auto cap_pe = ggml_ext_slice(gctx, pe, 3, 0, cap_len * 2);
+            auto img_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2);
+
+            auto img_adaln = ggml_concat(gctx, per_token(t_clean, img_len), per_token(t_noisy, img_len), 1);
+
+            for (int i = 0; i < config.num_refiner_layers; i++) {
+                auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(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<ZImage::JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
+                img        = block->forward(ctx, img, img_pe, nullptr, img_adaln);
+            }
+            if (sig != nullptr) {
+                auto sig_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]);
+                for (int i = 0; i < config.num_refiner_layers; i++) {
+                    auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["sigvq_refiner." + std::to_string(i)]);
+                    sig        = block->forward(ctx, sig, sig_pe, nullptr, nullptr);
+                }
+            }
+
+            auto seq = ggml_concat(gctx, cap, img, 1);
+
+            auto cap_adaln = ggml_concat(gctx, per_token(t_clean, cap_len), per_token(t_noisy, cap_len), 1);
+            auto seq_adaln = ggml_concat(gctx, cap_adaln, img_adaln, 1);
+            if (sig != nullptr) {
+                seq       = ggml_concat(gctx, seq, sig, 1);
+                seq_adaln = ggml_concat(gctx, seq_adaln, per_token(t_clean, sig_len), 1);
+            }
+
+            for (int i = 0; i < config.num_layers; i++) {
+                auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(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<float> pe_vec;
+
+        LLaDAImageRunner(ggml_backend_t backend,
+                         const String2TensorStorage& tensor_storage_map      = {},
+                         const std::string prefix                            = "",
+                         std::shared_ptr<RunnerWeightManager> 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<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            llada_image.get_param_tensors(tensors, prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& 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<int>(x->ne[1]),
+                                                   static_cast<int>(x->ne[0]),
+                                                   config.patch_size,
+                                                   static_cast<int>(x->ne[3]),
+                                                   static_cast<int>(context->ne[1]),
+                                                   ZImage::SEQ_MULTI_OF,
+                                                   config.theta,
+                                                   circular_y_enabled,
+                                                   circular_x_enabled,
+                                                   config.axes_dim);
+            int pos_len = static_cast<int>(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<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& 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<float>& x_tensor,
+                                      const sd::Tensor<float>& timesteps_tensor,
+                                      const sd::Tensor<float>& context_tensor,
+                                      const sd::Tensor<float>& semantic_tensor,
+                                      const sd::Tensor<float>& source_tensor) {
+            ggml_cgraph* gf        = new_graph_custom(LLADA_IMAGE_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
+            ggml_tensor* context   = make_input(context_tensor);
+            ggml_tensor* semantic  = make_optional_input(semantic_tensor);
+            ggml_tensor* source    = make_input(source_tensor);
+            GGML_ASSERT(x->ne[3] == 1);
+
+            pe_vec      = Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
+                                                        static_cast<int>(x->ne[0]),
+                                                        config.patch_size,
+                                                        static_cast<int>(context->ne[1]),
+                                                   semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
+                                                        ZImage::SEQ_MULTI_OF,
+                                                        config.theta,
+                                                        config.axes_dim);
+            int pos_len = static_cast<int>(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<float> 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<LLaDAImageDiffusionExtra>(&diffusion_params.extra);
+            bool has_semantic   = extra != nullptr && extra->semantic != nullptr && !extra->semantic->empty();
+            bool has_ref_latent = diffusion_params.ref_latents != nullptr && !diffusion_params.ref_latents->empty();
+            if (has_semantic && !has_ref_latent) {
+                LOG_WARN("llada_image: SigVQ features without a reference latent are not supported; falling back to text to image");
+            }
+            if (has_ref_latent) {
+                const auto& source = diffusion_params.ref_latents->front();
+                if (source.shape() != diffusion_params.x->shape()) {
+                    LOG_ERROR("llada_image: reference latent must match the target shape; use resize_vae_to_target=1");
+                    return {};
+                }
+                auto get_graph = [&]() -> ggml_cgraph* {
+                    return build_edit_graph(*diffusion_params.x,
+                                            *diffusion_params.timesteps,
+                                            tensor_or_empty(diffusion_params.context),
+                                            tensor_or_empty(extra != nullptr ? extra->semantic : nullptr),
+                                            source);
+                };
+                return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false),
+                                                       diffusion_params.x->dim());
+            }
+
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context));
+        }
+    };
+
+}  // namespace LLaDAImage
+
+#endif  // __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__

src/model/diffusion/model.hpp

@@ -39,6 +39,9 @@ const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
     {"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
     {"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}},
     {"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}},
+    // pass_to_vlm routes the reference image to the conditioner, which is where LLaDA-Image's
+    // SigVQ encoder lives; it does its own half-resolution resize.
+    {"llada_image", {true, true, Rope::RefIndexMode::FIXED, true, true, -1, RefImageResizeMode::NONE, -1, -1, true}},
     {"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}},
 };
 
@@ -131,6 +134,11 @@ struct HunyuanVideoDiffusionExtra {
     const sd::Tensor<float>* timestep_r = nullptr;
 };
 
+struct LLaDAImageDiffusionExtra {
+    // SigVQ semantic features of the reference image; present only in editing mode.
+    const sd::Tensor<float>* semantic = nullptr;
+};
+
 using DiffusionExtraParams = std::variant<std::monostate,
                                           UNetDiffusionExtra,
                                           SkipLayerDiffusionExtra,
@@ -143,7 +151,8 @@ using DiffusionExtraParams = std::variant<std::monostate,
                                           MiniMaxH3DiffusionExtra,
                                           MiniT2IDiffusionExtra,
                                           SenseNovaU1DiffusionExtra,
-                                          HunyuanVideoDiffusionExtra>;
+                                          HunyuanVideoDiffusionExtra,
+                                          LLaDAImageDiffusionExtra>;
 
 struct DiffusionParams {
     const sd::Tensor<float>* x                        = nullptr;

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<Linear>(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false);
-            float scale   = 1.f;
-            blocks["out"] = std::make_shared<Linear>(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<Linear>(hidden_size, num_heads * head_dim, false);
+                blocks["to_k"]     = std::make_shared<Linear>(hidden_size, num_kv_heads * head_dim, false);
+                blocks["to_v"]     = std::make_shared<Linear>(hidden_size, num_kv_heads * head_dim, false);
+                blocks["to_out.0"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
+            } else {
+                blocks["qkv"] = std::make_shared<Linear>(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false);
+                blocks["out"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
+            }
             if (qk_norm) {
-                blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim);
-                blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim);
+                blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-06f, norm_elementwise_affine);
+                blocks["k_norm"] = std::make_shared<RMSNorm>(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<Linear>(blocks["qkv"]);
-            auto out_proj   = std::dynamic_pointer_cast<Linear>(blocks["out"]);
+            auto out_proj   = std::dynamic_pointer_cast<Linear>(blocks[split_qkv ? "to_out.0" : "out"]);
+
+            if (split_qkv) {
+                auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
+                auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
+                auto v_proj = std::dynamic_pointer_cast<Linear>(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<RMSNorm>(blocks["q_norm"])->forward(ctx, q);
+                    k = std::dynamic_pointer_cast<RMSNorm>(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<Linear>(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<JointAttention>(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm);
+            blocks["attention"]       = std::make_shared<JointAttention>(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm, norm_elementwise_affine, split_qkv);
             blocks["feed_forward"]    = std::make_shared<FeedForward>(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier);
-            blocks["attention_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
-            blocks["ffn_norm1"]       = std::make_shared<RMSNorm>(hidden_size, norm_eps);
-            blocks["attention_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
-            blocks["ffn_norm2"]       = std::make_shared<RMSNorm>(hidden_size, norm_eps);
+            blocks["attention_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
+            blocks["ffn_norm1"]       = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
+            blocks["attention_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
+            blocks["ffn_norm2"]       = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
             if (modulation) {
                 blocks["adaLN_modulation.0"] = std::make_shared<Linear>(MIN(hidden_size, ADALN_EMBED_DIM), 4 * hidden_size);
             }

src/model/te/llada_image_te.hpp

@@ -0,0 +1,604 @@
+#ifndef __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__
+#define __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__
+
+#include <algorithm>
+#include <array>
+#include <cmath>
+
+#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<Linear>(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<Linear>(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<LayerNorm>(config.hidden_size, config.norm_eps, false);
+            blocks["norm_k"]     = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
+            blocks["cross_attn"] = std::make_shared<QueryAttention>(config.hidden_size, config.num_heads);
+            blocks["norm1"]      = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
+            blocks["mlp.fc1"]    = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, true);
+            blocks["mlp.fc2"]    = std::make_shared<Linear>(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<LayerNorm>(blocks["norm_q"]);
+            auto norm_k     = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_k"]);
+            auto cross_attn = std::dynamic_pointer_cast<QueryAttention>(blocks["cross_attn"]);
+            auto norm1      = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
+            auto fc1        = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
+            auto fc2        = std::dynamic_pointer_cast<Linear>(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<QueryFormerBlock>(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<QueryFormerBlock>(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<Linear>(config.hidden_size, config.hidden_size, true);
+            blocks["k_proj"]   = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
+            blocks["v_proj"]   = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
+            blocks["out_proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
+            blocks["q_norm"]   = std::make_shared<RMSNorm>(head_dim, config.norm_eps, false);
+            blocks["k_norm"]   = std::make_shared<RMSNorm>(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<Linear>(blocks["q_proj"]);
+            auto k_proj   = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
+            auto v_proj   = std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
+            auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
+            auto q_norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
+            auto k_norm   = std::dynamic_pointer_cast<RMSNorm>(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<TextProjectionAttention>(config);
+            blocks["layer_norm1"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps, false);
+            blocks["layer_norm2"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps, false);
+            blocks["mlp.fc1"]     = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, true);
+            blocks["mlp.fc2"]     = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto self_attn   = std::dynamic_pointer_cast<TextProjectionAttention>(blocks["self_attn"]);
+            auto layer_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["layer_norm1"]);
+            auto layer_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["layer_norm2"]);
+            auto fc1         = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
+            auto fc2         = std::dynamic_pointer_cast<Linear>(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<TextProjectionBlock>(config);
+            }
+            blocks["projector"] = std::make_shared<Linear>(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<TextProjectionBlock>(blocks["layers." + std::to_string(i)]);
+                x          = block->forward(ctx, x);
+            }
+            auto projector = std::dynamic_pointer_cast<Linear>(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<Linear>(config.hidden_size, config.hidden_size * 3, true);
+            blocks["proj"] = std::make_shared<Linear>(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<Linear>(blocks["qkv"]);
+            auto out_proj      = std::dynamic_pointer_cast<Linear>(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<LayerNorm>(config.hidden_size, config.norm_eps);
+            blocks["norm2"]   = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps);
+            blocks["attn"]    = std::make_shared<SigVQAttention>(config);
+            blocks["mlp.fc1"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, true);
+            blocks["mlp.fc2"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
+            auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
+            auto attn  = std::dynamic_pointer_cast<SigVQAttention>(blocks["attn"]);
+            auto fc1   = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
+            auto fc2   = std::dynamic_pointer_cast<Linear>(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<Conv2d>(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<SigVQBlock>(config);
+            }
+            blocks["vqmodel.quant_conv"]         = std::make_shared<Conv2d>(config.hidden_size,
+                                                                    config.codebook_embed_dim,
+                                                                    std::make_pair(1, 1));
+            blocks["prior_projector.net.0.proj"] = std::make_shared<Linear>(config.semantic_embed_dim, config.semantic_embed_dim, true);
+            blocks["prior_projector.net.2"]      = std::make_shared<Linear>(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<ggml_tensor*>& pos_idx,
+                             const std::vector<ggml_tensor*>& 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<Conv2d>(blocks["visual.patch_embed.proj"]);
+            auto quant_conv  = std::dynamic_pointer_cast<Conv2d>(blocks["vqmodel.quant_conv"]);
+            auto proj_0      = std::dynamic_pointer_cast<Linear>(blocks["prior_projector.net.0.proj"]);
+            auto proj_2      = std::dynamic_pointer_cast<Linear>(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<SigVQBlock>(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<RunnerWeightManager> 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<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
+            query_former.get_param_tensors(tensors, prefix);
+        }
+
+        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& 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<RunnerWeightManager> 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<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
+            text_projection.get_param_tensors(tensors, prefix);
+        }
+
+        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& 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<std::vector<int32_t>, 4> pos_idx_data;
+        std::array<std::vector<float>, 4> pos_weight_data;
+
+        SigVQRunner(ggml_backend_t backend,
+                    const String2TensorStorage& tensor_storage_map      = {},
+                    const std::string prefix                            = "",
+                    std::shared_ptr<RunnerWeightManager> 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<std::string, ggml_tensor*>& 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<int32_t>(std::min<int64_t>(std::max<int64_t>(v, 0), side - 1));
+            };
+
+            for (int64_t i = 0; i < grid_h; ++i) {
+                double src_h    = (static_cast<double>(i) + 0.5) * side / static_cast<double>(grid_h) - 0.5;
+                int64_t h_floor = static_cast<int64_t>(std::floor(src_h));
+                double dh       = src_h - static_cast<double>(h_floor);
+                for (int64_t j = 0; j < grid_w; ++j) {
+                    double src_w    = (static_cast<double>(j) + 0.5) * side / static_cast<double>(grid_w) - 0.5;
+                    int64_t w_floor = static_cast<int64_t>(std::floor(src_w));
+                    double dw       = src_w - static_cast<double>(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<int32_t>(side) + w0);
+                    pos_idx_data[1].push_back(h0 * static_cast<int32_t>(side) + w1);
+                    pos_idx_data[2].push_back(h1 * static_cast<int32_t>(side) + w0);
+                    pos_idx_data[3].push_back(h1 * static_cast<int32_t>(side) + w1);
+
+                    pos_weight_data[0].push_back(static_cast<float>((1.0 - dh) * (1.0 - dw)));
+                    pos_weight_data[1].push_back(static_cast<float>((1.0 - dh) * dw));
+                    pos_weight_data[2].push_back(static_cast<float>(dh * (1.0 - dw)));
+                    pos_weight_data[3].push_back(static_cast<float>(dh * dw));
+                }
+            }
+        }
+
+        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& 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<ggml_tensor*> pos_idx;
+                std::vector<ggml_tensor*> pos_weight;
+                for (int i = 0; i < 4; i++) {
+                    auto idx = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, static_cast<int64_t>(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<int64_t>(pos_weight_data[i].size()));
+                    set_backend_tensor_data(w, pos_weight_data[i].data());
+                    pos_idx.push_back(idx);
+                    pos_weight.push_back(w);
+                }
+
+                auto runner_ctx  = get_context();
+                ggml_tensor* out = sigvq.forward(&runner_ctx, x, pos_idx, pos_weight);
+                ggml_build_forward_expand(gf, out);
+                return gf;
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), 3);
+        }
+    };
+
+}  // namespace LLaDAImageTE
+
+#endif  // __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__

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<int> 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<MLP>(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<MLP>(shared_it->second);
+                out                 = ggml_add(gctx, out, shared_experts->forward(ctx, identity));
+            }
+
+            return out;
+        }
+    };
+
     struct GPTOSSMLP : public GGMLBlock {
     protected:
         int64_t hidden_size;
@@ -605,21 +832,31 @@ namespace LLM {
             }
             txt_token_end = image_embeds[i].first;
 
-            auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
-            if (input_embed == nullptr) {
-                input_embed = txt_embed;
-            } else {
-                input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1);
+            // An embed can sit flush against the previous one or at the very start/end of the
+            // sequence, leaving no text tokens to splice around it.
+            if (txt_token_end > txt_token_start) {
+                auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
+                if (input_embed == nullptr) {
+                    input_embed = txt_embed;
+                } else {
+                    input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1);
+                }
             }
 
-            input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1);
+            if (input_embed == nullptr) {
+                input_embed = image_embeds[i].second;
+            } else {
+                input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1);
+            }
         }
 
         txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1];
         txt_token_end   = raw_x->ne[1];
 
-        auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
-        input_embed          = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
+        if (txt_token_end > txt_token_start) {
+            auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
+            input_embed          = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
+        }
         GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]);
         return input_embed;
     }
@@ -1122,6 +1359,7 @@ namespace LLM {
         bool k_eq_v;
         bool v_norm;
         bool unscaled_attention;
+        bool qkv_fused;
         float rms_norm_eps;
         int rope_pairs;
 
@@ -1147,12 +1385,20 @@ namespace LLM {
               k_eq_v(global_layer && config.global_k_eq_v),
               v_norm(config.v_norm),
               unscaled_attention(config.unscaled_attention),
+              qkv_fused(config.qkv_fused),
               rms_norm_eps(config.rms_norm_eps),
               rope_pairs(0) {
-            blocks["q_proj"] = std::make_shared<Linear>(config.hidden_size, num_heads * head_dim, config.qkv_bias);
-            blocks["k_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
-            if (!k_eq_v) {
-                blocks["v_proj"] = std::make_shared<Linear>(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<Linear>(config.hidden_size, (num_heads + num_kv_heads * 2) * head_dim, config.qkv_bias);
+            } else {
+                blocks["q_proj"] = std::make_shared<Linear>(config.hidden_size, num_heads * head_dim, config.qkv_bias);
+                blocks["k_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
+                if (!k_eq_v) {
+                    blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
+                }
             }
             blocks["o_proj"] = std::make_shared<Linear>(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<int>(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<Linear>(blocks["q_proj"]);
-            auto k_proj     = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
-            auto v_proj     = k_eq_v ? nullptr : std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
             auto out_proj   = std::dynamic_pointer_cast<Linear>(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<Linear>(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<Linear>(blocks["q_proj"]);
+                auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
+                auto v_proj = k_eq_v ? nullptr : std::dynamic_pointer_cast<Linear>(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<int>(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<int>(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<Attention>(config, sliding_attention == 0);
             if (config.arch == LLMArch::GPT_OSS_20B) {
                 blocks["mlp"] = std::make_shared<GPTOSSMLP>(config);
+            } else if (config.arch == LLMArch::LLADA2_MOE && layer_index >= config.first_k_dense_replace) {
+                blocks["mlp"] = std::make_shared<LLaDA2MoEMLP>(config);
             } else {
                 blocks["mlp"] = std::make_shared<MLP>(config.hidden_size,
                                                       config.intermediate_size,
@@ -1485,6 +1779,10 @@ namespace LLM {
             if (arch == LLMArch::GPT_OSS_20B) {
                 auto mlp = std::dynamic_pointer_cast<GPTOSSMLP>(blocks["mlp"]);
                 x        = mlp->forward(ctx, x);
+            } else if (auto moe_mlp = std::dynamic_pointer_cast<LLaDA2MoEMLP>(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<MLP>(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<TextModel>(blocks["model"]);
+            return model->embed(ctx, input_ids);
+        }
+
         std::shared_ptr<VisionModel> vision_model() {
             GGML_ASSERT(enable_vision);
             return std::dynamic_pointer_cast<VisionModel>(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<float> compute_input_embeds(const int n_threads,
+                                               const sd::Tensor<int32_t>& 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<MistralTokenizer>();

src/model_loader.cpp

@@ -514,6 +514,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;
         }

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<std::pair<std::string, std::string>> 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<std::pair<std::string, std::string>> 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<std::string> 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<std::pair<std::string, std::string>> 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)) {

src/pipeline/diffusion_engine.cpp

@@ -101,6 +101,7 @@ const char* model_version_to_str[] = {
     "Krea2",
     "Mage Flow",
     "SenseNova U1.5",
+    "LLaDA-Image",
     "ESRGAN",
 };
 
@@ -1298,6 +1299,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)) {
@@ -1318,6 +1320,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;
             }
@@ -2399,6 +2403,9 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
                     condition.c_token_types.empty() ? nullptr : &condition.c_token_types,
                     condition.c_vinput_mask.empty() ? nullptr : &condition.c_vinput_mask,
                     condition.c_image_embeds.empty() ? nullptr : &condition.c_image_embeds};
+            } else if (sd_version_is_llada_image(version)) {
+                diffusion_params.extra = LLaDAImageDiffusionExtra{
+                    condition.extra_c_crossattns.empty() ? nullptr : &condition.extra_c_crossattns[0]};
             } else if (sd_version_is_minimax_h3(version)) {
                 diffusion_params.extra = MiniMaxH3DiffusionExtra{
                     condition.c_token_types.empty() ? nullptr : &condition.c_token_types,
@@ -2778,6 +2785,8 @@ std::string StableDiffusionGGML::get_default_ref_image_preset(SDVersion version)
         return "mage_flow";
     } else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) {
         return "z_image_omni";
+    } else if (sd_version_is_llada_image(version)) {
+        return "llada_image";
     } else if (sd_version_is_krea2(version)) {
         // have to make a choice between "krea2_edit" mode (for lbouaraba/krea2edit)
         // and "krea2_ostris_edit" (for krea2 ostris edit)

src/pipeline/image.cpp

@@ -475,8 +475,12 @@ namespace sd::pipeline {
                 }
                 condition_params.text            = request->negative_prompt;
                 condition_params.zero_out_masked = zero_out_masked;
-                uncond                           = sd->cond_stage_model->get_learned_condition(sd->n_threads,
-                                                                                               condition_params);
+                if (sd_version_is_llada_image(sd->version)) {
+                    // LLaDA-Image CFG keeps the source latent but drops its SigVQ features.
+                    condition_params.ref_images = nullptr;
+                }
+                uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
+                                                                     condition_params);
                 if (uncond.empty()) {
                     LOG_ERROR("failed to encode negative prompt");
                     return std::nullopt;

src/pipeline/model_builders.cpp

@@ -20,6 +20,7 @@
 #include "model/diffusion/krea2.hpp"
 #include "model/diffusion/lens.hpp"
 #include "model/diffusion/lingbot_video.hpp"
+#include "model/diffusion/llada_image.hpp"
 #include "model/diffusion/ltxv.hpp"
 #include "model/diffusion/mage_flow.hpp"
 #include "model/diffusion/minimax_h3.hpp"
@@ -370,6 +371,19 @@ namespace sd::model_builders {
                                                                       "model.diffusion_model",
                                                                       version,
                                                                       weight_manager);
+        } else if (sd_version_is_llada_image(version)) {
+            result.conditioner = std::make_shared<LLaDAImageEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                      tensor_storage_map,
+                                                                      "text_encoders.llm",
+                                                                      "queryformer",
+                                                                      "text_projection",
+                                                                      "sigvq",
+                                                                      weight_manager,
+                                                                      tokenizers);
+            result.diffusion   = std::make_shared<LLaDAImage::LLaDAImageRunner>(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<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                                tensor_storage_map,

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;
     }

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<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override {
+        std::vector<float> 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<float>(i) / static_cast<float>(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<Flux2Scheduler>(image_seq_len);
                 break;
             }
+            case LLADA_IMAGE_SCHEDULER: {
+                LOG_INFO("get_sigmas with LLaDA-Image scheduler");
+                scheduler = std::make_shared<LLaDAImageScheduler>(extra_sample_args);
+                break;
+            }
             case FLUX_SCHEDULER: {
                 LOG_INFO("get_sigmas with Flux scheduler");
                 scheduler = std::make_shared<FluxScheduler>(image_seq_len, extra_sample_args);

src/stable-diffusion.cpp

@@ -137,6 +137,7 @@ const char* scheduler_to_str[] = {
     "flux2",
     "flux",
     "beta",
+    "llada_image",
 };
 
 static_assert(SCHEDULER_COUNT == sizeof(scheduler_to_str) / sizeof(scheduler_to_str[0]),