Update handler.py

handler.py

@@ -5,11 +5,11 @@ Outputs: GIF, WebM, ZIP(frames)
 This version includes:
 - Defensive patches to avoid hosted runtime failing with:
     "`ffmpeg` is not a registered plugin name."
-- Robust frame extraction for shapes like (T,H,W,C) and (B,T,H,W,C)
+- Robust frame extraction for shapes like (T,H,W,C), (B,T,H,W,C), (T,C,H,W), (B,T,C,H,W)
 - Output encoders:
   - GIF: Pillow only (no ffmpeg)
   - ZIP: PNG frames zipped (no ffmpeg)
-  - WebM: imageio + imageio-ffmpeg (absolute ffmpeg path)
+  - WebM: imageio + imageio-ffmpeg via IMAGEIO_FFMPEG_EXE env var (NO executable= arg)
 
 IMPORTANT:
 - HF gateway often requires top-level { "inputs": {...} }.
@@ -87,7 +87,7 @@ def _patch_hf_toolkit_ffmpeg_plugin() -> dict:
         except Exception as e:
             diag["details"].append(f"failed wrapping {fn_name}: {e}")
 
-    # Some runtimes may import a module path for plugins; provide dummy module.
+    # Dummy module injection (helps if toolkit tries importing a plugin module)
     dummy_mod_name = "huggingface_inference_toolkit.plugins.ffmpeg"
     if dummy_mod_name not in sys.modules:
         dummy = types.ModuleType(dummy_mod_name)
@@ -119,7 +119,7 @@ import imageio
 try:
     import imageio_ffmpeg  # type: ignore
     _FFMPEG_EXE = imageio_ffmpeg.get_ffmpeg_exe()
-    # Force absolute path; some runtimes intercept plain "ffmpeg"
+    # Force absolute ffmpeg path via env var; do NOT use imageio executable= param (not supported here).
     os.environ["IMAGEIO_FFMPEG_EXE"] = _FFMPEG_EXE
 except Exception:
     _FFMPEG_EXE = ""
@@ -145,7 +145,7 @@ def _clamp_uint8_frame(frame: np.ndarray) -> np.ndarray:
     if not isinstance(frame, np.ndarray):
         frame = np.array(frame)
 
-    # squeeze batch dim
+    # squeeze batch dim (rare)
     if frame.ndim == 4 and frame.shape[0] == 1:
         frame = frame[0]
 
@@ -199,7 +199,11 @@ def _encode_gif(frames: List[np.ndarray], fps: int) -> bytes:
 
 def _encode_webm(frames: List[np.ndarray], fps: int, quality: str = "good") -> bytes:
     """
-    Encode WebM (VP9) via imageio. Requires ffmpeg through imageio-ffmpeg.
+    Encode WebM (VP9) via imageio.
+
+    IMPORTANT:
+    - Do NOT pass executable=...; HF's imageio build can reject that parameter.
+    - We rely on IMAGEIO_FFMPEG_EXE env var set at import time.
     """
     if not frames:
         raise ValueError("No frames to encode WebM.")
@@ -223,7 +227,6 @@ def _encode_webm(frames: List[np.ndarray], fps: int, quality: str = "good") -> b
             out_path,
             fps=max(1, fps),
             format="FFMPEG",
-            executable=_FFMPEG_EXE if _FFMPEG_EXE else None,
             codec="libvpx-vp9",
             ffmpeg_params=[
                 "-pix_fmt", "yuv420p",
@@ -344,7 +347,7 @@ class EndpointHandler:
         self.pipe = None
         self.init_error: Optional[str] = None
 
-        print("=== CUSTOM handler.py LOADED (updated frames shape handling) ===", flush=True)
+        print("=== CUSTOM handler.py LOADED (webm uses IMAGEIO_FFMPEG_EXE only) ===", flush=True)
         print(f"=== HF toolkit patch diag: {HF_TOOLKIT_PATCH_DIAG} ===", flush=True)
         print(f"=== imageio-ffmpeg exe: {_FFMPEG_EXE} ===", flush=True)
 
@@ -488,24 +491,21 @@ class EndpointHandler:
 
         frames: List[np.ndarray] = []
 
-        # 1) output.frames — can be list or ndarray/tensor-like
+        # 1) output.frames — may be list OR ndarray/tensor-like
         if hasattr(output, "frames") and getattr(output, "frames") is not None:
             frames_raw = getattr(output, "frames")
             arr = np.array(frames_raw)
 
-            # Common: (T, H, W, C)
+            # (T,H,W,C)
             if arr.ndim == 4:
                 frames = [arr[t] for t in range(arr.shape[0])]
-
-            # Sometimes: (B, T, H, W, C)
+            # (B,T,H,W,C)
             elif arr.ndim == 5:
                 arr = arr[0]
                 frames = [arr[t] for t in range(arr.shape[0])]
-
-            # Rare: list-of-frames already
+            # list-of-frames
             elif isinstance(frames_raw, list):
                 frames = [np.array(f) for f in frames_raw]
-
             else:
                 raise ValueError(f"Unexpected frames shape: {arr.shape}")
 
@@ -522,12 +522,12 @@ class EndpointHandler:
             except Exception:
                 arr = np.array(vids)
 
-            # shapes: (B,T,C,H,W) or (B,T,H,W,C) or (T,C,H,W) or (T,H,W,C)
+            # (B,T,C,H,W) or (B,T,H,W,C) or (T,C,H,W) or (T,H,W,C)
             if arr.ndim == 5:
                 arr = arr[0]
 
+            # (T,C,H,W) -> (T,H,W,C)
             if arr.ndim == 4 and arr.shape[1] in (1, 3, 4):
-                # (T,C,H,W) -> (T,H,W,C)
                 arr = np.transpose(arr, (0, 2, 3, 1))
 
             if arr.ndim != 4: