upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-07-30 06:40:41 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: add hunyuan video 1.5 support (#1795)

Browse Source
This commit is contained in:
leejet
2026-07-18 22:30:10 +08:00
committed by GitHub
parent b290693977
commit ea4e566ccf
19 changed files with 1970 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

103
src/conditioning/conditioner.hpp
View File

@@ -1782,6 +1782,7 @@ struct LLMEmbedder : public Conditioner {
SDVersion version;
std::shared_ptr<BPETokenizer> tokenizer;
std::shared_ptr<LLM::LLMRunner> llm;
std::shared_ptr<T5Runner> byt5;
LLMEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
@@ -1819,54 +1820,97 @@ struct LLMEmbedder : public Conditioner {
"text_encoders.llm",
enable_vision,
weight_manager);
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 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_stream_layers_enabled(bool enabled) override {
llm->set_stream_layers_enabled(enabled);
if (byt5) {
byt5->set_stream_layers_enabled(enabled);
}
}
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_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_done() override {
if (llm) {
llm->runner_done();
}
if (byt5) {
byt5->runner_done();
}
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
@@ -2060,7 +2104,24 @@ struct LLMEmbedder : public Conditioner {
int64_t t0 = ggml_time_ms();
RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode;
if (sd_version_is_lingbot_video(version)) {
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, "
@@ -2590,6 +2651,46 @@ struct LLMEmbedder : public Conditioner {
spell_quotes,
max_length);
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,
true,
true);
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],

13
src/model.h
View File

@@ -38,6 +38,7 @@ enum SDVersion {
VERSION_LINGBOT_VIDEO,
VERSION_QWEN_IMAGE,
VERSION_QWEN_IMAGE_LAYERED,
VERSION_HUNYUAN_VIDEO,
VERSION_ANIMA,
VERSION_FLUX2,
VERSION_FLUX2_KLEIN,
@@ -142,6 +143,13 @@ static inline bool sd_version_is_qwen_image(SDVersion version) {
return false;
}
static inline bool sd_version_is_hunyuan_video(SDVersion version) {
if (version == VERSION_HUNYUAN_VIDEO) {
return true;
}
return false;
}
static inline bool sd_version_is_anima(SDVersion version) {
if (version == VERSION_ANIMA) {
return true;
@@ -240,6 +248,10 @@ static inline bool sd_version_uses_wan_vae(SDVersion version) {
return false;
}
static inline bool sd_version_uses_hunyuan_video_vae(SDVersion version) {
return sd_version_is_hunyuan_video(version);
}
static inline bool sd_version_is_inpaint(SDVersion version) {
if (version == VERSION_SD1_INPAINT ||
version == VERSION_SD2_INPAINT ||
@@ -259,6 +271,7 @@ static inline bool sd_version_is_dit(SDVersion version) {
sd_version_is_wan(version) ||
sd_version_is_lingbot_video(version) ||
sd_version_is_qwen_image(version) ||
sd_version_is_hunyuan_video(version) ||
version == VERSION_HIDREAM_O1 ||
sd_version_is_anima(version) ||
sd_version_is_z_image(version) ||

27
src/model/common/rope.hpp
View File

@@ -535,6 +535,33 @@ namespace Rope {
return vid_ids_repeated;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_hunyuan_video_ids(int t,
int h,
int w,
int patch_t,
int patch_h,
int patch_w,
int bs,
int context_len) {
std::vector<std::vector<float>> txt_ids(bs * context_len, std::vector<float>(3, 0.0f));
auto img_ids = gen_vid_ids(t, h, w, patch_t, patch_h, patch_w, bs);
return concat_ids(txt_ids, img_ids, bs);
}
__STATIC_INLINE__ std::vector<float> gen_hunyuan_video_pe(int t,
int h,
int w,
int patch_t,
int patch_h,
int patch_w,
int bs,
int context_len,
float theta,
const std::vector<int>& axes_dim) {
auto ids = gen_hunyuan_video_ids(t, h, w, patch_t, patch_h, patch_w, bs, context_len);
return embed_nd(ids, bs, theta, axes_dim);
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int t,
int h,
int w,

8
src/model/diffusion/flux.hpp
View File

@@ -706,11 +706,13 @@ namespace Flux {
LastLayer(int64_t hidden_size,
int64_t patch_size,
int64_t out_channels,
bool prune_mod = false,
bool bias = true)
bool prune_mod = false,
bool bias = true,
int64_t patch_volume = 0)
: prune_mod(prune_mod) {
blocks["norm_final"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size, 1e-06f, false));
blocks["linear"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, patch_size * patch_size * out_channels, bias));
int64_t out_dim = (patch_volume > 0 ? patch_volume : patch_size * patch_size) * out_channels;
blocks["linear"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, out_dim, bias));
if (!prune_mod) {
blocks["adaLN_modulation.1"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, 2 * hidden_size, bias));
}

679
src/model/diffusion/hunyuan.hpp Normal file
View File

@@ -0,0 +1,679 @@
#ifndef __SD_MODEL_DIFFUSION_HUNYUAN_HPP__
#define __SD_MODEL_DIFFUSION_HUNYUAN_HPP__
#include <memory>
#include "model/common/block.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/mmdit.hpp"
#include "model/diffusion/wan.hpp"
#include "model_manager.h"
namespace Hunyuan {
constexpr int HUNYUAN_VIDEO_GRAPH_SIZE = 65536;
// Ref: https://github.com/huggingface/diffusers/pull/12696
struct IndividualTokenRefinerBlock : public GGMLBlock {
protected:
int64_t num_heads;
public:
IndividualTokenRefinerBlock(int64_t num_heads,
int64_t head_dim,
int64_t mlp_ratio = 4,
bool attn_bias = true)
: num_heads(num_heads) {
int64_t hidden_size = num_heads * head_dim;
blocks["self_attn.qkv"] = std::make_shared<Linear>(hidden_size, hidden_size * 3, attn_bias);
blocks["self_attn.proj"] = std::make_shared<Linear>(hidden_size, hidden_size, attn_bias);
blocks["norm1"] = std::make_shared<LayerNorm>(hidden_size, 1e-6f, true);
blocks["norm2"] = std::make_shared<LayerNorm>(hidden_size, 1e-6f, true);
blocks["mlp.0"] = std::make_shared<Linear>(hidden_size, hidden_size * mlp_ratio);
blocks["mlp.2"] = std::make_shared<Linear>(hidden_size * mlp_ratio, hidden_size);
// adaLN_modulation.0 is nn.SiLU()
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(hidden_size, hidden_size * 2);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* t_emb, ggml_tensor* mask) {
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
auto self_attn_qkv = std::dynamic_pointer_cast<Linear>(blocks["self_attn.qkv"]);
auto self_attn_proj = std::dynamic_pointer_cast<Linear>(blocks["self_attn.proj"]);
auto mlp_fc1 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
auto mlp_fc2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
auto adaLN_modulation_1 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
// self attn
auto qkv = self_attn_qkv->forward(ctx, norm1->forward(ctx, txt));
auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
auto q = qkv_vec[0];
auto k = qkv_vec[1];
auto v = qkv_vec[2];
auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled);
attn_out = self_attn_proj->forward(ctx, attn_out);
// adaLN_modulation
auto emb = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, t_emb));
auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0);
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, attn_out, mods[0]));
// mlp
auto mlp_out = mlp_fc1->forward(ctx, norm2->forward(ctx, txt));
mlp_out = ggml_silu_inplace(ctx->ggml_ctx, mlp_out);
mlp_out = mlp_fc2->forward(ctx, mlp_out);
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, mlp_out, mods[1]));
return txt;
}
};
struct IndividualTokenRefiner : public GGMLBlock {
protected:
int num_layers;
public:
IndividualTokenRefiner(int64_t num_heads,
int64_t head_dim,
int num_layers,
int64_t mlp_ratio = 4,
bool attn_bias = true)
: num_layers(num_layers) {
for (int i = 0; i < num_layers; i++) {
blocks["blocks." + std::to_string(i)] = std::make_shared<IndividualTokenRefinerBlock>(num_heads, head_dim, mlp_ratio, attn_bias);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* t_emb, ggml_tensor* mask) {
for (int i = 0; i < num_layers; i++) {
auto block = std::dynamic_pointer_cast<IndividualTokenRefinerBlock>(blocks["blocks." + std::to_string(i)]);
txt = block->forward(ctx, txt, t_emb, mask);
}
return txt;
}
};
struct TokenRefiner : public GGMLBlock {
public:
TokenRefiner(int64_t in_channels,
int64_t num_heads,
int64_t head_dim,
int num_layers,
int64_t mlp_ratio = 4,
bool attn_bias = true) {
int64_t hidden_size = num_heads * head_dim;
blocks["input_embedder"] = std::make_shared<Linear>(in_channels, hidden_size);
blocks["t_embedder"] = std::make_shared<Flux::MLPEmbedder>(256, hidden_size);
blocks["c_embedder"] = std::make_shared<Flux::MLPEmbedder>(in_channels, hidden_size);
blocks["individual_token_refiner"] = std::make_shared<IndividualTokenRefiner>(num_heads, head_dim, num_layers, mlp_ratio, attn_bias);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* timestep, ggml_tensor* mask) {
auto input_embedder = std::dynamic_pointer_cast<Linear>(blocks["input_embedder"]);
auto t_embedder = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["t_embedder"]);
auto c_embedder = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["c_embedder"]);
auto individual_token_refiner = std::dynamic_pointer_cast<IndividualTokenRefiner>(blocks["individual_token_refiner"]);
auto t_emb = t_embedder->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f));
auto h = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, txt, 1, 0, 2, 3));
auto pooled_projections = ggml_scale(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, h), 1.f / txt->ne[1]);
pooled_projections = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, pooled_projections, 1, 0, 2, 3));
auto c_emb = c_embedder->forward(ctx, pooled_projections);
t_emb = ggml_add(ctx->ggml_ctx, t_emb, c_emb);
txt = input_embedder->forward(ctx, txt);
txt = individual_token_refiner->forward(ctx, txt, t_emb, mask);
return txt;
}
};
struct ByT5Mapper : public UnaryBlock {
ByT5Mapper(int64_t in_dim, int64_t hidden_size) {
blocks["layernorm"] = std::make_shared<LayerNorm>(in_dim);
blocks["fc1"] = std::make_shared<Linear>(in_dim, 2048);
blocks["fc2"] = std::make_shared<Linear>(2048, 2048);
blocks["fc3"] = std::make_shared<Linear>(2048, hidden_size);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto layernorm = std::dynamic_pointer_cast<LayerNorm>(blocks["layernorm"]);
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
auto fc3 = std::dynamic_pointer_cast<Linear>(blocks["fc3"]);
x = fc1->forward(ctx, layernorm->forward(ctx, x));
x = ggml_ext_gelu(ctx->ggml_ctx, x);
x = fc2->forward(ctx, x);
x = ggml_ext_gelu(ctx->ggml_ctx, x);
return fc3->forward(ctx, x);
}
};
struct HunyuanVideoConfig {
std::tuple<int, int, int> patch_size = {1, 2, 2};
int64_t in_channels = 65;
int64_t out_channels = 32;
int64_t hidden_size = 2048;
int64_t vec_in_dim = 0;
int64_t context_in_dim = 3584;
int64_t vision_in_dim = 0;
float mlp_ratio = 4.0f;
int num_heads = 16;
int depth = 54;
int depth_single_blocks = 0;
bool qkv_bias = true;
bool guidance_embed = false;
bool use_byt5 = false;
bool use_cond_type_embedding = false;
bool use_meanflow = false;
bool use_meanflow_sum = false;
float theta = 256;
std::vector<int> axes_dim = {16, 56, 56};
int axes_dim_sum = 128;
int64_t patch_volume() const {
return static_cast<int64_t>(std::get<0>(patch_size)) * std::get<1>(patch_size) * std::get<2>(patch_size);
}
static HunyuanVideoConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix) {
HunyuanVideoConfig config;
config.depth = 0;
config.depth_single_blocks = 0;
bool inferred = false;
int64_t img_embed_dim = 0;
for (const auto& [name, storage] : tensor_storage_map) {
if (starts_with(name, prefix) && ends_with(name, "img_in.proj.bias")) {
img_embed_dim = storage.ne[0];
break;
}
}
for (const auto& [name, storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
auto update_depth = [&](const char* block_prefix, int* depth) {
size_t pos = name.find(block_prefix);
if (pos == std::string::npos) {
return;
}
pos += strlen(block_prefix);
size_t end = name.find('.', pos);
if (end != std::string::npos) {
*depth = std::max(*depth, atoi(name.substr(pos, end - pos).c_str()) + 1);
}
};
update_depth("double_blocks.", &config.depth);
update_depth("single_blocks.", &config.depth_single_blocks);
if (ends_with(name, "img_in.proj.weight") && storage.n_dims == 5) {
config.patch_size = {static_cast<int>(storage.ne[2]),
static_cast<int>(storage.ne[1]),
static_cast<int>(storage.ne[0])};
config.in_channels = storage.ne[3];
config.hidden_size = storage.ne[4];
inferred = true;
} else if (ends_with(name, "img_in.proj.weight") && storage.n_dims == 4) {
config.patch_size = {static_cast<int>(storage.ne[2]),
static_cast<int>(storage.ne[1]),
static_cast<int>(storage.ne[0])};
if (img_embed_dim > 0 && storage.ne[3] % img_embed_dim == 0) {
config.hidden_size = img_embed_dim;
config.in_channels = storage.ne[3] / img_embed_dim;
}
inferred = true;
} else if (ends_with(name, "txt_in.input_embedder.weight")) {
config.context_in_dim = storage.ne[0];
inferred = true;
} else if (ends_with(name, "vector_in.in_layer.weight")) {
config.vec_in_dim = storage.ne[0];
} else if (ends_with(name, "vision_in.proj.0.weight")) {
config.vision_in_dim = storage.ne[0];
} else if (ends_with(name, "double_blocks.0.img_attn.norm.key_norm.scale") ||
ends_with(name, "double_blocks.0.img_attn.norm.key_norm.weight")) {
config.num_heads = static_cast<int>(config.hidden_size / storage.ne[0]);
} else if (ends_with(name, "double_blocks.0.img_mlp.0.weight")) {
config.mlp_ratio = static_cast<float>(storage.ne[1]) / static_cast<float>(storage.ne[0]);
}
config.guidance_embed = config.guidance_embed || name.find("guidance_in.") != std::string::npos;
config.use_byt5 = config.use_byt5 || name.find("byt5_in.") != std::string::npos;
config.use_meanflow = config.use_meanflow || name.find("time_r_in.") != std::string::npos;
}
config.use_cond_type_embedding = tensor_storage_map.find(prefix + ".cond_type_embedding.weight") != tensor_storage_map.end();
config.use_meanflow_sum = config.vision_in_dim > 0;
auto final_iter = tensor_storage_map.find(prefix + ".final_layer.linear.weight");
if (final_iter != tensor_storage_map.end()) {
config.out_channels = final_iter->second.ne[1] / config.patch_volume();
}
config.qkv_bias = tensor_storage_map.find(prefix + ".double_blocks.0.img_attn.qkv.bias") != tensor_storage_map.end();
GGML_ASSERT(config.hidden_size % config.num_heads == 0);
GGML_ASSERT(config.hidden_size / config.num_heads == config.axes_dim_sum);
if (inferred) {
LOG_DEBUG("hunyuan video: depth = %d, single depth = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", hidden_size = %" PRId64 ", context_in_dim = %" PRId64 ", patch_size = %dx%dx%d",
config.depth,
config.depth_single_blocks,
config.in_channels,
config.out_channels,
config.hidden_size,
config.context_in_dim,
std::get<0>(config.patch_size),
std::get<1>(config.patch_size),
std::get<2>(config.patch_size));
}
return config;
}
};
class HunyuanVideoModel : public GGMLBlock {
protected:
HunyuanVideoConfig config;
void init_params(struct ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
if (config.use_cond_type_embedding) {
ggml_type type = get_type(prefix + "cond_type_embedding.weight", tensor_storage_map, GGML_TYPE_F16);
GGMLBlock::params["cond_type_embedding.weight"] = ggml_new_tensor_2d(ctx, type, config.hidden_size, 3);
}
}
public:
HunyuanVideoModel() {}
explicit HunyuanVideoModel(HunyuanVideoConfig config)
: config(std::move(config)) {
int64_t head_dim = this->config.hidden_size / this->config.num_heads;
blocks["txt_in"] = std::make_shared<TokenRefiner>(this->config.context_in_dim, this->config.num_heads, head_dim, 2);
blocks["img_in"] = std::make_shared<PatchEmbed>(static_cast<int64_t>(224) /*Not used*/,
this->config.patch_size,
this->config.in_channels,
this->config.hidden_size);
blocks["time_in"] = std::make_shared<Flux::MLPEmbedder>(256, this->config.hidden_size);
if (this->config.vec_in_dim > 0) {
blocks["vector_in"] = std::make_shared<Flux::MLPEmbedder>(this->config.vec_in_dim, this->config.hidden_size);
}
if (this->config.vision_in_dim > 0) {
blocks["vision_in"] = std::make_shared<WAN::MLPProj>(this->config.vision_in_dim, this->config.hidden_size);
}
if (this->config.guidance_embed) {
blocks["guidance_in"] = std::make_shared<Flux::MLPEmbedder>(256, this->config.hidden_size);
}
if (this->config.use_byt5) {
blocks["byt5_in"] = std::make_shared<ByT5Mapper>(1472, this->config.hidden_size);
}
if (this->config.use_meanflow) {
blocks["time_r_in"] = std::make_shared<Flux::MLPEmbedder>(256, this->config.hidden_size);
}
for (int i = 0; i < this->config.depth; i++) {
blocks["double_blocks." + std::to_string(i)] = std::make_shared<Flux::DoubleStreamBlock>(this->config.hidden_size,
this->config.num_heads,
this->config.mlp_ratio,
i,
this->config.qkv_bias);
}
for (int i = 0; i < this->config.depth_single_blocks; i++) {
blocks["single_blocks." + std::to_string(i)] = std::make_shared<Flux::SingleStreamBlock>(this->config.hidden_size,
this->config.num_heads,
this->config.mlp_ratio,
i,
0.f);
}
blocks["final_layer"] = std::make_shared<Flux::LastLayer>(this->config.hidden_size,
std::get<2>(this->config.patch_size),
this->config.out_channels,
false,
true,
this->config.patch_volume());
}
ggml_tensor* pad_to_patch_size(struct ggml_context* ctx,
ggml_tensor* x) {
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t T = x->ne[2];
int pt = std::get<0>(config.patch_size);
int ph = std::get<1>(config.patch_size);
int pw = std::get<2>(config.patch_size);
int pad_t = (pt - static_cast<int>(T % pt)) % pt;
int pad_h = (ph - static_cast<int>(H % ph)) % ph;
int pad_w = (pw - static_cast<int>(W % pw)) % pw;
x = ggml_pad(ctx, x, pad_w, pad_h, pad_t, 0); // [N*C, T + pad_t, H + pad_h, W + pad_w]
return x;
}
ggml_tensor* unpatchify(struct ggml_context* ctx,
ggml_tensor* x,
int64_t t_len,
int64_t h_len,
int64_t w_len) {
// x: [N, t_len*h_len*w_len, C*pt*ph*pw]
// return: [N*C, t_len*pt, h_len*ph, w_len*pw]
int64_t N = x->ne[3];
int64_t pt = std::get<0>(config.patch_size);
int64_t ph = std::get<1>(config.patch_size);
int64_t pw = std::get<2>(config.patch_size);
int64_t C = x->ne[0] / pt / ph / pw;
GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, ph, pt, w_len*pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N); // [N*C*t_len, h_len*ph, pt, w_len*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, pt, h_len*ph, w_len*pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N); // [N*C, t_len*pt, h_len*ph, w_len*pw]
return x;
}
ggml_tensor* add_condition_type(GGMLRunnerContext* ctx, ggml_tensor* x, int type) {
if (!config.use_cond_type_embedding) {
return x;
}
auto weight = GGMLBlock::params["cond_type_embedding.weight"];
auto row = ggml_view_1d(ctx->ggml_ctx,
weight,
weight->ne[0],
static_cast<size_t>(type) * weight->nb[1]);
auto target = ggml_new_tensor_3d(ctx->ggml_ctx, row->type, config.hidden_size, x->ne[1], x->ne[2]);
auto embed = ggml_repeat(ctx->ggml_ctx, row, target);
embed = ggml_cast(ctx->ggml_ctx, embed, x->type);
return ggml_add(ctx->ggml_ctx, x, embed);
}
ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
ggml_tensor* img,
ggml_tensor* txt,
ggml_tensor* timestep,
ggml_tensor* pe,
ggml_tensor* guidance = nullptr,
ggml_tensor* y = nullptr,
ggml_tensor* txt_byt5 = nullptr,
ggml_tensor* clip_fea = nullptr,
ggml_tensor* timestep_r = nullptr,
int64_t N = 1) {
// img: [N*C, T, H, W], C => in_dim
// txt: [N, L, text_dim]
// timestep: [N,] or [T]
// return: [N, t_len*h_len*w_len, out_dim*pt*ph*pw]
GGML_ASSERT(N == 1);
auto img_in = std::dynamic_pointer_cast<PatchEmbed>(blocks["img_in"]);
auto txt_in = std::dynamic_pointer_cast<TokenRefiner>(blocks["txt_in"]);
auto time_in = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["time_in"]);
auto final_layer = std::dynamic_pointer_cast<Flux::LastLayer>(blocks["final_layer"]);
img = img_in->forward(ctx, img); // [N*C, t_len*h_len*w_len, hidden_size]
txt = txt_in->forward(ctx, txt, timestep, nullptr); // [N, n_txt_token, hidden_size]
auto vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f));
if (config.use_meanflow && timestep_r != nullptr) {
auto time_r_in = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["time_r_in"]);
auto vec_r = time_r_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_r, 256, 10000, 1000.f));
vec = ggml_add(ctx->ggml_ctx, vec, vec_r);
if (!config.use_meanflow_sum) {
vec = ggml_scale(ctx->ggml_ctx, vec, 0.5f);
}
}
if (config.vec_in_dim > 0 && y != nullptr) {
auto vector_in = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["vector_in"]);
vec = ggml_add(ctx->ggml_ctx, vec, vector_in->forward(ctx, y));
}
if (config.guidance_embed && guidance != nullptr) {
auto guidance_in = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["guidance_in"]);
auto guidance_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 256, 10000, 1.f);
vec = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, guidance_emb));
}
txt = add_condition_type(ctx, txt, 0);
if (config.use_byt5 && txt_byt5 != nullptr) {
auto byt5_in = std::dynamic_pointer_cast<ByT5Mapper>(blocks["byt5_in"]);
txt_byt5 = add_condition_type(ctx, byt5_in->forward(ctx, txt_byt5), 1);
txt = config.use_cond_type_embedding ? ggml_concat(ctx->ggml_ctx, txt_byt5, txt, 1)
: ggml_concat(ctx->ggml_ctx, txt, txt_byt5, 1);
}
if (config.vision_in_dim > 0 && clip_fea != nullptr) {
auto vision_in = std::dynamic_pointer_cast<WAN::MLPProj>(blocks["vision_in"]);
clip_fea = add_condition_type(ctx, vision_in->forward(ctx, clip_fea), 2);
txt = ggml_concat(ctx->ggml_ctx, clip_fea, txt, 1);
}
for (int i = 0; i < config.depth; i++) {
auto block = std::dynamic_pointer_cast<Flux::DoubleStreamBlock>(blocks["double_blocks." + std::to_string(i)]);
auto img_txt = block->forward(ctx, img, txt, vec, pe, nullptr);
img = img_txt.first; // [N, n_img_token, hidden_size]
txt = img_txt.second; // [N, n_txt_token, hidden_size]
}
if (config.depth_single_blocks > 0) {
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_img_token, hidden_size]
for (int i = 0; i < config.depth_single_blocks; i++) {
auto block = std::dynamic_pointer_cast<Flux::SingleStreamBlock>(blocks["single_blocks." + std::to_string(i)]);
txt_img = block->forward(ctx, txt_img, vec, pe, nullptr);
}
txt_img = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, txt_img, 0, 2, 1, 3));
img = ggml_view_3d(ctx->ggml_ctx,
txt_img,
txt_img->ne[0],
txt_img->ne[1],
img->ne[1],
txt_img->nb[1],
txt_img->nb[2],
txt_img->nb[2] * txt->ne[1]);
img = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, img, 0, 2, 1, 3));
}
img = final_layer->forward(ctx, img, vec); // (N, t_len*h_len*w_len, out_channels * patch_size ** 3)
return img;
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
ggml_tensor* guidance = nullptr,
ggml_tensor* y = nullptr,
ggml_tensor* txt_byt5 = nullptr,
ggml_tensor* clip_fea = nullptr,
ggml_tensor* timestep_r = nullptr,
int64_t N = 1) {
// Forward pass of DiT.
// x: [N*C, T, H, W]
// timestep: [N,]
// context: [N, L, D]
// pe: [L, d_head/2, 2, 2]
// return: [N*C, T, H, W]
GGML_ASSERT(N == 1);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t T = x->ne[2];
x = pad_to_patch_size(ctx->ggml_ctx, x);
int64_t pt = std::get<0>(config.patch_size);
int64_t ph = std::get<1>(config.patch_size);
int64_t pw = std::get<2>(config.patch_size);
int64_t t_len = (T + pt - 1) / pt;
int64_t h_len = (H + ph - 1) / ph;
int64_t w_len = (W + pw - 1) / pw;
auto out = forward_orig(ctx, x, context, timestep, pe, guidance, y, txt_byt5, clip_fea, timestep_r, N);
out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
// slice
out = ggml_ext_slice(ctx->ggml_ctx, out, 2, 0, T); // [N*C, T, H + pad_h, W + pad_w]
out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w]
out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W); // [N*C, T, H, W]
return out;
}
};
struct HunyuanVideoRunner : public DiffusionModelRunner {
public:
HunyuanVideoConfig config;
HunyuanVideoModel hunyuan_video;
std::vector<float> pe_vec;
SDVersion version;
HunyuanVideoRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
SDVersion version = VERSION_HUNYUAN_VIDEO,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(HunyuanVideoConfig::detect_from_weights(tensor_storage_map, prefix)),
version(version) {
LOG_INFO("HunyuanVideo blocks: %d double, %d single", config.depth, config.depth_single_blocks);
hunyuan_video = HunyuanVideoModel(config);
hunyuan_video.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "hunyuan_video";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
hunyuan_video.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const sd::Tensor<float>& c_concat_tensor = {},
const sd::Tensor<float>& y_tensor = {},
const sd::Tensor<float>& guidance_tensor = {},
const sd::Tensor<float>& byt5_tensor = {},
const sd::Tensor<float>& vision_tensor = {},
const sd::Tensor<float>& timestep_r_tensor = {}) {
ggml_cgraph* gf = new_graph_custom(HUNYUAN_VIDEO_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
ggml_tensor* context = make_input(context_tensor);
ggml_tensor* c_concat = make_optional_input(c_concat_tensor);
ggml_tensor* y = make_optional_input(y_tensor);
ggml_tensor* guidance = make_optional_input(guidance_tensor);
ggml_tensor* byt5 = make_optional_input(byt5_tensor);
ggml_tensor* vision = make_optional_input(vision_tensor);
ggml_tensor* timestep_r = make_optional_input(timestep_r_tensor);
GGML_ASSERT(x->ne[3] == config.out_channels);
if (c_concat != nullptr) {
x = ggml_concat(compute_ctx, x, c_concat, 3);
}
GGML_ASSERT(x->ne[3] <= config.in_channels);
if (x->ne[3] < config.in_channels) {
x = ggml_pad(compute_ctx, x, 0, 0, 0, static_cast<int>(config.in_channels - x->ne[3]));
}
int text_len = static_cast<int>(context->ne[1]);
if (byt5 != nullptr) {
text_len += static_cast<int>(byt5->ne[1]);
}
if (vision != nullptr) {
text_len += static_cast<int>(vision->ne[1]);
}
pe_vec = Rope::gen_hunyuan_video_pe(static_cast<int>(x->ne[2]),
static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
std::get<0>(config.patch_size),
std::get<1>(config.patch_size),
std::get<2>(config.patch_size),
1,
text_len,
config.theta,
config.axes_dim);
int64_t pos_len = static_cast<int64_t>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_DEBUG("pos_len %d", pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
// pe->data = pe_vec.data();
// print_ggml_tensor(pe, true, "pe");
// pe->data = nullptr;
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
ggml_tensor* out = hunyuan_video.forward(&runner_ctx,
x,
timesteps,
context,
pe,
guidance,
y,
byt5,
vision,
timestep_r);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
const sd::Tensor<float>& c_concat = {},
const sd::Tensor<float>& y = {},
const sd::Tensor<float>& guidance = {},
const sd::Tensor<float>& byt5 = {},
const sd::Tensor<float>& vision = {},
const sd::Tensor<float>& timestep_r = {}) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, c_concat, y, guidance, byt5, vision, timestep_r);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
GGML_ASSERT(diffusion_params.context != nullptr);
const auto* extra = diffusion_extra_as<HunyuanVideoDiffusionExtra>(diffusion_params);
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
*diffusion_params.context,
tensor_or_empty(diffusion_params.c_concat),
tensor_or_empty(diffusion_params.y),
tensor_or_empty(extra->guidance),
tensor_or_empty(extra->byt5),
tensor_or_empty(extra->vision),
tensor_or_empty(extra->timestep_r));
}
};
} // namespace Hunyuan
#endif // __SD_MODEL_DIFFUSION_HUNYUAN_HPP__

