diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index 2f6f47eb8..35df3cc8a 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -184,6 +184,7 @@ class PipelineConfig: vae_config: VAEConfig = field(default_factory=VAEConfig) vae_precision: str = "fp32" vae_tiling: bool = True + vae_slicing: bool = False vae_sp: bool = True # Image encoder configuration @@ -470,6 +471,13 @@ class PipelineConfig: default=PipelineConfig.vae_tiling, help="Enable VAE tiling", ) + parser.add_argument( + f"--{prefix_with_dot}vae-slicing", + action=StoreBoolean, + dest=f"{prefix_with_dot.replace('-', '_')}vae_slicing", + default=PipelineConfig.vae_slicing, + help="Enable VAE slicing", + ) parser.add_argument( f"--{prefix_with_dot}vae-sp", action=StoreBoolean, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py index e2037097b..6a7d0446a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py @@ -54,7 +54,7 @@ class DiffusersExecutionStage(PipelineStage): def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Execute the diffusers pipeline.""" - kwargs = self._build_pipeline_kwargs(batch, server_args) + kwargs = self._build_pipeline_kwargs(batch) # Filter kwargs to only those supported by the pipeline, warn about ignored args kwargs, _ = self._filter_pipeline_kwargs(kwargs) @@ -82,8 +82,8 @@ class DiffusersExecutionStage(PipelineStage): return batch def _filter_pipeline_kwargs( - self, kwargs: dict, *, strict: bool = False - ) -> tuple[dict, list[str]]: + self, kwargs: dict[str, Any], *, strict: bool = False + ) -> tuple[dict[str, Any], list[str]]: """Filter kwargs to those accepted by the pipeline's __call__. Args: @@ -130,10 +130,7 @@ class DiffusersExecutionStage(PipelineStage): def _extract_output(self, output: Any) -> torch.Tensor | None: """Extract tensor output from pipeline result.""" for attr in ["images", "frames", "video", "sample", "pred_original_sample"]: - if not hasattr(output, attr): - continue - - data = getattr(output, attr) + data = getattr(output, attr, None) if data is None: continue @@ -166,7 +163,7 @@ class DiffusersExecutionStage(PipelineStage): tensor = tensor.permute(0, 4, 1, 2, 3) return tensor - if hasattr(data, "mode"): # PIL Image + if isinstance(data, Image.Image): return T.ToTensor()(data) if isinstance(data, list) and len(data) > 0: @@ -183,7 +180,7 @@ class DiffusersExecutionStage(PipelineStage): data = first first = data[0] - if hasattr(first, "mode"): # PIL images + if isinstance(first, Image.Image): tensors = [T.ToTensor()(img) for img in data] stacked = torch.stack(tensors) if len(tensors) > 1: @@ -259,7 +256,7 @@ class DiffusersExecutionStage(PipelineStage): return output - def _build_pipeline_kwargs(self, batch: Req, server_args: ServerArgs) -> dict: + def _build_pipeline_kwargs(self, batch: Req) -> dict[str, Any]: """Build kwargs dict for diffusers pipeline call.""" kwargs = {} @@ -316,7 +313,7 @@ class DiffusersExecutionStage(PipelineStage): component = getattr(self.diffusers_pipe, attr, None) if component is not None: try: - return next(component.parameters()).device + return str(next(component.parameters()).device) except StopIteration: pass return current_platform.device_type @@ -384,7 +381,9 @@ class DiffusersPipeline(ComposedPipelineBase): self.diffusers_pipe = self._load_diffusers_pipeline(model_path, server_args) self._detect_pipeline_type() - def _load_diffusers_pipeline(self, model_path: str, server_args: ServerArgs) -> Any: + def _load_diffusers_pipeline( + self, model_path: str, server_args: ServerArgs + ) -> DiffusionPipeline: """Load the diffusers pipeline. Optimizations applied: @@ -409,14 +408,10 @@ class DiffusersPipeline(ComposedPipelineBase): } # Add quantization config if provided (e.g., BitsAndBytesConfig for 4/8-bit) - config = server_args.pipeline_config - if config is not None: - quant_config = getattr(config, "quantization_config", None) - if quant_config is not None: - load_kwargs["quantization_config"] = quant_config - logger.info( - "Using quantization config: %s", type(quant_config).__name__ - ) + quant_config = getattr(server_args.pipeline_config, "quantization_config", None) + if quant_config is not None: + load_kwargs["quantization_config"] = quant_config + logger.info("Using quantization config: %s", type(quant_config).__name__) try: pipe = DiffusionPipeline.from_pretrained(model_path, **load_kwargs) @@ -483,27 +478,45 @@ class DiffusersPipeline(ComposedPipelineBase): logger.info("Loaded diffusers pipeline: %s", pipe.__class__.__name__) return pipe - def _apply_vae_optimizations(self, pipe: Any, server_args: ServerArgs) -> None: + def _apply_vae_optimizations( + self, pipe: DiffusionPipeline, server_args: ServerArgs + ) -> None: """Apply VAE memory optimizations (tiling, slicing) from pipeline config.""" config = server_args.pipeline_config - if config is None: - return # VAE slicing: decode latents slice-by-slice for lower peak memory # https://huggingface.co/docs/diffusers/optimization/memory#vae-slicing - if getattr(config, "vae_slicing", False): - if hasattr(pipe, "enable_vae_slicing"): + if config.vae_slicing: + if hasattr(pipe, "vae") and hasattr(pipe.vae, "enable_slicing"): + pipe.vae.enable_slicing() + logger.info("Enabled VAE slicing for lower memory usage") + elif hasattr(pipe, "enable_vae_slicing"): pipe.enable_vae_slicing() logger.info("Enabled VAE slicing for lower memory usage") + else: + logger.warning( + "VAE slicing is not available: neither " + "`pipe.vae.enable_slicing()` nor `pipe.enable_vae_slicing()` was found." + ) # VAE tiling: decode latents tile-by-tile for large images # https://huggingface.co/docs/diffusers/optimization/memory#vae-tiling - if getattr(config, "vae_tiling", False): - if hasattr(pipe, "enable_vae_tiling"): + if config.vae_tiling: + if hasattr(pipe, "vae") and hasattr(pipe.vae, "enable_tiling"): + pipe.vae.enable_tiling() + logger.info("Enabled VAE tiling for large image support") + elif hasattr(pipe, "enable_vae_tiling"): pipe.enable_vae_tiling() logger.info("Enabled VAE tiling for large image support") + else: + logger.warning( + "VAE tiling is not available: neither " + "`pipe.vae.enable_tiling()` nor `pipe.enable_vae_tiling()` was found." + ) - def _apply_attention_backend(self, pipe: Any, server_args: ServerArgs) -> None: + def _apply_attention_backend( + self, pipe: DiffusionPipeline, server_args: ServerArgs + ) -> None: """Apply attention backend setting from pipeline config or server_args. See: https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends @@ -511,6 +524,11 @@ class DiffusersPipeline(ComposedPipelineBase): """ backend = server_args.attention_backend + if backend is None: + backend = getattr( + server_args.pipeline_config, "diffusers_attention_backend", None + ) + if backend is None: return @@ -544,7 +562,9 @@ class DiffusersPipeline(ComposedPipelineBase): e, ) - def _apply_cache_dit(self, pipe: Any, server_args: ServerArgs) -> Any: + def _apply_cache_dit( + self, pipe: DiffusionPipeline, server_args: ServerArgs + ) -> DiffusionPipeline: """Enable cache-dit for diffusers pipeline if configured.""" cache_dit_config = server_args.cache_dit_config if not cache_dit_config: @@ -635,27 +655,23 @@ class DiffusersPipeline(ComposedPipelineBase): return pipe def _get_dtype(self, server_args: ServerArgs) -> torch.dtype: - """ - Determine the dtype to use for model loading. - """ dtype = ( torch.bfloat16 if torch.get_device_module().is_bf16_supported() else torch.float16 ) - if hasattr(server_args, "pipeline_config") and server_args.pipeline_config: - dit_precision = server_args.pipeline_config.dit_precision - if dit_precision == "fp16": - dtype = torch.float16 - elif dit_precision == "bf16": - dtype = torch.bfloat16 - elif dit_precision == "fp32": - dtype = torch.float32 + dit_precision = server_args.pipeline_config.dit_precision + if dit_precision == "fp16": + dtype = torch.float16 + elif dit_precision == "bf16": + dtype = torch.bfloat16 + elif dit_precision == "fp32": + dtype = torch.float32 return dtype - def _detect_pipeline_type(self): + def _detect_pipeline_type(self) -> None: """Detect if this is an image or video pipeline.""" pipe_class_name = self.diffusers_pipe.__class__.__name__.lower() video_indicators = ["video", "animat", "cogvideo", "wan", "hunyuan"] @@ -673,14 +689,14 @@ class DiffusersPipeline(ComposedPipelineBase): """Skip sglang's module loading - diffusers handles it.""" return {"diffusers_pipeline": self.diffusers_pipe} - def create_pipeline_stages(self, server_args: ServerArgs): + def create_pipeline_stages(self, server_args: ServerArgs) -> None: """Create the execution stage wrapping the diffusers pipeline.""" self.add_stage( - DiffusersExecutionStage(self.diffusers_pipe), "diffusers_execution" + stage_name="diffusers_execution", + stage=DiffusersExecutionStage(self.diffusers_pipe), ) - def initialize_pipeline(self, server_args: ServerArgs): - """Initialize the pipeline.""" + def initialize_pipeline(self, server_args: ServerArgs) -> None: pass def post_init(self) -> None: @@ -691,9 +707,7 @@ class DiffusersPipeline(ComposedPipelineBase): self.initialize_pipeline(self.server_args) self.create_pipeline_stages(self.server_args) - def add_stage( - self, stage: PipelineStage, stage_name: str | None = None - ) -> "DiffusersPipeline": + def add_stage(self, stage_name: str, stage: PipelineStage) -> None: """Add a stage to the pipeline.""" if stage_name is None: stage_name = self._infer_stage_name(stage)