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

if (ctx != nullptr) {

ggml_free(ctx);

}

}

const MimiCodecTransformerLayerWeights & weights) {

return {

weights.norm1,

weights.self_attn,

weights.layer_scale1,

weights.norm2,

weights.feed_forward,

weights.layer_scale2,

};

if (frames <= 0 || cache_steps < 0 || context <= 0 || current_end < 0) {

throw std::runtime_error("Mimi codec transformer attention mask received invalid dimensions");

}

std::vector<float> values(static_cast<size_t>(frames * (cache_steps + frames)), kMaskedAttentionBias);

for (int64_t query = 0; query < frames; ++query) {

const int64_t query_position = current_end + query;

for (int64_t slot = 0; slot < cache_steps; ++slot) {

int64_t key_position = -1;

if (current_end <= cache_steps) {

if (slot < current_end) {

key_position = slot;

}

} else {

const int64_t end_slot = current_end % cache_steps;

const int64_t delta = slot - end_slot;

key_position = delta < 0 ? current_end + delta : current_end + delta - cache_steps;

}

const int64_t distance = query_position - key_position;

if (key_position >= 0 && distance >= 0 && distance < context) {

values[static_cast<size_t>(query * (cache_steps + frames) + slot)] = 0.0F;

}

}

for (int64_t local_key = 0; local_key <= query; ++local_key) {

const int64_t key_position = current_end + local_key;

if (query_position - key_position < context) {

values[static_cast<size_t>(query * (cache_steps + frames) + cache_steps + local_key)] = 0.0F;

}

}

}

return values;

int64_t channels = 0;

int64_t history_frames = 0;

bool first = true;

std::vector<float> previous;

int64_t channels = 0;

int64_t partial_frames = 0;

std::vector<float> partial;

StreamingConvTranspose1dState encoder_rate_upsample;

MimiTransformerState transformer;

StreamingConv1dState input_projection;

std::vector<StreamingConvTranspose1dState> stage_upsamples;

std::vector<std::array<StreamingConv1dState, 2>> stage_residual_convs;

StreamingConv1dState output_projection;

StreamingConv1dState input_projection;

std::vector<std::array<StreamingConv1dState, 2>> stage_residual_convs;

std::vector<StreamingConv1dState> stage_downsamples;

StreamingConv1dState output_projection;

MimiTransformerState transformer;

StreamingConv1dState downsample;

ggml_gallocr_t & gallocr,

ggml_backend_buffer_t & params_buffer,

ggml_context *& context) {

if (gallocr != nullptr) {

ggml_gallocr_free(gallocr);

gallocr = nullptr;

}

if (params_buffer != nullptr) {

ggml_backend_buffer_free(params_buffer);

params_buffer = nullptr;

}

if (context != nullptr) {

ggml_free(context);

context = nullptr;

}

StreamingConv1dState state;

state.channels = channels;

state.history_frames = std::max<int64_t>(0, (kernel_size - 1) * dilation + 1 - stride);

state.previous.assign(static_cast<size_t>(state.channels * state.history_frames), 0.0F);

return state;

StreamingConvTranspose1dState state;

state.channels = channels;

state.partial_frames = std::max<int64_t>(0, kernel_size - stride);

state.partial.assign(static_cast<size_t>(state.channels * state.partial_frames), 0.0F);

return state;

MimiDecoderState state;

state.encoder_rate_upsample = make_streaming_convtranspose_state(config.hidden_size, 4, 2);

state.input_projection = make_streaming_conv1d_state(config.hidden_size, 7, 1, 1);

state.stage_upsamples.push_back(make_streaming_convtranspose_state(512, 16, 8));

state.stage_upsamples.push_back(make_streaming_convtranspose_state(256, 12, 6));

state.stage_upsamples.push_back(make_streaming_convtranspose_state(128, 10, 5));

state.stage_upsamples.push_back(make_streaming_convtranspose_state(64, 8, 4));

state.stage_residual_convs = {

std::array<StreamingConv1dState, 2>{

make_streaming_conv1d_state(512, 3, 1, 1),

make_streaming_conv1d_state(256, 1, 1, 1),

},

std::array<StreamingConv1dState, 2>{

make_streaming_conv1d_state(256, 3, 1, 1),

make_streaming_conv1d_state(128, 1, 1, 1),

},

std::array<StreamingConv1dState, 2>{

make_streaming_conv1d_state(128, 3, 1, 1),

make_streaming_conv1d_state(64, 1, 1, 1),

},

std::array<StreamingConv1dState, 2>{

make_streaming_conv1d_state(64, 3, 1, 1),

make_streaming_conv1d_state(32, 1, 1, 1),

},

};

state.output_projection = make_streaming_conv1d_state(64, 3, 1, 1);

return state;

MimiEncoderState state;

state.input_projection = make_streaming_conv1d_state(config.channels, 7, 1, 1);

state.stage_residual_convs = {

std::array<StreamingConv1dState, 2>{

make_streaming_conv1d_state(64, 3, 1, 1),

make_streaming_conv1d_state(32, 1, 1, 1),

},

std::array<StreamingConv1dState, 2>{

make_streaming_conv1d_state(128, 3, 1, 1),

make_streaming_conv1d_state(64, 1, 1, 1),

},

std::array<StreamingConv1dState, 2>{

make_streaming_conv1d_state(256, 3, 1, 1),

make_streaming_conv1d_state(128, 1, 1, 1),

},

std::array<StreamingConv1dState, 2>{

make_streaming_conv1d_state(512, 3, 1, 1),

make_streaming_conv1d_state(256, 1, 1, 1),

},

};

state.stage_downsamples.push_back(make_streaming_conv1d_state(64, 8, 4, 1));

state.stage_downsamples.push_back(make_streaming_conv1d_state(128, 10, 5, 1));

state.stage_downsamples.push_back(make_streaming_conv1d_state(256, 12, 6, 1));

state.stage_downsamples.push_back(make_streaming_conv1d_state(512, 16, 8, 1));

state.output_projection = make_streaming_conv1d_state(1024, 3, 1, 1);

state.downsample = make_streaming_conv1d_state(config.hidden_size, 4, 2, 1);

return state;

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

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

const modules::StreamingConv1dWeights & weights,

int64_t in_channels,

int64_t out_channels,

int64_t kernel_size,

int stride,

int dilation,

bool use_bias) {

auto full_input = input;

if (history.has_value()) {

full_input = modules::ConcatModule({2}).build(ctx, *history, input);

}

auto output = modules::Conv1dModule({

in_channels,

out_channels,

kernel_size,

stride,

0,

dilation,

use_bias,

}).build(ctx, full_input, weights);

std::optional<core::TensorValue> next_history = std::nullopt;

if (history.has_value()) {

next_history = modules::SliceModule({

2,

full_input.shape.dims[2] - history->shape.dims[2],

history->shape.dims[2],

}).build(ctx, full_input);

}

return {output, next_history};

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

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

const modules::ConvTranspose1dWeights & weights,

int64_t in_channels,

int64_t out_channels,

int64_t kernel_size,

int stride,

bool use_bias) {

auto raw = modules::ConvTranspose1dModule({

in_channels,

out_channels,

kernel_size,

stride,

0,

1,

use_bias,

}).build(ctx, input, weights);

if (!partial.has_value()) {

return {raw, std::nullopt};

}

const int64_t partial_frames = partial->shape.dims[2];

auto prefix = modules::SliceModule({2, 0, partial_frames}).build(ctx, raw);

prefix = modules::AddModule().build(ctx, prefix, *partial);

auto suffix = modules::SliceModule({2, partial_frames, raw.shape.dims[2] - partial_frames}).build(ctx, raw);

auto combined = modules::ConcatModule({2}).build(ctx, prefix, suffix);

auto output = modules::SliceModule({2, 0, raw.shape.dims[2] - partial_frames}).build(ctx, combined);

auto next_partial = modules::SliceModule({

2,

raw.shape.dims[2] - partial_frames,

partial_frames,

}).build(ctx, combined);

return {output, next_partial};

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

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

const modules::DepthwiseConvTranspose1dWeights & weights,

int64_t channels,

int64_t kernel_size,

int stride) {

auto raw = modules::DepthwiseConvTranspose1dModule({

channels,

kernel_size,

stride,

false,

}).build(ctx, input, weights);

if (!partial.has_value()) {

return {raw, std::nullopt};

}

const int64_t partial_frames = partial->shape.dims[2];

auto prefix = modules::SliceModule({2, 0, partial_frames}).build(ctx, raw);

prefix = modules::AddModule().build(ctx, prefix, *partial);

auto suffix = modules::SliceModule({2, partial_frames, raw.shape.dims[2] - partial_frames}).build(ctx, raw);

auto combined = modules::ConcatModule({2}).build(ctx, prefix, suffix);

auto output = modules::SliceModule({2, 0, raw.shape.dims[2] - partial_frames}).build(ctx, combined);

auto next_partial = modules::SliceModule({

2,

raw.shape.dims[2] - partial_frames,

partial_frames,

}).build(ctx, combined);

return {output, next_partial};

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

const MimiCodecResidualWeights & weights,

int64_t channels,

int64_t hidden_channels,

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

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

std::vector<core::TensorValue> & next_histories) {

auto x = modules::EluModule().build(ctx, input);

auto conv1 = build_stateful_conv1d(

ctx,

x,

conv1_history,

weights.conv1,

channels,

hidden_channels,

3,

1,

1,

weights.conv1.bias.has_value());

if (conv1.next_history.has_value()) {

next_histories.push_back(*conv1.next_history);

}

x = modules::EluModule().build(ctx, conv1.output);

auto conv2 = build_stateful_conv1d(

ctx,

x,

conv2_history,

weights.conv2,

hidden_channels,

channels,

1,

1,

1,

weights.conv2.bias.has_value());

if (conv2.next_history.has_value()) {

next_histories.push_back(*conv2.next_history);

}

return modules::ResidualAddModule().build(ctx, input, conv2.output);

core::ModuleBuildContext & ctx,

const core::TensorValue & input,

const SEANetResidualWeights & weights,

int64_t channels,

int64_t hidden_channels,

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

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

std::vector<core::TensorValue> & next_histories) {

auto x = modules::EluModule().build(ctx, input);

auto conv1 = build_stateful_conv1d(

ctx,

x,

conv1_history,

weights.conv1,

channels,

hidden_channels,

3,

1,

1,

weights.conv1.bias.has_value());

if (conv1.next_history.has_value()) {

next_histories.push_back(*conv1.next_history);

}

x = modules::EluModule().build(ctx, conv1.output);

auto conv2 = build_stateful_conv1d(

ctx,

x,

conv2_history,

weights.conv2,

hidden_channels,

channels,

1,

1,

1,

weights.conv2.bias.has_value());

if (conv2.next_history.has_value()) {

next_histories.push_back(*conv2.next_history);

}

return modules::ResidualAddModule().build(ctx, input, conv2.output);

const core::TensorValue & tensor,

const StreamingConv1dState & state,

const std::vector<float> & current,

int64_t channels,

int64_t frames,

modules::StreamingPadMode pad_mode) {

if (state.history_frames <= 0) {

return;

}

std::vector<float> history(static_cast<size_t>(channels * state.history_frames), 0.0F);

if (pad_mode == modules::StreamingPadMode::Replicate && state.first) {

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

std::fill_n(

history.data() + static_cast<size_t>(channel * state.history_frames),

static_cast<size_t>(state.history_frames),

current[static_cast<size_t>(channel * frames)]);

}

} else {

history = state.previous;

}

