feat: AnimateDiff SD 1.5 motion modules (v2 + v3) (#1784)

This commit is contained in:
fszontagh
2026-07-14 17:58:03 +02:00
committed by GitHub
parent 833369da84
commit c00a9e956d
18 changed files with 522 additions and 9 deletions

View File

@@ -0,0 +1,182 @@
#ifndef __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__
#define __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__
#include "core/ggml_extend.hpp"
#include "model/common/block.hpp"
// AnimateDiff (https://arxiv.org/abs/2307.04725) SD 1.5 motion modules.
namespace AnimateDiff {
struct MotionModuleConfig {
int max_frames = 32;
int64_t num_heads = 8;
int norm_num_groups = 32;
std::vector<int64_t> down_channels = {320, 640, 1280, 1280};
std::vector<int64_t> up_channels = {1280, 1280, 640, 320};
int num_down_motion_per_block = 2;
int num_up_motion_per_block = 3;
bool enable_mid_block = false;
int64_t mid_channels = 1280;
};
class TemporalAttention : public GGMLBlock {
protected:
int64_t channels;
int64_t num_heads;
int max_frames;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
params["pos_encoder.pe"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, channels, max_frames, 1);
}
public:
TemporalAttention(int64_t channels, int64_t num_heads, int max_frames)
: channels(channels), num_heads(num_heads), max_frames(max_frames) {
blocks["to_q"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, false));
blocks["to_k"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, false));
blocks["to_v"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, false));
blocks["to_out.0"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, true));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto to_q = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
auto to_k = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
auto to_v = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
int64_t C = x->ne[0];
int64_t F = x->ne[1];
auto pe = params["pos_encoder.pe"];
auto pe_f = (F == pe->ne[1])
? pe
: ggml_view_3d(ctx->ggml_ctx, pe, C, F, 1, pe->nb[1], pe->nb[2], 0);
auto x_pe = ggml_add(ctx->ggml_ctx, x, ggml_repeat(ctx->ggml_ctx, pe_f, x));
auto q = to_q->forward(ctx, x_pe);
auto k = to_k->forward(ctx, x_pe);
auto v = to_v->forward(ctx, x_pe);
auto a = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, (int)num_heads, nullptr, false);
return to_out->forward(ctx, a);
}
};
class TemporalTransformerBlock : public GGMLBlock {
public:
TemporalTransformerBlock(int64_t channels, int64_t num_heads, int max_frames) {
blocks["attention_blocks.0"] = std::make_shared<TemporalAttention>(channels, num_heads, max_frames);
blocks["attention_blocks.1"] = std::make_shared<TemporalAttention>(channels, num_heads, max_frames);
blocks["norms.0"] = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
blocks["norms.1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
blocks["ff"] = std::make_shared<FeedForward>(channels, channels, 4, FeedForward::Activation::GEGLU);
blocks["ff_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto attn0 = std::dynamic_pointer_cast<TemporalAttention>(blocks["attention_blocks.0"]);
auto attn1 = std::dynamic_pointer_cast<TemporalAttention>(blocks["attention_blocks.1"]);
auto norm0 = std::dynamic_pointer_cast<LayerNorm>(blocks["norms.0"]);
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norms.1"]);
auto ff = std::dynamic_pointer_cast<FeedForward>(blocks["ff"]);
auto ff_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["ff_norm"]);
auto r = x;
x = ggml_add(ctx->ggml_ctx, attn0->forward(ctx, norm0->forward(ctx, x)), r);
r = x;
x = ggml_add(ctx->ggml_ctx, attn1->forward(ctx, norm1->forward(ctx, x)), r);
r = x;
x = ggml_add(ctx->ggml_ctx, ff->forward(ctx, ff_norm->forward(ctx, x)), r);
return x;
}
};
class TemporalTransformer : public GGMLBlock {
public:
TemporalTransformer(int64_t channels, int64_t num_heads, int norm_num_groups, int max_frames) {
blocks["norm"] = std::shared_ptr<GGMLBlock>(new GroupNorm(norm_num_groups, channels));
blocks["proj_in"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, true));
blocks["transformer_blocks.0"] = std::make_shared<TemporalTransformerBlock>(channels, num_heads, max_frames);
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, true));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t num_frames) {
auto norm = std::dynamic_pointer_cast<GroupNorm>(blocks["norm"]);
auto proj_in = std::dynamic_pointer_cast<Linear>(blocks["proj_in"]);
auto tb0 = std::dynamic_pointer_cast<TemporalTransformerBlock>(blocks["transformer_blocks.0"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t C = x->ne[2];
GGML_ASSERT(x->ne[3] == num_frames);
auto residual = x;
auto h = norm->forward(ctx, x);
h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
h = ggml_reshape_3d(ctx->ggml_ctx, h, C, num_frames, W * H);
h = proj_in->forward(ctx, h);
h = tb0->forward(ctx, h);
h = proj_out->forward(ctx, h);
h = ggml_reshape_4d(ctx->ggml_ctx, h, C, num_frames, W, H);
h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
return ggml_add(ctx->ggml_ctx, h, residual);
}
};
class MotionModule : public GGMLBlock {
public:
MotionModule(int64_t channels, int64_t num_heads, int norm_num_groups, int max_frames) {
blocks["temporal_transformer"] = std::make_shared<TemporalTransformer>(channels, num_heads, norm_num_groups, max_frames);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t num_frames) {
auto tt = std::dynamic_pointer_cast<TemporalTransformer>(blocks["temporal_transformer"]);
return tt->forward(ctx, x, num_frames);
}
};
class AnimateDiffModel : public GGMLBlock {
public:
MotionModuleConfig config;
AnimateDiffModel(const MotionModuleConfig& cfg)
: config(cfg) {
for (int i = 0; i < static_cast<int>(cfg.down_channels.size()); ++i) {
int64_t ch = cfg.down_channels[i];
for (int j = 0; j < cfg.num_down_motion_per_block; ++j) {
blocks["down_blocks." + std::to_string(i) + ".motion_modules." + std::to_string(j)] =
std::make_shared<MotionModule>(ch, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames);
}
}
for (int i = 0; i < static_cast<int>(cfg.up_channels.size()); ++i) {
int64_t ch = cfg.up_channels[i];
for (int j = 0; j < cfg.num_up_motion_per_block; ++j) {
blocks["up_blocks." + std::to_string(i) + ".motion_modules." + std::to_string(j)] =
std::make_shared<MotionModule>(ch, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames);
}
}
if (cfg.enable_mid_block) {
blocks["mid_block.motion_modules.0"] =
std::make_shared<MotionModule>(cfg.mid_channels, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames);
}
}
std::shared_ptr<MotionModule> motion(const std::string& key) {
auto it = blocks.find(key);
if (it == blocks.end())
return nullptr;
return std::dynamic_pointer_cast<MotionModule>(it->second);
}
};
} // namespace AnimateDiff
#endif // __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__

