Bug Report: KeyError when getting audio duration from webm files

[holys]

Description

When processing webm audio files, the get_audio_duration function fails with a KeyError because the 'duration' field is not available in the audio stream metadata.

Error Message

Traceback (most recent call last):
  File "/srv/r8/monobase/cog/g1df235d5/cog0.13.6-python3.11/lib/python3.11/site-packages/cog/server/worker.py", line 646, in _handle_predict_error
    yield
  File "/srv/r8/monobase/cog/g1df235d5/cog0.13.6-python3.11/lib/python3.11/site-packages/cog/server/worker.py", line 538, in _predict
    result = predict(**payload)
  File "/src/predict.py", line 103, in predict
    audio_duration = get_audio_duration(audio_file)
  File "/src/predict.py", line 181, in get_audio_duration
    return float(stream['duration']) * 1000
KeyError: 'duration'

Current Behavior

The function attempts to read the duration directly from the audio stream metadata, but for some webm files, this information is only available in the format metadata section of the ffmpeg probe result.

Proposed Solution

def get_audio_duration(file_path):
   probe = ffmpeg.probe(file_path)
   
   # First try to get duration from audio stream
   stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'audio'), None)
   if stream and 'duration' in stream:
       return float(stream['duration']) * 1000
       
   # Fallback to format duration if stream duration is not available
   if 'format' in probe and 'duration' in probe['format']:
       return float(probe['format']['duration']) * 1000
       
   raise ValueError("Could not determine audio duration from file metadata")

PR: #12

Additional Context

This issue specifically occurs with webm files where the duration information is stored in the format metadata rather than the stream metadata.