update

Author: DN6

docs/source/en/using-diffusers/cli.md

@@ -141,10 +141,15 @@ Configure how the CLI loads model weights and custom pipeline code.
   pipelines. See [Attention backends](../optimization/attention_backends).
 - `--vae-tiling` / `--vae-slicing` — lower VAE decode VRAM. See
   [VAE tiling](../optimization/memory#vae-tiling) and [VAE slicing](../optimization/memory#vae-slicing).
-- `--compile [JSON]` — compile denoiser modules with [torch.compile](../optimization/fp16#torchcompile). The
-  CLI prefers [regional compilation](../optimization/fp16#regional-compilation) for modules with repeated
-  blocks. Bare `--compile` uses `fullgraph=true`. A JSON object is forwarded to `torch.compile`. Not supported
-  with `--context-parallel`.
+- `--compile [JSON]` — compile denoiser modules and the VAE decoder with
+  [torch.compile](../optimization/fp16#torchcompile). The CLI prefers
+  [regional compilation](../optimization/fp16#regional-compilation) for denoisers with repeated blocks; the VAE
+  decoder is always compiled whole. Bare `--compile` uses `fullgraph=true`. A JSON object is forwarded to
+  `torch.compile`. Not supported with `--context-parallel`.
+- `--compile-mode {regional,full}` — what `--compile` compiles. `regional` (default) compiles only the repeated
+  blocks, which keeps the first step fast. `full` compiles the whole denoiser. `torch.compile` modes that use CUDA
+  graphs (`max-autotune`, `reduce-overhead`) need `--compile-mode full`: repeated blocks share one compiled graph,
+  so with CUDA graphs each block's output is overwritten by the next block's run.
 - `--context-parallel` — Ulysses-style context parallelism on a DiT-based pipeline. Locally requires torchrun;
   under `--remote` the CLI wraps `torchrun --nproc-per-node=gpu` for you. See
   [Context parallelism](../training/distributed_inference#context-parallelism).
@@ -375,6 +380,14 @@ diffusers-cli generate --sandbox-id <id> --inputs '{"prompt": "a cat on the moon
 [`run --remote`](#remote-execution---remote). `--model` must be a Hub repo id. A local `--manifest` file is
 uploaded to the sandbox.
 
+With `--backend sglang` or `--backend vllm`, the sandbox starts from the engine's own image
+(`lmsysorg/sglang:latest` or `vllm/vllm-omni:latest`) unless `--image` names another, so the engine does not need
+installing:
+
+```bash
+diffusers-cli serve --model black-forest-labs/FLUX.2-klein-4B --backend sglang --remote --flavor a100-large
+```
+
 ### Manifests
 
 What a model accepts is described by its manifest. `serve` looks for one in this order:
@@ -425,8 +438,8 @@ diffusers-cli serve --model Qwen/Qwen-Image --backend sglang --backend-args '--n
 diffusers-cli serve --model Wan-AI/Wan2.2-T2V-A14B-Diffusers --backend vllm
 ```
 
-The pipeline flags (`--dtype`, `--cpu-offload`, `--compile`, ...) do not apply to these backends; pass the
-engine's own flags with `--backend-args`. Compared with the in-process backend, they serve image and video
+The pipeline flags (`--dtype`, `--cpu-offload`, `--compile`, ...) do not apply to these backends: they are
+ignored with a warning. Pass the engine's own flags with `--backend-args`. Compared with the in-process backend, they serve image and video
 outputs only, accept `image` as the only media input, expose the parameters both engines share (size, steps,
 guidance scale, seed, number of frames, fps, number of images), report no step progress, and apply no request
 adapters.

src/diffusers/commands/run.py

@@ -63,6 +63,14 @@
 
 
 ATTENTION_BACKEND_CHOICES = ("default", *sorted(b.value for b in _HUB_KERNELS_REGISTRY))
+COMPILE_MODE_CHOICES = ("regional", "full")
+# zlib level for PNG outputs. Pillow's default of 6 takes close to four times as long to encode a
+# 1024x1024 image for a file about 7% smaller; the output is lossless either way.
+PNG_COMPRESS_LEVEL = 1
+
+# `torch.compile` modes that replay CUDA graphs. Repeated blocks share one compiled graph, so each
+# block's output is overwritten when the next block replays it.
+_CUDA_GRAPH_COMPILE_MODES = ("max-autotune", "reduce-overhead")
 
 # Kwarg keys whose string value gets auto-loaded before being passed to the pipeline call.
 # Images resolve via `diffusers.utils.load_image` → PIL.Image.Image; videos resolve via
@@ -211,13 +219,23 @@ def _add_optimization_arguments(parser: ArgumentParser) -> None:
         default=None,
         metavar="JSON",
         help=(
-            "torch.compile every denoiser submodule on the pipeline. Accepts an optional JSON "
+            "torch.compile every denoiser submodule on the pipeline, and the VAE decoder. Accepts an optional JSON "
             'object of kwargs forwarded to `torch.compile`, e.g. \'{"mode": "max-autotune", '
             '"fullgraph": true}\'. Bare `--compile` uses `fullgraph=true`. Adds a one-time '
             "compilation cost on the first step but speeds up every subsequent step — worth it "
             "for multi-step generation (50+ steps)."
         ),
     )
+    parser.add_argument(
+        "--compile-mode",
+        choices=COMPILE_MODE_CHOICES,
+        default="regional",
+        help=(
+            "What --compile compiles. 'regional' (default) compiles only the repeated blocks of a denoiser "
+            "that declares them, which keeps the first step fast. 'full' compiles the whole denoiser, which "
+            "`torch.compile` modes that use CUDA graphs (max-autotune, reduce-overhead) require."
+        ),
+    )
 
 
 def _add_output_arguments(parser: ArgumentParser) -> None:
@@ -461,26 +479,44 @@ def _apply_optimizations(pipeline: Any, args: Namespace) -> None:
         if args.context_parallel:
             logger.warning("--compile is currently not supported with --context-parallel; skipping compile.")
         else:
-            _compile_denoiser(pipeline, args.compile)
+            _compile_denoiser(pipeline, args.compile, args.compile_mode)
+            _compile_vae_decoder(pipeline, args.compile)
+
+
+def _parse_compile_spec(compile_spec: str) -> dict[str, Any]:
+    try:
+        compile_kwargs = json.loads(compile_spec)
+    except json.JSONDecodeError as e:
+        raise SystemExit(f"--compile must be valid JSON: {e}") from e
+    if not isinstance(compile_kwargs, dict):
+        raise SystemExit("--compile must decode to a JSON object.")
+    return compile_kwargs
 
 
-def _compile_denoiser(pipeline: Any, compile_spec: str) -> None:
+def _compile_vae_decoder(pipeline: Any, compile_spec: str) -> None:
+    """Compile the decoder of the pipeline's VAE with the `--compile` kwargs.
+
+    The decoder runs once per call, so it has no repeated blocks to compile regionally and is always compiled whole.
+    """
+    decoder = getattr(getattr(pipeline, "vae", None), "decoder", None)
+    if not isinstance(decoder, torch.nn.Module):
+        return
+    pipeline.vae.decoder = torch.compile(decoder, **_parse_compile_spec(compile_spec))
+
+
+def _compile_denoiser(pipeline: Any, compile_spec: str, compile_mode: str = "regional") -> None:
     """Compile every `transformer*` and `unet*` submodule on the pipeline.
 
     `compile_spec` is the raw JSON string from `--compile` (`"{}"` for bare flag). Decoded into kwargs and forwarded
     verbatim to the compile call.
 
-    Prefers regional compilation via `module.compile_repeated_blocks(**kwargs)` — only compiles the repeated inner
-    blocks (the bulk of the compute), much faster first-step latency than compiling the whole module. Falls back to
-    full `torch.compile` if the model doesn't expose `_repeated_blocks`.
+    With `compile_mode="regional"`, prefers regional compilation via `module.compile_repeated_blocks(**kwargs)` — only
+    compiles the repeated inner blocks (the bulk of the compute), much faster first-step latency than compiling the
+    whole module. Falls back to full `torch.compile` if the model doesn't expose `_repeated_blocks`. With
+    `compile_mode="full"`, always compiles the whole module.
     """
 
-    try:
-        compile_kwargs = json.loads(compile_spec)
-    except json.JSONDecodeError as e:
-        raise SystemExit(f"--compile must be valid JSON: {e}") from e
-    if not isinstance(compile_kwargs, dict):
-        raise SystemExit("--compile must decode to a JSON object.")
+    compile_kwargs = _parse_compile_spec(compile_spec)
 
     for attr in dir(pipeline):
         if not any(attr.startswith(key) for key in _DENOISER_COMPONENT_KEYS):
@@ -489,11 +525,17 @@ def _compile_denoiser(pipeline: Any, compile_spec: str) -> None:
         if not isinstance(module, torch.nn.Module):
             continue
 
-        if getattr(module, "_repeated_blocks", None):
+        if compile_mode == "regional" and getattr(module, "_repeated_blocks", None):
+            if compile_kwargs.get("mode") in _CUDA_GRAPH_COMPILE_MODES:
+                raise SystemExit(
+                    f"--compile mode {compile_kwargs['mode']!r} uses CUDA graphs, which fail when the repeated "
+                    f"blocks of {type(module).__name__} are compiled separately. Add `--compile-mode full`, or "
+                    "use mode 'max-autotune-no-cudagraphs'."
+                )
             # Regional compile — only the repeated blocks. Mutates `module` in place.
             module.compile_repeated_blocks(**compile_kwargs)
         else:
-            # No regional metadata declared; fall back to compiling the whole module.
+            # Full compile was asked for, or no regional metadata is declared.
             setattr(pipeline, attr, torch.compile(module, **compile_kwargs))
 
 
@@ -535,6 +577,9 @@ def _load_lora(pipeline: Any, args: Namespace) -> None:
 
 
 def _load_pipeline(args: Namespace) -> Any:
+    if args.compile is None and args.compile_mode != "regional":
+        raise SystemExit(f"--compile-mode {args.compile_mode} only applies together with --compile.")
+
     # Detect modular repos by trying the standard config; `ModularPipeline` repos ship
     # `modular_model_index.json` instead of `model_index.json`, so `load_config` OSErrors.
     # A repo can also ship a `model_index.json` whose `_class_name` is a modular pipeline
@@ -800,7 +845,7 @@ def _save_videos(videos: list[Any], args: Namespace) -> list[str]:
                     arr = (np.clip(arr, 0.0, 1.0) * 255).round().astype(np.uint8)
                 frame = Image.fromarray(arr)
             frame_path = frames_dir / f"{i:04d}.png"
-            frame.save(frame_path)
+            frame.save(frame_path, compress_level=PNG_COMPRESS_LEVEL)
             saved.append(str(frame_path))
         try:
             export_to_video(frames, str(path), fps=args.fps)
@@ -837,14 +882,14 @@ def _save_output(value: Any, args: Namespace) -> list[str]:
             for arr, path in zip(value, paths):
                 if arr.dtype != np.uint8:
                     arr = (np.clip(arr, 0.0, 1.0) * 255).round().astype(np.uint8)
-                Image.fromarray(arr).save(path)
+                Image.fromarray(arr).save(path, compress_level=PNG_COMPRESS_LEVEL)
             return [str(p) for p in paths]
 
     pil_images = _as_pil_list(value)
     if pil_images is not None:
         paths = _resolve_output_paths(len(pil_images), args.output, ext="png")
         for img, path in zip(pil_images, paths):
-            img.save(path)
+            img.save(path, compress_level=PNG_COMPRESS_LEVEL)
         return [str(p) for p in paths]
 
     frames = _as_frame_sequence(value)

src/diffusers/commands/serve/__init__.py

@@ -33,6 +33,8 @@
 from .remote import add_remote_arguments, serve_remote
 
 
+logger = logging.get_logger("diffusers-cli/serve")
+
 DEFAULT_OUTPUT_DIR = str(Path.home() / ".diffusers" / "cli" / "serve" / "outputs")
 BACKEND_CHOICES = ("diffusers", "sglang", "vllm")
 
@@ -47,6 +49,7 @@
     "vae_tiling": False,
     "vae_slicing": False,
     "compile": None,
+    "compile_mode": "regional",
 }
 
 
@@ -212,14 +215,14 @@ def run(self) -> None:
         if args.backend == "diffusers" and args.backend_args is not None:
             raise SystemExit("--backend-args only applies to --backend sglang or vllm.")
         if args.backend != "diffusers":
-            rejected = [
+            ignored = [
                 "--" + name.replace("_", "-")
                 for name, unset in _DIFFUSERS_ONLY_FLAGS.items()
                 if getattr(args, name) != unset
             ]
-            if rejected:
-                raise SystemExit(
-                    f"{', '.join(rejected)} configure the in-process pipeline and have no effect with "
+            if ignored:
+                logger.warning(
+                    f"{', '.join(ignored)} configure the in-process pipeline and are ignored with "
                     f"--backend {args.backend}. Pass the engine's own flags through --backend-args instead."
                 )
 

src/diffusers/commands/serve/backends.py

@@ -41,7 +41,14 @@
 from ...modular_pipelines import ModularPipeline
 from ...utils import export_to_video, load_image, load_video, logging, numpy_to_pil
 from ...utils.constants import DIFFUSERS_REQUEST_TIMEOUT
-from ..run import _as_audio_arrays, _as_pil_list, _get_generator, _load_pipeline, _save_audio_arrays
+from ..run import (
+    PNG_COMPRESS_LEVEL,
+    _as_audio_arrays,
+    _as_pil_list,
+    _get_generator,
+    _load_pipeline,
+    _save_audio_arrays,
+)
 from .generations import Generation
 from .manifest import GMSError, ServedModel, blocks_signature
 
@@ -144,7 +151,10 @@ def _save_images(images: Any, directory: Path, name: str, media_type: str, rate:
         filename = f"{name}-{index}.{media_type.split('/')[1]}"
         if media_type == "image/jpeg":
             image = image.convert("RGB")
-        image.save(directory / filename)
+        if media_type == "image/png":
+            image.save(directory / filename, compress_level=PNG_COMPRESS_LEVEL)
+        else:
+            image.save(directory / filename)
         artifacts.append({"file": filename, "media_type": media_type, "width": image.width, "height": image.height})
     return artifacts
 

src/diffusers/commands/serve/remote.py

@@ -37,6 +37,8 @@
 SANDBOX_PORT = 8000
 _SANDBOX_MANIFEST_PATH = "/tmp/diffusers-cli/gms.json"
 _SERVE_DEPS = ("fastapi", "uvicorn")
+# Each engine's own image carries the engine, its kernels and the CUDA toolchain they are built with.
+_ENGINE_IMAGES = {"sglang": "lmsysorg/sglang:latest", "vllm": "vllm/vllm-omni:latest"}
 _READY_POLL_SECONDS = 5.0
 
 # Flags that say how the sandbox is set up, or that the sandbox server sets itself.
@@ -69,7 +71,10 @@ def add_remote_arguments(parser: ArgumentParser) -> None:
     parser.add_argument(
         "--image",
         default=None,
-        help=f"Sandbox image for --remote (defaults to {_DEFAULT_REMOTE_IMAGE!r}). Must provide torch + CUDA.",
+        help=(
+            f"Sandbox image for --remote. Defaults to {_DEFAULT_REMOTE_IMAGE!r}, or to the engine's own image "
+            f"with --backend sglang ({_ENGINE_IMAGES['sglang']!r}) or vllm ({_ENGINE_IMAGES['vllm']!r})."
+        ),
     )
     parser.add_argument(
         "--idle-timeout",
@@ -109,7 +114,7 @@ def serve_remote(args: Namespace, task: str) -> None:
     hf_token = args.token or get_token()
     logger.info(f"creating sandbox on flavor={args.flavor!r}...")
     create_kwargs: dict[str, Any] = {
-        "image": args.image or _DEFAULT_REMOTE_IMAGE,
+        "image": args.image or _ENGINE_IMAGES.get(args.backend, _DEFAULT_REMOTE_IMAGE),
         "flavor": args.flavor,
         "forward_hf_token": True,
         "token": hf_token,
@@ -133,6 +138,9 @@ def _stream(chunk: str) -> None:
         install_cmd = shlex.join(["uv", "pip", "install", "--system", "--break-system-packages", *dependencies])
         logger.info("installing dependencies in the sandbox...")
         sbx.run(install_cmd, on_stdout=_stream, on_stderr=_stream)
+        # An image can put another environment first on PATH, with its own older `diffusers-cli`
+        # (the SGLang image does), so the server is started with the interpreter the install went into.
+        python = sbx.run(["uv", "python", "find", "--system"]).stdout.strip()
 
         task_kwargs = _build_task_kwargs(args)
         for key in (*_LOCAL_KEYS, *_SANDBOX_SET_KEYS):
@@ -143,7 +151,8 @@ def _stream(chunk: str) -> None:
         task_kwargs["port"] = SANDBOX_PORT
         # Clients reach the server through the proxy, so output links must be built on the proxied address.
         task_kwargs["public_url"] = sbx.proxy_url_for(SANDBOX_PORT)
-        cli_argv = ["diffusers-cli", "--format", "quiet", *_kwargs_to_argv(task, task_kwargs)]
+        cli_module = [python, "-m", "diffusers.commands.diffusers_cli"]
+        cli_argv = [*cli_module, "--format", "quiet", *_kwargs_to_argv(task, task_kwargs)]
 
         threading.Thread(target=_announce_when_ready, args=(sbx, args, stopped), daemon=True).start()
         logger.info(f"starting the server in sandbox {sbx.id}; stop it with Ctrl-C.")

tests/others/test_cli_commands.py

@@ -34,6 +34,8 @@
     RunCommand,
     _build_task_kwargs,
     _collapse_frame_dirs,
+    _compile_denoiser,
+    _compile_vae_decoder,
     _download_outputs_from_sandbox,
     _kwargs_to_argv,
     _load_lora,
@@ -241,6 +243,41 @@ def test_compile_arg(self):
         assert compiled_blocks
         assert all(m._compiled_call_impl is not None for m in compiled_blocks)
 
+    def test_compile_mode_full_compiles_the_whole_denoiser(self, monkeypatch):
+        class Denoiser(torch.nn.Module):
+            _repeated_blocks = ["Block"]
+
+            def compile_repeated_blocks(self, **kwargs):
+                raise AssertionError("regional compilation must not run with compile_mode='full'")
+
+        compiled = {}
+        monkeypatch.setattr(torch, "compile", lambda module, **kwargs: compiled.update(kwargs) or "compiled")
+        pipeline = SimpleNamespace(transformer=Denoiser())
+        _compile_denoiser(pipeline, '{"mode": "max-autotune"}', "full")
+        assert pipeline.transformer == "compiled"
+        assert compiled == {"mode": "max-autotune"}
+
+    def test_regional_compile_rejects_cuda_graph_modes(self):
+        class Denoiser(torch.nn.Module):
+            _repeated_blocks = ["Block"]
+
+        pipeline = SimpleNamespace(transformer=Denoiser())
+        with pytest.raises(SystemExit, match="uses CUDA graphs.*--compile-mode full"):
+            _compile_denoiser(pipeline, '{"mode": "max-autotune"}', "regional")
+
+    def test_compile_vae_decoder_compiles_the_decoder(self, monkeypatch):
+        compiled = {}
+        monkeypatch.setattr(torch, "compile", lambda module, **kwargs: compiled.update(kwargs) or "compiled")
+        pipeline = SimpleNamespace(vae=SimpleNamespace(decoder=torch.nn.Conv2d(4, 3, 3)))
+        _compile_vae_decoder(pipeline, '{"fullgraph": true}')
+        assert pipeline.vae.decoder == "compiled"
+        assert compiled == {"fullgraph": True}
+
+    def test_compile_vae_decoder_skips_a_pipeline_without_one(self):
+        pipeline = SimpleNamespace(vae=None)
+        _compile_vae_decoder(pipeline, "{}")
+        assert pipeline.vae is None
+
     @require_torch_gpu
     # `--attention-backend` only exposes Hub-hosted kernels, all of which need `kernels>=0.12`.
     @require_kernels_version_greater_or_equal("0.12")