From 46d0152345f5a955247a559b2a653cb755a401d7 Mon Sep 17 00:00:00 2001 From: noctrex Date: Sun, 30 Aug 2026 17:55:47 +0300 Subject: [PATCH 1/5] feat: add Wan2.2 S2V (speech-to-video) support Audio-driven video generation for Wan2.2-S2V-14B: a reference portrait plus a driving WAV produce a lip-synced video whose output container carries the driving audio. - detect S2V checkpoints by audio-module key presence; uses the wan2.1 VAE - wav2vec2 audio encoder model class, backend module, --audio-encoder flag - S2V modules in the Wan DiT: causal audio encoder, motion encoder, audio injector (AdaLN + cross-attention), reference-latent tokens - audio windowing/bucketing to the latent frame count; uncond uses silence - new --audio flag; mux the driving track into avi/webm (wav sidecar) - skip ComfyUI fp8_scaled .scale_input keys at parse time - docs: downloads, S2V example, sampler and 24 GB VRAM guidance --- docs/wan.md | 67 +++ examples/common/common.cpp | 21 + examples/common/common.h | 1 + examples/common/media_io.cpp | 34 ++ examples/common/media_io.h | 18 + include/stable-diffusion.h | 1 + src/conditioning/audio_processing.hpp | 228 ++++++++ src/core/ggml_extend_backend.cpp | 6 + src/core/ggml_extend_backend.h | 1 + src/model.h | 3 +- src/model/audio/wav2vec2.hpp | 746 ++++++++++++++++++++++++++ src/model/common/rope.hpp | 8 +- src/model/diffusion/model.hpp | 2 + src/model/diffusion/wan.hpp | 420 +++++++++++++-- src/model_component.h | 3 + src/model_io/safetensors_io.cpp | 8 + src/model_loader.cpp | 9 + src/pipeline/diffusion_engine.cpp | 59 +- src/pipeline/diffusion_engine.h | 4 + src/pipeline/generation.h | 1 + src/pipeline/model_builders.cpp | 11 + src/pipeline/model_builders.h | 2 + src/pipeline/video.cpp | 155 ++++++ src/stable-diffusion.cpp | 2 + 24 files changed, 1770 insertions(+), 40 deletions(-) create mode 100644 src/conditioning/audio_processing.hpp create mode 100644 src/model/audio/wav2vec2.hpp diff --git a/docs/wan.md b/docs/wan.md index 02b46d62a..ebe4ebe0a 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,66 @@ +### 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. +- Use `--diffusion-fa` together with `--offload-to-cpu --vae-tiling`. Without + flash attention the S2V attention compute buffer alone needs ~40 GB VRAM at + 480x832; `--offload-to-cpu` keeps the model weights in system RAM, and + `--vae-tiling` keeps the video VAE decode (77-81 frames) within VRAM. +- `--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. +- The output video carries the driving audio track: it is muxed into `.avi` + / `.webm` outputs, truncated to the video duration. Other container types + fall back to a `.wav` sidecar file next to the video. +- Approximate VRAM use at 640x368, 13 frames: + + | checkpoint | VRAM | + |------------|------| + | wan2.2_s2v_14B_int8_convrot.safetensors | 17.6 GB | + | wan2.2_s2v-14B-Q8_0.gguf | 18.4 GB | + | wan2.2_s2v-14B-Q4_K_M.gguf | 13.4 GB | + + Q4_K_M is the lowest-memory option; bf16 and fp8_scaled need more than + 24 GB VRAM. +- 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. +- fp8_scaled and GGUF checkpoints are supported. The int8_convrot checkpoint + is available at + https://huggingface.co/noctrex/Wan2.2-S2V-14B-int8_convrot + ([int8_convrot](int8_convrot.md)); int8 convrot is supported on the CUDA, + Vulkan and ROCm backends. + ### Wan2.2 T2V A14B T2I ``` diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 335a6226f..400dd867d 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -432,6 +432,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)", @@ -858,6 +863,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" @@ -921,6 +927,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(); @@ -1471,6 +1478,16 @@ ArgOptions SDGenerationParams::get_options() { return 1; }; + auto on_audio_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + // S2V driving track shares the ref-audio transport (Wan2.2 S2V reads + // sd_vid_gen_params.ref_audios[0]). + 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; @@ -1660,6 +1677,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 ef084dfd4..92d03feed 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/common/media_io.cpp b/examples/common/media_io.cpp index 721aa6641..58b49a14c 100644 --- a/examples/common/media_io.cpp +++ b/examples/common/media_io.cpp @@ -1,4 +1,5 @@ #include "media_io.h" +#include "conditioning/audio_processing.hpp" #include "log.h" #include "resource_owners.hpp" @@ -10,6 +11,7 @@ #include #include #include +#include #include #include @@ -1506,3 +1508,35 @@ bool load_wav_from_file(const std::string& path, } return true; } + +// DSP core lives in src/conditioning/audio_processing.hpp (shared with +// stable-diffusion.cpp); these wrappers keep the examples-side API. +std::vector downmix_to_mono(const float* interleaved_samples, + uint64_t sample_count, + uint32_t channels) { + return AudioProcessing::downmix_to_mono(interleaved_samples, sample_count, channels); +} + +std::vector resample_audio(const float* samples, + uint64_t sample_count, + uint32_t orig_sample_rate, + uint32_t target_sample_rate) { + return AudioProcessing::resample_audio(samples, sample_count, orig_sample_rate, target_sample_rate); +} + +bool load_wav_from_file_mono(const std::string& path, + std::vector& mono_samples, + uint32_t target_sample_rate) { + std::vector interleaved; + uint32_t sample_rate = 0; + uint32_t channels = 0; + if (!load_wav_from_file(path, interleaved, sample_rate, channels)) { + return false; + } + const uint64_t frames = channels > 0 ? interleaved.size() / channels : 0; + mono_samples = downmix_to_mono(interleaved.data(), frames, channels); + if (sample_rate != target_sample_rate && !mono_samples.empty()) { + mono_samples = resample_audio(mono_samples.data(), mono_samples.size(), sample_rate, target_sample_rate); + } + return true; +} diff --git a/examples/common/media_io.h b/examples/common/media_io.h index 52b9221a2..0c2230674 100644 --- a/examples/common/media_io.h +++ b/examples/common/media_io.h @@ -115,4 +115,22 @@ bool load_wav_from_file(const std::string& path, uint32_t& sample_rate, uint32_t& channels); +// Downmix interleaved samples to mono by averaging channels. Returns empty vector on invalid input. +std::vector downmix_to_mono(const float* interleaved_samples, + uint64_t sample_count, + uint32_t channels); + +// Band-limited sinc resampler matching torchaudio.functional.resample +// (sinc_interp_hann, lowpass_filter_width 6, rolloff 0.99). 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); + +// Convenience: load WAV, downmix to mono and resample to the target sample rate. +bool load_wav_from_file_mono(const std::string& path, + std::vector& mono_samples, + uint32_t target_sample_rate = 16000); + #endif // __MEDIA_IO_H__ diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 08bc27aee..62ce26a41 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -207,6 +207,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; diff --git a/src/conditioning/audio_processing.hpp b/src/conditioning/audio_processing.hpp new file mode 100644 index 000000000..88eb7a731 --- /dev/null +++ b/src/conditioning/audio_processing.hpp @@ -0,0 +1,228 @@ +#ifndef __SD_CONDITIONING_AUDIO_PROCESSING_HPP__ +#define __SD_CONDITIONING_AUDIO_PROCESSING_HPP__ + +// Wan2.2-S2V audio windowing, ported from ComfyUI comfy_extras/nodes_wan.py +// (linear_interpolation + get_audio_embed_bucket_fps, m=0). +// +// Input is the wav2vec2 hidden states stacked per layer [num_layers, in_frames, dim] +// at the encoder frame rate (50 Hz). The frames are interpolated to video_rate +// (30 Hz), bucketed to fps (16) frames with zero padding past the audio end, and +// split into chunks of batch_frames = latent_t * 4 frames (one per diffusion chunk). + +#include +#include +#include +#include +#include +#include + +namespace AudioProcessing { + + // Chunk/padding math of get_audio_embed_bucket_fps (m=0). + struct BucketPlan { + int audio_frames; // input frames at video_rate (30 Hz) + int batch_frames; // pixel frames per chunk (latent_t * 4) + int video_rate; // timeline rate of audio_frames (30 Hz) + int fps; // bucket frame rate (16 fps) + int num_chunks; // ComfyUI num_repeat + int bucket_frames; // total bucket frames = num_chunks * batch_frames + int padded_audio_frames; // audio_frames plus zero padding applied + }; + + inline BucketPlan plan_buckets(int audio_frames, int batch_frames, int video_rate = 30, int fps = 16) { + 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; + // min_batch_num = int(audio_frame_num / (batch_frames * scale)) + 1 + plan.num_chunks = static_cast(audio_frames / (batch_frames * scale)) + 1; + plan.bucket_frames = plan.num_chunks * batch_frames; + // padd_audio_num = ceil(bucket_frames / fps * video_rate) - audio_frame_num + plan.padded_audio_frames = static_cast( + std::ceil(plan.bucket_frames / static_cast(fps) * video_rate)); + return plan; + } + + // Bucket frame index (fps timeline) -> source frame index (video_rate timeline). + // get_sample_indices with fixed_start=0 reduces to round-half-even(i * video_rate / fps), + // matching numpy's default rounding. + inline int bucket_source_frame(int bucket_frame, int video_rate = 30, int fps = 16) { + return static_cast(std::nearbyint(static_cast(bucket_frame) * video_rate / fps)); + } + + // torch.nn.functional.interpolate size computation: output_len = int(in_len / input_fps * output_fps) + inline int interpolated_frame_count(int in_frames, int input_fps = 50, int output_fps = 30) { + return static_cast(in_frames / static_cast(input_fps) * output_fps); + } + + // torch.nn.functional.interpolate(mode='linear', align_corners=True) along the frame + // dimension. in: [num_layers, in_frames, dim], out: [num_layers, out_frames, dim]. + inline 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; + } + + // Polyphase FIR from torchaudio.functional.resample: sinc interpolated at output + // phases, hann-windowed over lowpass_filter_width zero crossings, anti-aliased by + // clamping the sinc argument to +-lowpass_filter_width after scaling by min(rate)*rolloff. + // Returns the input unchanged when sample rates are equal, and an empty vector on + // invalid input. + inline 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; + } + + // Downmix interleaved samples to mono by averaging channels. Returns an empty + // vector on invalid input. + inline 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; + } + + // Full bucketing: stacked encoder states [num_layers, in_frames, dim] at input_fps -> + // bucket frames [bucket_frames, num_layers, dim] at fps, with zero frames past the + // audio end. Chunk c occupies rows [c * batch_frames, (c + 1) * batch_frames). + // Returns an empty vector on invalid input; the applied plan is stored in *plan_out. + inline 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) { + 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; // zero padding past the audio end + } + std::copy_n(interpolated.data() + static_cast(src) * num_layers * dim, + static_cast(num_layers) * dim, + buckets.data() + static_cast(frame) * num_layers * dim); + } + return buckets; + } + +} // namespace AudioProcessing + +#endif // __SD_CONDITIONING_AUDIO_PROCESSING_HPP__ diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index 9c4cfa42a..1782180c5 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; } @@ -962,6 +966,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 01652cbfb..7ee98e693 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 f59f1de7d..ae4319732 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, @@ -129,7 +130,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 000000000..cc542b8dd --- /dev/null +++ b/src/model/audio/wav2vec2.hpp @@ -0,0 +1,746 @@ +#ifndef __SD_MODEL_AUDIO_WAV2VEC2_HPP__ +#define __SD_MODEL_AUDIO_WAV2VEC2_HPP__ + +#include + +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "model/common/ggml_block.hpp" +#include "model.h" + +/*================================================ Wav2Vec2 audio encoder ================================================*/ + +// Port of ComfyUI comfy/audio_encoders/wav2vec2.py (HF wav2vec2 inference subset). +// Checkpoint: wav2vec2-large-english (embed_dim 1024, 24 layers, 16 heads, +// conv_norm/conv_bias true, stable layer norm). Keys carry a "wav2vec2." prefix; +// load the storage map with that prefix so names line up with the block tree. + +struct Wav2Vec2Config { + int64_t embed_dim = 1024; + int64_t conv_dim = 512; + int num_heads = 16; + int num_layers = 24; + bool conv_norm = true; + bool conv_bias = true; + bool do_normalize = true; + bool do_stable_layer_norm = true; + + // Mirrors comfy/audio_encoders/audio_encoders.py:49-76: embed_dim inferred from + // the shape of encoder.layer_norm.bias. + 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) { // large + config.embed_dim = 1024; + config.num_heads = 16; + config.num_layers = 24; + config.conv_norm = true; + config.conv_bias = true; + config.do_normalize = true; + config.do_stable_layer_norm = true; + } else if (config.embed_dim == 768) { // base + config.embed_dim = 768; + config.num_heads = 12; + config.num_layers = 12; + config.conv_norm = false; + 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; + } +}; + +// Conv1d with optional groups (weight [kernel, in_channels, out_channels]). For +// groups == channels == out_channels this stays depthwise-capable via ggml; general +// groups run per-group ggml_conv_1d on channel views and concat. +struct Wav2Vec2Conv1d : public UnaryBlock { + int64_t in_channels; + int64_t out_channels; + int64_t groups; + int kernel_size; + int stride; + int padding; + int dilation; + bool bias; + std::string prefix; + + Wav2Vec2Conv1d(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) + : in_channels(in_channels), + out_channels(out_channels), + groups(groups), + kernel_size(kernel_size), + stride(stride), + padding(padding), + dilation(dilation), + bias(bias) {} + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + this->prefix = prefix; + 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); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + // x: [L, in_channels, N] + ggml_tensor* w = params["weight"]; + ggml_tensor* b = bias ? params["bias"] : nullptr; + + GGML_ASSERT(x->ne[1] == in_channels); + + if (groups == 1) { + // F32 im2col + F32 weight keeps the conv core in full precision (see + // conv_1d_f32_im2col); mul_mat's mixed-type path requires an F32 + // right-hand side, so cast F16 checkpoint weights once. + ggml_tensor* w32 = w->type == GGML_TYPE_F32 ? w : ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F32); + ggml_tensor* out = conv_1d_f32_im2col(ctx->ggml_ctx, w32, x, stride, padding, dilation); + if (b != nullptr) { + b = ggml_reshape_3d(ctx->ggml_ctx, b, 1, out_channels, 1); + out = ggml_add_inplace(ctx->ggml_ctx, out, b); + } + return out; + } + return grouped_conv_1d(ctx, x, w, groups, stride, padding, b); + } + + // Grouped conv via per-group channel views; output channels concatenated in order. + // ggml_conv_1d hardcodes an F16 im2col destination, which rounds the input + // patches even for F32 weights; wav2vec2 carries large activation outliers + // that this rounding amplifies through 24 layers, so the grouped path uses + // ggml_conv_1d's exact composition with an F32 im2col instead (the pos_conv + // weight is always F32 here - it is recomputed in-graph from weight_norm). + static ggml_tensor* conv_1d_f32_im2col(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, int s0, int p0, int d0) { + ggml_tensor* im2col = ggml_im2col(ctx, w, x, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F32); + ggml_tensor* result = ggml_mul_mat(ctx, + ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[2] * im2col->ne[1]), + ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1], w->ne[2])); + return ggml_reshape_3d(ctx, result, im2col->ne[1], w->ne[2], im2col->ne[2]); + } + + static ggml_tensor* grouped_conv_1d(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* w, + int64_t groups, + int stride, + int padding, + ggml_tensor* b = nullptr) { + const int64_t ic_g = x->ne[1] / groups; + const int64_t oc_g = w->ne[2] / groups; + // ggml's im2col kernels (the backend of ggml_conv_1d) read input planes + // with flat indexing and only honor nb[1] for the channel offset. A + // non-contiguous x (e.g. a permuted view with time stride 4*C) therefore + // reads wrong elements through the per-group views below, because + // ggml_view_3d forces nb[0] to the type size and hides the true stride. + // Materialize once; contiguous inputs pass through untouched. + if (x->nb[0] != ggml_type_size(x->type) || x->nb[1] != (size_t)x->ne[0] * x->nb[0]) { + x = ggml_cont(ctx->ggml_ctx, x); + } + if (getenv("WAV2VEC2_DEBUG_CONV") != nullptr) { + printf("[dbg] x: ne=[%lld,%lld,%lld] nb=[%zu,%zu,%zu]\n", + (long long)x->ne[0], (long long)x->ne[1], (long long)x->ne[2], + x->nb[0], x->nb[1], x->nb[2]); + printf("[dbg] w: ne=[%lld,%lld,%lld] nb=[%zu,%zu,%zu] type=%d\n", + (long long)w->ne[0], (long long)w->ne[1], (long long)w->ne[2], + w->nb[0], w->nb[1], w->nb[2], (int)w->type); + } + ggml_tensor* acc = nullptr; + for (int64_t i = 0; i < groups; ++i) { + ggml_tensor* x_i = ggml_view_3d(ctx->ggml_ctx, x, + x->ne[0], ic_g, x->ne[2], + x->nb[1], x->nb[2], + i * ic_g * x->nb[1]); + ggml_tensor* w_i = ggml_view_3d(ctx->ggml_ctx, w, + w->ne[0], ic_g, oc_g, + w->nb[1], w->nb[2], + i * oc_g * w->nb[2]); + ggml_tensor* out_i = conv_1d_f32_im2col(ctx->ggml_ctx, w_i, x_i, stride, padding, 1); + if (b != nullptr) { + ggml_tensor* b_i = ggml_view_1d(ctx->ggml_ctx, b, oc_g, i * oc_g * b->nb[0]); + b_i = ggml_reshape_3d(ctx->ggml_ctx, b_i, 1, oc_g, 1); + out_i = ggml_add_inplace(ctx->ggml_ctx, out_i, b_i); + } + acc = (acc == nullptr) ? out_i : ggml_concat(ctx->ggml_ctx, acc, out_i, 1); + } + return acc; + } +}; + +// Conv + (optional) per-channel norm + GELU. ComfyUI LayerNormConv applies the +// LayerNorm over channels for every frame; equivalent to normalizing over the +// channel axis after transposing. GroupNorm variant mirrors LayerGroupNormConv. +struct Wav2Vec2ConvLayer : public UnaryBlock { + Wav2Vec2ConvLayer(int64_t in_channels, + int64_t out_channels, + int kernel_size, + int stride, + bool bias, + bool use_layer_norm) { + blocks["conv"] = std::shared_ptr(new Wav2Vec2Conv1d(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias)); + if (use_layer_norm) { + blocks["layer_norm"] = std::shared_ptr(new LayerNorm(out_channels)); + } + use_layer_norm_ = use_layer_norm; + channels_ = out_channels; + } + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + GGMLBlock::init_params(ctx, tensor_storage_map, prefix); + if (!use_layer_norm_) { + // LayerGroupNormConv: GroupNorm(num_groups=out_channels, affine=true) + params["layer_norm.weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels_); + params["layer_norm.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels_); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + // x: [L, C, N] + x = std::dynamic_pointer_cast(blocks["conv"])->forward(ctx, x); + if (use_layer_norm_) { + auto ln = std::dynamic_pointer_cast(blocks["layer_norm"]); + // [L, C, N] -> [C, L, N] -> LayerNorm over C -> back + x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); + x = ln->forward(ctx, x); + x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); + } else { + x = ggml_cont(ctx->ggml_ctx, x); + x = ggml_group_norm(ctx->ggml_ctx, x, (int)channels_, 1e-05f); + x = ggml_add(ctx->ggml_ctx, x, ggml_reshape_3d(ctx->ggml_ctx, params["layer_norm.bias"], 1, channels_, 1)); + x = ggml_mul(ctx->ggml_ctx, x, ggml_reshape_3d(ctx->ggml_ctx, params["layer_norm.weight"], 1, channels_, 1)); + } + return ggml_ext_gelu(ctx->ggml_ctx, x, true); + } + +private: + bool use_layer_norm_; + int64_t channels_; +}; + +struct Wav2Vec2ConvFeatureEncoder : public GGMLBlock { + Wav2Vec2ConvFeatureEncoder(const Wav2Vec2Config& config) { + // kernel sizes (10,3,3,3,3,2,2), strides (5,2,2,2,2,2,2); conv0 maps the + // single waveform channel to conv_dim and always has a bias. + 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) { + bool bias = (i == 0) ? true : config.conv_bias; + blocks["conv_layers." + std::to_string(i)] = std::shared_ptr( + new Wav2Vec2ConvLayer(in_channels, config.conv_dim, kernels[i], strides[i], bias, config.conv_norm)); + in_channels = config.conv_dim; + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [L, 1, N] waveform; returns [conv_dim, L', N] + 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); // [conv_dim, L', N] + } +}; + +struct Wav2Vec2FeatureProjection : public UnaryBlock { + Wav2Vec2FeatureProjection(const Wav2Vec2Config& config) { + blocks["layer_norm"] = std::shared_ptr(new LayerNorm(config.conv_dim)); + blocks["projection"] = std::shared_ptr(new Linear(config.conv_dim, config.embed_dim)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [conv_dim, L', N] -> [embed_dim, L', N] + 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; + } +}; + +struct Wav2Vec2PositionalConvEmbedding : public UnaryBlock { + int64_t embed_dim; + int64_t groups; + int kernel_size; + enum NormConvention { + NORM_PER_KERNEL_TAP, // g [1, 1, k]: norm over all channels per kernel tap + NORM_PER_OUT_KERNEL, // g [out, 1, k]: norm over in/g per (out, kernel tap) + NORM_PER_CHANNEL_PAIR, // g [out, ic_g, 1]: norm over kernel per channel pair + }; + NormConvention norm_convention; + bool legacy_key_names; // weight_g/weight_v vs parametrizations.weight.original0/1 + + Wav2Vec2PositionalConvEmbedding(const Wav2Vec2Config& config) + : embed_dim(config.embed_dim), groups(16), kernel_size(128), norm_convention(NORM_PER_KERNEL_TAP), legacy_key_names(true) {} + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + // prefix arrives with a trailing dot, e.g. "...encoder.pos_conv_embed."; + // the weight-norm params live under the inner "conv" module. + const std::string base = prefix + "conv."; + const int64_t ic_g = embed_dim / groups; + ggml_type wtype = GGML_TYPE_F16; + auto key_g = base + "weight_g"; + auto key_v = base + "weight_v"; + legacy_key_names = tensor_storage_map.find(key_g) != tensor_storage_map.end(); + if (!legacy_key_names) { + key_g = base + "parametrizations.weight.original0"; + key_v = base + "parametrizations.weight.original1"; + } + // The stored g shape decides the norm convention (files exist with + // parametrizations-style [1, 1, k] g under legacy key names). + auto it_g = tensor_storage_map.find(key_g); + GGML_ASSERT(it_g != tensor_storage_map.end() && it_g->second.ne[1] == 1); + if (it_g->second.ne[2] == 1) { + norm_convention = NORM_PER_KERNEL_TAP; + } else if (it_g->second.ne[1] == 1 && it_g->second.ne[2] == embed_dim) { + norm_convention = NORM_PER_OUT_KERNEL; + } else if (it_g->second.ne[0] == 1 && it_g->second.ne[1] == ic_g && it_g->second.ne[2] == embed_dim) { + norm_convention = NORM_PER_CHANNEL_PAIR; + } else { + GGML_ABORT("wav2vec2: unsupported weight-norm g shape"); + } + // Param keys must include the inner "conv." segment: get_param_tensors + // joins the block prefix with the param key verbatim. + const std::string rel_g = key_g.substr(prefix.size()); // "conv.weight_g" / "conv.parametrizations..." + const std::string rel_v = key_v.substr(prefix.size()); + params[rel_g] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, + it_g->second.ne[0], it_g->second.ne[1], it_g->second.ne[2]); + params[rel_v] = ggml_new_tensor_3d(ctx, get_type(key_v, tensor_storage_map, GGML_TYPE_F16), + kernel_size, ic_g, embed_dim); + if (tensor_storage_map.find(base + "bias") != tensor_storage_map.end()) { + params["conv.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, embed_dim); + } + } + + // Per-tap (or per-convention) norm of v, in-graph. + ggml_tensor* norm(GGMLRunnerContext* ctx) { + const int64_t ic_g = embed_dim / groups; + const char* key_v = legacy_key_names ? "conv.weight_v" : "conv.parametrizations.weight.original1"; + ggml_tensor* v = params[key_v]; + v = ggml_cast(ctx->ggml_ctx, v, GGML_TYPE_F32); + ggml_tensor* sq = ggml_mul(ctx->ggml_ctx, v, v); + ggml_tensor* norm; + if (norm_convention == NORM_PER_KERNEL_TAP) { + // sq [k, ic_g, out]: reduce channels per kernel tap -> [k, 1, 1] + // ggml permute axes are destinations of the source dims (inverse of + // torch's permute), so (2, 0, 1) yields [ic_g, out, k]. + ggml_tensor* pt = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, sq, 2, 0, 1, 3)); // [ic_g, out, k] + ggml_tensor* s = ggml_sum_rows(ctx->ggml_ctx, ggml_reshape_2d(ctx->ggml_ctx, pt, ic_g * embed_dim, kernel_size)); + norm = ggml_sqrt(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, s, kernel_size, 1, 1)); + } else if (norm_convention == NORM_PER_OUT_KERNEL) { + // sq [k, ic_g, out] -> [ic_g, k, out]: reduce ic_g per (k, out) -> [k, 1, out] + ggml_tensor* pt = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, sq, 1, 0, 2, 3)); + norm = ggml_sqrt(ctx->ggml_ctx, + ggml_reshape_3d(ctx->ggml_ctx, + ggml_sum_rows(ctx->ggml_ctx, ggml_reshape_2d(ctx->ggml_ctx, pt, ic_g, kernel_size * embed_dim)), + kernel_size, 1, embed_dim)); + } else { + // sq [k, ic_g, out]: reduce kernel per channel pair -> [1, ic_g, out] + norm = ggml_sqrt(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, sq)); + } + return norm; + } + + // Recombines the weight-normalized conv weight: weight = v * g / norm(v). + ggml_tensor* weight(GGMLRunnerContext* ctx) { + const char* key_g = legacy_key_names ? "conv.weight_g" : "conv.parametrizations.weight.original0"; + const char* key_v = legacy_key_names ? "conv.weight_v" : "conv.parametrizations.weight.original1"; + ggml_tensor* g = params[key_g]; + ggml_tensor* v = params[key_v]; + v = ggml_cast(ctx->ggml_ctx, v, GGML_TYPE_F32); + return ggml_mul(ctx->ggml_ctx, v, ggml_div(ctx->ggml_ctx, g, norm(ctx))); + } + + // Debug: raw v^2 (pre-reduction) and the transposed copy feeding the reduction. + ggml_tensor* sq_dump(GGMLRunnerContext* ctx) { + const char* key_v = legacy_key_names ? "conv.weight_v" : "conv.parametrizations.weight.original1"; + ggml_tensor* v = params[key_v]; + v = ggml_cast(ctx->ggml_ctx, v, GGML_TYPE_F32); + return ggml_mul(ctx->ggml_ctx, v, v); + } + + ggml_tensor* pt_dump(GGMLRunnerContext* ctx) { + return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, sq_dump(ctx), 2, 0, 1, 3)); + } + + // Debug: grouped conv output after drop-last, before GELU (with_bias=false + // also skips the bias add). + ggml_tensor* conv_raw(GGMLRunnerContext* ctx, ggml_tensor* x, bool with_bias) { + ggml_tensor* b = with_bias && params.count("conv.bias") > 0 ? params["conv.bias"] : nullptr; + ggml_tensor* t = conv_input(ctx, x); + t = Wav2Vec2Conv1d::grouped_conv_1d(ctx, t, weight(ctx), groups, 1, kernel_size / 2, b); + t = ggml_view_3d(ctx->ggml_ctx, t, t->ne[0] - 1, t->ne[1], t->ne[2], t->nb[1], t->nb[2], 0); + return t; // [L, embed_dim, N] + } + + // The materialized [L, embed_dim, N] plane-contiguous tensor that + // grouped_conv_1d actually consumes (cont of the permuted projection + // output; im2col requires this layout). + ggml_tensor* conv_input(GGMLRunnerContext* ctx, ggml_tensor* x) { + return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + } + + // Debug: group-0-only conv (no bias), post drop-last. + ggml_tensor* conv_group0(GGMLRunnerContext* ctx, ggml_tensor* x) { + const int64_t ic_g = embed_dim / groups; + ggml_tensor* t = conv_input(ctx, x); + ggml_tensor* w = weight(ctx); + ggml_tensor* x_0 = ggml_view_3d(ctx->ggml_ctx, t, t->ne[0], ic_g, t->ne[2], t->nb[1], t->nb[2], 0); + ggml_tensor* w_0 = ggml_view_3d(ctx->ggml_ctx, w, w->ne[0], ic_g, w->ne[2] / groups, w->nb[1], w->nb[2], 0); + ggml_tensor* out = Wav2Vec2Conv1d::conv_1d_f32_im2col(ctx->ggml_ctx, w_0, x_0, 1, kernel_size / 2, 1); + out = ggml_view_3d(ctx->ggml_ctx, out, out->ne[0] - 1, out->ne[1], out->ne[2], out->nb[1], out->nb[2], 0); + return out; // [L, ic_g, N] + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [embed_dim, L, N]; returns [embed_dim, L, N] (last sample dropped) + ggml_tensor* w = weight(ctx); + ggml_tensor* b = params.count("conv.bias") > 0 ? params["conv.bias"] : nullptr; + ggml_tensor* t = conv_input(ctx, x); + t = Wav2Vec2Conv1d::grouped_conv_1d(ctx, t, w, groups, 1, kernel_size / 2, b); + // GELU must be out-of-place and precede the drop-last view: an in-place + // op on a view writes through a buffer the graph allocator does not + // reserve for it, so reuse of the conv output's storage silently + // corrupts the result depending on the surrounding graph shape. + t = ggml_ext_gelu(ctx->ggml_ctx, t, false); + t = ggml_view_3d(ctx->ggml_ctx, t, t->ne[0] - 1, t->ne[1], t->ne[2], t->nb[1], t->nb[2], 0); + return ggml_permute(ctx->ggml_ctx, t, 1, 0, 2, 3); // [embed_dim, L, N] + } +}; + +struct Wav2Vec2FeedForward : public UnaryBlock { + Wav2Vec2FeedForward(const Wav2Vec2Config& config) { + blocks["intermediate_dense"] = std::shared_ptr(new Linear(config.embed_dim, config.embed_dim * 4)); + blocks["output_dense"] = std::shared_ptr(new Linear(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 Wav2Vec2TransformerEncoderLayer : public UnaryBlock { + bool do_stable_layer_norm; + + Wav2Vec2TransformerEncoderLayer(const Wav2Vec2Config& config) + : do_stable_layer_norm(config.do_stable_layer_norm) { + blocks["attention"] = std::shared_ptr(new MultiheadAttention(config.embed_dim, config.num_heads, true, true)); + blocks["layer_norm"] = std::shared_ptr(new LayerNorm(config.embed_dim)); + blocks["feed_forward"] = std::shared_ptr(new Wav2Vec2FeedForward(config)); + blocks["final_layer_norm"] = std::shared_ptr(new LayerNorm(config.embed_dim)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [embed_dim, L, N] + 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 Wav2Vec2TransformerEncoder : public GGMLBlock { + int num_layers; + bool do_stable_layer_norm; + + Wav2Vec2TransformerEncoder(const Wav2Vec2Config& config) + : num_layers(config.num_layers), do_stable_layer_norm(config.do_stable_layer_norm) { + blocks["pos_conv_embed"] = std::shared_ptr(new Wav2Vec2PositionalConvEmbedding(config)); + for (int i = 0; i < config.num_layers; ++i) { + blocks["layers." + std::to_string(i)] = std::shared_ptr(new Wav2Vec2TransformerEncoderLayer(config)); + } + blocks["layer_norm"] = std::shared_ptr(new LayerNorm(config.embed_dim)); + } + + ggml_tensor* pos_conv(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); + return pc->forward(ctx, x); + } + + ggml_tensor* pos_conv_weight(GGMLRunnerContext* ctx) { + auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); + return pc->weight(ctx); + } + + ggml_tensor* pos_conv_norm(GGMLRunnerContext* ctx) { + auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); + return pc->norm(ctx); + } + + ggml_tensor* pos_conv_sq(GGMLRunnerContext* ctx) { + auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); + return pc->sq_dump(ctx); + } + + ggml_tensor* pos_conv_pt(GGMLRunnerContext* ctx) { + auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); + return pc->pt_dump(ctx); + } + + ggml_tensor* pos_conv_raw(GGMLRunnerContext* ctx, ggml_tensor* x, bool with_bias) { + auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); + return pc->conv_raw(ctx, x, with_bias); + } + + ggml_tensor* pos_conv_input(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); + return pc->conv_input(ctx, x); + } + + ggml_tensor* pos_conv_group0(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); + return pc->conv_group0(ctx, x); + } + + // Returns the final hidden state [embed_dim, L, N]. When all_layers != nullptr, + // it is filled with the 24 pre-layer states plus the final state, concatenated + // along a new trailing axis: [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::shared_ptr(new Wav2Vec2ConvFeatureEncoder(config)); + blocks["feature_projection"] = std::shared_ptr(new Wav2Vec2FeatureProjection(config)); + blocks["encoder"] = std::shared_ptr(new Wav2Vec2TransformerEncoder(config)); + } + + // waveform: [L, 1, N] (already channel-mixed). Returns the final hidden state; + // all_layers receives the per-layer stack when non-null. + 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; + } + + // Debug front-end stages (see Wav2Vec2ModelRunner::build_graph): + // 1 = feature extractor, 2 = + projection, 3 = + pos_conv and residual add, + // 4 = the recomputed pos_conv weight tensor, 5 = the pos_conv weight norm. + ggml_tensor* forward_front(GGMLRunnerContext* ctx, ggml_tensor* x, int stage) { + 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"]); + + if (stage == 11) { + x = feature_extractor->forward(ctx, x); + x = feature_projection->forward(ctx, x); + return encoder->pos_conv_group0(ctx, x); + } + if (stage == 10) { + x = feature_extractor->forward(ctx, x); + x = feature_projection->forward(ctx, x); + return encoder->pos_conv_input(ctx, x); + } + if (stage == 9 || stage == 8) { + x = feature_extractor->forward(ctx, x); + x = feature_projection->forward(ctx, x); + return encoder->pos_conv_raw(ctx, x, stage == 8); + } + if (stage == 7) { + return encoder->pos_conv_pt(ctx); + } + if (stage == 6) { + return encoder->pos_conv_sq(ctx); + } + if (stage == 5) { + return encoder->pos_conv_norm(ctx); + } + if (stage == 4) { + return encoder->pos_conv_weight(ctx); + } + x = feature_extractor->forward(ctx, x); + if (stage == 1) { + return x; + } + x = feature_projection->forward(ctx, x); + if (stage == 2) { + return x; + } + return ggml_add(ctx->ggml_ctx, x, encoder->pos_conv(ctx, x)); + } +}; + +class Wav2Vec2ModelRunner : public GGMLRunner { +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), weight_prefix(prefix) { + config = Wav2Vec2Config::detect_from_weights(tensor_storage_map, prefix); + model = Wav2Vec2Model(config); + // GGMLBlock::init/get_param_tensors append their own '.'; the loader-facing + // prefix convention carries 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 { + char buf[256]; + snprintf(buf, sizeof(buf), "wav2vec2: embed_dim %" PRId64 ", %d layers, %d heads%s", + config.embed_dim, config.num_layers, config.num_heads, + config.do_stable_layer_norm ? ", stable-ln" : ""); + return std::string(buf); + } + + 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); + } + + // Allocates all param tensors on the backend's default buffer. Standalone + // harnesses need this before loading weights; the full pipeline routes + // allocation through ModelManager instead. + bool alloc_params_on_backend(ggml_backend_t backend) { + params_buffer = ggml_backend_alloc_ctx_tensors_from_buft(params_ctx, ggml_backend_get_default_buffer_type(backend)); + return params_buffer != nullptr; + } + + // waveform: [L, 1, 1] mono; already host-normalized. stage 0 (default) returns + // [embed_dim, L', num_layers + 1] with the per-layer stack (last slice = final + // hidden state). Stages 1..3 expose front-end intermediates for parity debugging: + // 1 = feature extractor [conv_dim, L', 1], 2 = + projection [embed_dim, L', 1], + // 3 = + pos_conv and residual add [embed_dim, L', 1]. + ggml_cgraph* build_graph(const sd::Tensor& waveform_tensor, int stage = 0) { + ggml_cgraph* gf = ggml_new_graph(compute_ctx); + ggml_tensor* waveform = make_input(waveform_tensor); + + auto runner_ctx = get_context(); + + if (stage == 0) { + ggml_tensor* all_layers = nullptr; + model.forward(&runner_ctx, waveform, &all_layers); + GGML_ASSERT(all_layers != nullptr); + ggml_build_forward_expand(gf, all_layers); + } else { + // stages produce views (permutes); materialize them so the flat dump + // reflects logical [C, T, N] order + ggml_tensor* out = ggml_cont(runner_ctx.ggml_ctx, model.forward_front(&runner_ctx, waveform, stage)); + ggml_build_forward_expand(gf, out); + } + return gf; + } + + sd::Tensor compute(const int n_threads, const std::vector& mono_waveform, int stage = 0) { + 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, stage); + }; + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); + } + +private: + Wav2Vec2Config config; + ggml_backend_buffer_t params_buffer = nullptr; + + // torch: (x - x.mean()) / torch.sqrt(x.var() + 1e-7); var is population variance. + 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); + } + } +}; + +#endif // __SD_MODEL_AUDIO_WAV2VEC2_HPP__ diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index 4f2580394..d5299755f 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -809,7 +809,8 @@ namespace Rope { return embed_nd(ids, bs, static_cast(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage); } - // Generate wan positional embeddings + // Generate wan positional embeddings; t_offset shifts the temporal ids + // (S2V reference latent uses t_start = max(30, t + 9)). __STATIC_INLINE__ std::vector gen_wan_pe(int t, int h, int w, @@ -818,8 +819,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 0e447cd06..cecf4cc3b 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: stacked wav2vec2 hidden states, [num_layers=25, frames=T_latent*4, dim=1024] + const sd::Tensor* audio_embed = nullptr; }; struct HiDreamO1DiffusionExtra { diff --git a/src/model/diffusion/wan.hpp b/src/model/diffusion/wan.hpp index 687f28380..53793270e 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 @@ -8,6 +9,7 @@ #include "core/ggml_extend_backend.h" #include "core/ggml_tensor_utils.h" +#include "model/audio/wav2vec2.hpp" #include "model/common/block.hpp" #include "model/common/rope.hpp" #include "model/diffusion/flux.hpp" @@ -33,11 +35,17 @@ 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; + // Wan2.2-S2V audio conditioning + int64_t audio_dim = 1024; // wav2vec2-large hidden size + int num_audio_token = 4; // motion tokens per frame; 1 learned padding token appended at use site + 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 +82,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; } @@ -528,10 +540,246 @@ namespace WAN { } }; + //================================================ Wan2.2-S2V audio conditioning ================================================= + // Ported from ComfyUI comfy/ldm/wan/model.py: CausalConv1d, MotionEncoder_tc, + // CausalAudioEncoder and AudioInjector_WAN. Only the adain_mode="attn_norm" + // configuration exists in the Wan2.2-S2V-14B checkpoint (injector_pre_norm_* + // have no keys), so the pre-norm fallback path is intentionally absent. + + class WanCausalConv1d : public GGMLBlock { + protected: + 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::shared_ptr(new Wav2Vec2Conv1d(in_dim, out_dim, kernel_size, stride)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [T, C] -> [T', C_out]; replicate left-pad (k-1, 0) along T + 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 { + protected: + 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 { + // Elementwise-consumed param (ggml_repeat/concat against F32 + // activations): hardcode F32; the loader converts the storage type. + params["padding_tokens"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_dim); + } + + // conv -> affine-free LayerNorm over channels -> SiLU. The norm runs on + // the permuted [C, T] form; to_conv_layout materializes the [T', C] + // form the next conv needs (im2col requires contiguous time rows), + // otherwise the token-layout [C, T'] is returned directly. + 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); // [T', C_out] + x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); // [C_out, T'] + 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)); // [T', C_out] + } + 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::shared_ptr(new WanCausalConv1d(in_dim, hidden_dim / 4 * num_token)); + if (need_global) { + blocks["conv1_global"] = std::shared_ptr(new WanCausalConv1d(in_dim, hidden_dim / 4)); + } + blocks["norm1"] = std::shared_ptr(new LayerNorm(hidden_dim / 4, 1e-6f, false)); + blocks["conv2"] = std::shared_ptr(new WanCausalConv1d(hidden_dim / 4, hidden_dim / 2, 3, 2)); + blocks["norm2"] = std::shared_ptr(new LayerNorm(hidden_dim / 2, 1e-6f, false)); + blocks["conv3"] = std::shared_ptr(new WanCausalConv1d(hidden_dim / 2, hidden_dim, 3, 2)); + blocks["norm3"] = std::shared_ptr(new LayerNorm(hidden_dim, 1e-6f, false)); + if (need_global) { + blocks["final_linear"] = std::shared_ptr(new Linear(hidden_dim, hidden_dim)); + } + } + + // x: [T, in_dim]. Returns (local [dim, num_token+1, T/4], global [dim, T/4]); + // torch groups the conv1_local channels into num_token batches and the + // group index becomes the per-frame token index. + std::pair forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto local = std::dynamic_pointer_cast(blocks["conv1_local"])->forward(ctx, x); // [T, hidden] + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + std::vector tokens; + for (auto& group : ggml_ext_chunk(ctx->ggml_ctx, local, num_token, 1)) { + // rearrange 'b (n c) t -> (b n) t c': per-group norm over c + ggml_tensor* s = ggml_permute(ctx->ggml_ctx, group, 1, 0, 2, 3); // [C, T] view + 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)); // [T, C] + s = conv_norm_silu(ctx, s, "conv2", "norm2", true); // [T2, hidden/2] + s = conv_norm_silu(ctx, s, "conv3", "norm3", false); // [hidden, T3] + tokens.push_back(ggml_reshape_3d(ctx->ggml_ctx, s, s->ne[0], 1, s->ne[1])); // [dim, 1, T3] + } + auto padding = ggml_reshape_3d(ctx->ggml_ctx, params["padding_tokens"], hidden_dim, 1, 1); + padding = ggml_repeat(ctx->ggml_ctx, padding, tokens[0]); // [dim, 1, T3] + tokens.push_back(padding); + ggml_tensor* local_out = ggml_ext_vec_concat(ctx->ggml_ctx, tokens, 1); // [dim, num_token+1, T3] + + if (!need_global) { + return {local_out, nullptr}; + } + // global branch shares conv2/conv3/norms with the local branch + ggml_tensor* g = conv_norm_silu(ctx, x, "conv1_global", "norm1", true); // [T, hidden/4] + g = conv_norm_silu(ctx, g, "conv2", "norm2", true); // [T2, hidden/2] + g = conv_norm_silu(ctx, g, "conv3", "norm3", false); // [hidden, T3] + g = std::dynamic_pointer_cast(blocks["final_linear"])->forward(ctx, g); + return {local_out, g}; + } + }; + + class WanCausalAudioEncoder : public GGMLBlock { + protected: + int64_t audio_dim; + int64_t dim; + int num_token; + int num_layers; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + // Elementwise-consumed param (silu/mul/div in forward): hardcode F32; + // the loader converts the storage type. Checkpoint shape is [1, 25, + // 1, 1] (torch order, leading 1s), so the reversed storage ne is + // [1, 1, 25, 1]; mirror it for raw-bytes loading. forward reshapes + // to [1, 1, L] either way. + 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) + : audio_dim(audio_dim), dim(dim), num_token(num_token), num_layers(num_layers) { + blocks["encoder"] = std::shared_ptr(new WanMotionEncoder(audio_dim, dim, num_token, true)); + } + + // features: [audio_dim, T*4, num_layers] stacked wav2vec2 hidden states. + // SiLU-weighted sum over the layer axis, then MotionEncoder. + // Returns (local [dim, num_token+1, T], global [dim, T]). + std::pair forward(GGMLRunnerContext* ctx, ggml_tensor* features) { + auto weights = ggml_silu(ctx->ggml_ctx, params["weights"]); // [L] + 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)); + // reduce over the layer axis: layers -> ne0 for sum_rows + // (torch-style 3-cycle; plain ggml_permute axes are destinations, + // the inverse mapping - see the pos_conv weight-norm bug) + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); // [L, dim, F] + x = ggml_sum_rows(ctx->ggml_ctx, x); // [1, dim, F] + x = ggml_reshape_2d(ctx->ggml_ctx, x, x->ne[1], x->ne[2]); // [dim, F] + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [F, dim] + return std::dynamic_pointer_cast(blocks["encoder"])->forward(ctx, x); + } + }; + + class WanAudioInjector : public GGMLBlock { + protected: + 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::shared_ptr(new WanT2VCrossAttention(dim, num_heads, qk_norm, eps)); + blocks["injector_adain_layers." + std::to_string(i) + ".linear"] = + std::shared_ptr(new Linear(dim, dim * 2)); + } + // AdaLayerNorm norm is affine-free (no checkpoint keys); eps 1e-5 per + // upstream AdaLayerNorm, distinct from the model-wide 1e-6. + blocks["adain_norm"] = std::shared_ptr(new LayerNorm(dim, 1e-5f, false)); + } + + // x: [dim, n_token] full sequence; audio_local: [dim, num_token+1, T]; + // audio_global: [dim, T]. The residual is applied to the first seq_len + // (video) tokens only; trailing reference tokens pass through unchanged. + // Batching over T at ne[2] mirrors torch's rearrange to (b t) n c. + 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); // [dim, T] + temb = adain_linear->forward(ctx, temb); // [2*dim, T] + auto shift = ggml_ext_slice(ctx->ggml_ctx, temb, 0, 0, dim); // [dim, T] + auto scale = ggml_ext_slice(ctx->ggml_ctx, temb, 0, dim, dim * 2); // [dim, T] + 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); // [dim, seq_len] view + auto h = ggml_reshape_3d(ctx->ggml_ctx, x_vid, dim, n_tok, T); // [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 * (1 + scale) + h = ggml_add(ctx->ggml_ctx, h, shift); + + auto res = injector->forward(ctx, h, audio_local, 0); // [dim, n_tok, T] + res = ggml_reshape_2d(ctx->ggml_ctx, res, dim, seq_len); + + auto x_head = ggml_add(ctx->ggml_ctx, x_vid, res); // out-of-place add on the view + 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; + } + }; + class Wan : public GGMLBlock { 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 +802,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 +843,15 @@ namespace WAN { blocks["vace_patch_embedding"] = std::shared_ptr(new Conv3d(config.vace_in_dim, config.dim, config.patch_size, config.patch_size)); } + + // s2v audio conditioning (checkpoint-contained modules) + if (config.model_type == "s2v") { + blocks["casual_audio_encoder"] = std::shared_ptr(new WanCausalAudioEncoder(config.audio_dim, config.dim, config.num_audio_token)); + blocks["audio_injector"] = std::shared_ptr(new WanAudioInjector(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 +899,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: [audio_dim, T*4, 25] stacked wav2vec2 hidden states (S2V) + // reference_latent: [N*C, T_ref, H, W] (S2V) + // 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 +933,41 @@ 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] + // s2v: audio conditioning + reference latent append + 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); // one pixel-frame per latent frame * 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; // [dim, num_audio_token+1, T] + audio_global = audio_emb.second; // [dim, T] + 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"]; // [dim, 3] + 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); // [N*dim, t_ref_len, h_len, w_len] + 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 frames run at timestep 0 (upstream t concat with zeros) + 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 +1012,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 +1034,15 @@ namespace WAN { c_skip = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength); x = ggml_add(ctx->ggml_ctx, x, c_skip); } + + // s2v: AdaIN from the global audio token, then cross-attention + // against the per-frame audio tokens; residual on video tokens only + 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 +1059,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 +1093,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) { + // drop the trailing reference tokens (upstream unpatchify slices to prod(grid_sizes)) + 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 +1158,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 +1213,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 +1225,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); - - 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); + 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); + if (ref_latent != nullptr) { + // s2v reference latent: temporal offset t_start = max(30, T + 9) + 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 +1274,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 +1292,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 +1307,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 +1321,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_component.h b/src/model_component.h index 96077a4a4..b3e2a55ba 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 807c915aa..fab8db28d 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -247,6 +247,14 @@ bool read_safetensors_file(const std::string& file_path, std::string dtype = tensor_info["dtype"]; nlohmann::json shape = tensor_info["shape"]; + // ComfyUI fp8_scaled checkpoints carry a per-module scale_input + // (activation scale, which cancels when activations are f16/f32). + // Drop it here; .scale_weight -> .weight_scale aliasing happens in + // convert_tensor_name. + 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 8325a6063..42795306a 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -424,6 +424,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; @@ -510,6 +511,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 shares patch_embedding shape with T2V-14B; key presence is the only reliable signal + is_s2v = true; + } if (tensor_storage.name.find("model.diffusion_model.patch_embedder.weight") != std::string::npos) { return VERSION_LINGBOT_VIDEO; } @@ -573,6 +579,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 7dba74b8f..82cd25d6c 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -28,11 +28,13 @@ #include "model_manager.h" #include "stable-diffusion.h" +#include "conditioning/audio_processing.hpp" #include "conditioning/conditioner.hpp" #include "core/backend_fit.h" #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" @@ -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", @@ -132,7 +135,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}}, @@ -800,6 +803,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, @@ -1000,6 +1010,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)); @@ -1009,11 +1020,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() { @@ -1111,6 +1126,13 @@ 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) { + // HF wav2vec2 inference leftovers: lm_head is a pretraining head, + // masked_spec_embed only applies to masked pretraining. + 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"); @@ -1737,6 +1759,30 @@ 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 = AudioProcessing::downmix_to_mono(audio.data, audio.sample_count, audio.channels); + if (mono.empty()) { + LOG_ERROR("audio mono downmix failed"); + return {}; + } + mono = AudioProcessing::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 = {}; @@ -1792,6 +1838,12 @@ 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; } @@ -2300,7 +2352,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; // wan2.2 s2v driving audio std::shared_ptr control_net; std::shared_ptr ip_adapter; sd::Tensor ip_adapter_tokens; @@ -363,6 +365,8 @@ class StableDiffusionGGML { 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 1210a8cba..17015af44 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; diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index 8d174e8d7..932d31d02 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" @@ -233,6 +234,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 73d134b3b..a1389e7eb 100644 --- a/src/pipeline/model_builders.h +++ b/src/pipeline/model_builders.h @@ -15,6 +15,7 @@ struct DiffusionModelRunner; struct VAE; struct AudioVAERunner; struct ControlNet; +class Wav2Vec2ModelRunner; struct GenerationExtension; struct GenerationExtensionInitContext; namespace IPAdapter { @@ -37,6 +38,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; // wan2.2 s2v driving audio }; struct VAEOptions { diff --git a/src/pipeline/video.cpp b/src/pipeline/video.cpp index cc979c0fa..7d326f179 100644 --- a/src/pipeline/video.cpp +++ b/src/pipeline/video.cpp @@ -5,6 +5,7 @@ #include #include +#include "conditioning/audio_processing.hpp" #include "core/rng.hpp" #include "core/rng_philox.hpp" #include "diffusion_engine.h" @@ -420,6 +421,50 @@ namespace sd::pipeline { return audio; } + // Wan2.2 S2V: build the first per-chunk audio window. Input: wav2vec2 stacked + // states [embed_dim, in_frames, num_layers]; output: [embed_dim, batch_frames, + // num_layers] as expected by WanDiffusionExtra.audio_embed, frames bucketed at + // 16 fps with zero padding past the audio end (nodes_wan.py + // get_audio_embed_bucket_fps, m=0, frame_offset=0). + 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 {}; + } + // [embed_dim, in_frames, num_layers] -> layer-first [num_layers, in_frames, embed_dim] + 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); + } + } + AudioProcessing::BucketPlan plan; + std::vector buckets = AudioProcessing::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 {}; + } + // Window rows [0, batch_frames): [frame, num_layers, dim] -> [dim, batch_frames, num_layers] + 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 +1078,56 @@ 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()) { + // ComfyUI WanSoundImageToVideo: the ref image is VAE-encoded and + // appended as reference_latents; the video latent itself stays + // unconstrained (no first-frame conditioning). + 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); // [W', H', 1, C, 1] + if (encoded_ref.empty()) { + LOG_ERROR("failed to encode S2V reference image"); + return std::nullopt; + } + // forward_orig consumes a 4d reference latent [N*C, t_ref, H, W] + 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 +1179,15 @@ namespace sd::pipeline { latents.keyframe_indices); } } + if (sd->version == VERSION_WAN2_2_S2V) { + // ComfyUI: positive gets the real audio window and the ref latent; + // negative gets audio * 0 while KEEPING the same ref latent (wan is + // excluded from ref-latent img cfg). + 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 +1200,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(); @@ -1725,6 +1835,35 @@ 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) { + // S2V does not generate audio; the driving track is conditioning only. + // Hand a copy back so the output container carries the same audio + // (the CLI muxes audio_out into avi/webm). + 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 +1913,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 +1926,21 @@ namespace sd::pipeline { if (frames_out != nullptr) { *frames_out = result; } + if (sd->version == VERSION_WAN2_2_S2V && generated_audio != nullptr) { + // The model conditioned on the first chunk window only; keep the muxed + // track aligned with the decoded video duration. + int fps = sd_vid_gen_params->fps > 0 ? sd_vid_gen_params->fps : 16; + 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/stable-diffusion.cpp b/src/stable-diffusion.cpp index cda44e92b..6bc72431b 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -352,6 +352,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" @@ -387,6 +388,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), From 6b945a759cb02ae59f5e15e26b2257e064f659dc Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 13 Sep 2026 21:28:56 +0800 Subject: [PATCH 2/5] correct S2V audio bucket layout and output frame rate --- docs/wan.md | 4 ++++ examples/cli/main.cpp | 8 ++++---- examples/server/async_jobs.cpp | 5 ++--- include/stable-diffusion.h | 5 ++++- src/conditioning/audio_processing.hpp | 8 +++++--- src/pipeline/generation.h | 3 ++- src/pipeline/request.cpp | 5 ++++- src/pipeline/video.cpp | 11 +++++++++-- src/stable-diffusion.cpp | 11 +++++++++-- 9 files changed, 43 insertions(+), 17 deletions(-) diff --git a/docs/wan.md b/docs/wan.md index ebe4ebe0a..a61503ce4 100644 --- a/docs/wan.md +++ b/docs/wan.md @@ -129,6 +129,10 @@ Notes: 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. - The output video carries the driving audio track: it is muxed into `.avi` / `.webm` outputs, truncated to the video duration. Other container types fall back to a `.wav` sidecar file next to the video. diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 22b26da79..a3cc1dde8 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/server/async_jobs.cpp b/examples/server/async_jobs.cpp index 1933dec17..8c4adf277 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 62ce26a41..12209d331 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -516,11 +516,14 @@ 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); +// fps_out may be NULL; otherwise it receives the effective frame rate before preview callbacks. +// Use this frame rate when encoding the returned frames. 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/audio_processing.hpp b/src/conditioning/audio_processing.hpp index 88eb7a731..4f10418a7 100644 --- a/src/conditioning/audio_processing.hpp +++ b/src/conditioning/audio_processing.hpp @@ -216,9 +216,11 @@ namespace AudioProcessing { if (src >= plan.audio_frames) { continue; // zero padding past the audio end } - std::copy_n(interpolated.data() + static_cast(src) * num_layers * dim, - static_cast(num_layers) * dim, - buckets.data() + static_cast(frame) * num_layers * dim); + 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; } diff --git a/src/pipeline/generation.h b/src/pipeline/generation.h index 17015af44..34aa000aa 100644 --- a/src/pipeline/generation.h +++ b/src/pipeline/generation.h @@ -60,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/request.cpp b/src/pipeline/request.cpp index aa78f7e6b..2874bbae9 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 7d326f179..31e88eff4 100644 --- a/src/pipeline/video.cpp +++ b/src/pipeline/video.cpp @@ -1532,10 +1532,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); } @@ -1547,6 +1551,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) { @@ -1929,7 +1936,7 @@ namespace sd::pipeline { if (sd->version == VERSION_WAN2_2_S2V && generated_audio != nullptr) { // The model conditioned on the first chunk window only; keep the muxed // track aligned with the decoded video duration. - int fps = sd_vid_gen_params->fps > 0 ? sd_vid_gen_params->fps : 16; + 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)", diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 6bc72431b..d19d444f7 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -724,8 +724,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; } @@ -741,10 +745,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) { From e3fbf5c4f4826a44889552fcb60e4704ef3885b2 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 13 Sep 2026 21:50:19 +0800 Subject: [PATCH 3/5] separate audio preprocessing from Wan S2V conditioning --- examples/common/media_io.cpp | 34 ---- examples/common/media_io.h | 18 -- src/conditioning/audio_processing.hpp | 230 -------------------------- src/conditioning/wan_audio.cpp | 113 +++++++++++++ src/conditioning/wan_audio.h | 35 ++++ src/pipeline/diffusion_engine.cpp | 6 +- src/pipeline/video.cpp | 16 +- src/runtime/audio_processing.cpp | 99 +++++++++++ src/runtime/audio_processing.h | 23 +++ 9 files changed, 281 insertions(+), 293 deletions(-) delete mode 100644 src/conditioning/audio_processing.hpp create mode 100644 src/conditioning/wan_audio.cpp create mode 100644 src/conditioning/wan_audio.h create mode 100644 src/runtime/audio_processing.cpp create mode 100644 src/runtime/audio_processing.h diff --git a/examples/common/media_io.cpp b/examples/common/media_io.cpp index 58b49a14c..721aa6641 100644 --- a/examples/common/media_io.cpp +++ b/examples/common/media_io.cpp @@ -1,5 +1,4 @@ #include "media_io.h" -#include "conditioning/audio_processing.hpp" #include "log.h" #include "resource_owners.hpp" @@ -11,7 +10,6 @@ #include #include #include -#include #include #include @@ -1508,35 +1506,3 @@ bool load_wav_from_file(const std::string& path, } return true; } - -// DSP core lives in src/conditioning/audio_processing.hpp (shared with -// stable-diffusion.cpp); these wrappers keep the examples-side API. -std::vector downmix_to_mono(const float* interleaved_samples, - uint64_t sample_count, - uint32_t channels) { - return AudioProcessing::downmix_to_mono(interleaved_samples, sample_count, channels); -} - -std::vector resample_audio(const float* samples, - uint64_t sample_count, - uint32_t orig_sample_rate, - uint32_t target_sample_rate) { - return AudioProcessing::resample_audio(samples, sample_count, orig_sample_rate, target_sample_rate); -} - -bool load_wav_from_file_mono(const std::string& path, - std::vector& mono_samples, - uint32_t target_sample_rate) { - std::vector interleaved; - uint32_t sample_rate = 0; - uint32_t channels = 0; - if (!load_wav_from_file(path, interleaved, sample_rate, channels)) { - return false; - } - const uint64_t frames = channels > 0 ? interleaved.size() / channels : 0; - mono_samples = downmix_to_mono(interleaved.data(), frames, channels); - if (sample_rate != target_sample_rate && !mono_samples.empty()) { - mono_samples = resample_audio(mono_samples.data(), mono_samples.size(), sample_rate, target_sample_rate); - } - return true; -} diff --git a/examples/common/media_io.h b/examples/common/media_io.h index 0c2230674..52b9221a2 100644 --- a/examples/common/media_io.h +++ b/examples/common/media_io.h @@ -115,22 +115,4 @@ bool load_wav_from_file(const std::string& path, uint32_t& sample_rate, uint32_t& channels); -// Downmix interleaved samples to mono by averaging channels. Returns empty vector on invalid input. -std::vector downmix_to_mono(const float* interleaved_samples, - uint64_t sample_count, - uint32_t channels); - -// Band-limited sinc resampler matching torchaudio.functional.resample -// (sinc_interp_hann, lowpass_filter_width 6, rolloff 0.99). 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); - -// Convenience: load WAV, downmix to mono and resample to the target sample rate. -bool load_wav_from_file_mono(const std::string& path, - std::vector& mono_samples, - uint32_t target_sample_rate = 16000); - #endif // __MEDIA_IO_H__ diff --git a/src/conditioning/audio_processing.hpp b/src/conditioning/audio_processing.hpp deleted file mode 100644 index 4f10418a7..000000000 --- a/src/conditioning/audio_processing.hpp +++ /dev/null @@ -1,230 +0,0 @@ -#ifndef __SD_CONDITIONING_AUDIO_PROCESSING_HPP__ -#define __SD_CONDITIONING_AUDIO_PROCESSING_HPP__ - -// Wan2.2-S2V audio windowing, ported from ComfyUI comfy_extras/nodes_wan.py -// (linear_interpolation + get_audio_embed_bucket_fps, m=0). -// -// Input is the wav2vec2 hidden states stacked per layer [num_layers, in_frames, dim] -// at the encoder frame rate (50 Hz). The frames are interpolated to video_rate -// (30 Hz), bucketed to fps (16) frames with zero padding past the audio end, and -// split into chunks of batch_frames = latent_t * 4 frames (one per diffusion chunk). - -#include -#include -#include -#include -#include -#include - -namespace AudioProcessing { - - // Chunk/padding math of get_audio_embed_bucket_fps (m=0). - struct BucketPlan { - int audio_frames; // input frames at video_rate (30 Hz) - int batch_frames; // pixel frames per chunk (latent_t * 4) - int video_rate; // timeline rate of audio_frames (30 Hz) - int fps; // bucket frame rate (16 fps) - int num_chunks; // ComfyUI num_repeat - int bucket_frames; // total bucket frames = num_chunks * batch_frames - int padded_audio_frames; // audio_frames plus zero padding applied - }; - - inline BucketPlan plan_buckets(int audio_frames, int batch_frames, int video_rate = 30, int fps = 16) { - 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; - // min_batch_num = int(audio_frame_num / (batch_frames * scale)) + 1 - plan.num_chunks = static_cast(audio_frames / (batch_frames * scale)) + 1; - plan.bucket_frames = plan.num_chunks * batch_frames; - // padd_audio_num = ceil(bucket_frames / fps * video_rate) - audio_frame_num - plan.padded_audio_frames = static_cast( - std::ceil(plan.bucket_frames / static_cast(fps) * video_rate)); - return plan; - } - - // Bucket frame index (fps timeline) -> source frame index (video_rate timeline). - // get_sample_indices with fixed_start=0 reduces to round-half-even(i * video_rate / fps), - // matching numpy's default rounding. - inline int bucket_source_frame(int bucket_frame, int video_rate = 30, int fps = 16) { - return static_cast(std::nearbyint(static_cast(bucket_frame) * video_rate / fps)); - } - - // torch.nn.functional.interpolate size computation: output_len = int(in_len / input_fps * output_fps) - inline int interpolated_frame_count(int in_frames, int input_fps = 50, int output_fps = 30) { - return static_cast(in_frames / static_cast(input_fps) * output_fps); - } - - // torch.nn.functional.interpolate(mode='linear', align_corners=True) along the frame - // dimension. in: [num_layers, in_frames, dim], out: [num_layers, out_frames, dim]. - inline 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; - } - - // Polyphase FIR from torchaudio.functional.resample: sinc interpolated at output - // phases, hann-windowed over lowpass_filter_width zero crossings, anti-aliased by - // clamping the sinc argument to +-lowpass_filter_width after scaling by min(rate)*rolloff. - // Returns the input unchanged when sample rates are equal, and an empty vector on - // invalid input. - inline 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; - } - - // Downmix interleaved samples to mono by averaging channels. Returns an empty - // vector on invalid input. - inline 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; - } - - // Full bucketing: stacked encoder states [num_layers, in_frames, dim] at input_fps -> - // bucket frames [bucket_frames, num_layers, dim] at fps, with zero frames past the - // audio end. Chunk c occupies rows [c * batch_frames, (c + 1) * batch_frames). - // Returns an empty vector on invalid input; the applied plan is stored in *plan_out. - inline 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) { - 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; // zero padding past the audio end - } - 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 AudioProcessing - -#endif // __SD_CONDITIONING_AUDIO_PROCESSING_HPP__ diff --git a/src/conditioning/wan_audio.cpp b/src/conditioning/wan_audio.cpp new file mode 100644 index 000000000..18b705df1 --- /dev/null +++ b/src/conditioning/wan_audio.cpp @@ -0,0 +1,113 @@ +#include "wan_audio.h" + +#include +#include +#include + +// Input is the wav2vec2 hidden states stacked per layer [num_layers, in_frames, dim] +// at the encoder frame rate (50 Hz). The frames are interpolated to video_rate +// (30 Hz), bucketed to fps (16) frames with zero padding past the audio end, and +// split into chunks of batch_frames = latent_t * 4 frames (one per diffusion chunk). + +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; + // min_batch_num = int(audio_frame_num / (batch_frames * scale)) + 1 + plan.num_chunks = static_cast(audio_frames / (batch_frames * scale)) + 1; + plan.bucket_frames = plan.num_chunks * batch_frames; + // padd_audio_num = ceil(bucket_frames / fps * video_rate) - audio_frame_num + plan.padded_audio_frames = static_cast( + std::ceil(plan.bucket_frames / static_cast(fps) * video_rate)); + return plan; + } + + // Bucket frame index (fps timeline) -> source frame index (video_rate timeline). + // get_sample_indices with fixed_start=0 reduces to round-half-even(i * video_rate / fps), + // matching numpy's default rounding. + 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)); + } + + // torch.nn.functional.interpolate size computation: output_len = int(in_len / input_fps * output_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); + } + + // torch.nn.functional.interpolate(mode='linear', align_corners=True) along the frame + // dimension. in: [num_layers, in_frames, dim], out: [num_layers, out_frames, dim]. + 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; // zero padding past the audio end + } + 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 000000000..e679608cd --- /dev/null +++ b/src/conditioning/wan_audio.h @@ -0,0 +1,35 @@ +#ifndef __SD_CONDITIONING_WAN_AUDIO_H__ +#define __SD_CONDITIONING_WAN_AUDIO_H__ + +#include + +namespace sd::wan_audio { + + // Chunk/padding math of get_audio_embed_bucket_fps (m=0). + struct BucketPlan { + int audio_frames; // input frames at video_rate (30 Hz) + int batch_frames; // pixel frames per chunk (latent_t * 4) + int video_rate; // timeline rate of audio_frames (30 Hz) + int fps; // bucket frame rate (16 fps) + int num_chunks; // number of audio chunks, including trailing padding + int bucket_frames; // total bucket frames = num_chunks * batch_frames + int padded_audio_frames; // audio_frames plus zero padding applied + }; + + // Full bucketing: stacked encoder states [num_layers, in_frames, dim] at input_fps -> + // bucket frames [bucket_frames, num_layers, dim] at fps, with zero frames past the + // audio end. Chunk c occupies rows [c * batch_frames, (c + 1) * batch_frames). + // Returns an empty vector on invalid input; the applied plan is stored in *plan_out. + 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/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 82cd25d6c..ad56d7b1b 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -28,7 +28,6 @@ #include "model_manager.h" #include "stable-diffusion.h" -#include "conditioning/audio_processing.hpp" #include "conditioning/conditioner.hpp" #include "core/backend_fit.h" #include "extensions/generation_extension.h" @@ -41,6 +40,7 @@ #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" @@ -1770,12 +1770,12 @@ sd::Tensor StableDiffusionGGML::get_audio_embedding(const sd_audio_t& aud LOG_ERROR("invalid driving audio"); return {}; } - auto mono = AudioProcessing::downmix_to_mono(audio.data, audio.sample_count, audio.channels); + 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 = AudioProcessing::resample_audio(mono.data(), mono.size(), audio.sample_rate, 16000); + 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 {}; diff --git a/src/pipeline/video.cpp b/src/pipeline/video.cpp index 31e88eff4..397d69e50 100644 --- a/src/pipeline/video.cpp +++ b/src/pipeline/video.cpp @@ -5,7 +5,7 @@ #include #include -#include "conditioning/audio_processing.hpp" +#include "conditioning/wan_audio.h" #include "core/rng.hpp" #include "core/rng_philox.hpp" #include "diffusion_engine.h" @@ -443,13 +443,13 @@ namespace sd::pipeline { layer_first.data() + (static_cast(l) * in_frames + f) * embed_dim); } } - AudioProcessing::BucketPlan plan; - std::vector buckets = AudioProcessing::build_audio_buckets(layer_first.data(), - static_cast(num_layers), - static_cast(in_frames), - static_cast(embed_dim), - static_cast(batch_frames), - &plan); + 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 {}; } diff --git a/src/runtime/audio_processing.cpp b/src/runtime/audio_processing.cpp new file mode 100644 index 000000000..9477d0dde --- /dev/null +++ b/src/runtime/audio_processing.cpp @@ -0,0 +1,99 @@ +#include "audio_processing.h" + +#include +#include +#include +#include + +namespace sd::audio { + + // Polyphase FIR from torchaudio.functional.resample: sinc interpolated at output + // phases, hann-windowed over lowpass_filter_width zero crossings, anti-aliased by + // clamping the sinc argument to +-lowpass_filter_width after scaling by min(rate)*rolloff. + 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 000000000..345289c55 --- /dev/null +++ b/src/runtime/audio_processing.h @@ -0,0 +1,23 @@ +#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); + + // Downmix interleaved samples to mono by averaging channels. Returns 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__ From aaf4c1368ce8cc1c6c44ab5e3bbb975c190c4d15 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 13 Sep 2026 23:02:03 +0800 Subject: [PATCH 4/5] refactor: simplify Wav2Vec2 and Wan S2V audio modules --- src/core/ggml_extend.cpp | 77 ++- src/core/ggml_extend.h | 15 +- src/model/audio/wav2vec2.hpp | 1057 ++++++++++------------------- src/model/common/ggml_block.hpp | 57 +- src/model/diffusion/wan.hpp | 241 +------ src/model/diffusion/wan_audio.hpp | 216 ++++++ src/pipeline/diffusion_engine.h | 6 +- src/pipeline/model_builders.cpp | 8 +- src/pipeline/model_builders.h | 6 +- 9 files changed, 743 insertions(+), 940 deletions(-) create mode 100644 src/model/diffusion/wan_audio.hpp diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 7cf585392..79e3ea9bf 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 585d4582c..6f3fe89ee 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/model/audio/wav2vec2.hpp b/src/model/audio/wav2vec2.hpp index cc542b8dd..e0a3fa017 100644 --- a/src/model/audio/wav2vec2.hpp +++ b/src/model/audio/wav2vec2.hpp @@ -1,746 +1,415 @@ #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/common/ggml_block.hpp" #include "model.h" +#include "model/common/ggml_block.hpp" -/*================================================ Wav2Vec2 audio encoder ================================================*/ - -// Port of ComfyUI comfy/audio_encoders/wav2vec2.py (HF wav2vec2 inference subset). -// Checkpoint: wav2vec2-large-english (embed_dim 1024, 24 layers, 16 heads, -// conv_norm/conv_bias true, stable layer norm). Keys carry a "wav2vec2." prefix; -// load the storage map with that prefix so names line up with the block tree. - -struct Wav2Vec2Config { - int64_t embed_dim = 1024; - int64_t conv_dim = 512; - int num_heads = 16; - int num_layers = 24; - bool conv_norm = true; - bool conv_bias = true; - bool do_normalize = true; - bool do_stable_layer_norm = true; - - // Mirrors comfy/audio_encoders/audio_encoders.py:49-76: embed_dim inferred from - // the shape of encoder.layer_norm.bias. - 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) { // large - config.embed_dim = 1024; - config.num_heads = 16; - config.num_layers = 24; - config.conv_norm = true; - config.conv_bias = true; - config.do_normalize = true; - config.do_stable_layer_norm = true; - } else if (config.embed_dim == 768) { // base - config.embed_dim = 768; - config.num_heads = 12; - config.num_layers = 12; - config.conv_norm = false; - 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; - } -}; - -// Conv1d with optional groups (weight [kernel, in_channels, out_channels]). For -// groups == channels == out_channels this stays depthwise-capable via ggml; general -// groups run per-group ggml_conv_1d on channel views and concat. -struct Wav2Vec2Conv1d : public UnaryBlock { - int64_t in_channels; - int64_t out_channels; - int64_t groups; - int kernel_size; - int stride; - int padding; - int dilation; - bool bias; - std::string prefix; - - Wav2Vec2Conv1d(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) - : in_channels(in_channels), - out_channels(out_channels), - groups(groups), - kernel_size(kernel_size), - stride(stride), - padding(padding), - dilation(dilation), - bias(bias) {} - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - this->prefix = prefix; - 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); - } - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { - // x: [L, in_channels, N] - ggml_tensor* w = params["weight"]; - ggml_tensor* b = bias ? params["bias"] : nullptr; - - GGML_ASSERT(x->ne[1] == in_channels); - - if (groups == 1) { - // F32 im2col + F32 weight keeps the conv core in full precision (see - // conv_1d_f32_im2col); mul_mat's mixed-type path requires an F32 - // right-hand side, so cast F16 checkpoint weights once. - ggml_tensor* w32 = w->type == GGML_TYPE_F32 ? w : ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F32); - ggml_tensor* out = conv_1d_f32_im2col(ctx->ggml_ctx, w32, x, stride, padding, dilation); - if (b != nullptr) { - b = ggml_reshape_3d(ctx->ggml_ctx, b, 1, out_channels, 1); - out = ggml_add_inplace(ctx->ggml_ctx, out, b); +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; } - return out; - } - return grouped_conv_1d(ctx, x, w, groups, stride, padding, b); - } - - // Grouped conv via per-group channel views; output channels concatenated in order. - // ggml_conv_1d hardcodes an F16 im2col destination, which rounds the input - // patches even for F32 weights; wav2vec2 carries large activation outliers - // that this rounding amplifies through 24 layers, so the grouped path uses - // ggml_conv_1d's exact composition with an F32 im2col instead (the pos_conv - // weight is always F32 here - it is recomputed in-graph from weight_norm). - static ggml_tensor* conv_1d_f32_im2col(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, int s0, int p0, int d0) { - ggml_tensor* im2col = ggml_im2col(ctx, w, x, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F32); - ggml_tensor* result = ggml_mul_mat(ctx, - ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[2] * im2col->ne[1]), - ggml_reshape_2d(ctx, w, w->ne[0] * w->ne[1], w->ne[2])); - return ggml_reshape_3d(ctx, result, im2col->ne[1], w->ne[2], im2col->ne[2]); - } - - static ggml_tensor* grouped_conv_1d(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* w, - int64_t groups, - int stride, - int padding, - ggml_tensor* b = nullptr) { - const int64_t ic_g = x->ne[1] / groups; - const int64_t oc_g = w->ne[2] / groups; - // ggml's im2col kernels (the backend of ggml_conv_1d) read input planes - // with flat indexing and only honor nb[1] for the channel offset. A - // non-contiguous x (e.g. a permuted view with time stride 4*C) therefore - // reads wrong elements through the per-group views below, because - // ggml_view_3d forces nb[0] to the type size and hides the true stride. - // Materialize once; contiguous inputs pass through untouched. - if (x->nb[0] != ggml_type_size(x->type) || x->nb[1] != (size_t)x->ne[0] * x->nb[0]) { - x = ggml_cont(ctx->ggml_ctx, x); - } - if (getenv("WAV2VEC2_DEBUG_CONV") != nullptr) { - printf("[dbg] x: ne=[%lld,%lld,%lld] nb=[%zu,%zu,%zu]\n", - (long long)x->ne[0], (long long)x->ne[1], (long long)x->ne[2], - x->nb[0], x->nb[1], x->nb[2]); - printf("[dbg] w: ne=[%lld,%lld,%lld] nb=[%zu,%zu,%zu] type=%d\n", - (long long)w->ne[0], (long long)w->ne[1], (long long)w->ne[2], - w->nb[0], w->nb[1], w->nb[2], (int)w->type); - } - ggml_tensor* acc = nullptr; - for (int64_t i = 0; i < groups; ++i) { - ggml_tensor* x_i = ggml_view_3d(ctx->ggml_ctx, x, - x->ne[0], ic_g, x->ne[2], - x->nb[1], x->nb[2], - i * ic_g * x->nb[1]); - ggml_tensor* w_i = ggml_view_3d(ctx->ggml_ctx, w, - w->ne[0], ic_g, oc_g, - w->nb[1], w->nb[2], - i * oc_g * w->nb[2]); - ggml_tensor* out_i = conv_1d_f32_im2col(ctx->ggml_ctx, w_i, x_i, stride, padding, 1); - if (b != nullptr) { - ggml_tensor* b_i = ggml_view_1d(ctx->ggml_ctx, b, oc_g, i * oc_g * b->nb[0]); - b_i = ggml_reshape_3d(ctx->ggml_ctx, b_i, 1, oc_g, 1); - out_i = ggml_add_inplace(ctx->ggml_ctx, out_i, b_i); + 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; } - acc = (acc == nullptr) ? out_i : ggml_concat(ctx->ggml_ctx, acc, out_i, 1); + return config; } - return acc; - } -}; - -// Conv + (optional) per-channel norm + GELU. ComfyUI LayerNormConv applies the -// LayerNorm over channels for every frame; equivalent to normalizing over the -// channel axis after transposing. GroupNorm variant mirrors LayerGroupNormConv. -struct Wav2Vec2ConvLayer : public UnaryBlock { - Wav2Vec2ConvLayer(int64_t in_channels, - int64_t out_channels, - int kernel_size, - int stride, - bool bias, - bool use_layer_norm) { - blocks["conv"] = std::shared_ptr(new Wav2Vec2Conv1d(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias)); - if (use_layer_norm) { - blocks["layer_norm"] = std::shared_ptr(new LayerNorm(out_channels)); + }; + + 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); } - use_layer_norm_ = use_layer_norm; - channels_ = out_channels; - } - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - GGMLBlock::init_params(ctx, tensor_storage_map, prefix); - if (!use_layer_norm_) { - // LayerGroupNormConv: GroupNorm(num_groups=out_channels, affine=true) - params["layer_norm.weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels_); - params["layer_norm.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels_); + + 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 { - // x: [L, C, N] - x = std::dynamic_pointer_cast(blocks["conv"])->forward(ctx, x); - if (use_layer_norm_) { - auto ln = std::dynamic_pointer_cast(blocks["layer_norm"]); - // [L, C, N] -> [C, L, N] -> LayerNorm over C -> back + + 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 = ln->forward(ctx, x); + x = layer_norm->forward(ctx, x); x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); - } else { - x = ggml_cont(ctx->ggml_ctx, x); - x = ggml_group_norm(ctx->ggml_ctx, x, (int)channels_, 1e-05f); - x = ggml_add(ctx->ggml_ctx, x, ggml_reshape_3d(ctx->ggml_ctx, params["layer_norm.bias"], 1, channels_, 1)); - x = ggml_mul(ctx->ggml_ctx, x, ggml_reshape_3d(ctx->ggml_ctx, params["layer_norm.weight"], 1, channels_, 1)); + return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x)); } - return ggml_ext_gelu(ctx->ggml_ctx, x, true); - } - -private: - bool use_layer_norm_; - int64_t channels_; -}; - -struct Wav2Vec2ConvFeatureEncoder : public GGMLBlock { - Wav2Vec2ConvFeatureEncoder(const Wav2Vec2Config& config) { - // kernel sizes (10,3,3,3,3,2,2), strides (5,2,2,2,2,2,2); conv0 maps the - // single waveform channel to conv_dim and always has a bias. - 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) { - bool bias = (i == 0) ? true : config.conv_bias; - blocks["conv_layers." + std::to_string(i)] = std::shared_ptr( - new Wav2Vec2ConvLayer(in_channels, config.conv_dim, kernels[i], strides[i], bias, config.conv_norm)); - in_channels = config.conv_dim; + }; + + 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) { - // x: [L, 1, N] waveform; returns [conv_dim, L', N] - 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); + 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 expects [N, C, H, W]; insert H = 1 for the temporal input. + 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)); } - return ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); // [conv_dim, L', N] - } -}; - -struct Wav2Vec2FeatureProjection : public UnaryBlock { - Wav2Vec2FeatureProjection(const Wav2Vec2Config& config) { - blocks["layer_norm"] = std::shared_ptr(new LayerNorm(config.conv_dim)); - blocks["projection"] = std::shared_ptr(new Linear(config.conv_dim, config.embed_dim)); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - // x: [conv_dim, L', N] -> [embed_dim, L', N] - 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; - } -}; - -struct Wav2Vec2PositionalConvEmbedding : public UnaryBlock { - int64_t embed_dim; - int64_t groups; - int kernel_size; - enum NormConvention { - NORM_PER_KERNEL_TAP, // g [1, 1, k]: norm over all channels per kernel tap - NORM_PER_OUT_KERNEL, // g [out, 1, k]: norm over in/g per (out, kernel tap) - NORM_PER_CHANNEL_PAIR, // g [out, ic_g, 1]: norm over kernel per channel pair }; - NormConvention norm_convention; - bool legacy_key_names; // weight_g/weight_v vs parametrizations.weight.original0/1 - - Wav2Vec2PositionalConvEmbedding(const Wav2Vec2Config& config) - : embed_dim(config.embed_dim), groups(16), kernel_size(128), norm_convention(NORM_PER_KERNEL_TAP), legacy_key_names(true) {} - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - // prefix arrives with a trailing dot, e.g. "...encoder.pos_conv_embed."; - // the weight-norm params live under the inner "conv" module. - const std::string base = prefix + "conv."; - const int64_t ic_g = embed_dim / groups; - ggml_type wtype = GGML_TYPE_F16; - auto key_g = base + "weight_g"; - auto key_v = base + "weight_v"; - legacy_key_names = tensor_storage_map.find(key_g) != tensor_storage_map.end(); - if (!legacy_key_names) { - key_g = base + "parametrizations.weight.original0"; - key_v = base + "parametrizations.weight.original1"; + + 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; + } } - // The stored g shape decides the norm convention (files exist with - // parametrizations-style [1, 1, k] g under legacy key names). - auto it_g = tensor_storage_map.find(key_g); - GGML_ASSERT(it_g != tensor_storage_map.end() && it_g->second.ne[1] == 1); - if (it_g->second.ne[2] == 1) { - norm_convention = NORM_PER_KERNEL_TAP; - } else if (it_g->second.ne[1] == 1 && it_g->second.ne[2] == embed_dim) { - norm_convention = NORM_PER_OUT_KERNEL; - } else if (it_g->second.ne[0] == 1 && it_g->second.ne[1] == ic_g && it_g->second.ne[2] == embed_dim) { - norm_convention = NORM_PER_CHANNEL_PAIR; - } else { - GGML_ABORT("wav2vec2: unsupported weight-norm g shape"); + + 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); } - // Param keys must include the inner "conv." segment: get_param_tensors - // joins the block prefix with the param key verbatim. - const std::string rel_g = key_g.substr(prefix.size()); // "conv.weight_g" / "conv.parametrizations..." - const std::string rel_v = key_v.substr(prefix.size()); - params[rel_g] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, - it_g->second.ne[0], it_g->second.ne[1], it_g->second.ne[2]); - params[rel_v] = ggml_new_tensor_3d(ctx, get_type(key_v, tensor_storage_map, GGML_TYPE_F16), - kernel_size, ic_g, embed_dim); - if (tensor_storage_map.find(base + "bias") != tensor_storage_map.end()) { - params["conv.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, embed_dim); + }; + + 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); } - } - - // Per-tap (or per-convention) norm of v, in-graph. - ggml_tensor* norm(GGMLRunnerContext* ctx) { - const int64_t ic_g = embed_dim / groups; - const char* key_v = legacy_key_names ? "conv.weight_v" : "conv.parametrizations.weight.original1"; - ggml_tensor* v = params[key_v]; - v = ggml_cast(ctx->ggml_ctx, v, GGML_TYPE_F32); - ggml_tensor* sq = ggml_mul(ctx->ggml_ctx, v, v); - ggml_tensor* norm; - if (norm_convention == NORM_PER_KERNEL_TAP) { - // sq [k, ic_g, out]: reduce channels per kernel tap -> [k, 1, 1] - // ggml permute axes are destinations of the source dims (inverse of - // torch's permute), so (2, 0, 1) yields [ic_g, out, k]. - ggml_tensor* pt = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, sq, 2, 0, 1, 3)); // [ic_g, out, k] - ggml_tensor* s = ggml_sum_rows(ctx->ggml_ctx, ggml_reshape_2d(ctx->ggml_ctx, pt, ic_g * embed_dim, kernel_size)); - norm = ggml_sqrt(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, s, kernel_size, 1, 1)); - } else if (norm_convention == NORM_PER_OUT_KERNEL) { - // sq [k, ic_g, out] -> [ic_g, k, out]: reduce ic_g per (k, out) -> [k, 1, out] - ggml_tensor* pt = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, sq, 1, 0, 2, 3)); - norm = ggml_sqrt(ctx->ggml_ctx, - ggml_reshape_3d(ctx->ggml_ctx, - ggml_sum_rows(ctx->ggml_ctx, ggml_reshape_2d(ctx->ggml_ctx, pt, ic_g, kernel_size * embed_dim)), - kernel_size, 1, embed_dim)); - } else { - // sq [k, ic_g, out]: reduce kernel per channel pair -> [1, ic_g, out] - norm = ggml_sqrt(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, sq)); + + 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; } - return norm; - } - - // Recombines the weight-normalized conv weight: weight = v * g / norm(v). - ggml_tensor* weight(GGMLRunnerContext* ctx) { - const char* key_g = legacy_key_names ? "conv.weight_g" : "conv.parametrizations.weight.original0"; - const char* key_v = legacy_key_names ? "conv.weight_v" : "conv.parametrizations.weight.original1"; - ggml_tensor* g = params[key_g]; - ggml_tensor* v = params[key_v]; - v = ggml_cast(ctx->ggml_ctx, v, GGML_TYPE_F32); - return ggml_mul(ctx->ggml_ctx, v, ggml_div(ctx->ggml_ctx, g, norm(ctx))); - } - - // Debug: raw v^2 (pre-reduction) and the transposed copy feeding the reduction. - ggml_tensor* sq_dump(GGMLRunnerContext* ctx) { - const char* key_v = legacy_key_names ? "conv.weight_v" : "conv.parametrizations.weight.original1"; - ggml_tensor* v = params[key_v]; - v = ggml_cast(ctx->ggml_ctx, v, GGML_TYPE_F32); - return ggml_mul(ctx->ggml_ctx, v, v); - } - - ggml_tensor* pt_dump(GGMLRunnerContext* ctx) { - return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, sq_dump(ctx), 2, 0, 1, 3)); - } - - // Debug: grouped conv output after drop-last, before GELU (with_bias=false - // also skips the bias add). - ggml_tensor* conv_raw(GGMLRunnerContext* ctx, ggml_tensor* x, bool with_bias) { - ggml_tensor* b = with_bias && params.count("conv.bias") > 0 ? params["conv.bias"] : nullptr; - ggml_tensor* t = conv_input(ctx, x); - t = Wav2Vec2Conv1d::grouped_conv_1d(ctx, t, weight(ctx), groups, 1, kernel_size / 2, b); - t = ggml_view_3d(ctx->ggml_ctx, t, t->ne[0] - 1, t->ne[1], t->ne[2], t->nb[1], t->nb[2], 0); - return t; // [L, embed_dim, N] - } - - // The materialized [L, embed_dim, N] plane-contiguous tensor that - // grouped_conv_1d actually consumes (cont of the permuted projection - // output; im2col requires this layout). - ggml_tensor* conv_input(GGMLRunnerContext* ctx, ggml_tensor* x) { - return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); - } - - // Debug: group-0-only conv (no bias), post drop-last. - ggml_tensor* conv_group0(GGMLRunnerContext* ctx, ggml_tensor* x) { - const int64_t ic_g = embed_dim / groups; - ggml_tensor* t = conv_input(ctx, x); - ggml_tensor* w = weight(ctx); - ggml_tensor* x_0 = ggml_view_3d(ctx->ggml_ctx, t, t->ne[0], ic_g, t->ne[2], t->nb[1], t->nb[2], 0); - ggml_tensor* w_0 = ggml_view_3d(ctx->ggml_ctx, w, w->ne[0], ic_g, w->ne[2] / groups, w->nb[1], w->nb[2], 0); - ggml_tensor* out = Wav2Vec2Conv1d::conv_1d_f32_im2col(ctx->ggml_ctx, w_0, x_0, 1, kernel_size / 2, 1); - out = ggml_view_3d(ctx->ggml_ctx, out, out->ne[0] - 1, out->ne[1], out->ne[2], out->nb[1], out->nb[2], 0); - return out; // [L, ic_g, N] - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - // x: [embed_dim, L, N]; returns [embed_dim, L, N] (last sample dropped) - ggml_tensor* w = weight(ctx); - ggml_tensor* b = params.count("conv.bias") > 0 ? params["conv.bias"] : nullptr; - ggml_tensor* t = conv_input(ctx, x); - t = Wav2Vec2Conv1d::grouped_conv_1d(ctx, t, w, groups, 1, kernel_size / 2, b); - // GELU must be out-of-place and precede the drop-last view: an in-place - // op on a view writes through a buffer the graph allocator does not - // reserve for it, so reuse of the conv output's storage silently - // corrupts the result depending on the surrounding graph shape. - t = ggml_ext_gelu(ctx->ggml_ctx, t, false); - t = ggml_view_3d(ctx->ggml_ctx, t, t->ne[0] - 1, t->ne[1], t->ne[2], t->nb[1], t->nb[2], 0); - return ggml_permute(ctx->ggml_ctx, t, 1, 0, 2, 3); // [embed_dim, L, N] - } -}; - -struct Wav2Vec2FeedForward : public UnaryBlock { - Wav2Vec2FeedForward(const Wav2Vec2Config& config) { - blocks["intermediate_dense"] = std::shared_ptr(new Linear(config.embed_dim, config.embed_dim * 4)); - blocks["output_dense"] = std::shared_ptr(new Linear(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 Wav2Vec2TransformerEncoderLayer : public UnaryBlock { - bool do_stable_layer_norm; - - Wav2Vec2TransformerEncoderLayer(const Wav2Vec2Config& config) - : do_stable_layer_norm(config.do_stable_layer_norm) { - blocks["attention"] = std::shared_ptr(new MultiheadAttention(config.embed_dim, config.num_heads, true, true)); - blocks["layer_norm"] = std::shared_ptr(new LayerNorm(config.embed_dim)); - blocks["feed_forward"] = std::shared_ptr(new Wav2Vec2FeedForward(config)); - blocks["final_layer_norm"] = std::shared_ptr(new LayerNorm(config.embed_dim)); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - // x: [embed_dim, L, N] - 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))); + }; + + 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_); + } } - return x; - } -}; - -struct Wav2Vec2TransformerEncoder : public GGMLBlock { - int num_layers; - bool do_stable_layer_norm; - - Wav2Vec2TransformerEncoder(const Wav2Vec2Config& config) - : num_layers(config.num_layers), do_stable_layer_norm(config.do_stable_layer_norm) { - blocks["pos_conv_embed"] = std::shared_ptr(new Wav2Vec2PositionalConvEmbedding(config)); - for (int i = 0; i < config.num_layers; ++i) { - blocks["layers." + std::to_string(i)] = std::shared_ptr(new Wav2Vec2TransformerEncoderLayer(config)); + + 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); } - blocks["layer_norm"] = std::shared_ptr(new LayerNorm(config.embed_dim)); - } - - ggml_tensor* pos_conv(GGMLRunnerContext* ctx, ggml_tensor* x) { - auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); - return pc->forward(ctx, x); - } - - ggml_tensor* pos_conv_weight(GGMLRunnerContext* ctx) { - auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); - return pc->weight(ctx); - } - - ggml_tensor* pos_conv_norm(GGMLRunnerContext* ctx) { - auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); - return pc->norm(ctx); - } - - ggml_tensor* pos_conv_sq(GGMLRunnerContext* ctx) { - auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); - return pc->sq_dump(ctx); - } - - ggml_tensor* pos_conv_pt(GGMLRunnerContext* ctx) { - auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); - return pc->pt_dump(ctx); - } - - ggml_tensor* pos_conv_raw(GGMLRunnerContext* ctx, ggml_tensor* x, bool with_bias) { - auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); - return pc->conv_raw(ctx, x, with_bias); - } - - ggml_tensor* pos_conv_input(GGMLRunnerContext* ctx, ggml_tensor* x) { - auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); - return pc->conv_input(ctx, x); - } - - ggml_tensor* pos_conv_group0(GGMLRunnerContext* ctx, ggml_tensor* x) { - auto pc = std::dynamic_pointer_cast(blocks["pos_conv_embed"]); - return pc->conv_group0(ctx, x); - } - - // Returns the final hidden state [embed_dim, L, N]. When all_layers != nullptr, - // it is filled with the 24 pre-layer states plus the final state, concatenated - // along a new trailing axis: [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); + }; + + 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); } - x = ggml_add(ctx->ggml_ctx, x, pos_conv_embed->forward(ctx, x)); - if (!do_stable_layer_norm) { - x = layer_norm->forward(ctx, x); + 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; } - for (int i = 0; i < num_layers; ++i) { - if (all_layers != nullptr) { - collected.push_back(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))); } - 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); + return 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); + }; + + 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); } - *all_layers = stack; + blocks["layer_norm"] = std::make_shared(config.embed_dim); } - return x; - } -}; - -struct Wav2Vec2Model : public GGMLBlock { - Wav2Vec2Config config; - - Wav2Vec2Model() = default; - Wav2Vec2Model(const Wav2Vec2Config& config_) : config(config_) { - blocks["feature_extractor"] = std::shared_ptr(new Wav2Vec2ConvFeatureEncoder(config)); - blocks["feature_projection"] = std::shared_ptr(new Wav2Vec2FeatureProjection(config)); - blocks["encoder"] = std::shared_ptr(new Wav2Vec2TransformerEncoder(config)); - } - - // waveform: [L, 1, N] (already channel-mixed). Returns the final hidden state; - // all_layers receives the per-layer stack when non-null. - 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; - } - - // Debug front-end stages (see Wav2Vec2ModelRunner::build_graph): - // 1 = feature extractor, 2 = + projection, 3 = + pos_conv and residual add, - // 4 = the recomputed pos_conv weight tensor, 5 = the pos_conv weight norm. - ggml_tensor* forward_front(GGMLRunnerContext* ctx, ggml_tensor* x, int stage) { - 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"]); - - if (stage == 11) { - x = feature_extractor->forward(ctx, x); - x = feature_projection->forward(ctx, x); - return encoder->pos_conv_group0(ctx, x); + + // 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; } - if (stage == 10) { - x = feature_extractor->forward(ctx, x); - x = feature_projection->forward(ctx, x); - return encoder->pos_conv_input(ctx, 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); } - if (stage == 9 || stage == 8) { + + 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); - return encoder->pos_conv_raw(ctx, x, stage == 8); - } - if (stage == 7) { - return encoder->pos_conv_pt(ctx); - } - if (stage == 6) { - return encoder->pos_conv_sq(ctx); - } - if (stage == 5) { - return encoder->pos_conv_norm(ctx); - } - if (stage == 4) { - return encoder->pos_conv_weight(ctx); - } - x = feature_extractor->forward(ctx, x); - if (stage == 1) { + x = encoder->forward(ctx, x, all_layers); return x; } - x = feature_projection->forward(ctx, x); - if (stage == 2) { - 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()); } - return ggml_add(ctx->ggml_ctx, x, encoder->pos_conv(ctx, x)); - } -}; - -class Wav2Vec2ModelRunner : public GGMLRunner { -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), weight_prefix(prefix) { - config = Wav2Vec2Config::detect_from_weights(tensor_storage_map, prefix); - model = Wav2Vec2Model(config); - // GGMLBlock::init/get_param_tensors append their own '.'; the loader-facing - // prefix convention carries one. - std::string block_prefix = weight_prefix; - if (!block_prefix.empty() && block_prefix.back() == '.') { - block_prefix.pop_back(); + + std::string get_desc() override { + return "wav2vec2"; } - model.init(params_ctx, tensor_storage_map, block_prefix); - LOG_INFO("%s", get_desc().c_str()); - } - - std::string get_desc() override { - char buf[256]; - snprintf(buf, sizeof(buf), "wav2vec2: embed_dim %" PRId64 ", %d layers, %d heads%s", - config.embed_dim, config.num_layers, config.num_heads, - config.do_stable_layer_norm ? ", stable-ln" : ""); - return std::string(buf); - } - - void get_param_tensors(std::map& tensors) { - std::string block_prefix = weight_prefix; - if (!block_prefix.empty() && block_prefix.back() == '.') { - block_prefix.pop_back(); + + 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); } - model.get_param_tensors(tensors, block_prefix); - } - - // Allocates all param tensors on the backend's default buffer. Standalone - // harnesses need this before loading weights; the full pipeline routes - // allocation through ModelManager instead. - bool alloc_params_on_backend(ggml_backend_t backend) { - params_buffer = ggml_backend_alloc_ctx_tensors_from_buft(params_ctx, ggml_backend_get_default_buffer_type(backend)); - return params_buffer != nullptr; - } - - // waveform: [L, 1, 1] mono; already host-normalized. stage 0 (default) returns - // [embed_dim, L', num_layers + 1] with the per-layer stack (last slice = final - // hidden state). Stages 1..3 expose front-end intermediates for parity debugging: - // 1 = feature extractor [conv_dim, L', 1], 2 = + projection [embed_dim, L', 1], - // 3 = + pos_conv and residual add [embed_dim, L', 1]. - ggml_cgraph* build_graph(const sd::Tensor& waveform_tensor, int stage = 0) { - ggml_cgraph* gf = ggml_new_graph(compute_ctx); - ggml_tensor* waveform = make_input(waveform_tensor); - - auto runner_ctx = get_context(); - - if (stage == 0) { + + // Normalized mono input [L, 1, 1]; returns [embed_dim, L', num_layers + 1]. + 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); - } else { - // stages produce views (permutes); materialize them so the flat dump - // reflects logical [C, T, N] order - ggml_tensor* out = ggml_cont(runner_ctx.ggml_ctx, model.forward_front(&runner_ctx, waveform, stage)); - ggml_build_forward_expand(gf, out); + return gf; } - return gf; - } - - sd::Tensor compute(const int n_threads, const std::vector& mono_waveform, int stage = 0) { - 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, stage); - }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); - } - -private: - Wav2Vec2Config config; - ggml_backend_buffer_t params_buffer = nullptr; - - // torch: (x - x.mean()) / torch.sqrt(x.var() + 1e-7); var is population variance. - 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; + + 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)); } - 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); + + private: + // Match waveform normalization with population variance and epsilon 1e-7. + 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 cc0f69419..20ed8f013 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -366,6 +366,61 @@ class Embedding : public UnaryBlock { } }; +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; @@ -764,7 +819,7 @@ class GroupNorm : public GGMLBlock { 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/diffusion/wan.hpp b/src/model/diffusion/wan.hpp index 53793270e..d1415c7df 100644 --- a/src/model/diffusion/wan.hpp +++ b/src/model/diffusion/wan.hpp @@ -9,7 +9,6 @@ #include "core/ggml_extend_backend.h" #include "core/ggml_tensor_utils.h" -#include "model/audio/wav2vec2.hpp" #include "model/common/block.hpp" #include "model/common/rope.hpp" #include "model/diffusion/flux.hpp" @@ -277,6 +276,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] @@ -540,235 +546,6 @@ namespace WAN { } }; - //================================================ Wan2.2-S2V audio conditioning ================================================= - // Ported from ComfyUI comfy/ldm/wan/model.py: CausalConv1d, MotionEncoder_tc, - // CausalAudioEncoder and AudioInjector_WAN. Only the adain_mode="attn_norm" - // configuration exists in the Wan2.2-S2V-14B checkpoint (injector_pre_norm_* - // have no keys), so the pre-norm fallback path is intentionally absent. - - class WanCausalConv1d : public GGMLBlock { - protected: - 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::shared_ptr(new Wav2Vec2Conv1d(in_dim, out_dim, kernel_size, stride)); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - // x: [T, C] -> [T', C_out]; replicate left-pad (k-1, 0) along T - 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 { - protected: - 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 { - // Elementwise-consumed param (ggml_repeat/concat against F32 - // activations): hardcode F32; the loader converts the storage type. - params["padding_tokens"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_dim); - } - - // conv -> affine-free LayerNorm over channels -> SiLU. The norm runs on - // the permuted [C, T] form; to_conv_layout materializes the [T', C] - // form the next conv needs (im2col requires contiguous time rows), - // otherwise the token-layout [C, T'] is returned directly. - 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); // [T', C_out] - x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); // [C_out, T'] - 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)); // [T', C_out] - } - 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::shared_ptr(new WanCausalConv1d(in_dim, hidden_dim / 4 * num_token)); - if (need_global) { - blocks["conv1_global"] = std::shared_ptr(new WanCausalConv1d(in_dim, hidden_dim / 4)); - } - blocks["norm1"] = std::shared_ptr(new LayerNorm(hidden_dim / 4, 1e-6f, false)); - blocks["conv2"] = std::shared_ptr(new WanCausalConv1d(hidden_dim / 4, hidden_dim / 2, 3, 2)); - blocks["norm2"] = std::shared_ptr(new LayerNorm(hidden_dim / 2, 1e-6f, false)); - blocks["conv3"] = std::shared_ptr(new WanCausalConv1d(hidden_dim / 2, hidden_dim, 3, 2)); - blocks["norm3"] = std::shared_ptr(new LayerNorm(hidden_dim, 1e-6f, false)); - if (need_global) { - blocks["final_linear"] = std::shared_ptr(new Linear(hidden_dim, hidden_dim)); - } - } - - // x: [T, in_dim]. Returns (local [dim, num_token+1, T/4], global [dim, T/4]); - // torch groups the conv1_local channels into num_token batches and the - // group index becomes the per-frame token index. - std::pair forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - auto local = std::dynamic_pointer_cast(blocks["conv1_local"])->forward(ctx, x); // [T, hidden] - auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); - std::vector tokens; - for (auto& group : ggml_ext_chunk(ctx->ggml_ctx, local, num_token, 1)) { - // rearrange 'b (n c) t -> (b n) t c': per-group norm over c - ggml_tensor* s = ggml_permute(ctx->ggml_ctx, group, 1, 0, 2, 3); // [C, T] view - 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)); // [T, C] - s = conv_norm_silu(ctx, s, "conv2", "norm2", true); // [T2, hidden/2] - s = conv_norm_silu(ctx, s, "conv3", "norm3", false); // [hidden, T3] - tokens.push_back(ggml_reshape_3d(ctx->ggml_ctx, s, s->ne[0], 1, s->ne[1])); // [dim, 1, T3] - } - auto padding = ggml_reshape_3d(ctx->ggml_ctx, params["padding_tokens"], hidden_dim, 1, 1); - padding = ggml_repeat(ctx->ggml_ctx, padding, tokens[0]); // [dim, 1, T3] - tokens.push_back(padding); - ggml_tensor* local_out = ggml_ext_vec_concat(ctx->ggml_ctx, tokens, 1); // [dim, num_token+1, T3] - - if (!need_global) { - return {local_out, nullptr}; - } - // global branch shares conv2/conv3/norms with the local branch - ggml_tensor* g = conv_norm_silu(ctx, x, "conv1_global", "norm1", true); // [T, hidden/4] - g = conv_norm_silu(ctx, g, "conv2", "norm2", true); // [T2, hidden/2] - g = conv_norm_silu(ctx, g, "conv3", "norm3", false); // [hidden, T3] - g = std::dynamic_pointer_cast(blocks["final_linear"])->forward(ctx, g); - return {local_out, g}; - } - }; - - class WanCausalAudioEncoder : public GGMLBlock { - protected: - int64_t audio_dim; - int64_t dim; - int num_token; - int num_layers; - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - // Elementwise-consumed param (silu/mul/div in forward): hardcode F32; - // the loader converts the storage type. Checkpoint shape is [1, 25, - // 1, 1] (torch order, leading 1s), so the reversed storage ne is - // [1, 1, 25, 1]; mirror it for raw-bytes loading. forward reshapes - // to [1, 1, L] either way. - 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) - : audio_dim(audio_dim), dim(dim), num_token(num_token), num_layers(num_layers) { - blocks["encoder"] = std::shared_ptr(new WanMotionEncoder(audio_dim, dim, num_token, true)); - } - - // features: [audio_dim, T*4, num_layers] stacked wav2vec2 hidden states. - // SiLU-weighted sum over the layer axis, then MotionEncoder. - // Returns (local [dim, num_token+1, T], global [dim, T]). - std::pair forward(GGMLRunnerContext* ctx, ggml_tensor* features) { - auto weights = ggml_silu(ctx->ggml_ctx, params["weights"]); // [L] - 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)); - // reduce over the layer axis: layers -> ne0 for sum_rows - // (torch-style 3-cycle; plain ggml_permute axes are destinations, - // the inverse mapping - see the pos_conv weight-norm bug) - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); // [L, dim, F] - x = ggml_sum_rows(ctx->ggml_ctx, x); // [1, dim, F] - x = ggml_reshape_2d(ctx->ggml_ctx, x, x->ne[1], x->ne[2]); // [dim, F] - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [F, dim] - return std::dynamic_pointer_cast(blocks["encoder"])->forward(ctx, x); - } - }; - - class WanAudioInjector : public GGMLBlock { - protected: - 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::shared_ptr(new WanT2VCrossAttention(dim, num_heads, qk_norm, eps)); - blocks["injector_adain_layers." + std::to_string(i) + ".linear"] = - std::shared_ptr(new Linear(dim, dim * 2)); - } - // AdaLayerNorm norm is affine-free (no checkpoint keys); eps 1e-5 per - // upstream AdaLayerNorm, distinct from the model-wide 1e-6. - blocks["adain_norm"] = std::shared_ptr(new LayerNorm(dim, 1e-5f, false)); - } - - // x: [dim, n_token] full sequence; audio_local: [dim, num_token+1, T]; - // audio_global: [dim, T]. The residual is applied to the first seq_len - // (video) tokens only; trailing reference tokens pass through unchanged. - // Batching over T at ne[2] mirrors torch's rearrange to (b t) n c. - 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); // [dim, T] - temb = adain_linear->forward(ctx, temb); // [2*dim, T] - auto shift = ggml_ext_slice(ctx->ggml_ctx, temb, 0, 0, dim); // [dim, T] - auto scale = ggml_ext_slice(ctx->ggml_ctx, temb, 0, dim, dim * 2); // [dim, T] - 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); // [dim, seq_len] view - auto h = ggml_reshape_3d(ctx->ggml_ctx, x_vid, dim, n_tok, T); // [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 * (1 + scale) - h = ggml_add(ctx->ggml_ctx, h, shift); - - auto res = injector->forward(ctx, h, audio_local, 0); // [dim, n_tok, T] - res = ggml_reshape_2d(ctx->ggml_ctx, res, dim, seq_len); - - auto x_head = ggml_add(ctx->ggml_ctx, x_vid, res); // out-of-place add on the view - 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; - } - }; - class Wan : public GGMLBlock { protected: WanConfig config; @@ -846,8 +623,8 @@ namespace WAN { // s2v audio conditioning (checkpoint-contained modules) if (config.model_type == "s2v") { - blocks["casual_audio_encoder"] = std::shared_ptr(new WanCausalAudioEncoder(config.audio_dim, config.dim, config.num_audio_token)); - blocks["audio_injector"] = std::shared_ptr(new WanAudioInjector(config.dim, config.num_heads, (int)config.audio_inject_layers.size(), config.qk_norm, config.eps)); + 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; } diff --git a/src/model/diffusion/wan_audio.hpp b/src/model/diffusion/wan_audio.hpp new file mode 100644 index 000000000..eda16a03c --- /dev/null +++ b/src/model/diffusion/wan_audio.hpp @@ -0,0 +1,216 @@ +#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}; + } + // The global branch shares conv2, conv3 and norms with the local branch. + 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); + } + // AdaLayerNorm is affine-free and uses eps=1e-5, unlike the 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/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h index 76f163c7c..5cad06f8c 100644 --- a/src/pipeline/diffusion_engine.h +++ b/src/pipeline/diffusion_engine.h @@ -27,7 +27,9 @@ struct LoraModel; struct ConditionerParams; struct SDCondition; struct RefImageParams; -class Wav2Vec2ModelRunner; +namespace Wav2Vec2 { + class Wav2Vec2ModelRunner; +} extern const char* model_version_to_str[]; @@ -64,7 +66,7 @@ class StableDiffusionGGML { std::shared_ptr first_stage_model; std::shared_ptr preview_vae; std::shared_ptr audio_vae_model; - std::shared_ptr audio_encoder; // wan2.2 s2v driving audio + std::shared_ptr audio_encoder; // wan2.2 s2v driving audio std::shared_ptr control_net; std::shared_ptr ip_adapter; sd::Tensor ip_adapter_tokens; diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index 932d31d02..7f5718e02 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -239,10 +239,10 @@ namespace sd::model_builders { 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); + 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; diff --git a/src/pipeline/model_builders.h b/src/pipeline/model_builders.h index a1389e7eb..aeaed5767 100644 --- a/src/pipeline/model_builders.h +++ b/src/pipeline/model_builders.h @@ -15,7 +15,9 @@ struct DiffusionModelRunner; struct VAE; struct AudioVAERunner; struct ControlNet; -class Wav2Vec2ModelRunner; +namespace Wav2Vec2 { + class Wav2Vec2ModelRunner; +} struct GenerationExtension; struct GenerationExtensionInitContext; namespace IPAdapter { @@ -38,7 +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; // wan2.2 s2v driving audio + std::shared_ptr audio_encoder; // wan2.2 s2v driving audio }; struct VAEOptions { From 3e8b9e23414622134474059def8cf0d6d93e1bc1 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 13 Sep 2026 23:19:13 +0800 Subject: [PATCH 5/5] simplify Wan S2V comments and usage notes --- docs/wan.md | 22 ---------------------- examples/common/common.cpp | 2 -- include/stable-diffusion.h | 3 +-- src/conditioning/wan_audio.cpp | 22 ++++++---------------- src/conditioning/wan_audio.h | 21 +++++++++------------ src/model/audio/wav2vec2.hpp | 4 +--- src/model/common/rope.hpp | 3 +-- src/model/diffusion/model.hpp | 2 +- src/model/diffusion/wan.hpp | 29 ++++++++++++----------------- src/model/diffusion/wan_audio.hpp | 3 +-- src/model_io/safetensors_io.cpp | 5 +---- src/model_loader.cpp | 2 +- src/pipeline/diffusion_engine.cpp | 8 ++------ src/pipeline/diffusion_engine.h | 2 +- src/pipeline/model_builders.h | 2 +- src/pipeline/video.cpp | 26 ++++++-------------------- src/runtime/audio_processing.cpp | 4 +--- src/runtime/audio_processing.h | 3 +-- 18 files changed, 46 insertions(+), 117 deletions(-) diff --git a/docs/wan.md b/docs/wan.md index a61503ce4..c7cf1dd9e 100644 --- a/docs/wan.md +++ b/docs/wan.md @@ -120,10 +120,6 @@ Notes: - 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. -- Use `--diffusion-fa` together with `--offload-to-cpu --vae-tiling`. Without - flash attention the S2V attention compute buffer alone needs ~40 GB VRAM at - 480x832; `--offload-to-cpu` keeps the model weights in system RAM, and - `--vae-tiling` keeps the video VAE decode (77-81 frames) within VRAM. - `--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: @@ -133,19 +129,6 @@ Notes: 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. -- The output video carries the driving audio track: it is muxed into `.avi` - / `.webm` outputs, truncated to the video duration. Other container types - fall back to a `.wav` sidecar file next to the video. -- Approximate VRAM use at 640x368, 13 frames: - - | checkpoint | VRAM | - |------------|------| - | wan2.2_s2v_14B_int8_convrot.safetensors | 17.6 GB | - | wan2.2_s2v-14B-Q8_0.gguf | 18.4 GB | - | wan2.2_s2v-14B-Q4_K_M.gguf | 13.4 GB | - - Q4_K_M is the lowest-memory option; bf16 and fp8_scaled need more than - 24 GB VRAM. - 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 @@ -159,11 +142,6 @@ Notes: ``` Expect some quality/dynamics loss compared to the full 20-step run. -- fp8_scaled and GGUF checkpoints are supported. The int8_convrot checkpoint - is available at - https://huggingface.co/noctrex/Wan2.2-S2V-14B-int8_convrot - ([int8_convrot](int8_convrot.md)); int8 convrot is supported on the CUDA, - Vulkan and ROCm backends. ### Wan2.2 T2V A14B T2I diff --git a/examples/common/common.cpp b/examples/common/common.cpp index fec7e1e64..04bd89ccc 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -1526,8 +1526,6 @@ ArgOptions SDGenerationParams::get_options() { if (++index >= argc) { return -1; } - // S2V driving track shares the ref-audio transport (Wan2.2 S2V reads - // sd_vid_gen_params.ref_audios[0]). ref_audio_paths.push_back(argv[index]); return 1; }; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index ce4dfbd56..31a87e85d 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -522,8 +522,7 @@ 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); -// fps_out may be NULL; otherwise it receives the effective frame rate before preview callbacks. -// Use this frame rate when encoding the returned frames. +// 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, diff --git a/src/conditioning/wan_audio.cpp b/src/conditioning/wan_audio.cpp index 18b705df1..d4d59ca12 100644 --- a/src/conditioning/wan_audio.cpp +++ b/src/conditioning/wan_audio.cpp @@ -4,11 +4,6 @@ #include #include -// Input is the wav2vec2 hidden states stacked per layer [num_layers, in_frames, dim] -// at the encoder frame rate (50 Hz). The frames are interpolated to video_rate -// (30 Hz), bucketed to fps (16) frames with zero padding past the audio end, and -// split into chunks of batch_frames = latent_t * 4 frames (one per diffusion chunk). - namespace sd::wan_audio { static BucketPlan plan_buckets(int audio_frames, int batch_frames, int video_rate, int fps) { @@ -18,29 +13,24 @@ namespace sd::wan_audio { plan.video_rate = video_rate; plan.fps = fps; const double scale = static_cast(video_rate) / fps; - // min_batch_num = int(audio_frame_num / (batch_frames * scale)) + 1 - plan.num_chunks = static_cast(audio_frames / (batch_frames * scale)) + 1; - plan.bucket_frames = plan.num_chunks * batch_frames; - // padd_audio_num = ceil(bucket_frames / fps * video_rate) - audio_frame_num + // 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; } - // Bucket frame index (fps timeline) -> source frame index (video_rate timeline). - // get_sample_indices with fixed_start=0 reduces to round-half-even(i * video_rate / fps), - // matching numpy's default rounding. + // 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)); } - // torch.nn.functional.interpolate size computation: output_len = int(in_len / input_fps * output_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); } - // torch.nn.functional.interpolate(mode='linear', align_corners=True) along the frame - // dimension. in: [num_layers, in_frames, dim], out: [num_layers, out_frames, dim]. + // Match PyTorch linear interpolation with align_corners=True. static std::vector linear_interpolate_frames(const std::vector& in, int num_layers, int in_frames, @@ -99,7 +89,7 @@ namespace sd::wan_audio { 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; // zero padding past the audio end + continue; } for (int layer = 0; layer < num_layers; ++layer) { std::copy_n(interpolated.data() + (static_cast(layer) * audio_frames + src) * dim, diff --git a/src/conditioning/wan_audio.h b/src/conditioning/wan_audio.h index e679608cd..1b245ed5c 100644 --- a/src/conditioning/wan_audio.h +++ b/src/conditioning/wan_audio.h @@ -5,21 +5,18 @@ namespace sd::wan_audio { - // Chunk/padding math of get_audio_embed_bucket_fps (m=0). struct BucketPlan { - int audio_frames; // input frames at video_rate (30 Hz) - int batch_frames; // pixel frames per chunk (latent_t * 4) - int video_rate; // timeline rate of audio_frames (30 Hz) - int fps; // bucket frame rate (16 fps) - int num_chunks; // number of audio chunks, including trailing padding - int bucket_frames; // total bucket frames = num_chunks * batch_frames - int padded_audio_frames; // audio_frames plus zero padding applied + 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; }; - // Full bucketing: stacked encoder states [num_layers, in_frames, dim] at input_fps -> - // bucket frames [bucket_frames, num_layers, dim] at fps, with zero frames past the - // audio end. Chunk c occupies rows [c * batch_frames, (c + 1) * batch_frames). - // Returns an empty vector on invalid input; the applied plan is stored in *plan_out. + // [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, diff --git a/src/model/audio/wav2vec2.hpp b/src/model/audio/wav2vec2.hpp index e0a3fa017..3f3c969ef 100644 --- a/src/model/audio/wav2vec2.hpp +++ b/src/model/audio/wav2vec2.hpp @@ -99,7 +99,7 @@ namespace Wav2Vec2 { 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 expects [N, C, H, W]; insert H = 1 for the temporal input. + // 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]); @@ -364,7 +364,6 @@ namespace Wav2Vec2 { model.get_param_tensors(tensors, block_prefix); } - // Normalized mono input [L, 1, 1]; returns [embed_dim, L', num_layers + 1]. ggml_cgraph* build_graph(const sd::Tensor& waveform_tensor) { ggml_cgraph* gf = ggml_new_graph(compute_ctx); ggml_tensor* waveform = make_input(waveform_tensor); @@ -390,7 +389,6 @@ namespace Wav2Vec2 { } private: - // Match waveform normalization with population variance and epsilon 1e-7. static void normalize(float* x, int64_t n) { double mean = 0.0; for (int64_t i = 0; i < n; ++i) { diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index 19dcad0fa..778bbf426 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -809,8 +809,7 @@ namespace Rope { return embed_nd(ids, bs, static_cast(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage); } - // Generate wan positional embeddings; t_offset shifts the temporal ids - // (S2V reference latent uses t_start = max(30, t + 9)). + // Generate wan positional embeddings __STATIC_INLINE__ std::vector gen_wan_pe(int t, int h, int w, diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 6c7564162..a4b3c38fd 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -69,7 +69,7 @@ struct AnimaDiffusionExtra { struct WanDiffusionExtra { const sd::Tensor* vace_context = nullptr; float vace_strength = 1.f; - // S2V: stacked wav2vec2 hidden states, [num_layers=25, frames=T_latent*4, dim=1024] + // S2V audio, sd::Tensor layout: [dim, T_latent*4, layers]. const sd::Tensor* audio_embed = nullptr; }; diff --git a/src/model/diffusion/wan.hpp b/src/model/diffusion/wan.hpp index ead41cfd7..90fd1285b 100644 --- a/src/model/diffusion/wan.hpp +++ b/src/model/diffusion/wan.hpp @@ -34,9 +34,8 @@ namespace WAN { int vace_layers = 0; int64_t vace_in_dim = 96; std::map vace_layers_mapping = {}; - // Wan2.2-S2V audio conditioning - int64_t audio_dim = 1024; // wav2vec2-large hidden size - int num_audio_token = 4; // motion tokens per frame; 1 learned padding token appended at use site + 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"; @@ -621,7 +620,6 @@ namespace WAN { blocks["vace_patch_embedding"] = std::shared_ptr(new Conv3d(config.vace_in_dim, config.dim, config.patch_size, config.patch_size)); } - // s2v audio conditioning (checkpoint-contained modules) 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); @@ -686,8 +684,8 @@ namespace WAN { // vace_context: [N*vace_in_dim, T, H, W] // timestep: [N,] or [T] // context: [N, L, text_dim] - // audio_embed: [audio_dim, T*4, 25] stacked wav2vec2 hidden states (S2V) - // reference_latent: [N*C, T_ref, H, W] (S2V) + // 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); @@ -710,36 +708,35 @@ 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] - // s2v: audio conditioning + reference latent append 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); // one pixel-frame per latent frame * 4 + 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; // [dim, num_audio_token+1, T] - audio_global = audio_emb.second; // [dim, T] + 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"]; // [dim, 3] + 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); // [N*dim, t_ref_len, h_len, w_len] + 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 frames run at timestep 0 (upstream t concat with zeros) + // 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); } @@ -812,8 +809,6 @@ namespace WAN { x = ggml_add(ctx->ggml_ctx, x, c_skip); } - // s2v: AdaIN from the global audio token, then cross-attention - // against the per-frame audio tokens; residual on video tokens only if (audio_injector != nullptr) { auto inject_iter = config.audio_inject_mapping.find(i); if (inject_iter != config.audio_inject_mapping.end()) { @@ -873,7 +868,7 @@ namespace WAN { 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) { - // drop the trailing reference tokens (upstream unpatchify slices to prod(grid_sizes)) + // Exclude reference tokens from the generated video. out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, t_len * h_len * w_len); } @@ -1015,7 +1010,7 @@ namespace WAN { config.theta, config.axes_dim); if (ref_latent != nullptr) { - // s2v reference latent: temporal offset t_start = max(30, T + 9) + // 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]), diff --git a/src/model/diffusion/wan_audio.hpp b/src/model/diffusion/wan_audio.hpp index eda16a03c..28f31bda9 100644 --- a/src/model/diffusion/wan_audio.hpp +++ b/src/model/diffusion/wan_audio.hpp @@ -104,7 +104,6 @@ namespace WAN { if (!need_global_) { return {local_out, nullptr}; } - // The global branch shares conv2, conv3 and norms with the local branch. 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); @@ -167,7 +166,7 @@ namespace WAN { blocks["injector_adain_layers." + std::to_string(i) + ".linear"] = std::make_shared(dim, dim * 2); } - // AdaLayerNorm is affine-free and uses eps=1e-5, unlike the attention norms. + // S2V AdaLayerNorm uses its own epsilon, independent of attention norms. blocks["adain_norm"] = std::make_shared(dim, 1e-5f, false); } diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index fab8db28d..d8b8dc510 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -247,10 +247,7 @@ bool read_safetensors_file(const std::string& file_path, std::string dtype = tensor_info["dtype"]; nlohmann::json shape = tensor_info["shape"]; - // ComfyUI fp8_scaled checkpoints carry a per-module scale_input - // (activation scale, which cancels when activations are f16/f32). - // Drop it here; .scale_weight -> .weight_scale aliasing happens in - // convert_tensor_name. + // ComfyUI FP8 activation scales cancel when inference uses F16/F32 activations. if (ends_with(name, ".scale_input")) { continue; } diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 53980b56d..7b7db3a09 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -527,7 +527,7 @@ SDVersion ModelLoader::get_sd_version() const { } 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 shares patch_embedding shape with T2V-14B; key presence is the only reliable signal + // 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) { diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index b4f25b55f..ecad31a53 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -1137,8 +1137,7 @@ bool StableDiffusionGGML::validate_and_load_runners() { ignore_tensors.insert("model.diffusion_model.__index_timestep_zero__"); if (audio_encoder != nullptr) { - // HF wav2vec2 inference leftovers: lm_head is a pretraining head, - // masked_spec_embed only applies to masked pretraining. + // These wav2vec2 tensors are unused during feature extraction. ignore_tensors.insert("wav2vec2.lm_head."); ignore_tensors.insert("wav2vec2.masked_spec_embed"); } @@ -1777,8 +1776,7 @@ 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"); @@ -1857,8 +1855,6 @@ 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]); diff --git a/src/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h index 5cad06f8c..40a5fc64c 100644 --- a/src/pipeline/diffusion_engine.h +++ b/src/pipeline/diffusion_engine.h @@ -66,7 +66,7 @@ class StableDiffusionGGML { std::shared_ptr first_stage_model; std::shared_ptr preview_vae; std::shared_ptr audio_vae_model; - std::shared_ptr audio_encoder; // wan2.2 s2v driving audio + std::shared_ptr audio_encoder; std::shared_ptr control_net; std::shared_ptr ip_adapter; sd::Tensor ip_adapter_tokens; diff --git a/src/pipeline/model_builders.h b/src/pipeline/model_builders.h index aeaed5767..5edc75680 100644 --- a/src/pipeline/model_builders.h +++ b/src/pipeline/model_builders.h @@ -40,7 +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; // wan2.2 s2v driving audio + std::shared_ptr audio_encoder; }; struct VAEOptions { diff --git a/src/pipeline/video.cpp b/src/pipeline/video.cpp index 397d69e50..7b9c8bde1 100644 --- a/src/pipeline/video.cpp +++ b/src/pipeline/video.cpp @@ -421,11 +421,7 @@ namespace sd::pipeline { return audio; } - // Wan2.2 S2V: build the first per-chunk audio window. Input: wav2vec2 stacked - // states [embed_dim, in_frames, num_layers]; output: [embed_dim, batch_frames, - // num_layers] as expected by WanDiffusionExtra.audio_embed, frames bucketed at - // 16 fps with zero padding past the audio end (nodes_wan.py - // get_audio_embed_bucket_fps, m=0, frame_offset=0). + // 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]; @@ -433,7 +429,6 @@ namespace sd::pipeline { if (embed_dim <= 0 || in_frames <= 0 || num_layers <= 0 || batch_frames <= 0) { return {}; } - // [embed_dim, in_frames, num_layers] -> layer-first [num_layers, in_frames, embed_dim] 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) { @@ -453,7 +448,7 @@ namespace sd::pipeline { if (buckets.empty() || plan.bucket_frames < batch_frames) { return {}; } - // Window rows [0, batch_frames): [frame, num_layers, dim] -> [dim, batch_frames, num_layers] + // 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) { @@ -1089,20 +1084,17 @@ namespace sd::pipeline { } int64_t t1 = ggml_time_ms(); if (!start_image.empty()) { - // ComfyUI WanSoundImageToVideo: the ref image is VAE-encoded and - // appended as reference_latents; the video latent itself stays - // unconstrained (no first-frame conditioning). 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); // [W', H', 1, C, 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; } - // forward_orig consumes a 4d reference latent [N*C, t_ref, H, W] + // 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], @@ -1180,9 +1172,6 @@ namespace sd::pipeline { } } if (sd->version == VERSION_WAN2_2_S2V) { - // ComfyUI: positive gets the real audio window and the ref latent; - // negative gets audio * 0 while KEEPING the same ref latent (wan is - // excluded from ref-latent img cfg). embeds.cond.c_ref_images = latents.ref_latents; if (!latents.s2v_audio_embed.empty()) { embeds.cond.c_ref_audios = {latents.s2v_audio_embed}; @@ -1843,9 +1832,7 @@ namespace sd::pipeline { sd_audio_t* generated_audio = nullptr; if (sd->version == VERSION_WAN2_2_S2V && sd_vid_gen_params->ref_audios_count > 0) { - // S2V does not generate audio; the driving track is conditioning only. - // Hand a copy back so the output container carries the same audio - // (the CLI muxes audio_out into avi/webm). + // 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) { @@ -1934,8 +1921,7 @@ namespace sd::pipeline { *frames_out = result; } if (sd->version == VERSION_WAN2_2_S2V && generated_audio != nullptr) { - // The model conditioned on the first chunk window only; keep the muxed - // track aligned with the decoded video duration. + // 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); diff --git a/src/runtime/audio_processing.cpp b/src/runtime/audio_processing.cpp index 9477d0dde..fe02b6ea2 100644 --- a/src/runtime/audio_processing.cpp +++ b/src/runtime/audio_processing.cpp @@ -7,9 +7,7 @@ namespace sd::audio { - // Polyphase FIR from torchaudio.functional.resample: sinc interpolated at output - // phases, hann-windowed over lowpass_filter_width zero crossings, anti-aliased by - // clamping the sinc argument to +-lowpass_filter_width after scaling by min(rate)*rolloff. + // Match torchaudio's Hann-windowed sinc resampler. std::vector resample_audio(const float* samples, uint64_t sample_count, uint32_t orig_sample_rate, diff --git a/src/runtime/audio_processing.h b/src/runtime/audio_processing.h index 345289c55..cdda0c463 100644 --- a/src/runtime/audio_processing.h +++ b/src/runtime/audio_processing.h @@ -12,8 +12,7 @@ namespace sd::audio { uint32_t orig_sample_rate, uint32_t target_sample_rate); - // Downmix interleaved samples to mono by averaging channels. Returns an empty - // vector on invalid input. + // 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);