| 35 | } // namespace |
| 36 | |
| 37 | VAEEncoderWeights load_vae_encoder_weights( |
| 38 | const AceStepAssets & assets, |
| 39 | ggml_backend_t backend, |
| 40 | core::BackendType backend_type, |
| 41 | size_t weight_context_bytes, |
| 42 | assets::TensorStorageType storage_type) { |
| 43 | const auto & config = assets.config.vae; |
| 44 | if (config.encoder_hidden_size <= 0 || |
| 45 | config.decoder_channels <= 0 || |
| 46 | config.channel_multiples.empty() || config.downsampling_ratios.empty()) { |
| 47 | throw std::runtime_error("ACE-Step VAE config is incomplete for encoder loading"); |
| 48 | } |
| 49 | const auto & source = *assets.vae_weights; |
| 50 | VAEEncoderWeights weights = {}; |
| 51 | weights.store = std::make_shared<core::BackendWeightStore>( |
| 52 | backend, |
| 53 | backend_type, |
| 54 | "ace_step.vae.encoder.weights", |
| 55 | weight_context_bytes); |
| 56 | |
| 57 | const int64_t encoder_hidden = config.decoder_channels; |
| 58 | const int64_t encoder_output_channels = config.encoder_hidden_size; |
| 59 | |
| 60 | weights.encoder_conv1 = load_weight_norm_conv1d( |
| 61 | *weights.store, |
| 62 | source, |
| 63 | "encoder.conv1", |
| 64 | encoder_hidden, |
| 65 | config.audio_channels, |
| 66 | 7, |
| 67 | 1, |
| 68 | 3, |
| 69 | 1, |
| 70 | storage_type, |
| 71 | true); |
| 72 | |
| 73 | weights.encoder_blocks.reserve(config.downsampling_ratios.size()); |
| 74 | for (size_t i = 0; i < config.downsampling_ratios.size(); ++i) { |
| 75 | const int64_t stride = config.downsampling_ratios[i]; |
| 76 | const int64_t in_channels = encoder_hidden * (i == 0 ? 1 : config.channel_multiples[i - 1]); |
| 77 | const int64_t out_channels = encoder_hidden * config.channel_multiples[i]; |
| 78 | const std::string prefix = "encoder.block." + std::to_string(i); |
| 79 | EncoderBlockWeights block = {}; |
| 80 | block.res1 = load_residual_unit(*weights.store, source, prefix + ".res_unit1", in_channels, storage_type); |
| 81 | block.res2 = load_residual_unit(*weights.store, source, prefix + ".res_unit2", in_channels, storage_type); |
| 82 | block.res3 = load_residual_unit(*weights.store, source, prefix + ".res_unit3", in_channels, storage_type); |
| 83 | block.snake = load_snake_exact(*weights.store, source, prefix + ".snake1", in_channels); |
| 84 | block.conv = load_weight_norm_conv1d( |
| 85 | *weights.store, |
| 86 | source, |
| 87 | prefix + ".conv1", |
| 88 | out_channels, |
| 89 | in_channels, |
| 90 | 2 * stride, |
| 91 | static_cast<int>(stride), |
| 92 | static_cast<int>((stride + 1) / 2), |
| 93 | 1, |
| 94 | storage_type, |
no test coverage detected