diff --git a/docs/cookbook/diffusion/Qwen-Image/Qwen-Image-Edit.mdx b/docs/cookbook/diffusion/Qwen-Image/Qwen-Image-Edit.mdx index 101f697db..5479fa733 100644 --- a/docs/cookbook/diffusion/Qwen-Image/Qwen-Image-Edit.mdx +++ b/docs/cookbook/diffusion/Qwen-Image/Qwen-Image-Edit.mdx @@ -44,6 +44,44 @@ See [Performance Optimization](/docs/sglang-diffusion/performance-optimization) - `--ulysses-degree`: The degree of DeepSpeed-Ulysses-style SP in USP - `--ring-degree`: The degree of ring attention-style SP in USP +### 3.3 Decompose an image into layers on H200 + +`Qwen/Qwen-Image-Layered` returns separate RGBA images. For this model, +`--num-frames 4` requests four output layers. The CLI saves all four PNGs, +and `DiffGenerator.generate()` returns one result per layer. + +On Linux with NVIDIA CUDA and two H200 GPUs, you can run the conditional and +unconditional branches on separate GPUs: + +```bash Command +CUDA_VISIBLE_DEVICES=0,1 sglang generate \ + --model-path Qwen/Qwen-Image-Layered \ + --num-gpus 2 \ + --cfg-parallel-size 2 \ + --tp-size 1 \ + --ulysses-degree 1 \ + --quality lossless \ + --enable-torch-compile false \ + --warmup-mode request \ + --width 640 --height 640 --num-frames 4 \ + --num-inference-steps 50 --guidance-scale 4.0 --seed 42 \ + --image-path https://raw.githubusercontent.com/QwenLM/Qwen-Image-Layered/main/assets/test_images/4.png \ + --prompt "a high quality, cute halloween themed illustration, consistent style and lighting" \ + --output-path outputs/qwen-layered \ + --save-output +``` + +For a single H200, set `CUDA_VISIBLE_DEVICES=0`, `--num-gpus 1`, and +`--cfg-parallel-size 1`. Both configurations use eager execution. Layered +does not currently support breakable CUDA graph; enabling BCG falls back to +eager execution. + +The Layered CFG policy gathers the branch predictions before applying the +single-GPU arithmetic order, preserving BF16 rounding and alpha values in the +validated fixed-seed example. Each GPU still holds a full DiT replica, so CFG +parallelism reduces request latency without reducing the model memory needed +on each GPU. + ## 4. API Usage For complete API documentation, please refer to the [official API usage guide](/docs/sglang-diffusion/api/openai_api). diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py index ed333b448..7b39cfcdc 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py @@ -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 diff --git a/python/sglang/multimodal_gen/configs/sample/qwenimage.py b/python/sglang/multimodal_gen/configs/sample/qwenimage.py index 8e0b09033..b7b5dc691 100644 --- a/python/sglang/multimodal_gen/configs/sample/qwenimage.py +++ b/python/sglang/multimodal_gen/configs/sample/qwenimage.py @@ -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 diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index e6950716e..f3e079bfd 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -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.""" diff --git a/python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py b/python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py index be2f750cd..fb3c000d7 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py +++ b/python/sglang/multimodal_gen/runtime/distributed/cfg_policy.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 712ba7d1b..d55bd454f 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -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: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index 2209dd108..d097cb226 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -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: diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 20be558a4..ed0532e1a 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -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, 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 64e43b569..4fa862cc9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -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, diff --git a/python/sglang/multimodal_gen/test/unit/test_qwen_image_layered.py b/python/sglang/multimodal_gen/test/unit/test_qwen_image_layered.py index 54edf15f2..eabb0ed44 100644 --- a/python/sglang/multimodal_gen/test/unit/test_qwen_image_layered.py +++ b/python/sglang/multimodal_gen/test/unit/test_qwen_image_layered.py @@ -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()