From c90ae4e22253739c7bfcacdf89334be3258346b8 Mon Sep 17 00:00:00 2001 From: leejet Date: Sat, 6 Sep 2025 13:39:09 +0800 Subject: [PATCH] add wan vace t2v support --- diffusion_model.hpp | 140 ++++++++++++++-------------- ggml_extend.hpp | 32 +++++++ stable-diffusion.cpp | 213 +++++++++++++++++++++++-------------------- wan.hpp | 202 +++++++++++++++++++++++++++++++++------- 4 files changed, 385 insertions(+), 202 deletions(-) diff --git a/diffusion_model.hpp b/diffusion_model.hpp index 312266e..80f81f8 100644 --- a/diffusion_model.hpp +++ b/diffusion_model.hpp @@ -6,22 +6,28 @@ #include "unet.hpp" #include "wan.hpp" +struct DiffusionParams { + struct ggml_tensor* x = NULL; + struct ggml_tensor* timesteps = NULL; + struct ggml_tensor* context = NULL; + struct ggml_tensor* c_concat = NULL; + struct ggml_tensor* y = NULL; + struct ggml_tensor* guidance = NULL; + std::vector ref_latents = {}; + int num_video_frames = -1; + std::vector controls = {}; + float control_strength = 0.f; + struct ggml_tensor* vace_context = NULL; + float vace_strength = 1.f; + std::vector skip_layers = {}; +}; + struct DiffusionModel { virtual std::string get_desc() = 0; virtual void compute(int n_threads, - struct ggml_tensor* x, - struct ggml_tensor* timesteps, - struct ggml_tensor* context, - struct ggml_tensor* c_concat, - struct ggml_tensor* y, - struct ggml_tensor* guidance, - std::vector ref_latents = {}, - int num_video_frames = -1, - std::vector controls = {}, - float control_strength = 0.f, - struct ggml_tensor** output = NULL, - struct ggml_context* output_ctx = NULL, - std::vector skip_layers = std::vector()) = 0; + DiffusionParams diffusion_params, + struct ggml_tensor** output = NULL, + struct ggml_context* output_ctx = NULL) = 0; virtual void alloc_params_buffer() = 0; virtual void free_params_buffer() = 0; virtual void free_compute_buffer() = 0; @@ -70,21 +76,18 @@ struct UNetModel : public DiffusionModel { } void compute(int n_threads, - struct ggml_tensor* x, - struct ggml_tensor* timesteps, - struct ggml_tensor* context, - struct ggml_tensor* c_concat, - struct ggml_tensor* y, - struct ggml_tensor* guidance, - std::vector ref_latents = {}, - int num_video_frames = -1, - std::vector controls = {}, - float control_strength = 0.f, - struct ggml_tensor** output = NULL, - struct ggml_context* output_ctx = NULL, - std::vector skip_layers = std::vector()) { - (void)skip_layers; // SLG doesn't work with UNet models - return unet.compute(n_threads, x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, output, output_ctx); + DiffusionParams diffusion_params, + struct ggml_tensor** output = NULL, + struct ggml_context* output_ctx = NULL) { + return unet.compute(n_threads, + diffusion_params.x, + diffusion_params.timesteps, + diffusion_params.context, + diffusion_params.c_concat, + diffusion_params.y, + diffusion_params.num_video_frames, + diffusion_params.controls, + diffusion_params.control_strength, output, output_ctx); } }; @@ -126,20 +129,17 @@ struct MMDiTModel : public DiffusionModel { } void compute(int n_threads, - struct ggml_tensor* x, - struct ggml_tensor* timesteps, - struct ggml_tensor* context, - struct ggml_tensor* c_concat, - struct ggml_tensor* y, - struct ggml_tensor* guidance, - std::vector ref_latents = {}, - int num_video_frames = -1, - std::vector controls = {}, - float control_strength = 0.f, - struct ggml_tensor** output = NULL, - struct ggml_context* output_ctx = NULL, - std::vector skip_layers = std::vector()) { - return mmdit.compute(n_threads, x, timesteps, context, y, output, output_ctx, skip_layers); + DiffusionParams diffusion_params, + struct ggml_tensor** output = NULL, + struct ggml_context* output_ctx = NULL) { + return mmdit.compute(n_threads, + diffusion_params.x, + diffusion_params.timesteps, + diffusion_params.context, + diffusion_params.y, + output, + output_ctx, + diffusion_params.skip_layers); } }; @@ -184,20 +184,20 @@ struct FluxModel : public DiffusionModel { } void compute(int n_threads, - struct ggml_tensor* x, - struct ggml_tensor* timesteps, - struct ggml_tensor* context, - struct ggml_tensor* c_concat, - struct ggml_tensor* y, - struct ggml_tensor* guidance, - std::vector ref_latents = {}, - int num_video_frames = -1, - std::vector controls = {}, - float control_strength = 0.f, - struct ggml_tensor** output = NULL, - struct ggml_context* output_ctx = NULL, - std::vector skip_layers = std::vector()) { - return flux.compute(n_threads, x, timesteps, context, c_concat, y, guidance, ref_latents, output, output_ctx, skip_layers); + DiffusionParams diffusion_params, + struct ggml_tensor** output = NULL, + struct ggml_context* output_ctx = NULL) { + return flux.compute(n_threads, + diffusion_params.x, + diffusion_params.timesteps, + diffusion_params.context, + diffusion_params.c_concat, + diffusion_params.y, + diffusion_params.guidance, + diffusion_params.ref_latents, + output, + output_ctx, + diffusion_params.skip_layers); } }; @@ -243,20 +243,20 @@ struct WanModel : public DiffusionModel { } void compute(int n_threads, - struct ggml_tensor* x, - struct ggml_tensor* timesteps, - struct ggml_tensor* context, - struct ggml_tensor* c_concat, - struct ggml_tensor* y, - struct ggml_tensor* guidance, - std::vector ref_latents = {}, - int num_video_frames = -1, - std::vector controls = {}, - float control_strength = 0.f, - struct ggml_tensor** output = NULL, - struct ggml_context* output_ctx = NULL, - std::vector skip_layers = std::vector()) { - return wan.compute(n_threads, x, timesteps, context, y, c_concat, NULL, output, output_ctx); + DiffusionParams diffusion_params, + struct ggml_tensor** output = NULL, + struct ggml_context* output_ctx = NULL) { + return wan.compute(n_threads, + diffusion_params.x, + diffusion_params.timesteps, + diffusion_params.context, + diffusion_params.y, + diffusion_params.c_concat, + NULL, + diffusion_params.vace_context, + diffusion_params.vace_strength, + output, + output_ctx); } }; diff --git a/ggml_extend.hpp b/ggml_extend.hpp index 0965784..cdd60c5 100644 --- a/ggml_extend.hpp +++ b/ggml_extend.hpp @@ -223,6 +223,38 @@ __STATIC_INLINE__ void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_ } } +__STATIC_INLINE__ void ggml_tensor_iter( + ggml_tensor* tensor, + const std::function& fn) { + int64_t n0 = tensor->ne[0]; + int64_t n1 = tensor->ne[1]; + int64_t n2 = tensor->ne[2]; + int64_t n3 = tensor->ne[3]; + + for (int64_t i3 = 0; i3 < n3; i3++) { + for (int64_t i2 = 0; i2 < n2; i2++) { + for (int64_t i1 = 0; i1 < n1; i1++) { + for (int64_t i0 = 0; i0 < n0; i0++) { + fn(tensor, i0, i1, i2, i3); + } + } + } + } +} + +__STATIC_INLINE__ void ggml_tensor_iter( + ggml_tensor* tensor, + const std::function& fn) { + int64_t n0 = tensor->ne[0]; + int64_t n1 = tensor->ne[1]; + int64_t n2 = tensor->ne[2]; + int64_t n3 = tensor->ne[3]; + + for (int64_t i = 0; i < ggml_nelements(tensor); i++) { + fn(tensor, i); + } +} + __STATIC_INLINE__ ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path) { std::ifstream file(file_path, std::ios::binary); if (!file.is_open()) { diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index 64164a2..6d7f17a 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -775,7 +775,12 @@ public: int64_t t0 = ggml_time_ms(); struct ggml_tensor* out = ggml_dup_tensor(work_ctx, x_t); - diffusion_model->compute(n_threads, x_t, timesteps, c, concat, NULL, NULL, {}, -1, {}, 0.f, &out); + DiffusionParams diffusion_params; + diffusion_params.x = x_t; + diffusion_params.timesteps = timesteps; + diffusion_params.context = c; + diffusion_params.c_concat = concat; + diffusion_model->compute(n_threads, diffusion_params, &out); diffusion_model->free_compute_buffer(); double result = 0.f; @@ -1032,7 +1037,9 @@ public: int start_merge_step, SDCondition id_cond, std::vector ref_latents = {}, - ggml_tensor* denoise_mask = nullptr) { + ggml_tensor* denoise_mask = NULL, + ggml_tensor* vace_context = NULL, + float vace_strength = 1.f) { std::vector skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count); float cfg_scale = guidance.txt_cfg; @@ -1116,32 +1123,30 @@ public: // GGML_ASSERT(0); } + DiffusionParams diffusion_params; + diffusion_params.x = noised_input; + diffusion_params.timesteps = timesteps; + diffusion_params.guidance = guidance_tensor; + diffusion_params.ref_latents = ref_latents; + diffusion_params.controls = controls; + diffusion_params.control_strength = control_strength; + diffusion_params.vace_context = vace_context; + diffusion_params.vace_strength = vace_strength; + if (start_merge_step == -1 || step <= start_merge_step) { // cond + diffusion_params.context = cond.c_crossattn; + diffusion_params.c_concat = cond.c_concat; + diffusion_params.y = cond.c_vector; work_diffusion_model->compute(n_threads, - noised_input, - timesteps, - cond.c_crossattn, - cond.c_concat, - cond.c_vector, - guidance_tensor, - ref_latents, - -1, - controls, - control_strength, + diffusion_params, &out_cond); } else { + diffusion_params.context = id_cond.c_crossattn; + diffusion_params.c_concat = cond.c_concat; + diffusion_params.y = id_cond.c_vector; work_diffusion_model->compute(n_threads, - noised_input, - timesteps, - id_cond.c_crossattn, - cond.c_concat, - id_cond.c_vector, - guidance_tensor, - ref_latents, - -1, - controls, - control_strength, + diffusion_params, &out_cond); } @@ -1152,34 +1157,23 @@ public: control_net->compute(n_threads, noised_input, control_hint, timesteps, uncond.c_crossattn, uncond.c_vector); controls = control_net->controls; } + diffusion_params.controls = controls; + diffusion_params.context = uncond.c_crossattn; + diffusion_params.c_concat = uncond.c_concat; + diffusion_params.y = uncond.c_vector; work_diffusion_model->compute(n_threads, - noised_input, - timesteps, - uncond.c_crossattn, - uncond.c_concat, - uncond.c_vector, - guidance_tensor, - ref_latents, - -1, - controls, - control_strength, + diffusion_params, &out_uncond); negative_data = (float*)out_uncond->data; } float* img_cond_data = NULL; if (has_img_cond) { + diffusion_params.context = img_cond.c_crossattn; + diffusion_params.c_concat = img_cond.c_concat; + diffusion_params.y = img_cond.c_vector; work_diffusion_model->compute(n_threads, - noised_input, - timesteps, - img_cond.c_crossattn, - img_cond.c_concat, - img_cond.c_vector, - guidance_tensor, - ref_latents, - -1, - controls, - control_strength, + diffusion_params, &out_img_cond); img_cond_data = (float*)out_img_cond->data; } @@ -1190,20 +1184,13 @@ public: if (is_skiplayer_step) { LOG_DEBUG("Skipping layers at step %d\n", step); // skip layer (same as conditionned) + diffusion_params.context = cond.c_crossattn; + diffusion_params.c_concat = cond.c_concat; + diffusion_params.y = cond.c_vector; + diffusion_params.skip_layers = skip_layers; work_diffusion_model->compute(n_threads, - noised_input, - timesteps, - cond.c_crossattn, - cond.c_concat, - cond.c_vector, - guidance_tensor, - ref_latents, - -1, - controls, - control_strength, - &out_skip, - NULL, - skip_layers); + diffusion_params, + &out_skip); skip_layer_data = (float*)out_skip->data; } float* vec_denoised = (float*)denoised->data; @@ -2412,7 +2399,7 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s } struct ggml_init_params params; - params.mem_size = static_cast(200 * 1024) * 1024; // 200 MB + params.mem_size = static_cast(1024 * 1024) * 1024; // 1G params.mem_size += width * height * frames * 3 * sizeof(float) * 2; params.mem_buffer = NULL; params.no_alloc = false; @@ -2440,6 +2427,7 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s ggml_tensor* clip_vision_output = NULL; ggml_tensor* concat_latent = NULL; ggml_tensor* denoise_mask = NULL; + ggml_tensor* vace_context = NULL; 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-FLF2V-14B") { @@ -2469,23 +2457,17 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s int64_t t1 = ggml_time_ms(); ggml_tensor* image = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, frames, 3); - for (int i3 = 0; i3 < image->ne[3]; i3++) { // channels - for (int i2 = 0; i2 < image->ne[2]; i2++) { - for (int i1 = 0; i1 < image->ne[1]; i1++) { // height - for (int i0 = 0; i0 < image->ne[0]; i0++) { // width - float value = 0.5f; - if (i2 == 0 && sd_vid_gen_params->init_image.data) { // start image - value = *(sd_vid_gen_params->init_image.data + i1 * width * 3 + i0 * 3 + i3); - value /= 255.f; - } else if (i2 == frames - 1 && sd_vid_gen_params->end_image.data) { - value = *(sd_vid_gen_params->end_image.data + i1 * width * 3 + i0 * 3 + i3); - value /= 255.f; - } - ggml_tensor_set_f32(image, value, i0, i1, i2, i3); - } - } + ggml_tensor_iter(image, [&](ggml_tensor* image, int64_t i0, int64_t i1, int64_t i2, int64_t i3) { + float value = 0.5f; + if (i2 == 0 && sd_vid_gen_params->init_image.data) { // start image + value = *(sd_vid_gen_params->init_image.data + i1 * width * 3 + i0 * 3 + i3); + value /= 255.f; + } else if (i2 == frames - 1 && sd_vid_gen_params->end_image.data) { + value = *(sd_vid_gen_params->end_image.data + i1 * width * 3 + i0 * 3 + i3); + value /= 255.f; } - } + ggml_tensor_set_f32(image, value, i0, i1, i2, i3); + }); concat_latent = sd_ctx->sd->encode_first_stage(work_ctx, image); // [b*c, t, h/8, w/8] @@ -2500,21 +2482,15 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s concat_latent->ne[1], concat_latent->ne[2], 4); // [b*4, t, w/8, h/8] - for (int i3 = 0; i3 < concat_mask->ne[3]; i3++) { - for (int i2 = 0; i2 < concat_mask->ne[2]; i2++) { - for (int i1 = 0; i1 < concat_mask->ne[1]; i1++) { - for (int i0 = 0; i0 < concat_mask->ne[0]; i0++) { - float value = 0.0f; - if (i2 == 0 && sd_vid_gen_params->init_image.data) { // start image - value = 1.0f; - } else if (i2 == frames - 1 && sd_vid_gen_params->end_image.data && i3 == 3) { - value = 1.0f; - } - ggml_tensor_set_f32(concat_mask, value, i0, i1, i2, i3); - } - } + ggml_tensor_iter(concat_mask, [&](ggml_tensor* concat_mask, int64_t i0, int64_t i1, int64_t i2, int64_t i3) { + float value = 0.0f; + if (i2 == 0 && sd_vid_gen_params->init_image.data) { // start image + value = 1.0f; + } else if (i2 == frames - 1 && sd_vid_gen_params->end_image.data && i3 == 3) { + value = 1.0f; } - } + ggml_tensor_set_f32(concat_mask, value, i0, i1, i2, i3); + }); concat_latent = ggml_tensor_concat(work_ctx, concat_mask, concat_latent, 3); // [b*(c+4), t, h/8, w/8] } else if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.2-TI2V-5B" && sd_vid_gen_params->init_image.data) { @@ -2533,24 +2509,59 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s sd_ctx->sd->process_latent_out(init_latent); - for (int i3 = 0; i3 < init_image_latent->ne[3]; i3++) { - for (int i2 = 0; i2 < init_image_latent->ne[2]; i2++) { - for (int i1 = 0; i1 < init_image_latent->ne[1]; i1++) { - for (int i0 = 0; i0 < init_image_latent->ne[0]; i0++) { - float value = ggml_tensor_get_f32(init_image_latent, i0, i1, i2, i3); - ggml_tensor_set_f32(init_latent, value, i0, i1, i2, i3); - if (i3 == 0) { - ggml_tensor_set_f32(denoise_mask, 0.f, i0, i1, i2, i3); - } - } - } + ggml_tensor_iter(init_image_latent, [&](ggml_tensor* t, int64_t i0, int64_t i1, int64_t i2, int64_t i3) { + float value = ggml_tensor_get_f32(t, i0, i1, i2, i3); + ggml_tensor_set_f32(init_latent, value, i0, i1, i2, i3); + if (i3 == 0) { + ggml_tensor_set_f32(denoise_mask, 0.f, i0, i1, i2, i3); } - } + }); sd_ctx->sd->process_latent_in(init_latent); int64_t t2 = ggml_time_ms(); LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); + } else if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-VACE-1.3B" || + sd_ctx->sd->diffusion_model->get_desc() == "Wan2.x-VACE-14B") { + LOG_INFO("VACE"); + ggml_tensor* control_video = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, frames, 3); + ggml_set_f32(control_video, 0.5f); + ggml_tensor* mask = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, frames, 1); + ggml_set_f32(mask, 1.0f); + ggml_tensor* inactive = ggml_dup_tensor(work_ctx, control_video); + ggml_tensor* reactive = ggml_dup_tensor(work_ctx, control_video); + + ggml_tensor_iter(control_video, [&](ggml_tensor* t, int64_t i0, int64_t i1, int64_t i2, int64_t i3) { + float control_video_value = ggml_tensor_get_f32(t, i0, i1, i2, i3) - 0.5f; + float mask_value = ggml_tensor_get_f32(mask, i0, i1, i2, 0); + float inactive_value = (control_video_value * (1.f - mask_value)) + 0.5f; + float reactive_value = (control_video_value * mask_value) + 0.5f; + + ggml_tensor_set_f32(inactive, inactive_value, i0, i1, i2, i3); + ggml_tensor_set_f32(reactive, reactive_value, i0, i1, i2, i3); + }); + + inactive = sd_ctx->sd->encode_first_stage(work_ctx, inactive); // [b*c, t, h/8, w/8] + reactive = sd_ctx->sd->encode_first_stage(work_ctx, reactive); // [b*c, t, h/8, w/8] + + sd_ctx->sd->process_latent_in(inactive); + sd_ctx->sd->process_latent_in(reactive); + + vace_context = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, inactive->ne[0], inactive->ne[1], inactive->ne[2], 96); // [b*96, t, h/8, w/8] + ggml_tensor_iter(vace_context, [&](ggml_tensor* vace_context, int64_t i0, int64_t i1, int64_t i2, int64_t i3) { + float value; + if (i3 < 16) { + value = ggml_tensor_get_f32(inactive, i0, i1, i2, i3); + } else if (i3 >= 16 && i3 < 32) { + value = ggml_tensor_get_f32(reactive, i0, i1, i2, i3); + } else { // mask + int64_t vae_stride = 8; + int64_t mask_height_index = i1 * vae_stride + (i3 - 32) / vae_stride; + int64_t mask_width_index = i0 * vae_stride + (i3 - 32) % vae_stride; + value = ggml_tensor_get_f32(mask, mask_width_index, mask_height_index, i2, 0); + } + ggml_tensor_set_f32(vace_context, value, i0, i1, i2, i3); + }); } if (init_latent == NULL) { @@ -2630,7 +2641,8 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s -1, {}, {}, - denoise_mask); + denoise_mask, + vace_context); int64_t sampling_end = ggml_time_ms(); LOG_INFO("sampling(high noise) completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); @@ -2662,7 +2674,8 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s -1, {}, {}, - denoise_mask); + denoise_mask, + vace_context); int64_t sampling_end = ggml_time_ms(); LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); diff --git a/wan.hpp b/wan.hpp index d385cac..2809a22 100644 --- a/wan.hpp +++ b/wan.hpp @@ -1528,12 +1528,12 @@ namespace WAN { blocks["ffn.2"] = std::shared_ptr(new Linear(ffn_dim, dim)); } - struct ggml_tensor* forward(struct ggml_context* ctx, - struct ggml_tensor* x, - struct ggml_tensor* e, - struct ggml_tensor* pe, - struct ggml_tensor* context, - int64_t context_img_len = 257) { + virtual struct ggml_tensor* forward(struct ggml_context* ctx, + struct ggml_tensor* x, + struct ggml_tensor* e, + struct ggml_tensor* pe, + struct ggml_tensor* context, + int64_t context_img_len = 257) { // x: [N, n_token, dim] // e: [N, 6, dim] or [N, T, 6, dim] // context: [N, context_img_len + context_txt_len, dim] @@ -1579,6 +1579,58 @@ namespace WAN { } }; + class VaceWanAttentionBlock : public WanAttentionBlock { + protected: + int block_id; + void init_params(struct ggml_context* ctx, const String2GGMLType& tensor_types = {}, const std::string prefix = "") { + enum ggml_type wtype = get_type(prefix + "weight", tensor_types, GGML_TYPE_F32); + params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 6, 1); + } + + public: + VaceWanAttentionBlock(bool t2v_cross_attn, + int64_t dim, + int64_t ffn_dim, + int64_t num_heads, + bool qk_norm = true, + bool cross_attn_norm = false, + float eps = 1e-6, + int block_id = 0, + bool flash_attn = false) + : WanAttentionBlock(t2v_cross_attn, dim, ffn_dim, num_heads, qk_norm, cross_attn_norm, eps, flash_attn), block_id(block_id) { + if (block_id == 0) { + blocks["before_proj"] = std::shared_ptr(new Linear(dim, dim)); + } + blocks["after_proj"] = std::shared_ptr(new Linear(dim, dim)); + } + + std::pair forward(struct ggml_context* ctx, + struct ggml_tensor* c, + struct ggml_tensor* x, + struct ggml_tensor* e, + struct ggml_tensor* pe, + struct ggml_tensor* context, + int64_t context_img_len = 257) { + // x: [N, n_token, dim] + // e: [N, 6, dim] or [N, T, 6, dim] + // context: [N, context_img_len + context_txt_len, dim] + // return [N, n_token, dim] + if (block_id == 0) { + auto before_proj = std::dynamic_pointer_cast(blocks["before_proj"]); + + c = before_proj->forward(ctx, c); + c = ggml_add(ctx, c, x); + } + + auto after_proj = std::dynamic_pointer_cast(blocks["after_proj"]); + + c = WanAttentionBlock::forward(ctx, c, e, pe, context, context_img_len); + auto c_skip = after_proj->forward(ctx, c); + + return {c_skip, c}; + } + }; + class Head : public GGMLBlock { protected: int dim; @@ -1675,22 +1727,25 @@ namespace WAN { }; struct WanParams { - std::string model_type = "t2v"; - std::tuple patch_size = {1, 2, 2}; - int64_t text_len = 512; - int64_t in_dim = 16; - int64_t dim = 2048; - int64_t ffn_dim = 8192; - int64_t freq_dim = 256; - int64_t text_dim = 4096; - int64_t out_dim = 16; - int64_t num_heads = 16; - int64_t num_layers = 32; - bool qk_norm = true; - bool cross_attn_norm = true; - float eps = 1e-6; - int64_t flf_pos_embed_token_number = 0; - int theta = 10000; + std::string model_type = "t2v"; + std::tuple patch_size = {1, 2, 2}; + int64_t text_len = 512; + int64_t in_dim = 16; + int64_t dim = 2048; + int64_t ffn_dim = 8192; + int64_t freq_dim = 256; + int64_t text_dim = 4096; + int64_t out_dim = 16; + int64_t num_heads = 16; + int64_t num_layers = 32; + int64_t vace_layers = 0; + int64_t vace_in_dim = 96; + std::map vace_layers_mapping = {}; + bool qk_norm = true; + bool cross_attn_norm = true; + float eps = 1e-6; + int64_t flf_pos_embed_token_number = 0; + int theta = 10000; // wan2.1 1.3B: 1536/12, wan2.1/2.2 14B: 5120/40, wan2.2 5B: 3074/24 std::vector axes_dim = {44, 42, 42}; int64_t axes_dim_sum = 128; @@ -1741,6 +1796,31 @@ namespace WAN { if (params.model_type == "i2v") { blocks["img_emb"] = std::shared_ptr(new MLPProj(1280, params.dim, params.flf_pos_embed_token_number)); } + + // vace + if (params.vace_layers > 0) { + for (int i = 0; i < params.vace_layers; i++) { + auto block = std::shared_ptr(new VaceWanAttentionBlock(params.model_type == "t2v", + params.dim, + params.ffn_dim, + params.num_heads, + params.qk_norm, + params.cross_attn_norm, + params.eps, + i, + params.flash_attn)); + blocks["vace_blocks." + std::to_string(i)] = block; + } + + int step = params.num_layers / params.vace_layers; + int n = 0; + for (int i = 0; i < params.num_layers; i += step) { + this->params.vace_layers_mapping[i] = n; + n++; + } + + blocks["vace_patch_embedding"] = std::shared_ptr(new Conv3d(params.vace_in_dim, params.dim, params.patch_size, params.patch_size)); + } } struct ggml_tensor* pad_to_patch_size(struct ggml_context* ctx, @@ -1789,9 +1869,12 @@ namespace WAN { struct ggml_tensor* timestep, struct ggml_tensor* context, struct ggml_tensor* pe, - struct ggml_tensor* clip_fea = NULL, - int64_t N = 1) { + struct ggml_tensor* clip_fea = NULL, + struct ggml_tensor* vace_context = NULL, + float vace_strength = 1.f, + int64_t N = 1) { // x: [N*C, T, H, W], C => in_dim + // vace_context: [N*vace_in_dim, T, H, W] // timestep: [N,] or [T] // context: [N, L, text_dim] // return: [N, t_len*h_len*w_len, out_dim*pt*ph*pw] @@ -1839,10 +1922,35 @@ namespace WAN { context_img_len = clip_fea->ne[1]; // 257 } + // vace_patch_embedding + ggml_tensor* c = NULL; + if (params.vace_layers > 0) { + auto vace_patch_embedding = std::dynamic_pointer_cast(blocks["vace_patch_embedding"]); + + c = vace_patch_embedding->forward(ctx, vace_context); // [N*dim, t_len, h_len, w_len] + c = ggml_reshape_3d(ctx, c, c->ne[0] * c->ne[1] * c->ne[2], c->ne[3] / N, N); // [N, dim, t_len*h_len*w_len] + c = ggml_nn_cont(ctx, ggml_torch_permute(ctx, c, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim] + } + + auto x_orig = x; + for (int i = 0; i < params.num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); x = block->forward(ctx, x, e0, pe, context, context_img_len); + + auto iter = params.vace_layers_mapping.find(i); + if (iter != params.vace_layers_mapping.end()) { + int n = iter->second; + + auto vace_block = std::dynamic_pointer_cast(blocks["vace_blocks." + std::to_string(n)]); + + auto result = vace_block->forward(ctx, c, x_orig, e0, pe, context, context_img_len); + auto c_skip = result.first; + c = result.second; + c_skip = ggml_scale(ctx, c_skip, vace_strength); + x = ggml_add(ctx, x, c_skip); + } } x = head->forward(ctx, x, e); // [N, t_len*h_len*w_len, pt*ph*pw*out_dim] @@ -1857,6 +1965,8 @@ namespace WAN { struct ggml_tensor* pe, struct ggml_tensor* clip_fea = NULL, struct ggml_tensor* time_dim_concat = NULL, + struct ggml_tensor* vace_context = NULL, + float vace_strength = 1.f, int64_t N = 1) { // Forward pass of DiT. // x: [N*C, T, H, W] @@ -1885,7 +1995,7 @@ namespace WAN { t_len = ((x->ne[2] + (std::get<0>(params.patch_size) / 2)) / std::get<0>(params.patch_size)); } - auto out = forward_orig(ctx, x, timestep, context, pe, clip_fea, N); // [N, t_len*h_len*w_len, pt*ph*pw*C] + auto out = forward_orig(ctx, x, timestep, context, pe, clip_fea, vace_context, vace_strength, N); // [N, t_len*h_len*w_len, pt*ph*pw*C] out = unpatchify(ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w] @@ -1920,7 +2030,19 @@ namespace WAN { std::string tensor_name = pair.first; if (tensor_name.find(prefix) == std::string::npos) continue; - size_t pos = tensor_name.find("blocks."); + size_t pos = tensor_name.find("vace_blocks."); + if (pos != std::string::npos) { + tensor_name = tensor_name.substr(pos); // remove prefix + auto items = split_string(tensor_name, '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + if (block_index + 1 > wan_params.vace_layers) { + wan_params.vace_layers = block_index + 1; + } + } + continue; + } + pos = tensor_name.find("blocks."); if (pos != std::string::npos) { tensor_name = tensor_name.substr(pos); // remove prefix auto items = split_string(tensor_name, '.'); @@ -1930,6 +2052,7 @@ namespace WAN { wan_params.num_layers = block_index + 1; } } + continue; } if (tensor_name.find("img_emb") != std::string::npos) { wan_params.model_type = "i2v"; @@ -1951,7 +2074,11 @@ namespace WAN { wan_params.out_dim = 48; wan_params.text_len = 512; } else { - desc = "Wan2.1-T2V-1.3B"; + if (wan_params.vace_layers > 0) { + desc = "Wan2.1-VACE-1.3B"; + } else { + desc = "Wan2.1-T2V-1.3B"; + } wan_params.dim = 1536; wan_params.eps = 1e-06; wan_params.ffn_dim = 8960; @@ -1967,7 +2094,11 @@ namespace WAN { desc = "Wan2.2-I2V-14B"; wan_params.in_dim = 36; } else { - desc = "Wan2.x-T2V-14B"; + if (wan_params.vace_layers > 0) { + desc = "Wan2.x-VACE-14B"; + } else { + desc = "Wan2.x-T2V-14B"; + } wan_params.in_dim = 16; } } else { @@ -2008,7 +2139,9 @@ namespace WAN { struct ggml_tensor* context, struct ggml_tensor* clip_fea = NULL, struct ggml_tensor* c_concat = NULL, - struct ggml_tensor* time_dim_concat = NULL) { + struct ggml_tensor* time_dim_concat = NULL, + struct ggml_tensor* vace_context = NULL, + float vace_strength = 1.f) { struct ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, WAN_GRAPH_SIZE, false); x = to_backend(x); @@ -2017,6 +2150,7 @@ namespace WAN { clip_fea = to_backend(clip_fea); c_concat = to_backend(c_concat); time_dim_concat = to_backend(time_dim_concat); + vace_context = to_backend(vace_context); pe_vec = Rope::gen_wan_pe(x->ne[2], x->ne[1], @@ -2045,7 +2179,9 @@ namespace WAN { context, pe, clip_fea, - time_dim_concat); + time_dim_concat, + vace_context, + vace_strength); ggml_build_forward_expand(gf, out); @@ -2059,10 +2195,12 @@ namespace WAN { struct ggml_tensor* clip_fea = NULL, struct ggml_tensor* c_concat = NULL, struct ggml_tensor* time_dim_concat = NULL, + struct ggml_tensor* vace_context = NULL, + float vace_strength = 1.f, struct ggml_tensor** output = NULL, struct ggml_context* output_ctx = NULL) { auto get_graph = [&]() -> struct ggml_cgraph* { - return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat); + return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength); }; GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx); @@ -2100,7 +2238,7 @@ namespace WAN { struct ggml_tensor* out = NULL; int t0 = ggml_time_ms(); - compute(8, x, timesteps, context, NULL, NULL, NULL, &out, work_ctx); + compute(8, x, timesteps, context, NULL, NULL, NULL, NULL, 1.f, &out, work_ctx); int t1 = ggml_time_ms(); print_ggml_tensor(out);