diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index 5affda3b..e84a1b96 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -6,7 +6,9 @@ body: - type: markdown attributes: value: | - Please use this template and include as many details as possible to help us reproduce and fix the issue. + Before submitting a bug report, please read the [Troubleshooting guide](https://github.com/leejet/stable-diffusion.cpp/blob/master/docs/troubleshooting.md) and try the steps relevant to your problem. + + If the problem persists, complete this form and include what you tried and the results, along with enough details to help us reproduce and fix the issue. - type: textarea id: commit attributes: diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 00000000..36c4fd9e --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,4 @@ +contact_links: + - name: Troubleshooting + url: https://github.com/leejet/stable-diffusion.cpp/blob/master/docs/troubleshooting.md + about: Read the troubleshooting guide first. If the problem persists, submit a bug report. diff --git a/README.md b/README.md index 6ba81d89..d99348cb 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) @@ -147,6 +148,7 @@ For runtime and parameter backend placement, see the [backend selection guide](. ## More Guides +- [Troubleshooting](./docs/troubleshooting.md) - [Backend selection](./docs/backend.md) - [RPC](./docs/rpc.md) - [LoRA](./docs/lora.md) diff --git a/docs/backend.md b/docs/backend.md index 429e659e..c14e8094 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -193,8 +193,9 @@ weights, compute buffers and caches must still fit the runner's capacity checks. Offloading weights does not guarantee that every resolution or frame count will fit, and auto-fit does not change a component to CPU computation solely because its full weights exceed VRAM. -If a VAE decode fails, auto-fit retries with spatial tiling; supported video -decoders try temporal tiling first and can then add spatial tiling. +If a VAE decode fails, decoding retries with spatial tiling even when `--auto-fit` +is off; supported video decoders try temporal tiling first and can then add +spatial tiling. Spatial retries use half-size tiles along each latent dimension. ## Modules diff --git a/docs/sensenova_u1.md b/docs/sensenova_u1.md new file mode 100644 index 00000000..fd19c4d3 --- /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/docs/troubleshooting.md b/docs/troubleshooting.md new file mode 100644 index 00000000..22007db2 --- /dev/null +++ b/docs/troubleshooting.md @@ -0,0 +1,45 @@ +# Troubleshooting + +## Completely black or white images or videos / NaNs + +Some ggml backends can encounter numerical overflow during inference, producing +NaN (not-a-number) values. This can result in completely black or white images or videos. +Whether it happens can depend on the backend, device, model, and weight format. + +Known overflow issues have been addressed as far as possible, but the maintainer +has limited hardware and cannot test every combination. Some cases may therefore +still need a manual workaround. + +These options are supported by both `sd-cli` and `sd-server`. If you encounter +this problem, add them to your CLI generation command or server startup command: + +```sh +--linear-scale 0.0078125 --attn-scale 0.0078125 +``` + +For `sd-server`, restart the server after changing these startup options. Run the +same prompt and seed again to see whether the output recovers. If the problem +persists, try smaller positive values, for example: + +```sh +--linear-scale 0.00390625 --attn-scale 0.00390625 +``` + +These options reduce intermediate values and compensate afterwards to preserve +the intended output scale: + +- `--linear-scale` scales Linear inputs before matrix multiplication and rescales + the result. +- `--attn-scale` scales attention keys and values (K/V). It takes effect only in + the Flash Attention path, where `--fa` or `--diffusion-fa` is enabled and the + backend supports it. + +The two values can be set independently and apply across model components. The +default `0` preserves each model's built-in settings; `1` explicitly disables the +corresponding scaling. Overrides must be finite positive values. C API users can +set `linear_scale` and `attn_scale` in `sd_ctx_params_t`. + +If the problem persists after trying the relevant steps above, +[submit a bug report](https://github.com/leejet/stable-diffusion.cpp/issues/new?template=bug_report.yml). +Include your full command, backend and hardware, model and weight format, logs, +and the scale values you tried with their results. diff --git a/examples/cli/README.md b/examples/cli/README.md index 9df3bb0b..fe47122b 100644 --- a/examples/cli/README.md +++ b/examples/cli/README.md @@ -22,3 +22,6 @@ Metadata mode inspects PNG/JPEG container metadata without loading any model: ./bin/sd-cli -M metadata --image ./output.png --metadata-raw ./bin/sd-cli -M metadata --image ./output.png --metadata-all ``` + +For completely black or white images or videos, NaNs, and the `--linear-scale` / +`--attn-scale` workaround, see [Troubleshooting](../../docs/troubleshooting.md). diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 335a6226..cdf21bc7 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -302,8 +302,12 @@ bool parse_options(int argc, const char** argv, const std::vector& o invalid_arg = true; return; } - *option.target = std::stoi(argv[i]); - found_arg = true; + try { + *option.target = std::stoi(argv[i]); + } catch (const std::invalid_argument&) { + invalid_arg = true; + } + found_arg = true; })) break; @@ -312,8 +316,12 @@ bool parse_options(int argc, const char** argv, const std::vector& o invalid_arg = true; return; } - *option.target = std::stof(argv[i]); - found_arg = true; + try { + *option.target = std::stof(argv[i]); + } catch (const std::invalid_argument&) { + invalid_arg = true; + } + found_arg = true; })) break; @@ -337,7 +345,8 @@ bool parse_options(int argc, const char** argv, const std::vector& o if (invalid_arg) { if (!valid) { - LOG_ERROR("error: invalid parameter for argument: %s", arg.c_str()); + LOG_ERROR("error: invalid parameter for argument \"%s\": \"%s\"", + arg.c_str(), (i >= argc) ? "" : argv[i]); } return false; } @@ -350,6 +359,25 @@ bool parse_options(int argc, const char** argv, const std::vector& o return true; } +static int parse_scale_override(int argc, const char** argv, int index, float& scale) { + if (++index >= argc) { + return -1; + } + try { + size_t end = 0; + const std::string value = argv[index]; + float parsed = std::stof(value, &end); + if (end != value.size() || !std::isfinite(parsed) || parsed < 0.f || + (parsed > 0.f && !std::isfinite(1.f / parsed))) { + return -1; + } + scale = parsed; + } catch (const std::exception&) { + return -1; + } + return 1; +} + ArgOptions SDContextParams::get_options() { ArgOptions options; options.string_options = { @@ -678,6 +706,18 @@ ArgOptions SDContextParams::get_options() { }; options.manual_options = { + {"", + "--linear-scale", + "linear input scale override (float, default: 0 = model default, 1 = no scaling)", + [this](int argc, const char** argv, int index) { + return parse_scale_override(argc, argv, index, linear_scale); + }}, + {"", + "--attn-scale", + "flash-attention K/V scale override (float, default: 0 = model default, 1 = no scaling); requires --fa or --diffusion-fa", + [this](int argc, const char** argv, int index) { + return parse_scale_override(argc, argv, index, attn_scale); + }}, {"", "--auto-fit", "on|off (default: on). Use one GPU for diffusion/te/vae computation and place weights on that GPU, " @@ -886,6 +926,8 @@ std::string SDContextParams::to_string() const { << " vae_on_cpu: " << (vae_on_cpu ? "true" : "false") << ",\n" << " flash_attn: " << (flash_attn ? "true" : "false") << ",\n" << " diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n" + << " linear_scale: " << linear_scale << ",\n" + << " attn_scale: " << attn_scale << ",\n" << " diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n" << " vae_conv_direct: " << (vae_conv_direct ? "true" : "false") << ",\n" << " prediction: " << sd_prediction_name(prediction) << ",\n" @@ -939,6 +981,8 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.enable_mmap = enable_mmap; sd_ctx_params.flash_attn = flash_attn; sd_ctx_params.diffusion_flash_attn = diffusion_flash_attn; + sd_ctx_params.linear_scale = linear_scale; + sd_ctx_params.attn_scale = attn_scale; sd_ctx_params.tae_preview_only = taesd_preview; sd_ctx_params.diffusion_conv_direct = diffusion_conv_direct; sd_ctx_params.vae_conv_direct = vae_conv_direct; diff --git a/examples/common/common.h b/examples/common/common.h index ef084dfd..660fef4f 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -175,6 +175,8 @@ struct SDContextParams { lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO; bool force_sdxl_vae_conv_scale = false; + float linear_scale = 0.f; + float attn_scale = 0.f; float flow_shift = INFINITY; ArgOptions get_options(); diff --git a/examples/server/README.md b/examples/server/README.md index 7622eb40..c59fc8ae 100644 --- a/examples/server/README.md +++ b/examples/server/README.md @@ -129,3 +129,6 @@ For detailed command-line arguments, run: ```bash ./bin/sd-server -h ``` + +For completely black or white images or videos, NaNs, and the `--linear-scale` / +`--attn-scale` startup options, see [Troubleshooting](../../docs/troubleshooting.md). diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 08bc27ae..cf69b48d 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 }; @@ -240,6 +241,8 @@ typedef struct { const char* rpc_servers; const char* model_args; bool disable_segmented_compute; // Force monolithic graph execution even when automatic graph cutting would fit memory better + float linear_scale; // Override linear input scaling; 0 keeps the model default + float attn_scale; // Override flash-attention K/V scaling; 0 keeps the model default } sd_ctx_params_t; typedef struct { @@ -493,6 +496,9 @@ SD_API void free_sd_audio(sd_audio_t* audio); SD_API void sd_sample_params_init(sd_sample_params_t* sample_params); SD_API char* sd_sample_params_to_str(const sd_sample_params_t* sample_params); +// Requires a loaded context; returns a static string owned by the library, or "Unknown". +SD_API const char* sd_get_model_version_name(const sd_ctx_t* sd_ctx); + SD_API enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx); SD_API enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_method_t sample_method); diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index f0a7286f..4dc46a22 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -16,6 +16,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; @@ -149,6 +150,7 @@ public: virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) {} virtual void get_layer_split_param_tensors(std::map& tensors) {} virtual void set_flash_attention_enabled(bool enabled) = 0; + virtual void set_scale_overrides(float linear_scale, float attn_scale) {} virtual void set_weight_adapter(const std::shared_ptr& adapter) {} virtual void runner_end() {} }; @@ -231,6 +233,13 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { } } + void set_scale_overrides(float linear_scale, float attn_scale) override { + text_model->set_scale_overrides(linear_scale, attn_scale); + if (sd_version_is_sdxl(version)) { + text_model2->set_scale_overrides(linear_scale, attn_scale); + } + } + void set_weight_adapter(const std::shared_ptr& adapter) override { text_model->set_weight_adapter(adapter); if (sd_version_is_sdxl(version)) { @@ -736,6 +745,18 @@ struct SD3CLIPEmbedder : public Conditioner { } } + void set_scale_overrides(float linear_scale, float attn_scale) override { + if (clip_l) { + clip_l->set_scale_overrides(linear_scale, attn_scale); + } + if (clip_g) { + clip_g->set_scale_overrides(linear_scale, attn_scale); + } + if (t5) { + t5->set_scale_overrides(linear_scale, attn_scale); + } + } + void set_weight_adapter(const std::shared_ptr& adapter) override { if (clip_l) { clip_l->set_weight_adapter(adapter); @@ -1106,6 +1127,15 @@ struct FluxCLIPEmbedder : public Conditioner { } } + void set_scale_overrides(float linear_scale, float attn_scale) override { + if (clip_l) { + clip_l->set_scale_overrides(linear_scale, attn_scale); + } + if (t5) { + t5->set_scale_overrides(linear_scale, attn_scale); + } + } + void set_weight_adapter(const std::shared_ptr& adapter) override { if (clip_l) { clip_l->set_weight_adapter(adapter); @@ -1374,6 +1404,12 @@ struct T5CLIPEmbedder : public Conditioner { } } + void set_scale_overrides(float linear_scale, float attn_scale) override { + if (t5) { + t5->set_scale_overrides(linear_scale, attn_scale); + } + } + void set_weight_adapter(const std::shared_ptr& adapter) override { if (t5) { t5->set_weight_adapter(adapter); @@ -1582,6 +1618,12 @@ struct MiniT2IConditioner : public Conditioner { } } + void set_scale_overrides(float linear_scale, float attn_scale) override { + if (t5) { + t5->set_scale_overrides(linear_scale, attn_scale); + } + } + void set_weight_adapter(const std::shared_ptr& adapter) override { if (t5) { t5->set_weight_adapter(adapter); @@ -1629,6 +1671,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; @@ -1678,6 +1785,10 @@ struct AnimaConditioner : public Conditioner { llm->set_flash_attention_enabled(enabled); } + void set_scale_overrides(float linear_scale, float attn_scale) override { + llm->set_scale_overrides(linear_scale, attn_scale); + } + void set_weight_adapter(const std::shared_ptr& adapter) override { llm->set_weight_adapter(adapter); } @@ -1882,6 +1993,13 @@ struct LLMEmbedder : public Conditioner { } } + void set_scale_overrides(float linear_scale, float attn_scale) override { + llm->set_scale_overrides(linear_scale, attn_scale); + if (byt5) { + byt5->set_scale_overrides(linear_scale, attn_scale); + } + } + void set_weight_adapter(const std::shared_ptr& adapter) override { if (llm) { llm->set_weight_adapter(adapter); @@ -2971,6 +3089,11 @@ struct LTXAVEmbedder : public Conditioner { projector->set_flash_attention_enabled(enabled); } + void set_scale_overrides(float linear_scale, float attn_scale) override { + llm->set_scale_overrides(linear_scale, attn_scale); + projector->set_scale_overrides(linear_scale, attn_scale); + } + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { llm->set_max_graph_vram_bytes(max_vram_bytes); projector->set_max_graph_vram_bytes(max_vram_bytes); diff --git a/src/core/backend_fit.cpp b/src/core/backend_fit.cpp index 0f2dd67c..82150ab9 100644 --- a/src/core/backend_fit.cpp +++ b/src/core/backend_fit.cpp @@ -389,7 +389,9 @@ namespace sd::backend_fit { tiling_params.temporal_tiling = true; retry_mode = tiling_params.enabled ? "spatial+temporal" : "temporal"; } else if (!tiling_params.enabled) { - tiling_params.enabled = true; + tiling_params.enabled = true; + tiling_params.rel_size_x = 0.5f; + tiling_params.rel_size_y = 0.5f; if (tiling_params.tile_size_x <= 0) { tiling_params.tile_size_x = 256; } @@ -401,7 +403,7 @@ namespace sd::backend_fit { return false; } - LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling", + LOG_WARN("VAE decode failed (likely out of memory); retrying with %s tiling", retry_mode); return true; } diff --git a/src/core/ggml_graph_cut.cpp b/src/core/ggml_graph_cut.cpp index 79710058..a6e2d99e 100644 --- a/src/core/ggml_graph_cut.cpp +++ b/src/core/ggml_graph_cut.cpp @@ -482,6 +482,16 @@ namespace sd::ggml_graph_cut { return ggml_nbytes(cache_src); } + static bool can_ignore_op_params(ggml_op op) { + // Exempt only parameters that cannot affect graph layout or backend allocation size. + switch (op) { + case GGML_OP_SCALE: + return true; + default: + return false; + } + } + std::vector graph_layout(ggml_cgraph* graph, bool include_bindings) { std::vector tensors; std::unordered_map indices; @@ -530,8 +540,10 @@ namespace sd::ggml_graph_cut { for (auto source : tensor->src) { signature.push_back(source == nullptr ? 0 : indices.at(source)); } - for (int value : tensor->op_params) { - signature.push_back(static_cast(value)); + if (!can_ignore_op_params(tensor->op)) { + for (int value : tensor->op_params) { + signature.push_back(static_cast(value)); + } } } return signature; diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 22afc47f..82fff07a 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -3,6 +3,7 @@ #include #include +#include "core/ggml_extend.h" #include "core/ggml_extend_backend.h" #include "core/ggml_runner.h" #include "core/ggml_tensor_utils.h" @@ -12,6 +13,21 @@ using namespace sd; +ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + int64_t n_head, + ggml_tensor* mask, + bool skip_reshape, + bool flash_attn, + float kv_scale) { + if (ctx->attn_scale > 0.f) { + kv_scale = ctx->attn_scale; + } + return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale); +} + void GGMLRunner::alloc_params_ctx() { ggml_init_params params; params.mem_size = static_cast(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead()); @@ -511,6 +527,8 @@ GGMLRunnerContext GGMLRunner::get_context() { runner_ctx.ggml_ctx = compute_ctx; runner_ctx.backend = runtime_backend; runner_ctx.flash_attn_enabled = flash_attn_enabled; + runner_ctx.linear_scale = linear_scale; + runner_ctx.attn_scale = attn_scale; runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled; runner_ctx.circular_x_enabled = circular_x_enabled; runner_ctx.circular_y_enabled = circular_y_enabled; diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index 7960f5aa..5dc3fa5a 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -68,6 +68,8 @@ struct GGMLRunnerContext { ggml_backend_t backend = nullptr; ggml_context* ggml_ctx = nullptr; bool flash_attn_enabled = false; + float linear_scale = 0.f; + float attn_scale = 0.f; bool conv2d_direct_enabled = false; bool circular_x_enabled = false; bool circular_y_enabled = false; @@ -113,6 +115,16 @@ struct GGMLRunnerContext { } }; +ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + int64_t n_head, + ggml_tensor* mask = nullptr, + bool skip_reshape = false, + bool flash_attn = false, + float kv_scale = 1.f); + struct GGMLRunner { private: std::map logged_compute_bytes_; @@ -163,6 +175,8 @@ protected: const std::string final_result_name = "ggml_runner_final_result_tensor"; bool flash_attn_enabled = false; + float linear_scale = 0.f; + float attn_scale = 0.f; bool conv2d_direct_enabled = false; bool circular_x_enabled = false; bool circular_y_enabled = false; @@ -323,6 +337,11 @@ public: flash_attn_enabled = enabled; } + void set_scale_overrides(float linear_scale, float attn_scale) { + this->linear_scale = linear_scale; + this->attn_scale = attn_scale; + } + void set_conv2d_direct_enabled(bool enabled) { conv2d_direct_enabled = enabled; } diff --git a/src/extensions/photomaker_extension.cpp b/src/extensions/photomaker_extension.cpp index 05bb46c0..a15e751e 100644 --- a/src/extensions/photomaker_extension.cpp +++ b/src/extensions/photomaker_extension.cpp @@ -135,6 +135,7 @@ struct PhotoMakerExtension : public GenerationExtension { pm_version, 20.f, ctx.model_manager); + pmid_model->set_scale_overrides(ctx.params->linear_scale, ctx.params->attn_scale); if (pm_version == PM_VERSION_2) { LOG_INFO("using PhotoMaker Version 2"); } diff --git a/src/model.h b/src/model.h index f59f1de7..4b3ed316 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/adapter/ip_adapter.hpp b/src/model/adapter/ip_adapter.hpp index 2cc660b5..9f8dc9f6 100644 --- a/src/model/adapter/ip_adapter.hpp +++ b/src/model/adapter/ip_adapter.hpp @@ -95,7 +95,7 @@ namespace IPAdapter { int64_t L = kv->ne[1]; ggml_tensor* k = ggml_cont(ctx->ggml_ctx, ggml_view_3d(ctx->ggml_ctx, kv, dim, L, N, kv->nb[1], kv->nb[2], 0)); ggml_tensor* v = ggml_cont(ctx->ggml_ctx, ggml_view_3d(ctx->ggml_ctx, kv, dim, L, N, kv->nb[1], kv->nb[2], dim * kv->nb[0])); - ggml_tensor* attn = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, heads, nullptr, false, false); + ggml_tensor* attn = ggml_ext_attention_ext(ctx, q, k, v, heads, nullptr, false, false); attn = to_out->forward(ctx, attn); latents = ggml_add(ctx->ggml_ctx, latents, attn); diff --git a/src/model/adapter/pulid.hpp b/src/model/adapter/pulid.hpp index 65b3064a..3ba84fb6 100644 --- a/src/model/adapter/pulid.hpp +++ b/src/model/adapter/pulid.hpp @@ -63,12 +63,11 @@ public: k = ggml_cont(ctx->ggml_ctx, k); v = ggml_cont(ctx->ggml_ctx, v); - ggml_tensor* attn_out = ggml_ext_attention_ext( - ctx->ggml_ctx, ctx->backend, - q, k, v, - heads, - /*mask=*/nullptr, - /*diag_mask_inf=*/false); + ggml_tensor* attn_out = ggml_ext_attention_ext(ctx, + q, k, v, + heads, + /*mask=*/nullptr, + /*diag_mask_inf=*/false); ggml_tensor* out = to_out->forward(ctx, attn_out); return out; diff --git a/src/model/common/block.hpp b/src/model/common/block.hpp index 09db9250..5a5e6615 100644 --- a/src/model/common/block.hpp +++ b/src/model/common/block.hpp @@ -380,14 +380,14 @@ public: if (xtra_dim) { context->ne[0] = 320; // reset dim to orig } - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, inner_dim] + x = ggml_ext_attention_ext(ctx, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, inner_dim] if (has_ip && ctx->ip_context != nullptr && ctx->ip_scale != 0.0f) { auto to_k_ip = std::dynamic_pointer_cast(blocks["to_k_ip"]); auto to_v_ip = std::dynamic_pointer_cast(blocks["to_v_ip"]); auto k_ip = to_k_ip->forward(ctx, ctx->ip_context); auto v_ip = to_v_ip->forward(ctx, ctx->ip_context); - auto x_ip = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_ip, v_ip, n_head, nullptr, false, ctx->flash_attn_enabled); + auto x_ip = ggml_ext_attention_ext(ctx, q, k_ip, v_ip, n_head, nullptr, false, ctx->flash_attn_enabled); x = ggml_add(ctx->ggml_ctx, x, ggml_scale(ctx->ggml_ctx, x_ip, ctx->ip_scale)); } diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index cc0f6941..cc5470c3 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -206,6 +206,7 @@ public: ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* w = params["weight"]; + const float scale = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale; ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr; if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) { bool supports_fp8_matmul = false; @@ -309,6 +310,7 @@ public: __STATIC_INLINE__ bool support_get_rows(ggml_type wtype) { switch (wtype) { case GGML_TYPE_F16: + case GGML_TYPE_BF16: case GGML_TYPE_Q8_0: case GGML_TYPE_Q5_1: case GGML_TYPE_Q5_0: @@ -869,7 +871,7 @@ public: v = v_proj->forward(ctx, x); } - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, false); // [N, n_token, embed_dim] + x = ggml_ext_attention_ext(ctx, q, k, v, n_head, mask, false); // [N, n_token, embed_dim] x = out_proj->forward(ctx, x); // [N, n_token, embed_dim] return x; diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index 4f258039..447c3a23 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -1024,7 +1024,7 @@ namespace Rope { q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved); // [N*n_head, L, d_head] k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved); // [N*n_head, L, d_head] - auto x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, true, ctx->flash_attn_enabled, kv_scale); // [N, L, n_head*d_head] + auto x = ggml_ext_attention_ext(ctx, q, k, v, n_head, mask, true, ctx->flash_attn_enabled, kv_scale); // [N, L, n_head*d_head] return x; } }; // namespace Rope diff --git a/src/model/diffusion/anima.hpp b/src/model/diffusion/anima.hpp index 930eface..a8977c96 100644 --- a/src/model/diffusion/anima.hpp +++ b/src/model/diffusion/anima.hpp @@ -237,8 +237,7 @@ namespace Anima { } auto q_rope = Rope::apply_rope(ctx->ggml_ctx, q4, pe_q, false); auto k_rope = Rope::apply_rope(ctx->ggml_ctx, k4, pe_k, false); - attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, - ctx->backend, + attn_out = ggml_ext_attention_ext(ctx, q_rope, k_rope, v4, @@ -249,8 +248,7 @@ namespace Anima { } else { auto q_flat = ggml_reshape_3d(ctx->ggml_ctx, q4, head_dim * num_heads, L_q, N); auto k_flat = ggml_reshape_3d(ctx->ggml_ctx, k4, head_dim * num_heads, L_k, N); - attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, - ctx->backend, + attn_out = ggml_ext_attention_ext(ctx, q_flat, k_flat, v, diff --git a/src/model/diffusion/animatediff.hpp b/src/model/diffusion/animatediff.hpp index 6a94e21d..eb0ffba3 100644 --- a/src/model/diffusion/animatediff.hpp +++ b/src/model/diffusion/animatediff.hpp @@ -61,7 +61,7 @@ namespace AnimateDiff { auto k = to_k->forward(ctx, x_pe); auto v = to_v->forward(ctx, x_pe); - auto a = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, (int)num_heads, nullptr, false); + auto a = ggml_ext_attention_ext(ctx, q, k, v, (int)num_heads, nullptr, false); return to_out->forward(ctx, a); } }; diff --git a/src/model/diffusion/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp index c4a5df19..81bf142a 100644 --- a/src/model/diffusion/ernie_image.hpp +++ b/src/model/diffusion/ernie_image.hpp @@ -183,7 +183,7 @@ namespace ErnieImage { k = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, k, 0, 2, 1, 3)); // [N, heads, S, head_dim] k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]); - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled); // [N, S, hidden_size] + x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled); // [N, S, hidden_size] x = to_out_0->forward(ctx, x); return x; } diff --git a/src/model/diffusion/hidream_o1.hpp b/src/model/diffusion/hidream_o1.hpp index 677f81e7..9e5a901f 100644 --- a/src/model/diffusion/hidream_o1.hpp +++ b/src/model/diffusion/hidream_o1.hpp @@ -504,6 +504,10 @@ namespace HiDreamO1 { vision_runner->set_flash_attention_enabled(enabled); } + void set_scale_overrides(float linear_scale, float attn_scale) override { + vision_runner->set_scale_overrides(linear_scale, attn_scale); + } + void set_weight_adapter(const std::shared_ptr& adapter) override { vision_runner->set_weight_adapter(adapter); } diff --git a/src/model/diffusion/hunyuan.hpp b/src/model/diffusion/hunyuan.hpp index c111674f..c106b9cd 100644 --- a/src/model/diffusion/hunyuan.hpp +++ b/src/model/diffusion/hunyuan.hpp @@ -54,7 +54,7 @@ namespace Hunyuan { auto k = qkv_vec[1]; auto v = qkv_vec[2]; - auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled); + auto attn_out = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled); attn_out = self_attn_proj->forward(ctx, attn_out); // adaLN_modulation diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp index 41760581..84ecfe0b 100644 --- a/src/model/diffusion/krea2.hpp +++ b/src/model/diffusion/krea2.hpp @@ -232,8 +232,7 @@ namespace Krea2 { q = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, q), head_dim_ * heads, Lq, N); k = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, k), head_dim_ * kv_heads, Lk, N); v = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, v), head_dim_ * kv_heads, Lk, N); - return ggml_ext_attention_ext(ctx->ggml_ctx, - ctx->backend, + return ggml_ext_attention_ext(ctx, q, k, v, diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp index 118606e4..31a65afa 100644 --- a/src/model/diffusion/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -709,8 +709,7 @@ namespace LTXV { k = apply_hidden_rope(ctx->ggml_ctx, k, k_pe, heads, dim_head, rope_interleaved); } - auto out = ggml_ext_attention_ext(ctx->ggml_ctx, - ctx->backend, + auto out = ggml_ext_attention_ext(ctx, q, k, v, diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 3e2b69db..80c29ab5 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -215,8 +215,7 @@ namespace MiniMaxH3 { q = attention_layout(ctx->ggml_ctx, q); k = attention_layout(ctx->ggml_ctx, k); } - auto out = ggml_ext_attention_ext(ctx->ggml_ctx, - ctx->backend, + auto out = ggml_ext_attention_ext(ctx, q, k, v, diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index 454ef4b3..42599c67 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -365,8 +365,8 @@ public: ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { auto qkv = pre_attention(ctx, x); - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] - x = post_attention(ctx, x); // [N, n_token, dim] + x = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] + x = post_attention(ctx, x); // [N, n_token, dim] return x; } }; @@ -587,8 +587,8 @@ public: auto qkv2 = std::get<1>(qkv_intermediates); auto intermediates = std::get<2>(qkv_intermediates); - auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] - auto attn2_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] + auto attn_out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] + auto attn2_out = ggml_ext_attention_ext(ctx, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] x = post_attention_x(ctx, attn_out, attn2_out, @@ -604,7 +604,7 @@ public: auto qkv = qkv_intermediates.first; auto intermediates = qkv_intermediates.second; - auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] + auto attn_out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] x = post_attention(ctx, attn_out, intermediates[0], @@ -648,7 +648,7 @@ block_mixing(GGMLRunnerContext* ctx, qkv.push_back(ggml_concat(ctx->ggml_ctx, context_qkv[i], x_qkv[i], 1)); } - auto attn = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_context + n_token, hidden_size] + auto attn = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_context + n_token, hidden_size] auto context_attn = ggml_view_3d(ctx->ggml_ctx, attn, @@ -680,7 +680,7 @@ block_mixing(GGMLRunnerContext* ctx, } if (x_block->self_attn) { - auto attn2 = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, hidden_size] + auto attn2 = ggml_ext_attention_ext(ctx, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, hidden_size] x = x_block->post_attention_x(ctx, x_attn, diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 0e447cd0..9d611c4e 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 00000000..929f5d0d --- /dev/null +++ b/src/model/diffusion/sensenova_u1.h @@ -0,0 +1,846 @@ +#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.h" +#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); + 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), + 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/diffusion/wan.hpp b/src/model/diffusion/wan.hpp index 687f2838..363beeeb 100644 --- a/src/model/diffusion/wan.hpp +++ b/src/model/diffusion/wan.hpp @@ -193,7 +193,7 @@ namespace WAN { k = norm_k->forward(ctx, k); auto v = v_proj->forward(ctx, context); // [N, n_context, dim] - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] + x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] x = o_proj->forward(ctx, x); // [N, n_token, dim] return x; @@ -255,8 +255,8 @@ namespace WAN { k_img = norm_k_img->forward(ctx, k_img); auto v_img = v_img_proj->forward(ctx, context_img); // [N, context_img_len, dim] - auto img_x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_img, v_img, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] + auto img_x = ggml_ext_attention_ext(ctx, q, k_img, v_img, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] + x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] x = ggml_add(ctx->ggml_ctx, x, img_x); diff --git a/src/model/te/clip.hpp b/src/model/te/clip.hpp index 441a26c2..cac613ed 100644 --- a/src/model/te/clip.hpp +++ b/src/model/te/clip.hpp @@ -145,6 +145,10 @@ protected: params["position_embedding.weight"] = ggml_new_tensor_2d(ctx, position_wtype, embed_dim, num_positions); } + enum ggml_op param_usage_op(const std::string& name) const override { + return name == "token_embedding.weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE; + } + public: CLIPEmbeddings(int64_t embed_dim, int64_t vocab_size = 49408, diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index 4f1db3bc..de01e5de 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -1359,7 +1359,7 @@ namespace LLM { x = ggml_ext_cont(ctx->ggml_ctx, kqv); x = ggml_reshape_3d(ctx->ggml_ctx, x, head_dim * num_heads, n_token, N); } else { - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, false); // [N, n_token, hidden_size] + x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, false); // [N, n_token, hidden_size] } x = out_proj->forward(ctx, x); // [N, n_token, hidden_size] diff --git a/src/model/te/t5.hpp b/src/model/te/t5.hpp index f89488d3..adc68aa0 100644 --- a/src/model/te/t5.hpp +++ b/src/model/te/t5.hpp @@ -251,7 +251,7 @@ public: k = ggml_ext_scale(ctx->ggml_ctx, k, ::sqrtf(static_cast(d_head)), true); - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask); // [N, n_token, d_head * n_head] + x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask); // [N, n_token, d_head * n_head] x = out_proj->forward(ctx, x); // [N, n_token, model_dim] return {x, past_bias}; diff --git a/src/model/vae/auto_encoder_kl.hpp b/src/model/vae/auto_encoder_kl.hpp index 116a8321..115ff345 100644 --- a/src/model/vae/auto_encoder_kl.hpp +++ b/src/model/vae/auto_encoder_kl.hpp @@ -142,7 +142,7 @@ public: v = ggml_reshape_3d(ctx->ggml_ctx, v, c, h * w, n); // [N, h * w, in_channels] } - h_ = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); + h_ = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); if (use_linear) { h_ = proj_out->forward(ctx, h_); // [N, h * w, in_channels] diff --git a/src/model/vae/hunyuan_vae.hpp b/src/model/vae/hunyuan_vae.hpp index 099a15b7..1aab8ac5 100644 --- a/src/model/vae/hunyuan_vae.hpp +++ b/src/model/vae/hunyuan_vae.hpp @@ -193,7 +193,7 @@ namespace Hunyuan { v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b); // [b, c, t*h*w] v = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [b, t*h*w, c] - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [b, t*h*w, c] + x = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [b, t*h*w, c] x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [b, c, t*h*w] x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b); // [b*c, t, h, w] diff --git a/src/model/vae/mage_vae.hpp b/src/model/vae/mage_vae.hpp index 5567aec2..38ff89c4 100644 --- a/src/model/vae/mage_vae.hpp +++ b/src/model/vae/mage_vae.hpp @@ -253,7 +253,7 @@ namespace MageVAE { q = to_patches(ctx->ggml_ctx, q); k = to_patches(ctx->ggml_ctx, k); v = to_patches(ctx->ggml_ctx, v); - h = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); + h = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); h = from_patches(ctx->ggml_ctx, h, np, batch, hp, wp); if (pad_h > 0) { h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, height); diff --git a/src/model/vae/minimax_h3_audio_vae.hpp b/src/model/vae/minimax_h3_audio_vae.hpp index 926d2fbe..a4fc1447 100644 --- a/src/model/vae/minimax_h3_audio_vae.hpp +++ b/src/model/vae/minimax_h3_audio_vae.hpp @@ -174,8 +174,7 @@ namespace MiniMaxH3 { auto mask = ggml_diag_mask_inf(ctx->ggml_ctx, ggml_ext_zeros(ctx->ggml_ctx, sequence, sequence, 1, 1), 0); - auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, - ctx->backend, + auto attn_out = ggml_ext_attention_ext(ctx, q, k, v, diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index ec65fa37..77756643 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -291,8 +291,7 @@ namespace MiniMaxH3VAE { k = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f); q = apply_partial_rope(ctx->ggml_ctx, q, pe); k = apply_partial_rope(ctx->ggml_ctx, k, pe); - auto out = ggml_ext_attention_ext(ctx->ggml_ctx, - ctx->backend, + auto out = ggml_ext_attention_ext(ctx, q, k, v, diff --git a/src/model/vae/vae.hpp b/src/model/vae/vae.hpp index 03cfd5e2..2c952213 100644 --- a/src/model/vae/vae.hpp +++ b/src/model/vae/vae.hpp @@ -166,7 +166,7 @@ public: 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/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index a632d66b..a56392e8 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -615,8 +615,8 @@ namespace WAN { auto v = qkv_vec[2]; v = ggml_reshape_3d(ctx->ggml_ctx, v, h * w, c, n); // [t, c, h * w] - v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [t, h * w, c] - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [t, h * w, c] + v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [t, h * w, c] + x = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [t, h * w, c] x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [t, c, h * w] x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, c, n); // [t, c, h, w] diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 8325a606..6ab370c3 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -69,6 +69,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) { @@ -185,6 +187,15 @@ bool ModelLoader::parse_file(const std::string& file_path, const std::string& pr } parsed_dependencies_.push_back(stamp); 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 parse_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)) { @@ -463,6 +474,9 @@ SDVersion ModelLoader::get_sd_version() const { if (tensor_storage.name.find("net.img_embedder.proj1.weight") != std::string::npos) { return VERSION_MINIT2I; } + if (tensor_storage.name.find("language_model.model.layers.0.self_attn.q_proj_mot_gen.weight") != std::string::npos) { + return VERSION_SENSENOVA_U1_5; + } if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) { auto img_in = tensor_storage_map.find("model.diffusion_model.img_in.weight"); if (img_in != tensor_storage_map.end() && img_in->second.ne[0] == 128) { diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index a8797874..5d51e241 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -96,9 +96,13 @@ const char* model_version_to_str[] = { "SeFi-Image", "Krea2", "Mage Flow", + "SenseNova U1.5", "ESRGAN", }; +static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]), + "\nnumber of elements in model_version_to_str[] != VERSION_COUNT"); + void calculate_alphas_cumprod(float* alphas_cumprod, float linear_start = 0.00085f, float linear_end = 0.0120f, @@ -843,6 +847,12 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf } bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) { + for (float scale : {sd_ctx_params->linear_scale, sd_ctx_params->attn_scale}) { + if (!std::isfinite(scale) || scale < 0.f || (scale > 0.f && !std::isfinite(1.f / scale))) { + LOG_ERROR("scale overrides must be finite positive values, or 0 to keep model defaults"); + return false; + } + } auto configuration = std::make_unique(*sd_ctx_params); n_threads = sd_ctx_params->n_threads; enable_mmap = sd_ctx_params->enable_mmap; @@ -1306,6 +1316,9 @@ bool StableDiffusionGGML::build_denoiser() { 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; } @@ -1351,6 +1364,11 @@ bool StableDiffusionGGML::build_denoiser() { 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; @@ -2333,6 +2351,9 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr StableDiffusionGGML::decode_first_stage(const sd::Tensordecode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); const bool prefer_temporal_tiling = decode_video && first_stage_model->can_temporal_tile_decode(); while (decoded.empty() && - auto_fit_enabled && sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) { decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); } diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index 69c48742..aada0f45 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -459,6 +459,9 @@ namespace sd::pipeline { // 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->version)) { + auto* sensenova_conditioner = static_cast(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->version) && diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index 8d174e8d..da5a469e 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -27,6 +27,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" @@ -306,6 +307,12 @@ namespace sd::model_builders { tensor_storage_map, "model.diffusion_model.model.net", weight_manager); + } else if (sd_version_is_sensenova_u1(version)) { + result.conditioner = std::make_shared(); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "", + weight_manager); } else if (sd_version_is_anima(version)) { result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), tensor_storage_map, @@ -396,6 +403,21 @@ namespace sd::model_builders { "ip_adapter", weight_manager); } + if (result.conditioner) { + result.conditioner->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale); + } + if (result.diffusion) { + result.diffusion->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale); + } + if (result.high_noise_diffusion) { + result.high_noise_diffusion->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale); + } + if (result.clip_vision) { + result.clip_vision->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale); + } + if (result.ip_adapter) { + result.ip_adapter->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale); + } runners = std::move(result); return true; } @@ -493,7 +515,7 @@ namespace sd::model_builders { } }; - 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"); result.vae = std::make_shared(version, ctx.backends.runtime_backend(SDBackendModule::VAE), @@ -531,6 +553,15 @@ namespace sd::model_builders { result.preview->set_conv2d_direct_enabled(true); } } + if (result.vae) { + result.vae->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale); + } + if (result.preview) { + result.preview->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale); + } + if (result.audio) { + result.audio->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale); + } runners = std::move(result); return true; } @@ -552,6 +583,7 @@ namespace sd::model_builders { LOG_INFO("Using Conv2d direct in the control net"); control_net->set_conv2d_direct_enabled(true); } + control_net->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale); runner = std::move(control_net); return true; } diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index 981669c2..ec89406d 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -1488,6 +1488,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 cda44e92..beb899cd 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -156,6 +156,7 @@ const char* prediction_to_str[] = { "flux_flow", "sefi_flow", "minit2i_flow", + "sensenova_u1_flow", }; const char* sd_prediction_name(enum prediction_t prediction) { @@ -322,6 +323,8 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { sd_ctx_params->eager_load = false; sd_ctx_params->enable_mmap = false; sd_ctx_params->diffusion_flash_attn = false; + sd_ctx_params->linear_scale = 0.f; + sd_ctx_params->attn_scale = 0.f; sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; sd_ctx_params->backend = nullptr; sd_ctx_params->params_backend = nullptr; @@ -373,6 +376,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "auto_fit: %s\n" "flash_attn: %s\n" "diffusion_flash_attn: %s\n" + "linear_scale: %g\n" + "attn_scale: %g\n" "vae_format: %s\n", SAFE_STR(sd_ctx_params->model_path), SAFE_STR(sd_ctx_params->clip_l_path), @@ -408,6 +413,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { BOOL_STR(sd_ctx_params->auto_fit), BOOL_STR(sd_ctx_params->flash_attn), BOOL_STR(sd_ctx_params->diffusion_flash_attn), + sd_ctx_params->linear_scale, + sd_ctx_params->attn_scale, sd_vae_format_name(sd_ctx_params->vae_format)); return buf; @@ -695,6 +702,13 @@ SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx) { return sd_ctx->sd->control_net != nullptr; } +const char* sd_get_model_version_name(const sd_ctx_t* sd_ctx) { + if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_ctx->sd->version >= VERSION_COUNT) { + return "Unknown"; + } + return model_version_to_str[sd_ctx->sd->version]; +} + enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) { return sd::pipeline::default_sample_method(sd_ctx != nullptr ? sd_ctx->sd : nullptr); } diff --git a/src/tokenizers/qwen2_tokenizer.cpp b/src/tokenizers/qwen2_tokenizer.cpp index a0c2e3bd..3ceeba43 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 04e92c2c..12739466 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 00000000..32e259b9 --- /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 00000000..31778ec6 --- /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__