diffusion: fix wan-2.2-TI2V and support sp (#12926)

This commit is contained in:
Mick
2025-11-10 14:37:57 +08:00
committed by GitHub
parent 90401cf7d2
commit e123648b36
12 changed files with 296 additions and 89 deletions
@@ -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 = []
+37 -20
View File
@@ -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",