From 64950d8f9735e0e0f652092b695d21e3ca34f310 Mon Sep 17 00:00:00 2001 From: wxy <1908865287@qq.com> Date: Sat, 7 Feb 2026 22:08:42 +0800 Subject: [PATCH] [diffusion] feat: support saving videos directly on the server to avoid the overhead of tensor transfer (#18253) --- .../configs/sample/sampling_params.py | 11 ++- .../entrypoints/diffusion_generator.py | 93 +++++++++++-------- .../runtime/entrypoints/http_server.py | 41 +++----- .../runtime/entrypoints/openai/utils.py | 63 ++++++------- .../runtime/entrypoints/utils.py | 68 +++++++++++++- .../runtime/managers/gpu_worker.py | 19 +++- .../runtime/pipelines_core/schedule_batch.py | 3 +- 7 files changed, 188 insertions(+), 110 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index f43625178..699cfd855 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -153,6 +153,8 @@ class SamplingParams: # if True, suppress verbose logging for this request suppress_logs: bool = False + return_file_paths_only: bool = True + def _set_output_file_ext(self): # add extension if needed if not any( @@ -738,6 +740,12 @@ class SamplingParams: "Default: true. Examples: --adjust-frames, --adjust-frames true, --adjust-frames false." ), ) + parser.add_argument( + "--return-file-paths-only", + action=StoreBoolean, + default=SamplingParams.return_file_paths_only, + help="If set, output file will be saved early to get a performance boost, while output tensors will not be returned.", + ) return parser @classmethod @@ -807,9 +815,6 @@ class SamplingParams: n_tokens = -1 return n_tokens - def output_file_path(self): - return os.path.join(self.output_path, self.output_file_name) - @dataclass class CacheParams: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index b43e08368..3c397f486 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -14,12 +14,8 @@ import time from typing import Any, List, Union import numpy as np -import torch -from sglang.multimodal_gen.configs.sample.sampling_params import ( - DataType, - SamplingParams, -) +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( ListLorasReq, MergeLoraWeightsReq, @@ -28,8 +24,8 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( format_lora_message, ) from sglang.multimodal_gen.runtime.entrypoints.utils import ( - post_process_sample, prepare_request, + save_outputs, ) from sglang.multimodal_gen.runtime.launch_server import launch_server from sglang.multimodal_gen.runtime.pipelines_core import Req @@ -229,49 +225,64 @@ class DiffGenerator: if output_batch.error: raise Exception(f"{output_batch.error}") - if output_batch.output is None: + if ( + output_batch.output is None + and output_batch.output_file_paths is None + ): logger.error( "Received empty output from scheduler for prompt %d", request_idx + 1, ) continue audio_sample_rate = output_batch.audio_sample_rate - for output_idx, sample in enumerate(output_batch.output): - num_outputs = len(output_batch.output) - audio = output_batch.audio - if req.data_type == DataType.VIDEO: - if isinstance(audio, torch.Tensor) and audio.ndim >= 2: - audio = ( - audio[output_idx] - if audio.shape[0] > output_idx - else None - ) - elif isinstance(audio, np.ndarray) and audio.ndim >= 2: - audio = ( - audio[output_idx] - if audio.shape[0] > output_idx - else None - ) - if audio is not None and not ( - isinstance(sample, (tuple, list)) and len(sample) == 2 - ): - sample = (sample, audio) - frames = post_process_sample( - sample, - fps=req.fps, - save_output=req.save_output, - # TODO: output file path for req should be determined - save_file_path=req.output_file_path( - num_outputs, output_idx - ), - data_type=req.data_type, - audio_sample_rate=audio_sample_rate, - ) + if req.save_output and req.return_file_paths_only: + for output_idx, output_path in enumerate( + output_batch.output_file_paths + ): + result_item: dict[str, Any] = { + "samples": None, + "frames": None, + "audio": None, + "prompts": req.prompt, + "size": (req.height, req.width, req.num_frames), + "generation_time": timer.duration, + "peak_memory_mb": output_batch.peak_memory_mb, + "timings": ( + output_batch.timings.to_dict() + if output_batch.timings + else {} + ), + "trajectory": output_batch.trajectory_latents, + "trajectory_timesteps": output_batch.trajectory_timesteps, + "trajectory_decoded": output_batch.trajectory_decoded, + "prompt_index": output_idx, + "output_file_path": output_path, + } + results.append(result_item) + continue + + samples_out: list[Any] = [] + audios_out: list[Any] = [] + frames_out: list[Any] = [] + save_outputs( + output_batch.output, + req.data_type, + req.fps, + req.save_output, + lambda idx: req.output_file_path(len(output_batch.output), idx), + audio=output_batch.audio, + audio_sample_rate=audio_sample_rate, + samples_out=samples_out, + audios_out=audios_out, + frames_out=frames_out, + ) + + for output_idx in range(len(samples_out)): result_item: dict[str, Any] = { - "samples": sample, - "frames": frames, - "audio": audio, + "samples": samples_out[output_idx], + "frames": frames_out[output_idx], + "audio": audios_out[output_idx], "prompts": req.prompt, "size": (req.height, req.width, req.num_frames), "generation_time": timer.duration, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index 960b96d73..20d2c3df9 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -16,8 +16,8 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( VertexGenerateReqInput, ) from sglang.multimodal_gen.runtime.entrypoints.utils import ( - post_process_sample, prepare_request, + save_outputs, ) from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args @@ -114,33 +114,22 @@ async def forward_to_scheduler(req_obj, sp): """Forwards request to scheduler and processes the result.""" try: response = await async_scheduler_client.forward(req_obj) - if response.output is None: + if response.output is None and response.output_file_paths is None: raise RuntimeError("Model generation returned no output.") - output_file_path = sp.output_file_path() - sample = response.output[0] - try: - audio = response.audio - except AttributeError: - audio = None - if isinstance(audio, torch.Tensor) and audio.ndim >= 2: - audio = audio[0] - if audio is not None and not ( - isinstance(sample, (tuple, list)) and len(sample) == 2 - ): - sample = (sample, audio) - post_process_sample( - sample=sample, - data_type=sp.data_type, - fps=sp.fps or 24, - save_output=True, - save_file_path=output_file_path, - audio_sample_rate=( - response.audio_sample_rate - if hasattr(response, "audio_sample_rate") - else None - ), - ) + if response.output_file_paths: + output_file_path = response.output_file_paths[0] + else: + output_file_path = sp.output_file_path() + save_outputs( + [response.output[0]], + sp.data_type, + sp.fps, + True, + lambda _idx: output_file_path, + audio=response.audio, + audio_sample_rate=response.audio_sample_rate, + ) if hasattr(response, "model_dump"): data = response.model_dump() diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py index b561cf8a4..a52643c55 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/utils.py @@ -7,11 +7,10 @@ import time from typing import Any, List, Optional, Union import httpx -import torch from fastapi import UploadFile from sglang.multimodal_gen.configs.sample.sampling_params import DataType -from sglang.multimodal_gen.runtime.entrypoints.utils import post_process_sample +from sglang.multimodal_gen.runtime.entrypoints.utils import save_outputs from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.scheduler_client import AsyncSchedulerClient from sglang.multimodal_gen.runtime.utils.logging_utils import ( @@ -211,51 +210,43 @@ async def process_generation_batch( with log_generation_timer(logger, batch.prompt): result = await scheduler_client.forward([batch]) - if result.output is None: + if result.output is None and result.output_file_paths is None: error_msg = result.error or "Unknown error" raise RuntimeError( f"Model generation returned no output. Error from scheduler: {error_msg}" ) save_file_path_list = [] - audio_sample_rate = result.audio_sample_rate - if batch.data_type == DataType.VIDEO: - for idx, output in enumerate(result.output): - save_file_path = str( - os.path.join(batch.output_path, batch.output_file_name) - ) - sample = result.output[idx] - audio = result.audio - if isinstance(audio, torch.Tensor) and audio.ndim >= 2: - audio = audio[idx] if audio.shape[0] > idx else None - if audio is not None and not ( - isinstance(sample, (tuple, list)) and len(sample) == 2 - ): - sample = (sample, audio) - post_process_sample( - sample, - batch.data_type, - batch.fps, - batch.save_output, - save_file_path, - audio_sample_rate=audio_sample_rate, - ) - save_file_path_list.append(save_file_path) + # If output_file_paths is provided, use it instead of output. + if result.output_file_paths: + save_file_path_list = result.output_file_paths else: - for idx, output in enumerate(result.output): - save_file_path = str( - os.path.join( - batch.output_path, f"sample_{idx}_" + batch.output_file_name - ) - ) - post_process_sample( - output, + audio_sample_rate = result.audio_sample_rate + if batch.data_type == DataType.VIDEO: + save_file_path_list = save_outputs( + result.output, batch.data_type, batch.fps, batch.save_output, - save_file_path, + lambda _idx: str( + os.path.join(batch.output_path, batch.output_file_name) + ), + audio=result.audio, + audio_sample_rate=audio_sample_rate, + ) + else: + save_file_path_list = save_outputs( + result.output, + batch.data_type, + batch.fps, + batch.save_output, + lambda idx: str( + os.path.join( + batch.output_path, + f"sample_{idx}_" + batch.output_file_name, + ) + ), audio_sample_rate=audio_sample_rate, ) - save_file_path_list.append(save_file_path) total_time = time.perf_counter() - total_start_time log_batch_completion(logger, 1, total_time) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index 5edd3f97a..9544d2b0f 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -12,7 +12,7 @@ import os import shutil import subprocess import tempfile -from typing import Any, Optional +from typing import Any, Callable, Optional, Sequence import imageio import numpy as np @@ -234,6 +234,72 @@ def prepare_request( return req +def attach_audio_to_video_sample( + sample: Any, + audio: Any, + output_idx: int, +) -> Any: + """Attach per-sample audio for video outputs when available.""" + if audio is None: + return sample + if isinstance(audio, torch.Tensor) and audio.ndim >= 2: + audio = audio[output_idx] if audio.shape[0] > output_idx else None + elif isinstance(audio, np.ndarray) and audio.ndim >= 2: + audio = audio[output_idx] if audio.shape[0] > output_idx else None + + if audio is not None and not ( + isinstance(sample, (tuple, list)) and len(sample) == 2 + ): + return (sample, audio) + return sample + + +def save_outputs( + outputs: Sequence[Any], + data_type: DataType, + fps: int, + save_output: bool, + build_output_path: Callable[[int], str], + *, + audio: Any = None, + audio_sample_rate: Optional[int] = None, + samples_out: Optional[list[Any]] = None, + audios_out: Optional[list[Any]] = None, + frames_out: Optional[list[Any]] = None, +) -> list[str]: + """Save outputs to files and return the list of file paths.""" + output_paths: list[str] = [] + for idx, output in enumerate(outputs): + save_file_path = build_output_path(idx) + sample = output + if data_type == DataType.VIDEO: + sample = attach_audio_to_video_sample(sample, audio, idx) + frames = post_process_sample( + sample, + data_type, + fps, + save_output, + save_file_path, + audio_sample_rate=audio_sample_rate, + ) + if samples_out is not None: + samples_out.append(sample) + if audios_out is not None: + if data_type == DataType.VIDEO: + audio_item = audio + if isinstance(audio, torch.Tensor) and audio.ndim >= 2: + audio_item = audio[idx] if audio.shape[0] > idx else None + elif isinstance(audio, np.ndarray) and audio.ndim >= 2: + audio_item = audio[idx] if audio.shape[0] > idx else None + audios_out.append(audio_item) + else: + audios_out.append(audio) + if frames_out is not None: + frames_out.append(frames) + output_paths.append(save_file_path) + return output_paths + + def post_process_sample( sample: Any, data_type: DataType, diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 9af2e1efc..2cefd7761 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -18,6 +18,7 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_cfg_group, get_tp_group, ) +from sglang.multimodal_gen.runtime.entrypoints.utils import save_outputs from sglang.multimodal_gen.runtime.pipelines_core import ( ComposedPipelineBase, LoRAPipeline, @@ -186,6 +187,21 @@ class GPUWorker: duration_ms = (time.monotonic() - start_time) * 1000 output_batch.timings.total_duration_ms = duration_ms + # Save output to file and return file path only if requested. Avoid the serialization + # and deserialization overhead between scheduler_client and gpu_worker. + if req.save_output and req.return_file_paths_only: + output_paths = save_outputs( + output_batch.output, + req.data_type, + req.fps, + True, + lambda idx: req.output_file_path(len(output_batch.output), idx), + audio=output_batch.audio, + audio_sample_rate=output_batch.audio_sample_rate, + ) + output_batch.output_file_paths = output_paths + output_batch.output = None + # TODO: extract to avoid duplication if req.perf_dump_path is not None or envs.SGLANG_DIFFUSION_STAGE_LOGGING: # Avoid logging warmup perf records that share the same request_id. @@ -198,8 +214,7 @@ class GPUWorker: if output_batch is None: output_batch = OutputBatch() output_batch.error = f"Error executing request {req.request_id}: {e}" - finally: - return output_batch + return output_batch def get_can_stay_resident_components( self, remaining_gpu_mem_gb: float diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 7fabc52bf..4fde4c3e9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -100,7 +100,6 @@ class Req: raw_audio_latent_shape: tuple[int, ...] | None = None # Audio Parameters - fps: float = 24.0 generate_audio: bool = True raw_latent_shape: torch.Tensor | None = None @@ -295,6 +294,7 @@ class Req: width: {target_width} height: {target_height} num_frames: {self.num_frames} + fps: {self.fps} prompt: {self.prompt} neg_prompt: {self.negative_prompt} seed: {self.seed} @@ -325,6 +325,7 @@ class OutputBatch: trajectory_latents: torch.Tensor | None = None trajectory_decoded: list[torch.Tensor] | None = None error: str | None = None + output_file_paths: list[str] | None = None # logged timings info, directly from Req.timings timings: Optional["RequestTimings"] = None