skip median filtering when the input is too small

whisper/model.py

@@ -1,16 +1,17 @@
+import base64
+import gzip
 from dataclasses import dataclass
 from typing import Dict
 from typing import Iterable, Optional
-import gzip
-import base64
+
 import numpy as np
 import torch
 import torch.nn.functional as F
 from torch import Tensor
 from torch import nn
 
-from .transcribe import transcribe as transcribe_function
 from .decoding import detect_language as detect_language_function, decode as decode_function
+from .transcribe import transcribe as transcribe_function
 
 
 @dataclass

whisper/timing.py

@@ -1,4 +1,3 @@
-import time
 import subprocess
 import warnings
 from dataclasses import dataclass
@@ -18,7 +17,13 @@
 
 def median_filter(x: torch.Tensor, filter_width: int):
  """Apply a median filter of width `filter_width` along the last dimension of `x`"""
- if (ndim := x.ndim) <= 2: # `F.pad` does not support 1D or 2D inputs for reflect padding
+ pad_width = filter_width // 2
+ if x.shape[-1] <= pad_width:
+ # F.pad requires the padding width to be smaller than the input dimension
+ return x
+
+ if (ndim := x.ndim) <= 2:
+ # `F.pad` does not support 1D or 2D inputs for reflect padding but supports 3D and 4D
  x = x[None, None, :]
 
  assert filter_width > 0 and filter_width % 2 == 1, "`filter_width` should be an odd number"