feat: add moss-image support (#2148)

README.md

@@ -46,6 +46,7 @@ API and command-line option may change frequently.***
     - [FLUX.1-dev/FLUX.1-schnell](./docs/flux.md)
     - [FLUX.2-dev/FLUX.2-klein](./docs/flux2.md)
     - [FLUX.3 Image Klein / Klein Base](./docs/flux3.md)
+    - [MOSS-Image-Base / MOSS-Image-SFT](./docs/moss_image.md)
     - [Lens](./docs/lens.md)
     - [Chroma](./docs/chroma.md)
     - [Chroma1-Radiance](./docs/chroma_radiance.md)

docs/moss_image.md

@@ -0,0 +1,34 @@
+# How to Use
+
+MOSS-Image uses a 5B diffusion transformer, the FLUX.2 VAE, and Qwen3.5 9B as the LLM text encoder. Both Base and SFT are supported. Prompt enhancement is not included.
+
+## Download weights
+
+- Download MOSS-Image Base or SFT
+    - safetensors: https://huggingface.co/Comfy-Org/MOSS-Image/tree/main/diffusion_models
+- Download vae
+    - safetensors: https://huggingface.co/Comfy-Org/MOSS-Image/tree/main/vae
+- Download Qwen3.5 9B
+    - gguf: https://huggingface.co/unsloth/Qwen3.5-9B-GGUF/tree/main
+- Download `tokenizer.json`
+    - https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base/tree/main/tokenizer
+
+Pass the external `tokenizer.json` with `--tokenizer`; Qwen3.5 uses a different vocabulary from the embedded Qwen tokenizer.
+
+## Examples
+
+### MOSS-Image Base
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\moss_image_base_int8_convrot_attn.safetensors --llm ..\models\text_encoders\Qwen3.5-9B-Q4_K_M.gguf --vae ..\models\vae\flux2_ae.safetensors --tokenizer ..\models\tokenizers\tokenizer_qwen3.5.json --cfg-scale 3.0 -p "a lovely cat holding a sign says 'moss-image.cpp'" -v --offload-to-cpu -H 1024 -W 1024
+```
+
+### MOSS-Image SFT
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\moss_image_sft_int8_convrot_attn.safetensors --llm ..\models\text_encoders\Qwen3.5-9B-Q4_K_M.gguf --vae ..\models\vae\flux2_ae.safetensors --tokenizer ..\models\tokenizers\tokenizer_qwen3.5.json --cfg-scale 3.0 -p "a lovely cat holding a sign says 'moss-image.cpp'" -v --offload-to-cpu -H 1024 -W 1024
+```
+
+<img src="../assets/moss/sft.png" />
+
+Use dimensions divisible by 16. The default discrete scheduler uses the upstream time shift of 5. Use `--flow-shift` to override it. The negative CFG branch skips DiT layer 22 (zero-based), matching the upstream pipeline.

src/conditioning/conditioner.hpp

@@ -1995,6 +1995,9 @@ struct LLMEmbedder : public Conditioner {
             if (sd_version_is_pid(version)) {
                 throw std::runtime_error("PiD requires an external Gemma 2 tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer");
             }
+            if (version == VERSION_MOSS_IMAGE) {
+                throw std::runtime_error("MOSS-Image requires the official Qwen3.5 tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer");
+            }
         }
         LLM::LLMArch arch = LLM::LLMArch::QWEN2_5_VL;
         if (version == VERSION_FLUX2) {
@@ -2018,6 +2021,8 @@ struct LLMEmbedder : public Conditioner {
             arch = LLM::LLMArch::QWEN3_VL;
         } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
             arch = LLM::LLMArch::QWEN3;
+        } else if (version == VERSION_MOSS_IMAGE) {
+            arch = LLM::LLMArch::QWEN3_5;
         }
         llm        = std::make_shared<LLM::LLMRunner>(arch,
                                                backend,
@@ -2032,6 +2037,8 @@ struct LLMEmbedder : public Conditioner {
             pad_id = 199999;
         } else if (arch == LLM::LLMArch::GEMMA2_2B) {
             pad_id = 0;
+        } else if (arch == LLM::LLMArch::QWEN3_5) {
+            pad_id = 248044;
         }
         tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, pad_id);
         if (!tokenizer) {
@@ -3039,6 +3046,11 @@ struct LLMEmbedder : public Conditioner {
 
                 prompt += "<|im_end|>\n<|im_start|>assistant\n";
             }
+        } else if (version == VERSION_MOSS_IMAGE) {
+            prompt_template_encode_start_idx = 0;
+            max_length                       = 4096;
+            prompt                           = "<|im_start|>user\n" + conditioner_params.text + "<|im_end|>\n<|im_start|>assistant\n<think>\n";
+            prompt_attn_range                = {0, 0};
         } else if (sd_version_is_flux3(version)) {
             prompt_template_encode_start_idx = 0;
             max_length                       = 4096;

src/core/ggml_extend.cpp

@@ -344,22 +344,15 @@ ggml_tensor* ggml_ext_pad(ggml_context* ctx,
 }
 
 static ggml_tensor* conv_1d(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, int s0, int p0, int d0, bool force_prec_f32) {
-    ggml_tensor* result;
-    if (force_prec_f32) {
-        ggml_tensor* patches = ggml_im2col(ctx, w, x, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F32);
-        result               = ggml_mul_mat(ctx,
-                                            ggml_reshape_2d(ctx, patches, patches->ne[0], patches->ne[2] * patches->ne[1]),
-                                            ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1], w->ne[2]));
-        result               = ggml_reshape_3d(ctx, result, patches->ne[1], w->ne[2], patches->ne[2]);
-    } else {
-        result = ggml_conv_1d(ctx, w, x, s0, p0, d0);
+    if (!force_prec_f32) {
+        return ggml_conv_1d(ctx, w, x, s0, p0, d0);
     }
-    if (x->ne[2] > 1) {
-        // mul_mat packs positions and batches before output channels: [OL, N, OC].
-        result = ggml_reshape_3d(ctx, result, result->ne[0], x->ne[2], w->ne[2]);
-        result = ggml_cont(ctx, ggml_permute(ctx, result, 0, 2, 1, 3));
-    }
-    return result;
+    ggml_tensor* patches = ggml_im2col(ctx, w, x, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F32);
+    ggml_tensor* result  = ggml_mul_mat(ctx,
+                                        ggml_reshape_2d(ctx, patches, patches->ne[0], patches->ne[2] * patches->ne[1]),
+                                        ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1], w->ne[2]));
+    result               = ggml_reshape_3d(ctx, result, patches->ne[1], patches->ne[2], w->ne[2]);
+    return ggml_cont(ctx, ggml_permute(ctx, result, 0, 2, 1, 3));
 }
 
 ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx,

