From 0bd72f075aa6aef5076309d6bda7ed2ac7e6e5b1 Mon Sep 17 00:00:00 2001 From: George <35490284+noctrex@users.noreply.github.com> Date: Sun, 13 Sep 2026 18:33:50 +0300 Subject: [PATCH] feat: add Wan2.2 S2V (audio+img-to-video) support (#1925) Co-authored-by: leejet --- docs/wan.md | 49 ++++ examples/cli/main.cpp | 8 +- examples/common/common.cpp | 19 ++ examples/common/common.h | 1 + examples/server/async_jobs.cpp | 5 +- include/stable-diffusion.h | 5 +- src/conditioning/wan_audio.cpp | 103 ++++++++ src/conditioning/wan_audio.h | 32 +++ src/core/ggml_extend.cpp | 77 +++++- src/core/ggml_extend.h | 15 +- src/core/ggml_extend_backend.cpp | 6 + src/core/ggml_extend_backend.h | 1 + src/model.h | 3 +- src/model/audio/wav2vec2.hpp | 413 ++++++++++++++++++++++++++++++ src/model/common/ggml_block.hpp | 57 ++++- src/model/common/rope.hpp | 5 +- src/model/diffusion/model.hpp | 2 + src/model/diffusion/wan.hpp | 190 +++++++++++--- src/model/diffusion/wan_audio.hpp | 215 ++++++++++++++++ src/model_component.h | 3 + src/model_io/safetensors_io.cpp | 5 + src/model_loader.cpp | 9 + src/pipeline/diffusion_engine.cpp | 55 +++- src/pipeline/diffusion_engine.h | 6 + src/pipeline/generation.h | 4 +- src/pipeline/model_builders.cpp | 11 + src/pipeline/model_builders.h | 4 + src/pipeline/request.cpp | 5 +- src/pipeline/video.cpp | 150 ++++++++++- src/runtime/audio_processing.cpp | 97 +++++++ src/runtime/audio_processing.h | 22 ++ src/stable-diffusion.cpp | 13 +- 32 files changed, 1533 insertions(+), 57 deletions(-) create mode 100644 src/conditioning/wan_audio.cpp create mode 100644 src/conditioning/wan_audio.h create mode 100644 src/model/audio/wav2vec2.hpp create mode 100644 src/model/diffusion/wan_audio.hpp create mode 100644 src/runtime/audio_processing.cpp create mode 100644 src/runtime/audio_processing.h diff --git a/docs/wan.md b/docs/wan.md index 02b46d62..c7cf1dd9 100644 --- a/docs/wan.md +++ b/docs/wan.md @@ -34,6 +34,10 @@ - Wan2.2 I2V A14B - safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/tree/main/split_files/diffusion_models - gguf: https://huggingface.co/QuantStack/Wan2.2-I2V-A14B-GGUF/tree/main + - Wan2.2 S2V 14B + - safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/tree/main/split_files/diffusion_models + - gguf: https://huggingface.co/QuantStack/Wan2.2-S2V-14B-GGUF/tree/main + - int8_convrot safetensors: https://huggingface.co/noctrex/Wan2.2-S2V-14B-int8_convrot - Download vae - wan_2.1_vae (for all the wan model except Wan2.2 TI2V 5B) - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors @@ -49,6 +53,9 @@ - Download clip_vison_h (for Wan2.1 I2V/FLF2V only) - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/clip_vision/clip_vision_h.safetensors +- Download audio_encoder (for Wan2.2 S2V only) + - safetensors: https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/blob/main/split_files/audio_encoders/wav2vec2_large_english_fp16.safetensors + ## Examples @@ -94,6 +101,48 @@ +### Wan2.2 S2V 14B + +Audio-driven video (speech-to-video). The reference image (`-i`) is the speaker +portrait, `--audio` is the driving audio track and `--audio-encoder` is the +wav2vec2 audio encoder. Wan2.2 S2V requires the wan_2.1 vae (16 channel), not +the wan2.2 vae. + +``` +.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\wan2.2_s2v-14B-Q8_0.gguf --audio-encoder ..\models\audio_encoders\wav2vec2_large_english_fp16.safetensors --vae ..\models\vae\wan_2.1_vae.safetensors --t5xxl ..\models\text_encoders\umt5-xxl-encoder-Q8_0.gguf -p "a person is talking" --cfg-scale 6.0 --steps 20 --sampling-method euler -v -n "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" -W 832 -H 480 --diffusion-fa --offload-to-cpu --vae-tiling --video-frames 81 -i ..\assets\cat_with_sd_cpp_42.png --audio .\input\speech.wav --flow-shift 3.0 +``` + +Notes: + +- Recommended settings: `--sampling-method euler --steps 20 --cfg-scale 6.0`. + `dpm++2m` produces heavy artifacts on S2V. 4 steps with the lightning LoRA + (below) is the fast option. +- Resolutions: width and height must be multiples of 16; the examples use + multiples of 64. 832x480 is a fast starting point; generation cost scales + with pixel area. +- `--audio` accepts a WAV file; it is downmixed to mono and resampled to 16 kHz + internally. Audio longer than the video is truncated, video longer than the + audio is padded with silence. Pick `--video-frames` to match the audio: + roughly `audio_seconds * 16` frames, capped at one chunk (77-81 frames, + ~5 s at the model's 16 fps). 33, 77 and 81 map to clean latent frame counts. +- S2V always uses 16 fps. Other requested frame rates are automatically + changed to 16 with a warning, including the CLI and server video output. + `generate_video()` returns the actual frame rate through `fps_out`; C API + callers should use that value when encoding the output video. +- One generation covers the first S2V chunk window (`--video-frames` frames). + Long-video chunked extend mode is not implemented yet. +- Speed: the lightx2v lightning LoRA works with S2V at 4 steps and + `--cfg-scale 1.0`. Use the **low_noise** variant; + the high_noise variant produces artifacts on S2V: + + ``` + --lora-model-dir ..\models\loras + -p "..." + --cfg-scale 1.0 --steps 4 + ``` + + Expect some quality/dynamics loss compared to the full 20-step run. + ### Wan2.2 T2V A14B T2I ``` diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 22b26da7..a3cc1dde 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -419,7 +419,8 @@ void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy, LOG_ERROR("save preview image to '%s' failed", path.string().c_str()); } } else { - if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, cli_params->preview_fps, cli_params->compression_quality) != 0) { + int fps = cli_params->preview_method == PREVIEW_PROJ ? cli_params->preview_fps / 4 : cli_params->preview_fps; + if (create_video_from_sd_images(cli_params->preview_path.c_str(), image, frame_count, fps, cli_params->compression_quality) != 0) { LOG_ERROR("save preview video to '%s' failed", cli_params->preview_path.c_str()); } } @@ -687,8 +688,6 @@ int main(int argc, const char* argv[]) { } } cli_params.preview_fps = gen_params.fps; - if (cli_params.preview_method == PREVIEW_PROJ) - cli_params.preview_fps /= 4; sd_set_preview_callback(step_callback, cli_params.preview_method, @@ -951,9 +950,10 @@ int main(int argc, const char* argv[]) { } else if (cli_params.mode == VID_GEN) { sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t(); sd_image_t* generated_video = nullptr; - if (!generate_video(sd_ctx.get(), &vid_gen_params, &generated_video, &num_results, &generated_audio)) { + if (!generate_video(sd_ctx.get(), &vid_gen_params, &generated_video, &num_results, &generated_audio, &cli_params.preview_fps)) { generated_video = nullptr; } + gen_params.fps = cli_params.preview_fps; results.adopt(generated_video, num_results); } diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 2357dcbb..04bd89cc 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -460,6 +460,11 @@ ArgOptions SDContextParams::get_options() { "path to standalone LTX audio vae model", 0, &audio_vae_path}, + {"", + "--audio-encoder", + "path to wav2vec2 audio encoder model (Wan2.2 S2V)", + 0, + &audio_encoder_path}, {"", "--taesd", "path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)", @@ -898,6 +903,7 @@ std::string SDContextParams::to_string() const { << " vae_path: \"" << vae_path << "\",\n" << " vae_format: \"" << vae_format << "\",\n" << " audio_vae_path: \"" << audio_vae_path << "\",\n" + << " audio_encoder_path: \"" << audio_encoder_path << "\",\n" << " taesd_path: \"" << taesd_path << "\",\n" << " esrgan_path: \"" << esrgan_path << "\",\n" << " control_net_path: \"" << control_net_path << "\",\n" @@ -963,6 +969,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.embeddings_connectors_path = embeddings_connectors_path.c_str(); sd_ctx_params.vae_path = vae_path.c_str(); sd_ctx_params.audio_vae_path = audio_vae_path.c_str(); + sd_ctx_params.audio_encoder_path = audio_encoder_path.c_str(); sd_ctx_params.taesd_path = taesd_path.c_str(); sd_ctx_params.control_net_path = control_net_path.c_str(); sd_ctx_params.ip_adapter_path = ip_adapter_path.c_str(); @@ -1515,6 +1522,14 @@ ArgOptions SDGenerationParams::get_options() { return 1; }; + auto on_audio_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + ref_audio_paths.push_back(argv[index]); + return 1; + }; + auto on_cache_mode_arg = [&](int argc, const char** argv, int index) { if (++index >= argc) { return -1; @@ -1704,6 +1719,10 @@ ArgOptions SDGenerationParams::get_options() { "--ref-audio", "standalone WAV reference for MiniMax-H3 Ref2VA (can be used multiple times)", on_ref_audio_arg}, + {"", + "--audio", + "driving audio track (Wan2.2 S2V; can be used once)", + on_audio_arg}, {"", "--cache-mode", "caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT Chebyshev+Taylor forecasting)", diff --git a/examples/common/common.h b/examples/common/common.h index 660fef4f..8a33edcf 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -131,6 +131,7 @@ struct SDContextParams { std::string vae_path; std::string vae_format = "auto"; std::string audio_vae_path; + std::string audio_encoder_path; std::string taesd_path; std::string esrgan_path; std::string control_net_path; diff --git a/examples/server/async_jobs.cpp b/examples/server/async_jobs.cpp index 1933dec1..8c4adf27 100644 --- a/examples/server/async_jobs.cpp +++ b/examples/server/async_jobs.cpp @@ -245,7 +245,7 @@ bool execute_vid_gen_job(ServerRuntime& runtime, { std::lock_guard lock(*runtime.sd_ctx_mutex); sd_image_t* raw_results = nullptr; - if (!generate_video(runtime.sd_ctx, ¶ms, &raw_results, &num_results, &generated_audio)) { + if (!generate_video(runtime.sd_ctx, ¶ms, &raw_results, &num_results, &generated_audio, &output_fps)) { raw_results = nullptr; } results.adopt(raw_results, num_results); @@ -261,7 +261,7 @@ bool execute_vid_gen_job(ServerRuntime& runtime, std::vector video_bytes = create_video_from_sd_images_to_vector(job.vid_gen.output_format, results.data(), num_results, - job.vid_gen.gen_params.fps, + output_fps, job.vid_gen.output_compression, generated_audio); free_sd_audio(generated_audio); @@ -273,7 +273,6 @@ bool execute_vid_gen_job(ServerRuntime& runtime, output_media_b64 = base64_encode(video_bytes); output_media_mime_type = video_mime_type(job.vid_gen.output_format); output_frame_count = num_results; - output_fps = job.vid_gen.gen_params.fps; return true; } diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index cf69b48d..31a87e85 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -208,6 +208,7 @@ typedef struct { const char* embeddings_connectors_path; const char* vae_path; const char* audio_vae_path; + const char* audio_encoder_path; const char* taesd_path; const char* control_net_path; const char* ip_adapter_path; @@ -521,11 +522,13 @@ enum sd_cancel_mode_t { SD_API void sd_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode); SD_API void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params); +// If non-NULL, fps_out receives the effective encoding frame rate before preview callbacks. SD_API bool generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params, sd_image_t** frames_out, int* num_frames_out, - sd_audio_t** audio_out); + sd_audio_t** audio_out, + int* fps_out); typedef struct upscaler_ctx_t upscaler_ctx_t; diff --git a/src/conditioning/wan_audio.cpp b/src/conditioning/wan_audio.cpp new file mode 100644 index 00000000..d4d59ca1 --- /dev/null +++ b/src/conditioning/wan_audio.cpp @@ -0,0 +1,103 @@ +#include "wan_audio.h" + +#include +#include +#include + +namespace sd::wan_audio { + + static BucketPlan plan_buckets(int audio_frames, int batch_frames, int video_rate, int fps) { + BucketPlan plan; + plan.audio_frames = audio_frames; + plan.batch_frames = batch_frames; + plan.video_rate = video_rate; + plan.fps = fps; + const double scale = static_cast(video_rate) / fps; + // Keep a trailing chunk even when audio ends on a chunk boundary. + plan.num_chunks = static_cast(audio_frames / (batch_frames * scale)) + 1; + plan.bucket_frames = plan.num_chunks * batch_frames; + plan.padded_audio_frames = static_cast( + std::ceil(plan.bucket_frames / static_cast(fps) * video_rate)); + return plan; + } + + // Match NumPy's round-half-even sampling. + static int bucket_source_frame(int bucket_frame, int video_rate, int fps) { + return static_cast(std::nearbyint(static_cast(bucket_frame) * video_rate / fps)); + } + + static int interpolated_frame_count(int in_frames, int input_fps, int output_fps) { + return static_cast(in_frames / static_cast(input_fps) * output_fps); + } + + // Match PyTorch linear interpolation with align_corners=True. + static std::vector linear_interpolate_frames(const std::vector& in, + int num_layers, + int in_frames, + int dim, + int out_frames) { + std::vector out(static_cast(num_layers) * out_frames * dim, 0.0f); + if (in.empty() || in_frames <= 0 || out_frames <= 0 || num_layers <= 0 || dim <= 0) { + return out; + } + const double scale = out_frames > 1 ? static_cast(in_frames - 1) / (out_frames - 1) : 0.0; + for (int layer = 0; layer < num_layers; ++layer) { + for (int out_i = 0; out_i < out_frames; ++out_i) { + const double pos = out_i * scale; + const int src0 = static_cast(pos); + const int src1 = std::min(src0 + 1, in_frames - 1); + const float frac = static_cast(pos - src0); + const float* in_row = &in[(static_cast(layer) * in_frames + src0) * dim]; + const float* in_next = &in[(static_cast(layer) * in_frames + src1) * dim]; + float* out_row = &out[(static_cast(layer) * out_frames + out_i) * dim]; + for (int d = 0; d < dim; ++d) { + out_row[d] = in_row[d] * (1.0f - frac) + in_next[d] * frac; + } + } + } + return out; + } + + std::vector build_audio_buckets(const float* stacked_states, + int num_layers, + int in_frames, + int dim, + int batch_frames, + BucketPlan* plan_out, + int input_fps, + int video_rate, + int fps) { + if (stacked_states == nullptr || num_layers <= 0 || in_frames <= 0 || dim <= 0 || batch_frames <= 0) { + return {}; + } + const int audio_frames = interpolated_frame_count(in_frames, input_fps, video_rate); + if (audio_frames <= 0) { + return {}; + } + const std::vector interpolated = + linear_interpolate_frames(std::vector(stacked_states, + stacked_states + static_cast(num_layers) * in_frames * dim), + num_layers, + in_frames, + dim, + audio_frames); + const BucketPlan plan = plan_buckets(audio_frames, batch_frames, video_rate, fps); + if (plan_out != nullptr) { + *plan_out = plan; + } + std::vector buckets(static_cast(plan.bucket_frames) * num_layers * dim, 0.0f); + for (int frame = 0; frame < plan.bucket_frames; ++frame) { + const int src = bucket_source_frame(frame, video_rate, fps); + if (src >= plan.audio_frames) { + continue; + } + for (int layer = 0; layer < num_layers; ++layer) { + std::copy_n(interpolated.data() + (static_cast(layer) * audio_frames + src) * dim, + static_cast(dim), + buckets.data() + (static_cast(frame) * num_layers + layer) * dim); + } + } + return buckets; + } + +} // namespace sd::wan_audio diff --git a/src/conditioning/wan_audio.h b/src/conditioning/wan_audio.h new file mode 100644 index 00000000..1b245ed5 --- /dev/null +++ b/src/conditioning/wan_audio.h @@ -0,0 +1,32 @@ +#ifndef __SD_CONDITIONING_WAN_AUDIO_H__ +#define __SD_CONDITIONING_WAN_AUDIO_H__ + +#include + +namespace sd::wan_audio { + + struct BucketPlan { + int audio_frames; // frames at video_rate + int batch_frames; // latent_t * 4 + int video_rate; + int fps; // bucket frame rate + int num_chunks; // includes trailing padding + int bucket_frames; + int padded_audio_frames; + }; + + // [layers, frames, dim] at input_fps -> [bucket_frames, layers, dim] at fps. + // Pads past the audio end; returns an empty vector on invalid input. + std::vector build_audio_buckets(const float* stacked_states, + int num_layers, + int in_frames, + int dim, + int batch_frames, + BucketPlan* plan_out = nullptr, + int input_fps = 50, + int video_rate = 30, + int fps = 16); + +} // namespace sd::wan_audio + +#endif // __SD_CONDITIONING_WAN_AUDIO_H__ diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 7cf58539..79e3ea9b 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -325,6 +325,76 @@ ggml_tensor* ggml_ext_pad(ggml_context* ctx, return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y); } +static ggml_tensor* conv_1d(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, int s0, int p0, int d0, bool force_prec_f32) { + ggml_tensor* result; + if (force_prec_f32) { + ggml_tensor* patches = ggml_im2col(ctx, w, x, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F32); + result = ggml_mul_mat(ctx, + ggml_reshape_2d(ctx, patches, patches->ne[0], patches->ne[2] * patches->ne[1]), + ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1], w->ne[2])); + result = ggml_reshape_3d(ctx, result, patches->ne[1], w->ne[2], patches->ne[2]); + } else { + result = ggml_conv_1d(ctx, w, x, s0, p0, d0); + } + if (x->ne[2] > 1) { + // mul_mat packs positions and batches before output channels: [OL, N, OC]. + result = ggml_reshape_3d(ctx, result, result->ne[0], x->ne[2], w->ne[2]); + result = ggml_cont(ctx, ggml_permute(ctx, result, 0, 2, 1, 3)); + } + return result; +} + +ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + int s0, + int p0, + int d0, + int64_t groups, + bool force_prec_f32) { + GGML_ASSERT(s0 > 0 && p0 >= 0 && d0 > 0 && groups > 0); + GGML_ASSERT(x->type == GGML_TYPE_F32 && x->ne[3] == 1 && w->ne[3] == 1); + GGML_ASSERT(x->ne[1] % groups == 0 && w->ne[2] % groups == 0); + GGML_ASSERT(w->ne[1] == x->ne[1] / groups); + GGML_ASSERT(b == nullptr || (b->type == GGML_TYPE_F32 && ggml_is_vector(b) && b->ne[0] == w->ne[2])); + + // im2col requires contiguous time rows; group views must retain the real channel and batch strides. + if (!ggml_is_contiguous(x)) { + x = ggml_cont(ctx, x); + } + if (force_prec_f32 && w->type != GGML_TYPE_F32) { + w = ggml_cast(ctx, w, GGML_TYPE_F32); + } + if (!ggml_is_contiguous(w)) { + w = ggml_cont(ctx, w); + } + + ggml_tensor* result = nullptr; + if (groups == 1) { + result = conv_1d(ctx, x, w, s0, p0, d0, force_prec_f32); + } else { + const int64_t ic_g = x->ne[1] / groups; + const int64_t oc_g = w->ne[2] / groups; + std::vector outputs; + outputs.reserve(groups); + for (int64_t group = 0; group < groups; ++group) { + ggml_tensor* x_i = ggml_view_3d(ctx, x, x->ne[0], ic_g, x->ne[2], x->nb[1], x->nb[2], group * ic_g * x->nb[1]); + ggml_tensor* w_i = ggml_view_3d(ctx, w, w->ne[0], ic_g, oc_g, w->nb[1], w->nb[2], group * oc_g * w->nb[2]); + outputs.push_back(conv_1d(ctx, x_i, w_i, s0, p0, d0, force_prec_f32)); + } + result = ggml_ext_vec_concat(ctx, outputs, 1); + } + if (b != nullptr) { + if (!ggml_is_contiguous(b)) { + b = ggml_cont(ctx, b); + } + b = ggml_reshape_3d(ctx, b, 1, w->ne[2], 1); + result = ggml_add_inplace(ctx, result, b); + } + return result; +} + ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, @@ -683,17 +753,16 @@ ggml_tensor* ggml_ext_group_norm(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, - int num_groups) { + int num_groups, + float eps) { if (ggml_n_dims(x) >= 3 && w != nullptr && b != nullptr) { w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1); b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1); } - const float eps = 1e-6f; // default eps parameter - x = ggml_group_norm(ctx, x, num_groups, eps); + x = ggml_group_norm(ctx, x, num_groups, eps); if (w != nullptr && b != nullptr) { x = ggml_mul_inplace(ctx, x, w); - // b = ggml_repeat(ctx, b, x); x = ggml_add_inplace(ctx, x, b); } return x; diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index 585d4582..6f3fe89e 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -103,6 +103,18 @@ ggml_tensor* ggml_ext_pad(ggml_context* ctx, bool circular_x = false, bool circular_y = false); +// ggml layout: x [L, IC, N], w [K, IC/groups, OC], b [OC], result [OL, OC, N]. +// force_prec_f32 keeps both input patches and weights in F32. +ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + int s0 = 1, + int p0 = 0, + int d0 = 1, + int64_t groups = 1, + bool force_prec_f32 = false); + // w: [OC，IC, KH, KW] // x: [N, IC, IH, IW] // b: [OC,] @@ -219,7 +231,8 @@ ggml_tensor* ggml_ext_group_norm(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, - int num_groups = 32); + int num_groups = 32, + float eps = 1e-6f); ggml_tensor* ggml_ext_timestep_embedding( ggml_context* ctx, diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index 10179a8e..b85fc5cb 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -87,6 +87,10 @@ static bool parse_backend_module(const std::string& raw_name, SDBackendModule* m *module = SDBackendModule::DETECTOR; return true; } + if (name == "audioencoder" || name == "audio") { + *module = SDBackendModule::AUDIO_ENCODER; + return true; + } return false; } @@ -968,6 +972,8 @@ const char* sd_backend_module_name(SDBackendModule module) { return "upscaler"; case SDBackendModule::DETECTOR: return "detector"; + case SDBackendModule::AUDIO_ENCODER: + return "audio_encoder"; } return "unknown"; } diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index 3430ccc5..a66bf51c 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -21,6 +21,7 @@ enum class SDBackendModule { PHOTOMAKER, UPSCALER, DETECTOR, + AUDIO_ENCODER, }; struct SDBackendAssignment { diff --git a/src/model.h b/src/model.h index 4b3ed316..7a8bc757 100644 --- a/src/model.h +++ b/src/model.h @@ -35,6 +35,7 @@ enum SDVersion { VERSION_WAN2, VERSION_WAN2_2_I2V, VERSION_WAN2_2_TI2V, + VERSION_WAN2_2_S2V, VERSION_LINGBOT_VIDEO, VERSION_QWEN_IMAGE, VERSION_QWEN_IMAGE_LAYERED, @@ -130,7 +131,7 @@ static inline bool sd_version_is_minimax_h3(SDVersion version) { } static inline bool sd_version_is_wan(SDVersion version) { - if (version == VERSION_WAN2 || version == VERSION_WAN2_2_I2V || version == VERSION_WAN2_2_TI2V) { + if (version == VERSION_WAN2 || version == VERSION_WAN2_2_I2V || version == VERSION_WAN2_2_TI2V || version == VERSION_WAN2_2_S2V) { return true; } return false; diff --git a/src/model/audio/wav2vec2.hpp b/src/model/audio/wav2vec2.hpp new file mode 100644 index 00000000..3f3c969e --- /dev/null +++ b/src/model/audio/wav2vec2.hpp @@ -0,0 +1,413 @@ +#ifndef __SD_MODEL_AUDIO_WAV2VEC2_HPP__ +#define __SD_MODEL_AUDIO_WAV2VEC2_HPP__ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "model.h" +#include "model/common/ggml_block.hpp" + +namespace Wav2Vec2 { + + struct Wav2Vec2Config { + int64_t embed_dim = 1024; + int64_t conv_dim = 512; + int num_heads = 16; + int num_layers = 24; + std::string feat_extract_norm = "layer"; + bool conv_bias = true; + bool do_normalize = true; + bool do_stable_layer_norm = true; + + static Wav2Vec2Config detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + Wav2Vec2Config config; + auto it = tensor_storage_map.find(prefix + "encoder.layer_norm.bias"); + if (it == tensor_storage_map.end()) { + LOG_WARN("wav2vec2: %sencoder.layer_norm.bias not found, using large defaults", prefix.c_str()); + return config; + } + config.embed_dim = it->second.ne[0]; + if (config.embed_dim == 1024) { + config.embed_dim = 1024; + config.num_heads = 16; + config.num_layers = 24; + config.feat_extract_norm = "layer"; + config.conv_bias = true; + config.do_normalize = true; + config.do_stable_layer_norm = true; + } else if (config.embed_dim == 768) { + config.embed_dim = 768; + config.num_heads = 12; + config.num_layers = 12; + config.feat_extract_norm = "group"; + config.conv_bias = false; + config.do_normalize = false; + config.do_stable_layer_norm = false; + } else { + LOG_WARN("wav2vec2: unsupported embed_dim %" PRId64 ", using large defaults", config.embed_dim); + config.embed_dim = 1024; + } + return config; + } + }; + + struct Wav2Vec2NoLayerNormConvLayer : public UnaryBlock { + Wav2Vec2NoLayerNormConvLayer(int64_t in_channels, int64_t out_channels, int kernel_size, int stride, bool bias) { + blocks["conv"] = std::make_shared(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto conv = std::dynamic_pointer_cast(blocks["conv"]); + x = conv->forward(ctx, x); + return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x)); + } + }; + + struct Wav2Vec2LayerNormConvLayer : public UnaryBlock { + Wav2Vec2LayerNormConvLayer(int64_t in_channels, int64_t out_channels, int kernel_size, int stride, bool bias) { + blocks["conv"] = std::make_shared(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias, true); + blocks["layer_norm"] = std::make_shared(out_channels); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto conv = std::dynamic_pointer_cast(blocks["conv"]); + auto layer_norm = std::dynamic_pointer_cast(blocks["layer_norm"]); + x = conv->forward(ctx, x); + // LayerNorm normalizes channels: [N, C, L] -> [N, L, C]. + x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); + x = layer_norm->forward(ctx, x); + x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); + return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x)); + } + }; + + struct Wav2Vec2GroupNormConvLayer : public UnaryBlock { + Wav2Vec2GroupNormConvLayer(int64_t in_channels, int64_t out_channels, int kernel_size, int stride, bool bias) { + blocks["conv"] = std::make_shared(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias, true); + blocks["layer_norm"] = std::make_shared((int)out_channels, out_channels, 1e-05f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto conv = std::dynamic_pointer_cast(blocks["conv"]); + auto layer_norm = std::dynamic_pointer_cast(blocks["layer_norm"]); + x = conv->forward(ctx, x); + // ggml GroupNorm needs [N, C, H, W], with H=1 for audio. + x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]); + x = layer_norm->forward(ctx, x); + x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[2], x->ne[3]); + return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x)); + } + }; + + struct Wav2Vec2FeatureEncoder : public UnaryBlock { + Wav2Vec2FeatureEncoder(const Wav2Vec2Config& config) { + GGML_ASSERT(config.feat_extract_norm == "layer" || config.feat_extract_norm == "group"); + const int kernels[7] = {10, 3, 3, 3, 3, 2, 2}; + const int strides[7] = {5, 2, 2, 2, 2, 2, 2}; + int64_t in_channels = 1; + for (int i = 0; i < 7; ++i) { + const std::string name = "conv_layers." + std::to_string(i); + if (config.feat_extract_norm == "layer") { + blocks[name] = std::make_shared(in_channels, config.conv_dim, kernels[i], strides[i], config.conv_bias); + } else if (i == 0) { + blocks[name] = std::make_shared(in_channels, config.conv_dim, kernels[i], strides[i], config.conv_bias); + } else { + blocks[name] = std::make_shared(in_channels, config.conv_dim, kernels[i], strides[i], config.conv_bias); + } + in_channels = config.conv_dim; + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + for (int i = 0; i < 7; ++i) { + auto conv = std::dynamic_pointer_cast(blocks["conv_layers." + std::to_string(i)]); + x = conv->forward(ctx, x); + } + return ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); + } + }; + + struct Wav2Vec2FeatureProjection : public UnaryBlock { + Wav2Vec2FeatureProjection(const Wav2Vec2Config& config) { + blocks["layer_norm"] = std::make_shared(config.conv_dim); + blocks["projection"] = std::make_shared(config.conv_dim, config.embed_dim); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto ln = std::dynamic_pointer_cast(blocks["layer_norm"]); + auto projection = std::dynamic_pointer_cast(blocks["projection"]); + x = ln->forward(ctx, x); + x = projection->forward(ctx, x); + return x; + } + }; + + class Wav2Vec2PositionalConvEmbedding : public UnaryBlock { + private: + int64_t embed_dim_; + static constexpr int groups_ = 16; + static constexpr int kernel_size_ = 128; + std::string weight_g_name_; + std::string weight_v_name_; + + ggml_tensor* weight(GGMLRunnerContext* ctx) { + auto g = params[weight_g_name_]; + auto v = ggml_cast(ctx->ggml_ctx, params[weight_v_name_], GGML_TYPE_F32); + auto squared = ggml_mul(ctx->ggml_ctx, v, v); + // PyTorch weight_norm(dim=2) reduces both channel axes, retaining each kernel tap. + squared = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, squared, 2, 0, 1, 3)); + squared = ggml_reshape_2d(ctx->ggml_ctx, squared, embed_dim_ / groups_ * embed_dim_, kernel_size_); + auto norm = ggml_sqrt(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, squared)); + norm = ggml_reshape_3d(ctx->ggml_ctx, norm, kernel_size_, 1, 1); + return ggml_mul(ctx->ggml_ctx, v, ggml_div(ctx->ggml_ctx, g, norm)); + } + + public: + Wav2Vec2PositionalConvEmbedding(const Wav2Vec2Config& config) + : embed_dim_(config.embed_dim) { + GGML_ASSERT(embed_dim_ > 0 && embed_dim_ % groups_ == 0); + } + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + bool legacy = tensor_storage_map.count(prefix + "conv.weight_g") > 0; + weight_g_name_ = legacy ? "conv.weight_g" : "conv.parametrizations.weight.original0"; + weight_v_name_ = legacy ? "conv.weight_v" : "conv.parametrizations.weight.original1"; + auto g = tensor_storage_map.find(prefix + weight_g_name_); + auto v = tensor_storage_map.find(prefix + weight_v_name_); + GGML_ASSERT(g != tensor_storage_map.end() && v != tensor_storage_map.end()); + GGML_ASSERT(g->second.ne[0] == kernel_size_ && g->second.ne[1] == 1 && g->second.ne[2] == 1 && g->second.ne[3] == 1); + GGML_ASSERT(v->second.ne[0] == kernel_size_ && v->second.ne[1] == embed_dim_ / groups_ && v->second.ne[2] == embed_dim_ && v->second.ne[3] == 1); + + params[weight_g_name_] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kernel_size_, 1, 1); + params[weight_v_name_] = ggml_new_tensor_3d(ctx, get_type(prefix + weight_v_name_, tensor_storage_map, GGML_TYPE_F16), + kernel_size_, embed_dim_ / groups_, embed_dim_); + if (tensor_storage_map.count(prefix + "conv.bias") > 0) { + params["conv.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, embed_dim_); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto w = weight(ctx); + auto b = params.count("conv.bias") > 0 ? params["conv.bias"] : nullptr; + x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = ggml_ext_conv_1d(ctx->ggml_ctx, x, w, b, 1, kernel_size_ / 2, 1, groups_, true); + // Apply GELU out of place before cropping to keep graph buffer reuse safe. + x = ggml_gelu_erf(ctx->ggml_ctx, x); + x = ggml_view_3d(ctx->ggml_ctx, x, x->ne[0] - 1, x->ne[1], x->ne[2], x->nb[1], x->nb[2], 0); + return ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); + } + }; + + struct Wav2Vec2FeedForward : public UnaryBlock { + Wav2Vec2FeedForward(const Wav2Vec2Config& config) { + blocks["intermediate_dense"] = std::make_shared(config.embed_dim, config.embed_dim * 4); + blocks["output_dense"] = std::make_shared(config.embed_dim * 4, config.embed_dim); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto intermediate_dense = std::dynamic_pointer_cast(blocks["intermediate_dense"]); + auto output_dense = std::dynamic_pointer_cast(blocks["output_dense"]); + x = intermediate_dense->forward(ctx, x); + x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = output_dense->forward(ctx, x); + return x; + } + }; + + struct Wav2Vec2EncoderLayer : public UnaryBlock { + bool do_stable_layer_norm; + + Wav2Vec2EncoderLayer(const Wav2Vec2Config& config) + : do_stable_layer_norm(config.do_stable_layer_norm) { + blocks["attention"] = std::make_shared(config.embed_dim, config.num_heads, true, true); + blocks["layer_norm"] = std::make_shared(config.embed_dim); + blocks["feed_forward"] = std::make_shared(config); + blocks["final_layer_norm"] = std::make_shared(config.embed_dim); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto attention = std::dynamic_pointer_cast(blocks["attention"]); + auto layer_norm = std::dynamic_pointer_cast(blocks["layer_norm"]); + auto feed_forward = std::dynamic_pointer_cast(blocks["feed_forward"]); + auto final_layer_norm = std::dynamic_pointer_cast(blocks["final_layer_norm"]); + + ggml_tensor* residual = x; + if (do_stable_layer_norm) { + x = layer_norm->forward(ctx, x); + x = attention->forward(ctx, x); + x = ggml_add(ctx->ggml_ctx, residual, x); + x = ggml_add(ctx->ggml_ctx, x, feed_forward->forward(ctx, final_layer_norm->forward(ctx, x))); + } else { + x = attention->forward(ctx, x); + x = ggml_add(ctx->ggml_ctx, residual, x); + x = layer_norm->forward(ctx, x); + x = final_layer_norm->forward(ctx, ggml_add(ctx->ggml_ctx, x, feed_forward->forward(ctx, x))); + } + return x; + } + }; + + struct Wav2Vec2Encoder : public GGMLBlock { + int num_layers; + bool do_stable_layer_norm; + + Wav2Vec2Encoder(const Wav2Vec2Config& config) + : num_layers(config.num_layers), do_stable_layer_norm(config.do_stable_layer_norm) { + blocks["pos_conv_embed"] = std::make_shared(config); + for (int i = 0; i < config.num_layers; ++i) { + blocks["layers." + std::to_string(i)] = std::make_shared(config); + } + blocks["layer_norm"] = std::make_shared(config.embed_dim); + } + + // For N == 1, all_layers stacks pre-layer states and the final state as [embed_dim, L, num_layers + 1]. + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor** all_layers = nullptr) { + auto pos_conv_embed = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); + auto layer_norm = std::dynamic_pointer_cast(blocks["layer_norm"]); + + std::vector collected; + if (all_layers != nullptr) { + collected.reserve(num_layers + 1); + } + + x = ggml_add(ctx->ggml_ctx, x, pos_conv_embed->forward(ctx, x)); + if (!do_stable_layer_norm) { + x = layer_norm->forward(ctx, x); + } + for (int i = 0; i < num_layers; ++i) { + if (all_layers != nullptr) { + collected.push_back(x); + } + auto layer = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); + x = layer->forward(ctx, x); + } + if (do_stable_layer_norm) { + x = layer_norm->forward(ctx, x); + } + if (all_layers != nullptr) { + collected.push_back(x); + ggml_tensor* stack = collected[0]; + for (size_t i = 1; i < collected.size(); ++i) { + stack = ggml_concat(ctx->ggml_ctx, stack, collected[i], 2); + } + *all_layers = stack; + } + return x; + } + }; + + struct Wav2Vec2Model : public GGMLBlock { + Wav2Vec2Config config; + + Wav2Vec2Model() = default; + Wav2Vec2Model(const Wav2Vec2Config& config_) + : config(config_) { + blocks["feature_extractor"] = std::make_shared(config); + blocks["feature_projection"] = std::make_shared(config); + blocks["encoder"] = std::make_shared(config); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor** all_layers = nullptr) { + auto feature_extractor = std::dynamic_pointer_cast(blocks["feature_extractor"]); + auto feature_projection = std::dynamic_pointer_cast(blocks["feature_projection"]); + auto encoder = std::dynamic_pointer_cast(blocks["encoder"]); + + x = feature_extractor->forward(ctx, x); + x = feature_projection->forward(ctx, x); + x = encoder->forward(ctx, x, all_layers); + return x; + } + }; + + class Wav2Vec2ModelRunner : public GGMLRunner { + private: + Wav2Vec2Config config; + + public: + Wav2Vec2Model model; + std::string weight_prefix; + + Wav2Vec2ModelRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "wav2vec2.", + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), + config(Wav2Vec2Config::detect_from_weights(tensor_storage_map, prefix)), + model(config), + weight_prefix(prefix) { + // GGMLBlock appends its own separator; loader prefixes already include one. + std::string block_prefix = weight_prefix; + if (!block_prefix.empty() && block_prefix.back() == '.') { + block_prefix.pop_back(); + } + model.init(params_ctx, tensor_storage_map, block_prefix); + LOG_INFO("%s", get_desc().c_str()); + } + + std::string get_desc() override { + return "wav2vec2"; + } + + void get_param_tensors(std::map& tensors) { + std::string block_prefix = weight_prefix; + if (!block_prefix.empty() && block_prefix.back() == '.') { + block_prefix.pop_back(); + } + model.get_param_tensors(tensors, block_prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& waveform_tensor) { + ggml_cgraph* gf = ggml_new_graph(compute_ctx); + ggml_tensor* waveform = make_input(waveform_tensor); + auto runner_ctx = get_context(); + ggml_tensor* all_layers = nullptr; + model.forward(&runner_ctx, waveform, &all_layers); + GGML_ASSERT(all_layers != nullptr); + ggml_build_forward_expand(gf, all_layers); + return gf; + } + + sd::Tensor compute(const int n_threads, const std::vector& mono_waveform) { + GGML_ASSERT(!mono_waveform.empty()); + const int64_t num_samples = (int64_t)mono_waveform.size(); + sd::Tensor waveform({num_samples, 1, 1}); + std::copy(mono_waveform.begin(), mono_waveform.end(), waveform.data()); + normalize(waveform.data(), num_samples); + + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(waveform); + }; + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); + } + + private: + static void normalize(float* x, int64_t n) { + double mean = 0.0; + for (int64_t i = 0; i < n; ++i) { + mean += x[i]; + } + mean /= n; + double var = 0.0; + for (int64_t i = 0; i < n; ++i) { + const double d = x[i] - mean; + var += d * d; + } + var /= n; + const float scale = (float)(1.0 / std::sqrt(var + 1e-7)); + for (int64_t i = 0; i < n; ++i) { + x[i] = (float)((x[i] - mean) * scale); + } + } + }; + +} // namespace Wav2Vec2 + +#endif // __SD_MODEL_AUDIO_WAV2VEC2_HPP__ diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index cc5470c3..2aca72f4 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -368,6 +368,61 @@ public: } }; +class Conv1d : public UnaryBlock { +protected: + int64_t in_channels; + int64_t out_channels; + int64_t groups; + int kernel_size; + int stride; + int padding; + int dilation; + bool bias; + bool force_prec_f32; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F16); + params["weight"] = ggml_new_tensor_3d(ctx, wtype, kernel_size, in_channels / groups, out_channels); + if (bias) { + params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); + } + } + +public: + Conv1d(int64_t in_channels, + int64_t out_channels, + int kernel_size, + int stride = 1, + int padding = 0, + int dilation = 1, + int64_t groups = 1, + bool bias = true, + bool force_prec_f32 = false) + : in_channels(in_channels), + out_channels(out_channels), + groups(groups), + kernel_size(kernel_size), + stride(stride), + padding(padding), + dilation(dilation), + bias(bias), + force_prec_f32(force_prec_f32) { + GGML_ASSERT(in_channels > 0 && out_channels > 0 && groups > 0); + GGML_ASSERT(in_channels % groups == 0 && out_channels % groups == 0); + GGML_ASSERT(kernel_size > 0 && stride > 0 && padding >= 0 && dilation > 0); + } + + std::string get_desc() override { + return "Conv1d"; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + GGML_ASSERT(x->ne[1] == in_channels); + return ggml_ext_conv_1d(ctx->ggml_ctx, x, params["weight"], bias ? params["bias"] : nullptr, + stride, padding, dilation, groups, force_prec_f32); + } +}; + class Conv2d : public UnaryBlock { protected: int64_t in_channels; @@ -766,7 +821,7 @@ public: b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); } } - return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups); + return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups, eps); } }; diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index 447c3a23..778bbf42 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -818,8 +818,9 @@ namespace Rope { int pw, int bs, int theta, - const std::vector& axes_dim) { - std::vector> ids = gen_vid_ids(t, h, w, pt, ph, pw, bs); + const std::vector& axes_dim, + int t_offset = 0) { + std::vector> ids = gen_vid_ids(t, h, w, pt, ph, pw, bs, t_offset); return embed_nd(ids, bs, static_cast(theta), axes_dim); } diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 9d611c4e..a4b3c38f 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -69,6 +69,8 @@ struct AnimaDiffusionExtra { struct WanDiffusionExtra { const sd::Tensor* vace_context = nullptr; float vace_strength = 1.f; + // S2V audio, sd::Tensor layout: [dim, T_latent*4, layers]. + const sd::Tensor* audio_embed = nullptr; }; struct HiDreamO1DiffusionExtra { diff --git a/src/model/diffusion/wan.hpp b/src/model/diffusion/wan.hpp index 363beeeb..90fd1285 100644 --- a/src/model/diffusion/wan.hpp +++ b/src/model/diffusion/wan.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_DIFFUSION_WAN_HPP__ #define __SD_MODEL_DIFFUSION_WAN_HPP__ +#include #include #include #include @@ -33,11 +34,16 @@ namespace WAN { int vace_layers = 0; int64_t vace_in_dim = 96; std::map vace_layers_mapping = {}; - bool qk_norm = true; - bool cross_attn_norm = true; - float eps = 1e-6f; - int64_t flf_pos_embed_token_number = 0; - int theta = 10000; + int64_t audio_dim = 1024; + int num_audio_token = 4; // excludes the learned padding token + std::vector audio_inject_layers = {}; + std::map audio_inject_mapping = {}; // block index -> injector index + std::string adain_mode = "attn_norm"; + bool qk_norm = true; + bool cross_attn_norm = true; + float eps = 1e-6f; + int64_t flf_pos_embed_token_number = 0; + int theta = 10000; // wan2.1 1.3B: 1536/12, wan2.1/2.2 14B: 5120/40, wan2.2 5B: 3074/24 std::vector axes_dim = {44, 42, 42}; int64_t axes_dim_sum = 128; @@ -74,6 +80,10 @@ namespace WAN { if (name.find("img_emb") != std::string::npos) { config.model_type = "i2v"; } + if (name.find("audio_injector") != std::string::npos || name.find("casual_audio_encoder") != std::string::npos) { + config.model_type = "s2v"; + config.audio_inject_layers = {0, 4, 8, 12, 16, 20, 24, 27, 30, 33, 36, 39}; + } if (name.find("img_emb.emb_pos") != std::string::npos) { config.flf_pos_embed_token_number = 514; } @@ -265,6 +275,13 @@ namespace WAN { } }; +} // namespace WAN + +// Audio injection reuses WanT2VCrossAttention defined above. +#include "model/diffusion/wan_audio.hpp" + +namespace WAN { + static ggml_tensor* modulate_add(ggml_context* ctx, ggml_tensor* x, ggml_tensor* e) { // x: [N, n_token, dim] // e: [N, 1, dim] or [N, T, 1, dim] @@ -532,6 +549,13 @@ namespace WAN { protected: WanConfig config; + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + if (config.model_type == "s2v") { + enum ggml_type wtype = GGML_TYPE_F32; // elementwise add vs F32 activations + params["trainable_cond_mask.weight"] = ggml_new_tensor_2d(ctx, wtype, config.dim, 3); + } + } + public: Wan() {} Wan(WanConfig config) @@ -554,7 +578,7 @@ namespace WAN { // blocks for (int i = 0; i < config.num_layers; i++) { - auto block = std::shared_ptr(new WanAttentionBlock(config.model_type == "t2v", + auto block = std::shared_ptr(new WanAttentionBlock(config.model_type != "i2v", config.dim, config.ffn_dim, config.num_heads, @@ -595,6 +619,14 @@ namespace WAN { blocks["vace_patch_embedding"] = std::shared_ptr(new Conv3d(config.vace_in_dim, config.dim, config.patch_size, config.patch_size)); } + + if (config.model_type == "s2v") { + blocks["casual_audio_encoder"] = std::make_shared(config.audio_dim, config.dim, config.num_audio_token); + blocks["audio_injector"] = std::make_shared(config.dim, config.num_heads, (int)config.audio_inject_layers.size(), config.qk_norm, config.eps); + for (size_t i = 0; i < config.audio_inject_layers.size(); i++) { + config.audio_inject_mapping[config.audio_inject_layers[i]] = (int)i; + } + } } ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx, @@ -642,18 +674,24 @@ namespace WAN { ggml_tensor* timestep, ggml_tensor* context, ggml_tensor* pe, - ggml_tensor* clip_fea = nullptr, - ggml_tensor* vace_context = nullptr, - float vace_strength = 1.f, - int64_t N = 1) { + ggml_tensor* clip_fea = nullptr, + ggml_tensor* vace_context = nullptr, + float vace_strength = 1.f, + int64_t N = 1, + ggml_tensor* audio_embed = nullptr, + ggml_tensor* reference_latent = nullptr) { // 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] + // audio_embed: [layers, T*4, audio_dim] + // reference_latent: [N*C, T_ref, H, W] + // return: [N, (t_len [+ t_ref_len]) * h_len*w_len, out_dim*pt*ph*pw] GGML_ASSERT(N == 1); + int64_t T = x->ne[2]; + auto patch_embedding = std::dynamic_pointer_cast(blocks["patch_embedding"]); auto text_embedding_0 = std::dynamic_pointer_cast(blocks["text_embedding.0"]); @@ -670,6 +708,40 @@ namespace WAN { x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3] / N, N); // [N, dim, t_len*h_len*w_len] x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim] + ggml_tensor* audio_local = nullptr; + ggml_tensor* audio_global = nullptr; + int64_t seq_len = x->ne[1]; + int64_t t_ref_len = 0; + if (config.model_type == "s2v") { + if (audio_embed != nullptr) { + GGML_ASSERT(audio_embed->ne[1] == T * 4); + auto audio_encoder = std::dynamic_pointer_cast(blocks["casual_audio_encoder"]); + auto audio_emb = audio_encoder->forward(ctx, audio_embed); + audio_local = audio_emb.first; + audio_global = audio_emb.second; + GGML_ASSERT(audio_local->ne[2] == T); + } + + // video tokens get cond_mask[0], reference tokens cond_mask[1] + auto cond_mask = params["trainable_cond_mask.weight"]; + auto cm0 = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 0, 1), config.dim, 1, 1); + x = ggml_add(ctx->ggml_ctx, x, cm0); + + if (reference_latent != nullptr) { + t_ref_len = reference_latent->ne[2]; + auto ref = patch_embedding->forward(ctx, reference_latent); + ref = ggml_reshape_3d(ctx->ggml_ctx, ref, ref->ne[0] * ref->ne[1] * ref->ne[2], ref->ne[3] / N, N); + ref = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ref, 1, 0, 2, 3)); // [N, t_ref*h_len*w_len, dim] + auto cm1 = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 1, 2), config.dim, 1, 1); + ref = ggml_add(ctx->ggml_ctx, ref, cm1); + x = ggml_concat(ctx->ggml_ctx, x, ref, 1); + + // Reference tokens use timestep 0. + GGML_ASSERT(timestep->ne[0] == T); + timestep = ggml_ext_pad(ctx->ggml_ctx, timestep, (int)t_ref_len, 0, 0, 0); + } + } + // time_embedding auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, config.freq_dim); e = time_embedding_0->forward(ctx, e); @@ -714,6 +786,11 @@ namespace WAN { auto x_orig = x; + std::shared_ptr audio_injector; + if (audio_local != nullptr) { + audio_injector = std::dynamic_pointer_cast(blocks["audio_injector"]); + } + for (int i = 0; i < config.num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); @@ -731,6 +808,13 @@ namespace WAN { c_skip = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength); x = ggml_add(ctx->ggml_ctx, x, c_skip); } + + if (audio_injector != nullptr) { + auto inject_iter = config.audio_inject_mapping.find(i); + if (inject_iter != config.audio_inject_mapping.end()) { + x = audio_injector->forward(ctx, x, seq_len, T, inject_iter->second, audio_local, audio_global); + } + } sd::ggml_graph_cut::mark_graph_cut(x, "wan.blocks." + std::to_string(i), "x"); if (c != nullptr) { sd::ggml_graph_cut::mark_graph_cut(c, "wan.blocks." + std::to_string(i), "c"); @@ -747,11 +831,13 @@ namespace WAN { ggml_tensor* timestep, ggml_tensor* context, ggml_tensor* pe, - ggml_tensor* clip_fea = nullptr, - ggml_tensor* time_dim_concat = nullptr, - ggml_tensor* vace_context = nullptr, - float vace_strength = 1.f, - int64_t N = 1) { + ggml_tensor* clip_fea = nullptr, + ggml_tensor* time_dim_concat = nullptr, + ggml_tensor* vace_context = nullptr, + float vace_strength = 1.f, + int64_t N = 1, + ggml_tensor* audio_embed = nullptr, + ggml_tensor* reference_latent = nullptr) { // Forward pass of DiT. // x: [N*C, T, H, W] // timestep: [N,] @@ -779,7 +865,12 @@ namespace WAN { t_len = ((x->ne[2] + (std::get<0>(config.patch_size) / 2)) / std::get<0>(config.patch_size)); } - 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] + auto out = forward_orig(ctx, x, timestep, context, pe, clip_fea, vace_context, vace_strength, N, audio_embed, reference_latent); // [N, (t_len [+t_ref]) *h_len*w_len, pt*ph*pw*C] + + if (reference_latent != nullptr) { + // Exclude reference tokens from the generated video. + out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, t_len * h_len * w_len); + } out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w] @@ -839,7 +930,10 @@ namespace WAN { config.text_len = 512; } } else if (config.num_layers == 40) { - if (config.model_type == "t2v") { + if (version == VERSION_WAN2_2_S2V) { + desc = "Wan2.2-S2V-14B"; + config.in_dim = 16; + } else if (config.model_type == "t2v") { if (version == VERSION_WAN2_2_I2V) { desc = "Wan2.2-I2V-14B"; config.in_dim = 36; @@ -891,7 +985,9 @@ namespace WAN { const sd::Tensor& c_concat_tensor = {}, const sd::Tensor& time_dim_concat_tensor = {}, const sd::Tensor& vace_context_tensor = {}, - float vace_strength = 1.f) { + float vace_strength = 1.f, + const sd::Tensor& audio_embed_tensor = {}, + const sd::Tensor& ref_latent_tensor = {}) { ggml_cgraph* gf = new_graph_custom(WAN_GRAPH_SIZE); ggml_tensor* x = make_input(x_tensor); @@ -901,16 +997,33 @@ namespace WAN { ggml_tensor* c_concat = make_optional_input(c_concat_tensor); ggml_tensor* time_dim_concat = make_optional_input(time_dim_concat_tensor); ggml_tensor* vace_context = make_optional_input(vace_context_tensor); + ggml_tensor* audio_embed = make_optional_input(audio_embed_tensor); + ggml_tensor* ref_latent = make_optional_input(ref_latent_tensor); - pe_vec = Rope::gen_wan_pe(static_cast(x->ne[2]), - static_cast(x->ne[1]), - static_cast(x->ne[0]), - std::get<0>(config.patch_size), - std::get<1>(config.patch_size), - std::get<2>(config.patch_size), - 1, - config.theta, - config.axes_dim); + pe_vec = Rope::gen_wan_pe(static_cast(x->ne[2]), + static_cast(x->ne[1]), + static_cast(x->ne[0]), + std::get<0>(config.patch_size), + std::get<1>(config.patch_size), + std::get<2>(config.patch_size), + 1, + config.theta, + config.axes_dim); + if (ref_latent != nullptr) { + // Match S2V's reference-frame temporal offset. + int t_start = std::max(30, static_cast(x->ne[2]) + 9); + auto ref_pe = Rope::gen_wan_pe(static_cast(ref_latent->ne[2]), + static_cast(ref_latent->ne[1]), + static_cast(ref_latent->ne[0]), + std::get<0>(config.patch_size), + std::get<1>(config.patch_size), + std::get<2>(config.patch_size), + 1, + config.theta, + config.axes_dim, + t_start); + pe_vec.insert(pe_vec.end(), ref_pe.begin(), ref_pe.end()); + } int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); // LOG_VERBOSE("pos_len %d", pos_len); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); @@ -933,7 +1046,10 @@ namespace WAN { clip_fea, time_dim_concat, vace_context, - vace_strength); + vace_strength, + 1, + audio_embed, + ref_latent); ggml_build_forward_expand(gf, out); @@ -948,9 +1064,11 @@ namespace WAN { const sd::Tensor& c_concat = {}, const sd::Tensor& time_dim_concat = {}, const sd::Tensor& vace_context = {}, - float vace_strength = 1.f) { + float vace_strength = 1.f, + const sd::Tensor& audio_embed = {}, + const sd::Tensor& ref_latent = {}) { auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength); + return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength, audio_embed, ref_latent); }; return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); @@ -961,6 +1079,12 @@ namespace WAN { GGML_ASSERT(diffusion_params.x != nullptr); GGML_ASSERT(diffusion_params.timesteps != nullptr); const auto* extra = diffusion_extra_as(diffusion_params); + static const std::vector> no_ref_latents; + const auto& ref_latents = config.model_type == "s2v" && diffusion_params.ref_latents != nullptr + ? *diffusion_params.ref_latents + : no_ref_latents; + const sd::Tensor empty_tensor; + const sd::Tensor& ref_latent = ref_latents.empty() ? empty_tensor : ref_latents[0]; return compute(n_threads, *diffusion_params.x, *diffusion_params.timesteps, @@ -969,7 +1093,9 @@ namespace WAN { tensor_or_empty(diffusion_params.c_concat), sd::Tensor(), tensor_or_empty(extra->vace_context), - extra->vace_strength); + extra->vace_strength, + tensor_or_empty(extra->audio_embed), + ref_latent); } void test() { diff --git a/src/model/diffusion/wan_audio.hpp b/src/model/diffusion/wan_audio.hpp new file mode 100644 index 00000000..28f31bda --- /dev/null +++ b/src/model/diffusion/wan_audio.hpp @@ -0,0 +1,215 @@ +#ifndef __SD_MODEL_DIFFUSION_WAN_AUDIO_HPP__ +#define __SD_MODEL_DIFFUSION_WAN_AUDIO_HPP__ + +#include +#include +#include +#include +#include + +#include "model/common/ggml_block.hpp" + +namespace WAN { + + class WanCausalConv1d : public UnaryBlock { + private: + int kernel_size_; + + public: + WanCausalConv1d(int64_t in_dim, + int64_t out_dim, + int kernel_size = 3, + int stride = 1) + : kernel_size_(kernel_size) { + blocks["conv"] = std::make_shared(in_dim, out_dim, kernel_size, stride, 0, 1, 1, true, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + // Replicate the first sample for causal left padding. + if (kernel_size_ > 1) { + auto first = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, 1); + for (int i = 0; i < kernel_size_ - 1; i++) { + x = ggml_concat(ctx->ggml_ctx, first, x, 0); + } + } + return std::dynamic_pointer_cast(blocks["conv"])->forward(ctx, x); + } + }; + + class WanMotionEncoder : public GGMLBlock { + private: + int64_t hidden_dim_; + int num_token_; + bool need_global_; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + // The padding token is combined with F32 activations. + params["padding_tokens"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_dim_); + } + + ggml_tensor* conv_norm_silu(GGMLRunnerContext* ctx, + ggml_tensor* x, + const std::string& conv_key, + const std::string& norm_key, + bool to_conv_layout) { + x = std::dynamic_pointer_cast(blocks[conv_key])->forward(ctx, x); + x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); + x = std::dynamic_pointer_cast(blocks[norm_key])->forward(ctx, x); + x = ggml_silu(ctx->ggml_ctx, x); + if (to_conv_layout) { + x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + } + return x; + } + + public: + WanMotionEncoder(int64_t in_dim, + int64_t hidden_dim, + int num_token, + bool need_global = true) + : hidden_dim_(hidden_dim), num_token_(num_token), need_global_(need_global) { + blocks["conv1_local"] = std::make_shared(in_dim, hidden_dim / 4 * num_token); + if (need_global) { + blocks["conv1_global"] = std::make_shared(in_dim, hidden_dim / 4); + } + blocks["norm1"] = std::make_shared(hidden_dim / 4, 1e-6f, false); + blocks["conv2"] = std::make_shared(hidden_dim / 4, hidden_dim / 2, 3, 2); + blocks["norm2"] = std::make_shared(hidden_dim / 2, 1e-6f, false); + blocks["conv3"] = std::make_shared(hidden_dim / 2, hidden_dim, 3, 2); + blocks["norm3"] = std::make_shared(hidden_dim, 1e-6f, false); + if (need_global) { + blocks["final_linear"] = std::make_shared(hidden_dim, hidden_dim); + } + } + + std::pair forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto local = std::dynamic_pointer_cast(blocks["conv1_local"])->forward(ctx, x); + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + std::vector tokens; + // Each token group is normalized independently over channels. + for (auto& group : ggml_ext_chunk(ctx->ggml_ctx, local, num_token_, 1)) { + ggml_tensor* s = ggml_permute(ctx->ggml_ctx, group, 1, 0, 2, 3); + s = norm1->forward(ctx, s); + s = ggml_silu(ctx->ggml_ctx, s); + s = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 0, 2, 3)); + s = conv_norm_silu(ctx, s, "conv2", "norm2", true); + s = conv_norm_silu(ctx, s, "conv3", "norm3", false); + tokens.push_back(ggml_reshape_3d(ctx->ggml_ctx, s, s->ne[0], 1, s->ne[1])); + } + auto padding = ggml_reshape_3d(ctx->ggml_ctx, params["padding_tokens"], hidden_dim_, 1, 1); + padding = ggml_repeat(ctx->ggml_ctx, padding, tokens[0]); + tokens.push_back(padding); + ggml_tensor* local_out = ggml_ext_vec_concat(ctx->ggml_ctx, tokens, 1); + + if (!need_global_) { + return {local_out, nullptr}; + } + ggml_tensor* g = conv_norm_silu(ctx, x, "conv1_global", "norm1", true); + g = conv_norm_silu(ctx, g, "conv2", "norm2", true); + g = conv_norm_silu(ctx, g, "conv3", "norm3", false); + g = std::dynamic_pointer_cast(blocks["final_linear"])->forward(ctx, g); + return {local_out, g}; + } + }; + + class WanCausalAudioEncoder : public GGMLBlock { + private: + int num_layers_; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + // Preserve the checkpoint shape for loading; layer mixing requires F32. + auto it = tensor_storage_map.find(prefix + "weights"); + if (it != tensor_storage_map.end()) { + params["weights"] = ggml_new_tensor(ctx, GGML_TYPE_F32, it->second.n_dims, it->second.ne); + } else { + params["weights"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_layers_); + } + } + + public: + WanCausalAudioEncoder(int64_t audio_dim, + int64_t dim, + int num_token, + int num_layers = 25) + : num_layers_(num_layers) { + blocks["encoder"] = std::make_shared(audio_dim, dim, num_token, true); + } + + // features: [layers, frames, audio_dim]; outputs: [T, tokens+1, dim] and [T, dim]. + std::pair forward(GGMLRunnerContext* ctx, ggml_tensor* features) { + auto weights = ggml_silu(ctx->ggml_ctx, params["weights"]); + auto x = ggml_mul(ctx->ggml_ctx, features, ggml_reshape_3d(ctx->ggml_ctx, weights, 1, 1, num_layers_)); + x = ggml_div(ctx->ggml_ctx, x, ggml_sum(ctx->ggml_ctx, weights)); + // Move the layer axis to ggml dimension 0 for reduction. + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); + x = ggml_sum_rows(ctx->ggml_ctx, x); + x = ggml_reshape_2d(ctx->ggml_ctx, x, x->ne[1], x->ne[2]); + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + return std::dynamic_pointer_cast(blocks["encoder"])->forward(ctx, x); + } + }; + + class WanAudioInjector : public GGMLBlock { + private: + int64_t dim_; + + public: + WanAudioInjector(int64_t dim, + int64_t num_heads, + int count, + bool qk_norm = true, + float eps = 1e-6f) + : dim_(dim) { + for (int i = 0; i < count; i++) { + blocks["injector." + std::to_string(i)] = + std::make_shared(dim, num_heads, qk_norm, eps); + blocks["injector_adain_layers." + std::to_string(i) + ".linear"] = + std::make_shared(dim, dim * 2); + } + // S2V AdaLayerNorm uses its own epsilon, independent of attention norms. + blocks["adain_norm"] = std::make_shared(dim, 1e-5f, false); + } + + // Inject into the video prefix; trailing reference tokens pass through unchanged. + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + int64_t seq_len, + int64_t T, + int injector_id, + ggml_tensor* audio_local, + ggml_tensor* audio_global) { + int64_t n_tok = seq_len / T; + int64_t n_token = x->ne[1]; + + auto adain_linear = std::dynamic_pointer_cast(blocks["injector_adain_layers." + std::to_string(injector_id) + ".linear"]); + auto injector = std::dynamic_pointer_cast(blocks["injector." + std::to_string(injector_id)]); + auto adain_norm = std::dynamic_pointer_cast(blocks["adain_norm"]); + + auto temb = ggml_silu(ctx->ggml_ctx, audio_global); + temb = adain_linear->forward(ctx, temb); + auto shift = ggml_ext_slice(ctx->ggml_ctx, temb, 0, 0, dim_); + auto scale = ggml_ext_slice(ctx->ggml_ctx, temb, 0, dim_, dim_ * 2); + shift = ggml_reshape_3d(ctx->ggml_ctx, shift, dim_, 1, T); + scale = ggml_reshape_3d(ctx->ggml_ctx, scale, dim_, 1, T); + + auto x_vid = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, seq_len); + auto h = ggml_reshape_3d(ctx->ggml_ctx, x_vid, dim_, n_tok, T); + h = adain_norm->forward(ctx, h); + h = ggml_add(ctx->ggml_ctx, h, ggml_mul(ctx->ggml_ctx, h, scale)); + h = ggml_add(ctx->ggml_ctx, h, shift); + + auto res = injector->forward(ctx, h, audio_local, 0); + res = ggml_reshape_2d(ctx->ggml_ctx, res, dim_, seq_len); + + auto x_head = ggml_add(ctx->ggml_ctx, x_vid, res); + if (seq_len < n_token) { + auto x_tail = ggml_ext_slice(ctx->ggml_ctx, x, 1, seq_len, n_token); + return ggml_concat(ctx->ggml_ctx, x_head, x_tail, 1); + } + return x_head; + } + }; + +} // namespace WAN + +#endif // __SD_MODEL_DIFFUSION_WAN_AUDIO_HPP__ diff --git a/src/model_component.h b/src/model_component.h index 96077a4a..b3e2a55b 100644 --- a/src/model_component.h +++ b/src/model_component.h @@ -10,6 +10,7 @@ enum class ModelComponent { VAE, PreviewVAE, AudioVAE, + AudioEncoder, ControlNet, PhotoMaker, PuLID, @@ -38,6 +39,8 @@ inline const char* model_component_name(ModelComponent component) { return "preview VAE"; case ModelComponent::AudioVAE: return "audio VAE"; + case ModelComponent::AudioEncoder: + return "audio encoder"; case ModelComponent::ControlNet: return "ControlNet"; case ModelComponent::PhotoMaker: diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index 807c915a..d8b8dc51 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -247,6 +247,11 @@ bool read_safetensors_file(const std::string& file_path, std::string dtype = tensor_info["dtype"]; nlohmann::json shape = tensor_info["shape"]; + // ComfyUI FP8 activation scales cancel when inference uses F16/F32 activations. + if (ends_with(name, ".scale_input")) { + continue; + } + size_t begin = tensor_info["data_offsets"][0].get(); size_t end = tensor_info["data_offsets"][1].get(); if (begin > end || end > file_size_ - data_start) { diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 6ab370c3..7b7db3a0 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -435,6 +435,7 @@ SDVersion ModelLoader::get_sd_version() const { bool is_flux2 = false; bool has_single_block_47 = false; bool is_wan = false; + bool is_s2v = false; int64_t patch_embedding_channels = 0; bool has_img_emb = false; bool has_middle_block_1 = false; @@ -524,6 +525,11 @@ SDVersion ModelLoader::get_sd_version() const { if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) { is_wan = true; } + if (tensor_storage.name.find("casual_audio_encoder.weights") != std::string::npos || + tensor_storage.name.find("audio_injector.injector.0.q.weight") != std::string::npos) { + // S2V and T2V-14B share patch_embedding shapes. + is_s2v = true; + } if (tensor_storage.name.find("model.diffusion_model.patch_embedder.weight") != std::string::npos) { return VERSION_LINGBOT_VIDEO; } @@ -587,6 +593,9 @@ SDVersion ModelLoader::get_sd_version() const { } if (is_wan) { LOG_VERBOSE("patch_embedding_channels %d", patch_embedding_channels); + if (is_s2v) { + return VERSION_WAN2_2_S2V; + } if (patch_embedding_channels == 184320 && !has_img_emb) { return VERSION_WAN2_2_I2V; } diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 18414017..ecad31a5 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -33,12 +33,14 @@ #include "extensions/generation_extension.h" #include "model/adapter/ip_adapter.hpp" #include "model/adapter/lora.hpp" +#include "model/audio/wav2vec2.hpp" #include "model/diffusion/animatediff.hpp" #include "model/diffusion/control.hpp" #include "model/diffusion/model.hpp" #include "model/vae/audio_vae.hpp" #include "model/vae/ltx_vae.hpp" #include "model/vae/vae.hpp" +#include "runtime/audio_processing.h" #include "runtime/denoiser.hpp" #include "runtime/guidance.h" #include "runtime/preview_interval.h" @@ -74,6 +76,7 @@ const char* model_version_to_str[] = { "Wan 2.x", "Wan 2.2 I2V", "Wan 2.2 TI2V", + "Wan 2.2 S2V", "LingBot Video", "Qwen Image", "Qwen Image Layered", @@ -136,7 +139,7 @@ StableDiffusionGGML::~StableDiffusionGGML() = default; const std::map>& StableDiffusionGGML::runner_components() { static const std::map> components{ - {RunnerGroup::Core, {ModelComponent::Conditioner, ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion, ModelComponent::CLIPVision, ModelComponent::IPAdapter}}, + {RunnerGroup::Core, {ModelComponent::Conditioner, ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion, ModelComponent::CLIPVision, ModelComponent::IPAdapter, ModelComponent::AudioEncoder}}, {RunnerGroup::VAE, {ModelComponent::VAE, ModelComponent::PreviewVAE, ModelComponent::AudioVAE}}, {RunnerGroup::ControlNet, {ModelComponent::ControlNet}}, {RunnerGroup::Extensions, {ModelComponent::PhotoMaker, ModelComponent::PuLID}}, @@ -804,6 +807,13 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf } } + if (strlen(SAFE_STR(sd_ctx_params->audio_encoder_path)) > 0) { + LOG_INFO("loading audio encoder (wav2vec2) from '%s'", sd_ctx_params->audio_encoder_path); + if (!model_loader.init_from_file(sd_ctx_params->audio_encoder_path, "wav2vec2.")) { + LOG_WARN("loading audio encoder weights from '%s' failed", sd_ctx_params->audio_encoder_path); + } + } + if (strlen(SAFE_STR(sd_ctx_params->motion_module_path)) > 0) { LOG_INFO("loading motion module (AnimateDiff) from '%s'", sd_ctx_params->motion_module_path); if (!model_loader.init_from_file(sd_ctx_params->motion_module_path, @@ -1010,6 +1020,7 @@ bool StableDiffusionGGML::build_core_runners() { high_noise_diffusion_model = std::move(runners.high_noise_diffusion); clip_vision = std::move(runners.clip_vision); ip_adapter = std::move(runners.ip_adapter); + audio_encoder = std::move(runners.audio_encoder); cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::TE)); diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION)); @@ -1019,11 +1030,15 @@ bool StableDiffusionGGML::build_core_runners() { if (clip_vision) { clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CLIP_VISION)); } + if (audio_encoder) { + audio_encoder->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::AUDIO_ENCODER)); + } return register_runner_params(ModelComponent::Conditioner, cond_stage_model, SDBackendModule::TE) && register_runner_params(ModelComponent::Diffusion, diffusion_model, SDBackendModule::DIFFUSION) && register_runner_params(ModelComponent::HighNoiseDiffusion, high_noise_diffusion_model, SDBackendModule::DIFFUSION) && register_runner_params(ModelComponent::CLIPVision, clip_vision, SDBackendModule::CLIP_VISION) && - register_runner_params(ModelComponent::IPAdapter, ip_adapter, SDBackendModule::DIFFUSION); + register_runner_params(ModelComponent::IPAdapter, ip_adapter, SDBackendModule::DIFFUSION) && + register_runner_params(ModelComponent::AudioEncoder, audio_encoder, SDBackendModule::AUDIO_ENCODER); } bool StableDiffusionGGML::build_vae_runners() { @@ -1121,6 +1136,12 @@ bool StableDiffusionGGML::validate_and_load_runners() { ignore_tensors.insert("model.diffusion_model.__32x32__"); ignore_tensors.insert("model.diffusion_model.__index_timestep_zero__"); + if (audio_encoder != nullptr) { + // These wav2vec2 tensors are unused during feature extraction. + ignore_tensors.insert("wav2vec2.lm_head."); + ignore_tensors.insert("wav2vec2.masked_spec_embed"); + } + if (audio_vae_model) { if (!sd_version_is_minimax_h3(version)) { ignore_tensors.insert("audio_vae.encoder"); @@ -1755,6 +1776,29 @@ sd::Tensor StableDiffusionGGML::get_clip_vision_output(const sd::Tensor StableDiffusionGGML::get_audio_embedding(const sd_audio_t& audio) { + if (audio_encoder == nullptr) { + LOG_ERROR("audio encoder model is not loaded"); + return {}; + } + if (audio.data == nullptr || audio.sample_count == 0 || audio.channels == 0 || audio.sample_rate == 0) { + LOG_ERROR("invalid driving audio"); + return {}; + } + auto mono = sd::audio::downmix_to_mono(audio.data, audio.sample_count, audio.channels); + if (mono.empty()) { + LOG_ERROR("audio mono downmix failed"); + return {}; + } + mono = sd::audio::resample_audio(mono.data(), mono.size(), audio.sample_rate, 16000); + if (mono.empty()) { + LOG_ERROR("audio resample to 16 kHz failed"); + return {}; + } + return audio_encoder->compute(n_threads, mono); +} + void StableDiffusionGGML::compute_ip_adapter_tokens(const sd_image_t& image, float strength) { ip_adapter_tokens = {}; ip_adapter_uncond_tokens = {}; @@ -1810,6 +1854,10 @@ std::vector StableDiffusionGGML::process_timesteps(const std::vectorget_desc() == "Wan2.2-S2V-14B") { + int64_t frame_count = init_latent.shape()[2]; + return std::vector(static_cast(frame_count), timesteps[0]); } else { return timesteps; } @@ -2318,7 +2366,8 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr first_stage_model; std::shared_ptr preview_vae; std::shared_ptr audio_vae_model; + std::shared_ptr audio_encoder; std::shared_ptr control_net; std::shared_ptr ip_adapter; sd::Tensor ip_adapter_tokens; @@ -363,6 +367,8 @@ public: int clip_skip = -1, bool zero_out_masked = false); + sd::Tensor get_audio_embedding(const sd_audio_t& audio); + void compute_ip_adapter_tokens(const sd_image_t& image, float strength); std::vector process_timesteps(const std::vector& timesteps, diff --git a/src/pipeline/generation.h b/src/pipeline/generation.h index 1210a8cb..34aa000a 100644 --- a/src/pipeline/generation.h +++ b/src/pipeline/generation.h @@ -28,6 +28,7 @@ namespace sd::pipeline { sd::Tensor denoise_mask; sd::Tensor clip_vision_output; sd::Tensor vace_context; + sd::Tensor s2v_audio_embed; int64_t ref_image_num = 0; int64_t video_conditioning_frame_count = 0; int64_t video_target_frame_count = 0; @@ -59,7 +60,8 @@ namespace sd::pipeline { const sd_vid_gen_params_t* sd_vid_gen_params, sd_image_t** frames_out, int* num_frames_out, - sd_audio_t** audio_out); + sd_audio_t** audio_out, + int* fps_out); sd::Tensor upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd, const char* model_path, diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index da5a469e..23d6dfa9 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -8,6 +8,7 @@ #include "core/util.h" #include "extensions/generation_extension.h" #include "model/adapter/ip_adapter.hpp" +#include "model/audio/wav2vec2.hpp" #include "model/diffusion/anima.hpp" #include "model/diffusion/boogu.hpp" #include "model/diffusion/control.hpp" @@ -234,6 +235,16 @@ namespace sd::model_builders { tensor_storage_map, weight_manager); } + if (version == VERSION_WAN2_2_S2V && + tensor_storage_map.count("wav2vec2.encoder.layer_norm.bias") > 0) { + if (!ensure_backend_pair(ctx.backends, SDBackendModule::AUDIO_ENCODER)) { + return false; + } + result.audio_encoder = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::AUDIO_ENCODER), + tensor_storage_map, + "wav2vec2.", + weight_manager); + } } else if (sd_version_is_lingbot_video(version)) { bool enable_vision = false; for (const auto& [name, _] : tensor_storage_map) { diff --git a/src/pipeline/model_builders.h b/src/pipeline/model_builders.h index 73d134b3..5edc7568 100644 --- a/src/pipeline/model_builders.h +++ b/src/pipeline/model_builders.h @@ -15,6 +15,9 @@ struct DiffusionModelRunner; struct VAE; struct AudioVAERunner; struct ControlNet; +namespace Wav2Vec2 { + class Wav2Vec2ModelRunner; +} struct GenerationExtension; struct GenerationExtensionInitContext; namespace IPAdapter { @@ -37,6 +40,7 @@ namespace sd::model_builders { std::shared_ptr high_noise_diffusion; std::shared_ptr clip_vision; std::shared_ptr ip_adapter; + std::shared_ptr audio_encoder; }; struct VAEOptions { diff --git a/src/pipeline/request.cpp b/src/pipeline/request.cpp index aa78f7e6..2874bbae 100644 --- a/src/pipeline/request.cpp +++ b/src/pipeline/request.cpp @@ -161,7 +161,10 @@ namespace sd::pipeline { frames = sd->align_video_frames(requested_frames); clip_skip = sd_vid_gen_params->clip_skip; fps = std::max(1, sd_vid_gen_params->fps); - if (sd_version_is_minimax_h3(sd->version) && fps != 24) { + if (sd->version == VERSION_WAN2_2_S2V && sd_vid_gen_params->fps != 16) { + LOG_WARN("Wan2.2 S2V uses 16 fps; overriding requested fps %d", sd_vid_gen_params->fps); + fps = 16; + } else if (sd_version_is_minimax_h3(sd->version) && fps != 24) { LOG_WARN("MiniMax-H3 uses 24 fps; overriding requested fps %d", fps); fps = 24; } diff --git a/src/pipeline/video.cpp b/src/pipeline/video.cpp index cc979c0f..7b9c8bde 100644 --- a/src/pipeline/video.cpp +++ b/src/pipeline/video.cpp @@ -5,6 +5,7 @@ #include #include +#include "conditioning/wan_audio.h" #include "core/rng.hpp" #include "core/rng_philox.hpp" #include "diffusion_engine.h" @@ -420,6 +421,45 @@ namespace sd::pipeline { return audio; } + // Build the first 16 fps audio window, zero-padding past the track end. + static sd::Tensor build_s2v_audio_window(const sd::Tensor& stacked, int64_t batch_frames) { + const int64_t embed_dim = stacked.shape()[0]; + const int64_t in_frames = stacked.shape()[1]; + const int64_t num_layers = stacked.shape()[2]; + if (embed_dim <= 0 || in_frames <= 0 || num_layers <= 0 || batch_frames <= 0) { + return {}; + } + std::vector layer_first(static_cast(num_layers) * in_frames * embed_dim); + for (int64_t l = 0; l < num_layers; ++l) { + for (int64_t f = 0; f < in_frames; ++f) { + const float* src = stacked.data() + l * embed_dim * in_frames + f * embed_dim; + std::copy_n(src, + static_cast(embed_dim), + layer_first.data() + (static_cast(l) * in_frames + f) * embed_dim); + } + } + sd::wan_audio::BucketPlan plan; + std::vector buckets = sd::wan_audio::build_audio_buckets(layer_first.data(), + static_cast(num_layers), + static_cast(in_frames), + static_cast(embed_dim), + static_cast(batch_frames), + &plan); + if (buckets.empty() || plan.bucket_frames < batch_frames) { + return {}; + } + // Reorder frame-major buckets into sd::Tensor's [dim, frame, layer] layout. + sd::Tensor window({embed_dim, batch_frames, num_layers}); + for (int64_t f = 0; f < batch_frames; ++f) { + for (int64_t l = 0; l < num_layers; ++l) { + const float* src = buckets.data() + (static_cast(f) * num_layers + l) * embed_dim; + float* dst = window.data() + l * embed_dim * batch_frames + f * embed_dim; + std::copy_n(src, static_cast(embed_dim), dst); + } + } + return window; + } + static std::optional prepare_video_generation_latents(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params, GenerationRequest* request) { @@ -1033,6 +1073,53 @@ namespace sd::pipeline { latents.vace_context = sd::ops::concat(vace_context, mask_context, 3); // [b, 2*c + vae_scale_factor*vae_scale_factor, t + 1 or t, h/vae_scale_factor, w/vae_scale_factor] int64_t t2 = ggml_time_ms(); LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1); + } else if (sd->diffusion_model->get_desc() == "Wan2.2-S2V-14B") { + LOG_INFO("S2V"); + if (!end_image.empty()) { + LOG_WARN("Wan2.2 S2V ignores end_image"); + } + if (sd_vid_gen_params->ref_audios_count > 1) { + LOG_ERROR("Wan2.2 S2V supports a single driving audio track"); + return std::nullopt; + } + int64_t t1 = ggml_time_ms(); + if (!start_image.empty()) { + auto ref_img = start_image.reshape({start_image.shape()[0], + start_image.shape()[1], + 1, + start_image.shape()[2], + 1}); + auto encoded_ref = sd->encode_first_stage(ref_img); + if (encoded_ref.empty()) { + LOG_ERROR("failed to encode S2V reference image"); + return std::nullopt; + } + // Wan consumes reference latents in 4D. + latents.ref_latents.push_back(encoded_ref.reshape({encoded_ref.shape()[0], + encoded_ref.shape()[1], + encoded_ref.shape()[2], + encoded_ref.shape()[3]})); + } + if (sd_vid_gen_params->ref_audios_count == 1) { + if (sd->audio_encoder == nullptr) { + LOG_ERROR("S2V audio conditioning requires --audio-encoder (wav2vec2)"); + return std::nullopt; + } + auto stacked = sd->get_audio_embedding(sd_vid_gen_params->ref_audios[0]); + if (stacked.empty()) { + LOG_ERROR("failed to compute wav2vec2 embedding for driving audio"); + return std::nullopt; + } + int64_t latent_t = sd->video_frames_to_latent_frames(request->frames); + int64_t batch_frames = latent_t * 4; + latents.s2v_audio_embed = build_s2v_audio_window(stacked, batch_frames); + if (latents.s2v_audio_embed.empty()) { + LOG_ERROR("failed to build S2V audio window"); + return std::nullopt; + } + } + int64_t t2 = ggml_time_ms(); + LOG_INFO("s2v conditioning prepared, taking %" PRId64 " ms", t2 - t1); } if (latents.init_latent.empty()) { @@ -1084,6 +1171,12 @@ namespace sd::pipeline { latents.keyframe_indices); } } + if (sd->version == VERSION_WAN2_2_S2V) { + embeds.cond.c_ref_images = latents.ref_latents; + if (!latents.s2v_audio_embed.empty()) { + embeds.cond.c_ref_audios = {latents.s2v_audio_embed}; + } + } if (request.use_uncond) { condition_params.text = request.negative_prompt; embeds.uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads, @@ -1096,6 +1189,12 @@ namespace sd::pipeline { embeds.uncond.c_reference_blocks = latents.minimax_reference_blocks; embeds.uncond.c_position_ids = embeds.cond.c_position_ids; } + if (sd->version == VERSION_WAN2_2_S2V) { + embeds.uncond.c_ref_images = latents.ref_latents; + if (!latents.s2v_audio_embed.empty()) { + embeds.uncond.c_ref_audios = {sd::Tensor::zeros_like(latents.s2v_audio_embed)}; + } + } } int64_t t1 = ggml_time_ms(); @@ -1422,10 +1521,14 @@ namespace sd::pipeline { const sd_vid_gen_params_t* sd_vid_gen_params, sd_image_t** frames_out, int* num_frames_out, - sd_audio_t** audio_out) { + sd_audio_t** audio_out, + int* fps_out) { if (sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd->version)) { LOG_INFO("AnimateDiff dispatch: %d frames, %dx%d", sd_vid_gen_params->video_frames, sd_vid_gen_params->width, sd_vid_gen_params->height); + if (fps_out != nullptr) { + *fps_out = std::max(1, sd_vid_gen_params->fps); + } return generate_animatediff_video(sd, sd_vid_gen_params, frames_out, num_frames_out); } @@ -1437,6 +1540,9 @@ namespace sd::pipeline { sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params; sd->apply_circular_axes(sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y); GenerationRequest request(sd, sd_vid_gen_params); + if (fps_out != nullptr) { + *fps_out = request.fps; + } bool latent_upscale_enabled = request.hires.enabled; GenerationRequest hires_request = request; if (latent_upscale_enabled) { @@ -1725,6 +1831,33 @@ namespace sd::pipeline { LOG_INFO("generating latent video completed, taking %.2fs", (latent_end - latent_start) * 1.0f / 1000); sd_audio_t* generated_audio = nullptr; + if (sd->version == VERSION_WAN2_2_S2V && sd_vid_gen_params->ref_audios_count > 0) { + // Return the driving track for muxing with the generated video. + const sd_audio_t& driving = sd_vid_gen_params->ref_audios[0]; + generated_audio = (sd_audio_t*)malloc(sizeof(sd_audio_t)); + if (generated_audio != nullptr) { + generated_audio->sample_rate = driving.sample_rate; + generated_audio->channels = driving.channels; + generated_audio->sample_count = driving.sample_count; + generated_audio->data = (float*)malloc(sizeof(float) * driving.sample_count * driving.channels); + if (generated_audio->data == nullptr) { + free(generated_audio); + generated_audio = nullptr; + } else { + memcpy(generated_audio->data, + driving.data, + sizeof(float) * driving.sample_count * driving.channels); + } + } + if (generated_audio != nullptr) { + LOG_DEBUG("s2v output audio: %u Hz, %u channels, %llu samples", + generated_audio->sample_rate, + generated_audio->channels, + (unsigned long long)generated_audio->sample_count); + } else { + LOG_DEBUG("s2v output audio copy failed (out of memory)"); + } + } if ((sd_version_is_ltxav(sd->version) || sd_version_is_minimax_h3(sd->version)) && latents.audio_length > 0 && sd->audio_vae_model != nullptr) { @@ -1774,6 +1907,7 @@ namespace sd::pipeline { return false; } auto result = decode_video_outputs(sd, latent_upscale_enabled ? hires_request : request, final_latent, num_frames_out); + LOG_DEBUG("decode_video_outputs returned %s", result == nullptr ? "nullptr (failed)" : "frames"); if (result == nullptr) { free_sd_audio(generated_audio); return false; @@ -1786,6 +1920,20 @@ namespace sd::pipeline { if (frames_out != nullptr) { *frames_out = result; } + if (sd->version == VERSION_WAN2_2_S2V && generated_audio != nullptr) { + // Limit the driving track to the generated video's duration. + int fps = request.fps; + uint64_t video_frames = num_frames_out != nullptr ? (uint64_t)*num_frames_out : 0; + uint64_t want_samples = (uint64_t)((double)video_frames / fps * generated_audio->sample_rate); + LOG_DEBUG("s2v audio truncate: %llu samples -> %llu (video %llu frames @ %d fps)", + (unsigned long long)generated_audio->sample_count, + (unsigned long long)want_samples, + (unsigned long long)video_frames, + fps); + if (want_samples > 0 && want_samples < generated_audio->sample_count) { + generated_audio->sample_count = want_samples; + } + } if (audio_out != nullptr) { *audio_out = generated_audio; } else { diff --git a/src/runtime/audio_processing.cpp b/src/runtime/audio_processing.cpp new file mode 100644 index 00000000..fe02b6ea --- /dev/null +++ b/src/runtime/audio_processing.cpp @@ -0,0 +1,97 @@ +#include "audio_processing.h" + +#include +#include +#include +#include + +namespace sd::audio { + + // Match torchaudio's Hann-windowed sinc resampler. + std::vector resample_audio(const float* samples, + uint64_t sample_count, + uint32_t orig_sample_rate, + uint32_t target_sample_rate) { + if (samples == nullptr || sample_count == 0 || orig_sample_rate == 0 || target_sample_rate == 0) { + return {}; + } + if (orig_sample_rate == target_sample_rate) { + return std::vector(samples, samples + sample_count); + } + + constexpr int kLowpassFilterWidth = 6; + constexpr double kRolloff = 0.99; + constexpr double kPi = 3.14159265358979323846; + + const uint64_t gcd = std::gcd(static_cast(orig_sample_rate), + static_cast(target_sample_rate)); + const int64_t orig_freq = static_cast(orig_sample_rate / gcd); + const int64_t new_freq = static_cast(target_sample_rate / gcd); + const double base_freq = static_cast(std::min(orig_freq, new_freq)) * kRolloff; + const int64_t width = static_cast(std::ceil(kLowpassFilterWidth * orig_freq / base_freq)); + const int64_t kernel_size = 2 * width + orig_freq; + + std::vector kernel(static_cast(new_freq) * kernel_size); + for (int64_t j = 0; j < new_freq; ++j) { + for (int64_t i = 0; i < kernel_size; ++i) { + double t = -static_cast(j) / new_freq + static_cast(i - width) / orig_freq; + t *= base_freq; + t = std::clamp(t, -static_cast(kLowpassFilterWidth), static_cast(kLowpassFilterWidth)); + const double cos_arg = std::cos(t * kPi / kLowpassFilterWidth / 2); + const double window = cos_arg * cos_arg; + double s = t * kPi; + const double sinc = (s == 0.0) ? 1.0 : std::sin(s) / s; + kernel[j * kernel_size + i] = sinc * window * (base_freq / orig_freq); + } + } + + const uint64_t num_phases = static_cast(sample_count / orig_freq) + 1; + const uint64_t target_length = (static_cast(new_freq) * sample_count + + static_cast(orig_freq) - 1) / + static_cast(orig_freq); + std::vector out(target_length); + for (uint64_t phase = 0; phase < num_phases; ++phase) { + const int64_t src_base = static_cast(phase * orig_freq) - width; + for (int64_t j = 0; j < new_freq; ++j) { + const uint64_t out_index = phase * new_freq + j; + if (out_index >= target_length) { + break; + } + const double* k = &kernel[j * kernel_size]; + double acc = 0.0; + for (int64_t i = 0; i < kernel_size; ++i) { + const int64_t src = src_base + i; + if (src >= 0 && src < static_cast(sample_count)) { + acc += samples[src] * k[i]; + } + } + out[out_index] = static_cast(acc); + } + } + return out; + } + + std::vector downmix_to_mono(const float* interleaved_samples, + uint64_t sample_count, + uint32_t channels) { + std::vector mono; + if (interleaved_samples == nullptr || sample_count == 0 || channels == 0) { + return mono; + } + mono.resize(static_cast(sample_count)); + if (channels == 1) { + std::memcpy(mono.data(), interleaved_samples, static_cast(sample_count) * sizeof(float)); + return mono; + } + const float scale = 1.0f / static_cast(channels); + for (uint64_t i = 0; i < sample_count; ++i) { + float sum = 0.0f; + for (uint32_t c = 0; c < channels; ++c) { + sum += interleaved_samples[i * channels + c]; + } + mono[static_cast(i)] = sum * scale; + } + return mono; + } + +} // namespace sd::audio diff --git a/src/runtime/audio_processing.h b/src/runtime/audio_processing.h new file mode 100644 index 00000000..cdda0c46 --- /dev/null +++ b/src/runtime/audio_processing.h @@ -0,0 +1,22 @@ +#ifndef __SD_RUNTIME_AUDIO_PROCESSING_H__ +#define __SD_RUNTIME_AUDIO_PROCESSING_H__ + +#include +#include + +namespace sd::audio { + + // Returns the input unchanged when sample rates are equal, and an empty vector on invalid input. + std::vector resample_audio(const float* samples, + uint64_t sample_count, + uint32_t orig_sample_rate, + uint32_t target_sample_rate); + + // Average interleaved channels; return an empty vector on invalid input. + std::vector downmix_to_mono(const float* interleaved_samples, + uint64_t sample_count, + uint32_t channels); + +} // namespace sd::audio + +#endif // __SD_RUNTIME_AUDIO_PROCESSING_H__ diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index beb899cd..4fb8771a 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -355,6 +355,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "embeddings_connectors_path: %s\n" "vae_path: %s\n" "audio_vae_path: %s\n" + "audio_encoder_path: %s\n" "taesd_path: %s\n" "control_net_path: %s\n" "photo_maker_path: %s\n" @@ -392,6 +393,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { SAFE_STR(sd_ctx_params->embeddings_connectors_path), SAFE_STR(sd_ctx_params->vae_path), SAFE_STR(sd_ctx_params->audio_vae_path), + SAFE_STR(sd_ctx_params->audio_encoder_path), SAFE_STR(sd_ctx_params->taesd_path), SAFE_STR(sd_ctx_params->control_net_path), SAFE_STR(sd_ctx_params->photo_maker_path), @@ -736,8 +738,12 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params, sd_image_t** frames_out, int* num_frames_out, - sd_audio_t** audio_out) { + sd_audio_t** audio_out, + int* fps_out) { if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) { + if (fps_out != nullptr) { + *fps_out = 0; + } return false; } @@ -753,10 +759,13 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd); if (!execution.ready) { + if (fps_out != nullptr) { + *fps_out = 0; + } return false; } - return sd::pipeline::generate_video(sd_ctx->sd, sd_vid_gen_params, frames_out, num_frames_out, audio_out); + return sd::pipeline::generate_video(sd_ctx->sd, sd_vid_gen_params, frames_out, num_frames_out, audio_out, fps_out); } SD_API void free_sd_images(sd_image_t* result_images, int num_images) {