[diffusion] http-server: support vertex generate pathway (#15348)

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
 

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)