From c90ae4e22253739c7bfcacdf89334be3258346b8 Mon Sep 17 00:00:00 2001
From: leejet <leejet714@gmail.com>
Date: Sat, 6 Sep 2025 13:39:09 +0800
Subject: [PATCH 1/5] add wan vace t2v support

---
 diffusion_model.hpp  | 140 ++++++++++++++--------------
 ggml_extend.hpp      |  32 +++++++
 stable-diffusion.cpp | 215 +++++++++++++++++++++++--------------------
 wan.hpp              | 202 +++++++++++++++++++++++++++++++++-------
 4 files changed, 386 insertions(+), 203 deletions(-)

diff --git a/diffusion_model.hpp b/diffusion_model.hpp
index 312266e8..80f81f8e 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<ggml_tensor*> ref_latents     = {};
+    int num_video_frames                      = -1;
+    std::vector<struct ggml_tensor*> controls = {};
+    float control_strength                    = 0.f;
+    struct ggml_tensor* vace_context          = NULL;
+    float vace_strength                       = 1.f;
+    std::vector<int> 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<ggml_tensor*> ref_latents     = {},
-                         int num_video_frames                      = -1,
-                         std::vector<struct ggml_tensor*> controls = {},
-                         float control_strength                    = 0.f,
-                         struct ggml_tensor** output               = NULL,
-                         struct ggml_context* output_ctx           = NULL,
-                         std::vector<int> skip_layers              = std::vector<int>())             = 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<ggml_tensor*> ref_latents     = {},
-                 int num_video_frames                      = -1,
-                 std::vector<struct ggml_tensor*> controls = {},
-                 float control_strength                    = 0.f,
-                 struct ggml_tensor** output               = NULL,
-                 struct ggml_context* output_ctx           = NULL,
-                 std::vector<int> skip_layers              = std::vector<int>()) {
-        (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<ggml_tensor*> ref_latents     = {},
-                 int num_video_frames                      = -1,
-                 std::vector<struct ggml_tensor*> controls = {},
-                 float control_strength                    = 0.f,
-                 struct ggml_tensor** output               = NULL,
-                 struct ggml_context* output_ctx           = NULL,
-                 std::vector<int> skip_layers              = std::vector<int>()) {
-        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<ggml_tensor*> ref_latents     = {},
-                 int num_video_frames                      = -1,
-                 std::vector<struct ggml_tensor*> controls = {},
-                 float control_strength                    = 0.f,
-                 struct ggml_tensor** output               = NULL,
-                 struct ggml_context* output_ctx           = NULL,
-                 std::vector<int> skip_layers              = std::vector<int>()) {
-        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<ggml_tensor*> ref_latents     = {},
-                 int num_video_frames                      = -1,
-                 std::vector<struct ggml_tensor*> controls = {},
-                 float control_strength                    = 0.f,
-                 struct ggml_tensor** output               = NULL,
-                 struct ggml_context* output_ctx           = NULL,
-                 std::vector<int> skip_layers              = std::vector<int>()) {
-        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 09657847..cdd60c5a 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<void(ggml_tensor*, int64_t, int64_t, int64_t, int64_t)>& 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<void(ggml_tensor*, int64_t)>& 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 64164a2f..6d7f17a5 100644
--- a/stable-diffusion.cpp
+++ b/stable-diffusion.cpp
@@ -775,7 +775,12 @@ class StableDiffusionGGML {
 
         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 @@ class StableDiffusionGGML {
                         int start_merge_step,
                         SDCondition id_cond,
                         std::vector<ggml_tensor*> ref_latents = {},
-                        ggml_tensor* denoise_mask             = nullptr) {
+                        ggml_tensor* denoise_mask             = NULL,
+                        ggml_tensor* vace_context             = NULL,
+                        float vace_strength                   = 1.f) {
         std::vector<int> skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count);
 
         float cfg_scale     = guidance.txt_cfg;
@@ -1116,32 +1123,30 @@ class StableDiffusionGGML {
                 // 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 @@ class StableDiffusionGGML {
                     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 @@ class StableDiffusionGGML {
             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<size_t>(200 * 1024) * 1024;  // 200 MB
+    params.mem_size = static_cast<size_t>(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 d385cac7..2809a226 100644
--- a/wan.hpp
+++ b/wan.hpp
@@ -1528,12 +1528,12 @@ namespace WAN {
             blocks["ffn.2"] = std::shared_ptr<GGMLBlock>(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<GGMLBlock>(new Linear(dim, dim));
+            }
+            blocks["after_proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim));
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> 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<Linear>(blocks["before_proj"]);
+
+                c = before_proj->forward(ctx, c);
+                c = ggml_add(ctx, c, x);
+            }
+
+            auto after_proj = std::dynamic_pointer_cast<Linear>(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<int, int, int> 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<int, int, int> 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<int, int> 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<int> 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<GGMLBlock>(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<GGMLBlock>(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<GGMLBlock>(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<Conv3d>(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<WanAttentionBlock>(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<VaceWanAttentionBlock>(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);

From 53aeb555bd1de21242691aecd3aa6de97ccc067b Mon Sep 17 00:00:00 2001
From: leejet <leejet714@gmail.com>
Date: Mon, 8 Sep 2025 23:02:38 +0800
Subject: [PATCH 2/5] add --vace-strength option

---
 diffusion_model.hpp   |  2 +-
 examples/cli/main.cpp | 12 ++++++++----
 stable-diffusion.cpp  | 23 +++++++++++++----------
 stable-diffusion.h    |  1 +
 wan.hpp               | 14 +++++++-------
 5 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/diffusion_model.hpp b/diffusion_model.hpp
index 896bed45..995a6a0f 100644
--- a/diffusion_model.hpp
+++ b/diffusion_model.hpp
@@ -14,7 +14,7 @@ struct DiffusionParams {
     struct ggml_tensor* y                     = NULL;
     struct ggml_tensor* guidance              = NULL;
     std::vector<ggml_tensor*> ref_latents     = {};
-    bool increase_ref_index = false;
+    bool increase_ref_index                   = false;
     int num_video_frames                      = -1;
     std::vector<struct ggml_tensor*> controls = {};
     float control_strength                    = 0.f;
diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
index 098c98cc..7779db26 100644
--- a/examples/cli/main.cpp
+++ b/examples/cli/main.cpp
@@ -91,10 +91,10 @@ struct SDParams {
     std::vector<int> high_noise_skip_layers = {7, 8, 9};
     sd_sample_params_t high_noise_sample_params;
 
-    float moe_boundary = 0.875f;
-
-    int video_frames = 1;
-    int fps          = 16;
+    float moe_boundary  = 0.875f;
+    int video_frames    = 1;
+    int fps             = 16;
+    float vace_strength = 1.f;
 
     float strength             = 0.75f;
     float control_strength     = 0.9f;
@@ -186,6 +186,7 @@ void print_params(SDParams params) {
     printf("    chroma_use_t5_mask:                %s\n", params.chroma_use_t5_mask ? "true" : "false");
     printf("    chroma_t5_mask_pad:                %d\n", params.chroma_t5_mask_pad);
     printf("    video_frames:                      %d\n", params.video_frames);
+    printf("    vace_strength:                     %.2f\n", params.vace_strength);
     printf("    fps:                               %d\n", params.fps);
     free(sample_params_str);
     free(high_noise_sample_params_str);
@@ -288,6 +289,7 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --moe-boundary BOUNDARY            timestep boundary for Wan2.2 MoE model. (default: 0.875)\n");
     printf("                                     only enabled if `--high-noise-steps` is set to -1\n");
     printf("  --flow-shift SHIFT                 shift value for Flow models like SD3.x or WAN (default: auto)\n");
+    printf("  --vace-strength                    wan vace strength\n");
     printf("  -v, --verbose                      print extra info\n");
 }
 
@@ -523,6 +525,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         {"", "--control-strength", "", &params.control_strength},
         {"", "--moe-boundary", "", &params.moe_boundary},
         {"", "--flow-shift", "", &params.flow_shift},
+        {"", "--vace-strength", "", &params.vace_strength},
     };
 
     options.bool_options = {
@@ -1244,6 +1247,7 @@ int main(int argc, const char* argv[]) {
             params.strength,
             params.seed,
             params.video_frames,
+            params.vace_strength,
         };
 
         results = generate_video(sd_ctx, &vid_gen_params, &num_results);
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
index d3c0a577..d2bf8173 100644
--- a/stable-diffusion.cpp
+++ b/stable-diffusion.cpp
@@ -1119,15 +1119,15 @@ class StableDiffusionGGML {
             }
 
             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.x                  = noised_input;
+            diffusion_params.timesteps          = timesteps;
+            diffusion_params.guidance           = guidance_tensor;
+            diffusion_params.ref_latents        = ref_latents;
             diffusion_params.increase_ref_index = increase_ref_index;
-            diffusion_params.controls         = controls;
-            diffusion_params.control_strength = control_strength;
-            diffusion_params.vace_context     = vace_context;
-            diffusion_params.vace_strength    = vace_strength;
+            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
@@ -1728,6 +1728,7 @@ void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params) {
     sd_vid_gen_params->seed                                  = -1;
     sd_vid_gen_params->video_frames                          = 6;
     sd_vid_gen_params->moe_boundary                          = 0.875f;
+    sd_vid_gen_params->vace_strength                         = 1.f;
 }
 
 struct sd_ctx_t {
@@ -2644,7 +2645,8 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
                                  {},
                                  false,
                                  denoise_mask,
-                                 vace_context);
+                                 vace_context,
+                                 sd_vid_gen_params->vace_strength);
 
         int64_t sampling_end = ggml_time_ms();
         LOG_INFO("sampling(high noise) completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000);
@@ -2678,7 +2680,8 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
                                           {},
                                           false,
                                           denoise_mask,
-                                          vace_context);
+                                          vace_context,
+                                          sd_vid_gen_params->vace_strength);
 
         int64_t sampling_end = ggml_time_ms();
         LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000);
diff --git a/stable-diffusion.h b/stable-diffusion.h
index 34b0d149..57aad811 100644
--- a/stable-diffusion.h
+++ b/stable-diffusion.h
@@ -211,6 +211,7 @@ typedef struct {
     float strength;
     int64_t seed;
     int video_frames;
+    float vace_strength;
 } sd_vid_gen_params_t;
 
 typedef struct sd_ctx_t sd_ctx_t;
diff --git a/wan.hpp b/wan.hpp
index 03bad4d6..f4cf8ab8 100644
--- a/wan.hpp
+++ b/wan.hpp
@@ -1533,12 +1533,12 @@ namespace WAN {
         }
 
         virtual struct ggml_tensor* forward(struct ggml_context* ctx,
-                                    ggml_backend_t backend,
-                                    struct ggml_tensor* x,
-                                    struct ggml_tensor* e,
-                                    struct ggml_tensor* pe,
-                                    struct ggml_tensor* context,
-                                    int64_t context_img_len = 257) {
+                                            ggml_backend_t backend,
+                                            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]
@@ -1610,7 +1610,7 @@ namespace WAN {
         }
 
         std::pair<ggml_tensor*, ggml_tensor*> forward(struct ggml_context* ctx,
-            ggml_backend_t backend,
+                                                      ggml_backend_t backend,
                                                       struct ggml_tensor* c,
                                                       struct ggml_tensor* x,
                                                       struct ggml_tensor* e,

From 64c5d8ea8f4cfb35ee67a891bf34e48ce571815c Mon Sep 17 00:00:00 2001
From: leejet <leejet714@gmail.com>
Date: Wed, 10 Sep 2025 21:52:52 +0800
Subject: [PATCH 3/5] add vace i2v support

---
 ggml                 |  2 +-
 ggml_extend.hpp      |  3 ++
 stable-diffusion.cpp | 68 +++++++++++++++++++++++++++++++++++++-------
 3 files changed, 62 insertions(+), 11 deletions(-)

diff --git a/ggml b/ggml
index 5fdc78ff..c46da318 160000
--- a/ggml
+++ b/ggml
@@ -1 +1 @@
-Subproject commit 5fdc78fff274094e2a1b155928131983362d8a71
+Subproject commit c46da318b9b6730806196ef7fff67c8160a74c8e
diff --git a/ggml_extend.hpp b/ggml_extend.hpp
index c6a0b7f9..25835b10 100644
--- a/ggml_extend.hpp
+++ b/ggml_extend.hpp
@@ -1543,6 +1543,7 @@ struct GGMLRunner {
             ggml_backend_tensor_copy(t, offload_t);
             std::swap(t->buffer, offload_t->buffer);
             std::swap(t->data, offload_t->data);
+            std::swap(t->extra, offload_t->extra);
 
             t         = ggml_get_next_tensor(params_ctx, t);
             offload_t = ggml_get_next_tensor(offload_ctx, offload_t);
@@ -1573,8 +1574,10 @@ struct GGMLRunner {
         while (t != NULL && offload_t != NULL) {
             t->buffer         = offload_t->buffer;
             t->data           = offload_t->data;
+            t->extra          = offload_t->extra;
             offload_t->buffer = NULL;
             offload_t->data   = NULL;
+            offload_t->extra  = NULL;
 
             t         = ggml_get_next_tensor(params_ctx, t);
             offload_t = ggml_get_next_tensor(offload_ctx, offload_t);
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
index d2bf8173..1e220ba7 100644
--- a/stable-diffusion.cpp
+++ b/stable-diffusion.cpp
@@ -2402,7 +2402,6 @@ 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<size_t>(1024 * 1024) * 1024;  // 1G
-    params.mem_size += width * height * frames * 3 * sizeof(float) * 2;
     params.mem_buffer = NULL;
     params.no_alloc   = false;
     // LOG_DEBUG("mem_size %u ", params.mem_size);
@@ -2430,6 +2429,7 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
     ggml_tensor* concat_latent      = NULL;
     ggml_tensor* denoise_mask       = NULL;
     ggml_tensor* vace_context       = NULL;
+    int64_t ref_image_num           = 0;  // for vace
     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") {
@@ -2526,6 +2526,20 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
     } 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");
+        int64_t t1                    = ggml_time_ms();
+        ggml_tensor* ref_image_latent = NULL;
+        if (sd_vid_gen_params->init_image.data) {
+            ggml_tensor* ref_img = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
+            sd_image_to_tensor(sd_vid_gen_params->init_image.data, ref_img);
+            ref_img = ggml_reshape_4d(work_ctx, ref_img, width, height, 1, 3);
+
+            ref_image_latent = sd_ctx->sd->encode_first_stage(work_ctx, ref_img);  // [b*c, 1, h/16, w/16]
+            sd_ctx->sd->process_latent_in(ref_image_latent);
+            auto zero_latent = ggml_dup_tensor(work_ctx, ref_image_latent);
+            ggml_set_f32(zero_latent, 0.f);
+            ref_image_latent = ggml_tensor_concat(work_ctx, ref_image_latent, zero_latent, 3);  // [b*2*c, 1, h/16, w/16]
+        }
+
         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);
@@ -2549,23 +2563,43 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
         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]
+        int64_t length = inactive->ne[2];
+        if (ref_image_latent) {
+            length += 1;
+            frames        = (length - 1) * 4 + 1;
+            ref_image_num = 1;
+        }
+        vace_context = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, inactive->ne[0], inactive->ne[1], length, 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);
+            if (i3 < 32) {
+                if (ref_image_latent && i2 == 0) {
+                    value = ggml_tensor_get_f32(ref_image_latent, i0, i1, 0, i3);
+                } else {
+                    if (i3 < 16) {
+                        value = ggml_tensor_get_f32(inactive, i0, i1, i2 - ref_image_num, i3);
+                    } else {
+                        value = ggml_tensor_get_f32(reactive, i0, i1, i2 - ref_image_num, 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);
+                if (ref_image_latent && i2 == 0) {
+                    value = 0.f;
+                } else {
+                    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 - ref_image_num, 0);
+                }
             }
             ggml_tensor_set_f32(vace_context, value, i0, i1, i2, i3);
         });
+        int64_t t2 = ggml_time_ms();
+        LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
     }
 
+    print_ggml_tensor(vace_context);
+
     if (init_latent == NULL) {
         init_latent = generate_init_latent(sd_ctx, work_ctx, width, height, frames, true);
     }
@@ -2690,6 +2724,20 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
         }
     }
 
+    if (ref_image_num > 0) {
+        ggml_tensor* trim_latent = ggml_new_tensor_4d(work_ctx,
+                                                      GGML_TYPE_F32,
+                                                      final_latent->ne[0],
+                                                      final_latent->ne[1],
+                                                      final_latent->ne[2] - ref_image_num,
+                                                      final_latent->ne[3]);
+        ggml_tensor_iter(trim_latent, [&](ggml_tensor* trim_latent, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+            float value = ggml_tensor_get_f32(final_latent, i0, i1, i2 + ref_image_num, i3);
+            ggml_tensor_set_f32(trim_latent, value, i0, i1, i2, i3);
+        });
+        final_latent = trim_latent;
+    }
+
     int64_t t4 = ggml_time_ms();
     LOG_INFO("generating latent video completed, taking %.2fs", (t4 - t2) * 1.0f / 1000);
     struct ggml_tensor* vid = sd_ctx->sd->decode_first_stage(work_ctx, final_latent, true);

From e751ae6d6f35fdd735abf3236650ecae9b6474f4 Mon Sep 17 00:00:00 2001
From: leejet <leejet714@gmail.com>
Date: Fri, 12 Sep 2025 00:32:16 +0800
Subject: [PATCH 4/5] fix the processing of vace_context

---
 ggml                 |  2 +-
 ggml_extend.hpp      | 15 +++++++++++++++
 stable-diffusion.cpp |  4 +---
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/ggml b/ggml
index c46da318..5fdc78ff 160000
--- a/ggml
+++ b/ggml
@@ -1 +1 @@
-Subproject commit c46da318b9b6730806196ef7fff67c8160a74c8e
+Subproject commit 5fdc78fff274094e2a1b155928131983362d8a71
diff --git a/ggml_extend.hpp b/ggml_extend.hpp
index 25835b10..390fc49b 100644
--- a/ggml_extend.hpp
+++ b/ggml_extend.hpp
@@ -255,6 +255,21 @@ __STATIC_INLINE__ void ggml_tensor_iter(
     }
 }
 
+
+__STATIC_INLINE__ void ggml_tensor_diff(
+    ggml_tensor* a,
+    ggml_tensor* b,
+    float gap = 0.1f) {
+    GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));
+    ggml_tensor_iter(a, [&] (ggml_tensor* a, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+        float a_value = ggml_tensor_get_f32(a, i0, i1, i2, i3);
+        float b_value = ggml_tensor_get_f32(b, i0, i1, i2, i3);
+        if (abs(a_value - b_value) > gap) {
+            LOG_WARN("[%ld, %ld, %ld, %ld] %f %f", i3, i2, i1, i0, a_value, b_value);
+        }
+    });
+}
+
 __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 1e220ba7..e2695084 100644
--- a/stable-diffusion.cpp
+++ b/stable-diffusion.cpp
@@ -2579,7 +2579,7 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
                     if (i3 < 16) {
                         value = ggml_tensor_get_f32(inactive, i0, i1, i2 - ref_image_num, i3);
                     } else {
-                        value = ggml_tensor_get_f32(reactive, i0, i1, i2 - ref_image_num, i3);
+                        value = ggml_tensor_get_f32(reactive, i0, i1, i2 - ref_image_num, i3 - 16);
                     }
                 }
             } else {  // mask
@@ -2598,8 +2598,6 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
         LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
     }
 
-    print_ggml_tensor(vace_context);
-
     if (init_latent == NULL) {
         init_latent = generate_init_latent(sd_ctx, work_ctx, width, height, frames, true);
     }

From f68ce0582abac6f6d41c81fc77965ca48efef9fd Mon Sep 17 00:00:00 2001
From: leejet <leejet714@gmail.com>
Date: Sat, 13 Sep 2025 16:08:56 +0800
Subject: [PATCH 5/5] add vace v2v support

---
 examples/cli/main.cpp | 75 +++++++++++++++++++++++++++++++++++++------
 ggml_extend.hpp       | 57 ++++++++++++--------------------
 preprocessing.hpp     | 17 +++++-----
 stable-diffusion.cpp  | 24 ++++++++------
 stable-diffusion.h    | 16 ++++-----
 upscaler.cpp          |  2 +-
 6 files changed, 118 insertions(+), 73 deletions(-)

diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
index 7779db26..97dd0109 100644
--- a/examples/cli/main.cpp
+++ b/examples/cli/main.cpp
@@ -35,6 +35,8 @@
 #define SAFE_STR(s) ((s) ? (s) : "")
 #define BOOL_STR(b) ((b) ? "true" : "false")
 
+namespace fs = std::filesystem;
+
 const char* modes_str[] = {
     "img_gen",
     "vid_gen",
@@ -75,6 +77,7 @@ struct SDParams {
     std::string mask_image_path;
     std::string control_image_path;
     std::vector<std::string> ref_image_paths;
+    std::string control_video_path;
     bool increase_ref_index = false;
 
     std::string prompt;
@@ -158,6 +161,7 @@ void print_params(SDParams params) {
     for (auto& path : params.ref_image_paths) {
         printf("        %s\n", path.c_str());
     };
+    printf("    control_video_path:                %s\n", params.control_video_path.c_str());
     printf("    increase_ref_index:                %s\n", params.increase_ref_index ? "true" : "false");
     printf("    offload_params_to_cpu:             %s\n", params.offload_params_to_cpu ? "true" : "false");
     printf("    clip_on_cpu:                       %s\n", params.clip_on_cpu ? "true" : "false");
@@ -178,7 +182,7 @@ void print_params(SDParams params) {
     printf("    flow_shift:                        %.2f\n", params.flow_shift);
     printf("    strength(img2img):                 %.2f\n", params.strength);
     printf("    rng:                               %s\n", sd_rng_type_name(params.rng_type));
-    printf("    seed:                              %ld\n", params.seed);
+    printf("    seed:                              %zd\n", params.seed);
     printf("    batch_count:                       %d\n", params.batch_count);
     printf("    vae_tiling:                        %s\n", params.vae_tiling ? "true" : "false");
     printf("    upscale_repeats:                   %d\n", params.upscale_repeats);
@@ -226,6 +230,9 @@ void print_usage(int argc, const char* argv[]) {
     printf("  -i, --end-img [IMAGE]              path to the end image, required by flf2v\n");
     printf("  --control-image [IMAGE]            path to image condition, control net\n");
     printf("  -r, --ref-image [PATH]             reference image for Flux Kontext models (can be used multiple times) \n");
+    printf("  --control-video [PATH]             path to control video frames, It must be a directory path.");
+    printf("                                     The video frames inside should be stored as images in lexicographical (character) order\n");
+    printf("                                     For example, if the control video path is `frames`, the directory contain images such as 00.png, 01.png, … etc.\n");
     printf("  --increase-ref-index               automatically increase the indices of references images based on the order they are listed (starting with 1).\n");
     printf("  -o, --output OUTPUT                path to write result image to (default: ./output.png)\n");
     printf("  -p, --prompt [PROMPT]              the prompt to render\n");
@@ -484,6 +491,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         {"", "--input-id-images-dir", "", &params.input_id_images_path},
         {"", "--mask", "", &params.mask_image_path},
         {"", "--control-image", "", &params.control_image_path},
+        {"", "--control-video", "", &params.control_video_path},
         {"-o", "--output", "", &params.output_path},
         {"-p", "--prompt", "", &params.prompt},
         {"-n", "--negative-prompt", "", &params.negative_prompt},
@@ -1062,6 +1070,7 @@ int main(int argc, const char* argv[]) {
     sd_image_t control_image = {(uint32_t)params.width, (uint32_t)params.height, 3, NULL};
     sd_image_t mask_image    = {(uint32_t)params.width, (uint32_t)params.height, 1, NULL};
     std::vector<sd_image_t> ref_images;
+    std::vector<sd_image_t> control_frames;
 
     auto release_all_resources = [&]() {
         free(init_image.data);
@@ -1073,6 +1082,11 @@ int main(int argc, const char* argv[]) {
             ref_image.data = NULL;
         }
         ref_images.clear();
+        for (auto frame : control_frames) {
+            free(frame.data);
+            frame.data = NULL;
+        }
+        control_frames.clear();
     };
 
     if (params.init_image_path.size() > 0) {
@@ -1131,14 +1145,12 @@ int main(int argc, const char* argv[]) {
             return 1;
         }
         if (params.canny_preprocess) {  // apply preprocessor
-            control_image.data = preprocess_canny(control_image.data,
-                                                  control_image.width,
-                                                  control_image.height,
-                                                  0.08f,
-                                                  0.08f,
-                                                  0.8f,
-                                                  1.0f,
-                                                  false);
+            preprocess_canny(control_image,
+                             0.08f,
+                             0.08f,
+                             0.8f,
+                             1.0f,
+                             false);
         }
     }
 
@@ -1160,6 +1172,48 @@ int main(int argc, const char* argv[]) {
         }
     }
 
+    if (!params.control_video_path.empty()) {
+        std::string dir = params.control_video_path;
+
+        if (!fs::exists(dir) || !fs::is_directory(dir)) {
+            fprintf(stderr, "'%s' is not a valid directory\n", dir.c_str());
+            release_all_resources();
+            return 1;
+        }
+
+        for (const auto& entry : fs::directory_iterator(dir)) {
+            if (!entry.is_regular_file())
+                continue;
+
+            std::string path = entry.path().string();
+            std::string ext  = entry.path().extension().string();
+            std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
+
+            if (ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".bmp") {
+                if (params.verbose) {
+                    printf("load control frame %zu from '%s'\n", control_frames.size(), path.c_str());
+                }
+                int width             = 0;
+                int height            = 0;
+                uint8_t* image_buffer = load_image(path.c_str(), width, height, params.width, params.height);
+                if (image_buffer == NULL) {
+                    fprintf(stderr, "load image from '%s' failed\n", path.c_str());
+                    release_all_resources();
+                    return 1;
+                }
+
+                control_frames.push_back({(uint32_t)params.width,
+                                          (uint32_t)params.height,
+                                          3,
+                                          image_buffer});
+
+                if (control_frames.size() >= params.video_frames) {
+                    break;
+                }
+            }
+        }
+    }
+
     if (params.mode == VID_GEN) {
         vae_decode_only = false;
     }
@@ -1239,6 +1293,8 @@ int main(int argc, const char* argv[]) {
             params.clip_skip,
             init_image,
             end_image,
+            control_frames.data(),
+            (int)control_frames.size(),
             params.width,
             params.height,
             params.sample_params,
@@ -1290,7 +1346,6 @@ int main(int argc, const char* argv[]) {
 
     // create directory if not exists
     {
-        namespace fs            = std::filesystem;
         const fs::path out_path = params.output_path;
         if (const fs::path out_dir = out_path.parent_path(); !out_dir.empty()) {
             std::error_code ec;
diff --git a/ggml_extend.hpp b/ggml_extend.hpp
index 390fc49b..a2a7435e 100644
--- a/ggml_extend.hpp
+++ b/ggml_extend.hpp
@@ -173,6 +173,14 @@ __STATIC_INLINE__ ggml_fp16_t ggml_tensor_get_f16(const ggml_tensor* tensor, int
     return *(ggml_fp16_t*)((char*)(tensor->data) + i * tensor->nb[3] + j * tensor->nb[2] + k * tensor->nb[1] + l * tensor->nb[0]);
 }
 
+__STATIC_INLINE__ float sd_image_get_f32(sd_image_t image, int iw, int ih, int ic, bool scale = true) {
+    float value = *(image.data + ih * image.width * image.channel + iw * image.channel + ic);
+    if (scale) {
+        value /= 255.f;
+    }
+    return value;
+}
+
 static struct ggml_tensor* get_tensor_from_graph(struct ggml_cgraph* gf, const char* name) {
     struct ggml_tensor* res = NULL;
     for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
@@ -255,13 +263,12 @@ __STATIC_INLINE__ void ggml_tensor_iter(
     }
 }
 
-
 __STATIC_INLINE__ void ggml_tensor_diff(
     ggml_tensor* a,
     ggml_tensor* b,
     float gap = 0.1f) {
     GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));
-    ggml_tensor_iter(a, [&] (ggml_tensor* a, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+    ggml_tensor_iter(a, [&](ggml_tensor* a, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
         float a_value = ggml_tensor_get_f32(a, i0, i1, i2, i3);
         float b_value = ggml_tensor_get_f32(b, i0, i1, i2, i3);
         if (abs(a_value - b_value) > gap) {
@@ -401,42 +408,18 @@ __STATIC_INLINE__ uint8_t* sd_tensor_to_image(struct ggml_tensor* input, int idx
     return image_data;
 }
 
-__STATIC_INLINE__ void sd_image_to_tensor(const uint8_t* image_data,
-                                          struct ggml_tensor* output,
+__STATIC_INLINE__ void sd_image_to_tensor(sd_image_t image,
+                                          ggml_tensor* tensor,
                                           bool scale = true) {
-    int64_t width    = output->ne[0];
-    int64_t height   = output->ne[1];
-    int64_t channels = output->ne[2];
-    GGML_ASSERT(channels == 3 && output->type == GGML_TYPE_F32);
-    for (int iy = 0; iy < height; iy++) {
-        for (int ix = 0; ix < width; ix++) {
-            for (int k = 0; k < channels; k++) {
-                float value = *(image_data + iy * width * channels + ix * channels + k);
-                if (scale) {
-                    value /= 255.f;
-                }
-                ggml_tensor_set_f32(output, value, ix, iy, k);
-            }
-        }
-    }
-}
-
-__STATIC_INLINE__ void sd_mask_to_tensor(const uint8_t* image_data,
-                                         struct ggml_tensor* output,
-                                         bool scale = true) {
-    int64_t width    = output->ne[0];
-    int64_t height   = output->ne[1];
-    int64_t channels = output->ne[2];
-    GGML_ASSERT(channels == 1 && output->type == GGML_TYPE_F32);
-    for (int iy = 0; iy < height; iy++) {
-        for (int ix = 0; ix < width; ix++) {
-            float value = *(image_data + iy * width * channels + ix);
-            if (scale) {
-                value /= 255.f;
-            }
-            ggml_tensor_set_f32(output, value, ix, iy);
-        }
-    }
+    GGML_ASSERT(image.width == tensor->ne[0]);
+    GGML_ASSERT(image.height == tensor->ne[1]);
+    GGML_ASSERT(image.channel == tensor->ne[2]);
+    GGML_ASSERT(1 == tensor->ne[3]);
+    GGML_ASSERT(tensor->type == GGML_TYPE_F32);
+    ggml_tensor_iter(tensor, [&](ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+        float value = sd_image_get_f32(image, i0, i1, i2, scale);
+        ggml_tensor_set_f32(tensor, value, i0, i1, i2, i3);
+    });
 }
 
 __STATIC_INLINE__ void sd_apply_mask(struct ggml_tensor* image_data,
diff --git a/preprocessing.hpp b/preprocessing.hpp
index 4ea1dbab..08df4a77 100644
--- a/preprocessing.hpp
+++ b/preprocessing.hpp
@@ -162,7 +162,7 @@ void threshold_hystersis(struct ggml_tensor* img, float high_threshold, float lo
     }
 }
 
-uint8_t* preprocess_canny(uint8_t* img, int width, int height, float high_threshold, float low_threshold, float weak, float strong, bool inverse) {
+bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold, float weak, float strong, bool inverse) {
     struct ggml_init_params params;
     params.mem_size               = static_cast<size_t>(10 * 1024 * 1024);  // 10
     params.mem_buffer             = NULL;
@@ -171,7 +171,7 @@ uint8_t* preprocess_canny(uint8_t* img, int width, int height, float high_thresh
 
     if (!work_ctx) {
         LOG_ERROR("ggml_init() failed");
-        return NULL;
+        return false;
     }
 
     float kX[9] = {
@@ -192,8 +192,8 @@ uint8_t* preprocess_canny(uint8_t* img, int width, int height, float high_thresh
     struct ggml_tensor* sf_ky = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 3, 3, 1, 1);
     memcpy(sf_ky->data, kY, ggml_nbytes(sf_ky));
     gaussian_kernel(gkernel);
-    struct ggml_tensor* image      = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
-    struct ggml_tensor* image_gray = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 1, 1);
+    struct ggml_tensor* image      = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 3, 1);
+    struct ggml_tensor* image_gray = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 1, 1);
     struct ggml_tensor* iX         = ggml_dup_tensor(work_ctx, image_gray);
     struct ggml_tensor* iY         = ggml_dup_tensor(work_ctx, image_gray);
     struct ggml_tensor* G          = ggml_dup_tensor(work_ctx, image_gray);
@@ -209,8 +209,8 @@ uint8_t* preprocess_canny(uint8_t* img, int width, int height, float high_thresh
     non_max_supression(image_gray, G, tetha);
     threshold_hystersis(image_gray, high_threshold, low_threshold, weak, strong);
     // to RGB channels
-    for (int iy = 0; iy < height; iy++) {
-        for (int ix = 0; ix < width; ix++) {
+    for (int iy = 0; iy < img.height; iy++) {
+        for (int ix = 0; ix < img.width; ix++) {
             float gray = ggml_tensor_get_f32(image_gray, ix, iy);
             gray       = inverse ? 1.0f - gray : gray;
             ggml_tensor_set_f32(image, gray, ix, iy);
@@ -218,10 +218,11 @@ uint8_t* preprocess_canny(uint8_t* img, int width, int height, float high_thresh
             ggml_tensor_set_f32(image, gray, ix, iy, 2);
         }
     }
-    free(img);
     uint8_t* output = sd_tensor_to_image(image);
+    free(img.data);
+    img.data = output;
     ggml_free(work_ctx);
-    return output;
+    return true;
 }
 
 #endif  // __PREPROCESSING_HPP__
\ No newline at end of file
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
index e2695084..807b624a 100644
--- a/stable-diffusion.cpp
+++ b/stable-diffusion.cpp
@@ -952,7 +952,7 @@ class StableDiffusionGGML {
                     free(resized_image.data);
                     resized_image.data = NULL;
                 } else {
-                    sd_image_to_tensor(init_image.data, init_img);
+                    sd_image_to_tensor(init_image, init_img);
                 }
                 if (augmentation_level > 0.f) {
                     struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, init_img);
@@ -1947,7 +1947,7 @@ sd_image_t* generate_image_internal(sd_ctx_t* sd_ctx,
     struct ggml_tensor* image_hint = NULL;
     if (control_image.data != NULL) {
         image_hint = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
-        sd_image_to_tensor(control_image.data, image_hint);
+        sd_image_to_tensor(control_image, image_hint);
     }
 
     // Sample
@@ -2208,8 +2208,8 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g
         ggml_tensor* init_img = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
         ggml_tensor* mask_img = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 1, 1);
 
-        sd_mask_to_tensor(sd_img_gen_params->mask_image.data, mask_img);
-        sd_image_to_tensor(sd_img_gen_params->init_image.data, init_img);
+        sd_image_to_tensor(sd_img_gen_params->mask_image, mask_img);
+        sd_image_to_tensor(sd_img_gen_params->init_image, init_img);
 
         if (sd_version_is_inpaint(sd_ctx->sd->version)) {
             int64_t mask_channels = 1;
@@ -2300,7 +2300,7 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g
                                               sd_img_gen_params->ref_images[i].height,
                                               3,
                                               1);
-        sd_image_to_tensor(sd_img_gen_params->ref_images[i].data, img);
+        sd_image_to_tensor(sd_img_gen_params->ref_images[i], img);
 
         ggml_tensor* latent = NULL;
         if (sd_ctx->sd->use_tiny_autoencoder) {
@@ -2401,7 +2401,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<size_t>(1024 * 1024) * 1024;  // 1G
+    params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1G
     params.mem_buffer = NULL;
     params.no_alloc   = false;
     // LOG_DEBUG("mem_size %u ", params.mem_size);
@@ -2500,7 +2500,7 @@ 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* init_img = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
-        sd_image_to_tensor(sd_vid_gen_params->init_image.data, init_img);
+        sd_image_to_tensor(sd_vid_gen_params->init_image, init_img);
         init_img = ggml_reshape_4d(work_ctx, init_img, width, height, 1, 3);
 
         auto init_image_latent = sd_ctx->sd->encode_first_stage(work_ctx, init_img);  // [b*c, 1, h/16, w/16]
@@ -2530,7 +2530,7 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
         ggml_tensor* ref_image_latent = NULL;
         if (sd_vid_gen_params->init_image.data) {
             ggml_tensor* ref_img = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, width, height, 3, 1);
-            sd_image_to_tensor(sd_vid_gen_params->init_image.data, ref_img);
+            sd_image_to_tensor(sd_vid_gen_params->init_image, ref_img);
             ref_img = ggml_reshape_4d(work_ctx, ref_img, width, height, 1, 3);
 
             ref_image_latent = sd_ctx->sd->encode_first_stage(work_ctx, ref_img);  // [b*c, 1, h/16, w/16]
@@ -2541,7 +2541,13 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
         }
 
         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_iter(control_video, [&](ggml_tensor* control_video, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+            float value = 0.5f;
+            if (i2 < sd_vid_gen_params->control_frames_size) {
+                value = sd_image_get_f32(sd_vid_gen_params->control_frames[i2], i0, i1, i3);
+            }
+            ggml_tensor_set_f32(control_video, value, i0, i1, i2, i3);
+        });
         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);
diff --git a/stable-diffusion.h b/stable-diffusion.h
index 57aad811..3abe195d 100644
--- a/stable-diffusion.h
+++ b/stable-diffusion.h
@@ -203,6 +203,8 @@ typedef struct {
     int clip_skip;
     sd_image_t init_image;
     sd_image_t end_image;
+    sd_image_t* control_frames;
+    int control_frames_size;
     int width;
     int height;
     sd_sample_params_t sample_params;
@@ -267,14 +269,12 @@ SD_API bool convert(const char* input_path,
                     enum sd_type_t output_type,
                     const char* tensor_type_rules);
 
-SD_API uint8_t* preprocess_canny(uint8_t* img,
-                                 int width,
-                                 int height,
-                                 float high_threshold,
-                                 float low_threshold,
-                                 float weak,
-                                 float strong,
-                                 bool inverse);
+SD_API bool preprocess_canny(sd_image_t image,
+                             float high_threshold,
+                             float low_threshold,
+                             float weak,
+                             float strong,
+                             bool inverse);
 
 #ifdef __cplusplus
 }
diff --git a/upscaler.cpp b/upscaler.cpp
index 4ab0b73c..652453ae 100644
--- a/upscaler.cpp
+++ b/upscaler.cpp
@@ -82,7 +82,7 @@ struct UpscalerGGML {
         }
         LOG_DEBUG("upscale work buffer size: %.2f MB", params.mem_size / 1024.f / 1024.f);
         ggml_tensor* input_image_tensor = ggml_new_tensor_4d(upscale_ctx, GGML_TYPE_F32, input_image.width, input_image.height, 3, 1);
-        sd_image_to_tensor(input_image.data, input_image_tensor);
+        sd_image_to_tensor(input_image, input_image_tensor);
 
         ggml_tensor* upscaled = ggml_new_tensor_4d(upscale_ctx, GGML_TYPE_F32, output_width, output_height, 3, 1);
         auto on_tiling        = [&](ggml_tensor* in, ggml_tensor* out, bool init) {