[diffusion] feat: allows quality adjustment of generated images/videos (#17937)

This commit is contained in:
HuangJi
2026-02-13 11:56:20 +08:00
committed by GitHub
parent 6555b2a71d
commit f4d80f9d42
9 changed files with 103 additions and 4 deletions

View File

@@ -97,6 +97,8 @@ class SamplingParams:
prompt_path: str | None = None
output_path: str | None = None
output_file_name: str | None = None
output_quality: str | None = "default"
output_compression: int | None = None
# Batch info
num_outputs_per_prompt: int = 1
@@ -207,6 +209,12 @@ class SamplingParams:
if self.height is None:
self.height_not_provided = True
# Handle output_quality to output_compression conversion
if self.output_compression is None and self.output_quality is not None:
self.output_compression = self._adjust_output_quality(
self.output_quality, self.data_type
)
self._validate()
# Allow env var to override num_inference_steps (for faster CI testing on AMD)
@@ -214,6 +222,13 @@ class SamplingParams:
if env_steps is not None and self.num_inference_steps is not None:
self.num_inference_steps = int(env_steps)
def _adjust_output_quality(self, output_quality: str, data_type: DataType) -> int:
"""Convert output_quality string to compression level."""
output_quality_mapper = {"maximum": 100, "high": 90, "medium": 55, "low": 35}
if output_quality == "default":
return 50 if data_type == DataType.VIDEO else 75
return output_quality_mapper.get(output_quality)
def _validate(self):
"""
check if the sampling params is correct by itself
@@ -567,6 +582,18 @@ class SamplingParams:
default=SamplingParams.output_file_name,
help="Name of the output file",
)
parser.add_argument(
"--output-quality",
type=str,
default=SamplingParams.output_quality,
help="Output quality setting (default, low, medium, high, maximum)",
)
parser.add_argument(
"--output-compression",
type=int,
default=SamplingParams.output_compression,
help="Output compression level (0-100, higher means better quality but larger file size)",
)
parser.add_argument(
"--num-outputs-per-prompt",
type=int,

View File

@@ -277,6 +277,7 @@ class DiffGenerator:
samples_out=samples_out,
audios_out=audios_out,
frames_out=frames_out,
output_compression=req.output_compression,
)
for output_idx in range(len(samples_out)):

View File

@@ -22,6 +22,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.stores import IMAGE_STORE
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
_parse_size,
add_common_data_to_response,
adjust_output_quality,
merge_image_input_list,
process_generation_batch,
save_image_to_path,
@@ -63,6 +64,7 @@ def _build_sampling_params_from_request(
negative_prompt: Optional[str] = None,
enable_teacache: Optional[bool] = None,
num_frames: int = 1,
output_compression: Optional[int] = None,
) -> SamplingParams:
if size is None:
width, height = None, None
@@ -90,6 +92,11 @@ def _build_sampling_params_from_request(
**({"guidance_scale": guidance_scale} if guidance_scale is not None else {}),
**({"negative_prompt": negative_prompt} if negative_prompt is not None else {}),
**({"true_cfg_scale": true_cfg_scale} if true_cfg_scale is not None else {}),
**(
{"output_compression": output_compression}
if output_compression is not None
else {}
),
)
if num_inference_steps is not None:
@@ -122,11 +129,16 @@ async def generations(
true_cfg_scale=request.true_cfg_scale,
negative_prompt=request.negative_prompt,
enable_teacache=request.enable_teacache,
output_compression=request.output_compression,
)
batch = prepare_request(
server_args=get_global_server_args(),
sampling_params=sampling,
)
if batch.output_compression is None:
batch.output_compression = adjust_output_quality(
request.output_quality, batch.data_type
)
# Add diffusers_kwargs if provided
if request.diffusers_kwargs:
batch.extra["diffusers_kwargs"] = request.diffusers_kwargs
@@ -217,6 +229,8 @@ async def edits(
guidance_scale: Optional[float] = Form(None),
true_cfg_scale: Optional[float] = Form(None),
num_inference_steps: Optional[int] = Form(None),
output_quality: Optional[str] = Form("default"),
output_compression: Optional[int] = Form(None),
enable_teacache: Optional[bool] = Form(False),
num_frames: int = Form(1),
):
@@ -264,12 +278,16 @@ async def edits(
num_inference_steps=num_inference_steps,
enable_teacache=enable_teacache,
num_frames=num_frames,
output_compression=output_compression,
)
batch = prepare_request(
server_args=get_global_server_args(),
sampling_params=sampling,
)
if batch.output_compression is None:
batch.output_compression = adjust_output_quality(
output_quality, batch.data_type
)
save_file_path_list, result = await process_generation_batch(
async_scheduler_client, batch
)

View File

@@ -43,6 +43,8 @@ class ImageGenerationsRequest(BaseModel):
seed: Optional[int] = 1024
generator_device: Optional[str] = "cuda"
negative_prompt: Optional[str] = None
output_quality: Optional[str] = "default"
output_compression: Optional[int] = None
enable_teacache: Optional[bool] = False
diffusers_kwargs: Optional[Dict[str, Any]] = None # kwargs for diffusers backend
@@ -88,6 +90,8 @@ class VideoGenerationsRequest(BaseModel):
)
negative_prompt: Optional[str] = None
enable_teacache: Optional[bool] = False
output_quality: Optional[str] = "default"
output_compression: Optional[int] = None
output_path: Optional[str] = None
diffusers_kwargs: Optional[Dict[str, Any]] = None # kwargs for diffusers backend

View File

@@ -21,6 +21,8 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import (
logger = init_logger(__name__)
OUTPUT_QUALITY_MAPPER = {"maximum": 100, "high": 90, "medium": 55, "low": 35}
@dataclasses.dataclass
class SetLoraReq:
@@ -237,6 +239,7 @@ async def process_generation_batch(
),
audio=result.audio,
audio_sample_rate=audio_sample_rate,
output_compression=batch.output_compression,
)
else:
save_file_path_list = save_outputs(
@@ -251,6 +254,7 @@ async def process_generation_batch(
)
),
audio_sample_rate=audio_sample_rate,
output_compression=batch.output_compression,
)
total_time = time.perf_counter() - total_start_time
@@ -302,3 +306,9 @@ def add_common_data_to_response(
response["id"] = request_id
return response
def adjust_output_quality(output_quality: str, data_type: DataType = None) -> int:
if output_quality == "default":
return 50 if data_type == DataType.VIDEO else 75
return OUTPUT_QUALITY_MAPPER.get(output_quality, None)

View File

@@ -32,6 +32,7 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.stores import VIDEO_STORE
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
_parse_size,
add_common_data_to_response,
adjust_output_quality,
merge_image_input_list,
process_generation_batch,
save_image_to_path,
@@ -88,6 +89,8 @@ def _build_sampling_params_from_request(
sampling_kwargs["enable_teacache"] = request.enable_teacache
if request.output_path is not None:
sampling_kwargs["output_path"] = request.output_path
if request.output_compression is not None:
sampling_kwargs["output_compression"] = request.output_compression
sampling_params = SamplingParams.from_user_sampling_params_args(
model_path=server_args.model_path,
server_args=server_args,
@@ -173,6 +176,8 @@ async def create_video(
guidance_scale: Optional[float] = Form(None),
num_inference_steps: Optional[int] = Form(None),
enable_teacache: Optional[bool] = Form(False),
output_quality: Optional[str] = Form("default"),
output_compression: Optional[int] = Form(None),
extra_body: Optional[str] = Form(None),
):
content_type = request.headers.get("content-type", "").lower()
@@ -232,6 +237,8 @@ async def create_video(
negative_prompt=negative_prompt,
num_inference_steps=num_inference_steps,
enable_teacache=enable_teacache,
output_compression=output_compression,
output_quality=output_quality,
**(
{"guidance_scale": guidance_scale} if guidance_scale is not None else {}
),
@@ -293,6 +300,10 @@ async def create_video(
server_args=server_args,
sampling_params=sampling_params,
)
if batch.output_compression is None:
batch.output_compression = adjust_output_quality(
req.output_quality, batch.data_type
)
# Add diffusers_kwargs if provided
if req.diffusers_kwargs:
batch.extra["diffusers_kwargs"] = req.diffusers_kwargs

View File

@@ -266,6 +266,7 @@ def save_outputs(
samples_out: Optional[list[Any]] = None,
audios_out: Optional[list[Any]] = None,
frames_out: Optional[list[Any]] = None,
output_compression: Optional[int] = None,
) -> list[str]:
"""Save outputs to files and return the list of file paths."""
output_paths: list[str] = []
@@ -281,6 +282,7 @@ def save_outputs(
save_output,
save_file_path,
audio_sample_rate=audio_sample_rate,
output_compression=output_compression,
)
if samples_out is not None:
samples_out.append(sample)
@@ -307,6 +309,7 @@ def post_process_sample(
save_output: bool = True,
save_file_path: Optional[str] = None,
audio_sample_rate: Optional[int] = None,
output_compression: Optional[int] = None,
):
"""
Process sample output and save video if necessary
@@ -352,8 +355,9 @@ def post_process_sample(
if save_file_path:
os.makedirs(os.path.dirname(save_file_path), exist_ok=True)
if data_type == DataType.VIDEO:
# TODO: make this configurable
quality = 5
quality = (
output_compression / 10 if output_compression is not None else 5
)
imageio.mimsave(
save_file_path,
frames,
@@ -372,7 +376,7 @@ def post_process_sample(
)
else:
quality = 75
quality = output_compression if output_compression is not None else 75
if len(frames) > 1:
for i, image in enumerate(frames):
parts = save_file_path.rsplit(".", 1)

View File

@@ -227,6 +227,7 @@ class GPUWorker:
lambda idx: req.output_file_path(len(output_batch.output), idx),
audio=output_batch.audio,
audio_sample_rate=output_batch.audio_sample_rate,
output_compression=req.output_compression,
)
output_batch.output_file_paths = output_paths
output_batch.output = None