32
src/model/diffusion/ltxv.hpp
View File

@@ -800,7 +800,7 @@ namespace LTXV {
auto gate_mlp = mods[5];
auto x_norm = rms_norm(ctx->ggml_ctx, x);
x_norm = modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa);
x_norm = LTXV::modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa);
auto msa = attn1->forward(ctx, x_norm, nullptr, self_attention_mask, pe);
x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, msa, gate_msa));
@@ -810,12 +810,12 @@ namespace LTXV {
auto gate_q = mods[8];
auto q = rms_norm(ctx->ggml_ctx, x);
q = modulate(ctx->ggml_ctx, q, shift_q, scale_q);
q = LTXV::modulate(ctx->ggml_ctx, q, shift_q, scale_q);
auto context_mod = context;
if (prompt_timestep != nullptr) {
auto prompt_mods = get_prompt_scale_shift_values(ctx, prompt_timestep);
context_mod = modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]);
context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]);
}
auto mca = attn2->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
@@ -826,7 +826,7 @@ namespace LTXV {
}
auto y = rms_norm(ctx->ggml_ctx, x);
y = modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp);
y = LTXV::modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp);
auto mlp_out = ff->forward(ctx, y);
x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, mlp_out, gate_mlp));
return x;
@@ -1177,11 +1177,11 @@ namespace LTXV {
if (cross_attention_adaln) {
auto q_mods = get_ada_values(ctx, table, timestep, dim, 9, 6, 3);
auto q = rms_norm(ctx->ggml_ctx, x);
q = modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]);
q = LTXV::modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]);
auto context_mod = context;
if (prompt_timestep != nullptr && prompt_table != nullptr) {
auto p_mods = get_ada_values(ctx, prompt_table, prompt_timestep, dim, 2);
context_mod = modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]);
context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]);
}
auto out = attn->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
return apply_gate(ctx->ggml_ctx, out, q_mods[2]);
@@ -1228,7 +1228,7 @@ namespace LTXV {
auto v_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6);
auto v_norm = rms_norm(ctx->ggml_ctx, vx);
v_norm = modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]);
v_norm = LTXV::modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]);
auto v_sa = attn1->forward(ctx, v_norm, nullptr, self_attention_mask, v_pe);
vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_sa, v_mods[2]));
auto v_txt = apply_text_cross_attention(ctx,
@@ -1246,7 +1246,7 @@ namespace LTXV {
if (run_ax) {
auto a_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6);
auto a_norm = rms_norm(ctx->ggml_ctx, ax);
a_norm = modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]);
a_norm = LTXV::modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]);
auto a_sa = audio_attn1->forward(ctx, a_norm, nullptr, nullptr, a_pe);
ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_sa, a_mods[2]));
auto a_txt = apply_text_cross_attention(ctx,
@@ -1269,8 +1269,8 @@ namespace LTXV {
auto a2v_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
auto a2v_audio = get_ada_values(ctx, a2v_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
auto a2v_video = get_ada_values(ctx, a2v_video_table, v_cross_scale_shift_timestep, v_dim, 4);
auto vx_scaled = modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]);
auto ax_scaled = modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]);
auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
auto a2v_out = audio_to_video_attn->forward(ctx, vx_scaled, ax_scaled, nullptr, v_cross_pe, a_cross_pe);
auto a2v_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 4, 5);
auto a2v_gate = get_ada_values(ctx, a2v_gate_table, v_cross_gate_timestep, v_dim, 1)[0];
@@ -1282,8 +1282,8 @@ namespace LTXV {
auto v2a_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
auto v2a_audio = get_ada_values(ctx, v2a_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
auto v2a_video = get_ada_values(ctx, v2a_video_table, v_cross_scale_shift_timestep, v_dim, 4);
auto ax_scaled = modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
auto vx_scaled = modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]);
auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]);
auto v2a_out = video_to_audio_attn->forward(ctx, ax_scaled, vx_scaled, nullptr, a_cross_pe, v_cross_pe);
auto v2a_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 4, 5);
auto v2a_gate = get_ada_values(ctx, v2a_gate_table, a_cross_gate_timestep, a_dim, 1)[0];
@@ -1291,14 +1291,14 @@ namespace LTXV {
}
auto a_ff_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6, 3, 3);
auto ax_scaled = rms_norm(ctx->ggml_ctx, ax);
ax_scaled = modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]);
ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]);
auto a_ff_out = audio_ff->forward(ctx, ax_scaled);
ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_ff_out, a_ff_mods[2]));
}
auto v_ff_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6, 3, 3);
auto vx_scaled = rms_norm(ctx->ggml_ctx, vx);
vx_scaled = modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]);
vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]);
auto v_ff_out = ff->forward(ctx, vx_scaled);
vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_ff_out, v_ff_mods[2]));
@@ -1657,14 +1657,14 @@ namespace LTXV {
auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, config.hidden_size);
vx = norm_out->forward(ctx, vx);
vx = modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]);
vx = LTXV::modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]);
vx = proj_out->forward(ctx, vx);
vx = unpatchify_video(ctx, vx, width, height, frames);
if (ax != nullptr && audio_time > 0) {
auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, config.audio_hidden_size);
ax = audio_norm_out->forward(ctx, ax);
ax = modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]);
ax = LTXV::modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]);
ax = audio_proj_out->forward(ctx, ax);
ax = unpatchify_audio(ctx, ax, audio_time);
}

