[diffusion] feat: progressive resolution growing for image and video models (#27524)

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,

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

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,
         )
 
 

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):

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

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

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
 

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

python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

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

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),
+        )

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),
+        )

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),
+        )

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)

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}")

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

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),
+        )

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:

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)

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()

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:

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):