Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 36 additions & 57 deletions src/model/diffusion/minimax_h3.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -123,25 +123,6 @@ namespace MiniMaxH3 {
return to_shift * base / (1.f + (to_shift - 1.f) * base);
}

static float time_shift_slope(float sigma, float from_shift, float to_shift) {
float base = sigma / (from_shift + sigma * (1.f - from_shift));
float a = 1.f + (from_shift - 1.f) * base;
float b = 1.f + (to_shift - 1.f) * base;
return to_shift * a * a / (from_shift * b * b);
}

static float time_shift_step_scale(float sigma,
float next_sigma,
float from_shift,
float to_shift) {
if (!std::isfinite(next_sigma) || next_sigma < 0.f || next_sigma == sigma) {
return time_shift_slope(sigma, from_shift, to_shift);
}
float shifted_sigma = time_shift_sigma(sigma, from_shift, to_shift);
float shifted_next_sigma = time_shift_sigma(next_sigma, from_shift, to_shift);
return (shifted_sigma - shifted_next_sigma) / (sigma - next_sigma);
}

struct TimeEmbedder : public GGMLBlock {
TimeEmbedder(int64_t input_dim, int64_t hidden_dim, int64_t output_dim) {
blocks["proj_in"] = std::make_shared<Linear>(input_dim, hidden_dim, true, true);
Expand Down Expand Up @@ -606,8 +587,7 @@ namespace MiniMaxH3 {
const std::vector<TokenModulationSpan>& segments,
const std::vector<SequenceSegment>& sequence_segments,
const TokenModulationSpan& video_segment,
const TokenModulationSpan& audio_segment,
float audio_slope) {
const TokenModulationSpan& audio_segment) {
auto video_proj = std::dynamic_pointer_cast<Linear>(blocks["video_patch_proj"]);
auto audio_proj = std::dynamic_pointer_cast<Linear>(blocks["audio_patch_proj"]);

Expand Down Expand Up @@ -727,7 +707,7 @@ namespace MiniMaxH3 {
audio->ne[2]);
audio_out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, audio_out, 1, 2, 0, 3));
video_out = ggml_ext_scale(ctx->ggml_ctx, video_out, -1.f);
audio_out = ggml_ext_scale(ctx->ggml_ctx, audio_out, -audio_slope);
audio_out = ggml_ext_scale(ctx->ggml_ctx, audio_out, -1.f);
return {video_out, audio_out};
}
};
Expand Down Expand Up @@ -1045,17 +1025,16 @@ namespace MiniMaxH3 {
const std::vector<MiniMaxH3ReferenceBlock>& reference_blocks,
int audio_length,
float video_shift,
float audio_shift,
float next_video_sigma) {
float audio_shift) {
auto split = split_av_latents(packed, audio_length);
video_input_cache = std::move(split.first);
audio_input_cache = std::move(split.second);
GGML_ASSERT(!audio_input_cache.empty());
GGML_ASSERT(!context_tensor.empty());

auto video = make_input(video_input_cache);
auto audio = make_input(audio_input_cache);
auto context = make_input(context_tensor);
auto video = make_input(video_input_cache);
auto audio_carrier = make_input(audio_input_cache);
auto context = make_input(context_tensor);
std::vector<ggml_tensor*> condition_inputs;
condition_inputs.reserve(condition_videos.size());
for (const auto& condition : condition_videos) {
Expand All @@ -1067,21 +1046,26 @@ namespace MiniMaxH3 {
audio_condition_inputs.push_back(make_input(condition));
}

float sigma_v = std::clamp(timestep[0] / 1000.f, 1e-6f, 1.f);
float t_v = 1.f - sigma_v;
float t_a = 1.f - time_shift_sigma(sigma_v, video_shift, audio_shift);
auto layout = build_layout(context_tensor.shape()[1],
video_input_cache.shape()[2],
video_input_cache.shape()[1],
video_input_cache.shape()[0],
audio_length,
condition_videos,
condition_audios,
keyframe_indices,
reference_blocks,
text_tags,
t_v,
t_a);
float sigma_v = std::clamp(timestep[0] / 1000.f, 1e-6f, 1.f);
float sigma_a = time_shift_sigma(sigma_v, video_shift, audio_shift);
float audio_scale = video_shift / audio_shift;
float t_v = 1.f - sigma_v;
float t_a = 1.f - sigma_a;
// The sampler carries c_a = (sigma_v / sigma_a) * x_a so the packed
// latent follows one sigma schedule. Restore x_a for the H3 network.
auto audio = ggml_ext_scale(compute_ctx, audio_carrier, sigma_a / sigma_v);
auto layout = build_layout(context_tensor.shape()[1],
video_input_cache.shape()[2],
video_input_cache.shape()[1],
video_input_cache.shape()[0],
audio_length,
condition_videos,
condition_audios,
keyframe_indices,
reference_blocks,
text_tags,
t_v,
t_a);

position_input_cache = sd::Tensor<float>(
{3, static_cast<int64_t>(layout.positions.size() / 3)},
Expand Down Expand Up @@ -1142,19 +1126,15 @@ namespace MiniMaxH3 {
layout.segments,
layout.sequence_segments,
layout.video_segment,
layout.audio_segment,
// The generic Euler sampler advances the packed tensor by
// `next_video_sigma - sigma_v`. For that sampler, scale H3's
// audio velocity by the exact ratio of the independent audio
// step. The derivative approximation substantially oversteps
// at low step counts (the Turbo use case). Retain the local
// slope for samplers that make extra/intermediate evaluations.
time_shift_step_scale(sigma_v,
next_video_sigma,
video_shift,
audio_shift));
auto merged = merge_av_latents(compute_ctx, output.first, output.second);
auto graph = new_graph_custom(H3_GRAPH_SIZE);
layout.audio_segment);
// Convert the model's audio velocity to d(c_a) / d(sigma_v).
output.second = ggml_add(compute_ctx,
ggml_ext_scale(compute_ctx, audio, 1.f - audio_scale),
ggml_ext_scale(compute_ctx,
output.second,
1.f + (audio_scale - 1.f) * sigma_a));
auto merged = merge_av_latents(compute_ctx, output.first, output.second);
auto graph = new_graph_custom(H3_GRAPH_SIZE);
ggml_build_forward_expand(graph, merged);
return graph;
}
Expand Down Expand Up @@ -1184,8 +1164,7 @@ namespace MiniMaxH3 {
reference_blocks,
extra->audio_length,
extra->video_sigma_shift,
extra->audio_sigma_shift,
extra->next_video_sigma);
extra->audio_sigma_shift);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
n_threads,
Expand Down
2 changes: 0 additions & 2 deletions src/model/diffusion/model.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -108,8 +108,6 @@ struct MiniMaxH3DiffusionExtra {
int audio_length = 0;
float video_sigma_shift = 12.f;
float audio_sigma_shift = 3.f;
// Negative when the outer sampler is not a single-evaluation Euler step.
float next_video_sigma = -1.f;
};

