Compare commits

...

2 Commits

Author SHA1 Message Date
KenForever
3590aa8d62
feat: add MiniT2I support (#1683) 2026-07-02 00:46:51 +08:00
stduhpf
1a13107bac
feat: add imatrix support (#633) 2026-07-01 23:26:37 +08:00
23 changed files with 1667 additions and 44 deletions

View File

@ -48,6 +48,7 @@ API and command-line option may change frequently.***
- [PiD](./docs/pid.md) - [PiD](./docs/pid.md)
- [LongCat Image](./docs/longcat_image.md) - [LongCat Image](./docs/longcat_image.md)
- [Z-Image](./docs/z_image.md) - [Z-Image](./docs/z_image.md)
- [MiniT2I](./docs/minit2i.md)
- [Ovis-Image](./docs/ovis_image.md) - [Ovis-Image](./docs/ovis_image.md)
- [Anima](./docs/anima.md) - [Anima](./docs/anima.md)
- [ERNIE-Image](./docs/ernie_image.md) - [ERNIE-Image](./docs/ernie_image.md)

59
docs/imatrix.md Normal file
View File

@ -0,0 +1,59 @@
# Importance Matrix (imatrix) Quantization
## What is an Importance Matrix?
Quantization reduces the precision of a model's weights, decreasing its size and computational requirements. However, this can lead to a loss of quality. An importance matrix helps mitigate this by identifying which weights are *most* important for the model's performance. During quantization, these important weights are preserved with higher precision, while less important weights are quantized more aggressively. This allows for better overall quality at a given quantization level.
This originates from work done with language models in [llama.cpp](https://github.com/ggml-org/llama.cpp/blob/master/tools/imatrix/README.md).
## Usage
The imatrix feature involves two main steps: *training* the matrix and *using* it during quantization.
### Training the Importance Matrix
To generate an imatrix, run stable-diffusion.cpp with the `--imat-out` flag, specifying the output filename. This process runs alongside normal image generation.
```bash
sd.exe [same exact parameters as normal generation] --imat-out imatrix.dat
```
* **`[same exact parameters as normal generation]`**: Use the same command-line arguments you would normally use for image generation (e.g., prompt, dimensions, sampling method, etc.).
* **`--imat-out imatrix.dat`**: Specifies the output file for the generated imatrix.
You can generate multiple images at once using the `-b` flag to speed up the training process.
### Continuing Training an Existing Matrix
If you want to refine an existing imatrix, use the `--imat-in` flag *in addition* to `--imat-out`. This will load the existing matrix and continue training it.
```bash
sd.exe [same exact parameters as normal generation] --imat-out imatrix.dat --imat-in imatrix.dat
```
With that, you can train and refine the imatrix while generating images like you'd normally do.
### Using Multiple Matrices
You can load and merge multiple imatrices together:
```bash
sd.exe [same exact parameters as normal generation] --imat-out imatrix.dat --imat-in imatrix.dat --imat-in imatrix2.dat
```
### Quantizing with an Importance Matrix
To quantize a model using a trained imatrix, use the `-M convert` option (or equivalent quantization command) and the `--imat-in` flag, specifying the imatrix file.
```bash
sd.exe -M convert [same exact parameters as normal quantization] --imat-in imatrix.dat
```
* **`[same exact parameters as normal quantization]`**: Use the same command-line arguments you would normally use for quantization (e.g., target quantization method, input/output filenames).
* **`--imat-in imatrix.dat`**: Specifies the imatrix file to use during quantization. You can specify multiple `--imat-in` flags to combine multiple matrices.
## Important Considerations
* The quality of the imatrix depends on the prompts and settings used during training. Use prompts and settings representative of the types of images you intend to generate for the best results.
* Experiment with different training parameters (e.g., number of images, prompt variations) to optimize the imatrix for your specific use case.
* The performance impact of training an imatrix during image generation or using an imatrix for quantization is negligible.
* Using already quantized models to train the imatrix seems to be working fine.

48
docs/minit2i.md Normal file
View File

@ -0,0 +1,48 @@
# How to Use
MiniT2I uses a MiniT2I diffusion transformer and `google/flan-t5-large` as the text encoder.
## Download weights
- Download MiniT2I diffusion model
- safetensors: https://huggingface.co/MiniT2I/MiniT2I/tree/main/minit2i-b-16/transformer (`diffusion_pytorch_model.safetensors`)
- Download flan-t5-large text encoder
- safetensors: https://huggingface.co/google/flan-t5-large/tree/main (`model.safetensors`)
## Examples
### Mac Metal
```
./bin/sd-cli \
--backend metal \
--diffusion-model ../models/minit2i/diffusion_pytorch_model.safetensors \
--t5xxl ../models/flan-t5-large/model.safetensors \
--prompt "a cat" \
--steps 100 \
--cfg-scale 6 \
--width 512 \
--height 512 \
--seed 42 \
--sampling-method euler \
--rng cpu \
--output minit2i_metal.png \
--threads 8
```
### CUDA with diffusion flash attention
```
./bin/sd-cli \
--diffusion-model ../models/minit2i/diffusion_pytorch_model.safetensors \
--t5xxl ../models/flan-t5-large/model.safetensors \
--prompt "a cat" \
--steps 100 \
--cfg-scale 6 \
--width 512 \
--height 512 \
--seed 42 \
--sampling-method euler \
--diffusion-fa \
--output minit2i_cuda.png
```

View File

@ -1,6 +1,7 @@
#include <stdio.h> #include <stdio.h>
#include <string.h> #include <string.h>
#include <time.h> #include <time.h>
#include <algorithm>
#include <cctype> #include <cctype>
#include <filesystem> #include <filesystem>
#include <functional> #include <functional>
@ -53,6 +54,9 @@ struct SDCliParams {
bool metadata_brief = false; bool metadata_brief = false;
bool metadata_all = false; bool metadata_all = false;
std::string imatrix_out;
std::vector<std::string> imatrix_in;
bool normal_exit = false; bool normal_exit = false;
ArgOptions get_options() { ArgOptions get_options() {
@ -79,6 +83,11 @@ struct SDCliParams {
"path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp", "path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp",
0, 0,
&preview_path}, &preview_path},
{"",
"--imat-out",
"compute the imatrix for this run and save it to the provided path",
0,
&imatrix_out},
}; };
options.int_options = { options.int_options = {
@ -179,6 +188,14 @@ struct SDCliParams {
return -1; return -1;
}; };
auto on_imatrix_in_arg = [&](int argc, const char** argv, int index) {
if (++index >= argc) {
return -1;
}
imatrix_in.push_back(argv[index]);
return 1;
};
options.manual_options = { options.manual_options = {
{"-M", {"-M",
"--mode", "--mode",
@ -192,6 +209,10 @@ struct SDCliParams {
"--help", "--help",
"show this help message and exit", "show this help message and exit",
on_help_arg}, on_help_arg},
{"",
"--imat-in",
"load an imatrix file for quantization or continued collection; can be specified multiple times",
on_imatrix_in_arg},
}; };
return options; return options;
@ -253,6 +274,7 @@ struct SDCliParams {
<< " preview_fps: " << preview_fps << ",\n" << " preview_fps: " << preview_fps << ",\n"
<< " taesd_preview: " << (taesd_preview ? "true" : "false") << ",\n" << " taesd_preview: " << (taesd_preview ? "true" : "false") << ",\n"
<< " preview_noisy: " << (preview_noisy ? "true" : "false") << ",\n" << " preview_noisy: " << (preview_noisy ? "true" : "false") << ",\n"
<< " imatrix_out: \"" << imatrix_out << "\",\n"
<< " metadata_raw: " << (metadata_raw ? "true" : "false") << ",\n" << " metadata_raw: " << (metadata_raw ? "true" : "false") << ",\n"
<< " metadata_brief: " << (metadata_brief ? "true" : "false") << ",\n" << " metadata_brief: " << (metadata_brief ? "true" : "false") << ",\n"
<< " metadata_all: " << (metadata_all ? "true" : "false") << "\n" << " metadata_all: " << (metadata_all ? "true" : "false") << "\n"
@ -605,8 +627,27 @@ int main(int argc, const char* argv[]) {
LOG_DEBUG("%s", ctx_params.to_string().c_str()); LOG_DEBUG("%s", ctx_params.to_string().c_str());
LOG_DEBUG("%s", gen_params.to_string().c_str()); LOG_DEBUG("%s", gen_params.to_string().c_str());
if (!cli_params.imatrix_out.empty()) {
if (fs::exists(cli_params.imatrix_out) &&
std::find(cli_params.imatrix_in.begin(), cli_params.imatrix_in.end(), cli_params.imatrix_out) == cli_params.imatrix_in.end()) {
LOG_WARN("imatrix file '%s' already exists and will be overwritten", cli_params.imatrix_out.c_str());
}
enable_imatrix_collection();
}
for (const auto& in_file : cli_params.imatrix_in) {
LOG_INFO("loading imatrix from '%s'", in_file.c_str());
if (!load_imatrix(in_file.c_str())) {
LOG_WARN("failed to load imatrix from '%s'", in_file.c_str());
}
}
if (cli_params.mode == CONVERT) { if (cli_params.mode == CONVERT) {
bool success = convert(ctx_params.model_path.c_str(), bool success = convert_with_components(ctx_params.model_path.c_str(),
ctx_params.clip_l_path.c_str(),
ctx_params.clip_g_path.c_str(),
ctx_params.t5xxl_path.c_str(),
ctx_params.diffusion_model_path.c_str(),
ctx_params.vae_path.c_str(), ctx_params.vae_path.c_str(),
cli_params.output_path.c_str(), cli_params.output_path.c_str(),
ctx_params.wtype, ctx_params.wtype,
@ -833,6 +874,11 @@ int main(int argc, const char* argv[]) {
return 1; return 1;
} }
if (!cli_params.imatrix_out.empty()) {
LOG_INFO("saving imatrix to '%s'", cli_params.imatrix_out.c_str());
save_imatrix(cli_params.imatrix_out.c_str());
}
free_sd_audio(generated_audio); free_sd_audio(generated_audio);
return 0; return 0;

View File

@ -710,7 +710,18 @@ bool SDContextParams::resolve(SDMode mode) {
} }
bool SDContextParams::validate(SDMode mode) { bool SDContextParams::validate(SDMode mode) {
if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) { if (mode == CONVERT) {
const bool has_convert_input = model_path.length() != 0 ||
clip_l_path.length() != 0 ||
clip_g_path.length() != 0 ||
t5xxl_path.length() != 0 ||
diffusion_model_path.length() != 0 ||
vae_path.length() != 0;
if (!has_convert_input) {
LOG_ERROR("error: convert mode needs at least one model input path\n");
return false;
}
} else if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) {
LOG_ERROR("error: the following arguments are required: model_path/diffusion_model\n"); LOG_ERROR("error: the following arguments are required: model_path/diffusion_model\n");
return false; return false;
} }

View File

@ -84,6 +84,7 @@ enum prediction_t {
FLOW_PRED, FLOW_PRED,
FLUX_FLOW_PRED, FLUX_FLOW_PRED,
SEFI_FLOW_PRED, SEFI_FLOW_PRED,
MINIT2I_FLOW_PRED,
PREDICTION_COUNT PREDICTION_COUNT
}; };
@ -407,14 +408,17 @@ typedef struct {
} sd_vid_gen_params_t; } sd_vid_gen_params_t;
typedef struct sd_ctx_t sd_ctx_t; typedef struct sd_ctx_t sd_ctx_t;
struct ggml_tensor;
typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data); typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data);
typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data); typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data);
typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data); typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data);
typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void* user_data);
SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data); SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data);
SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data); SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data);
SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data); SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data);
SD_API void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data);
SD_API int32_t sd_get_num_physical_cores(); SD_API int32_t sd_get_num_physical_cores();
SD_API const char* sd_get_system_info(); SD_API const char* sd_get_system_info();
SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx); SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx);
@ -503,6 +507,17 @@ SD_API bool convert(const char* input_path,
const char* tensor_type_rules, const char* tensor_type_rules,
bool convert_name); bool convert_name);
SD_API bool convert_with_components(const char* model_path,
const char* clip_l_path,
const char* clip_g_path,
const char* t5xxl_path,
const char* diffusion_model_path,
const char* vae_path,
const char* output_path,
enum sd_type_t output_type,
const char* tensor_type_rules,
bool convert_name);
SD_API bool preprocess_canny(sd_image_t image, SD_API bool preprocess_canny(sd_image_t image,
float high_threshold, float high_threshold,
float low_threshold, float low_threshold,
@ -510,6 +525,11 @@ SD_API bool preprocess_canny(sd_image_t image,
float strong, float strong,
bool inverse); bool inverse);
SD_API bool load_imatrix(const char* imatrix_path);
SD_API void save_imatrix(const char* imatrix_path);
SD_API void enable_imatrix_collection(void);
SD_API void disable_imatrix_collection(void);
SD_API const char* sd_commit(void); SD_API const char* sd_commit(void);
SD_API const char* sd_version(void); SD_API const char* sd_version(void);

View File

@ -1378,6 +1378,101 @@ struct T5CLIPEmbedder : public Conditioner {
} }
}; };
struct MiniT2IConditioner : public Conditioner {
T5UniGramTokenizer tokenizer;
std::shared_ptr<T5Runner> t5;
size_t prompt_length = 256;
MiniT2IConditioner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr) {
bool use_t5 = false;
for (const auto& pair : tensor_storage_map) {
if (pair.first.find("text_encoders.t5xxl") != std::string::npos) {
use_t5 = true;
break;
}
}
if (!use_t5) {
LOG_WARN("IMPORTANT NOTICE: No MiniT2I T5 text encoder provided, cannot process prompts!");
return;
}
t5 = std::make_shared<T5Runner>(backend, tensor_storage_map, "text_encoders.t5xxl.transformer", false, weight_manager);
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
if (t5) {
t5->set_max_graph_vram_bytes(max_vram_bytes);
}
}
void set_stream_layers_enabled(bool enabled) override {
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_flash_attention_enabled(bool enabled) override {
if (t5) {
t5->set_flash_attention_enabled(enabled);
}
}
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
if (t5) {
t5->set_weight_adapter(adapter);
}
}
void runner_done() override {
if (t5) {
t5->runner_done();
}
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
SDCondition result;
if (!t5) {
result.c_crossattn = sd::Tensor<float>::zeros({1024, static_cast<int64_t>(prompt_length)});
result.c_vector = sd::Tensor<float>::zeros({static_cast<int64_t>(prompt_length)});
return result;
}
std::vector<int> tokens = tokenizer.encode(conditioner_params.text);
if (tokens.size() > prompt_length) {
tokens.resize(prompt_length);
}
std::vector<float> mask(tokens.size(), 1.0f);
while (tokens.size() < prompt_length) {
tokens.push_back(tokenizer.PAD_TOKEN_ID);
mask.push_back(0.0f);
}
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
std::vector<float> t5_mask(mask.size(), 0.0f);
for (size_t i = 0; i < mask.size(); ++i) {
t5_mask[i] = mask[i] > 0.0f ? 0.0f : -HUGE_VALF;
}
sd::Tensor<float> hidden_states = t5->compute(n_threads,
input_ids,
sd::Tensor<float>::from_vector(t5_mask),
false,
true,
true);
GGML_ASSERT(!hidden_states.empty());
result.c_crossattn = std::move(hidden_states);
result.c_vector = sd::Tensor<float>::from_vector(mask);
return result;
}
};
struct AnimaConditioner : public Conditioner { struct AnimaConditioner : public Conditioner {
std::shared_ptr<BPETokenizer> qwen_tokenizer; std::shared_ptr<BPETokenizer> qwen_tokenizer;
T5UniGramTokenizer t5_tokenizer; T5UniGramTokenizer t5_tokenizer;

View File

@ -76,29 +76,22 @@ static bool load_tensors_for_export(ModelLoader& model_loader,
return success; return success;
} }
bool convert(const char* input_path, static bool init_convert_path(ModelLoader& model_loader, const char* path, const char* prefix, bool& loaded_any) {
const char* vae_path, if (path == nullptr || strlen(path) == 0) {
return true;
}
if (!model_loader.init_from_file(path, prefix)) {
LOG_ERROR("init model loader from file failed: '%s'", path);
return false;
}
loaded_any = true;
return true;
}
static bool export_loaded_model(ModelLoader& model_loader,
const char* output_path, const char* output_path,
sd_type_t output_type, sd_type_t output_type,
const char* tensor_type_rules, const char* tensor_type_rules) {
bool convert_name) {
ModelLoader model_loader;
if (!model_loader.init_from_file(input_path)) {
LOG_ERROR("init model loader from file failed: '%s'", input_path);
return false;
}
if (vae_path != nullptr && strlen(vae_path) > 0) {
if (!model_loader.init_from_file(vae_path, "vae.")) {
LOG_ERROR("init model loader from file failed: '%s'", vae_path);
return false;
}
}
if (convert_name) {
model_loader.convert_tensors_name();
}
ggml_type type = sd_type_to_ggml_type(output_type); ggml_type type = sd_type_to_ggml_type(output_type);
bool output_is_safetensors = ends_with(output_path, ".safetensors"); bool output_is_safetensors = ends_with(output_path, ".safetensors");
TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules); TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
@ -136,3 +129,55 @@ bool convert(const char* input_path,
ggml_free(ggml_ctx); ggml_free(ggml_ctx);
return success; return success;
} }
bool convert_with_components(const char* model_path,
const char* clip_l_path,
const char* clip_g_path,
const char* t5xxl_path,
const char* diffusion_model_path,
const char* vae_path,
const char* output_path,
sd_type_t output_type,
const char* tensor_type_rules,
bool convert_name) {
ModelLoader model_loader;
bool loaded_any = false;
if (!init_convert_path(model_loader, model_path, "", loaded_any) ||
!init_convert_path(model_loader, clip_l_path, "text_encoders.clip_l.transformer.", loaded_any) ||
!init_convert_path(model_loader, clip_g_path, "text_encoders.clip_g.transformer.", loaded_any) ||
!init_convert_path(model_loader, t5xxl_path, "text_encoders.t5xxl.transformer.", loaded_any) ||
!init_convert_path(model_loader, diffusion_model_path, "model.diffusion_model.", loaded_any) ||
!init_convert_path(model_loader, vae_path, "vae.", loaded_any)) {
return false;
}
if (!loaded_any) {
LOG_ERROR("no input model path provided for convert");
return false;
}
if (convert_name) {
model_loader.convert_tensors_name();
}
return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules);
}
bool convert(const char* input_path,
const char* vae_path,
const char* output_path,
sd_type_t output_type,
const char* tensor_type_rules,
bool convert_name) {
return convert_with_components(input_path,
nullptr,
nullptr,
nullptr,
nullptr,
vae_path,
output_path,
output_type,
tensor_type_rules,
convert_name);
}

