void operator()(ggml_context * ctx) const noexcept {

if (ctx != nullptr) {

ggml_free(ctx);

}

}

if (input.type == GGML_TYPE_F32) {

return input;

}

return core::wrap_tensor(ggml_cast(ctx.ggml, input.tensor, GGML_TYPE_F32), input.shape, GGML_TYPE_F32);

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

int64_t heads,

int64_t dim) {

auto contiguous = ensure_contiguous(ctx, input);

return core::reshape_tensor(

ctx,

contiguous,

core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim}));

if (repeats == 1) {

return input;

}

std::vector<core::TensorValue> heads;

heads.reserve(static_cast<size_t>(input.shape.dims[1] * repeats));

for (int64_t head = 0; head < input.shape.dims[1]; ++head) {

auto one = modules::SliceModule({1, head, 1}).build(ctx, input);

for (int64_t rep = 0; rep < repeats; ++rep) {

heads.push_back(one);

}

}

auto output = heads.front();

for (size_t i = 1; i < heads.size(); ++i) {

output = modules::ConcatModule({1}).build(ctx, output, heads[i]);

}

return output;

core::ModuleBuildContext & ctx,

const core::TensorValue & q_heads,

const core::TensorValue & k_heads,

const core::TensorValue & v_heads,

int64_t dim,

const std::optional<core::TensorValue> & attention_mask,

core::BackendType backend_type) {

(void)backend_type;

return modules::ScaledDotProductAttentionModule({

dim,

modules::ScaledDotProductAttentionLowering::Explicit,

GGML_PREC_F32,

}).build(ctx, q_heads, k_heads, v_heads, attention_mask);

core::ModuleBuildContext & ctx,

const core::TensorValue & hidden_states,

const core::TensorValue & positions,

const AceStepDiTAttentionWeights & weights,

const AceStepDiffusionConfig & config,

const std::optional<core::TensorValue> & attention_mask,

core::BackendType backend_type,

const std::optional<PrecomputedCrossAttentionKV> & cross_attention_kv = std::nullopt,

const std::optional<core::TensorValue> & encoder_hidden_states = std::nullopt) {

const int64_t dim = ace_step_diffusion_attention_head_dim(config, "ACE-Step diffusion");

const int64_t kv_repeats = config.num_attention_heads / config.num_key_value_heads;

const bool is_cross = cross_attention_kv.has_value() || encoder_hidden_states.has_value();

auto q = modules::LinearModule({config.hidden_size, config.num_attention_heads * dim, false, GGML_PREC_F32})

.build(ctx, hidden_states, {weights.q_weight, std::nullopt});

core::TensorValue k;

core::TensorValue v;

if (cross_attention_kv.has_value()) {

k = cross_attention_kv->key;

v = cross_attention_kv->value;

} else {

const auto & kv_source = is_cross ? *encoder_hidden_states : hidden_states;

k = modules::LinearModule({config.hidden_size, config.num_key_value_heads * dim, false, GGML_PREC_F32})

.build(ctx, kv_source, {weights.k_weight, std::nullopt});

v = modules::LinearModule({config.hidden_size, config.num_key_value_heads * dim, false, GGML_PREC_F32})

.build(ctx, kv_source, {weights.v_weight, std::nullopt});

}

q = modules::RMSNormModule({dim, config.rms_norm_eps, true, false})

.build(ctx, reshape_heads(ctx, q, config.num_attention_heads, dim), {weights.q_norm, std::nullopt});

if (!cross_attention_kv.has_value()) {

k = modules::RMSNormModule({dim, config.rms_norm_eps, true, false})

.build(ctx, reshape_heads(ctx, k, config.num_key_value_heads, dim), {weights.k_norm, std::nullopt});

v = reshape_heads(ctx, v, config.num_key_value_heads, dim);

}

if (!is_cross) {

q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, q, positions);

k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, k, positions);

}

auto q_heads = ensure_contiguous(

ctx,

modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q));

auto k_heads = repeat_kv_heads(

ctx,

modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k),

kv_repeats);

auto v_heads = repeat_kv_heads(

ctx,

modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v),

