From 44e4999ab2f1c136c4dcafebbcfb079c45d827f2 Mon Sep 17 00:00:00 2001 From: Elizaveta Martirosian Date: Sat, 18 Jul 2026 06:55:28 +0300 Subject: [PATCH] [Diffusion] Use SGLang server for ERNIE-Image prompt enhancement (#31354) Co-authored-by: Elizaveta Martirosian Co-authored-by: Elizaveta Martirosian Co-authored-by: ronnie_zheng --- docs_new/docs/sglang-diffusion/api/cli.mdx | 1 + .../docs/sglang-diffusion/models_with_pe.mdx | 61 +++++++++++++++++++ .../loader/component_loaders/pe_loader.py | 32 ++++++++++ .../runtime/server_args/server_args.py | 11 ++++ 4 files changed, 105 insertions(+) create mode 100644 docs_new/docs/sglang-diffusion/models_with_pe.mdx diff --git a/docs_new/docs/sglang-diffusion/api/cli.mdx b/docs_new/docs/sglang-diffusion/api/cli.mdx index 3d7ae504f..3386f36ff 100644 --- a/docs_new/docs/sglang-diffusion/api/cli.mdx +++ b/docs_new/docs/sglang-diffusion/api/cli.mdx @@ -94,6 +94,7 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis - `--srt-encoder-url {HTTPADDRESS}`: address of SGLang srt server with AR model for GLM-Image like models - `--srt-encoder-timeout {SECONDS}`: Timeout in seconds for HTTP requests to the SGLang encoder server - `--srt-encoder-connection-timeout {SECONDS}`: TCP connection timeout in seconds for SGLang encoder server +- `--pe-server-url {HTTPADDRESS}`: url of SGLang server hosting the PE model (e.g., for ERNIE-Image) ### Sampling and output diff --git a/docs_new/docs/sglang-diffusion/models_with_pe.mdx b/docs_new/docs/sglang-diffusion/models_with_pe.mdx new file mode 100644 index 000000000..cce6f0351 --- /dev/null +++ b/docs_new/docs/sglang-diffusion/models_with_pe.mdx @@ -0,0 +1,61 @@ +--- +title: "Diffusion Models with Prompt Enhancement (PE)" +--- + +## Quick Start + +By default, the PE model is loaded by the diffusion model server, which may not provide optimal performance. For higher performance, the PE model can be deployed as a separate SGLang server. This document uses `baidu/ERNIE-Image` as an example. + +Run the model with the built-in Transformers PE implementation (default): + +```bash +# Terminal 1: launch server +sglang serve --model-path baidu/ERNIE-Image --port ${PORT} +``` +```bash +# Terminal 2: launch client +curl -X POST http://${HOST}:${PORT}/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "prompt": "This is a photograph depicting an urban street scene. Shot at eye level, it shows a covered pedestrian or commercial street. Slightly below the center of the frame, a cyclist rides away from the camera toward the background, appearing as a dark silhouette against backlighting with indistinct details. The ground is paved with regular square tiles, bisected by a prominent tactile paving strip running through the scene, whose raised textures are clearly visible under the light. Light streams in diagonally from the right side of the frame, creating a strong backlight effect with a distinct Tyndall effect—visible light beams illuminating dust or vapor in the air and casting long shadows across the street. Several pedestrians appear on the left side and in the distance, some with their backs to the camera and others walking sideways, all rendered as silhouettes or semi-silhouettes. The overall color palette is warm, dominated by golden yellows and dark browns, evoking the atmosphere of dusk or early morning.", + "height": 1264, + "width": 848, + "num_inference_steps": 50, + "guidance_scale": 4.0, + "use_pe": true + }' +``` + +Run the model with an SGLang-served PE model (high performance): + +```bash +# Terminal 1: launch SGLang PE model server +sglang serve --model-path /path/to/baidu/ERNIE-Image/pe/ --port ${PE_PORT} +``` +```bash +# Terminal 2: launch diffusion model server with PE server +sglang serve --model-path /path/to/baidu/ERNIE-Image/ --pe-server-url "http://${HOST}:${PE_PORT}" +``` +```bash +# Terminal 3: launch client +curl -X POST http://${HOST}:${PORT}/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "prompt": "This is a photograph depicting an urban street scene. Shot at eye level, it shows a covered pedestrian or commercial street. Slightly below the center of the frame, a cyclist rides away from the camera toward the background, appearing as a dark silhouette against backlighting with indistinct details. The ground is paved with regular square tiles, bisected by a prominent tactile paving strip running through the scene, whose raised textures are clearly visible under the light. Light streams in diagonally from the right side of the frame, creating a strong backlight effect with a distinct Tyndall effect—visible light beams illuminating dust or vapor in the air and casting long shadows across the street. Several pedestrians appear on the left side and in the distance, some with their backs to the camera and others walking sideways, all rendered as silhouettes or semi-silhouettes. The overall color palette is warm, dominated by golden yellows and dark browns, evoking the atmosphere of dusk or early morning.", + "height": 1264, + "width": 848, + "num_inference_steps": 50, + "guidance_scale": 4.0, + "use_pe": true + }' +``` + +## Support matrix + +| Model | Built-in PE | SGLang PE Server | +|-------|-------------|------------------| +| ERNIE-Image | ✅ | ✅ | + +## Ascend NPU Environment + +Check here. diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py index 33fe3a5c3..699436774 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/pe_loader.py @@ -2,6 +2,7 @@ import json import os +import requests import torch from transformers import AutoModelForCausalLM, AutoTokenizer @@ -79,6 +80,31 @@ class PEModelWrapper: return self +class SGLangPEModelWrapper: + + def __init__(self, model_url): + self.model_url = model_url.rstrip("/") + # Tokenizer is initialized separately during pipeline setup + self.pe_tokenizer = None + + def generate(self, prompt: str, sampling_params: dict) -> dict: + response = requests.post( + self.model_url + "/generate", + json={ + "text": prompt, + "sampling_params": sampling_params, + }, + ) + + response.raise_for_status() + + return response.json() + + def to(self, *args, **kwargs): + logger.debug("Ignoring .to() because PE model is served externally") + return self + + class PELoader(ComponentLoader): """Loader for prompt-enhancement causal LM (Ministral-3 based).""" @@ -88,6 +114,12 @@ class PELoader(ComponentLoader): def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ): + if server_args.pe_server_url is not None: + logger.info( + f"Using external SGLang server for PE: {server_args.pe_server_url}" + ) + return SGLangPEModelWrapper(server_args.pe_server_url) + logger.info("Loading PE model from %s ...", component_model_path) pe_tokenizer_dir = os.path.join( diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 0d382aca6..a722333c0 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -426,6 +426,9 @@ class ServerArgs(DisaggServerArgsMixin): srt_encoder_connect_timeout: int = 3.05 srt_encoder_timeout: int = 100 + # SGLang server for PE model inference + pe_server_url: str | None = None + @property def broker_port(self) -> int: return self.port + 1 @@ -1938,6 +1941,14 @@ class ServerArgs(DisaggServerArgsMixin): "Increase value if connection between diffusion server and AR model server is slow.", ) + # SGLang server for PE model inference + parser.add_argument( + "--pe-server-url", + type=str, + default=ServerArgs.pe_server_url, + help="URL of SGLang server for PE model", + ) + return parser def url(self):