From 05eb0bcc6131b11e07bbfd9ad8589f0d162f7247 Mon Sep 17 00:00:00 2001 From: Yashika Gandhi - Google <56970160+yashikagandhi-google@users.noreply.github.com> Date: Fri, 19 Dec 2025 09:47:17 -0800 Subject: [PATCH] [diffusion] http-server: support vertex generate pathway (#15348) --- .../runtime/entrypoints/http_server.py | 116 +++++++++++++++++- .../multimodal_gen/runtime/launch_server.py | 18 ++- 2 files changed, 132 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index e44e7b9a0..0fd8f75f5 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -1,12 +1,27 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import asyncio +import base64 +import os +import uuid from contextlib import asynccontextmanager +import torch from fastapi import APIRouter, FastAPI, Request +from fastapi.responses import ORJSONResponse +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.openai import image_api, video_api -from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.entrypoints.utils import ( + post_process_sample, + prepare_request, +) +from sglang.multimodal_gen.runtime.scheduler_client import scheduler_client +from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args +from sglang.srt.managers.io_struct import VertexGenerateReqInput + +DEFAULT_SEED = 1024 +VERTEX_ROUTE = os.environ.get("AIP_PREDICT_ROUTE", "/vertex_generate") @asynccontextmanager @@ -70,6 +85,104 @@ async def health_generate(): return {"status": "ok"} +def make_serializable(obj): + """Recursively converts Tensors to None for JSON serialization.""" + if isinstance(obj, torch.Tensor): + return None + if isinstance(obj, dict): + return {k: make_serializable(v) for k, v in obj.items()} + if isinstance(obj, list): + return [make_serializable(v) for v in obj] + return obj + + +def encode_video_to_base64(file_path: str): + if not os.path.exists(file_path): + return None + with open(file_path, "rb") as f: + return base64.b64encode(f.read()).decode("utf-8") + + +async def forward_to_scheduler(req_obj, sp): + """Forwards request to scheduler and processes the result.""" + try: + response = await scheduler_client.forward(req_obj) + if response.output is None: + raise RuntimeError("Model generation returned no output.") + + output_file_path = sp.output_file_path() + post_process_sample( + sample=response.output[0], + data_type=sp.data_type, + fps=sp.fps or 24, + save_output=True, + save_file_path=output_file_path, + ) + + if hasattr(response, "model_dump"): + data = response.model_dump() + else: + data = response if isinstance(response, dict) else vars(response) + + if output_file_path: + print(f"Processing output file: {output_file_path}") + b64_video = encode_video_to_base64(output_file_path) + + if b64_video: + data["output"] = b64_video + data.pop("video_data", None) + data.pop("video_tensor", None) + + return make_serializable(data) + + except Exception as e: + print(f"Error during generation: {e}") + return {"error": str(e)} + + +vertex_router = APIRouter() + + +@vertex_router.post(VERTEX_ROUTE) +async def vertex_generate(vertex_req: VertexGenerateReqInput): + if not vertex_req.instances: + return ORJSONResponse({"predictions": []}) + + server_args = get_global_server_args() + params = vertex_req.parameters or {} + + futures = [] + + for inst in vertex_req.instances: + rid = f"vertex_{uuid.uuid4()}" + + prompt = inst.get("prompt") or inst.get("text") + image_input = inst.get("image") or inst.get("image_url") + seed_val = params.get("seed", DEFAULT_SEED) + + sp = SamplingParams.from_user_sampling_params_args( + model_path=server_args.model_path, + request_id=rid, + prompt=prompt, + image_path=image_input, + num_frames=params.get("num_frames"), + fps=params.get("fps"), + width=params.get("width"), + height=params.get("height"), + guidance_scale=params.get("guidance_scale"), + seed=seed_val, + server_args=server_args, + save_output=params.get("save_output"), + ) + + backend_req = prepare_request(server_args, sampling_params=sp) + futures.append(forward_to_scheduler(backend_req, sp)) + + results = await asyncio.gather(*futures) + + return ORJSONResponse({"predictions": results}) + + def create_app(server_args: ServerArgs): """ Create and configure the FastAPI application instance. @@ -77,6 +190,7 @@ def create_app(server_args: ServerArgs): app = FastAPI(lifespan=lifespan) app.include_router(health_router) + app.include_router(vertex_router) from sglang.multimodal_gen.runtime.entrypoints.openai import common_api diff --git a/python/sglang/multimodal_gen/runtime/launch_server.py b/python/sglang/multimodal_gen/runtime/launch_server.py index 75b50fba2..6c69b0e25 100644 --- a/python/sglang/multimodal_gen/runtime/launch_server.py +++ b/python/sglang/multimodal_gen/runtime/launch_server.py @@ -1,17 +1,24 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo import multiprocessing as mp +import os +import sys import uvicorn from sglang.multimodal_gen.runtime.entrypoints.http_server import create_app from sglang.multimodal_gen.runtime.managers.gpu_worker import run_scheduler_process -from sglang.multimodal_gen.runtime.server_args import ServerArgs, set_global_server_args +from sglang.multimodal_gen.runtime.server_args import ( + ServerArgs, + prepare_server_args, + set_global_server_args, +) from sglang.multimodal_gen.runtime.utils.logging_utils import ( configure_logger, logger, suppress_other_loggers, ) +from sglang.srt.utils import kill_process_tree def launch_server(server_args: ServerArgs, launch_http_server: bool = True): @@ -152,3 +159,12 @@ def launch_http_server_only(server_args): port=server_args.port, reload=False, ) + + +if __name__ == "__main__": + server_args = prepare_server_args(sys.argv[1:]) + + try: + launch_server(server_args) + finally: + kill_process_tree(os.getpid(), include_parent=False)