WIP: initial multimodal-gen support (#12484)

Co-authored-by: yhyang201 <yhyang201@gmail.com>
Co-authored-by: yizhang2077 <1109276519@qq.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: JiLi <leege233@gmail.com>
Co-authored-by: CHEN Xi <78632976+RubiaCx@users.noreply.github.com>
Co-authored-by: laixin <xielx@shanghaitech.edu.cn>
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
Co-authored-by: jzhang38 <a1286225768@gmail.com>
Co-authored-by: BrianChen1129 <yongqichcd@gmail.com>
Co-authored-by: Kevin Lin <42618777+kevin314@users.noreply.github.com>
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
Co-authored-by: rlsu9 <r3su@ucsd.edu>
Co-authored-by: Jinzhe Pan <48981407+eigensystem@users.noreply.github.com>
Co-authored-by: foreverpiano <pianoqwz@qq.com>
Co-authored-by: RandNMR73 <notomatthew31@gmail.com>
Co-authored-by: PorridgeSwim <yz3883@columbia.edu>
Co-authored-by: Jiali Chen <90408393+gary-chenjl@users.noreply.github.com>

python/sglang/multimodal_gen/test/__init__.py

@@ -0,0 +1 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

python/sglang/multimodal_gen/test/cli/test_generate_common.py

@@ -0,0 +1,105 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+"""
+    Common generate cli test, one test for image and video each
+"""
+
+import os
+import unittest
+from pathlib import Path
+
+from PIL import Image
+
+from sglang.multimodal_gen.test.test_utils import (
+    TestCLIBase,
+    check_image_size,
+    is_mp4,
+    run_command,
+)
+
+
+class TestGenerate(TestCLIBase):
+    model_path = "black-forest-labs/FLUX.1-dev"
+    launch_file_name = "launch_flux.json"
+    output_name = "FLUX.1-dev, single gpu"
+    ext = "jpg"
+
+    def test_generate_with_config(self):
+        test_dir = Path(__file__).parent
+        config_path = (
+            (test_dir / ".." / "test_files" / self.launch_file_name)
+            .resolve()
+            .as_posix()
+        )
+        command = [
+            "sgl_diffusion",
+            "generate",
+            f"--config={config_path}",
+        ]
+        duration = run_command(command)
+
+        self.assertIsNotNone(duration, f"Run command failed: {command}")
+
+        # verify
+        self.verify_image(self.output_name)
+
+    def test_generate_multiple_outputs(self):
+        command = [
+            "sglang",
+            "generate",
+            "--prompt='A curious raccoon'",
+            "--output-path=outputs",
+            f"--model-path={self.model_path}",
+            "--save-output",
+            f"--output-file-name={self.output_name}",
+            "--num-outputs-per-prompt=2",
+            "--width=720",
+            "--height=720",
+        ]
+        duration = run_command(command)
+        self.assertIsNotNone(duration, f"Run command failed: {command}")
+
+        self.verify_image(f"{self.output_name}_0.{self.ext}")
+        self.verify_image(f"{self.output_name}_1.{self.ext}")
+
+    def verify_image(self, output_name):
+        path = os.path.join("outputs", output_name)
+        with Image.open(path) as image:
+            check_image_size(self, image, 720, 720)
+
+    def verify_video(self, output_name):
+        path = os.path.join("outputs", output_name)
+        with open(path, "rb") as f:
+            header = f.read(12)
+            assert is_mp4(header)
+
+
+class TestWanGenerate(TestGenerate):
+    model_path = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    launch_file_name = "launch_wan.json"
+    output_name = "Wan2.1-T2V-1.3B-Diffusers, single gpu"
+    ext = "mp4"
+
+    def test_generate_multiple_outputs(self):
+        command = [
+            "sglang",
+            "generate",
+            "--prompt='A curious raccoon'",
+            "--output-path=outputs",
+            f"--model-path={self.model_path}",
+            "--save-output",
+            f"--output-file-name={self.output_name}",
+            "--num-outputs-per-prompt=2",
+            "--width=720",
+            "--height=720",
+        ]
+        duration = run_command(command)
+        self.assertIsNotNone(duration, f"Run command failed: {command}")
+
+        self.verify_video(f"{self.output_name}_0.{self.ext}")
+        # FIXME: second video is a meaningless output
+        self.verify_video(f"{self.output_name}_1.{self.ext}")
+
+
+if __name__ == "__main__":
+    unittest.main()

python/sglang/multimodal_gen/test/cli/test_generate_t2i_perf.py

@@ -0,0 +1,70 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import unittest
+from pathlib import Path
+
+from sglang.multimodal_gen.configs.sample.base import DataType
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.test.test_utils import TestGenerateBase
+
+logger = init_logger(__name__)
+
+
+class TestFlux_T2V(TestGenerateBase):
+    model_path = "black-forest-labs/FLUX.1-dev"
+    extra_args = []
+    data_type: DataType = DataType.IMAGE
+    thresholds = {
+        "test_single_gpu": 6.90 * 1.05,
+    }
+
+
+class TestQwenImage(TestGenerateBase):
+    model_path = "Qwen/Qwen-Image"
+    extra_args = []
+    data_type: DataType = DataType.IMAGE
+    thresholds = {
+        "test_single_gpu": 11.7 * 1.05,
+    }
+
+
+class TestQwenImageEdit(TestGenerateBase):
+    model_path = "Qwen/Qwen-Image-Edit"
+    extra_args = []
+    data_type: DataType = DataType.IMAGE
+    thresholds = {
+        "test_single_gpu": 43.5 * 1.05,
+    }
+
+    prompt: str | None = (
+        "Change the rabbit's color to purple, with a flash light background."
+    )
+
+    def setUp(self):
+        test_dir = Path(__file__).parent
+        img_path = (test_dir / ".." / "test_files" / "rabbit.jpg").resolve().as_posix()
+        self.base_command = [
+            "sglang",
+            "generate",
+            "--text-encoder-cpu-offload",
+            "--pin-cpu-memory",
+            f"--prompt='{self.prompt}'",
+            "--save-output",
+            "--log-level=debug",
+            f"--width={self.width}",
+            f"--height={self.height}",
+            f"--output-path={self.output_path}",
+        ] + [f"--image-path={img_path}"]
+
+    def test_single_gpu(self):
+        self._run_test(
+            name=f"{self.model_name()}, single gpu",
+            args=None,
+            model_path=self.model_path,
+            test_key="test_single_gpu",
+        )
+
+
+if __name__ == "__main__":
+    del TestGenerateBase
+    unittest.main()

python/sglang/multimodal_gen/test/cli/test_generate_t2v_perf.py

@@ -0,0 +1,68 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import unittest
+
+from sglang.multimodal_gen.configs.sample.base import DataType
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.test.test_utils import TestGenerateBase
+
+logger = init_logger(__name__)
+
+
+class TestFastWan2_1_T2V(TestGenerateBase):
+    model_path = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers"
+    extra_args = ["--attention-backend=video_sparse_attn"]
+    data_type: DataType = DataType.VIDEO
+    thresholds = {
+        "test_single_gpu": 13.0,
+        "test_cfg_parallel": 15.0,
+        "test_usp": 15.0,
+        "test_mixed": 15.0,
+    }
+
+
+class TestFastWan2_2_T2V(TestGenerateBase):
+    model_path = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
+    extra_args = []
+    data_type: DataType = DataType.VIDEO
+    thresholds = {
+        "test_single_gpu": 25.0,
+        "test_cfg_parallel": 30.0,
+        "test_usp": 30.0,
+        "test_mixed": 30.0,
+    }
+
+
+class TestWan2_1_T2V(TestGenerateBase):
+    model_path = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    extra_args = []
+    data_type: DataType = DataType.VIDEO
+    thresholds = {
+        "test_single_gpu": 76.0,
+        "test_cfg_parallel": 46.5 * 1.05,
+        "test_usp": 22.5,
+        "test_mixed": 26.5,
+    }
+
+
+class TestWan2_2_T2V(TestGenerateBase):
+    model_path = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+    extra_args = []
+    data_type: DataType = DataType.VIDEO
+    thresholds = {
+        "test_single_gpu": 865,
+        "test_cfg_parallel": 446,
+        "test_usp": 124,
+        "test_mixed": 159,
+    }
+
+    def test_mixed(self):
+        pass
+
+    def test_cfg_parallel(self):
+        pass
+
+
+if __name__ == "__main__":
+    del TestGenerateBase
+    unittest.main()

python/sglang/multimodal_gen/test/cli/test_generate_ti2v_perf.py

@@ -0,0 +1,62 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import unittest
+
+from sglang.multimodal_gen.configs.sample.base import DataType
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.test.test_utils import TestGenerateBase
+
+logger = init_logger(__name__)
+
+
+class TestGenerateTI2VBase(TestGenerateBase):
+    data_type: DataType = DataType.VIDEO
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_command = [
+            "sglang",
+            "generate",
+            f'--prompt="Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline\'s intricate details and the refreshing atmosphere of the seaside."',
+            "--image-path=https://github.com/Wan-Video/Wan2.2/blob/990af50de458c19590c245151197326e208d7191/examples/i2v_input.JPG?raw=true",
+            "--save-output",
+            "--log-level=debug",
+            f"--output-path={cls.output_path}",
+        ] + cls.extra_args
+
+    def test_single_gpu(self):
+        pass
+
+    def test_cfg_parallel(self):
+        pass
+
+    def test_mixed(self):
+        pass
+
+
+class TestWan2_1_I2V_14B_480P(TestGenerateTI2VBase):
+    model_path = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+    extra_args = ["--attention-backend=video_sparse_attn"]
+    thresholds = {
+        "test_single_gpu": 13.0,
+        "test_cfg_parallel": 191.7 * 1.05,
+        "test_usp": 15.0,
+        "test_mixed": 15.0,
+    }
+
+
+class TestWan2_2_TI2V_5B(TestGenerateTI2VBase):
+    model_path = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+    # FIXME: doesn't work with vsa at the moment
+    # extra_args = ["--attention-backend=video_sparse_attn"]
+    thresholds = {
+        "test_single_gpu": 13.0,
+        "test_cfg_parallel": 191.7 * 1.05,
+        "test_usp": 387.6 * 1.05,
+        "test_mixed": 15.0,
+    }
+
+
+if __name__ == "__main__":
+    del TestGenerateTI2VBase, TestGenerateBase
+    unittest.main()

python/sglang/multimodal_gen/test/cli/test_serve.py

@@ -0,0 +1,287 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+import asyncio
+import base64
+import subprocess
+import time
+import unittest
+from pathlib import Path
+
+from openai import OpenAI
+
+from sglang.multimodal_gen.runtime.utils.common import kill_process_tree
+from sglang.multimodal_gen.test.test_utils import is_mp4, is_png, wait_for_port
+
+
+def wait_for_video_completion(client, video_id, timeout=300, check_interval=3):
+    start = time.time()
+    video = client.videos.retrieve(video_id)
+
+    while video.status not in ("completed", "failed"):
+        time.sleep(check_interval)
+        video = client.videos.retrieve(video_id)
+        assert time.time() - start < timeout, "video generate timeout"
+
+    return video
+
+
+class TestVideoHttpServer(unittest.TestCase):
+    model_name = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    timeout = 120
+    extra_args = []
+
+    def _create_wait_and_download(
+        self, client: OpenAI, prompt: str, size: str
+    ) -> bytes:
+
+        video = client.videos.create(prompt=prompt, size=size)
+        video_id = video.id
+        self.assertEqual(video.status, "queued")
+
+        video = wait_for_video_completion(client, video_id, timeout=self.timeout)
+        self.assertEqual(video.status, "completed", "video generate failed")
+
+        response = client.videos.download_content(
+            video_id=video_id,
+        )
+        content = response.read()
+        return content
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_command = [
+            "sglang",
+            "serve",
+            "--model-path",
+            f"{cls.model_name}",
+            "--port",
+            "30010",
+        ]
+
+        process = subprocess.Popen(
+            cls.base_command + cls.extra_args,
+            # stdout=subprocess.PIPE,
+            # stderr=subprocess.PIPE,
+            text=True,
+            bufsize=1,
+        )
+        cls.pid = process.pid
+        wait_for_port(host="127.0.0.1", port=30010)
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.pid)
+
+    def test_http_server_basic(self):
+        client = OpenAI(
+            api_key="sk-proj-1234567890", base_url="http://localhost:30010/v1"
+        )
+        content = self._create_wait_and_download(
+            client, "A calico cat playing a piano on stage", "832x480"
+        )
+        self.assertTrue(is_mp4(content))
+
+    def test_concurrent_requests(self):
+        client = OpenAI(
+            api_key="sk-proj-1234567890", base_url="http://localhost:30010/v1"
+        )
+
+        num_requests = 2
+
+        async def generate_and_check_video(prompt, size):
+            content = await asyncio.to_thread(
+                self._create_wait_and_download, client, prompt, size
+            )
+            self.assertTrue(is_mp4(content))
+
+        async def send_concurrent_requests():
+            tasks = [
+                generate_and_check_video(
+                    "A dog playing a piano on stage",
+                    "832x480",
+                )
+                for _ in range(num_requests)
+            ]
+            await asyncio.gather(*tasks)
+
+        asyncio.run(send_concurrent_requests())
+
+
+class TestFastWan2_1HttpServer(TestVideoHttpServer):
+    model_name = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers"
+
+
+class TestFastWan2_2HttpServer(TestVideoHttpServer):
+    model_name = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
+
+
+class TestImage2VideoHttpServer(unittest.TestCase):
+    model_name = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
+    timeout = 1200
+    extra_args = []
+
+    def _create_wait_and_download(
+        self, client: OpenAI, prompt: str, size: str
+    ) -> bytes:
+
+        image_path = "https://github.com/Wan-Video/Wan2.2/blob/990af50de458c19590c245151197326e208d7191/examples/i2v_input.JPG?raw=true"
+        image_path = Path(image_path)
+        video = client.videos.create(
+            prompt=prompt,
+            input_reference=image_path,
+            size=size,
+            seconds=10,
+            extra_body={"fps": 16, "num_frames": 125},
+        )
+        # TODO: Some combinations of num_frames and fps may cause errors and need further investigation.
+        video_id = video.id
+        self.assertEqual(video.status, "queued")
+
+        video = wait_for_video_completion(client, video_id, timeout=self.timeout)
+        self.assertEqual(video.status, "completed", "video generate failed")
+
+        response = client.videos.download_content(
+            video_id=video_id,
+        )
+        content = response.read()
+        return content
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_command = [
+            "sgl-diffusion",
+            "serve",
+            "--model-path",
+            f"{cls.model_name}",
+            "--port",
+            "30010",
+        ]
+
+        process = subprocess.Popen(
+            cls.base_command + cls.extra_args,
+            # stdout=subprocess.PIPE,
+            # stderr=subprocess.PIPE,
+            text=True,
+            bufsize=1,
+        )
+        cls.pid = process.pid
+        wait_for_port(host="127.0.0.1", port=30010)
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.pid)
+
+    def test_http_server_basic(self):
+        client = OpenAI(
+            api_key="sk-proj-1234567890", base_url="http://localhost:30010/v1"
+        )
+        content = self._create_wait_and_download(
+            client, "A girl is fighting a monster.", "832x480"
+        )
+        self.assertTrue(is_mp4(content))
+
+    def test_concurrent_requests(self):
+        client = OpenAI(
+            api_key="sk-proj-1234567890", base_url="http://localhost:30010/v1"
+        )
+
+        num_requests = 2
+
+        async def generate_and_check_video(prompt, size):
+            content = await asyncio.to_thread(
+                self._create_wait_and_download, client, prompt, size
+            )
+            self.assertTrue(is_mp4(content))
+
+        async def send_concurrent_requests():
+            tasks = [
+                generate_and_check_video(
+                    "A dog playing a piano on stage",
+                    "832x480",
+                )
+                for _ in range(num_requests)
+            ]
+            await asyncio.gather(*tasks)
+
+        asyncio.run(send_concurrent_requests())
+
+
+class TestImageHttpServer(unittest.TestCase):
+
+    def _create_wait_and_download(
+        self, client: OpenAI, prompt: str, size: str
+    ) -> bytes:
+        img = client.images.generate(
+            model="gpt-image-1",
+            prompt=prompt,
+            n=1,
+            size=size,
+            response_format="b64_json",
+            output_format="png",
+        )
+        image_bytes = base64.b64decode(img.data[0].b64_json)
+        return image_bytes
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_command = [
+            "sglang",
+            "serve",
+            "--model-path",
+            "Qwen/Qwen-Image",
+            "--port",
+            "30020",
+        ]
+
+        process = subprocess.Popen(
+            cls.base_command,
+            # stdout=subprocess.PIPE,
+            # stderr=subprocess.PIPE,
+            text=True,
+            bufsize=1,
+        )
+        cls.pid = process.pid
+        wait_for_port(host="127.0.0.1", port=30020)
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.pid)
+
+    def test_http_server_basic(self):
+        client = OpenAI(
+            api_key="sk-proj-1234567890", base_url="http://localhost:30020/v1"
+        )
+        content = self._create_wait_and_download(
+            client, "A calico cat playing a piano on stage", "832x480"
+        )
+        self.assertTrue(is_png(content))
+
+    def test_concurrent_requests(self):
+        client = OpenAI(
+            api_key="sk-proj-1234567890", base_url="http://localhost:30020/v1"
+        )
+
+        num_requests = 2
+
+        async def generate_and_check_image(prompt, size):
+            content = await asyncio.to_thread(
+                self._create_wait_and_download, client, prompt, size
+            )
+            self.assertTrue(is_png(content))
+
+        async def send_concurrent_requests():
+            tasks = [
+                generate_and_check_image(
+                    "A dog playing a piano on stage",
+                    "832x480",
+                )
+                for _ in range(num_requests)
+            ]
+            await asyncio.gather(*tasks)
+
+        asyncio.run(send_concurrent_requests())
+
+
+if __name__ == "__main__":
+    # del TestPerform·anceBase
+    unittest.main()