core::write_tensor_f32(tensor, history);

const core::TensorValue & tensor,

StreamingConv1dState & state) {

if (state.history_frames <= 0) {

state.first = false;

return;

}

state.previous = core::read_tensor_f32(tensor.tensor);

state.first = false;

const core::TensorValue & tensor,

const StreamingConvTranspose1dState & state) {

if (state.partial_frames > 0) {

core::write_tensor_f32(tensor, state.partial);

}

const core::TensorValue & tensor,

StreamingConvTranspose1dState & state,

const std::vector<float> & bias_values) {

if (state.partial_frames <= 0) {

return;

}

state.partial = core::read_tensor_f32(tensor.tensor);

if (!bias_values.empty()) {

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

const float bias = bias_values[static_cast<size_t>(channel)];

for (int64_t frame = 0; frame < state.partial_frames; ++frame) {

state.partial[static_cast<size_t>(channel * state.partial_frames + frame)] -= bias;

}

}

}

MimiEncoderPreTransformerGraph(

ggml_backend_t backend,

core::BackendType backend_type,

int threads,

size_t graph_context_bytes,

std::shared_ptr<const MimiCodecWeights> weights,

MimiCodecConfig config)

: backend_(backend),

threads_(threads),

weights_(std::move(weights)),

config_(std::move(config)) {

if (weights_ == nullptr) {

throw std::runtime_error("Mimi codec encoder pre-transformer graph requires weights");

}

ggml_ctx_ = ggml_init({graph_context_bytes, nullptr, true});

if (ggml_ctx_ == nullptr) {

throw std::runtime_error("failed to initialize Mimi codec encoder pre-transformer graph context");

}

core::ModuleBuildContext ctx{ggml_ctx_, "mimi_codec.encoder_pre_transformer", backend_type};

input_ = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, config_.channels, kMimiFrameSamples}));

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

