mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-10-19 05:59:21 +08:00
Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b5d812008e | ||
|
|
c79d24bdbd | ||
|
|
1b042838d9 | ||
|
|
ead6bf521b | ||
|
|
9beb6aca69 | ||
|
|
12b6fbff28 |
@@ -65,6 +65,7 @@ API and command-line option may change frequently.***
|
||||
- Video Models
|
||||
- [Wan2.1/Wan2.2](./docs/wan.md)
|
||||
- [LTX-2.3](./docs/ltx2.md)
|
||||
- [LingBot-Video](./docs/lingbot_video.md)
|
||||
- [PhotoMaker](./docs/photo_maker.md) support.
|
||||
- Control Net support with SD 1.5
|
||||
- LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora)
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
# How to Use
|
||||
|
||||
Lingbot Video uses a Lingbot diffusion transformer, the Wan2.1 VAE, and Qwen3-VL 4B as the LLM text encoder.
|
||||
|
||||
## Download weights
|
||||
|
||||
- Download lingbot-video-dense-1.3b
|
||||
- safetensors: https://huggingface.co/robbyant/lingbot-video-dense-1.3b/tree/main/transformer
|
||||
- Download lingbot-video-moe-30b-a3b
|
||||
- safetensors: https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b/tree/main/transformer
|
||||
- Download vae
|
||||
- safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors
|
||||
- Download Qwen3-VL 4B
|
||||
- safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders
|
||||
- gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main
|
||||
|
||||
## Notes
|
||||
|
||||
- Use `-M vid_gen`.
|
||||
- T2V uses the text prompt only.
|
||||
- I2V uses `-i` as the first video frame. The same image is also passed to the
|
||||
Qwen3-VL prompt enhancer when vision weights are available.
|
||||
- Video frames are aligned to Wan-style temporal compression, so use frame counts
|
||||
like 33, 49, or 81.
|
||||
|
||||
## Examples
|
||||
|
||||
### LingBot-Video T2V
|
||||
|
||||
```sh
|
||||
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\lingbot-video-dens-1.3b.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\models\vae\wan_2.1_vae.safetensors -p '{"caption":{"comprehensive_description":"A lovely cat sits comfortably on a soft cushion near a sunlit window, looking calm, gentle, and adorable. The cat has soft fluffy fur, bright expressive eyes, small rounded ears, delicate whiskers, and a relaxed posture. Warm daylight falls across the cat from one side, creating soft highlights on the fur and a cozy glow around the scene. The background is softly blurred, showing hints of a peaceful indoor home environment with warm tones and gentle natural light. The overall atmosphere is cute, tender, serene, and photorealistic, emphasizing the cat''s charming appearance, soft texture, and affectionate presence.","camera_info":{"color":"Warm","frame_size":"Close Up","shot_type_angle":"Eye level","lens_size":"Medium Lens","composition":"Centered balanced","lighting":"Soft light","lighting_type":"Daylight"},"world_knowledge":[],"prominent_elements":[{"name":"lovely cat","description":"A cute and gentle domestic cat with soft fluffy fur, expressive eyes, and a calm relaxed presence.","location":"center of the frame","relative_size":"large","shape_and_color":"Small animal body with rounded face, triangular ears, bright eyes, and soft fur in warm natural tones","texture":"soft, fluffy, silky","appearance_details":"The cat has clean well-groomed fur, delicate whiskers, small ears, a cute nose, and bright attentive eyes. Its expression appears calm, affectionate, and slightly curious.","relationship":"Acts as the main subject and emotional focal point of the scene.","orientation":"facing the camera","pose":"sitting comfortably with a relaxed posture","expression":"gentle, adorable, calm, slightly curious","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"cat eyes","description":"Bright expressive eyes that give the cat a sweet and affectionate appearance.","location":"upper center of the cat face","relative_size":"small","shape_and_color":"Round almond-like eyes with glossy reflections","texture":"clear, glossy, reflective","appearance_details":"The eyes catch the soft daylight, creating small natural highlights that make the cat look vivid and alive.","relationship":"Enhance the emotional charm and cuteness of the cat.","orientation":"looking toward the camera","pose":"","expression":"soft and attentive","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"soft cushion","description":"A comfortable cushion or blanket where the cat is resting.","location":"bottom portion of the frame","relative_size":"medium","shape_and_color":"Soft rounded fabric surface in light warm neutral tones","texture":"plush, fabric, cozy","appearance_details":"The cushion gently supports the cat and adds a comfortable home-like feeling to the scene.","relationship":"Provides a cozy resting place for the cat.","orientation":"horizontal","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"sunlit indoor background","description":"A softly blurred indoor background with warm daylight and peaceful home atmosphere.","location":"behind the cat, filling the upper and side areas of the frame","relative_size":"large","shape_and_color":"Soft abstract shapes in warm beige, cream, and pale golden tones","texture":"soft, blurry, bokeh-like","appearance_details":"The background is intentionally out of focus, keeping attention on the cat while creating a cozy and serene mood.","relationship":"Provides a warm and gentle environment that supports the cute domestic scene.","orientation":"upright","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":"","is_cluster":true,"number_of_objects":"numerous"}]}}' -n '{"universal_negative":{"visual_quality":["low quality","worst quality","blurry","pixelated","jpeg artifacts","low resolution","unstable color","color flicker","underexposed","overexposed","invisible subject","subject hidden in darkness"],"artistic_style":["painting","illustration","drawing","cartoon","3d render","cgi","sketch","digital art"],"composition_and_content":["text","watermark","signature","logo","subtitles","pillarboxed","side bars","portrait image in landscape frame"],"temporal_and_motion_stability":["flickering","jittery","motion blur","temporal inconsistency","warping","morphing","incoherent motion","unnatural movement","static object with sudden jump","frame-to-frame inconsistency"],"material_and_structure":["plastic-like glass","unrealistic texture","deformed bottle","liquid freezing improperly","distorted reflections"]}}' --diffusion-fa --offload-to-cpu --cfg-scale 3 --video-frames 33 -v
|
||||
```
|
||||
@@ -428,6 +428,11 @@ SD_API const char* sd_get_system_info();
|
||||
SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx);
|
||||
SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx);
|
||||
|
||||
// ControlNet hot-swap APIs are not safe to call while generation is in flight.
|
||||
SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path);
|
||||
SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx);
|
||||
SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx);
|
||||
|
||||
SD_API const char* sd_type_name(enum sd_type_t type);
|
||||
SD_API enum sd_type_t str_to_sd_type(const char* str);
|
||||
SD_API const char* sd_rng_type_name(enum rng_type_t rng_type);
|
||||
|
||||
@@ -1797,7 +1797,7 @@ struct LLMEmbedder : public Conditioner {
|
||||
arch = LLM::LLMArch::GPT_OSS_20B;
|
||||
} else if (sd_version_is_pid(version)) {
|
||||
arch = LLM::LLMArch::GEMMA2_2B;
|
||||
} else if (sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) {
|
||||
} else if (sd_version_is_lingbot_video(version) || sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) {
|
||||
arch = LLM::LLMArch::QWEN3_VL;
|
||||
} else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
|
||||
arch = LLM::LLMArch::QWEN3;
|
||||
@@ -2009,7 +2009,84 @@ struct LLMEmbedder : public Conditioner {
|
||||
|
||||
int64_t t0 = ggml_time_ms();
|
||||
|
||||
if (sd_version_is_qwen_image(version)) {
|
||||
if (sd_version_is_lingbot_video(version)) {
|
||||
const int pad_token = 151643;
|
||||
const std::string prompt_prefix =
|
||||
"<|im_start|>system\nGiven a user input that may include a text prompt alone, "
|
||||
"a text prompt with an image reference, or a text prompt with a video reference "
|
||||
"or a video reference alone, generate an \"Enhanced prompt\" that provides detailed "
|
||||
"visual descriptions suitable for video generation. Evaluate the level of detail "
|
||||
"in the user's input: if it is simple, enrich it by adding specifics about colors, "
|
||||
"shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal "
|
||||
"progression, and spatial relationships to create vivid, concrete, and temporally "
|
||||
"coherent scenes to create vivid and concrete scenes. Please generate only the "
|
||||
"enhanced description for the prompt below and avoid including any additional "
|
||||
"commentary or evaluations:<|im_end|>\n<|im_start|>user\n";
|
||||
|
||||
auto prefix_tokens = tokenizer->encode(prompt_prefix, nullptr);
|
||||
prompt_template_encode_start_idx = 0;
|
||||
for (int token : prefix_tokens) {
|
||||
if (token != pad_token) {
|
||||
prompt_template_encode_start_idx++;
|
||||
}
|
||||
}
|
||||
LOG_DEBUG("prompt_template_encode_start_idx %d", prompt_template_encode_start_idx);
|
||||
|
||||
prompt = prompt_prefix;
|
||||
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
|
||||
LOG_INFO("LingBotVideoI2VPipeline");
|
||||
const std::string placeholder = "<|image_pad|>";
|
||||
std::string img_prompt;
|
||||
|
||||
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
|
||||
const auto& image = (*conditioner_params.ref_images)[i];
|
||||
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
|
||||
int height = static_cast<int>(image.shape()[1]);
|
||||
int width = static_cast<int>(image.shape()[0]);
|
||||
int min_pixels = static_cast<int>(4 * factor * factor);
|
||||
int max_pixels = static_cast<int>(16384 * factor * factor);
|
||||
int h_bar = std::max(static_cast<int>(factor), static_cast<int>(std::round(height / factor) * factor));
|
||||
int w_bar = std::max(static_cast<int>(factor), static_cast<int>(std::round(width / factor) * factor));
|
||||
|
||||
if (std::max(height, width) > 200 * std::min(height, width)) {
|
||||
LOG_WARN("LingBotVideo image aspect ratio is very large: %dx%d", width, height);
|
||||
}
|
||||
if (h_bar * w_bar > max_pixels) {
|
||||
double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
|
||||
h_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
|
||||
w_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
|
||||
} else if (h_bar * w_bar < min_pixels) {
|
||||
double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
|
||||
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
|
||||
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
|
||||
}
|
||||
|
||||
LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
|
||||
auto resized_image = clip_preprocess(image, w_bar, h_bar);
|
||||
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
|
||||
GGML_ASSERT(!image_embed.empty());
|
||||
|
||||
std::string image_prefix = prompt + img_prompt + "<|vision_start|>";
|
||||
int image_embed_idx = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
|
||||
image_embeds.emplace_back(image_embed_idx, image_embed);
|
||||
|
||||
img_prompt += "<|vision_start|>";
|
||||
int64_t num_image_tokens = image_embed.shape()[1];
|
||||
img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
|
||||
for (int j = 0; j < num_image_tokens; j++) {
|
||||
img_prompt += placeholder;
|
||||
}
|
||||
img_prompt += "<|vision_end|>";
|
||||
}
|
||||
prompt += img_prompt;
|
||||
}
|
||||
|
||||
prompt += conditioner_params.text;
|
||||
prompt_attn_range = {0, 0};
|
||||
prompt += "<|im_end|>\n<|im_start|>assistant\n";
|
||||
} else if (sd_version_is_qwen_image(version)) {
|
||||
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
|
||||
LOG_INFO("QwenImageEditPlusPipeline");
|
||||
prompt_template_encode_start_idx = 64;
|
||||
@@ -2141,6 +2218,41 @@ struct LLMEmbedder : public Conditioner {
|
||||
out_layers = {2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35};
|
||||
|
||||
prompt = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n";
|
||||
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
|
||||
std::string img_prompt = "";
|
||||
const std::string placeholder = "<|image_pad|>";
|
||||
|
||||
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
|
||||
const auto& image = (*conditioner_params.ref_images)[i];
|
||||
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
|
||||
int height = static_cast<int>(image.shape()[1]);
|
||||
int width = static_cast<int>(image.shape()[0]);
|
||||
double beta = std::sqrt((384.0 * 384.0) / (static_cast<double>(height) * static_cast<double>(width)));
|
||||
int h_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
|
||||
int w_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
|
||||
|
||||
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
|
||||
|
||||
auto resized_image = clip_preprocess(image, w_bar, h_bar);
|
||||
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
|
||||
GGML_ASSERT(!image_embed.empty());
|
||||
|
||||
std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
|
||||
int image_embed_idx = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
|
||||
image_embeds.emplace_back(image_embed_idx, image_embed);
|
||||
|
||||
img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
|
||||
int64_t num_image_tokens = image_embed.shape()[1];
|
||||
img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
|
||||
for (int j = 0; j < num_image_tokens; j++) {
|
||||
img_prompt += placeholder;
|
||||
}
|
||||
img_prompt += "<|vision_end|>";
|
||||
}
|
||||
prompt += img_prompt;
|
||||
}
|
||||
|
||||
prompt_attn_range.first = static_cast<int>(prompt.size());
|
||||
prompt += conditioner_params.text;
|
||||
|
||||
+10
-1
@@ -35,6 +35,7 @@ enum SDVersion {
|
||||
VERSION_WAN2,
|
||||
VERSION_WAN2_2_I2V,
|
||||
VERSION_WAN2_2_TI2V,
|
||||
VERSION_LINGBOT_VIDEO,
|
||||
VERSION_QWEN_IMAGE,
|
||||
VERSION_QWEN_IMAGE_LAYERED,
|
||||
VERSION_ANIMA,
|
||||
@@ -127,6 +128,13 @@ static inline bool sd_version_is_wan(SDVersion version) {
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_lingbot_video(SDVersion version) {
|
||||
if (version == VERSION_LINGBOT_VIDEO) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_qwen_image(SDVersion version) {
|
||||
if (version == VERSION_QWEN_IMAGE || version == VERSION_QWEN_IMAGE_LAYERED) {
|
||||
return true;
|
||||
@@ -226,7 +234,7 @@ static inline bool sd_version_uses_flux2_vae(SDVersion version) {
|
||||
}
|
||||
|
||||
static inline bool sd_version_uses_wan_vae(SDVersion version) {
|
||||
if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
|
||||
if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
@@ -249,6 +257,7 @@ static inline bool sd_version_is_dit(SDVersion version) {
|
||||
sd_version_is_ltxav(version) ||
|
||||
sd_version_is_sd3(version) ||
|
||||
sd_version_is_wan(version) ||
|
||||
sd_version_is_lingbot_video(version) ||
|
||||
sd_version_is_qwen_image(version) ||
|
||||
version == VERSION_HIDREAM_O1 ||
|
||||
sd_version_is_anima(version) ||
|
||||
|
||||
@@ -294,7 +294,7 @@ public:
|
||||
|
||||
auto net_0 = std::dynamic_pointer_cast<UnaryBlock>(blocks["net.0"]);
|
||||
auto net_2 = std::dynamic_pointer_cast<Linear>(blocks["net.2"]);
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
net_2->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
|
||||
@@ -759,6 +759,40 @@ namespace Rope {
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lingbot_video_ids(int t,
|
||||
int h,
|
||||
int w,
|
||||
int pt,
|
||||
int ph,
|
||||
int pw,
|
||||
int bs,
|
||||
int context_len) {
|
||||
auto vid_ids_repeated = gen_vid_ids(t, h, w, pt, ph, pw, bs, context_len + 1);
|
||||
|
||||
std::vector<std::vector<float>> txt_ids(bs * context_len, std::vector<float>(3, 0.0f));
|
||||
for (int i = 0; i < bs; ++i) {
|
||||
for (int j = 0; j < context_len; ++j) {
|
||||
txt_ids[i * context_len + j][0] = static_cast<float>(j + 1);
|
||||
}
|
||||
}
|
||||
|
||||
return concat_ids(vid_ids_repeated, txt_ids, bs);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_lingbot_video_pe(int t,
|
||||
int h,
|
||||
int w,
|
||||
int pt,
|
||||
int ph,
|
||||
int pw,
|
||||
int bs,
|
||||
int context_len,
|
||||
int theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_lingbot_video_ids(t, h, w, pt, ph, pw, bs, context_len);
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen2vl_ids(int grid_h,
|
||||
int grid_w,
|
||||
int merge_size,
|
||||
|
||||
@@ -199,7 +199,7 @@ namespace Boogu {
|
||||
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
|
||||
auto linear_3 = std::dynamic_pointer_cast<Linear>(blocks["linear_3"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
linear_2->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
@@ -259,7 +259,7 @@ namespace Boogu {
|
||||
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
|
||||
auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
to_out_0->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
@@ -383,7 +383,7 @@ namespace Boogu {
|
||||
auto instruct_out = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_out"]);
|
||||
auto img_out = std::dynamic_pointer_cast<Linear>(blocks["processor.img_out"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
to_out_0->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
|
||||
+33
-20
@@ -104,14 +104,15 @@ namespace DiT {
|
||||
return x;
|
||||
}
|
||||
|
||||
inline ggml_tensor* patchify(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
int pt,
|
||||
int ph,
|
||||
int pw,
|
||||
int64_t N = 1) {
|
||||
inline ggml_tensor* patchify_3d(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
int pt,
|
||||
int ph,
|
||||
int pw,
|
||||
int64_t N = 1,
|
||||
bool patch_last = true) {
|
||||
// x: [N*C, T, H, W]
|
||||
// return: [N, h*w, C*pt*ph*pw]
|
||||
// return: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
|
||||
int64_t C = x->ne[3] / N;
|
||||
int64_t T = x->ne[2];
|
||||
int64_t H = x->ne[1];
|
||||
@@ -123,15 +124,20 @@ namespace DiT {
|
||||
GGML_ASSERT(C * N == x->ne[3]);
|
||||
GGML_ASSERT(t_len * pt == T && h_len * ph == H && w_len * pw == W);
|
||||
|
||||
x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N); // [N*C*t_len, pt, h_len*ph, w_len*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, h_len*ph, pt, w_len*pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N); // [N*C*t_len*h_len, ph, pt, w_len*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt, ph, w_len*pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt*ph, w_len, pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, w_len, pt*ph, pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N); // [N, C, t_len*h_len*w_len, pt*ph*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, t_len*h_len*w_len, C, pt*ph*pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1); // [N, t_len*h_len*w_len, C*pt*ph*pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N); // [N*C*t_len, pt, h_len*ph, w_len*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, h_len*ph, pt, w_len*pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N); // [N*C*t_len*h_len, ph, pt, w_len*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt, ph, w_len*pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt*ph, w_len, pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, w_len, pt*ph, pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N); // [N, C, t_len*h_len*w_len, pt*ph*pw]
|
||||
if (patch_last) {
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, t_len*h_len*w_len, C, pt*ph*pw]
|
||||
} else {
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [N, t_len*h_len*w_len, pt*ph*pw, C]
|
||||
}
|
||||
// [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
|
||||
x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1);
|
||||
return x;
|
||||
}
|
||||
|
||||
@@ -142,16 +148,23 @@ namespace DiT {
|
||||
int64_t w_len,
|
||||
int pt,
|
||||
int ph,
|
||||
int pw) {
|
||||
// x: [N, t_len*h_len*w_len, C*pt*ph*pw]
|
||||
int pw,
|
||||
bool patch_last = true) {
|
||||
// x: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, pt*ph*pw*C]
|
||||
// return: [N*C, t_len*pt, h_len*ph, w_len*pw]
|
||||
int64_t N = x->ne[2];
|
||||
int64_t C = x->ne[0] / pt / ph / pw;
|
||||
|
||||
GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
|
||||
|
||||
x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
|
||||
if (patch_last) {
|
||||
x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
|
||||
} else {
|
||||
x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
|
||||
}
|
||||
|
||||
x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw]
|
||||
|
||||
+126
-27
@@ -267,7 +267,7 @@ namespace Krea2 {
|
||||
auto knorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["qknorm.knorm"]);
|
||||
auto wo = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
wo->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
@@ -421,29 +421,88 @@ namespace Krea2 {
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* vec,
|
||||
ggml_tensor* pe) {
|
||||
ggml_tensor* pe,
|
||||
ggml_tensor* vec_refs = nullptr,
|
||||
int64_t ref_start = -1) {
|
||||
auto mod = std::dynamic_pointer_cast<KreaDoubleSharedModulation>(blocks["mod"]);
|
||||
auto prenorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["prenorm"]);
|
||||
auto postnorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["postnorm"]);
|
||||
auto attn = std::dynamic_pointer_cast<KreaAttention>(blocks["attn"]);
|
||||
auto mlp = std::dynamic_pointer_cast<KreaSwiGLU>(blocks["mlp"]);
|
||||
|
||||
auto mods = mod->forward(ctx, vec);
|
||||
auto attn_input = Flux::modulate(ctx->ggml_ctx,
|
||||
prenorm->forward(ctx, x),
|
||||
mods[1],
|
||||
mods[0],
|
||||
true);
|
||||
auto attn_out = attn->forward(ctx, attn_input, pe);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2]));
|
||||
if (ref_start >= 0 && vec_refs) {
|
||||
// same as normal, but since vec is different for refs and the rest, needs a lot of views and concats
|
||||
auto mods_main = mod->forward(ctx, vec);
|
||||
auto mods_refs = mod->forward(ctx, vec_refs);
|
||||
|
||||
auto mlp_input = Flux::modulate(ctx->ggml_ctx,
|
||||
postnorm->forward(ctx, x),
|
||||
mods[4],
|
||||
mods[3],
|
||||
true);
|
||||
auto mlp_out = mlp->forward(ctx, mlp_input);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mods[5]));
|
||||
int64_t D = x->ne[0];
|
||||
int64_t N = x->ne[1];
|
||||
int64_t B = x->ne[2];
|
||||
size_t nb1 = x->nb[1];
|
||||
size_t nb2 = x->nb[2];
|
||||
|
||||
int64_t len_main = ref_start;
|
||||
int64_t len_refs = N - ref_start;
|
||||
|
||||
auto pre_x = prenorm->forward(ctx, x);
|
||||
|
||||
auto pre_x_main = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_main, B, nb1, nb2, 0);
|
||||
auto pre_x_refs = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_refs, B, nb1, nb2, len_main * nb1);
|
||||
|
||||
auto attn_in_main = Flux::modulate(ctx->ggml_ctx, pre_x_main, mods_main[1], mods_main[0], true);
|
||||
auto attn_in_refs = Flux::modulate(ctx->ggml_ctx, pre_x_refs, mods_refs[1], mods_refs[0], true);
|
||||
|
||||
auto attn_input = ggml_concat(ctx->ggml_ctx, attn_in_main, attn_in_refs, 1);
|
||||
|
||||
auto attn_out = attn->forward(ctx, attn_input, pe);
|
||||
|
||||
auto attn_out_main = ggml_view_3d(ctx->ggml_ctx, attn_out, D, len_main, B, attn_out->nb[1], attn_out->nb[2], 0);
|
||||
auto attn_out_refs = ggml_view_3d(ctx->ggml_ctx, attn_out, D, len_refs, B, attn_out->nb[1], attn_out->nb[2], len_main * attn_out->nb[1]);
|
||||
|
||||
auto res_main = ggml_mul(ctx->ggml_ctx, attn_out_main, mods_main[2]);
|
||||
auto res_refs = ggml_mul(ctx->ggml_ctx, attn_out_refs, mods_refs[2]);
|
||||
|
||||
auto attn_res = ggml_concat(ctx->ggml_ctx, res_main, res_refs, 1);
|
||||
|
||||
x = ggml_add(ctx->ggml_ctx, x, attn_res);
|
||||
|
||||
auto post_x = postnorm->forward(ctx, x);
|
||||
|
||||
auto post_x_main = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_main, B, post_x->nb[1], post_x->nb[2], 0);
|
||||
auto post_x_refs = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_refs, B, post_x->nb[1], post_x->nb[2], len_main * post_x->nb[1]);
|
||||
|
||||
auto mlp_in_main = Flux::modulate(ctx->ggml_ctx, post_x_main, mods_main[4], mods_main[3], true);
|
||||
auto mlp_in_refs = Flux::modulate(ctx->ggml_ctx, post_x_refs, mods_refs[4], mods_refs[3], true);
|
||||
|
||||
auto mlp_input = ggml_concat(ctx->ggml_ctx, mlp_in_main, mlp_in_refs, 1);
|
||||
auto mlp_out = mlp->forward(ctx, mlp_input);
|
||||
|
||||
auto mlp_out_main = ggml_view_3d(ctx->ggml_ctx, mlp_out, D, len_main, B, mlp_out->nb[1], mlp_out->nb[2], 0);
|
||||
auto mlp_out_refs = ggml_view_3d(ctx->ggml_ctx, mlp_out, D, len_refs, B, mlp_out->nb[1], mlp_out->nb[2], len_main * mlp_out->nb[1]);
|
||||
|
||||
auto mlp_res_main = ggml_mul(ctx->ggml_ctx, mlp_out_main, mods_main[5]);
|
||||
auto mlp_res_refs = ggml_mul(ctx->ggml_ctx, mlp_out_refs, mods_refs[5]);
|
||||
|
||||
auto mlp_res = ggml_concat(ctx->ggml_ctx, mlp_res_main, mlp_res_refs, 1);
|
||||
x = ggml_add(ctx->ggml_ctx, x, mlp_res);
|
||||
} else {
|
||||
auto mods = mod->forward(ctx, vec);
|
||||
auto attn_input = Flux::modulate(ctx->ggml_ctx,
|
||||
prenorm->forward(ctx, x),
|
||||
mods[1],
|
||||
mods[0],
|
||||
true);
|
||||
auto attn_out = attn->forward(ctx, attn_input, pe);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2]));
|
||||
|
||||
auto mlp_input = Flux::modulate(ctx->ggml_ctx,
|
||||
postnorm->forward(ctx, x),
|
||||
mods[4],
|
||||
mods[3],
|
||||
true);
|
||||
auto mlp_out = mlp->forward(ctx, mlp_input);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mods[5]));
|
||||
}
|
||||
return x;
|
||||
}
|
||||
|
||||
@@ -555,7 +614,8 @@ namespace Krea2 {
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* timestep,
|
||||
ggml_tensor* context,
|
||||
ggml_tensor* pe) {
|
||||
ggml_tensor* pe,
|
||||
std::vector<ggml_tensor*> ref_latents = {}) {
|
||||
int64_t W = x->ne[0];
|
||||
int64_t H = x->ne[1];
|
||||
int64_t N = x->ne[3];
|
||||
@@ -570,6 +630,13 @@ namespace Krea2 {
|
||||
|
||||
auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size, true);
|
||||
int64_t img_len = img->ne[1];
|
||||
if (ref_latents.size() > 0) {
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size, true);
|
||||
img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
|
||||
}
|
||||
}
|
||||
int64_t ref_len = img->ne[1] - img_len;
|
||||
img = first->forward(ctx, img);
|
||||
|
||||
auto t = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.timestep_dim), 10000, 1000.f);
|
||||
@@ -577,19 +644,30 @@ namespace Krea2 {
|
||||
t = ggml_reshape_3d(ctx->ggml_ctx, t, t->ne[0], 1, t->ne[1]);
|
||||
auto tvec = tproj->forward(ctx, t);
|
||||
|
||||
ggml_tensor* tvec_0 = nullptr;
|
||||
if (ref_latents.size() > 0) {
|
||||
// "index_timestep_zero" mode: use timestep = 0 for ref latents
|
||||
auto timestep_0 = ggml_scale(ctx->ggml_ctx, timestep, 0.0f);
|
||||
auto t_0 = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_0, static_cast<int>(config.timestep_dim), 10000, 1000.f);
|
||||
t_0 = tmlp->forward(ctx, t_0);
|
||||
t_0 = ggml_reshape_3d(ctx->ggml_ctx, t_0, t_0->ne[0], 1, t_0->ne[1]);
|
||||
tvec_0 = tproj->forward(ctx, t_0);
|
||||
}
|
||||
|
||||
auto txt = txtfusion->forward(ctx, context);
|
||||
txt = txtmlp->forward(ctx, txt);
|
||||
int64_t txt_len = txt->ne[1];
|
||||
|
||||
auto hidden_states = ggml_concat(ctx->ggml_ctx, txt, img, 1);
|
||||
int64_t ref_start = hidden_states->ne[1] - ref_len;
|
||||
for (int i = 0; i < config.layers; ++i) {
|
||||
auto block = std::dynamic_pointer_cast<KreaSingleStreamBlock>(blocks["blocks." + std::to_string(i)]);
|
||||
hidden_states = block->forward(ctx, hidden_states, tvec, pe);
|
||||
hidden_states = block->forward(ctx, hidden_states, tvec, pe, tvec_0, ref_start);
|
||||
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "krea2.blocks." + std::to_string(i), "hidden_states");
|
||||
}
|
||||
|
||||
hidden_states = last->forward(ctx, hidden_states, t);
|
||||
hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, txt_len, txt_len + img_len);
|
||||
hidden_states = last->forward(ctx, hidden_states, t);
|
||||
hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, true);
|
||||
return hidden_states;
|
||||
}
|
||||
@@ -601,10 +679,16 @@ namespace Krea2 {
|
||||
int bs,
|
||||
int context_len,
|
||||
float theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
const std::vector<int>& axes_dim,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
Rope::RefIndexMode ref_index_mode) {
|
||||
auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {});
|
||||
auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false);
|
||||
auto ids = Rope::concat_ids(txt_ids, img_ids, bs);
|
||||
if (ref_latents.size() > 0) {
|
||||
auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0);
|
||||
ids = Rope::concat_ids(ids, refs_ids, bs);
|
||||
}
|
||||
return Rope::embed_nd(ids, bs, theta, axes_dim);
|
||||
}
|
||||
|
||||
@@ -633,7 +717,9 @@ namespace Krea2 {
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
|
||||
const sd::Tensor<float>& timesteps_tensor,
|
||||
const sd::Tensor<float>& context_tensor) {
|
||||
const sd::Tensor<float>& context_tensor,
|
||||
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
|
||||
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) {
|
||||
ggml_cgraph* gf = new_graph_custom(KREA2_GRAPH_SIZE);
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
||||
@@ -641,19 +727,27 @@ namespace Krea2 {
|
||||
GGML_ASSERT(!context_tensor.empty());
|
||||
ggml_tensor* context = make_input(context_tensor);
|
||||
|
||||
std::vector<ggml_tensor*> ref_latents;
|
||||
ref_latents.reserve(ref_latents_tensor.size());
|
||||
for (const auto& ref_latent_tensor : ref_latents_tensor) {
|
||||
ref_latents.push_back(make_input(ref_latent_tensor));
|
||||
}
|
||||
|
||||
pe_vec = gen_krea2_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
config.theta,
|
||||
config.axes_dim);
|
||||
config.axes_dim,
|
||||
ref_latents,
|
||||
ref_index_mode);
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
set_backend_tensor_data(pe, pe_vec.data());
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe);
|
||||
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents);
|
||||
ggml_build_forward_expand(gf, out);
|
||||
return gf;
|
||||
}
|
||||
@@ -661,9 +755,11 @@ namespace Krea2 {
|
||||
sd::Tensor<float> compute(int n_threads,
|
||||
const sd::Tensor<float>& x,
|
||||
const sd::Tensor<float>& timesteps,
|
||||
const sd::Tensor<float>& context) {
|
||||
const sd::Tensor<float>& context,
|
||||
const std::vector<sd::Tensor<float>>& ref_latents = {},
|
||||
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) {
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_graph(x, timesteps, context);
|
||||
return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
|
||||
};
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
|
||||
}
|
||||
@@ -672,10 +768,13 @@ namespace Krea2 {
|
||||
const DiffusionParams& diffusion_params) override {
|
||||
GGML_ASSERT(diffusion_params.x != nullptr);
|
||||
GGML_ASSERT(diffusion_params.timesteps != nullptr);
|
||||
static const std::vector<sd::Tensor<float>> empty_ref_latents;
|
||||
return compute(n_threads,
|
||||
*diffusion_params.x,
|
||||
*diffusion_params.timesteps,
|
||||
tensor_or_empty(diffusion_params.context));
|
||||
tensor_or_empty(diffusion_params.context),
|
||||
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
|
||||
diffusion_params.ref_index_mode);
|
||||
}
|
||||
};
|
||||
} // namespace Krea2
|
||||
|
||||
@@ -0,0 +1,692 @@
|
||||
#ifndef __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__
|
||||
#define __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__
|
||||
|
||||
#include <algorithm>
|
||||
#include <cinttypes>
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "core/util.h"
|
||||
#include "model/common/rope.hpp"
|
||||
#include "model/diffusion/dit.hpp"
|
||||
#include "model/diffusion/flux.hpp"
|
||||
#include "model/diffusion/model.hpp"
|
||||
#include "model/diffusion/qwen_image.hpp"
|
||||
|
||||
namespace LingBotVideo {
|
||||
constexpr int LINGBOT_VIDEO_GRAPH_SIZE = 65536;
|
||||
|
||||
struct LingBotVideoConfig {
|
||||
int patch_t = 1;
|
||||
int patch_h = 2;
|
||||
int patch_w = 2;
|
||||
int64_t in_channels = 16;
|
||||
int64_t out_channels = 16;
|
||||
int64_t hidden_size = 2048;
|
||||
int64_t num_attention_heads = 16;
|
||||
int64_t depth = 24;
|
||||
int64_t intermediate_size = 6144;
|
||||
int64_t text_dim = 2560;
|
||||
int64_t freq_dim = 256;
|
||||
float norm_eps = 1e-6f;
|
||||
int rope_theta = 256;
|
||||
std::vector<int> axes_dim = {32, 48, 48};
|
||||
int axes_dim_sum = 128;
|
||||
bool qkv_bias = false;
|
||||
bool out_bias = true;
|
||||
bool patch_embed_bias = true;
|
||||
bool timestep_mlp_bias = true;
|
||||
int64_t num_experts = 0;
|
||||
int64_t num_experts_per_tok = 8;
|
||||
int64_t moe_intermediate_size = 512;
|
||||
int64_t decoder_sparse_step = 1;
|
||||
int64_t n_shared_experts = 0;
|
||||
bool norm_topk_prob = true;
|
||||
float routed_scaling_factor = 1.0f;
|
||||
int64_t n_group = 0;
|
||||
int64_t topk_group = 0;
|
||||
std::set<int> sparse_layers;
|
||||
|
||||
static LingBotVideoConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& prefix) {
|
||||
LingBotVideoConfig config;
|
||||
config.depth = 0;
|
||||
|
||||
for (const auto& [name, tensor_storage] : tensor_storage_map) {
|
||||
if (!starts_with(name, prefix)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (ends_with(name, "patch_embedder.weight") && tensor_storage.n_dims == 2) {
|
||||
int64_t patch_dim = tensor_storage.ne[0];
|
||||
config.hidden_size = tensor_storage.ne[1];
|
||||
int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w;
|
||||
if (patch_dim % patch_volume == 0) {
|
||||
config.in_channels = patch_dim / patch_volume;
|
||||
}
|
||||
} else if (ends_with(name, "text_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
|
||||
config.text_dim = tensor_storage.ne[0];
|
||||
} else if (ends_with(name, "time_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
|
||||
config.freq_dim = tensor_storage.ne[0];
|
||||
} else if (ends_with(name, "blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
|
||||
int64_t head_dim = tensor_storage.ne[0];
|
||||
if (head_dim > 0) {
|
||||
config.num_attention_heads = config.hidden_size / head_dim;
|
||||
}
|
||||
} else if (name.find(".attn.to_q.bias") != std::string::npos) {
|
||||
config.qkv_bias = true;
|
||||
} else if (name.find(".ffn.gate_proj.weight") != std::string::npos && tensor_storage.n_dims == 2) {
|
||||
config.intermediate_size = tensor_storage.ne[1];
|
||||
} else if (name.find(".ffn.experts.w1") != std::string::npos && tensor_storage.n_dims == 3) {
|
||||
config.num_experts = tensor_storage.ne[2];
|
||||
config.moe_intermediate_size = tensor_storage.ne[1];
|
||||
} else if (name.find(".ffn.shared_experts.gate_proj.weight") != std::string::npos && tensor_storage.n_dims == 2) {
|
||||
if (config.moe_intermediate_size > 0) {
|
||||
config.n_shared_experts = tensor_storage.ne[1] / config.moe_intermediate_size;
|
||||
}
|
||||
} else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
|
||||
int64_t out_dim = tensor_storage.ne[1];
|
||||
int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w;
|
||||
config.out_channels = patch_volume > 0 ? out_dim / patch_volume : config.out_channels;
|
||||
}
|
||||
|
||||
size_t block_pos = name.find("blocks.");
|
||||
if (block_pos != std::string::npos) {
|
||||
auto items = split_string(name.substr(block_pos), '.');
|
||||
if (items.size() > 1) {
|
||||
int block_index = atoi(items[1].c_str());
|
||||
if (block_index + 1 > config.depth) {
|
||||
config.depth = block_index + 1;
|
||||
}
|
||||
if (name.find("blocks." + std::to_string(block_index) + ".ffn.experts.w1") != std::string::npos) {
|
||||
config.sparse_layers.insert(block_index);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (config.depth == 0) {
|
||||
config.depth = 24;
|
||||
}
|
||||
config.axes_dim_sum = 0;
|
||||
for (int axis_dim : config.axes_dim) {
|
||||
config.axes_dim_sum += axis_dim;
|
||||
}
|
||||
if (!config.sparse_layers.empty()) {
|
||||
config.num_experts = 128;
|
||||
config.num_experts_per_tok = 8;
|
||||
config.moe_intermediate_size = 768;
|
||||
config.decoder_sparse_step = 1;
|
||||
config.n_shared_experts = 1;
|
||||
config.norm_topk_prob = true;
|
||||
config.n_group = 4;
|
||||
config.topk_group = 2;
|
||||
config.routed_scaling_factor = 2.5f;
|
||||
}
|
||||
LOG_DEBUG("lingbot_video: depth = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64 ", text_dim = %" PRId64 ", experts = %" PRId64 ", experts_per_tok = %" PRId64 ", n_group = %" PRId64 ", topk_group = %" PRId64 ", route_scale = %.2f, sparse_layers = %zu",
|
||||
config.depth,
|
||||
config.hidden_size,
|
||||
config.num_attention_heads,
|
||||
config.text_dim,
|
||||
config.num_experts,
|
||||
config.num_experts_per_tok,
|
||||
config.n_group,
|
||||
config.topk_group,
|
||||
config.routed_scaling_factor,
|
||||
config.sparse_layers.size());
|
||||
return config;
|
||||
}
|
||||
};
|
||||
|
||||
struct LingBotVideoTextEmbedder : public GGMLBlock {
|
||||
LingBotVideoTextEmbedder(int64_t text_dim,
|
||||
int64_t hidden_size,
|
||||
float eps = 1e-6f) {
|
||||
blocks["norm"] = std::make_shared<RMSNorm>(text_dim, eps);
|
||||
blocks["linear_1"] = std::make_shared<Linear>(text_dim, hidden_size, true);
|
||||
blocks["linear_2"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
|
||||
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
|
||||
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
|
||||
|
||||
x = norm->forward(ctx, x);
|
||||
x = linear_1->forward(ctx, x);
|
||||
x = ggml_silu_inplace(ctx->ggml_ctx, x);
|
||||
x = linear_2->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct LingBotVideoAttention : public GGMLBlock {
|
||||
int64_t num_heads;
|
||||
int64_t head_dim;
|
||||
|
||||
LingBotVideoAttention(int64_t hidden_size,
|
||||
int64_t num_heads,
|
||||
bool qkv_bias = false,
|
||||
bool out_bias = true,
|
||||
float eps = 1e-6f)
|
||||
: num_heads(num_heads),
|
||||
head_dim(hidden_size / num_heads) {
|
||||
int64_t inner_dim = num_heads * head_dim;
|
||||
blocks["to_q"] = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
|
||||
blocks["to_k"] = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
|
||||
blocks["to_v"] = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
|
||||
blocks["norm_q"] = std::make_shared<RMSNorm>(head_dim, eps);
|
||||
blocks["norm_k"] = std::make_shared<RMSNorm>(head_dim, eps);
|
||||
blocks["to_out"] = std::make_shared<Linear>(inner_dim, hidden_size, out_bias);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* pe,
|
||||
ggml_tensor* attention_mask = nullptr) {
|
||||
// x: [N, video_tokens + text_tokens, hidden_size]
|
||||
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 norm_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
|
||||
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
|
||||
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out"]);
|
||||
|
||||
int64_t S = x->ne[1];
|
||||
int64_t N = x->ne[2];
|
||||
|
||||
auto q = to_q->forward(ctx, x);
|
||||
auto k = to_k->forward(ctx, x);
|
||||
auto v = to_v->forward(ctx, x);
|
||||
|
||||
q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, S, N);
|
||||
k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, S, N);
|
||||
v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, S, N);
|
||||
|
||||
q = norm_q->forward(ctx, q);
|
||||
k = norm_k->forward(ctx, k);
|
||||
|
||||
x = Rope::attention(ctx, q, k, v, pe, attention_mask);
|
||||
x = to_out->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct LingBotVideoMLP : public UnaryBlock {
|
||||
LingBotVideoMLP(int64_t hidden_size,
|
||||
int64_t intermediate_size) {
|
||||
blocks["gate_proj"] = std::make_shared<Linear>(hidden_size, intermediate_size, false);
|
||||
blocks["up_proj"] = std::make_shared<Linear>(hidden_size, intermediate_size, false);
|
||||
blocks["down_proj"] = std::make_shared<Linear>(intermediate_size, hidden_size, false);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
|
||||
auto up_proj = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
|
||||
auto down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
|
||||
|
||||
auto gate = gate_proj->forward(ctx, x);
|
||||
gate = ggml_silu_inplace(ctx->ggml_ctx, gate);
|
||||
auto up = up_proj->forward(ctx, x);
|
||||
x = ggml_mul(ctx->ggml_ctx, gate, up);
|
||||
x = down_proj->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct LingBotVideoSparseMoeBlock : public UnaryBlock {
|
||||
int64_t hidden_size;
|
||||
int64_t intermediate_size;
|
||||
int64_t num_experts;
|
||||
int64_t num_experts_per_tok;
|
||||
bool has_shared_experts;
|
||||
bool norm_topk_prob;
|
||||
float routed_scaling_factor;
|
||||
int64_t n_group;
|
||||
int64_t topk_group;
|
||||
std::vector<float> group_expert_mask_vec;
|
||||
bool has_correction_bias = false;
|
||||
|
||||
LingBotVideoSparseMoeBlock(const LingBotVideoConfig& config)
|
||||
: hidden_size(config.hidden_size),
|
||||
intermediate_size(config.moe_intermediate_size),
|
||||
num_experts(config.num_experts),
|
||||
num_experts_per_tok(config.num_experts_per_tok),
|
||||
has_shared_experts(config.n_shared_experts > 0),
|
||||
norm_topk_prob(config.norm_topk_prob),
|
||||
routed_scaling_factor(config.routed_scaling_factor),
|
||||
n_group(config.n_group),
|
||||
topk_group(config.topk_group) {
|
||||
if (n_group > 1) {
|
||||
GGML_ASSERT(num_experts % n_group == 0);
|
||||
int64_t experts_per_group = num_experts / n_group;
|
||||
group_expert_mask_vec.assign(static_cast<size_t>(num_experts * n_group), 0.f);
|
||||
for (int64_t group = 0; group < n_group; ++group) {
|
||||
int64_t expert_begin = group * experts_per_group;
|
||||
int64_t expert_end = expert_begin + experts_per_group;
|
||||
for (int64_t expert = expert_begin; expert < expert_end; ++expert) {
|
||||
group_expert_mask_vec[static_cast<size_t>(group * num_experts + expert)] = 1.f;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (has_shared_experts) {
|
||||
blocks["shared_experts"] = std::make_shared<LingBotVideoMLP>(hidden_size,
|
||||
intermediate_size * config.n_shared_experts);
|
||||
}
|
||||
}
|
||||
|
||||
void init_params(ggml_context* ctx,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "") override {
|
||||
ggml_type router_type = get_type(prefix + "router.weight", tensor_storage_map, GGML_TYPE_F32);
|
||||
ggml_type w1_type = get_type(prefix + "experts.w1", tensor_storage_map, GGML_TYPE_F32);
|
||||
ggml_type w2_type = get_type(prefix + "experts.w2", tensor_storage_map, GGML_TYPE_F32);
|
||||
ggml_type w3_type = get_type(prefix + "experts.w3", tensor_storage_map, GGML_TYPE_F32);
|
||||
params["router.weight"] = ggml_new_tensor_2d(ctx, router_type, hidden_size, num_experts);
|
||||
if (tensor_storage_map.find(prefix + "router.e_score_correction_bias") != tensor_storage_map.end()) {
|
||||
params["router.e_score_correction_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts);
|
||||
has_correction_bias = true;
|
||||
}
|
||||
params["experts.w1"] = ggml_new_tensor_3d(ctx, w1_type, hidden_size, intermediate_size, num_experts);
|
||||
params["experts.w2"] = ggml_new_tensor_3d(ctx, w2_type, intermediate_size, hidden_size, num_experts);
|
||||
params["experts.w3"] = ggml_new_tensor_3d(ctx, w3_type, hidden_size, intermediate_size, num_experts);
|
||||
}
|
||||
|
||||
ggml_tensor* expert_linear(GGMLRunnerContext* ctx,
|
||||
const std::string& weight_name,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* selected_experts) {
|
||||
return ggml_mul_mat_id(ctx->ggml_ctx, params[weight_name], x, selected_experts);
|
||||
}
|
||||
|
||||
ggml_tensor* select_experts(GGMLRunnerContext* ctx, ggml_tensor* choice_scores) {
|
||||
ggml_context* gctx = ctx->ggml_ctx;
|
||||
if (n_group <= 1 || topk_group <= 0) {
|
||||
return ggml_argsort_top_k(gctx, choice_scores, static_cast<int>(num_experts_per_tok));
|
||||
}
|
||||
|
||||
GGML_ASSERT(choice_scores->ne[0] == num_experts);
|
||||
GGML_ASSERT(num_experts % n_group == 0);
|
||||
GGML_ASSERT(topk_group > 0 && topk_group <= n_group);
|
||||
GGML_ASSERT(!group_expert_mask_vec.empty());
|
||||
|
||||
const int64_t n_token_total = choice_scores->ne[1];
|
||||
const int64_t experts_per_group = num_experts / n_group;
|
||||
const int group_score_k = 2;
|
||||
GGML_ASSERT(experts_per_group >= group_score_k);
|
||||
|
||||
ggml_tensor* grouped_scores = ggml_reshape_3d(gctx, choice_scores, experts_per_group, n_group, n_token_total);
|
||||
ggml_tensor* group_top_ids = ggml_argsort_top_k(gctx, grouped_scores, group_score_k);
|
||||
grouped_scores = ggml_reshape_3d(gctx, grouped_scores, 1, experts_per_group, n_group * n_token_total);
|
||||
group_top_ids = ggml_cont(gctx, group_top_ids);
|
||||
group_top_ids = ggml_reshape_2d(gctx, group_top_ids, group_score_k, n_group * n_token_total);
|
||||
|
||||
ggml_tensor* group_top_values = ggml_get_rows(gctx, grouped_scores, group_top_ids);
|
||||
group_top_values = ggml_reshape_3d(gctx, group_top_values, group_score_k, n_group, n_token_total);
|
||||
|
||||
ggml_tensor* group_scores = nullptr;
|
||||
for (int rank = 0; rank < group_score_k; ++rank) {
|
||||
ggml_tensor* value = ggml_view_3d(gctx,
|
||||
group_top_values,
|
||||
1,
|
||||
n_group,
|
||||
n_token_total,
|
||||
group_top_values->nb[1],
|
||||
group_top_values->nb[2],
|
||||
rank * group_top_values->nb[0]);
|
||||
group_scores = group_scores == nullptr ? value : ggml_add(gctx, group_scores, value);
|
||||
}
|
||||
group_scores = ggml_reshape_2d(gctx, group_scores, n_group, n_token_total);
|
||||
|
||||
ggml_tensor* selected_groups = ggml_argsort_top_k(gctx, group_scores, static_cast<int>(topk_group));
|
||||
selected_groups = ggml_cont(gctx, selected_groups);
|
||||
|
||||
ggml_tensor* group_expert_mask = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, num_experts, n_group, 1);
|
||||
ctx->bind_backend_tensor_data(group_expert_mask, group_expert_mask_vec.data());
|
||||
ggml_tensor* group_expert_mask_template = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, num_experts, n_group, n_token_total);
|
||||
group_expert_mask = ggml_repeat(gctx, group_expert_mask, group_expert_mask_template);
|
||||
|
||||
ggml_tensor* selected_group_masks = ggml_get_rows(gctx, group_expert_mask, selected_groups);
|
||||
ggml_tensor* selected_mask = nullptr;
|
||||
for (int64_t rank = 0; rank < topk_group; ++rank) {
|
||||
ggml_tensor* mask = ggml_view_3d(gctx,
|
||||
selected_group_masks,
|
||||
num_experts,
|
||||
1,
|
||||
n_token_total,
|
||||
selected_group_masks->nb[1],
|
||||
selected_group_masks->nb[2],
|
||||
rank * selected_group_masks->nb[1]);
|
||||
selected_mask = selected_mask == nullptr ? mask : ggml_add(gctx, selected_mask, mask);
|
||||
}
|
||||
selected_mask = ggml_reshape_2d(gctx, selected_mask, num_experts, n_token_total);
|
||||
|
||||
ggml_tensor* excluded_group_mask = ggml_sub(gctx, selected_mask, ggml_ext_ones_like(gctx, selected_mask));
|
||||
ggml_tensor* masked_scores = ggml_add(gctx, choice_scores, ggml_scale(gctx, excluded_group_mask, 1.0e9f));
|
||||
return ggml_argsort_top_k(gctx, masked_scores, static_cast<int>(num_experts_per_tok));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [N, tokens, hidden_size]
|
||||
GGML_ASSERT(num_experts > 0);
|
||||
GGML_ASSERT(num_experts_per_tok > 0 && num_experts_per_tok <= num_experts);
|
||||
|
||||
ggml_tensor* residual = x;
|
||||
const int64_t n_token = x->ne[1];
|
||||
const int64_t N = x->ne[2];
|
||||
const int64_t n_token_total = n_token * N;
|
||||
|
||||
ggml_tensor* router_logits = ggml_mul_mat(ctx->ggml_ctx, params["router.weight"], x);
|
||||
router_logits = ggml_reshape_2d(ctx->ggml_ctx, router_logits, num_experts, n_token_total);
|
||||
ggml_tensor* probs = ggml_sigmoid(ctx->ggml_ctx, router_logits);
|
||||
ggml_tensor* choice_scores = probs;
|
||||
if (has_correction_bias) {
|
||||
choice_scores = ggml_add(ctx->ggml_ctx, choice_scores, params["router.e_score_correction_bias"]);
|
||||
}
|
||||
|
||||
ggml_tensor* selected_experts = select_experts(ctx, choice_scores);
|
||||
ggml_tensor* weights = ggml_get_rows(ctx->ggml_ctx,
|
||||
ggml_reshape_3d(ctx->ggml_ctx, probs, 1, num_experts, n_token_total),
|
||||
selected_experts);
|
||||
weights = ggml_reshape_2d(ctx->ggml_ctx, weights, num_experts_per_tok, n_token_total);
|
||||
if (norm_topk_prob && num_experts_per_tok > 1) {
|
||||
auto weights_sum = ggml_sum_rows(ctx->ggml_ctx, weights);
|
||||
weights_sum = ggml_clamp(ctx->ggml_ctx, weights_sum, 6.103515625e-5f, INFINITY);
|
||||
weights = ggml_div(ctx->ggml_ctx, weights, weights_sum);
|
||||
}
|
||||
if (routed_scaling_factor != 1.0f) {
|
||||
weights = ggml_scale(ctx->ggml_ctx, weights, routed_scaling_factor);
|
||||
}
|
||||
weights = ggml_reshape_3d(ctx->ggml_ctx, weights, 1, num_experts_per_tok, n_token_total);
|
||||
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size, 1, n_token_total);
|
||||
auto gate = expert_linear(ctx, "experts.w1", x, selected_experts);
|
||||
gate = ggml_silu_inplace(ctx->ggml_ctx, gate);
|
||||
auto up = expert_linear(ctx, "experts.w3", x, selected_experts);
|
||||
auto act = ggml_mul(ctx->ggml_ctx, gate, up);
|
||||
auto out = expert_linear(ctx, "experts.w2", act, selected_experts);
|
||||
out = ggml_mul(ctx->ggml_ctx, out, weights);
|
||||
|
||||
ggml_tensor* summed = nullptr;
|
||||
for (int64_t i = 0; i < num_experts_per_tok; ++i) {
|
||||
auto expert_out = ggml_view_2d(ctx->ggml_ctx,
|
||||
out,
|
||||
hidden_size,
|
||||
n_token_total,
|
||||
out->nb[2],
|
||||
i * out->nb[1]);
|
||||
summed = summed == nullptr ? expert_out : ggml_add(ctx->ggml_ctx, summed, expert_out);
|
||||
}
|
||||
if (num_experts_per_tok == 1) {
|
||||
summed = ggml_cont(ctx->ggml_ctx, summed);
|
||||
}
|
||||
summed = ggml_reshape_3d(ctx->ggml_ctx, summed, hidden_size, n_token, N);
|
||||
|
||||
if (has_shared_experts) {
|
||||
auto shared_experts = std::dynamic_pointer_cast<LingBotVideoMLP>(blocks["shared_experts"]);
|
||||
summed = ggml_add(ctx->ggml_ctx, summed, shared_experts->forward(ctx, residual));
|
||||
}
|
||||
return summed;
|
||||
}
|
||||
};
|
||||
|
||||
struct LingBotVideoBlock : public GGMLBlock {
|
||||
int64_t hidden_size;
|
||||
|
||||
LingBotVideoBlock(const LingBotVideoConfig& config,
|
||||
bool sparse)
|
||||
: hidden_size(config.hidden_size) {
|
||||
blocks["norm1"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
|
||||
blocks["attn"] = std::make_shared<LingBotVideoAttention>(config.hidden_size,
|
||||
config.num_attention_heads,
|
||||
config.qkv_bias,
|
||||
config.out_bias,
|
||||
config.norm_eps);
|
||||
blocks["norm_post_attn"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
|
||||
blocks["norm2"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
|
||||
if (sparse) {
|
||||
blocks["ffn"] = std::make_shared<LingBotVideoSparseMoeBlock>(config);
|
||||
} else {
|
||||
blocks["ffn"] = std::make_shared<LingBotVideoMLP>(config.hidden_size, config.intermediate_size);
|
||||
}
|
||||
blocks["norm_post_ffn"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
|
||||
}
|
||||
|
||||
void init_params(ggml_context* ctx,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "") override {
|
||||
params["scale_shift_table"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size * 6, 1);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* temb6,
|
||||
ggml_tensor* pe,
|
||||
ggml_tensor* attention_mask = nullptr) {
|
||||
// x: [N, tokens, hidden_size], temb6: [N, tokens, 6 * hidden_size]
|
||||
auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
|
||||
auto attn = std::dynamic_pointer_cast<LingBotVideoAttention>(blocks["attn"]);
|
||||
auto norm_post_attn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_post_attn"]);
|
||||
auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
|
||||
auto ffn = std::dynamic_pointer_cast<UnaryBlock>(blocks["ffn"]);
|
||||
auto norm_post_ffn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_post_ffn"]);
|
||||
|
||||
auto table = ggml_reshape_3d(ctx->ggml_ctx, params["scale_shift_table"], hidden_size * 6, 1, 1);
|
||||
auto mod = ggml_add(ctx->ggml_ctx, temb6, table);
|
||||
auto mods = ggml_ext_chunk(ctx->ggml_ctx, mod, 6, 0);
|
||||
|
||||
auto shift_msa = mods[0];
|
||||
auto scale_msa = mods[1];
|
||||
auto gate_msa = ggml_tanh(ctx->ggml_ctx, mods[2]);
|
||||
auto shift_mlp = mods[3];
|
||||
auto scale_mlp = mods[4];
|
||||
auto gate_mlp = ggml_tanh(ctx->ggml_ctx, mods[5]);
|
||||
|
||||
auto attn_in = Flux::modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa, true);
|
||||
auto attn_out = attn->forward(ctx, attn_in, pe, attention_mask);
|
||||
attn_out = norm_post_attn->forward(ctx, attn_out);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_msa, attn_out));
|
||||
|
||||
auto ffn_in = Flux::modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp, true);
|
||||
auto ffn_out = ffn->forward(ctx, ffn_in);
|
||||
ffn_out = norm_post_ffn->forward(ctx, ffn_out);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_mlp, ffn_out));
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
class LingBotVideoModel : public GGMLBlock {
|
||||
public:
|
||||
LingBotVideoConfig config;
|
||||
|
||||
LingBotVideoModel() = default;
|
||||
LingBotVideoModel(LingBotVideoConfig config)
|
||||
: config(config) {
|
||||
int64_t patch_dim = config.patch_t * config.patch_h * config.patch_w * config.in_channels;
|
||||
blocks["patch_embedder"] = std::make_shared<Linear>(patch_dim, config.hidden_size, config.patch_embed_bias);
|
||||
blocks["time_embedder"] = std::make_shared<Qwen::TimestepEmbedding>(config.freq_dim,
|
||||
config.hidden_size,
|
||||
config.hidden_size,
|
||||
0,
|
||||
config.timestep_mlp_bias);
|
||||
blocks["time_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size, true);
|
||||
blocks["text_embedder"] = std::make_shared<LingBotVideoTextEmbedder>(config.text_dim,
|
||||
config.hidden_size,
|
||||
config.norm_eps);
|
||||
for (int i = 0; i < config.depth; i++) {
|
||||
bool sparse = config.sparse_layers.find(i) != config.sparse_layers.end();
|
||||
blocks["blocks." + std::to_string(i)] = std::make_shared<LingBotVideoBlock>(config, sparse);
|
||||
}
|
||||
blocks["norm_out"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
|
||||
blocks["norm_out_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 2 * config.hidden_size, true);
|
||||
blocks["proj_out"] = std::make_shared<Linear>(config.hidden_size,
|
||||
config.patch_t * config.patch_h * config.patch_w * config.out_channels,
|
||||
true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* timestep,
|
||||
ggml_tensor* context,
|
||||
ggml_tensor* pe) {
|
||||
// x: [N*C, T, H, W], context: [N, text_tokens, text_dim]
|
||||
GGML_ASSERT(context != nullptr);
|
||||
GGML_ASSERT(x->ne[3] == config.in_channels);
|
||||
GGML_ASSERT(x->ne[2] % config.patch_t == 0);
|
||||
GGML_ASSERT(x->ne[1] % config.patch_h == 0);
|
||||
GGML_ASSERT(x->ne[0] % config.patch_w == 0);
|
||||
|
||||
auto patch_embedder = std::dynamic_pointer_cast<Linear>(blocks["patch_embedder"]);
|
||||
auto time_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["time_embedder"]);
|
||||
auto time_modulation = std::dynamic_pointer_cast<Linear>(blocks["time_modulation.1"]);
|
||||
auto text_embedder = std::dynamic_pointer_cast<LingBotVideoTextEmbedder>(blocks["text_embedder"]);
|
||||
auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
|
||||
auto norm_out_modulation = std::dynamic_pointer_cast<Linear>(blocks["norm_out_modulation.1"]);
|
||||
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 T = x->ne[2];
|
||||
int64_t N = 1;
|
||||
int64_t t_len = T / config.patch_t;
|
||||
int64_t h_len = H / config.patch_h;
|
||||
int64_t w_len = W / config.patch_w;
|
||||
int64_t n_img = t_len * h_len * w_len;
|
||||
|
||||
auto img = DiT::patchify_3d(ctx->ggml_ctx, x, config.patch_t, config.patch_h, config.patch_w, N, false);
|
||||
img = patch_embedder->forward(ctx, img);
|
||||
|
||||
auto txt = text_embedder->forward(ctx, context);
|
||||
auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1);
|
||||
|
||||
auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx,
|
||||
timestep,
|
||||
static_cast<int>(config.freq_dim),
|
||||
10000,
|
||||
1.0f);
|
||||
auto t_emb = time_embedder->forward(ctx, timestep_proj);
|
||||
auto temb_template = ggml_new_tensor_3d(ctx->ggml_ctx, t_emb->type, t_emb->ne[0], hidden_states->ne[1], t_emb->ne[1]);
|
||||
auto temb_tokens = ggml_repeat(ctx->ggml_ctx,
|
||||
ggml_reshape_3d(ctx->ggml_ctx, t_emb, t_emb->ne[0], 1, t_emb->ne[1]),
|
||||
temb_template);
|
||||
auto temb6 = time_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens));
|
||||
|
||||
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.prelude", "hidden_states");
|
||||
|
||||
for (int i = 0; i < config.depth; i++) {
|
||||
auto block = std::dynamic_pointer_cast<LingBotVideoBlock>(blocks["blocks." + std::to_string(i)]);
|
||||
hidden_states = block->forward(ctx, hidden_states, temb6, pe);
|
||||
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.blocks." + std::to_string(i), "hidden_states");
|
||||
}
|
||||
|
||||
auto final_mods = ggml_ext_chunk(ctx->ggml_ctx,
|
||||
norm_out_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens)),
|
||||
2,
|
||||
0);
|
||||
hidden_states = norm_out->forward(ctx, hidden_states);
|
||||
hidden_states = Flux::modulate(ctx->ggml_ctx, hidden_states, final_mods[0], final_mods[1], true);
|
||||
hidden_states = proj_out->forward(ctx, hidden_states);
|
||||
hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);
|
||||
|
||||
auto out = DiT::unpatchify_3d(ctx->ggml_ctx,
|
||||
hidden_states,
|
||||
t_len,
|
||||
h_len,
|
||||
w_len,
|
||||
config.patch_t,
|
||||
config.patch_h,
|
||||
config.patch_w,
|
||||
false);
|
||||
return out;
|
||||
}
|
||||
};
|
||||
|
||||
struct LingBotVideoRunner : public DiffusionModelRunner {
|
||||
LingBotVideoConfig config;
|
||||
LingBotVideoModel lingbot_video;
|
||||
std::vector<float> pe_vec;
|
||||
|
||||
LingBotVideoRunner(ggml_backend_t backend,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "",
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
|
||||
const char* model_args = nullptr)
|
||||
: DiffusionModelRunner(backend, prefix, weight_manager),
|
||||
config(LingBotVideoConfig::detect_from_weights(tensor_storage_map, prefix)) {
|
||||
SD_UNUSED(model_args);
|
||||
|
||||
lingbot_video = LingBotVideoModel(config);
|
||||
lingbot_video.init(params_ctx, tensor_storage_map, prefix);
|
||||
}
|
||||
|
||||
std::string get_desc() override {
|
||||
return "lingbot_video";
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
|
||||
lingbot_video.get_param_tensors(tensors, prefix);
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
|
||||
const sd::Tensor<float>& timesteps_tensor,
|
||||
const sd::Tensor<float>& context_tensor) {
|
||||
ggml_cgraph* gf = new_graph_custom(LINGBOT_VIDEO_GRAPH_SIZE);
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
||||
GGML_ASSERT(x_tensor.dim() == 5);
|
||||
GGML_ASSERT(x->ne[3] == config.in_channels);
|
||||
GGML_ASSERT(!context_tensor.empty());
|
||||
ggml_tensor* context = make_input(context_tensor);
|
||||
|
||||
int64_t W = x->ne[0];
|
||||
int64_t H = x->ne[1];
|
||||
int64_t T = x->ne[2];
|
||||
int64_t N = x_tensor.shape()[4];
|
||||
GGML_ASSERT(N == 1);
|
||||
pe_vec = Rope::gen_lingbot_video_pe(static_cast<int>(T),
|
||||
static_cast<int>(H),
|
||||
static_cast<int>(W),
|
||||
config.patch_t,
|
||||
config.patch_h,
|
||||
config.patch_w,
|
||||
static_cast<int>(N),
|
||||
static_cast<int>(context->ne[1]),
|
||||
config.rope_theta,
|
||||
config.axes_dim);
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
set_backend_tensor_data(pe, pe_vec.data());
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* out = lingbot_video.forward(&runner_ctx, x, timesteps, context, pe);
|
||||
ggml_build_forward_expand(gf, out);
|
||||
return gf;
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(int n_threads,
|
||||
const sd::Tensor<float>& x,
|
||||
const sd::Tensor<float>& timesteps,
|
||||
const sd::Tensor<float>& context) {
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_graph(x, timesteps, context);
|
||||
};
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(int n_threads,
|
||||
const DiffusionParams& diffusion_params) override {
|
||||
GGML_ASSERT(diffusion_params.x != nullptr);
|
||||
GGML_ASSERT(diffusion_params.timesteps != nullptr);
|
||||
return compute(n_threads,
|
||||
*diffusion_params.x,
|
||||
*diffusion_params.timesteps,
|
||||
tensor_or_empty(diffusion_params.context));
|
||||
}
|
||||
};
|
||||
} // namespace LingBotVideo
|
||||
|
||||
#endif // __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__
|
||||
@@ -183,7 +183,7 @@ namespace Qwen {
|
||||
auto to_v = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
|
||||
auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
to_out_0->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
@@ -517,7 +517,7 @@ namespace Qwen {
|
||||
if (input->ne[3] == 1) {
|
||||
input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]);
|
||||
}
|
||||
return DiT::patchify(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
|
||||
return DiT::patchify_3d(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
|
||||
};
|
||||
|
||||
auto img = patchify_input(x);
|
||||
|
||||
@@ -150,6 +150,8 @@ namespace ZImage {
|
||||
|
||||
if (sd_backend_is(ctx->backend, "ROCm")) {
|
||||
out_proj->set_scale(1.f / 16.f);
|
||||
out_proj->set_force_prec_f32(true);
|
||||
qkv_proj->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
auto qkv = qkv_proj->forward(ctx, x); // [N, n_token, (num_heads + num_kv_heads*2)*head_dim]
|
||||
@@ -227,7 +229,7 @@ namespace ZImage {
|
||||
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
|
||||
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
w2->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
|
||||
@@ -3,15 +3,19 @@
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <exception>
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <ostream>
|
||||
#include <string>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "binary_io.h"
|
||||
#include "core/util.h"
|
||||
#include "json.hpp"
|
||||
|
||||
namespace fs = std::filesystem;
|
||||
|
||||
static constexpr size_t ST_HEADER_SIZE_LEN = 8;
|
||||
|
||||
static void set_error(std::string* error, const std::string& message) {
|
||||
@@ -20,6 +24,14 @@ static void set_error(std::string* error, const std::string& message) {
|
||||
}
|
||||
}
|
||||
|
||||
static std::string resolve_index_shard_path(const std::string& index_path, const std::string& shard_path) {
|
||||
fs::path shard_fs_path(shard_path);
|
||||
if (shard_fs_path.is_absolute()) {
|
||||
return shard_fs_path.lexically_normal().string();
|
||||
}
|
||||
return (fs::path(index_path).parent_path() / shard_fs_path).lexically_normal().string();
|
||||
}
|
||||
|
||||
bool is_safetensors_file(const std::string& file_path) {
|
||||
std::ifstream file(file_path, std::ios::binary);
|
||||
if (!file.is_open()) {
|
||||
@@ -230,6 +242,52 @@ bool read_safetensors_file(const std::string& file_path,
|
||||
return true;
|
||||
}
|
||||
|
||||
bool read_safetensors_index_file(const std::string& file_path,
|
||||
std::vector<std::string>& shard_paths,
|
||||
std::string* error) {
|
||||
shard_paths.clear();
|
||||
|
||||
std::ifstream file(file_path);
|
||||
if (!file.is_open()) {
|
||||
set_error(error, "failed to open '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
nlohmann::json index;
|
||||
try {
|
||||
index = nlohmann::json::parse(file);
|
||||
} catch (const std::exception&) {
|
||||
set_error(error, "parsing safetensors index failed: '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!index.is_object() || !index.contains("weight_map") || !index["weight_map"].is_object()) {
|
||||
set_error(error, "invalid safetensors index '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
std::unordered_set<std::string> seen_shard_paths;
|
||||
for (const auto& item : index["weight_map"].items()) {
|
||||
if (!item.value().is_string()) {
|
||||
set_error(error, "invalid shard path for tensor '" + item.key() + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
std::string shard_path = resolve_index_shard_path(file_path,
|
||||
item.value().get<std::string>());
|
||||
if (seen_shard_paths.insert(shard_path).second) {
|
||||
shard_paths.push_back(std::move(shard_path));
|
||||
}
|
||||
}
|
||||
|
||||
if (shard_paths.empty()) {
|
||||
set_error(error, "safetensors index has no tensors: '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool ggml_type_to_safetensors_dtype(ggml_type type, std::string* dtype) {
|
||||
switch (type) {
|
||||
case GGML_TYPE_F16:
|
||||
|
||||
@@ -11,6 +11,9 @@ bool is_safetensors_file(const std::string& file_path);
|
||||
bool read_safetensors_file(const std::string& file_path,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::string* error = nullptr);
|
||||
bool read_safetensors_index_file(const std::string& file_path,
|
||||
std::vector<std::string>& shard_paths,
|
||||
std::string* error = nullptr);
|
||||
bool write_safetensors_file(const std::string& file_path,
|
||||
const std::vector<TensorWriteInfo>& tensors,
|
||||
std::string* error = nullptr);
|
||||
|
||||
@@ -235,6 +235,9 @@ bool ModelLoader::init_from_file(const std::string& file_path, const std::string
|
||||
} else if (is_gguf_file(file_path)) {
|
||||
LOG_INFO("load %s using gguf format", file_path.c_str());
|
||||
return init_from_gguf_file(file_path, prefix);
|
||||
} else if (ends_with(file_path, ".json")) {
|
||||
LOG_INFO("load %s using safetensors index format", file_path.c_str());
|
||||
return init_from_safetensors_index_file(file_path, prefix);
|
||||
} else if (is_safetensors_file(file_path)) {
|
||||
LOG_INFO("load %s using safetensors format", file_path.c_str());
|
||||
return init_from_safetensors_file(file_path, prefix);
|
||||
@@ -339,6 +342,25 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
|
||||
return true;
|
||||
}
|
||||
|
||||
bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, const std::string& prefix) {
|
||||
LOG_DEBUG("init from safetensors index '%s', prefix = '%s'", file_path.c_str(), prefix.c_str());
|
||||
|
||||
std::vector<std::string> shard_paths;
|
||||
std::string error;
|
||||
if (!read_safetensors_index_file(file_path, shard_paths, &error)) {
|
||||
LOG_ERROR("%s", error.c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const std::string& shard_path : shard_paths) {
|
||||
if (!init_from_file(shard_path, prefix)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/*================================================= TorchLegacyModelLoader ==================================================*/
|
||||
|
||||
bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix) {
|
||||
@@ -511,6 +533,9 @@ SDVersion ModelLoader::get_sd_version() {
|
||||
if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
|
||||
is_wan = true;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.patch_embedder.weight") != std::string::npos) {
|
||||
return VERSION_LINGBOT_VIDEO;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.patch_embedding.weight") != std::string::npos) {
|
||||
patch_embedding_channels = tensor_storage.ne[3];
|
||||
}
|
||||
|
||||
@@ -43,6 +43,7 @@ protected:
|
||||
|
||||
bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = "");
|
||||
bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = "");
|
||||
bool init_from_safetensors_index_file(const std::string& file_path, const std::string& prefix = "");
|
||||
bool init_from_torch_zip_file(const std::string& file_path, const std::string& prefix = "");
|
||||
bool init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix = "");
|
||||
bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = "");
|
||||
|
||||
@@ -179,6 +179,102 @@ bool ModelManager::register_param_tensors(const std::string& desc,
|
||||
return true;
|
||||
}
|
||||
|
||||
bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* registered_tensor_size) {
|
||||
if (desc.empty()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
std::unordered_set<TensorState*> target_states;
|
||||
size_t released_size = 0;
|
||||
for (auto& state : tensor_states_) {
|
||||
if (state == nullptr || state->desc != desc) {
|
||||
continue;
|
||||
}
|
||||
if (state->active_prepare_count > 0) {
|
||||
LOG_ERROR("model manager cannot unregister active %s tensor '%s'",
|
||||
desc.c_str(),
|
||||
state->name.c_str());
|
||||
return false;
|
||||
}
|
||||
target_states.insert(state.get());
|
||||
if (state->tensor != nullptr) {
|
||||
released_size += ggml_nbytes(state->tensor);
|
||||
}
|
||||
}
|
||||
|
||||
if (target_states.empty()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
release_compute_staging_blocks(false);
|
||||
|
||||
std::vector<ParamsStorageBlock*> storage_blocks_to_release;
|
||||
std::unordered_set<TensorState*> affected_storage_states;
|
||||
for (const auto& block : params_storage_blocks_) {
|
||||
if (block == nullptr) {
|
||||
continue;
|
||||
}
|
||||
bool has_target_state = false;
|
||||
for (TensorState* state : block->states) {
|
||||
if (state != nullptr && target_states.count(state) > 0) {
|
||||
has_target_state = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!has_target_state) {
|
||||
continue;
|
||||
}
|
||||
storage_blocks_to_release.push_back(block.get());
|
||||
for (TensorState* state : block->states) {
|
||||
if (state != nullptr) {
|
||||
affected_storage_states.insert(state);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (TensorState* state : affected_storage_states) {
|
||||
if (state == nullptr) {
|
||||
continue;
|
||||
}
|
||||
if (state->active_prepare_count > 0 || state->staged_to_compute_backend) {
|
||||
LOG_ERROR("model manager cannot unregister %s while tensor '%s' is active",
|
||||
desc.c_str(),
|
||||
state->name.c_str());
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
for (ParamsStorageBlock* block : storage_blocks_to_release) {
|
||||
if (block != nullptr) {
|
||||
free_params_storage_block(*block);
|
||||
erase_params_storage_block(block);
|
||||
}
|
||||
}
|
||||
|
||||
for (auto it = tensor_states_by_name_.begin(); it != tensor_states_by_name_.end();) {
|
||||
if (target_states.count(it->second) > 0) {
|
||||
it = tensor_states_by_name_.erase(it);
|
||||
} else {
|
||||
++it;
|
||||
}
|
||||
}
|
||||
tensor_states_.erase(std::remove_if(tensor_states_.begin(),
|
||||
tensor_states_.end(),
|
||||
[&](const std::unique_ptr<TensorState>& s) {
|
||||
return s == nullptr || target_states.count(s.get()) > 0;
|
||||
}),
|
||||
tensor_states_.end());
|
||||
|
||||
if (registered_tensor_size != nullptr) {
|
||||
if (released_size > *registered_tensor_size) {
|
||||
*registered_tensor_size = 0;
|
||||
} else {
|
||||
*registered_tensor_size -= released_size;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool ModelManager::load_all_params_eagerly() {
|
||||
std::vector<TensorState*> all_states;
|
||||
all_states.reserve(tensor_states_.size());
|
||||
|
||||
@@ -134,6 +134,9 @@ public:
|
||||
bool allow_split_buffer = false,
|
||||
bool params_follow_compute_backend = false);
|
||||
|
||||
bool unregister_param_tensors(const std::string& desc,
|
||||
size_t* registered_tensor_size = nullptr);
|
||||
|
||||
template <typename Runner>
|
||||
bool register_runner_params(const std::string& desc,
|
||||
Runner& runner,
|
||||
|
||||
+148
-7
@@ -32,6 +32,7 @@
|
||||
#include "model/diffusion/ideogram4.hpp"
|
||||
#include "model/diffusion/krea2.hpp"
|
||||
#include "model/diffusion/lens.hpp"
|
||||
#include "model/diffusion/lingbot_video.hpp"
|
||||
#include "model/diffusion/ltxv.hpp"
|
||||
#include "model/diffusion/minit2i.hpp"
|
||||
#include "model/diffusion/mmdit.hpp"
|
||||
@@ -87,6 +88,7 @@ const char* model_version_to_str[] = {
|
||||
"Wan 2.x",
|
||||
"Wan 2.2 I2V",
|
||||
"Wan 2.2 TI2V",
|
||||
"LingBot Video",
|
||||
"Qwen Image",
|
||||
"Qwen Image Layered",
|
||||
"Anima",
|
||||
@@ -222,9 +224,13 @@ public:
|
||||
std::string split_mode_spec;
|
||||
bool auto_fit_enabled = false;
|
||||
|
||||
bool diffusion_conv_direct = false;
|
||||
|
||||
bool is_using_v_parameterization = false;
|
||||
bool is_using_edm_v_parameterization = false;
|
||||
|
||||
size_t control_net_params_mem_size = 0;
|
||||
|
||||
std::shared_ptr<ModelManager> model_manager;
|
||||
|
||||
std::shared_ptr<Denoiser> denoiser = std::make_shared<CompVisDenoiser>();
|
||||
@@ -494,6 +500,76 @@ public:
|
||||
params_follow_runtime);
|
||||
}
|
||||
|
||||
bool unload_control_net() {
|
||||
if (control_net == nullptr) {
|
||||
return true;
|
||||
}
|
||||
if (model_manager != nullptr) {
|
||||
if (!model_manager->unregister_param_tensors("ControlNet", &control_net_params_mem_size)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
control_net.reset();
|
||||
control_net_params_mem_size = 0;
|
||||
return true;
|
||||
}
|
||||
|
||||
bool load_control_net_from_file(const std::string& path) {
|
||||
if (path.empty()) {
|
||||
LOG_ERROR("sd_ctx_load_control_net: empty path");
|
||||
return false;
|
||||
}
|
||||
if (model_manager == nullptr) {
|
||||
LOG_ERROR("sd_ctx_load_control_net: model_manager not initialized");
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!unload_control_net()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
ModelLoader& shared_loader = model_manager->loader();
|
||||
if (!shared_loader.init_from_file(path)) {
|
||||
LOG_ERROR("sd_ctx_load_control_net: failed to load '%s'", path.c_str());
|
||||
return false;
|
||||
}
|
||||
shared_loader.convert_tensors_name();
|
||||
|
||||
if (!ensure_backend_pair(SDBackendModule::CONTROL_NET)) {
|
||||
LOG_ERROR("sd_ctx_load_control_net: control_net backend unavailable");
|
||||
return false;
|
||||
}
|
||||
|
||||
control_net = std::make_shared<ControlNet>(backend_for(SDBackendModule::CONTROL_NET),
|
||||
params_backend_for(SDBackendModule::CONTROL_NET),
|
||||
shared_loader.get_tensor_storage_map(),
|
||||
version,
|
||||
"",
|
||||
model_manager);
|
||||
if (diffusion_conv_direct) {
|
||||
LOG_INFO("Using Conv2d direct in the control net");
|
||||
control_net->set_conv2d_direct_enabled(true);
|
||||
}
|
||||
if (!register_runner_params("ControlNet",
|
||||
control_net,
|
||||
SDBackendModule::CONTROL_NET,
|
||||
&control_net_params_mem_size)) {
|
||||
LOG_ERROR("sd_ctx_load_control_net: register_runner_params failed");
|
||||
control_net.reset();
|
||||
control_net_params_mem_size = 0;
|
||||
return false;
|
||||
}
|
||||
if (!model_manager->validate_registered_tensors()) {
|
||||
LOG_ERROR("sd_ctx_load_control_net: registered tensors validation failed");
|
||||
unload_control_net();
|
||||
return false;
|
||||
}
|
||||
LOG_INFO("sd_ctx_load_control_net: loaded '%s' (%.2f MB)",
|
||||
path.c_str(),
|
||||
control_net_params_mem_size / 1024.0 / 1024.0);
|
||||
return true;
|
||||
}
|
||||
|
||||
bool init_backend() {
|
||||
std::string error;
|
||||
if (!backend_manager.init(backend_spec.c_str(),
|
||||
@@ -852,10 +928,12 @@ public:
|
||||
model_loader.process_model_files(enable_mmap, needs_writable_mmap);
|
||||
load_alphas_cumprod(model_loader);
|
||||
|
||||
diffusion_conv_direct = sd_ctx_params->diffusion_conv_direct;
|
||||
|
||||
size_t text_encoder_params_mem_size = 0;
|
||||
size_t unet_params_mem_size = 0;
|
||||
size_t vae_params_mem_size = 0;
|
||||
size_t control_net_params_mem_size = 0;
|
||||
control_net_params_mem_size = 0;
|
||||
size_t extension_params_mem_size = 0;
|
||||
|
||||
bool tae_preview_only = sd_ctx_params->tae_preview_only;
|
||||
@@ -905,7 +983,7 @@ public:
|
||||
tensor_storage_map,
|
||||
version,
|
||||
"",
|
||||
false,
|
||||
true,
|
||||
model_manager);
|
||||
diffusion_model = std::make_shared<Krea2::Krea2Runner>(backend_for(SDBackendModule::DIFFUSION),
|
||||
tensor_storage_map,
|
||||
@@ -1004,6 +1082,25 @@ public:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
} else if (sd_version_is_lingbot_video(version)) {
|
||||
bool enable_vision = false;
|
||||
for (const auto& [name, _] : tensor_storage_map) {
|
||||
if (starts_with(name, "text_encoders.llm.visual.")) {
|
||||
enable_vision = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
|
||||
tensor_storage_map,
|
||||
version,
|
||||
"",
|
||||
enable_vision,
|
||||
model_manager);
|
||||
diffusion_model = std::make_shared<LingBotVideo::LingBotVideoRunner>(backend_for(SDBackendModule::DIFFUSION),
|
||||
tensor_storage_map,
|
||||
"model.diffusion_model",
|
||||
model_manager,
|
||||
sd_ctx_params->model_args);
|
||||
} else if (sd_version_is_qwen_image(version)) {
|
||||
bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED;
|
||||
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
|
||||
@@ -1503,6 +1600,7 @@ public:
|
||||
}
|
||||
} else if (sd_version_is_sd3(version) ||
|
||||
sd_version_is_wan(version) ||
|
||||
sd_version_is_lingbot_video(version) ||
|
||||
sd_version_is_qwen_image(version) ||
|
||||
version == VERSION_HIDREAM_O1 ||
|
||||
sd_version_is_anima(version) ||
|
||||
@@ -2332,7 +2430,7 @@ public:
|
||||
sd::Tensor<float> timesteps_tensor({static_cast<int64_t>(timesteps_vec.size())}, timesteps_vec);
|
||||
sd::Tensor<float> guidance_tensor({1}, std::vector<float>{guidance.distilled_guidance});
|
||||
sd::Tensor<float> noised_input = x * c_in;
|
||||
if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version))) {
|
||||
if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) {
|
||||
noised_input = noised_input * denoise_mask + init_latent * (1.0f - denoise_mask);
|
||||
}
|
||||
|
||||
@@ -2578,7 +2676,7 @@ public:
|
||||
int get_diffusion_model_down_factor() {
|
||||
int down_factor = 8; // unet
|
||||
if (sd_version_is_dit(version)) {
|
||||
if (sd_version_is_wan(version)) {
|
||||
if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) {
|
||||
down_factor = 2;
|
||||
} else {
|
||||
down_factor = 1;
|
||||
@@ -2641,7 +2739,7 @@ public:
|
||||
int latent_frames = frames;
|
||||
if (sd_version_is_ltxav(version)) {
|
||||
latent_frames = ((frames - 1) / 8) + 1;
|
||||
} else if (sd_version_is_wan(version)) {
|
||||
} else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) {
|
||||
latent_frames = ((frames - 1) / 4) + 1;
|
||||
}
|
||||
return latent_frames;
|
||||
@@ -2654,7 +2752,7 @@ public:
|
||||
if (sd_version_is_ltxav(version)) {
|
||||
return (latent_frames - 1) * 8 + 1;
|
||||
}
|
||||
if (sd_version_is_wan(version)) {
|
||||
if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) {
|
||||
return (latent_frames - 1) * 4 + 1;
|
||||
}
|
||||
return latent_frames;
|
||||
@@ -3328,7 +3426,7 @@ struct sd_ctx_t {
|
||||
};
|
||||
|
||||
static bool sd_version_supports_video_generation(SDVersion version) {
|
||||
return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_ltxav(version);
|
||||
return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version);
|
||||
}
|
||||
|
||||
static bool sd_version_supports_image_generation(SDVersion version) {
|
||||
@@ -3429,6 +3527,27 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) {
|
||||
return sd_version_supports_video_generation(sd_ctx->sd->version);
|
||||
}
|
||||
|
||||
SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path) {
|
||||
if (sd_ctx == nullptr || sd_ctx->sd == nullptr || path == nullptr) {
|
||||
return false;
|
||||
}
|
||||
return sd_ctx->sd->load_control_net_from_file(path);
|
||||
}
|
||||
|
||||
SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx) {
|
||||
if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
|
||||
return false;
|
||||
}
|
||||
return sd_ctx->sd->unload_control_net();
|
||||
}
|
||||
|
||||
SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx) {
|
||||
if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
|
||||
return false;
|
||||
}
|
||||
return sd_ctx->sd->control_net != nullptr;
|
||||
}
|
||||
|
||||
enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) {
|
||||
if (sd_ctx != nullptr && sd_ctx->sd != nullptr) {
|
||||
if (sd_version_is_pid(sd_ctx->sd->version)) {
|
||||
@@ -5419,6 +5538,27 @@ static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd
|
||||
sd::ops::fill_slice(&latents.denoise_mask, 2, latents.init_latent.shape()[2] - 1, latents.init_latent.shape()[2], 0.0f);
|
||||
}
|
||||
|
||||
int64_t t2 = ggml_time_ms();
|
||||
LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
|
||||
} else if (sd_version_is_lingbot_video(sd_ctx->sd->version) && !start_image.empty()) {
|
||||
LOG_INFO("LingBot Video IMG2VID");
|
||||
|
||||
int64_t t1 = ggml_time_ms();
|
||||
auto init_img = start_image.reshape({start_image.shape()[0], start_image.shape()[1], 1, start_image.shape()[2], 1});
|
||||
auto init_image_latent = sd_ctx->sd->encode_first_stage(init_img);
|
||||
if (init_image_latent.empty()) {
|
||||
LOG_ERROR("failed to encode init video frame");
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
latents.init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->frames, true);
|
||||
sd::ops::slice_assign(&latents.init_latent, 2, 0, init_image_latent.shape()[2], init_image_latent);
|
||||
|
||||
latents.denoise_mask = sd::full<float>({latents.init_latent.shape()[0], latents.init_latent.shape()[1], latents.init_latent.shape()[2], 1, 1}, 1.f);
|
||||
sd::ops::fill_slice(&latents.denoise_mask, 2, 0, init_image_latent.shape()[2], 0.0f);
|
||||
|
||||
latents.ref_images.push_back(start_image);
|
||||
|
||||
int64_t t2 = ggml_time_ms();
|
||||
LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
|
||||
} else if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-VACE-1.3B" ||
|
||||
@@ -5524,6 +5664,7 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx,
|
||||
condition_params.clip_skip = request.clip_skip;
|
||||
condition_params.text = request.prompt;
|
||||
condition_params.zero_out_masked = true;
|
||||
condition_params.ref_images = &latents.ref_images;
|
||||
|
||||
int64_t prepare_start_ms = ggml_time_ms();
|
||||
embeds.cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
|
||||
|
||||
Reference in New Issue
Block a user