python/sglang/multimodal_gen/test/test_files/launch_flux.json

@@ -0,0 +1,11 @@
+{
+    "model_path": "black-forest-labs/FLUX.1-dev",
+    "prompt": "A beautiful woman in a red dress walking down a street",
+    "text_encoder_cpu_offload": true,
+    "pin_cpu_memory": true,
+    "save_output": true,
+    "width": 720,
+    "height": 720,
+    "output_path": "outputs",
+    "output_file_name": "FLUX.1-dev, single gpu"
+}

python/sglang/multimodal_gen/test/test_files/launch_wan.json

@@ -0,0 +1,11 @@
+{
+    "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    "prompt": "A beautiful woman in a red dress walking down a street",
+    "text_encoder_cpu_offload": true,
+    "pin_cpu_memory": true,
+    "save_output": true,
+    "width": 720,
+    "height": 720,
+    "output_path": "outputs",
+    "output_file_name": "Wan2.1-T2V-1.3B-Diffusers, single gpu"
+}

python/sglang/multimodal_gen/test/test_offline_api.py

@@ -0,0 +1,75 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+"""
+    Testing the performance of generate command of sgl_diffusion' CLI
+"""
+
+import unittest
+
+import torch
+
+from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class TestGeneratorAPIBase(unittest.TestCase):
+    # server args
+    server_kwargs = {}
+
+    # sampling
+    output_path: str = "outputs"
+
+    results = []
+
+    @classmethod
+    def setUpClass(cls):
+        cls.results = []
+
+    def verify_single_generation_result(self, result):
+        self.assertIsNotNone(result, "Generation failed")
+        self.assertTrue(
+            "samples" in result and isinstance(result["samples"], torch.Tensor),
+            f"Incorrect Generation result",
+        )
+
+    def _run_test(self, name, server_kwargs, test_key: str):
+        generator = DiffGenerator.from_pretrained(**server_kwargs)
+        result = generator.generate(prompt="A curious raccoon")
+        self.verify_single_generation_result(result)
+
+    def test_single_gpu(self):
+        self._run_test(
+            name=self.server_kwargs["model_path"],
+            server_kwargs=self.server_kwargs | dict(num_gpus=1),
+            test_key="test_single_gpu",
+        )
+
+    def test_cfg_parallel(self):
+        self._run_test(
+            name=self.server_kwargs["model_path"],
+            server_kwargs=self.server_kwargs
+            | dict(num_gpus=2, enable_cfg_parallel=True),
+            test_key="test_cfg_parallel",
+        )
+
+    def test_multiple_prompts(self):
+        generator = DiffGenerator.from_pretrained(
+            **self.server_kwargs | dict(num_gpus=2, enable_cfg_parallel=True)
+        )
+        prompts = ["A curious raccoon", "A curious cat"]
+        results = generator.generate(prompt=prompts)
+
+        self.assertEqual(len(results), len(prompts), "Some generation tasks fail")
+        for result in results:
+            self.verify_single_generation_result(result)
+
+
+class TestWan2_1_T2V(TestGeneratorAPIBase):
+    server_kwargs = {"model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"}
+
+
+if __name__ == "__main__":
+    del TestGeneratorAPIBase
+    unittest.main()