auto input_history = make_history_tensor(ctx, config_.channels, 6);

auto input_projection = build_stateful_conv1d(

ctx,

input_,

input_history,

weights_->encoder.input_projection,

config_.channels,

64,

7,

1,

1,

weights_->encoder.input_projection.bias.has_value());

next_histories_.push_back(*input_projection.next_history);

auto x = input_projection.output;

const int64_t channels[] = {64, 128, 256, 512};

const int64_t hidden_channels[] = {32, 64, 128, 256};

const int64_t downsample_kernels[] = {8, 10, 12, 16};

const int downsample_strides[] = {4, 5, 6, 8};

for (size_t stage = 0; stage < 4; ++stage) {

x = build_mimi_residual_block(

ctx,

x,

weights_->encoder.residual_blocks[stage],

channels[stage],

hidden_channels[stage],

make_history_tensor(ctx, channels[stage], 2),

make_history_tensor(ctx, hidden_channels[stage], 0),

residual_histories);

x = modules::EluModule().build(ctx, x);

auto downsample = build_stateful_conv1d(

ctx,

x,

make_history_tensor(ctx, channels[stage], downsample_kernels[stage] - downsample_strides[stage]),

weights_->encoder.downsample_layers[stage],

channels[stage],

channels[stage] * 2,

downsample_kernels[stage],

downsample_strides[stage],

1,

weights_->encoder.downsample_layers[stage].bias.has_value());

x = downsample.output;

next_histories_.push_back(*downsample.next_history);

}

