DarthAffe/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-09-24 20:20:37 +00:00
Code Issues Packages Projects Releases Wiki Activity
stable-diffusion.cpp/src/conditioning/conditioner.hpp

3714 lines
173 KiB
C++
Raw Blame History

#ifndef __SD_CONDITIONING_CONDITIONER_HPP__
#define __SD_CONDITIONING_CONDITIONER_HPP__
#include <cinttypes>
#include <cmath>
#include <iomanip>
#include <limits>
#include <optional>
#include <sstream>
#include <stdexcept>
#include "core/ggml_tensor_utils.h"
#include "core/tensor_ggml.hpp"
#include "core/util.h"
#include "model/diffusion/model.hpp"
#include "model/te/clip.hpp"
#include "model/te/llada_image_te.hpp"
#include "model/te/llm.hpp"
#include "model/te/t5.hpp"
#include "model_loader.h"
#include "tokenizers/sensenova_u1_tokenizer.h"
#include "tokenizers/tokenizer_config.h"
struct SDCondition {
sd::Tensor<float> c_crossattn;
sd::Tensor<float> c_vector;
sd::Tensor<float> c_concat;
sd::Tensor<int32_t> c_t5_ids;
sd::Tensor<float> c_t5_weights;
sd::Tensor<int32_t> c_input_ids;
sd::Tensor<int32_t> c_position_ids;
sd::Tensor<int32_t> c_token_types;
sd::Tensor<int32_t> c_vinput_mask;
std::vector<std::pair<int, sd::Tensor<float>>> c_image_embeds;
std::vector<sd::Tensor<float>> c_ref_images;
std::vector<sd::Tensor<float>> c_ref_audios;
std::vector<MiniMaxH3ReferenceBlock> c_reference_blocks;
std::vector<sd::Tensor<float>> extra_c_crossattns;
SDCondition() = default;
SDCondition(sd::Tensor<float> c_crossattn,
sd::Tensor<float> c_vector,
sd::Tensor<float> c_concat)
: c_crossattn(std::move(c_crossattn)), c_vector(std::move(c_vector)), c_concat(std::move(c_concat)) {}
bool empty() const {
if (!c_crossattn.empty() || !c_vector.empty() || !c_concat.empty() ||
!c_t5_ids.empty() || !c_t5_weights.empty() ||
!c_input_ids.empty() || !c_position_ids.empty() ||
!c_token_types.empty() || !c_vinput_mask.empty()) {
return false;
}
for (const auto& image_embed : c_image_embeds) {
if (!image_embed.second.empty()) {
return false;
}
}
for (const auto& tensor : c_ref_images) {
if (!tensor.empty()) {
return false;
}
}
for (const auto& tensor : c_ref_audios) {
if (!tensor.empty()) {
return false;
}
}
for (const auto& tensor : extra_c_crossattns) {
if (!tensor.empty()) {
return false;
}
}
return true;
}
};
enum class MiniMaxH3PresentationKind {
IMAGE,
VIDEO,
AUDIO,
};
struct MiniMaxH3PresentationItem {
MiniMaxH3PresentationKind kind = MiniMaxH3PresentationKind::IMAGE;
std::vector<sd::Tensor<float>> frames;
std::vector<float> timestamps;
};
static inline sd::Tensor<float> apply_token_weights(sd::Tensor<float> hidden_states,
const std::vector<float>& weights) {
if (hidden_states.empty()) {
return hidden_states;
}
bool all_one = true;
for (float weight : weights) {
if (weight != 1.0f) {
all_one = false;
break;
}
}
if (all_one) {
return hidden_states;
}
if (hidden_states.dim() == 1) {
hidden_states.unsqueeze_(1);
}
GGML_ASSERT(static_cast<size_t>(hidden_states.shape()[1]) == weights.size());
float original_mean = hidden_states.mean();
auto chunk_weights = sd::Tensor<float>::from_vector(weights);
chunk_weights.reshape_({1, static_cast<int64_t>(weights.size())});
hidden_states *= chunk_weights;
float new_mean = hidden_states.mean();
if (std::isfinite(original_mean) && std::isfinite(new_mean) && new_mean != 0.0f) {
hidden_states *= (original_mean / new_mean);
}
return hidden_states;
}
struct ConditionerParams {
std::string text;
int clip_skip = -1;
int width = -1;
int height = -1;
bool zero_out_masked = false;
const std::vector<sd::Tensor<float>>* ref_images = nullptr; // for qwen image edit
const std::vector<MiniMaxH3PresentationItem>* minimax_h3_references = nullptr;
RefImageParams ref_image_params;
};
struct Conditioner {
virtual ~Conditioner() = default;
public:
virtual SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) = 0;
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {}
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
virtual void set_graph_cut_layer_split_enabled(bool enabled) {}
virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {}
virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& 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<WeightAdapter>& adapter) {}
virtual void runner_end() {}
};
// ldm.modules.encoders.modules.FrozenCLIPEmbedder
// Ref: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/cad87bf4e3e0b0a759afa94e933527c3123d59bc/modules/sd_hijack_clip.py#L283
struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
SDVersion version = VERSION_SD1;
std::shared_ptr<Tokenizer> tokenizer;
std::shared_ptr<CLIPTextModelRunner> text_model;
std::shared_ptr<CLIPTextModelRunner> text_model2;
std::map<std::string, std::string> embedding_map;
int32_t num_custom_embeddings = 0;
std::vector<uint8_t> token_embed_custom;
std::vector<uint8_t> token_embed_custom2;
std::map<std::string, std::pair<int, int>> embedding_pos_map;
FrozenCLIPEmbedderWithCustomWords(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::map<std::string, std::string>& orig_embedding_map,
SDVersion version = VERSION_SD1,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {})
: version(version) {
const int pad_id = sd_version_is_sd2(version) ? 0 : 49407;
tokenizer = tokenizers.create(TokenizerConfig::MAIN, 49408, pad_id, false, true);
if (!tokenizer) {
tokenizer = std::make_shared<CLIPTokenizer>(pad_id);
}
for (const auto& kv : orig_embedding_map) {
std::string name = normalize_embedding_name(kv.first);
embedding_map[name] = kv.second;
tokenizer->add_special_token(name);
}
bool force_clip_f32 = !embedding_map.empty();
if (sd_version_is_sd1(version)) {
text_model = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "cond_stage_model.transformer.text_model", OPENAI_CLIP_VIT_L_14, true, force_clip_f32, weight_manager);
} else if (sd_version_is_sd2(version)) {
text_model = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "cond_stage_model.transformer.text_model", OPEN_CLIP_VIT_H_14, true, force_clip_f32, weight_manager);
} else if (sd_version_is_sdxl(version)) {
text_model = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "cond_stage_model.transformer.text_model", OPENAI_CLIP_VIT_L_14, false, force_clip_f32, weight_manager);
text_model2 = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "cond_stage_model.1.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false, force_clip_f32, weight_manager);
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
text_model->get_param_tensors(tensors, "cond_stage_model.transformer.text_model");
if (sd_version_is_sdxl(version)) {
text_model2->get_param_tensors(tensors, "cond_stage_model.1.transformer.text_model");
}
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
text_model->set_max_graph_vram_bytes(max_vram_bytes);
if (sd_version_is_sdxl(version)) {
text_model2->set_max_graph_vram_bytes(max_vram_bytes);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
text_model->set_runtime_backends(backends);
if (sd_version_is_sdxl(version)) {
text_model2->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
text_model->set_graph_cut_layer_split_enabled(enabled);
if (sd_version_is_sdxl(version)) {
text_model2->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
text_model->set_graph_cut_layer_split_backend_vram_limits(limits);
if (sd_version_is_sdxl(version)) {
text_model2->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void set_flash_attention_enabled(bool enabled) override {
text_model->set_flash_attention_enabled(enabled);
if (sd_version_is_sdxl(version)) {
text_model2->set_flash_attention_enabled(enabled);
}
}
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<WeightAdapter>& adapter) override {
text_model->set_weight_adapter(adapter);
if (sd_version_is_sdxl(version)) {
text_model2->set_weight_adapter(adapter);
}
}
void runner_end() override {
text_model->runner_end();
if (sd_version_is_sdxl(version)) {
text_model2->runner_end();
}
}
bool load_embedding(std::string embd_name, std::string embd_path, std::vector<int32_t>& bpe_tokens) {
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(embd_path)) {
LOG_ERROR("embedding '%s' failed", embd_name.c_str());
return false;
}
auto push_ids = [&](int pos_start, int pos_end, bool cached) {
for (int i = pos_start; i < pos_end; i++) {
bpe_tokens.push_back(text_model->model.vocab_size + i);
}
if (!cached) {
LOG_VERBOSE("embedding '%s' applied: %i token(s), custom embeddings: %i", embd_name.c_str(), pos_end - pos_start, num_custom_embeddings);
}
};
auto iter = embedding_pos_map.find(embd_name);
if (iter != embedding_pos_map.end()) {
LOG_VERBOSE("embedding already read in: %s", embd_name.c_str());
push_ids(iter->second.first, iter->second.second, true);
return true;
}
ggml_init_params params;
params.mem_size = 100 * 1024 * 1024; // max for custom embeddings 100 MB
params.mem_buffer = nullptr;
params.no_alloc = false;
auto ggml_ctx_deleter = [](ggml_context* ctx) { ggml_free(ctx); };
auto embd_ctx = std::unique_ptr<ggml_context, decltype(ggml_ctx_deleter)>(ggml_init(params), ggml_ctx_deleter);
if (!embd_ctx.get()) {
LOG_ERROR("ggml_init failed when loading embeddings file");
return false;
}
ggml_tensor* embd = nullptr;
ggml_tensor* embd2 = nullptr;
ggml_type embd_type = text_model->model.get_token_embed_weight()->type;
ggml_type embd2_type = text_model2 ? text_model2->model.get_token_embed_weight()->type : embd_type;
int64_t hidden_size = text_model->model.hidden_size;
int64_t hidden_size2 = text_model2 ? text_model2->model.hidden_size : 0;
auto on_load = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) {
if (tensor_storage.ne[0] == hidden_size) {
embd = ggml_new_tensor_2d(embd_ctx.get(), embd_type, hidden_size, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1);
if (embd == nullptr) {
return false;
}
*dst_tensor = embd;
} else if (text_model2) {
if (tensor_storage.ne[0] == hidden_size2) {
embd2 = ggml_new_tensor_2d(embd_ctx.get(), embd2_type, hidden_size2, tensor_storage.n_dims > 1 ? tensor_storage.ne[1] : 1);
if (embd2 == nullptr) {
return false;
}
*dst_tensor = embd2;
} else {
LOG_VERBOSE("embedding skipped, wrong hidden size, got %i, expected %i or %i", tensor_storage.ne[0], hidden_size, hidden_size2);
}
} else {
LOG_VERBOSE("embedding skipped, wrong hidden size, got %i, expected %i", tensor_storage.ne[0], hidden_size);
}
return true;
};
model_loader.set_n_threads(1);
if (!model_loader.load_tensors(on_load)) {
LOG_ERROR("embedding '%s' failed", embd_name.c_str());
return false;
}
if (!embd && !embd2) {
LOG_WARN("embedding '%s' has no usable tensor", embd_name.c_str());
return false;
}
int pos_start = num_custom_embeddings;
int64_t embd_rows = embd ? embd->ne[1] : 0;
int64_t embd2_rows = embd2 ? embd2->ne[1] : 0;
if (embd_rows < embd2_rows) {
LOG_WARN("embedding '%s' has fewer rows for text model 1, zero-padding", embd_name.c_str());
} else if (text_model2 && embd2_rows < embd_rows) {
LOG_WARN("embedding '%s' has fewer rows for text model 2, zero-padding", embd_name.c_str());
}
int64_t rows = std::max(embd_rows, embd2_rows);
size_t embd_bytes = hidden_size * ggml_type_size(embd_type);
token_embed_custom.resize(token_embed_custom.size() + embd_bytes * rows);
if (embd) {
memcpy((void*)(token_embed_custom.data() + embd_bytes * num_custom_embeddings),
embd->data, embd_bytes * embd_rows);
}
if (text_model2) {
size_t embd2_bytes = hidden_size2 * ggml_type_size(embd2_type);
token_embed_custom2.resize(token_embed_custom2.size() + embd2_bytes * rows);
if (embd2) {
memcpy((void*)(token_embed_custom2.data() + embd2_bytes * num_custom_embeddings),
embd2->data, embd2_bytes * embd2_rows);
}
}
num_custom_embeddings += (int)rows;
int pos_end = num_custom_embeddings;
push_ids(pos_start, pos_end, false);
embedding_pos_map[embd_name] = std::pair{pos_start, pos_end};
return true;
}
static std::string normalize_embedding_name(std::string name) {
std::transform(name.begin(), name.end(), name.begin(), [](unsigned char c) { return std::tolower(c); });
return name;
}
bool append_embedding_tokens(std::string str, std::vector<int32_t>& bpe_tokens) {
std::string name = normalize_embedding_name(std::move(str));
auto iter = embedding_map.find(name);
if (iter == embedding_map.end()) {
return false;
}
return load_embedding(name, iter->second, bpe_tokens);
}
bool convert_token_to_id(const std::string& text, std::vector<int>& tokens) {
auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
return append_embedding_tokens(str, bpe_tokens);
};
return tokenizer->encode(text, tokens, on_new_token_cb);
}
bool decode(const std::vector<int>& tokens, std::string& text) {
return tokenizer->decode(tokens, text);
}
std::pair<std::vector<int>, std::vector<float>> tokenize(std::string text,
size_t min_length = 0,
size_t max_length = 0,
bool allow_overflow_expand = true) {
auto parsed_attention = parse_prompt_attention(text);
{
std::stringstream ss;
ss << "[";
for (const auto& item : parsed_attention) {
ss << "['" << item.first << "', " << item.second << "], ";
}
ss << "]";
LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
}
auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
return append_embedding_tokens(str, bpe_tokens);
};
std::vector<int> tokens;
std::vector<float> weights;
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
float curr_weight = item.second;
if (curr_text == "BREAK" && curr_weight == -1.0f) {
// Pad token array up to chunk size at this point.
// TODO: This is a hardcoded chunk_len, like in stable-diffusion.cpp, make it a parameter for the future?
// Also, this is 75 instead of 77 to leave room for BOS and EOS tokens.
size_t current_size = tokens.size();
size_t padding_size = (75 - (current_size % 75)) % 75; // Ensure no negative padding
if (padding_size > 0) {
LOG_VERBOSE("BREAK token encountered, padding current chunk by %zu tokens.", padding_size);
tokens.insert(tokens.end(), padding_size, tokenizer->EOS_TOKEN_ID);
weights.insert(weights.end(), padding_size, 1.0f);
}
continue; // Skip to the next item after handling BREAK
}
std::vector<int> curr_tokens;
if (!tokenizer->encode(curr_text, curr_tokens, on_new_token_cb)) {
return {};
}
tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
weights.insert(weights.end(), curr_tokens.size(), curr_weight);
}
tokenizer->pad_tokens(tokens, &weights, nullptr, min_length, max_length, allow_overflow_expand);
// for (int i = 0; i < tokens.size(); i++) {
// std::cout << tokens[i] << ":" << weights[i] << ", ";
// }
// std::cout << std::endl;
return {tokens, weights};
}
SDCondition get_learned_condition_common(int n_threads,
std::vector<int>& tokens,
std::vector<float>& weights,
int clip_skip,
int width,
int height,
bool zero_out_masked = false) {
int64_t t0 = ggml_time_ms();
sd::Tensor<float> hidden_states; // [n_token, hidden_size] or [n_token, hidden_size + hidden_size2]
sd::Tensor<float> pooled;
if (clip_skip <= 0) {
clip_skip = (sd_version_is_sd2(version) || sd_version_is_sdxl(version)) ? 2 : 1;
}
size_t chunk_len = 77;
size_t chunk_count = tokens.size() / chunk_len;
for (int chunk_idx = 0; chunk_idx < chunk_count; chunk_idx++) {
std::vector<int> chunk_tokens(tokens.begin() + chunk_idx * chunk_len,
tokens.begin() + (chunk_idx + 1) * chunk_len);
std::vector<float> chunk_weights(weights.begin() + chunk_idx * chunk_len,
weights.begin() + (chunk_idx + 1) * chunk_len);
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(chunk_tokens.size())}, chunk_tokens);
sd::Tensor<int32_t> input_ids2;
size_t max_token_idx = 0;
if (sd_version_is_sdxl(version)) {
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), tokenizer->EOS_TOKEN_ID);
if (it != chunk_tokens.end()) {
std::fill(std::next(it), chunk_tokens.end(), 0);
}
max_token_idx = std::min<size_t>(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1);
input_ids2 = sd::Tensor<int32_t>({static_cast<int64_t>(chunk_tokens.size())}, chunk_tokens);
// for (int i = 0; i < chunk_tokens.size(); i++) {
// printf("%d ", chunk_tokens[i]);
// }
// printf("\n");
}
{
auto chunk_hidden_states = text_model->compute(n_threads,
input_ids,
num_custom_embeddings,
token_embed_custom.data(),
max_token_idx,
false,
clip_skip,
false);
GGML_ASSERT(!chunk_hidden_states.empty());
if (sd_version_is_sdxl(version)) {
auto chunk_hidden_states2 = text_model2->compute(n_threads,
input_ids2,
num_custom_embeddings,
token_embed_custom2.data(),
max_token_idx,
false,
clip_skip,
false);
GGML_ASSERT(!chunk_hidden_states2.empty());
chunk_hidden_states = sd::ops::concat(chunk_hidden_states, chunk_hidden_states2, 0);
if (chunk_idx == 0) {
pooled = text_model2->compute(n_threads,
input_ids2,
num_custom_embeddings,
token_embed_custom2.data(),
max_token_idx,
true,
clip_skip,
false);
GGML_ASSERT(!pooled.empty());
}
}
int64_t t1 = ggml_time_ms();
LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
chunk_hidden_states = apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
if (zero_out_masked) {
chunk_hidden_states.fill_(0.0f);
}
if (!hidden_states.empty()) {
hidden_states = sd::ops::concat(hidden_states, chunk_hidden_states, 1);
} else {
hidden_states = std::move(chunk_hidden_states);
}
}
}
sd::Tensor<float> vec;
if (sd_version_is_sdxl(version)) {
int out_dim = 256;
int adm_in_channels = 2816;
GGML_ASSERT(!pooled.empty());
vec = sd::Tensor<float>({adm_in_channels});
vec.fill_(0.0f);
size_t offset = 0;
std::copy(pooled.values().begin(), pooled.values().end(), vec.values().begin());
offset += pooled.values().size();
auto append_embedding = [&](const std::vector<float>& timesteps) {
sd::Tensor<float> embedding;
set_timestep_embedding(timesteps, &embedding, out_dim);
std::copy(embedding.values().begin(), embedding.values().end(), vec.values().begin() + static_cast<int64_t>(offset));
offset += embedding.values().size();
};
append_embedding({static_cast<float>(height), static_cast<float>(width)});
append_embedding({0.0f, 0.0f});
append_embedding({static_cast<float>(height), static_cast<float>(width)});
GGML_ASSERT(offset == vec.values().size());
}
SDCondition result;
if (!hidden_states.empty()) {
result.c_crossattn = std::move(hidden_states);
}
if (!vec.empty()) {
result.c_vector = std::move(vec);
}
return result;
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
auto tokens_and_weights = tokenize(conditioner_params.text, text_model->model.n_token, text_model->model.n_token, true);
if (tokens_and_weights.first.empty()) {
return {};
}
std::vector<int>& tokens = tokens_and_weights.first;
std::vector<float>& weights = tokens_and_weights.second;
return get_learned_condition_common(n_threads,
tokens,
weights,
conditioner_params.clip_skip,
conditioner_params.width,
conditioner_params.height,
conditioner_params.zero_out_masked);
}
};
struct FrozenCLIPVisionEmbedder : public GGMLRunner {
CLIPVisionModelProjection vision_model;
std::string weight_prefix = "cond_stage_model.transformer";
FrozenCLIPVisionEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager) {
bool proj_in = false;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, weight_prefix)) {
continue;
}
if (contains(name, "self_attn.in_proj")) {
proj_in = true;
break;
}
}
vision_model = CLIPVisionModelProjection(OPEN_CLIP_VIT_H_14, false, proj_in);
vision_model.init(params_ctx, tensor_storage_map, weight_prefix);
}
std::string get_desc() override {
return "clip_vision";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
vision_model.get_param_tensors(tensors, weight_prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& pixel_values_tensor, bool return_pooled, int clip_skip) {
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
ggml_tensor* pixel_values = make_input(pixel_values_tensor);
auto runner_ctx = get_context();
ggml_tensor* hidden_states = vision_model.forward(&runner_ctx, pixel_values, return_pooled, clip_skip);
ggml_build_forward_expand(gf, hidden_states);
return gf;
}
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<float>& pixel_values,
bool return_pooled,
int clip_skip) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(pixel_values, return_pooled, clip_skip);
};
return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
}
};
struct SD3CLIPEmbedder : public Conditioner {
std::shared_ptr<Tokenizer> clip_l_tokenizer;
std::shared_ptr<Tokenizer> clip_g_tokenizer;
T5UniGramTokenizer t5_tokenizer;
std::shared_ptr<CLIPTextModelRunner> clip_l;
std::shared_ptr<CLIPTextModelRunner> clip_g;
std::shared_ptr<T5Runner> t5;
SD3CLIPEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {}) {
bool use_clip_l = false;
bool use_clip_g = false;
bool use_t5 = false;
for (auto pair : tensor_storage_map) {
if (pair.first.find("text_encoders.clip_l") != std::string::npos) {
use_clip_l = true;
} else if (pair.first.find("text_encoders.clip_g") != std::string::npos) {
use_clip_g = true;
} else if (pair.first.find("text_encoders.t5xxl") != std::string::npos) {
use_t5 = true;
}
}
if (!use_clip_l && !use_clip_g && !use_t5) {
LOG_WARN("IMPORTANT NOTICE: No text encoders provided, cannot process prompts!");
return;
}
if (use_clip_l) {
clip_l_tokenizer = tokenizers.create(TokenizerConfig::CLIP_L, 49408, 49407, false, true);
if (!clip_l_tokenizer) {
clip_l_tokenizer = std::make_shared<CLIPTokenizer>();
}
clip_l = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, false, false, weight_manager);
}
if (use_clip_g) {
clip_g_tokenizer = tokenizers.create(TokenizerConfig::CLIP_G, 49408, 0, false, true);
if (!clip_g_tokenizer) {
clip_g_tokenizer = std::make_shared<CLIPTokenizer>(0);
}
clip_g = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "text_encoders.clip_g.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false, false, weight_manager);
}
if (use_t5) {
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 (clip_l) {
clip_l->get_param_tensors(tensors, "text_encoders.clip_l.transformer.text_model");
}
if (clip_g) {
clip_g->get_param_tensors(tensors, "text_encoders.clip_g.transformer.text_model");
}
if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
if (clip_l) {
clip_l->set_max_graph_vram_bytes(max_vram_bytes);
}
if (clip_g) {
clip_g->set_max_graph_vram_bytes(max_vram_bytes);
}
if (t5) {
t5->set_max_graph_vram_bytes(max_vram_bytes);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (clip_l) {
clip_l->set_runtime_backends(backends);
}
if (clip_g) {
clip_g->set_runtime_backends(backends);
}
if (t5) {
t5->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_graph_cut_layer_split_enabled(enabled);
}
if (clip_g) {
clip_g->set_graph_cut_layer_split_enabled(enabled);
}
if (t5) {
t5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (clip_l) {
clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
}
if (clip_g) {
clip_g->set_graph_cut_layer_split_backend_vram_limits(limits);
}
if (t5) {
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_flash_attention_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_flash_attention_enabled(enabled);
}
if (clip_g) {
clip_g->set_flash_attention_enabled(enabled);
}
if (t5) {
t5->set_flash_attention_enabled(enabled);
}
}
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<WeightAdapter>& adapter) override {
if (clip_l) {
clip_l->set_weight_adapter(adapter);
}
if (clip_g) {
clip_g->set_weight_adapter(adapter);
}
if (t5) {
t5->set_weight_adapter(adapter);
}
}
void runner_end() override {
if (clip_l) {
clip_l->runner_end();
}
if (clip_g) {
clip_g->runner_end();
}
if (t5) {
t5->runner_end();
}
}
std::vector<std::pair<std::vector<int>, std::vector<float>>> tokenize(std::string text,
size_t min_length = 0,
size_t max_length = 0,
bool allow_overflow_expand = true) {
auto parsed_attention = parse_prompt_attention(text);
{
std::stringstream ss;
ss << "[";
for (const auto& item : parsed_attention) {
ss << "['" << item.first << "', " << item.second << "], ";
}
ss << "]";
LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
}
auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
return false;
};
std::vector<int> clip_l_tokens;
std::vector<float> clip_l_weights;
std::vector<int> clip_g_tokens;
std::vector<float> clip_g_weights;
std::vector<int> t5_tokens;
std::vector<float> t5_weights;
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
float curr_weight = item.second;
if (clip_l) {
std::vector<int> curr_tokens;
if (!clip_l_tokenizer->encode(curr_text, curr_tokens, on_new_token_cb)) {
return {};
}
clip_l_tokens.insert(clip_l_tokens.end(), curr_tokens.begin(), curr_tokens.end());
clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight);
}
if (clip_g) {
std::vector<int> curr_tokens;
if (!clip_g_tokenizer->encode(curr_text, curr_tokens, on_new_token_cb)) {
return {};
}
clip_g_tokens.insert(clip_g_tokens.end(), curr_tokens.begin(), curr_tokens.end());
clip_g_weights.insert(clip_g_weights.end(), curr_tokens.size(), curr_weight);
}
if (t5) {
std::vector<int> curr_tokens;
if (!t5_tokenizer.encode(curr_text, curr_tokens)) {
return {};
}
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
}
}
if (clip_l) {
clip_l_tokenizer->pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, min_length, max_length, allow_overflow_expand);
}
if (clip_g) {
clip_g_tokenizer->pad_tokens(clip_g_tokens, &clip_g_weights, nullptr, min_length, max_length, allow_overflow_expand);
}
if (t5) {
t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr, min_length, max_length, true);
}
// for (int i = 0; i < clip_l_tokens.size(); i++) {
// std::cout << clip_l_tokens[i] << ":" << clip_l_weights[i] << ", ";
// }
// std::cout << std::endl;
// for (int i = 0; i < clip_g_tokens.size(); i++) {
// std::cout << clip_g_tokens[i] << ":" << clip_g_weights[i] << ", ";
// }
// std::cout << std::endl;
// for (int i = 0; i < t5_tokens.size(); i++) {
// std::cout << t5_tokens[i] << ":" << t5_weights[i] << ", ";
// }
// std::cout << std::endl;
return {{clip_l_tokens, clip_l_weights}, {clip_g_tokens, clip_g_weights}, {t5_tokens, t5_weights}};
}
SDCondition get_learned_condition_common(int n_threads,
std::vector<std::pair<std::vector<int>, std::vector<float>>> token_and_weights,
int clip_skip,
bool zero_out_masked = false) {
auto& clip_l_tokens = token_and_weights[0].first;
auto& clip_l_weights = token_and_weights[0].second;
auto& clip_g_tokens = token_and_weights[1].first;
auto& clip_g_weights = token_and_weights[1].second;
auto& t5_tokens = token_and_weights[2].first;
auto& t5_weights = token_and_weights[2].second;
if (clip_skip <= 0) {
clip_skip = 2;
}
size_t chunk_len = 77;
int64_t t0 = ggml_time_ms();
sd::Tensor<float> hidden_states;
sd::Tensor<float> pooled;
size_t chunk_count = std::max(std::max(clip_l_tokens.size(), clip_g_tokens.size()), t5_tokens.size()) / chunk_len;
for (int chunk_idx = 0; chunk_idx < chunk_count; chunk_idx++) {
// clip_l
sd::Tensor<float> chunk_hidden_states_l;
sd::Tensor<float> pooled_l;
if (clip_l) {
std::vector<int> chunk_tokens(clip_l_tokens.begin() + chunk_idx * chunk_len,
clip_l_tokens.begin() + (chunk_idx + 1) * chunk_len);
std::vector<float> chunk_weights(clip_l_weights.begin() + chunk_idx * chunk_len,
clip_l_weights.begin() + (chunk_idx + 1) * chunk_len);
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(chunk_tokens.size())}, chunk_tokens);
size_t max_token_idx = 0;
chunk_hidden_states_l = clip_l->compute(n_threads,
input_ids,
0,
nullptr,
max_token_idx,
false,
clip_skip,
false);
GGML_ASSERT(!chunk_hidden_states_l.empty());
chunk_hidden_states_l = ::apply_token_weights(std::move(chunk_hidden_states_l), chunk_weights);
if (chunk_idx == 0) {
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_l_tokenizer->EOS_TOKEN_ID);
max_token_idx = std::min<size_t>(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1);
pooled_l = clip_l->compute(n_threads,
input_ids,
0,
nullptr,
max_token_idx,
true,
clip_skip,
false);
GGML_ASSERT(!pooled_l.empty());
}
} else {
chunk_hidden_states_l = sd::Tensor<float>::zeros({768, static_cast<int64_t>(chunk_len), 1});
if (chunk_idx == 0) {
pooled_l = sd::Tensor<float>::zeros({768, 1});
}
}
// clip_g
sd::Tensor<float> chunk_hidden_states_g;
sd::Tensor<float> pooled_g;
if (clip_g) {
std::vector<int> chunk_tokens(clip_g_tokens.begin() + chunk_idx * chunk_len,
clip_g_tokens.begin() + (chunk_idx + 1) * chunk_len);
std::vector<float> chunk_weights(clip_g_weights.begin() + chunk_idx * chunk_len,
clip_g_weights.begin() + (chunk_idx + 1) * chunk_len);
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(chunk_tokens.size())}, chunk_tokens);
size_t max_token_idx = 0;
chunk_hidden_states_g = clip_g->compute(n_threads,
input_ids,
0,
nullptr,
max_token_idx,
false,
clip_skip,
false);
GGML_ASSERT(!chunk_hidden_states_g.empty());
chunk_hidden_states_g = ::apply_token_weights(std::move(chunk_hidden_states_g), chunk_weights);
if (chunk_idx == 0) {
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_g_tokenizer->EOS_TOKEN_ID);
max_token_idx = std::min<size_t>(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1);
pooled_g = clip_g->compute(n_threads,
input_ids,
0,
nullptr,
max_token_idx,
true,
clip_skip,
false);
GGML_ASSERT(!pooled_g.empty());
}
} else {
chunk_hidden_states_g = sd::Tensor<float>::zeros({1280, static_cast<int64_t>(chunk_len), 1});
if (chunk_idx == 0) {
pooled_g = sd::Tensor<float>::zeros({1280, 1});
}
}
// t5
sd::Tensor<float> chunk_hidden_states_t5;
if (t5) {
std::vector<int> chunk_tokens(t5_tokens.begin() + chunk_idx * chunk_len,
t5_tokens.begin() + (chunk_idx + 1) * chunk_len);
std::vector<float> chunk_weights(t5_weights.begin() + chunk_idx * chunk_len,
t5_weights.begin() + (chunk_idx + 1) * chunk_len);
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(chunk_tokens.size())}, chunk_tokens);
chunk_hidden_states_t5 = t5->compute(n_threads,
input_ids,
sd::Tensor<float>(),
false);
GGML_ASSERT(!chunk_hidden_states_t5.empty());
chunk_hidden_states_t5 = ::apply_token_weights(std::move(chunk_hidden_states_t5), chunk_weights);
} else {
chunk_hidden_states_t5 = sd::Tensor<float>::zeros({4096, static_cast<int64_t>(chunk_len), 1});
}
sd::Tensor<float> chunk_hidden_states_lg = sd::ops::concat(chunk_hidden_states_l, chunk_hidden_states_g, 0);
if (chunk_hidden_states_lg.shape()[0] < 4096) {
auto pad_shape = chunk_hidden_states_lg.shape();
pad_shape[0] = 4096 - chunk_hidden_states_lg.shape()[0];
chunk_hidden_states_lg = sd::ops::concat(chunk_hidden_states_lg,
sd::Tensor<float>::zeros(pad_shape),
0);
}
sd::Tensor<float> chunk_hidden_states = sd::ops::concat(chunk_hidden_states_lg,
chunk_hidden_states_t5,
1); // [n_token*2, 4096]
if (chunk_idx == 0) {
pooled = sd::ops::concat(pooled_l, pooled_g, 0); // [768 + 1280]
}
int64_t t1 = ggml_time_ms();
LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
if (zero_out_masked) {
chunk_hidden_states.fill_(0.0f);
}
if (!hidden_states.empty()) {
hidden_states = sd::ops::concat(hidden_states, chunk_hidden_states, 1);
} else {
hidden_states = std::move(chunk_hidden_states);
}
}
SDCondition result;
result.c_crossattn = std::move(hidden_states);
result.c_vector = std::move(pooled);
return result;
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
auto tokens_and_weights = tokenize(conditioner_params.text, 77, 77, true);
if (tokens_and_weights.empty()) {
return {};
}
return get_learned_condition_common(n_threads,
tokens_and_weights,
conditioner_params.clip_skip,
conditioner_params.zero_out_masked);
}
};
struct FluxCLIPEmbedder : public Conditioner {
std::shared_ptr<Tokenizer> clip_l_tokenizer;
T5UniGramTokenizer t5_tokenizer;
std::shared_ptr<CLIPTextModelRunner> clip_l;
std::shared_ptr<T5Runner> t5;
size_t chunk_len = 256;
FluxCLIPEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {}) {
bool use_clip_l = false;
bool use_t5 = false;
for (auto pair : tensor_storage_map) {
if (pair.first.find("text_encoders.clip_l") != std::string::npos) {
use_clip_l = true;
} else if (pair.first.find("text_encoders.t5xxl") != std::string::npos) {
use_t5 = true;
}
}
if (!use_clip_l && !use_t5) {
LOG_WARN("IMPORTANT NOTICE: No text encoders provided, cannot process prompts!");
return;
}
if (use_clip_l) {
auto slot = tokenizers.has(TokenizerConfig::CLIP_L) ? TokenizerConfig::CLIP_L : TokenizerConfig::MAIN;
clip_l_tokenizer = tokenizers.create(slot, 49408, 49407, false, true);
if (!clip_l_tokenizer) {
clip_l_tokenizer = std::make_shared<CLIPTokenizer>();
}
clip_l = std::make_shared<CLIPTextModelRunner>(backend, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, true, false, weight_manager);
} else {
LOG_WARN("clip_l text encoder not found! Prompt adherence might be degraded.");
}
if (use_t5) {
t5 = std::make_shared<T5Runner>(backend, tensor_storage_map, "text_encoders.t5xxl.transformer", false, weight_manager);
} else {
LOG_WARN("t5xxl text encoder not found! Prompt adherence might be degraded.");
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (clip_l) {
clip_l->get_param_tensors(tensors, "text_encoders.clip_l.transformer.text_model");
}
if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
if (clip_l) {
clip_l->set_max_graph_vram_bytes(max_vram_bytes);
}
if (t5) {
t5->set_max_graph_vram_bytes(max_vram_bytes);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (clip_l) {
clip_l->set_runtime_backends(backends);
}
if (t5) {
t5->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_graph_cut_layer_split_enabled(enabled);
}
if (t5) {
t5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (clip_l) {
clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
}
if (t5) {
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_flash_attention_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_flash_attention_enabled(enabled);
}
if (t5) {
t5->set_flash_attention_enabled(enabled);
}
}
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<WeightAdapter>& adapter) override {
if (clip_l) {
clip_l->set_weight_adapter(adapter);
}
if (t5) {
t5->set_weight_adapter(adapter);
}
}
void runner_end() override {
if (clip_l) {
clip_l->runner_end();
}
if (t5) {
t5->runner_end();
}
}
std::vector<std::pair<std::vector<int>, std::vector<float>>> tokenize(std::string text,
size_t min_length = 0,
size_t max_length = 0) {
auto parsed_attention = parse_prompt_attention(text);
{
std::stringstream ss;
ss << "[";
for (const auto& item : parsed_attention) {
ss << "['" << item.first << "', " << item.second << "], ";
}
ss << "]";
LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
}
auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
return false;
};
std::vector<int> clip_l_tokens;
std::vector<float> clip_l_weights;
std::vector<int> t5_tokens;
std::vector<float> t5_weights;
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
float curr_weight = item.second;
if (clip_l) {
std::vector<int> curr_tokens;
if (!clip_l_tokenizer->encode(curr_text, curr_tokens, on_new_token_cb)) {
return {};
}
clip_l_tokens.insert(clip_l_tokens.end(), curr_tokens.begin(), curr_tokens.end());
clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight);
}
if (t5) {
std::vector<int> curr_tokens;
if (!t5_tokenizer.encode(curr_text, curr_tokens)) {
return {};
}
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
}
}
if (clip_l) {
clip_l_tokenizer->pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, 77, 77, true);
}
if (t5) {
t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr, min_length, max_length, true);
}
// for (int i = 0; i < clip_l_tokens.size(); i++) {
// std::cout << clip_l_tokens[i] << ":" << clip_l_weights[i] << ", ";
// }
// std::cout << std::endl;
// for (int i = 0; i < t5_tokens.size(); i++) {
// std::cout << t5_tokens[i] << ":" << t5_weights[i] << ", ";
// }
// std::cout << std::endl;
return {{clip_l_tokens, clip_l_weights}, {t5_tokens, t5_weights}};
}
SDCondition get_learned_condition_common(int n_threads,
std::vector<std::pair<std::vector<int>, std::vector<float>>> token_and_weights,
int clip_skip,
bool zero_out_masked = false) {
auto& clip_l_tokens = token_and_weights[0].first;
auto& clip_l_weights = token_and_weights[0].second;
auto& t5_tokens = token_and_weights[1].first;
auto& t5_weights = token_and_weights[1].second;
if (clip_skip <= 0) {
clip_skip = 2;
}
int64_t t0 = ggml_time_ms();
sd::Tensor<float> hidden_states; // [N, n_token, 4096]
sd::Tensor<float> pooled; // [768,]
size_t chunk_count = std::max(clip_l_tokens.size() > 0 ? chunk_len : 0, t5_tokens.size()) / chunk_len;
for (int chunk_idx = 0; chunk_idx < chunk_count; chunk_idx++) {
// clip_l
if (chunk_idx == 0) {
if (clip_l) {
size_t chunk_len_l = 77;
std::vector<int> chunk_tokens(clip_l_tokens.begin(),
clip_l_tokens.begin() + chunk_len_l);
std::vector<float> chunk_weights(clip_l_weights.begin(),
clip_l_weights.begin() + chunk_len_l);
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(chunk_tokens.size())}, chunk_tokens);
size_t max_token_idx = 0;
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), clip_l_tokenizer->EOS_TOKEN_ID);
max_token_idx = std::min<size_t>(std::distance(chunk_tokens.begin(), it), chunk_tokens.size() - 1);
pooled = clip_l->compute(n_threads,
input_ids,
0,
nullptr,
max_token_idx,
true,
clip_skip,
false);
if (pooled.empty()) {
LOG_ERROR("Flux CLIP-L encoding failed");
return {};
}
} else {
pooled = sd::Tensor<float>::zeros({768});
}
}
// t5
sd::Tensor<float> chunk_hidden_states;
if (t5) {
std::vector<int> chunk_tokens(t5_tokens.begin() + chunk_idx * chunk_len,
t5_tokens.begin() + (chunk_idx + 1) * chunk_len);
std::vector<float> chunk_weights(t5_weights.begin() + chunk_idx * chunk_len,
t5_weights.begin() + (chunk_idx + 1) * chunk_len);
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(chunk_tokens.size())}, chunk_tokens);
chunk_hidden_states = t5->compute(n_threads,
input_ids,
sd::Tensor<float>(),
false);
if (chunk_hidden_states.empty()) {
LOG_ERROR("Flux T5 encoding failed at chunk %d/%zu", chunk_idx + 1, chunk_count);
return {};
}
chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
if (zero_out_masked) {
chunk_hidden_states.fill_(0.0f);
}
} else {
chunk_hidden_states = sd::Tensor<float>::zeros({4096, static_cast<int64_t>(chunk_len)});
}
int64_t t1 = ggml_time_ms();
LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
if (!hidden_states.empty()) {
hidden_states = sd::ops::concat(hidden_states, chunk_hidden_states, 1);
} else {
hidden_states = std::move(chunk_hidden_states);
}
}
SDCondition result;
result.c_crossattn = std::move(hidden_states);
result.c_vector = std::move(pooled);
return result;
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, chunk_len);
if (tokens_and_weights.empty()) {
return {};
}
return get_learned_condition_common(n_threads,
tokens_and_weights,
conditioner_params.clip_skip,
conditioner_params.zero_out_masked);
}
};
struct T5CLIPEmbedder : public Conditioner {
T5UniGramTokenizer t5_tokenizer;
std::shared_ptr<T5Runner> t5;
size_t chunk_len = 512;
bool use_mask = false;
int mask_pad = 0;
bool is_umt5 = false;
T5CLIPEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
bool use_mask = false,
int mask_pad = 0,
bool is_umt5 = false,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const char* model_args = nullptr)
: use_mask(use_mask), mask_pad(mask_pad), t5_tokenizer(is_umt5) {
for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
if (key == "chroma_use_t5_mask") {
bool parsed = false;
if (parse_strict_bool(value, parsed)) {
this->use_mask = parsed;
} else {
LOG_WARN("ignoring invalid Chroma T5 model arg '%s=%s'", key.c_str(), value.c_str());
}
} else if (key == "chroma_t5_mask_pad") {
int parsed = 0;
if (parse_strict_int(value, parsed)) {
this->mask_pad = parsed;
} else {
LOG_WARN("ignoring invalid Chroma T5 model arg '%s=%s'", key.c_str(), value.c_str());
}
}
}
bool use_t5 = false;
for (auto pair : tensor_storage_map) {
if (pair.first.find("text_encoders.t5xxl") != std::string::npos) {
use_t5 = true;
}
}
if (!use_t5) {
LOG_WARN("IMPORTANT NOTICE: No text encoders provided, cannot process prompts!");
return;
} else {
t5 = std::make_shared<T5Runner>(backend, tensor_storage_map, "text_encoders.t5xxl.transformer", is_umt5, 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_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (t5) {
t5->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (t5) {
t5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (t5) {
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_flash_attention_enabled(bool enabled) override {
if (t5) {
t5->set_flash_attention_enabled(enabled);
}
}
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<WeightAdapter>& adapter) override {
if (t5) {
t5->set_weight_adapter(adapter);
}
}
void runner_end() override {
if (t5) {
t5->runner_end();
}
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
size_t min_length = 0,
size_t max_length = 0) {
auto parsed_attention = parse_prompt_attention(text);
{
std::stringstream ss;
ss << "[";
for (const auto& item : parsed_attention) {
ss << "['" << item.first << "', " << item.second << "], ";
}
ss << "]";
LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
}
auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
return false;
};
std::vector<int> t5_tokens;
std::vector<float> t5_weights;
std::vector<float> t5_mask;
if (t5) {
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
float curr_weight = item.second;
std::vector<int> curr_tokens;
if (!t5_tokenizer.encode(curr_text, curr_tokens)) {
return {};
}
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
}
t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, &t5_mask, min_length, max_length, true);
for (auto& mask_value : t5_mask) {
mask_value = mask_value > 0.0f ? 0.0f : -HUGE_VALF;
}
}
return {t5_tokens, t5_weights, t5_mask};
}
void modify_mask_to_attend_padding(sd::Tensor<float>* mask, int max_seq_length, int num_extra_padding = 8) {
GGML_ASSERT(mask != nullptr);
float* mask_data = mask->data();
int num_pad = 0;
for (int64_t i = 0; i < max_seq_length; i++) {
if (num_pad >= num_extra_padding) {
break;
}
if (std::isinf(mask_data[i])) {
mask_data[i] = 0;
++num_pad;
}
}
// LOG_VERBOSE("PAD: %d", num_pad);
}
SDCondition get_learned_condition_common(int n_threads,
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> token_and_weights,
int clip_skip,
bool zero_out_masked = false) {
if (!t5) {
SDCondition result;
result.c_crossattn = sd::Tensor<float>::zeros({4096, 256});
result.c_vector = sd::Tensor<float>::full({256}, -HUGE_VALF);
return result;
}
auto& t5_tokens = std::get<0>(token_and_weights);
auto& t5_weights = std::get<1>(token_and_weights);
auto& t5_attn_mask_vec = std::get<2>(token_and_weights);
int64_t t0 = ggml_time_ms();
sd::Tensor<float> t5_attn_mask = sd::Tensor<float>::from_vector(t5_attn_mask_vec);
sd::Tensor<float> hidden_states;
size_t chunk_count = t5_tokens.size() / chunk_len;
for (int chunk_idx = 0; chunk_idx < chunk_count; chunk_idx++) {
// t5
std::vector<int> chunk_tokens(t5_tokens.begin() + chunk_idx * chunk_len,
t5_tokens.begin() + (chunk_idx + 1) * chunk_len);
std::vector<float> chunk_weights(t5_weights.begin() + chunk_idx * chunk_len,
t5_weights.begin() + (chunk_idx + 1) * chunk_len);
std::vector<float> chunk_mask(t5_attn_mask_vec.begin() + chunk_idx * chunk_len,
t5_attn_mask_vec.begin() + (chunk_idx + 1) * chunk_len);
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(chunk_tokens.size())}, chunk_tokens);
sd::Tensor<float> t5_attn_mask_chunk;
if (use_mask) {
t5_attn_mask_chunk = sd::Tensor<float>({static_cast<int64_t>(chunk_mask.size())}, chunk_mask);
}
auto chunk_hidden_states = t5->compute(n_threads,
input_ids,
t5_attn_mask_chunk,
false);
GGML_ASSERT(!chunk_hidden_states.empty());
chunk_hidden_states = apply_token_weights(std::move(chunk_hidden_states), chunk_weights);
if (zero_out_masked) {
auto chunk_mask_tensor = sd::Tensor<float>::from_vector(chunk_mask)
.reshape_({1, static_cast<int64_t>(chunk_mask.size())});
chunk_hidden_states.masked_fill_(chunk_mask_tensor < 0.0f, 0.0f);
}
int64_t t1 = ggml_time_ms();
LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
if (!hidden_states.empty()) {
hidden_states = sd::ops::concat(hidden_states, chunk_hidden_states, 1);
} else {
hidden_states = std::move(chunk_hidden_states);
}
}
modify_mask_to_attend_padding(&t5_attn_mask, static_cast<int>(t5_attn_mask.numel()), mask_pad);
SDCondition result;
result.c_crossattn = std::move(hidden_states);
result.c_vector = std::move(t5_attn_mask);
return result;
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, chunk_len);
if (std::get<0>(tokens_and_weights).empty()) {
return {};
}
return get_learned_condition_common(n_threads,
tokens_and_weights,
conditioner_params.clip_skip,
conditioner_params.zero_out_masked);
}
};
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_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
if (t5) {
t5->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (t5) {
t5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (t5) {
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (t5) {
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_flash_attention_enabled(bool enabled) override {
if (t5) {
t5->set_flash_attention_enabled(enabled);
}
}
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<WeightAdapter>& adapter) override {
if (t5) {
t5->set_weight_adapter(adapter);
}
}
void runner_end() override {
if (t5) {
t5->runner_end();
}
}
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;
if (!tokenizer.encode(conditioner_params.text, tokens)) {
return {};
}
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);
GGML_ASSERT(!hidden_states.empty());
result.c_crossattn = std::move(hidden_states);
result.c_vector = sd::Tensor<float>::from_vector(mask);
return result;
}
};
struct SenseNovaU1Conditioner : public Conditioner {
static constexpr size_t kMaxPromptTokens = 12288;
SenseNovaU1Tokenizer tokenizer;
void get_param_tensors(std::map<std::string, ggml_tensor*>& 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 "
"<think></think> 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 ? "<img>" : "<think>\n\n</think>\n\n<img>";
return query;
}
SDCondition tokenize_condition(const std::string& text, bool is_negative) {
std::vector<int> tokens;
if (!tokenizer.encode(build_query(text, is_negative), tokens)) {
return {};
}
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<int32_t>({static_cast<int64_t>(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<Tokenizer> qwen_tokenizer;
T5UniGramTokenizer t5_tokenizer;
std::shared_ptr<LLM::LLMRunner> llm;
AnimaConditioner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {}) {
llm = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::QWEN3,
backend,
tensor_storage_map,
"text_encoders.llm",
false,
weight_manager);
qwen_tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 151643);
if (!qwen_tokenizer) {
qwen_tokenizer = std::make_shared<Qwen2Tokenizer>();
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
llm->set_max_graph_vram_bytes(max_vram_bytes);
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
llm->set_graph_cut_layer_split_enabled(enabled);
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
void set_flash_attention_enabled(bool enabled) override {
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<WeightAdapter>& adapter) override {
llm->set_weight_adapter(adapter);
}
void runner_end() override {
llm->runner_end();
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<int>, std::vector<float>> tokenize(std::string text) {
auto parsed_attention = parse_prompt_attention(text);
{
std::stringstream ss;
ss << "[";
for (const auto& item : parsed_attention) {
ss << "['" << item.first << "', " << item.second << "], ";
}
ss << "]";
LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
}
std::vector<int> qwen_tokens;
std::vector<float> qwen_weights;
std::vector<int> t5_tokens;
std::vector<float> t5_weights;
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
std::vector<int> curr_tokens;
if (!qwen_tokenizer->tokenize(curr_text, curr_tokens, nullptr)) {
return {};
}
qwen_tokens.insert(qwen_tokens.end(), curr_tokens.begin(), curr_tokens.end());
// Anima uses uniform Qwen token weights.
qwen_weights.insert(qwen_weights.end(), curr_tokens.size(), 1.f);
}
if (qwen_tokens.empty()) {
qwen_tokens.push_back(151643); // qwen3 pad token
qwen_weights.push_back(1.f);
}
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
float curr_weight = item.second;
std::vector<int> curr_tokens;
if (!t5_tokenizer.encode(curr_text, curr_tokens)) {
return {};
}
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
}
t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr);
return {qwen_tokens, qwen_weights, t5_tokens, t5_weights};
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
int64_t t0 = ggml_time_ms();
auto tokenized = tokenize(conditioner_params.text);
auto& qwen_tokens = std::get<0>(tokenized);
auto& qwen_weights = std::get<1>(tokenized);
auto& t5_tokens = std::get<2>(tokenized);
auto& t5_weights = std::get<3>(tokenized);
if (qwen_tokens.empty()) {
return {};
}
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(qwen_tokens.size()), 1}, qwen_tokens);
auto hidden_states = llm->compute(n_threads,
input_ids,
sd::Tensor<float>(),
{},
{},
false,
false);
GGML_ASSERT(!hidden_states.empty());
hidden_states = apply_token_weights(std::move(hidden_states), qwen_weights);
auto t5_ids_tensor = sd::Tensor<int32_t>::from_vector(t5_tokens);
auto t5_weight_tensor = sd::Tensor<float>::from_vector(t5_weights);
int64_t t1 = ggml_time_ms();
LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
SDCondition result;
result.c_crossattn = std::move(hidden_states);
result.c_t5_ids = std::move(t5_ids_tensor);
result.c_t5_weights = std::move(t5_weight_tensor);
return result;
}
};
struct LLMEmbedder : public Conditioner {
SDVersion version;
std::shared_ptr<Tokenizer> tokenizer;
std::shared_ptr<LLM::LLMRunner> llm;
std::shared_ptr<T5Runner> byt5;
LLMEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
SDVersion version = VERSION_QWEN_IMAGE,
const std::string prefix = "",
bool enable_vision = false,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {})
: version(version) {
if (!tokenizers.has(TokenizerConfig::MAIN)) {
if (sd_version_is_lens(version)) {
throw std::runtime_error("Lens requires an external GPT-OSS tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer");
}
if (sd_version_is_pid(version)) {
throw std::runtime_error("PiD requires an external Gemma 2 tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer");
}
}
LLM::LLMArch arch = LLM::LLMArch::QWEN2_5_VL;
if (version == VERSION_FLUX2) {
arch = LLM::LLMArch::MISTRAL_SMALL_3_2;
} else if (sd_version_is_ernie_image(version)) {
arch = LLM::LLMArch::MINISTRAL_3_3B;
} else if (sd_version_is_lens(version)) {
arch = LLM::LLMArch::GPT_OSS_20B;
} else if (sd_version_is_pid(version)) {
arch = LLM::LLMArch::GEMMA2_2B;
} else if (version == VERSION_QWEN_IMAGE_2_1 ||
sd_version_is_lingbot_video(version) ||
sd_version_is_ideogram4(version) ||
sd_version_is_boogu_image(version) ||
sd_version_is_sefi_image(version) ||
sd_version_is_krea2(version) ||
sd_version_is_minimax_h3(version) ||
sd_version_is_mage_flow(version)) {
arch = LLM::LLMArch::QWEN3_VL;
} else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
arch = LLM::LLMArch::QWEN3;
}
llm = std::make_shared<LLM::LLMRunner>(arch,
backend,
tensor_storage_map,
"text_encoders.llm",
enable_vision,
weight_manager);
int pad_id = 151643;
if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
pad_id = 11;
} else if (arch == LLM::LLMArch::GPT_OSS_20B) {
pad_id = 199999;
} else if (arch == LLM::LLMArch::GEMMA2_2B) {
pad_id = 0;
}
tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, pad_id);
if (!tokenizer) {
if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
tokenizer = std::make_shared<MistralTokenizer>();
} else {
tokenizer = std::make_shared<Qwen2Tokenizer>();
}
}
if (sd_version_is_hunyuan_video(version)) {
const std::string byt5_prefix = "text_encoders.t5xxl.transformer";
for (const auto& [name, _] : tensor_storage_map) {
if (starts_with(name, byt5_prefix + ".")) {
byt5 = std::make_shared<T5Runner>(backend,
tensor_storage_map,
byt5_prefix,
false,
weight_manager);
break;
}
}
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
if (byt5) {
byt5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
llm->set_max_graph_vram_bytes(max_vram_bytes);
if (byt5) {
byt5->set_max_graph_vram_bytes(max_vram_bytes);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
if (byt5) {
byt5->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (llm) {
llm->set_graph_cut_layer_split_enabled(enabled);
}
if (byt5) {
byt5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (llm) {
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
}
if (byt5) {
byt5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
if (byt5) {
byt5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
if (byt5) {
byt5->set_flash_attention_enabled(enabled);
}
}
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<WeightAdapter>& adapter) override {
if (llm) {
llm->set_weight_adapter(adapter);
}
if (byt5) {
byt5->set_weight_adapter(adapter);
}
}
void runner_end() override {
if (llm) {
llm->runner_end();
}
if (byt5) {
byt5->runner_end();
}
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
const std::pair<int, int>& attn_range,
size_t min_length = 0,
size_t max_length = 100000000,
bool spell_quotes = false) {
std::vector<std::pair<std::string, float>> parsed_attention;
if (attn_range.first >= 0 && attn_range.second > 0) {
if (attn_range.first > 0) {
parsed_attention.emplace_back(text.substr(0, attn_range.first), 1.f);
}
if (attn_range.second - attn_range.first > 0) {
auto new_parsed_attention = parse_prompt_attention(text.substr(attn_range.first, attn_range.second - attn_range.first));
if (spell_quotes) {
new_parsed_attention = split_quotation_attention(new_parsed_attention);
}
parsed_attention.insert(parsed_attention.end(),
new_parsed_attention.begin(),
new_parsed_attention.end());
}
if (attn_range.second < text.size()) {
parsed_attention.emplace_back(text.substr(attn_range.second), 1.f);
}
} else {
parsed_attention.emplace_back(text, 1.f);
}
{
std::stringstream ss;
ss << "[";
for (const auto& item : parsed_attention) {
ss << "['" << item.first << "', " << item.second << "], ";
}
ss << "]";
LOG_VERBOSE("parse '%s' to %s", text.c_str(), ss.str().c_str());
}
std::vector<int> tokens;
std::vector<float> weights;
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
float curr_weight = item.second;
std::vector<int> curr_tokens;
if (!tokenizer->encode(curr_text, curr_tokens, nullptr)) {
return {};
}
tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
weights.insert(weights.end(), curr_tokens.size(), curr_weight);
}
std::vector<float> mask;
tokenizer->pad_tokens(tokens, &weights, &mask, min_length, max_length);
// for (int i = 0; i < tokens.size(); i++) {
// std::cout << tokens[i] << ":" << weights[i] << ", " << i << std::endl;
// }
// std::cout << std::endl;
return {tokens, weights, mask};
}
sd::Tensor<float> encode_prompt(int n_threads,
const std::string prompt,
const std::pair<int, int>& prompt_attn_range,
int min_length,
int hidden_states_min_length,
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
const std::set<int>& out_layers,
int prompt_template_encode_start_idx,
bool spell_quotes = false,
int max_length = 100000000,
const LLM::DeepStackImageEmbeds& deepstack_image_embeds = {},
const std::vector<LLM::ImageGrid>& image_grids = {}) {
auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes);
auto& tokens = std::get<0>(tokens_weights_mask);
auto& weights = std::get<1>(tokens_weights_mask);
auto& mask = std::get<2>(tokens_weights_mask);
if (tokens.empty()) {
return {};
}
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
sd::Tensor<float> attention_mask;
if (!mask.empty()) {
attention_mask = sd::Tensor<float>({static_cast<int64_t>(mask.size()), static_cast<int64_t>(mask.size())});
const float masked_attention_value = -std::numeric_limits<float>::max() / 4.0f;
for (size_t i1 = 0; i1 < mask.size(); ++i1) {
for (size_t i0 = 0; i0 < mask.size(); ++i0) {
float value = 0.0f;
if (mask[i0] == 0.0f) {
value += masked_attention_value;
}
if (i0 > i1) {
value += masked_attention_value;
}
attention_mask[static_cast<int64_t>(i0 + mask.size() * i1)] = value;
}
}
}
auto hidden_states = llm->compute(n_threads,
input_ids,
attention_mask,
image_embeds,
out_layers,
false,
false,
deepstack_image_embeds,
image_grids);
if (hidden_states.empty()) {
LOG_ERROR("LLM prompt encoding failed");
return {};
}
hidden_states = apply_token_weights(std::move(hidden_states), weights);
GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
int64_t zero_pad_len = 0;
if (hidden_states_min_length > 0) {
if (hidden_states.shape()[1] - prompt_template_encode_start_idx < hidden_states_min_length) {
zero_pad_len = hidden_states_min_length - hidden_states.shape()[1] + prompt_template_encode_start_idx;
}
}
sd::Tensor<float> new_hidden_states = sd::ops::slice(hidden_states,
1,
prompt_template_encode_start_idx,
hidden_states.shape()[1]);
if (zero_pad_len > 0) {
auto pad_shape = new_hidden_states.shape();
pad_shape[1] = zero_pad_len;
new_hidden_states = sd::ops::concat(new_hidden_states,
sd::Tensor<float>::zeros(std::move(pad_shape)),
1);
}
return new_hidden_states;
}
void resize_image_dims(int height, int width, int& h_bar, int& w_bar, int factor, int min_size, int max_size, RefImageResizeMode mode) {
if (min_size > 0 && min_size == max_size) {
if (mode == RefImageResizeMode::AREA) {
double beta = std::sqrt(static_cast<double>(min_size) / (static_cast<double>(height) * width));
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
} else if (mode == RefImageResizeMode::LONGEST_SIDE) {
int current_max_side = std::max(height, width);
double beta = static_cast<double>(min_size) / current_max_side;
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
}
return;
}
if (mode == RefImageResizeMode::AREA) {
double current_area = static_cast<double>(h_bar) * w_bar;
if (max_size > 0 && current_area > max_size) {
double beta = std::sqrt((static_cast<double>(height) * width) / static_cast<double>(max_size));
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
} else if (min_size > 0 && current_area < min_size) {
double beta = std::sqrt(static_cast<double>(min_size) / (static_cast<double>(height) * width));
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
}
} else if (mode == RefImageResizeMode::LONGEST_SIDE) {
int current_max_side = std::max(height, width);
if (max_size > 0 && current_max_side > max_size) {
double beta = static_cast<double>(max_size) / current_max_side;
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(height * beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(width * beta / factor)) * static_cast<int>(factor));
} else if (min_size > 0 && current_max_side < min_size) {
double beta = static_cast<double>(min_size) / current_max_side;
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
}
}
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
std::string prompt;
std::pair<int, int> prompt_attn_range;
std::vector<std::string> extra_prompts;
std::vector<std::pair<int, int>> extra_prompts_attn_range;
std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
LLM::DeepStackImageEmbeds deepstack_image_embeds;
std::vector<LLM::ImageGrid> image_grids;
int prompt_template_encode_start_idx = 34;
int min_length = 0; // pad tokens
int max_length = 100000000;
int hidden_states_min_length = 0; // zero pad hidden_states
bool spell_quotes = false;
std::set<int> out_layers;
int64_t t0 = ggml_time_ms();
RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode;
if (sd_version_is_minimax_h3(version)) {
prompt_template_encode_start_idx = 0;
out_layers = {50};
prompt_attn_range = {0, 0};
if (llm->enable_vision) {
const std::string placeholder = "<|image_pad|>";
const int patch_size = llm->config.vision.patch_size;
const int factor = patch_size * llm->config.vision.spatial_merge_size;
auto resize_for_vision = [&](const sd::Tensor<float>& image) {
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
int w_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
resize_image_dims(height,
width,
h_bar,
w_bar,
factor,
3136,
12845056,
RefImageResizeMode::AREA);
auto resized = sd::ops::interpolate(
image,
std::vector<int64_t>{w_bar, h_bar, image.shape()[2], image.shape()[3]});
for (int64_t i = 0; i < resized.numel(); ++i) {
resized[i] = std::clamp(resized[i], 0.f, 1.f) * 2.f - 1.f;
}
return resized;
};
auto add_vision_outputs = [&](std::vector<sd::Tensor<float>> image_outputs,
int grid_h,
int grid_w) {
GGML_ASSERT(image_outputs.size() == 4);
auto image_embed = std::move(image_outputs[0]);
prompt += "<|vision_start|>";
std::vector<int> prefix_tokens;
if (!tokenizer->encode(prompt, prefix_tokens, nullptr)) {
return false;
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
image_embeds.emplace_back(image_embed_idx, image_embed);
if (deepstack_image_embeds.empty()) {
deepstack_image_embeds.resize(image_outputs.size() - 1);
}
for (size_t layer = 0; layer < deepstack_image_embeds.size(); ++layer) {
deepstack_image_embeds[layer].emplace_back(image_embed_idx, std::move(image_outputs[layer + 1]));
}
image_grids.push_back({image_embed_idx,
static_cast<int>(image_embed.shape()[1]),
grid_h,
grid_w});
for (int64_t i = 0; i < image_embed.shape()[1]; ++i) {
prompt += placeholder;
}
prompt += "<|vision_end|>";
return true;
};
const auto* references = conditioner_params.minimax_h3_references;
if (references != nullptr && !references->empty()) {
int picture_index = 0;
int video_index = 0;
int audio_index = 0;
for (const auto& item : *references) {
if (item.kind == MiniMaxH3PresentationKind::AUDIO) {
prompt += "<Audio " + std::to_string(++audio_index) + ">: ";
continue;
}
if (item.kind == MiniMaxH3PresentationKind::IMAGE) {
GGML_ASSERT(item.frames.size() == 1);
auto resized = resize_for_vision(item.frames[0]);
prompt += "<Picture " + std::to_string(++picture_index) + ">: ";
if (!add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size)) {
return {};
}
continue;
}
GGML_ASSERT(!item.frames.empty());
prompt += "<Video " + std::to_string(++video_index) + ">: ";
for (size_t frame = 0; frame < item.frames.size(); frame += 2) {
size_t next = std::min(frame + 1, item.frames.size() - 1);
float t0 = frame < item.timestamps.size() ? item.timestamps[frame] : frame / 2.f;
float t1 = next < item.timestamps.size() ? item.timestamps[next] : next / 2.f;
std::ostringstream timestamp;
timestamp << '<' << std::fixed << std::setprecision(1) << (t0 + t1) * 0.5f << " seconds>";
prompt += timestamp.str();
auto first = resize_for_vision(item.frames[frame]);
auto second = resize_for_vision(item.frames[next]);
if (first.shape()[0] != second.shape()[0] || first.shape()[1] != second.shape()[1]) {
second = sd::ops::interpolate(second,
std::vector<int64_t>{first.shape()[0],
first.shape()[1],
second.shape()[2],
second.shape()[3]});
}
auto pair = sd::ops::concat(first.unsqueeze(2), second.unsqueeze(2), 2);
if (!add_vision_outputs(llm->encode_video_block_outputs(n_threads,
pair,
false),
static_cast<int>(first.shape()[1]) / patch_size,
static_cast<int>(first.shape()[0]) / patch_size)) {
return {};
}
}
}
} else if (conditioner_params.ref_images != nullptr) {
for (size_t i = 0; i < conditioner_params.ref_images->size(); ++i) {
auto resized = resize_for_vision((*conditioner_params.ref_images)[i]);
prompt += "<Picture " + std::to_string(i + 1) + ">: ";
if (!add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size)) {
return {};
}
}
}
}
prompt += conditioner_params.text;
} else if (sd_version_is_hunyuan_video(version)) {
prompt_template_encode_start_idx = 98;
out_layers = {26};
prompt =
"<|im_start|>system\nYou are a helpful assistant. Describe the video by detailing the following aspects:\n"
"1. The main content and theme of the video.\n"
"2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects.\n"
"3. Actions, events, behaviors temporal relationships, physical movement changes of the objects.\n"
"4. background environment, light, style and atmosphere.\n"
"5. camera angles, movements, and transitions used in the video.<|im_end|>\n"
"<|im_start|>user\n";
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n";
} else if (sd_version_is_lingbot_video(version)) {
const int pad_token = 151643;
const std::string prompt_prefix =
"<|im_start|>system\nGiven a user input that may include a text prompt alone, "
"a text prompt with an image reference, or a text prompt with a video reference "
"or a video reference alone, generate an \"Enhanced prompt\" that provides detailed "
"visual descriptions suitable for video generation. Evaluate the level of detail "
"in the user's input: if it is simple, enrich it by adding specifics about colors, "
"shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal "
"progression, and spatial relationships to create vivid, concrete, and temporally "
"coherent scenes to create vivid and concrete scenes. Please generate only the "
"enhanced description for the prompt below and avoid including any additional "
"commentary or evaluations:<|im_end|>\n<|im_start|>user\n";
std::vector<int> prefix_tokens;
if (!tokenizer->encode(prompt_prefix, prefix_tokens, nullptr)) {
return {};
}
prompt_template_encode_start_idx = 0;
for (int token : prefix_tokens) {
if (token != pad_token) {
prompt_template_encode_start_idx++;
}
}
LOG_VERBOSE("prompt_template_encode_start_idx %d", prompt_template_encode_start_idx);
prompt = prompt_prefix;
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
LOG_INFO("LingBotVideoI2VPipeline");
const std::string placeholder = "<|image_pad|>";
std::string img_prompt;
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels = static_cast<int>(4 * factor * factor);
} else {
min_pixels = static_cast<int>(2 * factor);
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels = static_cast<int>(16384 * factor * factor);
} else {
max_pixels = static_cast<int>(128 * factor);
}
}
int h_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(height) / factor) * factor));
int w_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(width) / factor) * factor));
if (std::max(height, width) > 200 * std::min(height, width)) {
LOG_WARN("LingBotVideo image aspect ratio is very large: %dx%d", width, height);
}
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_VERBOSE("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt + img_prompt + "<|vision_start|>";
std::vector<int> prefix_tokens;
if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) {
return {};
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
image_embeds.emplace_back(image_embed_idx, image_embed);
img_prompt += "<|vision_start|>";
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
for (int j = 0; j < num_image_tokens; j++) {
img_prompt += placeholder;
}
img_prompt += "<|vision_end|>";
}
prompt += img_prompt;
}
prompt += conditioner_params.text;
prompt_attn_range = {0, 0};
prompt += "<|im_end|>\n<|im_start|>assistant\n";
} else if (version == VERSION_QWEN_IMAGE_2_1) {
if (!llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
LOG_ERROR("Qwen Image 2.1 editing requires Qwen3-VL vision weights; provide --llm_vision or a combined encoder");
return {};
}
prompt = "<|im_start|>system\nComprehend and analyze the provided prompt.<|im_end|>\n";
std::vector<int> system_tokens;
if (!tokenizer->encode(prompt, system_tokens, nullptr)) {
return {};
}
prompt_template_encode_start_idx = static_cast<int>(system_tokens.size());
out_layers = {static_cast<int>(llm->config.num_layers)};
prompt += "<|im_start|>user\n";
if (llm->enable_vision && conditioner_params.ref_images != nullptr) {
for (size_t i = 0; i < conditioner_params.ref_images->size(); ++i) {
const auto& image = (*conditioner_params.ref_images)[i];
int64_t width = image.shape()[0];
int64_t height = image.shape()[1];
int64_t pixels = width * height;
if (width % 32 != 0 || height % 32 != 0) {
LOG_ERROR("Qwen Image 2.1 reference dimensions must be multiples of 32");
return {};
}
auto rgb = sd::Tensor<float>({width, height, 3, 1});
for (int64_t p = 0; p < pixels; ++p) {
float alpha = image.shape()[2] == 4 ? image[p + 3 * pixels] : 1.f;
for (int c = 0; c < 3; ++c) {
rgb[p + c * pixels] = 2.f * (image[p + c * pixels] * alpha + 1.f - alpha) - 1.f;
}
}
auto outputs = llm->encode_image_outputs(n_threads, rgb, false);
if (outputs.empty()) {
return {};
}
prompt += (i == 0 ? "" : " ") + std::string("<image") + std::to_string(i + 1) + "><|vision_start|>";
std::vector<int> prefix_tokens;
if (!tokenizer->encode(prompt, prefix_tokens, nullptr)) {
return {};
}
int index = static_cast<int>(prefix_tokens.size());
int count = static_cast<int>(outputs[0].shape()[1]);
image_embeds.emplace_back(index, std::move(outputs[0]));
if (deepstack_image_embeds.empty()) {
deepstack_image_embeds.resize(outputs.size() - 1);
}
for (size_t layer = 1; layer < outputs.size(); ++layer) {
deepstack_image_embeds[layer - 1].emplace_back(index, std::move(outputs[layer]));
}
image_grids.push_back({index, count,
static_cast<int>(height) / llm->config.vision.patch_size,
static_cast<int>(width) / llm->config.vision.patch_size});
for (int j = 0; j < count; ++j) {
prompt += "<|image_pad|>";
}
prompt += "<|vision_end|>";
}
}
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text.empty() ? " " : conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n";
} else if (sd_version_is_qwen_image(version) || sd_version_is_mage_flow(version)) {
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
LOG_INFO("%s", sd_version_is_mage_flow(version) ? "MageFlowEditPipeline" : "QwenImageEditPlusPipeline");
prompt_template_encode_start_idx = 64;
int image_embed_idx = 64 + 6;
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
min_pixels = sd_version_is_mage_flow(version) ? -1 : 384;
if (min_pixels > 0 && resize_mode == RefImageResizeMode::AREA) {
min_pixels *= min_pixels;
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
max_pixels = sd_version_is_mage_flow(version) ? 384 : 560;
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels *= max_pixels;
}
}
std::string placeholder = "<|image_pad|>";
std::string img_prompt;
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
int w_bar = static_cast<int>(std::round(static_cast<double>(width) / factor) * factor);
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
image_embeds.emplace_back(image_embed_idx, image_embed);
image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
img_prompt += (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>";
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(num_image_tokens * placeholder.size());
for (int j = 0; j < num_image_tokens; j++) {
img_prompt += placeholder;
}
img_prompt += "<|vision_end|>";
}
prompt = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n";
prompt += img_prompt;
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n";
} else {
prompt_template_encode_start_idx = 34;
prompt = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n";
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n";
}
if (sd_version_is_mage_flow(version)) {
max_length = 2048 + prompt_template_encode_start_idx;
}
} else if (sd_version_is_boogu_image(version)) {
prompt_template_encode_start_idx = 0;
const std::string t2i_system_prompt =
"You are a helpful assistant that generates high-quality images based on user instructions. The instructions are as follows.";
const std::string edit_system_prompt =
"Describe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.";
const bool has_ref_images = llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty();
const bool text_empty = conditioner_params.text.find_first_not_of(" \t\r\n") == std::string::npos;
if (has_ref_images) {
LOG_INFO("BooguImageEditPipeline");
const std::string prompt_prefix = "<|im_start|>system\n" + edit_system_prompt + "<|im_end|>\n<|im_start|>user\n";
std::string img_prompt;
const std::string placeholder = "<|image_pad|>";
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
min_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels *= min_pixels;
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
max_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels *= max_pixels;
}
}
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = std::max(factor,
static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
int w_bar = std::max(factor,
static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt_prefix + img_prompt + "<|vision_start|>";
std::vector<int> prefix_tokens;
if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) {
return {};
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
image_embeds.emplace_back(image_embed_idx, image_embed);
img_prompt += "<|vision_start|>";
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
for (int j = 0; j < num_image_tokens; j++) {
img_prompt += placeholder;
}
img_prompt += "<|vision_end|>";
}
prompt = prompt_prefix + img_prompt;
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n";
} else {
const std::string& system_prompt = text_empty ? edit_system_prompt : t2i_system_prompt;
prompt = "<|im_start|>system\n" + system_prompt + "<|im_end|>\n<|im_start|>user\n";
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n";
}
} else if (sd_version_is_krea2(version)) {
prompt_template_encode_start_idx = 34;
out_layers = {2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35};
prompt = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n";
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
std::string img_prompt = "";
const std::string placeholder = "<|image_pad|>";
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
min_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels *= min_pixels;
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
max_pixels = 1024;
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels *= max_pixels;
}
}
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = std::max(factor,
static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
int w_bar = std::max(factor,
static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
std::vector<int> prefix_tokens;
if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) {
return {};
}
int image_embed_idx = static_cast<int>(prefix_tokens.size());
image_embeds.emplace_back(image_embed_idx, image_embed);
img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
for (int j = 0; j < num_image_tokens; j++) {
img_prompt += placeholder;
}
img_prompt += "<|vision_end|>";
}
prompt += img_prompt;
}
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n";
} else if (sd_version_is_longcat(version)) {
spell_quotes = true;
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
LOG_INFO("LongCatEditPipeline");
prompt_template_encode_start_idx = 67;
min_length = 512 + prompt_template_encode_start_idx;
int image_embed_idx = 36 + 6;
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
min_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels *= min_pixels;
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
max_pixels = 560;
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels *= max_pixels;
}
}
std::string placeholder = "<|image_pad|>";
std::string img_prompt;
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
int w_bar = static_cast<int>(std::round(static_cast<double>(width) / factor) * factor);
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_VERBOSE("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false);
GGML_ASSERT(!image_embed.empty());
image_embeds.emplace_back(image_embed_idx, image_embed);
image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
img_prompt += "<|vision_start|>";
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(num_image_tokens * placeholder.size());
for (int j = 0; j < num_image_tokens; j++) {
img_prompt += placeholder;
}
img_prompt += "<|vision_end|>";
}
prompt = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n";
prompt += img_prompt;
} else {
prompt_template_encode_start_idx = 36;
min_length = 512 + prompt_template_encode_start_idx;
prompt = "<|im_start|>system\nAs an image captioning expert, generate a descriptive text prompt based on an image content, suitable for input to a text-to-image model.<|im_end|>\n<|im_start|>user\n";
}
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n";
} else if (version == VERSION_FLUX2) {
prompt_template_encode_start_idx = 0;
hidden_states_min_length = 512;
out_layers = {10, 20, 30};
prompt = "[SYSTEM_PROMPT]You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object\nattribution and actions without speculation.[/SYSTEM_PROMPT][INST]";
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "[/INST]";
} else if (sd_version_is_ideogram4(version)) {
prompt_template_encode_start_idx = 0;
out_layers = {1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 36};
prompt = "<|im_start|>user\n";
prompt += conditioner_params.text;
prompt += "<|im_end|>\n<|im_start|>assistant\n";
prompt_attn_range = {0, 0};
} else if (sd_version_is_ernie_image(version)) {
prompt_template_encode_start_idx = 0;
out_layers = {25}; // -2
prompt_attn_range.first = 0;
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
} else if (sd_version_is_lens(version)) {
prompt_template_encode_start_idx = 97;
min_length = 0;
max_length = 512;
out_layers = {6, 12, 18, 24};
prompt =
"<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.\n"
"Knowledge cutoff: 2024-06\n"
"Current date: 2026-05-26\n" // fix for current date
"\n"
"Reasoning: medium\n"
"\n"
"# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>developer<|message|># Instructions\n"
"\n"
"Describe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background.\n"
"\n"
"<|end|><|start|>user<|message|>";
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|end|><|start|>assistant<|channel|>analysis<|message|>Need to generate one image according to the description.<|end|><|start|>assistant<|channel|>final<|message|>";
} else if (sd_version_is_z_image(version)) {
prompt_template_encode_start_idx = 0;
out_layers = {35}; // -2
if (conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
LOG_INFO("ZImageOmniPipeline");
prompt = "<|im_start|>user\n<|vision_start|>";
for (int i = 0; i < conditioner_params.ref_images->size() - 1; i++) {
extra_prompts.push_back("<|vision_end|><|vision_start|>");
}
extra_prompts.push_back("<|vision_end|>" + conditioner_params.text + "<|im_end|>\n<|im_start|>assistant\n<|vision_start|>");
extra_prompts.push_back("<|vision_end|><|im_end|>");
} else {
prompt = "<|im_start|>user\n";
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n";
}
} else if (version == VERSION_FLUX2_KLEIN) {
prompt_template_encode_start_idx = 0;
min_length = 512;
out_layers = {9, 18, 27};
prompt = "<|im_start|>user\n";
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
} else if (sd_version_is_sefi_image(version)) {
prompt_template_encode_start_idx = 0;
min_length = 1024;
out_layers = {9, 18, 27};
prompt = "<|im_start|>user\n";
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n";
} else if (version == VERSION_OVIS_IMAGE) {
prompt_template_encode_start_idx = 28;
min_length = prompt_template_encode_start_idx + 256;
prompt = "<|im_start|>user\nDescribe the image by detailing the color, quantity, text, shape, size, texture, spatial relationships of the objects and background:";
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += " " + conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
} else if (sd_version_is_pid(version)) {
constexpr int pixeldit_max_length = 300;
const std::string chi_prompt =
"Given a user prompt, generate an \"Enhanced prompt\" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:\n"
"- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.\n"
"- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.\n"
"Here are examples of how to transform or refine prompts:\n"
"- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.\n"
"- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.\n"
"Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:\n"
"User Prompt: ";
auto chi_tokens = std::get<0>(tokenize(chi_prompt, {0, 0}));
if (chi_tokens.empty()) {
return {};
}
size_t num_chi_tokens = chi_tokens.size();
max_length = (int)num_chi_tokens + pixeldit_max_length - 2;
min_length = max_length;
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += " " + conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
auto hidden_states = encode_prompt(n_threads,
prompt,
prompt_attn_range,
min_length,
0,
image_embeds,
out_layers,
0,
false,
max_length);
if (hidden_states.empty()) {
return {};
}
if (hidden_states.shape()[1] > pixeldit_max_length) {
auto bos = sd::ops::slice(hidden_states, 1, 0, 1);
auto tail = sd::ops::slice(hidden_states,
1,
hidden_states.shape()[1] - (pixeldit_max_length - 1),
hidden_states.shape()[1]);
hidden_states = sd::ops::concat(bos, tail, 1);
}
int64_t t1 = ggml_time_ms();
LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
SDCondition result;
result.c_crossattn = std::move(hidden_states);
return result;
} else {
GGML_ABORT("unknown version %d", version);
}
auto hidden_states = encode_prompt(n_threads,
prompt,
prompt_attn_range,
min_length,
hidden_states_min_length,
image_embeds,
out_layers,
prompt_template_encode_start_idx,
spell_quotes,
max_length,
deepstack_image_embeds,
image_grids);
if (hidden_states.empty()) {
return {};
}
std::vector<sd::Tensor<float>> extra_hidden_states_vec;
if (sd_version_is_hunyuan_video(version) && byt5) {
std::vector<std::string> quoted_texts;
auto collect_quoted = [&](const std::string& open, const std::string& close) {
size_t begin = 0;
while ((begin = conditioner_params.text.find(open, begin)) != std::string::npos) {
size_t content_begin = begin + open.size();
size_t end = conditioner_params.text.find(close, content_begin);
if (end == std::string::npos) {
break;
}
quoted_texts.push_back(conditioner_params.text.substr(content_begin, end - content_begin));
begin = end + close.size();
}
};
collect_quoted("\"", "\"");
collect_quoted("\xE2\x80\x98", "\xE2\x80\x99");
collect_quoted("\xE2\x80\x9C", "\xE2\x80\x9D");
if (!quoted_texts.empty()) {
std::string byt5_text;
for (const auto& text : quoted_texts) {
byt5_text += "Text \"" + text + "\". ";
}
std::vector<int> tokens;
tokens.reserve(byt5_text.size() + 1);
for (unsigned char byte : byt5_text) {
tokens.push_back(static_cast<int>(byte) + 3);
}
tokens.push_back(1);
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
auto byt5_hidden_states = byt5->compute(n_threads,
input_ids,
sd::Tensor<float>(),
false);
GGML_ASSERT(!byt5_hidden_states.empty());
extra_hidden_states_vec.push_back(std::move(byt5_hidden_states));
}
}
for (int i = 0; i < extra_prompts.size(); i++) {
auto extra_hidden_states = encode_prompt(n_threads,
extra_prompts[i],
extra_prompts_attn_range[i],
min_length,
hidden_states_min_length,
image_embeds,
out_layers,
prompt_template_encode_start_idx,
spell_quotes,
max_length);
if (extra_hidden_states.empty()) {
return {};
}
extra_hidden_states_vec.push_back(std::move(extra_hidden_states));
}
int64_t t1 = ggml_time_ms();
LOG_VERBOSE("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
SDCondition result;
result.c_crossattn = std::move(hidden_states);
result.extra_c_crossattns = std::move(extra_hidden_states_vec);
if (version == VERSION_QWEN_IMAGE_2_1) {
auto slots = sd::Tensor<int32_t>::zeros({result.c_crossattn.shape()[1]});
for (size_t i = 0; i < image_embeds.size(); ++i) {
int64_t begin = image_embeds[i].first - prompt_template_encode_start_idx;
int64_t end = begin + image_embeds[i].second.shape()[1];
if (begin < 0 || end > slots.numel()) {
LOG_ERROR("Qwen Image 2.1 image slots exceed the encoded prompt");
return {};
}
for (int64_t j = begin; j < end; ++j) {
slots[j] = static_cast<int32_t>(i + 1);
}
}
result.c_token_types = std::move(slots);
}
if (sd_version_is_minimax_h3(version)) {
std::vector<int32_t> tags(static_cast<size_t>(result.c_crossattn.shape()[1]), 1);
for (const auto& [index, image_embed] : image_embeds) {
int64_t begin = std::max<int64_t>(0, index - 1);
int64_t end = std::min<int64_t>(static_cast<int64_t>(tags.size()),
index + image_embed.shape()[1] + 1);
std::fill(tags.begin() + begin, tags.begin() + end, 0);
}
int64_t tag_count = static_cast<int64_t>(tags.size());
result.c_token_types = sd::Tensor<int32_t>({tag_count}, std::move(tags));
}
return result;
}
};
struct LTXAVTextProjection : public GGMLBlock {
static constexpr int64_t kHiddenSize = 3840;
static constexpr int64_t kNumStates = 49;
bool dual_projection = false;
LTXAVTextProjection(bool dual_projection = false)
: dual_projection(dual_projection) {
if (dual_projection) {
blocks["video_aggregate_embed"] = std::make_shared<Linear>(kHiddenSize * kNumStates, 4096, true);
blocks["audio_aggregate_embed"] = std::make_shared<Linear>(kHiddenSize * kNumStates, 2048, true);
} else {
blocks["projection"] = std::make_shared<Linear>(kHiddenSize * kNumStates, kHiddenSize, false);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
if (!dual_projection) {
auto projection = std::dynamic_pointer_cast<Linear>(blocks["projection"]);
return projection->forward(ctx, x);
}
auto video_projection = std::dynamic_pointer_cast<Linear>(blocks["video_aggregate_embed"]);
auto audio_projection = std::dynamic_pointer_cast<Linear>(blocks["audio_aggregate_embed"]);
auto video_in = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(4096.f / static_cast<float>(kHiddenSize)));
auto audio_in = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(2048.f / static_cast<float>(kHiddenSize)));
auto video = video_projection->forward(ctx, video_in);
auto audio = audio_projection->forward(ctx, audio_in);
return ggml_concat(ctx->ggml_ctx, video, audio, 0);
}
};
struct LTXAVTextProjectionRunner : public GGMLRunner {
LTXAVTextProjection model;
LTXAVTextProjectionRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string& prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager),
model(tensor_storage_map.find(prefix + ".video_aggregate_embed.weight") != tensor_storage_map.end()) {
model.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "ltxav_text_projection";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
model.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
auto x = make_input(x_tensor);
auto runner_ctx = get_context();
auto out = model.forward(&runner_ctx, x);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
bool auto_runner_end = true) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x);
};
return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_runner_end));
}
};
// LLaDA-Image's text path is a three-stage pipeline rather than a single encoder pass:
// the token embeddings feed a QueryFormer whose 256 queries are appended to the backbone
// input, and the backbone's final hidden states are projected to the denoiser's caption dim.
// Ref: LLaDAImagePipeline._encode_text.
struct LLaDAImageEmbedder : public Conditioner {
std::shared_ptr<Tokenizer> tokenizer;
std::shared_ptr<LLM::LLMRunner> llm;
std::shared_ptr<LLaDAImageTE::QueryFormerRunner> query_former;
std::shared_ptr<LLaDAImageTE::TextProjectionRunner> text_projection;
std::shared_ptr<LLaDAImageTE::SigVQRunner> sigvq;
std::string llm_prefix;
std::string query_former_prefix;
std::string text_projection_prefix;
std::string sigvq_prefix;
LLaDAImageEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string& llm_prefix = "text_encoders.llm",
const std::string& query_former_prefix = "queryformer",
const std::string& text_projection_prefix = "text_projection",
const std::string& sigvq_prefix = "sigvq",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {})
: llm_prefix(llm_prefix),
query_former_prefix(query_former_prefix),
text_projection_prefix(text_projection_prefix),
sigvq_prefix(sigvq_prefix) {
if (!tokenizers.has(TokenizerConfig::MAIN)) {
throw std::runtime_error("LLaDA-Image requires an external LLaDA2 tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer");
}
llm = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::LLADA2_MOE,
backend,
tensor_storage_map,
llm_prefix,
false,
weight_manager);
// <|endoftext|> doubles as the pad token in LLaDA2's tokenizer.json.
tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 156892);
query_former = std::make_shared<LLaDAImageTE::QueryFormerRunner>(backend,
tensor_storage_map,
query_former_prefix,
weight_manager);
text_projection = std::make_shared<LLaDAImageTE::TextProjectionRunner>(backend,
tensor_storage_map,
text_projection_prefix,
weight_manager);
// SigVQ is only present when the user supplies the editing weights.
for (const auto& [name, _] : tensor_storage_map) {
if (starts_with(name, sigvq_prefix + ".")) {
sigvq = std::make_shared<LLaDAImageTE::SigVQRunner>(backend,
tensor_storage_map,
sigvq_prefix,
weight_manager);
break;
}
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, llm_prefix);
query_former->get_param_tensors(tensors, query_former_prefix);
text_projection->get_param_tensors(tensors, text_projection_prefix);
if (sigvq != nullptr) {
sigvq->get_param_tensors(tensors, sigvq_prefix);
}
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
query_former->set_flash_attention_enabled(enabled);
text_projection->set_flash_attention_enabled(enabled);
if (sigvq != nullptr) {
sigvq->set_flash_attention_enabled(enabled);
}
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
llm->set_max_graph_vram_bytes(max_vram_bytes);
query_former->set_max_graph_vram_bytes(max_vram_bytes);
text_projection->set_max_graph_vram_bytes(max_vram_bytes);
if (sigvq != nullptr) {
sigvq->set_max_graph_vram_bytes(max_vram_bytes);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
llm->set_graph_cut_layer_split_enabled(enabled);
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, llm_prefix);
}
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
llm->set_weight_adapter(adapter);
query_former->set_weight_adapter(adapter);
text_projection->set_weight_adapter(adapter);
if (sigvq != nullptr) {
sigvq->set_weight_adapter(adapter);
}
}
void runner_end() override {
llm->runner_end();
query_former->runner_end();
text_projection->runner_end();
if (sigvq != nullptr) {
sigvq->runner_end();
}
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
const int64_t num_queries = 256;
const bool has_ref_images = conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty();
if (has_ref_images && sigvq == nullptr) {
LOG_ERROR("LLaDA-Image editing requires connectors with SigVQ weights");
return {};
}
std::string text = conditioner_params.text;
while (!text.empty() && std::isspace(static_cast<unsigned char>(text.front()))) {
text.erase(text.begin());
}
while (!text.empty() && std::isspace(static_cast<unsigned char>(text.back()))) {
text.pop_back();
}
std::string prompt = text.empty()
? "<role>HUMAN</role> Generate an image.\n<role>ASSISTANT</role>\n<IMAGE1>"
: "<role>HUMAN</role> Generate an image: " + text + "\n<role>ASSISTANT</role>\n<IMAGE1>";
std::vector<int> tokens;
if (!tokenizer->encode(prompt, tokens, nullptr)) {
return {};
}
int64_t n_text = static_cast<int64_t>(tokens.size());
GGML_ASSERT(n_text > 0);
sd::Tensor<int32_t> text_ids({n_text}, std::vector<int32_t>(tokens.begin(), tokens.end()));
auto inputs_embeds = llm->compute_input_embeds(n_threads, text_ids);
auto query_embeds = query_former->compute(n_threads, inputs_embeds);
// splice_image_embeds() replaces tokens in place, so the query slots have to exist in
// input_ids; their ids are irrelevant because the embeddings are overwritten.
std::vector<int32_t> padded(tokens.begin(), tokens.end());
padded.resize(static_cast<size_t>(n_text + num_queries), tokenizer->PAD_TOKEN_ID);
int64_t n_total = static_cast<int64_t>(padded.size());
sd::Tensor<int32_t> input_ids({n_total}, padded);
// Bidirectional everywhere except that the text tokens must not see the appended
// queries, matching backbone_attention_mask[:, :, :text_length, text_length:] = min.
const float mask_min = std::numeric_limits<float>::lowest() / 4.0f;
sd::Tensor<float> attention_mask({n_total, n_total});
for (int64_t i1 = 0; i1 < n_total; ++i1) {
for (int64_t i0 = 0; i0 < n_total; ++i0) {
float value = (i1 < n_text && i0 >= n_text) ? mask_min : 0.0f;
attention_mask[i0 + n_total * i1] = value;
}
}
LLM::ImageEmbeds image_embeds;
image_embeds.emplace_back(static_cast<int>(n_text), query_embeds);
std::set<int> out_layers = {static_cast<int>(llm->config.num_layers) + 1};
auto hidden_states = llm->compute(n_threads,
input_ids,
attention_mask,
image_embeds,
out_layers);
SDCondition result;
result.c_crossattn = text_projection->compute(n_threads, hidden_states);
// Editing: SigVQ sees the reference at half the output resolution, as in
// LLaDAImagePipeline._encode_source_image.
if (has_ref_images) {
const auto& ref = conditioner_params.ref_images->front();
auto resized = sd::ops::interpolate(ref,
{conditioner_params.width / 2,
conditioner_params.height / 2,
ref.shape()[2],
ref.shape()[3]},
sd::ops::InterpolateMode::Bilinear);
resized = resized * 2.f - 1.f;
auto semantic = sigvq->compute(n_threads, resized);
if (semantic.empty()) {
return {};
}
result.extra_c_crossattns.push_back(std::move(semantic));
}
return result;
}
};
struct LTXAVEmbedder : public Conditioner {
static constexpr int64_t kHiddenSize = 3840;
static constexpr int64_t kNumStates = 49;
static constexpr int64_t kMinLength = 1024;
std::shared_ptr<Tokenizer> tokenizer;
std::shared_ptr<LLM::LLMRunner> llm;
std::shared_ptr<LTXAVTextProjectionRunner> projector;
std::string projector_prefix;
bool dual_projection = false;
// Gemma 4 keeps a per-layer output scalar that no Gemma 3 checkpoint has, and widens its
// full-attention heads to 512 so their q_proj is twice a sliding layer's.
static LLM::LLMArch detect_gemma_arch(const String2TensorStorage& tensor_storage_map,
const std::string& llm_prefix) {
if (tensor_storage_map.find(llm_prefix + ".model.layers.0.layer_scalar") != tensor_storage_map.end()) {
return LLM::LLMArch::GEMMA4_12B;
}
auto global_q = tensor_storage_map.find(llm_prefix + ".model.layers.5.self_attn.q_proj.weight");
auto sliding_q = tensor_storage_map.find(llm_prefix + ".model.layers.0.self_attn.q_proj.weight");
if (global_q != tensor_storage_map.end() &&
sliding_q != tensor_storage_map.end() &&
global_q->second.ne[1] == sliding_q->second.ne[1] * 2) {
return LLM::LLMArch::GEMMA4_12B;
}
return LLM::LLMArch::GEMMA3_12B;
}
LTXAVEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string& llm_prefix = "text_encoders.llm",
const std::string& projector_prefix = "text_embedding_projection",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {})
: projector_prefix(projector_prefix) {
LLM::LLMArch arch = detect_gemma_arch(tensor_storage_map, llm_prefix);
LOG_INFO("ltxav text encoder: %s", arch == LLM::LLMArch::GEMMA4_12B ? "gemma 4" : "gemma 3");
llm = std::make_shared<LLM::LLMRunner>(arch,
backend,
tensor_storage_map,
llm_prefix,
false,
weight_manager);
tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 0, true);
if (!tokenizer) {
tokenizer = std::make_shared<GemmaTokenizer>();
}
dual_projection = tensor_storage_map.find(projector_prefix + ".video_aggregate_embed.weight") != tensor_storage_map.end();
projector = std::make_shared<LTXAVTextProjectionRunner>(backend,
tensor_storage_map,
projector_prefix,
weight_manager);
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
projector->get_param_tensors(tensors, projector_prefix);
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
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);
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
llm->set_graph_cut_layer_split_enabled(enabled);
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
llm->set_weight_adapter(adapter);
projector->set_weight_adapter(adapter);
}
void runner_end() override {
llm->runner_end();
projector->runner_end();
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
const std::pair<int, int>& attn_range) {
std::vector<std::pair<std::string, float>> parsed_attention;
if (attn_range.first >= 0 && attn_range.second > 0) {
if (attn_range.first > 0) {
parsed_attention.emplace_back(text.substr(0, attn_range.first), 1.f);
}
if (attn_range.second - attn_range.first > 0) {
auto new_parsed_attention = parse_prompt_attention(text.substr(attn_range.first, attn_range.second - attn_range.first));
parsed_attention.insert(parsed_attention.end(), new_parsed_attention.begin(), new_parsed_attention.end());
}
if (static_cast<size_t>(attn_range.second) < text.size()) {
parsed_attention.emplace_back(text.substr(attn_range.second), 1.f);
}
} else {
parsed_attention.emplace_back(text, 1.f);
}
std::vector<int> tokens;
std::vector<float> weights;
for (const auto& item : parsed_attention) {
std::vector<int> curr_tokens;
if (!tokenizer->encode(item.first, curr_tokens, nullptr)) {
return {};
}
tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
weights.insert(weights.end(), curr_tokens.size(), item.second);
}
std::vector<float> mask;
tokenizer->pad_tokens(tokens, &weights, &mask, kMinLength);
return {tokens, weights, mask};
}
sd::Tensor<float> encode_prompt(int n_threads,
const std::string& prompt,
const std::pair<int, int>& prompt_attn_range) {
auto tokens_weights_mask = tokenize(prompt, prompt_attn_range);
auto& tokens = std::get<0>(tokens_weights_mask);
auto& weights = std::get<1>(tokens_weights_mask);
auto& mask = std::get<2>(tokens_weights_mask);
if (tokens.empty()) {
return {};
}
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, std::vector<int32_t>(tokens.begin(), tokens.end()));
sd::Tensor<float> attention_mask;
if (!mask.empty()) {
const float mask_min = std::numeric_limits<float>::lowest() / 4.0f;
attention_mask = sd::Tensor<float>({static_cast<int64_t>(mask.size()), static_cast<int64_t>(mask.size())});
for (size_t i1 = 0; i1 < mask.size(); ++i1) {
for (size_t i0 = 0; i0 < mask.size(); ++i0) {
float value = 0.0f;
if (mask[i0] == 0.0f) {
value += mask_min;
}
if (i0 > i1) {
value += mask_min;
}
attention_mask[static_cast<int64_t>(i0 + mask.size() * i1)] = value;
}
}
}
auto hidden_states = llm->compute(n_threads,
input_ids,
attention_mask,
{},
{},
true,
false);
GGML_ASSERT(!hidden_states.empty());
hidden_states = apply_token_weights(std::move(hidden_states), weights);
int64_t valid_tokens = 0;
for (float value : mask) {
valid_tokens += static_cast<int64_t>(value > 0.0f);
}
GGML_ASSERT(valid_tokens > 0);
hidden_states = sd::ops::slice(hidden_states,
1,
hidden_states.shape()[1] - valid_tokens,
hidden_states.shape()[1]);
hidden_states.reshape_({kHiddenSize, kNumStates, valid_tokens});
hidden_states = hidden_states.permute({1, 0, 2});
if (dual_projection) {
for (int64_t state_idx = 0; state_idx < kNumStates; ++state_idx) {
for (int64_t token_idx = 0; token_idx < valid_tokens; ++token_idx) {
double sq_sum = 0.0;
for (int64_t hidden_idx = 0; hidden_idx < kHiddenSize; ++hidden_idx) {
float value = hidden_states.index(state_idx, hidden_idx, token_idx);
sq_sum += static_cast<double>(value) * static_cast<double>(value);
}
float inv_rms = 1.0f / std::sqrt(static_cast<float>(sq_sum / static_cast<double>(kHiddenSize)) + 1e-6f);
for (int64_t hidden_idx = 0; hidden_idx < kHiddenSize; ++hidden_idx) {
hidden_states.index(state_idx, hidden_idx, token_idx) *= inv_rms;
}
}
}
} else {
for (int64_t state_idx = 0; state_idx < kNumStates; ++state_idx) {
double sum = 0.0;
float min_value = std::numeric_limits<float>::infinity();
float max_value = -std::numeric_limits<float>::infinity();
for (int64_t token_idx = 0; token_idx < valid_tokens; ++token_idx) {
for (int64_t hidden_idx = 0; hidden_idx < kHiddenSize; ++hidden_idx) {
float value = hidden_states.index(state_idx, hidden_idx, token_idx);
sum += value;
min_value = std::min(min_value, value);
max_value = std::max(max_value, value);
}
}
float mean_value = static_cast<float>(sum / static_cast<double>(kHiddenSize * valid_tokens));
float denom = max_value - min_value + 1e-6f;
float scale_value = 8.0f / denom;
for (int64_t token_idx = 0; token_idx < valid_tokens; ++token_idx) {
for (int64_t hidden_idx = 0; hidden_idx < kHiddenSize; ++hidden_idx) {
float value = hidden_states.index(state_idx, hidden_idx, token_idx);
hidden_states.index(state_idx, hidden_idx, token_idx) = (value - mean_value) * scale_value;
}
}
}
}
hidden_states.reshape_({kNumStates * kHiddenSize, valid_tokens});
return projector->compute(n_threads, hidden_states, false);
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
int64_t t0 = ggml_time_ms();
std::string prompt;
std::pair<int, int> prompt_attn_range;
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
auto hidden_states = encode_prompt(n_threads, prompt, prompt_attn_range);
if (hidden_states.empty()) {
return {};
}
int64_t t1 = ggml_time_ms();
LOG_VERBOSE("computing LTXAV condition graph completed, taking %" PRId64 " ms", t1 - t0);
SDCondition result;
result.c_crossattn = std::move(hidden_states);
return result;
}
};
#endif // __SD_CONDITIONING_CONDITIONER_HPP__
Reference in New Issue View Git Blame Copy Permalink