python/sglang/multimodal_gen/test/test_utils.py

@@ -0,0 +1,260 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+import os
+import shlex
+import socket
+import subprocess
+import sys
+import time
+import unittest
+
+from PIL import Image
+
+from sglang.multimodal_gen.configs.sample.base import DataType
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def run_command(command):
+    """Runs a command and returns the execution time and status."""
+    print(f"Running command: {' '.join(command)}")
+
+    duration = None
+    with subprocess.Popen(
+        command,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        encoding="utf-8",
+    ) as process:
+        for line in process.stdout:
+            sys.stdout.write(line)
+            if "Pixel data generated" in line:
+                words = line.split(" ")
+                duration = float(words[-2])
+
+    if process.returncode == 0:
+        return duration
+    else:
+        print(f"Command failed with exit code {process.returncode}")
+        return None
+
+
+def probe_port(host="127.0.0.1", port=30010, timeout=2.0) -> bool:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.settimeout(timeout)
+        try:
+            s.connect((host, port))
+            return True
+        except OSError:
+            return False
+
+
+def is_mp4(data):
+    idx = data.find(b"ftyp")
+    return 0 <= idx <= 32
+
+
+def is_png(data):
+    # PNG files start with: 89 50 4E 47 0D 0A 1A 0A
+    return data.startswith(b"\x89PNG\r\n\x1a\n")
+
+
+def wait_for_port(host="127.0.0.1", port=30010, deadline=300.0, interval=0.5):
+    end = time.time() + deadline
+    last_err = None
+    while time.time() < end:
+        if probe_port(host, port, timeout=interval):
+            return True
+        time.sleep(interval)
+    raise TimeoutError(f"Port {host}:{port} not ready. Last error: {last_err}")
+
+
+def check_image_size(ut, image, width, height):
+    # check image size
+    ut.assertEqual(image.size, (width, height))
+
+
+class TestCLIBase(unittest.TestCase):
+    model_path: str = None
+    extra_args = []
+    data_type: DataType = None
+    # tested on h100
+    thresholds = {}
+
+    width: int = 720
+    height: int = 720
+    output_path: str = "outputs"
+
+    base_command = [
+        "sglang",
+        "generate",
+        "--text-encoder-cpu-offload",
+        "--pin-cpu-memory",
+        "--prompt='A curious raccoon'",
+        "--save-output",
+        "--log-level=debug",
+        f"--width={width}",
+        f"--height={height}",
+        f"--output-path={output_path}",
+    ]
+
+    results = []
+
+    @classmethod
+    def setUpClass(cls):
+        cls.results = []
+
+    def _run_command(self, name, model_path: str, test_key: str = "", args=[]):
+        command = (
+            self.base_command
+            + [f"--model-path={model_path}"]
+            + shlex.split(args or "")
+            + [f"--output-file-name={name}"]
+            + self.extra_args
+        )
+        duration = run_command(command)
+        status = "Success" if duration else "Failed"
+
+        duration_str = f"{duration:.4f}s" if duration else "NA"
+        self.__class__.results.append(
+            {"name": name, "key": test_key, "duration": duration_str, "status": status}
+        )
+
+        return name, duration, status
+
+
+class TestGenerateBase(TestCLIBase):
+    model_path: str = None
+    extra_args = []
+    data_type: DataType = None
+    # tested on h100
+    thresholds = {}
+
+    width: int = 720
+    height: int = 720
+    output_path: str = "outputs"
+    image_path: str | None = None
+    prompt: str | None = "A curious raccoon"
+
+    base_command = [
+        "sglang",
+        "generate",
+        # "--text-encoder-cpu-offload",
+        # "--pin-cpu-memory",
+        f"--prompt='{prompt}'",
+        "--save-output",
+        "--log-level=debug",
+        f"--width={width}",
+        f"--height={height}",
+        f"--output-path={output_path}",
+    ]
+
+    results = []
+
+    @classmethod
+    def setUpClass(cls):
+        cls.results = []
+
+    @classmethod
+    def tearDownClass(cls):
+        # Print markdown table
+        print("\n## Test Results\n")
+        print("| Test Case                      | Duration | Status  |")
+        print("|--------------------------------|----------|---------|")
+        test_keys = ["test_single_gpu", "test_cfg_parallel", "test_usp", "test_mixed"]
+        test_key_to_order = {
+            test_key: order for order, test_key in enumerate(test_keys)
+        }
+
+        ordered_results: list[dict] = [{}] * len(test_keys)
+
+        for result in cls.results:
+            order = test_key_to_order[result["key"]]
+            ordered_results[order] = result
+
+        for result in ordered_results:
+            if not result:
+                continue
+            status = (
+                result["status"] and result["duration"] <= cls.thresholds[result["key"]]
+            )
+            print(f"| {result['name']:<30} | {result['duration']:<8} | {status:<7} |")
+        print()
+        durations = [result["duration"] for result in cls.results]
+        print(" | ".join([""] + durations + [""]))
+
+    def _run_test(self, name, args, model_path: str, test_key: str):
+        time_threshold = self.thresholds[test_key]
+        name, duration, status = self._run_command(
+            name, args=args, model_path=model_path, test_key=test_key
+        )
+        self.verify(status, name, duration, time_threshold)
+
+    def verify(self, status, name, duration, time_threshold):
+        print("-" * 80)
+        print("\n" * 3)
+
+        # test task status
+        self.assertEqual(status, "Success", f"{name} command failed")
+        self.assertIsNotNone(duration, f"Could not parse duration for {name}")
+        self.assertLessEqual(
+            duration,
+            time_threshold,
+            f"{name} failed with {duration:.4f}s > {time_threshold}s",
+        )
+
+        # test output file
+        path = os.path.join(
+            self.output_path, f"{name}.{self.data_type.get_default_extension()}"
+        )
+        self.assertTrue(os.path.exists(path), f"Output file not exist for {path}")
+        if self.data_type == DataType.IMAGE:
+            with Image.open(path) as image:
+                check_image_size(self, image, self.width, self.height)
+        logger.info(f"{name} passed in {duration:.4f}s (threshold: {time_threshold}s)")
+
+    def model_name(self):
+        return self.model_path.split("/")[-1]
+
+    def test_single_gpu(self):
+        """single gpu"""
+        self._run_test(
+            name=f"{self.model_name()}, single gpu",
+            args=None,
+            model_path=self.model_path,
+            test_key="test_single_gpu",
+        )
+
+    def test_cfg_parallel(self):
+        """cfg parallel"""
+        if self.data_type == DataType.IMAGE:
+            return
+        self._run_test(
+            name=f"{self.model_name()}, cfg parallel",
+            args="--num-gpus 2 --enable-cfg-parallel",
+            model_path=self.model_path,
+            test_key="test_cfg_parallel",
+        )
+
+    def test_usp(self):
+        """usp"""
+        if self.data_type == DataType.IMAGE:
+            return
+        self._run_test(
+            name=f"{self.model_name()}, usp",
+            args="--num-gpus 4 --ulysses-degree=2 --ring-degree=2",
+            model_path=self.model_path,
+            test_key="test_usp",
+        )
+
+    def test_mixed(self):
+        """mixed"""
+        if self.data_type == DataType.IMAGE:
+            return
+        self._run_test(
+            name=f"{self.model_name()}, mixed",
+            args="--num-gpus 4 --ulysses-degree=2 --ring-degree=1 --enable-cfg-parallel",
+            model_path=self.model_path,
+            test_key="test_mixed",
+        )