next_histories_.insert(next_histories_.end(), residual_histories.begin(), residual_histories.end());

x = modules::EluModule().build(ctx, x);

auto output_projection = build_stateful_conv1d(

ctx,

x,

make_history_tensor(ctx, 1024, 2),

weights_->encoder.output_projection,

1024,

config_.hidden_size,

3,

1,

1,

weights_->encoder.output_projection.bias.has_value());

next_histories_.push_back(*output_projection.next_history);

output_ = output_projection.output;

build_runtime(ctx);

}

~MimiEncoderPreTransformerGraph() {

release_graph_runtime(galloc_, params_buffer_, ggml_ctx_);

}

std::vector<float> run(const std::vector<float> & input, MimiEncoderState & state) const {

if (static_cast<int64_t>(input.size()) != config_.channels * kMimiFrameSamples) {

throw std::runtime_error("Mimi codec encoder pre-transformer input size mismatch");

}

core::write_tensor_f32(input_, input);

size_t history = 0;

write_streaming_history(history_tensors_[history++], state.input_projection, input, config_.channels, kMimiFrameSamples, modules::StreamingPadMode::Constant);

for (size_t stage = 0; stage < 4; ++stage) {

write_streaming_history(history_tensors_[history++], state.stage_residual_convs[stage][0], {}, 0, 0, modules::StreamingPadMode::Constant);

write_streaming_history(history_tensors_[history++], state.stage_downsamples[stage], {}, 0, 0, modules::StreamingPadMode::Constant);

}

write_streaming_history(history_tensors_[history++], state.output_projection, {}, 0, 0, modules::StreamingPadMode::Constant);

if (core::compute_backend_graph(backend_, graph_) != GGML_STATUS_SUCCESS) {

throw std::runtime_error("Mimi codec encoder pre-transformer graph compute failed");

}

size_t output_history = 0;

read_streaming_history(next_histories_[output_history++], state.input_projection);

for (size_t stage = 0; stage < 4; ++stage) {

read_streaming_history(next_histories_[output_history++], state.stage_downsamples[stage]);

}

for (size_t stage = 0; stage < 4; ++stage) {

read_streaming_history(next_histories_[output_history++], state.stage_residual_convs[stage][0]);

}

read_streaming_history(next_histories_[output_history++], state.output_projection);

return core::read_tensor_f32(output_.tensor);

}