src/model.h

@@ -67,6 +67,7 @@ enum SDVersion {
     VERSION_Z_IMAGE_L2P,
     VERSION_IRIS,
     VERSION_FLUX3,
+    VERSION_MOSS_IMAGE,
     VERSION_COUNT,
 };
 
@@ -291,7 +292,8 @@ static inline bool sd_version_uses_flux2_vae(SDVersion version) {
         sd_version_is_lens(version) ||
         sd_version_is_ideogram4(version) ||
         sd_version_is_sefi_image(version) ||
-        sd_version_is_llada_image(version)) {
+        sd_version_is_llada_image(version) ||
+        version == VERSION_MOSS_IMAGE) {
         return true;
     }
     return false;
@@ -348,7 +350,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_mage_flow(version) ||
         sd_version_is_sensenova_u1(version) ||
         sd_version_is_pixart(version) ||
-        version == VERSION_IRIS) {
+        version == VERSION_IRIS ||
+        version == VERSION_MOSS_IMAGE) {
         return true;
     }
     return false;

src/model/common/ggml_block.hpp

@@ -13,8 +13,8 @@
 #include <vector>
 
 #include "core/ggml_extend.h"
-#include "core/ggml_runner.h"
 #include "core/ggml_extend_backend.h"
+#include "core/ggml_runner.h"
 #include "model.h"
 
 class GGMLBlock {

src/model/diffusion/moss_image.hpp

@@ -0,0 +1,262 @@
+#ifndef __SD_MODEL_DIFFUSION_MOSS_IMAGE_HPP__
+#define __SD_MODEL_DIFFUSION_MOSS_IMAGE_HPP__
+
+#include <algorithm>
+#include <cinttypes>
+
+#include "model/common/ggml_block.hpp"
+#include "model/common/rope.hpp"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/model.hpp"
+#include "model_loader.h"
+
+namespace MOSSImage {
+    struct MOSSImageConfig {
+        int64_t dim      = 3072;
+        int64_t ffn_dim  = 8192;
+        int64_t text_dim = 4096;
+        int64_t channels = 128;
+        int64_t head_dim = 128;
+        int64_t depth    = 42;
+        int64_t freq_dim = 256;
+
+        static MOSSImageConfig detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) {
+            MOSSImageConfig config;
+            int64_t depth = 0;
+            for (const auto& [name, tensor] : weights) {
+                if (!starts_with(name, prefix + ".")) {
+                    continue;
+                }
+                auto key = name.substr(prefix.size() + 1);
+                if (key == "video_embedder.weight") {
+                    config.channels = tensor.ne[0];
+                    config.dim      = tensor.ne[1];
+                } else if (key == "text_embedder.proj.weight") {
+                    config.text_dim = tensor.ne[0];
+                } else if (key == "layers.0.proj.ffn_down.weight") {
+                    config.ffn_dim = tensor.ne[0];
+                } else if (key == "layers.0.q_norm.weight") {
+                    config.head_dim = tensor.ne[0];
+                } else if (key == "time_embedding.0.weight") {
+                    config.freq_dim = tensor.ne[0];
+                }
+                if (starts_with(key, "layers.")) {
+                    depth = std::max<int64_t>(depth, std::stoi(key.substr(7)) + 1);
+                }
+            }
+            if (depth) {
+                config.depth = depth;
+                LOG_VERBOSE("moss_image: dim = %" PRId64 ", ffn_dim = %" PRId64 ", depth = %" PRId64 ", head_dim = %" PRId64,
+                            config.dim, config.ffn_dim, config.depth, config.head_dim);
+            }
+            return config;
+        }
+    };
+
+    inline ggml_tensor* situ(ggml_context* c, ggml_tensor* hidden, ggml_tensor* gate) {
+        auto a = ggml_scale(c, ggml_tanh(c, ggml_scale(c, gate, 0.25f)), 4.f);
+        a      = ggml_mul(c, a, ggml_sigmoid(c, gate));
+        auto b = ggml_scale(c, ggml_tanh(c, ggml_scale(c, hidden, 1.f / 25.f)), 25.f);
+        return ggml_mul(c, a, b);
+    }
+
+    struct TextBlock : public GGMLBlock {
+        MOSSImageConfig config;
+
+        TextBlock(const MOSSImageConfig& config)
+            : config(config) {
+            blocks["norm1"]    = std::make_shared<RMSNorm>(config.dim, 1e-6f);
+            blocks["norm2"]    = std::make_shared<RMSNorm>(config.dim, 1e-6f);
+            blocks["q_norm"]   = std::make_shared<RMSNorm>(config.head_dim, 1e-6f);
+            blocks["k_norm"]   = std::make_shared<RMSNorm>(config.head_dim, 1e-6f);
+            blocks["qkv"]      = std::make_shared<Linear>(config.dim, config.dim * 3, true);
+            blocks["proj"]     = std::make_shared<Linear>(config.dim, config.dim, true);
+            blocks["ffn_up"]   = std::make_shared<Linear>(config.dim, config.ffn_dim, true);
+            blocks["ffn_gate"] = std::make_shared<Linear>(config.dim, config.ffn_dim, true);
+            blocks["ffn_down"] = std::make_shared<Linear>(config.ffn_dim, config.dim, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe) {
+            auto c      = ctx->ggml_ctx;
+            auto linear = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, h); };
+            auto norm   = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast<RMSNorm>(blocks[name])->forward(ctx, h); };
+            auto y      = norm("norm1", x);
+            auto qkv    = linear("qkv", y);
+            auto split  = [&](int i) {
+                return ggml_reshape_4d(c, ggml_ext_slice(c, qkv, 0, i * config.dim, (i + 1) * config.dim), config.head_dim, config.dim / config.head_dim, x->ne[1], x->ne[2]);
+            };
+            auto q = norm("q_norm", split(0));
+            auto k = norm("k_norm", split(1));
+            auto v = split(2);
+            y      = Rope::attention(ctx, q, k, v, pe, nullptr, 1.f, false);
+            y      = linear("proj", y);
+            x      = ggml_add(c, x, y);
+            y      = norm("norm2", x);
+            y      = situ(c, linear("ffn_up", y), linear("ffn_gate", y));
+            y      = linear("ffn_down", y);
+            return ggml_add(c, x, y);
+        }
+    };
+
+    struct Block : public GGMLBlock {
+        MOSSImageConfig config;
+
+        Block(const MOSSImageConfig& config)
+            : config(config) {
+            blocks["norm1"]            = std::make_shared<RMSNorm>(config.dim, 1e-6f);
+            blocks["norm2"]            = std::make_shared<RMSNorm>(config.dim, 1e-6f);
+            blocks["q_norm"]           = std::make_shared<RMSNorm>(config.head_dim, 1e-6f);
+            blocks["k_norm"]           = std::make_shared<RMSNorm>(config.head_dim, 1e-6f);
+            blocks["proj.qkv"]         = std::make_shared<Linear>(config.dim, config.dim * 3, false);
+            blocks["proj.o"]           = std::make_shared<Linear>(config.dim, config.dim, false);
+            blocks["proj.ffn_up_gate"] = std::make_shared<Linear>(config.dim, config.ffn_dim * 2, false);
+            blocks["proj.ffn_down"]    = std::make_shared<Linear>(config.ffn_dim, config.dim, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, ggml_tensor* mod) {
+            auto c      = ctx->ggml_ctx;
+            auto linear = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, h); };
+            auto norm   = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast<RMSNorm>(blocks[name])->forward(ctx, h); };
+            auto part   = [&](int i) { return ggml_ext_slice(c, mod, 0, i * config.dim, (i + 1) * config.dim); };
+            auto y      = norm("norm1", x);
+            y           = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(1))), part(0));
+            auto qkv    = linear("proj.qkv", y);
+            auto split  = [&](int i) {
+                return ggml_reshape_4d(c, ggml_ext_slice(c, qkv, 0, i * config.dim, (i + 1) * config.dim), config.head_dim, config.dim / config.head_dim, x->ne[1], x->ne[2]);
+            };
+            auto q       = norm("q_norm", split(0));
+            auto k       = norm("k_norm", split(1));
+            auto v       = split(2);
+            y            = Rope::attention(ctx, q, k, v, pe, nullptr, 1.f, false);
+            y            = linear("proj.o", y);
+            y            = ggml_mul(c, y, part(2));
+            x            = ggml_add(c, x, y);
+            y            = norm("norm2", x);
+            y            = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(4))), part(3));
+            auto up_gate = linear("proj.ffn_up_gate", y);
+            y            = situ(c, ggml_ext_slice(c, up_gate, 0, 0, config.ffn_dim), ggml_ext_slice(c, up_gate, 0, config.ffn_dim, 2 * config.ffn_dim));
+            y            = linear("proj.ffn_down", y);
+            y            = ggml_mul(c, y, part(5));
+            return ggml_add(c, x, y);
+        }
+    };
+
+    struct Model : public GGMLBlock {
+        MOSSImageConfig config;
+
+        Model(const MOSSImageConfig& config)
+            : config(config) {
+            blocks["video_embedder"]     = std::make_shared<Linear>(config.channels, config.dim, true);
+            blocks["text_embedder.proj"] = std::make_shared<Linear>(config.text_dim, config.dim, true);
+            for (int i = 0; i < 2; ++i) {
+                blocks["text_embedder.blocks." + std::to_string(i)] = std::make_shared<TextBlock>(config);
+            }
+            blocks["time_embedding.0"]        = std::make_shared<Linear>(config.freq_dim, config.dim, true);
+            blocks["time_embedding.2"]        = std::make_shared<Linear>(config.dim, config.dim, true);
+            blocks["shared_block_modulation"] = std::make_shared<Linear>(config.dim, 6 * config.dim, true);
+            blocks["video_head_modulation"]   = std::make_shared<Linear>(config.dim, 2 * config.dim, true);
+            for (int64_t i = 0; i < config.depth; ++i) {
+                blocks["layers." + std::to_string(i)] = std::make_shared<Block>(config);
+            }
+            blocks["final_norm_video"]   = std::make_shared<RMSNorm>(config.dim, 1e-6f);
+            blocks["final_linear_video"] = std::make_shared<Linear>(config.dim, config.channels, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image, ggml_tensor* time, ggml_tensor* text, ggml_tensor* pe, const std::vector<int>& skip) {
+            auto c            = ctx->ggml_ctx;
+            auto linear       = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, h); };
+            int64_t image_len = image->ne[0] * image->ne[1];
+            auto x            = linear("video_embedder", DiT::patchify(c, image, 1, 1));
+            text              = linear("text_embedder.proj", text);
+            auto text_pe      = ggml_ext_slice(c, pe, 3, image_len, image_len + text->ne[1]);
+            for (int i = 0; i < 2; ++i) {
+                text = std::dynamic_pointer_cast<TextBlock>(blocks["text_embedder.blocks." + std::to_string(i)])->forward(ctx, text, text_pe);
+            }
+            time      = linear("time_embedding.2", ggml_silu(c, linear("time_embedding.0", time)));
+            auto mod  = linear("shared_block_modulation", time);
+            auto head = linear("video_head_modulation", time);
+            x         = ggml_concat(c, x, text, 1);
+            for (int64_t i = 0; i < config.depth; ++i) {
+                if (std::find(skip.begin(), skip.end(), i) == skip.end()) {
+                    x = std::dynamic_pointer_cast<Block>(blocks["layers." + std::to_string(i)])->forward(ctx, x, pe, mod);
+                    sd::ggml_graph_cut::mark_graph_cut(x, "moss_image.layers." + std::to_string(i), "x");
+                }
+            }
+            x          = ggml_ext_slice(c, x, 1, 0, image_len);
+            x          = std::dynamic_pointer_cast<RMSNorm>(blocks["final_norm_video"])->forward(ctx, x);
+            auto shift = ggml_ext_slice(c, head, 0, 0, config.dim);
+            auto scale = ggml_ext_slice(c, head, 0, config.dim, 2 * config.dim);
+            x          = ggml_add(c, ggml_add(c, x, ggml_mul(c, x, scale)), shift);
+            x          = linear("final_linear_video", x);
+            return DiT::unpatchify(c, x, image->ne[1], image->ne[0], 1, 1);
+        }
+    };
+
+    struct MOSSImageRunner : public DiffusionModelRunner {
+        MOSSImageConfig config;
+        Model model;
+        std::vector<float> pe_data, time_data;
+
+        MOSSImageRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager), config(MOSSImageConfig::detect_from_weights(weights, prefix)), model(config) {
+            model.init(params_ctx, weights, prefix);
+        }
+
+        std::string get_desc() override { return "MOSS-Image"; }
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            model.get_param_tensors(tensors, prefix);
+        }
+
+        sd::Tensor<float> compute(int n_threads, const DiffusionParams& inputs) override {
+            if (!inputs.x || !inputs.context || !inputs.timesteps || inputs.context->empty()) {
+                LOG_ERROR("MOSS-Image requires latents, timesteps and text features");
+                return {};
+            }
+            std::vector<int> skip;
+            if (auto extra = std::get_if<SkipLayerDiffusionExtra>(&inputs.extra)) {
+                if (extra->skip_layers) {
+                    skip = *extra->skip_layers;
+                }
+            }
+            auto graph = [&]() {
+                auto gf   = new_graph_custom(65536);
+                auto x    = make_input(*inputs.x);
+                auto text = make_input(*inputs.context);
+                GGML_ASSERT(x->ne[3] == 1 && x->ne[2] == config.channels && text->ne[0] == config.text_dim);
+                int64_t h = x->ne[1], w = x->ne[0], nt = text->ne[1];
+                Rope::Embedding embedding;
+                embedding.positions.append_image(static_cast<int>(h), static_cast<int>(w));
+                embedding.positions.images.back().height_axis = 2;
+                embedding.positions.images.back().width_axis  = 3;
+                embedding.positions.append_tokens(nt);
+                for (int64_t token = 0; token < h * w + nt; ++token) {
+                    float pos[4] = {0.f, float(token - h * w), 0.f, 0.f};
+                    if (token < h * w) {
+                        pos[1] = float(nt);
+                        pos[2] = float(token / w) - float(h - 1) / 2.f;
+                        pos[3] = float(token % w) - float(w - 1) / 2.f;
+                    }
+                    embedding.ids.push_back({pos[0], pos[1], pos[2], pos[3]});
+                }
+                embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, std::vector<int>(4, static_cast<int>(config.head_dim / 4)), embedding.layout, &embedding.frequencies);
+                pe_data          = finish_rope_pe(std::move(embedding));
+                time_data.resize(config.freq_dim);
+                for (int64_t i = 0; i < config.freq_dim / 2; ++i) {
+                    double angle                       = inputs.timesteps->values()[0] * std::pow(10000.0, -double(i) / double(config.freq_dim / 2));
+                    time_data[i]                       = float(std::cos(angle));
+                    time_data[i + config.freq_dim / 2] = float(std::sin(angle));
+                }
+                auto pe   = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, h * w + nt);
+                auto time = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.freq_dim, 1, 1);
+                set_backend_tensor_data(pe, pe_data.data());
+                set_backend_tensor_data(time, time_data.data());
+                auto ctx = get_context();
+                ggml_build_forward_expand(gf, model.forward(&ctx, x, time, text, pe, skip));
+                return gf;
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(graph, n_threads, false), inputs.x->dim());
+        }
+    };
+}  // namespace MOSSImage
+
+#endif  // __SD_MODEL_DIFFUSION_MOSS_IMAGE_HPP__