python/sglang/multimodal_gen/test/utils.py

@@ -0,0 +1,162 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import json
+import os
+
+import numpy as np
+import torch
+from pytorch_msssim import ms_ssim, ssim
+from torchvision.io import read_video
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
+    """
+    Compute SSIM between two videos.
+
+    Args:
+        video1_path: Path to the first video.
+        video2_path: Path to the second video.
+        use_ms_ssim: Whether to use Multi-Scale Structural Similarity(MS-SSIM) instead of SSIM.
+    """
+    print(f"Computing SSIM between {video1_path} and {video2_path}...")
+    if not os.path.exists(video1_path):
+        raise FileNotFoundError(f"Video1 not found: {video1_path}")
+    if not os.path.exists(video2_path):
+        raise FileNotFoundError(f"Video2 not found: {video2_path}")
+
+    frames1, _, _ = read_video(video1_path, pts_unit="sec", output_format="TCHW")
+    frames2, _, _ = read_video(video2_path, pts_unit="sec", output_format="TCHW")
+
+    # Ensure same number of frames
+    min_frames = min(frames1.shape[0], frames2.shape[0])
+    frames1 = frames1[:min_frames]
+    frames2 = frames2[:min_frames]
+
+    frames1 = frames1.float() / 255.0
+    frames2 = frames2.float() / 255.0
+
+    if torch.cuda.is_available():
+        frames1 = frames1.cuda()
+        frames2 = frames2.cuda()
+
+    ssim_values = []
+
+    # Process each frame individually
+    for i in range(min_frames):
+        img1 = frames1[i : i + 1]
+        img2 = frames2[i : i + 1]
+
+        with torch.no_grad():
+            if use_ms_ssim:
+                value = ms_ssim(img1, img2, data_range=1.0)
+            else:
+                value = ssim(img1, img2, data_range=1.0)
+
+            ssim_values.append(value.item())
+
+    if ssim_values:
+        mean_ssim = np.mean(ssim_values)
+        min_ssim = np.min(ssim_values)
+        max_ssim = np.max(ssim_values)
+        min_frame_idx = np.argmin(ssim_values)
+        max_frame_idx = np.argmax(ssim_values)
+
+        print(f"Mean SSIM: {mean_ssim:.4f}")
+        print(f"Min SSIM: {min_ssim:.4f} (at frame {min_frame_idx})")
+        print(f"Max SSIM: {max_ssim:.4f} (at frame {max_frame_idx})")
+
+        return mean_ssim, min_ssim, max_ssim
+    else:
+        print("No SSIM values calculated")
+        return 0, 0, 0
+
+
+def compare_folders(reference_folder, generated_folder, use_ms_ssim=True):
+    """
+    Compare videos with the same filename between reference_folder and generated_folder
+
+    Example usage:
+        results = compare_folders(reference_folder, generated_folder,
+                              args.use_ms_ssim)
+        for video_name, ssim_value in results.items():
+            if ssim_value is not None:
+                print(
+                    f"{video_name}: {ssim_value[0]:.4f}, Min SSIM: {ssim_value[1]:.4f}, Max SSIM: {ssim_value[2]:.4f}"
+                )
+            else:
+                print(f"{video_name}: Error during comparison")
+
+        valid_ssims = [v for v in results.values() if v is not None]
+        if valid_ssims:
+            avg_ssim = np.mean([v[0] for v in valid_ssims])
+            print(f"\nAverage SSIM across all videos: {avg_ssim:.4f}")
+        else:
+            print("\nNo valid SSIM values to average")
+    """
+
+    reference_videos = [f for f in os.listdir(reference_folder) if f.endswith(".mp4")]
+
+    results = {}
+
+    for video_name in reference_videos:
+        ref_path = os.path.join(reference_folder, video_name)
+        gen_path = os.path.join(generated_folder, video_name)
+
+        if os.path.exists(gen_path):
+            print(f"\nComparing {video_name}...")
+            try:
+                ssim_value = compute_video_ssim_torchvision(
+                    ref_path, gen_path, use_ms_ssim
+                )
+                results[video_name] = ssim_value
+            except Exception as e:
+                print(f"Error comparing {video_name}: {e}")
+                results[video_name] = None
+        else:
+            print(f"\nSkipping {video_name} - no matching file in generated folder")
+
+    return results
+
+
+def write_ssim_results(
+    output_dir, ssim_values, reference_path, generated_path, num_inference_steps, prompt
+):
+    """
+    Write SSIM results to a JSON file in the same directory as the generated videos.
+    """
+    try:
+        logger.info(f"Attempting to write SSIM results to directory: {output_dir}")
+
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir, exist_ok=True)
+
+        mean_ssim, min_ssim, max_ssim = ssim_values
+
+        result = {
+            "mean_ssim": mean_ssim,
+            "min_ssim": min_ssim,
+            "max_ssim": max_ssim,
+            "reference_video": reference_path,
+            "generated_video": generated_path,
+            "parameters": {
+                "num_inference_steps": num_inference_steps,
+                "prompt": prompt,
+            },
+        }
+
+        test_name = f"steps{num_inference_steps}_{prompt[:100]}"
+        result_file = os.path.join(output_dir, f"{test_name}_ssim.json")
+        logger.info(f"Writing JSON results to: {result_file}")
+        with open(result_file, "w") as f:
+            json.dump(result, f, indent=2)
+
+        logger.info(f"SSIM results written to {result_file}")
+        return True
+    except Exception as e:
+        logger.error(f"ERROR writing SSIM results: {str(e)}")
+        return False