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 35812157c..fb790bb3e 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -412,6 +412,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)", @@ -817,6 +822,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" @@ -879,6 +885,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(); @@ -1428,6 +1435,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; @@ -1617,6 +1634,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 34b4a013b..47a7ec5c0 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -130,6 +130,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 9ebd0b5a1..62d58ffba 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 d5cda6d9a..98c7704b7 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -206,6 +206,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 a83166438..74ac1c425 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 d5498e8f4..e3cda0efc 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..fbe5b8da1 --- /dev/null +++ b/src/model/audio/wav2vec2.hpp @@ -0,0 +1,742 @@ +#ifndef __SD_MODEL_AUDIO_WAV2VEC2_HPP__ +#define __SD_MODEL_AUDIO_WAV2VEC2_HPP__ + +#include "core/ggml_extend.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, true, 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 c36cc7e9c..15a6963f0 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -805,7 +805,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, @@ -814,8 +815,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 070ca53d4..7d37d1b7b 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 9a907dcfa..8de841ccc 100644 --- a/src/model/diffusion/wan.hpp +++ b/src/model/diffusion/wan.hpp @@ -1,10 +1,12 @@ #ifndef __SD_MODEL_DIFFUSION_WAN_HPP__ #define __SD_MODEL_DIFFUSION_WAN_HPP__ +#include #include #include #include +#include "model/audio/wav2vec2.hpp" #include "model/common/block.hpp" #include "model/common/rope.hpp" #include "model/diffusion/flux.hpp" @@ -30,11 +32,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; @@ -71,6 +79,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; } @@ -525,10 +537,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) @@ -551,7 +799,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, @@ -592,6 +840,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, @@ -639,18 +896,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"]); @@ -667,6 +930,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); @@ -711,6 +1009,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)]); @@ -728,6 +1031,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"); @@ -744,11 +1056,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,] @@ -776,7 +1090,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] @@ -836,7 +1155,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; @@ -888,7 +1210,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); @@ -898,16 +1222,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_DEBUG("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); @@ -930,7 +1271,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); @@ -945,9 +1289,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, false, false), x.dim()); @@ -958,6 +1304,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, @@ -966,7 +1318,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_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index 1d1269488..060a26e4f 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -250,6 +250,14 @@ bool read_safetensors_file(const std::string& file_path, continue; } + // 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 dc2bb4bf0..3980fe533 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -398,6 +398,7 @@ SDVersion ModelLoader::get_sd_version() { 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; @@ -484,6 +485,11 @@ SDVersion ModelLoader::get_sd_version() { 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; } @@ -547,6 +553,9 @@ SDVersion ModelLoader::get_sd_version() { } if (is_wan) { LOG_DEBUG("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/stable-diffusion.cpp b/src/stable-diffusion.cpp index 54ff82928..1875b7e69 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -19,11 +19,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/anima.hpp" #include "model/diffusion/animatediff.hpp" #include "model/diffusion/boogu.hpp" @@ -102,6 +104,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", @@ -233,6 +236,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 control_net; std::shared_ptr ip_adapter; sd::Tensor ip_adapter_tokens; @@ -821,6 +825,13 @@ class StableDiffusionGGML { } } + 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, @@ -1198,6 +1209,22 @@ class StableDiffusionGGML { return false; } } + if (version == VERSION_WAN2_2_S2V && + tensor_storage_map.count("wav2vec2.encoder.layer_norm.bias") > 0) { + if (!ensure_backend_pair(SDBackendModule::AUDIO_ENCODER)) { + return false; + } + audio_encoder = std::make_shared(backend_for(SDBackendModule::AUDIO_ENCODER), + tensor_storage_map, + "wav2vec2.", + model_manager); + audio_encoder->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::AUDIO_ENCODER)); + if (!register_runner_params("audio encoder", + audio_encoder, + SDBackendModule::AUDIO_ENCODER)) { + return false; + } + } } else if (sd_version_is_lingbot_video(version)) { bool enable_vision = false; for (const auto& [name, _] : tensor_storage_map) { @@ -1656,6 +1683,13 @@ class StableDiffusionGGML { 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"); @@ -2188,6 +2222,30 @@ class StableDiffusionGGML { return output; } + // Driving audio for Wan2.2 S2V: downmix to mono, resample to 16 kHz, then + // wav2vec2 stacked states [embed_dim, frames, num_layers + 1] at 50 Hz. + sd::Tensor 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 compute_ip_adapter_tokens(const sd_image_t& image, float strength) { ip_adapter_tokens = {}; ip_adapter_uncond_tokens = {}; @@ -2243,6 +2301,12 @@ class StableDiffusionGGML { } } return new_timesteps; + } + // Wan2.2 S2V runs per-frame timesteps (ComfyUI repeats t across frames); + // reference tokens are zeroed in-model. + if (diffusion_model->get_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; } @@ -2759,7 +2823,8 @@ class StableDiffusionGGML { condition.c_t5_weights.empty() ? nullptr : &condition.c_t5_weights}; } else if (sd_version_is_wan(version)) { diffusion_params.extra = WanDiffusionExtra{vace_context.empty() ? nullptr : &vace_context, - vace_strength}; + vace_strength, + condition.c_ref_audios.empty() ? nullptr : &condition.c_ref_audios[0]}; } else if (sd_version_is_hunyuan_video(version)) { diffusion_params.extra = HunyuanVideoDiffusionExtra{ &guidance_tensor, @@ -3609,6 +3674,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" @@ -3643,6 +3709,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), @@ -4424,6 +4491,7 @@ struct ImageGenerationLatents { 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; @@ -5897,6 +5965,50 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx, return true; } +// 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(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params, GenerationRequest* request) { @@ -6510,6 +6622,56 @@ static std::optional prepare_video_generation_latents(sd 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_ctx->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_ctx->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_ctx->sd->audio_encoder == nullptr) { + LOG_ERROR("S2V audio conditioning requires --audio-encoder (wav2vec2)"); + return std::nullopt; + } + auto stacked = sd_ctx->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_ctx->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()) { @@ -6561,6 +6723,15 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx, latents.keyframe_indices); } } + if (sd_ctx->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_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, @@ -6573,6 +6744,12 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx, embeds.uncond.c_reference_blocks = latents.minimax_reference_blocks; embeds.uncond.c_position_ids = embeds.cond.c_position_ids; } + if (sd_ctx->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(); @@ -7185,6 +7362,35 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, LOG_INFO("generating latent video completed, taking %.2fs", (latent_end - latent_start) * 1.0f / 1000); sd_audio_t* generated_audio = nullptr; + if (sd_ctx->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_ctx->sd->version) || sd_version_is_minimax_h3(sd_ctx->sd->version)) && latents.audio_length > 0 && sd_ctx->sd->audio_vae_model != nullptr) { @@ -7234,6 +7440,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, return false; } auto result = decode_video_outputs(sd_ctx, 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; @@ -7246,6 +7453,21 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, if (frames_out != nullptr) { *frames_out = result; } + if (sd_ctx->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 {