std::optional<core::TensorValue> make_history_tensor(

core::ModuleBuildContext & ctx,

int64_t channels,

int64_t frames) {

if (frames <= 0) {

return std::nullopt;

}

history_tensors_.push_back(core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, channels, frames})));

return history_tensors_.back();

}

void build_runtime(core::ModuleBuildContext &) {

if (core::is_host_backend(backend_)) {

params_buffer_ = ggml_backend_alloc_ctx_tensors(ggml_ctx_, backend_);

if (params_buffer_ == nullptr) {

release_graph_runtime(galloc_, params_buffer_, ggml_ctx_);

throw std::runtime_error("Mimi codec encoder pre-transformer context tensor allocation failed");

}

}

core::set_backend_threads(backend_, threads_);

graph_ = ggml_new_graph_custom(ggml_ctx_, 65536, false);

ggml_build_forward_expand(graph_, output_.tensor);

for (const auto & history : next_histories_) {

ggml_build_forward_expand(graph_, history.tensor);

}

galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));

if (galloc_ == nullptr || !ggml_gallocr_reserve(galloc_, graph_) || !ggml_gallocr_alloc_graph(galloc_, graph_)) {

release_graph_runtime(galloc_, params_buffer_, ggml_ctx_);

throw std::runtime_error("Mimi codec encoder pre-transformer graph allocation failed");

}

}

ggml_backend_t backend_ = nullptr;

int threads_ = 1;

std::shared_ptr<const MimiCodecWeights> weights_;

MimiCodecConfig config_;

ggml_context * ggml_ctx_ = nullptr;

ggml_backend_buffer_t params_buffer_ = nullptr;

ggml_cgraph * graph_ = nullptr;

ggml_gallocr_t galloc_ = nullptr;

core::TensorValue input_;

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

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

core::TensorValue output_;

MimiEncoderPostTransformerGraph(

ggml_backend_t backend,

core::BackendType backend_type,

int threads,

size_t graph_context_bytes,

std::shared_ptr<const MimiCodecWeights> weights,

MimiCodecConfig config,

int64_t frames)

: backend_(backend),

threads_(threads),

weights_(std::move(weights)),

config_(std::move(config)),

frames_(frames) {

if (weights_ == nullptr || frames_ <= 0) {

throw std::runtime_error("Mimi codec encoder post-transformer graph requires weights and frames");

}

ggml_ctx_ = ggml_init({graph_context_bytes, nullptr, true});

if (ggml_ctx_ == nullptr) {

throw std::runtime_error("failed to initialize Mimi codec encoder post-transformer graph context");

}

core::ModuleBuildContext ctx{ggml_ctx_, "mimi_codec.encoder_post_transformer", backend_type};

input_ = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, config_.hidden_size, frames_}));

downsample_history_ = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, config_.hidden_size, 2}));

auto downsample = build_stateful_conv1d(

ctx,

input_,

downsample_history_,

weights_->encoder.downsample_conv,

config_.hidden_size,

config_.hidden_size,

4,

2,

1,

weights_->encoder.downsample_conv.bias.has_value());

next_downsample_history_ = *downsample.next_history;

if (downsample.output.shape.dims[2] != 1) {

throw std::runtime_error("Mimi codec encoder post-transformer expected one latent frame");

}

auto latent = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, downsample.output);

latent = modules::ReshapeModule({core::TensorShape::from_dims({1, config_.hidden_size})}).build(ctx, latent);

const auto semantic_weight = core::reshape_tensor(

ctx,

weights_->quantizer.semantic_input_proj.weight,

core::TensorShape::from_dims({config_.latent_size, config_.hidden_size}));

semantic_ = modules::LinearModule({

config_.hidden_size,

config_.latent_size,

weights_->quantizer.semantic_input_proj.bias.has_value(),

GGML_PREC_F32,

}).build(ctx, latent, modules::LinearWeights{semantic_weight, weights_->quantizer.semantic_input_proj.bias});

