mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2025-12-13 05:48:56 +00:00
add wan vace t2v support
This commit is contained in:
leejet
parent
1c07fb6fb1
commit
c90ae4e222
@ -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);
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@ -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()) {
|
||||
|
||||
@ -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<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 @@ 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<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);
|
||||
|
||||
202
wan.hpp
202
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);
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user