View File

@@ -6,6 +6,7 @@
#include "model.h"
#include "model/common/block.hpp"
#include "model/diffusion/animatediff.hpp"
#include "model/diffusion/model.hpp"
/*==================================================== UnetModel =====================================================*/
@@ -29,6 +30,8 @@ struct UNetConfig {
bool tiny_unet = false;
int model_channels = 320;
int adm_in_channels = 2816; // only for VERSION_SDXL/SVD
bool enable_animatediff = false;
bool animatediff_has_mid_block = false;
static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
@@ -84,6 +87,13 @@ struct UNetConfig {
return &it->second;
};
if (find_weight("motion_module.down_blocks.0.motion_modules.0.temporal_transformer.proj_in.weight") != nullptr) {
config.enable_animatediff = true;
if (find_weight("motion_module.mid_block.motion_modules.0.temporal_transformer.proj_in.weight") != nullptr) {
config.animatediff_has_mid_block = true;
}
}
if (const TensorStorage* input = find_weight("input_blocks.0.0.weight")) {
if (input->n_dims == 4) {
config.in_channels = static_cast<int>(input->ne[2]);
@@ -473,6 +483,12 @@ public:
blocks["out.0"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(ch)); // ch == model_channels
// out_1 is nn.SiLU()
blocks["out.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(model_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
if (this->config.enable_animatediff) {
AnimateDiff::MotionModuleConfig mm_cfg;
mm_cfg.enable_mid_block = this->config.animatediff_has_mid_block;
blocks["motion_module"] = std::make_shared<AnimateDiff::AnimateDiffModel>(mm_cfg);
}
}
ggml_tensor* resblock_forward(std::string name,
@@ -583,6 +599,42 @@ public:
ggml_set_name(h, "bench-start");
hs.push_back(h);
auto motion_root = config.enable_animatediff && num_video_frames > 1
? std::dynamic_pointer_cast<AnimateDiff::AnimateDiffModel>(blocks["motion_module"])
: nullptr;
auto apply_motion_input = [&](int input_block_idx, ggml_tensor* h_in) -> ggml_tensor* {
if (!motion_root)
return h_in;
int di = (input_block_idx - 1) / 3;
int mj = (input_block_idx - 1) % 3;
if (di < 0 || di >= (int)channel_mult.size() || mj < 0 || mj >= num_res_blocks)
return h_in;
auto mm = motion_root->motion("down_blocks." + std::to_string(di) + ".motion_modules." + std::to_string(mj));
if (!mm)
return h_in;
return mm->forward(ctx, h_in, num_video_frames);
};
auto apply_motion_output = [&](int output_block_idx, ggml_tensor* h_in) -> ggml_tensor* {
if (!motion_root)
return h_in;
int ui = output_block_idx / 3;
int mj = output_block_idx % 3;
if (ui < 0 || ui >= (int)channel_mult.size() || mj < 0 || mj > num_res_blocks)
return h_in;
auto mm = motion_root->motion("up_blocks." + std::to_string(ui) + ".motion_modules." + std::to_string(mj));
if (!mm)
return h_in;
return mm->forward(ctx, h_in, num_video_frames);
};
auto apply_motion_mid = [&](ggml_tensor* h_in) -> ggml_tensor* {
if (!motion_root)
return h_in;
auto mm = motion_root->motion("mid_block.motion_modules.0");
if (!mm)
return h_in;
return mm->forward(ctx, h_in, num_video_frames);
};
// input block 1-11
size_t len_mults = channel_mult.size();
int input_block_idx = 0;
@@ -597,6 +649,7 @@ public:
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
h = attention_layer_forward(name, ctx, h, context, num_video_frames); // [N, mult*model_channels, h, w]
}
h = apply_motion_input(input_block_idx, h);
sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h");
hs.push_back(h);
}
@@ -624,6 +677,7 @@ public:
h = attention_layer_forward("middle_block.1", ctx, h, context, num_video_frames); // [N, 4*model_channels, h/8, w/8]
h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames); // [N, 4*model_channels, h/8, w/8]
}
h = apply_motion_mid(h);
}
sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h");
if (controls.size() > 0) {
@@ -660,6 +714,8 @@ public:
up_sample_idx++;
}
h = apply_motion_output(output_block_idx, h);
if (i > 0 && j == num_res_blocks) {
if (tiny_unet) {
output_block_idx++;

View File

@@ -304,6 +304,12 @@ std::string convert_diffusers_unet_to_original_sd1(std::string name) {
}
}
static const std::vector<std::pair<std::string, std::string>> name_map{
{"to_out.weight", "to_out.0.weight"},
{"to_out.bias", "to_out.0.bias"},
};
replace_with_name_map(result, name_map);
return result;
}
@@ -1343,13 +1349,25 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
// diffusion model
{
bool matched = false;
for (const auto& prefix : diffuison_model_prefix_vec) {
if (starts_with(name, prefix)) {
name = convert_diffusion_model_name(name.substr(prefix.size()), prefix, version);
name = prefix + name;
name = convert_diffusion_model_name(name.substr(prefix.size()), prefix, version);
name = prefix + name;
matched = true;
break;
}
}
if (is_lora && !matched && !diffuison_model_prefix_vec.empty()) {
if (starts_with(name, "down_blocks.") || starts_with(name, "up_blocks.") ||
starts_with(name, "mid_block.") || starts_with(name, "conv_in.") ||
starts_with(name, "conv_out.") || starts_with(name, "time_embedding.") ||
starts_with(name, "conv_norm_out.")) {
const std::string& canonical_prefix = diffuison_model_prefix_vec.front();
name = convert_diffusion_model_name(name, canonical_prefix, version);
name = canonical_prefix + name;
}
}
}
// cond_stage_model