View File

@ -2473,7 +2473,10 @@ protected:
sd_backend_cpu_set_n_threads(runtime_backend, n_threads); sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
} }
ggml_status status = ggml_backend_graph_compute(runtime_backend, gf); ggml_status status = sd_backend_graph_compute_with_eval_callback(runtime_backend,
gf,
sd_get_backend_eval_callback(),
sd_get_backend_eval_callback_data());
if (status != GGML_STATUS_SUCCESS) { if (status != GGML_STATUS_SUCCESS) {
LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
return std::nullopt; return std::nullopt;

View File

@ -9,6 +9,7 @@
#include <vector> #include <vector>
#include "core/util.h" #include "core/util.h"
#include "ggml/src/ggml-impl.h"
#include "stable-diffusion.h" #include "stable-diffusion.h"
static std::string trim_copy(const std::string& value) { static std::string trim_copy(const std::string& value) {
@ -110,7 +111,67 @@ static std::string resolve_first_device_by_type(enum ggml_backend_dev_type type)
if (dev == nullptr) { if (dev == nullptr) {
return ""; return "";
} }
return ggml_backend_dev_name(dev); const char* dev_name = ggml_backend_dev_name(dev);
if (dev_name != nullptr && dev_name[0] != '\0') {
return dev_name;
}
ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
const char* reg_name = reg != nullptr ? ggml_backend_reg_name(reg) : nullptr;
return reg_name != nullptr ? reg_name : "";
}
static ggml_backend_dev_t resolve_first_device_by_registry_name(const std::string& name) {
std::string lower = lower_copy(trim_copy(name));
if (lower == "metal") {
lower = "mtl";
}
if (lower.empty()) {
return nullptr;
}
const size_t device_count = ggml_backend_dev_count();
for (size_t i = 0; i < device_count; ++i) {
ggml_backend_dev_t dev = ggml_backend_dev_get(i);
ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
if (reg == nullptr) {
continue;
}
const char* reg_name = ggml_backend_reg_name(reg);
if (reg_name != nullptr && lower_copy(reg_name) == lower) {
return dev;
}
}
return nullptr;
}
static ggml_backend_dev_t resolve_device_by_name(const std::string& name) {
const std::string lower = lower_copy(trim_copy(name));
if (lower.empty()) {
return nullptr;
}
const size_t device_count = ggml_backend_dev_count();
for (size_t i = 0; i < device_count; ++i) {
ggml_backend_dev_t dev = ggml_backend_dev_get(i);
const char* dev_name = ggml_backend_dev_name(dev);
if (dev_name != nullptr && lower_copy(dev_name) == lower) {
return dev;
}
}
return nullptr;
}
static std::string backend_device_name(ggml_backend_dev_t dev) {
if (dev == nullptr) {
return "";
}
const char* name = ggml_backend_dev_name(dev);
if (name != nullptr && name[0] != '\0') {
return name;
}
ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
const char* reg_name = reg != nullptr ? ggml_backend_reg_name(reg) : nullptr;
return reg_name != nullptr ? reg_name : "";
} }
static ggml_backend_buffer_t ggml_backend_tensor_buffer(const struct ggml_tensor* tensor) { static ggml_backend_buffer_t ggml_backend_tensor_buffer(const struct ggml_tensor* tensor) {
@ -296,6 +357,10 @@ std::string sd_backend_resolve_name(const std::string& name) {
return resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU); return resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU);
} }
if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(requested)) {
return backend_device_name(dev);
}
const size_t device_count = ggml_backend_dev_count(); const size_t device_count = ggml_backend_dev_count();
for (size_t i = 0; i < device_count; ++i) { for (size_t i = 0; i < device_count; ++i) {
ggml_backend_dev_t dev = ggml_backend_dev_get(i); ggml_backend_dev_t dev = ggml_backend_dev_get(i);
@ -328,7 +393,20 @@ static ggml_backend_t init_named_backend(const std::string& name) {
return ggml_backend_init_best(); return ggml_backend_init_best();
} }
if (ggml_backend_dev_t dev = resolve_device_by_name(name)) {
return ggml_backend_dev_init(dev, nullptr);
}
if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(name)) {
return ggml_backend_dev_init(dev, nullptr);
}
std::string resolved = sd_backend_resolve_name(name); std::string resolved = sd_backend_resolve_name(name);
if (ggml_backend_dev_t dev = resolve_device_by_name(resolved)) {
return ggml_backend_dev_init(dev, nullptr);
}
if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(resolved)) {
return ggml_backend_dev_init(dev, nullptr);
}
if (resolved.empty()) { if (resolved.empty()) {
return nullptr; return nullptr;
} }
@ -364,6 +442,68 @@ bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
return false; return false;
} }
static ggml_cgraph sd_ggml_graph_view(ggml_cgraph* cgraph0, int i0, int i1) {
ggml_cgraph cgraph = {
/*.size =*/0,
/*.n_nodes =*/i1 - i0,
/*.n_leafs =*/0,
/*.nodes =*/cgraph0->nodes + i0,
/*.grads =*/nullptr,
/*.grad_accs =*/nullptr,
/*.leafs =*/nullptr,
/*.use_counts =*/cgraph0->use_counts,
/*.visited_hash_set =*/cgraph0->visited_hash_set,
/*.order =*/cgraph0->order,
/*.uid =*/0,
};
return cgraph;
}
ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
ggml_cgraph* gf,
sd_graph_eval_callback_t callback_eval,
void* callback_eval_user_data) {
if (callback_eval == nullptr) {
return ggml_backend_graph_compute(backend, gf);
}
ggml_status status = GGML_STATUS_SUCCESS;
const int n_nodes = ggml_graph_n_nodes(gf);
bool stopped = false;
for (int j0 = 0; j0 < n_nodes; ++j0) {
ggml_tensor* t = ggml_graph_node(gf, j0);
bool need = callback_eval(t, true, callback_eval_user_data);
int j1 = j0;
while (!need && j1 < n_nodes - 1) {
t = ggml_graph_node(gf, ++j1);
need = callback_eval(t, true, callback_eval_user_data);
}
ggml_cgraph gv = sd_ggml_graph_view(gf, j0, j1 + 1);
status = ggml_backend_graph_compute_async(backend, &gv);
if (status != GGML_STATUS_SUCCESS) {
break;
}
ggml_backend_synchronize(backend);
if (need && !callback_eval(t, false, callback_eval_user_data)) {
stopped = true;
break;
}
j0 = j1;
}
ggml_backend_synchronize(backend);
if (stopped && status == GGML_STATUS_SUCCESS) {
status = GGML_STATUS_ABORTED;
}
return status;
}
const char* sd_get_system_info() { const char* sd_get_system_info() {
static std::string cache_info = []() -> std::string { static std::string cache_info = []() -> std::string {
ggml_backend_load_all_once(); ggml_backend_load_all_once();
@ -599,7 +739,7 @@ bool SDBackendManager::validate(std::string* error) const {
} }
return false; return false;
} }
if (!sd_backend_resolve_name(name).empty()) { if (!sd_backend_resolve_name(name).empty() || resolve_first_device_by_registry_name(name) != nullptr) {
return true; return true;
} }
if (error != nullptr) { if (error != nullptr) {

View File

@ -9,6 +9,7 @@
#include "ggml-backend.h" #include "ggml-backend.h"
#include "ggml.h" #include "ggml.h"
#include "stable-diffusion.h"
enum class SDBackendModule { enum class SDBackendModule {
DIFFUSION, DIFFUSION,
@ -71,6 +72,10 @@ bool sd_backend_is(ggml_backend_t backend, const std::string& name);
bool sd_backend_is_cpu(ggml_backend_t backend); bool sd_backend_is_cpu(ggml_backend_t backend);
ggml_backend_t sd_backend_cpu_init(); ggml_backend_t sd_backend_cpu_init();
bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads); bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
ggml_cgraph* gf,
sd_graph_eval_callback_t callback_eval,
void* callback_eval_user_data);
std::string sd_backend_resolve_name(const std::string& name); std::string sd_backend_resolve_name(const std::string& name);
const char* sd_backend_module_name(SDBackendModule module); const char* sd_backend_module_name(SDBackendModule module);
void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value); void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);

View File

@ -346,6 +346,9 @@ int sd_preview_interval = 1;
bool sd_preview_denoised = true; bool sd_preview_denoised = true;
bool sd_preview_noisy = false; bool sd_preview_noisy = false;
static sd_graph_eval_callback_t sd_backend_eval_cb = nullptr;
static void* sd_backend_eval_cb_data = nullptr;
std::u32string utf8_to_utf32(const std::string& utf8_str) { std::u32string utf8_to_utf32(const std::string& utf8_str) {
std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter; std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
return converter.from_bytes(utf8_str); return converter.from_bytes(utf8_str);
@ -629,6 +632,11 @@ void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode, int interval, b
sd_preview_noisy = noisy; sd_preview_noisy = noisy;
} }
void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data) {
sd_backend_eval_cb = cb;
sd_backend_eval_cb_data = data;
}
sd_preview_cb_t sd_get_preview_callback() { sd_preview_cb_t sd_get_preview_callback() {
return sd_preview_cb; return sd_preview_cb;
} }
@ -649,6 +657,14 @@ bool sd_should_preview_noisy() {
return sd_preview_noisy; return sd_preview_noisy;
} }
sd_graph_eval_callback_t sd_get_backend_eval_callback() {
return sd_backend_eval_cb;
}
void* sd_get_backend_eval_callback_data() {
return sd_backend_eval_cb_data;
}
sd_progress_cb_t sd_get_progress_callback() { sd_progress_cb_t sd_get_progress_callback() {
return sd_progress_cb; return sd_progress_cb;
} }

