diff --git a/README.md b/README.md index 6ba81d89a..cb0efe2d1 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,7 @@ API and command-line option may change frequently.*** - [LongCat Image](./docs/longcat_image.md) - [Z-Image](./docs/z_image.md) - [MiniT2I](./docs/minit2i.md) + - [SenseNova U1.5](./docs/sensenova_u1.md) - [Ovis-Image](./docs/ovis_image.md) - [Anima](./docs/anima.md) - [ERNIE-Image](./docs/ernie_image.md) diff --git a/docs/sensenova_u1.md b/docs/sensenova_u1.md new file mode 100644 index 000000000..fd19c4d38 --- /dev/null +++ b/docs/sensenova_u1.md @@ -0,0 +1,46 @@ +# How to Use + +SenseNova U1.5 is an 8B MoT model that performs diffusion directly in RGB pixel +space. It does not require a separate text encoder or VAE. + +## Download weights + +- Download SenseNova U1.5 8B MoT + - safetensors: https://huggingface.co/sensenova/SenseNova-U1.5-8B-MoT + +Pass the complete downloaded repository directory to `--model`. The directory +must contain `model.safetensors.index.json`, every referenced Safetensors shard, +and the tokenizer files. + +## Examples + +### CUDA + +```bash +./bin/sd-cli \ + --model /path/to/SenseNova-U1.5-8B-MoT \ + --prompt "a red cube on a white background" \ + --width 2048 \ + --height 2048 \ + --steps 50 \ + --cfg-scale 4 \ + --flow-shift 3 \ + --seed 42 \ + --sampling-method euler \ + --rng cuda \ + --fa \ + --output output.png +``` + +## Notes + +- To match the official non-thinking text-to-image pipeline, use 50 Euler + steps, CFG 4, flow shift 3, seed 42, CUDA RNG, and an empty negative prompt. +- Width and height must be multiples of 32. The trained 1:1 resolution is + 2048x2048; lower resolutions are useful for smoke tests but are outside the + training buckets. +- The SenseNova prompt template and unconditional prompt are built + automatically. +- This implementation supports non-thinking text-to-image generation. Image + editing, visual understanding, interleaved generation, and thinking-mode + prompt expansion are not implemented. diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index d5cda6d9a..c0971f4f4 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -92,6 +92,7 @@ enum prediction_t { FLUX_FLOW_PRED, SEFI_FLOW_PRED, MINIT2I_FLOW_PRED, + SENSENOVA_U1_FLOW_PRED, PREDICTION_COUNT }; diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index 8968676b3..ea527d253 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -14,6 +14,7 @@ #include "model/te/llm.hpp" #include "model/te/t5.hpp" #include "model_loader.h" +#include "tokenizers/sensenova_u1_tokenizer.h" struct SDCondition { sd::Tensor c_crossattn; @@ -1667,6 +1668,71 @@ struct MiniT2IConditioner : public Conditioner { } }; +struct SenseNovaU1Conditioner : public Conditioner { + static constexpr size_t kMaxPromptTokens = 12288; + SenseNovaU1Tokenizer tokenizer; + + void get_param_tensors(std::map& tensors) override { + SD_UNUSED(tensors); + } + + void set_flash_attention_enabled(bool enabled) override { + SD_UNUSED(enabled); + } + + static std::string build_query(const std::string& text, bool is_negative) { + static const std::string kSystemMessage = + "You are an image generation and editing assistant that accurately understands and executes user intent.\n\n" + "You support two modes:\n\n1. Think Mode:\nIf the task requires reasoning, you MUST start with a " + " block. Put all reasoning inside the block using plain text. DO NOT include any image tags. " + "Keep it reasonable and directly useful for producing the final image.\n\n2. Non-Think Mode:\nIf no reasoning " + "is needed, directly produce the final image.\n\nTask Types:\n\nA. Text-to-Image Generation:\n- Generate a " + "high-quality image based on the user's description.\n- Ensure visual clarity, semantic consistency, and " + "completeness.\n- DO NOT introduce elements that contradict or override the user's intent.\n\nB. Image Editing:\n" + "- Use the provided image(s) as input or reference for modification or transformation.\n- The result can be an " + "edited image or a new image based on the reference(s).\n- Preserve all unspecified attributes unless explicitly " + "changed.\n\nGeneral Rules:\n- For any visible text in the image, follow the language specified for the rendered " + "text in the user's description, not the language of the prompt. If no language is specified, use the user's input " + "language."; + + std::string query; + if (!is_negative) { + query += "<|im_start|>system\n"; + query += kSystemMessage; + query += "<|im_end|>\n"; + } + query += "<|im_start|>user\n"; + query += text; + query += "<|im_end|>\n<|im_start|>assistant\n"; + query += is_negative ? "" : "\n\n\n\n"; + return query; + } + + SDCondition tokenize_condition(const std::string& text, bool is_negative) { + auto tokens = tokenizer.encode(build_query(text, is_negative)); + if (tokens.empty() || tokens.size() > kMaxPromptTokens) { + LOG_ERROR("SenseNova U1.5 prompt token count %zu is outside [1, %zu]", + tokens.size(), + kMaxPromptTokens); + return {}; + } + + SDCondition result; + result.c_input_ids = sd::Tensor({static_cast(tokens.size())}, tokens); + return result; + } + + SDCondition get_learned_condition(int n_threads, + const ConditionerParams& conditioner_params) override { + SD_UNUSED(n_threads); + return tokenize_condition(conditioner_params.text, false); + } + + SDCondition get_unconditional_condition(const std::string& text) { + return tokenize_condition(text, true); + } +}; + struct AnimaConditioner : public Conditioner { std::shared_ptr qwen_tokenizer; T5UniGramTokenizer t5_tokenizer; diff --git a/src/model.h b/src/model.h index f59f1de7d..4b3ed316f 100644 --- a/src/model.h +++ b/src/model.h @@ -57,6 +57,7 @@ enum SDVersion { VERSION_SEFI_IMAGE, VERSION_KREA2, VERSION_MAGE_FLOW, + VERSION_SENSENOVA_U1_5, VERSION_ESRGAN, VERSION_COUNT, }; @@ -237,6 +238,10 @@ static inline bool sd_version_is_mage_flow(SDVersion version) { return version == VERSION_MAGE_FLOW; } +static inline bool sd_version_is_sensenova_u1(SDVersion version) { + return version == VERSION_SENSENOVA_U1_5; +} + static inline bool sd_version_uses_flux_vae(SDVersion version) { if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) { return true; @@ -295,7 +300,8 @@ static inline bool sd_version_is_dit(SDVersion version) { sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version) || - sd_version_is_mage_flow(version)) { + sd_version_is_mage_flow(version) || + sd_version_is_sensenova_u1(version)) { return true; } return false; diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 070ca53d4..5c8156ae9 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -114,6 +114,10 @@ struct MiniT2IDiffusionExtra { const sd::Tensor* mask = nullptr; }; +struct SenseNovaU1DiffusionExtra { + const sd::Tensor* input_ids = nullptr; +}; + struct HunyuanVideoDiffusionExtra { const sd::Tensor* guidance = nullptr; const sd::Tensor* byt5 = nullptr; @@ -131,6 +135,7 @@ using DiffusionExtraParams = std::variant; struct DiffusionParams { diff --git a/src/model/diffusion/sensenova_u1.h b/src/model/diffusion/sensenova_u1.h new file mode 100644 index 000000000..bc1e524bf --- /dev/null +++ b/src/model/diffusion/sensenova_u1.h @@ -0,0 +1,851 @@ +#ifndef __SD_MODEL_DIFFUSION_SENSENOVA_U1_H__ +#define __SD_MODEL_DIFFUSION_SENSENOVA_U1_H__ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/ggml_extend.hpp" +#include "model/diffusion/dit.hpp" +#include "model/diffusion/model.hpp" +#include "model/te/llm.hpp" +#include "model_loader.h" + +namespace SenseNovaU1 { + constexpr int SENSENOVA_U1_GRAPH_SIZE = 327680; + + struct SenseNovaU1Config { + int64_t hidden_size = 4096; + int64_t intermediate_size = 12288; + int64_t num_layers = 42; + int64_t num_heads = 32; + int64_t num_kv_heads = 8; + int64_t head_dim = 128; + int64_t vocab_size = 151936; + int64_t max_position_embeddings = 262144; + int64_t max_position_embeddings_hw = 10000; + int64_t vision_hidden_size = 1024; + int64_t patch_size = 16; + int64_t vision_downsample_factor = 2; + int64_t in_channels = 3; + int64_t timestep_embedding_size = 256; + float rms_norm_eps = 1e-6f; + float rope_theta = 5000000.f; + float rope_theta_hw = 10000.f; + float noise_scale_base_image_seq_len = 64.f; + float noise_scale_max_value = 16.f; + float t_eps = 0.02f; + bool add_noise_scale_embedding = true; + + int64_t image_token_stride() const { + return patch_size * vision_downsample_factor; + } + + static SenseNovaU1Config detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix) { + SenseNovaU1Config config; + config.num_layers = 0; + const std::string root = prefix.empty() ? "" : prefix + "."; + + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, root)) { + continue; + } + if (ends_with(name, "language_model.model.embed_tokens.weight") && tensor_storage.n_dims == 2) { + config.hidden_size = tensor_storage.ne[0]; + config.vocab_size = tensor_storage.ne[1]; + } else if (ends_with(name, "language_model.model.layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) { + config.intermediate_size = tensor_storage.ne[1]; + } else if (ends_with(name, "language_model.model.layers.0.self_attn.q_proj.weight") && tensor_storage.n_dims == 2) { + config.num_heads = tensor_storage.ne[1] / config.head_dim; + } else if (ends_with(name, "language_model.model.layers.0.self_attn.k_proj.weight") && tensor_storage.n_dims == 2) { + config.num_kv_heads = tensor_storage.ne[1] / config.head_dim; + } else if (ends_with(name, "fm_modules.vision_model_mot_gen.embeddings.patch_embedding.weight") && tensor_storage.n_dims == 4) { + config.patch_size = tensor_storage.ne[0]; + config.in_channels = tensor_storage.ne[2]; + config.vision_hidden_size = tensor_storage.ne[3]; + } else if (ends_with(name, "fm_modules.vision_model_mot_gen.embeddings.dense_embedding.weight") && tensor_storage.n_dims == 4) { + config.vision_downsample_factor = tensor_storage.ne[0]; + } + + const std::string layer_prefix = root + "language_model.model.layers."; + if (starts_with(name, layer_prefix)) { + const char* index_begin = name.c_str() + layer_prefix.size(); + config.num_layers = std::max(config.num_layers, std::strtoll(index_begin, nullptr, 10) + 1); + } + } + + if (config.num_layers == 0) { + config.num_layers = 42; + } + config.add_noise_scale_embedding = tensor_storage_map.find(root + "fm_modules.noise_scale_embedder.mlp.0.weight") != tensor_storage_map.end(); + + LOG_DEBUG("sensenova-u1.5: layers=%" PRId64 ", hidden=%" PRId64 ", intermediate=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", patch=%" PRId64 "x%" PRId64, + config.num_layers, + config.hidden_size, + config.intermediate_size, + config.num_heads, + config.num_kv_heads, + config.patch_size, + config.vision_downsample_factor); + return config; + } + }; + + class StorageConv2d : public Conv2d { + protected: + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + this->prefix = prefix; + ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F16); + params["weight"] = ggml_new_tensor_4d(ctx, + wtype, + kernel_size.second, + kernel_size.first, + in_channels, + out_channels); + if (bias) { + params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); + } + } + + public: + StorageConv2d(int64_t in_channels, + int64_t out_channels, + std::pair kernel_size, + std::pair stride = {1, 1}, + std::pair padding = {0, 0}, + bool bias = true) + : Conv2d(in_channels, + out_channels, + kernel_size, + stride, + padding, + {1, 1}, + bias) {} + }; + + struct TimestepEmbedder : public GGMLBlock { + int64_t frequency_embedding_size; + + TimestepEmbedder(int64_t hidden_size, int64_t frequency_embedding_size = 256) + : frequency_embedding_size(frequency_embedding_size) { + blocks["mlp.0"] = std::make_shared(frequency_embedding_size, hidden_size, true); + blocks["mlp.2"] = std::make_shared(hidden_size, hidden_size, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timesteps) { + auto mlp_0 = std::dynamic_pointer_cast(blocks["mlp.0"]); + auto mlp_2 = std::dynamic_pointer_cast(blocks["mlp.2"]); + auto x = ggml_ext_timestep_embedding(ctx->ggml_ctx, + timesteps, + static_cast(frequency_embedding_size), + 10000.f, + 1.f); + x = mlp_0->forward(ctx, x); + x = ggml_silu_inplace(ctx->ggml_ctx, x); + return mlp_2->forward(ctx, x); + } + }; + + inline ggml_tensor* apply_vision_rope(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* position_x, + ggml_tensor* position_y, + float theta, + int max_position) { + GGML_ASSERT(x->ne[0] % 2 == 0); + // ggml_rope_ext addresses positions through ne[2]. The vision + // embeddings arrive as [hidden, tokens, batch], so add the singleton + // head axis used by the RoPE kernel: [hidden, 1, tokens, batch]. + x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]); + const int64_t half = x->ne[0] / 2; + auto x_part = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, half); + auto y_part = ggml_ext_slice(ctx->ggml_ctx, x, 0, half, x->ne[0]); + x_part = ggml_rope_ext(ctx->ggml_ctx, + x_part, + position_x, + nullptr, + static_cast(half), + GGML_ROPE_TYPE_NORMAL, + max_position, + theta, + 1.f, + 0.f, + 1.f, + 32.f, + 1.f); + y_part = ggml_rope_ext(ctx->ggml_ctx, + y_part, + position_y, + nullptr, + static_cast(half), + GGML_ROPE_TYPE_NORMAL, + max_position, + theta, + 1.f, + 0.f, + 1.f, + 32.f, + 1.f); + return ggml_concat(ctx->ggml_ctx, x_part, y_part, 0); + } + + struct VisionEmbeddings : public GGMLBlock { + SenseNovaU1Config config; + + explicit VisionEmbeddings(const SenseNovaU1Config& config) + : config(config) { + blocks["patch_embedding"] = std::make_shared(config.in_channels, + config.vision_hidden_size, + std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}, + std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}, + std::pair{0, 0}, + true); + blocks["dense_embedding"] = std::make_shared(config.vision_hidden_size, + config.hidden_size, + std::pair{static_cast(config.vision_downsample_factor), static_cast(config.vision_downsample_factor)}, + std::pair{static_cast(config.vision_downsample_factor), static_cast(config.vision_downsample_factor)}, + std::pair{0, 0}, + true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* image, + ggml_tensor* position_x, + ggml_tensor* position_y) { + auto patch_embedding = std::dynamic_pointer_cast(blocks["patch_embedding"]); + auto dense_embedding = std::dynamic_pointer_cast(blocks["dense_embedding"]); + + auto x = patch_embedding->forward(ctx, image); + x = ggml_gelu_erf(ctx->ggml_ctx, x); + + const int64_t grid_w = x->ne[0]; + const int64_t grid_h = x->ne[1]; + const int64_t batch = x->ne[3]; + x = ggml_reshape_3d(ctx->ggml_ctx, x, grid_w * grid_h, x->ne[2], batch); + x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = apply_vision_rope(ctx, + x, + position_x, + position_y, + config.rope_theta_hw, + static_cast(config.max_position_embeddings_hw)); + x = ggml_reshape_4d(ctx->ggml_ctx, x, config.vision_hidden_size, grid_w, grid_h, batch); + x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); + x = dense_embedding->forward(ctx, x); + + const int64_t token_w = x->ne[0]; + const int64_t token_h = x->ne[1]; + x = ggml_reshape_3d(ctx->ggml_ctx, x, token_w * token_h, x->ne[2], x->ne[3]); + return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + } + }; + + inline ggml_tensor* pixel_shuffle(GGMLRunnerContext* ctx, + ggml_tensor* x, + int upscale_factor) { + GGML_ASSERT(upscale_factor > 0); + const int64_t h = x->ne[1]; + const int64_t w = x->ne[0]; + GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0); + x = ggml_ext_cont(ctx->ggml_ctx, + ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); + x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]); + return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true); + } + + struct PixelDecoder : public GGMLBlock { + explicit PixelDecoder(const SenseNovaU1Config& config) { + blocks["conv1"] = std::make_shared(config.hidden_size / 4, + 1024, + std::pair{3, 3}, + std::pair{1, 1}, + std::pair{1, 1}, + true); + blocks["conv2"] = std::make_shared(256, + 192, + std::pair{3, 3}, + std::pair{1, 1}, + std::pair{1, 1}, + true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); + x = pixel_shuffle(ctx, x, 2); + x = conv1->forward(ctx, x); + x = ggml_gelu_erf(ctx->ggml_ctx, x); + x = pixel_shuffle(ctx, x, 2); + x = conv2->forward(ctx, x); + return pixel_shuffle(ctx, x, 8); + } + }; + + enum class Branch { + UNDERSTANDING, + GENERATION, + }; + + struct Attention : public GGMLBlock { + SenseNovaU1Config config; + int layer_index; + + Attention(const SenseNovaU1Config& config, int layer_index) + : config(config), layer_index(layer_index) { + blocks["q_proj"] = std::make_shared(config.hidden_size, config.num_heads * config.head_dim, false); + blocks["k_proj"] = std::make_shared(config.hidden_size, config.num_kv_heads * config.head_dim, false); + blocks["v_proj"] = std::make_shared(config.hidden_size, config.num_kv_heads * config.head_dim, false); + blocks["o_proj"] = std::make_shared(config.num_heads * config.head_dim, config.hidden_size, false); + blocks["q_proj_mot_gen"] = std::make_shared(config.hidden_size, config.num_heads * config.head_dim, false); + blocks["k_proj_mot_gen"] = std::make_shared(config.hidden_size, config.num_kv_heads * config.head_dim, false); + blocks["v_proj_mot_gen"] = std::make_shared(config.hidden_size, config.num_kv_heads * config.head_dim, false); + blocks["o_proj_mot_gen"] = std::make_shared(config.num_heads * config.head_dim, config.hidden_size, false); + + const int64_t axis_dim = config.head_dim / 2; + blocks["q_norm"] = std::make_shared(axis_dim, config.rms_norm_eps); + blocks["k_norm"] = std::make_shared(axis_dim, config.rms_norm_eps); + blocks["q_norm_hw"] = std::make_shared(axis_dim, config.rms_norm_eps); + blocks["k_norm_hw"] = std::make_shared(axis_dim, config.rms_norm_eps); + blocks["q_norm_mot_gen"] = std::make_shared(axis_dim, config.rms_norm_eps); + blocks["k_norm_mot_gen"] = std::make_shared(axis_dim, config.rms_norm_eps); + blocks["q_norm_hw_mot_gen"] = std::make_shared(axis_dim, config.rms_norm_eps); + blocks["k_norm_hw_mot_gen"] = std::make_shared(axis_dim, config.rms_norm_eps); + } + + ggml_tensor* apply_axis_rope(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* positions, + int dimensions, + float theta, + int max_position) { + return ggml_rope_ext(ctx->ggml_ctx, + x, + positions, + nullptr, + dimensions, + GGML_ROPE_TYPE_NEOX, + max_position, + theta, + 1.f, + 0.f, + 1.f, + 32.f, + 1.f); + } + + ggml_tensor* normalize_and_rotate(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* position_t, + ggml_tensor* position_h, + ggml_tensor* position_w, + const std::string& norm_name, + const std::string& norm_hw_name) { + const int64_t temporal_dim = config.head_dim / 2; + const int64_t spatial_dim = config.head_dim - temporal_dim; + const int64_t axis_dim = spatial_dim / 2; + + auto temporal = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, temporal_dim); + auto spatial = ggml_ext_slice(ctx->ggml_ctx, x, 0, temporal_dim, config.head_dim); + temporal = std::dynamic_pointer_cast(blocks[norm_name])->forward(ctx, temporal); + spatial = std::dynamic_pointer_cast(blocks[norm_hw_name])->forward(ctx, spatial); + + auto height = ggml_ext_slice(ctx->ggml_ctx, spatial, 0, 0, axis_dim); + auto width = ggml_ext_slice(ctx->ggml_ctx, spatial, 0, axis_dim, spatial_dim); + temporal = apply_axis_rope(ctx, + temporal, + position_t, + static_cast(temporal_dim), + config.rope_theta, + static_cast(config.max_position_embeddings)); + height = apply_axis_rope(ctx, + height, + position_h, + static_cast(axis_dim), + config.rope_theta_hw, + static_cast(config.max_position_embeddings_hw)); + width = apply_axis_rope(ctx, + width, + position_w, + static_cast(axis_dim), + config.rope_theta_hw, + static_cast(config.max_position_embeddings_hw)); + return ggml_concat(ctx->ggml_ctx, + ggml_concat(ctx->ggml_ctx, temporal, height, 0), + width, + 0); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* position_t, + ggml_tensor* position_h, + ggml_tensor* position_w, + ggml_tensor* attention_mask, + Branch branch, + const std::string& cache_prefix) { + const bool generation = branch == Branch::GENERATION; + const std::string suffix = generation ? "_mot_gen" : ""; + auto q_proj = std::dynamic_pointer_cast(blocks["q_proj" + suffix]); + auto k_proj = std::dynamic_pointer_cast(blocks["k_proj" + suffix]); + auto v_proj = std::dynamic_pointer_cast(blocks["v_proj" + suffix]); + auto o_proj = std::dynamic_pointer_cast(blocks["o_proj" + suffix]); + + const int64_t n_tokens = x->ne[1]; + const int64_t batch = x->ne[2]; + auto q = ggml_reshape_4d(ctx->ggml_ctx, + q_proj->forward(ctx, x), + config.head_dim, + config.num_heads, + n_tokens, + batch); + auto k = ggml_reshape_4d(ctx->ggml_ctx, + k_proj->forward(ctx, x), + config.head_dim, + config.num_kv_heads, + n_tokens, + batch); + auto v = ggml_reshape_4d(ctx->ggml_ctx, + v_proj->forward(ctx, x), + config.head_dim, + config.num_kv_heads, + n_tokens, + batch); + + q = normalize_and_rotate(ctx, + q, + position_t, + position_h, + position_w, + "q_norm" + suffix, + "q_norm_hw" + suffix); + k = normalize_and_rotate(ctx, + k, + position_t, + position_h, + position_w, + "k_norm" + suffix, + "k_norm_hw" + suffix); + + const std::string layer_cache = cache_prefix + "." + std::to_string(layer_index); + if (generation) { + auto prefix_k = ctx->load_cache_tensor(layer_cache + ".k"); + auto prefix_v = ctx->load_cache_tensor(layer_cache + ".v"); + GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr); + k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 2); + v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2); + } else { + // Keep dedicated graph outputs alive until the runner copies them + // into its persistent cache buffer after graph execution. + auto cache_k = ggml_dup_tensor(ctx->ggml_ctx, k); + cache_k = ggml_cpy(ctx->ggml_ctx, k, cache_k); + ggml_set_output(cache_k); + auto cache_v = ggml_dup_tensor(ctx->ggml_ctx, v); + cache_v = ggml_cpy(ctx->ggml_ctx, v, cache_v); + ggml_set_output(cache_v); + ctx->persist_cache_tensor(layer_cache + ".k", cache_k); + ctx->persist_cache_tensor(layer_cache + ".v", cache_v); + } + + q = ggml_cont(ctx->ggml_ctx, + ggml_ext_torch_permute(ctx->ggml_ctx, q, 0, 2, 1, 3)); + q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]); + k = ggml_cont(ctx->ggml_ctx, + ggml_ext_torch_permute(ctx->ggml_ctx, k, 0, 2, 1, 3)); + k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]); + + auto out = ggml_ext_attention_ext(ctx->ggml_ctx, + ctx->backend, + q, + k, + v, + config.num_heads, + attention_mask, + true, + ctx->flash_attn_enabled); + return o_proj->forward(ctx, out); + } + }; + + struct TransformerBlock : public GGMLBlock { + TransformerBlock(const SenseNovaU1Config& config, int layer_index) { + blocks["self_attn"] = std::make_shared(config, layer_index); + blocks["mlp"] = std::make_shared(config.hidden_size, config.intermediate_size, false); + blocks["mlp_mot_gen"] = std::make_shared(config.hidden_size, config.intermediate_size, false); + blocks["input_layernorm"] = std::make_shared(config.hidden_size, config.rms_norm_eps); + blocks["input_layernorm_mot_gen"] = std::make_shared(config.hidden_size, config.rms_norm_eps); + blocks["post_attention_layernorm"] = std::make_shared(config.hidden_size, config.rms_norm_eps); + blocks["post_attention_layernorm_mot_gen"] = std::make_shared(config.hidden_size, config.rms_norm_eps); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* position_t, + ggml_tensor* position_h, + ggml_tensor* position_w, + ggml_tensor* attention_mask, + Branch branch, + const std::string& cache_prefix) { + const bool generation = branch == Branch::GENERATION; + auto input_norm = std::dynamic_pointer_cast( + blocks[generation ? "input_layernorm_mot_gen" : "input_layernorm"]); + auto post_norm = std::dynamic_pointer_cast( + blocks[generation ? "post_attention_layernorm_mot_gen" : "post_attention_layernorm"]); + auto attention = std::dynamic_pointer_cast(blocks["self_attn"]); + auto mlp = std::dynamic_pointer_cast(blocks[generation ? "mlp_mot_gen" : "mlp"]); + + auto residual = x; + x = input_norm->forward(ctx, x); + x = attention->forward(ctx, + x, + position_t, + position_h, + position_w, + attention_mask, + branch, + cache_prefix); + x = ggml_add_inplace(ctx->ggml_ctx, x, residual); + + residual = x; + x = post_norm->forward(ctx, x); + x = mlp->forward(ctx, x); + return ggml_add_inplace(ctx->ggml_ctx, x, residual); + } + }; + + struct TextModel : public GGMLBlock { + SenseNovaU1Config config; + + explicit TextModel(const SenseNovaU1Config& config) + : config(config) { + blocks["embed_tokens"] = std::make_shared(config.vocab_size, config.hidden_size); + for (int i = 0; i < config.num_layers; ++i) { + blocks["layers." + std::to_string(i)] = std::make_shared(config, i); + } + blocks["norm"] = std::make_shared(config.hidden_size, config.rms_norm_eps); + blocks["norm_mot_gen"] = std::make_shared(config.hidden_size, config.rms_norm_eps); + } + + ggml_tensor* embed(GGMLRunnerContext* ctx, ggml_tensor* input_ids) { + return std::dynamic_pointer_cast(blocks["embed_tokens"])->forward(ctx, input_ids); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* position_t, + ggml_tensor* position_h, + ggml_tensor* position_w, + ggml_tensor* attention_mask, + Branch branch, + const std::string& cache_prefix) { + for (int i = 0; i < config.num_layers; ++i) { + auto layer = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); + x = layer->forward(ctx, + x, + position_t, + position_h, + position_w, + attention_mask, + branch, + cache_prefix); + } + auto norm = std::dynamic_pointer_cast( + blocks[branch == Branch::GENERATION ? "norm_mot_gen" : "norm"]); + return norm->forward(ctx, x); + } + }; + + struct SenseNovaU1Model : public GGMLBlock { + SenseNovaU1Config config; + + explicit SenseNovaU1Model(const SenseNovaU1Config& config) + : config(config) { + blocks["language_model.model"] = std::make_shared(config); + blocks["fm_modules.vision_model_mot_gen.embeddings"] = std::make_shared(config); + blocks["fm_modules.timestep_embedder"] = std::make_shared(config.hidden_size, + config.timestep_embedding_size); + if (config.add_noise_scale_embedding) { + blocks["fm_modules.noise_scale_embedder"] = std::make_shared(config.hidden_size, + config.timestep_embedding_size); + } + blocks["fm_modules.fm_head"] = std::make_shared(config); + } + + std::shared_ptr text_model() { + return std::dynamic_pointer_cast(blocks["language_model.model"]); + } + + std::shared_ptr vision_embeddings() { + return std::dynamic_pointer_cast(blocks["fm_modules.vision_model_mot_gen.embeddings"]); + } + + std::shared_ptr timestep_embedder() { + return std::dynamic_pointer_cast(blocks["fm_modules.timestep_embedder"]); + } + + std::shared_ptr noise_scale_embedder() { + if (!config.add_noise_scale_embedding) { + return nullptr; + } + return std::dynamic_pointer_cast(blocks["fm_modules.noise_scale_embedder"]); + } + + std::shared_ptr pixel_decoder() { + return std::dynamic_pointer_cast(blocks["fm_modules.fm_head"]); + } + }; + + struct SenseNovaU1Runner : public DiffusionModelRunner { + SenseNovaU1Config config; + SenseNovaU1Model model; + std::unordered_set cached_prefix_hashes; + std::vector position_t_vec; + std::vector position_h_vec; + std::vector position_w_vec; + std::vector attention_mask_vec; + std::vector noise_scale_vec; + + SenseNovaU1Runner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string& prefix = "", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(SenseNovaU1Config::detect_from_weights(tensor_storage_map, prefix)), + model(config) { + model.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "SenseNova U1.5"; + } + + void get_param_tensors(std::map& tensors, + const std::string& prefix) override { + model.get_param_tensors(tensors, prefix); + } + + static uint64_t hash_input_ids(const sd::Tensor& input_ids) { + uint64_t hash = 1469598103934665603ULL; + for (int32_t token : input_ids.values()) { + uint32_t value = static_cast(token); + for (int byte = 0; byte < 4; ++byte) { + hash ^= static_cast(value & 0xffU); + hash *= 1099511628211ULL; + value >>= 8; + } + } + hash ^= static_cast(input_ids.numel()); + hash *= 1099511628211ULL; + return hash; + } + + static std::string cache_prefix(uint64_t hash) { + return "snu15." + std::to_string(hash); + } + + ggml_tensor* make_position_tensor(const std::vector& values, + const std::string& name) { + auto tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, values.size()); + ggml_set_name(tensor, name.c_str()); + set_backend_tensor_data(tensor, values.data()); + return tensor; + } + + ggml_cgraph* build_prefix_graph(const sd::Tensor& input_ids_tensor, + const std::string& prefix_cache) { + ggml_cgraph* graph = new_graph_custom(SENSENOVA_U1_GRAPH_SIZE); + ggml_tensor* ids = make_input(input_ids_tensor); + const int64_t length = input_ids_tensor.numel(); + + position_t_vec.resize(length); + position_h_vec.assign(length, 0); + position_w_vec.assign(length, 0); + for (int64_t i = 0; i < length; ++i) { + position_t_vec[i] = static_cast(i); + } + auto position_t = make_position_tensor(position_t_vec, "snu15.prefix.position_t"); + auto position_h = make_position_tensor(position_h_vec, "snu15.prefix.position_h"); + auto position_w = make_position_tensor(position_w_vec, "snu15.prefix.position_w"); + + attention_mask_vec.assign(static_cast(length * length), 0.f); + for (int64_t query = 0; query < length; ++query) { + for (int64_t key = query + 1; key < length; ++key) { + attention_mask_vec[static_cast(query * length + key)] = -INFINITY; + } + } + auto attention_mask = ggml_new_tensor_2d(compute_ctx, + GGML_TYPE_F32, + length, + length); + ggml_set_name(attention_mask, "snu15.prefix.attention_mask"); + set_backend_tensor_data(attention_mask, attention_mask_vec.data()); + + auto runner_ctx = get_context(); + auto text_model = model.text_model(); + auto hidden = text_model->embed(&runner_ctx, ids); + hidden = text_model->forward(&runner_ctx, + hidden, + position_t, + position_h, + position_w, + attention_mask, + Branch::UNDERSTANDING, + prefix_cache); + ggml_build_forward_expand(graph, hidden); + return graph; + } + + bool ensure_prefix_cache(int n_threads, + const sd::Tensor& input_ids, + std::string* prefix_cache) { + const uint64_t hash = hash_input_ids(input_ids); + *prefix_cache = cache_prefix(hash); + if (cached_prefix_hashes.find(hash) != cached_prefix_hashes.end() && + get_cache_tensor_by_name(*prefix_cache + ".0.k") != nullptr) { + return true; + } + + if (cached_prefix_hashes.size() >= 2) { + free_cache_ctx_and_buffer(); + cached_prefix_hashes.clear(); + } + auto get_graph = [&]() { + return build_prefix_graph(input_ids, *prefix_cache); + }; + auto result = GGMLRunner::compute(get_graph, + n_threads, + false, + true, + true, + true); + if (!result.has_value()) { + LOG_ERROR("SenseNova U1.5 prefix cache computation failed"); + return false; + } + cached_prefix_hashes.insert(hash); + return true; + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timestep_tensor, + const std::string& prefix_cache, + int64_t prefix_length) { + ggml_cgraph* graph = new_graph_custom(SENSENOVA_U1_GRAPH_SIZE); + ggml_tensor* x = make_input(x_tensor); + ggml_tensor* t = make_input(timestep_tensor); + GGML_ASSERT(x->ne[3] == 1); + GGML_ASSERT(x->ne[0] % config.image_token_stride() == 0); + GGML_ASSERT(x->ne[1] % config.image_token_stride() == 0); + + const int64_t grid_w = x->ne[0] / config.patch_size; + const int64_t grid_h = x->ne[1] / config.patch_size; + const int64_t token_w = grid_w / config.vision_downsample_factor; + const int64_t token_h = grid_h / config.vision_downsample_factor; + const int64_t tokens = token_w * token_h; + + position_h_vec.resize(grid_w * grid_h); + position_w_vec.resize(grid_w * grid_h); + for (int64_t index = 0; index < grid_w * grid_h; ++index) { + position_h_vec[index] = static_cast(index / grid_w); + position_w_vec[index] = static_cast(index % grid_w); + } + auto vision_position_x = make_position_tensor(position_w_vec, "snu15.vision.position_x"); + auto vision_position_y = make_position_tensor(position_h_vec, "snu15.vision.position_y"); + + auto runner_ctx = get_context(); + auto hidden = model.vision_embeddings()->forward(&runner_ctx, + x, + vision_position_x, + vision_position_y); + auto time_embedding = model.timestep_embedder()->forward(&runner_ctx, t); + time_embedding = ggml_reshape_3d(compute_ctx, time_embedding, config.hidden_size, 1, 1); + hidden = ggml_add(compute_ctx, hidden, time_embedding); + + if (config.add_noise_scale_embedding) { + const float image_tokens = static_cast(tokens); + const float noise_scale = std::min(config.noise_scale_max_value, + std::sqrt(image_tokens / config.noise_scale_base_image_seq_len)); + noise_scale_vec = {noise_scale / config.noise_scale_max_value}; + auto noise_scale_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, 1); + ggml_set_name(noise_scale_tensor, "snu15.noise_scale"); + set_backend_tensor_data(noise_scale_tensor, noise_scale_vec.data()); + auto noise_embedding = model.noise_scale_embedder()->forward(&runner_ctx, noise_scale_tensor); + noise_embedding = ggml_reshape_3d(compute_ctx, noise_embedding, config.hidden_size, 1, 1); + hidden = ggml_add(compute_ctx, hidden, noise_embedding); + } + + position_t_vec.assign(tokens, static_cast(prefix_length)); + position_h_vec.resize(tokens); + position_w_vec.resize(tokens); + for (int64_t index = 0; index < tokens; ++index) { + position_h_vec[index] = static_cast(index / token_w); + position_w_vec[index] = static_cast(index % token_w); + } + auto position_t = make_position_tensor(position_t_vec, "snu15.image.position_t"); + auto position_h = make_position_tensor(position_h_vec, "snu15.image.position_h"); + auto position_w = make_position_tensor(position_w_vec, "snu15.image.position_w"); + + hidden = model.text_model()->forward(&runner_ctx, + hidden, + position_t, + position_h, + position_w, + nullptr, + Branch::GENERATION, + prefix_cache); + hidden = ggml_reshape_4d(compute_ctx, + hidden, + config.hidden_size, + token_w, + token_h, + x->ne[3]); + hidden = ggml_cont(compute_ctx, ggml_permute(compute_ctx, hidden, 2, 0, 1, 3)); + auto x_prediction = model.pixel_decoder()->forward(&runner_ctx, hidden); + + const float timestep = timestep_tensor.values()[0]; + const float denom = std::max(1.f - timestep, config.t_eps); + auto velocity = ggml_scale(compute_ctx, + ggml_sub(compute_ctx, x_prediction, x), + 1.f / denom); + ggml_build_forward_expand(graph, velocity); + return graph; + } + + sd::Tensor compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timestep, + const sd::Tensor& input_ids) { + std::string prefix_cache; + if (!ensure_prefix_cache(n_threads, input_ids, &prefix_cache)) { + return {}; + } + auto get_graph = [&]() { + return build_graph(x, timestep, prefix_cache, input_ids.numel()); + }; + return restore_trailing_singleton_dims( + GGMLRunner::compute(get_graph, n_threads, false, false, false), + x.dim()); + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + const auto* extra = diffusion_extra_as(diffusion_params); + GGML_ASSERT(extra->input_ids != nullptr); + return compute(n_threads, + *diffusion_params.x, + *diffusion_params.timesteps, + *extra->input_ids); + } + }; +} // namespace SenseNovaU1 + +#endif // __SD_MODEL_DIFFUSION_SENSENOVA_U1_H__ diff --git a/src/model/vae/vae.hpp b/src/model/vae/vae.hpp index f3adb0ccc..a7b21087a 100644 --- a/src/model/vae/vae.hpp +++ b/src/model/vae/vae.hpp @@ -165,7 +165,7 @@ struct VAE : public GGMLRunner { scale_factor = 16; } else if (sd_version_uses_flux2_vae(version)) { scale_factor = 16; - } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) { + } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version) || sd_version_is_sensenova_u1(version)) { scale_factor = 1; } return scale_factor; diff --git a/src/model_loader.cpp b/src/model_loader.cpp index dc2bb4bf0..a24d5b17c 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -67,6 +67,8 @@ const char* unused_tensors[] = { // "v_pred", // Used to detect SDXL vpred models "text_encoders.llm.output.weight", "text_encoders.llm.lm_head.", + "language_model.lm_head.", + "vision_model.", }; bool is_unused_tensor(const std::string& name) { @@ -170,6 +172,15 @@ void ModelLoader::set_n_threads(int n_threads) { bool ModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) { if (is_directory(file_path)) { + const std::string diffusers_index_path = path_join(file_path, "model_index.json"); + const std::string diffusers_unet_path = path_join(file_path, "unet/diffusion_pytorch_model.safetensors"); + const bool has_diffusers_layout = file_exists(diffusers_index_path) || file_exists(diffusers_unet_path); + + const std::string safetensors_index_path = path_join(file_path, "model.safetensors.index.json"); + if (!has_diffusers_layout && file_exists(safetensors_index_path)) { + LOG_INFO("load %s using root safetensors index", file_path.c_str()); + return init_from_safetensors_index_file(safetensors_index_path, prefix); + } LOG_INFO("load %s using diffusers format", file_path.c_str()); return init_from_diffusers_file(file_path, prefix); } else if (is_gguf_file(file_path)) { @@ -387,7 +398,20 @@ bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const s return true; } +static bool has_sensenova_u1_signature(const String2TensorStorage& tensors) { + // Require independent markers for the generation MoT branch, the vision + // input path, and the pixel-flow output head to avoid matching a plain or + // partially exported Qwen checkpoint. + return tensors.find("language_model.model.layers.0.self_attn.q_proj_mot_gen.weight") != tensors.end() && + tensors.find("fm_modules.vision_model_mot_gen.embeddings.patch_embedding.weight") != tensors.end() && + tensors.find("fm_modules.fm_head.conv1.weight") != tensors.end(); +} + SDVersion ModelLoader::get_sd_version() { + if (has_sensenova_u1_signature(tensor_storage_map)) { + return VERSION_SENSENOVA_U1_5; + } + TensorStorage token_embedding_weight, input_block_weight, context_ebedding_weight; bool has_multiple_encoders = false; diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index b6f1843ec..88347670e 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -1486,6 +1486,82 @@ struct MiniT2IFlowDenoiser : public Denoiser { } }; +// SenseNova U1.5 integrates velocity over t=0..1 while the generic sampler +// integrates over descending sigma. With sigma=1-t, returning +// denoised=x+sigma*v makes the generic Euler derivative exactly -v, so the +// descending-sigma update is identical to the official ascending-time update. +struct SenseNovaU1FlowDenoiser : public DiscreteFlowDenoiser { + explicit SenseNovaU1FlowDenoiser(float shift = 3.f) + : DiscreteFlowDenoiser(shift) {} + + float sigma_min() override { + return 0.f; + } + + float sigma_max() override { + return 1.f; + } + + float sigma_to_t(float sigma) override { + return 1.f - sigma; + } + + float t_to_sigma(float t) override { + float sigma = 1.f - t; + return shift * sigma / (1.f + (shift - 1.f) * sigma); + } + + std::vector get_scalings(float sigma) override { + return {1.f, sigma, 1.f}; + } + + sd::Tensor noise_scaling(float sigma, + const sd::Tensor& noise, + const sd::Tensor& latent) override { + SD_UNUSED(sigma); + SD_UNUSED(latent); + GGML_ASSERT(noise.dim() >= 2); + const float token_w = static_cast(noise.shape()[0]) / 32.f; + const float token_h = static_cast(noise.shape()[1]) / 32.f; + const float noise_scale = std::min(16.f, std::sqrt((token_w * token_h) / 64.f)); + return noise * noise_scale; + } + + sd::Tensor inverse_noise_scaling(float sigma, + const sd::Tensor& latent) override { + SD_UNUSED(sigma); + return latent; + } + + float noise_level_to_sigma(float noise_level) override { + SD_UNUSED(noise_level); + return 1.f; + } + + std::vector get_sigmas(uint32_t n, + int image_seq_len, + scheduler_t scheduler_type, + SDVersion version, + const char* extra_sample_args = nullptr) override { + SD_UNUSED(image_seq_len); + SD_UNUSED(scheduler_type); + SD_UNUSED(version); + SD_UNUSED(extra_sample_args); + std::vector sigmas; + sigmas.reserve(n + 1); + if (n == 0) { + sigmas.push_back(0.f); + return sigmas; + } + for (uint32_t i = 0; i <= n; ++i) { + const float t = static_cast(i) / static_cast(n); + sigmas.push_back(t_to_sigma(t)); + } + sigmas.back() = 0.f; + return sigmas; + } +}; + typedef std::function&, float, int)> denoise_cb_t; static std::pair get_ancestral_step(float sigma_from, diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index d2193e14c..f139e3375 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -44,6 +44,7 @@ #include "model/diffusion/model.hpp" #include "model/diffusion/pid.hpp" #include "model/diffusion/qwen_image.hpp" +#include "model/diffusion/sensenova_u1.h" #include "model/diffusion/unet.hpp" #include "model/diffusion/wan.hpp" #include "model/diffusion/z_image.hpp" @@ -124,6 +125,7 @@ const char* model_version_to_str[] = { "SeFi-Image", "Krea2", "Mage Flow", + "SenseNova U1.5", "ESRGAN", }; @@ -1271,6 +1273,12 @@ class StableDiffusionGGML { tensor_storage_map, "model.diffusion_model.model.net", model_manager); + } else if (sd_version_is_sensenova_u1(version)) { + cond_stage_model = std::make_shared(); + diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), + tensor_storage_map, + "", + model_manager); } else if (sd_version_is_anima(version)) { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), tensor_storage_map, @@ -1488,7 +1496,7 @@ class StableDiffusionGGML { } }; - if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) { + if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version) || sd_version_is_sensenova_u1(version)) { LOG_INFO("using FakeVAE"); first_stage_model = std::make_shared(version, backend_for(SDBackendModule::VAE), @@ -1840,6 +1848,9 @@ class StableDiffusionGGML { pred_type = SEFI_FLOW_PRED; } else if (sd_version_is_minit2i(version)) { pred_type = MINIT2I_FLOW_PRED; + } else if (sd_version_is_sensenova_u1(version)) { + pred_type = SENSENOVA_U1_FLOW_PRED; + default_flow_shift = 3.f; } else { pred_type = EPS_PRED; } @@ -1885,6 +1896,11 @@ class StableDiffusionGGML { denoiser = std::make_shared(); break; } + case SENSENOVA_U1_FLOW_PRED: { + LOG_INFO("running in SenseNova U1.5 FLOW mode"); + denoiser = std::make_shared(default_flow_shift); + break; + } default: { LOG_ERROR("Unknown predition type %i", pred_type); return false; @@ -2790,6 +2806,9 @@ class StableDiffusionGGML { } else if (sd_version_is_minit2i(version)) { diffusion_params.extra = MiniT2IDiffusionExtra{ condition.c_vector.empty() ? nullptr : &condition.c_vector}; + } else if (sd_version_is_sensenova_u1(version)) { + diffusion_params.extra = SenseNovaU1DiffusionExtra{ + condition.c_input_ids.empty() ? nullptr : &condition.c_input_ids}; } else { diffusion_params.extra = std::monostate{}; } @@ -2954,7 +2973,9 @@ class StableDiffusionGGML { int get_diffusion_model_down_factor() { int down_factor = 8; // unet if (sd_version_is_dit(version)) { - if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version)) { + if (sd_version_is_sensenova_u1(version)) { + down_factor = 32; + } else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version)) { down_factor = 2; } else { down_factor = 1; @@ -2980,6 +3001,8 @@ class StableDiffusionGGML { latent_channel = 3; } else if (sd_version_is_minit2i(version)) { latent_channel = 3; + } else if (sd_version_is_sensenova_u1(version)) { + latent_channel = 3; } else if (sd_version_is_pid(version)) { latent_channel = 3; } else if (sd_version_is_sefi_image(version)) { @@ -3394,6 +3417,7 @@ const char* prediction_to_str[] = { "flux_flow", "sefi_flow", "minit2i_flow", + "sensenova_u1_flow", }; const char* sd_prediction_name(enum prediction_t prediction) { @@ -5271,6 +5295,10 @@ static std::optional prepare_image_generation_embeds(sd_c // states with a zeroed prompt mask, so no extra text encode is needed. uncond.c_crossattn = cond.c_crossattn; uncond.c_vector = sd::Tensor::zeros_like(cond.c_vector); + } else if (sd_version_is_sensenova_u1(sd_ctx->sd->version)) { + auto* sensenova_conditioner = static_cast( + sd_ctx->sd->cond_stage_model.get()); + uncond = sensenova_conditioner->get_unconditional_condition(request->negative_prompt); } else { bool zero_out_masked = false; if (sd_version_is_sdxl(sd_ctx->sd->version) && diff --git a/src/tokenizers/qwen2_tokenizer.cpp b/src/tokenizers/qwen2_tokenizer.cpp index 79e683e7b..6bc21ed2b 100644 --- a/src/tokenizers/qwen2_tokenizer.cpp +++ b/src/tokenizers/qwen2_tokenizer.cpp @@ -45,16 +45,8 @@ void Qwen2Tokenizer::load_from_merges(const std::string& merges_utf8_str) { bpe_len = rank; } -Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str) { - UNK_TOKEN = "<|endoftext|>"; - EOS_TOKEN = "<|endoftext|>"; - PAD_TOKEN = "<|endoftext|>"; - - UNK_TOKEN_ID = 151643; - EOS_TOKEN_ID = 151643; - PAD_TOKEN_ID = 151643; - - special_tokens = { +static const std::vector& qwen2_special_tokens() { + static const std::vector tokens = { "<|endoftext|>", "<|im_start|>", "<|im_end|>", @@ -87,6 +79,24 @@ Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str) { "<|bot_token|>", "<|tms_token|>", }; + return tokens; +} + +Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str) + : Qwen2Tokenizer(merges_utf8_str, qwen2_special_tokens()) { +} + +Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str, + const std::vector& special_tokens_override) { + UNK_TOKEN = "<|endoftext|>"; + EOS_TOKEN = "<|endoftext|>"; + PAD_TOKEN = "<|endoftext|>"; + + UNK_TOKEN_ID = 151643; + EOS_TOKEN_ID = 151643; + PAD_TOKEN_ID = 151643; + + special_tokens = special_tokens_override; if (merges_utf8_str.size() > 0) { load_from_merges(merges_utf8_str); diff --git a/src/tokenizers/qwen2_tokenizer.h b/src/tokenizers/qwen2_tokenizer.h index 04e92c2c3..127394665 100644 --- a/src/tokenizers/qwen2_tokenizer.h +++ b/src/tokenizers/qwen2_tokenizer.h @@ -2,12 +2,15 @@ #define __SD_TOKENIZERS_QWEN2_TOKENIZER_H__ #include +#include #include "bpe_tokenizer.h" class Qwen2Tokenizer : public BPETokenizer { protected: void load_from_merges(const std::string& merges_utf8_str); + Qwen2Tokenizer(const std::string& merges_utf8_str, + const std::vector& special_tokens_override); public: explicit Qwen2Tokenizer(const std::string& merges_utf8_str = ""); diff --git a/src/tokenizers/sensenova_u1_tokenizer.cpp b/src/tokenizers/sensenova_u1_tokenizer.cpp new file mode 100644 index 000000000..32e259b9c --- /dev/null +++ b/src/tokenizers/sensenova_u1_tokenizer.cpp @@ -0,0 +1,44 @@ +#include "sensenova_u1_tokenizer.h" + +#include + +static const std::vector& sensenova_u1_special_tokens() { + static const std::vector tokens = { + "<|endoftext|>", + "<|im_start|>", + "<|im_end|>", + "<|object_ref_start|>", + "<|object_ref_end|>", + "<|box_start|>", + "<|box_end|>", + "<|quad_start|>", + "<|quad_end|>", + "<|vision_start|>", + "<|vision_end|>", + "<|vision_pad|>", + "<|image_pad|>", + "<|video_pad|>", + "", + "", + "<|fim_prefix|>", + "<|fim_middle|>", + "<|fim_suffix|>", + "<|fim_pad|>", + "<|repo_name|>", + "<|file_sep|>", + "", + "", + "", + "", + "", + "", + "", + }; + return tokens; +} + +SenseNovaU1Tokenizer::SenseNovaU1Tokenizer(const std::string& merges_utf8_str) + : Qwen2Tokenizer(merges_utf8_str, sensenova_u1_special_tokens()) { + EOS_TOKEN = "<|im_end|>"; + EOS_TOKEN_ID = 151645; +} diff --git a/src/tokenizers/sensenova_u1_tokenizer.h b/src/tokenizers/sensenova_u1_tokenizer.h new file mode 100644 index 000000000..31778ec61 --- /dev/null +++ b/src/tokenizers/sensenova_u1_tokenizer.h @@ -0,0 +1,13 @@ +#ifndef __SD_TOKENIZERS_SENSENOVA_U1_TOKENIZER_H__ +#define __SD_TOKENIZERS_SENSENOVA_U1_TOKENIZER_H__ + +#include + +#include "qwen2_tokenizer.h" + +class SenseNovaU1Tokenizer : public Qwen2Tokenizer { +public: + explicit SenseNovaU1Tokenizer(const std::string& merges_utf8_str = ""); +}; + +#endif // __SD_TOKENIZERS_SENSENOVA_U1_TOKENIZER_H__