struct MiniT2IDiffusionExtra {
Expand Down
44 changes: 44 additions & 0 deletions src/runtime/denoiser.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -1043,6 +1043,16 @@ struct Denoiser {
const sd::Tensor<float>& latent) = 0;
virtual float noise_level_to_sigma(float noise_level) = 0;

virtual sd::Tensor<float> process_latent_in(const sd::Tensor<float>& latent) {
// An empty result means the original latent can be used unchanged.
SD_UNUSED(latent);
return {};
}

virtual sd::Tensor<float> process_latent_out(sd::Tensor<float> latent) {
return latent;
}

virtual std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) {
auto bound_t_to_sigma = std::bind(&Denoiser::t_to_sigma, this, std::placeholders::_1);
std::shared_ptr<SigmaScheduler> scheduler;
Expand Down Expand Up @@ -1286,6 +1296,40 @@ struct DiscreteFlowDenoiser : public Denoiser {
}
};

struct H3AVFlowDenoiser : public DiscreteFlowDenoiser {
int64_t video_channels;
float audio_shift;

H3AVFlowDenoiser(float shift, float audio_shift, int64_t video_channels)
: DiscreteFlowDenoiser(shift),
video_channels(video_channels),
audio_shift(audio_shift) {
GGML_ASSERT(shift > 0.f && audio_shift > 0.f && video_channels > 0);
}

sd::Tensor<float> process_latent_in(const sd::Tensor<float>& latent) override {
return scale_audio(latent, shift / audio_shift);
}

sd::Tensor<float> process_latent_out(sd::Tensor<float> latent) override {
auto transformed = scale_audio(latent, audio_shift / shift);
if (transformed.empty()) {
return latent;
}
return transformed;
}

private:
sd::Tensor<float> scale_audio(const sd::Tensor<float>& latent, float scale) const {
if (scale == 1.f || latent.dim() < 4 || latent.shape()[3] <= video_channels) {
return {};
}
auto video = sd::ops::slice(latent, 3, 0, video_channels);
auto audio = sd::ops::slice(latent, 3, video_channels, latent.shape()[3]) * scale;
return sd::ops::concat(video, audio, 3);
}
};

struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
FluxFlowDenoiser() = default;

