| 320 | } |
| 321 | |
| 322 | VAEDecoderWeights load_vae_decoder_weights( |
| 323 | const AceStepAssets & assets, |
| 324 | ggml_backend_t backend, |
| 325 | core::BackendType backend_type, |
| 326 | size_t weight_context_bytes, |
| 327 | assets::TensorStorageType storage_type) { |
| 328 | const auto & config = assets.config.vae; |
| 329 | if (config.decoder_channels <= 0 || |
| 330 | config.decoder_input_channels <= 0 || |
| 331 | config.channel_multiples.empty() || config.downsampling_ratios.empty()) { |
| 332 | throw std::runtime_error("ACE-Step VAE config is incomplete for decoder loading"); |
| 333 | } |
| 334 | const auto & source = *assets.vae_weights; |
| 335 | VAEDecoderWeights weights = {}; |
| 336 | weights.store = std::make_shared<core::BackendWeightStore>( |
| 337 | backend, |
| 338 | backend_type, |
| 339 | "ace_step.vae.decoder.weights", |
| 340 | weight_context_bytes); |
| 341 | |
| 342 | const int64_t deepest_channels = config.decoder_channels * config.channel_multiples.back(); |
| 343 | weights.conv1 = load_weight_norm_conv1d( |
| 344 | *weights.store, |
| 345 | source, |
| 346 | "decoder.conv1", |
| 347 | deepest_channels, |
| 348 | config.decoder_input_channels, |
| 349 | 7, |
| 350 | 1, |
| 351 | 3, |
| 352 | 1, |
| 353 | storage_type, |
| 354 | true); |
| 355 | |
| 356 | const size_t num_blocks = config.downsampling_ratios.size(); |
| 357 | weights.blocks.reserve(num_blocks); |
| 358 | for (size_t i = 0; i < num_blocks; ++i) { |
| 359 | const int64_t stride = config.downsampling_ratios[num_blocks - 1 - i]; |
| 360 | const int64_t in_channels = config.decoder_channels * config.channel_multiples[num_blocks - i - 1]; |
| 361 | const int64_t out_channels = config.decoder_channels * (i + 1 < num_blocks ? config.channel_multiples[num_blocks - i - 2] : 1); |
| 362 | const std::string prefix = "decoder.block." + std::to_string(i); |
| 363 | DecoderBlockWeights block = {}; |
| 364 | block.snake = load_snake_exact(*weights.store, source, prefix + ".snake1", in_channels); |
| 365 | block.conv_t = load_weight_norm_conv_transpose1d( |
| 366 | *weights.store, |
| 367 | source, |
| 368 | prefix + ".conv_t1", |
| 369 | in_channels, |
| 370 | out_channels, |
| 371 | 2 * stride, |
| 372 | static_cast<int>(stride), |
| 373 | static_cast<int>((stride + 1) / 2), |
| 374 | 1, |
| 375 | storage_type, |
| 376 | true); |
| 377 | block.res1 = load_residual_unit(*weights.store, source, prefix + ".res_unit1", out_channels, storage_type); |
| 378 | block.res2 = load_residual_unit(*weights.store, source, prefix + ".res_unit2", out_channels, storage_type); |
| 379 | block.res3 = load_residual_unit(*weights.store, source, prefix + ".res_unit3", out_channels, storage_type); |
no test coverage detected