diff --git a/docs_new/docs/sglang-diffusion/support_new_models.mdx b/docs_new/docs/sglang-diffusion/support_new_models.mdx index 766e24b82..efdf665eb 100644 --- a/docs_new/docs/sglang-diffusion/support_new_models.mdx +++ b/docs_new/docs/sglang-diffusion/support_new_models.mdx @@ -1,285 +1,156 @@ --- -title: "How to Support New Diffusion Models" -metatags: - description: "This document explains how to add support for new diffusion models in SGLang Diffusion." +title: "Support New Diffusion Models" +description: "A concise implementation guide for adding diffusion model families to SGLang-Diffusion." --- -This document explains how to add support for new diffusion models in SGLang Diffusion. +Use this guide as a triage flow for finding the smallest change that can +support a model. Most new model work should touch a small set of files, even +though the runtime is split into separate folders. -## Architecture Overview +## Read the Code in This Order -SGLang Diffusion is engineered for both performance and flexibility, built upon a pipeline architecture. This -design allows developers to construct pipelines for various diffusion models while keeping the core generation -loop standardized for optimization. +The files are split by runtime responsibility. For a new model, read the +request path first: -At its core, the architecture revolves around two key concepts, as highlighted in our [blog post](https://lmsys.org/blog/2025-11-07-sglang-diffusion/#architecture): +1. `registry.py` chooses the model family, sampling params, and pipeline config. +2. `configs/pipeline_configs/{model}.py` defines model-specific denoising and + decoding behavior. +3. `runtime/pipelines/{model}.py` wires modules into stages. +4. `runtime/pipelines_core/stages/` runs the shared stage logic. +5. `runtime/models/` contains native model components only when the architecture + cannot be reused. -- **`ComposedPipeline`**: This class orchestrates a series of `PipelineStage`s to define the complete generation process for a specific model. It acts as the main entry point for a model and manages the data flow between the different stages of the diffusion process. -- **`PipelineStage`**: Each stage is a modular component that encapsulates a function within the diffusion process. Examples include prompt encoding, the denoising loop, or VAE decoding. +That is the dependency direction. Avoid making a model PR that requires readers +to jump between folders in a different order. -### Two Pipeline Styles +`runtime/models/` owns modeling code: checkpoint-defined neural modules, +architecture wrappers, and weight-loading or forward-path details that are +intrinsic to one model family. Reusable serving infrastructure belongs in +SGLang-Diffusion runtime folders such as `runtime/cache/`, +`runtime/distributed/`, `runtime/utils/`, or shared pipeline stages. This +includes cache managers, graph runners, process-group transport, request +utilities, and common action-policy helpers. Model packages may call these +helpers. Keep ownership in shared runtime folders unless the code is truly +architecture-specific. -SGLang Diffusion supports two pipeline composition styles. Both are valid; choose the one that best fits your model. +## Start With the Smallest Change -#### Style A: Hybrid Monolithic Pipeline (Recommended Default) +Before adding files, decide which path fits the model. -The recommended default for most new models. Uses a three-stage structure: +| Situation | What to do | +| --- | --- | +| A new checkpoint uses an existing native family | Add the Hugging Face path and, if needed, a small `SamplingParams` or `PipelineConfig` variant. Reuse the existing pipeline and modules. | +| The model has a new native DiT/UNet architecture | Add a native SGLang pipeline and the missing model components. Keep denoising and decoding on the shared stages unless measured behavior requires model-specific logic. | +| The model is long-tail or you only need compatibility first | Prefer the Diffusers backend for compatibility-first support. Add native support later if performance or deployment needs justify it. | -``` -BeforeDenoisingStage (model-specific) → DenoisingStage (standard) → DecodingStage (standard) +Do not add a folder just to mirror the Diffusers repository layout. Add a new +file only when an existing pipeline, stage, module, config, or sampler cannot +express the behavior clearly. + +## Minimal File Map + +The source tree is split by runtime responsibility. That split is useful for +optimization. Keep new model PRs focused on the files required by model +behavior. + +| Area | Add or edit when | Typical file | +| --- | --- | --- | +| Registry | Always, unless extending an already registered family | `python/sglang/multimodal_gen/registry.py` | +| Runtime parameters | The request schema differs from existing models | `configs/sample/{model}.py` | +| Pipeline config | Denoising, decoding, precision, position encoding, or CFG hooks differ | `configs/pipeline_configs/{model}.py` | +| Pipeline wiring | The model needs a new stage layout or module list | `runtime/pipelines/{model}.py` | +| DiT/UNet module | The denoising network is new | `runtime/models/dits/{model}.py` | +| dVLA or policy module | The checkpoint defines a new action-policy architecture | `runtime/models/{family}/modeling_*.py` or a task-specific model subfolder | +| Shared runtime infrastructure | Cache, CUDA graph, distributed transfer, request utilities, or action-policy helpers can be reused by future models | `runtime/cache/`, `runtime/distributed/`, `runtime/utils/`, `runtime/pipelines_core/stages/` | +| Model component config | A model component has static architecture config | `configs/models/dits/{model}.py`, `configs/models/vaes/{model}.py` | +| Model-specific stage | Pre-processing is too custom for the standard stages | `runtime/pipelines_core/stages/model_specific_stages/{model}.py` | +| Encoder, VAE, scheduler | No existing implementation can be reused | `runtime/models/encoders/`, `runtime/models/vaes/`, `runtime/models/schedulers/` | + +For a new native architecture, the common minimum is: + +1. `registry.py` +2. `configs/sample/{model}.py` +3. `configs/pipeline_configs/{model}.py` +4. `runtime/pipelines/{model}.py` +5. `runtime/models/dits/{model}.py` + +Every extra file should map to model behavior that existing code cannot express +clearly. + +For dVLA or other non-image diffusion policies, keep the same ownership rule. +The policy network, VLM/action expert modules, checkpoint mapping, and +model-specific forward code belong under `runtime/models/`. Prefix caches, +request-local contexts, denoising graph runners, OpenPI-compatible transport, +and prefix/action process-group utilities should be shared SGLang-Diffusion +runtime infrastructure when they are useful beyond the first model. + +## Read the Reference First + +Use the model's Diffusers pipeline, official implementation, or +`model_index.json` as the source of truth. Write down: + +- Which modules must be loaded: tokenizer, text encoder, image encoder, + transformer, scheduler, VAE, processor, and any extra adapters. +- The prompt and image encoding flow. +- Latent shape, packing, scale, shift, dtype, and device rules. +- Timestep and sigma schedule. +- The exact `forward()` kwargs expected by the denoising network. +- VAE decode rules and output post-processing. + +If the new model is close to Flux, Qwen-Image, GLM-Image, Wan, HunyuanVideo, or +LTX, extend that implementation before starting from an empty file. + +## Choose a Pipeline Shape + +SGLang-Diffusion uses `ComposedPipelineBase` to wire stages together. Most +native pipelines should use one of these two shapes. + +| Shape | Use when | Layout | +| --- | --- | --- | +| Standard stages | Text/image encoding, latent prep, denoising, and decoding match existing helpers | `add_standard_t2i_stages()`, `add_standard_ti2i_stages()`, or a similar helper | +| Model-specific pre-processing | The reference pipeline has custom captioning, image conditioning, latent packing, or timestep preparation | `{Model}BeforeDenoisingStage -> DenoisingStage -> DecodingStage` | + +Prefer standard stages when possible. Use a model-specific +`BeforeDenoisingStage` when trying to force the model into shared stages would +create many conditionals. + +## Implement the Pieces + +### 1. Sampling Params + +Create request parameters only for values users can set at runtime. + +```python +# python/sglang/multimodal_gen/configs/sample/my_model.py +from dataclasses import dataclass + +from sglang.multimodal_gen.configs.sample.sampling_params import ImageSamplingParams + + +@dataclass +class MyModelSamplingParams(ImageSamplingParams): + guidance_scale: float = 4.0 + num_inference_steps: int = 28 ``` - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
StageOwnershipResponsibility
{Model}BeforeDenoisingStageModel-specificAll pre-processing: input validation, text/image encoding, latent preparation, timestep computation
DenoisingStageFramework-standardThe denoising loop (DiT/UNet forward passes), shared across all models
DecodingStageFramework-standardVAE decoding from latent space to pixel space, shared across all models
+### 2. Pipeline Config -**Why recommended?** Modern diffusion models often have highly heterogeneous pre-processing requirements — different text encoders, different latent formats, different conditioning mechanisms. The Hybrid approach keeps pre-processing isolated per model, avoids fragile shared stages with excessive conditional logic, and lets developers port Diffusers reference code quickly. - -#### Style B: Modular Composition Style - -Uses the framework's fine-grained standard stages (`TextEncodingStage`, `LatentPreparationStage`, `TimestepPreparationStage`, etc.) to build the pipeline by composition. Convenience methods like `add_standard_t2i_stages()` and `add_standard_ti2i_stages()` make this very concise. - -This style is appropriate when: -- **The new model's pre-processing can largely reuse existing stages** — e.g., a model that uses standard CLIP/T5 text encoding + standard latent preparation with minimal customization. -- **A model-specific optimization needs to be extracted as a standalone stage** — e.g., a specialized encoding or conditioning step that benefits from being a separate stage for profiling, parallelism control, or reuse across multiple pipeline variants. - -#### How to Choose - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
SituationRecommended Style
Model has unique/complex pre-processing (VLM captioning, AR token generation, custom latent packing, etc.)Hybrid — consolidate into a BeforeDenoisingStage
Model fits neatly into standard text-to-image or text+image-to-image patternModular — use add_standard_t2i_stages() / add_standard_ti2i_stages()
Porting a Diffusers pipeline with many custom stepsHybrid — copy the __call__ logic into a single stage
Adding a variant of an existing model that shares most logicModular — reuse existing stages, customize via PipelineConfig callbacks
A specific pre-processing step needs special parallelism or profiling isolationModular — extract that step as a dedicated stage
- -## Key Components for Implementation - -To add support for a new diffusion model, you will need to define or configure the following components: - -1. **`PipelineConfig`**: A dataclass holding static configurations for your model pipeline — precision settings, model architecture parameters, and callback methods used by the standard `DenoisingStage` and `DecodingStage`. Each model has its own subclass. - -2. **`SamplingParams`**: A dataclass defining runtime generation parameters — `prompt`, `negative_prompt`, `guidance_scale`, `num_inference_steps`, `seed`, `height`, `width`, etc. - -3. **Pre-processing stage(s)**: Either a single model-specific `{Model}BeforeDenoisingStage` (Hybrid style) or a combination of standard stages (Modular style). See [Two Pipeline Styles](#two-pipeline-styles) above. - -4. **`ComposedPipeline`**: A class that wires together your pre-processing stage(s) with the standard `DenoisingStage` and `DecodingStage`. See base definitions: - - [`ComposedPipelineBase`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py) - - [`PipelineStage`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py) - - [Central registry](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/registry.py) - -5. **Modules (model components)**: Each pipeline references modules loaded from the model repository (e.g., Diffusers `model_index.json`): - - `text_encoder`: Encodes text prompts into embeddings. - - `tokenizer`: Tokenizes raw text input for the text encoder(s). - - `processor`: Preprocesses images and extracts features; often used in image-to-image tasks. - - `image_encoder`: Specialized image feature extractor. - - `dit/transformer`: The core denoising network (DiT/UNet architecture) operating in latent space. - - `scheduler`: Controls the timestep schedule and denoising dynamics. - - `vae`: Variational Autoencoder for encoding/decoding between pixel space and latent space. - -## Pipeline Stages Reference - -### Core Stages (used by all pipelines) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
Stage ClassDescription
DenoisingStageExecutes the main denoising loop, iteratively applying the model (DiT/UNet) to refine the latents.
DecodingStageDecodes the final latent tensor back into pixel space using the VAE.
DmdDenoisingStageA specialized denoising stage for DMD model architectures.
CausalDMDDenoisingStageA specialized causal denoising stage for specific video models.
- -### Pre-processing Stages (for Modular Composition Style) - -The following fine-grained stages can be composed to build the pre-processing portion of a pipeline. They are best suited for models whose pre-processing largely fits the standard patterns. If your model requires significant customization, consider the Hybrid style with a single `BeforeDenoisingStage` instead. - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
Stage ClassDescription
InputValidationStageValidates user-provided SamplingParams.
TextEncodingStageEncodes text prompts into embeddings using one or more text encoders.
ImageEncodingStageEncodes input images into embeddings, often used in image-to-image tasks.
ImageVAEEncodingStageEncodes an input image into latent space using the VAE.
TimestepPreparationStagePrepares the scheduler's timesteps for the diffusion process.
LatentPreparationStageCreates the initial noisy latent tensor that will be denoised.
- -## Implementation Guide - -### Step 1: Obtain and Study the Reference Implementation - -Before writing any code, obtain the model's original implementation or Diffusers pipeline code: -- The model's Diffusers pipeline source (e.g., the `pipeline_*.py` file from the `diffusers` library or HuggingFace repo) -- Or the model's official reference implementation (e.g., from the model author's GitHub repo) -- Or the HuggingFace model ID to look up `model_index.json` and the associated pipeline class - -Once you have the reference code, study it thoroughly: - -1. Find the model's `model_index.json` to identify required modules. -2. Read the Diffusers pipeline's `__call__` method to understand: - - How text prompts are encoded - - How latents are prepared (shape, dtype, scaling) - - How timesteps/sigmas are computed - - What conditioning kwargs the DiT expects - - How the denoising loop works - - How VAE decoding is done - -### Step 2: Evaluate Reuse of Existing Pipelines and Stages - -Before creating any new files, check whether an existing pipeline or stage can be reused or extended. Only create new pipelines/stages when the existing ones would need substantial structural changes or when no architecturally similar implementation exists. - -- **Compare against existing pipelines** (Flux, Wan, Qwen-Image, GLM-Image, HunyuanVideo, LTX, etc.). If the new model shares most of its structure with an existing one, prefer adding a new config variant or reusing existing stages. -- **Check existing stages** in `runtime/pipelines_core/stages/` and `stages/model_specific_stages/`. -- **Check existing model components** — many models share VAEs (e.g., `AutoencoderKL`), text encoders (CLIP, T5), and schedulers. Reuse these directly. - -### Step 3: Implement Model Components - -Adapt the model's core components: - -- **DiT/Transformer**: Implement in [`runtime/models/dits/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/dits/) -- **Encoders**: Implement in [`runtime/models/encoders/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/encoders/) -- **VAEs**: Implement in [`runtime/models/vaes/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/vaes/) -- **Schedulers**: Implement in [`runtime/models/schedulers/`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/models/schedulers/) if needed - -Use SGLang's fused kernels where possible (see `LayerNormScaleShift`, `RMSNormScaleShift`, `apply_qk_norm`, etc.). - -**Tensor Parallel (TP) and Sequence Parallel (SP)**: For multi-GPU deployment, it is recommended to add TP/SP support to the DiT model. This can be done incrementally after the single-GPU implementation is verified. Reference implementations: -- **Wan model** (`runtime/models/dits/wanvideo.py`) — Full TP + SP: `ColumnParallelLinear`/`RowParallelLinear` for attention, sequence dimension sharding via `get_sp_world_size()` -- **Qwen-Image model** (`runtime/models/dits/qwen_image.py`) — SP via `USPAttention` (Ulysses + Ring Attention) - -### Step 4: Create Configs - -- **DiT Config**: `configs/models/dits/{model_name}.py` -- **VAE Config**: `configs/models/vaes/{model_name}.py` -- **SamplingParams**: `configs/sample/{model_name}.py` - -### Step 5: Create PipelineConfig - -The `PipelineConfig` provides callbacks that the standard `DenoisingStage` and `DecodingStage` use: +`PipelineConfig` is where shared denoising and decoding stages get model-specific +callbacks. ```python # python/sglang/multimodal_gen/configs/pipeline_configs/my_model.py +from dataclasses import dataclass, field + @dataclass class MyModelPipelineConfig(ImagePipelineConfig): task_type: ModelTaskType = ModelTaskType.T2I - vae_precision: str = "bf16" should_use_guidance: bool = True - dit_config: DiTConfig = field(default_factory=MyModelDitConfig) + dit_config: DiTConfig = field(default_factory=MyModelDiTConfig) vae_config: VAEConfig = field(default_factory=MyModelVAEConfig) - def get_freqs_cis(self, batch, device, rotary_emb, dtype): - """Prepare rotary position embeddings for the DiT.""" - ... - def prepare_pos_cond_kwargs(self, batch, latent_model_input, t, **kwargs): - """Build positive conditioning kwargs for each denoising step.""" return { "hidden_states": latent_model_input, "encoder_hidden_states": batch.prompt_embeds[0], @@ -287,315 +158,173 @@ class MyModelPipelineConfig(ImagePipelineConfig): } def prepare_neg_cond_kwargs(self, batch, latent_model_input, t, **kwargs): - """Build negative conditioning kwargs for CFG.""" return { "hidden_states": latent_model_input, "encoder_hidden_states": batch.negative_prompt_embeds[0], "timestep": t, } - - def get_decode_scale_and_shift(self): - """Return (scale, shift) for latent denormalization before VAE decode.""" - ... ``` -### Step 6: Implement Pre-processing +Make these kwargs match the denoising module's `forward()` signature exactly. -Choose based on your model's needs (see [How to Choose](#how-to-choose)): +### 3. Pipeline Wiring -#### Option A: BeforeDenoisingStage (Hybrid Style) - -Create a single stage that handles all pre-processing. Best when the model has custom/complex pre-processing logic. - -```python -# python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/my_model.py - -class MyModelBeforeDenoisingStage(PipelineStage): - """Monolithic pre-processing stage for MyModel. - - Consolidates: input validation, text/image encoding, latent - preparation, and timestep computation. - """ - - def __init__(self, vae, text_encoder, tokenizer, transformer, scheduler): - super().__init__() - self.vae = vae - self.text_encoder = text_encoder - self.tokenizer = tokenizer - self.transformer = transformer - self.scheduler = scheduler - - @torch.no_grad() - def forward(self, batch: Req, server_args: ServerArgs) -> Req: - device = get_local_torch_device() - - # 1. Encode prompt (model-specific logic) - prompt_embeds, negative_prompt_embeds = self._encode_prompt(...) - - # 2. Prepare latents - latents = self._prepare_latents(...) - - # 3. Prepare timesteps - timesteps, sigmas = self._prepare_timesteps(...) - - # 4. Populate batch for DenoisingStage - batch.prompt_embeds = [prompt_embeds] - batch.negative_prompt_embeds = [negative_prompt_embeds] - batch.latents = latents - batch.timesteps = timesteps - batch.num_inference_steps = len(timesteps) - batch.sigmas = sigmas.tolist() - batch.generator = generator - batch.raw_latent_shape = latents.shape - return batch -``` - -#### Option B: Standard Stages (Modular Style) - -Skip creating a custom stage entirely — configure via `PipelineConfig` callbacks and use framework helpers. Best when the model fits standard patterns. - -(This option has no separate stage file; the pipeline class in Step 7 calls `add_standard_t2i_stages()` directly.) - -**Key batch fields that `DenoisingStage` expects** (regardless of which option you choose): - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
FieldTypeDescription
batch.latentstorch.TensorInitial noisy latent tensor
batch.timestepstorch.TensorTimestep schedule
batch.num_inference_stepsintNumber of denoising steps
batch.sigmaslist[float]Sigma schedule (must be a Python list, not numpy)
batch.prompt_embedslist[torch.Tensor]Positive prompt embeddings (wrapped in a list)
batch.negative_prompt_embedslist[torch.Tensor]Negative prompt embeddings (wrapped in a list)
batch.generatortorch.GeneratorRNG generator for reproducibility
batch.raw_latent_shapetupleOriginal latent shape before any packing
- -### Step 7: Define the Pipeline Class - -#### Hybrid Style +Use the standard helper when the model fits it. ```python # python/sglang/multimodal_gen/runtime/pipelines/my_model.py - class MyModelPipeline(LoRAPipeline, ComposedPipelineBase): - pipeline_name = "MyModelPipeline" # Must match model_index.json _class_name + pipeline_name = "MyModelPipeline" _required_config_modules = [ - "text_encoder", "tokenizer", "vae", "transformer", "scheduler", + "text_encoder", + "tokenizer", + "transformer", + "scheduler", + "vae", + ] + + def create_pipeline_stages(self, server_args: ServerArgs): + self.add_standard_t2i_stages() + + +EntryClass = [MyModelPipeline] +``` + +Use a model-specific pre-processing stage when the reference pipeline cannot be +cleanly expressed by standard helpers. + +```python +class MyModelPipeline(LoRAPipeline, ComposedPipelineBase): + pipeline_name = "MyModelPipeline" + + _required_config_modules = [ + "text_encoder", + "tokenizer", + "transformer", + "scheduler", + "vae", ] def create_pipeline_stages(self, server_args: ServerArgs): - # 1. Monolithic pre-processing (model-specific) self.add_stage( MyModelBeforeDenoisingStage( - vae=self.get_module("vae"), text_encoder=self.get_module("text_encoder"), tokenizer=self.get_module("tokenizer"), transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), - ), + vae=self.get_module("vae"), + ) ) - - # 2. Standard denoising loop (framework-provided) self.add_stage( DenoisingStage( transformer=self.get_module("transformer"), scheduler=self.get_module("scheduler"), - ), + ) ) - - # 3. Standard VAE decoding (framework-provided) self.add_standard_decoding_stage() EntryClass = [MyModelPipeline] ``` -#### Modular Style +### 4. Optional Before-Denoising Stage + +A `BeforeDenoisingStage` should populate the batch fields consumed by +`DenoisingStage`. ```python -# python/sglang/multimodal_gen/runtime/pipelines/my_model.py +class MyModelBeforeDenoisingStage(PipelineStage): + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + prompt_embeds, negative_prompt_embeds = self._encode_prompt(batch) + latents = self._prepare_latents(batch) + timesteps, sigmas = self._prepare_timesteps(batch) -class MyModelPipeline(LoRAPipeline, ComposedPipelineBase): - pipeline_name = "MyModelPipeline" - - _required_config_modules = [ - "text_encoder", "tokenizer", "vae", "transformer", "scheduler", - ] - - def create_pipeline_stages(self, server_args: ServerArgs): - # All pre-processing + denoising + decoding in one call - self.add_standard_t2i_stages( - prepare_extra_timestep_kwargs=[prepare_mu], # model-specific hooks - ) - - -EntryClass = [MyModelPipeline] + batch.prompt_embeds = [prompt_embeds] + batch.negative_prompt_embeds = [negative_prompt_embeds] + batch.latents = latents + batch.timesteps = timesteps + batch.num_inference_steps = len(timesteps) + batch.sigmas = sigmas.tolist() + batch.raw_latent_shape = latents.shape + return batch ``` -### Step 8: Register the Model +Required fields for `DenoisingStage`: -Register your configs in [`registry.py`](https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/registry.py): +| Field | Notes | +| --- | --- | +| `batch.latents` | Initial latent tensor, including any packing required by the model. | +| `batch.timesteps` | Timestep tensor in the exact order used by the reference pipeline. | +| `batch.sigmas` | Python list when the scheduler expects sigma values. | +| `batch.prompt_embeds` | Positive embeddings, wrapped in a list. | +| `batch.negative_prompt_embeds` | Negative embeddings, wrapped in a list when CFG is used. | +| `batch.num_inference_steps` | Number of denoising iterations. | +| `batch.raw_latent_shape` | Original latent shape before packing, if decode needs it. | + +### 5. Denoising Module + +Add a file under `runtime/models/dits/` only when the architecture is new. Reuse +existing encoders, VAEs, schedulers, normalization layers, and fused kernels +whenever possible. + +For multi-GPU serving, add TP/SP support after the single-GPU path is correct. +Useful references: + +- `runtime/models/dits/wanvideo.py` for TP plus SP. +- `runtime/models/dits/qwen_image.py` for USP attention. + +### 6. Registry + +Register the family once the sampling params and pipeline config exist. ```python register_configs( model_family="my_model", sampling_param_cls=MyModelSamplingParams, pipeline_config_cls=MyModelPipelineConfig, - hf_model_paths=["org/my-model-name"], + hf_model_paths=["org/my-model"], ) ``` -The `EntryClass` in your pipeline file is automatically discovered by the registry — no additional registration needed for the pipeline class itself. +The pipeline file is discovered through its `EntryClass`; do not add a second +pipeline registry unless the existing registry requires it. -### Step 9: Verify Output Quality +## Verify the Port -After implementation, verify that the generated output is not noise. A noisy or garbled output is the most common sign of an incorrect implementation. Common causes include: +Use one deterministic prompt and seed while comparing with the reference +implementation. -- Incorrect latent scale/shift factors -- Wrong timestep/sigma schedule (order, dtype, or value range) -- Mismatched conditioning kwargs -- Rotary embedding style mismatch (`is_neox_style`) +1. Run a single-GPU smoke test and check that the output contains coherent + content. +2. Compare latent scale and shift, timestep order, sigma values, and conditioning + kwargs against Diffusers or the official implementation. +3. Verify VAE decode and post-processing separately from denoising. +4. If the model supports LoRA, CFG parallelism, TP, SP, or disaggregation, test + each feature explicitly. +5. Add or update docs, examples, or the compatibility matrix when users need a + new launch command. -Debug by comparing intermediate tensor values against the Diffusers reference pipeline with the same seed. +Common failure points: -## Reference Implementations +- Wrong latent scale or shift. +- Reversed or dtype-mismatched timesteps. +- Missing negative embeddings when CFG is enabled. +- Conditioning kwarg names mismatched with the DiT `forward()`. +- Rotary embedding shape or style mismatch. +- Decoding packed latents without restoring `raw_latent_shape`. -### Hybrid Style +## PR Checklist - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
ModelPipelineBeforeDenoisingStagePipelineConfig
GLM-Imageruntime/pipelines/glm_image.pystages/model_specific_stages/glm_image.pyconfigs/pipeline_configs/glm_image.py
Qwen-Image-Layeredruntime/pipelines/qwen_image.pystages/model_specific_stages/qwen_image_layered.pyconfigs/pipeline_configs/qwen_image.py
- -### Modular Style - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
ModelPipelineNotes
Qwen-Image (T2I)runtime/pipelines/qwen_image.pyUses add_standard_t2i_stages()
Qwen-Image-Editruntime/pipelines/qwen_image.pyUses add_standard_ti2i_stages()
Fluxruntime/pipelines/flux.pyUses add_standard_t2i_stages() with custom prepare_mu
Wanruntime/pipelines/wan_pipeline.pyUses add_standard_ti2v_stages()
- -## Checklist - -Before submitting your implementation, verify: - -**Common (both styles):** -- [ ] **Pipeline file** at `runtime/pipelines/{model_name}.py` with `EntryClass` -- [ ] **PipelineConfig** at `configs/pipeline_configs/{model_name}.py` -- [ ] **SamplingParams** at `configs/sample/{model_name}.py` -- [ ] **DiT model** at `runtime/models/dits/{model_name}.py` -- [ ] **Model configs** (DiT, VAE) at `configs/models/dits/` and `configs/models/vaes/` -- [ ] **Registry entry** in `registry.py` via `register_configs()` -- [ ] `pipeline_name` matches Diffusers `model_index.json` `_class_name` -- [ ] `_required_config_modules` lists all modules from `model_index.json` -- [ ] `PipelineConfig` callbacks (`prepare_pos_cond_kwargs`, etc.) match the DiT's `forward()` signature -- [ ] Uses framework-standard `DenoisingStage` and `DecodingStage` (not custom denoising loops) -- [ ] **TP/SP support** considered for DiT model (recommended; reference `wanvideo.py` for TP+SP, `qwen_image.py` for USPAttention) -- [ ] **Output quality verified** — generated images/videos are not noise; compared against Diffusers reference output - -**Hybrid style only:** -- [ ] **BeforeDenoisingStage** at `stages/model_specific_stages/{model_name}.py` -- [ ] `BeforeDenoisingStage.forward()` populates all batch fields required by `DenoisingStage` +- [ ] Reused an existing family, stage, module, scheduler, or VAE wherever + possible. +- [ ] Kept the new-model touch surface small and justified any extra files. +- [ ] Added `SamplingParams`, `PipelineConfig`, pipeline wiring, DiT module, and + registry entry when native support is needed. +- [ ] Confirmed `pipeline_name` matches the Diffusers `model_index.json` + `_class_name` when applicable. +- [ ] Confirmed `_required_config_modules` matches the model repo. +- [ ] Verified image or video quality against a reference output. +- [ ] Tested multi-GPU paths if the PR claims TP, SP, CFG parallelism, or + distributed serving support.