Function load_s3_flow_encoder_weights

src/models/chatterbox/s3gen_flow.cpp:1731–1783

    state_->decoder_backend = {};
}

std::shared_ptr<const S3FlowEncoderWeights> load_s3_flow_encoder_weights(
    const engine::assets::TensorSource & source,
    const engine::core::ExecutionContext & execution_context,
    engine::assets::TensorStorageType weight_storage_type) {
    auto weights = std::make_shared<S3FlowEncoderWeights>();
    weights->execution_context = &execution_context;
    weights->store = std::make_shared<engine::core::BackendWeightStore>(
        execution_context.backend(),
        execution_context.backend_type(),
        "chatterbox.s3_flow_encoder.weights",
        1024ull * 1024ull * 1024ull);
    weights->input_embedding_tensor = weights->store->load_tensor(
        source,
        "flow.input_embedding.weight",
        weight_storage_type,
        {6561, 512});
    weights->speaker_affine = load_flow_linear(*weights->store, source, "flow.spk_embed_affine_layer", 80, 192, true, weight_storage_type);
    weights->encoder_proj = load_flow_linear(*weights->store, source, "flow.encoder_proj", 80, 512, true, weight_storage_type);
    weights->embed_linear = load_flow_linear(*weights->store, source, "flow.encoder.embed.out.0", 512, 512, true, weight_storage_type);
    weights->embed_norm = load_flow_layer_norm(*weights->store, source, "flow.encoder.embed.out.1", 512);
    weights->prelook_conv1 =
        load_flow_conv1d(*weights->store, source, "flow.encoder.pre_lookahead_layer.conv1", 512, 512, 4, 1, 0, weight_storage_type);
    weights->prelook_conv2 =
        load_flow_conv1d(*weights->store, source, "flow.encoder.pre_lookahead_layer.conv2", 512, 512, 3, 1, 0, weight_storage_type);
    auto load_layer = [&](const std::string & prefix) {
        S3FlowEncoderWeights::EncoderLayerWeights layer;
        layer.norm_mha = load_flow_layer_norm(*weights->store, source, prefix + ".norm_mha", 512);
        layer.attn.q = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_q", 512, 512, true, weight_storage_type);
        layer.attn.k = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_k", 512, 512, true, weight_storage_type);
        layer.attn.v = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_v", 512, 512, true, weight_storage_type);
        layer.attn.out = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_out", 512, 512, true, weight_storage_type);
        layer.attn.pos = load_flow_linear(*weights->store, source, prefix + ".self_attn.linear_pos", 512, 512, false, weight_storage_type);
        layer.attn.pos_bias_u_tensor = weights->store->load_f32_tensor(source, prefix + ".self_attn.pos_bias_u", {8, 64});
        layer.attn.pos_bias_v_tensor = weights->store->load_f32_tensor(source, prefix + ".self_attn.pos_bias_v", {8, 64});
        layer.norm_ff = load_flow_layer_norm(*weights->store, source, prefix + ".norm_ff", 512);
        layer.ff.w1 = load_flow_linear(*weights->store, source, prefix + ".feed_forward.w_1", 2048, 512, true, weight_storage_type);
        layer.ff.w2 = load_flow_linear(*weights->store, source, prefix + ".feed_forward.w_2", 512, 2048, true, weight_storage_type);
        return layer;
    };
    for (int i = 0; i < 6; ++i) {
        weights->encoders.push_back(load_layer("flow.encoder.encoders." + std::to_string(i)));
    }
    weights->up_layer_conv = load_flow_conv1d(*weights->store, source, "flow.encoder.up_layer.conv", 512, 512, 5, 1, 0, weight_storage_type);
    weights->up_embed_linear = load_flow_linear(*weights->store, source, "flow.encoder.up_embed.out.0", 512, 512, true, weight_storage_type);
    weights->up_embed_norm = load_flow_layer_norm(*weights->store, source, "flow.encoder.up_embed.out.1", 512);
    for (int i = 0; i < 4; ++i) {
        weights->up_encoders.push_back(load_layer("flow.encoder.up_encoders." + std::to_string(i)));
    }
    weights->after_norm = load_flow_layer_norm(*weights->store, source, "flow.encoder.after_norm", 512);
    weights->store->upload();
    return weights;
}

S3FlowEncoderOutputs compute_s3_flow_encoder_forward(
    S3FlowSessionCache & cache,