diff --git a/docs/docs/sglang-diffusion/environment_variables.mdx b/docs/docs/sglang-diffusion/environment_variables.mdx index 44dde40f8..6c61676bf 100644 --- a/docs/docs/sglang-diffusion/environment_variables.mdx +++ b/docs/docs/sglang-diffusion/environment_variables.mdx @@ -18,6 +18,11 @@ description: "Configure SGLang diffusion behavior with environment variables." + + SGLANG_EXTERNAL_MODEL_PACKAGE + not set + Installed package that registers out-of-tree diffusion pipelines and component models. The package is imported once in every process. + SGLANG_DIFFUSION_TARGET_DEVICE cuda diff --git a/docs/docs/sglang-diffusion/support_new_models.mdx b/docs/docs/sglang-diffusion/support_new_models.mdx index 5e4de1c42..114623367 100644 --- a/docs/docs/sglang-diffusion/support_new_models.mdx +++ b/docs/docs/sglang-diffusion/support_new_models.mdx @@ -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. diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index 3ff7b41af..4da5b3f84 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -24,6 +24,7 @@ if TYPE_CHECKING: SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD: str = "fork" SGLANG_DIFFUSION_TARGET_DEVICE: str = "cuda" SGLANG_DIFFUSION_PLATFORM_OVERRIDE: str = "" + SGLANG_EXTERNAL_MODEL_PACKAGE: str = "" MAX_JOBS: str | None = None NVCC_THREADS: str | None = None CMAKE_BUILD_TYPE: str | None = None @@ -209,6 +210,9 @@ environment_variables: dict[str, Callable[[], Any]] = { "SGLANG_DIFFUSION_PLATFORM_OVERRIDE": _lazy_str( "SGLANG_DIFFUSION_PLATFORM_OVERRIDE", "" ), + # Import an installed package that registers out-of-tree diffusion models + # and pipelines. This is shared with the SRT model plugin mechanism. + "SGLANG_EXTERNAL_MODEL_PACKAGE": _lazy_str("SGLANG_EXTERNAL_MODEL_PACKAGE", ""), # Enables torch profiler if set. Path to the directory where torch profiler # traces are saved. Note that it must be an absolute path. "SGLANG_DIFFUSION_TORCH_PROFILER_DIR": _lazy_path( diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 72fcc21ed..25288bea8 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -196,6 +196,9 @@ from sglang.multimodal_gen.configs.sample.zimage import ( from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) +from sglang.multimodal_gen.runtime.utils.external_model_package import ( + load_external_model_package, +) from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( maybe_download_model_index, ) @@ -206,6 +209,7 @@ logger = init_logger(__name__) # --- Part 1: Pipeline Discovery --- _PIPELINE_REGISTRY: Dict[str, Type[ComposedPipelineBase]] = {} +_BUILTIN_PIPELINES_DISCOVERED = False # Registry for pipeline configuration classes (for safetensors files without model_index.json) # Maps pipeline_class_name -> (PipelineConfig class, SamplingParams class) @@ -219,11 +223,13 @@ def _discover_and_register_pipelines(): finds modules with an 'EntryClass' attribute, and maps the class's 'pipeline_name' to the class itself in a global registry. """ - if _PIPELINE_REGISTRY: # run only once + global _BUILTIN_PIPELINES_DISCOVERED + if _BUILTIN_PIPELINES_DISCOVERED: return package_name = "sglang.multimodal_gen.runtime.pipelines" package = importlib.import_module(package_name) + _BUILTIN_PIPELINES_DISCOVERED = True for _, module_name, ispkg in pkgutil.walk_packages( package.__path__, package.__name__ + "." @@ -276,6 +282,11 @@ def _discover_and_register_pipelines(): ) +def _ensure_registry_initialized() -> None: + _discover_and_register_pipelines() + load_external_model_package() + + def get_pipeline_config_classes( pipeline_class_name: str, ) -> Tuple[Type[PipelineConfig], Type[Any]] | None: @@ -283,10 +294,23 @@ def get_pipeline_config_classes( Get the configuration classes for a pipeline. """ # Ensure pipelines are discovered first - _discover_and_register_pipelines() + _ensure_registry_initialized() return _PIPELINE_CONFIG_REGISTRY.get(pipeline_class_name) +def get_pipeline_class( + pipeline_class_name: str, +) -> Type[ComposedPipelineBase] | None: + """Get a registered pipeline class by name.""" + _ensure_registry_initialized() + return _PIPELINE_REGISTRY.get(pipeline_class_name) + + +def get_registered_pipeline_names() -> List[str]: + _ensure_registry_initialized() + return list(_PIPELINE_REGISTRY) + + # --- Part 2: Config Registration --- @dataclasses.dataclass class ConfigInfo: @@ -328,7 +352,7 @@ def register_configs( pipeline_config_cls: Type[PipelineConfig], hf_model_paths: Optional[List[str]] = None, model_detectors: Optional[List[Callable[[str], bool]]] = None, -): +) -> str: """ Registers configuration classes for a new model family. """ @@ -349,6 +373,67 @@ def register_configs( if model_detectors: for detector in model_detectors: _MODEL_NAME_DETECTORS.append((model_id, detector)) + return model_id + + +def register_pipeline( + pipeline_cls: Type[ComposedPipelineBase], + *, + sampling_param_cls: Any, + pipeline_config_cls: Type[PipelineConfig], + hf_model_paths: Optional[List[str]] = None, + model_detectors: Optional[List[Callable[[str], bool]]] = None, + overwrite: bool = False, +) -> None: + """Register an out-of-tree native diffusion pipeline and its configs.""" + _discover_and_register_pipelines() + if not issubclass(pipeline_cls, ComposedPipelineBase): + raise TypeError("pipeline_cls must inherit from ComposedPipelineBase") + if not issubclass(pipeline_config_cls, PipelineConfig): + raise TypeError("pipeline_config_cls must inherit from PipelineConfig") + + pipeline_name = pipeline_cls.pipeline_name + existing_pipeline = _PIPELINE_REGISTRY.get(pipeline_name) + if existing_pipeline is not None and not overwrite: + raise ValueError( + f"Pipeline '{pipeline_name}' is already registered; pass overwrite=True to replace it" + ) + for model_path in hf_model_paths or []: + if model_path in _MODEL_HF_PATH_TO_NAME and not overwrite: + raise ValueError(f"Model path '{model_path}' is already registered") + registered_pipeline = KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS.get( + model_path.lower() + ) + if registered_pipeline is not None and not overwrite: + raise ValueError( + f"Model path '{model_path}' is already registered for pipeline " + f"'{registered_pipeline}'" + ) + + _PIPELINE_REGISTRY[pipeline_name] = pipeline_cls + _PIPELINE_CONFIG_REGISTRY[pipeline_name] = ( + pipeline_config_cls, + sampling_param_cls, + ) + config_id = register_configs( + sampling_param_cls=sampling_param_cls, + pipeline_config_cls=pipeline_config_cls, + hf_model_paths=hf_model_paths, + model_detectors=None if overwrite else model_detectors, + ) + if overwrite and model_detectors: + _MODEL_NAME_DETECTORS[:0] = [ + (config_id, detector) for detector in model_detectors + ] + for model_path in hf_model_paths or []: + KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS[model_path.lower()] = pipeline_name + _get_config_info.cache_clear() + get_model_info.cache_clear() + logger.info( + "Registered external diffusion pipeline '%s' from %s", + pipeline_name, + pipeline_cls.__module__, + ) def get_model_short_name(model_id: str) -> str: @@ -367,6 +452,7 @@ def _normalize_hf_cache_path(path: str) -> str: def has_registered_diffusion_model_path(model_path: str) -> bool: + _ensure_registry_initialized() all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True) if model_path in _MODEL_HF_PATH_TO_NAME: @@ -396,6 +482,7 @@ def _get_config_info( """ Gets the ConfigInfo for a given model path using mappings and detectors. """ + _ensure_registry_initialized() all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True) # 0. Explicit model_id override: match by short name @@ -580,7 +667,7 @@ def get_model_info( # For AUTO or SGLANG backend, try native implementation first # 1. Discover all available pipeline classes and cache them - _discover_and_register_pipelines() + _ensure_registry_initialized() # Detect quantized models and fallback to diffusers is_quantized = any(q in model_path.lower() for q in ["-4bit", "-awq", "-gptq"]) diff --git a/python/sglang/multimodal_gen/runtime/models/registry.py b/python/sglang/multimodal_gen/runtime/models/registry.py index 3af49cb6d..cb44d9fb6 100644 --- a/python/sglang/multimodal_gen/runtime/models/registry.py +++ b/python/sglang/multimodal_gen/runtime/models/registry.py @@ -19,6 +19,9 @@ from typing import NoReturn, TypeVar, cast import cloudpickle from torch import nn +from sglang.multimodal_gen.runtime.utils.external_model_package import ( + load_external_model_package, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) @@ -242,7 +245,6 @@ class _LazyRegisteredModel(_BaseRegisteredModel): """ module_name: str - component_name: str class_name: str # Performed in another process to avoid initializing CUDA @@ -289,10 +291,12 @@ class _ModelRegistry: registered_models: dict[str, _BaseRegisteredModel] = field(default_factory=dict) def get_supported_archs(self) -> Set[str]: + load_external_model_package() return self.registered_models.keys() def resolve_by_alias(self, alias: str) -> type[nn.Module] | None: """Resolve a model class by its alias (external module path).""" + load_external_model_package() if alias in _ALIAS_TO_MODEL: canonical_name = _ALIAS_TO_MODEL[alias] return self._try_load_model_cls(canonical_name) @@ -304,7 +308,7 @@ class _ModelRegistry: model_cls: type[nn.Module] | str, ) -> None: """ - Register an external model to be used in vLLM. + Register an external model to be used in SGLang-Diffusion. :code:`model_cls` can be either: @@ -328,7 +332,9 @@ class _ModelRegistry: msg = "Expected a string in the format `:`" raise ValueError(msg) - model = _LazyRegisteredModel(*split_str) + model = _LazyRegisteredModel( + module_name=split_str[0], class_name=split_str[1] + ) else: model = _RegisteredModel.from_model_cls(model_cls) @@ -364,6 +370,7 @@ class _ModelRegistry: self, architectures: str | list[str], ) -> list[str]: + load_external_model_package() if isinstance(architectures, str): architectures = [architectures] if not architectures: @@ -417,7 +424,6 @@ ModelRegistry = _ModelRegistry( { model_arch: _LazyRegisteredModel( module_name=f"sglang.multimodal_gen.runtime.models.{component_name}.{mod_relname}", - component_name=component_name, class_name=cls_name, ) for model_arch, ( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py index 063ff1d3c..adbc60ce1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py @@ -48,20 +48,18 @@ def build_pipeline( # Check if pipeline class is explicitly specified if server_args.pipeline_class_name: from sglang.multimodal_gen.registry import ( - _PIPELINE_REGISTRY, - _discover_and_register_pipelines, + get_pipeline_class, + get_registered_pipeline_names, ) - _discover_and_register_pipelines() + available_pipelines = get_registered_pipeline_names() logger.info(f"Requested pipeline_class_name: {server_args.pipeline_class_name}") - logger.info( - f"Available pipelines in registry: {list(_PIPELINE_REGISTRY.keys())}" - ) - pipeline_cls = _PIPELINE_REGISTRY.get(server_args.pipeline_class_name) + logger.info(f"Available pipelines in registry: {available_pipelines}") + pipeline_cls = get_pipeline_class(server_args.pipeline_class_name) if pipeline_cls is None: raise ValueError( f"Pipeline class '{server_args.pipeline_class_name}' not found in registry. " - f"Available pipelines: {list(_PIPELINE_REGISTRY.keys())}" + f"Available pipelines: {available_pipelines}" ) logger.info( f"✓ Using explicitly specified pipeline: {server_args.pipeline_class_name} (class: {pipeline_cls.__name__})" diff --git a/python/sglang/multimodal_gen/runtime/utils/external_model_package.py b/python/sglang/multimodal_gen/runtime/utils/external_model_package.py new file mode 100644 index 000000000..c72863175 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/utils/external_model_package.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 + +import importlib +from functools import lru_cache + +from sglang.multimodal_gen import envs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +@lru_cache(maxsize=1) +def load_external_model_package() -> None: + """Import the configured out-of-tree model package once per process.""" + package_name = envs.SGLANG_EXTERNAL_MODEL_PACKAGE + if not package_name: + return + + logger.info("Loading external model package: %s", package_name) + importlib.import_module(package_name) diff --git a/python/sglang/multimodal_gen/test/unit/test_external_model_package.py b/python/sglang/multimodal_gen/test/unit/test_external_model_package.py new file mode 100644 index 000000000..a02f0c0aa --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_external_model_package.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +import subprocess +import sys +import textwrap + + +def test_external_diffusion_package_registration(tmp_path): + package_dir = tmp_path / "external_diffusion_plugin" + package_dir.mkdir() + (package_dir / "__init__.py").write_text( + textwrap.dedent(""" + from sglang.multimodal_gen.registry import register_pipeline + from sglang.multimodal_gen.runtime.models.registry import ModelRegistry + + from .plugin import ( + ExternalPipeline, + ExternalPipelineConfig, + ExternalSamplingParams, + ) + + ModelRegistry.register_model( + "ExternalTransformer", + "external_diffusion_plugin.plugin:ExternalTransformer", + ) + register_pipeline( + ExternalPipeline, + sampling_param_cls=ExternalSamplingParams, + pipeline_config_cls=ExternalPipelineConfig, + hf_model_paths=["external-org/external-checkpoint"], + model_detectors=[lambda value: "external-checkpoint" in value], + ) + """), + encoding="utf-8", + ) + (package_dir / "plugin.py").write_text( + textwrap.dedent(""" + from torch import nn + + from sglang.multimodal_gen.configs.pipeline_configs.base import ( + PipelineConfig, + ) + from sglang.multimodal_gen.configs.sample.sampling_params import ( + SamplingParams, + ) + from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, + ) + + + class ExternalTransformer(nn.Module): + pass + + + class ExternalPipelineConfig(PipelineConfig): + pass + + + class ExternalSamplingParams(SamplingParams): + pass + + + class ExternalPipeline(ComposedPipelineBase): + pipeline_name = "ExternalPipeline" + + def create_pipeline_stages(self, server_args): + pass + """), + encoding="utf-8", + ) + + model_dir = tmp_path / "external-checkpoint" + model_dir.mkdir() + + script = textwrap.dedent(""" + import sys + + from sglang.cli.utils import get_is_diffusion_model + from sglang.multimodal_gen.registry import ( + get_model_info, + get_pipeline_config_classes, + ) + from sglang.multimodal_gen.runtime.models.registry import ModelRegistry + + assert get_is_diffusion_model(sys.argv[1]) + model_cls, architecture = ModelRegistry.resolve_model_cls( + "ExternalTransformer" + ) + assert model_cls.__name__ == "ExternalTransformer" + assert architecture == "ExternalTransformer" + + model_info = get_model_info(sys.argv[1], backend="sglang") + assert model_info.pipeline_cls.__name__ == "ExternalPipeline" + assert model_info.pipeline_config_cls.__name__ == "ExternalPipelineConfig" + assert model_info.sampling_param_cls.__name__ == "ExternalSamplingParams" + + pipeline_config_cls, sampling_param_cls = get_pipeline_config_classes( + "ExternalPipeline" + ) + assert pipeline_config_cls is model_info.pipeline_config_cls + assert sampling_param_cls is model_info.sampling_param_cls + """) + env = os.environ.copy() + env["SGLANG_EXTERNAL_MODEL_PACKAGE"] = "external_diffusion_plugin" + env["PYTHONPATH"] = os.pathsep.join( + path for path in (str(tmp_path), env.get("PYTHONPATH")) if path + ) + result = subprocess.run( + [sys.executable, "-c", script, str(model_dir)], + check=False, + env=env, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr