Skip to content

Commit 516b85b

Browse files
rmatifrmatif
committed
add anima
1 parent c9cd497 commit 516b85b

8 files changed

Lines changed: 1037 additions & 12 deletions

File tree

‎src/anima.hpp‎

Lines changed: 774 additions & 0 deletions
Large diffs are not rendered by default.

‎src/conditioner.hpp‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1641,6 +1641,142 @@ struct T5CLIPEmbedder : public Conditioner {
16411641
}
16421642
};
16431643

1644+
struct AnimaConditioner : public Conditioner {
1645+
std::shared_ptr<LLM::BPETokenizer> qwen_tokenizer;
1646+
T5UniGramTokenizer t5_tokenizer;
1647+
std::shared_ptr<LLM::LLMRunner> llm;
1648+
1649+
AnimaConditioner(ggml_backend_t backend,
1650+
bool offload_params_to_cpu,
1651+
const String2TensorStorage& tensor_storage_map = {}) {
1652+
qwen_tokenizer = std::make_shared<LLM::Qwen2Tokenizer>();
1653+
llm = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::QWEN3,
1654+
backend,
1655+
offload_params_to_cpu,
1656+
tensor_storage_map,
1657+
"text_encoders.llm",
1658+
false);
1659+
}
1660+
1661+
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
1662+
llm->get_param_tensors(tensors, "text_encoders.llm");
1663+
}
1664+
1665+
void alloc_params_buffer() override {
1666+
llm->alloc_params_buffer();
1667+
}
1668+
1669+
void free_params_buffer() override {
1670+
llm->free_params_buffer();
1671+
}
1672+
1673+
size_t get_params_buffer_size() override {
1674+
return llm->get_params_buffer_size();
1675+
}
1676+
1677+
void set_flash_attention_enabled(bool enabled) override {
1678+
llm->set_flash_attention_enabled(enabled);
1679+
}
1680+
1681+
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
1682+
llm->set_weight_adapter(adapter);
1683+
}
1684+
1685+
std::tuple<std::vector<int>, std::vector<float>, std::vector<int>, std::vector<float>> tokenize(std::string text) {
1686+
auto parsed_attention = parse_prompt_attention(text);
1687+
1688+
{
1689+
std::stringstream ss;
1690+
ss << "[";
1691+
for (const auto& item : parsed_attention) {
1692+
ss << "['" << item.first << "', " << item.second << "], ";
1693+
}
1694+
ss << "]";
1695+
LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str());
1696+
}
1697+
1698+
std::vector<int> qwen_tokens;
1699+
std::vector<float> qwen_weights;
1700+
std::vector<int> t5_tokens;
1701+
std::vector<float> t5_weights;
1702+
1703+
for (const auto& item : parsed_attention) {
1704+
const std::string& curr_text = item.first;
1705+
std::vector<int> curr_tokens = qwen_tokenizer->tokenize(curr_text, nullptr);
1706+
qwen_tokens.insert(qwen_tokens.end(), curr_tokens.begin(), curr_tokens.end());
1707+
// Anima uses uniform Qwen token weights.
1708+
qwen_weights.insert(qwen_weights.end(), curr_tokens.size(), 1.f);
1709+
}
1710+
if (qwen_tokens.empty()) {
1711+
qwen_tokens.push_back(151643); // qwen3 pad token
1712+
qwen_weights.push_back(1.f);
1713+
}
1714+
1715+
for (const auto& item : parsed_attention) {
1716+
const std::string& curr_text = item.first;
1717+
float curr_weight = item.second;
1718+
std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
1719+
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
1720+
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
1721+
}
1722+
1723+
return {qwen_tokens, qwen_weights, t5_tokens, t5_weights};
1724+
}
1725+
1726+
SDCondition get_learned_condition(ggml_context* work_ctx,
1727+
int n_threads,
1728+
const ConditionerParams& conditioner_params) override {
1729+
int64_t t0 = ggml_time_ms();
1730+
1731+
auto tokenized = tokenize(conditioner_params.text);
1732+
auto& qwen_tokens = std::get<0>(tokenized);
1733+
auto& qwen_weights = std::get<1>(tokenized);
1734+
auto& t5_tokens = std::get<2>(tokenized);
1735+
auto& t5_weights = std::get<3>(tokenized);
1736+
1737+
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, qwen_tokens);
1738+
1739+
struct ggml_tensor* hidden_states = nullptr; // [N, n_token, 1024]
1740+
llm->compute(n_threads,
1741+
input_ids,
1742+
nullptr,
1743+
{},
1744+
{},
1745+
&hidden_states,
1746+
work_ctx);
1747+
1748+
{
1749+
auto tensor = hidden_states;
1750+
float original_mean = ggml_ext_tensor_mean(tensor);
1751+
for (int i2 = 0; i2 < tensor->ne[2]; i2++) {
1752+
for (int i1 = 0; i1 < tensor->ne[1]; i1++) {
1753+
for (int i0 = 0; i0 < tensor->ne[0]; i0++) {
1754+
float value = ggml_ext_tensor_get_f32(tensor, i0, i1, i2);
1755+
value *= qwen_weights[i1];
1756+
ggml_ext_tensor_set_f32(tensor, value, i0, i1, i2);
1757+
}
1758+
}
1759+
}
1760+
float new_mean = ggml_ext_tensor_mean(tensor);
1761+
if (new_mean != 0.f) {
1762+
ggml_ext_tensor_scale_inplace(tensor, (original_mean / new_mean));
1763+
}
1764+
}
1765+
1766+
struct ggml_tensor* t5_ids_tensor = nullptr;
1767+
struct ggml_tensor* t5_weight_tensor = nullptr;
1768+
if (!t5_tokens.empty()) {
1769+
t5_ids_tensor = vector_to_ggml_tensor_i32(work_ctx, t5_tokens);
1770+
t5_weight_tensor = vector_to_ggml_tensor(work_ctx, t5_weights);
1771+
}
1772+
1773+
int64_t t1 = ggml_time_ms();
1774+
LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
1775+
1776+
return {hidden_states, t5_weight_tensor, t5_ids_tensor};
1777+
}
1778+
};
1779+
16441780
struct LLMEmbedder : public Conditioner {
16451781
SDVersion version;
16461782
std::shared_ptr<LLM::BPETokenizer> tokenizer;

‎src/diffusion_model.hpp‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#ifndef __DIFFUSION_MODEL_H__
22
#define __DIFFUSION_MODEL_H__
33

4+
#include "anima.hpp"
45
#include "flux.hpp"
56
#include "mmdit.hpp"
67
#include "qwen_image.hpp"
@@ -242,6 +243,72 @@ struct FluxModel : public DiffusionModel {
242243
}
243244
};
244245

