diff --git a/README.md b/README.md index 877cf7b9..33b1ccb1 100644 --- a/README.md +++ b/README.md @@ -64,6 +64,7 @@ API and command-line option may change frequently.*** - [HiDream-O1-Image](./docs/hidream_o1_image.md) - [Ideogram4](./docs/ideogram4.md) - [LLaDA-Image](./docs/llada_image.md) + - [Ming-Image Design](./docs/ming_image.md) - [PixArt](./docs/pixart.md) - [Image Edit Models](./docs/edit.md) - [FLUX.1-Kontext-dev](./docs/kontext.md) diff --git a/docs/ming_image.md b/docs/ming_image.md new file mode 100644 index 00000000..6077548c --- /dev/null +++ b/docs/ming_image.md @@ -0,0 +1,24 @@ +# Ming-Image + +[Ming-Image](https://github.com/inclusionAI/Ming-Image) 0.1 Design uses a 6B diffusion transformer (DiT), Ling-mini-2.0 for text conditioning, and the Ming-Image VAE. Text-to-image generation with RGBA output is supported. + +## Download weights + +- Download Ming-Image 0.1 Design DiT + - safetensors: https://huggingface.co/Comfy-Org/Ming-Image/tree/main/diffusion_models +- Download Ling-mini-2.0 BF16 + - safetensors: https://huggingface.co/Comfy-Org/Ming-Image/tree/main/text_encoders +- Download Ming-Image VAE + - safetensors: https://huggingface.co/Comfy-Org/Ming-Image/tree/main/vae +- Download Ling tokenizer + - tokenizer.json: https://huggingface.co/inclusionAI/Ming-Image-0.1-Design/blob/main/mllm/tokenizer.json + +The example below uses `ming_image_0.1_design_bf16.safetensors` for the DiT. You can also use `ming_image_0.1_design_int8_convrot.safetensors` with [INT8 convrot support](int8_convrot.md). Use the BF16 text encoder. + +## Text-to-image + +Pass the Ling `tokenizer.json` with `--tokenizer` and use the matching Ming-Image VAE. Image dimensions must be multiples of 16. Save the output as PNG to preserve the alpha channel. + +```bash +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\ming_image_0.1_design_bf16.safetensors --llm ..\models\text_encoders\ming_image_0.1_ling_mini_2.0_bf16.safetensors --vae ..\models\vae\ming_image_vae_bf16.safetensors --tokenizer ..\models\text_encoders\tokenizer.json -p "A cheerful orange cat sticker, transparent background" --width 1024 --height 1024 --steps 12 --cfg-scale 1 --sampling-method euler --diffusion-fa -v --offload-to-cpu -o ming_image.png +``` diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index 6b750f93..7796a40c 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -16,6 +16,7 @@ #include "model/te/clip.hpp" #include "model/te/llada_image_te.hpp" #include "model/te/llm.hpp" +#include "model/te/ming_image_te.hpp" #include "model/te/t5.hpp" #include "model_loader.h" #include "tokenizers/sensenova_u1_tokenizer.h" @@ -3171,6 +3172,65 @@ struct LLMEmbedder : public Conditioner { } }; +struct MingImageEmbedder : public Conditioner { + std::shared_ptr tokenizer; + std::shared_ptr text_model; + const std::string prefix = "text_encoders.llm"; + + MingImageEmbedder(ggml_backend_t backend, const String2TensorStorage& tensors, std::shared_ptr weight_manager, const TokenizerConfig& tokenizers) { + if (!tokenizers.has(TokenizerConfig::MAIN)) { + throw std::runtime_error("Ming-Image requires the Ling tokenizer.json; pass --tokenizer FILE"); + } + text_model = std::make_shared(backend, tensors, prefix, weight_manager); + tokenizer = tokenizers.create(TokenizerConfig::MAIN, text_model->config.backbone.vocab_size, 156895); + } + + void get_param_tensors(std::map& tensors) override { + text_model->get_param_tensors(tensors, prefix); + } + void get_param_tensor_ops(std::map& ops) override { + text_model->get_param_tensor_ops(ops); + } + void get_layer_split_param_tensors(std::map& tensors) override { + text_model->get_param_tensors(tensors, prefix); + } + void set_flash_attention_enabled(bool enabled) override { text_model->set_flash_attention_enabled(enabled); } + void set_max_graph_vram_bytes(size_t bytes) override { text_model->set_max_graph_vram_bytes(bytes); } + void set_runtime_backends(const std::vector& backends) override { text_model->set_runtime_backends(backends); } + void set_graph_cut_layer_split_enabled(bool enabled) override { text_model->set_graph_cut_layer_split_enabled(enabled); } + void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) override { text_model->set_graph_cut_layer_split_backend_vram_limits(limits); } + void set_scale_overrides(float linear, float attention) override { text_model->set_scale_overrides(linear, attention); } + void set_weight_adapter(const std::shared_ptr& adapter) override { text_model->set_weight_adapter(adapter); } + void runner_end() override { text_model->runner_end(); } + + SDCondition get_learned_condition(int n_threads, const ConditionerParams& input) override { + if (input.ref_images != nullptr && !input.ref_images->empty()) { + LOG_ERROR("Ming-Image currently supports text-to-image only"); + return {}; + } + std::string prompt = + "SYSTEM你是一个友好的AI助手。\n\ndetailed thinking off<|role_end|>" + "HUMAN" + + input.text + "<|role_end|>ASSISTANT"; + std::vector tokens; + if (!tokenizer->encode(prompt, tokens, nullptr)) { + return {}; + } + if (tokens.size() + 258 > 32768) { + LOG_ERROR("Ming-Image prompt exceeds the text encoder context length"); + return {}; + } + auto output = text_model->compute(n_threads, tokens); + if (output.empty()) { + return {}; + } + SDCondition result; + result.c_crossattn = sd::ops::slice(sd::ops::slice(output, 1, 0, 256), 0, 0, text_model->config.caption_dim); + result.extra_c_crossattns.push_back(sd::ops::slice(output, 1, 256, output.shape()[1])); + return result; + } +}; + struct LTXAVTextProjection : public GGMLBlock { static constexpr int64_t kHiddenSize = 3840; static constexpr int64_t kNumStates = 49; diff --git a/src/model.h b/src/model.h index 055feb94..277d7b93 100644 --- a/src/model.h +++ b/src/model.h @@ -63,6 +63,7 @@ enum SDVersion { VERSION_LLADA_IMAGE, VERSION_ESRGAN, VERSION_PIXART, + VERSION_MING_IMAGE, VERSION_COUNT, }; @@ -280,7 +281,7 @@ static inline bool sd_version_uses_flux2_vae(SDVersion version) { } static inline bool sd_version_uses_wan_vae(SDVersion version) { - if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) { + if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version) || version == VERSION_MING_IMAGE) { return true; } return false; @@ -314,6 +315,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) || + version == VERSION_MING_IMAGE || sd_version_is_llada_image(version) || sd_version_is_boogu_image(version) || sd_version_is_ernie_image(version) || diff --git a/src/model/diffusion/ming_image.hpp b/src/model/diffusion/ming_image.hpp new file mode 100644 index 00000000..f19be990 --- /dev/null +++ b/src/model/diffusion/ming_image.hpp @@ -0,0 +1,135 @@ +#ifndef __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__ +#define __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__ + +#include "z_image.hpp" + +namespace MingImage { + struct MingImageConfig : ZImage::ZImageConfig { + bool split_qkv = true; + + static MingImageConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) { + MingImageConfig config; + static_cast(config) = ZImage::ZImageConfig::detect_from_weights(tensors, prefix); + config.split_qkv = tensors.count(prefix + ".layers.0.attention.qkv.weight") == 0; + return config; + } + }; + + class MingImageModel : public GGMLBlock { + MingImageConfig config; + + public: + explicit MingImageModel(const MingImageConfig& config) + : config(config) { + blocks["x_embedder"] = std::make_shared(config.patch_size * config.patch_size * config.in_channels, config.hidden_size); + blocks["t_embedder"] = std::make_shared(1024, 256, std::min(config.hidden_size, 256)); + blocks["cap_embedder.0"] = std::make_shared(config.cap_feat_dim, config.norm_eps); + blocks["cap_embedder.1"] = std::make_shared(config.cap_feat_dim, config.hidden_size); + auto add_blocks = [&](const std::string& prefix, int64_t count, bool modulation) { + for (int64_t i = 0; i < count; ++i) { + blocks[prefix + std::to_string(i)] = std::make_shared( + static_cast(i), 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, true, config.split_qkv, 1e-5f); + } + }; + add_blocks("noise_refiner.", config.num_refiner_layers, true); + add_blocks("context_refiner.", config.num_refiner_layers, false); + add_blocks("layers.", config.num_layers, true); + blocks["final_layer"] = std::make_shared(config.hidden_size, config.patch_size, config.out_channels); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, ggml_tensor* direct, ggml_tensor* pe) { + auto gctx = ctx->ggml_ctx; + const int64_t width = x->ne[0], height = x->ne[1]; + auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size, false); + img = std::dynamic_pointer_cast(blocks["x_embedder"])->forward(ctx, img); + auto txt = std::dynamic_pointer_cast(blocks["cap_embedder.0"])->forward(ctx, context); + txt = std::dynamic_pointer_cast(blocks["cap_embedder.1"])->forward(ctx, txt); + txt = ggml_concat(gctx, txt, direct, 1); + auto t = std::dynamic_pointer_cast(blocks["t_embedder"])->forward(ctx, timestep); + const int64_t n_txt = txt->ne[1], n_img = img->ne[1]; + auto txt_pe = ggml_ext_slice(gctx, pe, 3, 0, n_txt); + auto img_pe = ggml_ext_slice(gctx, pe, 3, n_txt, n_txt + n_img); + for (int64_t i = 0; i < config.num_refiner_layers; ++i) { + txt = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)])->forward(ctx, txt, txt_pe); + sd::ggml_graph_cut::mark_graph_cut(txt, "ming_image.context_refiner." + std::to_string(i), "txt"); + } + for (int64_t i = 0; i < config.num_refiner_layers; ++i) { + img = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)])->forward(ctx, img, img_pe, nullptr, t); + sd::ggml_graph_cut::mark_graph_cut(img, "ming_image.noise_refiner." + std::to_string(i), "img"); + } + auto combined = ggml_concat(gctx, txt, img, 1); + for (int64_t i = 0; i < config.num_layers; ++i) { + combined = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)])->forward(ctx, combined, pe, nullptr, t); + sd::ggml_graph_cut::mark_graph_cut(combined, "ming_image.layers." + std::to_string(i), "combined"); + } + img = ggml_ext_slice(gctx, combined, 1, n_txt, n_txt + n_img); + img = std::dynamic_pointer_cast(blocks["final_layer"])->forward(ctx, img, t); + img = DiT::unpatchify_and_crop(gctx, img, height, width, config.patch_size, config.patch_size, false); + return ggml_scale(gctx, img, -1.f); + } + }; + + struct MingImageRunner : DiffusionModelRunner { + MingImageConfig config; + MingImageModel model; + std::vector pe_values; + + MingImageRunner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(MingImageConfig::detect_from_weights(tensors, prefix)), + model(config) { + model.init(params_ctx, tensors, prefix); + } + + std::string get_desc() override { return "ming_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 { + const auto* extra = diffusion_extra_as(inputs); + if (inputs.ref_latents != nullptr && !inputs.ref_latents->empty()) { + LOG_ERROR("Ming-Image reference-image conditioning is not supported"); + return {}; + } + if (inputs.context == nullptr || extra->direct_context == nullptr) { + LOG_ERROR("Ming-Image requires both query and direct text conditions"); + return {}; + } + auto graph = [&]() { + auto gf = new_graph_custom(ZImage::Z_IMAGE_GRAPH_SIZE); + auto x = make_input(*inputs.x); + auto t = make_input(*inputs.timesteps); + auto context = make_input(*inputs.context); + auto direct = make_input(*extra->direct_context); + GGML_ASSERT(x->ne[3] == 1); + const int64_t n_txt = context->ne[1] + direct->ne[1]; + const int64_t n_img = ((x->ne[0] + config.patch_size - 1) / config.patch_size) * + ((x->ne[1] + config.patch_size - 1) / config.patch_size); + auto padded = finish_rope_pe(Rope::gen_z_image_pe( + static_cast(x->ne[1]), static_cast(x->ne[0]), config.patch_size, 1, + static_cast(n_txt), ZImage::SEQ_MULTI_OF, {}, Rope::RefIndexMode::FIXED, + config.theta, config.axes_dim)); + // Zero-masked alignment tokens cannot affect valid queries. Omit them while + // retaining the padded caption length used to position image tokens. + const size_t stride = config.axes_dim_sum * 2; + const int64_t padded_txt = n_txt + Rope::bound_mod(static_cast(n_txt), ZImage::SEQ_MULTI_OF); + pe_values.assign(padded.begin(), padded.begin() + n_txt * stride); + pe_values.insert(pe_values.end(), padded.begin() + padded_txt * stride, + padded.begin() + (padded_txt + n_img) * stride); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, n_txt + n_img); + set_backend_tensor_data(pe, pe_values.data()); + auto ctx = get_context(); + auto out = model.forward(&ctx, x, t, context, direct, pe); + ggml_build_forward_expand(gf, out); + return gf; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(graph, n_threads, false), inputs.x->dim()); + } + }; +} + +#endif // __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__ diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index c82b9d71..957d575c 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -141,6 +141,10 @@ struct LLaDAImageDiffusionExtra { const sd::Tensor* semantic = nullptr; }; +struct MingImageDiffusionExtra { + const sd::Tensor* direct_context = nullptr; +}; + using DiffusionExtraParams = std::variant; + LLaDAImageDiffusionExtra, + MingImageDiffusionExtra>; struct DiffusionParams { const sd::Tensor* x = nullptr; diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index d654bd38..2b330a1a 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -140,7 +140,8 @@ namespace ZImage { int64_t num_kv_heads, bool qk_norm, bool norm_elementwise_affine = true, - bool split_qkv = false) + bool split_qkv = false, + float qk_norm_eps = 1e-6f) : 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) { @@ -153,8 +154,8 @@ namespace ZImage { blocks["out"] = std::make_shared(num_heads * head_dim, hidden_size, false, false, false, scale); } if (qk_norm) { - blocks["q_norm"] = std::make_shared(head_dim, 1e-06f, norm_elementwise_affine); - blocks["k_norm"] = std::make_shared(head_dim, 1e-06f, norm_elementwise_affine); + blocks["q_norm"] = std::make_shared(head_dim, qk_norm_eps, norm_elementwise_affine); + blocks["k_norm"] = std::make_shared(head_dim, qk_norm_eps, norm_elementwise_affine); } } @@ -318,9 +319,10 @@ namespace ZImage { bool qk_norm, bool modulation = true, bool norm_elementwise_affine = true, - bool split_qkv = false) + bool split_qkv = false, + float qk_norm_eps = 1e-6f) : modulation(modulation) { - blocks["attention"] = std::make_shared(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm, norm_elementwise_affine, split_qkv); + blocks["attention"] = std::make_shared(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm, norm_elementwise_affine, split_qkv, qk_norm_eps); blocks["feed_forward"] = std::make_shared(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier); blocks["attention_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); blocks["ffn_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index da59cfab..20521cc7 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -50,6 +50,8 @@ namespace LLM { GEMMA4_12B, GPT_OSS_20B, LLADA2_MOE, + BAILING_MOE, + QWEN2, ARCH_COUNT, }; @@ -64,6 +66,8 @@ namespace LLM { "gemma4_12b", "gpt_oss_20b", "llada2_moe", + "bailing_moe", + "qwen2", }; enum class MLPActivation { @@ -225,7 +229,7 @@ namespace LLM { config.intermediate_size = 9216; config.num_layers = 26; config.vocab_size = 256000; - } else if (arch == LLMArch::LLADA2_MOE) { + } else if (arch == LLMArch::LLADA2_MOE || arch == LLMArch::BAILING_MOE) { config.head_dim = 128; config.num_heads = 16; config.num_kv_heads = 4; @@ -240,7 +244,7 @@ namespace LLM { config.max_position_embeddings = 16384; config.rope_thetas = {600000.f}; config.qkv_fused = true; - config.bidirectional = true; + config.bidirectional = arch == LLMArch::LLADA2_MOE; config.partial_rotary = 0.5f; config.num_experts = 256; config.num_experts_per_tok = 8; @@ -250,6 +254,18 @@ namespace LLM { config.n_group = 8; config.topk_group = 4; config.routed_scaling_factor = 2.5f; + if (arch == LLMArch::BAILING_MOE) { + config.vocab_size = 157184; + config.max_position_embeddings = 32768; + } + } else if (arch == LLMArch::QWEN2) { + config.hidden_size = 1536; + config.intermediate_size = 8960; + config.num_heads = 12; + config.num_kv_heads = 2; + config.vocab_size = 151936; + config.max_position_embeddings = 32768; + config.bidirectional = true; } else if (arch == LLMArch::GPT_OSS_20B) { config.head_dim = 64; config.num_heads = 64; @@ -471,6 +487,8 @@ namespace LLM { int64_t n_group; int64_t topk_group; float routed_scaling_factor; + bool image_router; + bool fused_experts = false; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, @@ -488,6 +506,10 @@ namespace LLM { // 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); + if (image_router) { + params["image_gate.weight"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_experts); + params["image_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); @@ -506,8 +528,14 @@ namespace LLM { } }; - 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); + fused_experts = tensor_storage_map.count(prefix + "experts.gate_up_proj.weight") != 0; + if (fused_experts) { + auto type = supported_type(get_type(prefix + "experts.gate_up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size); + declare_experts("experts.gate_up_proj.weight", type, hidden_size, 2 * moe_intermediate_size); + } else { + 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); } @@ -519,7 +547,8 @@ namespace LLM { 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) { + routed_scaling_factor(config.routed_scaling_factor), + image_router(config.arch == LLMArch::BAILING_MOE) { if (config.num_shared_experts > 0) { blocks["shared_experts"] = std::make_shared(config.hidden_size, config.moe_intermediate_size * config.num_shared_experts, @@ -582,7 +611,7 @@ namespace LLM { return ggml_mul_mat_id(ctx->ggml_ctx, w, x, selected_experts); } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* image_mask = nullptr) { // 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); @@ -596,10 +625,21 @@ namespace LLM { 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] + auto bias = params["gate.expert_bias"]; + if (image_router) { + GGML_ASSERT(image_mask != nullptr); + auto mask = ggml_reshape_2d(gctx, image_mask, 1, n_token_total); + auto inverse_mask = ggml_scale_bias(gctx, mask, -1.f, 1.f); + auto image_logits = ggml_reshape_2d(gctx, ggml_mul_mat(gctx, params["image_gate.weight"], x), num_experts, n_token_total); + logits = ggml_add(gctx, ggml_mul(gctx, logits, inverse_mask), ggml_mul(gctx, image_logits, mask)); + bias = ggml_add(gctx, + ggml_mul(gctx, ggml_repeat_4d(gctx, bias, num_experts, n_token_total, 1, 1), inverse_mask), + ggml_mul(gctx, ggml_repeat_4d(gctx, params["image_gate.expert_bias"], num_experts, n_token_total, 1, 1), mask)); + } + auto scores = ggml_sigmoid(gctx, logits); // The bias steers selection only; the combine weights come from the unbiased scores. - auto routing = ggml_add(gctx, scores, params["gate.expert_bias"]); + auto routing = ggml_add(gctx, scores, 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] @@ -614,12 +654,17 @@ namespace LLM { 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); + auto xf = ggml_reshape_3d(gctx, x, hidden_size, 1, n_token_total); + ggml_tensor* activated; + if (fused_experts) { + activated = ggml_swiglu(gctx, expert_linear(ctx, "experts.gate_up_proj.weight", xf, selected_experts)); + } else { + auto gate = expert_linear(ctx, "experts.gate_proj.weight", xf, selected_experts); + auto up = expert_linear(ctx, "experts.up_proj.weight", xf, selected_experts); + 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) { @@ -1428,7 +1473,9 @@ namespace LLM { ggml_tensor* x, ggml_tensor* input_pos, ggml_tensor* attention_mask = nullptr, - int rope_index = 0) { + int rope_index = 0, + ggml_tensor* rope_cos = nullptr, + ggml_tensor* rope_sin = nullptr) { // x: [N, n_token, hidden_size] int64_t n_token = x->ne[1]; int64_t N = x->ne[2]; @@ -1471,13 +1518,28 @@ namespace LLM { v = ggml_rms_norm(ctx->ggml_ctx, v, rms_norm_eps); } - if (arch == LLMArch::MISTRAL_SMALL_3_2) { + if (rope_cos != nullptr) { + GGML_ASSERT(rope_sin != nullptr); + // Bailing video RoPE interleaves spatial frequencies within a partial NEOX head. + auto rotate = [&](ggml_tensor* input) { + auto gctx = ctx->ggml_ctx; + int64_t half = rope_cos->ne[0]; + auto first = ggml_ext_slice(gctx, input, 0, 0, half); + auto second = ggml_ext_slice(gctx, input, 0, half, 2 * half); + auto left = ggml_sub(gctx, ggml_mul(gctx, first, rope_cos), ggml_mul(gctx, second, rope_sin)); + auto right = ggml_add(gctx, ggml_mul(gctx, second, rope_cos), ggml_mul(gctx, first, rope_sin)); + auto rotated = ggml_concat(gctx, left, right, 0); + return ggml_concat(gctx, rotated, ggml_ext_slice(gctx, input, 0, 2 * half, input->ne[0]), 0); + }; + q = rotate(q); + k = rotate(k); + } else if (arch == LLMArch::MISTRAL_SMALL_3_2) { q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); } else if (arch == LLMArch::MINISTRAL_3_3B) { q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 262144, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 262144, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); - } else if (arch == LLMArch::QWEN3) { + } else if (arch == LLMArch::QWEN3 || arch == LLMArch::QWEN2) { q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f); } else if (arch == LLMArch::GPT_OSS_20B) { @@ -1724,7 +1786,7 @@ namespace LLM { blocks["self_attn"] = std::make_shared(config, sliding_attention == 0); if (config.arch == LLMArch::GPT_OSS_20B) { blocks["mlp"] = std::make_shared(config); - } else if (config.arch == LLMArch::LLADA2_MOE && layer_index >= config.first_k_dense_replace) { + } else if ((config.arch == LLMArch::LLADA2_MOE || config.arch == LLMArch::BAILING_MOE) && layer_index >= config.first_k_dense_replace) { blocks["mlp"] = std::make_shared(config); } else { blocks["mlp"] = std::make_shared(config.hidden_size, @@ -1746,7 +1808,10 @@ namespace LLM { ggml_tensor* x, ggml_tensor* input_pos, ggml_tensor* attention_mask = nullptr, - ggml_tensor* sliding_attention_mask = nullptr) { + ggml_tensor* sliding_attention_mask = nullptr, + ggml_tensor* image_mask = nullptr, + 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 input_layernorm = std::dynamic_pointer_cast(blocks["input_layernorm"]); @@ -1768,7 +1833,7 @@ namespace LLM { auto residual = x; x = input_layernorm->forward(ctx, x); - x = self_attn->forward(ctx, x, input_pos, block_attention_mask, rope_index); + 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); } @@ -1782,7 +1847,7 @@ namespace LLM { } else if (auto moe_mlp = std::dynamic_pointer_cast(blocks["mlp"])) { // LLaDA2 is dense for the first first_k_dense_replace layers and MoE afterwards, // so the block type varies per layer rather than per arch. - x = moe_mlp->forward(ctx, x); + x = moe_mlp->forward(ctx, x, image_mask); } else { auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); x = mlp->forward(ctx, x); @@ -1804,10 +1869,11 @@ namespace LLM { protected: int64_t num_layers; LLMConfig config; + std::string graph_cut_prefix; public: - TextModel(const LLMConfig& config) - : num_layers(config.num_layers), config(config) { + TextModel(const LLMConfig& config, const std::string& graph_cut_prefix = "llm.text") + : num_layers(config.num_layers), config(config), graph_cut_prefix(graph_cut_prefix) { blocks["embed_tokens"] = std::shared_ptr(new Embedding(config.vocab_size, config.hidden_size)); for (int i = 0; i < num_layers; i++) { blocks["layers." + std::to_string(i)] = std::shared_ptr(new TransformerBlock(config, i)); @@ -1831,7 +1897,10 @@ namespace LLM { std::set out_layers, const std::vector>>& deepstack_image_embeds = {}, ggml_tensor* sliding_attention_mask = nullptr, - bool return_all_hidden_states = false) { + bool return_all_hidden_states = false, + ggml_tensor* image_mask = nullptr, + ggml_tensor* rope_cos = nullptr, + ggml_tensor* rope_sin = nullptr) { auto norm = config.final_norm ? std::dynamic_pointer_cast(blocks["norm"]) : nullptr; std::vector intermediate_outputs; @@ -1843,18 +1912,18 @@ namespace LLM { intermediate_outputs.push_back(x); } - sd::ggml_graph_cut::mark_graph_cut(x, "llm.text.prelude", "x"); + sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".prelude", "x"); for (int i = 0; i < num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); - x = block->forward(ctx, x, input_pos, attention_mask, sliding_attention_mask); + x = block->forward(ctx, x, input_pos, attention_mask, sliding_attention_mask, image_mask, rope_cos, rope_sin); if (i < static_cast(deepstack_image_embeds.size())) { x = add_deepstack_image_embeds(ctx, x, deepstack_image_embeds[static_cast(i)]); } if (return_all_hidden_states || out_layers.size() > 1) { x = ggml_cont(ctx->ggml_ctx, x); } - sd::ggml_graph_cut::mark_graph_cut(x, "llm.text.layers." + std::to_string(i), "x"); + sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".layers." + std::to_string(i), "x"); if (return_all_hidden_states) { if (i + 1 < num_layers) { intermediate_outputs.push_back(x); diff --git a/src/model/te/ming_image_te.hpp b/src/model/te/ming_image_te.hpp new file mode 100644 index 00000000..ab3b8e6d --- /dev/null +++ b/src/model/te/ming_image_te.hpp @@ -0,0 +1,159 @@ +#ifndef __SD_MODEL_TE_MING_IMAGE_TE_HPP__ +#define __SD_MODEL_TE_MING_IMAGE_TE_HPP__ + +#include "llm.hpp" + +namespace MingImageTE { + struct MingImageTEConfig { + LLM::LLMConfig backbone; + LLM::LLMConfig connector; + int64_t num_queries = 256; + int64_t caption_dim = 2560; + int64_t diffusion_dim = 3840; + + static MingImageTEConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) { + MingImageTEConfig config; + for (const auto& entry : tensors) { + if (starts_with(entry.first, prefix + ".backbone.") && + contains(entry.first, ".mlp.experts.") && entry.second.type == GGML_TYPE_I8) { + throw std::runtime_error("Ming-Image INT8/W4A8 text encoder experts are not supported; use the BF16 text encoder"); + } + } + bool vision = false; + config.backbone = LLM::LLMConfig::detect_from_weights(tensors, prefix + ".backbone.", LLM::LLMArch::BAILING_MOE, vision); + config.connector = LLM::LLMConfig::detect_from_weights(tensors, prefix + ".connector.", LLM::LLMArch::QWEN2, vision); + const auto query = tensors.find(prefix + ".query_tokens_dict.16x16"); + const auto projection = tensors.find(prefix + ".proj_out.weight"); + const auto direct = tensors.find(prefix + ".proj_directvlm.1.weight"); + if (query == tensors.end() || projection == tensors.end() || direct == tensors.end()) { + throw std::runtime_error("Ming-Image requires the learned queries, connector and both condition projections"); + } + config.num_queries = query->second.ne[1]; + config.caption_dim = projection->second.ne[1]; + config.diffusion_dim = direct->second.ne[1]; + if (config.num_queries != 256 || config.caption_dim != 2560 || config.diffusion_dim != 3840 || + config.backbone.num_layers != 20 || config.backbone.hidden_size != 2048 || + config.connector.num_layers != 28 || config.connector.hidden_size != 1536) { + throw std::runtime_error("unsupported Ming-Image text encoder configuration"); + } + LOG_VERBOSE("ming_image_te: queries = %" PRId64 ", caption_dim = %" PRId64 ", diffusion_dim = %" PRId64, + config.num_queries, config.caption_dim, config.diffusion_dim); + return config; + } + }; + + struct ConnectorModel : LLM::TextModel { + explicit ConnectorModel(const LLM::LLMConfig& config) + : LLM::TextModel(config, "ming_image.connector") { + blocks.erase("embed_tokens"); + } + }; + + class MingImageTextModel : public GGMLBlock { + MingImageTEConfig config; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensors = {}, const std::string prefix = "") override { + params["query_tokens_dict.16x16"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.backbone.hidden_size, config.num_queries); + } + + public: + explicit MingImageTextModel(const MingImageTEConfig& config) + : config(config) { + blocks["backbone"] = std::make_shared(config.backbone, "ming_image.backbone"); + blocks["connector"] = std::make_shared(config.connector); + blocks["proj_in"] = std::make_shared(config.backbone.hidden_size, config.connector.hidden_size); + blocks["proj_out"] = std::make_shared(config.connector.hidden_size, config.caption_dim); + blocks["proj_directvlm.0"] = std::make_shared(config.backbone.hidden_size * 3, 1e-5f); + blocks["proj_directvlm.1"] = std::make_shared(config.backbone.hidden_size * 3, config.diffusion_dim); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* ids, int64_t prompt_length, ggml_tensor* mask, ggml_tensor* image_mask, ggml_tensor* cos, ggml_tensor* sin, ggml_tensor* connector_positions) { + auto gctx = ctx->ggml_ctx; + auto backbone = std::dynamic_pointer_cast(blocks["backbone"]); + auto connector = std::dynamic_pointer_cast(blocks["connector"]); + auto x = backbone->embed(ctx, ids); + const int64_t query_start = prompt_length + 1; + auto before = ggml_ext_slice(gctx, x, 1, 0, query_start); + auto after = ggml_ext_slice(gctx, x, 1, query_start + config.num_queries, x->ne[1]); + x = ggml_concat(gctx, ggml_concat(gctx, before, params["query_tokens_dict.16x16"], 1), after, 1); + // HF hidden_states[20] includes final RMSNorm; sd.cpp selects it as num_layers + 1. + x = backbone->forward_embeds(ctx, x, nullptr, mask, {5, 12, 21}, {}, nullptr, false, image_mask, cos, sin); + auto direct = ggml_ext_slice(gctx, x, 1, 0, prompt_length); + direct = std::dynamic_pointer_cast(blocks["proj_directvlm.0"])->forward(ctx, direct); + direct = std::dynamic_pointer_cast(blocks["proj_directvlm.1"])->forward(ctx, direct); + auto queries = ggml_ext_slice(gctx, x, 0, config.backbone.hidden_size * 2, config.backbone.hidden_size * 3); + queries = ggml_ext_slice(gctx, queries, 1, query_start, query_start + config.num_queries); + queries = std::dynamic_pointer_cast(blocks["proj_in"])->forward(ctx, queries); + queries = connector->forward_embeds(ctx, queries, connector_positions, nullptr, {}); + queries = std::dynamic_pointer_cast(blocks["proj_out"])->forward(ctx, queries); + queries = ggml_pad(gctx, queries, static_cast(config.diffusion_dim - config.caption_dim), 0, 0, 0); + return ggml_concat(gctx, queries, direct, 1); + } + }; + + struct MingImageTextRunner : GGMLRunner { + MingImageTEConfig config; + MingImageTextModel model; + + MingImageTextRunner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), config(MingImageTEConfig::detect_from_weights(tensors, prefix)), model(config) { + model.init(params_ctx, tensors, prefix); + } + + std::string get_desc() override { return "ming_image_text"; } + + void get_param_tensors(std::map& tensors, const std::string& prefix) { + model.get_param_tensors(tensors, prefix); + } + + void get_param_tensor_ops(std::map& ops) { + model.get_param_tensor_ops(ops); + } + + sd::Tensor compute(int n_threads, const std::vector& tokens) { + const int64_t prompt_length = tokens.size(); + const int64_t query_start = prompt_length + 1; + const int64_t total = prompt_length + config.num_queries + 2; + std::vector ids(tokens.begin(), tokens.end()); + ids.push_back(157158); + ids.insert(ids.end(), config.num_queries, 157157); + ids.push_back(157159); + auto input = sd::Tensor({total}, std::move(ids)); + sd::Tensor attention_mask({total, total}); + sd::Tensor image_mask({1, total}); + sd::Tensor cos({32, 1, total}), sin({32, 1, total}); + std::vector positions(config.num_queries); + std::iota(positions.begin(), positions.end(), 0); + auto connector_positions = sd::Tensor({config.num_queries}, std::move(positions)); + for (int64_t token = 0; token < total; ++token) { + const bool query = token >= query_start && token < query_start + config.num_queries; + image_mask[token] = query ? 1.f : 0.f; + for (int64_t key = 0; key < total; ++key) { + attention_mask[key + token * total] = key > token ? -INFINITY : 0.f; + } + // A 16x16 query bank is represented upstream as a [1, 2, 512] image grid, + // then spatially merged and centered to [1, 1, 256]. + int64_t temporal = query ? query_start : (token == total - 1 ? query_start + 1 : token); + int64_t width = query ? token - 127 : temporal; + for (int j = 0; j < 32; ++j) { + int64_t position = query && j < 24 && j % 2 ? width : temporal; + float frequency = 1.f / std::pow(600000.f, static_cast(2 * j) / 64.f); + float angle = static_cast(position) * frequency; + cos[token * 32 + j] = std::cos(angle); + sin[token * 32 + j] = std::sin(angle); + } + } + auto graph = [&]() { + auto gf = new_graph_custom(LLM::LLM_GRAPH_SIZE); + auto ctx = get_context(); + auto out = model.forward(&ctx, make_input(input), prompt_length, make_input(attention_mask), + make_input(image_mask), make_input(cos), make_input(sin), make_input(connector_positions)); + ggml_build_forward_expand(gf, out); + return gf; + }; + return take_or_empty(GGMLRunner::compute(graph, n_threads)); + } + }; +} + +#endif // __SD_MODEL_TE_MING_IMAGE_TE_HPP__ diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 58cff816..9dd30fdf 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -1084,7 +1084,7 @@ namespace WAN { _conv_num = 34; _enc_conv_num = 26; - } else if (version == VERSION_QWEN_IMAGE_LAYERED) { + } else if (version == VERSION_QWEN_IMAGE_LAYERED || version == VERSION_MING_IMAGE) { input_channels = 4; } @@ -1423,11 +1423,17 @@ namespace WAN { } sd::Tensor diffusion_to_vae_latents(const sd::Tensor& latents) override { + if (version == VERSION_MING_IMAGE) { + return latents / 8.0064f; + } auto [mean_tensor, std_tensor] = get_latents_mean_std(latents); return (latents * std_tensor) / scale_factor + mean_tensor; } sd::Tensor vae_to_diffusion_latents(const sd::Tensor& latents) override { + if (version == VERSION_MING_IMAGE) { + return latents * 8.0064f; + } auto [mean_tensor, std_tensor] = get_latents_mean_std(latents); return ((latents - mean_tensor) * scale_factor) / std_tensor; } diff --git a/src/model_loader.cpp b/src/model_loader.cpp index d96b597f..bcab0ac7 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -524,6 +524,9 @@ SDVersion ModelLoader::get_sd_version() const { return VERSION_LLADA_IMAGE; } if (tensor_storage.name.find("model.diffusion_model.cap_embedder.0.weight") != std::string::npos) { + if (tensor_storage_map.find("text_encoders.llm.connector.layers.0.self_attn.q_proj.weight") != tensor_storage_map.end()) { + return VERSION_MING_IMAGE; + } return VERSION_Z_IMAGE; } if (tensor_storage.name.find("double_stream_layers.0.img_instruct_attn.processor.img_to_q.weight") != std::string::npos) { diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index bff51c95..0daeda66 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -105,7 +105,25 @@ std::string convert_open_clip_to_hf_clip_name(std::string 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 std::string convert_ming_image_te_name(std::string name) { + if (name == "llm.query_tokens") { + name = "llm.query_tokens_dict.16x16"; + } + static const std::vector> name_map = { + {"thinker.", "backbone."}, + {"attention.", "self_attn."}, + {"self_attn.dense.", "self_attn.o_proj."}, + {"gate.proj.", "gate."}, + }; + replace_with_name_map(name, name_map); + return name; +} + +std::string convert_cond_stage_model_name(std::string name, std::string prefix, SDVersion version) { + if (version == VERSION_MING_IMAGE && prefix == "text_encoders." && starts_with(name, "llm.")) { + name = convert_ming_image_te_name(name); + } + static const std::vector> clip_name_map{ {"transformer.text_projection.weight", "transformer.text_model.text_projection"}, {"model.text_projection.weight", "transformer.text_model.text_projection"}, @@ -994,6 +1012,17 @@ static std::string convert_diffusers_dit_to_original_pixart(std::string name) { return prefix + name; } +static std::string convert_ming_image_dit_name(std::string name) { + static const std::vector> name_map = { + {"all_x_embedder.2-1.", "x_embedder."}, + {"all_final_layer.2-1.", "final_layer."}, + {"attention.norm_q.", "attention.q_norm."}, + {"attention.norm_k.", "attention.k_norm."}, + }; + replace_with_name_map(name, name_map); + return name; +} + std::string convert_diffusion_model_name(std::string name, std::string prefix, SDVersion version) { if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) { name = convert_diffusers_unet_to_original_sd1(name); @@ -1005,6 +1034,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S name = convert_diffusers_dit_to_original_flux(name); } else if (sd_version_is_hunyuan_video(version)) { name = convert_hunyuan_video_to_original_flux(name); + } else if (version == VERSION_MING_IMAGE) { + name = convert_ming_image_dit_name(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)) { @@ -1668,7 +1699,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) { { for (const auto& prefix : cond_stage_model_prefix_vec) { if (starts_with(name, prefix)) { - name = convert_cond_stage_model_name(name.substr(prefix.size()), prefix); + name = convert_cond_stage_model_name(name.substr(prefix.size()), prefix, version); name = prefix + name; break; } diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index aacbec99..359b27b8 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -106,6 +106,7 @@ const char* model_version_to_str[] = { "LLaDA-Image", "ESRGAN", "PixArt", + "Ming-Image", }; static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]), @@ -1210,6 +1211,11 @@ bool StableDiffusionGGML::validate_and_load_runners() { LOG_VERBOSE("validating model metadata"); std::set ignore_tensors; + if (version == VERSION_MING_IMAGE) { + ignore_tensors.insert("text_encoders.llm.vision."); + ignore_tensors.insert("text_encoders.llm.linear_proj."); + ignore_tensors.insert("text_encoders.llm.backbone.lm_head."); + } if (use_tae && !tae_preview_only) { ignore_tensors.insert("first_stage_model."); } @@ -1370,6 +1376,7 @@ bool StableDiffusionGGML::build_denoiser() { sd_version_is_anima(version) || sd_version_is_ernie_image(version) || sd_version_is_z_image(version) || + version == VERSION_MING_IMAGE || sd_version_is_llada_image(version) || sd_version_is_boogu_image(version) || sd_version_is_pid(version) || @@ -1391,6 +1398,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 (version == VERSION_MING_IMAGE) { + default_flow_shift = INFINITY; } else if (sd_version_is_llada_image(version)) { default_flow_shift = 1.0f; // unused: LLADA_IMAGE_SCHEDULER builds a fixed grid } else { @@ -1451,6 +1460,8 @@ bool StableDiffusionGGML::build_denoiser() { } else if (sd_version_is_minimax_h3(version)) { LOG_INFO("running in MiniMax H3 AV FLOW mode"); denoiser = std::make_shared(default_flow_shift, 3.f, get_latent_channel()); + } else if (version == VERSION_MING_IMAGE) { + denoiser = std::make_shared(); } else { LOG_INFO("running in FLOW mode"); denoiser = std::make_shared(); @@ -2161,7 +2172,7 @@ std::vector StableDiffusionGGML::prepare_sample_timesteps(float sigma, if (version == VERSION_HIDREAM_O1) { return std::vector{1.0f - (t / static_cast(TIMESTEPS))}; } - if (sd_version_is_z_image(version) || sd_version_is_ideogram4(version)) { + if (sd_version_is_z_image(version) || sd_version_is_ideogram4(version) || version == VERSION_MING_IMAGE) { return std::vector{1000.f - t}; } return std::vector{t}; @@ -2514,6 +2525,9 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr::zeros_like(cond.c_vector); + } else if (sd->version == VERSION_MING_IMAGE) { + uncond.c_crossattn = sd::Tensor::zeros_like(cond.c_crossattn); + for (const auto& extra : cond.extra_c_crossattns) { + uncond.extra_c_crossattns.push_back(sd::Tensor::zeros_like(extra)); + } } else if (sd_version_is_sensenova_u1(sd->version)) { auto* sensenova_conditioner = static_cast(sd->cond_stage_model.get()); uncond = sensenova_conditioner->get_unconditional_condition(request->negative_prompt); diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index afc6ee8e..f392c364 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -24,6 +24,7 @@ #include "model/diffusion/llada_image.hpp" #include "model/diffusion/ltxv.hpp" #include "model/diffusion/mage_flow.hpp" +#include "model/diffusion/ming_image.hpp" #include "model/diffusion/minimax_h3.hpp" #include "model/diffusion/minit2i.hpp" #include "model/diffusion/mmdit.hpp" @@ -374,6 +375,11 @@ namespace sd::model_builders { tensor_storage_map, "model.diffusion_model", weight_manager); + } else if (version == VERSION_MING_IMAGE) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, weight_manager, tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, "model.diffusion_model", weight_manager); } else if (sd_version_is_z_image(version)) { result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), tensor_storage_map, diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index 15c12608..02ebe57e 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -1369,6 +1369,42 @@ struct DiscreteFlowDenoiser : public Denoiser { } }; +struct MingImageFlowDenoiser : DiscreteFlowDenoiser { + float resolution_shift = std::exp(1.35f); + + MingImageFlowDenoiser() + : DiscreteFlowDenoiser(INFINITY) {} + + float sigma_min() override { return 0.f; } + float sigma_max() override { return 1.f; } + + float t_to_sigma(float t) override { + return time_snr_shift(std::isfinite(shift) ? shift : resolution_shift, (t + 1.f) / 1000.f); + } + + std::vector get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler, SDVersion version, const char* extra_sample_args = nullptr) override { + const float tokens = static_cast(image_seq_len) / 4.f; + const float mu = tokens >= 4096.f ? 1.35f : 0.5f + (tokens - 256.f) * (1.15f - 0.5f) / (4096.f - 256.f); + resolution_shift = std::exp(mu); + if (scheduler != DISCRETE_SCHEDULER) { + return DiscreteFlowDenoiser::get_sigmas(n, image_seq_len, scheduler, version, extra_sample_args); + } + if (n == 0) { + return {}; + } + const float factor = std::isfinite(shift) ? shift : resolution_shift; + std::vector sigmas; + sigmas.reserve(n + 1); + for (uint32_t i = 0; i < n; ++i) { + const float t = n == 1 ? 1.f : 1.f - static_cast(i) / static_cast(n - 1); + sigmas.push_back(time_snr_shift(factor, t)); + } + // Upstream sets sigma_min to zero before linspace, then appends the terminal zero. + sigmas.push_back(0.f); + return sigmas; + } +}; + struct H3AVFlowDenoiser : public DiscreteFlowDenoiser { int64_t video_channels; float audio_shift; @@ -1783,7 +1819,10 @@ static sd::Tensor sample_euler(denoise_cb_t model, const std::vector& sigmas) { int steps = static_cast(sigmas.size()) - 1; for (int i = 0; i < steps; i++) { - float sigma = sigmas[i]; + float sigma = sigmas[i]; + if (sigma == sigmas[i + 1]) { + continue; + } auto denoised_opt = model(x, sigma, i + 1); if (denoised_opt.pred.empty()) { return {};