From 3c2f4c7bbe0c76ed900db2569d4b3802ed7fd24d Mon Sep 17 00:00:00 2001 From: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Date: Thu, 29 Jan 2026 21:10:11 +0800 Subject: [PATCH] [diffusion] model: sync with upstream z-Image (#17822) --- .../configs/pipeline_configs/zimage.py | 12 +++++++++ .../configs/sample/sampling_params.py | 13 ++++++++++ .../multimodal_gen/configs/sample/zimage.py | 13 +++++++++- python/sglang/multimodal_gen/registry.py | 19 +++++++++++--- .../pipelines_core/stages/denoising.py | 26 +++++++++++++++++++ 5 files changed, 79 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py index 29ba9e4f4..534c02da6 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py @@ -317,3 +317,15 @@ class ZImagePipelineConfig(ImagePipelineConfig): batch, ), } + + def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): + return { + "freqs_cis": self.get_freqs_cis( + batch.prompt_embeds[0], + batch.width, + batch.height, + device, + rotary_emb, + batch, + ), + } diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 4afae9505..f43625178 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -126,6 +126,7 @@ class SamplingParams: guidance_scale_2: float = None true_cfg_scale: float = None # for CFG vs guidance distillation (e.g., QwenImage) guidance_rescale: float = 0.0 + cfg_normalization: float | bool = 0.0 boundary_ratio: float | None = None # TeaCache parameters @@ -275,6 +276,11 @@ class SamplingParams: "guidance_rescale", self.guidance_rescale, allow_none=False ) + if self.cfg_normalization is None: + self.cfg_normalization = 0.0 + elif isinstance(self.cfg_normalization, bool): + self.cfg_normalization = 1.0 if self.cfg_normalization else 0.0 + if self.boundary_ratio is not None: if isinstance(self.boundary_ratio, bool) or not isinstance( self.boundary_ratio, (int, float) @@ -646,6 +652,13 @@ class SamplingParams: default=SamplingParams.guidance_rescale, help="Guidance rescale factor", ) + parser.add_argument( + "--cfg-normalization", + type=float, + default=SamplingParams.cfg_normalization, # type: ignore[arg-type] + dest="cfg_normalization", + help=("CFG renormalization factor (for Z-Image). "), + ) parser.add_argument( "--boundary-ratio", type=float, diff --git a/python/sglang/multimodal_gen/configs/sample/zimage.py b/python/sglang/multimodal_gen/configs/sample/zimage.py index f4530e977..77a9dabf9 100644 --- a/python/sglang/multimodal_gen/configs/sample/zimage.py +++ b/python/sglang/multimodal_gen/configs/sample/zimage.py @@ -8,7 +8,7 @@ from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams @dataclass -class ZImageSamplingParams(SamplingParams): +class ZImageTurboSamplingParams(SamplingParams): num_inference_steps: int = 9 num_frames: int = 1 @@ -18,6 +18,7 @@ class ZImageSamplingParams(SamplingParams): # fps: int = 24 guidance_scale: float = 0.0 + cfg_normalization: float | bool = False teacache_params: TeaCacheParams = field( default_factory=lambda: TeaCacheParams( @@ -31,3 +32,13 @@ class ZImageSamplingParams(SamplingParams): ], ) ) + + +@dataclass +class ZImageSamplingParams(SamplingParams): + num_inference_steps: int = 50 + + num_frames: int = 1 + negative_prompt: str = " " + guidance_scale: float = 5.0 + cfg_normalization: float | bool = True diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 6d239c8f9..10c769936 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -98,7 +98,10 @@ from sglang.multimodal_gen.configs.sample.wan import ( WanT2V_1_3B_SamplingParams, WanT2V_14B_SamplingParams, ) -from sglang.multimodal_gen.configs.sample.zimage import ZImageSamplingParams +from sglang.multimodal_gen.configs.sample.zimage import ( + ZImageSamplingParams, + ZImageTurboSamplingParams, +) from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) @@ -598,12 +601,22 @@ def _register_configs(): ], ) register_configs( - sampling_param_cls=ZImageSamplingParams, + sampling_param_cls=ZImageTurboSamplingParams, pipeline_config_cls=ZImagePipelineConfig, hf_model_paths=[ "Tongyi-MAI/Z-Image-Turbo", ], - model_detectors=[lambda hf_id: "z-image" in hf_id.lower()], + model_detectors=[lambda hf_id: "z-image-turbo" in hf_id.lower()], + ) + register_configs( + sampling_param_cls=ZImageSamplingParams, + pipeline_config_cls=ZImagePipelineConfig, + hf_model_paths=[ + "Tongyi-MAI/Z-Image", + ], + model_detectors=[ + lambda hf_id: "z-image" in hf_id.lower() and "turbo" not in hf_id.lower() + ], ) # Qwen-Image register_configs( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index ce4c8a86f..0fab36a7c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -1316,6 +1316,21 @@ class DenoisingStage(PipelineStage): noise_pred = cfg_model_parallel_all_reduce(partial) + if batch.cfg_normalization and float(batch.cfg_normalization) > 0: + factor = float(batch.cfg_normalization) + pred_f = noise_pred.float() + new_norm = torch.linalg.vector_norm(pred_f) + if cfg_rank == 0: + cond_f = noise_pred_cond.float() + ori_norm = torch.linalg.vector_norm(cond_f) + else: + ori_norm = torch.empty_like(new_norm) + ori_norm = get_cfg_group().broadcast(ori_norm, src=0) + max_norm = ori_norm * factor + + if new_norm > max_norm: + noise_pred = noise_pred * (max_norm / new_norm) + # Guidance rescale: broadcast std(cond) from rank 0, compute std(cfg) locally if batch.guidance_rescale > 0.0: std_cfg = noise_pred.std( @@ -1343,6 +1358,17 @@ class DenoisingStage(PipelineStage): noise_pred_cond - noise_pred_uncond ) + if batch.cfg_normalization and float(batch.cfg_normalization) > 0: + factor = float(batch.cfg_normalization) + cond_f = noise_pred_cond.float() + pred_f = noise_pred.float() + ori_norm = torch.linalg.vector_norm(cond_f) + new_norm = torch.linalg.vector_norm(pred_f) + max_norm = ori_norm * factor + + if new_norm > max_norm: + noise_pred = noise_pred * (max_norm / new_norm) + if batch.guidance_rescale > 0.0: noise_pred = self.rescale_noise_cfg( noise_pred,