94
src/model/diffusion/mmdit.hpp
View File

@@ -136,11 +136,15 @@ struct MMDiTConfig {
};
struct PatchEmbed : public GGMLBlock {
// 2D Image to Patch Embedding
// 2D/3D Image to Patch Embedding
protected:
bool is_3d;
bool flatten;
bool dynamic_img_pad;
int patch_size;
int patch_t;
int patch_h;
int patch_w;
int64_t embed_dim;
public:
PatchEmbed(int64_t img_size = 224,
@@ -149,42 +153,90 @@ public:
int64_t embed_dim = 1536,
bool bias = true,
bool flatten = true,
bool dynamic_img_pad = true)
: patch_size(patch_size),
bool dynamic_img_pad = true,
bool is_3d = false)
: patch_t(is_3d ? patch_size : 1),
patch_h(patch_size),
patch_w(patch_size),
embed_dim(embed_dim),
flatten(flatten),
dynamic_img_pad(dynamic_img_pad) {
dynamic_img_pad(dynamic_img_pad),
is_3d(is_3d) {
// img_size is always None
// patch_size is always 2
// in_chans is always 16
// norm_layer is always False
// strict_img_size is always true, but not used
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_chans,
embed_dim,
{patch_size, patch_size},
{patch_size, patch_size},
{0, 0},
{1, 1},
bias));
if (is_3d) {
blocks["proj"] = std::make_shared<Conv3d>(in_chans,
embed_dim,
std::tuple{patch_size, patch_size, patch_size},
std::tuple{patch_size, patch_size, patch_size},
std::tuple{0, 0, 0},
std::tuple{1, 1, 1},
bias);
} else {
blocks["proj"] = std::make_shared<Conv2d>(in_chans,
embed_dim,
std::pair{patch_size, patch_size},
std::pair{patch_size, patch_size},
std::pair{0, 0},
std::pair{1, 1},
bias);
}
}
PatchEmbed(int64_t img_size,
std::tuple<int, int, int> patch_size,
int64_t in_chans,
int64_t embed_dim,
bool bias = true,
bool flatten = true,
bool dynamic_img_pad = true)
: patch_t(std::get<0>(patch_size)),
patch_h(std::get<1>(patch_size)),
patch_w(std::get<2>(patch_size)),
embed_dim(embed_dim),
flatten(flatten),
dynamic_img_pad(dynamic_img_pad),
is_3d(true) {
SD_UNUSED(img_size);
blocks["proj"] = std::make_shared<Conv3d>(in_chans,
embed_dim,
patch_size,
patch_size,
std::tuple{0, 0, 0},
std::tuple{1, 1, 1},
bias);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, C, H, W]
// return: [N, H*W, embed_dim]
auto proj = std::dynamic_pointer_cast<Conv2d>(blocks["proj"]);
// x: [N, C, H, W] or [N*C, T, H, W]
// return: [N, h_len*w_len, embed_dim] or [N, t_len*h_len*w_len, embed_dim]
auto proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["proj"]);
if (dynamic_img_pad) {
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int pad_h = (patch_size - H % patch_size) % patch_size;
int pad_w = (patch_size - W % patch_size) % patch_size;
x = ggml_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0); // TODO: reflect pad mode
int pad_t = 0;
int pad_h = (patch_h - static_cast<int>(H % patch_h)) % patch_h;
int pad_w = (patch_w - static_cast<int>(W % patch_w)) % patch_w;
if (is_3d) {
int64_t T = x->ne[2];
pad_t = (patch_t - static_cast<int>(T % patch_t)) % patch_t;
}
x = ggml_pad(ctx->ggml_ctx, x, pad_w, pad_h, pad_t, 0); // TODO: reflect pad mode
}
x = proj->forward(ctx, x);
x = proj->forward(ctx, x); // [N, C, h_len, w_len] or [N*C, t_len, h_len, w_len]
if (flatten) {
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
if (is_3d) {
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], embed_dim, x->ne[3] / embed_dim); // [N, C, t_len*h_len*w_len]
} else {
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]); // [N, C, h_len*w_len]
}
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [N, h_len*w_len, C]
}
return x;
}

10
src/model/diffusion/model.hpp
View File

@@ -87,6 +87,13 @@ struct MiniT2IDiffusionExtra {
const sd::Tensor<float>* mask = nullptr;
};
struct HunyuanVideoDiffusionExtra {
const sd::Tensor<float>* guidance = nullptr;
const sd::Tensor<float>* byt5 = nullptr;
const sd::Tensor<float>* vision = nullptr;
const sd::Tensor<float>* timestep_r = nullptr;
};
using DiffusionExtraParams = std::variant<std::monostate,
UNetDiffusionExtra,
SkipLayerDiffusionExtra,
@@ -95,7 +102,8 @@ using DiffusionExtraParams = std::variant<std::monostate,
WanDiffusionExtra,
HiDreamO1DiffusionExtra,
LTXAVDiffusionExtra,
MiniT2IDiffusionExtra>;
MiniT2IDiffusionExtra,
HunyuanVideoDiffusionExtra>;
struct DiffusionParams {
const sd::Tensor<float>* x = nullptr;

4
src/model/te/t5.hpp
View File

@@ -18,6 +18,7 @@
struct T5Config {
int64_t num_layers = 24;
int64_t model_dim = 4096;
int64_t inner_dim = 4096;
int64_t ff_dim = 10240;
int64_t num_heads = 64;
int64_t vocab_size = 32128;
@@ -53,6 +54,7 @@ struct T5Config {
if (q->n_dims == 2) {
config.model_dim = q->ne[0];
int64_t inner_dim = q->ne[1];
config.inner_dim = inner_dim;
// Flan-T5/T5 uses d_kv=64 for common sizes.
if (inner_dim % 64 == 0) {
config.num_heads = inner_dim / 64;
@@ -357,7 +359,7 @@ public:
: config(config) {
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new T5Stack(config.num_layers,
config.model_dim,
config.model_dim,
config.inner_dim,
config.ff_dim,
config.num_heads,
config.relative_attention));

834
src/model/vae/hunyuan_vae.hpp Normal file
View File

@@ -0,0 +1,834 @@
#ifndef __SD_MODEL_VAE_HUNYUAN_VAE_HPP__
#define __SD_MODEL_VAE_HUNYUAN_VAE_HPP__
#include <algorithm>
#include <cmath>
#include <map>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "model/vae/wan_vae.hpp"
#include "model_manager.h"
namespace Hunyuan {
constexpr int HUNYUAN_VIDEO_VAE_GRAPH_SIZE = 65536;
constexpr int HUNYUAN_VIDEO_VAE_GRAPH_SIZE_PER_LATENT_FRAME = 8192;
constexpr int HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE = 1;
struct TemporalConvCarry {
const std::vector<ggml_tensor*>* input = nullptr;
std::vector<ggml_tensor*>* output = nullptr;
size_t input_index = 0;
bool is_continuation() const {
return input != nullptr;
}
ggml_tensor* take() {
GGML_ASSERT(input != nullptr && input_index < input->size());
return (*input)[input_index++];
}
void push(ggml_tensor* tensor) {
if (output != nullptr) {
output->push_back(tensor);
}
}
void finish() const {
GGML_ASSERT(input == nullptr || input_index == input->size());
}
};
static ggml_tensor* repeat_interleave_channels(GGMLRunnerContext* ctx,
ggml_tensor* x,
int64_t repeats,
int64_t width,
int64_t height,
int64_t frames) {
GGML_ASSERT(repeats > 0);
GGML_ASSERT(width * height * frames == x->ne[0] * x->ne[1] * x->ne[2]);
int64_t channels = x->ne[3];
if (repeats == 1) {
return ggml_reshape_4d(ctx->ggml_ctx, x, width, height, frames, channels);
}
x = ggml_reshape_3d(ctx->ggml_ctx, x, width * height * frames, 1, channels);
auto target = ggml_new_tensor_3d(ctx->ggml_ctx, x->type, width * height * frames, repeats, channels);
x = ggml_repeat(ctx->ggml_ctx, x, target);
return ggml_reshape_4d(ctx->ggml_ctx, x, width, height, frames, channels * repeats);
}
class CausalConv3d : public GGMLBlock {
protected:
std::tuple<int, int, int> kernel_size;
public:
CausalConv3d(int64_t in_channels,
int64_t out_channels,
std::tuple<int, int, int> kernel_size,
std::tuple<int, int, int> stride = {1, 1, 1},
std::tuple<int, int, int> padding = {0, 0, 0},
std::tuple<int, int, int> dilation = {1, 1, 1},
bool bias = true)
: kernel_size(kernel_size) {
blocks["conv"] = std::make_shared<Conv3d>(in_channels, out_channels, kernel_size, stride, padding, dilation, bias);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
TemporalConvCarry* carry = nullptr) {
// x: [N*IC, ID, IH, IW]
// result: x: [N*OC, OD, OH, OW]
// assert N == 1
auto conv = std::dynamic_pointer_cast<Conv3d>(blocks["conv"]);
int pad_w = std::get<2>(kernel_size) / 2;
int pad_h = std::get<1>(kernel_size) / 2;
int pad_t = std::get<0>(kernel_size) - 1;
std::vector<ggml_tensor*> temporal_frames;
temporal_frames.reserve(x->ne[2] + pad_t);
if (pad_t > 0) {
if (carry != nullptr && carry->is_continuation()) {
auto previous = carry->take();
GGML_ASSERT(previous->ne[2] <= pad_t);
for (int64_t frame = 0; frame < previous->ne[2]; frame++) {
temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, previous, 2, frame, frame + 1));
}
for (int64_t frame = previous->ne[2]; frame < pad_t; frame++) {
temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1));
}
} else {
auto first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
for (int frame = 0; frame < pad_t; frame++) {
temporal_frames.push_back(first);
}
}
}
for (int64_t frame = 0; frame < x->ne[2]; frame++) {
temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + 1));
}
if (pad_t > 0 && carry != nullptr && carry->output != nullptr) {
ggml_tensor* next = nullptr;
for (int frame = pad_t; frame > 0; frame--) {
auto item = temporal_frames[temporal_frames.size() - frame];
next = next == nullptr ? item : ggml_concat(ctx->ggml_ctx, next, item, 2);
}
carry->push(ggml_cont(ctx->ggml_ctx, next));
}
ggml_tensor* padded = nullptr;
for (auto frame : temporal_frames) {
padded = padded == nullptr ? frame : ggml_concat(ctx->ggml_ctx, padded, frame, 2);
}
auto replicate_pad = [&](ggml_tensor* input, int dim, int left, int right) {
if (left > 0) {
auto first = ggml_ext_slice(ctx->ggml_ctx, input, dim, 0, 1);
for (int i = 0; i < left; i++) {
input = ggml_concat(ctx->ggml_ctx, first, input, dim);
}
}
if (right > 0) {
auto last = ggml_ext_slice(ctx->ggml_ctx, input, dim, input->ne[dim] - 1, input->ne[dim]);
for (int i = 0; i < right; i++) {
input = ggml_concat(ctx->ggml_ctx, input, last, dim);
}
}
return input;
};
padded = replicate_pad(padded, 0, pad_w, pad_w);
padded = replicate_pad(padded, 1, pad_h, pad_h);
return conv->forward(ctx, padded);
}
};
class AttnBlock : public UnaryBlock {
protected:
int64_t in_channels;
public:
AttnBlock(int64_t in_channels)
: in_channels(in_channels) {
blocks["norm"] = std::make_shared<WAN::RMS_norm>(in_channels);
blocks["q"] = std::make_shared<Conv3d>(in_channels, in_channels, std::tuple{1, 1, 1});
blocks["k"] = std::make_shared<Conv3d>(in_channels, in_channels, std::tuple{1, 1, 1});
blocks["v"] = std::make_shared<Conv3d>(in_channels, in_channels, std::tuple{1, 1, 1});
blocks["proj_out"] = std::make_shared<Conv3d>(in_channels, in_channels, std::tuple{1, 1, 1});
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x) override {
// x: [b*c, t, h, w]
auto norm = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm"]);
auto q_proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["q"]);
auto k_proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["k"]);
auto v_proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["v"]);
auto proj_out = std::dynamic_pointer_cast<UnaryBlock>(blocks["proj_out"]);
const int64_t b = x->ne[3] / in_channels;
auto identity = x;
x = norm->forward(ctx, x);
const int64_t c = x->ne[3] / b;
const int64_t t = x->ne[2];
const int64_t h = x->ne[1];
const int64_t w = x->ne[0];
auto q = q_proj->forward(ctx, x); // [b*c, t, h, w]
auto k = k_proj->forward(ctx, x); // [b*c, t, h, w]
auto v = v_proj->forward(ctx, x); // [b*c, t, h, w]
q = ggml_reshape_3d(ctx->ggml_ctx, q, w * h * t, c, b); // [b, c, t*h*w]
q = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 1, 0, 2, 3)); // [b, t*h*w, c]
k = ggml_reshape_3d(ctx->ggml_ctx, k, w * h * t, c, b); // [b, c, t*h*w]
k = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 1, 0, 2, 3)); // [b, t*h*w, c]
v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b); // [b, c, t*h*w]
v = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [b, t*h*w, c]
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [b, t*h*w, c]
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [b, c, t*h*w]
x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b); // [b*c, t, h, w]
x = proj_out->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, x, identity);
return x;
}
};
class ResnetBlock : public UnaryBlock {
protected:
int64_t in_channels;
int64_t out_channels;
public:
ResnetBlock(int64_t in_channels,
int64_t out_channels)
: in_channels(in_channels),
out_channels(out_channels) {
blocks["norm1"] = std::make_shared<WAN::RMS_norm>(in_channels);
blocks["conv1"] = std::make_shared<CausalConv3d>(in_channels, out_channels, std::tuple{3, 3, 3});
blocks["norm2"] = std::make_shared<WAN::RMS_norm>(out_channels);
blocks["conv2"] = std::make_shared<CausalConv3d>(out_channels, out_channels, std::tuple{3, 3, 3});
if (out_channels != in_channels) {
blocks["nin_shortcut"] = std::make_shared<CausalConv3d>(in_channels, out_channels, std::tuple{1, 1, 1});
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
return forward(ctx, x, nullptr);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
TemporalConvCarry* carry) {
// x: [B*IC, IT, OH, OW]
// return: [B*OC, OT, OH, OW]
auto norm1 = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm1"]);
auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
auto norm2 = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm2"]);
auto conv2 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv2"]);
auto h = x;
h = norm1->forward(ctx, h);
h = ggml_silu_inplace(ctx->ggml_ctx, h); // swish
h = conv1->forward(ctx, h, carry);
h = norm2->forward(ctx, h);
h = ggml_silu_inplace(ctx->ggml_ctx, h); // swish
// dropout, skip for inference
h = conv2->forward(ctx, h, carry);
// skip connection
if (out_channels != in_channels) {
auto nin_shortcut = std::dynamic_pointer_cast<CausalConv3d>(blocks["nin_shortcut"]);
x = nin_shortcut->forward(ctx, x); // [B*OC, OT, OH, OW]
}
h = ggml_add(ctx->ggml_ctx, h, x);
return h; // [B*OC, OT, OH, OW]
}
};
class Upsample : public GGMLBlock {
protected:
int64_t in_channels;
int64_t out_channels;
int64_t factor_t;
int64_t factor_s;
int64_t factor;
int64_t repeats;
public:
Upsample(int64_t in_channels, int64_t out_channels, bool add_temporal_upsample)
: in_channels(in_channels), out_channels(out_channels) {
if (add_temporal_upsample) {
factor_t = 2;
} else {
factor_t = 1;
}
factor_s = 2;
factor = factor_t * factor_s * factor_s;
GGML_ASSERT(out_channels * factor % in_channels == 0);
repeats = out_channels * factor / in_channels;
blocks["conv"] = std::make_shared<CausalConv3d>(in_channels, out_channels * factor, std::tuple{3, 3, 3});
}
static ggml_tensor* _pixel_shuffle_3d(GGMLRunnerContext* ctx,
ggml_tensor* x,
int64_t factor_t,
int64_t factor_s,
int64_t B = 1) {
// x: [B*factor*C, T, H, W]
// return: [B*C, T*factor_t, H*factor_s, W*factor_s]
GGML_ASSERT(B == 1);
int64_t factor = factor_t * factor_s * factor_s;
int64_t C = x->ne[3] / factor;
int64_t T = x->ne[2];
int64_t H = x->ne[1];
int64_t W = x->ne[0];
x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, C, factor); // [factor, C, T*H, W]
x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // [C, factor, T*H, W]
x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor_s, factor_s * factor_t * C); // [C*factor_t*factor_s, factor_s, T*H, W]
x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); // [C*factor_t*factor_s, T*H, W, factor_s]
x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, H * T, factor_s, factor_t * C); // [C*factor_t, factor_s, T*H, W*factor_s]
x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // [C*factor_t, T*H, factor_s, W*factor_s]
x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, T, factor_t, C); // [C, factor_t, T, H*factor_s*W*factor_s]
x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // [C, T, factor_t, H*factor_s*W*factor_s]
x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s * H, factor_t * T, C); // [C, T*factor_t, H*factor_s, W*factor_s]
return x;
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
TemporalConvCarry* carry = nullptr) {
// x: [B*IC, T, H, W]
// return: [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s]
const int64_t B = x->ne[3] / in_channels;
GGML_ASSERT(B == 1);
auto conv = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv"]);
const bool continuation = carry != nullptr && carry->is_continuation();
auto h = conv->forward(ctx, x, carry); // [B*factor*OC, T, H, W]
ggml_tensor* shortcut = nullptr;
if (factor_t == 2 && !continuation) {
auto h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1); // [B*factor*OC, 1, H, W]
h_first = _pixel_shuffle_3d(ctx, h_first, 1, factor_s, B); // [B*2*OC, 1, H*factor_s, W*factor_s]
h_first = ggml_ext_slice(ctx->ggml_ctx, h_first, 3, 0, out_channels); // [B*OC, 1, H*factor_s, W*factor_s]
auto x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
x_first = repeat_interleave_channels(ctx, x_first, repeats / 2, x->ne[0], x->ne[1], 1);
x_first = _pixel_shuffle_3d(ctx, x_first, 1, factor_s, B);
if (x->ne[2] == 1) {
return ggml_add(ctx->ggml_ctx, h_first, x_first);
}
auto h_next = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]); // [B*factor*OC, T - 1, H, W]
h_next = _pixel_shuffle_3d(ctx, h_next, factor_t, factor_s, B); // [B*OC, (T - 1)*factor_t, H*factor_s, W*factor_s]
h = ggml_concat(ctx->ggml_ctx, h_first, h_next, 2); // [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s]
auto x_next = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);
x_next = repeat_interleave_channels(ctx, x_next, repeats, x->ne[0], x->ne[1], x->ne[2] - 1);
x_next = _pixel_shuffle_3d(ctx, x_next, factor_t, factor_s, B);
shortcut = ggml_concat(ctx->ggml_ctx, x_first, x_next, 2); // [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s]
} else {
h = _pixel_shuffle_3d(ctx, h, factor_t, factor_s, B);
shortcut = repeat_interleave_channels(ctx, x, repeats, x->ne[0], x->ne[1], x->ne[2]);
shortcut = _pixel_shuffle_3d(ctx, shortcut, factor_t, factor_s, B); // [B*OC, T*factor_t, H*factor_s, W*factor_s]
}
return ggml_add(ctx->ggml_ctx, h, shortcut);
}
};
static ggml_tensor* pixel_unshuffle_3d(GGMLRunnerContext* ctx,
ggml_tensor* x,
int64_t factor_t,
int64_t factor_s) {
GGML_ASSERT(x->ne[0] % factor_s == 0);
GGML_ASSERT(x->ne[1] % factor_s == 0);
GGML_ASSERT(x->ne[2] % factor_t == 0);
int64_t W = x->ne[0] / factor_s;
int64_t H = x->ne[1] / factor_s;
int64_t T = x->ne[2] / factor_t;
int64_t C = x->ne[3];
int64_t factor = factor_t * factor_s * factor_s;
x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, factor_t, T, C);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s, H * T, factor_t * C);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s, W, H * T, factor_s * factor_t * C);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor, C);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, C * factor);
}
static ggml_tensor* mean_channel_groups(GGMLRunnerContext* ctx,
ggml_tensor* x,
int64_t group_size) {
GGML_ASSERT(group_size > 0);
GGML_ASSERT(x->ne[3] % group_size == 0);
if (group_size == 1) {
return x;
}
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t T = x->ne[2];
int64_t spatial = W * H * T;
int64_t groups = x->ne[3] / group_size;
x = ggml_reshape_3d(ctx->ggml_ctx, x, spatial, group_size, groups);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
x = ggml_sum_rows(ctx->ggml_ctx, x);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, groups);
return ggml_scale(ctx->ggml_ctx, x, 1.f / static_cast<float>(group_size));
}
class Downsample : public GGMLBlock {
protected:
int64_t in_channels;
int64_t out_channels;
int64_t factor_t;
int64_t factor_s = 2;
int64_t factor;
int64_t group_size;
public:
Downsample(int64_t in_channels, int64_t out_channels, bool add_temporal_downsample)
: in_channels(in_channels),
out_channels(out_channels),
factor_t(add_temporal_downsample ? 2 : 1),
factor(factor_t * factor_s * factor_s),
group_size(factor * in_channels / out_channels) {
GGML_ASSERT(out_channels % factor == 0);
GGML_ASSERT(factor * in_channels % out_channels == 0);
blocks["conv"] = std::make_shared<CausalConv3d>(in_channels,
out_channels / factor,
std::tuple{3, 3, 3});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto conv = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv"]);
auto h = conv->forward(ctx, x);
ggml_tensor* h_first = nullptr;
ggml_tensor* x_first = nullptr;
if (factor_t == 2) {
h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1);
h_first = pixel_unshuffle_3d(ctx, h_first, 1, factor_s);
h_first = ggml_concat(ctx->ggml_ctx, h_first, h_first, 3);
x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
x_first = pixel_unshuffle_3d(ctx, x_first, 1, factor_s);
x_first = mean_channel_groups(ctx, x_first, group_size / 2);
if (x->ne[2] == 1) {
return ggml_add(ctx->ggml_ctx, h_first, x_first);
}
h = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]);
x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);
}
GGML_ASSERT(h->ne[2] % factor_t == 0);
h = pixel_unshuffle_3d(ctx, h, factor_t, factor_s);
x = pixel_unshuffle_3d(ctx, x, factor_t, factor_s);
x = mean_channel_groups(ctx, x, group_size);
if (factor_t == 2) {
h = ggml_concat(ctx->ggml_ctx, h_first, h, 2);
x = ggml_concat(ctx->ggml_ctx, x_first, x, 2);
}
return ggml_add(ctx->ggml_ctx, h, x);
}
};
class MidBlock : public UnaryBlock {
protected:
int64_t in_channels;
int num_layers;
bool add_attention;
public:
MidBlock(int64_t in_channels,
int num_layers = 1,
bool add_attention = true)
: in_channels(in_channels),
num_layers(num_layers),
add_attention(add_attention) {
blocks["block_1"] = std::make_shared<ResnetBlock>(in_channels, in_channels);
for (int i = 0; i < num_layers; i++) {
if (add_attention) {
blocks["attn_" + std::to_string(i + 1)] = std::make_shared<AttnBlock>(in_channels);
}
blocks["block_" + std::to_string(i + 2)] = std::make_shared<ResnetBlock>(in_channels, in_channels);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
// x: [B*C, T, H, W]
// return: [B*C, T, H, W]
auto block_1 = std::dynamic_pointer_cast<ResnetBlock>(blocks["block_1"]);
x = block_1->forward(ctx, x);
for (int i = 0; i < num_layers; i++) {
if (add_attention) {
auto block = std::dynamic_pointer_cast<AttnBlock>(blocks["attn_" + std::to_string(i + 1)]);
x = block->forward(ctx, x);
}
auto block = std::dynamic_pointer_cast<ResnetBlock>(blocks["block_" + std::to_string(i + 2)]);
x = block->forward(ctx, x);
}
return x;
}
};
class UpBlock : public UnaryBlock {
protected:
int num_layers;
int64_t upsample_out_channels;
public:
UpBlock(int64_t in_channels,
int64_t out_channels,
int num_layers = 1,
int64_t upsample_out_channels = 0,
bool add_temporal_upsample = true)
: num_layers(num_layers),
upsample_out_channels(upsample_out_channels) {
for (int i = 0; i < num_layers; i++) {
int64_t IC = i == 0 ? in_channels : out_channels;
blocks["block." + std::to_string(i)] = std::make_shared<ResnetBlock>(IC, out_channels);
}
if (upsample_out_channels > 0) {
blocks["upsample"] = std::make_shared<Upsample>(out_channels, upsample_out_channels, add_temporal_upsample);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
return forward(ctx, x, nullptr);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
TemporalConvCarry* carry) {
// x: [B*IC, T, H, W]
// return: [B*OC, T, H, W] or [B*OC, T, H*2, W*2] or [B*OC, T*2, H*2, W*2]
for (int i = 0; i < num_layers; i++) {
auto block = std::dynamic_pointer_cast<ResnetBlock>(blocks["block." + std::to_string(i)]);
x = block->forward(ctx, x, carry);
}
if (upsample_out_channels > 0) {
auto upsample = std::dynamic_pointer_cast<Upsample>(blocks["upsample"]);
x = upsample->forward(ctx, x, carry);
}
return x;
}
};
class DownBlock : public UnaryBlock {
protected:
int num_layers;
int64_t downsample_out_channels;
public:
DownBlock(int64_t in_channels,
int64_t out_channels,
int num_layers,
int64_t downsample_out_channels = 0,
bool add_temporal_downsample = false)
: num_layers(num_layers),
downsample_out_channels(downsample_out_channels) {
for (int i = 0; i < num_layers; i++) {
int64_t IC = i == 0 ? in_channels : out_channels;
blocks["block." + std::to_string(i)] = std::make_shared<ResnetBlock>(IC, out_channels);
}
if (downsample_out_channels > 0) {
blocks["downsample"] = std::make_shared<Downsample>(out_channels,
downsample_out_channels,
add_temporal_downsample);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
for (int i = 0; i < num_layers; i++) {
auto block = std::dynamic_pointer_cast<ResnetBlock>(blocks["block." + std::to_string(i)]);
x = block->forward(ctx, x);
}
if (downsample_out_channels > 0) {
auto downsample = std::dynamic_pointer_cast<Downsample>(blocks["downsample"]);
x = downsample->forward(ctx, x);
}
return x;
}
};
class Encoder : public GGMLBlock {
protected:
int64_t z_channels;
std::vector<int64_t> block_out_channels;
public:
Encoder(int64_t in_channels = 3,
int64_t z_channels = 32,
std::vector<int64_t> block_out_channels = {128, 256, 512, 1024, 1024},
int layers_per_block = 2,
int spatial_compression_ratio = 16,
int temporal_compression_ratio = 4,
bool downsample_match_channel = true)
: z_channels(z_channels),
block_out_channels(std::move(block_out_channels)) {
blocks["conv_in"] = std::make_shared<CausalConv3d>(in_channels,
this->block_out_channels[0],
std::tuple{3, 3, 3});
int spatial_depth = static_cast<int>(std::log2(static_cast<double>(spatial_compression_ratio)));
int temporal_start = static_cast<int>(std::log2(static_cast<double>(spatial_compression_ratio / temporal_compression_ratio)));
int64_t channels = this->block_out_channels[0];
for (int i = 0; i < static_cast<int>(this->block_out_channels.size()); i++) {
int64_t out_channels = this->block_out_channels[i];
if (i < spatial_depth) {
int64_t next_channels = downsample_match_channel ? this->block_out_channels[i + 1] : out_channels;
blocks["down." + std::to_string(i)] = std::make_shared<DownBlock>(channels,
out_channels,
layers_per_block,
next_channels,
i >= temporal_start);
channels = next_channels;
} else {
blocks["down." + std::to_string(i)] = std::make_shared<DownBlock>(channels,
out_channels,
layers_per_block);
channels = out_channels;
}
}
blocks["mid"] = std::make_shared<MidBlock>(channels);
blocks["norm_out"] = std::make_shared<WAN::RMS_norm>(channels);
blocks["conv_out"] = std::make_shared<CausalConv3d>(channels,
z_channels * 2,
std::tuple{3, 3, 3});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto conv_in = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_in"]);
auto mid = std::dynamic_pointer_cast<MidBlock>(blocks["mid"]);
auto norm_out = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm_out"]);
auto conv_out = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_out"]);
x = conv_in->forward(ctx, x);
for (int i = 0; i < static_cast<int>(block_out_channels.size()); i++) {
auto down = std::dynamic_pointer_cast<DownBlock>(blocks["down." + std::to_string(i)]);
x = down->forward(ctx, x);
}
x = mid->forward(ctx, x);
auto shortcut = mean_channel_groups(ctx, x, x->ne[3] / (z_channels * 2));
x = norm_out->forward(ctx, x);
x = ggml_silu_inplace(ctx->ggml_ctx, x);
x = conv_out->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, x, shortcut);
return ggml_ext_slice(ctx->ggml_ctx, x, 3, 0, z_channels);
}
};
class Decoder : public GGMLBlock {
protected:
int64_t repeats;
std::vector<int64_t> block_out_channels;
public:
Decoder(int64_t in_channels = 32,
int64_t out_channels = 3,
std::vector<int64_t> block_out_channels = {1024, 1024, 512, 256, 128},
int layers_per_block = 2,
int spatial_compression_ratio = 16,
int temporal_compression_ratio = 4,
bool upsample_match_channel = true)
: block_out_channels(std::move(block_out_channels)) {
repeats = this->block_out_channels[0] / in_channels;
blocks["conv_in"] = std::make_shared<CausalConv3d>(in_channels, this->block_out_channels[0], std::tuple{3, 3, 3});
blocks["mid"] = std::make_shared<MidBlock>(this->block_out_channels[0]);
int64_t IC = this->block_out_channels[0];
for (int i = 0; i < this->block_out_channels.size(); i++) {
int64_t OC = this->block_out_channels[i];
bool add_spatial_upsample = i < std::log2(static_cast<double>(spatial_compression_ratio));
bool add_temporal_upsample = i < std::log2(static_cast<double>(temporal_compression_ratio));
if (add_spatial_upsample || add_temporal_upsample) {
int64_t upsample_out_channels = upsample_match_channel ? this->block_out_channels[i + 1] : OC;
blocks["up." + std::to_string(i)] = std::make_shared<UpBlock>(IC, OC, layers_per_block + 1, upsample_out_channels, add_temporal_upsample);
IC = upsample_out_channels;
} else {
blocks["up." + std::to_string(i)] = std::make_shared<UpBlock>(IC, OC, layers_per_block + 1, 0, false);
}
}
blocks["norm_out"] = std::make_shared<WAN::RMS_norm>(this->block_out_channels.back());
blocks["conv_out"] = std::make_shared<CausalConv3d>(this->block_out_channels.back(), out_channels, std::tuple{3, 3, 3});
}
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
auto conv_in = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_in"]);
auto mid_block = std::dynamic_pointer_cast<MidBlock>(blocks["mid"]);
auto norm_out = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm_out"]);
auto conv_out = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_out"]);
auto h = conv_in->forward(ctx, z);
auto shortcut = repeat_interleave_channels(ctx, z, repeats, z->ne[0], z->ne[1], z->ne[2]);
h = ggml_add(ctx->ggml_ctx, h, shortcut);
h = mid_block->forward(ctx, h);
ggml_tensor* output = nullptr;
std::vector<ggml_tensor*> carry_input;
const int64_t frames = h->ne[2];
for (int64_t start = 0; start < frames; start += HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE) {
const int64_t end = std::min(start + HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE, frames);
auto chunk = ggml_ext_slice(ctx->ggml_ctx, h, 2, start, end);
std::vector<ggml_tensor*> carry_output;
TemporalConvCarry carry{
start == 0 ? nullptr : &carry_input,
end == frames ? nullptr : &carry_output,
};
for (int i = 0; i < block_out_channels.size(); i++) {
auto up_block = std::dynamic_pointer_cast<UpBlock>(blocks["up." + std::to_string(i)]);
chunk = up_block->forward(ctx, chunk, &carry);
}
chunk = norm_out->forward(ctx, chunk);
chunk = ggml_silu_inplace(ctx->ggml_ctx, chunk); // nonlinearity/swish
chunk = conv_out->forward(ctx, chunk, &carry);
carry.finish();
output = output == nullptr ? chunk : ggml_concat(ctx->ggml_ctx, output, chunk, 2);
carry_input = std::move(carry_output);
}
return output;
}
};
class HunyuanVideoVAERunner : public VAE {
protected:
bool decode_only;
Encoder encoder;
Decoder decoder;
public:
HunyuanVideoVAERunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
bool decode_only,
SDVersion version,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: VAE(version, backend, prefix, weight_manager),
decode_only(decode_only ||
tensor_storage_map.find(prefix + ".encoder.conv_in.conv.weight") == tensor_storage_map.end()) {
if (!this->decode_only) {
encoder.init(params_ctx, tensor_storage_map, prefix + ".encoder");
}
decoder.init(params_ctx, tensor_storage_map, prefix + ".decoder");
}
std::string get_desc() override {
return "hunyuan_video_vae";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (!decode_only) {
encoder.get_param_tensors(tensors, weight_prefix + ".encoder");
}
decoder.get_param_tensors(tensors, weight_prefix + ".decoder");
}
int get_encoder_output_channels(int input_channels) override {
SD_UNUSED(input_channels);
return 32;
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output,
std::shared_ptr<RNG> rng) override {
SD_UNUSED(rng);
return vae_output;
}
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
return latents / 1.03682f;
}
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
return latents * 1.03682f;
}
ggml_cgraph* build_graph(const sd::Tensor<float>& input_tensor, bool decode_graph) {
size_t graph_size = HUNYUAN_VIDEO_VAE_GRAPH_SIZE;
if (decode_graph) {
graph_size = std::max(graph_size,
HUNYUAN_VIDEO_VAE_GRAPH_SIZE_PER_LATENT_FRAME *
static_cast<size_t>(input_tensor.shape()[2]));
}
ggml_cgraph* gf = new_graph_custom(graph_size);
ggml_tensor* input = make_input(input_tensor);
auto runner_ctx = get_context();
ggml_tensor* output = decode_graph ? decoder.forward(&runner_ctx, input)
: encoder.forward(&runner_ctx, input);
ggml_build_forward_expand(gf, output);
return gf;
}
sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& input,
bool decode_graph) override {
if (!decode_graph && decode_only) {
LOG_ERROR("Hunyuan Video VAE encoder weights are not available");
return {};
}
sd::Tensor<float> expanded;
if (input.dim() == 4) {
expanded = input.unsqueeze(2);
}
const auto& graph_input = expanded.empty() ? input : expanded;
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(graph_input, decode_graph);
};
auto output = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
n_threads,
true,
true,
true),
graph_input.dim());
if (!output.empty() && input.dim() == 4) {
output.squeeze_(2);
}
return output;
}
};
} // namespace Hunyuan
#endif // __SD_MODEL_VAE_HUNYUAN_VAE_HPP__