View File

@ -98,6 +98,9 @@ int sd_get_preview_interval();
bool sd_should_preview_denoised(); bool sd_should_preview_denoised();
bool sd_should_preview_noisy(); bool sd_should_preview_noisy();
sd_graph_eval_callback_t sd_get_backend_eval_callback();
void* sd_get_backend_eval_callback_data();
// test if the backend is a specific one, e.g. "CUDA", "ROCm", "Vulkan" etc. // test if the backend is a specific one, e.g. "CUDA", "ROCm", "Vulkan" etc.
bool sd_backend_is(ggml_backend_t backend, const std::string& name); bool sd_backend_is(ggml_backend_t backend, const std::string& name);

View File

@ -46,6 +46,7 @@ enum SDVersion {
VERSION_OVIS_IMAGE, VERSION_OVIS_IMAGE,
VERSION_ERNIE_IMAGE, VERSION_ERNIE_IMAGE,
VERSION_LENS, VERSION_LENS,
VERSION_MINIT2I,
VERSION_LONGCAT, VERSION_LONGCAT,
VERSION_PID, VERSION_PID,
VERSION_IDEOGRAM4, VERSION_IDEOGRAM4,
@ -174,6 +175,13 @@ static inline bool sd_version_is_lens(SDVersion version) {
return false; return false;
} }
static inline bool sd_version_is_minit2i(SDVersion version) {
if (version == VERSION_MINIT2I) {
return true;
}
return false;
}
static inline bool sd_version_is_pid(SDVersion version) { static inline bool sd_version_is_pid(SDVersion version) {
if (version == VERSION_PID) { if (version == VERSION_PID) {
return true; return true;
@ -247,6 +255,7 @@ static inline bool sd_version_is_dit(SDVersion version) {
sd_version_is_boogu_image(version) || sd_version_is_boogu_image(version) ||
sd_version_is_ernie_image(version) || sd_version_is_ernie_image(version) ||
sd_version_is_lens(version) || sd_version_is_lens(version) ||
sd_version_is_minit2i(version) ||
sd_version_is_longcat(version) || sd_version_is_longcat(version) ||
sd_version_is_pid(version) || sd_version_is_pid(version) ||
sd_version_is_ideogram4(version) || sd_version_is_ideogram4(version) ||

View File

@ -0,0 +1,611 @@
#ifndef __SD_MODEL_DIFFUSION_MINIT2I_HPP__
#define __SD_MODEL_DIFFUSION_MINIT2I_HPP__
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <string>
#include <vector>
#include "core/ggml_extend.hpp"
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/model.hpp"
#include "model_loader.h"
namespace MiniT2I {
constexpr int MINIT2I_GRAPH_SIZE = 196608;
struct MiniT2IConfig {
int64_t image_size = 512;
int64_t patch_size = 16;
int64_t in_channels = 3;
int64_t txt_input_size = 1024;
int64_t hidden_size = 768;
int64_t txt_hidden_size = 768;
int64_t cond_vec_size = 768;
int64_t depth_double = 17;
int64_t txt_preamble_depth = 2;
int64_t num_heads = 12;
int64_t head_dim = 64;
float mlp_ratio = 2.6667f;
int64_t pca_channels = 128;
int64_t prompt_length = 256;
int64_t n_T = 100;
float cfg_interval_start = 0.0f;
float cfg_interval_end = 1.0f;
static MiniT2IConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
MiniT2IConfig config;
config.depth_double = 0;
config.txt_preamble_depth = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "img_embedder.proj1.weight") && tensor_storage.n_dims == 4) {
config.patch_size = tensor_storage.ne[0];
config.in_channels = tensor_storage.ne[2];
config.pca_channels = tensor_storage.ne[3];
} else if (ends_with(name, "img_embedder.proj2.weight") && tensor_storage.n_dims == 4) {
config.pca_channels = tensor_storage.ne[2];
config.hidden_size = tensor_storage.ne[3];
} else if (ends_with(name, "txt_embedder.weight") && tensor_storage.n_dims == 2) {
config.txt_input_size = tensor_storage.ne[0];
config.txt_hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "pooled_embedder.weight") && tensor_storage.n_dims == 2) {
config.cond_vec_size = tensor_storage.ne[1];
} else if (ends_with(name, "double_blocks.0.img_qkv.weight") && tensor_storage.n_dims == 2) {
int64_t inner3 = tensor_storage.ne[1];
int64_t inner = inner3 / 3;
config.hidden_size = tensor_storage.ne[0];
if (config.hidden_size == 768) {
config.num_heads = 12;
config.head_dim = 64;
} else if (config.hidden_size == 1248) {
config.num_heads = 24;
config.head_dim = 52;
} else if (inner > 0) {
config.head_dim = 64;
config.num_heads = std::max<int64_t>(1, inner / config.head_dim);
}
} else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) {
int64_t patch_area = config.patch_size * config.patch_size;
config.hidden_size = tensor_storage.ne[0];
config.in_channels = patch_area > 0 ? tensor_storage.ne[1] / patch_area : config.in_channels;
} else if (ends_with(name, "mask_token") && tensor_storage.n_dims >= 2) {
config.prompt_length = tensor_storage.ne[1];
}
size_t pos = name.find("double_blocks.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int64_t idx = atoi(items[1].c_str());
config.depth_double = std::max<int64_t>(config.depth_double, idx + 1);
}
}
pos = name.find("txt_preamble_blocks.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int64_t idx = atoi(items[1].c_str());
config.txt_preamble_depth = std::max<int64_t>(config.txt_preamble_depth, idx + 1);
}
}
}
if (config.depth_double <= 0) {
config.depth_double = config.hidden_size == 1248 ? 23 : 17;
}
if (config.txt_preamble_depth <= 0) {
config.txt_preamble_depth = 2;
}
if (config.head_dim <= 0 || config.num_heads <= 0) {
config.head_dim = config.hidden_size == 1248 ? 52 : 64;
config.num_heads = config.hidden_size / config.head_dim;
}
LOG_DEBUG("minit2i: hidden_size=%" PRId64 ", txt_hidden_size=%" PRId64 ", heads=%" PRId64 ", head_dim=%" PRId64 ", double_blocks=%" PRId64 ", txt_blocks=%" PRId64 ", patch=%" PRId64 ", in_channels=%" PRId64,
config.hidden_size,
config.txt_hidden_size,
config.num_heads,
config.head_dim,
config.depth_double,
config.txt_preamble_depth,
config.patch_size,
config.in_channels);
return config;
}
};
inline std::vector<float> make_2d_sincos_pos_embed(int grid_size, int dim) {
GGML_ASSERT(dim % 4 == 0);
int half_dim = dim / 2;
int quarter = half_dim / 2;
std::vector<float> out(static_cast<size_t>(grid_size) * grid_size * dim);
std::vector<float> omega(quarter);
for (int i = 0; i < quarter; ++i) {
omega[i] = 1.0f / std::pow(10000.0f, static_cast<float>(i) / static_cast<float>(quarter));
}
for (int y = 0; y < grid_size; ++y) {
for (int x = 0; x < grid_size; ++x) {
size_t base = static_cast<size_t>(y * grid_size + x) * dim;
for (int i = 0; i < quarter; ++i) {
float ay = y * omega[i];
float ax = x * omega[i];
out[base + i] = std::sin(ax);
out[base + quarter + i] = std::cos(ax);
out[base + half_dim + i] = std::sin(ay);
out[base + half_dim + quarter + i] = std::cos(ay);
}
}
}
return out;
}
inline std::vector<float> make_text_rope(int length, int head_dim) {
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), head_dim, 10000.f));
}
inline std::vector<float> make_vision_rope(int side, int head_dim) {
GGML_ASSERT(head_dim % 4 == 0);
int dim = head_dim / 2;
int quarter = dim / 2;
int length = side * side;
std::vector<float> out(static_cast<size_t>(length) * (head_dim / 2) * 4);
std::vector<float> freqs(quarter);
for (int i = 0; i < quarter; ++i) {
freqs[i] = 1.0f / std::pow(10000.0f, static_cast<float>(2 * i) / static_cast<float>(dim));
}
for (int y = 0; y < side; ++y) {
for (int x = 0; x < side; ++x) {
int pos = y * side + x;
size_t base = static_cast<size_t>(pos) * (head_dim / 2) * 4;
for (int i = 0; i < quarter; ++i) {
float ay = y * freqs[i];
float ax = x * freqs[i];
float angles[2] = {ay, ax};
for (int axis = 0; axis < 2; ++axis) {
int j = axis * quarter + i;
out[base + 4 * j] = std::cos(angles[axis]);
out[base + 4 * j + 1] = -std::sin(angles[axis]);
out[base + 4 * j + 2] = std::sin(angles[axis]);
out[base + 4 * j + 3] = std::cos(angles[axis]);
}
}
}
}
return out;
}
struct SwiGLUMlp : public GGMLBlock {
SwiGLUMlp(int64_t in_features, int64_t hidden_features) {
int64_t hidden_dim = ((hidden_features + 7) / 8) * 8;
blocks["w1"] = std::make_shared<Linear>(in_features, hidden_dim, false);
blocks["w3"] = std::make_shared<Linear>(in_features, hidden_dim, false);
blocks["w2"] = std::make_shared<Linear>(hidden_dim, in_features, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
auto gate = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
auto up = w3->forward(ctx, x);
return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, up));
}
};
struct BottleneckPatchEmbed : public GGMLBlock {
int64_t patch_size;
BottleneckPatchEmbed(int64_t patch_size, int64_t in_channels, int64_t pca_channels, int64_t hidden_size)
: patch_size(patch_size) {
blocks["proj1"] = std::make_shared<Conv2d>(in_channels,
pca_channels,
std::pair<int, int>{static_cast<int>(patch_size), static_cast<int>(patch_size)},
std::pair<int, int>{static_cast<int>(patch_size), static_cast<int>(patch_size)},
std::pair<int, int>{0, 0},
std::pair<int, int>{1, 1},
false);
blocks["proj2"] = std::make_shared<Conv2d>(pca_channels,
hidden_size,
std::pair<int, int>{1, 1},
std::pair<int, int>{1, 1},
std::pair<int, int>{0, 0},
std::pair<int, int>{1, 1},
true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto proj1 = std::dynamic_pointer_cast<Conv2d>(blocks["proj1"]);
auto proj2 = std::dynamic_pointer_cast<Conv2d>(blocks["proj2"]);
x = proj1->forward(ctx, x);
x = proj2->forward(ctx, x);
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
return x;
}
};
struct TimestepEmbedder : public GGMLBlock {
int frequency_embedding_size;
TimestepEmbedder(int64_t hidden_size, int frequency_embedding_size = 256)
: frequency_embedding_size(frequency_embedding_size) {
blocks["mlp.0"] = std::make_shared<Linear>(frequency_embedding_size, hidden_size, true, true);
blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, hidden_size, true, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1.0f);
t_emb = mlp_0->forward(ctx, t_emb);
t_emb = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
return mlp_2->forward(ctx, t_emb);
}
};
inline std::vector<ggml_tensor*> split_qkv(ggml_context* ctx, ggml_tensor* qkv, int64_t num_heads, int64_t head_dim) {
int64_t N = qkv->ne[2];
int64_t L = qkv->ne[1];
auto q = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], 0);
auto k = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], qkv->nb[0] * head_dim * num_heads);
auto v = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N,
qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], qkv->nb[0] * head_dim * num_heads * 2);
return {q, k, v};
}
struct PlainTextTransformerBlock : public GGMLBlock {
int64_t num_heads;
int64_t head_dim;
PlainTextTransformerBlock(int64_t hidden_size, int64_t num_heads, int64_t head_dim, float mlp_ratio)
: num_heads(num_heads), head_dim(head_dim) {
int64_t inner_dim = num_heads * head_dim;
blocks["norm1"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["norm2"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["qkv"] = std::make_shared<Linear>(hidden_size, inner_dim * 3, true);
blocks["attn_proj"] = std::make_shared<Linear>(inner_dim, hidden_size, true);
blocks["mlp"] = std::make_shared<SwiGLUMlp>(hidden_size, static_cast<int64_t>(hidden_size * mlp_ratio));
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* pe) {
auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto attn_proj = std::dynamic_pointer_cast<Linear>(blocks["attn_proj"]);
auto mlp = std::dynamic_pointer_cast<SwiGLUMlp>(blocks["mlp"]);
auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
auto qkv = split_qkv(ctx->ggml_ctx, qkv_proj->forward(ctx, norm1->forward(ctx, txt)), num_heads, head_dim);
auto q = q_norm->forward(ctx, qkv[0]);
auto k = k_norm->forward(ctx, qkv[1]);
auto v = qkv[2];
auto out = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false);
txt = ggml_add(ctx->ggml_ctx, txt, attn_proj->forward(ctx, out));
txt = ggml_add(ctx->ggml_ctx, txt, mlp->forward(ctx, norm2->forward(ctx, txt)));
return txt;
}
};
struct DoubleStreamDiTBlock : public GGMLBlock {
int64_t num_heads;
int64_t head_dim;
DoubleStreamDiTBlock(int64_t hidden_size, int64_t txt_hidden_size, int64_t num_heads, int64_t head_dim, float mlp_ratio)
: num_heads(num_heads), head_dim(head_dim) {
int64_t inner_dim = num_heads * head_dim;
blocks["img_norm1"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["img_norm2"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["txt_norm1"] = std::make_shared<RMSNorm>(txt_hidden_size, 1e-6f);
blocks["txt_norm2"] = std::make_shared<RMSNorm>(txt_hidden_size, 1e-6f);
blocks["img_qkv"] = std::make_shared<Linear>(hidden_size, inner_dim * 3, true);
blocks["txt_qkv"] = std::make_shared<Linear>(txt_hidden_size, inner_dim * 3, true);
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["img_attn_proj"] = std::make_shared<Linear>(inner_dim, hidden_size, true);
blocks["txt_attn_proj"] = std::make_shared<Linear>(inner_dim, txt_hidden_size, true);
blocks["img_mlp"] = std::make_shared<SwiGLUMlp>(hidden_size, static_cast<int64_t>(hidden_size * mlp_ratio));
blocks["txt_mlp"] = std::make_shared<SwiGLUMlp>(txt_hidden_size, static_cast<int64_t>(txt_hidden_size * mlp_ratio));
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* img,
ggml_tensor* txt,
ggml_tensor* pe) {
auto img_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm1"]);
auto img_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm2"]);
auto txt_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm1"]);
auto txt_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm2"]);
auto img_qkv_p = std::dynamic_pointer_cast<Linear>(blocks["img_qkv"]);
auto txt_qkv_p = std::dynamic_pointer_cast<Linear>(blocks["txt_qkv"]);
auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
auto img_proj = std::dynamic_pointer_cast<Linear>(blocks["img_attn_proj"]);
auto txt_proj = std::dynamic_pointer_cast<Linear>(blocks["txt_attn_proj"]);
auto img_mlp = std::dynamic_pointer_cast<SwiGLUMlp>(blocks["img_mlp"]);
auto txt_mlp = std::dynamic_pointer_cast<SwiGLUMlp>(blocks["txt_mlp"]);
int64_t li = img->ne[1];
int64_t lt = txt->ne[1];
auto img_qkv = split_qkv(ctx->ggml_ctx, img_qkv_p->forward(ctx, img_norm1->forward(ctx, img)), num_heads, head_dim);
auto txt_qkv = split_qkv(ctx->ggml_ctx, txt_qkv_p->forward(ctx, txt_norm1->forward(ctx, txt)), num_heads, head_dim);
auto q = ggml_concat(ctx->ggml_ctx, q_norm->forward(ctx, txt_qkv[0]), q_norm->forward(ctx, img_qkv[0]), 2);
auto k = ggml_concat(ctx->ggml_ctx, k_norm->forward(ctx, txt_qkv[1]), k_norm->forward(ctx, img_qkv[1]), 2);
auto v = ggml_concat(ctx->ggml_ctx, txt_qkv[2], img_qkv[2], 2);
auto out = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false);
auto out_txt = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, lt);
auto out_img = ggml_ext_slice(ctx->ggml_ctx, out, 1, lt, lt + li);
img = ggml_add(ctx->ggml_ctx, img, img_proj->forward(ctx, out_img));
txt = ggml_add(ctx->ggml_ctx, txt, txt_proj->forward(ctx, out_txt));
img = ggml_add(ctx->ggml_ctx, img, img_mlp->forward(ctx, img_norm2->forward(ctx, img)));
txt = ggml_add(ctx->ggml_ctx, txt, txt_mlp->forward(ctx, txt_norm2->forward(ctx, txt)));
return {img, txt};
}
};
struct FinalLayer : public GGMLBlock {
FinalLayer(int64_t hidden_size, int64_t patch_size, int64_t out_channels) {
blocks["norm_final"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["linear"] = std::make_shared<Linear>(hidden_size, patch_size * patch_size * out_channels, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm_final = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_final"]);
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
return linear->forward(ctx, norm_final->forward(ctx, x));
}
};
struct MMJiT : public GGMLBlock {
MiniT2IConfig config;
MMJiT(const MiniT2IConfig& config)
: config(config) {
blocks["img_embedder"] = std::make_shared<BottleneckPatchEmbed>(config.patch_size, config.in_channels, config.pca_channels, config.hidden_size);
blocks["txt_embedder"] = std::make_shared<Linear>(config.txt_input_size, config.txt_hidden_size, false);
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>(config.cond_vec_size);
blocks["pooled_embedder"] = std::make_shared<Linear>(config.txt_input_size, config.cond_vec_size, false);
for (int64_t i = 0; i < config.txt_preamble_depth; ++i) {
blocks["txt_preamble_blocks." + std::to_string(i)] = std::make_shared<PlainTextTransformerBlock>(config.txt_hidden_size, config.num_heads, config.head_dim, config.mlp_ratio);
}
for (int64_t i = 0; i < config.depth_double; ++i) {
blocks["double_blocks." + std::to_string(i)] = std::make_shared<DoubleStreamDiTBlock>(config.hidden_size, config.txt_hidden_size, config.num_heads, config.head_dim, config.mlp_ratio);
}
blocks["final_layer"] = std::make_shared<FinalLayer>(config.hidden_size, config.patch_size, config.in_channels);
}
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
enum ggml_type wtype = get_type(prefix + "mask_token", tensor_storage_map, GGML_TYPE_F32);
params["mask_token"] = ggml_new_tensor_3d(ctx, wtype, config.txt_input_size, 1, 1);
}
ggml_tensor* apply_text_mask(GGMLRunnerContext* ctx, ggml_tensor* context, ggml_tensor* mask) {
if (mask == nullptr) {
return context;
}
mask = ggml_reshape_3d(ctx->ggml_ctx, mask, 1, mask->ne[0], mask->ne[1]);
mask = ggml_repeat(ctx->ggml_ctx, mask, context);
auto keep = ggml_mul(ctx->ggml_ctx, context, mask);
auto inv = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, mask), mask);
auto mask_token = ggml_repeat(ctx->ggml_ctx, params["mask_token"], context);
return ggml_add(ctx->ggml_ctx, keep, ggml_mul(ctx->ggml_ctx, mask_token, inv));
}
ggml_tensor* pool_context(GGMLRunnerContext* ctx, ggml_tensor* context) {
int64_t dim = context->ne[0];
int64_t len = context->ne[1];
int64_t N = context->ne[2];
auto x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, context, 1, 0, 2, 3));
x = ggml_reshape_3d(ctx->ggml_ctx, x, len, dim, N);
x = ggml_mean(ctx->ggml_ctx, x);
x = ggml_reshape_2d(ctx->ggml_ctx, x, dim, N);
return x;
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* img,
ggml_tensor* context,
ggml_tensor* mask,
ggml_tensor* pos_embed,
ggml_tensor* txt_pe,
ggml_tensor* joint_pe) {
auto img_embedder = std::dynamic_pointer_cast<BottleneckPatchEmbed>(blocks["img_embedder"]);
auto txt_embedder = std::dynamic_pointer_cast<Linear>(blocks["txt_embedder"]);
auto final_layer = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
int64_t W = img->ne[0];
int64_t H = img->ne[1];
int64_t hp = H / config.patch_size;
int64_t wp = W / config.patch_size;
context = apply_text_mask(ctx, context, mask);
auto x = img_embedder->forward(ctx, img);
x = ggml_add(ctx->ggml_ctx, x, pos_embed);
auto txt = txt_embedder->forward(ctx, context);
for (int64_t i = 0; i < config.txt_preamble_depth; ++i) {
auto block = std::dynamic_pointer_cast<PlainTextTransformerBlock>(blocks["txt_preamble_blocks." + std::to_string(i)]);
txt = block->forward(ctx, txt, txt_pe);
sd::ggml_graph_cut::mark_graph_cut(txt, "minit2i.txt_preamble_blocks." + std::to_string(i), "txt");
}
for (int64_t i = 0; i < config.depth_double; ++i) {
auto block = std::dynamic_pointer_cast<DoubleStreamDiTBlock>(blocks["double_blocks." + std::to_string(i)]);
auto out = block->forward(ctx, x, txt, joint_pe);
x = out.first;
txt = out.second;
sd::ggml_graph_cut::mark_graph_cut(x, "minit2i.double_blocks." + std::to_string(i), "x");
sd::ggml_graph_cut::mark_graph_cut(txt, "minit2i.double_blocks." + std::to_string(i), "txt");
}
auto combined = ggml_concat(ctx->ggml_ctx, txt, x, 1);
auto out = final_layer->forward(ctx, combined);
auto img_out = ggml_ext_slice(ctx->ggml_ctx, out, 1, txt->ne[1], txt->ne[1] + x->ne[1]);
return DiT::unpatchify(ctx->ggml_ctx, img_out, hp, wp, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), false);
}
};
struct MiniT2IRunner : public DiffusionModelRunner {
MiniT2IConfig config;
MMJiT model;
ggml_context* position_cache_ctx = nullptr;
ggml_backend_buffer_t position_cache_buffer = nullptr;
ggml_tensor* cached_pos_embed = nullptr;
ggml_tensor* cached_txt_pe = nullptr;
ggml_tensor* cached_joint_pe = nullptr;
int64_t cached_img_side = -1;
int64_t cached_txt_len = -1;
int64_t cached_hidden_size = -1;
int64_t cached_head_dim = -1;
MiniT2IRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(MiniT2IConfig::detect_from_weights(tensor_storage_map, this->prefix)),
model(config) {
model.init(params_ctx, tensor_storage_map, this->prefix);
}
~MiniT2IRunner() override {
free_position_cache();
}
std::string get_desc() override {
return "MiniT2I";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
model.get_param_tensors(tensors, prefix);
}
void free_position_cache() {
if (position_cache_buffer != nullptr) {
ggml_backend_buffer_free(position_cache_buffer);
position_cache_buffer = nullptr;
}
if (position_cache_ctx != nullptr) {
ggml_free(position_cache_ctx);
position_cache_ctx = nullptr;
}
cached_pos_embed = nullptr;
cached_txt_pe = nullptr;
cached_joint_pe = nullptr;
cached_img_side = -1;
cached_txt_len = -1;
cached_hidden_size = -1;
cached_head_dim = -1;
}
void ensure_position_cache(int64_t img_side, int64_t txt_len) {
if (cached_img_side == img_side &&
cached_txt_len == txt_len &&
cached_hidden_size == config.hidden_size &&
cached_head_dim == config.head_dim &&
cached_pos_embed != nullptr &&
cached_txt_pe != nullptr &&
cached_joint_pe != nullptr) {
return;
}
free_position_cache();
auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast<int>(img_side), static_cast<int>(config.hidden_size));
auto txt_pe_vec = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim));
auto img_pe_vec = make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim));
auto joint_pe_vec = txt_pe_vec;
joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end());
ggml_init_params params;
params.mem_size = static_cast<size_t>(3 * ggml_tensor_overhead());
params.mem_buffer = nullptr;
params.no_alloc = true;
position_cache_ctx = ggml_init(params);
GGML_ASSERT(position_cache_ctx != nullptr);
cached_pos_embed = ggml_new_tensor_3d(position_cache_ctx, GGML_TYPE_F32, config.hidden_size, img_side * img_side, 1);
ggml_set_name(cached_pos_embed, "minit2i.pos_embed");
cached_txt_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len);
ggml_set_name(cached_txt_pe, "minit2i.txt_pe");
cached_joint_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len + img_side * img_side);
ggml_set_name(cached_joint_pe, "minit2i.joint_pe");
position_cache_buffer = ggml_backend_alloc_ctx_tensors(position_cache_ctx, runtime_backend);
GGML_ASSERT(position_cache_buffer != nullptr);
ggml_backend_buffer_set_usage(position_cache_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
ggml_backend_tensor_set(cached_pos_embed, pos_embed_vec.data(), 0, ggml_nbytes(cached_pos_embed));
ggml_backend_tensor_set(cached_txt_pe, txt_pe_vec.data(), 0, ggml_nbytes(cached_txt_pe));
ggml_backend_tensor_set(cached_joint_pe, joint_pe_vec.data(), 0, ggml_nbytes(cached_joint_pe));
ggml_backend_synchronize(runtime_backend);
cached_img_side = img_side;
cached_txt_len = txt_len;
cached_hidden_size = config.hidden_size;
cached_head_dim = config.head_dim;
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const sd::Tensor<float>& mask_tensor) {
ggml_cgraph* gf = new_graph_custom(MINIT2I_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* context = make_input(context_tensor);
ggml_tensor* mask = make_input(mask_tensor);
SD_UNUSED(timesteps_tensor);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t img_side = H / config.patch_size;
int64_t txt_len = context->ne[1];
ensure_position_cache(img_side, txt_len);
auto runner_ctx = get_context();
auto out = model.forward(&runner_ctx, x, context, mask, cached_pos_embed, cached_txt_pe, cached_joint_pe);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
const sd::Tensor<float>& mask) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, mask);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
GGML_ASSERT(diffusion_params.context != nullptr);
const auto* extra = diffusion_extra_as<MiniT2IDiffusionExtra>(diffusion_params);
GGML_ASSERT(extra->mask != nullptr);
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
*diffusion_params.context,
*extra->mask);
}
};
} // namespace MiniT2I
#endif // __SD_MODEL_DIFFUSION_MINIT2I_HPP__

