How do I export a video with audio from separated stream?

[Saxamos]

Overview

I want to make modification on the images only on a video. I am using pyav to load and save it.

Expected behavior

The output video must be a video with both video and audio stream.

Actual behavior

I only manage to output either an audio or a video (without audio) file.

My code look like this:

container = av.open("my_input_video.mp4")
frames = [some_transformation(img.to_ndarray(format="rgb24")) for img in container.decode(container.streams.video[0])]

output = av.open("my_output_video.mp4", "w")

audio_stream = output.add_stream("aac", 44100)
video_stream = output.add_stream("h264", 25)
video_stream.bit_rate = 1081899
video_stream.height = 1500
video_stream.width = 844

for frame in frames:
    av_frame = av.VideoFrame.from_ndarray(frame, format="rgb24")
    packet = video_stream.encode(av_frame)
    output.mux(packet)

for audio_frame in container.decode(container.streams.audio[0]):
    packet = audio_stream.encode(audio_frame)
    output.mux(packet)

output.close()

With the previous code, the output is a video without audio and I got the following traceback:

Traceback (most recent call last):
  File "/Users/XXX/export.py", line 30, in save_reframed_video
    packet = audio_stream.encode(audio_frame)
  File "av/stream.pyx", line 152, in av.stream.Stream.encode
  File "av/codec/context.pyx", line 466, in av.codec.context.CodecContext.encode
  File "av/audio/codeccontext.pyx", line 49, in av.audio.codeccontext.AudioCodecContext._prepare_frames_for_encode
  File "av/audio/fifo.pyx", line 88, in av.audio.fifo.AudioFifo.write
ValueError: Frame.pts (1934336) != expected (0); fix or set to None.
Traceback (most recent call last):
  File "av/container/output.pyx", line 25, in av.container.output.close_output
TypeError: 'NoneType' object is not iterable
Exception ignored in: 'av.container.output.OutputContainer.__dealloc__'
Traceback (most recent call last):
  File "av/container/output.pyx", line 25, in av.container.output.close_output
TypeError: 'NoneType' object is not iterable

Process finished with exit code 1

Now I've tried to change the order of stream (starting with the audio) and I only got an audio file without video. In this order there is no error displayed, and I can see 2 streams in the output when I'm in debug mode.