246+
struct AnimaModel : public DiffusionModel {
247+
std::string prefix;
248+
Anima::AnimaRunner anima;
249+
250+
AnimaModel(ggml_backend_t backend,
251+
bool offload_params_to_cpu,
252+
const String2TensorStorage& tensor_storage_map = {},
253+
const std::string prefix = "model.diffusion_model")
254+
: prefix(prefix), anima(backend, offload_params_to_cpu, tensor_storage_map, prefix) {
255+
}
256+
257+
std::string get_desc() override {
258+
return anima.get_desc();
259+
}
260+
261+
void alloc_params_buffer() override {
262+
anima.alloc_params_buffer();
263+
}
264+
265+
void free_params_buffer() override {
266+
anima.free_params_buffer();
267+
}
268+
269+
void free_compute_buffer() override {
270+
anima.free_compute_buffer();
271+
}
272+
273+
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
274+
anima.get_param_tensors(tensors, prefix);
275+
}
276+
277+
size_t get_params_buffer_size() override {
278+
return anima.get_params_buffer_size();
279+
}
280+
281+
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
282+
anima.set_weight_adapter(adapter);
283+
}
284+
285+
int64_t get_adm_in_channels() override {
286+
return 768;
287+
}
288+
289+
void set_flash_attention_enabled(bool enabled) {
290+
anima.set_flash_attention_enabled(enabled);
291+
}
292+
293+
void set_circular_axes(bool circular_x, bool circular_y) override {
294+
anima.set_circular_axes(circular_x, circular_y);
295+
}
296+
297+
bool compute(int n_threads,
298+
DiffusionParams diffusion_params,
299+
struct ggml_tensor** output = nullptr,
300+
struct ggml_context* output_ctx = nullptr) override {
301+
return anima.compute(n_threads,
302+
diffusion_params.x,
303+
diffusion_params.timesteps,
304+
diffusion_params.context,
305+
diffusion_params.c_concat,
306+
diffusion_params.y,
307+
output,
308+
output_ctx);
309+
}
310+
};
311+
245312
struct WanModel : public DiffusionModel {
246313
std::string prefix;
247314
WAN::WanRunner wan;

‎src/model.cpp‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1057,6 +1057,9 @@ SDVersion ModelLoader::get_sd_version() {
10571057
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
10581058
return VERSION_QWEN_IMAGE;
10591059
}
1060+
if (tensor_storage.name.find("model.diffusion_model.net.llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {
1061+
return VERSION_ANIMA;
1062+
}
10601063
if (tensor_storage.name.find("model.diffusion_model.double_stream_modulation_img.lin.weight") != std::string::npos) {
10611064
is_flux2 = true;
10621065
}

‎src/model.h‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ enum SDVersion {
4545
VERSION_WAN2_2_I2V,
4646
VERSION_WAN2_2_TI2V,
4747
VERSION_QWEN_IMAGE,
48+
VERSION_ANIMA,
4849
VERSION_FLUX2,
4950
VERSION_FLUX2_KLEIN,
5051
VERSION_Z_IMAGE,
@@ -122,6 +123,13 @@ static inline bool sd_version_is_qwen_image(SDVersion version) {
122123
return false;
123124
}
124125

126+
static inline bool sd_version_is_anima(SDVersion version) {
127+
if (version == VERSION_ANIMA) {
128+
return true;
129+
}
130+
return false;
131+
}
132+
125133
static inline bool sd_version_is_z_image(SDVersion version) {
126134
if (version == VERSION_Z_IMAGE) {
127135
return true;
@@ -146,6 +154,7 @@ static inline bool sd_version_is_dit(SDVersion version) {
146154
sd_version_is_sd3(version) ||
147155
sd_version_is_wan(version) ||
148156
sd_version_is_qwen_image(version) ||
157+
sd_version_is_anima(version) ||
149158
sd_version_is_z_image(version)) {
150159
return true;
151160
}

‎src/name_conversion.cpp‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1094,6 +1094,14 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
10941094
}
10951095
}
10961096

1097+
if (is_lora && sd_version_is_anima(version)) {
1098+
static const std::string anima_diffusion_prefix = "model.diffusion_model.";
1099+
static const std::string anima_net_prefix = "model.diffusion_model.net.";
1100+
if (starts_with(name, anima_diffusion_prefix) && !starts_with(name, anima_net_prefix)) {
1101+
name = anima_net_prefix + name.substr(anima_diffusion_prefix.size());
1102+
}
1103+
}
1104+
10971105
// cond_stage_model
10981106
{
10991107
for (const auto& prefix : cond_stage_model_prefix_vec) {

‎src/rope.hpp‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ namespace Rope {
4343

4444
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
4545
int dim,
46-
int theta,
46+
float theta,
4747
const std::vector<int>& axis_wrap_dims = {}) {
4848
assert(dim % 2 == 0);
4949
int half_dim = dim / 2;
@@ -167,7 +167,7 @@ namespace Rope {
167167

168168
__STATIC_INLINE__ std::vector<float> embed_nd(const std::vector<std::vector<float>>& ids,
169169
int bs,
170-
int theta,
170+
const std::vector<float>& axis_thetas,
171171
const std::vector<int>& axes_dim,
172172
const std::vector<std::vector<int>>& wrap_dims = {}) {
173173
std::vector<std::vector<float>> trans_ids = transpose(ids);
@@ -188,8 +188,12 @@ namespace Rope {
188188
if (!wrap_dims.empty() && i < (int)wrap_dims.size()) {
189189
axis_wrap_dims = wrap_dims[i];
190190
}
191+
float axis_theta = 10000.0f;
192+
if (!axis_thetas.empty()) {
193+
axis_theta = axis_thetas[std::min(i, axis_thetas.size() - 1)];
194+
}
191195
std::vector<std::vector<float>> rope_emb =
192-
rope(trans_ids[i], axes_dim[i], theta, axis_wrap_dims); // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
196+
rope(trans_ids[i], axes_dim[i], axis_theta, axis_wrap_dims); // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
193197
for (int b = 0; b < bs; ++b) {
194198
for (int j = 0; j < pos_len; ++j) {
195199
for (int k = 0; k < rope_emb[0].size(); ++k) {
@@ -203,6 +207,15 @@ namespace Rope {
203207
return flatten(emb);
204208
}
205209

210+
__STATIC_INLINE__ std::vector<float> embed_nd(const std::vector<std::vector<float>>& ids,
211+
int bs,
212+
float theta,
213+
const std::vector<int>& axes_dim,
214+
const std::vector<std::vector<int>>& wrap_dims = {}) {
215+
std::vector<float> axis_thetas(axes_dim.size(), theta);
216+
return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims);
217+
}
218+
206219
__STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
207220
int bs,
208221
int axes_dim_num,

0 commit comments

Comments
 (0)