kv_repeats);

k_heads = ensure_contiguous(ctx, k_heads);

v_heads = ensure_contiguous(ctx, v_heads);

auto context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask, backend_type);

context = ensure_contiguous(ctx, context);

// The attention width is heads x head_dim, which only equals hidden_size when

// head_dim was derived from it. XL states head_dim outright (32 x 128 against

// a hidden size of 2560), so o_proj is the one rectangular projection here.

const int64_t attention_size = config.num_attention_heads * dim;

context = core::reshape_tensor(

ctx,

context,

core::TensorShape::from_dims({hidden_states.shape.dims[0], hidden_states.shape.dims[1], attention_size}));

return modules::LinearModule({attention_size, config.hidden_size, false, GGML_PREC_F32})

.build(ctx, context, {weights.out_weight, std::nullopt});

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

const core::TensorValue & norm_weight,

const core::TensorValue & one,

const core::TensorValue & shift,

const core::TensorValue & scale,

float rms_eps) {

auto hidden = modules::RMSNormModule({input.shape.dims[input.shape.rank - 1], rms_eps, true, false})

.build(ctx, input, {norm_weight, std::nullopt});

auto scaled = core::wrap_tensor(ggml_add(ctx.ggml, scale.tensor, one.tensor), scale.shape, GGML_TYPE_F32);

hidden = modules::MulModule{}.build(ctx, hidden, scaled);

return modules::AddModule{}.build(ctx, hidden, shift);

core::ModuleBuildContext & ctx,

const core::TensorValue & conditioning,

int64_t seq_len) {

auto reshaped = core::reshape_tensor(

ctx,

conditioning,

core::TensorShape::from_dims({conditioning.shape.dims[0], 1, conditioning.shape.dims[1]}));

return modules::RepeatModule({core::TensorShape::from_dims({conditioning.shape.dims[0], seq_len, conditioning.shape.dims[1]})})

.build(ctx, reshaped);

core::ModuleBuildContext & ctx,

ggml_tensor * freqs_tensor,

const core::TensorValue & timestep) {

const int64_t batch_size = timestep.shape.dims[0];

auto freqs = core::wrap_tensor(freqs_tensor, core::TensorShape::from_dims({128}), GGML_TYPE_F32);

auto expanded_freqs = core::reshape_tensor(ctx, freqs, core::TensorShape::from_dims({1, 128}));

expanded_freqs = modules::RepeatModule({core::TensorShape::from_dims({batch_size, 128})}).build(ctx, expanded_freqs);

auto expanded_time = modules::RepeatModule({core::TensorShape::from_dims({batch_size, 128})}).build(ctx, timestep);

return modules::MulModule{}.build(ctx, expanded_time, expanded_freqs);

core::ModuleBuildContext & ctx,

ggml_tensor * freqs_tensor,

const core::TensorValue & timestep,

const AceStepTimeEmbeddingWeights & weights,

int64_t hidden_size) {

auto scaled_t = core::wrap_tensor(ggml_scale(ctx.ggml, timestep.tensor, 1000.0F), timestep.shape, GGML_TYPE_F32);

auto args = build_time_frequency_tensor(ctx, freqs_tensor, scaled_t);

auto cos_part = core::wrap_tensor(ggml_cos(ctx.ggml, args.tensor), args.shape, GGML_TYPE_F32);

auto sin_part = core::wrap_tensor(ggml_sin(ctx.ggml, args.tensor), args.shape, GGML_TYPE_F32);

auto embedding = modules::ConcatModule({1}).build(ctx, cos_part, sin_part);

auto temb = modules::LinearModule({256, hidden_size, true, GGML_PREC_F32}).build(ctx, embedding, weights.fc1);

temb = modules::SiluModule{}.build(ctx, temb);

temb = modules::LinearModule({hidden_size, hidden_size, true, GGML_PREC_F32}).build(ctx, temb, weights.fc2);

auto proj_in = modules::SiluModule{}.build(ctx, temb);

auto timestep_proj =

modules::LinearModule({hidden_size, hidden_size * 6, true, GGML_PREC_F32}).build(ctx, proj_in, weights.time_proj);

timestep_proj = core::reshape_tensor(

ctx,

timestep_proj,

core::TensorShape::from_dims({timestep_proj.shape.dims[0], int64_t{6}, hidden_size}));

return TimeEmbeddingOutputs{temb, timestep_proj};

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

const modules::LinearWeights & gate_proj,

const modules::LinearWeights & up_proj,

const modules::LinearWeights & down_proj,

const AceStepDiffusionConfig & config) {

auto gate = modules::LinearModule({config.hidden_size, config.intermediate_size, false, GGML_PREC_F32})

.build(ctx, input, gate_proj);

gate = modules::SiluModule{}.build(ctx, gate);

auto up = modules::LinearModule({config.hidden_size, config.intermediate_size, false, GGML_PREC_F32})

.build(ctx, input, up_proj);

auto ff = modules::MulModule{}.build(ctx, gate, up);

return modules::LinearModule({config.intermediate_size, config.hidden_size, false, GGML_PREC_F32})

.build(ctx, ff, down_proj);

core::ModuleBuildContext & ctx,

const core::TensorValue & conditioning,

int64_t plane,

int64_t hidden_size,

int64_t seq_len) {

auto slice = modules::SliceModule({1, plane, 1}).build(ctx, conditioning);

slice = ensure_contiguous(ctx, slice);

slice = core::reshape_tensor(

ctx,

slice,

core::TensorShape::from_dims({conditioning.shape.dims[0], hidden_size}));

return expand_conditioning(ctx, slice, seq_len);

core::ModuleBuildContext & ctx,

const core::TensorValue & hidden_states,

const core::TensorValue & positions,

const core::TensorValue & timestep_proj,

const core::TensorValue & one,

const std::optional<core::TensorValue> & self_attention_mask,

const core::TensorValue & encoder_attention_mask,

const AceStepDiTLayerWeights & weights,

const AceStepDiffusionConfig & config,

core::BackendType backend_type,

const std::optional<PrecomputedCrossAttentionKV> & cross_attention_kv,

const std::optional<core::TensorValue> & encoder_hidden_states = std::nullopt) {

const int64_t hidden_size = config.hidden_size;

const int64_t seq_len = hidden_states.shape.dims[1];

auto scale_shift_table = modules::RepeatModule({timestep_proj.shape}).build(ctx, ensure_f32(ctx, weights.scale_shift_table));

auto modulation = modules::AddModule{}.build(ctx, scale_shift_table, timestep_proj);

auto shift_msa = slice_conditioning_plane(ctx, modulation, 0, hidden_size, seq_len);

auto scale_msa = slice_conditioning_plane(ctx, modulation, 1, hidden_size, seq_len);

auto gate_msa = slice_conditioning_plane(ctx, modulation, 2, hidden_size, seq_len);

auto c_shift_msa = slice_conditioning_plane(ctx, modulation, 3, hidden_size, seq_len);

auto c_scale_msa = slice_conditioning_plane(ctx, modulation, 4, hidden_size, seq_len);

auto c_gate_msa = slice_conditioning_plane(ctx, modulation, 5, hidden_size, seq_len);

auto norm_hidden = apply_modulated_rms_norm(

ctx,

hidden_states,

weights.self_attn_norm,

one,

shift_msa,

scale_msa,

config.rms_norm_eps);

auto attn_out = build_attention(

ctx,

norm_hidden,

positions,

weights.self_attn,

config,

self_attention_mask,

backend_type);

attn_out = modules::MulModule{}.build(ctx, attn_out, gate_msa);

auto x = modules::AddModule{}.build(ctx, hidden_states, attn_out);

auto cross_norm = modules::RMSNormModule({hidden_size, config.rms_norm_eps, true, false})

.build(ctx, x, {weights.cross_attn_norm, std::nullopt});

auto cross_out = build_attention(

ctx,

cross_norm,

positions,

weights.cross_attn,

config,

encoder_attention_mask,

backend_type,

cross_attention_kv,

encoder_hidden_states);

x = modules::AddModule{}.build(ctx, x, cross_out);

auto mlp_norm = apply_modulated_rms_norm(

ctx,

x,

weights.mlp_norm,

one,

c_shift_msa,

c_scale_msa,

config.rms_norm_eps);

auto ff = build_mlp(ctx, mlp_norm, weights.mlp_gate, weights.mlp_up, weights.mlp_down, config);

ff = modules::MulModule{}.build(ctx, ff, c_gate_msa);

return modules::AddModule{}.build(ctx, x, ff);

int64_t seq_len,

const std::vector<int32_t> & encoder_attention_mask,

int64_t valid_seq_len) {

const int64_t encoder_tokens = static_cast<int64_t>(encoder_attention_mask.size());

if (seq_len <= 0 || encoder_tokens <= 0 || valid_seq_len <= 0 || valid_seq_len > seq_len) {

throw std::runtime_error("ACE-Step diffusion cross attention mask shape is invalid");

}

std::vector<float> values(static_cast<size_t>(seq_len * encoder_tokens), 0.0F);

const float masked = std::numeric_limits<float>::lowest();

for (int64_t q = 0; q < seq_len; ++q) {

for (int64_t k = 0; k < encoder_tokens; ++k) {

if (q >= valid_seq_len) {

values[static_cast<size_t>(q * encoder_tokens + k)] = masked;

}

}

}

return values;

std::vector<float> values(static_cast<size_t>(tokens * tokens), 0.0F);

const float masked = std::numeric_limits<float>::lowest();

for (int64_t q = 0; q < tokens; ++q) {

for (int64_t k = 0; k < tokens; ++k) {

if (std::llabs(q - k) > sliding_window) {

values[static_cast<size_t>(q * tokens + k)] = masked;

}

}

}

return values;

if (seq_len <= 0 || valid_seq_len <= 0 || valid_seq_len > seq_len) {

throw std::runtime_error("ACE-Step diffusion self attention padding mask is invalid");

}

std::vector<float> values(static_cast<size_t>(seq_len * seq_len), 0.0F);

const float masked = std::numeric_limits<float>::lowest();

for (int64_t q = 0; q < valid_seq_len; ++q) {

for (int64_t k = valid_seq_len; k < seq_len; ++k) {

values[static_cast<size_t>(q * seq_len + k)] = masked;

}

}

return values;

std::vector<float> freqs(128, 0.0F);

for (int64_t i = 0; i < 128; ++i) {

freqs[static_cast<size_t>(i)] =

std::exp(-std::log(10000.0F) * static_cast<float>(i) / 128.0F);

}

return freqs;

return {

1.0F, 0.95454544F, 0.93333334F, 0.9F, 0.875F,

0.85714287F, 0.83333331F, 0.76923078F, 0.75F,

0.66666669F, 0.64285713F, 0.625F, 0.54545456F,

0.5F, 0.4F, 0.375F, 0.3F, 0.25F, 0.22222222F, 0.125F

};

if (shift == 1.0F) {

return {1.0F, 0.875F, 0.75F, 0.625F, 0.5F, 0.375F, 0.25F, 0.125F};

}

if (shift == 2.0F) {

return {1.0F, 0.93333334F, 0.85714287F, 0.76923078F, 0.66666669F, 0.54545456F, 0.4F, 0.22222222F};

}

return {1.0F, 0.95454544F, 0.9F, 0.83333331F, 0.75F, 0.64285713F, 0.5F, 0.3F};

constexpr float kValid[] = {1.0F, 2.0F, 3.0F};

float best = kValid[0];

float best_distance = std::abs(shift - best);

for (float candidate : kValid) {

const float distance = std::abs(shift - candidate);

if (distance < best_distance) {

best = candidate;

best_distance = distance;

}

}

return best;

float best = valid.front();

float best_distance = std::abs(t - best);

for (float candidate : valid) {

const float distance = std::abs(t - candidate);

if (distance < best_distance) {

best = candidate;

best_distance = distance;

}

}

return best;

std::vector<float> schedule;

if (!options.timesteps.empty()) {

schedule = options.timesteps;

while (!schedule.empty() && schedule.back() == 0.0F) {

schedule.pop_back();

}

if (is_turbo && schedule.size() > 20) {

schedule.resize(20);

}

if (!schedule.empty() && is_turbo) {

const auto valid = valid_timesteps();

for (float & t : schedule) {

t = nearest_valid_timestep(t, valid);

}

return schedule;

}

if (!schedule.empty()) {

return schedule;

}

}

if (options.num_inference_steps > 0) {

const int64_t n = is_turbo ? std::min<int64_t>(options.num_inference_steps, 20) : options.num_inference_steps;

schedule.reserve(static_cast<size_t>(n));

for (int64_t i = 0; i < n; ++i) {

float t = 1.0F - static_cast<float>(i) / static_cast<float>(n);

if (options.shift != 1.0F) {

t = options.shift * t / (1.0F + (options.shift - 1.0F) * t);

}

schedule.push_back(t);

}

return schedule;

}

return shift_schedule(nearest_valid_shift(options.shift));

return engine::sampling::generate_torch_cuda_randn(

count,

seed,

engine::sampling::TorchRandnPrecision::Float32,

start_index);

const std::vector<float> & x,

float t,

uint32_t seed,

uint64_t & noise_offset,

const std::vector<float> * noise_override = nullptr) {

std::vector<float> sampled_noise;

const std::vector<float> * noise = noise_override;

if (noise == nullptr) {

sampled_noise = gaussian_noise(x.size(), seed, noise_offset);

noise_offset += static_cast<uint64_t>(x.size());

noise = &sampled_noise;

}

if (noise->size() != x.size()) {

throw std::runtime_error("ACE-Step diffusion renoise input/noise size mismatch");

}

return engine::sampling::renoise(x, *noise, t);

std::vector<float> & velocity,

const std::vector<float> & hidden,

int64_t frames,

int64_t channels,

float threshold) {

const size_t count = static_cast<size_t>(frames * channels);

if (velocity.size() != count || hidden.size() != count) {

throw std::runtime_error("ACE-Step diffusion velocity clamp shape mismatch");

}

engine::sampling::clamp_velocity_norm(velocity, hidden, threshold);

std::vector<float> & velocity,

const std::vector<float> & previous_velocity,

float factor) {

if (factor <= 0.0F || previous_velocity.empty()) {

return;

}

for (size_t i = 0; i < velocity.size(); ++i) {

velocity[i] = (1.0F - factor) * velocity[i] + factor * previous_velocity[i];

}

if (value < 0.0F) {

throw std::runtime_error("ACE-Step diffusion expected a nonnegative value for Python-style rounding");

}

const double x = static_cast<double>(value);

const double floor_value = std::floor(x);

const double fraction = x - floor_value;

int64_t rounded = static_cast<int64_t>(floor_value);

if (fraction > 0.5) {

++rounded;

} else if (fraction == 0.5 && (rounded % 2) != 0) {

++rounded;

}

return rounded;

const std::vector<float> & pred_cond,

const std::vector<float> & pred_uncond,

float guidance_scale) {

return engine::sampling::cfg_guidance(pred_cond, pred_uncond, guidance_scale);

const std::vector<float> & pred_cond,

const std::vector<float> & pred_uncond,

float guidance_scale,

int64_t frames,

int64_t channels,

std::vector<float> & momentum) {

return engine::sampling::apg_guidance(pred_cond, pred_uncond, guidance_scale, frames, channels, momentum);

const std::vector<float> & latents,

const std::vector<float> & pred_cond,

const std::vector<float> & pred_uncond,

float sigma,

float guidance_scale,

int64_t frames,

int64_t channels) {

return engine::sampling::adg_guidance(latents, pred_cond, pred_uncond, sigma, guidance_scale, frames, channels);

if (latents.frames <= 0 || latents.channels <= 0 || padded_frames < latents.frames) {

throw std::runtime_error("ACE-Step diffusion latent padding request is invalid");

}

std::vector<float> padded(static_cast<size_t>(padded_frames * latents.channels), 0.0F);

std::copy(latents.values.begin(), latents.values.end(), padded.begin());

return padded;

if (latents.channels != expected_channels) {

throw std::runtime_error("ACE-Step diffusion context latent channel count mismatch");

}

return pad_latent_values(latents, padded_frames);

const AceStepEncoderConditioning & conditioning,

int64_t padded_tokens,

int64_t hidden_size) {

if (conditioning.hidden_size != hidden_size || padded_tokens < conditioning.tokens) {

throw std::runtime_error("ACE-Step diffusion encoder hidden padding request is invalid");

}

std::vector<float> padded(static_cast<size_t>(padded_tokens * hidden_size), 0.0F);

std::copy(conditioning.values.begin(), conditioning.values.end(), padded.begin());

return padded;

const std::vector<float> & null_condition_emb,

int64_t padded_tokens,

int64_t hidden_size) {

if (padded_tokens <= 0 || hidden_size <= 0 ||

static_cast<int64_t>(null_condition_emb.size()) != hidden_size) {

throw std::runtime_error("ACE-Step diffusion null condition embedding shape mismatch");

}

std::vector<float> out(static_cast<size_t>(padded_tokens * hidden_size), 0.0F);

for (int64_t token = 0; token < padded_tokens; ++token) {

std::copy(

null_condition_emb.begin(),

null_condition_emb.end(),

out.begin() + static_cast<std::ptrdiff_t>(token * hidden_size));

}

return out;

if (batch_size <= 0) {

throw std::runtime_error("ACE-Step diffusion batch size is invalid");

}

if (batch_size == 1) {

return values;

}

std::vector<float> out;

out.reserve(values.size() * static_cast<size_t>(batch_size));

for (int64_t batch = 0; batch < batch_size; ++batch) {

out.insert(out.end(), values.begin(), values.end());

}

return out;

const AceStepEncoderConditioning & conditioning,

int64_t padded_tokens) {

if (static_cast<int64_t>(conditioning.attention_mask.size()) != conditioning.tokens || padded_tokens < conditioning.tokens) {

throw std::runtime_error("ACE-Step diffusion encoder attention mask padding request is invalid");

}

std::vector<int32_t> padded(static_cast<size_t>(padded_tokens), 0);

std::copy(conditioning.attention_mask.begin(), conditioning.attention_mask.end(), padded.begin());

return padded;

std::vector<float> & hidden,

const std::vector<float> & clean_src_latents,

const std::vector<int32_t> & repaint_mask,

float next_timestep,

const std::vector<float> & noise,

int64_t channels) {

if (hidden.size() != clean_src_latents.size() || hidden.size() != noise.size()) {

throw std::runtime_error("ACE-Step repaint injection latent shape mismatch");

}

if (channels <= 0 ||

hidden.size() % static_cast<size_t>(channels) != 0 ||

repaint_mask.size() != hidden.size() / static_cast<size_t>(channels)) {

throw std::runtime_error("ACE-Step repaint injection mask shape mismatch");

}

for (size_t frame = 0; frame < repaint_mask.size(); ++frame) {

if (repaint_mask[frame] != 0) {

continue;

}

const size_t offset = frame * static_cast<size_t>(channels);

for (int64_t channel = 0; channel < channels; ++channel) {

const size_t index = offset + static_cast<size_t>(channel);

hidden[index] = next_timestep * noise[index] + (1.0F - next_timestep) * clean_src_latents[index];

}

}

std::vector<float> & generated,

const std::vector<float> & clean_src_latents,

const std::vector<int32_t> & repaint_mask,

int64_t crossfade_frames,

int64_t channels) {

const auto soft_mask = build_soft_repaint_mask(repaint_mask, crossfade_frames);

engine::sampling::blend_by_mask_in_place(generated, clean_src_latents, soft_mask, channels);

if (!dcw_mode_is_valid(options.dcw_mode)) {

throw std::runtime_error("ACE-Step diffusion dcw_mode is invalid");

}

if (!options.dcw_enabled) {

return false;

}

if (reference_skips_wavelet_dcw && options.dcw_mode != "pix") {

// In the validated index-tts reference environment, the 2B Base path

// logs wavelet DCW as unavailable and leaves the latent unchanged.

return false;

}

if (options.dcw_mode == "double") {

return options.dcw_scaler != 0.0F || options.dcw_high_scaler != 0.0F;

}

return options.dcw_scaler != 0.0F;

const char * name = nullptr;

std::vector<float> dec_lo;

std::vector<float> dec_hi;

std::vector<float> low;

std::vector<float> high;

const WaveletSpec * spec = nullptr;

int64_t frames = 0;

int64_t channels = 0;

const std::vector<float> & values,

int64_t frames,

int64_t channels,

const WaveletSpec & spec) {

const int64_t filter_size = static_cast<int64_t>(spec.dec_lo.size());

const int64_t band_frames = (frames + filter_size - 1) / 2;

WaveletBands bands;

bands.low.assign(static_cast<size_t>(band_frames * channels), 0.0F);

bands.high.assign(static_cast<size_t>(band_frames * channels), 0.0F);

bands.spec = &spec;

bands.frames = frames;

bands.channels = channels;

for (int64_t pair = 0; pair < band_frames; ++pair) {

for (int64_t channel = 0; channel < channels; ++channel) {

const size_t index = static_cast<size_t>(pair * channels + channel);

float low = 0.0F;

float high = 0.0F;

for (int64_t k = 0; k < filter_size; ++k) {

const int64_t src_index = 2 * pair + 1 - k;

if (src_index < 0 || src_index >= frames) {

continue;

}

const float value = values[static_cast<size_t>(src_index * channels + channel)];

low += spec.dec_lo[static_cast<size_t>(k)] * value;

high += spec.dec_hi[static_cast<size_t>(k)] * value;

}

bands.low[index] = low;

bands.high[index] = high;

}

}

return bands;

if (bands.spec == nullptr) {

throw std::runtime_error("ACE-Step diffusion wavelet IDWT requires valid filter spec");

}

const int64_t filter_size = static_cast<int64_t>(bands.spec->dec_lo.size());

const int64_t band_frames = static_cast<int64_t>(bands.low.size() / bands.channels);

std::vector<float> out(static_cast<size_t>(bands.frames * bands.channels), 0.0F);

for (int64_t pair = 0; pair < band_frames; ++pair) {

for (int64_t channel = 0; channel < bands.channels; ++channel) {

const size_t index = static_cast<size_t>(pair * bands.channels + channel);

for (int64_t k = 0; k < filter_size; ++k) {

const int64_t dst_index = 2 * pair + 1 - k;

if (dst_index < 0 || dst_index >= bands.frames) {

continue;

}

out[static_cast<size_t>(dst_index * bands.channels + channel)] +=

bands.low[index] * bands.spec->dec_lo[static_cast<size_t>(k)] +

bands.high[index] * bands.spec->dec_hi[static_cast<size_t>(k)];

}

}

}

return out;

std::vector<float> & hidden,

const std::vector<float> & denoised,

int64_t frames,

int64_t channels,

float t_curr,

bool reference_skips_wavelet_dcw,

const AceStepGenerationOptions & options) {

if (!dcw_is_active(options, reference_skips_wavelet_dcw)) {

return;

}

if (hidden.size() < denoised.size()) {

throw std::runtime_error("ACE-Step diffusion DCW hidden/denoised size mismatch");

}

if (options.dcw_mode == "pix") {

if (options.dcw_scaler == 0.0F) {

return;

}

for (size_t i = 0; i < denoised.size(); ++i) {

hidden[i] = hidden[i] + options.dcw_scaler * (hidden[i] - denoised[i]);

}

return;

}

const auto & spec = lookup_wavelet_spec(options.dcw_wavelet);

auto x_bands = wavelet_dwt_prefix(hidden, frames, channels, spec);

const auto y_bands = wavelet_dwt_prefix(denoised, frames, channels, spec);

const float low_scaler = t_curr * options.dcw_scaler;

const float high_scaler = (1.0F - t_curr) *

(options.dcw_mode == "double" ? options.dcw_high_scaler : options.dcw_scaler);

if (options.dcw_mode == "low" || options.dcw_mode == "double") {

for (size_t i = 0; i < x_bands.low.size(); ++i) {

x_bands.low[i] = x_bands.low[i] + low_scaler * (x_bands.low[i] - y_bands.low[i]);

}

}

if (options.dcw_mode == "high" || options.dcw_mode == "double") {

for (size_t i = 0; i < x_bands.high.size(); ++i) {

x_bands.high[i] = x_bands.high[i] + high_scaler * (x_bands.high[i] - y_bands.high[i]);

}

}

const auto corrected = wavelet_idwt_prefix(x_bands);

std::copy(corrected.begin(), corrected.end(), hidden.begin());

if (noise_file.empty()) {

return gaussian_noise(count, seed);

}

auto values = io::read_f32_file(noise_file);

if (values.size() != count) {

throw std::runtime_error(

"ACE-Step diffusion noise file size mismatch: expected " +

std::to_string(count) + " floats, got " + std::to_string(values.size()));

}

return values;

NoiseSchedule schedule;

if (noise_file.empty()) {

schedule.initial_noise = gaussian_noise(count, seed);

return schedule;

}

auto values = io::read_f32_file(noise_file);

if (values.empty() || values.size() % count != 0) {

throw std::runtime_error(

"ACE-Step diffusion noise file size mismatch: expected a positive multiple of " +

std::to_string(count) + " floats, got " + std::to_string(values.size()));

}

const size_t chunks = values.size() / count;

schedule.initial_noise.assign(values.begin(), values.begin() + static_cast<std::ptrdiff_t>(count));

schedule.renoise_noises.reserve(chunks > 0 ? chunks - 1 : 0);

for (size_t chunk = 1; chunk < chunks; ++chunk) {

const auto begin = values.begin() + static_cast<std::ptrdiff_t>(chunk * count);

schedule.renoise_noises.emplace_back(begin, begin + static_cast<std::ptrdiff_t>(count));

}

return schedule;

public:

CrossAttentionCacheGraph(

std::shared_ptr<const AceStepAssets> assets,

ggml_backend_t backend,

core::BackendType backend_type,

int threads,

std::shared_ptr<const AceStepDiffusionWeights> weights,

int64_t encoder_tokens_capacity,

size_t graph_arena_bytes)

: assets_(std::move(assets)),

backend_(backend),

backend_type_(backend_type),

threads_(threads),

weights_(std::move(weights)),

encoder_tokens_(encoder_tokens_capacity) {

build(graph_arena_bytes);

}

~CrossAttentionCacheGraph() {

if (backend_ != nullptr && graph_ != nullptr) {

engine::core::release_backend_graph_resources(backend_, graph_);

}

if (buffer_ != nullptr) {

ggml_backend_buffer_free(buffer_);

}

}

bool can_run(int64_t encoder_tokens) const noexcept {

return encoder_tokens <= encoder_tokens_;

}

int64_t encoder_token_capacity() const noexcept { return encoder_tokens_; }

struct Output {

std::vector<std::vector<float>> keys;

std::vector<std::vector<float>> values;

};

Output run(const std::vector<float> & encoder_hidden_states) const {

const auto & config = assets_->config.diffusion;

// Still in the encoder's width here: the condition embedder inside the

// graph is what lifts it to the DiT's.

const int64_t encoder_hidden_size = assets_->config.encoder.hidden_size;

if (static_cast<int64_t>(encoder_hidden_states.size()) != encoder_tokens_ * encoder_hidden_size) {

throw std::runtime_error("ACE-Step diffusion cross-attention cache encoder hidden state shape mismatch");

}

core::write_tensor_f32(encoder_value_, encoder_hidden_states);

core::set_backend_threads(backend_, threads_);

const ggml_status status = engine::core::compute_backend_graph(backend_, graph_);

ggml_backend_synchronize(backend_);

if (status != GGML_STATUS_SUCCESS) {

throw std::runtime_error("ACE-Step diffusion cross-attention cache graph compute failed");

}

Output out;

out.keys.resize(key_outputs_.size());

out.values.resize(value_outputs_.size());