diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 9190ac4d4..fa7ff7a9c 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -337,28 +337,33 @@ int BetaCudaDeviceInterface::streamPropertyChange(CUVIDEOFORMAT* videoFormat) { // Moral equivalent of avcodec_send_packet(). Here, we pass the AVPacket down to // the NVCUVID parser. int BetaCudaDeviceInterface::sendPacket(ReferenceAVPacket& packet) { + TORCH_CHECK( + packet.get() && packet->data && packet->size > 0, + "sendPacket received an empty packet, this is unexpected, please report."); + + applyBSF(packet); + CUVIDSOURCEDATAPACKET cuvidPacket = {}; + cuvidPacket.payload = packet->data; + cuvidPacket.payload_size = packet->size; + cuvidPacket.flags = CUVID_PKT_TIMESTAMP; + cuvidPacket.timestamp = packet->pts; - if (packet.get() && packet->data && packet->size > 0) { - applyBSF(packet); + return sendCuvidPacket(cuvidPacket); +} - // Regular packet with data - cuvidPacket.payload = packet->data; - cuvidPacket.payload_size = packet->size; - cuvidPacket.flags = CUVID_PKT_TIMESTAMP; - cuvidPacket.timestamp = packet->pts; +int BetaCudaDeviceInterface::sendEOFPacket() { + CUVIDSOURCEDATAPACKET cuvidPacket = {}; + cuvidPacket.flags = CUVID_PKT_ENDOFSTREAM; + eofSent_ = true; - } else { - // End of stream packet - cuvidPacket.flags = CUVID_PKT_ENDOFSTREAM; - eofSent_ = true; - } + return sendCuvidPacket(cuvidPacket); +} +int BetaCudaDeviceInterface::sendCuvidPacket( + CUVIDSOURCEDATAPACKET& cuvidPacket) { CUresult result = cuvidParseVideoData(videoParser_, &cuvidPacket); - if (result != CUDA_SUCCESS) { - return AVERROR_EXTERNAL; - } - return AVSUCCESS; + return result == CUDA_SUCCESS ? AVSUCCESS : AVERROR_EXTERNAL; } void BetaCudaDeviceInterface::applyBSF(ReferenceAVPacket& packet) { @@ -551,17 +556,7 @@ UniqueAVFrame BetaCudaDeviceInterface::convertCudaFrameToAVFrame( void BetaCudaDeviceInterface::flush() { isFlushing_ = true; - // TODONVDEC P0: simplify flushing and "eofSent_" logic. We should just have a - // "sendEofPacket()" function that does the right thing, instead of setting - // CUVID_PKT_ENDOFSTREAM in different places. - if (!eofSent_) { - CUVIDSOURCEDATAPACKET cuvidPacket = {}; - cuvidPacket.flags = CUVID_PKT_ENDOFSTREAM; - CUresult result = cuvidParseVideoData(videoParser_, &cuvidPacket); - if (result == CUDA_SUCCESS) { - eofSent_ = true; - } - } + sendEOFPacket(); isFlushing_ = false; diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.h b/src/torchcodec/_core/BetaCudaDeviceInterface.h index 9c452c48c..d4b9fc0a5 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.h +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.h @@ -52,6 +52,7 @@ class BetaCudaDeviceInterface : public DeviceInterface { } int sendPacket(ReferenceAVPacket& packet) override; + int sendEOFPacket() override; int receiveFrame(UniqueAVFrame& avFrame) override; void flush() override; @@ -61,6 +62,7 @@ class BetaCudaDeviceInterface : public DeviceInterface { int frameReadyInDisplayOrder(CUVIDPARSERDISPINFO* dispInfo); private: + int sendCuvidPacket(CUVIDSOURCEDATAPACKET& cuvidPacket); // Apply bitstream filter, modifies packet in-place void applyBSF(ReferenceAVPacket& packet); void initializeBSF( diff --git a/src/torchcodec/_core/DeviceInterface.h b/src/torchcodec/_core/DeviceInterface.h index 284d83378..35984da26 100644 --- a/src/torchcodec/_core/DeviceInterface.h +++ b/src/torchcodec/_core/DeviceInterface.h @@ -100,6 +100,14 @@ class DeviceInterface { return AVERROR(ENOSYS); } + // Send an EOF packet to flush the decoder + // Returns AVSUCCESS on success, or other AVERROR on failure + virtual int sendEOFPacket() { + TORCH_CHECK( + false, "Send EOF packet not implemented for this device interface"); + return AVERROR(ENOSYS); + } + // Moral equivalent of avcodec_receive_frame() // Returns AVSUCCESS on success, AVERROR(EAGAIN) if no frame ready, // AVERROR_EOF if end of stream, or other AVERROR on failure diff --git a/src/torchcodec/_core/SingleStreamDecoder.cpp b/src/torchcodec/_core/SingleStreamDecoder.cpp index 53c15cd33..27d46b4d4 100644 --- a/src/torchcodec/_core/SingleStreamDecoder.cpp +++ b/src/torchcodec/_core/SingleStreamDecoder.cpp @@ -1220,12 +1220,7 @@ UniqueAVFrame SingleStreamDecoder::decodeAVFrame( if (status == AVERROR_EOF) { // End of file reached. We must drain the decoder if (deviceInterface_->canDecodePacketDirectly()) { - // TODONVDEC P0: Re-think this. This should be simpler. - AutoAVPacket eofAutoPacket; - ReferenceAVPacket eofPacket(eofAutoPacket); - eofPacket->data = nullptr; - eofPacket->size = 0; - status = deviceInterface_->sendPacket(eofPacket); + status = deviceInterface_->sendEOFPacket(); } else { status = avcodec_send_packet( streamInfo.codecContext.get(),