chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Yuhao Yang
2026-01-29 21:10:11 +08:00
committed by GitHub
parent 30adf78f82
commit 3c2f4c7bbe
5 changed files with 79 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
View File

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