feat: add Wan2.2 S2V (audio+img-to-video) support (#1925)

Co-authored-by: leejet <leejet714@gmail.com>

src/pipeline/diffusion_engine.cpp

@@ -33,12 +33,14 @@
 #include "extensions/generation_extension.h"
 #include "model/adapter/ip_adapter.hpp"
 #include "model/adapter/lora.hpp"
+#include "model/audio/wav2vec2.hpp"
 #include "model/diffusion/animatediff.hpp"
 #include "model/diffusion/control.hpp"
 #include "model/diffusion/model.hpp"
 #include "model/vae/audio_vae.hpp"
 #include "model/vae/ltx_vae.hpp"
 #include "model/vae/vae.hpp"
+#include "runtime/audio_processing.h"
 #include "runtime/denoiser.hpp"
 #include "runtime/guidance.h"
 #include "runtime/preview_interval.h"
@@ -74,6 +76,7 @@ const char* model_version_to_str[] = {
     "Wan 2.x",
     "Wan 2.2 I2V",
     "Wan 2.2 TI2V",
+    "Wan 2.2 S2V",
     "LingBot Video",
     "Qwen Image",
     "Qwen Image Layered",
@@ -136,7 +139,7 @@ StableDiffusionGGML::~StableDiffusionGGML() = default;
 
 const std::map<StableDiffusionGGML::RunnerGroup, std::set<ModelComponent>>& StableDiffusionGGML::runner_components() {
     static const std::map<RunnerGroup, std::set<ModelComponent>> components{
-        {RunnerGroup::Core, {ModelComponent::Conditioner, ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion, ModelComponent::CLIPVision, ModelComponent::IPAdapter}},
+        {RunnerGroup::Core, {ModelComponent::Conditioner, ModelComponent::Diffusion, ModelComponent::HighNoiseDiffusion, ModelComponent::CLIPVision, ModelComponent::IPAdapter, ModelComponent::AudioEncoder}},
         {RunnerGroup::VAE, {ModelComponent::VAE, ModelComponent::PreviewVAE, ModelComponent::AudioVAE}},
         {RunnerGroup::ControlNet, {ModelComponent::ControlNet}},
         {RunnerGroup::Extensions, {ModelComponent::PhotoMaker, ModelComponent::PuLID}},
@@ -804,6 +807,13 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf
         }
     }
 
+    if (strlen(SAFE_STR(sd_ctx_params->audio_encoder_path)) > 0) {
+        LOG_INFO("loading audio encoder (wav2vec2) from '%s'", sd_ctx_params->audio_encoder_path);
+        if (!model_loader.init_from_file(sd_ctx_params->audio_encoder_path, "wav2vec2.")) {
+            LOG_WARN("loading audio encoder weights from '%s' failed", sd_ctx_params->audio_encoder_path);
+        }
+    }
+
     if (strlen(SAFE_STR(sd_ctx_params->motion_module_path)) > 0) {
         LOG_INFO("loading motion module (AnimateDiff) from '%s'", sd_ctx_params->motion_module_path);
         if (!model_loader.init_from_file(sd_ctx_params->motion_module_path,
@@ -1010,6 +1020,7 @@ bool StableDiffusionGGML::build_core_runners() {
     high_noise_diffusion_model = std::move(runners.high_noise_diffusion);
     clip_vision                = std::move(runners.clip_vision);
     ip_adapter                 = std::move(runners.ip_adapter);
+    audio_encoder              = std::move(runners.audio_encoder);
 
     cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::TE));
     diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION));
@@ -1019,11 +1030,15 @@ bool StableDiffusionGGML::build_core_runners() {
     if (clip_vision) {
         clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CLIP_VISION));
     }
+    if (audio_encoder) {
+        audio_encoder->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::AUDIO_ENCODER));
+    }
     return register_runner_params(ModelComponent::Conditioner, cond_stage_model, SDBackendModule::TE) &&
            register_runner_params(ModelComponent::Diffusion, diffusion_model, SDBackendModule::DIFFUSION) &&
            register_runner_params(ModelComponent::HighNoiseDiffusion, high_noise_diffusion_model, SDBackendModule::DIFFUSION) &&
            register_runner_params(ModelComponent::CLIPVision, clip_vision, SDBackendModule::CLIP_VISION) &&
