diff --git a/CMakeLists.txt b/CMakeLists.txt index 3c2b7b4..2c42716 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -17,6 +17,7 @@ if (MSVC) $<$:/utf-8> $<$,$>:/MP> $<$:/utf-8> + $<$,$>:/bigobj> ) endif() diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 1cc7a7a..2a59882 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -924,6 +924,7 @@ int main(int argc, const char* argv[]) { } sd_img_gen_params_t img_gen_params{}; + bool video_gen_ok = true; const bool use_img_gen_params = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER; if (use_img_gen_params) { img_gen_params = gen_params.to_sd_img_gen_params_t(); @@ -942,13 +943,17 @@ int main(int argc, const char* argv[]) { } else if (cli_params.mode == VID_GEN) { sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t(); sd_image_t* generated_video = nullptr; - if (!generate_video(sd_ctx.get(), &vid_gen_params, &generated_video, &num_results, &generated_audio)) { + video_gen_ok = generate_video(sd_ctx.get(), &vid_gen_params, &generated_video, &num_results, &generated_audio); + if (!video_gen_ok) { generated_video = nullptr; } results.adopt(generated_video, num_results); } if (!results) { + if (cli_params.mode == VID_GEN && !gen_params.dump_conditioning_path.empty()) { + return video_gen_ok ? 0 : 1; + } LOG_ERROR("generate failed"); return 1; } diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 3581215..ea3950f 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -957,6 +957,11 @@ ArgOptions SDGenerationParams::get_options() { "path to the init image", 0, &init_image_path}, + {"", + "--dump-conditioning", + "compute the text conditioning, write it to a safetensors file (compatible with ComfyUI-LTXVideo's LTXVLoadConditioning cache node) at this path, and exit before loading the diffusion model or VAE", + 0, + &dump_conditioning_path}, {"", "--end-img", "path to the end image, required by flf2v", @@ -2661,6 +2666,7 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() { params.lora_count = static_cast(lora_vec.size()); params.prompt = prompt.c_str(); params.negative_prompt = negative_prompt.c_str(); + params.dump_conditioning_path = dump_conditioning_path.empty() ? nullptr : dump_conditioning_path.c_str(); params.clip_skip = clip_skip; params.init_image = init_image.get(); params.end_image = end_image.get(); diff --git a/examples/common/common.h b/examples/common/common.h index 34b4a01..6376429 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -241,6 +241,8 @@ struct SDGenerationParams { std::string ref_image_args; + std::string dump_conditioning_path; + std::string pm_id_images_dir; std::string pm_id_embed_path; float pm_style_strength = 20.f; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index bab62ba..e196d5d 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -431,6 +431,7 @@ typedef struct { sd_hires_params_t hires; bool circular_x; bool circular_y; + const char* dump_conditioning_path; } sd_vid_gen_params_t; typedef struct sd_ctx_t sd_ctx_t; diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index 69bcaa1..6b0a681 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -425,9 +425,18 @@ static bool ggml_type_to_safetensors_dtype(ggml_type type, std::string* dtype) { bool write_safetensors_file(const std::string& file_path, const std::vector& tensors, - std::string* error) { + std::string* error, + const std::map* metadata) { nlohmann::ordered_json header = nlohmann::ordered_json::object(); + if (metadata != nullptr && !metadata->empty()) { + nlohmann::ordered_json metadata_json = nlohmann::ordered_json::object(); + for (const auto& [key, value] : *metadata) { + metadata_json[key] = value; + } + header["__metadata__"] = metadata_json; + } + uint64_t data_offset = 0; for (const TensorWriteInfo& write_tensor : tensors) { ggml_tensor* tensor = write_tensor.tensor; diff --git a/src/model_io/safetensors_io.h b/src/model_io/safetensors_io.h index 4291b54..0133b4c 100644 --- a/src/model_io/safetensors_io.h +++ b/src/model_io/safetensors_io.h @@ -18,7 +18,8 @@ bool read_safetensors_index_file(const std::string& file_path, std::string* error = nullptr); bool write_safetensors_file(const std::string& file_path, const std::vector& tensors, - std::string* error = nullptr); + std::string* error = nullptr, + const std::map* metadata = nullptr); class SafetensorsStreamingWriter : public StreamingModelWriter { public: diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 9d7b6c8..4696445 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -15,6 +16,7 @@ #include "core/rng_mt19937.hpp" #include "core/rng_philox.hpp" #include "core/util.h" +#include "model_io/safetensors_io.h" #include "model_loader.h" #include "model_manager.h" #include "stable-diffusion.h" @@ -6506,6 +6508,75 @@ static std::optional prepare_video_generation_latents(sd return latents; } +// Writes the LTXAV text conditioning to a safetensors file compatible with +// ComfyUI-LTXVideo's LTXVLoadConditioning cache node: a `conditioning_data_0` +// tensor, with a `non_tensor_options` metadata entry for any non-tensor +// conditioning options. +// +// sd.cpp's LTXAV text projection only runs the video/audio_aggregate_embed +// linear layers (conditioner.hpp's LTXAVTextProjection); it does not run the +// separate video/audio_embeddings_connector transformer stage that +// ComfyUI's av_model.py applies afterwards (Embeddings1DConnector: 2 RoPE +// self-attention blocks + learnable register-token padding to >=1024 +// tokens). Those connector weights live in the diffusion checkpoint, not +// the text encoder, and ComfyUI's own diffusion model already loads and +// runs that stage internally whenever conditioning carries +// unprocessed_ltxav_embeds=true. So rather than reimplementing the +// connector here, mark the dump as unprocessed and let ComfyUI finish the +// job with its own correctly-loaded weights. No attention_mask tensor is +// written: the connector rebuilds its own mask internally once it pads the +// sequence with register tokens, and a stale one sized for our raw token +// count crashes downstream cross-attention with a shape mismatch. +static bool dump_ltxav_conditioning_to_safetensors(const SDCondition& cond, const std::string& path, std::string* error) { + if (cond.c_crossattn.empty() || cond.c_crossattn.dim() != 2) { + if (error != nullptr) { + *error = "dump-conditioning: expected a 2D LTXAV crossattn tensor, got dim=" + + std::to_string(cond.c_crossattn.dim()); + } + return false; + } + + const int64_t hidden_dim = cond.c_crossattn.shape()[0]; + const int64_t seq_len = cond.c_crossattn.shape()[1]; + + ggml_init_params ggml_params{}; + ggml_params.mem_size = static_cast(hidden_dim * seq_len) * sizeof(float) + + 2 * ggml_tensor_overhead() + 4096; + ggml_params.mem_buffer = nullptr; + ggml_params.no_alloc = false; + ggml_context* ctx = ggml_init(ggml_params); + if (ctx == nullptr) { + if (error != nullptr) { + *error = "dump-conditioning: failed to allocate scratch context"; + } + return false; + } + + ggml_tensor* crossattn_2d = sd::make_ggml_tensor(ctx, cond.c_crossattn); + ggml_tensor* crossattn_3d = ggml_reshape_3d(ctx, crossattn_2d, hidden_dim, seq_len, 1); + ggml_set_name(crossattn_3d, "conditioning_data_0"); + + std::vector tensors; + + TensorWriteInfo crossattn_info; + crossattn_info.tensor = crossattn_3d; + crossattn_info.n_dims = 3; + crossattn_info.ne[0] = hidden_dim; + crossattn_info.ne[1] = seq_len; + crossattn_info.ne[2] = 1; + tensors.push_back(crossattn_info); + + std::map metadata; + metadata["num_conditionings"] = "1"; + metadata["dtype"] = "float32"; + metadata["created_at"] = "stable-diffusion.cpp --dump-conditioning"; + metadata["non_tensor_options"] = "{\"0:unprocessed_ltxav_embeds\": true}"; + + bool ok = write_safetensors_file(path, tensors, error, &metadata); + ggml_free(ctx); + return ok; +} + static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params, const GenerationRequest& request, @@ -6909,6 +6980,21 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, sd_vid_gen_params, request, latents); + + if (sd_vid_gen_params->dump_conditioning_path != nullptr && + strlen(sd_vid_gen_params->dump_conditioning_path) > 0) { + std::string dump_error; + bool dump_ok = dump_ltxav_conditioning_to_safetensors(embeds.cond, + sd_vid_gen_params->dump_conditioning_path, + &dump_error); + if (!dump_ok) { + LOG_ERROR("dump-conditioning failed: %s", dump_error.c_str()); + } else { + LOG_INFO("dumped conditioning to '%s'", sd_vid_gen_params->dump_conditioning_path); + } + return dump_ok; + } + if (latent_upscale_enabled) { LOG_INFO("generate_video %dx%dx%d -> LTX latent spatial upscale", request.width,