View File

@@ -24,6 +24,7 @@
#include "extensions/generation_extension.h"
#include "model/adapter/lora.hpp"
#include "model/diffusion/anima.hpp"
#include "model/diffusion/animatediff.hpp"
#include "model/diffusion/boogu.hpp"
#include "model/diffusion/control.hpp"
#include "model/diffusion/ernie_image.hpp"
@@ -63,6 +64,10 @@
const char* sd_vae_format_name(enum sd_vae_format_t format);
static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback);
static bool sd_version_supports_animatediff(SDVersion version) {
return version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX;
}
const char* model_version_to_str[] = {
"SD 1.x",
"SD 1.x Inpaint",
@@ -212,6 +217,8 @@ public:
std::vector<std::shared_ptr<GenerationExtension>> generation_extensions;
std::vector<std::shared_ptr<LoraModel>> runtime_lora_models;
bool apply_lora_immediately = false;
bool animatediff_loaded = false;
int animatediff_num_frames = 0;
std::string taesd_path;
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr};
@@ -808,6 +815,16 @@ public:
}
}
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,
"model.diffusion_model.motion_module.")) {
LOG_WARN("loading motion module from '%s' failed", sd_ctx_params->motion_module_path);
} else {
animatediff_loaded = true;
}
}
if (strlen(SAFE_STR(sd_ctx_params->control_net_path)) > 0) {
if (!model_loader.init_from_file(sd_ctx_params->control_net_path)) {
LOG_ERROR("init control net model loader from file failed: '%s'", sd_ctx_params->control_net_path);
@@ -2487,7 +2504,11 @@ public:
diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images);
if (sd_version_is_unet(version)) {
diffusion_params.extra = UNetDiffusionExtra{-1, &controls, control_strength};
int nvf = -1;
if (animatediff_loaded && noised_input.dim() >= 4 && noised_input.shape()[3] > 1) {
nvf = static_cast<int>(noised_input.shape()[3]);
}
diffusion_params.extra = UNetDiffusionExtra{nvf, &controls, control_strength};
} else if (sd_version_is_sd3(version)) {
diffusion_params.extra = SkipLayerDiffusionExtra{local_skip_layers};
} else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
@@ -3643,6 +3664,9 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) {
if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
return false;
}
if (sd_ctx->sd->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->sd->version)) {
return true;
}
return sd_version_supports_video_generation(sd_ctx->sd->version);
}
@@ -4659,6 +4683,14 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
}
}
if (sd_ctx->sd->animatediff_num_frames > 1 &&
init_latent.dim() >= 4 && init_latent.shape()[3] == 1) {
int n_frames = sd_ctx->sd->animatediff_num_frames;
std::vector<int64_t> shape(init_latent.shape().begin(), init_latent.shape().end());
shape[3] = n_frames;
init_latent = sd::Tensor<float>(std::move(shape)); // zero-filled batch of N frames; per-frame noise is generated later via randn_like.
}
if (!control_image_tensor.empty()) {
control_latent = sd_ctx->sd->encode_first_stage(control_image_tensor);
if (control_latent.empty()) {
@@ -4967,6 +4999,24 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
if (cancelled) {
break;
}
} else if (sd_ctx->sd->animatediff_num_frames > 1 &&
final_latents[i].dim() >= 4 &&
final_latents[i].shape()[3] == sd_ctx->sd->animatediff_num_frames) {
int n_frames = sd_ctx->sd->animatediff_num_frames;
for (int f = 0; f < n_frames; ++f) {
if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
LOG_ERROR("cancelling latent decodings");
cancelled = true;
break;
}
sd::Tensor<float> frame_latent = sd::ops::slice(final_latents[i], 3, f, f + 1);
sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(frame_latent);
if (image.empty()) {
LOG_ERROR("decode_first_stage failed for AnimateDiff frame %d/%d", f + 1, n_frames);
return nullptr;
}
decoded_images.push_back(std::move(image));
}
} else {
sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(final_latents[i]);
if (image.empty()) {
@@ -6057,6 +6107,47 @@ static bool apply_ltxv_refine_image_conditioning(sd_ctx_t* sd_ctx,
return true;
}
static bool generate_animatediff_video(sd_ctx_t* sd_ctx,
const sd_vid_gen_params_t* sd_vid_gen_params,
sd_image_t** frames_out,
int* num_frames_out) {
int n_frames = sd_vid_gen_params->video_frames;
if (n_frames < 1) {
LOG_ERROR("AnimateDiff: --video-frames must be >= 1");
return false;
}
if (n_frames > 32) {
LOG_WARN("AnimateDiff motion modules have a 32-frame positional-encoding context; capping to 32");
n_frames = 32;
}
sd_img_gen_params_t img_gen_params;
sd_img_gen_params_init(&img_gen_params);
img_gen_params.loras = sd_vid_gen_params->loras;
img_gen_params.lora_count = sd_vid_gen_params->lora_count;
img_gen_params.prompt = sd_vid_gen_params->prompt;
img_gen_params.negative_prompt = sd_vid_gen_params->negative_prompt;
img_gen_params.clip_skip = sd_vid_gen_params->clip_skip;
img_gen_params.width = sd_vid_gen_params->width;
img_gen_params.height = sd_vid_gen_params->height;
img_gen_params.sample_params = sd_vid_gen_params->sample_params;
img_gen_params.strength = sd_vid_gen_params->strength;
img_gen_params.seed = sd_vid_gen_params->seed;
img_gen_params.batch_count = 1;
img_gen_params.control_strength = 1.0f;
img_gen_params.vae_tiling_params = sd_vid_gen_params->vae_tiling_params;
img_gen_params.cache = sd_vid_gen_params->cache;
img_gen_params.hires = sd_vid_gen_params->hires;
img_gen_params.qwen_image_layers = 0;
img_gen_params.circular_x = sd_vid_gen_params->circular_x;
img_gen_params.circular_y = sd_vid_gen_params->circular_y;
sd_ctx->sd->animatediff_num_frames = n_frames;
bool ok = generate_image(sd_ctx, &img_gen_params, frames_out, num_frames_out);
sd_ctx->sd->animatediff_num_frames = 0;
return ok;
}
SD_API bool generate_video(sd_ctx_t* sd_ctx,
const sd_vid_gen_params_t* sd_vid_gen_params,
sd_image_t** frames_out,
@@ -6071,14 +6162,20 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
if (audio_out != nullptr) {
*audio_out = nullptr;
}
if (num_frames_out != nullptr) {
*num_frames_out = 0;
}
if (sd_ctx->sd->animatediff_loaded && sd_version_supports_animatediff(sd_ctx->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);
return generate_animatediff_video(sd_ctx, sd_vid_gen_params, frames_out, num_frames_out);
}
sd_ctx->sd->reset_cancel_flag();
const RefImageParams ref_image_params;
if (num_frames_out != nullptr) {
*num_frames_out = 0;
}
int64_t t0 = ggml_time_ms();
sd_ctx->sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params;
apply_circular_axes_to_diffusion(sd_ctx, sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y);