# SPDX-License-Identifier: Apache-2.0 """ Diffusers backend pipeline wrapper. This module provides a wrapper that allows running any diffusers-supported model through sglang's infrastructure using vanilla diffusers pipelines. """ import argparse import inspect import re import warnings from typing import Any import numpy as np import torch import torchvision.transforms as T from diffusers import DiffusionPipeline from PIL import Image from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( ComponentResidencyStrategy, get_global_component_residency_manager, ) from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import ( PipelineExecutor, ) from sglang.multimodal_gen.runtime.pipelines_core.executors.sync_executor import ( SyncExecutor, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages import PipelineStage from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.precision import resolve_precision from sglang.multimodal_gen.runtime.utils.vision import load_image as load_vision_image logger = init_logger(__name__) class DiffusersExecutionStage(PipelineStage): """Pipeline stage that wraps diffusers pipeline execution.""" def __init__(self, diffusers_pipe: DiffusionPipeline): super().__init__() self.diffusers_pipe = diffusers_pipe def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Execute the diffusers pipeline.""" 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) # Request tensor output for cleaner handling if "output_type" not in kwargs: kwargs["output_type"] = "pt" with torch.no_grad(), warnings.catch_warnings(record=True): warnings.simplefilter("always") try: output = self.diffusers_pipe(**kwargs) except TypeError as e: # Some pipelines don't support output_type="pt" if "output_type" in str(e): kwargs.pop("output_type", None) output = self.diffusers_pipe(**kwargs) else: raise batch.output = self._extract_output(output) if batch.output is not None: batch.output = self._postprocess_output(batch.output) return batch def _filter_pipeline_kwargs( 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: kwargs: Arguments to filter strict: If True, raise ValueError on unsupported args; otherwise warn Returns: Tuple of (filtered_kwargs, ignored_keys) """ try: sig = inspect.signature(self.diffusers_pipe.__call__) except (ValueError, TypeError): return kwargs, [] params = sig.parameters accepts_var_kwargs = any( p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values() ) if accepts_var_kwargs: return kwargs, [] valid = set(params.keys()) - {"self"} filtered = {} ignored = [] for k, v in kwargs.items(): if k in valid: filtered[k] = v else: ignored.append(k) if ignored: pipe_name = type(self.diffusers_pipe).__name__ msg = ( f"Pipeline '{pipe_name}' does not support: {', '.join(sorted(ignored))}. " "These arguments will be ignored." ) if strict: raise ValueError(msg) logger.warning(msg) return filtered, ignored 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"]: data = getattr(output, attr, None) if data is None: continue result = self._convert_to_tensor(data) if result is not None: logger.debug( "Extracted output from '%s': shape=%s, dtype=%s", attr, result.shape, result.dtype, ) return result logger.warning("Could not extract output from pipeline result") return None def _convert_to_tensor(self, data: Any) -> torch.Tensor | None: """Convert various data formats to a tensor.""" if isinstance(data, torch.Tensor): return data if isinstance(data, np.ndarray): tensor = torch.from_numpy(data).float() if tensor.max() > 1.0: tensor = tensor / 255.0 # (B, H, W, C) -> (B, C, H, W) or (B, T, H, W, C) -> (B, C, T, H, W) if tensor.ndim == 4: tensor = tensor.permute(0, 3, 1, 2) elif tensor.ndim == 5: tensor = tensor.permute(0, 4, 1, 2, 3) return tensor if isinstance(data, Image.Image): return T.ToTensor()(data) if isinstance(data, list) and len(data) > 0: return self._convert_list_to_tensor(data) return None def _convert_list_to_tensor(self, data: list) -> torch.Tensor | None: """Convert a list of items to a tensor.""" first = data[0] # Nested list (e.g., [[frame1, frame2, ...]] for video batches) if isinstance(first, list) and len(first) > 0: data = first first = data[0] if isinstance(first, Image.Image): tensors = [T.ToTensor()(img) for img in data] stacked = torch.stack(tensors) if len(tensors) > 1: return stacked.permute(1, 0, 2, 3) # (T, C, H, W) -> (C, T, H, W) return stacked[0] if isinstance(first, torch.Tensor): stacked = torch.stack(data) if len(data) > 1: return stacked.permute(1, 0, 2, 3) return stacked[0] if isinstance(first, np.ndarray): tensors = [torch.from_numpy(arr).float() for arr in data] if tensors[0].max() > 1.0: tensors = [t / 255.0 for t in tensors] if tensors[0].ndim == 3: tensors = [t.permute(2, 0, 1) for t in tensors] stacked = torch.stack(tensors) if len(data) > 1: return stacked.permute(1, 0, 2, 3) return stacked[0] return None def _postprocess_output(self, output: torch.Tensor) -> torch.Tensor: """Post-process output tensor to ensure valid values and correct shape.""" output = output.cpu().float() # Handle NaN or Inf values if torch.isnan(output).any() or torch.isinf(output).any(): logger.warning("Output contains invalid values, fixing...") output = torch.nan_to_num(output, nan=0.5, posinf=1.0, neginf=0.0) # Normalize to [0, 1] range if needed min_val, max_val = output.min().item(), output.max().item() if min_val < -0.5 or max_val > 1.5: output = (output + 1) / 2 output = output.clamp(0, 1) # Ensure correct shape for downstream processing output = self._fix_output_shape(output) logger.debug("Final output tensor shape: %s", output.shape) return output def _fix_output_shape(self, output: torch.Tensor) -> torch.Tensor: """Fix tensor shape for downstream processing. Expected: (B, C, H, W) for images or (B, C, T, H, W) for videos. """ if output.dim() == 5: # Video: (B, T, C, H, W) -> (B, C, T, H, W) return output.permute(0, 2, 1, 3, 4) if output.dim() == 4: if output.shape[0] == 1 or output.shape[1] in [1, 3, 4]: return output # Already (B, C, H, W) # (T, C, H, W) -> (1, C, T, H, W) return output.unsqueeze(0).permute(0, 2, 1, 3, 4) if output.dim() == 3: c, h, w = output.shape if c > 4 and w <= 4: output = output.permute(2, 0, 1) if output.shape[0] == 1: output = output.repeat(3, 1, 1) return output.unsqueeze(0) if output.dim() == 2: return output.unsqueeze(0).repeat(3, 1, 1).unsqueeze(0) return output def _build_pipeline_kwargs(self, batch: Req) -> dict[str, Any]: """Build kwargs dict for diffusers pipeline call.""" kwargs = {} if batch.prompt is not None: kwargs["prompt"] = batch.prompt if batch.negative_prompt: kwargs["negative_prompt"] = batch.negative_prompt if batch.num_inference_steps is not None: kwargs["num_inference_steps"] = batch.num_inference_steps if batch.guidance_scale is not None: kwargs["guidance_scale"] = batch.guidance_scale if batch.true_cfg_scale is not None: kwargs["true_cfg_scale"] = batch.true_cfg_scale if batch.height is not None: kwargs["height"] = batch.height if batch.width is not None: kwargs["width"] = batch.width if batch.num_frames is not None and batch.num_frames > 1: kwargs["num_frames"] = batch.num_frames # Generator for reproducibility if batch.generator is not None: kwargs["generator"] = batch.generator elif batch.seed is not None: device = self._get_generator_device(batch) kwargs["generator"] = torch.Generator(device=device).manual_seed(batch.seed) # Image input for img2img or inpainting image = self._load_input_image(batch) if image is not None: kwargs["image"] = image if batch.num_outputs_per_prompt > 1: kwargs["num_images_per_prompt"] = batch.num_outputs_per_prompt # Extra diffusers-specific kwargs if batch.extra: diffusers_kwargs = batch.extra.get("diffusers_kwargs", {}) if diffusers_kwargs: kwargs.update(diffusers_kwargs) return kwargs def _get_generator_device(self, batch: Req) -> str: """Resolve RNG device consistently with the non-diffusers path. Diffusers CPU offload can temporarily park modules on CPU, but that should not silently switch a CUDA request to CPU RNG, otherwise the same seed produces different outputs depending on runtime placement. """ if batch.generator_device == "cpu": return "cpu" return current_platform.device_type def _load_input_image(self, batch: Req) -> Image.Image | None: """Load input image from batch.""" # Check for PIL image in condition_image or pixel_values if batch.condition_image is not None and isinstance( batch.condition_image, Image.Image ): return batch.condition_image if batch.pixel_values is not None and isinstance( batch.pixel_values, Image.Image ): return batch.pixel_values if not batch.image_path: return None if isinstance(batch.image_path, list): batch.image_path = batch.image_path[0] try: image = load_vision_image(batch.image_path) return image.convert("RGB") except Exception as e: logger.error("Failed to load image from %s: %s", batch.image_path, e) return None class DiffusersPipeline(ComposedPipelineBase): """ Pipeline wrapper that uses vanilla diffusers pipelines. This allows running any diffusers-supported model through sglang's infrastructure without requiring native sglang implementation. """ pipeline_name = "DiffusersPipeline" is_video_pipeline = False _required_config_modules: list[str] = [] def __init__( self, model_path: str, server_args: ServerArgs, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None, executor: PipelineExecutor | None = None, ): self.server_args = server_args self.model_path = model_path self._stages: list[PipelineStage] = [] self._stage_name_mapping: dict[str, PipelineStage] = {} self.modules: dict[str, Any] = {} self.memory_usages: dict[str, float] = {} self.component_residency_strategies: dict[str, ComponentResidencyStrategy] = {} self.component_residency_manager = None self.post_init_called = False self.executor = executor or SyncExecutor(server_args=server_args) self._cache_dit_enabled = False logger.info("Loading diffusers pipeline from %s", model_path) 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 ) -> DiffusionPipeline: """Load the diffusers pipeline. Optimizations applied: - device_map: Loads models directly to GPU, warming up CUDA caching allocator to avoid small tensor allocations during inference. - Parallel shard loading: When using device_map with accelerate, model shards are loaded in parallel for faster initialization. """ original_model_path = model_path # Keep original for custom_pipeline model_path = maybe_download_model(model_path, force_diffusers_model=True) self.model_path = model_path dtype = self._get_dtype(server_args) logger.info("Loading diffusers pipeline with dtype=%s", dtype) # Build common kwargs for from_pretrained load_kwargs = { "torch_dtype": dtype, "trust_remote_code": server_args.trust_remote_code, "revision": server_args.revision, } # Add quantization config if provided (e.g., BitsAndBytesConfig for 4/8-bit) 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) except AttributeError as e: if "has no attribute" in str(e): # Custom pipeline class not in diffusers - try loading with custom_pipeline logger.info( "Pipeline class not found in diffusers, trying custom_pipeline from repo..." ) try: custom_kwargs = { **load_kwargs, "custom_pipeline": original_model_path, } custom_kwargs["trust_remote_code"] = True pipe = DiffusionPipeline.from_pretrained( model_path, **custom_kwargs ) except Exception as e2: match = re.search(r"has no attribute (\w+)", str(e)) class_name = match.group(1) if match else "unknown" raise RuntimeError( f"Pipeline class '{class_name}' not found in diffusers and no custom pipeline.py in repo. " f"Try: pip install --upgrade diffusers (some pipelines require latest version). " f"Original error: {e}" ) from e2 else: raise except Exception as e: # Only retry with float32 for dtype-related errors if "dtype" in str(e).lower() or "float" in str(e).lower(): logger.warning( "Failed with dtype=%s, falling back to float32: %s", dtype, e ) load_kwargs["torch_dtype"] = torch.float32 pipe = DiffusionPipeline.from_pretrained(model_path, **load_kwargs) else: raise # Use CPU offload (all-or-nothing in diffusers) if any component offload is requested. any_offload = ( server_args.dit_cpu_offload or server_args.text_encoder_cpu_offload or server_args.image_encoder_cpu_offload or server_args.vae_cpu_offload ) if any_offload: device = get_local_torch_device() gpu_id = device.index if device.index is not None else 0 pipe.enable_model_cpu_offload(gpu_id=gpu_id) logger.info( "Enabled model CPU offload for diffusers pipeline (gpu_id=%d)", gpu_id ) else: pipe = pipe.to(get_local_torch_device()) # Apply VAE memory optimizations from pipeline config self._apply_vae_optimizations(pipe, server_args) # Apply attention backend if specified self._apply_attention_backend(pipe, server_args) # Apply cache-dit acceleration if configured pipe = self._apply_cache_dit(pipe, server_args) # Apply torch.compile if enabled and supported pipe = self._apply_torch_compile(pipe, server_args) logger.info("Loaded diffusers pipeline: %s", pipe.__class__.__name__) return pipe 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 # VAE slicing: decode latents slice-by-slice for lower peak memory # https://huggingface.co/docs/diffusers/optimization/memory#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 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: 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 Available backends: flash, _flash_3_hub, sage, xformers, native, etc. """ backend = server_args.attention_backend if backend is None: backend = getattr( server_args.pipeline_config, "diffusers_attention_backend", None ) if backend is None: return backend = backend.lower() sglang_backends = {e.name.lower() for e in AttentionBackendEnum} | { "fa3", "fa4", } if backend in sglang_backends: logger.debug( "Skipping diffusers attention backend '%s' because it matches a " "SGLang backend name. Use diffusers backend names when running " "the diffusers backend.", backend, ) return for component_name in ["transformer", "unet"]: component = getattr(pipe, component_name, None) if component is not None and hasattr(component, "set_attention_backend"): try: component.set_attention_backend(backend) logger.info( "Set attention backend '%s' on %s", backend, component_name ) except Exception as e: logger.warning( "Failed to set attention backend '%s' on %s: %s", backend, component_name, e, ) 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: return pipe try: import cache_dit except ImportError as e: raise RuntimeError( "cache-dit is required for --cache-dit-config. " "Install it with `pip install cache-dit`." ) from e if not hasattr(cache_dit, "load_configs"): raise RuntimeError( "cache-dit>=1.2.0 is required for --cache-dit-config. " "Please upgrade cache-dit." ) try: cache_options = cache_dit.load_configs(cache_dit_config) except Exception as e: raise ValueError( "Failed to load cache-dit config. Provide a YAML/JSON path (or a dict " "supported by cache-dit>=1.2.0)." ) from e try: pipe = cache_dit.enable_cache(pipe, **cache_options) except Exception: # cache-dit is an external integration and can raise a variety of errors. logger.exception("Failed to enable cache-dit for diffusers pipeline") raise logger.info("Enabled cache-dit for diffusers pipeline") self._cache_dit_enabled = True return pipe def _apply_torch_compile(self, pipe: Any, server_args: ServerArgs) -> Any: """Apply torch.compile to the pipeline if configured and supported.""" if not server_args.enable_torch_compile: return pipe # check if the pipeline has 'transformer' or 'unet' components which are # typically the most expensive parts to compile. 'transformer_2' for some # video pipelines, e.g, Wan 2.2 series, also check for that. compilable_components = ["transformer", "transformer_2", "unet"] if not any(hasattr(pipe, comp) for comp in compilable_components): logger.warning( "Pipeline does not have 'transformer' or 'unet' components. " "torch.compile may not provide significant benefits and could increase latency." ) return pipe if self._cache_dit_enabled: try: import cache_dit if hasattr(cache_dit, "set_compile_configs"): cache_dit.set_compile_configs() except Exception as e: logger.warning( f"Failed to set torch_compile configs for cache-dit: {e}" ) for comp in compilable_components: if hasattr(pipe, comp): try: component = getattr(pipe, comp) repeated_blocks = getattr(component, "_repeated_blocks", None) if ( isinstance(component, torch.nn.Module) and repeated_blocks and hasattr(component, "compile_repeated_blocks") ): # Regional compilation: compile a single instance of each # repeated transformer block and let inductor's cache reuse # it for all repeats, instead of compiling the whole DiT as # one graph component.compile_repeated_blocks() elif isinstance(component, torch.nn.Module) and hasattr( component, "compile" ): # Prefer in-place compilation if supported. According to PyTorch documentation: # https://docs.pytorch.org/docs/stable/generated/torch.compile.html component.compile() else: compiled_component = torch.compile(component) setattr(pipe, comp, compiled_component) logger.info( f"Applied torch.compile to {comp} component of the pipeline" ) except Exception as e: logger.warning(f"Failed to apply torch.compile to {comp}: {e}") return pipe def _get_dtype(self, server_args: ServerArgs) -> torch.dtype: """ Determine the dtype to use for model loading. """ if hasattr(server_args, "pipeline_config") and server_args.pipeline_config: return resolve_precision(server_args, "dit", precision_attr="dit_precision") # precision-constraint: legacy fallback for callers without pipeline_config; # prefer explicit dit_precision policy when available. return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 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"] self.is_video_pipeline = any(ind in pipe_class_name for ind in video_indicators) logger.debug( "Detected pipeline type: %s", "video" if self.is_video_pipeline else "image", ) def load_modules( self, server_args: ServerArgs, loaded_modules: dict[str, torch.nn.Module] | None = None, ) -> dict[str, Any]: """Skip sglang's module loading - diffusers handles it.""" return {"diffusers_pipeline": self.diffusers_pipe} def create_pipeline_stages(self, server_args: ServerArgs) -> None: """Create the execution stage wrapping the diffusers pipeline.""" self.add_stage( stage_name="diffusers_execution", stage=DiffusersExecutionStage(self.diffusers_pipe), ) def initialize_pipeline(self, server_args: ServerArgs) -> None: pass def post_init(self) -> None: """Post initialization hook.""" if self.post_init_called: return self.post_init_called = True self.initialize_pipeline(self.server_args) self.create_pipeline_stages(self.server_args) 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) if stage_name in self._stage_name_mapping: raise ValueError(f"Duplicate stage name detected: {stage_name}") stage.set_registered_stage_name(stage_name) stage.set_profile_stage_name(self._profile_stage_name(stage, stage_name)) self._stages.append(stage) self._stage_name_mapping[stage_name] = stage return self @property def stages(self) -> list[PipelineStage]: """List of stages in the pipeline.""" return self._stages @torch.no_grad() def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Execute the pipeline on the given batch.""" if not self.post_init_called: self.post_init() self.component_residency_manager = get_global_component_residency_manager( self, server_args ) self.executor.component_residency_manager = self.component_residency_manager return self.executor.execute_with_profiling(self.stages, batch, server_args) @classmethod def from_pretrained( cls, model_path: str, device: str | None = None, torch_dtype: torch.dtype | None = None, pipeline_config: str | PipelineConfig | None = None, args: argparse.Namespace | None = None, required_config_modules: list[str] | None = None, loaded_modules: dict[str, torch.nn.Module] | None = None, **kwargs, ) -> "DiffusersPipeline": """Load a pipeline from a pretrained model using diffusers backend.""" kwargs["model_path"] = model_path server_args = ServerArgs.from_kwargs(**kwargs) pipe = cls( model_path, server_args, required_config_modules=required_config_modules, loaded_modules=loaded_modules, ) pipe.post_init() return pipe def get_module(self, module_name: str, default_value: Any = None) -> Any: """Get a module by name.""" if module_name == "diffusers_pipeline": return self.diffusers_pipe return self.modules.get(module_name, default_value) EntryClass = DiffusersPipeline