[diffusion] feat: implement upscaling (#19723)
This commit is contained in:
@@ -2136,6 +2136,202 @@
|
||||
"expected_e2e_ms": 716.81,
|
||||
"expected_avg_denoise_ms": 67.32,
|
||||
"expected_median_denoise_ms": 74.1
|
||||
},
|
||||
"flux_2_image_t2i_upscaling_4x": {
|
||||
"stages_ms": {
|
||||
"InputValidationStage": 0.15,
|
||||
"TextEncodingStage": 537.4,
|
||||
"ImageVAEEncodingStage": 0.01,
|
||||
"LatentPreparationStage": 1.11,
|
||||
"TimestepPreparationStage": 39.57,
|
||||
"DenoisingStage": 24738.35,
|
||||
"DecodingStage": 14.04
|
||||
},
|
||||
"denoise_step_ms": {
|
||||
"0": 70.71,
|
||||
"1": 476.99,
|
||||
"2": 502.01,
|
||||
"3": 483.39,
|
||||
"4": 500.65,
|
||||
"5": 487.43,
|
||||
"6": 502.39,
|
||||
"7": 484.76,
|
||||
"8": 498.5,
|
||||
"9": 489.98,
|
||||
"10": 499.22,
|
||||
"11": 490.47,
|
||||
"12": 498.98,
|
||||
"13": 491.42,
|
||||
"14": 495.17,
|
||||
"15": 492.24,
|
||||
"16": 494.69,
|
||||
"17": 491.68,
|
||||
"18": 497.13,
|
||||
"19": 493.63,
|
||||
"20": 495.29,
|
||||
"21": 496.19,
|
||||
"22": 496.52,
|
||||
"23": 496.31,
|
||||
"24": 493.64,
|
||||
"25": 494.35,
|
||||
"26": 493.27,
|
||||
"27": 495.52,
|
||||
"28": 493.06,
|
||||
"29": 494.66,
|
||||
"30": 494.08,
|
||||
"31": 496.01,
|
||||
"32": 494.79,
|
||||
"33": 495.81,
|
||||
"34": 493.94,
|
||||
"35": 495.56,
|
||||
"36": 493.5,
|
||||
"37": 495.98,
|
||||
"38": 495.82,
|
||||
"39": 496.7,
|
||||
"40": 495.29,
|
||||
"41": 496.84,
|
||||
"42": 495.67,
|
||||
"43": 495.32,
|
||||
"44": 496.48,
|
||||
"45": 496.03,
|
||||
"46": 495.65,
|
||||
"47": 498.27,
|
||||
"48": 496.44,
|
||||
"49": 496.79
|
||||
},
|
||||
"expected_e2e_ms": 25735.08,
|
||||
"expected_avg_denoise_ms": 486.1,
|
||||
"expected_median_denoise_ms": 495.42
|
||||
},
|
||||
"wan2_1_t2v_1.3b_upscaling_4x": {
|
||||
"stages_ms": {
|
||||
"InputValidationStage": 0.08,
|
||||
"TextEncodingStage": 1164.21,
|
||||
"LatentPreparationStage": 0.26,
|
||||
"TimestepPreparationStage": 3.86,
|
||||
"DenoisingStage": 10234.35,
|
||||
"DecodingStage": 499.44
|
||||
},
|
||||
"denoise_step_ms": {
|
||||
"0": 200.87,
|
||||
"1": 202.85,
|
||||
"2": 203.32,
|
||||
"3": 206.02,
|
||||
"4": 205.75,
|
||||
"5": 204.14,
|
||||
"6": 205.12,
|
||||
"7": 204.65,
|
||||
"8": 203.99,
|
||||
"9": 204.96,
|
||||
"10": 204.35,
|
||||
"11": 206.89,
|
||||
"12": 200.69,
|
||||
"13": 209.67,
|
||||
"14": 204.91,
|
||||
"15": 203.5,
|
||||
"16": 206.73,
|
||||
"17": 202.43,
|
||||
"18": 205.92,
|
||||
"19": 204.61,
|
||||
"20": 211.47,
|
||||
"21": 197.43,
|
||||
"22": 203.58,
|
||||
"23": 205.82,
|
||||
"24": 204.01,
|
||||
"25": 205.06,
|
||||
"26": 204.86,
|
||||
"27": 206.03,
|
||||
"28": 200.78,
|
||||
"29": 206.99,
|
||||
"30": 206.58,
|
||||
"31": 202.84,
|
||||
"32": 204.51,
|
||||
"33": 204.19,
|
||||
"34": 202.89,
|
||||
"35": 204.55,
|
||||
"36": 205.03,
|
||||
"37": 204.2,
|
||||
"38": 203.92,
|
||||
"39": 204.9,
|
||||
"40": 203.24,
|
||||
"41": 204.21,
|
||||
"42": 205.76,
|
||||
"43": 205.32,
|
||||
"44": 202.63,
|
||||
"45": 205.67,
|
||||
"46": 204.55,
|
||||
"47": 202.89,
|
||||
"48": 205.29,
|
||||
"49": 203.87
|
||||
},
|
||||
"expected_e2e_ms": 12021.58,
|
||||
"expected_avg_denoise_ms": 204.49,
|
||||
"expected_median_denoise_ms": 204.55
|
||||
},
|
||||
"wan2_1_t2v_1.3b_frame_interp_2x_upscaling_4x": {
|
||||
"stages_ms": {
|
||||
"InputValidationStage": 0.03,
|
||||
"TextEncodingStage": 1089.94,
|
||||
"LatentPreparationStage": 0.12,
|
||||
"TimestepPreparationStage": 1.99,
|
||||
"DenoisingStage": 8617.56,
|
||||
"DecodingStage": 469.68
|
||||
},
|
||||
"denoise_step_ms": {
|
||||
"0": 122.51,
|
||||
"1": 172.31,
|
||||
"2": 169.58,
|
||||
"3": 171.76,
|
||||
"4": 171.5,
|
||||
"5": 174.24,
|
||||
"6": 175.85,
|
||||
"7": 171.2,
|
||||
"8": 172.0,
|
||||
"9": 172.13,
|
||||
"10": 171.22,
|
||||
"11": 174.8,
|
||||
"12": 174.01,
|
||||
"13": 172.28,
|
||||
"14": 172.77,
|
||||
"15": 173.52,
|
||||
"16": 172.19,
|
||||
"17": 175.12,
|
||||
"18": 172.78,
|
||||
"19": 175.1,
|
||||
"20": 171.54,
|
||||
"21": 173.38,
|
||||
"22": 171.61,
|
||||
"23": 174.14,
|
||||
"24": 174.06,
|
||||
"25": 172.32,
|
||||
"26": 173.08,
|
||||
"27": 173.94,
|
||||
"28": 173.32,
|
||||
"29": 174.3,
|
||||
"30": 173.63,
|
||||
"31": 172.21,
|
||||
"32": 174.4,
|
||||
"33": 173.25,
|
||||
"34": 173.54,
|
||||
"35": 175.12,
|
||||
"36": 172.93,
|
||||
"37": 172.76,
|
||||
"38": 174.73,
|
||||
"39": 174.46,
|
||||
"40": 172.66,
|
||||
"41": 174.58,
|
||||
"42": 173.9,
|
||||
"43": 174.88,
|
||||
"44": 172.35,
|
||||
"45": 173.52,
|
||||
"46": 175.94,
|
||||
"47": 172.88,
|
||||
"48": 174.97,
|
||||
"49": 172.94
|
||||
},
|
||||
"expected_e2e_ms": 10425.77,
|
||||
"expected_avg_denoise_ms": 172.28,
|
||||
"expected_median_denoise_ms": 173.29
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -871,9 +871,7 @@ def get_generate_fn(
|
||||
req_background = None # Not specified in current request
|
||||
|
||||
# Build extra_body for optional features
|
||||
extra_body = {}
|
||||
if sampling_params.enable_teacache:
|
||||
extra_body["enable_teacache"] = True
|
||||
extra_body = dict(sampling_params.extras)
|
||||
|
||||
response = client.images.with_raw_response.generate(
|
||||
model=model_path,
|
||||
@@ -898,6 +896,13 @@ def get_generate_fn(
|
||||
|
||||
# Validate output file
|
||||
expected_width, expected_height = parse_dimensions(output_size)
|
||||
if (
|
||||
sampling_params.extras.get("enable_upscaling")
|
||||
and expected_width
|
||||
and expected_height
|
||||
):
|
||||
expected_width *= sampling_params.extras.get("upscaling_scale", 4)
|
||||
expected_height *= sampling_params.extras.get("upscaling_scale", 4)
|
||||
validate_image_file(
|
||||
tmp_path,
|
||||
expected_filename,
|
||||
@@ -947,8 +952,7 @@ def get_generate_fn(
|
||||
|
||||
# Build extra_body for optional features
|
||||
extra_body = {"num_frames": sampling_params.num_frames}
|
||||
if sampling_params.enable_teacache:
|
||||
extra_body["enable_teacache"] = True
|
||||
extra_body.update(sampling_params.extras)
|
||||
|
||||
images = [open(image_path, "rb") for image_path in image_paths]
|
||||
try:
|
||||
@@ -1076,22 +1080,18 @@ def get_generate_fn(
|
||||
pytest.skip(f"{case_id}: no text prompt configured")
|
||||
|
||||
# Build extra_body for optional features
|
||||
extra_body = {}
|
||||
if sampling_params.enable_teacache:
|
||||
extra_body["enable_teacache"] = True
|
||||
extra_body = dict(sampling_params.extras)
|
||||
if sampling_params.num_frames:
|
||||
extra_body["num_frames"] = sampling_params.num_frames
|
||||
if sampling_params.enable_frame_interpolation:
|
||||
extra_body["enable_frame_interpolation"] = True
|
||||
extra_body["frame_interpolation_exp"] = (
|
||||
sampling_params.frame_interpolation_exp
|
||||
)
|
||||
|
||||
# Compute expected output frame count for validation
|
||||
expected_frame_count = None
|
||||
if sampling_params.enable_frame_interpolation and sampling_params.num_frames:
|
||||
if (
|
||||
sampling_params.extras.get("enable_frame_interpolation")
|
||||
and sampling_params.num_frames
|
||||
):
|
||||
n = sampling_params.num_frames
|
||||
exp = sampling_params.frame_interpolation_exp
|
||||
exp = sampling_params.extras.get("frame_interpolation_exp", 1)
|
||||
expected_frame_count = (n - 1) * (2**exp) + 1
|
||||
|
||||
return _create_and_download_video(
|
||||
@@ -1118,9 +1118,7 @@ def get_generate_fn(
|
||||
pytest.skip(f"{case_id}: file missing: {image_path}")
|
||||
|
||||
# Build extra_body for optional features
|
||||
extra_body = {}
|
||||
if sampling_params.enable_teacache:
|
||||
extra_body["enable_teacache"] = True
|
||||
extra_body = dict(sampling_params.extras)
|
||||
|
||||
with image_path.open("rb") as fh:
|
||||
return _create_and_download_video(
|
||||
@@ -1140,8 +1138,7 @@ def get_generate_fn(
|
||||
|
||||
# Build extra_body for optional features
|
||||
extra_body = {"reference_url": sampling_params.image_path}
|
||||
if sampling_params.enable_teacache:
|
||||
extra_body["enable_teacache"] = True
|
||||
extra_body.update(sampling_params.extras)
|
||||
|
||||
return _create_and_download_video(
|
||||
client,
|
||||
@@ -1170,9 +1167,7 @@ def get_generate_fn(
|
||||
pytest.skip(f"{case_id}: file missing: {image_path}")
|
||||
|
||||
# Build extra_body for optional features
|
||||
extra_body = {}
|
||||
if sampling_params.enable_teacache:
|
||||
extra_body["enable_teacache"] = True
|
||||
extra_body = dict(sampling_params.extras)
|
||||
|
||||
with image_path.open("rb") as fh:
|
||||
return _create_and_download_video(
|
||||
@@ -1186,6 +1181,7 @@ def get_generate_fn(
|
||||
extra_body={
|
||||
"fps": sampling_params.fps,
|
||||
"num_frames": sampling_params.num_frames,
|
||||
**extra_body,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@@ -233,12 +233,9 @@ class DiffusionSamplingParams:
|
||||
|
||||
num_outputs_per_prompt: int = 1
|
||||
|
||||
# TeaCache acceleration
|
||||
enable_teacache: bool = False
|
||||
|
||||
# Frame interpolation
|
||||
enable_frame_interpolation: bool = False
|
||||
frame_interpolation_exp: int = 1 # 1 = 2×, 2 = 4×
|
||||
# Additional request-level parameters (e.g. enable_teacache, enable_upscaling, …)
|
||||
# merged directly into the OpenAI extra_body dict.
|
||||
extras: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -494,6 +491,19 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
|
||||
),
|
||||
MULTI_FRAME_I2I_sampling_params,
|
||||
),
|
||||
# Upscaling (Real-ESRGAN 4×) for T2I
|
||||
DiffusionTestCase(
|
||||
"flux_2_image_t2i_upscaling_4x",
|
||||
DiffusionServerArgs(
|
||||
model_path="black-forest-labs/FLUX.2-dev",
|
||||
modality="image",
|
||||
),
|
||||
DiffusionSamplingParams(
|
||||
prompt="Doraemon is eating dorayaki",
|
||||
output_size="1024x1024",
|
||||
extras={"enable_upscaling": True, "upscaling_scale": 4},
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
HUNYUAN3D_SHAPE_sampling_params = DiffusionSamplingParams(
|
||||
@@ -536,10 +546,10 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
|
||||
),
|
||||
DiffusionSamplingParams(
|
||||
prompt=T2V_PROMPT,
|
||||
enable_teacache=True,
|
||||
extras={"enable_teacache": True},
|
||||
),
|
||||
),
|
||||
# Frame interpolation correctness (2× / exp=1)
|
||||
# Frame interpolation (2× / exp=1)
|
||||
# Uses the same 1.3B model already in the suite;
|
||||
DiffusionTestCase(
|
||||
"wan2_1_t2v_1.3b_frame_interp_2x",
|
||||
@@ -550,8 +560,40 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
|
||||
),
|
||||
DiffusionSamplingParams(
|
||||
prompt=T2V_PROMPT,
|
||||
enable_frame_interpolation=True,
|
||||
frame_interpolation_exp=1,
|
||||
extras={"enable_frame_interpolation": True, "frame_interpolation_exp": 1},
|
||||
),
|
||||
),
|
||||
# Upscaling (Real-ESRGAN 4×)
|
||||
# Uses the same 1.3B model already in the suite;
|
||||
DiffusionTestCase(
|
||||
"wan2_1_t2v_1.3b_upscaling_4x",
|
||||
DiffusionServerArgs(
|
||||
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
modality="video",
|
||||
custom_validator="video",
|
||||
),
|
||||
DiffusionSamplingParams(
|
||||
prompt=T2V_PROMPT,
|
||||
extras={"enable_upscaling": True, "upscaling_scale": 4},
|
||||
),
|
||||
),
|
||||
# Combined: Frame interpolation (2×) + Upscaling (4×)
|
||||
# Verifies that both post-processing steps compose correctly.
|
||||
DiffusionTestCase(
|
||||
"wan2_1_t2v_1.3b_frame_interp_2x_upscaling_4x",
|
||||
DiffusionServerArgs(
|
||||
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
|
||||
modality="video",
|
||||
custom_validator="video",
|
||||
),
|
||||
DiffusionSamplingParams(
|
||||
prompt=T2V_PROMPT,
|
||||
extras={
|
||||
"enable_frame_interpolation": True,
|
||||
"frame_interpolation_exp": 1,
|
||||
"enable_upscaling": True,
|
||||
"upscaling_scale": 4,
|
||||
},
|
||||
),
|
||||
),
|
||||
# LoRA test case for single transformer + merge/unmerge API test
|
||||
|
||||
Reference in New Issue
Block a user