From 3067b3f050503c7a3aab71b8a5bacfdfee36ed6b Mon Sep 17 00:00:00 2001 From: Jinyan Chen <93358689+liz-badada@users.noreply.github.com> Date: Tue, 2 Dec 2025 18:28:59 +0800 Subject: [PATCH] [diffusion] chore: improve model info registration and searching strategy (#14281) Co-authored-by: Jinyan Chen Co-authored-by: Mick --- .../configs/pipeline_configs/base.py | 13 +- .../configs/sample/sampling_params.py | 9 +- python/sglang/multimodal_gen/registry.py | 136 ++++++++---------- .../runtime/pipelines_core/__init__.py | 2 - 4 files changed, 65 insertions(+), 95 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 43477ccd6..5ffc92783 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -489,18 +489,9 @@ class PipelineConfig: # 1. Get the pipeline config class from the registry model_info = get_model_info(model_path) - # 2. Instantiate PipelineConfig - if model_info is None: - # The error is already logged in get_model_info. - # We raise an exception here to stop the execution. - raise ValueError( - f"Failed to get model info for '{model_path}'. " - "Please check the model path and ensure it is registered correctly." - ) - pipeline_config = model_info.pipeline_config_cls() - # 3. Load PipelineConfig from a json file or a PipelineConfig object if provided + # 2. Load PipelineConfig from a json file or a PipelineConfig object if provided if isinstance(pipeline_config_or_path, str): pipeline_config.load_from_json(pipeline_config_or_path) kwargs[prefix_with_dot + "pipeline_config_path"] = pipeline_config_or_path @@ -509,7 +500,7 @@ class PipelineConfig: elif isinstance(pipeline_config_or_path, dict): pipeline_config.update_pipeline_config(pipeline_config_or_path) - # 4. Update PipelineConfig from CLI arguments if provided + # 3. Update PipelineConfig from CLI arguments if provided kwargs[prefix_with_dot + "model_path"] = model_path pipeline_config.update_config_from_dict(kwargs, config_cli_prefix) return pipeline_config diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index e73dd8de5..ef691b1ae 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -299,14 +299,7 @@ class SamplingParams: from sglang.multimodal_gen.registry import get_model_info model_info = get_model_info(model_path) - if model_info is not None: - sampling_params: SamplingParams = model_info.sampling_param_cls(**kwargs) - else: - logger.warning( - "Couldn't find an optimal sampling param for %s. Using the default sampling param.", - model_path, - ) - sampling_params = cls(**kwargs) + sampling_params: SamplingParams = model_info.sampling_param_cls(**kwargs) return sampling_params @staticmethod diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 26b4980c4..1e2fafeb4 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -124,42 +124,38 @@ class ConfigInfo: _CONFIG_REGISTRY: Dict[str, ConfigInfo] = {} # Mappings from Hugging Face model paths to our internal model names -_MODEL_PATH_TO_NAME: Dict[str, str] = {} +_MODEL_HF_PATH_TO_NAME: Dict[str, str] = {} # Detectors to identify model families from paths or class names _MODEL_NAME_DETECTORS: List[Tuple[str, Callable[[str], bool]]] = [] def register_configs( - model_name: str, sampling_param_cls: Any, pipeline_config_cls: Type[PipelineConfig], - model_paths: Optional[List[str]] = None, + hf_model_paths: Optional[List[str]] = None, model_detectors: Optional[List[Callable[[str], bool]]] = None, ): """ Registers configuration classes for a new model family. """ - if model_name in _CONFIG_REGISTRY: - logger.warning( - f"Config for model '{model_name}' is already registered and will be overwritten." - ) + model_id = str(len(_CONFIG_REGISTRY)) - _CONFIG_REGISTRY[model_name] = ConfigInfo( + _CONFIG_REGISTRY[model_id] = ConfigInfo( sampling_param_cls=sampling_param_cls, pipeline_config_cls=pipeline_config_cls, ) - if model_paths: - for path in model_paths: - if path in _MODEL_PATH_TO_NAME: + if hf_model_paths: + for path in hf_model_paths: + if path in _MODEL_HF_PATH_TO_NAME: logger.warning( - f"Model path '{path}' is already mapped to '{_MODEL_PATH_TO_NAME[path]}' and will be overwritten by '{model_name}'." + f"Model path '{path}' is already mapped to '{_MODEL_HF_PATH_TO_NAME[path]}' and will be overwritten by '{model_id}'." ) - _MODEL_PATH_TO_NAME[path] = model_name + _MODEL_HF_PATH_TO_NAME[path] = model_id if model_detectors: for detector in model_detectors: - _MODEL_NAME_DETECTORS.append((model_name, detector)) + _MODEL_NAME_DETECTORS.append((model_id, detector)) def _get_config_info(model_path: str) -> Optional[ConfigInfo]: @@ -167,18 +163,21 @@ def _get_config_info(model_path: str) -> Optional[ConfigInfo]: Gets the ConfigInfo for a given model path using mappings and detectors. """ # 1. Exact match - if model_path in _MODEL_PATH_TO_NAME: - model_name = _MODEL_PATH_TO_NAME[model_path] - logger.debug(f"Resolved model name '{model_name}' from exact path match.") - return _CONFIG_REGISTRY.get(model_name) + if model_path in _MODEL_HF_PATH_TO_NAME: + model_id = _MODEL_HF_PATH_TO_NAME[model_path] + logger.debug(f"Resolved model path '{model_path}' from exact path match.") + return _CONFIG_REGISTRY.get(model_id) - # 2. Partial match: find the best (longest) match against all registered model names. + # 2. Partial match: find the best (longest) match against all registered model hf paths. cleaned_model_path = re.sub(r"--", "/", model_path.lower()) - all_model_names = sorted(_CONFIG_REGISTRY.keys(), key=len, reverse=True) - for model_name in all_model_names: - if model_name in cleaned_model_path: - logger.debug(f"Resolved model name '{model_name}' from partial path match.") - return _CONFIG_REGISTRY.get(model_name) + all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True) + for model_hf_path in all_model_hf_paths: + if model_hf_path.lower() in cleaned_model_path: + logger.debug( + f"Resolved model name '{model_hf_path}' from partial path match." + ) + model_id = _MODEL_HF_PATH_TO_NAME[model_hf_path] + return _CONFIG_REGISTRY.get(model_id) # 3. Use detectors if os.path.exists(model_path): @@ -188,14 +187,23 @@ def _get_config_info(model_path: str) -> Optional[ConfigInfo]: pipeline_name = config.get("_class_name", "").lower() - for model_name, detector in _MODEL_NAME_DETECTORS: + matched_model_names = [] + for model_id, detector in _MODEL_NAME_DETECTORS: if detector(model_path.lower()) or detector(pipeline_name): logger.debug( - f"Resolved model name '{model_name}' using a registered detector." + f"Matched model name '{model_id}' using a registered detector." ) - return _CONFIG_REGISTRY.get(model_name) + matched_model_names += [model_id] - return None + if len(matched_model_names) >= 1: + if len(matched_model_names) > 1: + logger.warning( + f"More than one model name is matched, using the first matched" + ) + model_id = matched_model_names[0] + return _CONFIG_REGISTRY.get(model_id) + else: + raise RuntimeError(f"No model info found for model path: {model_path}") # --- Part 3: Main Resolver --- @@ -256,7 +264,7 @@ def get_model_info(model_path: str) -> Optional[ModelInfo]: logger.error( f"Could not resolve configuration for model '{model_path}'. " "It is not a registered model path or detected by any registered model family detectors. " - f"Known model paths: {list(_MODEL_PATH_TO_NAME.keys())}" + f"Known model paths: {list(_MODEL_HF_PATH_TO_NAME.keys())}" ) return None @@ -275,160 +283,140 @@ def get_model_info(model_path: str) -> Optional[ModelInfo]: def _register_configs(): # Hunyuan register_configs( - model_name="hunyuan", sampling_param_cls=HunyuanSamplingParams, pipeline_config_cls=HunyuanConfig, - model_paths=[ + hf_model_paths=[ "hunyuanvideo-community/HunyuanVideo", ], - model_detectors=[lambda id: "hunyuan" in id.lower()], + model_detectors=[lambda hf_id: "hunyuan" in hf_id.lower()], ) register_configs( - model_name="fasthunyuan", sampling_param_cls=FastHunyuanSamplingParam, pipeline_config_cls=FastHunyuanConfig, - model_paths=[ + hf_model_paths=[ "FastVideo/FastHunyuan-diffusers", ], ) # StepVideo register_configs( - model_name="stepvideo", sampling_param_cls=StepVideoT2VSamplingParams, pipeline_config_cls=StepVideoT2VConfig, - model_paths=[ + hf_model_paths=[ "FastVideo/stepvideo-t2v-diffusers", ], - model_detectors=[lambda id: "stepvideo" in id.lower()], + model_detectors=[lambda hf_id: "stepvideo" in hf_id.lower()], ) # Wan register_configs( - model_name="wan-t2v-1.3b", sampling_param_cls=WanT2V_1_3B_SamplingParams, pipeline_config_cls=WanT2V480PConfig, - model_paths=[ + hf_model_paths=[ "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", ], - model_detectors=[lambda id: "wanpipeline" in id.lower()], + model_detectors=[lambda hf_id: "wanpipeline" in hf_id.lower()], ) register_configs( - model_name="wan-t2v-14b", sampling_param_cls=WanT2V_14B_SamplingParams, pipeline_config_cls=WanT2V720PConfig, - model_paths=[ + hf_model_paths=[ "Wan-AI/Wan2.1-T2V-14B-Diffusers", ], ) register_configs( - model_name="wan-i2v-14b-480p", sampling_param_cls=WanI2V_14B_480P_SamplingParam, pipeline_config_cls=WanI2V480PConfig, - model_paths=[ + hf_model_paths=[ "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", ], - model_detectors=[lambda id: "wanimagetovideo" in id.lower()], + model_detectors=[lambda hf_id: "wanimagetovideo" in hf_id.lower()], ) register_configs( - model_name="wan-i2v-14b-720p", sampling_param_cls=WanI2V_14B_720P_SamplingParam, pipeline_config_cls=WanI2V720PConfig, - model_paths=[ + hf_model_paths=[ "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers", ], ) register_configs( - model_name="wan-fun-1.3b-inp", sampling_param_cls=Wan2_1_Fun_1_3B_InP_SamplingParams, pipeline_config_cls=WanI2V480PConfig, - model_paths=[ + hf_model_paths=[ "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers", ], ) register_configs( - model_name="wan-ti2v-5b", sampling_param_cls=Wan2_2_TI2V_5B_SamplingParam, pipeline_config_cls=Wan2_2_TI2V_5B_Config, - model_paths=[ + hf_model_paths=[ "Wan-AI/Wan2.2-TI2V-5B-Diffusers", ], ) register_configs( - model_name="fastwan-ti2v-5b", sampling_param_cls=Wan2_2_TI2V_5B_SamplingParam, pipeline_config_cls=FastWan2_2_TI2V_5B_Config, - model_paths=[ + hf_model_paths=[ "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", "FastVideo/FastWan2.2-TI2V-5B-Diffusers", ], ) register_configs( - model_name="wan-t2v-a14b", sampling_param_cls=Wan2_2_T2V_A14B_SamplingParam, pipeline_config_cls=Wan2_2_T2V_A14B_Config, - model_paths=[ - "Wan-AI/Wan2.2-T2V-A14B-Diffusers", - ], + hf_model_paths=["Wan-AI/Wan2.2-T2V-A14B-Diffusers"], ) register_configs( - model_name="wan-i2v-a14b", sampling_param_cls=Wan2_2_I2V_A14B_SamplingParam, pipeline_config_cls=Wan2_2_I2V_A14B_Config, - model_paths=[ - "Wan-AI/Wan2.2-I2V-A14B-Diffusers", - ], + hf_model_paths=["Wan-AI/Wan2.2-I2V-A14B-Diffusers"], ) register_configs( - model_name="fast-wan-t2v-1.3b", sampling_param_cls=FastWanT2V480PConfig, pipeline_config_cls=FastWan2_1_T2V_480P_Config, - model_paths=[ + hf_model_paths=[ "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", ], ) # FLUX register_configs( - model_name="flux", sampling_param_cls=FluxSamplingParams, pipeline_config_cls=FluxPipelineConfig, - model_paths=[ + hf_model_paths=[ "black-forest-labs/FLUX.1-dev", ], - model_detectors=[lambda id: "flux" in id.lower()], + model_detectors=[lambda hf_id: "flux.1" in hf_id.lower()], ) register_configs( - model_name="flux-2", sampling_param_cls=FluxSamplingParams, pipeline_config_cls=Flux2PipelineConfig, - model_paths=[ + hf_model_paths=[ "black-forest-labs/FLUX.2-dev", ], - model_detectors=[lambda id: "flux.2" in id.lower()], + model_detectors=[lambda hf_id: "flux.2" in hf_id.lower()], ) register_configs( - model_name="Z-image", sampling_param_cls=ZImageSamplingParams, pipeline_config_cls=ZImagePipelineConfig, - model_paths=[ + hf_model_paths=[ "Tongyi-MAI/Z-Image-Turbo", ], - model_detectors=[lambda id: "z-image" in id.lower()], + model_detectors=[lambda hf_id: "z-image" in hf_id.lower()], ) # Qwen-Image register_configs( - model_name="qwen-image", sampling_param_cls=QwenImageSamplingParams, pipeline_config_cls=QwenImagePipelineConfig, + hf_model_paths=["Qwen/Qwen-Image"], ) register_configs( - model_name="qwen-image-edit", sampling_param_cls=QwenImageSamplingParams, pipeline_config_cls=QwenImageEditPipelineConfig, + hf_model_paths=["Qwen/Qwen-Image-Edit"], ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py index feda424e5..32243af35 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py @@ -43,8 +43,6 @@ def build_pipeline( """ model_path = server_args.model_path model_info = get_model_info(model_path) - if model_info is None: - raise ValueError(f"Unsupported model: {model_path}") pipeline_cls = model_info.pipeline_cls