11
src/model/vae/tae.hpp
View File

@@ -528,6 +528,9 @@ public:
if (version == VERSION_WAN2_2_TI2V) {
z_channels = 48;
patch = 2;
} else if (sd_version_is_hunyuan_video(version)) {
z_channels = 32;
patch = 2;
} else if (sd_version_is_ltxav(version)) {
z_channels = 128;
patch = 4;
@@ -542,12 +545,12 @@ public:
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
auto decoder = std::dynamic_pointer_cast<TinyVideoDecoder>(blocks["decoder"]);
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
// (W, H, C, T) -> (W, H, T, C)
z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 1, 3, 2));
}
auto result = decoder->forward(ctx, z);
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
// (W, H, T, C) -> (W, H, C, T)
result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
}
@@ -556,7 +559,7 @@ public:
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
// (W, H, T, C) -> (W, H, C, T)
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
}
@@ -569,7 +572,7 @@ public:
}
}
x = encoder->forward(ctx, x);
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
// (W, H, C, T) -> (W, H, T, C)
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
}

4
src/model/vae/vae.hpp
View File

@@ -74,7 +74,7 @@ public:
int scale_factor = 8;
if (version == VERSION_LTXAV) {
scale_factor = 32;
} else if (version == VERSION_WAN2_2_TI2V) {
} else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version)) {
scale_factor = 16;
} else if (sd_version_uses_flux2_vae(version)) {
scale_factor = 16;
@@ -136,7 +136,7 @@ public:
// Image VAE encode is more sensitive to tile boundary context than decode.
// Keep the smaller legacy factor for video VAEs, but default image encode
// tiles to 64 latent pixels so a 512px SD image is encoded as one tile.
const float encode_tile_factor = (sd_version_is_wan(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f;
const float encode_tile_factor = (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f;
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor);
LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
output = tiled_compute(input,

3
src/model_loader.cpp
View File

@@ -503,6 +503,9 @@ SDVersion ModelLoader::get_sd_version() {
}
return VERSION_QWEN_IMAGE;
}
if (tensor_storage.name.find("model.diffusion_model.txt_in.individual_token_refiner.blocks.0.adaLN_modulation.1.weight") != std::string::npos) {
return VERSION_HUNYUAN_VIDEO;
}
if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {
return VERSION_ANIMA;
}

71
src/name_conversion.cpp
View File

@@ -664,6 +664,72 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
return name;
}
std::string convert_hunyuan_video_to_original_flux(std::string name) {
int num_layers = 54;
int num_single_layers = 0;
static std::unordered_map<std::string, std::string> hy_name_map;
if (hy_name_map.empty()) {
// --- double transformer blocks ---
for (int i = 0; i < num_layers; ++i) {
std::string block_prefix = "double_blocks." + std::to_string(i) + ".";
std::string dst_prefix = "double_blocks." + std::to_string(i) + ".";
hy_name_map[block_prefix + "img_mod.linear"] = dst_prefix + "img_mod.lin";
hy_name_map[block_prefix + "txt_mod.linear"] = dst_prefix + "txt_mod.lin";
// attn
hy_name_map[block_prefix + "img_attn_qkv"] = dst_prefix + "img_attn.qkv";
hy_name_map[block_prefix + "txt_attn_qkv"] = dst_prefix + "txt_attn.qkv";
// norm
hy_name_map[block_prefix + "img_attn_q_norm.weight"] = dst_prefix + "img_attn.norm.query_norm.scale";
hy_name_map[block_prefix + "img_attn_k_norm.weight"] = dst_prefix + "img_attn.norm.key_norm.scale";
hy_name_map[block_prefix + "txt_attn_q_norm.weight"] = dst_prefix + "txt_attn.norm.query_norm.scale";
hy_name_map[block_prefix + "txt_attn_k_norm.weight"] = dst_prefix + "txt_attn.norm.key_norm.scale";
// ff
hy_name_map[block_prefix + "img_mlp.fc1"] = dst_prefix + "img_mlp.0";
hy_name_map[block_prefix + "img_mlp.fc2"] = dst_prefix + "img_mlp.2";
hy_name_map[block_prefix + "txt_mlp.fc1"] = dst_prefix + "txt_mlp.0";
hy_name_map[block_prefix + "txt_mlp.fc2"] = dst_prefix + "txt_mlp.2";
// output projections
hy_name_map[block_prefix + "img_attn_proj"] = dst_prefix + "img_attn.proj";
hy_name_map[block_prefix + "txt_attn_proj"] = dst_prefix + "txt_attn.proj";
}
}
hy_name_map["time_in.mlp.0"] = "time_in.in_layer";
hy_name_map["time_in.mlp.2"] = "time_in.out_layer";
hy_name_map["time_r_in.mlp.0"] = "time_r_in.in_layer";
hy_name_map["time_r_in.mlp.2"] = "time_r_in.out_layer";
hy_name_map["vector_in.mlp.0"] = "vector_in.in_layer";
hy_name_map["vector_in.mlp.2"] = "vector_in.out_layer";
hy_name_map["guidance_in.mlp.0"] = "guidance_in.in_layer";
hy_name_map["guidance_in.mlp.2"] = "guidance_in.out_layer";
hy_name_map["txt_in.c_embedder.linear_1"] = "txt_in.c_embedder.in_layer";
hy_name_map["txt_in.c_embedder.linear_2"] = "txt_in.c_embedder.out_layer";
hy_name_map["txt_in.t_embedder.mlp.0"] = "txt_in.t_embedder.in_layer";
hy_name_map["txt_in.t_embedder.mlp.2"] = "txt_in.t_embedder.out_layer";
replace_with_prefix_map(name, hy_name_map);
static const std::vector<std::pair<std::string, std::string>> generic_name_map = {
{"_attn_qkv.", "_attn.qkv."},
{"_attn_proj.", "_attn.proj."},
{"mlp.fc1.", "mlp.0."},
{"mlp.fc2.", "mlp.2."},
{".modulation.linear.", ".modulation.lin."},
};
replace_with_name_map(name, generic_name_map);
return name;
}
std::string convert_diffusers_dit_to_original_lumina2(std::string name) {
int num_layers = 30;
int num_refiner_layers = 2;
@@ -807,6 +873,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
name = convert_diffusers_dit_to_original_sd3(name);
} else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
name = convert_diffusers_dit_to_original_flux(name);
} else if (sd_version_is_hunyuan_video(version)) {
name = convert_hunyuan_video_to_original_flux(name);
} else if (sd_version_is_z_image(version)) {
name = convert_diffusers_dit_to_original_lumina2(name);
} else if (sd_version_is_anima(version)) {
@@ -980,6 +1048,9 @@ std::string convert_diffusers_to_original_wan_vae(std::string name) {
}
std::string convert_first_stage_model_name(std::string name, std::string prefix, SDVersion version) {
if (sd_version_is_hunyuan_video(version)) {
return name;
}
if (sd_version_uses_wan_vae(version)) {
return convert_diffusers_to_original_wan_vae(name);
}

105
src/stable-diffusion.cpp
View File

@@ -30,6 +30,7 @@
#include "model/diffusion/ernie_image.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/hidream_o1.hpp"
#include "model/diffusion/hunyuan.hpp"
#include "model/diffusion/ideogram4.hpp"
#include "model/diffusion/krea2.hpp"
#include "model/diffusion/lens.hpp"
@@ -46,6 +47,7 @@
#include "model/upscaler/esrgan.hpp"
#include "model/upscaler/ltx_latent_upscaler.hpp"
#include "model/vae/auto_encoder_kl.hpp"
#include "model/vae/hunyuan_vae.hpp"
#include "model/vae/ltx_audio_vae.hpp"
#include "model/vae/ltx_vae.hpp"
#include "model/vae/tae.hpp"
@@ -96,6 +98,7 @@ const char* model_version_to_str[] = {
"LingBot Video",
"Qwen Image",
"Qwen Image Layered",
"Hunyuan Video",
"Anima",
"Flux.2",
"Flux.2 klein",
@@ -1064,6 +1067,18 @@ public:
tensor_storage_map,
"model.diffusion_model",
model_manager);
} else if (sd_version_is_hunyuan_video(version)) {
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
tensor_storage_map,
version,
"",
false,
model_manager);
diffusion_model = std::make_shared<Hunyuan::HunyuanVideoRunner>(backend_for(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
version,
model_manager);
} else if (sd_version_is_wan(version)) {
cond_stage_model = std::make_shared<T5CLIPEmbedder>(backend_for(SDBackendModule::TE),
tensor_storage_map,
@@ -1269,7 +1284,7 @@ public:
}
auto create_tae = [&](bool decode_only) -> std::shared_ptr<VAE> {
if (sd_version_uses_wan_vae(version) || sd_version_is_ltxav(version)) {
if (sd_version_uses_wan_vae(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
return std::make_shared<TinyVideoAutoEncoder>(backend_for(SDBackendModule::VAE),
tensor_storage_map,
"decoder",
@@ -1306,6 +1321,13 @@ public:
false,
version,
model_manager);
} else if (sd_version_uses_hunyuan_video_vae(vae_version)) {
return std::make_shared<Hunyuan::HunyuanVideoVAERunner>(backend_for(SDBackendModule::VAE),
tensor_storage_map,
"first_stage_model",
false,
vae_version,
model_manager);
} else if (sd_version_uses_wan_vae(vae_version)) {
return std::make_shared<WAN::WanVAERunner>(backend_for(SDBackendModule::VAE),
tensor_storage_map,
@@ -1617,6 +1639,7 @@ public:
}
} else if (sd_version_is_sd3(version) ||
sd_version_is_wan(version) ||
sd_version_is_hunyuan_video(version) ||
sd_version_is_lingbot_video(version) ||
sd_version_is_qwen_image(version) ||
version == VERSION_HIDREAM_O1 ||
@@ -1629,6 +1652,8 @@ public:
pred_type = FLOW_PRED;
if (sd_version_is_wan(version)) {
default_flow_shift = 5.f;
} else if (sd_version_is_hunyuan_video(version)) {
default_flow_shift = 7.f;
} else if (sd_version_is_ernie_image(version)) {
default_flow_shift = 4.f;
} else if (sd_version_is_pid(version)) {
@@ -2446,6 +2471,10 @@ public:
sd::Tensor<float> timesteps_tensor({static_cast<int64_t>(timesteps_vec.size())}, timesteps_vec);
sd::Tensor<float> guidance_tensor({1}, std::vector<float>{guidance.distilled_guidance});
sd::Tensor<float> hunyuan_timestep_r_tensor;
if (sd_version_is_hunyuan_video(version) && step + 1 < sigmas.size()) {
hunyuan_timestep_r_tensor = sd::Tensor<float>::from_vector({sigmas[step + 1]});
}
sd::Tensor<float> noised_input = x * c_in;
if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) {
noised_input = noised_input * denoise_mask + init_latent * (1.0f - denoise_mask);
@@ -2520,6 +2549,12 @@ public:
} else if (sd_version_is_wan(version)) {
diffusion_params.extra = WanDiffusionExtra{vace_context.empty() ? nullptr : &vace_context,
vace_strength};
} else if (sd_version_is_hunyuan_video(version)) {
diffusion_params.extra = HunyuanVideoDiffusionExtra{
&guidance_tensor,
condition.extra_c_crossattns.empty() ? nullptr : &condition.extra_c_crossattns[0],
condition.c_vector.empty() ? nullptr : &condition.c_vector,
hunyuan_timestep_r_tensor.empty() ? nullptr : &hunyuan_timestep_r_tensor};
} else if (version == VERSION_HIDREAM_O1) {
diffusion_params.extra = HiDreamO1DiffusionExtra{
condition.c_input_ids.empty() ? nullptr : &condition.c_input_ids,
@@ -2713,6 +2748,8 @@ public:
latent_channel = 128;
} else if (version == VERSION_WAN2_2_TI2V) {
latent_channel = 48;
} else if (sd_version_is_hunyuan_video(version)) {
latent_channel = 32;
} else if (version == VERSION_HIDREAM_O1) {
latent_channel = 3;
} else if (version == VERSION_CHROMA_RADIANCE) {
@@ -2760,7 +2797,7 @@ public:
int latent_frames = frames;
if (sd_version_is_ltxav(version)) {
latent_frames = ((frames - 1) / 8) + 1;
} else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) {
} else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_hunyuan_video(version)) {
latent_frames = ((frames - 1) / 4) + 1;
}
return latent_frames;
@@ -2773,7 +2810,7 @@ public:
if (sd_version_is_ltxav(version)) {
return (latent_frames - 1) * 8 + 1;
}
if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) {
if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_hunyuan_video(version)) {
return (latent_frames - 1) * 4 + 1;
}
return latent_frames;
@@ -3570,7 +3607,7 @@ struct sd_ctx_t {
};
static bool sd_version_supports_video_generation(SDVersion version) {
return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version);
return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version);
}
static bool sd_version_supports_image_generation(SDVersion version) {
@@ -5627,6 +5664,66 @@ static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd
}
}
if (sd_version_is_hunyuan_video(sd_ctx->sd->version) &&
(!start_image.empty() || !end_image.empty())) {
LOG_INFO("Hunyuan Video IMG2VID");
int64_t t1 = ggml_time_ms();
auto concat_latent = sd_ctx->sd->generate_init_latent(request->width,
request->height,
request->frames,
true);
auto encode_condition_frame = [&](const sd::Tensor<float>& image,
int64_t latent_frame,
const char* name) -> bool {
auto encoded = sd_ctx->sd->encode_first_stage(image);
if (encoded.empty()) {
LOG_ERROR("failed to encode Hunyuan Video %s conditioning frame", name);
return false;
}
if (encoded.dim() == 4) {
encoded.unsqueeze_(2);
}
if (encoded.dim() != 5 ||
encoded.shape()[0] != concat_latent.shape()[0] ||
encoded.shape()[1] != concat_latent.shape()[1] ||
encoded.shape()[3] != concat_latent.shape()[3]) {
LOG_ERROR("invalid Hunyuan Video %s conditioning latent shape", name);
return false;
}
sd::ops::slice_assign(&concat_latent,
2,
latent_frame,
latent_frame + 1,
sd::ops::slice(encoded, 2, 0, 1));
return true;
};
if (!start_image.empty() && !encode_condition_frame(start_image, 0, "start")) {
return std::nullopt;
}
if (!end_image.empty() &&
!encode_condition_frame(end_image, concat_latent.shape()[2] - 1, "end")) {
return std::nullopt;
}
sd::Tensor<float> concat_mask = sd::zeros<float>({concat_latent.shape()[0],
concat_latent.shape()[1],
concat_latent.shape()[2],
1,
1});
if (!start_image.empty()) {
sd::ops::fill_slice(&concat_mask, 2, 0, 1, 1.0f);
}
if (!end_image.empty()) {
sd::ops::fill_slice(&concat_mask, 2, concat_mask.shape()[2] - 1, concat_mask.shape()[2], 1.0f);
}
latents.concat_latent = sd::ops::concat(concat_latent, concat_mask, 3);
int64_t t2 = ggml_time_ms();
LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
}
if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-I2V-14B" ||
sd_ctx->sd->diffusion_model->get_desc() == "Wan2.2-I2V-14B" ||
sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-I2V-1.3B" ||