[diffusion] model: sync with upstream z-Image (#17822)

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

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,

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

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(

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,