diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index cd353ffbc..96fa94386 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -19,6 +19,7 @@ from typing import Any, Callable, List, Optional, Sequence, Union import imageio import numpy as np import torch +from PIL import Image try: import scipy.io.wavfile as scipy_wavfile @@ -643,11 +644,30 @@ def post_process_sample( indexed_path = f"{parts[0]}_{i}.{parts[1]}" else: indexed_path = f"{save_file_path}_{i}" - imageio.imwrite(indexed_path, image, quality=quality) + _save_image_frame( + indexed_path, image, quality, output_compression + ) else: - imageio.imwrite(save_file_path, frames[0], quality=quality) + _save_image_frame( + save_file_path, frames[0], quality, output_compression + ) logger.info(f"Output saved to {CYAN}{save_file_path}{RESET}") else: logger.info(f"No output path provided, output not saved") return frames + + +def _save_image_frame( + path: str, frame: np.ndarray, quality: int | None, output_compression: int | None +) -> None: + ext = os.path.splitext(path)[1].lower() + if ext == ".png": + compress_level = 1 + if output_compression is not None and output_compression != 75: + compress_level = max(0, min(9, round(output_compression / 100 * 9))) + if frame.ndim == 3 and frame.shape[-1] == 1: + frame = frame[..., 0] + Image.fromarray(frame).save(path, format="PNG", compress_level=compress_level) + else: + imageio.imwrite(path, frame, quality=quality) diff --git a/python/sglang/multimodal_gen/test/unit/test_output_saving.py b/python/sglang/multimodal_gen/test/unit/test_output_saving.py new file mode 100644 index 000000000..c3113331b --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_output_saving.py @@ -0,0 +1,68 @@ +import numpy as np +import pytest +from PIL import Image + +import sglang.multimodal_gen.runtime.entrypoints.utils as output_utils +from sglang.multimodal_gen.configs.sample.sampling_params import DataType +from sglang.multimodal_gen.runtime.entrypoints.utils import post_process_sample + + +def _rgb_frame() -> np.ndarray: + return np.array( + [ + [[0, 32, 255], [64, 128, 192], [255, 224, 16]], + [[9, 17, 33], [127, 128, 129], [240, 12, 88]], + ], + dtype=np.uint8, + ) + + +@pytest.mark.parametrize("output_compression", [None, 0, 75]) +def test_png_output_saving_preserves_pixels(tmp_path, output_compression): + frame = _rgb_frame() + output_path = tmp_path / f"sample_{output_compression}.png" + + frames = post_process_sample( + frame, + DataType.IMAGE, + fps=1, + save_file_path=str(output_path), + output_compression=output_compression, + ) + + assert output_path.exists() + np.testing.assert_array_equal(frames[0], frame) + np.testing.assert_array_equal(np.array(Image.open(output_path)), frame) + + +@pytest.mark.parametrize( + ("output_compression", "expected_compress_level"), [(None, 1), (0, 0), (75, 1)] +) +def test_png_output_saving_uses_fast_pillow_path( + tmp_path, monkeypatch, output_compression, expected_compress_level +): + frame = _rgb_frame() + output_path = tmp_path / f"sample_{output_compression}.png" + + def fail_imageio_imwrite(*args, **kwargs): + raise AssertionError("PNG output should use Pillow's PNG fast path") + + original_save = Image.Image.save + save_calls = [] + + def save_spy(self, fp, format=None, **params): + save_calls.append((format, params.get("compress_level"))) + return original_save(self, fp, format=format, **params) + + monkeypatch.setattr(output_utils.imageio, "imwrite", fail_imageio_imwrite) + monkeypatch.setattr(Image.Image, "save", save_spy) + + post_process_sample( + frame, + DataType.IMAGE, + fps=1, + save_file_path=str(output_path), + output_compression=output_compression, + ) + + assert save_calls == [("PNG", expected_compress_level)]