const auto acoustic_weight = core::reshape_tensor(

ctx,

weights_->quantizer.acoustic_input_proj.weight,

core::TensorShape::from_dims({config_.latent_size, config_.hidden_size}));

acoustic_ = modules::LinearModule({

config_.hidden_size,

config_.latent_size,

weights_->quantizer.acoustic_input_proj.bias.has_value(),

GGML_PREC_F32,

}).build(ctx, latent, modules::LinearWeights{acoustic_weight, weights_->quantizer.acoustic_input_proj.bias});

build_runtime();

}

~MimiEncoderPostTransformerGraph() {

release_graph_runtime(galloc_, params_buffer_, ggml_ctx_);

}

std::pair<std::vector<float>, std::vector<float>> run(

const std::vector<float> & input,

MimiEncoderState & state) const {

if (static_cast<int64_t>(input.size()) != config_.hidden_size * frames_) {

throw std::runtime_error("Mimi codec encoder post-transformer input size mismatch");

}

core::write_tensor_f32(input_, input);

write_streaming_history(

downsample_history_,

state.downsample,

input,

config_.hidden_size,

frames_,

modules::StreamingPadMode::Replicate);

if (core::compute_backend_graph(backend_, graph_) != GGML_STATUS_SUCCESS) {

throw std::runtime_error("Mimi codec encoder post-transformer graph compute failed");

}

read_streaming_history(next_downsample_history_, state.downsample);

return {core::read_tensor_f32(semantic_.tensor), core::read_tensor_f32(acoustic_.tensor)};

}

void build_runtime() {

if (core::is_host_backend(backend_)) {

params_buffer_ = ggml_backend_alloc_ctx_tensors(ggml_ctx_, backend_);

if (params_buffer_ == nullptr) {

release_graph_runtime(galloc_, params_buffer_, ggml_ctx_);

throw std::runtime_error("Mimi codec encoder post-transformer context tensor allocation failed");

}

}

core::set_backend_threads(backend_, threads_);

graph_ = ggml_new_graph_custom(ggml_ctx_, 32768, false);

ggml_build_forward_expand(graph_, semantic_.tensor);

ggml_build_forward_expand(graph_, acoustic_.tensor);

ggml_build_forward_expand(graph_, next_downsample_history_.tensor);

galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));

if (galloc_ == nullptr || !ggml_gallocr_reserve(galloc_, graph_) || !ggml_gallocr_alloc_graph(galloc_, graph_)) {

release_graph_runtime(galloc_, params_buffer_, ggml_ctx_);

throw std::runtime_error("Mimi codec encoder post-transformer graph allocation failed");

}

}

ggml_backend_t backend_ = nullptr;

int threads_ = 1;

std::shared_ptr<const MimiCodecWeights> weights_;

MimiCodecConfig config_;

int64_t frames_ = 0;

ggml_context * ggml_ctx_ = nullptr;

ggml_backend_buffer_t params_buffer_ = nullptr;

ggml_cgraph * graph_ = nullptr;

ggml_gallocr_t galloc_ = nullptr;

core::TensorValue input_;

core::TensorValue downsample_history_;

core::TensorValue next_downsample_history_;

core::TensorValue semantic_;

core::TensorValue acoustic_;

MimiDecoderPreTransformerGraph(

ggml_backend_t backend,

core::BackendType backend_type,

int threads,

size_t graph_context_bytes,

std::shared_ptr<const MimiCodecWeights> weights,

MimiCodecConfig config)

: backend_(backend),

threads_(threads),

weights_(std::move(weights)),

config_(std::move(config)) {

if (weights_ == nullptr) {

throw std::runtime_error("Mimi codec decoder pre-transformer graph requires weights");

}

ggml_ctx_ = ggml_init({graph_context_bytes, nullptr, true});

if (ggml_ctx_ == nullptr) {

throw std::runtime_error("failed to initialize Mimi codec decoder pre-transformer graph context");

}

core::ModuleBuildContext ctx{ggml_ctx_, "mimi_codec.decoder_pre_transformer", backend_type};

input_ = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, config_.hidden_size, 1}));

partial_ = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, config_.hidden_size, 2}));

auto upsample = build_stateful_depthwise_convtranspose1d(

ctx,

input_,

partial_,

weights_->decoder.upsample_conv,

config_.hidden_size,

4,

2);

