diff --git a/README.md b/README.md index ab67af4e3..340ba5b2d 100644 --- a/README.md +++ b/README.md @@ -65,6 +65,7 @@ API and command-line option may change frequently.*** - Video Models - [Wan2.1/Wan2.2](./docs/wan.md) - [LTX-2.3](./docs/ltx2.md) + - [HunyuanVideo 1.5](./docs/hunyuan_video.md) - [LingBot-Video](./docs/lingbot_video.md) - [PhotoMaker](./docs/photo_maker.md) support. - Control Net support with SD 1.5 diff --git a/assets/huanyuan_video/hy1.5_t2v.mp4 b/assets/huanyuan_video/hy1.5_t2v.mp4 new file mode 100644 index 000000000..76a960fbf Binary files /dev/null and b/assets/huanyuan_video/hy1.5_t2v.mp4 differ diff --git a/assets/hunyuan_video/hy1.5_t2v.mp4 b/assets/hunyuan_video/hy1.5_t2v.mp4 new file mode 100644 index 000000000..76a960fbf Binary files /dev/null and b/assets/hunyuan_video/hy1.5_t2v.mp4 differ diff --git a/docs/hunyuan_video.md b/docs/hunyuan_video.md new file mode 100644 index 000000000..2b05e0f71 --- /dev/null +++ b/docs/hunyuan_video.md @@ -0,0 +1,24 @@ +# HunyuanVideo 1.5 + +HunyuanVideo 1.5 uses a HunyuanVideo diffusion transformer, a causal video VAE, Qwen2.5-VL 7B for the main text conditioning, +and ByT5 Small GlyphXL for glyph-aware text conditioning. + +## Download weights + +- Download HunyuanVideo 1.5 + - safetensors: https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/tree/main/split_files/diffusion_models +- Download vae + - safetensors: https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/tree/main/split_files/vae +- Download qwen_2.5_vl 7b + - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/tree/main/split_files/text_encoders + - gguf: https://huggingface.co/mradermacher/Qwen2.5-VL-7B-Instruct-GGUF/tree/main +- Download byt5 small glyphxl + - safetensros: https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/tree/main/split_files/text_encoders + +## Text-to-video example + +```shell +.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\hunyuanvideo1.5_720p_t2v_fp16.safetensors --vae ..\models\vae\hunyuanvideo15_vae_fp16.safetensors --llm ..\models\text_encoders\qwen_2.5_vl_7b.safetensors --t5xxl ..\models\text_encoders\byt5_small_glyphxl_fp16.safetensors -p "a lovely cat" --cfg-scale 6.0 --sampling-method euler -v -W 1280 -H 720 --offload-to-cpu --diffusion-fa --video-frames 33 --vae-tiling +``` + + \ No newline at end of file diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index 4dc4e0cd4..a0bc49b37 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -1782,6 +1782,7 @@ struct LLMEmbedder : public Conditioner { SDVersion version; std::shared_ptr tokenizer; std::shared_ptr llm; + std::shared_ptr 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(backend, + tensor_storage_map, + byt5_prefix, + false, + weight_manager); + break; + } + } + } } void get_param_tensors(std::map& 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& 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& 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& 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& 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, std::vector> 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(prompt.size()); + prompt += conditioner_params.text; + prompt_attn_range.second = static_cast(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> extra_hidden_states_vec; + if (sd_version_is_hunyuan_video(version) && byt5) { + std::vector 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 tokens; + tokens.reserve(byt5_text.size() + 1); + for (unsigned char byte : byt5_text) { + tokens.push_back(static_cast(byte) + 3); + } + tokens.push_back(1); + sd::Tensor input_ids({static_cast(tokens.size())}, tokens); + auto byt5_hidden_states = byt5->compute(n_threads, + input_ids, + sd::Tensor(), + 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], diff --git a/src/model.h b/src/model.h index cb2b2039e..67523656a 100644 --- a/src/model.h +++ b/src/model.h @@ -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) || diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index d4142c181..100f96ee0 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -535,6 +535,33 @@ namespace Rope { return vid_ids_repeated; } + __STATIC_INLINE__ std::vector> 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> txt_ids(bs * context_len, std::vector(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 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& 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> gen_qwen_image_ids(int t, int h, int w, diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index edd7a89b8..375576c14 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -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(new LayerNorm(hidden_size, 1e-06f, false)); - blocks["linear"] = std::shared_ptr(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(new Linear(hidden_size, out_dim, bias)); if (!prune_mod) { blocks["adaLN_modulation.1"] = std::shared_ptr(new Linear(hidden_size, 2 * hidden_size, bias)); } diff --git a/src/model/diffusion/hunyuan.hpp b/src/model/diffusion/hunyuan.hpp new file mode 100644 index 000000000..067d92aaf --- /dev/null +++ b/src/model/diffusion/hunyuan.hpp @@ -0,0 +1,679 @@ +#ifndef __SD_MODEL_DIFFUSION_HUNYUAN_HPP__ +#define __SD_MODEL_DIFFUSION_HUNYUAN_HPP__ + +#include + +#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(hidden_size, hidden_size * 3, attn_bias); + blocks["self_attn.proj"] = std::make_shared(hidden_size, hidden_size, attn_bias); + + blocks["norm1"] = std::make_shared(hidden_size, 1e-6f, true); + blocks["norm2"] = std::make_shared(hidden_size, 1e-6f, true); + + blocks["mlp.0"] = std::make_shared(hidden_size, hidden_size * mlp_ratio); + blocks["mlp.2"] = std::make_shared(hidden_size * mlp_ratio, hidden_size); + + // adaLN_modulation.0 is nn.SiLU() + blocks["adaLN_modulation.1"] = std::make_shared(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(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto self_attn_qkv = std::dynamic_pointer_cast(blocks["self_attn.qkv"]); + auto self_attn_proj = std::dynamic_pointer_cast(blocks["self_attn.proj"]); + auto mlp_fc1 = std::dynamic_pointer_cast(blocks["mlp.0"]); + auto mlp_fc2 = std::dynamic_pointer_cast(blocks["mlp.2"]); + auto adaLN_modulation_1 = std::dynamic_pointer_cast(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(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(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(in_channels, hidden_size); + blocks["t_embedder"] = std::make_shared(256, hidden_size); + blocks["c_embedder"] = std::make_shared(in_channels, hidden_size); + blocks["individual_token_refiner"] = std::make_shared(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(blocks["input_embedder"]); + auto t_embedder = std::dynamic_pointer_cast(blocks["t_embedder"]); + auto c_embedder = std::dynamic_pointer_cast(blocks["c_embedder"]); + auto individual_token_refiner = std::dynamic_pointer_cast(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(in_dim); + blocks["fc1"] = std::make_shared(in_dim, 2048); + blocks["fc2"] = std::make_shared(2048, 2048); + blocks["fc3"] = std::make_shared(2048, hidden_size); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto layernorm = std::dynamic_pointer_cast(blocks["layernorm"]); + auto fc1 = std::dynamic_pointer_cast(blocks["fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); + auto fc3 = std::dynamic_pointer_cast(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 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 axes_dim = {16, 56, 56}; + int axes_dim_sum = 128; + + int64_t patch_volume() const { + return static_cast(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(storage.ne[2]), + static_cast(storage.ne[1]), + static_cast(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(storage.ne[2]), + static_cast(storage.ne[1]), + static_cast(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(config.hidden_size / storage.ne[0]); + } else if (ends_with(name, "double_blocks.0.img_mlp.0.weight")) { + config.mlp_ratio = static_cast(storage.ne[1]) / static_cast(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(this->config.context_in_dim, this->config.num_heads, head_dim, 2); + blocks["img_in"] = std::make_shared(static_cast(224) /*Not used*/, + this->config.patch_size, + this->config.in_channels, + this->config.hidden_size); + blocks["time_in"] = std::make_shared(256, this->config.hidden_size); + if (this->config.vec_in_dim > 0) { + blocks["vector_in"] = std::make_shared(this->config.vec_in_dim, this->config.hidden_size); + } + if (this->config.vision_in_dim > 0) { + blocks["vision_in"] = std::make_shared(this->config.vision_in_dim, this->config.hidden_size); + } + if (this->config.guidance_embed) { + blocks["guidance_in"] = std::make_shared(256, this->config.hidden_size); + } + if (this->config.use_byt5) { + blocks["byt5_in"] = std::make_shared(1472, this->config.hidden_size); + } + if (this->config.use_meanflow) { + blocks["time_r_in"] = std::make_shared(256, this->config.hidden_size); + } + + for (int i = 0; i < this->config.depth; i++) { + blocks["double_blocks." + std::to_string(i)] = std::make_shared(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(this->config.hidden_size, + this->config.num_heads, + this->config.mlp_ratio, + i, + 0.f); + } + + blocks["final_layer"] = std::make_shared(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(T % pt)) % pt; + int pad_h = (ph - static_cast(H % ph)) % ph; + int pad_w = (pw - static_cast(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(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(blocks["img_in"]); + auto txt_in = std::dynamic_pointer_cast(blocks["txt_in"]); + auto time_in = std::dynamic_pointer_cast(blocks["time_in"]); + auto final_layer = std::dynamic_pointer_cast(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(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(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(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(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(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(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(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 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 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& tensors, const std::string& prefix) override { + hunyuan_video.get_param_tensors(tensors, prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor, + const sd::Tensor& c_concat_tensor = {}, + const sd::Tensor& y_tensor = {}, + const sd::Tensor& guidance_tensor = {}, + const sd::Tensor& byt5_tensor = {}, + const sd::Tensor& vision_tensor = {}, + const sd::Tensor& 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(config.in_channels - x->ne[3])); + } + + int text_len = static_cast(context->ne[1]); + if (byt5 != nullptr) { + text_len += static_cast(byt5->ne[1]); + } + if (vision != nullptr) { + text_len += static_cast(vision->ne[1]); + } + pe_vec = Rope::gen_hunyuan_video_pe(static_cast(x->ne[2]), + static_cast(x->ne[1]), + static_cast(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(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 compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context, + const sd::Tensor& c_concat = {}, + const sd::Tensor& y = {}, + const sd::Tensor& guidance = {}, + const sd::Tensor& byt5 = {}, + const sd::Tensor& vision = {}, + const sd::Tensor& 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(get_graph, n_threads, false, false, false), x.dim()); + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + GGML_ASSERT(diffusion_params.context != nullptr); + const auto* extra = diffusion_extra_as(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__ diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp index b89ff32c6..39c633747 100644 --- a/src/model/diffusion/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -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); } diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index d8e76dfb9..6731b5fb5 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -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 @@ struct PatchEmbed : public GGMLBlock { 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(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(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(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 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(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(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(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(H % patch_h)) % patch_h; + int pad_w = (patch_w - static_cast(W % patch_w)) % patch_w; + if (is_3d) { + int64_t T = x->ne[2]; + pad_t = (patch_t - static_cast(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; } diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 6e04f9f1f..d8ee3d26f 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -87,6 +87,13 @@ struct MiniT2IDiffusionExtra { const sd::Tensor* mask = nullptr; }; +struct HunyuanVideoDiffusionExtra { + const sd::Tensor* guidance = nullptr; + const sd::Tensor* byt5 = nullptr; + const sd::Tensor* vision = nullptr; + const sd::Tensor* timestep_r = nullptr; +}; + using DiffusionExtraParams = std::variant; + MiniT2IDiffusionExtra, + HunyuanVideoDiffusionExtra>; struct DiffusionParams { const sd::Tensor* x = nullptr; diff --git a/src/model/te/t5.hpp b/src/model/te/t5.hpp index 6d2326f94..90c781e1d 100644 --- a/src/model/te/t5.hpp +++ b/src/model/te/t5.hpp @@ -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 @@ struct T5 : public GGMLBlock { : config(config) { blocks["encoder"] = std::shared_ptr(new T5Stack(config.num_layers, config.model_dim, - config.model_dim, + config.inner_dim, config.ff_dim, config.num_heads, config.relative_attention)); diff --git a/src/model/vae/hunyuan_vae.hpp b/src/model/vae/hunyuan_vae.hpp new file mode 100644 index 000000000..b16126181 --- /dev/null +++ b/src/model/vae/hunyuan_vae.hpp @@ -0,0 +1,834 @@ +#ifndef __SD_MODEL_VAE_HUNYUAN_VAE_HPP__ +#define __SD_MODEL_VAE_HUNYUAN_VAE_HPP__ + +#include +#include +#include +#include +#include +#include +#include +#include + +#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* input = nullptr; + std::vector* 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 kernel_size; + + public: + CausalConv3d(int64_t in_channels, + int64_t out_channels, + std::tuple kernel_size, + std::tuple stride = {1, 1, 1}, + std::tuple padding = {0, 0, 0}, + std::tuple dilation = {1, 1, 1}, + bool bias = true) + : kernel_size(kernel_size) { + blocks["conv"] = std::make_shared(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(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 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(in_channels); + blocks["q"] = std::make_shared(in_channels, in_channels, std::tuple{1, 1, 1}); + blocks["k"] = std::make_shared(in_channels, in_channels, std::tuple{1, 1, 1}); + blocks["v"] = std::make_shared(in_channels, in_channels, std::tuple{1, 1, 1}); + blocks["proj_out"] = std::make_shared(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(blocks["norm"]); + auto q_proj = std::dynamic_pointer_cast(blocks["q"]); + auto k_proj = std::dynamic_pointer_cast(blocks["k"]); + auto v_proj = std::dynamic_pointer_cast(blocks["v"]); + auto proj_out = std::dynamic_pointer_cast(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(in_channels); + blocks["conv1"] = std::make_shared(in_channels, out_channels, std::tuple{3, 3, 3}); + + blocks["norm2"] = std::make_shared(out_channels); + blocks["conv2"] = std::make_shared(out_channels, out_channels, std::tuple{3, 3, 3}); + + if (out_channels != in_channels) { + blocks["nin_shortcut"] = std::make_shared(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(blocks["norm1"]); + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto conv2 = std::dynamic_pointer_cast(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(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(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(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(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(in_channels, + out_channels / factor, + std::tuple{3, 3, 3}); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto conv = std::dynamic_pointer_cast(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(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(in_channels); + } + blocks["block_" + std::to_string(i + 2)] = std::make_shared(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(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(blocks["attn_" + std::to_string(i + 1)]); + x = block->forward(ctx, x); + } + auto block = std::dynamic_pointer_cast(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(IC, out_channels); + } + if (upsample_out_channels > 0) { + blocks["upsample"] = std::make_shared(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(blocks["block." + std::to_string(i)]); + x = block->forward(ctx, x, carry); + } + if (upsample_out_channels > 0) { + auto upsample = std::dynamic_pointer_cast(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(IC, out_channels); + } + if (downsample_out_channels > 0) { + blocks["downsample"] = std::make_shared(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(blocks["block." + std::to_string(i)]); + x = block->forward(ctx, x); + } + if (downsample_out_channels > 0) { + auto downsample = std::dynamic_pointer_cast(blocks["downsample"]); + x = downsample->forward(ctx, x); + } + return x; + } + }; + + class Encoder : public GGMLBlock { + protected: + int64_t z_channels; + std::vector block_out_channels; + + public: + Encoder(int64_t in_channels = 3, + int64_t z_channels = 32, + std::vector 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(in_channels, + this->block_out_channels[0], + std::tuple{3, 3, 3}); + + int spatial_depth = static_cast(std::log2(static_cast(spatial_compression_ratio))); + int temporal_start = static_cast(std::log2(static_cast(spatial_compression_ratio / temporal_compression_ratio))); + int64_t channels = this->block_out_channels[0]; + for (int i = 0; i < static_cast(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(channels, + out_channels, + layers_per_block, + next_channels, + i >= temporal_start); + channels = next_channels; + } else { + blocks["down." + std::to_string(i)] = std::make_shared(channels, + out_channels, + layers_per_block); + channels = out_channels; + } + } + + blocks["mid"] = std::make_shared(channels); + blocks["norm_out"] = std::make_shared(channels); + blocks["conv_out"] = std::make_shared(channels, + z_channels * 2, + std::tuple{3, 3, 3}); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto conv_in = std::dynamic_pointer_cast(blocks["conv_in"]); + auto mid = std::dynamic_pointer_cast(blocks["mid"]); + auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); + auto conv_out = std::dynamic_pointer_cast(blocks["conv_out"]); + + x = conv_in->forward(ctx, x); + for (int i = 0; i < static_cast(block_out_channels.size()); i++) { + auto down = std::dynamic_pointer_cast(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 block_out_channels; + + public: + Decoder(int64_t in_channels = 32, + int64_t out_channels = 3, + std::vector 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(in_channels, this->block_out_channels[0], std::tuple{3, 3, 3}); + blocks["mid"] = std::make_shared(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(spatial_compression_ratio)); + bool add_temporal_upsample = i < std::log2(static_cast(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(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(IC, OC, layers_per_block + 1, 0, false); + } + } + + blocks["norm_out"] = std::make_shared(this->block_out_channels.back()); + blocks["conv_out"] = std::make_shared(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(blocks["conv_in"]); + auto mid_block = std::dynamic_pointer_cast(blocks["mid"]); + auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); + auto conv_out = std::dynamic_pointer_cast(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 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 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(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 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& 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 vae_output_to_latents(const sd::Tensor& vae_output, + std::shared_ptr rng) override { + SD_UNUSED(rng); + return vae_output; + } + + sd::Tensor diffusion_to_vae_latents(const sd::Tensor& latents) override { + return latents / 1.03682f; + } + + sd::Tensor vae_to_diffusion_latents(const sd::Tensor& latents) override { + return latents * 1.03682f; + } + + ggml_cgraph* build_graph(const sd::Tensor& 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(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 _compute(const int n_threads, + const sd::Tensor& input, + bool decode_graph) override { + if (!decode_graph && decode_only) { + LOG_ERROR("Hunyuan Video VAE encoder weights are not available"); + return {}; + } + + sd::Tensor 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(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__ diff --git a/src/model/vae/tae.hpp b/src/model/vae/tae.hpp index 558c0544c..710729fd7 100644 --- a/src/model/vae/tae.hpp +++ b/src/model/vae/tae.hpp @@ -528,6 +528,9 @@ class TAEHV : public GGMLBlock { 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 @@ class TAEHV : public GGMLBlock { ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) { auto decoder = std::dynamic_pointer_cast(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 @@ class TAEHV : public GGMLBlock { ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) { auto encoder = std::dynamic_pointer_cast(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 @@ class TAEHV : public GGMLBlock { } } 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)); } diff --git a/src/model/vae/vae.hpp b/src/model/vae/vae.hpp index b97aecee0..a462c8728 100644 --- a/src/model/vae/vae.hpp +++ b/src/model/vae/vae.hpp @@ -74,7 +74,7 @@ struct VAE : public GGMLRunner { 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 @@ struct VAE : public GGMLRunner { // 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, diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 853ed1440..2f9e68cac 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -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; } diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 3aa7880e2..904e12e86 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -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 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> 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); } diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 7429ab72d..0bddba2bb 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -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 @@ class StableDiffusionGGML { tensor_storage_map, "model.diffusion_model", model_manager); + } else if (sd_version_is_hunyuan_video(version)) { + cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + model_manager); + diffusion_model = std::make_shared(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(backend_for(SDBackendModule::TE), tensor_storage_map, @@ -1269,7 +1284,7 @@ class StableDiffusionGGML { } auto create_tae = [&](bool decode_only) -> std::shared_ptr { - 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(backend_for(SDBackendModule::VAE), tensor_storage_map, "decoder", @@ -1306,6 +1321,13 @@ class StableDiffusionGGML { false, version, model_manager); + } else if (sd_version_uses_hunyuan_video_vae(vae_version)) { + return std::make_shared(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(backend_for(SDBackendModule::VAE), tensor_storage_map, @@ -1617,6 +1639,7 @@ class StableDiffusionGGML { } } 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 @@ class StableDiffusionGGML { 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 @@ class StableDiffusionGGML { sd::Tensor timesteps_tensor({static_cast(timesteps_vec.size())}, timesteps_vec); sd::Tensor guidance_tensor({1}, std::vector{guidance.distilled_guidance}); + sd::Tensor hunyuan_timestep_r_tensor; + if (sd_version_is_hunyuan_video(version) && step + 1 < sigmas.size()) { + hunyuan_timestep_r_tensor = sd::Tensor::from_vector({sigmas[step + 1]}); + } sd::Tensor 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 @@ class StableDiffusionGGML { } 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 @@ class StableDiffusionGGML { 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 @@ class StableDiffusionGGML { 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 @@ class StableDiffusionGGML { 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 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& 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 concat_mask = sd::zeros({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" ||