[diffusion] feat: support out-of-tree models and pipelines (#35713)

This commit is contained in:
Mick
2026-08-21 00:33:34 +08:00
committed by GitHub
parent 7f8f030000
commit be373395b4
8 changed files with 297 additions and 16 deletions
@@ -18,6 +18,11 @@ description: "Configure SGLang diffusion behavior with environment variables."
</tr>
</thead>
<tbody>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_EXTERNAL_MODEL_PACKAGE</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>not set</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Installed package that registers out-of-tree diffusion pipelines and component models. The package is imported once in every process.</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DIFFUSION_TARGET_DEVICE</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>cuda</code></td>
@@ -33,6 +33,51 @@ utilities, and common action-policy helpers. Model packages may call these
helpers. Keep ownership in shared runtime folders unless the code is truly
architecture-specific.
## Out-of-Tree Models and Pipelines
An installed package can register native component models and a pipeline
without modifying SGLang-Diffusion. Register them in the package's
`__init__.py`:
```python
from sglang.multimodal_gen.registry import register_pipeline
from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
from .configs import CustomPipelineConfig, CustomSamplingParams
from .pipeline import CustomPipeline
ModelRegistry.register_model(
"CustomTransformer2DModel",
"custom_diffusion.models:CustomTransformer2DModel",
)
register_pipeline(
CustomPipeline,
sampling_param_cls=CustomSamplingParams,
pipeline_config_cls=CustomPipelineConfig,
hf_model_paths=["my-org/custom-diffusion-model"],
model_detectors=[lambda value: "custom-diffusion" in value.lower()],
)
```
Install the package in the server environment and set the same environment
variable used by SRT plugins:
```bash
pip install -e /path/to/custom-diffusion
SGLANG_EXTERNAL_MODEL_PACKAGE=custom_diffusion \
sglang serve --model-path my-org/custom-diffusion-model
```
Notes:
- The string form of `register_model` keeps component imports lazy.
- `hf_model_paths` also supports checkpoints without `model_index.json`. Other
Diffusers checkpoints can select the pipeline through `_class_name`.
- For a standalone safetensors file, pass `--pipeline CustomPipeline`.
- Set the environment variable before startup. Each process imports the package
once. Use `overwrite=True` only to intentionally replace a built-in pipeline.
## Start With the Smallest Change
Before adding files, decide which path fits the model.