diff --git a/CMakeLists.txt b/CMakeLists.txt index f15fdbf28..a003f400a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -312,7 +312,7 @@ set(CMAKE_POLICY_DEFAULT_CMP0077 NEW) if (NOT SD_USE_SYSTEM_GGML) # see https://github.com/ggerganov/ggml/pull/682 - add_definitions(-DGGML_MAX_NAME=128) + add_definitions(-DGGML_MAX_NAME=160) endif() # deps diff --git a/assets/animatediff/v2_coast.gif b/assets/animatediff/v2_coast.gif new file mode 100644 index 000000000..275f4023f Binary files /dev/null and b/assets/animatediff/v2_coast.gif differ diff --git a/assets/animatediff/v2_house.gif b/assets/animatediff/v2_house.gif new file mode 100644 index 000000000..6ef98f969 Binary files /dev/null and b/assets/animatediff/v2_house.gif differ diff --git a/assets/animatediff/v2_man.gif b/assets/animatediff/v2_man.gif new file mode 100644 index 000000000..3cb6d98a4 Binary files /dev/null and b/assets/animatediff/v2_man.gif differ diff --git a/assets/animatediff/v2_rabbit.gif b/assets/animatediff/v2_rabbit.gif new file mode 100644 index 000000000..814a1a587 Binary files /dev/null and b/assets/animatediff/v2_rabbit.gif differ diff --git a/assets/animatediff/v3_coast.gif b/assets/animatediff/v3_coast.gif new file mode 100644 index 000000000..9e38def31 Binary files /dev/null and b/assets/animatediff/v3_coast.gif differ diff --git a/assets/animatediff/v3_house.gif b/assets/animatediff/v3_house.gif new file mode 100644 index 000000000..9d867304a Binary files /dev/null and b/assets/animatediff/v3_house.gif differ diff --git a/assets/animatediff/v3_man.gif b/assets/animatediff/v3_man.gif new file mode 100644 index 000000000..7d7094fa8 Binary files /dev/null and b/assets/animatediff/v3_man.gif differ diff --git a/assets/animatediff/v3_rabbit.gif b/assets/animatediff/v3_rabbit.gif new file mode 100644 index 000000000..f35277062 Binary files /dev/null and b/assets/animatediff/v3_rabbit.gif differ diff --git a/assets/animatediff/v3_rabbit_domain_lora.gif b/assets/animatediff/v3_rabbit_domain_lora.gif new file mode 100644 index 000000000..63847e158 Binary files /dev/null and b/assets/animatediff/v3_rabbit_domain_lora.gif differ diff --git a/docs/animatediff.md b/docs/animatediff.md new file mode 100644 index 000000000..b26b6aaee --- /dev/null +++ b/docs/animatediff.md @@ -0,0 +1,152 @@ +# AnimateDiff (SD 1.5) + +AnimateDiff adds motion to a frozen Stable Diffusion 1.5 checkpoint by +injecting a temporal-attention module at 20 UNet slots. The base SD 1.5 +model, VAE, and text encoder are unchanged; only the motion module produces +the temporal residual that turns a batch of independent frames into a +coherent animation. Reference: Guo et al., "AnimateDiff: Animate Your +Personalized Text-to-Image Diffusion Models without Specific Tuning" +(https://arxiv.org/abs/2307.04725). + +## Download weights + +- Motion module (v3, recommended) + - fp16 safetensors: https://huggingface.co/conrevo/AnimateDiff-A1111/resolve/main/motion_module/mm_sd15_v3.safetensors + - original checkpoint: https://huggingface.co/guoyww/animatediff/resolve/main/v3_sd15_mm.ckpt +- SD 1.5 base model + - any SD 1.5 checkpoint works. `realisticVisionV60B1` and `toonyou_beta3` + are the ones used in guoyww's reference configs. +- Domain Adapter LoRA (optional, v3 only, sharpens the base UNet's output + toward the motion module's trained distribution) + - ckpt: https://huggingface.co/guoyww/animatediff/resolve/main/v3_sd15_adapter.ckpt + - place under your `--lora-model-dir` and reference in the prompt as + ``. + +The motion module is `~836 MB` and loads alongside the SD 1.5 UNet via +`--motion-module`. + +## Motion module versions + +Per [animatediff.net/models](https://animatediff.net/models): + +| Module | Base | Native res | Character | +|---------------------|------|------------|-----------| +| `mm_sd_v14.ckpt` | 1.5 | 256x256 | earliest, more jittery | +| `mm_sd_v15.ckpt` | 1.5 | 256x256 | improved stability over v1.4 | +| `mm_sd_v15_v2.ckpt` | 1.5 | 384x384 | significantly better motion dynamics | +| `v3_sd15_mm.ckpt` | 1.5 | 512x512 | smoothest, highest quality; pairs with a Domain Adapter LoRA | +| `mm_sdxl_v10_beta` | SDXL | 512x512 | experimental, not yet supported here | + +Match your `-H -W` to the module's native resolution for best results. v3 is +trained at 512x512 - going smaller (e.g. 384x384) still works but the motion +character is closer to v2. + +## Examples + +Generate an 8-frame animation at 512x512, seed 42, 20 steps. The sampler / +scheduler / CFG values below match what mm_sd15_v3 was trained with; using +SD 1.5 defaults (euler_a, low CFG) produces noise-like output. + +``` +.\bin\Release\sd-cli.exe -M vid_gen \ + --model ..\models\checkpoints\realisticVisionV60B1.safetensors \ + --motion-module ..\models\animatediff\mm_sd15_v3.safetensors \ + --offload-to-cpu --diffusion-fa \ + -p "a red apple on a wooden table" \ + --cfg-scale 8.0 --sampling-method euler --scheduler discrete \ + -H 512 -W 512 --video-frames 8 --fps 8 --steps 20 -s 42 \ + -o out.avi +``` + +Generate at the motion module's native 16-frame context (recommended for +best temporal quality). Needs more VRAM at 512x512, so drop to 384x384 or +use layer streaming: + +``` +.\bin\Release\sd-cli.exe -M vid_gen \ + --model ..\models\checkpoints\realisticVisionV60B1.safetensors \ + --motion-module ..\models\animatediff\mm_sd15_v3.safetensors \ + --offload-to-cpu --diffusion-fa \ + -p "photo of coastline, rocks, storm weather, wind, waves, lightning" \ + --cfg-scale 8.0 --sampling-method euler --scheduler discrete \ + -H 384 -W 384 --video-frames 16 --fps 8 --steps 20 -s 42 \ + -o out.avi +``` + +Low-VRAM streaming (verified with a 2 GiB cap on RTX 3060): + +``` +.\bin\Release\sd-cli.exe -M vid_gen \ + --model ..\models\checkpoints\realisticVisionV60B1.safetensors \ + --motion-module ..\models\animatediff\mm_sd15_v3.safetensors \ + --max-vram 2.0 --stream-layers --diffusion-fa \ + -p "photo of coastline, rocks, storm weather, wind, waves, lightning" \ + --cfg-scale 8.0 --sampling-method euler --scheduler discrete \ + -H 384 -W 384 --video-frames 8 --fps 8 --steps 20 -s 42 \ + -o out.avi +``` + +## Reference-quality reproduction + +Using guoyww's official reference configs on this impl (RealisticVision v6.0 +base + `mm_sd15_v3` or `mm_sd_v15_v2` + native resolution + 16 frames + euler ++ 25 steps + CFG 8 + linear beta schedule) reproduces the reference +AnimateDiff output style. + +### v3 (512x512, `mm_sd15_v3`) + +| Prompt | Sample | +|---------------------------------------|--------| +| B&W man on stormy coastline | | +| Close-up rabbit macro shot | | +| Coastline, storm, waves, lightning | | +| Old house, storm, forest, night | | + +### v2 (384x384, `mm_sd_v15_v2.ckpt`) + +| Prompt | Sample | +|---------------------------------------|--------| +| B&W man on stormy coastline | | +| Close-up rabbit macro shot | | +| Coastline, storm, waves, lightning | | +| Old house, storm, forest, night | | + +Motion is strong for scenes with motion cues in the prompt (storm/waves/wind) +and subtle for static subjects (close-up macro), matching reference behavior. +v2 has an additional motion module at the UNet middle block that v3 dropped; +this impl auto-detects the topology from the checkpoint. + +### v3 + Domain Adapter LoRA + +Attaching the `v3_sd15_adapter` LoRA sharpens the base UNet output toward +the training distribution the motion module was fine-tuned against. Same +prompt, seed, config as above: + + + +Individual fur strands, glowing inner-ear, and richer forest detail become +visible compared to the no-LoRA rendering. + +``` +sd-cli -M vid_gen --model realisticVisionV60B1.safetensors \ + --motion-module mm_sd15_v3.safetensors \ + --lora-model-dir ./loras \ + -p "close up photo of a rabbit ..." ... +``` + +## Notes + +- The motion module was trained at `video_length=16`. Running with + `--video-frames 16` gives the best coherence; F=8 works but shows a shorter + motion arc. Frame counts up to 32 are supported by the positional encoding + but exceed the trained regime and produce more static output. +- At `--video-frames 1` the motion module is skipped entirely and the output + is bit-identical to `-M img_gen`. This avoids the single-token + temporal-attention degeneracy that would otherwise emit an untrained-magnitude + residual on a single-frame sample. +- The base UNet is frozen, so character identity and style follow the SD 1.5 + checkpoint you pass to `--model`. LoRAs and prompt weighting attach to the + base model in the usual way. +- No mid_block motion module in v3. `mm_sdxl_v10_beta` (SDXL variant) is not + supported yet. +- Output is written as MJPEG AVI. Use `--fps` to set playback speed. diff --git a/examples/common/common.cpp b/examples/common/common.cpp index ed23a155e..5042fd112 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -423,6 +423,11 @@ ArgOptions SDContextParams::get_options() { "path to control net model", 0, &control_net_path}, + {"", + "--motion-module", + "path to AnimateDiff motion module (SD 1.5); enables video generation on --video-frames > 1", + 0, + &motion_module_path}, {"", "--embd-dir", "embeddings directory", @@ -672,7 +677,7 @@ ArgOptions SDContextParams::get_options() { } void SDContextParams::build_embedding_map() { - static const std::vector valid_ext = {".gguf", ".safetensors", ".pt"}; + static const std::vector valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"}; if (!fs::exists(embedding_dir) || !fs::is_directory(embedding_dir)) { return; @@ -867,6 +872,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.audio_vae_path = audio_vae_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.motion_module_path = motion_module_path.c_str(); sd_ctx_params.embeddings = embedding_vec.data(); sd_ctx_params.embedding_count = static_cast(embedding_vec.size()); sd_ctx_params.photo_maker_path = photo_maker_path.c_str(); @@ -2032,7 +2038,7 @@ void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_d return; } static const std::regex re(R"(]+):([^>]+)>)"); - static const std::vector valid_ext = {".gguf", ".safetensors", ".pt"}; + static const std::vector valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"}; std::smatch m; std::string tmp = prompt; diff --git a/examples/common/common.h b/examples/common/common.h index 824c6d931..094d1e00b 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -132,6 +132,7 @@ struct SDContextParams { std::string taesd_path; std::string esrgan_path; std::string control_net_path; + std::string motion_module_path; std::string embedding_dir; std::string photo_maker_path; std::string pulid_weights_path; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 1dc61291c..9a8842621 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -199,6 +199,7 @@ typedef struct { const char* audio_vae_path; const char* taesd_path; const char* control_net_path; + const char* motion_module_path; const sd_embedding_t* embeddings; uint32_t embedding_count; const char* photo_maker_path; diff --git a/src/model/diffusion/animatediff.hpp b/src/model/diffusion/animatediff.hpp new file mode 100644 index 000000000..d470544fe --- /dev/null +++ b/src/model/diffusion/animatediff.hpp @@ -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 down_channels = {320, 640, 1280, 1280}; + std::vector 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(new Linear(channels, channels, false)); + blocks["to_k"] = std::shared_ptr(new Linear(channels, channels, false)); + blocks["to_v"] = std::shared_ptr(new Linear(channels, channels, false)); + blocks["to_out.0"] = std::shared_ptr(new Linear(channels, channels, true)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto to_q = std::dynamic_pointer_cast(blocks["to_q"]); + auto to_k = std::dynamic_pointer_cast(blocks["to_k"]); + auto to_v = std::dynamic_pointer_cast(blocks["to_v"]); + auto to_out = std::dynamic_pointer_cast(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(channels, num_heads, max_frames); + blocks["attention_blocks.1"] = std::make_shared(channels, num_heads, max_frames); + blocks["norms.0"] = std::shared_ptr(new LayerNorm(channels)); + blocks["norms.1"] = std::shared_ptr(new LayerNorm(channels)); + blocks["ff"] = std::make_shared(channels, channels, 4, FeedForward::Activation::GEGLU); + blocks["ff_norm"] = std::shared_ptr(new LayerNorm(channels)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto attn0 = std::dynamic_pointer_cast(blocks["attention_blocks.0"]); + auto attn1 = std::dynamic_pointer_cast(blocks["attention_blocks.1"]); + auto norm0 = std::dynamic_pointer_cast(blocks["norms.0"]); + auto norm1 = std::dynamic_pointer_cast(blocks["norms.1"]); + auto ff = std::dynamic_pointer_cast(blocks["ff"]); + auto ff_norm = std::dynamic_pointer_cast(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(new GroupNorm(norm_num_groups, channels)); + blocks["proj_in"] = std::shared_ptr(new Linear(channels, channels, true)); + blocks["transformer_blocks.0"] = std::make_shared(channels, num_heads, max_frames); + blocks["proj_out"] = std::shared_ptr(new Linear(channels, channels, true)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t num_frames) { + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto proj_in = std::dynamic_pointer_cast(blocks["proj_in"]); + auto tb0 = std::dynamic_pointer_cast(blocks["transformer_blocks.0"]); + auto proj_out = std::dynamic_pointer_cast(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(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(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(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(ch, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames); + } + } + for (int i = 0; i < static_cast(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(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(cfg.mid_channels, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames); + } + } + + std::shared_ptr motion(const std::string& key) { + auto it = blocks.find(key); + if (it == blocks.end()) + return nullptr; + return std::dynamic_pointer_cast(it->second); + } + }; + +} // namespace AnimateDiff + +#endif // __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__ diff --git a/src/model/diffusion/unet.hpp b/src/model/diffusion/unet.hpp index 253b3b4b2..118ed6327 100644 --- a/src/model/diffusion/unet.hpp +++ b/src/model/diffusion/unet.hpp @@ -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(input->ne[2]); @@ -473,6 +483,12 @@ class UnetModelBlock : public GGMLBlock { blocks["out.0"] = std::shared_ptr(new GroupNorm32(ch)); // ch == model_channels // out_1 is nn.SiLU() blocks["out.2"] = std::shared_ptr(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(mm_cfg); + } } ggml_tensor* resblock_forward(std::string name, @@ -583,6 +599,42 @@ class UnetModelBlock : public GGMLBlock { ggml_set_name(h, "bench-start"); hs.push_back(h); + + auto motion_root = config.enable_animatediff && num_video_frames > 1 + ? std::dynamic_pointer_cast(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 @@ class UnetModelBlock : public GGMLBlock { 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 @@ class UnetModelBlock : public GGMLBlock { 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 @@ class UnetModelBlock : public GGMLBlock { up_sample_idx++; } + h = apply_motion_output(output_block_idx, h); + if (i > 0 && j == num_res_blocks) { if (tiny_unet) { output_block_idx++; diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 6382c33e8..3aa7880e2 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -304,6 +304,12 @@ std::string convert_diffusers_unet_to_original_sd1(std::string name) { } } + static const std::vector> 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 diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index fd6a93c68..10fdde11c 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -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 @@ class StableDiffusionGGML { std::vector> generation_extensions; std::vector> 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 @@ class StableDiffusionGGML { } } + 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 @@ class StableDiffusionGGML { 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(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 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 shape(init_latent.shape().begin(), init_latent.shape().end()); + shape[3] = n_frames; + init_latent = sd::Tensor(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 frame_latent = sd::ops::slice(final_latents[i], 3, f, f + 1); + sd::Tensor 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 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);