Expand Down
32 changes: 18 additions & 14 deletions src/stable-diffusion.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1860,6 +1860,9 @@ class StableDiffusionGGML {
if (sd_version_is_ltxav(version)) {
LOG_INFO("running in LTXAV FLOW mode");
denoiser = std::make_shared<FluxFlowDenoiser>();
} else if (sd_version_is_minimax_h3(version)) {
LOG_INFO("running in MiniMax H3 AV FLOW mode");
denoiser = std::make_shared<H3AVFlowDenoiser>(default_flow_shift, 3.f, get_latent_channel());
} else {
LOG_INFO("running in FLOW mode");
denoiser = std::make_shared<DiscreteFlowDenoiser>();
Expand Down Expand Up @@ -2617,10 +2620,14 @@ class StableDiffusionGGML {
int64_t last_progress_us = ggml_time_us();
SamplePreviewContext preview = prepare_sample_preview_context();

sd::Tensor<float> x_t = !noise.empty()
? denoiser->noise_scaling(sigmas[0], noise, init_latent)
: init_latent;
sd::Tensor<float> denoised = x_t;
sd::Tensor<float> processed_init_latent = denoiser->process_latent_in(init_latent);
const sd::Tensor<float>& sampling_init_latent = processed_init_latent.empty()
? init_latent
: processed_init_latent;
sd::Tensor<float> x_t = !noise.empty()
? denoiser->noise_scaling(sigmas[0], noise, sampling_init_latent)
: sampling_init_latent;
sd::Tensor<float> denoised = x_t;

auto denoise = [&](const sd::Tensor<float>& x, float sigma, int step) -> sd::guidance::GuiderOutput {
if (get_cancel_flag() == SD_CANCEL_ALL) {
Expand All @@ -2643,10 +2650,10 @@ class StableDiffusionGGML {
std::vector<float> timesteps_vec = base_timesteps_vec;
sd::Tensor<float> audio_timesteps_tensor;
if (sd_version_is_ltxav(version) && !denoise_mask.empty()) {
timesteps_vec = process_ltxav_video_timesteps(base_timesteps_vec, init_latent, denoise_mask);
timesteps_vec = process_ltxav_video_timesteps(base_timesteps_vec, sampling_init_latent, denoise_mask);
audio_timesteps_tensor = sd::Tensor<float>({static_cast<int64_t>(base_timesteps_vec.size())}, base_timesteps_vec);
} else {
timesteps_vec = process_timesteps(timesteps_vec, init_latent, denoise_mask, step);
timesteps_vec = process_timesteps(timesteps_vec, sampling_init_latent, denoise_mask, step);
}
const std::vector<float>& scaling_timesteps_vec = (sd_version_is_ltxav(version) && !denoise_mask.empty())
? base_timesteps_vec
Expand All @@ -2661,13 +2668,13 @@ class StableDiffusionGGML {
}
sd::Tensor<float> noised_input = x * c_in;
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);
noised_input = noised_input * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
}

if (cache_runtime.spectrum_enabled && cache_runtime.spectrum.should_predict()) {
cache_runtime.spectrum.predict(&denoised);
if (!denoise_mask.empty()) {
denoised = denoised * denoise_mask + init_latent * (1.0f - denoise_mask);
denoised = denoised * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
}
if (sd_should_preview_denoised() && preview.callback != nullptr) {
if (step % sd_get_preview_interval() == 0) {
Expand Down Expand Up @@ -2765,11 +2772,7 @@ class StableDiffusionGGML {
condition.c_reference_blocks.empty() ? nullptr : &condition.c_reference_blocks,
audio_length,
std::isfinite(active_flow_shift) ? active_flow_shift : 12.f,
3.f,
method == EULER_SAMPLE_METHOD && step > 0 &&
static_cast<size_t>(step) < sigmas.size()
? sigmas[step]
: -1.f};
3.f};
} else if (sd_version_is_ltxav(version)) {
diffusion_params.extra = LTXAVDiffusionExtra{
nullptr,
Expand Down Expand Up @@ -2894,7 +2897,7 @@ class StableDiffusionGGML {
cache_runtime.spectrum.update(denoised);
}
if (!denoise_mask.empty()) {
denoised = denoised * denoise_mask + init_latent * (1.0f - denoise_mask);
denoised = denoised * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
}
if (sd_should_preview_denoised() && preview.callback != nullptr) {
if (step % sd_get_preview_interval() == 0) {
Expand Down Expand Up @@ -2924,6 +2927,7 @@ class StableDiffusionGGML {
if (inverse_noise_scaling) {
x0 = denoiser->inverse_noise_scaling(sigmas[sigmas.size() - 1], x0);
}
x0 = denoiser->process_latent_out(std::move(x0));

if (control_net) {
control_net->free_control_ctx();
Expand Down
Loading