minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Diffusion] Return Qwen-Image-Layered outputs and preserve CFG2 rounding (#38549)

Browse Source 
Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
Xiaoyu Zhang
2026-09-13 09:15:47 +08:00
committed by GitHub
co-authored by Mick Qian
parent 3e035a3513
commit 23bc4c6ed9
10 changed files with 341 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py
+4
View File
@@ -26,6 +26,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config impo
from sglang.multimodal_gen.configs.post_training.pipeline_configs import (
QwenImageRolloutPipelineMixin,
)
from sglang.multimodal_gen.runtime.distributed.cfg_policy import CFGPolicy
from sglang.multimodal_gen.runtime.utils.condition_expansion import (
PromptToSampleBatchExpander,
)
@@ -778,6 +779,9 @@ class QwenImageEditPlus_2511_PipelineConfig(QwenImageEditPlusPipelineConfig):
class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig):
resolution: int = 640
vae_precision: str = "bf16"
cfg_policy: CFGPolicy = field(
default_factory=lambda: CFGPolicy(parallel_uses_serial_arithmetic=True)
)
# promoting the auxiliary components regresses first-request latency
supports_auto_residency: bool = False
python/sglang/multimodal_gen/configs/sample/qwenimage.py
+4
View File
@@ -40,3 +40,7 @@ class QwenImageLayeredSamplingParams(QwenImageSamplingParams):
num_inference_steps: int = 50
cfg_normalize: bool = True
use_en_prompt: bool = True
@property
def num_samples_per_request(self) -> int:
return self.num_frames
python/sglang/multimodal_gen/configs/sample/sampling_params.py
+5
View File
@@ -507,6 +507,11 @@ class SamplingParams:
return frozenset()
@property
def num_samples_per_request(self) -> int:
"""Number of final samples produced by one expanded scheduler request."""
return 1
@classmethod
def default_image_output_format(cls) -> str | None:
"""Return a model-owned default format for the image API, if any."""
python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py
+4 -1
View File
@@ -43,6 +43,9 @@ class CFGPolicy:
"""
branches: list[CFGBranch] = field(default_factory=list)
# Gather predictions before combining when a model needs the same bf16
# rounding as serial CFG. The default retains legacy WAN all-reduce outputs.
parallel_uses_serial_arithmetic: bool = False
def build(
self,
@@ -83,7 +86,7 @@ class CFGPolicy:
return predictions[0]
pos_t = _wrap(predictions[0])
neg_t = _wrap(predictions[1])
if cfg_parallel:
if cfg_parallel and not self.parallel_uses_serial_arithmetic:
# Match the old CFG-parallel calculation: multiply the positive
# prediction by cfg_scale and the negative prediction by
# (1 - cfg_scale) before adding them. The serial CFG formula is
python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py
+30 -14
View File
@@ -30,6 +30,7 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
GenerationResult,
expand_request_outputs,
format_lora_message,
map_request_outputs,
prepare_request,
save_outputs,
)
@@ -234,8 +235,9 @@ class DiffGenerator:
) -> GenerationResult | list[GenerationResult] | None:
"""Generate image(s)/video(s) based on the given prompt(s).
Returns a single GenerationResult for a single prompt, a list for
multiple prompts, or None when every request failed.
Returns one GenerationResult per final sample, including each layer
of a layered image. Returns a single result without a list wrapper,
or None when every request failed.
"""
# 1. prepare requests
prompts = self._resolve_prompts(
@@ -307,7 +309,9 @@ class DiffGenerator:
global_output_index = 0
for requests in request_groups:
output_requests = []
try:
output_requests = map_request_outputs(requests)
timer_prompt = [req.prompt for req in requests]
logger.info("Processing %d grouped request(s)", len(requests))
with ExitStack() as stack:
@@ -332,10 +336,11 @@ class DiffGenerator:
if requests[0].save_output and requests[0].return_file_paths_only:
output_file_paths = output_batch.output_file_paths or []
self._validate_output_count(
len(output_file_paths), len(requests)
len(output_file_paths), len(output_requests)
)
for idx, path in enumerate(output_file_paths):
req = requests[idx]
output_request = output_requests[idx]
req = output_request.request
if req.data_type == DataType.VIDEO:
req.sampling_params.validate_video_final_outputs(
[path], req
@@ -343,7 +348,10 @@ class DiffGenerator:
results.append(
GenerationResult(
**self._result_common(
req, output_batch, timer.duration, idx
req,
output_batch,
timer.duration,
output_request.request_index,
),
prompt_index=global_output_index + idx,
output_file_path=path,
@@ -352,14 +360,18 @@ class DiffGenerator:
elif requests[0].data_type == DataType.MESH:
output_file_paths = output_batch.output_file_paths or []
self._validate_output_count(
len(output_file_paths), len(requests)
len(output_file_paths), len(output_requests)
)
for idx, sample in enumerate(output_file_paths):
req = requests[idx]
output_request = output_requests[idx]
req = output_request.request
results.append(
GenerationResult(
**self._result_common(
req, output_batch, timer.duration, idx
req,
output_batch,
timer.duration,
output_request.request_index,
),
prompt_index=global_output_index + idx,
output_file_path=sample,
@@ -367,7 +379,7 @@ class DiffGenerator:
)
else:
self._validate_output_count(
len(output_batch.output), len(requests)
len(output_batch.output), len(output_requests)
)
samples_out: list[Any] = []
audios_out: list[Any] = []
@@ -377,7 +389,7 @@ class DiffGenerator:
requests[0].data_type,
requests[0].fps,
requests[0].save_output,
lambda idx: requests[idx].output_file_path(1, 0),
lambda idx: output_requests[idx].output_file_path(),
audio=output_batch.audio,
audio_sample_rate=output_batch.audio_sample_rate,
samples_out=samples_out,
@@ -400,8 +412,9 @@ class DiffGenerator:
)
for idx in range(len(samples_out)):
req = requests[idx]
output_file_path = req.output_file_path(1, 0)
output_request = output_requests[idx]
req = output_request.request
output_file_path = output_request.output_file_path()
if req.data_type == DataType.VIDEO and req.save_output:
req.sampling_params.validate_video_final_outputs(
[output_file_path], req
@@ -409,7 +422,10 @@ class DiffGenerator:
results.append(
GenerationResult(
**self._result_common(
req, output_batch, timer.duration, idx
req,
output_batch,
timer.duration,
output_request.request_index,
),
samples=samples_out[idx],
frames=frames_out[idx],
@@ -432,7 +448,7 @@ class DiffGenerator:
"Failed to clean up model-owned video request resources",
exc_info=True,
)
global_output_index += len(requests)
global_output_index += len(output_requests)
total_gen_time = time.perf_counter() - total_start_time
if self.server_args.batching_max_size > 1:
python/sglang/multimodal_gen/runtime/entrypoints/utils.py
+25
View File
@@ -219,6 +219,31 @@ class MaterializedOutput:
fps: int = 0
@dataclass(frozen=True)
class RequestOutput:
"""Map one final sample to its request, metrics and output filename."""
request: Req
request_index: int
sample_index: int
sample_count: int
def output_file_path(self):
return self.request.output_file_path(self.sample_count, self.sample_index)
def map_request_outputs(requests: list[Req]) -> list[RequestOutput]:
outputs = []
for request_index, req in enumerate(requests):
count = req.sampling_params.num_samples_per_request
if count < 1:
raise ValueError(f"num_samples_per_request must be positive, got {count}")
outputs.extend(
RequestOutput(req, request_index, index, count) for index in range(count)
)
return outputs
def _normalize_audio_to_numpy(audio: Any) -> np.ndarray | None:
"""Convert audio (torch / numpy) into a float32 numpy array in [-1, 1], best-effort."""
if audio is None:
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
+5 -3
View File
@@ -47,6 +47,7 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_ulysses_parallel_world_size,
)
from sglang.multimodal_gen.runtime.entrypoints.utils import (
map_request_outputs,
materialize_output_sample,
post_process_sample,
save_outputs,
@@ -1196,9 +1197,10 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
) -> None:
if not self.is_output_rank or output_batch.output is None:
return
if len(output_batch.output) != len(reqs):
output_requests = map_request_outputs(reqs)
if len(output_batch.output) != len(output_requests):
raise RuntimeError(
f"Expected {len(reqs)} grouped outputs, got {len(output_batch.output)}"
f"Expected {len(output_requests)} grouped outputs, got {len(output_batch.output)}"
)
first_req = reqs[0]
@@ -1207,7 +1209,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
first_req.data_type,
first_req.fps,
True,
lambda idx: reqs[idx].output_file_path(1, 0),
lambda idx: output_requests[idx].output_file_path(),
audio=output_batch.audio,
audio_sample_rate=output_batch.audio_sample_rate,
output_compression=first_req.output_compression,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+1
View File
@@ -2254,6 +2254,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
if (
len(cfg_policy.branches) == 2
and get_classifier_free_guidance_world_size() == 2
and not cfg_policy.parallel_uses_serial_arithmetic
):
return run_two_branch_cfg_parallel(
cfg_policy,
python/sglang/multimodal_gen/test/unit/test_qwen_image_layered.py
+225
View File
@@ -1,11 +1,30 @@
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
import torch
from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
QwenImageLayeredPipelineConfig,
QwenImagePipelineConfig,
)
from sglang.multimodal_gen.configs.sample.qwenimage import (
QwenImageLayeredSamplingParams,
)
from sglang.multimodal_gen.configs.sample.sampling_params import (
DataType,
SamplingParams,
)
from sglang.multimodal_gen.runtime.distributed.cfg_policy import CFGPolicy
from sglang.multimodal_gen.runtime.entrypoints import diffusion_generator as dg
from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
from sglang.multimodal_gen.runtime.entrypoints.utils import map_request_outputs
from sglang.multimodal_gen.runtime.managers import gpu_worker
from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage
from sglang.test.test_utils import CustomTestCase
class TestQwenImageLayeredPipelineConfig(unittest.TestCase):
@@ -43,5 +62,211 @@ class TestQwenImageLayeredPipelineConfig(unittest.TestCase):
self.assertEqual(unpacked.shape, (1, channels // 4, generated_layers, 80, 80))
class TestLayeredOutputMapping(CustomTestCase):
def _params(self, **kwargs):
return QwenImageLayeredSamplingParams(
prompt="layers",
num_frames=4,
output_path="/tmp",
output_file_name="layers.png",
**kwargs,
)
def _generate(self, params, response, save_mock=None):
generator = object.__new__(DiffGenerator)
generator.local_scheduler_process = []
generator.owns_scheduler_client = False
generator.server_args = SimpleNamespace(
model_path="Qwen/Qwen-Image-Layered",
prompt_file_path=None,
warmup_mode="off",
batching_max_size=1,
)
with (
patch.object(
SamplingParams, "from_user_sampling_params_args", return_value=params
),
patch.object(
dg,
"prepare_request",
side_effect=lambda **kw: Req(sampling_params=kw["sampling_params"]),
),
patch.object(
generator,
"_send_to_scheduler_and_wait_for_response",
return_value=response,
),
patch.object(dg, "save_outputs", side_effect=save_mock),
):
return generator.generate(
{"prompt": "layers", "output_file_name": "layers.png"}
)
def test_all_four_saved_layers_are_returned(self):
paths = [f"/tmp/layers_{i}.png" for i in range(4)]
result = self._generate(
self._params(save_output=True, return_file_paths_only=True),
OutputBatch(output_file_paths=paths),
)
self.assertEqual([r.output_file_path for r in result], paths)
self.assertEqual([r.prompt_index for r in result], [0, 1, 2, 3])
def test_grouped_layers_keep_their_request_metrics(self):
paths = [
f"/tmp/layers_{draw}_{layer}.png" for draw in range(2) for layer in range(4)
]
metrics = [
SimpleNamespace(to_dict=lambda index=index: {"request": index})
for index in range(2)
]
result = self._generate(
self._params(
save_output=True, return_file_paths_only=True, num_outputs_per_prompt=2
),
OutputBatch(output_file_paths=paths, metrics_list=metrics),
)
self.assertEqual([r.output_file_path for r in result], paths)
self.assertEqual([r.metrics["request"] for r in result], [0] * 4 + [1] * 4)
self.assertEqual([r.prompt_index for r in result], list(range(8)))
def test_returned_samples_use_unique_layer_filenames(self):
samples = [np.full((2, 2, 4), index, dtype=np.uint8) for index in range(4)]
paths = []
def save(outputs, data_type, fps, should_save, build_path, **kwargs):
paths.extend(build_path(index) for index in range(len(outputs)))
kwargs["samples_out"].extend(samples)
kwargs["frames_out"].extend([None] * 4)
kwargs["audios_out"].extend([None] * 4)
result = self._generate(
self._params(save_output=False, return_file_paths_only=False),
OutputBatch(output=samples),
save_mock=save,
)
self.assertEqual(paths, [f"/tmp/layers_{i}.png" for i in range(4)])
self.assertEqual([r.output_file_path for r in result], paths)
for actual, expected in zip(result, samples):
self.assertIs(actual.samples, expected)
def test_partial_layer_response_is_rejected(self):
with self.assertLogs(dg.logger, level="ERROR") as logs:
result = self._generate(
self._params(save_output=True, return_file_paths_only=True),
OutputBatch(output_file_paths=["a.png", "b.png", "c.png"]),
)
self.assertIsNone(result)
self.assertTrue(
any("Expected 4 outputs, got 3" in line for line in logs.output)
)
def test_video_frames_do_not_expand_into_separate_results(self):
req = Req(
sampling_params=SamplingParams(data_type=DataType.VIDEO, num_frames=121)
)
outputs = map_request_outputs([req])
self.assertEqual(len(outputs), 1)
self.assertIs(outputs[0].request, req)
def test_worker_group_saves_each_layer_under_its_parent(self):
requests = [
Req(sampling_params=self._params(request_id=f"draw{i}")) for i in range(2)
]
for index, req in enumerate(requests):
req.output_file_name = f"draw{index}.png"
paths = []
def save(outputs, data_type, fps, should_save, build_path, **kwargs):
paths.extend(build_path(index) for index in range(len(outputs)))
return paths
batch = OutputBatch(output=[None] * 8)
with patch.object(gpu_worker, "save_outputs", side_effect=save):
GPUWorker._save_group_output_paths(
SimpleNamespace(is_output_rank=True), requests, batch
)
expected = [
f"/tmp/draw{draw}_{layer}.png" for draw in range(2) for layer in range(4)
]
self.assertEqual(paths, expected)
self.assertEqual(batch.output_file_paths, expected)
class TestLayeredCFGOrder(CustomTestCase):
def test_layered_parallel_uses_serial_formula_and_postprocess(self):
config = QwenImageLayeredPipelineConfig()
req = SimpleNamespace(
do_classifier_free_guidance=True,
true_cfg_scale=7.0,
cfg_normalization=0,
guidance_rescale=0,
cfg_normalize=True,
)
policy = config.cfg_policy.build(req, {}, {}, {})
self.assertTrue(policy.parallel_uses_serial_arithmetic)
self.assertFalse(
QwenImagePipelineConfig().cfg_policy.parallel_uses_serial_arithmetic
)
pos = torch.tensor([[1.0, 0.3]], dtype=torch.bfloat16)
neg = torch.tensor([[0.1, -0.2]], dtype=torch.bfloat16)
serial = policy.combine([pos, neg], req, 7.0, config)
parallel = policy.combine([pos, neg], req, 7.0, config, cfg_parallel=True)
self.assertTrue(torch.equal(serial, parallel))
self.assertFalse(torch.equal(serial, neg + 7.0 * (pos - neg)))
def test_dispatch_gathers_for_layered_and_preserves_legacy_fast_path(self):
stage = DenoisingStage.__new__(DenoisingStage)
batch = SimpleNamespace(
do_classifier_free_guidance=True,
cfg_normalization=0,
guidance_rescale=0,
)
config = SimpleNamespace(
get_classifier_free_guidance_scale=lambda batch, scale: scale,
postprocess_cfg_noise=lambda batch, pred, cond: pred,
)
pos = torch.tensor([1.0], dtype=torch.bfloat16)
neg = torch.tensor([0.1], dtype=torch.bfloat16)
module = "sglang.multimodal_gen.runtime.pipelines_core.stages.denoising"
for same_order in [False, True]:
with self.subTest(same_order=same_order):
policy = CFGPolicy(parallel_uses_serial_arithmetic=same_order).build(
batch, {}, {}, {}
)
with (
patch(
f"{module}.get_classifier_free_guidance_world_size",
return_value=2,
),
patch(
f"{module}.run_cfg_parallel", return_value=[pos, neg]
) as gather,
patch(
f"{module}.run_two_branch_cfg_parallel", return_value=pos
) as reduce,
):
result = stage._predict_noise_with_cfg(
current_model=None,
latent_model_input=pos,
timestep=torch.tensor(1),
batch=batch,
timestep_index=0,
attn_metadata=None,
target_dtype=torch.bfloat16,
current_guidance_scale=7.0,
cfg_policy=policy,
cfg_gate_state=None,
server_args=SimpleNamespace(
enable_cfg_parallel=True, pipeline_config=config
),
guidance=None,
latents=pos,
)
self.assertEqual(gather.call_count, int(same_order))
self.assertEqual(reduce.call_count, int(not same_order))
expected = neg + 7.0 * (pos - neg) if same_order else pos
self.assertTrue(torch.equal(result, expected))
if __name__ == "__main__":
unittest.main()