diffusion: fix wan-2.2-TI2V and support sp (#12926)
This commit is contained in:
@@ -20,6 +20,10 @@ class TestFastWan2_1_T2V(TestGenerateBase):
|
||||
"test_mixed": 15.0,
|
||||
}
|
||||
|
||||
# disabled for vsa
|
||||
def test_usp(self):
|
||||
pass
|
||||
|
||||
|
||||
class TestFastWan2_2_T2V(TestGenerateBase):
|
||||
model_path = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
|
||||
@@ -38,7 +42,7 @@ class TestWan2_1_T2V(TestGenerateBase):
|
||||
extra_args = []
|
||||
data_type: DataType = DataType.VIDEO
|
||||
thresholds = {
|
||||
"test_single_gpu": 76.0,
|
||||
"test_single_gpu": 76.0 * 1.05,
|
||||
"test_cfg_parallel": 46.5 * 1.05,
|
||||
"test_usp": 22.5,
|
||||
"test_mixed": 26.5,
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.multimodal_gen.configs.sample.base import DataType
|
||||
@@ -18,7 +16,8 @@ class TestGenerateTI2VBase(TestGenerateBase):
|
||||
"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",
|
||||
"--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}",
|
||||
@@ -36,12 +35,15 @@ class TestGenerateTI2VBase(TestGenerateBase):
|
||||
|
||||
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,
|
||||
"test_usp": 530.5 * 1.05,
|
||||
}
|
||||
|
||||
|
||||
class TestWan2_1_I2V_14B_720P(TestGenerateTI2VBase):
|
||||
model_path = "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers"
|
||||
thresholds = {
|
||||
"test_usp": 530.5 * 1.05,
|
||||
}
|
||||
|
||||
|
||||
@@ -50,13 +52,19 @@ class TestWan2_2_TI2V_5B(TestGenerateTI2VBase):
|
||||
# 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,
|
||||
"test_usp": 82.3 * 1.05,
|
||||
}
|
||||
|
||||
|
||||
# OOM
|
||||
# class TestWan2_2_I2V_A14B(TestGenerateTI2VBase):
|
||||
# model_path = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
|
||||
# # FIXME: doesn't work with vsa at the moment
|
||||
# thresholds = {
|
||||
# "test_usp": 66.3 * 1.05,
|
||||
# }
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
del TestGenerateTI2VBase, TestGenerateBase
|
||||
unittest.main()
|
||||
|
||||
@@ -19,7 +19,7 @@ class TestGeneratorAPIBase(unittest.TestCase):
|
||||
server_kwargs = {}
|
||||
|
||||
# sampling
|
||||
output_path: str = "outputs"
|
||||
output_path: str = "test_outputs"
|
||||
|
||||
results = []
|
||||
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
||||
import dataclasses
|
||||
import os
|
||||
import shlex
|
||||
import socket
|
||||
@@ -6,6 +7,7 @@ import subprocess
|
||||
import sys
|
||||
import time
|
||||
import unittest
|
||||
from typing import Optional
|
||||
|
||||
from PIL import Image
|
||||
|
||||
@@ -15,7 +17,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def run_command(command):
|
||||
def run_command(command) -> Optional[float]:
|
||||
"""Runs a command and returns the execution time and status."""
|
||||
print(f"Running command: {' '.join(command)}")
|
||||
|
||||
@@ -75,6 +77,18 @@ def check_image_size(ut, image, width, height):
|
||||
ut.assertEqual(image.size, (width, height))
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TestResult:
|
||||
name: str
|
||||
key: str
|
||||
duration: Optional[float]
|
||||
succeed: bool
|
||||
|
||||
@property
|
||||
def duration_str(self):
|
||||
return f"{self.duration:.4f}" if self.duration else "NA"
|
||||
|
||||
|
||||
class TestCLIBase(unittest.TestCase):
|
||||
model_path: str = None
|
||||
extra_args = []
|
||||
@@ -84,7 +98,7 @@ class TestCLIBase(unittest.TestCase):
|
||||
|
||||
width: int = 720
|
||||
height: int = 720
|
||||
output_path: str = "outputs"
|
||||
output_path: str = "test_outputs"
|
||||
|
||||
base_command = [
|
||||
"sglang",
|
||||
@@ -105,7 +119,7 @@ class TestCLIBase(unittest.TestCase):
|
||||
def setUpClass(cls):
|
||||
cls.results = []
|
||||
|
||||
def _run_command(self, name, model_path: str, test_key: str = "", args=[]):
|
||||
def _run_command(self, name: str, model_path: str, test_key: str = "", args=[]):
|
||||
command = (
|
||||
self.base_command
|
||||
+ [f"--model-path={model_path}"]
|
||||
@@ -115,11 +129,10 @@ class TestCLIBase(unittest.TestCase):
|
||||
)
|
||||
duration = run_command(command)
|
||||
status = "Success" if duration else "Failed"
|
||||
succeed = duration is not None
|
||||
|
||||
duration_str = f"{duration:.4f}s" if duration else "NA"
|
||||
self.__class__.results.append(
|
||||
{"name": name, "key": test_key, "duration": duration_str, "status": status}
|
||||
)
|
||||
duration = float(duration) if succeed else None
|
||||
self.results.append(TestResult(name, test_key, duration, succeed))
|
||||
|
||||
return name, duration, status
|
||||
|
||||
@@ -133,7 +146,7 @@ class TestGenerateBase(TestCLIBase):
|
||||
|
||||
width: int = 720
|
||||
height: int = 720
|
||||
output_path: str = "outputs"
|
||||
output_path: str = "test_outputs"
|
||||
image_path: str | None = None
|
||||
prompt: str | None = "A curious raccoon"
|
||||
|
||||
@@ -150,7 +163,7 @@ class TestGenerateBase(TestCLIBase):
|
||||
f"--output-path={output_path}",
|
||||
]
|
||||
|
||||
results = []
|
||||
results: list[TestResult] = []
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
@@ -167,24 +180,28 @@ class TestGenerateBase(TestCLIBase):
|
||||
test_key: order for order, test_key in enumerate(test_keys)
|
||||
}
|
||||
|
||||
ordered_results: list[dict] = [{}] * len(test_keys)
|
||||
|
||||
ordered_results: list[TestResult] = [None] * len(test_keys)
|
||||
for result in cls.results:
|
||||
order = test_key_to_order[result["key"]]
|
||||
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"]]
|
||||
"Succeed"
|
||||
if (
|
||||
result.succeed
|
||||
and float(result.duration) <= float(cls.thresholds[result.key])
|
||||
)
|
||||
else "Failed"
|
||||
)
|
||||
print(f"| {result['name']:<30} | {result['duration']:<8} | {status:<7} |")
|
||||
print(f"| {result.name:<30} | {result.duration_str:<8} | {status:<7} |")
|
||||
print()
|
||||
durations = [result["duration"] for result in cls.results]
|
||||
durations = [result.duration_str for result in cls.results]
|
||||
print(" | ".join([""] + durations + [""]))
|
||||
|
||||
def _run_test(self, name, args, model_path: str, test_key: str):
|
||||
def _run_test(self, name: str, 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
|
||||
@@ -220,7 +237,7 @@ class TestGenerateBase(TestCLIBase):
|
||||
def test_single_gpu(self):
|
||||
"""single gpu"""
|
||||
self._run_test(
|
||||
name=f"{self.model_name()}, single gpu",
|
||||
name=f"{self.model_name()}_single gpu",
|
||||
args=None,
|
||||
model_path=self.model_path,
|
||||
test_key="test_single_gpu",
|
||||
@@ -231,7 +248,7 @@ class TestGenerateBase(TestCLIBase):
|
||||
if self.data_type == DataType.IMAGE:
|
||||
return
|
||||
self._run_test(
|
||||
name=f"{self.model_name()}, cfg parallel",
|
||||
name=f"{self.model_name()}_cfg parallel",
|
||||
args="--num-gpus 2 --enable-cfg-parallel",
|
||||
model_path=self.model_path,
|
||||
test_key="test_cfg_parallel",
|
||||
@@ -242,7 +259,7 @@ class TestGenerateBase(TestCLIBase):
|
||||
if self.data_type == DataType.IMAGE:
|
||||
return
|
||||
self._run_test(
|
||||
name=f"{self.model_name()}, usp",
|
||||
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",
|
||||
@@ -253,7 +270,7 @@ class TestGenerateBase(TestCLIBase):
|
||||
if self.data_type == DataType.IMAGE:
|
||||
return
|
||||
self._run_test(
|
||||
name=f"{self.model_name()}, mixed",
|
||||
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",
|
||||
|
||||
Reference in New Issue
Block a user