[diffusion] app: add ComfyUI plugin support for SGLang-Diffusion (#15271)

Co-authored-by: niehen6174 <niehen.6174@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
Co-authored-by: niehen6174 <nihen6174@gmail.com>

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -2,6 +2,7 @@
 
 # SPDX-License-Identifier: Apache-2.0
 import json
+import os
 from collections.abc import Callable
 from dataclasses import asdict, dataclass, field, fields
 from enum import Enum, auto
@@ -518,15 +519,55 @@ class PipelineConfig:
         if model_path is None:
             raise ValueError("model_path is required in kwargs")
 
+        # Check if model_path is a safetensors file and pipeline_class_name is specified
+        pipeline_class_name = kwargs.get(
+            prefix_with_dot + "pipeline_class_name"
+        ) or kwargs.get("pipeline_class_name")
+        is_safetensors_file = os.path.isfile(model_path) and model_path.endswith(
+            ".safetensors"
+        )
+
         # 1. Get the pipeline config class from the registry
         from sglang.multimodal_gen.configs.pipeline_configs.flux import (
             Flux2PipelineConfig,
         )
+        from sglang.multimodal_gen.registry import get_pipeline_config_classes
 
-        model_info = get_model_info(model_path)
+        # If model_path is a safetensors file and pipeline_class_name is specified,
+        # try to get PipelineConfig from the registry first
+        if is_safetensors_file and pipeline_class_name:
+            config_classes = get_pipeline_config_classes(pipeline_class_name)
+            if config_classes is not None:
+                pipeline_config_cls, _ = config_classes
+                logger.info(
+                    f"Detected safetensors file with {pipeline_class_name}, "
+                    f"using {pipeline_config_cls.__name__} directly without model_index.json"
+                )
+            else:
+                model_info = get_model_info(model_path)
+                if model_info is None:
+                    from sglang.multimodal_gen.registry import (
+                        _PIPELINE_CONFIG_REGISTRY,
+                        _discover_and_register_pipelines,
+                    )
 
-        # 1.5. Adjust pipeline config for fine-tuned VAE if needed
-        pipeline_config_cls = model_info.pipeline_config_cls
+                    _discover_and_register_pipelines()
+                    available_pipelines = list(_PIPELINE_CONFIG_REGISTRY.keys())
+                    raise ValueError(
+                        f"Could not get model info for '{model_path}'. "
+                        f"If using a safetensors file, please specify a valid pipeline_class_name. "
+                        f"Available pipelines with config classes: {available_pipelines}"
+                    )
+                pipeline_config_cls = model_info.pipeline_config_cls
+        else:
+            model_info = get_model_info(model_path)
+            if model_info is None:
+                raise ValueError(
+                    f"Could not get model info for '{model_path}'. "
+                    f"If using a safetensors file, please specify pipeline_class_name"
+                )
+            # 1.5. Adjust pipeline config for fine-tuned VAE if needed
+            pipeline_config_cls = model_info.pipeline_config_cls
         vae_path = kwargs.get(prefix_with_dot + "vae_path") or kwargs.get("vae_path")
 
         # Check if this is a Flux2 model with fal/FLUX.2-Tiny-AutoEncoder

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -405,7 +405,8 @@ class SamplingParams:
                 )
                 self.num_frames = new_num_frames
 
-        self._set_output_file_name()
+        if not server_args.comfyui_mode:
+            self._set_output_file_name()
 
     @classmethod
     def from_pretrained(cls, model_path: str, **kwargs) -> "SamplingParams":
@@ -417,7 +418,45 @@ class SamplingParams:
 
     @staticmethod
     def from_user_sampling_params_args(model_path: str, server_args, *args, **kwargs):
-        sampling_params = SamplingParams.from_pretrained(model_path)
+        try:
+            sampling_params = SamplingParams.from_pretrained(model_path)
+        except (AttributeError, ValueError) as e:
+            # Handle safetensors files or other cases where model_index.json is not available
+            # Use appropriate SamplingParams based on pipeline_class_name from registry
+            if os.path.isfile(model_path) and model_path.endswith(".safetensors"):
+                # Determine which sampling params to use based on pipeline_class_name
+                pipeline_class_name = getattr(server_args, "pipeline_class_name", None)
+
+                # Try to get SamplingParams from registry
+                from sglang.multimodal_gen.registry import get_pipeline_config_classes
+
+                config_classes = (
+                    get_pipeline_config_classes(pipeline_class_name)
+                    if pipeline_class_name
+                    else None
+                )
+
+                if config_classes is not None:
+                    _, sampling_params_cls = config_classes
+                    try:
+                        sampling_params = sampling_params_cls()
+                        logger.info(
+                            f"Using {sampling_params_cls.__name__} for {pipeline_class_name} safetensors file (no model_index.json): %s",
+                            model_path,
+                        )
+                    except Exception as import_error:
+                        logger.warning(
+                            f"Failed to instantiate {sampling_params_cls.__name__}: {import_error}. "
+                            "Using default SamplingParams"
+                        )
+                        sampling_params = SamplingParams()
+                else:
+                    raise ValueError(
+                        f"Could not get pipeline config classes for {pipeline_class_name}"
+                    )
+            else:
+                # Re-raise if it's not a safetensors file issue
+                raise
 
         user_sampling_params = SamplingParams(*args, **kwargs)
         # TODO: refactor