chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] fix: add VAE tiling/slicing argument handling for diffusers backend (#17825)

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