[diffusion] feat: speed up png image output saving (#26947)

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)

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)]