[diffusion] feat: allows quality adjustment of generated images/videos (#17937)
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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)):
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user