[Bugfix] Preserve channel-last video tensor layout (#5418)
Signed-off-by: pxljs <1621352782@qq.com>
This commit is contained in:
1 parent
67aa6c55a4
commit
185e3d6df6
2 files changed
+21
-3
No files matched your search
@@ -48,6 +48,25 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch):
|
||||
assert mux_calls[0]["audio"] is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("frame_count", [3, 4])
|
||||
def test_channel_last_video_tensor_preserves_channel_sized_frame_count(frame_count):
|
||||
video = torch.arange(frame_count * 2 * 5 * 3, dtype=torch.uint8).reshape(frame_count, 2, 5, 3)
|
||||
|
||||
frames = video_api_utils._coerce_video_to_uint8_frames(video)
|
||||
|
||||
assert frames.shape == (frame_count, 2, 5, 3)
|
||||
np.testing.assert_array_equal(frames, video.numpy())
|
||||
|
||||
|
||||
def test_channel_first_video_tensor_is_converted_to_channel_last():
|
||||
video = torch.arange(3 * 5 * 2 * 6, dtype=torch.uint8).reshape(3, 5, 2, 6)
|
||||
|
||||
frames = video_api_utils._coerce_video_to_uint8_frames(video)
|
||||
|
||||
assert frames.shape == (5, 2, 6, 3)
|
||||
np.testing.assert_array_equal(frames, video.permute(1, 2, 3, 0).numpy())
|
||||
|
||||
|
||||
def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch):
|
||||
muxers = []
|
||||
|
||||
|
||||
@@ -308,9 +308,6 @@ def _normalize_video_tensor(video_tensor: torch.Tensor) -> np.ndarray:
|
||||
video_tensor = video_tensor.detach().cpu()
|
||||
if video_tensor.dim() == 5:
|
||||
raise ValueError("Batched video tensors are not supported for single-video encoding.")
|
||||
elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
|
||||
# [C, F, H, W] -> [F, H, W, C]
|
||||
video_tensor = video_tensor.permute(1, 2, 3, 0)
|
||||
|
||||
if video_tensor.is_floating_point():
|
||||
# Cast to float32 first: bf16 (e.g. SANA-WM's refiner output) has no
|
||||
@@ -329,6 +326,8 @@ def _normalize_single_video_array(video_array: np.ndarray) -> np.ndarray:
|
||||
|
||||
if video_array.ndim == 4:
|
||||
# Convert channel-first layouts to channel-last
|
||||
# Prefer an explicit channel-last dimension for ambiguous 3/4-frame
|
||||
# videos before interpreting a leading dimension as channels.
|
||||
if video_array.shape[0] in (3, 4) and video_array.shape[-1] not in (3, 4):
|
||||
video_array = np.transpose(video_array, (1, 2, 3, 0))
|
||||
elif video_array.shape[1] in (3, 4) and video_array.shape[-1] not in (3, 4):
|
||||
|
||||
Reference in new issue
Block a user