Function graph_run_downsampled_encoder

src/framework/audio/zipenhancer.cpp:795–825

}

core::TensorValue graph_run_downsampled_encoder(
    core::ModuleBuildContext & ctx,
    std::vector<GraphConstant> & constants,
    const core::TensorValue & input,
    const std::unordered_map<std::string, Param> & params,
    const std::string & root,
    int64_t ds) {
    const int64_t frames = input.shape.dims[2];
    const int64_t freq = input.shape.dims[3];
    auto t_seq = graph_transpose(ctx, input, {{2, 3, 0, 1}}, 4);
    t_seq = graph_reshape(ctx, t_seq, core::TensorShape::from_dims({frames, freq, kDense}));
    auto td = graph_downsample(ctx, t_seq, require_param(params, root + ".downsample_t.bias"));
    const int64_t d_frames = td.shape.dims[0];
    auto f_seq = graph_transpose(ctx, td, {{1, 0, 2, 3}}, 3);
    auto fd = graph_downsample(ctx, f_seq, require_param(params, root + ".downsample_f.bias"));
    const int64_t d_freq = fd.shape.dims[0];
    auto x = graph_transpose(ctx, fd, {{2, 1, 0, 3}}, 3);
    x = graph_reshape(ctx, x, core::TensorShape::from_dims({1, kDense, d_frames, d_freq}));
    x = graph_run_dual_encoder(ctx, constants, x, params, root + ".encoder");

    auto f_up = graph_transpose(ctx, x, {{3, 2, 0, 1}}, 4);
    f_up = graph_reshape(ctx, f_up, core::TensorShape::from_dims({d_freq, d_frames, kDense}));
    f_up = graph_upsample_slice(ctx, f_up, ds, freq);
    auto t_up = graph_transpose(ctx, f_up, {{1, 0, 2, 3}}, 3);
    t_up = graph_upsample_slice(ctx, t_up, ds, frames);
    auto orig = graph_transpose(ctx, input, {{2, 3, 0, 1}}, 4);
    orig = graph_reshape(ctx, orig, core::TensorShape::from_dims({frames, freq, kDense}));
    auto combined = graph_bypass(ctx, orig, t_up, require_param(params, root + ".out_combiner.bypass_scale"));
    auto out = graph_transpose(ctx, combined, {{2, 0, 1, 3}}, 3);
    return graph_reshape(ctx, out, core::TensorShape::from_dims({1, kDense, frames, freq}));
}

core::TensorValue graph_tsconformer(
    core::ModuleBuildContext & ctx,