diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index d0683166d..c3b9c0690 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -266,12 +267,24 @@ namespace MiniMaxH3 { blocks["final_norm"] = std::make_shared(config.hidden_size, config.final_norm_eps); } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + bool cut_after_last = true) { + auto final_norm = std::dynamic_pointer_cast(blocks["final_norm"]); for (int64_t i = 0; i < num_layers; ++i) { auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); x = block->forward(ctx, x); + const bool is_last = i + 1 == num_layers; + if (is_last) { + x = final_norm->forward(ctx, x); + } + if (!is_last || cut_after_last) { + sd::ggml_graph_cut::mark_graph_cut(x, + "minimax_h3.token_refiner.blocks." + std::to_string(i), + "hidden_states"); + } } - return std::dynamic_pointer_cast(blocks["final_norm"])->forward(ctx, x); + return num_layers == 0 ? final_norm->forward(ctx, x) : x; } }; @@ -526,14 +539,20 @@ namespace MiniMaxH3 { } } - ggml_tensor* refine_context(GGMLRunnerContext* ctx, ggml_tensor* context) { + ggml_tensor* refine_context(GGMLRunnerContext* ctx, + ggml_tensor* context, + bool cut_after_last_refiner = true) { if (context->ne[0] == config.hidden_size) { return context; } GGML_ASSERT(context->ne[0] == config.text_dim); auto condition_proj = std::dynamic_pointer_cast(blocks["condition_proj"]); auto token_refiner = std::dynamic_pointer_cast(blocks["token_refiner"]); - return token_refiner->forward(ctx, condition_proj->forward(ctx, context)); + auto projected = condition_proj->forward(ctx, context); + sd::ggml_graph_cut::mark_graph_cut(projected, + "minimax_h3.condition_proj", + "hidden_states"); + return token_refiner->forward(ctx, projected, cut_after_last_refiner); } ggml_tensor* time_embedding(GGMLRunnerContext* ctx, @@ -952,6 +971,42 @@ namespace MiniMaxH3 { } struct MiniMaxH3Runner : public DiffusionModelRunner { + struct RefinedContextCacheEntry { + const void* context_cache_identity = nullptr; + const sd::Tensor* source_context = nullptr; + const float* source_data = nullptr; + std::vector source_shape; + std::shared_ptr weight_adapter = nullptr; + ggml_context* refined_ctx = nullptr; + ggml_backend_buffer_t refined_buffer = nullptr; + ggml_tensor* refined = nullptr; + + ~RefinedContextCacheEntry() { + if (refined_buffer != nullptr) { + ggml_backend_buffer_free(refined_buffer); + } + if (refined_ctx != nullptr) { + ggml_free(refined_ctx); + } + } + + RefinedContextCacheEntry() = default; + RefinedContextCacheEntry(const RefinedContextCacheEntry&) = delete; + RefinedContextCacheEntry& operator=(const RefinedContextCacheEntry&) = delete; + + bool matches(const void* identity, + const sd::Tensor& context, + const std::shared_ptr& adapter) const { + return context_cache_identity == identity && + weight_adapter == adapter && + source_context == &context && + source_data == context.data() && + source_shape == context.shape(); + } + }; + + static constexpr size_t REFINED_CONTEXT_CACHE_CAPACITY = 4; + Config config; MiniMaxH3Transformer3DModel model; sd::Tensor video_input_cache; @@ -961,6 +1016,7 @@ namespace MiniMaxH3 { sd::Tensor curve_index_input_cache; sd::Tensor curve_upper_index_input_cache; sd::Tensor curve_fraction_input_cache; + std::vector> refined_context_cache; MiniMaxH3Runner(ggml_backend_t backend, const String2TensorStorage& tensors, @@ -981,6 +1037,95 @@ namespace MiniMaxH3 { model.get_param_tensors(tensors, prefix); } + std::unique_ptr create_refined_context_cache_entry( + const sd::Tensor& context, + const void* context_cache_identity) { + auto entry = std::make_unique(); + entry->context_cache_identity = context_cache_identity; + entry->source_context = &context; + entry->source_data = context.data(); + entry->source_shape = context.shape(); + entry->weight_adapter = weight_adapter; + + auto refined_shape = context.shape(); + refined_shape[0] = config.hidden_size; + ggml_init_params params; + params.mem_size = ggml_tensor_overhead(); + params.mem_buffer = nullptr; + params.no_alloc = true; + entry->refined_ctx = ggml_init(params); + GGML_ASSERT(entry->refined_ctx != nullptr); + entry->refined = ggml_new_tensor(entry->refined_ctx, + GGML_TYPE_F32, + static_cast(refined_shape.size()), + refined_shape.data()); + ggml_set_name(entry->refined, "minimax_h3.refined_context"); + entry->refined_buffer = ggml_backend_alloc_ctx_tensors(entry->refined_ctx, + runtime_backend); + GGML_ASSERT(entry->refined_buffer != nullptr); + ggml_backend_buffer_set_usage(entry->refined_buffer, + GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + return entry; + } + + ggml_cgraph* build_context_refinement_graph(const sd::Tensor& context, + ggml_tensor* refined_output) { + GGML_ASSERT(!context.empty() && context.shape()[0] == config.text_dim); + GGML_ASSERT(refined_output != nullptr && refined_output->ne[0] == config.hidden_size); + auto context_input = make_input(context); + auto runner_ctx = get_context(); + auto refined = model.refine_context(&runner_ctx, context_input, false); + // Refinement graph buffers are transient; persist only their final output. + auto output = ggml_cpy(runner_ctx.ggml_ctx, refined, refined_output); + auto graph = new_graph_custom(H3_GRAPH_SIZE); + ggml_build_forward_expand(graph, output); + return graph; + } + + ggml_tensor* get_refined_context(const sd::Tensor& context, + const void* context_cache_identity, + int n_threads) { + GGML_ASSERT(!context.empty()); + GGML_ASSERT(context_cache_identity != nullptr); + GGML_ASSERT(context.shape()[0] == config.text_dim || + context.shape()[0] == config.hidden_size); + + for (const auto& entry : refined_context_cache) { + if (entry->matches(context_cache_identity, context, weight_adapter)) { + return entry->refined; + } + } + + auto entry = create_refined_context_cache_entry(context, context_cache_identity); + if (context.shape()[0] == config.hidden_size) { + ggml_backend_tensor_set(entry->refined, + context.data(), + 0, + ggml_nbytes(entry->refined)); + ggml_backend_synchronize(runtime_backend); + } else { + auto get_graph = [&]() { + return build_context_refinement_graph(context, entry->refined); + }; + auto result = GGMLRunner::compute(get_graph, + n_threads, + false, + true, + true, + true); + if (!result.has_value()) { + return nullptr; + } + } + + auto refined = entry->refined; + if (refined_context_cache.size() == REFINED_CONTEXT_CACHE_CAPACITY) { + refined_context_cache.erase(refined_context_cache.begin()); + } + refined_context_cache.push_back(std::move(entry)); + return refined; + } + std::pair, sd::Tensor> split_av_latents(const sd::Tensor& packed, int audio_length) const { GGML_ASSERT(packed.dim() == 4 || packed.dim() == 5); @@ -1026,6 +1171,7 @@ namespace MiniMaxH3 { ggml_cgraph* build_graph(const sd::Tensor& packed, const sd::Tensor& timestep, const sd::Tensor& context_tensor, + ggml_tensor* refined_context, const std::vector>& condition_videos, const std::vector>& condition_audios, const sd::Tensor& text_tags, @@ -1039,10 +1185,12 @@ namespace MiniMaxH3 { audio_input_cache = std::move(split.second); GGML_ASSERT(!audio_input_cache.empty()); GGML_ASSERT(!context_tensor.empty()); + GGML_ASSERT(refined_context != nullptr && + refined_context->ne[0] == config.hidden_size); auto video = make_input(video_input_cache); auto audio = make_input(audio_input_cache); - auto context = make_input(context_tensor); + auto context = refined_context; std::vector condition_inputs; condition_inputs.reserve(condition_videos.size()); for (const auto& condition : condition_videos) { @@ -1151,10 +1299,20 @@ namespace MiniMaxH3 { ? empty_reference_blocks : *extra->reference_blocks; const sd::Tensor empty_int; + const void* context_cache_identity = params.context_cache_identity != nullptr + ? params.context_cache_identity + : params.context; + auto context = get_refined_context(*params.context, + context_cache_identity, + n_threads); + if (context == nullptr) { + return {}; + } auto get_graph = [&]() { return build_graph(*params.x, *params.timesteps, *params.context, + context, conditions, audio_conditions, extra->text_token_tags == nullptr ? empty_int : *extra->text_token_tags, @@ -1171,6 +1329,11 @@ namespace MiniMaxH3 { false), params.x->dim()); } + + protected: + void on_sampling_done() override { + refined_context_cache.clear(); + } }; } // namespace MiniMaxH3 diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 070ca53d4..f895e5eae 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -137,6 +137,7 @@ struct DiffusionParams { const sd::Tensor* x = nullptr; const sd::Tensor* timesteps = nullptr; const sd::Tensor* context = nullptr; + const void* context_cache_identity = nullptr; const sd::Tensor* c_concat = nullptr; const sd::Tensor* y = nullptr; const std::vector>* ref_latents = nullptr; @@ -160,6 +161,7 @@ static inline const sd::Tensor& tensor_or_empty(const sd::Tensor* tensor) struct DiffusionModelRunner : public GGMLRunner { protected: std::string prefix; + virtual void on_sampling_done() {} public: DiffusionModelRunner(ggml_backend_t backend, @@ -171,6 +173,11 @@ struct DiffusionModelRunner : public GGMLRunner { virtual sd::Tensor compute(int n_threads, const DiffusionParams& diffusion_params) = 0; + void sampling_done() { + runner_done(); + on_sampling_done(); + } + void get_param_tensors(std::map& tensors) { get_param_tensors(tensors, prefix); } diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 109a2a483..c7587b0e3 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -2531,6 +2531,16 @@ class StableDiffusionGGML { float frame_rate, const sd_cache_params_t* cache_params, const sd::Tensor& video_positions = {}) { + struct SamplingDoneOnExit { + DiffusionModelRunner* runner = nullptr; + ~SamplingDoneOnExit() { + if (runner != nullptr) { + runner->sampling_done(); + } + } + }; + SamplingDoneOnExit sample_diffusion_runner_done{work_diffusion_model.get()}; + struct RunnerDoneOnExit { GGMLRunner* runner = nullptr; ~RunnerDoneOnExit() { @@ -2539,8 +2549,6 @@ class StableDiffusionGGML { } } }; - RunnerDoneOnExit sample_diffusion_runner_done{work_diffusion_model.get()}; - RunnerDoneOnExit sample_control_runner_done{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; std::vector skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count); @@ -2712,10 +2720,11 @@ class StableDiffusionGGML { const std::vector* local_skip_layers = nullptr, const std::vector>* ref_latents_override = nullptr, bool use_uncond_ip = false) -> sd::Tensor { - diffusion_params.context = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn; - diffusion_params.c_concat = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat); - diffusion_params.y = condition.c_vector.empty() ? nullptr : &condition.c_vector; - diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images); + diffusion_params.context = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn; + diffusion_params.context_cache_identity = &condition; + diffusion_params.c_concat = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat); + diffusion_params.y = condition.c_vector.empty() ? nullptr : &condition.c_vector; + diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images); if (sd_version_is_unet(version)) { int nvf = -1;