output_ = upsample.output;

next_partial_ = *upsample.next_partial;

build_runtime();

}

~MimiDecoderPreTransformerGraph() {

release_graph_runtime(galloc_, params_buffer_, ggml_ctx_);

}

std::vector<float> run(const std::vector<float> & input, MimiDecoderState & state) const {

if (static_cast<int64_t>(input.size()) != config_.hidden_size) {

throw std::runtime_error("Mimi codec decoder pre-transformer input size mismatch");

}

core::write_tensor_f32(input_, input);

write_streaming_partial(partial_, state.encoder_rate_upsample);

if (core::compute_backend_graph(backend_, graph_) != GGML_STATUS_SUCCESS) {

throw std::runtime_error("Mimi codec decoder pre-transformer graph compute failed");

}

read_streaming_partial(next_partial_, state.encoder_rate_upsample, {});

return core::read_tensor_f32(output_.tensor);

}

void build_runtime() {

if (core::is_host_backend(backend_)) {

params_buffer_ = ggml_backend_alloc_ctx_tensors(ggml_ctx_, backend_);

if (params_buffer_ == nullptr) {

release_graph_runtime(galloc_, params_buffer_, ggml_ctx_);

throw std::runtime_error("Mimi codec decoder pre-transformer context tensor allocation failed");

}

}

core::set_backend_threads(backend_, threads_);

graph_ = ggml_new_graph_custom(ggml_ctx_, 32768, false);

ggml_build_forward_expand(graph_, output_.tensor);

ggml_build_forward_expand(graph_, next_partial_.tensor);

galloc_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));

if (galloc_ == nullptr || !ggml_gallocr_reserve(galloc_, graph_) || !ggml_gallocr_alloc_graph(galloc_, graph_)) {

release_graph_runtime(galloc_, params_buffer_, ggml_ctx_);

throw std::runtime_error("Mimi codec decoder pre-transformer graph allocation failed");

}

}

ggml_backend_t backend_ = nullptr;

int threads_ = 1;

std::shared_ptr<const MimiCodecWeights> weights_;

MimiCodecConfig config_;

ggml_context * ggml_ctx_ = nullptr;

ggml_backend_buffer_t params_buffer_ = nullptr;

ggml_cgraph * graph_ = nullptr;

ggml_gallocr_t galloc_ = nullptr;

core::TensorValue input_;

core::TensorValue partial_;

core::TensorValue next_partial_;

core::TensorValue output_;

MimiDecoderPostTransformerGraph(

ggml_backend_t backend,

core::BackendType backend_type,

int threads,

size_t graph_context_bytes,

std::shared_ptr<const MimiCodecWeights> weights,

MimiCodecConfig config,

int64_t frames)

: backend_(backend),

threads_(threads),

weights_(std::move(weights)),

config_(std::move(config)),

frames_(frames) {

if (weights_ == nullptr || frames_ <= 0) {

throw std::runtime_error("Mimi codec decoder post-transformer graph requires weights and frames");

}

ggml_ctx_ = ggml_init({graph_context_bytes, nullptr, true});

if (ggml_ctx_ == nullptr) {

throw std::runtime_error("failed to initialize Mimi codec decoder post-transformer graph context");

}

core::ModuleBuildContext ctx{ggml_ctx_, "mimi_codec.decoder_post_transformer", backend_type};

input_ = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, config_.hidden_size, frames_}));

auto input_projection = build_stateful_conv1d(

ctx,

input_,

make_history_tensor(ctx, config_.hidden_size, 6),

weights_->decoder.seanet.input_projection,

config_.hidden_size,

1024,

7,

1,

1,

weights_->decoder.seanet.input_projection.bias.has_value());

next_histories_.push_back(*input_projection.next_history);

auto x = input_projection.output;

const int64_t stage_in_channels[] = {1024, 512, 256, 128};

const int64_t stage_out_channels[] = {512, 256, 128, 64};

const int64_t stage_hidden_channels[] = {256, 128, 64, 32};

const int64_t stage_kernel_sizes[] = {16, 12, 10, 8};

const int stage_strides[] = {8, 6, 5, 4};

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

