DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 03:24:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: leejet <leejet714@gmail.com>
This commit is contained in:
Georgeandleejet authored and GitHub committed 2026-09-13 23:33:50 +08:00
1 parent 4a7da26b73
commit 0bd72f075a
32 files changed
+1533 -57
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/audio/wav2vec2.hpp
+413
View File
@@ -0,0 +1,413 @@
#ifndef __SD_MODEL_AUDIO_WAV2VEC2_HPP__
#define __SD_MODEL_AUDIO_WAV2VEC2_HPP__
#include <algorithm>
#include <cinttypes>
#include <cmath>
#include <cstdio>
#include <map>
#include <memory>
#include <string>
#include <vector>
#include "core/ggml_extend.h"
#include "core/ggml_runner.h"
#include "model.h"
#include "model/common/ggml_block.hpp"
namespace Wav2Vec2 {
struct Wav2Vec2Config {
int64_t embed_dim = 1024;
int64_t conv_dim = 512;
int num_heads = 16;
int num_layers = 24;
std::string feat_extract_norm = "layer";
bool conv_bias = true;
bool do_normalize = true;
bool do_stable_layer_norm = true;
static Wav2Vec2Config detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
Wav2Vec2Config config;
auto it = tensor_storage_map.find(prefix + "encoder.layer_norm.bias");
if (it == tensor_storage_map.end()) {
LOG_WARN("wav2vec2: %sencoder.layer_norm.bias not found, using large defaults", prefix.c_str());
return config;
}
config.embed_dim = it->second.ne[0];
if (config.embed_dim == 1024) {
config.embed_dim = 1024;
config.num_heads = 16;
config.num_layers = 24;
config.feat_extract_norm = "layer";
config.conv_bias = true;
config.do_normalize = true;
config.do_stable_layer_norm = true;
} else if (config.embed_dim == 768) {
config.embed_dim = 768;
config.num_heads = 12;
config.num_layers = 12;
config.feat_extract_norm = "group";
config.conv_bias = false;
config.do_normalize = false;
config.do_stable_layer_norm = false;
} else {
LOG_WARN("wav2vec2: unsupported embed_dim %" PRId64 ", using large defaults", config.embed_dim);
config.embed_dim = 1024;
}
return config;
}
};
struct Wav2Vec2NoLayerNormConvLayer : public UnaryBlock {
Wav2Vec2NoLayerNormConvLayer(int64_t in_channels, int64_t out_channels, int kernel_size, int stride, bool bias) {
blocks["conv"] = std::make_shared<Conv1d>(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto conv = std::dynamic_pointer_cast<Conv1d>(blocks["conv"]);
x = conv->forward(ctx, x);
return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
}
};
struct Wav2Vec2LayerNormConvLayer : public UnaryBlock {
Wav2Vec2LayerNormConvLayer(int64_t in_channels, int64_t out_channels, int kernel_size, int stride, bool bias) {
blocks["conv"] = std::make_shared<Conv1d>(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias, true);
blocks["layer_norm"] = std::make_shared<LayerNorm>(out_channels);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto conv = std::dynamic_pointer_cast<Conv1d>(blocks["conv"]);
auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
x = conv->forward(ctx, x);
// LayerNorm normalizes channels: [N, C, L] -> [N, L, C].
x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
x = layer_norm->forward(ctx, x);
x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
}
};
struct Wav2Vec2GroupNormConvLayer : public UnaryBlock {
Wav2Vec2GroupNormConvLayer(int64_t in_channels, int64_t out_channels, int kernel_size, int stride, bool bias) {
blocks["conv"] = std::make_shared<Conv1d>(in_channels, out_channels, kernel_size, stride, 0, 1, 1, bias, true);
blocks["layer_norm"] = std::make_shared<GroupNorm>((int)out_channels, out_channels, 1e-05f);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto conv = std::dynamic_pointer_cast<Conv1d>(blocks["conv"]);
auto layer_norm = std::dynamic_pointer_cast<GroupNorm>(blocks["layer_norm"]);
x = conv->forward(ctx, x);
// ggml GroupNorm needs [N, C, H, W], with H=1 for audio.
x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]);
x = layer_norm->forward(ctx, x);
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[2], x->ne[3]);
return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x));
}
};
struct Wav2Vec2FeatureEncoder : public UnaryBlock {
Wav2Vec2FeatureEncoder(const Wav2Vec2Config& config) {
GGML_ASSERT(config.feat_extract_norm == "layer" || config.feat_extract_norm == "group");
const int kernels[7] = {10, 3, 3, 3, 3, 2, 2};
const int strides[7] = {5, 2, 2, 2, 2, 2, 2};
int64_t in_channels = 1;
for (int i = 0; i < 7; ++i) {
const std::string name = "conv_layers." + std::to_string(i);
if (config.feat_extract_norm == "layer") {
blocks[name] = std::make_shared<Wav2Vec2LayerNormConvLayer>(in_channels, config.conv_dim, kernels[i], strides[i], config.conv_bias);
} else if (i == 0) {
blocks[name] = std::make_shared<Wav2Vec2GroupNormConvLayer>(in_channels, config.conv_dim, kernels[i], strides[i], config.conv_bias);
} else {
blocks[name] = std::make_shared<Wav2Vec2NoLayerNormConvLayer>(in_channels, config.conv_dim, kernels[i], strides[i], config.conv_bias);
}
in_channels = config.conv_dim;
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
for (int i = 0; i < 7; ++i) {
auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv_layers." + std::to_string(i)]);
x = conv->forward(ctx, x);
}
return ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
}
};
struct Wav2Vec2FeatureProjection : public UnaryBlock {
Wav2Vec2FeatureProjection(const Wav2Vec2Config& config) {
blocks["layer_norm"] = std::make_shared<LayerNorm>(config.conv_dim);
blocks["projection"] = std::make_shared<Linear>(config.conv_dim, config.embed_dim);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto ln = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
auto projection = std::dynamic_pointer_cast<Linear>(blocks["projection"]);
x = ln->forward(ctx, x);
x = projection->forward(ctx, x);
return x;
}
};
class Wav2Vec2PositionalConvEmbedding : public UnaryBlock {
private:
int64_t embed_dim_;
static constexpr int groups_ = 16;
static constexpr int kernel_size_ = 128;
std::string weight_g_name_;
std::string weight_v_name_;
ggml_tensor* weight(GGMLRunnerContext* ctx) {
auto g = params[weight_g_name_];
auto v = ggml_cast(ctx->ggml_ctx, params[weight_v_name_], GGML_TYPE_F32);
auto squared = ggml_mul(ctx->ggml_ctx, v, v);
// PyTorch weight_norm(dim=2) reduces both channel axes, retaining each kernel tap.
squared = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, squared, 2, 0, 1, 3));
squared = ggml_reshape_2d(ctx->ggml_ctx, squared, embed_dim_ / groups_ * embed_dim_, kernel_size_);
auto norm = ggml_sqrt(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, squared));
norm = ggml_reshape_3d(ctx->ggml_ctx, norm, kernel_size_, 1, 1);
return ggml_mul(ctx->ggml_ctx, v, ggml_div(ctx->ggml_ctx, g, norm));
}
public:
Wav2Vec2PositionalConvEmbedding(const Wav2Vec2Config& config)
: embed_dim_(config.embed_dim) {
GGML_ASSERT(embed_dim_ > 0 && embed_dim_ % groups_ == 0);
}
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
bool legacy = tensor_storage_map.count(prefix + "conv.weight_g") > 0;
weight_g_name_ = legacy ? "conv.weight_g" : "conv.parametrizations.weight.original0";
weight_v_name_ = legacy ? "conv.weight_v" : "conv.parametrizations.weight.original1";
auto g = tensor_storage_map.find(prefix + weight_g_name_);
auto v = tensor_storage_map.find(prefix + weight_v_name_);
GGML_ASSERT(g != tensor_storage_map.end() && v != tensor_storage_map.end());
GGML_ASSERT(g->second.ne[0] == kernel_size_ && g->second.ne[1] == 1 && g->second.ne[2] == 1 && g->second.ne[3] == 1);
GGML_ASSERT(v->second.ne[0] == kernel_size_ && v->second.ne[1] == embed_dim_ / groups_ && v->second.ne[2] == embed_dim_ && v->second.ne[3] == 1);
params[weight_g_name_] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kernel_size_, 1, 1);
params[weight_v_name_] = ggml_new_tensor_3d(ctx, get_type(prefix + weight_v_name_, tensor_storage_map, GGML_TYPE_F16),
kernel_size_, embed_dim_ / groups_, embed_dim_);
if (tensor_storage_map.count(prefix + "conv.bias") > 0) {
params["conv.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, embed_dim_);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto w = weight(ctx);
auto b = params.count("conv.bias") > 0 ? params["conv.bias"] : nullptr;
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
x = ggml_ext_conv_1d(ctx->ggml_ctx, x, w, b, 1, kernel_size_ / 2, 1, groups_, true);
// Apply GELU out of place before cropping to keep graph buffer reuse safe.
x = ggml_gelu_erf(ctx->ggml_ctx, x);
x = ggml_view_3d(ctx->ggml_ctx, x, x->ne[0] - 1, x->ne[1], x->ne[2], x->nb[1], x->nb[2], 0);
return ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
}
};
struct Wav2Vec2FeedForward : public UnaryBlock {
Wav2Vec2FeedForward(const Wav2Vec2Config& config) {
blocks["intermediate_dense"] = std::make_shared<Linear>(config.embed_dim, config.embed_dim * 4);
blocks["output_dense"] = std::make_shared<Linear>(config.embed_dim * 4, config.embed_dim);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto intermediate_dense = std::dynamic_pointer_cast<Linear>(blocks["intermediate_dense"]);
auto output_dense = std::dynamic_pointer_cast<Linear>(blocks["output_dense"]);
x = intermediate_dense->forward(ctx, x);
x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
x = output_dense->forward(ctx, x);
return x;
}
};
struct Wav2Vec2EncoderLayer : public UnaryBlock {
bool do_stable_layer_norm;
Wav2Vec2EncoderLayer(const Wav2Vec2Config& config)
: do_stable_layer_norm(config.do_stable_layer_norm) {
blocks["attention"] = std::make_shared<MultiheadAttention>(config.embed_dim, config.num_heads, true, true);
blocks["layer_norm"] = std::make_shared<LayerNorm>(config.embed_dim);
blocks["feed_forward"] = std::make_shared<Wav2Vec2FeedForward>(config);
blocks["final_layer_norm"] = std::make_shared<LayerNorm>(config.embed_dim);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto attention = std::dynamic_pointer_cast<MultiheadAttention>(blocks["attention"]);
auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
auto feed_forward = std::dynamic_pointer_cast<Wav2Vec2FeedForward>(blocks["feed_forward"]);
auto final_layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer_norm"]);
ggml_tensor* residual = x;
if (do_stable_layer_norm) {
x = layer_norm->forward(ctx, x);
x = attention->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, residual, x);
x = ggml_add(ctx->ggml_ctx, x, feed_forward->forward(ctx, final_layer_norm->forward(ctx, x)));
} else {
x = attention->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, residual, x);
x = layer_norm->forward(ctx, x);
x = final_layer_norm->forward(ctx, ggml_add(ctx->ggml_ctx, x, feed_forward->forward(ctx, x)));
}
return x;
}
};
struct Wav2Vec2Encoder : public GGMLBlock {
int num_layers;
bool do_stable_layer_norm;
Wav2Vec2Encoder(const Wav2Vec2Config& config)
: num_layers(config.num_layers), do_stable_layer_norm(config.do_stable_layer_norm) {
blocks["pos_conv_embed"] = std::make_shared<Wav2Vec2PositionalConvEmbedding>(config);
for (int i = 0; i < config.num_layers; ++i) {
blocks["layers." + std::to_string(i)] = std::make_shared<Wav2Vec2EncoderLayer>(config);
}
blocks["layer_norm"] = std::make_shared<LayerNorm>(config.embed_dim);
}
// For N == 1, all_layers stacks pre-layer states and the final state as [embed_dim, L, num_layers + 1].
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor** all_layers = nullptr) {
auto pos_conv_embed = std::dynamic_pointer_cast<Wav2Vec2PositionalConvEmbedding>(blocks["pos_conv_embed"]);
auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
std::vector<ggml_tensor*> 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<Wav2Vec2EncoderLayer>(blocks["layers." + std::to_string(i)]);
x = layer->forward(ctx, x);
}
if (do_stable_layer_norm) {
x = layer_norm->forward(ctx, x);
}
if (all_layers != nullptr) {
collected.push_back(x);
ggml_tensor* stack = collected[0];
for (size_t i = 1; i < collected.size(); ++i) {
stack = ggml_concat(ctx->ggml_ctx, stack, collected[i], 2);
}
*all_layers = stack;
}
return x;
}
};
struct Wav2Vec2Model : public GGMLBlock {
Wav2Vec2Config config;
Wav2Vec2Model() = default;
Wav2Vec2Model(const Wav2Vec2Config& config_)
: config(config_) {
blocks["feature_extractor"] = std::make_shared<Wav2Vec2FeatureEncoder>(config);
blocks["feature_projection"] = std::make_shared<Wav2Vec2FeatureProjection>(config);
blocks["encoder"] = std::make_shared<Wav2Vec2Encoder>(config);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor** all_layers = nullptr) {
auto feature_extractor = std::dynamic_pointer_cast<Wav2Vec2FeatureEncoder>(blocks["feature_extractor"]);
auto feature_projection = std::dynamic_pointer_cast<Wav2Vec2FeatureProjection>(blocks["feature_projection"]);
auto encoder = std::dynamic_pointer_cast<Wav2Vec2Encoder>(blocks["encoder"]);
x = feature_extractor->forward(ctx, x);
x = feature_projection->forward(ctx, x);
x = encoder->forward(ctx, x, all_layers);
return x;
}
};
class Wav2Vec2ModelRunner : public GGMLRunner {
private:
Wav2Vec2Config config;
public:
Wav2Vec2Model model;
std::string weight_prefix;
Wav2Vec2ModelRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "wav2vec2.",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager),
config(Wav2Vec2Config::detect_from_weights(tensor_storage_map, prefix)),
model(config),
weight_prefix(prefix) {
// GGMLBlock appends its own separator; loader prefixes already include one.
std::string block_prefix = weight_prefix;
if (!block_prefix.empty() && block_prefix.back() == '.') {
block_prefix.pop_back();
}
model.init(params_ctx, tensor_storage_map, block_prefix);
LOG_INFO("%s", get_desc().c_str());
}
std::string get_desc() override {
return "wav2vec2";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
std::string block_prefix = weight_prefix;
if (!block_prefix.empty() && block_prefix.back() == '.') {
block_prefix.pop_back();
}
model.get_param_tensors(tensors, block_prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& waveform_tensor) {
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
ggml_tensor* waveform = make_input(waveform_tensor);
auto runner_ctx = get_context();
ggml_tensor* all_layers = nullptr;
model.forward(&runner_ctx, waveform, &all_layers);
GGML_ASSERT(all_layers != nullptr);
ggml_build_forward_expand(gf, all_layers);
return gf;
}
sd::Tensor<float> compute(const int n_threads, const std::vector<float>& mono_waveform) {
GGML_ASSERT(!mono_waveform.empty());
const int64_t num_samples = (int64_t)mono_waveform.size();
sd::Tensor<float> waveform({num_samples, 1, 1});
std::copy(mono_waveform.begin(), mono_waveform.end(), waveform.data());
normalize(waveform.data(), num_samples);
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(waveform);
};
return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true));
}
private:
static void normalize(float* x, int64_t n) {
double mean = 0.0;
for (int64_t i = 0; i < n; ++i) {
mean += x[i];
}
mean /= n;
double var = 0.0;
for (int64_t i = 0; i < n; ++i) {
const double d = x[i] - mean;
var += d * d;
}
var /= n;
const float scale = (float)(1.0 / std::sqrt(var + 1e-7));
for (int64_t i = 0; i < n; ++i) {
x[i] = (float)((x[i] - mean) * scale);
}
}
};
} // namespace Wav2Vec2
#endif // __SD_MODEL_AUDIO_WAV2VEC2_HPP__
src/model/common/ggml_block.hpp
+56 -1
View File
@@ -368,6 +368,61 @@ public:
}
};
class Conv1d : public UnaryBlock {
protected:
int64_t in_channels;
int64_t out_channels;
int64_t groups;
int kernel_size;
int stride;
int padding;
int dilation;
bool bias;
bool force_prec_f32;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F16);
params["weight"] = ggml_new_tensor_3d(ctx, wtype, kernel_size, in_channels / groups, out_channels);
if (bias) {
params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels);
}
}
public:
Conv1d(int64_t in_channels,
int64_t out_channels,
int kernel_size,
int stride = 1,
int padding = 0,
int dilation = 1,
int64_t groups = 1,
bool bias = true,
bool force_prec_f32 = false)
: in_channels(in_channels),
out_channels(out_channels),
groups(groups),
kernel_size(kernel_size),
stride(stride),
padding(padding),
dilation(dilation),
bias(bias),
force_prec_f32(force_prec_f32) {
GGML_ASSERT(in_channels > 0 && out_channels > 0 && groups > 0);
GGML_ASSERT(in_channels % groups == 0 && out_channels % groups == 0);
GGML_ASSERT(kernel_size > 0 && stride > 0 && padding >= 0 && dilation > 0);
}
std::string get_desc() override {
return "Conv1d";
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
GGML_ASSERT(x->ne[1] == in_channels);
return ggml_ext_conv_1d(ctx->ggml_ctx, x, params["weight"], bias ? params["bias"] : nullptr,
stride, padding, dilation, groups, force_prec_f32);
}
};
class Conv2d : public UnaryBlock {
protected:
int64_t in_channels;
@@ -766,7 +821,7 @@ public:
b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
}
}
return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups);
return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups, eps);
}
};
src/model/common/rope.hpp
+3 -2
View File
@@ -818,8 +818,9 @@ namespace Rope {
int pw,
int bs,
int theta,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_vid_ids(t, h, w, pt, ph, pw, bs);
const std::vector<int>& axes_dim,
int t_offset = 0) {
std::vector<std::vector<float>> ids = gen_vid_ids(t, h, w, pt, ph, pw, bs, t_offset);
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
}
src/model/diffusion/model.hpp
+2
View File
@@ -69,6 +69,8 @@ struct AnimaDiffusionExtra {
struct WanDiffusionExtra {
const sd::Tensor<float>* vace_context = nullptr;
float vace_strength = 1.f;
// S2V audio, sd::Tensor layout: [dim, T_latent*4, layers].
const sd::Tensor<float>* audio_embed = nullptr;
};
struct HiDreamO1DiffusionExtra {
src/model/diffusion/wan.hpp
+158 -32
View File
@@ -1,6 +1,7 @@
#ifndef __SD_MODEL_DIFFUSION_WAN_HPP__
#define __SD_MODEL_DIFFUSION_WAN_HPP__
#include <algorithm>
#include <cinttypes>
#include <map>
#include <memory>
@@ -33,11 +34,16 @@ namespace WAN {
int vace_layers = 0;
int64_t vace_in_dim = 96;
std::map<int, int> vace_layers_mapping = {};
bool qk_norm = true;
bool cross_attn_norm = true;
float eps = 1e-6f;
int64_t flf_pos_embed_token_number = 0;
int theta = 10000;
int64_t audio_dim = 1024;
int num_audio_token = 4; // excludes the learned padding token
std::vector<int> audio_inject_layers = {};
std::map<int, int> 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<int> axes_dim = {44, 42, 42};
int64_t axes_dim_sum = 128;
@@ -74,6 +80,10 @@ namespace WAN {
if (name.find("img_emb") != std::string::npos) {
config.model_type = "i2v";
}
if (name.find("audio_injector") != std::string::npos || name.find("casual_audio_encoder") != std::string::npos) {
config.model_type = "s2v";
config.audio_inject_layers = {0, 4, 8, 12, 16, 20, 24, 27, 30, 33, 36, 39};
}
if (name.find("img_emb.emb_pos") != std::string::npos) {
config.flf_pos_embed_token_number = 514;
}
@@ -265,6 +275,13 @@ namespace WAN {
}
};
} // namespace WAN
// Audio injection reuses WanT2VCrossAttention defined above.
#include "model/diffusion/wan_audio.hpp"
namespace WAN {
static ggml_tensor* modulate_add(ggml_context* ctx, ggml_tensor* x, ggml_tensor* e) {
// x: [N, n_token, dim]
// e: [N, 1, dim] or [N, T, 1, dim]
@@ -532,6 +549,13 @@ namespace WAN {
protected:
WanConfig config;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
if (config.model_type == "s2v") {
enum ggml_type wtype = GGML_TYPE_F32; // elementwise add vs F32 activations
params["trainable_cond_mask.weight"] = ggml_new_tensor_2d(ctx, wtype, config.dim, 3);
}
}
public:
Wan() {}
Wan(WanConfig config)
@@ -554,7 +578,7 @@ namespace WAN {
// blocks
for (int i = 0; i < config.num_layers; i++) {
auto block = std::shared_ptr<GGMLBlock>(new WanAttentionBlock(config.model_type == "t2v",
auto block = std::shared_ptr<GGMLBlock>(new WanAttentionBlock(config.model_type != "i2v",
config.dim,
config.ffn_dim,
config.num_heads,
@@ -595,6 +619,14 @@ namespace WAN {
blocks["vace_patch_embedding"] = std::shared_ptr<GGMLBlock>(new Conv3d(config.vace_in_dim, config.dim, config.patch_size, config.patch_size));
}
if (config.model_type == "s2v") {
blocks["casual_audio_encoder"] = std::make_shared<WanCausalAudioEncoder>(config.audio_dim, config.dim, config.num_audio_token);
blocks["audio_injector"] = std::make_shared<WanAudioInjector>(config.dim, config.num_heads, (int)config.audio_inject_layers.size(), config.qk_norm, config.eps);
for (size_t i = 0; i < config.audio_inject_layers.size(); i++) {
config.audio_inject_mapping[config.audio_inject_layers[i]] = (int)i;
}
}
}
ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
@@ -642,18 +674,24 @@ namespace WAN {
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
ggml_tensor* clip_fea = nullptr,
ggml_tensor* vace_context = nullptr,
float vace_strength = 1.f,
int64_t N = 1) {
ggml_tensor* clip_fea = nullptr,
ggml_tensor* vace_context = nullptr,
float vace_strength = 1.f,
int64_t N = 1,
ggml_tensor* audio_embed = nullptr,
ggml_tensor* reference_latent = nullptr) {
// x: [N*C, T, H, W], C => in_dim
// vace_context: [N*vace_in_dim, T, H, W]
// timestep: [N,] or [T]
// context: [N, L, text_dim]
// return: [N, t_len*h_len*w_len, out_dim*pt*ph*pw]
// audio_embed: [layers, T*4, audio_dim]
// reference_latent: [N*C, T_ref, H, W]
// return: [N, (t_len [+ t_ref_len]) * h_len*w_len, out_dim*pt*ph*pw]
GGML_ASSERT(N == 1);
int64_t T = x->ne[2];
auto patch_embedding = std::dynamic_pointer_cast<Conv3d>(blocks["patch_embedding"]);
auto text_embedding_0 = std::dynamic_pointer_cast<Linear>(blocks["text_embedding.0"]);
@@ -670,6 +708,40 @@ namespace WAN {
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3] / N, N); // [N, dim, t_len*h_len*w_len]
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim]
ggml_tensor* audio_local = nullptr;
ggml_tensor* audio_global = nullptr;
int64_t seq_len = x->ne[1];
int64_t t_ref_len = 0;
if (config.model_type == "s2v") {
if (audio_embed != nullptr) {
GGML_ASSERT(audio_embed->ne[1] == T * 4);
auto audio_encoder = std::dynamic_pointer_cast<WanCausalAudioEncoder>(blocks["casual_audio_encoder"]);
auto audio_emb = audio_encoder->forward(ctx, audio_embed);
audio_local = audio_emb.first;
audio_global = audio_emb.second;
GGML_ASSERT(audio_local->ne[2] == T);
}
// video tokens get cond_mask[0], reference tokens cond_mask[1]
auto cond_mask = params["trainable_cond_mask.weight"];
auto cm0 = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 0, 1), config.dim, 1, 1);
x = ggml_add(ctx->ggml_ctx, x, cm0);
if (reference_latent != nullptr) {
t_ref_len = reference_latent->ne[2];
auto ref = patch_embedding->forward(ctx, reference_latent);
ref = ggml_reshape_3d(ctx->ggml_ctx, ref, ref->ne[0] * ref->ne[1] * ref->ne[2], ref->ne[3] / N, N);
ref = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ref, 1, 0, 2, 3)); // [N, t_ref*h_len*w_len, dim]
auto cm1 = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 1, 2), config.dim, 1, 1);
ref = ggml_add(ctx->ggml_ctx, ref, cm1);
x = ggml_concat(ctx->ggml_ctx, x, ref, 1);
// Reference tokens use timestep 0.
GGML_ASSERT(timestep->ne[0] == T);
timestep = ggml_ext_pad(ctx->ggml_ctx, timestep, (int)t_ref_len, 0, 0, 0);
}
}
// time_embedding
auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, config.freq_dim);
e = time_embedding_0->forward(ctx, e);
@@ -714,6 +786,11 @@ namespace WAN {
auto x_orig = x;
std::shared_ptr<WanAudioInjector> audio_injector;
if (audio_local != nullptr) {
audio_injector = std::dynamic_pointer_cast<WanAudioInjector>(blocks["audio_injector"]);
}
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<WanAttentionBlock>(blocks["blocks." + std::to_string(i)]);
@@ -731,6 +808,13 @@ namespace WAN {
c_skip = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength);
x = ggml_add(ctx->ggml_ctx, x, c_skip);
}
if (audio_injector != nullptr) {
auto inject_iter = config.audio_inject_mapping.find(i);
if (inject_iter != config.audio_inject_mapping.end()) {
x = audio_injector->forward(ctx, x, seq_len, T, inject_iter->second, audio_local, audio_global);
}
}
sd::ggml_graph_cut::mark_graph_cut(x, "wan.blocks." + std::to_string(i), "x");
if (c != nullptr) {
sd::ggml_graph_cut::mark_graph_cut(c, "wan.blocks." + std::to_string(i), "c");
@@ -747,11 +831,13 @@ namespace WAN {
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
ggml_tensor* clip_fea = nullptr,
ggml_tensor* time_dim_concat = nullptr,
ggml_tensor* vace_context = nullptr,
float vace_strength = 1.f,
int64_t N = 1) {
ggml_tensor* clip_fea = nullptr,
ggml_tensor* time_dim_concat = nullptr,
ggml_tensor* vace_context = nullptr,
float vace_strength = 1.f,
int64_t N = 1,
ggml_tensor* audio_embed = nullptr,
ggml_tensor* reference_latent = nullptr) {
// Forward pass of DiT.
// x: [N*C, T, H, W]
// timestep: [N,]
@@ -779,7 +865,12 @@ namespace WAN {
t_len = ((x->ne[2] + (std::get<0>(config.patch_size) / 2)) / std::get<0>(config.patch_size));
}
auto out = forward_orig(ctx, x, timestep, context, pe, clip_fea, vace_context, vace_strength, N); // [N, t_len*h_len*w_len, pt*ph*pw*C]
auto out = forward_orig(ctx, x, timestep, context, pe, clip_fea, vace_context, vace_strength, N, audio_embed, reference_latent); // [N, (t_len [+t_ref]) *h_len*w_len, pt*ph*pw*C]
if (reference_latent != nullptr) {
// Exclude reference tokens from the generated video.
out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, t_len * h_len * w_len);
}
out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
@@ -839,7 +930,10 @@ namespace WAN {
config.text_len = 512;
}
} else if (config.num_layers == 40) {
if (config.model_type == "t2v") {
if (version == VERSION_WAN2_2_S2V) {
desc = "Wan2.2-S2V-14B";
config.in_dim = 16;
} else if (config.model_type == "t2v") {
if (version == VERSION_WAN2_2_I2V) {
desc = "Wan2.2-I2V-14B";
config.in_dim = 36;
@@ -891,7 +985,9 @@ namespace WAN {
const sd::Tensor<float>& c_concat_tensor = {},
const sd::Tensor<float>& time_dim_concat_tensor = {},
const sd::Tensor<float>& vace_context_tensor = {},
float vace_strength = 1.f) {
float vace_strength = 1.f,
const sd::Tensor<float>& audio_embed_tensor = {},
const sd::Tensor<float>& ref_latent_tensor = {}) {
ggml_cgraph* gf = new_graph_custom(WAN_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
@@ -901,16 +997,33 @@ namespace WAN {
ggml_tensor* c_concat = make_optional_input(c_concat_tensor);
ggml_tensor* time_dim_concat = make_optional_input(time_dim_concat_tensor);
ggml_tensor* vace_context = make_optional_input(vace_context_tensor);
ggml_tensor* audio_embed = make_optional_input(audio_embed_tensor);
ggml_tensor* ref_latent = make_optional_input(ref_latent_tensor);
pe_vec = Rope::gen_wan_pe(static_cast<int>(x->ne[2]),
static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
std::get<0>(config.patch_size),
std::get<1>(config.patch_size),
std::get<2>(config.patch_size),
1,
config.theta,
config.axes_dim);
pe_vec = Rope::gen_wan_pe(static_cast<int>(x->ne[2]),
static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
std::get<0>(config.patch_size),
std::get<1>(config.patch_size),
std::get<2>(config.patch_size),
1,
config.theta,
config.axes_dim);
if (ref_latent != nullptr) {
// Match S2V's reference-frame temporal offset.
int t_start = std::max(30, static_cast<int>(x->ne[2]) + 9);
auto ref_pe = Rope::gen_wan_pe(static_cast<int>(ref_latent->ne[2]),
static_cast<int>(ref_latent->ne[1]),
static_cast<int>(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<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_VERBOSE("pos_len %d", pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
@@ -933,7 +1046,10 @@ namespace WAN {
clip_fea,
time_dim_concat,
vace_context,
vace_strength);
vace_strength,
1,
audio_embed,
ref_latent);
ggml_build_forward_expand(gf, out);
@@ -948,9 +1064,11 @@ namespace WAN {
const sd::Tensor<float>& c_concat = {},
const sd::Tensor<float>& time_dim_concat = {},
const sd::Tensor<float>& vace_context = {},
float vace_strength = 1.f) {
float vace_strength = 1.f,
const sd::Tensor<float>& audio_embed = {},
const sd::Tensor<float>& ref_latent = {}) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength, audio_embed, ref_latent);
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
@@ -961,6 +1079,12 @@ namespace WAN {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
const auto* extra = diffusion_extra_as<WanDiffusionExtra>(diffusion_params);
static const std::vector<sd::Tensor<float>> 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<float> empty_tensor;
const sd::Tensor<float>& ref_latent = ref_latents.empty() ? empty_tensor : ref_latents[0];
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
@@ -969,7 +1093,9 @@ namespace WAN {
tensor_or_empty(diffusion_params.c_concat),
sd::Tensor<float>(),
tensor_or_empty(extra->vace_context),
extra->vace_strength);
extra->vace_strength,
tensor_or_empty(extra->audio_embed),
ref_latent);
}
void test() {
src/model/diffusion/wan_audio.hpp
+215
View File
@@ -0,0 +1,215 @@
#ifndef __SD_MODEL_DIFFUSION_WAN_AUDIO_HPP__
#define __SD_MODEL_DIFFUSION_WAN_AUDIO_HPP__
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "model/common/ggml_block.hpp"
namespace WAN {
class WanCausalConv1d : public UnaryBlock {
private:
int kernel_size_;
public:
WanCausalConv1d(int64_t in_dim,
int64_t out_dim,
int kernel_size = 3,
int stride = 1)
: kernel_size_(kernel_size) {
blocks["conv"] = std::make_shared<Conv1d>(in_dim, out_dim, kernel_size, stride, 0, 1, 1, true, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
// Replicate the first sample for causal left padding.
if (kernel_size_ > 1) {
auto first = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, 1);
for (int i = 0; i < kernel_size_ - 1; i++) {
x = ggml_concat(ctx->ggml_ctx, first, x, 0);
}
}
return std::dynamic_pointer_cast<Conv1d>(blocks["conv"])->forward(ctx, x);
}
};
class WanMotionEncoder : public GGMLBlock {
private:
int64_t hidden_dim_;
int num_token_;
bool need_global_;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
// The padding token is combined with F32 activations.
params["padding_tokens"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_dim_);
}
ggml_tensor* conv_norm_silu(GGMLRunnerContext* ctx,
ggml_tensor* x,
const std::string& conv_key,
const std::string& norm_key,
bool to_conv_layout) {
x = std::dynamic_pointer_cast<WanCausalConv1d>(blocks[conv_key])->forward(ctx, x);
x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3);
x = std::dynamic_pointer_cast<LayerNorm>(blocks[norm_key])->forward(ctx, x);
x = ggml_silu(ctx->ggml_ctx, x);
if (to_conv_layout) {
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
}
return x;
}
public:
WanMotionEncoder(int64_t in_dim,
int64_t hidden_dim,
int num_token,
bool need_global = true)
: hidden_dim_(hidden_dim), num_token_(num_token), need_global_(need_global) {
blocks["conv1_local"] = std::make_shared<WanCausalConv1d>(in_dim, hidden_dim / 4 * num_token);
if (need_global) {
blocks["conv1_global"] = std::make_shared<WanCausalConv1d>(in_dim, hidden_dim / 4);
}
blocks["norm1"] = std::make_shared<LayerNorm>(hidden_dim / 4, 1e-6f, false);
blocks["conv2"] = std::make_shared<WanCausalConv1d>(hidden_dim / 4, hidden_dim / 2, 3, 2);
blocks["norm2"] = std::make_shared<LayerNorm>(hidden_dim / 2, 1e-6f, false);
blocks["conv3"] = std::make_shared<WanCausalConv1d>(hidden_dim / 2, hidden_dim, 3, 2);
blocks["norm3"] = std::make_shared<LayerNorm>(hidden_dim, 1e-6f, false);
if (need_global) {
blocks["final_linear"] = std::make_shared<Linear>(hidden_dim, hidden_dim);
}
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto local = std::dynamic_pointer_cast<WanCausalConv1d>(blocks["conv1_local"])->forward(ctx, x);
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
std::vector<ggml_tensor*> tokens;
// Each token group is normalized independently over channels.
for (auto& group : ggml_ext_chunk(ctx->ggml_ctx, local, num_token_, 1)) {
ggml_tensor* s = ggml_permute(ctx->ggml_ctx, group, 1, 0, 2, 3);
s = norm1->forward(ctx, s);
s = ggml_silu(ctx->ggml_ctx, s);
s = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 0, 2, 3));
s = conv_norm_silu(ctx, s, "conv2", "norm2", true);
s = conv_norm_silu(ctx, s, "conv3", "norm3", false);
tokens.push_back(ggml_reshape_3d(ctx->ggml_ctx, s, s->ne[0], 1, s->ne[1]));
}
auto padding = ggml_reshape_3d(ctx->ggml_ctx, params["padding_tokens"], hidden_dim_, 1, 1);
padding = ggml_repeat(ctx->ggml_ctx, padding, tokens[0]);
tokens.push_back(padding);
ggml_tensor* local_out = ggml_ext_vec_concat(ctx->ggml_ctx, tokens, 1);
if (!need_global_) {
return {local_out, nullptr};
}
ggml_tensor* g = conv_norm_silu(ctx, x, "conv1_global", "norm1", true);
g = conv_norm_silu(ctx, g, "conv2", "norm2", true);
g = conv_norm_silu(ctx, g, "conv3", "norm3", false);
g = std::dynamic_pointer_cast<Linear>(blocks["final_linear"])->forward(ctx, g);
return {local_out, g};
}
};
class WanCausalAudioEncoder : public GGMLBlock {
private:
int num_layers_;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
// Preserve the checkpoint shape for loading; layer mixing requires F32.
auto it = tensor_storage_map.find(prefix + "weights");
if (it != tensor_storage_map.end()) {
params["weights"] = ggml_new_tensor(ctx, GGML_TYPE_F32, it->second.n_dims, it->second.ne);
} else {
params["weights"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_layers_);
}
}
public:
WanCausalAudioEncoder(int64_t audio_dim,
int64_t dim,
int num_token,
int num_layers = 25)
: num_layers_(num_layers) {
blocks["encoder"] = std::make_shared<WanMotionEncoder>(audio_dim, dim, num_token, true);
}
// features: [layers, frames, audio_dim]; outputs: [T, tokens+1, dim] and [T, dim].
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* features) {
auto weights = ggml_silu(ctx->ggml_ctx, params["weights"]);
auto x = ggml_mul(ctx->ggml_ctx, features, ggml_reshape_3d(ctx->ggml_ctx, weights, 1, 1, num_layers_));
x = ggml_div(ctx->ggml_ctx, x, ggml_sum(ctx->ggml_ctx, weights));
// Move the layer axis to ggml dimension 0 for reduction.
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
x = ggml_sum_rows(ctx->ggml_ctx, x);
x = ggml_reshape_2d(ctx->ggml_ctx, x, x->ne[1], x->ne[2]);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
return std::dynamic_pointer_cast<WanMotionEncoder>(blocks["encoder"])->forward(ctx, x);
}
};
class WanAudioInjector : public GGMLBlock {
private:
int64_t dim_;
public:
WanAudioInjector(int64_t dim,
int64_t num_heads,
int count,
bool qk_norm = true,
float eps = 1e-6f)
: dim_(dim) {
for (int i = 0; i < count; i++) {
blocks["injector." + std::to_string(i)] =
std::make_shared<WanT2VCrossAttention>(dim, num_heads, qk_norm, eps);
blocks["injector_adain_layers." + std::to_string(i) + ".linear"] =
std::make_shared<Linear>(dim, dim * 2);
}
// S2V AdaLayerNorm uses its own epsilon, independent of attention norms.
blocks["adain_norm"] = std::make_shared<LayerNorm>(dim, 1e-5f, false);
}
// Inject into the video prefix; trailing reference tokens pass through unchanged.
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
int64_t seq_len,
int64_t T,
int injector_id,
ggml_tensor* audio_local,
ggml_tensor* audio_global) {
int64_t n_tok = seq_len / T;
int64_t n_token = x->ne[1];
auto adain_linear = std::dynamic_pointer_cast<Linear>(blocks["injector_adain_layers." + std::to_string(injector_id) + ".linear"]);
auto injector = std::dynamic_pointer_cast<WanT2VCrossAttention>(blocks["injector." + std::to_string(injector_id)]);
auto adain_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["adain_norm"]);
auto temb = ggml_silu(ctx->ggml_ctx, audio_global);
temb = adain_linear->forward(ctx, temb);
auto shift = ggml_ext_slice(ctx->ggml_ctx, temb, 0, 0, dim_);
auto scale = ggml_ext_slice(ctx->ggml_ctx, temb, 0, dim_, dim_ * 2);
shift = ggml_reshape_3d(ctx->ggml_ctx, shift, dim_, 1, T);
scale = ggml_reshape_3d(ctx->ggml_ctx, scale, dim_, 1, T);
auto x_vid = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, seq_len);
auto h = ggml_reshape_3d(ctx->ggml_ctx, x_vid, dim_, n_tok, T);
h = adain_norm->forward(ctx, h);
h = ggml_add(ctx->ggml_ctx, h, ggml_mul(ctx->ggml_ctx, h, scale));
h = ggml_add(ctx->ggml_ctx, h, shift);
auto res = injector->forward(ctx, h, audio_local, 0);
res = ggml_reshape_2d(ctx->ggml_ctx, res, dim_, seq_len);
auto x_head = ggml_add(ctx->ggml_ctx, x_vid, res);
if (seq_len < n_token) {
auto x_tail = ggml_ext_slice(ctx->ggml_ctx, x, 1, seq_len, n_token);
return ggml_concat(ctx->ggml_ctx, x_head, x_tail, 1);
}
return x_head;
}
};
} // namespace WAN
#endif // __SD_MODEL_DIFFUSION_WAN_AUDIO_HPP__