From b0cd533a963a53dafc5b904763cf838e3249d095 Mon Sep 17 00:00:00 2001 From: Brian Chao Date: Mon, 8 Jun 2026 20:44:07 +0800 Subject: [PATCH] [diffusion] feat: progressive resolution growing for image and video models (#27524) --- .../progressive_resolution.mdx | 293 +++++++++ .../configs/sample/sampling_params.py | 48 ++ .../multimodal_gen/runtime/pipelines/flux.py | 30 +- .../runtime/pipelines/flux_2.py | 4 + .../runtime/pipelines/qwen_image.py | 8 +- .../runtime/pipelines/wan_pipeline.py | 13 +- .../runtime/pipelines/zimage_pipeline.py | 9 +- .../pipelines_core/composed_pipeline_base.py | 79 ++- .../pipelines_core/stages/denoising.py | 6 +- .../stages/progressive_resolution/__init__.py | 1 + .../progressive_resolution/denoising.py | 622 ++++++++++++++++++ .../stages/progressive_resolution/flux.py | 131 ++++ .../stages/progressive_resolution/flux_2.py | 200 ++++++ .../progressive_resolution/qwen_image.py | 163 +++++ .../progressive_resolution/spectral_ops.py | 78 +++ .../stages/progressive_resolution/upsample.py | 108 +++ .../stages/progressive_resolution/wan.py | 187 ++++++ .../stages/progressive_resolution/zimage.py | 169 +++++ .../pipelines_core/stages/text_encoding.py | 13 +- .../multimodal_gen/test/unit/conftest.py | 56 ++ .../test_progressive.py | 358 ++++++++++ .../test/unit/test_cfg_parallel_warmup.py | 11 +- .../test/unit/test_input_validation.py | 11 +- 23 files changed, 2543 insertions(+), 55 deletions(-) create mode 100644 docs_new/docs/sglang-diffusion/progressive_resolution.mdx create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/__init__.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/flux.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/flux_2.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/qwen_image.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/spectral_ops.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/upsample.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/wan.py create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/zimage.py create mode 100644 python/sglang/multimodal_gen/test/unit/conftest.py create mode 100644 python/sglang/multimodal_gen/test/unit/progressive_resolution/test_progressive.py diff --git a/docs_new/docs/sglang-diffusion/progressive_resolution.mdx b/docs_new/docs/sglang-diffusion/progressive_resolution.mdx new file mode 100644 index 000000000..a7b71992e --- /dev/null +++ b/docs_new/docs/sglang-diffusion/progressive_resolution.mdx @@ -0,0 +1,293 @@ +--- +title: "Progressive Resolution Generation" +description: "Experimental spectral progressive resolution growing for selected SGLang Diffusion pipelines." +--- + +Progressive resolution growing is an experimental feature for selected SGLang Diffusion pipelines. It runs early denoising steps at a coarser latent resolution and spectrally upsamples the latent before the full-resolution steps. On the benchmark setup below, this reduces the quadratic attention cost of the DiT transformer and yields up to **1.63× speedup on FLUX.1**, **1.93× speedup on FLUX.2**, **2.33× speedup on Z-Image**, **2.78× speedup on Wan 2.1 T2V**, and **1.69× speedup on Qwen-Image**. + +This page is intentionally not linked from the main documentation navigation while the feature is still experimental. + +Based on [Spectral Progressive Diffusion (arXiv 2605.18736)](https://arxiv.org/abs/2605.18736). + +## Overview + +DiT attention is O(n²) in sequence length. Running the first N denoising steps at half the spatial resolution cuts the attention cost to ~6% for those steps. + +The transition point — how many steps to run at each resolution — is computed from the **Bayes-optimal frequency-activation criterion**: frequencies that cannot be resolved at the coarse scale are not denoised there. This makes the speedup lossless by construction. + +| Model | Full-res tokens | Half-res tokens | Token-step ratio | +|-------|----------------|----------------|-----------------| +| FLUX.1 1024×1024 | 4,096 | 1,024 | 4.0× | +| FLUX.2 1024×1024 | 4,096 | 1,024 | 4.0× | +| Z-Image 1024×1024 | 4,096 | 1,024 | 4.0× | +| Wan 2.1 T2V 480×832 (81 frames) | 6,240 | 1,560 | 4.0× | + +## Parameters + +| Parameter | CLI flag | Default | Description | +|-----------|----------|---------|-------------| +| `progressive_mode` | `--progressive-mode` | `"fullres"` | `"fullres"` disables (identical to standard generation). `"dct_rewind"` enables spectral upsample with scheduler rewind (recommended). `"dct"` enables upsample without rewind. | +| `progressive_levels` | `--progressive-levels` | `1` | Number of resolution halvings. `1` = one coarse stage (64×64 latent → 128×128). `2` = two coarse stages (32×32 → 64×64 → 128×128). | +| `progressive_delta` | `--progressive-delta` | `0.01` | Noise-dominated tolerance δ. Controls how many steps run at coarse resolution. Higher δ = more coarse steps = more speedup. | + +> **Tip:** Add `--dit-cpu-offload false` to keep the transformer GPU-resident. With CPU offload each step pays a fixed PCIe transfer cost regardless of sequence length, which dilutes the speedup. + +--- + +## FLUX.1 + +### Usage + +```bash +sglang generate \ + --model-path black-forest-labs/FLUX.1-dev \ + --prompt "A serene mountain lake at golden hour, photorealistic" \ + --num-inference-steps 50 \ + --dit-cpu-offload false \ + --progressive-mode dct_rewind \ + --progressive-levels 1 \ + --progressive-delta 0.05 +``` + +### Choosing delta + +| δ | Coarse steps (50 total) | Denoising speedup | +|---|------------------------|-------------------| +| `0.01` | 18 @ 64² + 32 @ 128² | **1.32×** | +| `0.05` | 28 @ 64² + 22 @ 128² | **1.63×** | + +For most prompts `0.05` is recommended — it gives the largest speedup with no visible degradation. + +### Benchmark + +Hardware: RTX A6000 48 GB, `--dit-cpu-offload false`. Timing = denoising loop only. + +| Config | Stage split | Denoise | Speedup | +|--------|-------------|---------|---------| +| Fullres (baseline) | 50 @ 128² latent | 36.65 s | 1.00× | +| dct_rewind L1 δ=0.01 | 18@64² + 32@128² | 27.67 s | **1.32×** | +| dct_rewind L1 δ=0.05 | 28@64² + 22@128² | 22.58 s | **1.62×** | +| dct_rewind L2 δ=0.01 | 10@32² + 8@64² + 32@128² | 26.48 s | **1.38×** | + +### Python API + +```python +from sglang.multimodal_gen import DiffGenerator + +gen = DiffGenerator.from_pretrained( + model_path="black-forest-labs/FLUX.1-dev", + dit_cpu_offload=False, +) +result = gen.generate(sampling_params_kwargs={ + "prompt": "A serene mountain lake at golden hour, photorealistic", + "num_inference_steps": 50, + "height": 1024, + "width": 1024, + "progressive_mode": "dct_rewind", + "progressive_levels": 1, + "progressive_delta": 0.05, +}) +``` + +--- + +## FLUX.2 + +Supports `FLUX.2-dev`, `FLUX.2-klein-4B`, and `FLUX.2-klein-9B`. + +### Usage + +```bash +sglang generate \ + --model-path black-forest-labs/FLUX.2-klein-4B \ + --prompt "A serene mountain lake at golden hour, photorealistic" \ + --num-inference-steps 30 \ + --dit-cpu-offload false \ + --progressive-mode dct_rewind \ + --progressive-levels 1 \ + --progressive-delta 0.10 +``` + +### Benchmark + +Hardware: RTX A6000 48 GB, `--dit-cpu-offload false`. Model: FLUX.2-klein-4B, 30 steps, 1024×1024. +Timing = denoising loop only, averaged across 10 diverse prompts. + +| Config | Stage split | Denoise | Speedup | +|--------|-------------|---------|---------| +| Fullres (baseline) | 30 @ 64² latent | 9.72 s | 1.00× | +| dct_rewind L1 δ=0.05 | 18@32² + 12@64² | 5.50 s | **1.77×** | +| dct_rewind L1 δ=0.10 | 20@32² + 10@64² | 5.03 s | **1.93×** | + +### Python API + +```python +from sglang.multimodal_gen import DiffGenerator + +gen = DiffGenerator.from_pretrained( + model_path="black-forest-labs/FLUX.2-klein-4B", + dit_cpu_offload=False, +) +result = gen.generate(sampling_params_kwargs={ + "prompt": "A serene mountain lake at golden hour, photorealistic", + "num_inference_steps": 30, + "progressive_mode": "dct_rewind", + "progressive_levels": 1, + "progressive_delta": 0.10, +}) +``` + +--- + +## Wan 2.1 T2V + +Supports `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` and `Wan-AI/Wan2.1-T2V-14B-Diffusers`. + +> **Note:** Progressive generation grows only the **spatial** H×W dimensions. The temporal dimension T (number of latent frames) is kept fixed across all stages. + +### Usage + +```bash +sglang generate \ + --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ + --prompt "A cheetah sprinting across the Serengeti at sunset, slow motion, photorealistic" \ + --num-inference-steps 50 \ + --num-frames 81 \ + --height 480 \ + --width 832 \ + --guidance-scale 5.0 \ + --flow-shift 5.0 \ + --dit-cpu-offload false \ + --progressive-mode dct_rewind \ + --progressive-levels 1 \ + --progressive-delta 0.05 +``` + +### Choosing delta + +| δ | Coarse steps (50 total) | Denoising speedup | +|---|------------------------|-------------------| +| `0.01` | 23 @ 30×52 + 27 @ 60×104 | **1.65×** | +| `0.02` | 27 @ 30×52 + 23 @ 60×104 | **1.86×** | +| `0.05` | 33 @ 30×52 + 17 @ 60×104 | **2.32×** | +| `0.10` | 37 @ 30×52 + 13 @ 60×104 | **2.78×** | + +For most prompts `0.05` is recommended. `0.10` provides maximum speedup but should be validated on motion-heavy scenes. + +### Python API + +```python +from sglang.multimodal_gen import DiffGenerator + +gen = DiffGenerator.from_pretrained( + model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + dit_cpu_offload=False, + flow_shift=5.0, +) +result = gen.generate(sampling_params_kwargs={ + "prompt": "A cheetah sprinting across the Serengeti at sunset, slow motion, photorealistic", + "num_inference_steps": 50, + "num_frames": 81, + "height": 480, + "width": 832, + "guidance_scale": 5.0, + "progressive_mode": "dct_rewind", + "progressive_levels": 1, + "progressive_delta": 0.05, +}) +``` + +--- + +## Z-Image + +Supports `Tongyi-MAI/Z-Image`. Z-Image uses the same VAE as FLUX.1 (`FluxVAEConfig`), so the power-law spectrum constants are identical. The progressive stage handles Z-Image's 5-D latent format `[B, C, 1, H, W]` with squeeze/unsqueeze hooks and recomputes caption+image RoPE positional embeddings on each stage transition. + +> **Note:** Always specify `--height 1024 --width 1024` (or another resolution where H_lat and W_lat are both divisible by 2). Z-Image's default resolution (360×640) produces a 45×80 latent where H=45 is not divisible by the patch size. + +### Usage + +```bash +# Standard fullres — unchanged behavior +sglang generate --model-path Tongyi-MAI/Z-Image \ + --prompt "A serene mountain lake at golden hour, photorealistic" \ + --height 1024 --width 1024 + +# Progressive dct_rewind L1 δ=0.10 → 2.33× denoising speedup +sglang generate --model-path Tongyi-MAI/Z-Image \ + --prompt "A serene mountain lake at golden hour, photorealistic" \ + --height 1024 --width 1024 \ + --num-inference-steps 50 \ + --dit-cpu-offload false \ + --progressive-mode dct_rewind \ + --progressive-levels 1 \ + --progressive-delta 0.10 +``` + +### Choosing delta + +| δ | Coarse steps (50 total) | Denoising speedup | +|---|------------------------|-------------------| +| `0.01` | 26 @ 64² + 24 @ 128² | **1.53×** | +| `0.05` | 35 @ 64² + 15 @ 128² | **2.03×** | +| `0.10` | 42 @ 64² + 8 @ 128² | **2.33×** | + +Z-Image achieves higher progressive speedups than FLUX.1 at the same δ because it uses dual CFG (two forward passes per step), doubling the absolute attention savings at coarse resolution. `0.10` is the recommended tradeoff. + +### Python API + +```python +from sglang.multimodal_gen import DiffGenerator + +gen = DiffGenerator.from_pretrained( + model_path="Tongyi-MAI/Z-Image", + dit_cpu_offload=False, +) +result = gen.generate(sampling_params_kwargs={ + "prompt": "A serene mountain lake at golden hour, photorealistic", + "num_inference_steps": 50, + "height": 1024, + "width": 1024, + "progressive_mode": "dct_rewind", + "progressive_levels": 1, + "progressive_delta": 0.10, +}) +``` + +--- + +## Qwen-Image + +Qwen-Image uses the same 2×2 patchify convention as FLUX.1 (in_channels=64, C=16), so the same progressive stage wires in with model-specific hooks for RoPE (`freqs_cis`) and spatial metadata (`img_shapes`). + +```bash +# Standard fullres — unchanged behavior +sglang generate --model-path Qwen/Qwen-Image \ + --prompt "A serene mountain lake at golden hour" + +# Progressive dct_rewind L1 δ=0.20 → 1.69× denoising speedup +sglang generate --model-path Qwen/Qwen-Image \ + --prompt "A serene mountain lake at golden hour" \ + --progressive-mode dct_rewind --progressive-levels 1 --progressive-delta 0.20 \ + --num-inference-steps 30 --dit-cpu-offload false +``` + +Hardware: RTX A6000 48 GB, `--dit-cpu-offload false`. Timing = denoising loop only. + +| Config | Stage split | Denoise | Speedup | +|--------|-------------|---------|---------| +| Fullres (baseline) | 30 @ 128² | 43.00 s | 1.00× | +| dct_rewind L1 δ=0.05 | 13@64² + 17@128² | 33.25 s | **1.29×** | +| dct_rewind L1 δ=0.10 | 16@64² + 14@128² | 33.86 s | **1.27×** | +| dct_rewind L1 δ=0.20 | 19@64² + 11@128² | 25.40 s | **1.69×** | + +## Limitations + +- **Sequence parallelism incompatible.** Cannot be combined with `--ulysses-degree` or `--ring-degree`. The stage raises a `RuntimeError` if SP is enabled. +- **torch.compile incompatible.** Compiled kernels have a fixed sequence length; the resolution transition causes a recompile or error. Use progressive without `--enable-torch-compile`. +- **Cache-DiT interaction is experimental.** The stage refreshes Cache-DiT context at resolution transitions, but quality and speedup should be benchmarked before relying on this combination. + +## References + +- [Spectral Progressive Diffusion (arXiv 2605.18736)](https://arxiv.org/abs/2605.18736) +- [SGLang Diffusion Performance Optimization](./performance-optimization) diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index f559fa5c5..7cbfc101c 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -167,6 +167,10 @@ class SamplingParams: cfg_normalization: float | bool = 0.0 boundary_ratio: float | None = None + progressive_mode: str = "fullres" + progressive_levels: int = 1 + progressive_delta: float = 0.01 + # TeaCache parameters enable_teacache: bool = False teacache_params: Any = ( @@ -379,6 +383,28 @@ class SamplingParams: f"num_inference_steps must be a positive int, got {self.num_inference_steps!r}" ) + if self.progressive_mode not in ("fullres", "dct", "dct_rewind"): + raise ValueError( + "progressive_mode must be one of 'fullres', 'dct', or " + f"'dct_rewind', got {self.progressive_mode!r}" + ) + if ( + isinstance(self.progressive_levels, bool) + or not isinstance(self.progressive_levels, int) + or self.progressive_levels <= 0 + ): + raise ValueError( + f"progressive_levels must be a positive int, got {self.progressive_levels!r}" + ) + if ( + isinstance(self.progressive_delta, bool) + or not isinstance(self.progressive_delta, (int, float)) + or not 0 < float(self.progressive_delta) < 1 + ): + raise ValueError( + f"progressive_delta must be in (0, 1), got {self.progressive_delta!r}" + ) + # Numeric hyperparams should not be NaN/Inf and should be within basic ranges. # Note: bool is a subclass of int; reject it explicitly to avoid silent surprises. def _finite_non_negative_float( @@ -738,6 +764,28 @@ class SamplingParams: help="", ) + # Progressive resolution growing (DCT spectral upsampling) + add_argument( + "--progressive-mode", + type=str, + dest="progressive_mode", + choices=["fullres", "dct", "dct_rewind"], + help="Progressive resolution mode. 'fullres' disables (default). " + "'dct_rewind' uses DCT-II upsample + scheduler sigma rewind (recommended).", + ) + add_argument( + "--progressive-levels", + type=int, + dest="progressive_levels", + help="Number of resolution halvings for progressive generation (default: 1).", + ) + add_argument( + "--progressive-delta", + type=float, + dest="progressive_delta", + help="Noise-dominated tolerance δ for stage-transition thresholds (default: 0.01).", + ) + add_argument( "--prompt", type=str, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/flux.py b/python/sglang/multimodal_gen/runtime/pipelines/flux.py index 80cd8fe8a..f96a54a91 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/flux.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/flux.py @@ -7,9 +7,8 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req -from sglang.multimodal_gen.runtime.pipelines_core.stages import ( - InputValidationStage, - TextEncodingStage, +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.flux import ( + FluxProgressiveDenoisingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -67,26 +66,13 @@ class FluxPipeline(LoRAPipeline, ComposedPipelineBase): ] def create_pipeline_stages(self, server_args: ServerArgs): - self.add_stage(InputValidationStage()) - - self.add_stage( - TextEncodingStage( - text_encoders=[ - self.get_module("text_encoder"), - self.get_module("text_encoder_2"), - ], - tokenizers=[ - self.get_module("tokenizer"), - self.get_module("tokenizer_2"), - ], - ), - "prompt_encoding_stage_primary", + self.add_standard_t2i_stages( + text_encoder_key=["text_encoder", "text_encoder_2"], + tokenizer_key=["tokenizer", "tokenizer_2"], + text_encoding_stage_name="prompt_encoding_stage_primary", + prepare_extra_timestep_kwargs=[prepare_mu], + progressive_denoising_stage_cls=FluxProgressiveDenoisingStage, ) - self.add_standard_timestep_preparation_stage(prepare_extra_kwargs=[prepare_mu]) - self.add_standard_latent_preparation_stage() - self.add_standard_denoising_stage() - self.add_standard_decoding_stage() - EntryClass = FluxPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines/flux_2.py b/python/sglang/multimodal_gen/runtime/pipelines/flux_2.py index 4910f6ef1..3f825a66b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/flux_2.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/flux_2.py @@ -7,6 +7,9 @@ from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline, Req from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.flux_2 import ( + Flux2ProgressiveDenoisingStage, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -56,6 +59,7 @@ class Flux2Pipeline(LoRAPipeline, ComposedPipelineBase): prompt_encoding="text", image_vae_stage_kwargs={"vae_image_processor": vae_image_processor}, prepare_extra_timestep_kwargs=[compute_empirical_mu], + progressive_denoising_stage_cls=Flux2ProgressiveDenoisingStage, ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py b/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py index ad1c0504a..f2d20ca94 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py @@ -12,6 +12,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.qwen_image_layered import ( QwenImageLayeredBeforeDenoisingStage, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.qwen_image import ( + QwenImageProgressiveDenoisingStage, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.utils import PRECISION_TO_TYPE @@ -64,7 +67,10 @@ class QwenImagePipeline(LoRAPipeline, ComposedPipelineBase): ] def create_pipeline_stages(self, server_args: ServerArgs): - self.add_standard_t2i_stages(prepare_extra_timestep_kwargs=[prepare_mu]) + self.add_standard_t2i_stages( + prepare_extra_timestep_kwargs=[prepare_mu], + progressive_denoising_stage_cls=QwenImageProgressiveDenoisingStage, + ) class QwenImageEditPipeline(LoRAPipeline, ComposedPipelineBase): diff --git a/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py index b52045754..9b3e63a80 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py @@ -15,6 +15,12 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ComposedPipelineBase, ) from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline +from sglang.multimodal_gen.runtime.pipelines_core.stages import ( + InputValidationStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.wan import ( + WanProgressiveDenoisingStage, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -43,7 +49,12 @@ class WanPipeline(LoRAPipeline, ComposedPipelineBase): ) def create_pipeline_stages(self, server_args: ServerArgs) -> None: - self.add_standard_t2i_stages() + self.add_stage(InputValidationStage()) + self.add_standard_text_encoding_stage() + self.add_standard_latent_preparation_stage() + self.add_standard_timestep_preparation_stage() + self.add_progressive_denoising_stage(WanProgressiveDenoisingStage) + self.add_standard_decoding_stage() EntryClass = WanPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py index 8f1f71474..f903b07b2 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py @@ -1,11 +1,13 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # SPDX-License-Identifier: Apache-2.0 - from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline, Req from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.zimage import ( + ZImageProgressiveDenoisingStage, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -55,7 +57,10 @@ class ZImagePipeline(LoRAPipeline, ComposedPipelineBase): ] def create_pipeline_stages(self, server_args: ServerArgs): - self.add_standard_t2i_stages(prepare_extra_timestep_kwargs=[prepare_mu]) + self.add_standard_t2i_stages( + prepare_extra_timestep_kwargs=[prepare_mu], + progressive_denoising_stage_cls=ZImageProgressiveDenoisingStage, + ) EntryClass = ZImagePipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 3bac67e46..efb88187b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -48,6 +48,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages import ( TextEncodingStage, TimestepPreparationStage, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.denoising import ( + ProgressiveDenoisingStage, + ProgressiveDenoisingStageRouter, +) from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( @@ -650,14 +654,24 @@ class ComposedPipelineBase(ABC): def add_standard_text_encoding_stage( self, - text_encoder_key: str = "text_encoder", - tokenizer_key: str = "tokenizer", + text_encoder_key: str | list[str] = "text_encoder", + tokenizer_key: str | list[str] = "tokenizer", + stage_name: str | None = None, ) -> "ComposedPipelineBase": + text_encoder_keys = ( + [text_encoder_key] + if isinstance(text_encoder_key, str) + else text_encoder_key + ) + tokenizer_keys = ( + [tokenizer_key] if isinstance(tokenizer_key, str) else tokenizer_key + ) return self.add_stage( TextEncodingStage( - text_encoders=[self.get_module(text_encoder_key)], - tokenizers=[self.get_module(tokenizer_key)], + text_encoders=[self.get_module(key) for key in text_encoder_keys], + tokenizers=[self.get_module(key) for key in tokenizer_keys], ), + stage_name, ) def add_standard_timestep_preparation_stage( @@ -718,6 +732,42 @@ class ComposedPipelineBase(ABC): stage_name, ) + def add_progressive_denoising_stage( + self, + progressive_stage_cls: type[ProgressiveDenoisingStage], + transformer_key: str = "transformer", + transformer_2_key: str | None = "transformer_2", + scheduler_key: str = "scheduler", + vae_key: str | None = "vae", + stage_name: str = "denoising_stage", + ) -> "ComposedPipelineBase": + + def create_stage() -> PipelineStage: + kwargs = { + "transformer": self.get_module(transformer_key), + "scheduler": self.get_module(scheduler_key), + "pipeline": self, + } + + if transformer_2_key: + transformer_2 = self.get_module(transformer_2_key, None) + if transformer_2 is not None: + kwargs["transformer_2"] = transformer_2 + + if vae_key: + kwargs["vae"] = self.get_module(vae_key, None) + + return ProgressiveDenoisingStageRouter( + standard_stage=DenoisingStage(**kwargs), + progressive_stage_factory=lambda: progressive_stage_cls(**kwargs), + ) + + return self.add_stage_factory( + RoleType.DENOISER, + create_stage, + stage_name, + ) + def add_standard_decoding_stage( self, vae_key: str = "vae", @@ -740,19 +790,30 @@ class ComposedPipelineBase(ABC): def add_standard_t2i_stages( self, include_input_validation: bool = True, + text_encoder_key: str | list[str] = "text_encoder", + tokenizer_key: str | list[str] = "tokenizer", + text_encoding_stage_name: str | None = None, prepare_extra_timestep_kwargs: list[Callable] | None = None, + progressive_denoising_stage_cls: type[ProgressiveDenoisingStage] | None = None, ) -> "ComposedPipelineBase": if include_input_validation: self.add_stage(InputValidationStage()) - self.add_standard_text_encoding_stage() + self.add_standard_text_encoding_stage( + text_encoder_key=text_encoder_key, + tokenizer_key=tokenizer_key, + stage_name=text_encoding_stage_name, + ) self.add_standard_latent_preparation_stage() self.add_standard_timestep_preparation_stage( prepare_extra_kwargs=prepare_extra_timestep_kwargs ) - self.add_standard_denoising_stage() + if progressive_denoising_stage_cls is None: + self.add_standard_denoising_stage() + else: + self.add_progressive_denoising_stage(progressive_denoising_stage_cls) self.add_standard_decoding_stage() return self @@ -770,6 +831,7 @@ class ComposedPipelineBase(ABC): image_vae_key: str = "vae", image_vae_stage_kwargs: dict[str, Any] | None = None, prepare_extra_timestep_kwargs: list[Callable] | None = None, + progressive_denoising_stage_cls: type[ProgressiveDenoisingStage] | None = None, ) -> "ComposedPipelineBase": if include_input_validation: self.add_stage( @@ -806,7 +868,10 @@ class ComposedPipelineBase(ABC): self.add_standard_timestep_preparation_stage( prepare_extra_kwargs=prepare_extra_timestep_kwargs ) - self.add_standard_denoising_stage() + if progressive_denoising_stage_cls is None: + self.add_standard_denoising_stage() + else: + self.add_progressive_denoising_stage(progressive_denoising_stage_cls) self.add_standard_decoding_stage() return self diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 89664611a..03a8b5cf0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -611,7 +611,11 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): scheduler = batch.scheduler assert scheduler is not None - boundary_timestep = self._handle_boundary_ratio(server_args, batch, scheduler) + boundary_timestep = ( + self._handle_boundary_ratio(server_args, batch, scheduler) + if self.transformer_2 is not None + else None + ) # Get timesteps and calculate warmup steps timesteps = batch.timesteps num_inference_steps = batch.num_inference_steps diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py new file mode 100644 index 000000000..cc9db15f3 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py @@ -0,0 +1,622 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Progressive-resolution denoising stage. + +Extends DenoisingStage with a multi-stage coarse-to-fine denoising loop: + Stage 1 runs at 1/(2^levels) of the full latent resolution. + Between stages, the latent is upsampled via the spectral method selected by + progressive_mode. + Stage N runs at full resolution. + +When progressive_mode == "fullres" (default), route the request to the standard +DenoisingStage instead of this stage. + +Supported progressive_mode values + "dct" : DCT-II embed, IDCT upsample, no scheduler rewind + "dct_rewind" : DCT upsample + gamma scaling + scheduler sigma rewind (paper §3) + +Extension hooks for model-specific subclasses + _unpack_latent(latent, h_lat, w_lat) → spatial [B, C, H, W] + _repack_latent(x_spatial, h_lat, w_lat, batch) → model-native latent + _on_resolution_change(ctx, batch, srv, h_px, w_px) → update resolution-dep. state +""" + +from __future__ import annotations + +import math +import time +from collections.abc import Callable, Sequence +from typing import Any + +import torch +from diffusers.utils.torch_utils import randn_tensor + +from sglang.multimodal_gen import envs +from sglang.multimodal_gen.runtime.cache.cache_dit_integration import ( + refresh_context_on_dual_transformer, + refresh_context_on_transformer, +) +from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType +from sglang.multimodal_gen.runtime.distributed import ( + get_local_torch_device, + get_sp_world_size, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( + DenoisingContext, + DenoisingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.upsample import ( + apply_upsample, +) +from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +PROGRESSIVE_MODES = frozenset({"dct", "dct_rewind"}) + + +def is_progressive_resolution_mode(mode: str | None) -> bool: + return (mode or "fullres") in PROGRESSIVE_MODES + + +def unpack_2x2_latent(latent: torch.Tensor, h_lat: int, w_lat: int) -> torch.Tensor: + batch_size, _seq_len, packed_channels = latent.shape + spatial_channels = packed_channels // 4 + x = latent.view(batch_size, h_lat // 2, w_lat // 2, spatial_channels, 2, 2) + x = x.permute(0, 3, 1, 4, 2, 5) + return x.reshape(batch_size, spatial_channels, h_lat, w_lat) + + +def pack_2x2_latent(x: torch.Tensor, h_lat: int, w_lat: int) -> torch.Tensor: + batch_size, spatial_channels = x.shape[:2] + x = x.view(batch_size, spatial_channels, h_lat // 2, 2, w_lat // 2, 2) + x = x.permute(0, 2, 4, 1, 3, 5) + return x.reshape(batch_size, (h_lat // 2) * (w_lat // 2), spatial_channels * 4) + + +def _P_omega(w: float, A: float, beta: float) -> float: + return A * abs(w) ** (-beta) + + +def _activation_time(P: float, delta: float) -> float: + denom = P * (1.0 + P - delta) + if denom <= 0 or delta >= 1.0 + P: + raise ValueError( + f"delta={delta} >= 1+P={1+P:.4f}; criterion trivially satisfied." + ) + return 1.0 / (1.0 + math.sqrt(delta / denom)) + + +def compute_stage_transitions( + delta: float, + n_levels: int, + A: float, + beta: float, + H_lat: int, + W_lat: int, +) -> dict[int, float]: + stage_sigmas: dict[int, float] = {1: 1.0} + num_stages = n_levels + 1 + for stage in range(2, num_stages + 1): + H_prev = H_lat // (2 ** (num_stages - stage + 1)) + W_prev = W_lat // (2 ** (num_stages - stage + 1)) + w = min(H_prev, W_prev) // 2 + stage_sigmas[stage] = _activation_time(_P_omega(w, A, beta), delta) + return stage_sigmas + + +def find_transition_steps( + scheduler_sigmas: torch.Tensor, + stage_sigmas: dict[int, float], + n_steps: int, +) -> dict[int, int]: + transition_steps: dict[int, int] = {} + sigmas_list = scheduler_sigmas.cpu().tolist() + for stage, threshold in stage_sigmas.items(): + if stage == 1: + continue + found = n_steps + for step_index in range(n_steps): + if sigmas_list[step_index] <= threshold: + found = step_index + break + transition_steps[stage] = found + return transition_steps + + +def reset_scheduler_at_step(scheduler: object, step_index: int) -> None: + if hasattr(scheduler, "model_outputs"): + solver_order = getattr( + getattr(scheduler, "config", None), + "solver_order", + len(scheduler.model_outputs), + ) + scheduler.model_outputs = [None] * solver_order + if hasattr(scheduler, "lower_order_nums"): + scheduler.lower_order_nums = 0 + if hasattr(scheduler, "last_sample"): + scheduler.last_sample = None + if hasattr(scheduler, "this_order"): + scheduler.this_order = 0 + if hasattr(scheduler, "timestep_list"): + solver_order = getattr( + getattr(scheduler, "config", None), + "solver_order", + len(scheduler.timestep_list), + ) + scheduler.timestep_list = [None] * solver_order + scheduler._step_index = step_index + + +class ProgressiveDenoisingStageRouter(PipelineStage): + def __init__( + self, + standard_stage: DenoisingStage, + progressive_stage_factory: Callable[[], DenoisingStage], + ) -> None: + super().__init__() + self.standard_stage = standard_stage + self._progressive_stage_factory = progressive_stage_factory + self._progressive_stage: DenoisingStage | None = None + + def _get_progressive_stage(self) -> DenoisingStage: + if self._progressive_stage is None: + stage = self._progressive_stage_factory() + if self._component_residency_manager is not None: + stage.set_component_residency_manager(self._component_residency_manager) + if self._registered_stage_name is not None: + stage.set_registered_stage_name(self._registered_stage_name) + if self._profile_stage_name is not None: + stage.set_profile_stage_name(self._profile_stage_name) + self._progressive_stage = stage + return self._progressive_stage + + @property + def role_affinity(self): + return RoleType.DENOISER + + @property + def parallelism_type(self): + return self.standard_stage.parallelism_type + + def set_component_residency_manager(self, manager) -> None: + super().set_component_residency_manager(manager) + self.standard_stage.set_component_residency_manager(manager) + if self._progressive_stage is not None: + self._progressive_stage.set_component_residency_manager(manager) + + def set_registered_stage_name(self, stage_name: str) -> None: + super().set_registered_stage_name(stage_name) + self.standard_stage.set_registered_stage_name(stage_name) + if self._progressive_stage is not None: + self._progressive_stage.set_registered_stage_name(stage_name) + + def set_profile_stage_name(self, stage_name: str) -> None: + super().set_profile_stage_name(stage_name) + self.standard_stage.set_profile_stage_name(stage_name) + if self._progressive_stage is not None: + self._progressive_stage.set_profile_stage_name(stage_name) + + def _active_profile_stage_name(self) -> str: + # keep progressive requests under the canonical perf baseline stage name + return "DenoisingStage" + + def component_uses(self, server_args: ServerArgs, stage_name: str | None = None): + return self.standard_stage.component_uses(server_args, stage_name) + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + mode = getattr(batch, "progressive_mode", "fullres") or "fullres" + if is_progressive_resolution_mode(mode): + return self._get_progressive_stage().forward(batch, server_args) + if mode == "fullres": + return self.standard_stage.forward(batch, server_args) + raise ValueError(f"Unsupported progressive_mode: {mode!r}") + + +def _get_scm_preset() -> str | None: + preset = envs.SGLANG_CACHE_DIT_SCM_PRESET + return None if (preset is None or preset == "none") else preset + + +class ProgressiveDenoisingStage(DenoisingStage): + """DenoisingStage extended with progressive resolution growing. + + Subclass and override _unpack_latent / _repack_latent / _on_resolution_change + for model-specific latent packing and positional-embedding updates. + + spectrum_A and spectrum_beta are the fitted power-law coefficients for + P(ω) = A * |ω|^{-β} describing the latent frequency spectrum. + """ + + def __init__( + self, + transformer, + scheduler, + pipeline=None, + transformer_2=None, + vae=None, + spectrum_A: float = 1.0, + spectrum_beta: float = 2.0, + ) -> None: + super().__init__(transformer, scheduler, pipeline, transformer_2, vae) + self._spectrum_A = spectrum_A + self._spectrum_beta = spectrum_beta + + # ------------------------------------------------------------------ + # Extension hooks (override in model-specific subclasses) + # ------------------------------------------------------------------ + + def _latent_scale_factor(self, server_args: ServerArgs) -> int: + """Pixel-to-latent scale factor used for spatial latent dimensions. + + Defaults to vae_scale_factor. Models that apply an extra patchification + step (e.g. FLUX.2 uses vae_scale_factor * 2) should override this. + """ + return server_args.pipeline_config.vae_config.arch_config.vae_scale_factor + + def _unpack_latent( + self, latent: torch.Tensor, h_lat: int, w_lat: int + ) -> torch.Tensor: + """Convert model-native latent → spatial [B, C, H_lat, W_lat].""" + return latent + + def _repack_latent( + self, + x_spatial: torch.Tensor, + h_lat: int, + w_lat: int, + batch: Req, + server_args: ServerArgs, + ) -> torch.Tensor: + """Convert spatial [B, C, H_lat, W_lat] → model-native latent.""" + return x_spatial + + def _on_resolution_change( + self, + ctx: DenoisingContext, + batch: Req, + server_args: ServerArgs, + new_h_pixel: int, + new_w_pixel: int, + ) -> None: + """Called after each stage transition. Update resolution-dependent state.""" + pass + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _prepare_resolution_pos_cond_kwargs( + self, + ctx: DenoisingContext, + batch: Req, + server_args: ServerArgs, + ) -> dict[str, Any]: + rotary_emb = self._get_transformer_attr("rotary_emb") + return server_args.pipeline_config.prepare_pos_cond_kwargs( + batch, + self.device, + rotary_emb, + dtype=ctx.target_dtype, + ) + + @staticmethod + def _update_cfg_branch_kwargs( + ctx: DenoisingContext, + updates: dict[str, Any | None], + ) -> None: + assert ctx.cfg_policy is not None + for branch in ctx.cfg_policy.branches: + for name, value in updates.items(): + if value is not None and name in branch.kwargs: + branch.kwargs[name] = value + + for name, value in updates.items(): + if value is not None and name in ctx.pos_cond_kwargs: + ctx.pos_cond_kwargs[name] = value + + @staticmethod + def _get_seed(batch: Req) -> int: + seeds = getattr(batch, "seeds", None) + if seeds: + return int(seeds[0]) + sp = getattr(batch, "sampling_params", None) + seed = getattr(sp, "seed", None) if sp is not None else None + return int(seed) if seed is not None else 42 + + @staticmethod + def _initial_noise_batch_size(batch: Req) -> int: + try: + return int(batch.batch_size) + except AttributeError: + prompt_embeds = getattr(batch, "prompt_embeds", None) + if prompt_embeds: + return int(prompt_embeds[0].shape[0]) + latents = getattr(batch, "latents", None) + if latents is not None: + return int(latents.shape[0]) + return 1 + + def _get_seeds(self, batch: Req, seed: int | Sequence[int]) -> list[int]: + batch_size = self._initial_noise_batch_size(batch) + if isinstance(seed, Sequence) and not isinstance(seed, (str, bytes)): + seeds = [int(item) for item in seed] + else: + batch_seeds = getattr(batch, "seeds", None) + if batch_seeds: + seeds = [int(item) for item in batch_seeds] + else: + seeds = [int(seed) + i for i in range(batch_size)] + if len(seeds) != batch_size: + raise ValueError( + "progressive seeds length must match batch size: " + f"{len(seeds)} vs {batch_size}" + ) + return seeds + + def _get_initial_noise_generator( + self, batch: Req, seed: int | Sequence[int], device: torch.device | str + ): + seeds = self._get_seeds(batch, seed) + generators = [ + torch.Generator(device=device).manual_seed(seed) for seed in seeds + ] + if len(generators) == 1: + return generators[0] + return generators + + def _generate_initial_noise( + self, + batch: Req, + server_args: ServerArgs, + h_lat: int, + w_lat: int, + seed: int | Sequence[int], + ) -> torch.Tensor: + """Generate low-res initial noise and return in model-native format.""" + device = get_local_torch_device() + C = server_args.pipeline_config.dit_config.arch_config.in_channels // 4 + dtype = server_args.pipeline_config.get_latent_dtype( + batch.prompt_embeds[0].dtype if batch.prompt_embeds else torch.bfloat16 + ) + noise_spatial = randn_tensor( + (self._initial_noise_batch_size(batch), C, h_lat, w_lat), + generator=self._get_initial_noise_generator(batch, seed, device), + device=device, + dtype=dtype, + ) + return self._repack_latent(noise_spatial, h_lat, w_lat, batch, server_args) + + def _run_stage_steps( + self, + ctx: DenoisingContext, + batch: Req, + server_args: ServerArgs, + timesteps_cpu: torch.Tensor, + start_step: int, + end_step: int, + ) -> None: + """Run denoising steps [start_step, end_step) using the parent infrastructure.""" + for step_index in range(start_step, end_step): + t_host = timesteps_cpu[step_index] + step = self._prepare_step_state( + ctx, batch, server_args, step_index, t_host, timesteps_cpu + ) + self._run_denoising_step(ctx, step, batch, server_args) + + # ------------------------------------------------------------------ + # Progressive forward + # ------------------------------------------------------------------ + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + mode = getattr(batch, "progressive_mode", "fullres") or "fullres" + + if mode not in PROGRESSIVE_MODES: + raise ValueError( + "ProgressiveDenoisingStage requires progressive_mode to be " + "'dct' or 'dct_rewind'. Route fullres requests to DenoisingStage." + ) + + if get_sp_world_size() > 1: + raise RuntimeError( + "Progressive resolution growing is not compatible with sequence " + "parallelism. Disable --ulysses-degree / --ring-degree or set " + "progressive_mode='fullres'." + ) + + levels = int(getattr(batch, "progressive_levels", 1)) + delta = float(getattr(batch, "progressive_delta", 0.01)) + seed = self._get_seed(batch) + seeds = self._get_seeds(batch, seed) + + latent_scale = self._latent_scale_factor(server_args) + H_lat = batch.height // latent_scale + W_lat = batch.width // latent_scale + downsample = 2**levels + init_h_lat = H_lat // downsample + init_w_lat = W_lat // downsample + + # Compute stage transitions from the power-law spectrum + stage_sigmas = compute_stage_transitions( + delta, levels, self._spectrum_A, self._spectrum_beta, H_lat, W_lat + ) + num_stages = len(stage_sigmas) + + logger.info( + "Progressive denoising: mode=%s levels=%d delta=%.3f initial=%dx%d", + mode, + levels, + delta, + init_h_lat, + init_w_lat, + ) + + # ── Prepare initial state ────────────────────────────────────────────── + # Save the full-res dimensions that were set by LatentPreparationStage. + orig_h, orig_w = batch.height, batch.width + + # Override batch with low-res initial noise; _prepare_denoising_loop + # reads batch.latents and batch.height/width to build freqs_cis. + batch.height = init_h_lat * latent_scale + batch.width = init_w_lat * latent_scale + batch.latents = self._generate_initial_noise( + batch, server_args, init_h_lat, init_w_lat, seed + ) + batch.raw_latent_shape = batch.latents.shape + + ctx = self._prepare_denoising_loop(batch, server_args) + self._before_denoising_loop(ctx, batch, server_args) + + scheduler = ctx.scheduler + n_steps = int(batch.num_inference_steps) + timesteps_cpu = ctx.timesteps.cpu() + + transition_steps = find_transition_steps( + scheduler.sigmas, stage_sigmas, n_steps + ) + rewind = mode.endswith("_rewind") + + # For rewind mode we patch scheduler.sigmas/timesteps and ctx.timesteps + # in-place at transition points. The scheduler tensors may be inference + # tensors (created inside torch.inference_mode), so clone them once now + # to obtain normal mutable tensors. timesteps_cpu is already a fresh + # CPU tensor from .cpu(), so no clone is needed there. + if rewind: + scheduler.sigmas = scheduler.sigmas.clone() + scheduler.timesteps = scheduler.timesteps.clone() + ctx.timesteps = ctx.timesteps.clone() + + denoising_start = time.time() + stage_start = 0 + cur_h_lat = init_h_lat + cur_w_lat = init_w_lat + + # ── Stage loop ──────────────────────────────────────────────────────── + # DenoisingStage.forward() wraps its denoising loop in torch.autocast; + # we bypass that path, so we must apply the same context here. + with torch.autocast( + device_type=current_platform.device_type, + dtype=ctx.target_dtype, + enabled=ctx.autocast_enabled, + ): + for stage in range(1, num_stages + 1): + stage_end = transition_steps.get(stage + 1, n_steps) + + logger.info( + "Stage %d/%d: %dx%d latent, steps [%d, %d)", + stage, + num_stages, + cur_h_lat, + cur_w_lat, + stage_start, + stage_end, + ) + + self._run_stage_steps( + ctx, batch, server_args, timesteps_cpu, stage_start, stage_end + ) + + if stage == num_stages: + break + + # ── Resolution transition ────────────────────────────────────── + sigma_t = float(scheduler.sigmas[stage_end]) + upsample_seed = [item + stage * 10_000 for item in seeds] + + # Unpack → spatial, upsample, repack + x_spatial = self._unpack_latent(ctx.latents, cur_h_lat, cur_w_lat) + + result = apply_upsample(x_spatial, sigma_t, upsample_seed, mode) + + if rewind: + x_spatial_up, t_eff = result + # Patch scheduler sigma/timestep at transition point for rewind + scheduler.sigmas[stage_end] = t_eff + scheduler.timesteps[stage_end] = t_eff * 1000 + ctx.timesteps[stage_end] = t_eff * 1000 + timesteps_cpu[stage_end] = t_eff * 1000 + logger.info( + " rewind: sigma=%.4f → t_eff=%.4f at step %d", + sigma_t, + t_eff, + stage_end, + ) + else: + x_spatial_up = result + + new_h_lat = cur_h_lat * 2 + new_w_lat = cur_w_lat * 2 + ctx.latents = self._repack_latent( + x_spatial_up, new_h_lat, new_w_lat, batch, server_args + ) + + # Update batch dimensions and model-specific state + new_h_pixel = new_h_lat * latent_scale + new_w_pixel = new_w_lat * latent_scale + batch.height = new_h_pixel + batch.width = new_w_pixel + self._on_resolution_change( + ctx, batch, server_args, new_h_pixel, new_w_pixel + ) + + reset_scheduler_at_step(scheduler, stage_end) + + # Refresh cache-dit context so its step counter and cached + # activations start clean for the new resolution. The coarse- + # stage activations have the wrong shape and would corrupt the + # residual-diff decision for the first full-res steps. + if self._cache_dit_enabled: + n_remaining = n_steps - stage_end + scm_preset = _get_scm_preset() + if self.transformer_2 is not None: + n_high = n_remaining // 2 + n_low = n_remaining - n_high + refresh_context_on_dual_transformer( + self.transformer, + self.transformer_2, + n_high, + n_low, + scm_preset=scm_preset, + ) + else: + refresh_context_on_transformer( + self.transformer, + n_remaining, + scm_preset=scm_preset, + ) + logger.info( + "cache-dit context refreshed at stage transition " + "(step %d, %d steps remaining)", + stage_end, + n_remaining, + ) + + cur_h_lat = new_h_lat + cur_w_lat = new_w_lat + stage_start = stage_end + + denoising_end = time.time() + if not ctx.is_warmup: + logger.info( + "Progressive denoising done in %.2fs (avg %.4fs/step)", + denoising_end - denoising_start, + (denoising_end - denoising_start) / max(n_steps, 1), + ) + + # raw_latent_shape was set to the low-res initial noise shape when we + # replaced batch.latents. Update it to the final full-res latent so + # maybe_unpad_latents in post_denoising_loop does not truncate tokens. + batch.raw_latent_shape = ctx.latents.shape + + # Ensure batch resolution reflects the final full-res output + batch.height = orig_h + batch.width = orig_w + + self._finish_active_component_use() + self._finalize_denoising_loop(ctx, batch, server_args) + return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/flux.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/flux.py new file mode 100644 index 000000000..df4ac601c --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/flux.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +FLUX.1-specific progressive-resolution denoising stage. + +Provides pack/unpack for FLUX's patchify format and updates the RoPE +positional embeddings (freqs_cis) when the latent resolution changes +between progressive stages. +""" + +from __future__ import annotations + +import torch + +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( + DenoisingContext, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.denoising import ( + ProgressiveDenoisingStage, + pack_2x2_latent, + unpack_2x2_latent, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +# Power-law spectrum constants for FLUX.1-dev VAE +# Fitted on Aesthetics-Train-V2 (105k images) +FLUX_SPECTRUM_A: float = 203.615097 +FLUX_SPECTRUM_BETA: float = 1.915461 + + +class FluxProgressiveDenoisingStage(ProgressiveDenoisingStage): + """FLUX-specific progressive denoising stage. + + Handles: + - FLUX patchify pack/unpack + - freqs_cis (RoPE image position embeddings) update on resolution change + - img_ids cache keyed on (h_lat, w_lat) to avoid redundant computation + """ + + def __init__(self, transformer, scheduler, pipeline=None, vae=None) -> None: + super().__init__( + transformer, + scheduler, + pipeline=pipeline, + vae=vae, + spectrum_A=FLUX_SPECTRUM_A, + spectrum_beta=FLUX_SPECTRUM_BETA, + ) + # Cache freqs_cis per latent resolution (h_lat, w_lat) to avoid + # redundant rotary embedding recomputation between requests. + self._freqs_cis_cache: dict[ + tuple[int, int], tuple[torch.Tensor, torch.Tensor] + ] = {} + + # ------------------------------------------------------------------ + # Pack / Unpack overrides + # ------------------------------------------------------------------ + + def _unpack_latent( + self, latent: torch.Tensor, h_lat: int, w_lat: int + ) -> torch.Tensor: + return unpack_2x2_latent(latent, h_lat, w_lat) + + def _repack_latent( + self, + x_spatial: torch.Tensor, + h_lat: int, + w_lat: int, + batch: Req, + server_args: ServerArgs, + ) -> torch.Tensor: + return pack_2x2_latent(x_spatial, h_lat, w_lat) + + # ------------------------------------------------------------------ + # Resolution-change hook + # ------------------------------------------------------------------ + + def _on_resolution_change( + self, + ctx: DenoisingContext, + batch: Req, + server_args: ServerArgs, + new_h_pixel: int, + new_w_pixel: int, + ) -> None: + """Recompute freqs_cis for the new resolution and update all CFG branches. + + CFGBranch.kwargs is a shallow copy made at build() time; updating + ctx.pos_cond_kwargs alone does NOT reach the transformer. We must + update branch.kwargs["freqs_cis"] directly in every branch. + """ + if ctx.cfg_policy is None: + return + + vae_scale_factor = ( + server_args.pipeline_config.vae_config.arch_config.vae_scale_factor + ) + new_h_lat = new_h_pixel // vae_scale_factor + new_w_lat = new_w_pixel // vae_scale_factor + key = (new_h_lat, new_w_lat) + + if key not in self._freqs_cis_cache: + new_pos_kwargs = self._prepare_resolution_pos_cond_kwargs( + ctx, batch, server_args + ) + freqs_cis = new_pos_kwargs.get("freqs_cis") + if freqs_cis is not None: + self._freqs_cis_cache[key] = freqs_cis + + cached = self._freqs_cis_cache.get(key) + if cached is None: + logger.warning( + "freqs_cis not available for %dx%d latent; skipping update", + new_h_lat, + new_w_lat, + ) + return + + self._update_cfg_branch_kwargs(ctx, {"freqs_cis": cached}) + + logger.info( + "Updated freqs_cis for %dx%d latent (pixel %dx%d) across %d branch(es)", + new_h_lat, + new_w_lat, + new_h_pixel, + new_w_pixel, + len(ctx.cfg_policy.branches), + ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/flux_2.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/flux_2.py new file mode 100644 index 000000000..e6fed75d7 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/flux_2.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +FLUX.2-specific progressive-resolution denoising stage. + +Provides pack/unpack for FLUX.2's simple row-major token format and updates +both batch.latent_ids and freqs_cis when the latent resolution changes +between progressive stages. + +FLUX.2 latent layout (before packing): + spatial: [B, C, H_lat, W_lat] where H_lat = H_pixel // (vae_scale_factor * 2) + packed: [B, H_lat * W_lat, C] (row-major reshape) + +This differs from FLUX.1 which uses a 2×2 patchification to interleave spatial +blocks into packed tokens. +""" + +from __future__ import annotations + +import torch +from diffusers.utils.torch_utils import randn_tensor + +from sglang.multimodal_gen.configs.pipeline_configs.flux import _prepare_latent_ids +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( + DenoisingContext, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.denoising import ( + ProgressiveDenoisingStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +# Power-law spectrum constants — using FLUX.1-dev VAE values as a placeholder +# until FLUX.2-specific coefficients are fitted. +# Fitted on Aesthetics-Train-V2 (105k images) for the FLUX.1-dev VAE. +FLUX_SPECTRUM_A: float = 203.615097 +FLUX_SPECTRUM_BETA: float = 1.915461 + + +def _flux2_unpack(latent: torch.Tensor, h_lat: int, w_lat: int) -> torch.Tensor: + """Packed [B, H_lat*W_lat, C] → spatial [B, C, H_lat, W_lat] (row-major).""" + B, _S, C = latent.shape + return latent.permute(0, 2, 1).reshape(B, C, h_lat, w_lat) + + +def _flux2_pack(x: torch.Tensor) -> torch.Tensor: + """Spatial [B, C, H_lat, W_lat] → packed [B, H_lat*W_lat, C] (row-major).""" + B, C, H, W = x.shape + return x.reshape(B, C, H * W).permute(0, 2, 1) + + +class Flux2ProgressiveDenoisingStage(ProgressiveDenoisingStage): + """FLUX.2-specific progressive denoising stage. + + Handles: + - FLUX.2 row-major pack/unpack + - latent_ids update on resolution change (needed for 4-D RoPE in FLUX.2) + - freqs_cis cache and branch update on resolution change + """ + + def __init__(self, transformer, scheduler, pipeline=None, vae=None) -> None: + super().__init__( + transformer, + scheduler, + pipeline=pipeline, + vae=vae, + spectrum_A=FLUX_SPECTRUM_A, + spectrum_beta=FLUX_SPECTRUM_BETA, + ) + self._freqs_cis_cache: dict[ + tuple[int, int], tuple[torch.Tensor, torch.Tensor] + ] = {} + + # ------------------------------------------------------------------ + # Scale factor override + # ------------------------------------------------------------------ + + def _latent_scale_factor(self, server_args: ServerArgs) -> int: + # FLUX.2 latent spatial dimensions are at 1/(vae_scale_factor * 2) of + # pixel resolution due to the extra patchification step. + return server_args.pipeline_config.vae_config.arch_config.vae_scale_factor * 2 + + # ------------------------------------------------------------------ + # Pack / Unpack overrides + # ------------------------------------------------------------------ + + def _unpack_latent( + self, latent: torch.Tensor, h_lat: int, w_lat: int + ) -> torch.Tensor: + return _flux2_unpack(latent, h_lat, w_lat) + + def _repack_latent( + self, + x_spatial: torch.Tensor, + h_lat: int, + w_lat: int, + batch: Req, + server_args: ServerArgs, + ) -> torch.Tensor: + return _flux2_pack(x_spatial) + + # ------------------------------------------------------------------ + # Initial noise generation + # ------------------------------------------------------------------ + + def _generate_initial_noise( + self, + batch: Req, + server_args: ServerArgs, + h_lat: int, + w_lat: int, + seed, + ) -> torch.Tensor: + """Generate low-res noise, set batch.latent_ids, and return packed latent. + + FLUX.2 uses in_channels directly (no //4) because the spatial latent + already incorporates the patchification channel expansion. + """ + device = get_local_torch_device() + C = server_args.pipeline_config.dit_config.arch_config.in_channels + dtype = server_args.pipeline_config.get_latent_dtype( + batch.prompt_embeds[0].dtype if batch.prompt_embeds else torch.bfloat16 + ) + noise_spatial = randn_tensor( + (self._initial_noise_batch_size(batch), C, h_lat, w_lat), + generator=self._get_initial_noise_generator(batch, seed, device), + device=device, + dtype=dtype, + ) + + # latent_ids are derived from the spatial shape; _prepare_denoising_loop + # will read batch.latent_ids when building freqs_cis. + latent_ids = _prepare_latent_ids(noise_spatial) + batch.latent_ids = latent_ids.to(device) + + return _flux2_pack(noise_spatial) + + # ------------------------------------------------------------------ + # Resolution-change hook + # ------------------------------------------------------------------ + + def _on_resolution_change( + self, + ctx: DenoisingContext, + batch: Req, + server_args: ServerArgs, + new_h_pixel: int, + new_w_pixel: int, + ) -> None: + """Update batch.latent_ids and freqs_cis for the new latent resolution. + + Called after the upsampled latent is stored in ctx.latents and + batch.height/width are updated to new_h_pixel/new_w_pixel. + """ + if ctx.cfg_policy is None: + return + + latent_scale = self._latent_scale_factor(server_args) + new_h_lat = new_h_pixel // latent_scale + new_w_lat = new_w_pixel // latent_scale + key = (new_h_lat, new_w_lat) + + # Update batch.latent_ids so that prepare_pos_cond_kwargs sees the + # correct grid coordinates for the upsampled resolution. + C = server_args.pipeline_config.dit_config.arch_config.in_channels + dummy = ctx.latents.new_zeros(1, C, new_h_lat, new_w_lat) + latent_ids = _prepare_latent_ids(dummy) + batch.latent_ids = latent_ids.to(ctx.latents.device) + + if key not in self._freqs_cis_cache: + new_pos_kwargs = self._prepare_resolution_pos_cond_kwargs( + ctx, batch, server_args + ) + freqs_cis = new_pos_kwargs.get("freqs_cis") + if freqs_cis is not None: + self._freqs_cis_cache[key] = freqs_cis + + cached = self._freqs_cis_cache.get(key) + if cached is None: + logger.warning( + "freqs_cis not available for %dx%d latent; skipping update", + new_h_lat, + new_w_lat, + ) + return + + self._update_cfg_branch_kwargs(ctx, {"freqs_cis": cached}) + + logger.info( + "Updated latent_ids and freqs_cis for %dx%d latent (pixel %dx%d) " + "across %d branch(es)", + new_h_lat, + new_w_lat, + new_h_pixel, + new_w_pixel, + len(ctx.cfg_policy.branches), + ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/qwen_image.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/qwen_image.py new file mode 100644 index 000000000..5c6c2cb07 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/qwen_image.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Qwen-Image-specific progressive-resolution denoising stage. + +Provides pack/unpack for Qwen-Image's patchify format and updates the RoPE +positional embeddings (freqs_cis) and image shape metadata (img_shapes) when +the latent resolution changes between progressive stages. + +Qwen-Image uses the same patchify convention as FLUX.1-dev: + - in_channels = 64, spatial channels C = in_channels // 4 = 16 + - 2×2 patchification → packed [B, S, 64] where S = (H_lat/2) * (W_lat/2) + +The Qwen DiT forward() uses both ``freqs_cis`` (RoPE) and ``img_shapes`` +(for build_modulate_index), so _on_resolution_change updates both. + +Extension points (from ProgressiveDenoisingStage base class): + _unpack_latent : [B, S, 64] → [B, 16, H_lat, W_lat] + _repack_latent : [B, 16, H_lat, W_lat] → [B, S, 64] + _on_resolution_change : update freqs_cis + img_shapes in every CFG branch +""" + +from __future__ import annotations + +import torch + +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( + DenoisingContext, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.denoising import ( + ProgressiveDenoisingStage, + pack_2x2_latent, + unpack_2x2_latent, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +# Power-law spectrum constants P(ω) = A·|ω|^{-β} for Qwen-Image VAE latents. +# TODO: fit these from Qwen-Image VAE latent statistics on a representative +# dataset (e.g. Aesthetics-Train-V2). Using FLUX.1-dev fitted values as +# a reasonable starting point — both VAEs compress 2-D images into a +# 16-channel spatial latent with similar frequency roll-off. +QWEN_IMAGE_SPECTRUM_A: float = 203.615097 +QWEN_IMAGE_SPECTRUM_BETA: float = 1.915461 + + +class QwenImageProgressiveDenoisingStage(ProgressiveDenoisingStage): + """Qwen-Image progressive denoising stage. + + Inherits the full coarse-to-fine denoising loop from + ProgressiveDenoisingStage and overrides three model-specific hooks: + + * _unpack_latent / _repack_latent — Qwen's 2×2 patchify format + * _on_resolution_change — update freqs_cis AND img_shapes in + every CFG branch so the Qwen DiT's + build_modulate_index sees the right + spatial dimensions at each stage + + When progressive_mode == "fullres" (the default) the stage delegates + entirely to DenoisingStage.forward(), so existing non-progressive + requests are completely unaffected. + """ + + def __init__( + self, + transformer, + scheduler, + pipeline=None, + vae=None, + ) -> None: + super().__init__( + transformer, + scheduler, + pipeline=pipeline, + vae=vae, + spectrum_A=QWEN_IMAGE_SPECTRUM_A, + spectrum_beta=QWEN_IMAGE_SPECTRUM_BETA, + ) + + # ------------------------------------------------------------------ + # Pack / Unpack overrides + # ------------------------------------------------------------------ + + def _unpack_latent( + self, latent: torch.Tensor, h_lat: int, w_lat: int + ) -> torch.Tensor: + return unpack_2x2_latent(latent, h_lat, w_lat) + + def _repack_latent( + self, + x_spatial: torch.Tensor, + h_lat: int, + w_lat: int, + batch: Req, + server_args: ServerArgs, + ) -> torch.Tensor: + return pack_2x2_latent(x_spatial, h_lat, w_lat) + + # ------------------------------------------------------------------ + # Resolution-change hook + # ------------------------------------------------------------------ + + def _on_resolution_change( + self, + ctx: DenoisingContext, + batch: Req, + server_args: ServerArgs, + new_h_pixel: int, + new_w_pixel: int, + ) -> None: + """Update freqs_cis and img_shapes for the new latent resolution. + + batch.height / batch.width are already set to new_h_pixel / new_w_pixel + by the base class before this hook fires, so prepare_pos_cond_kwargs + computes the correct RoPE cache and img_shapes for the new resolution. + + Both freqs_cis (RoPE) and img_shapes (build_modulate_index) are updated + in every CFG branch because CFGBranch.kwargs is a shallow copy made at + build() time — updating ctx.pos_cond_kwargs alone does not reach the + transformer. + """ + if ctx.cfg_policy is None: + return + + new_pos_kwargs = self._prepare_resolution_pos_cond_kwargs( + ctx, batch, server_args + ) + freqs_cis = new_pos_kwargs.get("freqs_cis") + img_shapes = new_pos_kwargs.get("img_shapes") + + if freqs_cis is None: + logger.warning( + "freqs_cis not available for pixel %dx%d; skipping update", + new_h_pixel, + new_w_pixel, + ) + return + + self._update_cfg_branch_kwargs( + ctx, + { + "freqs_cis": freqs_cis, + "img_shapes": img_shapes, + }, + ) + + vae_scale_factor = ( + server_args.pipeline_config.vae_config.arch_config.vae_scale_factor + ) + new_h_lat = new_h_pixel // vae_scale_factor + new_w_lat = new_w_pixel // vae_scale_factor + + logger.info( + "Updated freqs_cis + img_shapes for %dx%d latent (pixel %dx%d)" + " across %d branch(es)", + new_h_lat, + new_w_lat, + new_h_pixel, + new_w_pixel, + len(ctx.cfg_policy.branches), + ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/spectral_ops.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/spectral_ops.py new file mode 100644 index 000000000..27d3ad55d --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/spectral_ops.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +GPU DCT-II / IDCT-II via torch.fft — no CPU↔GPU transfers. + +Algorithm: Makhoul (1980) "A fast cosine transform in one and two dimensions", +adapted for PyTorch. Operates on the last two spatial dims; input can be any +shape (..., H, W). +""" + +import math + +import torch + +# --------------------------------------------------------------------------- +# 1-D DCT-II / IDCT-II (ortho-normalized, operates on last dim) +# --------------------------------------------------------------------------- + + +def dct_1d(x: torch.Tensor, norm: str = "ortho") -> torch.Tensor: + """1-D DCT-II via torch.fft. Input: (..., N). Output: same shape.""" + shape = x.shape + N = shape[-1] + x = x.reshape(-1, N) + + # Reorder: [x0, x2, x4, ..., xN-1, ..., x3, x1] + v = torch.cat([x[:, ::2], x[:, 1::2].flip(dims=[1])], dim=1) + + Vc = torch.fft.fft(v, dim=1) + + k = torch.arange(N, dtype=x.dtype, device=x.device) * (-math.pi / (2 * N)) + W = torch.exp(torch.complex(torch.zeros_like(k), k)) # e^{-i*pi*k/(2N)} + V = (Vc * W).real + + if norm == "ortho": + V[:, 0] /= math.sqrt(N) * 2 + V[:, 1:] /= math.sqrt(N / 2) * 2 + + return (2 * V).reshape(shape) + + +def idct_1d(X: torch.Tensor, norm: str = "ortho") -> torch.Tensor: + """1-D IDCT-II (= scaled DCT-III) via torch.fft. Input: (..., N).""" + shape = X.shape + N = shape[-1] + X_v = X.reshape(-1, N) / 2 + + if norm == "ortho": + X_v = X_v.clone() + X_v[:, 0] *= math.sqrt(N) * 2 + X_v[:, 1:] *= math.sqrt(N / 2) * 2 + + k = torch.arange(N, dtype=X.dtype, device=X.device) * (math.pi / (2 * N)) + W = torch.exp(torch.complex(torch.zeros_like(k), k)) # e^{i*pi*k/(2N)} + + # Build complex input for IFFT + V_t_i = torch.cat([X_v[:, :1] * 0, -X_v.flip(dims=[1])[:, :-1]], dim=1) + Vc = torch.complex(X_v, V_t_i) * W + + v = torch.fft.ifft(Vc, dim=1).real + x = torch.zeros_like(v) + x[:, ::2] = v[:, : N - (N // 2)] + x[:, 1::2] = v.flip(dims=[1])[:, : N // 2] + return x.reshape(shape) + + +# --------------------------------------------------------------------------- +# 2-D DCT-II / IDCT-II (separable: apply 1-D along H then W) +# --------------------------------------------------------------------------- + + +def dct_2d(x: torch.Tensor, norm: str = "ortho") -> torch.Tensor: + """2-D DCT-II on the last two dims of x (..., H, W).""" + return dct_1d(dct_1d(x, norm).transpose(-1, -2), norm).transpose(-1, -2) + + +def idct_2d(X: torch.Tensor, norm: str = "ortho") -> torch.Tensor: + """2-D IDCT-II on the last two dims of X (..., H, W).""" + return idct_1d(idct_1d(X, norm).transpose(-1, -2), norm).transpose(-1, -2) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/upsample.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/upsample.py new file mode 100644 index 000000000..6270c761f --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/upsample.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +GPU-native latent upsample operations for progressive resolution growing. + +All ops run entirely on GPU via torch.fft — no CPU↔GPU data movement. +Supported modes: "dct", "dct_rewind". + +Each function takes a spatial latent tensor (..., H, W) and returns a 2× larger +tensor (..., 2H, 2W). The rewind variant also returns t_eff. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch + +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.spectral_ops import ( + dct_2d, + idct_2d, +) + + +def dct_upsample_2d( + x: torch.Tensor, + sigma_t: float, + seed: int | Sequence[int], + rewind: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, float]: + """DCT-II 2× upsample: embed low-res coefficients top-left, noise-pad, IDCT. + + x: (..., H, W) spatial latent tensor. + sigma_t: current noise level (used to scale the high-freq padding noise). + seed: deterministic RNG seed for the noise padding. + rewind: if True, multiply by 2/(1+sigma_t) and return (result, t_eff). + + Matches the CPU reference in inference_progressive.py but runs fully on GPU. + """ + *leading, H, W = x.shape + H2, W2 = H * 2, W * 2 + + # 2-D DCT-II of the source (ortho-normalized, Parseval identity preserved). + # All intermediate computation stays in float32 to match the reference + # (inference_progressive.py uses scipy float32 throughout). bfloat16 has + # only 7 mantissa bits; quantising the DCT coefficients before IDCT would + # introduce mean absolute error ~0.8 against an output range of ±4. + X_low = dct_2d(x.float(), norm="ortho") # (..., H, W) float32 + + # Fill 2N×2N grid with float32 white Gaussian noise of variance sigma_t² + # per DCT bin, matching the reference's float32 noise path. + if isinstance(seed, Sequence) and not isinstance(seed, (str, bytes)): + if not leading or len(seed) != leading[0]: + batch_dim = leading[0] if leading else 0 + raise ValueError( + "seed list length must match leading batch dimension: " + f"{len(seed)} vs {batch_dim}" + ) + big = torch.cat( + [ + torch.randn( + 1, + *leading[1:], + H2, + W2, + generator=torch.Generator(device=x.device).manual_seed(int(item)), + dtype=torch.float32, + device=x.device, + ) + for item in seed + ], + dim=0, + ) + else: + generator = torch.Generator(device=x.device).manual_seed(int(seed)) + big = torch.randn( + *leading, H2, W2, generator=generator, dtype=torch.float32, device=x.device + ) + big = big * sigma_t + + # Embed low-res DCT coefficients in the top-left corner (no precision loss). + big[..., :H, :W] = X_low + + # 2-D IDCT-II → spatial domain, then cast back to original dtype. + result = idct_2d(big, norm="ortho").to(x.dtype) + + if rewind: + gamma = 1.0 + sigma_t + result = result * (2.0 / gamma) + t_eff = 2.0 * sigma_t / gamma + return result, t_eff + return result + + +def apply_upsample( + x: torch.Tensor, + sigma_t: float, + seed: int | Sequence[int], + mode: str, +) -> torch.Tensor | tuple[torch.Tensor, float]: + """Dispatch to the requested upsample function. + + Returns tensor for plain modes, (tensor, t_eff) for rewind modes. + """ + if mode == "dct": + return dct_upsample_2d(x, sigma_t, seed, rewind=False) + if mode == "dct_rewind": + return dct_upsample_2d(x, sigma_t, seed, rewind=True) + raise ValueError(f"Unsupported progressive upsample mode: {mode!r}") diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/wan.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/wan.py new file mode 100644 index 000000000..720b95f85 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/wan.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Wan video progressive-resolution denoising stage. + +Extends ProgressiveDenoisingStage for the Wan T2V video model: + - Latent format: [B, C, T, H, W] (already spatial — no pack/unpack required) + - Upsample: spatial H×W dims only; T (temporal frames) is fixed across all stages + - No RoPE / freqs_cis update needed (Wan T2V uses no spatial positional embeddings + that depend on H or W in the context) + +Power-law spectrum constants fitted on VChitect dataset (9050 videos), spatial +spectrum P(ω) = A * |ω|^(-β): + A = 219.484718 + β = 2.422687 +""" + +from __future__ import annotations + +import torch +from diffusers.utils.torch_utils import randn_tensor + +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.denoising import ( + ProgressiveDenoisingStage, + is_progressive_resolution_mode, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +# Power-law spectrum constants for WAN 2.1 VAE +# Fitted on VChitect (9050 videos): P(ω) = A * |ω|^(-β) +WAN_SPECTRUM_A: float = 219.484718 +WAN_SPECTRUM_BETA: float = 2.422687 + + +class WanProgressiveDenoisingStage(ProgressiveDenoisingStage): + """Wan T2V–specific progressive denoising stage. + + Differences from the FLUX progressive stage: + - Wan latent is [B, C, T, H, W] — no patchify pack/unpack needed. + - Progressive upsample grows only the spatial H×W plane; T stays fixed. + - Wan T2V has no spatial RoPE freqs_cis that depends on H/W, so + _on_resolution_change is a no-op. + - Initial noise must carry the temporal dimension T_lat. + """ + + def __init__(self, transformer, scheduler, pipeline=None, vae=None) -> None: + super().__init__( + transformer, + scheduler, + pipeline=pipeline, + vae=vae, + spectrum_A=WAN_SPECTRUM_A, + spectrum_beta=WAN_SPECTRUM_BETA, + ) + + # ------------------------------------------------------------------ + # Latent scale factor (WanVAEArchConfig uses spatial_compression_ratio) + # ------------------------------------------------------------------ + + def _latent_scale_factor(self, server_args: ServerArgs) -> int: + arch = server_args.pipeline_config.vae_config.arch_config + return getattr(arch, "vae_scale_factor", None) or getattr( + arch, "spatial_compression_ratio", 8 + ) + + # ------------------------------------------------------------------ + # Pack / Unpack overrides (Wan latent is already [B, C, T, H, W]) + # ------------------------------------------------------------------ + + def _unpack_latent( + self, latent: torch.Tensor, h_lat: int, w_lat: int + ) -> torch.Tensor: + return latent + + def _repack_latent( + self, + x_spatial: torch.Tensor, + h_lat: int, + w_lat: int, + batch: Req, + server_args: ServerArgs, + ) -> torch.Tensor: + return x_spatial + + # ------------------------------------------------------------------ + # Resolution-change hook (no-op for Wan T2V) + # ------------------------------------------------------------------ + + def _on_resolution_change( + self, + ctx, + batch: Req, + server_args: ServerArgs, + new_h_pixel: int, + new_w_pixel: int, + ) -> None: + """Wan T2V has no spatial positional embeddings that require updating.""" + pass + + # ------------------------------------------------------------------ + # Resolution alignment (Wan patch embedding requires even spatial dims) + # ------------------------------------------------------------------ + + @torch.no_grad() + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + """Snap batch.height / batch.width to even multiples before progressive loop. + + Wan's patch embedding is Conv3d(stride=(1,2,2)), so each progressive + stage latent must have even H and W. At L levels of downsampling the + initial latent is H_lat//(2^L) × W_lat//(2^L); if either is odd the + patchification fails with a tensor size mismatch. + + Fix: align batch.height/width down to the nearest multiple of + vae_scale_factor * 2^L * 2 (= vae_scale * align_unit) so that every + stage latent is guaranteed even. For 480p L=1 this is a no-op (60 is + already divisible by 4). For 720p L=1: 90→88 latent rows (704 px). + """ + mode = getattr(batch, "progressive_mode", "fullres") or "fullres" + if not is_progressive_resolution_mode(mode): + return super().forward(batch, server_args) + + levels = int(getattr(batch, "progressive_levels", 1)) + arch = server_args.pipeline_config.vae_config.arch_config + vae_scale = getattr(arch, "vae_scale_factor", None) or getattr( + arch, "spatial_compression_ratio", 8 + ) + # Each stage halves the spatial dims; Wan needs even dims at every stage. + # Required: H_lat divisible by 2^L * patch_spatial (= 2^L * 2). + align_pixels = vae_scale * (2**levels) * 2 + h_aligned = max((batch.height // align_pixels) * align_pixels, align_pixels) + w_aligned = max((batch.width // align_pixels) * align_pixels, align_pixels) + + if h_aligned != batch.height or w_aligned != batch.width: + logger.info( + "WanProgressiveDenoisingStage: aligning resolution %dx%d → %dx%d " + "so all progressive stage latents have even spatial dims (patch=2, L=%d)", + batch.height, + batch.width, + h_aligned, + w_aligned, + levels, + ) + batch.height = h_aligned + batch.width = w_aligned + + return super().forward(batch, server_args) + + # ------------------------------------------------------------------ + # Initial noise (must include the temporal dim T_lat) + # ------------------------------------------------------------------ + + def _generate_initial_noise( + self, + batch: Req, + server_args: ServerArgs, + h_lat: int, + w_lat: int, + seed, + ) -> torch.Tensor: + """Generate low-res initial noise [1, C, T_lat, h_lat, w_lat]. + + The base-class version generates 4-D [1, C, H, W] noise and uses + in_channels // 4 for the channel count (FLUX patchify convention). + Wan operates directly on the 5-D latent, so we override to: + - Use z_dim (= 16) as the correct latent channel count. + - Preserve T_lat from the original full-res latent in batch.latents, + since progressive upsample only grows spatial H×W. + """ + device = get_local_torch_device() + C = server_args.pipeline_config.vae_config.arch_config.z_dim + # batch.latents still holds the full-res latent from LatentPreparationStage + # at this call site, so shape[2] gives the fixed T_lat. + T_lat = batch.latents.shape[2] + dtype = server_args.pipeline_config.get_latent_dtype( + batch.prompt_embeds[0].dtype if batch.prompt_embeds else torch.bfloat16 + ) + noise = randn_tensor( + (self._initial_noise_batch_size(batch), C, T_lat, h_lat, w_lat), + generator=self._get_initial_noise_generator(batch, seed, device), + device=device, + dtype=dtype, + ) + return noise diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/zimage.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/zimage.py new file mode 100644 index 000000000..91535878f --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/zimage.py @@ -0,0 +1,169 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Z-Image-specific progressive-resolution denoising stage. + +Provides pack/unpack for Z-Image's 5-D latent format [B, C, F, H, W] and updates +the RoPE positional embeddings (freqs_cis) when the latent resolution changes +between progressive stages. +""" + +from __future__ import annotations + +import torch +from diffusers.utils.torch_utils import randn_tensor + +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import ( + DenoisingContext, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.denoising import ( + ProgressiveDenoisingStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +# Power-law spectrum constants for Z-Image. +# Z-Image uses the same VAE as FLUX.1-dev (FluxVAEConfig), so the spectrum +# constants fitted on Aesthetics-Train-V2 (105k images) apply directly. +ZIMAGE_SPECTRUM_A: float = 203.615097 +ZIMAGE_SPECTRUM_BETA: float = 1.915461 + + +def _zimage_unpack(latent: torch.Tensor) -> torch.Tensor: + """5-D latent [B, C, 1, H_lat, W_lat] → spatial [B, C, H_lat, W_lat].""" + return latent.squeeze(2) + + +def _zimage_repack(x: torch.Tensor) -> torch.Tensor: + """Spatial [B, C, H_lat, W_lat] → 5-D latent [B, C, 1, H_lat, W_lat].""" + return x.unsqueeze(2) + + +class ZImageProgressiveDenoisingStage(ProgressiveDenoisingStage): + """Z-Image-specific progressive denoising stage. + + Handles: + - Z-Image 5-D latent pack/unpack [B, C, 1, H, W] ↔ [B, C, H, W] + - freqs_cis (RoPE caption + image position embeddings) update on resolution change + """ + + def __init__(self, transformer, scheduler, pipeline=None, vae=None) -> None: + super().__init__( + transformer, + scheduler, + pipeline=pipeline, + vae=vae, + spectrum_A=ZIMAGE_SPECTRUM_A, + spectrum_beta=ZIMAGE_SPECTRUM_BETA, + ) + + # ------------------------------------------------------------------ + # Initial noise + # ------------------------------------------------------------------ + + def _generate_initial_noise( + self, + batch: Req, + server_args: ServerArgs, + h_lat: int, + w_lat: int, + seed, + ) -> torch.Tensor: + """Generate low-res initial noise in Z-Image's native 5-D format [B, C, 1, H, W]. + + The base class uses in_channels // 4 which is correct for FLUX (64 // 4 = 16), + but Z-Image's in_channels = 16 already refers to the spatial channel count. + We use it directly and return 5-D via _repack_latent. + """ + device = get_local_torch_device() + C = server_args.pipeline_config.dit_config.arch_config.in_channels + dtype = server_args.pipeline_config.get_latent_dtype( + batch.prompt_embeds[0].dtype if batch.prompt_embeds else torch.bfloat16 + ) + noise_spatial = randn_tensor( + (self._initial_noise_batch_size(batch), C, h_lat, w_lat), + generator=self._get_initial_noise_generator(batch, seed, device), + device=device, + dtype=dtype, + ) + return self._repack_latent(noise_spatial, h_lat, w_lat, batch, server_args) + + # ------------------------------------------------------------------ + # Pack / Unpack overrides + # ------------------------------------------------------------------ + + def _unpack_latent( + self, latent: torch.Tensor, h_lat: int, w_lat: int + ) -> torch.Tensor: + return _zimage_unpack(latent) + + def _repack_latent( + self, + x_spatial: torch.Tensor, + h_lat: int, + w_lat: int, + batch: Req, + server_args: ServerArgs, + ) -> torch.Tensor: + return _zimage_repack(x_spatial) + + # ------------------------------------------------------------------ + # Resolution-change hook + # ------------------------------------------------------------------ + + def _on_resolution_change( + self, + ctx: DenoisingContext, + batch: Req, + server_args: ServerArgs, + new_h_pixel: int, + new_w_pixel: int, + ) -> None: + """Recompute freqs_cis for the new resolution and update all CFG branches. + + Z-Image freqs_cis is a (cap_freqs_cis, x_freqs_cis) tuple. The image + position offsets depend on caption length, so the full tuple must be + recomputed rather than only the image portion. + + batch.height / batch.width are already updated to new_h_pixel / new_w_pixel + by the time this hook is called, so prepare_pos_cond_kwargs uses the + correct new resolution automatically. + + CFGBranch.kwargs is a shallow copy made at build() time; updating + ctx.pos_cond_kwargs alone does NOT reach the transformer. We must + update branch.kwargs["freqs_cis"] directly in every branch. + """ + if ctx.cfg_policy is None: + return + + new_pos_kwargs = self._prepare_resolution_pos_cond_kwargs( + ctx, batch, server_args + ) + freqs_cis = new_pos_kwargs.get("freqs_cis") + if freqs_cis is None: + logger.warning( + "freqs_cis not available for pixel %dx%d; skipping update", + new_h_pixel, + new_w_pixel, + ) + return + + vae_scale_factor = ( + server_args.pipeline_config.vae_config.arch_config.vae_scale_factor + ) + new_h_lat = new_h_pixel // vae_scale_factor + new_w_lat = new_w_pixel // vae_scale_factor + + self._update_cfg_branch_kwargs(ctx, {"freqs_cis": freqs_cis}) + + logger.info( + "Updated freqs_cis for %dx%d latent (pixel %dx%d) across %d branch(es)", + new_h_lat, + new_w_lat, + new_h_pixel, + new_w_pixel, + len(ctx.cfg_policy.branches), + ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py index e98b1ce93..be9948ecd 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py @@ -249,12 +249,21 @@ class TextEncodingStage(ConditionEncodingStage): ) -> None: assert batch.negative_prompt_embeds is not None - # a single negative prompt can be shared across positive prompts - target_batch_sizes = [pe.shape[0] for pe in prompt_embeds_list] + # a single negative prompt can be shared across positive prompts. + # 2-D embeddings (seq × dim, e.g. Z-Image single-prompt) carry no explicit + # batch dimension; treat them as batch=1. + target_batch_sizes = [ + 1 if pe.ndim == 2 else pe.shape[0] for pe in prompt_embeds_list + ] def align_negative_batch_dim( tensor: torch.Tensor, target_batch: int, name: str ) -> torch.Tensor: + # 2-D: seq × dim with no batch dim — implicitly batch=1. + if tensor.ndim == 2: + if target_batch > 1: + return tensor.unsqueeze(0).repeat(target_batch, 1, 1) + return tensor if tensor.shape[0] == target_batch: return tensor if tensor.shape[0] == 1 and target_batch > 1: diff --git a/python/sglang/multimodal_gen/test/unit/conftest.py b/python/sglang/multimodal_gen/test/unit/conftest.py new file mode 100644 index 000000000..d46be29db --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/conftest.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest + +from sglang.multimodal_gen.runtime import server_args as server_args_module +from sglang.multimodal_gen.runtime.server_args import set_global_server_args + + +def _make_unit_server_args(): + dit_config = SimpleNamespace( + hidden_size=64, + num_attention_heads=4, + boundary_ratio=None, + arch_config=SimpleNamespace(in_channels=16, patch_size=2), + ) + vae_config = SimpleNamespace( + vae_tiling=False, + arch_config=SimpleNamespace( + vae_scale_factor=8, + spatial_compression_ratio=8, + z_dim=16, + scale_factor_spatial=8, + ), + get_vae_scale_factor=lambda: 8, + ) + pipeline_config = SimpleNamespace( + dit_config=dit_config, + vae_config=vae_config, + dit_precision="bfloat16", + vae_precision="bfloat16", + get_latent_dtype=lambda dtype: dtype, + ) + return SimpleNamespace( + attention_backend=None, + attention_backend_config=None, + comfyui_mode=False, + disable_autocast=False, + enable_cfg_parallel=False, + enable_layerwise_nvtx_marker=False, + enable_torch_compile=False, + model_loaded={}, + model_paths={}, + pipeline_config=pipeline_config, + ) + + +@pytest.fixture(autouse=True) +def default_global_server_args(): + previous = server_args_module._global_server_args + set_global_server_args(_make_unit_server_args()) + try: + yield + finally: + set_global_server_args(previous) diff --git a/python/sglang/multimodal_gen/test/unit/progressive_resolution/test_progressive.py b/python/sglang/multimodal_gen/test/unit/progressive_resolution/test_progressive.py new file mode 100644 index 000000000..d0aac40c8 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/progressive_resolution/test_progressive.py @@ -0,0 +1,358 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Focused unit tests for the experimental progressive-resolution path.""" + +import argparse +import dataclasses +import unittest +from types import SimpleNamespace + +import torch + +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.denoising import ( + ProgressiveDenoisingStage, + ProgressiveDenoisingStageRouter, + compute_stage_transitions, + find_transition_steps, + is_progressive_resolution_mode, + reset_scheduler_at_step, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.flux import ( + _flux_pack, + _flux_unpack, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.flux_2 import ( + Flux2ProgressiveDenoisingStage, + _flux2_pack, + _flux2_unpack, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.qwen_image import ( + _qwen_image_pack, + _qwen_image_unpack, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.spectral_ops import ( + dct_2d, + idct_2d, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.upsample import ( + apply_upsample, + dct_upsample_2d, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.wan import ( + WanProgressiveDenoisingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.zimage import ( + _zimage_repack, + _zimage_unpack, +) + + +class _DummyDenoisingStage: + parallelism_type = None + + def __init__(self, route_name: str): + self.route_name = route_name + self.component_manager = None + self.registered_stage_name = None + self.profile_stage_name = None + + def set_component_residency_manager(self, manager): + self.component_manager = manager + + def set_registered_stage_name(self, stage_name: str): + self.registered_stage_name = stage_name + + def set_profile_stage_name(self, stage_name: str): + self.profile_stage_name = stage_name + + def component_uses(self, server_args, stage_name=None): + return [] + + def forward(self, batch, server_args): + batch.route_name = self.route_name + return batch + + +class TestProgressiveSamplingParams(unittest.TestCase): + def _parse_cli_kwargs(self, argv: list[str]) -> dict: + parser = argparse.ArgumentParser() + SamplingParams.add_cli_args(parser) + return SamplingParams.get_cli_args(parser.parse_args(argv)) + + def test_defaults_and_valid_modes(self): + params = SamplingParams() + self.assertEqual(params.progressive_mode, "fullres") + self.assertEqual(params.progressive_levels, 1) + self.assertAlmostEqual(params.progressive_delta, 0.01) + + for mode in ("fullres", "dct", "dct_rewind"): + with self.subTest(mode=mode): + self.assertEqual( + SamplingParams(progressive_mode=mode).progressive_mode, mode + ) + + def test_validation_rejects_invalid_values(self): + invalid_cases = [ + {"progressive_mode": "wavelet"}, + {"progressive_levels": 0}, + {"progressive_levels": True}, + {"progressive_delta": 0}, + {"progressive_delta": 1}, + ] + for kwargs in invalid_cases: + with self.subTest(kwargs=kwargs): + with self.assertRaises(ValueError): + SamplingParams(**kwargs) + + def test_fields_stay_in_batch_signature(self): + fields = {field.name: field for field in dataclasses.fields(SamplingParams)} + for name in ("progressive_mode", "progressive_levels", "progressive_delta"): + with self.subTest(field=name): + self.assertFalse(fields[name].metadata.get("batch_sig_exclude")) + + def test_cli_only_emits_explicit_progressive_args(self): + self.assertEqual(self._parse_cli_kwargs([]), {}) + + kwargs = self._parse_cli_kwargs( + [ + "--progressive-mode", + "dct_rewind", + "--progressive-levels", + "2", + "--progressive-delta", + "0.05", + ] + ) + + self.assertEqual(kwargs["progressive_mode"], "dct_rewind") + self.assertEqual(kwargs["progressive_levels"], 2) + self.assertAlmostEqual(kwargs["progressive_delta"], 0.05) + + +class TestProgressiveRouter(unittest.TestCase): + def test_fullres_uses_standard_stage_without_constructing_progressive_stage(self): + calls = [] + + def create_progressive_stage(): + calls.append(1) + return _DummyDenoisingStage("progressive") + + router = ProgressiveDenoisingStageRouter( + standard_stage=_DummyDenoisingStage("standard"), + progressive_stage_factory=create_progressive_stage, + ) + batch = SimpleNamespace(progressive_mode="fullres") + + out = router.forward(batch, SimpleNamespace()) + + self.assertEqual(out.route_name, "standard") + self.assertEqual(calls, []) + + def test_progressive_stage_is_lazy_and_reused(self): + calls = [] + + def create_progressive_stage(): + calls.append(1) + return _DummyDenoisingStage("progressive") + + router = ProgressiveDenoisingStageRouter( + standard_stage=_DummyDenoisingStage("standard"), + progressive_stage_factory=create_progressive_stage, + ) + batch = SimpleNamespace(progressive_mode="dct_rewind") + + router.forward(batch, SimpleNamespace()) + router.forward(batch, SimpleNamespace()) + + self.assertEqual(batch.route_name, "progressive") + self.assertEqual(len(calls), 1) + + def test_invalid_mode_raises(self): + router = ProgressiveDenoisingStageRouter( + standard_stage=_DummyDenoisingStage("standard"), + progressive_stage_factory=lambda: _DummyDenoisingStage("progressive"), + ) + + with self.assertRaises(ValueError): + router.forward( + SimpleNamespace(progressive_mode="wavelet"), SimpleNamespace() + ) + + def test_mode_predicate(self): + self.assertTrue(is_progressive_resolution_mode("dct")) + self.assertTrue(is_progressive_resolution_mode("dct_rewind")) + self.assertFalse(is_progressive_resolution_mode("fullres")) + self.assertFalse(is_progressive_resolution_mode(None)) + + +class TestStageTransitionHelpers(unittest.TestCase): + def test_compute_stage_transitions_returns_one_threshold_per_stage(self): + transitions = compute_stage_transitions( + delta=0.01, + n_levels=2, + A=203.615097, + beta=1.915461, + H_lat=128, + W_lat=128, + ) + + self.assertEqual(set(transitions), {1, 2, 3}) + self.assertEqual(transitions[1], 1.0) + self.assertTrue(0 < transitions[2] < 1) + self.assertTrue(0 < transitions[3] < 1) + + def test_find_transition_steps_maps_thresholds_to_scheduler_indices(self): + scheduler_sigmas = torch.tensor([1.0, 0.8, 0.5, 0.25, 0.1]) + transitions = find_transition_steps( + scheduler_sigmas, + {1: 1.0, 2: 0.5, 3: 0.2}, + n_steps=5, + ) + + self.assertEqual(transitions, {2: 2, 3: 4}) + + def test_reset_scheduler_clears_solver_state(self): + scheduler = SimpleNamespace( + config=SimpleNamespace(solver_order=2), + model_outputs=[torch.ones(1), torch.ones(1)], + lower_order_nums=1, + last_sample=torch.ones(1), + this_order=1, + timestep_list=[1, 2], + _step_index=0, + ) + + reset_scheduler_at_step(scheduler, 3) + + self.assertEqual(scheduler.model_outputs, [None, None]) + self.assertEqual(scheduler.lower_order_nums, 0) + self.assertIsNone(scheduler.last_sample) + self.assertEqual(scheduler.this_order, 0) + self.assertEqual(scheduler.timestep_list, [None, None]) + self.assertEqual(scheduler._step_index, 3) + + +class TestSpectralUpsample(unittest.TestCase): + def test_dct_roundtrip(self): + x = torch.randn(2, 3, 8, 10) + + reconstructed = idct_2d(dct_2d(x)) + + torch.testing.assert_close(reconstructed, x, rtol=1e-5, atol=1e-5) + + def test_apply_upsample_shapes_and_rewind_return(self): + x = torch.randn(2, 3, 4, 5) + + out = apply_upsample(x, sigma_t=0.25, seed=[1, 2], mode="dct") + rewind_out, t_eff = apply_upsample( + x, sigma_t=0.25, seed=[1, 2], mode="dct_rewind" + ) + + self.assertEqual(out.shape, (2, 3, 8, 10)) + self.assertEqual(rewind_out.shape, (2, 3, 8, 10)) + self.assertGreater(t_eff, 0.25) + self.assertEqual(out.dtype, x.dtype) + + def test_seed_list_is_batch_checked_and_deterministic(self): + x = torch.randn(2, 3, 4, 4) + + out1 = dct_upsample_2d(x, sigma_t=0.1, seed=[7, 8]) + out2 = dct_upsample_2d(x, sigma_t=0.1, seed=[7, 8]) + + torch.testing.assert_close(out1, out2) + with self.assertRaises(ValueError): + dct_upsample_2d(x, sigma_t=0.1, seed=[7]) + + def test_invalid_upsample_mode_raises(self): + with self.assertRaises(ValueError): + apply_upsample(torch.zeros(1, 1, 2, 2), 0.1, 0, "wavelet") + + +class TestProgressiveStageHelpers(unittest.TestCase): + def test_seed_helpers_support_batch_seed_lists(self): + stage = object.__new__(ProgressiveDenoisingStage) + + batch = SimpleNamespace(batch_size=2, seeds=[11, 12], sampling_params=None) + self.assertEqual(stage._get_seed(batch), 11) + self.assertEqual(stage._get_seeds(batch, seed=0), [11, 12]) + + batch = SimpleNamespace( + prompt_embeds=[torch.zeros(3, 4, 5)], + seeds=None, + sampling_params=SimpleNamespace(seed=20), + ) + self.assertEqual(stage._get_seeds(batch, seed=20), [20, 21, 22]) + + def test_seed_helper_rejects_wrong_seed_count(self): + stage = object.__new__(ProgressiveDenoisingStage) + batch = SimpleNamespace(batch_size=2, seeds=[1], sampling_params=None) + + with self.assertRaises(ValueError): + stage._get_seeds(batch, seed=0) + + def test_model_specific_latent_scale_factors(self): + flux2_stage = object.__new__(Flux2ProgressiveDenoisingStage) + wan_stage = object.__new__(WanProgressiveDenoisingStage) + + flux2_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + vae_config=SimpleNamespace( + arch_config=SimpleNamespace(vae_scale_factor=8) + ) + ) + ) + wan_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + vae_config=SimpleNamespace( + arch_config=SimpleNamespace(spatial_compression_ratio=8) + ) + ) + ) + + self.assertEqual(flux2_stage._latent_scale_factor(flux2_args), 16) + self.assertEqual(wan_stage._latent_scale_factor(wan_args), 8) + + +class TestLatentAdapters(unittest.TestCase): + def test_flux_and_qwen_patchify_roundtrip(self): + x = torch.arange(1 * 16 * 8 * 12, dtype=torch.float32).reshape(1, 16, 8, 12) + + for name, pack, unpack in ( + ("flux", _flux_pack, _flux_unpack), + ("qwen_image", _qwen_image_pack, _qwen_image_unpack), + ): + with self.subTest(adapter=name): + packed = pack(x, 8, 12) + self.assertEqual(packed.shape, (1, (8 // 2) * (12 // 2), 64)) + torch.testing.assert_close(unpack(packed, 8, 12), x) + + def test_flux2_row_major_roundtrip(self): + x = torch.arange(2 * 4 * 3 * 5, dtype=torch.float32).reshape(2, 4, 3, 5) + + packed = _flux2_pack(x) + + self.assertEqual(packed.shape, (2, 3 * 5, 4)) + torch.testing.assert_close(packed[0, 7], x[0, :, 1, 2]) + torch.testing.assert_close(_flux2_unpack(packed, 3, 5), x) + + def test_zimage_adds_and_removes_frame_dim(self): + latent = torch.randn(2, 16, 1, 8, 8) + + spatial = _zimage_unpack(latent) + + self.assertEqual(spatial.shape, (2, 16, 8, 8)) + torch.testing.assert_close(_zimage_repack(spatial), latent) + + def test_wan_latent_adapter_is_identity(self): + stage = object.__new__(WanProgressiveDenoisingStage) + latent = torch.randn(1, 16, 5, 8, 8) + + self.assertIs(stage._unpack_latent(latent, 8, 8), latent) + self.assertIs( + stage._repack_latent(latent, 8, 8, SimpleNamespace(), SimpleNamespace()), + latent, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py index 6c653366f..84b7be560 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py @@ -38,11 +38,6 @@ from sglang.multimodal_gen.runtime.server_warmup import ( should_include_warmup_image, ) -# Patch path for get_global_server_args used by Stage.__init__ -_GLOBAL_ARGS_PATCH = ( - "sglang.multimodal_gen.runtime.pipelines_core.stages.base.get_global_server_args" -) - def _make_bare_scheduler(enable_cfg_parallel: bool) -> Scheduler: """ @@ -75,11 +70,7 @@ def _make_bare_scheduler(enable_cfg_parallel: bool) -> Scheduler: def _make_input_validation_stage() -> InputValidationStage: - """Construct InputValidationStage with the global server-args patch - that existing tests in this suite use (see test_input_validation.py).""" - with patch(_GLOBAL_ARGS_PATCH) as m: - m.return_value = MagicMock() - return InputValidationStage() + return InputValidationStage() def _make_validation_server_args(enable_cfg_parallel: bool) -> MagicMock: diff --git a/python/sglang/multimodal_gen/test/unit/test_input_validation.py b/python/sglang/multimodal_gen/test/unit/test_input_validation.py index 13fedecb8..b97cc06dd 100644 --- a/python/sglang/multimodal_gen/test/unit/test_input_validation.py +++ b/python/sglang/multimodal_gen/test/unit/test_input_validation.py @@ -21,11 +21,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import InputValidationStage, ) -# Patch path for get_global_server_args used by Stage.__init__ -_GLOBAL_ARGS_PATCH = ( - "sglang.multimodal_gen.runtime.pipelines_core.stages.base.get_global_server_args" -) - def _make_batch(condition_image: Image.Image, width=None, height=None) -> Req: """Create a minimal Req with a condition image and optional user dimensions.""" @@ -109,8 +104,7 @@ class TestPreprocessConditionImageResolution(unittest.TestCase): """ def setUp(self): - with patch(_GLOBAL_ARGS_PATCH, return_value=MagicMock()): - self.stage = InputValidationStage() + self.stage = InputValidationStage() def _run(self, config, img_w, img_h, user_w=None, user_h=None): """Run preprocess_condition_image and return (batch.width, batch.height).""" @@ -238,8 +232,7 @@ class TestFlux2ConditionImagePreprocess(unittest.TestCase): class TestFlux2TI2ISizeResolution(unittest.TestCase): def setUp(self): - with patch(_GLOBAL_ARGS_PATCH, return_value=MagicMock()): - self.stage = InputValidationStage() + self.stage = InputValidationStage() self.config = _DummyTI2IConfig() def test_uses_condition_image_size_when_width_height_not_explicit(self):