for (size_t stage = 0; stage < 4; ++stage) {

x = modules::EluModule().build(ctx, x);

auto upsample = build_stateful_convtranspose1d(

ctx,

x,

make_partial_tensor(ctx, stage_out_channels[stage], stage_kernel_sizes[stage] - stage_strides[stage]),

weights_->decoder.seanet.stages[stage].upsample,

stage_in_channels[stage],

stage_out_channels[stage],

stage_kernel_sizes[stage],

stage_strides[stage],

weights_->decoder.seanet.stages[stage].upsample.bias.has_value());

next_partials_.push_back(*upsample.next_partial);

x = build_seanet_residual_block(

ctx,

upsample.output,

weights_->decoder.seanet.stages[stage].residual_blocks.front(),

stage_out_channels[stage],

stage_hidden_channels[stage],

make_history_tensor(ctx, stage_out_channels[stage], 2),

make_history_tensor(ctx, stage_hidden_channels[stage], 0),

residual_histories);

}

next_histories_.insert(next_histories_.end(), residual_histories.begin(), residual_histories.end());

x = modules::EluModule().build(ctx, x);

auto output_projection = build_stateful_conv1d(

ctx,

x,

make_history_tensor(ctx, 64, 2),

weights_->decoder.seanet.output_projection,

64,

config_.channels,

3,

1,

1,

weights_->decoder.seanet.output_projection.bias.has_value());

next_histories_.push_back(*output_projection.next_history);

output_ = output_projection.output;

build_runtime();

}

~MimiDecoderPostTransformerGraph() {

release_graph_runtime(galloc_, params_buffer_, ggml_ctx_);

}

std::vector<float> run(const std::vector<float> & input, MimiDecoderState & state) const {

if (static_cast<int64_t>(input.size()) != config_.hidden_size * frames_) {

throw std::runtime_error("Mimi codec decoder post-transformer input size mismatch");

}

core::write_tensor_f32(input_, input);

size_t history = 0;

write_streaming_history(history_tensors_[history++], state.input_projection, input, config_.hidden_size, frames_, modules::StreamingPadMode::Constant);

for (size_t stage = 0; stage < 4; ++stage) {

write_streaming_partial(partial_tensors_[stage], state.stage_upsamples[stage]);

write_streaming_history(history_tensors_[history++], state.stage_residual_convs[stage][0], {}, 0, 0, modules::StreamingPadMode::Constant);

}

write_streaming_history(history_tensors_[history++], state.output_projection, {}, 0, 0, modules::StreamingPadMode::Constant);

if (core::compute_backend_graph(backend_, graph_) != GGML_STATUS_SUCCESS) {

throw std::runtime_error("Mimi codec decoder post-transformer graph compute failed");

}

size_t output_history = 0;

read_streaming_history(next_histories_[output_history++], state.input_projection);

for (size_t stage = 0; stage < 4; ++stage) {

read_streaming_history(next_histories_[output_history++], state.stage_residual_convs[stage][0]);

}

read_streaming_history(next_histories_[output_history++], state.output_projection);

for (size_t stage = 0; stage < 4; ++stage) {

read_streaming_partial(next_partials_[stage], state.stage_upsamples[stage], weights_->decoder.seanet_upsample_bias_values[stage]);

}

auto output = core::read_tensor_f32(output_.tensor);

for (float & sample : output) {

sample = std::clamp(sample, -1.0F, 1.0F);

}

return output;

}

std::optional<core::TensorValue> make_history_tensor(

core::ModuleBuildContext & ctx,

int64_t channels,

int64_t frames) {

if (frames <= 0) {

return std::nullopt;

}

history_tensors_.push_back(core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, channels, frames})));

return history_tensors_.back();

}

std::optional<core::TensorValue> make_partial_tensor(

core::ModuleBuildContext & ctx,

int64_t channels,

int64_t frames) {

if (frames <= 0) {

return std::nullopt;

}

partial_tensors_.push_back(core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, channels, frames})));

return partial_tensors_.back();

}

void build_runtime() {

if (core::is_host_backend(backend_)) {

params_buffer_ = ggml_backend_alloc_ctx_tensors(ggml_ctx_, backend_);

if (params_buffer_ == nullptr) {

release_graph_runtime(galloc_, params_buffer_, ggml_ctx_);

throw std::runtime_error("Mimi codec decoder post-transformer context tensor allocation failed");

}

}

core::set_backend_threads(backend_, threads_);