-           register_runner_params(ModelComponent::IPAdapter, ip_adapter, SDBackendModule::DIFFUSION);
+           register_runner_params(ModelComponent::IPAdapter, ip_adapter, SDBackendModule::DIFFUSION) &&
+           register_runner_params(ModelComponent::AudioEncoder, audio_encoder, SDBackendModule::AUDIO_ENCODER);
 }
 
 bool StableDiffusionGGML::build_vae_runners() {
@@ -1121,6 +1136,12 @@ bool StableDiffusionGGML::validate_and_load_runners() {
     ignore_tensors.insert("model.diffusion_model.__32x32__");
     ignore_tensors.insert("model.diffusion_model.__index_timestep_zero__");
 
+    if (audio_encoder != nullptr) {
+        // These wav2vec2 tensors are unused during feature extraction.
+        ignore_tensors.insert("wav2vec2.lm_head.");
+        ignore_tensors.insert("wav2vec2.masked_spec_embed");
+    }
+
     if (audio_vae_model) {
         if (!sd_version_is_minimax_h3(version)) {
             ignore_tensors.insert("audio_vae.encoder");
@@ -1755,6 +1776,29 @@ sd::Tensor<float> StableDiffusionGGML::get_clip_vision_output(const sd::Tensor<f
     return output;
 }
 
+// Returns 50 Hz wav2vec2 states in sd::Tensor layout: [dim, frames, layers].
+sd::Tensor<float> StableDiffusionGGML::get_audio_embedding(const sd_audio_t& audio) {
+    if (audio_encoder == nullptr) {
+        LOG_ERROR("audio encoder model is not loaded");
+        return {};
+    }
+    if (audio.data == nullptr || audio.sample_count == 0 || audio.channels == 0 || audio.sample_rate == 0) {
+        LOG_ERROR("invalid driving audio");
+        return {};
+    }
+    auto mono = sd::audio::downmix_to_mono(audio.data, audio.sample_count, audio.channels);
+    if (mono.empty()) {
+        LOG_ERROR("audio mono downmix failed");
+        return {};
+    }
+    mono = sd::audio::resample_audio(mono.data(), mono.size(), audio.sample_rate, 16000);
+    if (mono.empty()) {
+        LOG_ERROR("audio resample to 16 kHz failed");
+        return {};
+    }
+    return audio_encoder->compute(n_threads, mono);
+}
+
 void StableDiffusionGGML::compute_ip_adapter_tokens(const sd_image_t& image, float strength) {
     ip_adapter_tokens        = {};
     ip_adapter_uncond_tokens = {};
@@ -1810,6 +1854,10 @@ std::vector<float> StableDiffusionGGML::process_timesteps(const std::vector<floa
             }
         }
         return new_timesteps;
+    }
+    if (diffusion_model->get_desc() == "Wan2.2-S2V-14B") {
+        int64_t frame_count = init_latent.shape()[2];
+        return std::vector<float>(static_cast<size_t>(frame_count), timesteps[0]);
     } else {
         return timesteps;
     }
@@ -2318,7 +2366,8 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
                                                              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,

src/pipeline/diffusion_engine.h

@@ -27,6 +27,9 @@ struct LoraModel;
 struct ConditionerParams;
 struct SDCondition;
 struct RefImageParams;
+namespace Wav2Vec2 {
+    class Wav2Vec2ModelRunner;
+}
 
 extern const char* model_version_to_str[];
 
@@ -63,6 +66,7 @@ public:
     std::shared_ptr<VAE> first_stage_model;
     std::shared_ptr<VAE> preview_vae;
     std::shared_ptr<AudioVAERunner> audio_vae_model;
+    std::shared_ptr<Wav2Vec2::Wav2Vec2ModelRunner> audio_encoder;
     std::shared_ptr<ControlNet> control_net;
     std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
     sd::Tensor<float> ip_adapter_tokens;
@@ -363,6 +367,8 @@ public:
                                              int clip_skip        = -1,
                                              bool zero_out_masked = false);
 
+    sd::Tensor<float> get_audio_embedding(const sd_audio_t& audio);
+
     void compute_ip_adapter_tokens(const sd_image_t& image, float strength);
 
     std::vector<float> process_timesteps(const std::vector<float>& timesteps,

src/pipeline/generation.h

@@ -28,6 +28,7 @@ namespace sd::pipeline {
         sd::Tensor<float> denoise_mask;
         sd::Tensor<float> clip_vision_output;
         sd::Tensor<float> vace_context;
+        sd::Tensor<float> s2v_audio_embed;
         int64_t ref_image_num                  = 0;
         int64_t video_conditioning_frame_count = 0;
         int64_t video_target_frame_count       = 0;
@@ -59,7 +60,8 @@ namespace sd::pipeline {
                         const sd_vid_gen_params_t* sd_vid_gen_params,
                         sd_image_t** frames_out,
                         int* num_frames_out,
-                        sd_audio_t** audio_out);
+                        sd_audio_t** audio_out,
+                        int* fps_out);
 
     sd::Tensor<float> upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd,
                                                        const char* model_path,

src/pipeline/model_builders.cpp

@@ -8,6 +8,7 @@
 #include "core/util.h"
 #include "extensions/generation_extension.h"
 #include "model/adapter/ip_adapter.hpp"
+#include "model/audio/wav2vec2.hpp"
 #include "model/diffusion/anima.hpp"
 #include "model/diffusion/boogu.hpp"
 #include "model/diffusion/control.hpp"
@@ -234,6 +235,16 @@ namespace sd::model_builders {
                                                                                 tensor_storage_map,
                                                                                 weight_manager);
             }
+            if (version == VERSION_WAN2_2_S2V &&
+                tensor_storage_map.count("wav2vec2.encoder.layer_norm.bias") > 0) {
+                if (!ensure_backend_pair(ctx.backends, SDBackendModule::AUDIO_ENCODER)) {
+                    return false;
+                }
+                result.audio_encoder = std::make_shared<Wav2Vec2::Wav2Vec2ModelRunner>(ctx.backends.runtime_backend(SDBackendModule::AUDIO_ENCODER),
+                                                                                       tensor_storage_map,
+                                                                                       "wav2vec2.",
+                                                                                       weight_manager);
+            }
         } else if (sd_version_is_lingbot_video(version)) {
             bool enable_vision = false;
             for (const auto& [name, _] : tensor_storage_map) {

src/pipeline/model_builders.h

@@ -15,6 +15,9 @@ struct DiffusionModelRunner;
 struct VAE;
 struct AudioVAERunner;
 struct ControlNet;
+namespace Wav2Vec2 {
+    class Wav2Vec2ModelRunner;
+}
 struct GenerationExtension;
 struct GenerationExtensionInitContext;
 namespace IPAdapter {
@@ -37,6 +40,7 @@ namespace sd::model_builders {
         std::shared_ptr<DiffusionModelRunner> high_noise_diffusion;
         std::shared_ptr<FrozenCLIPVisionEmbedder> clip_vision;
         std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
+        std::shared_ptr<Wav2Vec2::Wav2Vec2ModelRunner> audio_encoder;
     };
 
     struct VAEOptions {

src/pipeline/request.cpp

@@ -161,7 +161,10 @@ namespace sd::pipeline {
         frames           = sd->align_video_frames(requested_frames);
         clip_skip        = sd_vid_gen_params->clip_skip;
         fps              = std::max(1, sd_vid_gen_params->fps);
-        if (sd_version_is_minimax_h3(sd->version) && fps != 24) {
+        if (sd->version == VERSION_WAN2_2_S2V && sd_vid_gen_params->fps != 16) {
+            LOG_WARN("Wan2.2 S2V uses 16 fps; overriding requested fps %d", sd_vid_gen_params->fps);
+            fps = 16;
+        } else if (sd_version_is_minimax_h3(sd->version) && fps != 24) {
             LOG_WARN("MiniMax-H3 uses 24 fps; overriding requested fps %d", fps);
             fps = 24;
         }

src/pipeline/video.cpp

@@ -5,6 +5,7 @@
 #include <cstdlib>
 #include <optional>
 
+#include "conditioning/wan_audio.h"
 #include "core/rng.hpp"
 #include "core/rng_philox.hpp"
 #include "diffusion_engine.h"
@@ -420,6 +421,45 @@ namespace sd::pipeline {
         return audio;
     }
 
+    // Build the first 16 fps audio window, zero-padding past the track end.
+    static sd::Tensor<float> build_s2v_audio_window(const sd::Tensor<float>& stacked, int64_t batch_frames) {
+        const int64_t embed_dim  = stacked.shape()[0];
+        const int64_t in_frames  = stacked.shape()[1];
+        const int64_t num_layers = stacked.shape()[2];
+        if (embed_dim <= 0 || in_frames <= 0 || num_layers <= 0 || batch_frames <= 0) {
+            return {};
+        }
+        std::vector<float> layer_first(static_cast<size_t>(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<size_t>(embed_dim),
+                            layer_first.data() + (static_cast<size_t>(l) * in_frames + f) * embed_dim);
+            }
+        }
+        sd::wan_audio::BucketPlan plan;
+        std::vector<float> buckets = sd::wan_audio::build_audio_buckets(layer_first.data(),
+                                                                        static_cast<int>(num_layers),
+                                                                        static_cast<int>(in_frames),
+                                                                        static_cast<int>(embed_dim),
+                                                                        static_cast<int>(batch_frames),
+                                                                        &plan);
+        if (buckets.empty() || plan.bucket_frames < batch_frames) {
+            return {};
+        }
+        // Reorder frame-major buckets into sd::Tensor's [dim, frame, layer] layout.
+        sd::Tensor<float> 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<size_t>(f) * num_layers + l) * embed_dim;
+                float* dst       = window.data() + l * embed_dim * batch_frames + f * embed_dim;
+                std::copy_n(src, static_cast<size_t>(embed_dim), dst);
+            }
+        }
+        return window;
+    }
+
     static std::optional<ImageGenerationLatents> prepare_video_generation_latents(StableDiffusionGGML* sd,
                                                                                   const sd_vid_gen_params_t* sd_vid_gen_params,
                                                                                   GenerationRequest* request) {
@@ -1033,6 +1073,53 @@ namespace sd::pipeline {
             latents.vace_context = sd::ops::concat(vace_context, mask_context, 3);  // [b, 2*c + vae_scale_factor*vae_scale_factor, t + 1 or t, h/vae_scale_factor, w/vae_scale_factor]
             int64_t t2           = ggml_time_ms();
             LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
+        } else if (sd->diffusion_model->get_desc() == "Wan2.2-S2V-14B") {
+            LOG_INFO("S2V");
+            if (!end_image.empty()) {
+                LOG_WARN("Wan2.2 S2V ignores end_image");
+            }
+            if (sd_vid_gen_params->ref_audios_count > 1) {
+                LOG_ERROR("Wan2.2 S2V supports a single driving audio track");
+                return std::nullopt;
+            }
+            int64_t t1 = ggml_time_ms();
+            if (!start_image.empty()) {
+                auto ref_img     = start_image.reshape({start_image.shape()[0],
+                                                        start_image.shape()[1],
+                                                        1,
+                                                        start_image.shape()[2],
+                                                        1});
+                auto encoded_ref = sd->encode_first_stage(ref_img);
+                if (encoded_ref.empty()) {
+                    LOG_ERROR("failed to encode S2V reference image");
+                    return std::nullopt;
+                }
+                // Wan consumes reference latents in 4D.
+                latents.ref_latents.push_back(encoded_ref.reshape({encoded_ref.shape()[0],
+                                                                   encoded_ref.shape()[1],
+                                                                   encoded_ref.shape()[2],
+                                                                   encoded_ref.shape()[3]}));
+            }
+            if (sd_vid_gen_params->ref_audios_count == 1) {
+                if (sd->audio_encoder == nullptr) {
+                    LOG_ERROR("S2V audio conditioning requires --audio-encoder (wav2vec2)");
+                    return std::nullopt;
+                }
+                auto stacked = sd->get_audio_embedding(sd_vid_gen_params->ref_audios[0]);
+                if (stacked.empty()) {
+                    LOG_ERROR("failed to compute wav2vec2 embedding for driving audio");
+                    return std::nullopt;
+                }
+                int64_t latent_t        = sd->video_frames_to_latent_frames(request->frames);
+                int64_t batch_frames    = latent_t * 4;
+                latents.s2v_audio_embed = build_s2v_audio_window(stacked, batch_frames);
+                if (latents.s2v_audio_embed.empty()) {
+                    LOG_ERROR("failed to build S2V audio window");
+                    return std::nullopt;
+                }
+            }
+            int64_t t2 = ggml_time_ms();
+            LOG_INFO("s2v conditioning prepared, taking %" PRId64 " ms", t2 - t1);
         }
 
         if (latents.init_latent.empty()) {
@@ -1084,6 +1171,12 @@ namespace sd::pipeline {
                     latents.keyframe_indices);
             }
         }
+        if (sd->version == VERSION_WAN2_2_S2V) {
+            embeds.cond.c_ref_images = latents.ref_latents;
+            if (!latents.s2v_audio_embed.empty()) {
+                embeds.cond.c_ref_audios = {latents.s2v_audio_embed};
+            }
+        }
         if (request.use_uncond) {
             condition_params.text  = request.negative_prompt;
             embeds.uncond          = sd->cond_stage_model->get_learned_condition(sd->n_threads,
@@ -1096,6 +1189,12 @@ namespace sd::pipeline {
                 embeds.uncond.c_reference_blocks = latents.minimax_reference_blocks;
                 embeds.uncond.c_position_ids     = embeds.cond.c_position_ids;
             }
+            if (sd->version == VERSION_WAN2_2_S2V) {
+                embeds.uncond.c_ref_images = latents.ref_latents;
+                if (!latents.s2v_audio_embed.empty()) {
+                    embeds.uncond.c_ref_audios = {sd::Tensor<float>::zeros_like(latents.s2v_audio_embed)};
+                }
+            }
         }
 
         int64_t t1 = ggml_time_ms();
@@ -1422,10 +1521,14 @@ namespace sd::pipeline {
                         const sd_vid_gen_params_t* sd_vid_gen_params,
                         sd_image_t** frames_out,
                         int* num_frames_out,
-                        sd_audio_t** audio_out) {
+                        sd_audio_t** audio_out,
+                        int* fps_out) {
         if (sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd->version)) {
             LOG_INFO("AnimateDiff dispatch: %d frames, %dx%d",
                      sd_vid_gen_params->video_frames, sd_vid_gen_params->width, sd_vid_gen_params->height);
+            if (fps_out != nullptr) {
+                *fps_out = std::max(1, sd_vid_gen_params->fps);
+            }
             return generate_animatediff_video(sd, sd_vid_gen_params, frames_out, num_frames_out);
         }
 
@@ -1437,6 +1540,9 @@ namespace sd::pipeline {
         sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params;
         sd->apply_circular_axes(sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y);
         GenerationRequest request(sd, sd_vid_gen_params);
+        if (fps_out != nullptr) {
+            *fps_out = request.fps;
+        }
         bool latent_upscale_enabled     = request.hires.enabled;
         GenerationRequest hires_request = request;
         if (latent_upscale_enabled) {
@@ -1725,6 +1831,33 @@ namespace sd::pipeline {
         LOG_INFO("generating latent video completed, taking %.2fs", (latent_end - latent_start) * 1.0f / 1000);
 
         sd_audio_t* generated_audio = nullptr;
+        if (sd->version == VERSION_WAN2_2_S2V && sd_vid_gen_params->ref_audios_count > 0) {
+            // Return the driving track for muxing with the generated video.
+            const sd_audio_t& driving = sd_vid_gen_params->ref_audios[0];
+            generated_audio           = (sd_audio_t*)malloc(sizeof(sd_audio_t));
+            if (generated_audio != nullptr) {
+                generated_audio->sample_rate  = driving.sample_rate;
+                generated_audio->channels     = driving.channels;
+                generated_audio->sample_count = driving.sample_count;
+                generated_audio->data         = (float*)malloc(sizeof(float) * driving.sample_count * driving.channels);
+                if (generated_audio->data == nullptr) {
+                    free(generated_audio);
+                    generated_audio = nullptr;
+                } else {
+                    memcpy(generated_audio->data,
+                           driving.data,
+                           sizeof(float) * driving.sample_count * driving.channels);
+                }
+            }
+            if (generated_audio != nullptr) {
+                LOG_DEBUG("s2v output audio: %u Hz, %u channels, %llu samples",
+                          generated_audio->sample_rate,
+                          generated_audio->channels,
+                          (unsigned long long)generated_audio->sample_count);
+            } else {
+                LOG_DEBUG("s2v output audio copy failed (out of memory)");
+            }
+        }
         if ((sd_version_is_ltxav(sd->version) || sd_version_is_minimax_h3(sd->version)) &&
             latents.audio_length > 0 &&
             sd->audio_vae_model != nullptr) {
@@ -1774,6 +1907,7 @@ namespace sd::pipeline {
             return false;
         }
         auto result = decode_video_outputs(sd, latent_upscale_enabled ? hires_request : request, final_latent, num_frames_out);
+        LOG_DEBUG("decode_video_outputs returned %s", result == nullptr ? "nullptr (failed)" : "frames");
         if (result == nullptr) {
             free_sd_audio(generated_audio);
             return false;
@@ -1786,6 +1920,20 @@ namespace sd::pipeline {
         if (frames_out != nullptr) {
             *frames_out = result;
         }
+        if (sd->version == VERSION_WAN2_2_S2V && generated_audio != nullptr) {
+            // Limit the driving track to the generated video's duration.
+            int fps               = request.fps;
+            uint64_t video_frames = num_frames_out != nullptr ? (uint64_t)*num_frames_out : 0;
+            uint64_t want_samples = (uint64_t)((double)video_frames / fps * generated_audio->sample_rate);
+            LOG_DEBUG("s2v audio truncate: %llu samples -> %llu (video %llu frames @ %d fps)",
+                      (unsigned long long)generated_audio->sample_count,
+                      (unsigned long long)want_samples,
+                      (unsigned long long)video_frames,
+                      fps);
+            if (want_samples > 0 && want_samples < generated_audio->sample_count) {
+                generated_audio->sample_count = want_samples;
+            }
+        }
         if (audio_out != nullptr) {
             *audio_out = generated_audio;
         } else {