#include #include "Encoder.h" #include "StableABICompat.h" extern "C" { #include #include #include } namespace facebook::torchcodec { namespace { torch::stable::Tensor validate_samples(const torch::stable::Tensor& samples) { STD_TORCH_CHECK( samples.scalar_type() == kStableFloat32, "samples must have float32 dtype, got ", (samples.scalar_type())); STD_TORCH_CHECK( samples.device().type() == kStableCPU, "samples must be on CPU, got ", device_type_name(samples.device().type())); STD_TORCH_CHECK( samples.dim() == 2, "samples must have 2 dimensions, got ", samples.dim()); // We enforce this, but if we get user reports we should investigate whether // that's actually needed. int num_channels = static_cast(samples.sizes()[0]); STD_TORCH_CHECK( num_channels <= AV_NUM_DATA_POINTERS, "Trying to encode ", num_channels, " channels, but FFmpeg only supports ", AV_NUM_DATA_POINTERS, " channels per frame."); return torch::stable::contiguous(samples); } void validate_sample_rate(const AVCodec& av_codec, int sample_rate) { const int* supported_sample_rates = get_supported_sample_rates(av_codec); if (supported_sample_rates == nullptr) { return; } for (auto i = 0; supported_sample_rates[i] != 0; ++i) { if (sample_rate == supported_sample_rates[i]) { return; } } std::stringstream supported_rates; for (auto i = 0; supported_sample_rates[i] != 0; ++i) { if (i > 0) { supported_rates << ", "; } supported_rates << supported_sample_rates[i]; } STD_TORCH_CHECK( false, "invalid sample rate=", sample_rate, ". Supported sample rate values are: ", supported_rates.str()); } static const std::vector preferred_formats_order = { AV_SAMPLE_FMT_FLTP, AV_SAMPLE_FMT_FLT, AV_SAMPLE_FMT_DBLP, AV_SAMPLE_FMT_DBL, AV_SAMPLE_FMT_S64P, AV_SAMPLE_FMT_S64, AV_SAMPLE_FMT_S32P, AV_SAMPLE_FMT_S32, AV_SAMPLE_FMT_S16P, AV_SAMPLE_FMT_S16, AV_SAMPLE_FMT_U8P, AV_SAMPLE_FMT_U8}; AVSampleFormat find_best_output_sample_format(const AVCodec& av_codec) { const AVSampleFormat* supported_sample_formats = get_supported_output_sample_formats(av_codec); // Find a sample format that the encoder supports. We prefer using FLT[P], // since this is the format of the input samples. If FLTP isn't supported // then we'll need to convert the AVFrame's format. Our heuristic is to encode // into the format with the highest resolution. if (supported_sample_formats == nullptr) { // Can't really validate anything in this case, best we can do is hope that // FLTP is supported by the encoder. If not, FFmpeg will raise. return AV_SAMPLE_FMT_FLTP; } for (AVSampleFormat preferred_format : preferred_formats_order) { for (int i = 0; supported_sample_formats[i] != -1; ++i) { if (supported_sample_formats[i] == preferred_format) { return preferred_format; } } } // We should always find a match in preferred_formats_order, so we should // always return earlier. But in the event that a future FFmpeg version // defines an additional sample format that isn't in preferred_formats_order, // we fallback: return supported_sample_formats[0]; } void close_avio_context( AVFormatContext* av_format_context, AVIOContextHolder* avio_context_holder) { if (!av_format_context || !av_format_context->pb) { return; } if (av_format_context->pb->error == 0) { avio_flush(av_format_context->pb); } if (!avio_context_holder) { if (av_format_context->pb->error == 0) { avio_close(av_format_context->pb); } } av_format_context->pb = nullptr; } torch::stable::Tensor validate_frames( const torch::stable::Tensor& frames, const AVCodecContext* av_codec_context = nullptr, DeviceInterface* device_interface = nullptr) { STD_TORCH_CHECK( frames.scalar_type() == kStableUInt8, "frames must have uint8 dtype, got ", frames.scalar_type()); STD_TORCH_CHECK( frames.dim() == 4, "frames must have 4 dimensions (N, C, H, W), got ", frames.dim()); STD_TORCH_CHECK( frames.sizes()[1] == 3, "frame must have 3 channels (R, G, B), got ", frames.sizes()[1]); if (device_interface != nullptr) { auto& expected_device = device_interface->device(); auto frames_device = frames.device(); STD_TORCH_CHECK( frames_device == expected_device, "All frames must be on the same device. Expected ", device_type_name(expected_device.type()), ":", expected_device.index(), ", got ", device_type_name(frames_device.type()), ":", frames_device.index()); } if (av_codec_context) { STD_TORCH_CHECK( static_cast(frames.sizes()[2]) == av_codec_context->height && static_cast(frames.sizes()[3]) == av_codec_context->width, "All frames must have the same dimensions. Expected height=", av_codec_context->height, " width=", av_codec_context->width, ", got height=", frames.sizes()[2], " width=", frames.sizes()[3]); } return torch::stable::contiguous(frames); } void try_to_validate_codec_option( const AVCodec& av_codec, const char* option_name, const std::string& value) { if (!av_codec.priv_class) { return; } const AVOption* option = av_opt_find2( // Convert obj arg from const AVClass* const* to non-const void* // First cast to remove const, then cast to void* const_cast(static_cast(&av_codec.priv_class)), option_name, nullptr, 0, AV_OPT_SEARCH_FAKE_OBJ, nullptr); // If option is not found we cannot validate it, let FFmpeg handle it if (!option) { return; } // Validate if option is defined as a numeric type if (option->type == AV_OPT_TYPE_INT || option->type == AV_OPT_TYPE_INT64 || option->type == AV_OPT_TYPE_FLOAT || option->type == AV_OPT_TYPE_DOUBLE) { try { double numeric_value = std::stod(value); STD_TORCH_CHECK( numeric_value >= option->min && numeric_value <= option->max, option_name, "=", numeric_value, " is out of valid range [", option->min, ", ", option->max, "] for this codec. For more details, run 'ffmpeg -h encoder=", av_codec.name, "'"); } catch (const std::invalid_argument&) { STD_TORCH_CHECK( false, "Option ", option_name, " expects a numeric value but got '", value, "'"); } } } void sort_codec_options( const AVFormatContext* av_format_context, const std::map& extra_options, UniqueAVDictionary& codec_dict, UniqueAVDictionary& format_dict) { // Accepts a map of options as input, then sorts them into codec options and // format options. The sorted options are returned into two separate dicts. const AVClass* format_class = avformat_get_class(); const AVClass* muxer_class = av_format_context->oformat ? av_format_context->oformat->priv_class : nullptr; for (const auto& [key, value] : extra_options) { // Check if option is generic format option const AVOption* fmt_opt = av_opt_find2( &format_class, key.c_str(), nullptr, 0, AV_OPT_SEARCH_CHILDREN | AV_OPT_SEARCH_FAKE_OBJ, nullptr); // Check if option is muxer-specific option // (Returned from `ffmpeg -h muxer=mp4`) const AVOption* muxer_opt = nullptr; if (muxer_class) { muxer_opt = av_opt_find2( &muxer_class, key.c_str(), nullptr, 0, AV_OPT_SEARCH_FAKE_OBJ, nullptr); } if (fmt_opt || muxer_opt) { // Pass container-format options to formatDict to be used in // avformat_write_header av_dict_set(format_dict.get_address(), key.c_str(), value.c_str(), 0); } else { // By default, pass as codec option to be used in avcodec_open2 av_dict_set(codec_dict.get_address(), key.c_str(), value.c_str(), 0); } } } } // namespace MultiStreamEncoder::~MultiStreamEncoder() { close(); } MultiStreamEncoder::MultiStreamEncoder() { set_ffmpeg_log_level(); } void MultiStreamEncoder::open(std::string_view file_name) { STD_TORCH_CHECK(!closed_, "Cannot open after close() was called."); STD_TORCH_CHECK(!header_written_, "open() was already called."); AVFormatContext* av_format_context = nullptr; int status = avformat_alloc_output_context2( &av_format_context, nullptr, nullptr, file_name.data()); STD_TORCH_CHECK( av_format_context != nullptr, "Couldn't allocate AVFormatContext. ", "The destination file is ", file_name, ", check the desired extension? ", get_ffmpeg_error_string_from_error_code(status)); av_format_context_.reset(av_format_context); status = avio_open(&av_format_context_->pb, file_name.data(), AVIO_FLAG_WRITE); STD_TORCH_CHECK( status >= 0, "avio_open failed. The destination file is ", file_name, ", make sure it's a valid path? ", get_ffmpeg_error_string_from_error_code(status)); open_streams_and_write_header(); } void MultiStreamEncoder::open( std::string_view format_name, std::unique_ptr avio_context_holder) { STD_TORCH_CHECK(!closed_, "Cannot open after close() was called."); STD_TORCH_CHECK(!header_written_, "open() was already called."); avio_context_holder_ = std::move(avio_context_holder); // Map mkv -> matroska when used as format name format_name = (format_name == "mkv") ? "matroska" : format_name; AVFormatContext* av_format_context = nullptr; int status = avformat_alloc_output_context2( &av_format_context, nullptr, format_name.data(), nullptr); STD_TORCH_CHECK( av_format_context != nullptr, "Couldn't allocate AVFormatContext. ", "Check the desired format? Got format=", format_name, ". ", get_ffmpeg_error_string_from_error_code(status)); av_format_context_.reset(av_format_context); av_format_context_->pb = avio_context_holder_->get_avio_context(); open_streams_and_write_header(); } int MultiStreamEncoder::add_video_stream( int height, int width, double frame_rate, std::string device, std::optional codec, std::optional pixel_format, std::optional crf, std::optional preset, std::optional> extra_options) { STD_TORCH_CHECK(height > 0, "height must be > 0, got ", height); STD_TORCH_CHECK(width > 0, "width must be > 0, got ", width); STD_TORCH_CHECK(frame_rate > 0, "frame_rate must be > 0, got ", frame_rate); VideoStream video_stream; StableDevice stable_device(std::move(device)); video_stream.device_interface = create_device_interface( stable_device, stable_device.type() == kStableCUDA ? "ffmpeg" : "default"); video_stream.in_height = height; video_stream.in_width = width; video_stream.in_frame_rate = frame_rate; video_stream.options.codec = std::move(codec); video_stream.options.pixel_format = std::move(pixel_format); video_stream.options.crf = crf; video_stream.options.preset = std::move(preset); video_stream.options.extra_options = std::move(extra_options); video_streams_.push_back(std::move(video_stream)); return static_cast(video_streams_.size() - 1); } int MultiStreamEncoder::add_audio_stream( int sample_rate, int num_channels, std::optional bit_rate, std::optional out_num_channels, std::optional out_sample_rate) { STD_TORCH_CHECK( sample_rate > 0, "sample_rate must be > 0, got ", sample_rate); STD_TORCH_CHECK( num_channels > 0, "num_channels must be > 0, got ", num_channels); STD_TORCH_CHECK( num_channels <= AV_NUM_DATA_POINTERS, "Trying to encode ", num_channels, " channels, but FFmpeg only supports ", AV_NUM_DATA_POINTERS, " channels per frame."); AudioStream audio_stream; audio_stream.in_sample_rate = sample_rate; audio_stream.in_num_channels = num_channels; audio_stream.options.bit_rate = bit_rate; audio_stream.options.num_channels = out_num_channels; audio_stream.options.sample_rate = out_sample_rate; audio_streams_.push_back(std::move(audio_stream)); return static_cast(audio_streams_.size() - 1); } void MultiStreamEncoder::initialize_video_stream(VideoStream& video_stream) { const AVCodec* av_codec = nullptr; // If codec arg is provided, find codec using logic similar to FFmpeg: // https://github.com/FFmpeg/FFmpeg/blob/master/fftools/ffmpeg_opt.c#L804-L835 if (video_stream.options.codec.has_value()) { const std::string& codec = video_stream.options.codec.value(); // Try to find codec by name ("libx264", "libsvtav1") av_codec = avcodec_find_encoder_by_name(codec.c_str()); // Try to find by codec descriptor ("h264", "av1") if (!av_codec) { const AVCodecDescriptor* desc = avcodec_descriptor_get_by_name(codec.c_str()); if (desc) { av_codec = avcodec_find_encoder(desc->id); } } } else { STD_TORCH_CHECK( av_format_context_->oformat != nullptr, "Output format is null, unable to find default codec."); // Try to substitute the default codec with its hardware equivalent // This will return std::nullopt when device is CPU. auto hw_codec = video_stream.device_interface->find_codec( av_format_context_->oformat->video_codec, /*isDecoder=*/false); if (hw_codec.has_value()) { av_codec = hw_codec.value(); } if (!av_codec) { av_codec = avcodec_find_encoder(av_format_context_->oformat->video_codec); } } STD_TORCH_CHECK( av_codec != nullptr, "Video codec ", video_stream.options.codec.has_value() ? video_stream.options.codec.value() + " " : "", "not found. To see available codecs, run: ffmpeg -encoders"); AVCodecContext* av_codec_context = avcodec_alloc_context3(av_codec); STD_TORCH_CHECK( av_codec_context != nullptr, "Couldn't allocate codec context."); video_stream.av_codec_context.reset(av_codec_context); int out_height = video_stream.in_height; int out_width = video_stream.in_width; // Pixel-format selection is delegated to the device interface. AVPixelFormat out_pixel_format = video_stream.device_interface->get_encoding_pixel_format( *av_codec, video_stream.options.pixel_format); // Configure codec parameters video_stream.av_codec_context->codec_id = av_codec->id; video_stream.av_codec_context->width = out_width; video_stream.av_codec_context->height = out_height; video_stream.av_codec_context->pix_fmt = out_pixel_format; video_stream.av_codec_context->framerate = av_d2q(video_stream.in_frame_rate, INT_MAX); video_stream.av_codec_context->time_base = av_inv_q(video_stream.av_codec_context->framerate); // Set flag for containers that require extradata to be in the codec context if (av_format_context_->oformat->flags & AVFMT_GLOBALHEADER) { video_stream.av_codec_context->flags |= AV_CODEC_FLAG_GLOBAL_HEADER; } // Apply videoStreamOptions UniqueAVDictionary av_codec_options; if (video_stream.options.extra_options.has_value()) { for (const auto& [key, value] : video_stream.options.extra_options.value()) { try_to_validate_codec_option(*av_codec, key.c_str(), value); } sort_codec_options( av_format_context_.get(), video_stream.options.extra_options.value(), av_codec_options, av_format_options_); } if (video_stream.options.crf.has_value()) { std::string crf_value = std::to_string(video_stream.options.crf.value()); try_to_validate_codec_option(*av_codec, "crf", crf_value); av_dict_set(av_codec_options.get_address(), "crf", crf_value.c_str(), 0); } if (video_stream.options.preset.has_value()) { av_dict_set( av_codec_options.get_address(), "preset", video_stream.options.preset.value().c_str(), 0); } // Hardware setup is a no-op on CPU; HW devices override these hooks. video_stream.device_interface->register_hardware_device_with_codec( video_stream.av_codec_context.get()); video_stream.device_interface->setup_hardware_frame_context_for_encoding( video_stream.av_codec_context.get()); int status = avcodec_open2( video_stream.av_codec_context.get(), av_codec, av_codec_options.get_address()); STD_TORCH_CHECK( status == AVSUCCESS, "avcodec_open2 failed: ", get_ffmpeg_error_string_from_error_code(status)); video_stream.av_stream = avformat_new_stream(av_format_context_.get(), nullptr); STD_TORCH_CHECK( video_stream.av_stream != nullptr, "Couldn't create new stream."); // Set the stream time base to encode correct frame timestamps video_stream.av_stream->time_base = video_stream.av_codec_context->time_base; // Set the stream frame rate to store correct frame durations for some // containers (webm, mkv) video_stream.av_stream->r_frame_rate = video_stream.av_codec_context->framerate; status = avcodec_parameters_from_context( video_stream.av_stream->codecpar, video_stream.av_codec_context.get()); STD_TORCH_CHECK( status == AVSUCCESS, "avcodec_parameters_from_context failed: ", get_ffmpeg_error_string_from_error_code(status)); } void MultiStreamEncoder::initialize_audio_stream(AudioStream& audio_stream) { // We use the AVFormatContext's default codec for that // specific format/container. const AVCodec* av_codec = avcodec_find_encoder(av_format_context_->oformat->audio_codec); STD_TORCH_CHECK(av_codec != nullptr, "Codec not found"); AVCodecContext* av_codec_context = avcodec_alloc_context3(av_codec); STD_TORCH_CHECK( av_codec_context != nullptr, "Couldn't allocate codec context."); audio_stream.av_codec_context.reset(av_codec_context); auto desired_bit_rate = audio_stream.options.bit_rate; if (desired_bit_rate.has_value()) { STD_TORCH_CHECK( *desired_bit_rate >= 0, "bit_rate=", *desired_bit_rate, " must be >= 0."); } // bit_rate=None defaults to 0, which is what the FFmpeg CLI seems to use as // well when "-b:a" isn't specified. audio_stream.av_codec_context->bit_rate = desired_bit_rate.value_or(0); int out_num_channels = audio_stream.options.num_channels.value_or(audio_stream.in_num_channels); audio_stream.out_num_channels = out_num_channels; validate_num_channels(*av_codec, out_num_channels); set_default_channel_layout(audio_stream.av_codec_context, out_num_channels); int out_sample_rate = audio_stream.options.sample_rate.value_or(audio_stream.in_sample_rate); audio_stream.out_sample_rate = out_sample_rate; validate_sample_rate(*av_codec, out_sample_rate); audio_stream.av_codec_context->sample_rate = out_sample_rate; audio_stream.av_codec_context->time_base = AVRational{1, out_sample_rate}; // Input samples are expected to be FLTP. Not all encoders support FLTP, so we // may need to convert the samples into a supported output sample format, // which is what the `.sample_fmt` defines. audio_stream.av_codec_context->sample_fmt = find_best_output_sample_format(*av_codec); int status = avcodec_open2(audio_stream.av_codec_context.get(), av_codec, nullptr); STD_TORCH_CHECK( status == AVSUCCESS, "avcodec_open2 failed: ", get_ffmpeg_error_string_from_error_code(status)); // We're allocating the stream here. Streams are meant to be freed by // avformat_free_context(avFormatContext), which we call in the // avFormatContext_'s destructor. audio_stream.av_stream = avformat_new_stream(av_format_context_.get(), nullptr); STD_TORCH_CHECK( audio_stream.av_stream != nullptr, "Couldn't create new audio stream."); status = avcodec_parameters_from_context( audio_stream.av_stream->codecpar, audio_stream.av_codec_context.get()); STD_TORCH_CHECK( status == AVSUCCESS, "avcodec_parameters_from_context failed: ", get_ffmpeg_error_string_from_error_code(status)); // If a codec supports variable frame size, frame_size may not be defined, in // which case we default to 256 like torchaudio. audio_stream.frame_size = audio_stream.av_codec_context->frame_size > 0 ? audio_stream.av_codec_context->frame_size : 256; // We always create a FIFO so that addSamples() can be called multiple times // with various chunk sizes that are then buffered and encoded in frame_size // sized batches. auto av_audio_fifo = av_audio_fifo_alloc( audio_stream.av_codec_context->sample_fmt, out_num_channels, audio_stream.frame_size * 2); STD_TORCH_CHECK(av_audio_fifo != nullptr, "Couldn't create AVAudioFifo."); audio_stream.av_audio_fifo.reset(av_audio_fifo); } void MultiStreamEncoder::open_streams_and_write_header() { STD_TORCH_CHECK( !video_streams_.empty() || !audio_streams_.empty(), "Call addVideoStream() or addAudioStream() before open()."); for (auto& video_stream : video_streams_) { initialize_video_stream(video_stream); } for (auto& audio_stream : audio_streams_) { initialize_audio_stream(audio_stream); } int status = avformat_write_header( av_format_context_.get(), av_format_options_.get_address()); STD_TORCH_CHECK( status == AVSUCCESS, "Error in avformat_write_header: ", get_ffmpeg_error_string_from_error_code(status)); header_written_ = true; } void MultiStreamEncoder::add_frames( const torch::stable::Tensor& frames, int stream_index) { STD_TORCH_CHECK(!closed_, "Cannot add frames after close() was called."); STD_TORCH_CHECK(header_written_, "Call open() before addFrames()."); STD_TORCH_CHECK( stream_index >= 0 && stream_index < static_cast(video_streams_.size()), "Invalid stream index ", stream_index, ". Number of video streams: ", video_streams_.size()); auto& video_stream = video_streams_[stream_index]; auto validated_frames = validate_frames( frames, video_stream.av_codec_context.get(), video_stream.device_interface.get()); AutoAVPacket auto_av_packet; // TODO MultiStreamEncoder: Consider using accessor for potential performance // improvement int num_frames = static_cast(validated_frames.sizes()[0]); for (int i = 0; i < num_frames; ++i) { torch::stable::Tensor curr_frame = select_row(validated_frames, i); int frame_index = video_stream.num_encoded_frames + i; UniqueAVFrame av_frame = video_stream.device_interface->convert_tensor_to_av_frame_for_encoding( curr_frame, frame_index, video_stream.av_codec_context.get()); STD_TORCH_CHECK( av_frame != nullptr, "convertTensorToAVFrameForEncoding failed for frame ", frame_index, " on device: ", device_type_name(validated_frames.device().type())); encode_video_frame(auto_av_packet, av_frame, video_stream); } video_stream.num_encoded_frames += num_frames; } void MultiStreamEncoder::encode_video_frame( AutoAVPacket& auto_av_packet, const UniqueAVFrame& av_frame, VideoStream& video_stream) { auto status = avcodec_send_frame(video_stream.av_codec_context.get(), av_frame.get()); STD_TORCH_CHECK( status == AVSUCCESS, "Error while sending frame: ", get_ffmpeg_error_string_from_error_code(status)); while (status >= 0) { ReferenceAVPacket packet(auto_av_packet); status = avcodec_receive_packet( video_stream.av_codec_context.get(), packet.get()); if (status == AVERROR(EAGAIN) || status == AVERROR_EOF) { if (status == AVERROR_EOF) { // Flush remaining buffered packets status = av_interleaved_write_frame(av_format_context_.get(), nullptr); STD_TORCH_CHECK( status == AVSUCCESS, "Failed to flush packet: ", get_ffmpeg_error_string_from_error_code(status)); } return; } STD_TORCH_CHECK( status >= 0, "Error receiving packet: ", get_ffmpeg_error_string_from_error_code(status)); // The code below is borrowed from torchaudio: // https://github.com/pytorch/audio/blob/b6a3368a45aaafe05f1a6a9f10c68adc5e944d9e/src/libtorio/ffmpeg/stream_writer/encoder.cpp#L46 // Setting packet->duration to 1 allows the last frame to be properly // encoded, and needs to be set before calling av_packet_rescale_ts. if (packet->duration == 0) { packet->duration = 1; } av_packet_rescale_ts( packet.get(), video_stream.av_codec_context->time_base, video_stream.av_stream->time_base); packet->stream_index = video_stream.av_stream->index; status = av_interleaved_write_frame(av_format_context_.get(), packet.get()); STD_TORCH_CHECK( status == AVSUCCESS, "Error in av_interleaved_write_frame: ", get_ffmpeg_error_string_from_error_code(status)); } } void MultiStreamEncoder::add_samples( const torch::stable::Tensor& samples, int stream_index) { STD_TORCH_CHECK(!closed_, "Cannot add samples after close() was called."); STD_TORCH_CHECK(header_written_, "Call open() before addSamples()."); STD_TORCH_CHECK( stream_index >= 0 && stream_index < static_cast(audio_streams_.size()), "Invalid stream index ", stream_index, ". Number of audio streams: ", audio_streams_.size()); auto& audio_stream = audio_streams_[stream_index]; auto validated_samples = validate_samples(samples); STD_TORCH_CHECK( static_cast(validated_samples.sizes()[0]) == audio_stream.in_num_channels, "Expected ", audio_stream.in_num_channels, " channels, got ", validated_samples.sizes()[0]); encode_audio_samples(validated_samples, audio_stream); } void MultiStreamEncoder::encode_audio_samples( const torch::stable::Tensor& samples, AudioStream& audio_stream) { UniqueAVFrame av_frame = allocate_av_frame( audio_stream.frame_size, audio_stream.in_sample_rate, audio_stream.in_num_channels, AV_SAMPLE_FMT_FLTP); AutoAVPacket auto_av_packet; const uint8_t* psamples = static_cast(samples.const_data_ptr()); int num_samples = static_cast(samples.sizes()[1]); // per channel int num_encoded_samples = 0; // per channel int num_bytes_per_sample = static_cast(samples.element_size()); int num_bytes_per_channel = num_samples * num_bytes_per_sample; while (num_encoded_samples < num_samples) { int num_samples_to_encode = std::min(audio_stream.frame_size, num_samples - num_encoded_samples); int num_bytes_to_encode = num_samples_to_encode * num_bytes_per_sample; for (int ch = 0; ch < audio_stream.in_num_channels; ch++) { std::memcpy( av_frame->data[ch], psamples + ch * num_bytes_per_channel, num_bytes_to_encode); } psamples += num_bytes_to_encode; // Above, we set the AVFrame's .nb_samples to AVCodecContext.frame_size so // that the frame buffers are allocated to a big enough size. Here, we reset // it to the exact number of samples that need to be encoded, otherwise the // encoded frame would contain more samples than necessary and our results // wouldn't match the ffmpeg CLI. av_frame->nb_samples = num_samples_to_encode; UniqueAVFrame converted_av_frame = maybe_convert_audio_av_frame(av_frame, audio_stream); encode_audio_frame_through_fifo( auto_av_packet, converted_av_frame, audio_stream); num_encoded_samples += num_samples_to_encode; } STD_TORCH_CHECK( num_encoded_samples == num_samples, "Hmmmmmm something went wrong."); } UniqueAVFrame MultiStreamEncoder::maybe_convert_audio_av_frame( const UniqueAVFrame& av_frame, AudioStream& audio_stream) { if (static_cast(av_frame->format) == audio_stream.av_codec_context->sample_fmt && get_num_channels(*av_frame) == audio_stream.out_num_channels && av_frame->sample_rate == audio_stream.out_sample_rate) { // Note: the clone references the same underlying data, it's a cheap copy. return UniqueAVFrame(av_frame_clone(av_frame.get())); } if (!audio_stream.swr_context) { audio_stream.swr_context.reset(create_swr_context( static_cast(av_frame->format), audio_stream.av_codec_context->sample_fmt, av_frame->sample_rate, audio_stream.out_sample_rate, *av_frame, audio_stream.out_num_channels)); } // convertAudioAVFrameSamples uses avFrame's extended_data field, so we ensure // it's the same as data. This should always be the case since we validated // earlier that we have less than AV_NUM_DATA_POINTERS channels. STD_TORCH_CHECK( av_frame->data == av_frame->extended_data, "Codec context data and extended_data pointers differ, this is unexpected."); UniqueAVFrame converted_av_frame = convert_audio_av_frame_samples( audio_stream.swr_context, *av_frame, audio_stream.av_codec_context->sample_fmt, audio_stream.out_sample_rate, audio_stream.out_num_channels); if (av_frame->sample_rate == audio_stream.out_sample_rate) { STD_TORCH_CHECK( converted_av_frame->nb_samples == av_frame->nb_samples, "convertedAVFrame->nb_samples=", converted_av_frame->nb_samples, " differs from ", "avFrame->nb_samples=", av_frame->nb_samples, "This is unexpected, please report on the TorchCodec bug tracker."); } return converted_av_frame; } void MultiStreamEncoder::encode_audio_frame_through_fifo( AutoAVPacket& auto_av_packet, const UniqueAVFrame& av_frame, AudioStream& audio_stream, // flushFifo is only set to true in maybeFlushSwrAndFifo(), i.e. at the very // end of the encoding process when we're flushing buffers. We also want to // flush the FIFO so as to not leave any remaining samples in it. bool flush_fifo) { if (av_frame != nullptr) { int num_samples_written = av_audio_fifo_write( audio_stream.av_audio_fifo.get(), reinterpret_cast(av_frame->data), av_frame->nb_samples); STD_TORCH_CHECK( num_samples_written == av_frame->nb_samples, "Tried to write ", av_frame->nb_samples, " samples, but only wrote ", num_samples_written); } UniqueAVFrame newav_frame = allocate_av_frame( audio_stream.frame_size, audio_stream.av_codec_context->sample_rate, audio_stream.out_num_channels, audio_stream.av_codec_context->sample_fmt); // Explaining the while bound: // - if we're not flushing the FIFO, i.e. in most cases, we want to pull // exactly `frame_size` samples from the FIFO, so we have to stop before it // contains less than `frame_size` samples. // - if we're flushing the FIFO, we want to read from the FIFO until the very // last sample it contains. // // In both cases, for as long as we can, we're trying to pull exactly // `frame_size` samples from the FIFO and send each `frame_size`-sized avFrame // to encodeAudioFrame(). Only the very last avFrame of the encoding process // is allowed to contain less than frame_size samples. That only happens when // flushFifo is true. while (av_audio_fifo_size(audio_stream.av_audio_fifo.get()) >= (flush_fifo ? 1 : audio_stream.frame_size)) { int samples_to_read = std::min( av_audio_fifo_size(audio_stream.av_audio_fifo.get()), newav_frame->nb_samples); int num_samples_read = av_audio_fifo_read( audio_stream.av_audio_fifo.get(), reinterpret_cast(newav_frame->data), samples_to_read); STD_TORCH_CHECK( num_samples_read == samples_to_read, "Tried to read ", samples_to_read, " samples, but only read ", num_samples_read); newav_frame->nb_samples = num_samples_read; encode_audio_frame(auto_av_packet, newav_frame, audio_stream); } } void MultiStreamEncoder::encode_audio_frame( AutoAVPacket& auto_av_packet, const UniqueAVFrame& av_frame, AudioStream& audio_stream) { if (av_frame != nullptr) { av_frame->pts = audio_stream.last_encoded_av_frame_pts; audio_stream.last_encoded_av_frame_pts += av_frame->nb_samples; } auto status = avcodec_send_frame(audio_stream.av_codec_context.get(), av_frame.get()); STD_TORCH_CHECK( status == AVSUCCESS, "Error while sending frame: ", get_ffmpeg_error_string_from_error_code(status)); while (status >= 0) { ReferenceAVPacket packet(auto_av_packet); status = avcodec_receive_packet( audio_stream.av_codec_context.get(), packet.get()); if (status == AVERROR(EAGAIN) || status == AVERROR_EOF) { if (status == AVERROR_EOF) { // Flush the packets that were potentially buffered by // av_interleaved_write_frame(). See corresponding block in // TorchAudio: // https://github.com/pytorch/audio/blob/d60ce09e2c532d5bf2e05619e700ab520543465e/src/libtorio/ffmpeg/stream_writer/encoder.cpp#L21 status = av_interleaved_write_frame(av_format_context_.get(), nullptr); STD_TORCH_CHECK( status == AVSUCCESS, "Failed to flush packet: ", get_ffmpeg_error_string_from_error_code(status)); } return; } STD_TORCH_CHECK( status >= 0, "Error receiving packet: ", get_ffmpeg_error_string_from_error_code(status)); packet->stream_index = audio_stream.av_stream->index; av_packet_rescale_ts( packet.get(), audio_stream.av_codec_context->time_base, audio_stream.av_stream->time_base); status = av_interleaved_write_frame(av_format_context_.get(), packet.get()); STD_TORCH_CHECK( status == AVSUCCESS, "Error in av_interleaved_write_frame: ", get_ffmpeg_error_string_from_error_code(status)); } } void MultiStreamEncoder::maybe_flush_swr_and_fifo( AutoAVPacket& auto_av_packet, AudioStream& audio_stream) { // When sample conversion is involved, libswresample may have buffered some // samples that we need to flush into the FIFO before draining it. UniqueAVFrame swr_frame(nullptr); if (audio_stream.swr_context != nullptr) { int num_remaining_samples = // this is an upper bound swr_get_out_samples(audio_stream.swr_context.get(), 0); if (num_remaining_samples > 0) { swr_frame = allocate_av_frame( num_remaining_samples, audio_stream.out_sample_rate, audio_stream.out_num_channels, audio_stream.av_codec_context->sample_fmt); int actual_num_remaining_samples = swr_convert( audio_stream.swr_context.get(), swr_frame->data, swr_frame->nb_samples, nullptr, 0); swr_frame->nb_samples = actual_num_remaining_samples; } } // Flush any remaining swr samples into the FIFO, then drain it. encode_audio_frame_through_fifo( auto_av_packet, swr_frame, audio_stream, /*flushFifo=*/true); } void MultiStreamEncoder::flush_buffers() { for (auto& audio_stream : audio_streams_) { if (audio_stream.av_stream != nullptr) { AutoAVPacket audio_av_packet; maybe_flush_swr_and_fifo(audio_av_packet, audio_stream); encode_audio_frame(audio_av_packet, UniqueAVFrame(nullptr), audio_stream); } } for (auto& video_stream : video_streams_) { if (video_stream.av_stream != nullptr) { AutoAVPacket video_av_packet; encode_video_frame(video_av_packet, UniqueAVFrame(nullptr), video_stream); } } } void MultiStreamEncoder::close() { if (closed_) { return; } // TODO MultiStreamEncoder: Revisit if "closed_" flag is useful closed_ = true; if (header_written_) { flush_buffers(); int status = av_write_trailer(av_format_context_.get()); // av_write_trailer returns mfra atom size (positive) for fragmented // containers. All FFmpeg errors are negative, so positive is not an error. if (status > 0) { status = AVSUCCESS; } STD_TORCH_CHECK( status == AVSUCCESS, "Error in av_write_trailer: ", get_ffmpeg_error_string_from_error_code(status)); } close_avio_context(av_format_context_.get(), avio_context_holder_.get()); } } // namespace facebook::torchcodec