[diffusion] feat: implement upscaling (#19723)

This commit is contained in:
Junhao Liu
2026-03-09 02:06:40 +08:00
committed by GitHub
parent b77dd41db0
commit 7662b8b919
17 changed files with 1005 additions and 63 deletions
@@ -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