View File

@ -52,6 +52,10 @@ struct LTXAVDiffusionExtra {
const sd::Tensor<float>* video_positions = nullptr; const sd::Tensor<float>* video_positions = nullptr;
}; };
struct MiniT2IDiffusionExtra {
const sd::Tensor<float>* mask = nullptr;
};
using DiffusionExtraParams = std::variant<std::monostate, using DiffusionExtraParams = std::variant<std::monostate,
UNetDiffusionExtra, UNetDiffusionExtra,
SkipLayerDiffusionExtra, SkipLayerDiffusionExtra,
@ -59,7 +63,8 @@ using DiffusionExtraParams = std::variant<std::monostate,
AnimaDiffusionExtra, AnimaDiffusionExtra,
WanDiffusionExtra, WanDiffusionExtra,
HiDreamO1DiffusionExtra, HiDreamO1DiffusionExtra,
LTXAVDiffusionExtra>; LTXAVDiffusionExtra,
MiniT2IDiffusionExtra>;
struct DiffusionParams { struct DiffusionParams {
const sd::Tensor<float>* x = nullptr; const sd::Tensor<float>* x = nullptr;

View File

@ -26,13 +26,66 @@ struct T5Config {
static T5Config detect_from_weights(const String2TensorStorage& tensor_storage_map, static T5Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix, const std::string& prefix,
bool is_umt5 = false) { bool is_umt5 = false) {
(void)tensor_storage_map;
(void)prefix;
T5Config config; T5Config config;
if (is_umt5) { if (is_umt5) {
config.vocab_size = 256384; config.vocab_size = 256384;
config.relative_attention = false; config.relative_attention = false;
} }
auto find_tensor = [&](const std::string& suffix) -> const TensorStorage* {
auto it = tensor_storage_map.find(prefix + "." + suffix);
if (it != tensor_storage_map.end()) {
return &it->second;
}
it = tensor_storage_map.find(prefix + suffix);
if (it != tensor_storage_map.end()) {
return &it->second;
}
return nullptr;
};
if (const TensorStorage* shared = find_tensor("shared.weight")) {
if (shared->n_dims == 2) {
config.vocab_size = shared->ne[1];
config.model_dim = shared->ne[0];
}
}
if (const TensorStorage* q = find_tensor("encoder.block.0.layer.0.SelfAttention.q.weight")) {
if (q->n_dims == 2) {
config.model_dim = q->ne[0];
int64_t inner_dim = q->ne[1];
// Flan-T5/T5 uses d_kv=64 for common sizes.
if (inner_dim % 64 == 0) {
config.num_heads = inner_dim / 64;
}
}
}
if (const TensorStorage* wi = find_tensor("encoder.block.0.layer.1.DenseReluDense.wi_0.weight")) {
if (wi->n_dims == 2) {
config.model_dim = wi->ne[0];
config.ff_dim = wi->ne[1];
}
}
int64_t detected_layers = 0;
for (const auto& [name, _] : tensor_storage_map) {
std::string base = prefix;
if (!base.empty() && base.back() != '.') {
base += ".";
}
std::string layer_prefix = base + "encoder.block.";
if (!starts_with(name, layer_prefix)) {
continue;
}
size_t pos = layer_prefix.size();
size_t dot = name.find('.', pos);
if (dot == std::string::npos) {
continue;
}
int64_t layer = atoi(name.substr(pos, dot - pos).c_str());
detected_layers = std::max(detected_layers, layer + 1);
}
if (detected_layers > 0) {
config.num_layers = detected_layers;
}
return config; return config;
} }
}; };

