Method AceStepDitWeightsRuntime

src/models/ace_step/dit_weights_runtime.cpp:293–315

}  // namespace

AceStepDitWeightsRuntime::AceStepDitWeightsRuntime(
    std::shared_ptr<const AceStepAssets> assets,
    core::ExecutionContext & execution,
    assets::TensorStorageType storage_type,
    size_t weight_context_bytes)
    : assets_(std::move(assets)) {
    if (!assets_) {
        throw std::runtime_error("ACE-Step DiT weights runtime requires assets");
    }
    if (execution.backend() == nullptr) {
        throw std::runtime_error("ACE-Step DiT weights runtime backend is not initialized");
    }
    store_ = std::make_shared<core::BackendWeightStore>(
        execution.backend(),
        execution.backend_type(),
        "ace_step.dit.weights",
        weight_context_bytes);
    detokenizer_weights_ = load_detokenizer_weights(store_, *assets_, storage_type);
    condition_encoder_weights_ = load_condition_encoder_weights(store_, *assets_, storage_type);
    diffusion_weights_ = load_diffusion_weights(store_, *assets_, storage_type);
    store_->upload();
    assets_->dit_weights->release_storage();
}

const std::shared_ptr<const AceStepAssets> & AceStepDitWeightsRuntime::assets() const noexcept {
    return assets_;