[Diffusion] Use SGLang server for ERNIE-Image prompt enhancement (#31354)

Co-authored-by: Elizaveta Martirosian <you@example.com>
Co-authored-by: Elizaveta Martirosian <elizaveta.martirosian@gmail.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>

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
 

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
+
+<a href="https://github.com/sgl-project/sglang/tree/main/docs_new/docs/sglang-diffusion/models_with_ar.mdx#ascend-npu-env">Check here.</a>

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(

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):