diff --git a/README.md b/README.md
index abb2dcb20..4bba56c28 100644
--- a/README.md
+++ b/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)
diff --git a/assets/moss/sft.png b/assets/moss/sft.png
new file mode 100644
index 000000000..ede42d1f7
Binary files /dev/null and b/assets/moss/sft.png differ
diff --git a/docs/moss_image.md b/docs/moss_image.md
new file mode 100644
index 000000000..2784c2fc1
--- /dev/null
+++ b/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
+```
+
+
+
+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.
diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp
index d5fa6750d..7bbc38348 100644
--- a/src/conditioning/conditioner.hpp
+++ b/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(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\n";
+ prompt_attn_range = {0, 0};
} else if (sd_version_is_flux3(version)) {
prompt_template_encode_start_idx = 0;
max_length = 4096;
diff --git a/src/model.h b/src/model.h
index baf3a2a34..c8d803a66 100644
--- a/src/model.h
+++ b/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;
diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp
new file mode 100644
index 000000000..f1cbbc7f3
--- /dev/null
+++ b/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
+#include
+
+#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(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(config.dim, 1e-6f);
+ blocks["norm2"] = std::make_shared(config.dim, 1e-6f);
+ blocks["q_norm"] = std::make_shared(config.head_dim, 1e-6f);
+ blocks["k_norm"] = std::make_shared(config.head_dim, 1e-6f);
+ blocks["qkv"] = std::make_shared(config.dim, config.dim * 3, true);
+ blocks["proj"] = std::make_shared(config.dim, config.dim, true);
+ blocks["ffn_up"] = std::make_shared(config.dim, config.ffn_dim, true);
+ blocks["ffn_gate"] = std::make_shared(config.dim, config.ffn_dim, true);
+ blocks["ffn_down"] = std::make_shared(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(blocks[name])->forward(ctx, h); };
+ auto norm = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast(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(config.dim, 1e-6f);
+ blocks["norm2"] = std::make_shared(config.dim, 1e-6f);
+ blocks["q_norm"] = std::make_shared(config.head_dim, 1e-6f);
+ blocks["k_norm"] = std::make_shared(config.head_dim, 1e-6f);
+ blocks["proj.qkv"] = std::make_shared(config.dim, config.dim * 3, false);
+ blocks["proj.o"] = std::make_shared(config.dim, config.dim, false);
+ blocks["proj.ffn_up_gate"] = std::make_shared(config.dim, config.ffn_dim * 2, false);
+ blocks["proj.ffn_down"] = std::make_shared(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(blocks[name])->forward(ctx, h); };
+ auto norm = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast(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(config.channels, config.dim, true);
+ blocks["text_embedder.proj"] = std::make_shared(config.text_dim, config.dim, true);
+ for (int i = 0; i < 2; ++i) {
+ blocks["text_embedder.blocks." + std::to_string(i)] = std::make_shared(config);
+ }
+ blocks["time_embedding.0"] = std::make_shared(config.freq_dim, config.dim, true);
+ blocks["time_embedding.2"] = std::make_shared(config.dim, config.dim, true);
+ blocks["shared_block_modulation"] = std::make_shared(config.dim, 6 * config.dim, true);
+ blocks["video_head_modulation"] = std::make_shared(config.dim, 2 * config.dim, true);
+ for (int64_t i = 0; i < config.depth; ++i) {
+ blocks["layers." + std::to_string(i)] = std::make_shared(config);
+ }
+ blocks["final_norm_video"] = std::make_shared(config.dim, 1e-6f);
+ blocks["final_linear_video"] = std::make_shared(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& skip) {
+ auto c = ctx->ggml_ctx;
+ auto linear = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast(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(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(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(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 pe_data, time_data;
+
+ MOSSImageRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr 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& tensors, const std::string& prefix) override {
+ model.get_param_tensors(tensors, prefix);
+ }
+
+ sd::Tensor 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 skip;
+ if (auto extra = std::get_if(&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(h), static_cast(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(4, static_cast(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__
diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp
index a782ea7bd..8ea918025 100644
--- a/src/model/te/llm.hpp
+++ b/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 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(tensor.ne[0]);
+ } else if (ends_with(name, "linear_attn.norm.weight")) {
+ config.linear_head_dim = static_cast(tensor.ne[0]);
+ } else if (ends_with(name, "linear_attn.conv1d.weight")) {
+ config.linear_conv_kernel = static_cast(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(tensor.ne[1] / (2 * config.head_dim));
+ } else if (ends_with(name, "self_attn.k_proj.weight")) {
+ config.num_kv_heads = static_cast(tensor.ne[1] / config.head_dim);
+ } else if (ends_with(name, "linear_attn.in_proj_z.weight")) {
+ config.linear_value_heads = static_cast(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((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(hidden_size, conv_dim, false);
+ blocks["in_proj_z"] = std::make_shared(hidden_size, head_dim * value_heads, false);
+ blocks["in_proj_a"] = std::make_shared(hidden_size, value_heads, false);
+ blocks["in_proj_b"] = std::make_shared(hidden_size, value_heads, false);
+ blocks["out_proj"] = std::make_shared(head_dim * value_heads, hidden_size, false);
+ blocks["norm"] = std::make_shared(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(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(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(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(config.hidden_size, (num_heads + num_kv_heads * 2) * head_dim, config.qkv_bias);
} else {
- blocks["q_proj"] = std::make_shared(config.hidden_size, num_heads * head_dim, config.qkv_bias);
+ blocks["q_proj"] = std::make_shared(config.hidden_size, num_heads * head_dim * (arch == LLMArch::QWEN3_5 ? 2 : 1), config.qkv_bias);
blocks["k_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
if (!k_eq_v) {
blocks["v_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
@@ -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(config, sliding_attention == 0);
+ if (config.arch == LLMArch::QWEN3_5 && config.linear_attention_layers.count(layer_index)) {
+ blocks["linear_attn"] = std::make_shared(config);
+ } else {
+ blocks["self_attn"] = std::make_shared(config, sliding_attention == 0);
+ }
if (config.arch == LLMArch::GPT_OSS_20B) {
blocks["mlp"] = std::make_shared(config);
} else if ((config.arch == LLMArch::LLADA2_MOE || 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(blocks["self_attn"]);
+ auto self_attn = blocks.count("self_attn") ? std::dynamic_pointer_cast(blocks["self_attn"]) : nullptr;
auto input_layernorm = std::dynamic_pointer_cast(blocks["input_layernorm"]);
auto pre_ffw_norm = std::dynamic_pointer_cast(blocks[pre_ffw_norm_name]);
std::shared_ptr 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(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 ||
diff --git a/src/model_loader.cpp b/src/model_loader.cpp
index d881c7e58..8822b2c3f 100644
--- a/src/model_loader.cpp
+++ b/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;
diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp
index 8b286e813..a342eb82e 100644
--- a/src/name_conversion.cpp
+++ b/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> 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)) {
diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp
index daa15175e..70338bd99 100644
--- a/src/pipeline/diffusion_engine.cpp
+++ b/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 StableDiffusionGGML::sample(const std::shared_ptr StableDiffusionGGML::sample(const std::shared_ptr* uncond_skip_layers = nullptr;
+ const std::vector 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,
diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp
index 494567904..4b435170f 100644
--- a/src/pipeline/model_builders.cpp
+++ b/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(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(ctx.backends.runtime_backend(SDBackendModule::TE),
+ tensor_storage_map, version, "", false, weight_manager, tokenizers);
+ result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+ tensor_storage_map, "model.diffusion_model", weight_manager);
} else { // SD1.x SD2.x SDXL
std::map embbeding_map;
for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) {