View File

@ -78,7 +78,7 @@ public:
scale_factor = 16; scale_factor = 16;
} else if (sd_version_uses_flux2_vae(version)) { } else if (sd_version_uses_flux2_vae(version)) {
scale_factor = 16; scale_factor = 16;
} else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) { } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) {
scale_factor = 1; scale_factor = 1;
} }
return scale_factor; return scale_factor;

View File

@ -20,6 +20,7 @@
#include "model_io/torch_legacy_io.h" #include "model_io/torch_legacy_io.h"
#include "model_io/torch_zip_io.h" #include "model_io/torch_zip_io.h"
#include "model_loader.h" #include "model_loader.h"
#include "runtime/imatrix.h"
#include "stable-diffusion.h" #include "stable-diffusion.h"
#include "core/ggml_extend_backend.h" #include "core/ggml_extend_backend.h"
@ -156,7 +157,8 @@ void convert_tensor(void* src,
void* dst, void* dst,
ggml_type dst_type, ggml_type dst_type,
int nrows, int nrows,
int n_per_row) { int n_per_row,
std::vector<float> imatrix = {}) {
int n = nrows * n_per_row; int n = nrows * n_per_row;
if (src_type == dst_type) { if (src_type == dst_type) {
size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type); size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type);
@ -165,7 +167,7 @@ void convert_tensor(void* src,
if (dst_type == GGML_TYPE_F16) { if (dst_type == GGML_TYPE_F16) {
ggml_fp32_to_fp16_row((float*)src, (ggml_fp16_t*)dst, n); ggml_fp32_to_fp16_row((float*)src, (ggml_fp16_t*)dst, n);
} else { } else {
std::vector<float> imatrix(n_per_row, 1.0f); // dummy importance matrix imatrix.resize(n_per_row, 1.0f);
const float* im = imatrix.data(); const float* im = imatrix.data();
ggml_quantize_chunk(dst_type, (float*)src, dst, 0, nrows, n_per_row, im); ggml_quantize_chunk(dst_type, (float*)src, dst, 0, nrows, n_per_row, im);
} }
@ -195,7 +197,7 @@ void convert_tensor(void* src,
if (dst_type == GGML_TYPE_F16) { if (dst_type == GGML_TYPE_F16) {
ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n); ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n);
} else { } else {
std::vector<float> imatrix(n_per_row, 1.0f); // dummy importance matrix imatrix.resize(n_per_row, 1.0f);
const float* im = imatrix.data(); const float* im = imatrix.data();
ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, nrows, n_per_row, im); ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, nrows, n_per_row, im);
} }
@ -470,6 +472,9 @@ SDVersion ModelLoader::get_sd_version() {
tensor_storage_map.find("model.diffusion_model.transformer_blocks.0.img_mlp.w1.weight") != tensor_storage_map.end()) { tensor_storage_map.find("model.diffusion_model.transformer_blocks.0.img_mlp.w1.weight") != tensor_storage_map.end()) {
return VERSION_LENS; return VERSION_LENS;
} }
if (tensor_storage.name.find("net.img_embedder.proj1.weight") != std::string::npos) {
return VERSION_MINIT2I;
}
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) { if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
return VERSION_QWEN_IMAGE; return VERSION_QWEN_IMAGE;
} }
@ -970,6 +975,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
size_t total_tensors_processed = 0; size_t total_tensors_processed = 0;
const int64_t t_start = start_time; const int64_t t_start = start_time;
int last_n_threads = 1; int last_n_threads = 1;
SDVersion imatrix_version = (version_ == VERSION_COUNT) ? get_sd_version() : version_;
for (size_t file_index = 0; file_index < file_data.size(); ++file_index) { for (size_t file_index = 0; file_index < file_data.size(); ++file_index) {
auto& fdata = file_data[file_index]; auto& fdata = file_data[file_index];
@ -1154,12 +1160,15 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
failed = true; failed = true;
return; return;
} }
std::string processed_name = convert_tensor_name(tensor_storage.name, imatrix_version);
std::vector<float> imatrix = get_imatrix_collector().get_values(processed_name);
convert_tensor((void*)target_buf, convert_tensor((void*)target_buf,
tensor_storage.type, tensor_storage.type,
convert_buf, convert_buf,
dst_tensor->type, dst_tensor->type,
(int)tensor_storage.nelements() / (int)tensor_storage.ne[0], (int)tensor_storage.nelements() / (int)tensor_storage.ne[0],
(int)tensor_storage.ne[0]); (int)tensor_storage.ne[0],
std::move(imatrix));
} else { } else {
convert_buf = read_buf; convert_buf = read_buf;
} }

