From 162b18e5f3e80b48adc5191b0fd8d5df045cc7a7 Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Fri, 23 Oct 2026 12:41:10 -0300 Subject: [PATCH] feat: add image reference support to MOSS-Image --- docs/edit.md | 3 ++ src/conditioning/conditioner.hpp | 70 +++++++++++++++++++++++++-- src/model/diffusion/model.hpp | 1 + src/model/integrations/moss_image.cpp | 10 +++- src/model/te/llm.hpp | 12 +++-- src/name_conversion.cpp | 2 +- 6 files changed, 90 insertions(+), 8 deletions(-) diff --git a/docs/edit.md b/docs/edit.md index 1312854b71..21a1c5a7e1 100644 --- a/docs/edit.md +++ b/docs/edit.md @@ -21,6 +21,7 @@ Depending on the architecture, different models handle reference images differen | [**LLaDA-Image**](./llada_image.md#image-editing) | `llada_image` | | [**Fibo 2.0 Dev / Edit Turbo**](./fibo2.md#examples) | `fibo2` | | **Anima (Community Edit LoRAs)** | `cosmos_reference` | +| **MOSS-Image (native multimodal support)** | `moss_vision` | Stable-diffusion.spp also supports basic Unet-based editing models like instruct-pix2pix or CosXL-Edit. This document is not about those. @@ -59,6 +60,7 @@ The `--ref-image-args` argument accepts a comma-separated list of key-value pair | `mage_flow` | Mage-Flow-Edit | | `krea2_edit` | Specifically for [lbouaraba/krea2edit](https://huggingface.co/conradlocke/krea2-identity-edit). (or similar) | | `cosmos_reference` | For Anima | +| `moss_vision` | For MOSS-Image | | `default` | Uses the automatic detection based on model architecture. | --- @@ -106,6 +108,7 @@ For a technical overview of how each preset is configured, see the table below. | `krea2_ostris_edit`| Yes | `increase` | `area` | `force_ref_timestep_zero = true` | | `krea2_edit` | Yes | `increase` | `longest` | `vlm_size = 768` | | `cosmos_reference` | No | `fixed` | `none` | `resize_before_vae = false` | +| `moss_vision` | Yes | `increase` | `area` | | **Additional Default Notes:** - **VLM Input Sizes:** For most presets, `vlm_max_size` and `vlm_min_size` are set to `-1`, meaning the values are model-dependent and handled automatically. In `area` mode they represent pixel area; in `longest_side` mode they represent a side length in pixels. diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index e3832a1d8e..b7375ef7ef 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -3090,9 +3090,73 @@ struct LLMEmbedder : public Conditioner { } } 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}; + max_length = 0; + prompt = "<|im_start|>user\n"; + if (conditioner_params.ref_image_params.pass_to_vlm && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { + if (!llm->enable_vision) { + LOG_ERROR("MOSS-Image vision conditioning requires Qwen3.5 vision weights; provide --llm_vision or a combined encoder"); + return {}; + } + const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; + int min_pixels = conditioner_params.ref_image_params.vlm_min_size; + if (min_pixels <= 0) { + min_pixels = 2 * factor; + if (resize_mode == RefImageResizeMode::AREA) { + min_pixels *= min_pixels; + } + } + int max_pixels = conditioner_params.ref_image_params.vlm_max_size; + if (max_pixels <= 0) { + max_pixels = 1024; + if (resize_mode == RefImageResizeMode::AREA) { + max_pixels *= max_pixels; + } + } + for (size_t i = 0; i < conditioner_params.ref_images->size(); ++i) { + const auto& image = (*conditioner_params.ref_images)[i]; + int height = static_cast(image.shape()[1]); + int width = static_cast(image.shape()[0]); + int h_bar = std::max(factor, static_cast(std::round(static_cast(height) / factor)) * factor); + int w_bar = std::max(factor, static_cast(std::round(static_cast(width) / factor)) * factor); + resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode); + + LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar); + auto resized = image; + if (resized.shape()[0] != w_bar || resized.shape()[1] != h_bar) { + resized = sd::ops::interpolate(resized, + {w_bar, h_bar, resized.shape()[2], resized.shape()[3]}, + sd::ops::InterpolateMode::Bilinear, false, true); + } + const int64_t W = w_bar, H = h_bar, P = W * H; + auto rgb = sd::Tensor({W, H, 3, 1}); + for (int64_t p = 0; p < P; ++p) { + float alpha = resized.shape()[2] == 4 ? resized[p + 3 * P] : 1.f; + for (int c = 0; c < 3; ++c) { + float v = 2.f * (resized[p + c * P] * alpha + 1.f - alpha) - 1.f; + rgb[p + c * P] = std::clamp(v, -1.f, 1.f); + } + } + auto outputs = llm->encode_image_outputs(n_threads, rgb, false); + if (outputs.empty()) { + return {}; + } + prompt += (i == 0 ? "" : " ") + std::string("Picture ") + std::to_string(i + 1) + ": <|vision_start|>"; + size_t marker_start = prompt.size(); + prompt += "<|image_pad|><|vision_end|>"; + if (!expand_vision_marker(prompt, marker_start, std::move(outputs), + static_cast(H) / llm->config.vision.patch_size, + static_cast(W) / llm->config.vision.patch_size, + image_embeds, deepstack_image_embeds, image_grids)) { + return {}; + } + max_length += (H / factor) * (W / factor); + } + } + max_length += 4096; + prompt_attn_range.first = static_cast(prompt.size()); + prompt += conditioner_params.text; + prompt_attn_range.second = static_cast(prompt.size()); + prompt += "<|im_end|>\n<|im_start|>assistant\n\n"; } else if (sd_version_is_flux3(version)) { prompt_template_encode_start_idx = 0; max_length = 4096; diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 4ff91c416c..b89e09449c 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -65,6 +65,7 @@ const std::unordered_map REF_IMAGE_PRESETS = { // 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}}, + {"moss_vision", {true, false, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::AREA, -1, -1}}, }; struct UNetDiffusionExtra { diff --git a/src/model/integrations/moss_image.cpp b/src/model/integrations/moss_image.cpp index 4a85d169a5..5c96293bbb 100644 --- a/src/model/integrations/moss_image.cpp +++ b/src/model/integrations/moss_image.cpp @@ -7,8 +7,15 @@ static bool build_moss_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& runners) { + bool enable_vision = false; + for (const auto& [name, _] : ctx.tensor_storage_map) { + if (starts_with(name, "text_encoders.llm.visual.")) { + enable_vision = true; + break; + } + } runners.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - ctx.tensor_storage_map, ctx.version, "", false, + ctx.tensor_storage_map, ctx.version, "", enable_vision, ctx.weight_manager, tokenizers); runners.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), ctx.tensor_storage_map, "model.diffusion_model", ctx.weight_manager); @@ -20,6 +27,7 @@ const std::vector& moss_image_definitions() { ModelDefinition definition; definition.version = VERSION_MOSS_IMAGE; definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "moss_vision"; definition.name = "MOSS-Image"; definition.traits.is_dit = true; definition.traits.vae_kind = ModelVAEKind::FLUX2; diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index add8b6ff9a..1fbadf8ed6 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -387,6 +387,7 @@ namespace LLM { config.num_heads = 16; } if (arch == LLMArch::QWEN3_5) { + config.vision.arch = LLMVisionArch::QWEN3_VL; 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) { @@ -1800,7 +1801,7 @@ namespace LLM { 1.f, 32.f, 1.f); - } else if (arch == LLMArch::LLADA2_MOE || arch == LLMArch::QWEN3_5) { + } 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 = @@ -1832,6 +1833,12 @@ namespace LLM { 1.f, 32.f, 1.f); + } else if (arch == LLMArch::QWEN3_5) { + int sections[4] = {11, 11, 10, 0}; + int rotary_dim = rope_pairs * 2; + q = ggml_rope_multi(ctx->ggml_ctx, q, input_pos, nullptr, rotary_dim, sections, GGML_ROPE_TYPE_IMROPE, 262144, rope_thetas[0], 1.f, 0.f, 1.f, 32.f, 1.f); + k = ggml_rope_multi(ctx->ggml_ctx, k, input_pos, nullptr, rotary_dim, sections, GGML_ROPE_TYPE_IMROPE, 262144, 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); @@ -2516,7 +2523,6 @@ 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 || @@ -2534,7 +2540,7 @@ namespace LLM { input_pos_vec[2 * n_tokens + i] = i; input_pos_vec[3 * n_tokens + i] = 0; } - if (config.arch == LLMArch::QWEN3_VL && !image_grids.empty()) { + if ((config.arch == LLMArch::QWEN3_VL || config.arch == LLMArch::QWEN3_5) && !image_grids.empty()) { int offset = 0; for (const auto& grid : image_grids) { int end = grid.index + grid.size; diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 57a41c5c39..645c569b6f 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -1753,7 +1753,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) { } } - if (version == VERSION_FIBO2 || 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)) { + if (version == VERSION_FIBO2 || 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) || version == VERSION_MOSS_IMAGE) { const std::string hf_vision_prefix = "text_encoders.llm.model.visual."; if (starts_with(name, hf_vision_prefix)) { name = "text_encoders.llm.visual." + name.substr(hf_vision_prefix.size());