src/model/te/llm.hpp

@@ -43,6 +43,7 @@ namespace LLM {
         QWEN2_5_VL,
         QWEN3,
         QWEN3_VL,
+        QWEN3_5,
         MISTRAL_SMALL_3_2,
         MINISTRAL_3_3B,
         GEMMA3_12B,
@@ -59,6 +60,7 @@ namespace LLM {
         "qwen2.5vl",
         "qwen3",
         "qwen3vl",
+        "qwen3.5",
         "mistral_small3.2",
         "ministral3.3b",
         "gemma3_12b",
@@ -121,6 +123,7 @@ namespace LLM {
         bool qk_norm                    = false;
         bool final_norm                 = true;
         bool rms_norm_add               = false;
+        bool qwen35_gguf                = false;
         bool normalize_input            = false;
         int64_t vocab_size              = 152064;
         int64_t max_position_embeddings = 128000;
@@ -134,6 +137,11 @@ namespace LLM {
         bool qkv_fused              = false;
         bool bidirectional          = false;
         float partial_rotary        = 1.f;
+        int linear_head_dim         = 128;
+        int linear_key_heads        = 16;
+        int linear_value_heads      = 32;
+        int linear_conv_kernel      = 4;
+        std::set<int> linear_attention_layers;
 
         // DeepSeek-V3-style grouped-sigmoid MoE routing (LLaDA2)
         int64_t moe_intermediate_size = 0;
@@ -167,6 +175,19 @@ namespace LLM {
                 config.num_kv_heads = 8;
                 config.qkv_bias     = false;
                 config.rms_norm_eps = 1e-5f;
+            } else if (arch == LLMArch::QWEN3_5) {
+                config.hidden_size             = 4096;
+                config.intermediate_size       = 12288;
+                config.head_dim                = 256;
+                config.num_heads               = 16;
+                config.num_kv_heads            = 4;
+                config.qkv_bias                = false;
+                config.qk_norm                 = true;
+                config.rms_norm_add            = true;
+                config.partial_rotary          = 0.25f;
+                config.rope_thetas             = {10000000.f};
+                config.max_position_embeddings = 262144;
+                config.vocab_size              = 248320;
             } else if (arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) {
                 config.head_dim     = 128;
                 config.num_heads    = 32;
@@ -365,6 +386,41 @@ namespace LLM {
             if ((arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) && config.num_layers == 28) {
                 config.num_heads = 16;
             }
+            if (arch == LLMArch::QWEN3_5) {
+                config.qwen35_gguf  = tensor_storage_map.count(prefix + ".model.layers.0.linear_attn.A") > 0;
+                config.rms_norm_add = !config.qwen35_gguf;
+                for (const auto& [name, tensor] : tensor_storage_map) {
+                    if (!starts_with(name, prefix + ".")) {
+                        continue;
+                    }
+                    if (ends_with(name, "self_attn.q_norm.weight")) {
+                        config.head_dim = static_cast<int>(tensor.ne[0]);
+                    } else if (ends_with(name, "linear_attn.norm.weight")) {
+                        config.linear_head_dim = static_cast<int>(tensor.ne[0]);
+                    } else if (ends_with(name, "linear_attn.conv1d.weight")) {
+                        config.linear_conv_kernel = static_cast<int>(tensor.ne[0]);
+                        const auto pos            = name.find("layers.");
+                        config.linear_attention_layers.insert(std::stoi(name.substr(pos + 7)));
+                    }
+                }
+                for (const auto& [name, tensor] : tensor_storage_map) {
+                    if (!starts_with(name, prefix + ".")) {
+                        continue;
+                    }
+                    if (ends_with(name, "self_attn.q_proj.weight")) {
+                        config.num_heads = static_cast<int>(tensor.ne[1] / (2 * config.head_dim));
+                    } else if (ends_with(name, "self_attn.k_proj.weight")) {
+                        config.num_kv_heads = static_cast<int>(tensor.ne[1] / config.head_dim);
+                    } else if (ends_with(name, "linear_attn.in_proj_z.weight")) {
+                        config.linear_value_heads = static_cast<int>(tensor.ne[1] / config.linear_head_dim);
+                    }
+                }
+                for (const auto& [name, tensor] : tensor_storage_map) {
+                    if (starts_with(name, prefix + ".") && ends_with(name, "linear_attn.in_proj_qkv.weight")) {
+                        config.linear_key_heads = static_cast<int>((tensor.ne[1] / config.linear_head_dim - config.linear_value_heads) / 2);
+                    }
+                }
+            }
             if (arch == LLMArch::QWEN3_VL &&
                 (config.num_layers == 50 || config.num_layers == 64) &&
                 config.hidden_size == 5120) {
@@ -1390,6 +1446,81 @@ namespace LLM {
         }
     };
 
+    struct Qwen35DeltaNet : public UnaryBlock {
+        int head_dim;
+        int key_heads;
+        int value_heads;
+        int conv_dim;
+        int kernel;
+        bool gguf_layout;
+
+        Qwen35DeltaNet(const LLMConfig& config)
+            : head_dim(config.linear_head_dim), key_heads(config.linear_key_heads), value_heads(config.linear_value_heads), conv_dim(head_dim * (key_heads * 2 + value_heads)), kernel(config.linear_conv_kernel), gguf_layout(config.qwen35_gguf) {
+            auto hidden_size      = config.hidden_size;
+            blocks["in_proj_qkv"] = std::make_shared<Linear>(hidden_size, conv_dim, false);
+            blocks["in_proj_z"]   = std::make_shared<Linear>(hidden_size, head_dim * value_heads, false);
+            blocks["in_proj_a"]   = std::make_shared<Linear>(hidden_size, value_heads, false);
+            blocks["in_proj_b"]   = std::make_shared<Linear>(hidden_size, value_heads, false);
+            blocks["out_proj"]    = std::make_shared<Linear>(head_dim * value_heads, hidden_size, false);
+            blocks["norm"]        = std::make_shared<LLMRMSNorm>(head_dim, 1e-6f);
+        }
+
+        void init_params(ggml_context* ctx, const String2TensorStorage& = {}, std::string = "") override {
+            params["conv1d.weight"]             = gguf_layout ? ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kernel, conv_dim) : ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kernel, 1, conv_dim);
+            params[gguf_layout ? "A" : "A_log"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, value_heads);
+            params["dt_bias"]                   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, value_heads);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto c          = ctx->ggml_ctx;
+            const int64_t n = x->ne[1];
+            GGML_ASSERT(x->ne[2] == 1);
+            auto linear        = [&](const char* name) { return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x); };
+            auto qkv           = linear("in_proj_qkv");
+            auto padded        = ggml_pad_ext(c, qkv, 0, 0, kernel - 1, 0, 0, 0, 0, 0);
+            ggml_tensor* mixed = nullptr;
+            // Causal depthwise convolution, with zero state for every independent prompt.
+            for (int i = 0; i < kernel; ++i) {
+                auto w    = ggml_cont(c, ggml_view_2d(c, params["conv1d.weight"], 1, conv_dim, kernel * sizeof(float), i * sizeof(float)));
+                w         = ggml_reshape_1d(c, w, conv_dim);
+                auto part = ggml_mul(c, ggml_ext_slice(c, padded, 1, i, i + n), w);
+                mixed     = mixed ? ggml_add(c, mixed, part) : part;
+            }
+            mixed      = ggml_silu(c, mixed);
+            auto heads = [&](int start, int count) {
+                auto h = ggml_ext_slice(c, mixed, 0, start, start + head_dim * count);
+                return ggml_reshape_4d(c, h, head_dim, count, n, 1);
+            };
+            auto repeat_keys = [&](ggml_tensor* h) {
+                h = ggml_scale(c, ggml_rms_norm(c, h, 1e-6f / head_dim), 1.f / std::sqrt(float(head_dim)));
+                // llama.cpp reorders V heads and stores negative exp(A_log) and offset norms.
+                if (gguf_layout) {
+                    return h;
+                }
+                // ggml GDN broadcasts heads cyclically; Qwen uses adjacent repeated heads.
+                h          = ggml_reshape_4d(c, h, head_dim, 1, key_heads, n);
+                auto shape = ggml_new_tensor_4d(c, GGML_TYPE_F32, head_dim, value_heads / key_heads, key_heads, n);
+                return ggml_reshape_4d(c, ggml_repeat(c, h, shape), head_dim, value_heads, n, 1);
+            };
+            auto q     = repeat_keys(heads(0, key_heads));
+            auto k     = repeat_keys(heads(head_dim * key_heads, key_heads));
+            auto v     = heads(2 * head_dim * key_heads, value_heads);
+            auto beta  = ggml_reshape_4d(c, ggml_sigmoid(c, linear("in_proj_b")), 1, value_heads, n, 1);
+            auto decay = ggml_softplus(c, ggml_add(c, linear("in_proj_a"), params["dt_bias"]));
+            auto a     = gguf_layout ? params["A"] : ggml_neg(c, ggml_exp(c, params["A_log"]));
+            decay      = ggml_mul(c, decay, a);
+            decay      = ggml_reshape_4d(c, decay, 1, value_heads, n, 1);
+            auto state = ggml_ext_zeros(c, head_dim, head_dim, value_heads, 1);
+            auto out   = ggml_gated_delta_net(c, q, k, v, decay, beta, state, 1);
+            out        = ggml_view_3d(c, out, head_dim, value_heads, n, head_dim * sizeof(float), head_dim * value_heads * sizeof(float), 0);
+            out        = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"])->forward(ctx, out);
+            auto z     = ggml_reshape_3d(c, linear("in_proj_z"), head_dim, value_heads, n);
+            out        = ggml_mul(c, out, ggml_silu(c, z));
+            out        = ggml_reshape_3d(c, out, head_dim * value_heads, n, 1);
+            return std::dynamic_pointer_cast<Linear>(blocks["out_proj"])->forward(ctx, out);
+        }
+    };
+
     struct Attention : public GGMLBlock {
     protected:
         LLMArch arch;
@@ -1439,7 +1570,7 @@ namespace LLM {
                 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["q_proj"] = std::make_shared<Linear>(config.hidden_size, num_heads * head_dim * (arch == LLMArch::QWEN3_5 ? 2 : 1), 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);
@@ -1502,6 +1633,13 @@ namespace LLM {
                 v = k_eq_v ? k : v_proj->forward(ctx, x);  // [N, n_token, num_kv_heads*head_dim]
             }
 
+            ggml_tensor* output_gate = nullptr;
+            if (arch == LLMArch::QWEN3_5) {
+                q           = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim * 2, num_heads, n_token, N);
+                output_gate = ggml_ext_slice(ctx->ggml_ctx, q, 0, head_dim, head_dim * 2);
+                output_gate = ggml_reshape_3d(ctx->ggml_ctx, output_gate, head_dim * num_heads, n_token, N);
+                q           = ggml_ext_slice(ctx->ggml_ctx, q, 0, 0, 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]
             v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_kv_heads, n_token, N);  // [N, n_token, num_kv_heads, head_dim]
@@ -1658,7 +1796,7 @@ namespace LLM {
                                   1.f,
                                   32.f,
                                   1.f);
-            } else if (arch == LLMArch::LLADA2_MOE) {
+            } else if (arch == LLMArch::LLADA2_MOE || arch == LLMArch::QWEN3_5) {
                 // 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 =
@@ -1734,6 +1872,9 @@ namespace LLM {
                 x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled);  // [N, n_token, hidden_size]
             }
 
+            if (output_gate) {
+                x = ggml_mul(ctx->ggml_ctx, x, ggml_sigmoid(ctx->ggml_ctx, output_gate));
+            }
             x = out_proj->forward(ctx, x);  // [N, n_token, hidden_size]
             return x;
         }
@@ -1783,7 +1924,11 @@ namespace LLM {
             if (!config.sliding_attention.empty()) {
                 sliding_attention = config.sliding_attention[layer_index % config.sliding_attention.size()];
             }
-            blocks["self_attn"] = std::make_shared<Attention>(config, sliding_attention == 0);
+            if (config.arch == LLMArch::QWEN3_5 && config.linear_attention_layers.count(layer_index)) {
+                blocks["linear_attn"] = std::make_shared<Qwen35DeltaNet>(config);
+            } else {
+                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 || config.arch == LLMArch::BAILING_MOE) && layer_index >= config.first_k_dense_replace) {
@@ -1813,7 +1958,7 @@ namespace LLM {
                              ggml_tensor* rope_cos               = nullptr,
                              ggml_tensor* rope_sin               = nullptr) {
             // x: [N, n_token, hidden_size]
-            auto self_attn                                  = std::dynamic_pointer_cast<Attention>(blocks["self_attn"]);
+            auto self_attn                                  = blocks.count("self_attn") ? std::dynamic_pointer_cast<Attention>(blocks["self_attn"]) : nullptr;
             auto input_layernorm                            = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["input_layernorm"]);
             auto pre_ffw_norm                               = std::dynamic_pointer_cast<LLMRMSNorm>(blocks[pre_ffw_norm_name]);
             std::shared_ptr<LLMRMSNorm> post_attention_norm = nullptr;
@@ -1833,7 +1978,11 @@ namespace LLM {
 
             auto residual = x;
             x             = input_layernorm->forward(ctx, x);
-            x             = self_attn->forward(ctx, x, input_pos, block_attention_mask, rope_index, rope_cos, rope_sin);
+            if (blocks.count("linear_attn")) {
+                x = std::dynamic_pointer_cast<Qwen35DeltaNet>(blocks["linear_attn"])->forward(ctx, x);
+            } else {
+                x = self_attn->forward(ctx, x, input_pos, block_attention_mask, rope_index, rope_cos, rope_sin);
+            }
             if (post_attention_norm != nullptr) {
                 x = post_attention_norm->forward(ctx, x);
             }
@@ -2360,6 +2509,7 @@ namespace LLM {
             if (config.arch == LLMArch::MISTRAL_SMALL_3_2 ||
                 config.arch == LLMArch::MINISTRAL_3_3B ||
                 config.arch == LLMArch::QWEN3 ||
+                config.arch == LLMArch::QWEN3_5 ||
                 config.arch == LLMArch::GEMMA3_12B ||
                 config.arch == LLMArch::GEMMA4_12B ||
                 config.arch == LLMArch::GEMMA2_2B ||

src/model/vae/minimax_h3_audio_vae.hpp

@@ -363,33 +363,22 @@ namespace MiniMaxH3 {
             auto pre       = std::dynamic_pointer_cast<AudioAttentionProjection>(blocks["pre_block"]);
             auto mean_proj = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["mean_proj"]);
 
-            // GGML's batched conv1d storage interleaves the stream dimension
-            // with output channels. Subsequent layers then read stereo samples
-            // as adjacent feature channels. Run each mono stream independently,
-            // matching PyTorch's reshape(B*S, 1, samples), and concatenate only
-            // the completed normalized latents.
             const int64_t streams = waveform->ne[2] * waveform->ne[3];
             waveform              = ggml_reshape_3d(ctx->ggml_ctx,
                                                     waveform,
                                                     waveform->ne[0],
                                                     1,
                                                     streams);
-            ggml_tensor* stereo_z = nullptr;
-            for (int64_t stream = 0; stream < streams; ++stream) {
-                auto mono = ggml_ext_slice(ctx->ggml_ctx, waveform, 2, stream, stream + 1);
-                auto x    = encoder->forward(ctx, mono);
-                x         = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
-                x         = pre->forward(ctx, x);
-                x         = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
-                auto z    = mean_proj->forward(ctx, x);
+            auto x                = encoder->forward(ctx, waveform);
+            x                     = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x                     = pre->forward(ctx, x);
+            x                     = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            auto z                = mean_proj->forward(ctx, x);
 
-                auto mean = ggml_reshape_4d(ctx->ggml_ctx, params["latents_mean"], 1, kLatentChannels, 1, 1);
-                auto std  = ggml_reshape_4d(ctx->ggml_ctx, params["latents_std"], 1, kLatentChannels, 1, 1);
-                z         = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, z, mean), std);
-                z         = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 2, 1, 3));
-                stereo_z  = stereo_z == nullptr ? z : ggml_concat(ctx->ggml_ctx, stereo_z, z, 1);
-            }
-            return stereo_z;
+            auto mean = ggml_reshape_4d(ctx->ggml_ctx, params["latents_mean"], 1, kLatentChannels, 1, 1);
+            auto std  = ggml_reshape_4d(ctx->ggml_ctx, params["latents_std"], 1, kLatentChannels, 1, 1);
+            z         = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, z, mean), std);
+            return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 2, 1, 3));
         }
 
         ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent) {

src/model_loader.cpp

@@ -444,6 +444,11 @@ SDVersion ModelLoader::get_sd_version() const {
         return VERSION_FLUX3;
     }
 
+    if (tensor_storage_map.count("model.diffusion_model.video_embedder.weight") &&
+        tensor_storage_map.count("model.diffusion_model.shared_block_modulation.weight")) {
+        return VERSION_MOSS_IMAGE;
+    }
+
     bool has_multiple_encoders = false;
     bool is_unet               = false;
 

src/name_conversion.cpp

@@ -1663,6 +1663,26 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
 
     replace_with_prefix_map(name, prefix_map);
 
+    if (version == VERSION_MOSS_IMAGE) {
+        const std::string te_prefix = "text_encoders.llm.";
+        if (starts_with(name, te_prefix + "blk.")) {
+            static const std::vector<std::pair<std::string, std::string>> qwen35_name_map = {
+                {"attn_post_norm.", "post_attention_layernorm."},
+                {"post_attention_norm.", "post_attention_layernorm."},
+                {"attn_qkv.", "linear_attn.in_proj_qkv."},
+                {"attn_gate.", "linear_attn.in_proj_z."},
+                {"ssm_conv1d.", "linear_attn.conv1d."},
+                {"ssm_dt.bias", "linear_attn.dt_bias"},
+                {"ssm_norm.", "linear_attn.norm."},
+                {"ssm_out.", "linear_attn.out_proj."},
+                {"ssm_alpha.", "linear_attn.in_proj_a."},
+                {"ssm_beta.", "linear_attn.in_proj_b."},
+                {"ssm_a", "linear_attn.A"},
+            };
+            replace_with_name_map(name, qwen35_name_map);
+        }
+    }
+
     if (version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
         const std::string hf_vision_prefix = "text_encoders.llm.model.visual.";
         if (starts_with(name, hf_vision_prefix)) {

src/pipeline/diffusion_engine.cpp

@@ -112,6 +112,7 @@ const char* model_version_to_str[] = {
     "Z-Image L2P",
     "Iris-3B",
     "Flux.3 Image Klein",
+    "MOSS-Image",
 };
 
 static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]),
@@ -1490,9 +1491,10 @@ bool StableDiffusionGGML::build_denoiser() {
                    sd_version_is_boogu_image(version) ||
                    sd_version_is_pid(version) ||
                    sd_version_is_ideogram4(version) ||
-                   version == VERSION_IRIS) {
+                   version == VERSION_IRIS ||
+                   version == VERSION_MOSS_IMAGE) {
             pred_type = FLOW_PRED;
-            if (sd_version_is_wan(version)) {
+            if (sd_version_is_wan(version) || version == VERSION_MOSS_IMAGE) {
                 default_flow_shift = 5.f;
             } else if (sd_version_is_hunyuan_video(version)) {
                 default_flow_shift = 7.f;
@@ -2644,7 +2646,7 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
                     unet_extra.ip_scale   = ip_adapter_strength;
                 }
                 diffusion_params.extra = unet_extra;
-            } else if (sd_version_is_sd3(version)) {
+            } else if (sd_version_is_sd3(version) || version == VERSION_MOSS_IMAGE) {
                 diffusion_params.extra = SkipLayerDiffusionExtra{local_skip_layers};
             } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
                 diffusion_params.extra = FluxDiffusionExtra{&guidance_tensor,
@@ -2873,9 +2875,12 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
         if (uncond_out.empty() && !uncond.empty()) {
             if (!skip_uncond) {
                 const std::vector<int>* uncond_skip_layers = nullptr;
+                const std::vector<int> moss_skip_layers    = {22};
                 if (is_skiplayer_step && slg_uncond) {
                     LOG_VERBOSE("Skipping layers at uncond step %d\n", step);
                     uncond_skip_layers = &skip_layer_guidance.layers();
+                } else if (version == VERSION_MOSS_IMAGE) {
+                    uncond_skip_layers = &moss_skip_layers;
                 }
                 uncond_out = run_condition(uncond,
                                            uncond.c_concat.empty() ? nullptr : &uncond.c_concat,

src/pipeline/model_builders.cpp

@@ -31,6 +31,7 @@
 #include "model/diffusion/minit2i.hpp"
 #include "model/diffusion/mmdit.hpp"
 #include "model/diffusion/model.hpp"
+#include "model/diffusion/moss_image.hpp"
 #include "model/diffusion/pid.hpp"
 #include "model/diffusion/pixart.hpp"
 #include "model/diffusion/qwen_image.hpp"
@@ -468,6 +469,11 @@ namespace sd::model_builders {
                                                                tensor_storage_map, version, "", false, weight_manager, tokenizers);
             result.diffusion   = std::make_shared<Iris::IrisRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
                                                                   tensor_storage_map, "model.diffusion_model", weight_manager);
+        } else if (version == VERSION_MOSS_IMAGE) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map, version, "", false, weight_manager, tokenizers);
+            result.diffusion   = std::make_shared<MOSSImage::MOSSImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                            tensor_storage_map, "model.diffusion_model", weight_manager);
         } else {  // SD1.x SD2.x SDXL
             std::map<std::string, std::string> embbeding_map;
             for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) {