View File

@ -1338,6 +1338,68 @@ struct SefiFlowDenoiser : public FluxFlowDenoiser {
} }
}; };
// MiniT2I predicts x0 directly and integrates a linear flow ODE:
// x_{t+dt} = x_t + (x0 - x_t)/(1 - t) * dt, t in [0, 1), x0 = start = noise * 2.
// Mapping sigma = 1 - t makes the generic Euler update
// x += (x - denoised)/sigma * (sigma_next - sigma)
// exactly reproduce that step when denoised == x0. To make the generic
// `denoised = pred * c_out + x * c_skip` yield x0 from the model's raw x0
// prediction we use c_skip = 0, c_out = 1, c_in = 1. Sigmas run linearly 1 -> 0.
struct MiniT2IFlowDenoiser : public Denoiser {
float sigma_min() override {
return 0.0f;
}
float sigma_max() override {
return 1.0f;
}
float sigma_to_t(float sigma) override {
return 1.0f - sigma;
}
float t_to_sigma(float t) override {
return 1.0f - t;
}
std::vector<float> get_scalings(float sigma) override {
SD_UNUSED(sigma);
float c_skip = 0.0f;
float c_out = 1.0f;
float c_in = 1.0f;
return {c_skip, c_out, c_in};
}
sd::Tensor<float> noise_scaling(float sigma,
const sd::Tensor<float>& noise,
const sd::Tensor<float>& latent) override {
SD_UNUSED(sigma);
SD_UNUSED(latent);
// Sampling starts from x0_init = noise * 2 (see MiniT2I reference).
return noise * 2.0f;
}
sd::Tensor<float> inverse_noise_scaling(float sigma, const sd::Tensor<float>& latent) override {
SD_UNUSED(sigma);
return latent;
}
std::vector<float> 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);
// Uniform t schedule 0 -> 1 => sigma 1 -> 0, matching the reference loop.
std::vector<float> sigmas;
sigmas.reserve(n + 1);
for (uint32_t i = 0; i < n; ++i) {
sigmas.push_back(1.0f - static_cast<float>(i) / static_cast<float>(n));
}
sigmas.push_back(0.0f);
return sigmas;
}
};
typedef std::function<sd::guidance::GuiderOutput(const sd::Tensor<float>&, float, int)> denoise_cb_t; typedef std::function<sd::guidance::GuiderOutput(const sd::Tensor<float>&, float, int)> denoise_cb_t;
static std::pair<float, float> get_ancestral_step(float sigma_from, static std::pair<float, float> get_ancestral_step(float sigma_from,

308
src/runtime/imatrix.cpp Normal file
View File

@ -0,0 +1,308 @@
#include "runtime/imatrix.h"
/* Adapted from llama.cpp (credits: Kawrakow). */
#include "core/util.h"
#include "ggml-backend.h"
#include "ggml.h"
#include "stable-diffusion.h"
#include <cmath>
#include <cstdlib>
#include <cstring>
static IMatrixCollector imatrix_collector;
IMatrixCollector& get_imatrix_collector() {
return imatrix_collector;
}
// remove any prefix and suffixes from the name
// CUDA0#blk.0.attn_k.weight#0 => blk.0.attn_k.weight
static std::string filter_tensor_name(const char* name) {
std::string wname;
const char* p = strchr(name, '#');
if (p != NULL) {
p = p + 1;
const char* q = strchr(p, '#');
if (q != NULL) {
wname = std::string(p, q - p);
} else {
wname = p;
}
} else {
wname = name;
}
return wname;
}
bool IMatrixCollector::collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) {
GGML_UNUSED(user_data);
if (t == nullptr) {
return false;
}
if (t->op != GGML_OP_MUL_MAT && t->op != GGML_OP_MUL_MAT_ID) {
return false;
}
const struct ggml_tensor* src0 = t->src[0];
const struct ggml_tensor* src1 = t->src[1];
if (src0 == nullptr || src1 == nullptr) {
return false;
}
std::string wname = filter_tensor_name(src0->name);
// when ask is true, the scheduler wants to know if we are interested in data from this tensor
// if we return true, a follow-up call will be made with ask=false in which we can do the actual collection
if (ask) {
if (t->op == GGML_OP_MUL_MAT_ID) {
return true; // collect all indirect matrix multiplications
}
// why are small batches ignored (<16 tokens)?
// if (src1->ne[1] < 16 || src1->type != GGML_TYPE_F32) return false;
if (!(wname.substr(0, 6) == "model." || wname.substr(0, 17) == "cond_stage_model." || wname.substr(0, 14) == "text_encoders.")) {
return false;
}
return true;
}
std::lock_guard<std::mutex> lock(mutex_);
// copy the data from the GPU memory if needed
const bool is_host = src1->buffer == NULL || ggml_backend_buffer_is_host(src1->buffer);
if (!is_host) {
src1_data_.resize(ggml_nelements(src1));
ggml_backend_tensor_get(src1, src1_data_.data(), 0, ggml_nbytes(src1));
}
const float* data = is_host ? (const float*)src1->data : src1_data_.data();
// this has been adapted to the new format of storing merged experts in a single 3d tensor
// ref: https://github.com/ggml-org/llama.cpp/pull/6387
if (t->op == GGML_OP_MUL_MAT_ID) {
// ids -> [n_experts_used, n_tokens]
// src1 -> [cols, n_expert_used, n_tokens]
const ggml_tensor* ids = t->src[2];
const int n_as = static_cast<int>(src0->ne[2]);
const int n_ids = static_cast<int>(ids->ne[0]);
// the top-k selected expert ids are stored in the ids tensor
// for simplicity, always copy ids to host, because it is small
// take into account that ids is not contiguous!
GGML_ASSERT(ids->ne[1] == src1->ne[2]);
ids_.resize(ggml_nbytes(ids));
ggml_backend_tensor_get(ids, ids_.data(), 0, ggml_nbytes(ids));
auto& e = stats_[wname];
++e.ncall;
if (e.values.empty()) {
e.values.resize(src1->ne[0] * n_as, 0);
e.counts.resize(src1->ne[0] * n_as, 0);
} else if (e.values.size() != (size_t)src1->ne[0] * n_as) {
LOG_ERROR("inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0] * n_as);
exit(1); // GGML_ABORT("fatal error");
}
// loop over all possible experts, regardless if they are used or not in the batch
for (int ex = 0; ex < n_as; ++ex) {
size_t e_start = ex * src1->ne[0];
for (int idx = 0; idx < n_ids; ++idx) {
for (int row = 0; row < (int)src1->ne[2]; ++row) {
const int excur = *(const int32_t*)(ids_.data() + row * ids->nb[1] + idx * ids->nb[0]);
GGML_ASSERT(excur >= 0 && excur < n_as); // sanity check
if (excur != ex)
continue;
const int64_t i11 = idx % src1->ne[1];
const int64_t i12 = row;
const float* x = (const float*)((const char*)data + i11 * src1->nb[1] + i12 * src1->nb[2]);
for (int j = 0; j < (int)src1->ne[0]; ++j) {
e.values[e_start + j] += x[j] * x[j];
e.counts[e_start + j]++;
if (!std::isfinite(e.values[e_start + j])) {
LOG_ERROR("%f detected in %s\n", e.values[e_start + j], wname.c_str());
exit(1);
}
}
}
}
}
} else {
auto& e = stats_[wname];
if (e.values.empty()) {
e.values.resize(src1->ne[0], 0);
e.counts.resize(src1->ne[0], 0);
} else if (e.values.size() != (size_t)src1->ne[0]) {
LOG_WARN("inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0]);
exit(1); // GGML_ABORT("fatal error");
}
++e.ncall;
for (int row = 0; row < (int)src1->ne[1]; ++row) {
const float* x = data + row * src1->ne[0];
for (int j = 0; j < (int)src1->ne[0]; ++j) {
if (std::isfinite(x[j])) {
e.values[j] += x[j] * x[j];
e.counts[j]++;
if (!std::isfinite(e.values[j])) {
LOG_WARN("%f detected in %s\n", e.values[j], wname.c_str());
exit(1);
}
} else {
// Likely something from an attention mask?
}
}
}
}
return true;
}
bool load_imatrix(const char* imatrix_path) {
return imatrix_collector.load_imatrix(imatrix_path);
}
void save_imatrix(const char* imatrix_path) {
imatrix_collector.save_imatrix(imatrix_path);
}
static bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) {
return imatrix_collector.collect_imatrix(t, ask, user_data);
}
void enable_imatrix_collection() {
sd_set_backend_eval_callback(collect_imatrix, nullptr);
}
void disable_imatrix_collection() {
sd_set_backend_eval_callback(nullptr, nullptr);
}
void IMatrixCollector::save_imatrix(std::string fname, int ncall) const {
if (ncall > 0) {
fname += ".at_";
fname += std::to_string(ncall);
}
// avoid writing imatrix entries that do not have full data
// this can happen with MoE models where some of the experts end up not being exercised by the provided training data
int n_entries = 0;
std::vector<std::string> to_store;
for (const auto& kv : stats_) {
const int n_all = static_cast<int>(kv.second.counts.size());
if (n_all == 0) {
continue;
}
int n_zeros = 0;
for (const int c : kv.second.counts) {
if (c == 0) {
n_zeros++;
}
}
if (n_zeros == n_all) {
LOG_WARN("entry '%40s' has no data - skipping\n", kv.first.c_str());
continue;
}
if (n_zeros > 0) {
LOG_WARN("entry '%40s' has partial data (%.2f%%) - skipping\n", kv.first.c_str(), 100.0f * (n_all - n_zeros) / n_all);
continue;
}
n_entries++;
to_store.push_back(kv.first);
}
if (to_store.size() < stats_.size()) {
LOG_WARN("storing only %zu out of %zu entries\n", to_store.size(), stats_.size());
}
std::ofstream out(fname, std::ios::binary);
out.write((const char*)&n_entries, sizeof(n_entries));
for (const auto& name : to_store) {
const auto& stat = stats_.at(name);
int len = static_cast<int>(name.size());
out.write((const char*)&len, sizeof(len));
out.write(name.c_str(), len);
out.write((const char*)&stat.ncall, sizeof(stat.ncall));
int nval = static_cast<int>(stat.values.size());
out.write((const char*)&nval, sizeof(nval));
if (nval > 0) {
std::vector<float> tmp(nval);
for (int i = 0; i < nval; i++) {
tmp[i] = (stat.values[i] / static_cast<float>(stat.counts[i])) * static_cast<float>(stat.ncall);
}
out.write((const char*)tmp.data(), nval * sizeof(float));
}
}
// Write the number of call the matrix was computed with
out.write((const char*)&last_call_, sizeof(last_call_));
}
bool IMatrixCollector::load_imatrix(const char* fname) {
std::ifstream in(fname, std::ios::binary);
if (!in) {
LOG_ERROR("failed to open %s\n", fname);
return false;
}
int n_entries;
in.read((char*)&n_entries, sizeof(n_entries));
if (in.fail() || n_entries < 1) {
LOG_ERROR("no data in file %s\n", fname);
return false;
}
for (int i = 0; i < n_entries; ++i) {
int len;
in.read((char*)&len, sizeof(len));
std::vector<char> name_as_vec(len + 1);
in.read((char*)name_as_vec.data(), len);
if (in.fail()) {
LOG_ERROR("failed reading name for entry %d from %s\n", i + 1, fname);
return false;
}
name_as_vec[len] = 0;
std::string name{name_as_vec.data()};
auto& e = stats_[std::move(name)];
int ncall;
in.read((char*)&ncall, sizeof(ncall));
int nval;
in.read((char*)&nval, sizeof(nval));
if (in.fail() || nval < 1) {
LOG_ERROR("failed reading number of values for entry %d\n", i);
stats_ = {};
return false;
}
if (e.values.empty()) {
e.values.resize(nval, 0);
e.counts.resize(nval, 0);
}
std::vector<float> tmp(nval);
in.read((char*)tmp.data(), nval * sizeof(float));
if (in.fail()) {
LOG_ERROR("failed reading data for entry %d\n", i);
stats_ = {};
return false;
}
// Recreate the state as expected by save_imatrix(), and correct for weighted sum.
for (int i = 0; i < nval; i++) {
e.values[i] += tmp[i];
e.counts[i] += ncall;
}
e.ncall += ncall;
}
return true;
}

45
src/runtime/imatrix.h Normal file
View File

@ -0,0 +1,45 @@
#ifndef __SD_RUNTIME_IMATRIX_H__
#define __SD_RUNTIME_IMATRIX_H__
#include <fstream>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>
/* Adapted from llama.cpp (credits: Kawrakow). */
struct ggml_tensor;
struct IMatrixStats {
std::vector<float> values{};
std::vector<int> counts{};
int ncall = 0;
};
class IMatrixCollector {
private:
std::unordered_map<std::string, IMatrixStats> stats_ = {};
std::mutex mutex_;
int last_call_ = 0;
std::vector<float> src1_data_;
std::vector<char> ids_; // the expert ids from ggml_mul_mat_id
public:
IMatrixCollector() = default;
bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data);
void save_imatrix(std::string fname, int ncall = -1) const;
bool load_imatrix(const char* fname);
std::vector<float> get_values(const std::string& key) const {
auto it = stats_.find(key);
if (it != stats_.end()) {
return it->second.values;
} else {
return {};
}
}
};
IMatrixCollector& get_imatrix_collector();
#endif // __SD_RUNTIME_IMATRIX_H__

View File

@ -29,6 +29,7 @@
#include "model/diffusion/krea2.hpp" #include "model/diffusion/krea2.hpp"
#include "model/diffusion/lens.hpp" #include "model/diffusion/lens.hpp"
#include "model/diffusion/ltxv.hpp" #include "model/diffusion/ltxv.hpp"
#include "model/diffusion/minit2i.hpp"
#include "model/diffusion/mmdit.hpp" #include "model/diffusion/mmdit.hpp"
#include "model/diffusion/model.hpp" #include "model/diffusion/model.hpp"
#include "model/diffusion/pid.hpp" #include "model/diffusion/pid.hpp"
@ -93,6 +94,7 @@ const char* model_version_to_str[] = {
"Ovis Image", "Ovis Image",
"Ernie Image", "Ernie Image",
"Lens", "Lens",
"MiniT2I",
"Longcat-Image", "Longcat-Image",
"PiD", "PiD",
"Ideogram 4", "Ideogram 4",
@ -785,6 +787,14 @@ public:
tensor_storage_map, tensor_storage_map,
"model", "model",
model_manager); model_manager);
} else if (sd_version_is_minit2i(version)) {
cond_stage_model = std::make_shared<MiniT2IConditioner>(backend_for(SDBackendModule::TE),
tensor_storage_map,
model_manager);
diffusion_model = std::make_shared<MiniT2I::MiniT2IRunner>(backend_for(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model.model.net",
model_manager);
} else if (sd_version_is_anima(version)) { } else if (sd_version_is_anima(version)) {
cond_stage_model = std::make_shared<AnimaConditioner>(backend_for(SDBackendModule::TE), cond_stage_model = std::make_shared<AnimaConditioner>(backend_for(SDBackendModule::TE),
tensor_storage_map, tensor_storage_map,
@ -958,7 +968,7 @@ public:
} }
}; };
if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) { if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) {
LOG_INFO("using FakeVAE"); LOG_INFO("using FakeVAE");
first_stage_model = std::make_shared<FakeVAE>(version, first_stage_model = std::make_shared<FakeVAE>(version,
backend_for(SDBackendModule::VAE), backend_for(SDBackendModule::VAE),
@ -1299,6 +1309,8 @@ public:
} }
} else if (sd_version_is_sefi_image(version)) { } else if (sd_version_is_sefi_image(version)) {
pred_type = SEFI_FLOW_PRED; pred_type = SEFI_FLOW_PRED;
} else if (sd_version_is_minit2i(version)) {
pred_type = MINIT2I_FLOW_PRED;
} else { } else {
pred_type = EPS_PRED; pred_type = EPS_PRED;
} }
@ -1336,6 +1348,11 @@ public:
denoiser = std::make_shared<SefiFlowDenoiser>(); denoiser = std::make_shared<SefiFlowDenoiser>();
break; break;
} }
case MINIT2I_FLOW_PRED: {
LOG_INFO("running in MiniT2I FLOW mode");
denoiser = std::make_shared<MiniT2IFlowDenoiser>();
break;
}
default: { default: {
LOG_ERROR("Unknown predition type %i", pred_type); LOG_ERROR("Unknown predition type %i", pred_type);
return false; return false;
@ -2032,11 +2049,12 @@ public:
} }
int64_t last_progress_us = ggml_time_us(); int64_t last_progress_us = ggml_time_us();
SamplePreviewContext preview = prepare_sample_preview_context();
sd::Tensor<float> x_t = !noise.empty() sd::Tensor<float> x_t = !noise.empty()
? denoiser->noise_scaling(sigmas[0], noise, init_latent) ? denoiser->noise_scaling(sigmas[0], noise, init_latent)
: init_latent; : init_latent;
sd::Tensor<float> denoised = x_t; sd::Tensor<float> denoised = x_t;
SamplePreviewContext preview = prepare_sample_preview_context();
auto denoise = [&](const sd::Tensor<float>& x, float sigma, int step) -> sd::guidance::GuiderOutput { auto denoise = [&](const sd::Tensor<float>& x, float sigma, int step) -> sd::guidance::GuiderOutput {
if (get_cancel_flag() == SD_CANCEL_ALL) { if (get_cancel_flag() == SD_CANCEL_ALL) {
@ -2155,6 +2173,9 @@ public:
audio_length, audio_length,
frame_rate, frame_rate,
video_positions.empty() ? nullptr : &video_positions}; video_positions.empty() ? nullptr : &video_positions};
} else if (sd_version_is_minit2i(version)) {
diffusion_params.extra = MiniT2IDiffusionExtra{
condition.c_vector.empty() ? nullptr : &condition.c_vector};
} else { } else {
diffusion_params.extra = std::monostate{}; diffusion_params.extra = std::monostate{};
} }
@ -2335,6 +2356,8 @@ public:
latent_channel = 3; latent_channel = 3;
} else if (version == VERSION_CHROMA_RADIANCE) { } else if (version == VERSION_CHROMA_RADIANCE) {
latent_channel = 3; latent_channel = 3;
} else if (sd_version_is_minit2i(version)) {
latent_channel = 3;
} else if (sd_version_is_pid(version)) { } else if (sd_version_is_pid(version)) {
latent_channel = 3; latent_channel = 3;
} else if (sd_version_is_sefi_image(version)) { } else if (sd_version_is_sefi_image(version)) {
@ -2416,7 +2439,7 @@ public:
} }
sd::Tensor<float> decode_first_stage(const sd::Tensor<float>& x, bool decode_video = false) { sd::Tensor<float> decode_first_stage(const sd::Tensor<float>& x, bool decode_video = false) {
if (sd_version_is_pid(version)) { if (sd_version_is_pid(version) || sd_version_is_minit2i(version)) {
return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f); return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f);
} }
auto latents = first_stage_model->diffusion_to_vae_latents(x); auto latents = first_stage_model->diffusion_to_vae_latents(x);
@ -2591,6 +2614,7 @@ const char* prediction_to_str[] = {
"sd3_flow", "sd3_flow",
"flux_flow", "flux_flow",
"sefi_flow", "sefi_flow",
"minit2i_flow",
}; };
const char* sd_prediction_name(enum prediction_t prediction) { const char* sd_prediction_name(enum prediction_t prediction) {
@ -4224,6 +4248,11 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
if (request->use_uncond || request->use_high_noise_uncond) { if (request->use_uncond || request->use_high_noise_uncond) {
if (sd_version_is_ideogram4(sd_ctx->sd->version)) { if (sd_version_is_ideogram4(sd_ctx->sd->version)) {
uncond.c_vector = sd::Tensor<float>::from_vector({1.0f}); uncond.c_vector = sd::Tensor<float>::from_vector({1.0f});
} else if (sd_version_is_minit2i(sd_ctx->sd->version)) {
// MiniT2I derives the unconditional signal from the same T5 hidden
// states with a zeroed prompt mask, so no extra text encode is needed.
uncond.c_crossattn = cond.c_crossattn;
uncond.c_vector = sd::Tensor<float>::zeros_like(cond.c_vector);
} else { } else {
bool zero_out_masked = false; bool zero_out_masked = false;
if (sd_version_is_sdxl(sd_ctx->sd->version) && if (sd_version_is_sdxl(sd_ctx->sd->version) &&