From 260ae4a1c2116bb2df0fef04ea8f1e254c99cbd9 Mon Sep 17 00:00:00 2001 From: DN6 Date: Fri, 9 Oct 2026 11:02:21 +0530 Subject: [PATCH 1/4] add diffusers serve and generate commands --- .ai/skills/diffusers-cli/SKILL.md | 4 + docs/source/en/using-diffusers/cli.md | 144 ++++ setup.py | 3 + src/diffusers/commands/diffusers_cli.py | 4 + src/diffusers/commands/generate.py | 238 ++++++ src/diffusers/commands/serve/__init__.py | 281 +++++++ src/diffusers/commands/serve/app.py | 218 ++++++ src/diffusers/commands/serve/backends.py | 472 +++++++++++ src/diffusers/commands/serve/generations.py | 194 +++++ src/diffusers/commands/serve/manifest.py | 779 ++++++++++++++++++ src/diffusers/dependency_versions_table.py | 2 + tests/others/test_cli_commands.py | 2 +- tests/others/test_cli_serve.py | 826 ++++++++++++++++++++ 13 files changed, 3166 insertions(+), 1 deletion(-) create mode 100644 src/diffusers/commands/generate.py create mode 100644 src/diffusers/commands/serve/__init__.py create mode 100644 src/diffusers/commands/serve/app.py create mode 100644 src/diffusers/commands/serve/backends.py create mode 100644 src/diffusers/commands/serve/generations.py create mode 100644 src/diffusers/commands/serve/manifest.py create mode 100644 tests/others/test_cli_serve.py diff --git a/.ai/skills/diffusers-cli/SKILL.md b/.ai/skills/diffusers-cli/SKILL.md index dd9ffd66dabd..9baf0540855a 100644 --- a/.ai/skills/diffusers-cli/SKILL.md +++ b/.ai/skills/diffusers-cli/SKILL.md @@ -16,6 +16,8 @@ description: > | --------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `run` | Run any `DiffusionPipeline` or `ModularPipeline`. Forwards `--pipeline-kwargs` verbatim, saves output by detecting its runtime type, optionally runs on HF Jobs via `--remote`. | | `schema` | Print the input schema for a pipeline repo (kwarg names, types, defaults, descriptions). **No weights downloaded** — only the small index file. | +| `serve` | Serve a pipeline over HTTP with the Generative Media Spec API: model discovery, queued generations, polling or SSE progress. In-process, or behind SGLang / vLLM-Omni via `--backend`. | +| `generate` | Send a request to a Generative Media Spec server (such as `serve`) and save the outputs. Loads no pipeline; takes `--url`, `--inputs` and `--parameters` with the server's declared names. | | `custom_blocks` | Package a local `ModularPipelineBlocks` subclass for the Hub. | | `env` | Print versions of diffusers + torch + transformers + accelerate + safetensors + CUDA + GPU info. Use when investigating environment issues, dtype/precision support, or building bug reports. | @@ -33,6 +35,8 @@ The other commands are small enough that `diffusers-cli --help` is the ```bash diffusers-cli schema --help +diffusers-cli serve --help +diffusers-cli generate --help diffusers-cli custom_blocks --help diffusers-cli env --help ``` diff --git a/docs/source/en/using-diffusers/cli.md b/docs/source/en/using-diffusers/cli.md index 02be2bd2fd0e..d7e8d6523b38 100644 --- a/docs/source/en/using-diffusers/cli.md +++ b/docs/source/en/using-diffusers/cli.md @@ -21,6 +21,8 @@ specific language governing permissions and limitations under the License. | [`env`](#env) | Print environment info for bug reports. | | [`schema`](#schema) | Inspect a pipeline's `__call__` signature without downloading weights. | | [`run`](#run) | Run a pipeline locally or in a Hugging Face Sandbox. | +| [`serve`](#serve) | Serve a pipeline over HTTP with the Generative Media Spec API. | +| [`generate`](#generate) | Send a generation to a Generative Media Spec server and save its outputs. | | [`custom_blocks`](#customblocks) | Package a local `ModularPipelineBlocks` subclass for the Hub. | | [`fp16_safetensors`](#fp16safetensors) | Convert a checkpoint to fp16 `.safetensors`. | | [`skills`](#skills) | Install pre-authored skill bundles into your AI coding agent. | @@ -287,6 +289,148 @@ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \ hf sandbox kill ``` +## `serve` + +Serve a pipeline over HTTP. The server implements the Generative Media Spec (GMS, draft `gms/0.1`): clients +discover what the model accepts, create a generation, follow it while it runs, and read the results. It needs +the `serve` extra: + +```bash +pip install "diffusers[serve]" +diffusers-cli serve --model black-forest-labs/FLUX.1-dev --dtype bf16 +``` + +`serve` takes the same loading and optimization flags as [`run`](#run) (`--dtype`, `--device-map`, +`--cpu-offload`, `--compile`, `--lora`, ...), plus `--host` (default `127.0.0.1`) and `--port` (default `8000`). +The server has no authentication, so keep it on a trusted network. + +[Modular repos](../modular_diffusers/overview) are detected the same way as in `run` and served in-process. They +report no step progress, and a cancel request only takes effect once the pipeline call returns. + +| Flag | Effect | +|---|---| +| `--default-seed` | Seed used for requests that do not set one. Without it, each such request gets a random seed. | +| `--enable-cors` | Allow cross-origin requests from any origin, so a browser app on another host can call the server. | +| `--log-level` | Logging level for diffusers and the HTTP server (`debug`, `info`, `warning`, `error`, `critical`). | + +`GET /health` returns `{"status": "ok"}`, or `503` when an SGLang or vLLM-Omni backend has exited. Every response +carries an `x-request-id` header: the one sent with the request, or a generated one. + +Generations run one at a time. Creating one returns immediately with a `queued` generation; poll it, subscribe +to its events, or ask the server to hold the connection with `Prefer: wait=`: + +```bash +# What the model accepts: tasks, inputs, parameters with types and defaults. +curl localhost:8000/v1/models/black-forest-labs/FLUX.1-dev + +# Create a generation and wait up to 120 seconds for it. +curl localhost:8000/v1/generations -H 'Prefer: wait=120' -d '{ + "model": "black-forest-labs/FLUX.1-dev", + "inputs": {"prompt": "an astronaut riding a horse"}, + "parameters": {"width": 1024, "height": 1024, "steps": 28, "seed": 42} +}' + +# Follow step progress as server-sent events, or poll the generation. +curl -N localhost:8000/v1/generations//events +curl localhost:8000/v1/generations/ + +# Cancel a queued or running generation, or delete it together with its files. +curl -X POST localhost:8000/v1/generations//cancel +curl -X DELETE localhost:8000/v1/generations/ +``` + +A completed generation lists its outputs with a download URL each (or inline base64 with +`"response_format": "b64_json"`), the seed that was used, and the model revision. Outputs are kept under +`--output-dir` and expire after one hour. Images are PNG by default (`"output_format": "image/jpeg"` or +`"image/webp"` to change it), videos MP4, audio WAV. Media inputs are passed as +`{"type": "image", "url": "https://..."}` or `{"type": "image", "base64": "...", "media_type": "image/png"}`. + +Validation is strict: an unknown parameter, a value outside its declared range, or a missing input is rejected +with a JSON pointer to the offending field. + +```json +{"error": {"code": "invalid_request", "message": "task 'text_to_image' declares no parameter 'cfg'", "pointer": "/parameters/cfg"}} +``` + +### Manifests + +What a model accepts is described by its manifest. `serve` looks for one in this order: + +1. the file passed with `--manifest `; +2. a `gms.json` published in the model repo, served byte-for-byte; +3. a manifest derived from the pipeline's `__call__` signature. + +A derived manifest is a best-effort reading of the signature: scalar arguments become parameters, prompts and +media arguments become inputs, and arguments with a core GMS name are renamed (`num_inference_steps` becomes +`steps`, `num_images_per_prompt` becomes `num_outputs`, `generator` becomes `seed`). It carries no value ranges. +Write a `gms.json` when you need ranges, presets, or a duration in seconds. + +When a manifest uses a name that differs from the pipeline argument, map it per task under `x-diffusers`: + +```json +"x-diffusers": {"inputs": {"caption": "prompt"}, "parameters": {"cfg": "guidance_scale"}} +``` + +A task is only hosted when every input it requires maps to a pipeline argument. A parameter the pipeline cannot +set is rejected with `not_implemented` unless the request sends its default value. + +### Adapters + +A request can apply LoRAs for that generation only. They are loaded before the call and unloaded after it: + +```bash +curl localhost:8000/v1/generations -d '{ + "model": "black-forest-labs/FLUX.1-dev", + "inputs": {"prompt": "a tiny cat"}, + "adapters": [{"type": "lora", "path": "hf:alvdansen/littletinies", "scale": 0.8}] +}' +``` + +Request adapters are turned off when the server starts with `--compile` or `--lora`; such requests fail with +`not_implemented`. + +### SGLang and vLLM-Omni backends + +`--backend sglang` and `--backend vllm` run the model in [SGLang Diffusion](https://docs.sglang.io/) or +[vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni) instead of in-process. `serve` starts `sglang serve` or +`vllm serve --omni` as a child process on a local port, waits for it to become healthy, and translates GMS +requests to the engine's OpenAI-compatible image and video routes. Clients only ever talk to the GMS endpoints +of `serve`. The engine must be installed separately. + +```bash +diffusers-cli serve --model Qwen/Qwen-Image --backend sglang --backend-args '--num-gpus 2' +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 +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. + +## `generate` + +Send a generation to a server started with [`serve`](#serve), or to any other Generative Media Spec server, and +save its outputs. Unlike [`run`](#run), `generate` loads no pipeline: it reads the model's discovery document, +submits the request, polls the generation until it finishes, and downloads the files. + +```bash +diffusers-cli generate --url http://localhost:8000 --inputs '{"prompt": "a cat on the moon"}' + +diffusers-cli generate --url http://localhost:8000 \ + --inputs '{"prompt": "make the fur grey", "image": "cat.png"}' \ + --parameters '{"steps": 8, "seed": 0}' --output outputs/ +``` + +`--inputs` and `--parameters` use the names the server declares, not pipeline argument names (`steps`, not +`num_inference_steps`). An image, mask, video or audio input can be an http(s) URL or a local file path; local +files are sent inline as base64. `--task`, `--preset`, `--adapters` and `--output-format` map to the request +fields of the same name. `--model` is only needed when the server hosts more than one model. + +Outputs are saved to `~/.diffusers/cli/generate/outputs//` unless `--output` names a directory. +`--token` sends a bearer token in the `Authorization` header, for servers behind an authenticating proxy; no +token is sent by default. Interrupting the command cancels the generation on the server. + ## `custom_blocks` Package a local `ModularPipelineBlocks` subclass for upload to the Hub. Reads a Python file, AST-scans it for diff --git a/setup.py b/setup.py index 34981886d0f8..94620fb3c722 100644 --- a/setup.py +++ b/setup.py @@ -146,6 +146,8 @@ "opencv-python", "timm", "flashpack", + "fastapi", + "uvicorn", ] # this is a lookup table with items like: @@ -253,6 +255,7 @@ def run(self): extras["nvidia_modelopt"] = deps_list("nvidia_modelopt[hf]") extras["sdnq"] = deps_list("sdnq") extras["flashpack"] = deps_list("flashpack") +extras["serve"] = deps_list("fastapi", "uvicorn") extras["dev"] = extras["quality"] + extras["test"] + extras["training"] + extras["docs"] + extras["torch"] diff --git a/src/diffusers/commands/diffusers_cli.py b/src/diffusers/commands/diffusers_cli.py index 0e4f2c27fb64..1df91bd7ddb4 100644 --- a/src/diffusers/commands/diffusers_cli.py +++ b/src/diffusers/commands/diffusers_cli.py @@ -20,8 +20,10 @@ from .custom_blocks import CustomBlocksCommand from .env import EnvironmentCommand from .fp16_safetensors import FP16SafetensorsCommand +from .generate import GenerateCommand from .run import RunCommand from .schema import SchemaCommand +from .serve import ServeCommand from .skills import SkillsCommand @@ -47,8 +49,10 @@ def main(): EnvironmentCommand.register_subcommand(commands_parser) FP16SafetensorsCommand.register_subcommand(commands_parser) CustomBlocksCommand.register_subcommand(commands_parser) + GenerateCommand.register_subcommand(commands_parser) RunCommand.register_subcommand(commands_parser) SchemaCommand.register_subcommand(commands_parser) + ServeCommand.register_subcommand(commands_parser) SkillsCommand.register_subcommand(commands_parser) # Let's go diff --git a/src/diffusers/commands/generate.py b/src/diffusers/commands/generate.py new file mode 100644 index 000000000000..f3ea85d54e1d --- /dev/null +++ b/src/diffusers/commands/generate.py @@ -0,0 +1,238 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""`diffusers-cli generate` — send a generation to a Generative Media Spec (GMS) server and save its outputs. + +The client side of `diffusers-cli serve`: it loads no pipeline. It reads the model's discovery document, submits the +request, polls the generation until it finishes, and downloads the files. +""" + +from __future__ import annotations + +import base64 +import json +import mimetypes +import time +from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter, _SubParsersAction +from pathlib import Path +from typing import Any + +from huggingface_hub.cli._output import out +from huggingface_hub.utils import httpx + +from . import BaseDiffusersCLICommand + + +DEFAULT_OUTPUT_DIR = str(Path.home() / ".diffusers" / "cli" / "generate" / "outputs") +_MEDIA_TYPES = ("image", "mask", "video", "audio") +_POLL_SECONDS = 1.0 +_TIMEOUT_SECONDS = 60.0 + + +def _parse_json(raw: str | None, flag: str, expected: type) -> Any: + if raw is None: + return None + try: + parsed = json.loads(raw) + except json.JSONDecodeError as e: + raise SystemExit(f"{flag} must be valid JSON: {e}") from e + if not isinstance(parsed, expected): + raise SystemExit(f"{flag} must decode to a JSON {'object' if expected is dict else 'array'}.") + return parsed + + +def _media_value(name: str, media_type: str, value: Any) -> Any: + """Turn a URL or local path into a GMS media value; anything else is sent as given.""" + if isinstance(value, list): + return [_media_value(name, media_type, item) for item in value] + if not isinstance(value, str): + return value + if value.startswith(("http://", "https://")): + return {"type": media_type, "url": value} + + path = Path(value).expanduser() + if not path.is_file(): + raise SystemExit(f"Input {name!r} is neither an http(s) URL nor an existing file: {value!r}") + media = {"type": media_type, "base64": base64.b64encode(path.read_bytes()).decode()} + guessed = mimetypes.guess_type(path.name)[0] + if guessed is not None: + media["media_type"] = guessed + return media + + +class GenerateCommand(BaseDiffusersCLICommand): + task = "generate" + + @staticmethod + def register_subcommand(subparsers: _SubParsersAction) -> None: + epilog = ( + "Examples\n" + ' $ diffusers-cli generate --url http://localhost:8000 --inputs \'{"prompt": "a cat on the moon"}\'\n' + " $ diffusers-cli generate --url http://localhost:8000 \\\n" + ' --inputs \'{"prompt": "make the fur grey", "image": "cat.png"}\' \\\n' + ' --parameters \'{"steps": 8, "seed": 0}\' -o outputs/\n' + "\n" + "Learn more\n" + " Use `diffusers-cli --help` for more information about a command.\n" + " Read the documentation at https://huggingface.co/docs/diffusers\n" + ) + + parser: ArgumentParser = subparsers.add_parser( + "generate", + help="Send a generation to a Generative Media Spec server and save its outputs.", + usage="\n diffusers-cli generate --url --inputs [options]", + epilog=epilog, + formatter_class=RawDescriptionHelpFormatter, + ) + parser._optionals.title = "Options" + parser.add_argument("--url", required=True, help="Base URL of the server, e.g. http://localhost:8000.") + parser.add_argument( + "--model", "-m", default=None, help="Model id to request. Defaults to the server's only model." + ) + parser.add_argument( + "--inputs", + required=True, + metavar="JSON", + help=( + "JSON object of the task's inputs. An image, mask, video or audio input given as an http(s) URL " + "or a local file path is converted to a media value; local files are sent inline as base64." + ), + ) + parser.add_argument("--parameters", default=None, metavar="JSON", help="JSON object of task parameters.") + parser.add_argument( + "--task", default=None, help="Task to run. Defaults to the one the server infers from the inputs." + ) + parser.add_argument("--preset", default=None, help="Named preset of the task to start from.") + parser.add_argument( + "--adapters", + default=None, + metavar="JSON", + help='JSON array of adapters to apply, e.g. \'[{"type": "lora", "path": "hf:org/name", "scale": 0.8}]\'.', + ) + parser.add_argument("--output-format", default=None, help="Media type to encode outputs as, e.g. image/jpeg.") + parser.add_argument( + "--output", + "-o", + default=None, + help=f"Directory to save outputs in. Defaults to {DEFAULT_OUTPUT_DIR}//.", + ) + parser.add_argument( + "--token", + default=None, + help="Bearer token sent to the server in the Authorization header. No token is sent by default.", + ) + parser.set_defaults(func=GenerateCommand) + + def __init__(self, args: Namespace): + self.args = args + + def run(self) -> None: + args = self.args + inputs = _parse_json(args.inputs, "--inputs", dict) + parameters = _parse_json(args.parameters, "--parameters", dict) + adapters = _parse_json(args.adapters, "--adapters", list) + + with self._client() as client: + model = args.model or self._only_model(client) + document = self._request(client, "GET", f"/v1/models/{model}") + tasks = document["manifest"]["tasks"] + if args.task is not None and args.task not in tasks: + raise SystemExit(f"Model {model!r} declares no task {args.task!r}; its tasks are {list(tasks)}.") + + declared = [tasks[args.task]] if args.task is not None else list(tasks.values()) + input_types = {name: spec.get("type") for task in declared for name, spec in task["inputs"].items()} + body: dict[str, Any] = {"model": model, "inputs": {}} + for name, value in inputs.items(): + if input_types.get(name) in _MEDIA_TYPES: + value = _media_value(name, input_types[name], value) + body["inputs"][name] = value + optional = { + "task": args.task, + "preset": args.preset, + "parameters": parameters, + "adapters": adapters, + "output_format": args.output_format, + } + body.update({key: value for key, value in optional.items() if value is not None}) + + generation = self._request(client, "POST", "/v1/generations", json=body) + try: + while generation["status"] in ("queued", "generating"): + time.sleep(_POLL_SECONDS) + generation = self._request(client, "GET", f"/v1/generations/{generation['id']}") + except KeyboardInterrupt: + # Without this the server keeps generating for a client that has gone away. + self._request(client, "POST", f"/v1/generations/{generation['id']}/cancel") + raise + if generation["status"] != "complete": + raise SystemExit(f"Generation {generation['id']} ended with status {generation['status']!r}.") + + directory = Path(args.output or Path(DEFAULT_OUTPUT_DIR) / generation["id"]).expanduser() + directory.mkdir(parents=True, exist_ok=True) + saved = [] + for name, artifacts in generation["outputs"].items(): + for index, artifact in enumerate(artifacts): + path = directory / f"{name}-{index}.{artifact['media_type'].split('/')[1]}" + path.write_bytes(self._artifact_bytes(client, artifact)) + saved.append(str(path)) + + out.result( + self.task, + url=args.url, + model=model, + generation=generation["id"], + seed=generation.get("seed"), + outputs=saved, + ) + + def _client(self) -> httpx.Client: + headers = {"Authorization": f"Bearer {self.args.token}"} if self.args.token else {} + return httpx.Client(base_url=self.args.url, headers=headers, timeout=_TIMEOUT_SECONDS, follow_redirects=True) + + def _request(self, client: httpx.Client, method: str, path: str, **kwargs: Any) -> Any: + try: + response = client.request(method, path, **kwargs) + except httpx.HTTPError as e: + raise SystemExit(f"Could not reach {self.args.url}: {type(e).__name__}: {e}") from e + if response.status_code < 400: + return response.json() + + try: + error = response.json()["error"] + message = f"{error['code']}: {error['message']}" + if error.get("pointer"): + message += f" (at {error['pointer']})" + except (ValueError, KeyError, TypeError): + message = response.text + raise SystemExit(f"{method} {path} failed with HTTP {response.status_code}: {message}") + + def _only_model(self, client: httpx.Client) -> str: + models = [model["id"] for model in self._request(client, "GET", "/v1/models")["data"]] + if len(models) != 1: + raise SystemExit(f"{self.args.url} serves {len(models)} models ({models}); pick one with --model.") + return models[0] + + def _artifact_bytes(self, client: httpx.Client, artifact: dict[str, Any]) -> bytes: + if "base64" in artifact: + return base64.b64decode(artifact["base64"]) + try: + # The token is for the server; an output hosted elsewhere is fetched without it. + if httpx.URL(artifact["url"]).host == client.base_url.host: + response = client.get(artifact["url"]) + else: + response = httpx.get(artifact["url"], timeout=_TIMEOUT_SECONDS, follow_redirects=True) + response.raise_for_status() + except httpx.HTTPError as e: + raise SystemExit(f"Could not download output {artifact['url']}: {type(e).__name__}: {e}") from e + return response.content diff --git a/src/diffusers/commands/serve/__init__.py b/src/diffusers/commands/serve/__init__.py new file mode 100644 index 000000000000..e56c73a5c530 --- /dev/null +++ b/src/diffusers/commands/serve/__init__.py @@ -0,0 +1,281 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""`diffusers-cli serve` — serve a pipeline over HTTP with the Generative Media Spec (GMS) API. + +Generations run one at a time on a background worker: a request returns immediately with a queued generation that +clients poll or follow over server-sent events. The pipeline runs in-process by default, or behind an SGLang or +vLLM-Omni server with `--backend`. +""" + +from __future__ import annotations + +import inspect +from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter, _SubParsersAction +from pathlib import Path + +from huggingface_hub.cli._output import out + +from ...utils import logging +from .. import BaseDiffusersCLICommand +from ..run import _add_loading_arguments, _add_optimization_arguments + + +DEFAULT_OUTPUT_DIR = str(Path.home() / ".diffusers" / "cli" / "serve" / "outputs") +BACKEND_CHOICES = ("diffusers", "sglang", "vllm") + +# Flags that configure the in-process pipeline, paired with the value argparse gives them when unset. +_DIFFUSERS_ONLY_FLAGS = { + "device_map": None, + "dtype": "auto", + "variant": None, + "lora": None, + "cpu_offload": None, + "attention_backend": "default", + "vae_tiling": False, + "vae_slicing": False, + "compile": None, +} + + +def _resolve_pipeline(args: Namespace) -> tuple[type, inspect.Signature, str | None]: + """Return the pipeline class, the signature of its call, and the commit hash of a standard Hub repo.""" + import diffusers + + from .manifest import blocks_signature + + # A modular repo ships `modular_model_index.json` instead of `model_index.json`, so `load_config` raises. + try: + config, commit_hash = diffusers.DiffusionPipeline.load_config( + args.model, token=args.token, revision=args.revision, return_commit_hash=True + ) + except OSError: + config, commit_hash = None, None + + if config is not None: + class_name = config.get("_class_name") + pipeline_cls = getattr(diffusers, str(class_name), None) + if not isinstance(pipeline_cls, type): + raise SystemExit( + f"Pipeline class {class_name!r} declared in {diffusers.DiffusionPipeline.config_name} is not " + "exported by the installed diffusers." + ) + if not issubclass(pipeline_cls, diffusers.ModularPipeline): + return pipeline_cls, inspect.signature(pipeline_cls.__call__), commit_hash + + if args.backend != "diffusers": + raise SystemExit( + f"{args.model!r} is a modular pipeline, which --backend {args.backend} cannot serve. " + "Use --backend diffusers." + ) + try: + pipeline = diffusers.ModularPipeline.from_pretrained( + args.model, trust_remote_code=args.trust_remote_code, token=args.token, revision=args.revision + ) + except OSError as e: + raise SystemExit(f"Could not read a pipeline config for {args.model!r}: {e}") from e + return type(pipeline), blocks_signature(pipeline.blocks), commit_hash + + +class ServeCommand(BaseDiffusersCLICommand): + task = "serve" + + @staticmethod + def register_subcommand(subparsers: _SubParsersAction) -> None: + epilog = ( + "Examples\n" + " $ diffusers-cli serve -m black-forest-labs/FLUX.1-dev --dtype bf16\n" + " $ diffusers-cli serve -m black-forest-labs/FLUX.1-dev --dtype bf16 --compile --port 9000\n" + " $ diffusers-cli serve -m Wan-AI/Wan2.2-T2V-A14B-Diffusers --backend vllm\n" + " $ diffusers-cli serve -m Qwen/Qwen-Image --backend sglang --backend-args '--num-gpus 2'\n" + "\n" + " $ curl localhost:8000/v1/generations -H 'Prefer: wait=120' \\\n" + ' -d \'{"model": "black-forest-labs/FLUX.1-dev", "inputs": {"prompt": "a cat on the moon"}}\'\n' + "\n" + "Learn more\n" + " Use `diffusers-cli --help` for more information about a command.\n" + " Read the documentation at https://huggingface.co/docs/diffusers\n" + ) + + parser: ArgumentParser = subparsers.add_parser( + "serve", + help="Serve a diffusers pipeline over HTTP with the Generative Media Spec API.", + usage="\n diffusers-cli serve [options]", + epilog=epilog, + formatter_class=RawDescriptionHelpFormatter, + ) + parser._optionals.title = "Options" + _add_loading_arguments(parser) + _add_optimization_arguments(parser) + parser.add_argument( + "--backend", + choices=BACKEND_CHOICES, + default="diffusers", + help=( + "What runs the model. 'diffusers' loads the pipeline in this process. 'sglang' and 'vllm' start " + "`sglang serve` / `vllm serve --omni` as a child process and translate requests to its " + "OpenAI-compatible image and video routes." + ), + ) + parser.add_argument( + "--backend-args", + default=None, + metavar="ARGS", + help=( + "Extra arguments appended to the `sglang serve` / `vllm serve` command, as one quoted string " + "(e.g. '--num-gpus 2'). Only used with --backend sglang or vllm." + ), + ) + parser.add_argument("--host", default="127.0.0.1", help="Interface to bind. Defaults to 127.0.0.1.") + parser.add_argument("--port", type=int, default=8000, help="Port to bind. Defaults to 8000.") + parser.add_argument( + "--manifest", + default=None, + metavar="PATH", + help=( + "Path to a GMS manifest (`gms.json`) describing the model. Defaults to the `gms.json` published " + "in the model repo; when there is none, a manifest is derived from the pipeline's call signature." + ), + ) + parser.add_argument( + "--output-dir", + default=DEFAULT_OUTPUT_DIR, + help=f"Directory generated files are stored in until they expire. Defaults to {DEFAULT_OUTPUT_DIR}.", + ) + parser.add_argument( + "--fps", + type=int, + default=8, + help="FPS used to encode video outputs when the pipeline call takes no fps argument.", + ) + parser.add_argument( + "--sampling-rate", + type=int, + default=None, + help="Sample rate used to encode audio outputs. Defaults to 16000.", + ) + parser.add_argument( + "--default-seed", + type=int, + default=None, + help="Seed used for requests that do not set one. Defaults to a random seed per request.", + ) + parser.add_argument( + "--enable-cors", + action="store_true", + help="Allow cross-origin requests from any origin, so browser apps on other hosts can call the server.", + ) + parser.add_argument( + "--log-level", + choices=tuple(logging.get_log_levels_dict()), + default=None, + help="Logging level for diffusers and the HTTP server. Defaults to the current diffusers verbosity.", + ) + # `_load_pipeline` reads these `run` flags, which `serve` does not expose. + parser.set_defaults(func=ServeCommand, workflow=None, offload_margin=None) + + def __init__(self, args: Namespace): + self.args = args + + def run(self) -> None: + args = self.args + try: + import fastapi # noqa: F401 + import uvicorn + except ImportError as e: + raise SystemExit( + '`diffusers-cli serve` requires FastAPI and uvicorn. Install them with: pip install "diffusers[serve]"' + ) from e + + if args.log_level is not None: + logging.set_verbosity(logging.get_log_levels_dict()[args.log_level]) + + if args.context_parallel: + raise SystemExit( + "--context-parallel is not supported by `diffusers-cli serve`: under torchrun every rank would " + "start its own server on the same port." + ) + 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 = [ + "--" + 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 " + f"--backend {args.backend}. Pass the engine's own flags through --backend-args instead." + ) + + from .app import create_app + from .backends import ( + DiffusersBackend, + EngineBackend, + backend_capabilities, + backend_kwargs, + backend_output_types, + ) + from .generations import GenerationQueue + from .manifest import ServedModel, load_manifest + + pipeline_cls, signature, commit_hash = _resolve_pipeline(args) + kwarg_names = backend_kwargs(args.backend, signature) + manifest = load_manifest( + args.model, + pipeline_cls, + signature, + kwarg_names, + manifest_path=args.manifest, + revision=args.revision, + token=args.token, + ) + served = ServedModel( + manifest, + kwarg_names, + backend_output_types(args.backend), + backend_capabilities(args, pipeline_cls), + revision=commit_hash, + ) + if not served.tasks: + raise SystemExit( + f"None of the tasks {list(manifest.data['tasks'])} in the manifest for {args.model!r} can be " + f"served by {pipeline_cls.__name__} with --backend {args.backend}." + ) + + if args.backend == "diffusers": + backend = DiffusersBackend(args, served) + else: + backend = EngineBackend.launch(args, served) + + try: + app = create_app( + served, + GenerationQueue(backend, args.output_dir), + default_seed=args.default_seed, + enable_cors=args.enable_cors, + ) + out.result( + self.task, + url=f"http://{args.host}:{args.port}", + model=served.id, + backend=args.backend, + pipeline_class=pipeline_cls.__name__, + tasks=list(served.tasks), + capabilities=served.capabilities, + ) + uvicorn.run(app, host=args.host, port=args.port, log_level=args.log_level or "info") + finally: + backend.close() diff --git a/src/diffusers/commands/serve/app.py b/src/diffusers/commands/serve/app.py new file mode 100644 index 000000000000..05213fc302ce --- /dev/null +++ b/src/diffusers/commands/serve/app.py @@ -0,0 +1,218 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The Generative Media Spec HTTP surface, as a FastAPI app over one served model and one generation queue.""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import re +import uuid +from typing import Any + +from fastapi import FastAPI, Request +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import FileResponse, JSONResponse, Response, StreamingResponse +from starlette.exceptions import HTTPException + +from ...utils import logging +from .generations import Generation, GenerationQueue +from .manifest import GMSError, ServedModel + + +logger = logging.get_logger("diffusers-cli/serve") + +MAX_WAIT_SECONDS = 300.0 +REQUEST_ID_HEADER = "x-request-id" +_KEEP_ALIVE_SECONDS = 15.0 + + +def _timestamp(value: Any) -> str | None: + if value is None: + return None + return value.strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _render(generation: Generation, request: Request) -> dict[str, Any]: + outputs = None + if generation.status == "complete": + outputs = {} + for name, artifacts in generation.outputs.items(): + outputs[name] = [] + for artifact in artifacts: + rendered = {key: value for key, value in artifact.items() if key != "file"} + if generation.request.response_format == "b64_json": + content = (generation.directory / artifact["file"]).read_bytes() + rendered["base64"] = base64.b64encode(content).decode() + else: + url = request.url_for( + "get_generation_file", generation_id=generation.id, filename=artifact["file"] + ) + rendered["url"] = str(url) + outputs[name].append(rendered) + + return { + "id": generation.id, + "model": generation.model, + "revision": generation.revision, + "status": generation.status, + "created_at": _timestamp(generation.created_at), + "progress": generation.progress if generation.status == "generating" else None, + "outputs": outputs, + "expires_at": _timestamp(generation.expires_at), + "seed": generation.seed if generation.status == "complete" else None, + "usage": generation.usage, + } + + +def create_app( + served: ServedModel, + generations: GenerationQueue, + default_seed: int | None = None, + enable_cors: bool = False, +) -> FastAPI: + app = FastAPI(title="diffusers-cli serve", docs_url=None, redoc_url=None, openapi_url=None) + + @app.middleware("http") + async def request_id(request: Request, call_next) -> Response: + request.state.request_id = request.headers.get(REQUEST_ID_HEADER) or str(uuid.uuid4()) + response = await call_next(request) + response.headers[REQUEST_ID_HEADER] = request.state.request_id + return response + + if enable_cors: + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_methods=["*"], + allow_headers=["*"], + expose_headers=[REQUEST_ID_HEADER], + ) + logger.warning("--enable-cors allows requests from any origin. Not recommended for production.") + + @app.exception_handler(GMSError) + async def gms_error(request: Request, error: GMSError) -> JSONResponse: + return JSONResponse(error.to_dict(), status_code=error.status) + + @app.exception_handler(HTTPException) + async def http_error(request: Request, error: HTTPException) -> JSONResponse: + if error.status_code == 404: + code = "not_found" + elif error.status_code < 500: + code = "invalid_request" + else: + code = "internal" + return JSONResponse({"error": {"code": code, "message": str(error.detail)}}, status_code=error.status_code) + + @app.exception_handler(Exception) + async def internal_error(request: Request, error: Exception) -> JSONResponse: + message = f"{type(error).__name__}: {error}" + # An unhandled error is answered outside the middleware stack, so the request id is attached here. + headers = {REQUEST_ID_HEADER: request.state.request_id} + return JSONResponse({"error": {"code": "internal", "message": message}}, status_code=500, headers=headers) + + @app.get("/health") + async def health() -> JSONResponse: + failure = generations.failure() + if failure is not None: + return JSONResponse({"status": "unhealthy", "reason": failure}, status_code=503) + return JSONResponse({"status": "ok"}) + + @app.get("/v1/models") + async def list_models() -> dict[str, Any]: + return {"data": [served.summary()]} + + @app.get("/v1/models/{model_id:path}") + async def get_model(model_id: str) -> Response: + if model_id != served.id: + raise GMSError("not_found", f"model {model_id!r} is not served here; this server hosts {served.id!r}") + return Response(served.discovery_document(), media_type="application/json") + + @app.post("/v1/generations", status_code=201) + async def create_generation(request: Request) -> dict[str, Any]: + try: + body = await request.json() + except json.JSONDecodeError as e: + raise GMSError("invalid_request", f"request body is not valid JSON: {e}") from e + resolved = served.resolve(body) + seeded = served.seeded(resolved.task) + generation = Generation(served.id, served.revision, resolved, seeded=seeded, default_seed=default_seed) + generations.submit(generation) + + wait = re.fullmatch(r"\s*wait=(\d+)\s*", request.headers.get("prefer", "")) + if wait is not None: + await asyncio.to_thread(generation.wait, min(float(wait.group(1)), MAX_WAIT_SECONDS)) + if generation.error is not None: + raise generation.error + return _render(generation, request) + + @app.get("/v1/generations/{generation_id}") + async def get_generation(generation_id: str, request: Request) -> dict[str, Any]: + generation = generations.get(generation_id) + if generation.error is not None: + raise generation.error + return _render(generation, request) + + @app.delete("/v1/generations/{generation_id}", status_code=204) + async def delete_generation(generation_id: str) -> Response: + generations.delete(generation_id) + return Response(status_code=204) + + @app.post("/v1/generations/{generation_id}/cancel") + async def cancel_generation(generation_id: str, request: Request) -> dict[str, Any]: + generation = generations.cancel(generation_id) + if generation.error is not None: + raise generation.error + return _render(generation, request) + + @app.get("/v1/generations/{generation_id}/events") + async def get_generation_events(generation_id: str, request: Request) -> StreamingResponse: + generation = generations.get(generation_id) + last_event_id = request.headers.get("last-event-id", "") + start = int(last_event_id) if last_event_id.isdigit() and len(last_event_id) < 10 else 0 + + async def stream(): + index = start + while True: + events, done = await asyncio.to_thread(generation.wait_for_events, index, _KEEP_ALIVE_SECONDS) + for event, data in events: + index += 1 + payload = _render(generation, request) if event == "complete" else data + yield f"id: {index}\nevent: {event}\ndata: {json.dumps(payload)}\n\n" + if done: + return + if not events: + yield ": keep-alive\n\n" + + return StreamingResponse(stream(), media_type="text/event-stream", headers={"Cache-Control": "no-cache"}) + + @app.get("/v1/generations/{generation_id}/files/{filename}") + async def get_generation_file(generation_id: str, filename: str) -> FileResponse: + generation = generations.get(generation_id) + artifacts = [a for group in (generation.outputs or {}).values() for a in group if a["file"] == filename] + if not artifacts: + raise GMSError("not_found", f"generation {generation_id!r} has no file {filename!r}") + return FileResponse(generation.directory / filename, media_type=artifacts[0]["media_type"]) + + @app.get("/v1/generations/{generation_id}/outputs/{name}/stream") + async def stream_generation_output(generation_id: str, name: str) -> Response: + raise GMSError("not_implemented", "this server does not advertise the `stream` capability") + + @app.post("/v1/uploads") + async def create_upload() -> Response: + raise GMSError("not_implemented", "this server does not advertise the `presigned_upload` capability") + + return app diff --git a/src/diffusers/commands/serve/backends.py b/src/diffusers/commands/serve/backends.py new file mode 100644 index 000000000000..c0a779804dad --- /dev/null +++ b/src/diffusers/commands/serve/backends.py @@ -0,0 +1,472 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Backends that execute a resolved generation: an in-process pipeline, or an SGLang / vLLM-Omni server.""" + +from __future__ import annotations + +import base64 +import inspect +import io +import mimetypes +import os +import shlex +import shutil +import socket +import subprocess +import tempfile +import time +import wave +from argparse import Namespace +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +import numpy as np +import torch +from huggingface_hub.utils import httpx +from PIL import Image + +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 .generations import Generation +from .manifest import GMSError, ServedModel, blocks_signature + + +logger = logging.get_logger("diffusers-cli/serve") + +# Request fields both engines accept on their OpenAI-compatible image and video routes, named as the +# pipeline kwargs they correspond to. +_ENGINE_KWARGS = frozenset( + { + "prompt", + "negative_prompt", + "image", + "width", + "height", + "num_frames", + "fps", + "num_inference_steps", + "guidance_scale", + "guidance_scale_2", + "true_cfg_scale", + "generator", + "num_images_per_prompt", + } +) +_ENGINE_FIELDS = {"num_images_per_prompt": "n"} +_ENGINE_INSTALL = {"sglang": 'pip install "sglang[diffusion]"', "vllm": "pip install vllm-omni"} +_POLL_SECONDS = 1.0 + +_OUTPUT_FIELDS = {"image": ("images",), "video": ("frames", "videos"), "audio": ("audios", "audio")} + + +def backend_kwargs(backend: str, signature: inspect.Signature) -> set[str]: + names = set(signature.parameters) - {"self"} + if backend == "diffusers": + return names + return names & _ENGINE_KWARGS + + +def backend_output_types(backend: str) -> set[str]: + if backend == "diffusers": + return {"image", "video", "audio"} + return {"image", "video"} + + +def backend_capabilities(args: Namespace, pipeline_cls: type) -> list[str]: + capabilities = ["sse", "wait"] + if args.backend != "diffusers": + return capabilities + # Swapping adapter weights retraces a compiled denoiser on every request, and LoRAs passed at + # startup are part of the deployment rather than something a request may replace. + if hasattr(pipeline_cls, "load_lora_weights") and args.compile is None and not args.lora: + capabilities.append("adapters") + return capabilities + + +def _media_bytes(name: str, value: dict[str, Any], pointer: str) -> bytes: + try: + if "base64" in value: + return base64.b64decode(value["base64"], validate=True) + response = httpx.get(value["url"], follow_redirects=True, timeout=DIFFUSERS_REQUEST_TIMEOUT) + response.raise_for_status() + return response.content + except (ValueError, httpx.HTTPError) as e: + raise GMSError("invalid_request", f"could not read input {name!r}: {e}", pointer) from e + + +def _media_suffix(value: dict[str, Any], default: str) -> str: + if value.get("media_type"): + return mimetypes.guess_extension(value["media_type"]) or default + return Path(urlparse(value.get("url", "")).path).suffix or default + + +def _load_media(name: str, value: dict[str, Any], pointer: str) -> Any: + content = _media_bytes(name, value, pointer) + try: + if value["type"] in ("image", "mask"): + return load_image(Image.open(io.BytesIO(content))) + if value["type"] == "audio": + import torchaudio + + return torchaudio.load(io.BytesIO(content)) + with tempfile.NamedTemporaryFile(suffix=_media_suffix(value, ".mp4")) as file: + file.write(content) + file.flush() + return load_video(file.name) + except Exception as e: + raise GMSError("invalid_request", f"could not decode input {name!r}: {e}", pointer) from e + + +def _save_images(images: Any, directory: Path, name: str, media_type: str, rate: float) -> list[dict[str, Any]]: + if isinstance(images, np.ndarray): + images = numpy_to_pil(images) + images = _as_pil_list(images) + if images is None: + raise GMSError("internal", f"cannot encode output {name!r} as {media_type}: not a list of images") + + artifacts = [] + for index, image in enumerate(images): + filename = f"{name}-{index}.{media_type.split('/')[1]}" + if media_type == "image/jpeg": + image = image.convert("RGB") + image.save(directory / filename) + artifacts.append({"file": filename, "media_type": media_type, "width": image.width, "height": image.height}) + return artifacts + + +def _save_videos(videos: Any, directory: Path, name: str, media_type: str, rate: float) -> list[dict[str, Any]]: + artifacts = [] + for index, frames in enumerate(videos): + frames = list(frames) + filename = f"{name}-{index}.mp4" + export_to_video(frames, str(directory / filename), fps=rate) + first = frames[0] + width, height = first.size if isinstance(first, Image.Image) else (first.shape[1], first.shape[0]) + artifacts.append( + { + "file": filename, + "media_type": media_type, + "width": width, + "height": height, + "duration": round(len(frames) / rate, 2), + "fps": rate, + } + ) + return artifacts + + +def _save_audio(audios: Any, directory: Path, name: str, media_type: str, rate: float) -> list[dict[str, Any]]: + if isinstance(audios, torch.Tensor): + audios = audios.detach().to(torch.float32).cpu().numpy() + arrays = _as_audio_arrays(audios) + if arrays is None: + raise GMSError("internal", f"cannot encode output {name!r} as {media_type}: not an audio array") + + staging = directory / f"{name}-staging" + paths = _save_audio_arrays(arrays, int(rate), Namespace(output=str(staging) + os.sep)) + artifacts = [] + for index, path in enumerate(paths): + filename = f"{name}-{index}.wav" + Path(path).rename(directory / filename) + with wave.open(str(directory / filename), "rb") as file: + duration = round(file.getnframes() / file.getframerate(), 2) + artifacts.append({"file": filename, "media_type": media_type, "duration": duration, "sample_rate": int(rate)}) + staging.rmdir() + return artifacts + + +_SAVERS = {"image": _save_images, "video": _save_videos, "audio": _save_audio} + + +class DiffusersBackend: + def __init__(self, args: Namespace, served: ServedModel): + self._args = args + self._served = served + self.pipeline = _load_pipeline(args) + self._modular = isinstance(self.pipeline, ModularPipeline) + if self._modular: + parameters = blocks_signature(self.pipeline.blocks).parameters + else: + parameters = inspect.signature(type(self.pipeline).__call__).parameters + self._defaults = { + name: parameter.default + for name, parameter in parameters.items() + if parameter.default is not inspect.Parameter.empty + } + self._reports_steps = "callback_on_step_end" in parameters + + def close(self) -> None: + pass + + def failure(self) -> str | None: + return None + + def generate(self, generation: Generation) -> dict[str, list[dict[str, Any]]]: + request = generation.request + kwargs = dict(request.kwargs) + for kwarg, (name, value) in request.media.items(): + if isinstance(value, list): + loaded = [_load_media(name, item, f"/inputs/{name}/{index}") for index, item in enumerate(value)] + else: + loaded = [_load_media(name, value, f"/inputs/{name}")] + if self._served.tasks[request.task]["inputs"][name]["type"] == "audio": + if kwarg == "initial_audio_waveforms": + kwargs.setdefault("initial_audio_sampling_rate", loaded[0][1]) + loaded = [waveform for waveform, _ in loaded] + kwargs[kwarg] = loaded if isinstance(value, list) else loaded[0] + + if generation.seed is not None: + device = self.pipeline.device.type if hasattr(self.pipeline, "device") else "cpu" + kwargs["generator"] = _get_generator(generation.seed, device) + + if self._reports_steps: + + def on_step_end(pipeline, step, timestep, callback_kwargs): + total_steps = getattr(pipeline, "num_timesteps", None) + if total_steps: + generation.report_progress(step + 1, total_steps) + if generation.cancel_requested: + pipeline._interrupt = True + return callback_kwargs + + kwargs["callback_on_step_end"] = on_step_end + + try: + self._load_adapters(request.adapters) + result = self.pipeline(**kwargs) + finally: + if request.adapters: + self.pipeline.unload_lora_weights() + if generation.cancel_requested: + return {} + + outputs = {} + for name, spec in self._served.tasks[request.task]["outputs"].items(): + kind = spec["type"] + # A modular pipeline returns its `PipelineState`; a standard one returns an output object. + if self._modular: + fields = (result.get(field) for field in _OUTPUT_FIELDS.get(kind, ())) + else: + fields = (getattr(result, field, None) for field in _OUTPUT_FIELDS.get(kind, ())) + value = next((field for field in fields if field is not None), None) + if value is None: + continue + rate = self._fps(kwargs) if kind == "video" else (self._args.sampling_rate or 16000) + outputs[name] = _SAVERS[kind](value, generation.directory, name, request.output_formats[kind], rate) + if not outputs: + declared = list(self._served.tasks[request.task]["outputs"]) + raise GMSError( + "internal", + f"{type(self.pipeline).__name__} returned none of the outputs {declared} declared for task " + f"{request.task!r}", + ) + return outputs + + def _fps(self, kwargs: dict[str, Any]) -> float: + for kwarg in ("fps", "frame_rate"): + rate = kwargs.get(kwarg) or self._defaults.get(kwarg) + if rate: + return rate + return self._args.fps + + def _load_adapters(self, adapters: list[dict[str, Any]]) -> None: + if not adapters: + return + names = [] + weights = [] + for index, adapter in enumerate(adapters): + name = f"request_{index}" + try: + self.pipeline.load_lora_weights(adapter["path"].removeprefix("hf:"), adapter_name=name) + except Exception as e: + raise GMSError( + "invalid_request", + f"could not load adapter {adapter['path']!r}: {type(e).__name__}: {e}", + f"/adapters/{index}/path", + ) from e + scale = adapter.get("scale", 1.0) + targets = adapter.get("targets") + if targets: + valid_targets = self._served.adapters[adapter["type"]].get("targets", []) + weights.append({target: targets.get(target, scale) for target in valid_targets}) + else: + weights.append(scale) + names.append(name) + self.pipeline.set_adapters(names, adapter_weights=weights) + + +class EngineBackend: + def __init__(self, engine: str, served: ServedModel, client: Any, process: subprocess.Popen | None = None): + self._engine = engine + self._served = served + self._client = client + self._process = process + + @classmethod + def launch(cls, args: Namespace, served: ServedModel) -> EngineBackend: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + port = probe.getsockname()[1] + + if args.backend == "sglang": + command = ["sglang", "serve", "--model-path", args.model] + else: + command = ["vllm", "serve", args.model, "--omni"] + command += ["--host", "127.0.0.1", "--port", str(port)] + if args.revision: + command += ["--revision", args.revision] + command += shlex.split(args.backend_args or "") + + if shutil.which(command[0]) is None: + raise SystemExit( + f"--backend {args.backend} needs the `{command[0]}` executable on PATH. " + f"Install it with: {_ENGINE_INSTALL[args.backend]}" + ) + + logger.info(f"starting {args.backend} backend: {shlex.join(command)}") + process = subprocess.Popen(command) + client = httpx.Client(base_url=f"http://127.0.0.1:{port}", timeout=None) + backend = cls(args.backend, served, client, process) + try: + backend._wait_until_ready(command) + except BaseException: + backend.close() + raise + return backend + + def failure(self) -> str | None: + if self._process is None or self._process.poll() is None: + return None + return f"the {self._engine} backend exited with code {self._process.returncode}" + + def close(self) -> None: + self._client.close() + if self._process is None or self._process.poll() is not None: + return + self._process.terminate() + try: + self._process.wait(timeout=30) + except subprocess.TimeoutExpired: + self._process.kill() + self._process.wait() + + def generate(self, generation: Generation) -> dict[str, list[dict[str, Any]]]: + request = generation.request + outputs = {spec["type"]: name for name, spec in self._served.tasks[request.task]["outputs"].items()} + + fields = {_ENGINE_FIELDS.get(kwarg, kwarg): value for kwarg, value in request.kwargs.items()} + if generation.seed is not None: + fields["seed"] = generation.seed + width = fields.pop("width", None) + height = fields.pop("height", None) + if (width is None) != (height is None): + missing = "width" if width is None else "height" + raise GMSError( + "invalid_request", + f"the {self._engine} backend takes width and height together; {missing} is missing", + f"/parameters/{missing}", + ) + if width is not None: + fields["size"] = f"{width}x{height}" + + images = [] + if "image" in request.media: + name, value = request.media["image"] + for index, item in enumerate(value if isinstance(value, list) else [value]): + pointer = f"/inputs/{name}/{index}" if isinstance(value, list) else f"/inputs/{name}" + content = _media_bytes(name, item, pointer) + images.append((f"image-{index}{_media_suffix(item, '.png')}", content, item.get("media_type"))) + + if "video" in outputs: + return {outputs["video"]: self._generate_video(generation, fields, images)} + return {outputs["image"]: self._generate_images(generation, fields, images)} + + def _generate_images(self, generation: Generation, fields: dict[str, Any], images: list) -> list[dict[str, Any]]: + media_type = generation.request.output_formats["image"] + extension = media_type.split("/")[1] + fields = {**fields, "response_format": "b64_json", "output_format": extension} + if images: + response = self._client.post("/v1/images/edits", files=self._form(fields, "image", images)) + else: + response = self._client.post("/v1/images/generations", json=fields) + + artifacts = [] + for index, item in enumerate(self._json(response)["data"]): + content = base64.b64decode(item["b64_json"]) + filename = f"image-{index}.{extension}" + (generation.directory / filename).write_bytes(content) + width, height = Image.open(io.BytesIO(content)).size + artifacts.append({"file": filename, "media_type": media_type, "width": width, "height": height}) + return artifacts + + def _generate_video(self, generation: Generation, fields: dict[str, Any], images: list) -> list[dict[str, Any]]: + job = self._json(self._client.post("/v1/videos", files=self._form(fields, "input_reference", images[:1]))) + while job["status"] in ("queued", "in_progress"): + if generation.cancel_requested: + self._client.delete(f"/v1/videos/{job['id']}") + return [] + time.sleep(_POLL_SECONDS) + job = self._json(self._client.get(f"/v1/videos/{job['id']}")) + if job["status"] != "completed": + message = (job.get("error") or {}).get("message", f"video job ended with status {job['status']!r}") + raise GMSError("internal", f"{self._engine} backend: {message}") + + response = self._client.get(f"/v1/videos/{job['id']}/content") + if response.status_code >= 400: + self._json(response) + filename = "video-0.mp4" + (generation.directory / filename).write_bytes(response.content) + self._client.delete(f"/v1/videos/{job['id']}") + + artifact = {"file": filename, "media_type": "video/mp4"} + if job.get("fps"): + artifact["fps"] = job["fps"] + if job.get("duration_s"): + artifact["duration"] = job["duration_s"] + return [artifact] + + @staticmethod + def _form(fields: dict[str, Any], file_field: str, files: list) -> list: + # Both engines read these routes as multipart forms, so plain fields are sent as nameless + # parts: httpx falls back to urlencoding when a request carries no file parts. + parts = [] + for key, value in fields.items(): + text = str(value).lower() if isinstance(value, bool) else str(value) + parts.append((key, (None, text))) + for filename, content, media_type in files: + parts.append((file_field, (filename, content, media_type or "application/octet-stream"))) + return parts + + def _json(self, response: Any) -> Any: + if response.status_code >= 400: + code = "invalid_request" if response.status_code < 500 else "internal" + raise GMSError(code, f"{self._engine} backend returned {response.status_code}: {response.text[:500]}") + return response.json() + + def _wait_until_ready(self, command: list[str]) -> None: + while True: + if self._process.poll() is not None: + raise SystemExit( + f"`{shlex.join(command)}` exited with code {self._process.returncode} before becoming ready." + ) + try: + if self._client.get("/health", timeout=5.0).status_code == 200: + return + except httpx.HTTPError: + pass + time.sleep(2.0) diff --git a/src/diffusers/commands/serve/generations.py b/src/diffusers/commands/serve/generations.py new file mode 100644 index 000000000000..592ca965d54a --- /dev/null +++ b/src/diffusers/commands/serve/generations.py @@ -0,0 +1,194 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The generation resource and the single-worker queue that runs generations one at a time.""" + +from __future__ import annotations + +import queue +import secrets +import shutil +import threading +import time +import uuid +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Any + +from ...utils import logging +from .manifest import GMSError, ResolvedRequest + + +logger = logging.get_logger("diffusers-cli/serve") + +EXPIRY_SECONDS = 3600 + + +class Generation: + def __init__( + self, + model: str, + revision: str | None, + request: ResolvedRequest, + seeded: bool, + default_seed: int | None = None, + ): + self.id = f"gen_{uuid.uuid4().hex[:12]}" + self.model = model + self.revision = revision + self.request = request + # The spec has the server pick the seed when the request omits it, so the result is reproducible. + self.seed = None + if seeded and request.seed is not None: + self.seed = request.seed + elif seeded and default_seed is not None: + self.seed = default_seed + elif seeded: + self.seed = secrets.randbelow(2**32) + self.status = "queued" + self.created_at = datetime.now(timezone.utc) + self.expires_at: datetime | None = None + self.progress: dict[str, Any] | None = None + self.outputs: dict[str, list[dict[str, Any]]] | None = None + self.usage: dict[str, Any] | None = None + self.error: GMSError | None = None + self.cancel_requested = False + self.directory: Path | None = None + self.events: list[tuple[str, dict[str, Any] | None]] = [] + self.done = False + self._condition = threading.Condition() + + def report_progress(self, step: int, total_steps: int) -> None: + self.progress = {"step": step, "total_steps": total_steps, "fraction": round(step / total_steps, 4)} + self._emit("progress", self.progress) + + def wait_for_events(self, index: int, timeout: float) -> tuple[list[tuple[str, dict[str, Any] | None]], bool]: + with self._condition: + self._condition.wait_for(lambda: len(self.events) > index or self.done, timeout) + return self.events[index:], self.done + + def wait(self, timeout: float) -> None: + with self._condition: + self._condition.wait_for(lambda: self.done, timeout) + + def finish(self, status: str, error: GMSError | None = None) -> None: + with self._condition: + self.status = status + self.error = error + self.progress = None + self.expires_at = datetime.now(timezone.utc) + timedelta(seconds=EXPIRY_SECONDS) + if error is not None: + self.events.append(("error", {"code": error.code, "message": error.message})) + elif status == "complete": + self.events.append(("complete", None)) + self.done = True + self._condition.notify_all() + + def _emit(self, event: str, data: dict[str, Any] | None) -> None: + with self._condition: + self.events.append((event, data)) + self._condition.notify_all() + + +class GenerationQueue: + def __init__(self, backend: Any, output_dir: str): + self._backend = backend + self._output_dir = Path(output_dir) + self._generations: dict[str, Generation] = {} + self._lock = threading.Lock() + self._pending: queue.Queue[Generation] = queue.Queue() + threading.Thread(target=self._work, name="diffusers-serve-worker", daemon=True).start() + + def failure(self) -> str | None: + return self._backend.failure() + + def submit(self, generation: Generation) -> None: + self._prune() + with self._lock: + self._generations[generation.id] = generation + self._pending.put(generation) + + def get(self, generation_id: str) -> Generation: + self._prune() + with self._lock: + generation = self._generations.get(generation_id) + if generation is None: + raise GMSError("not_found", f"no generation with id {generation_id!r}") + return generation + + def cancel(self, generation_id: str) -> Generation: + generation = self.get(generation_id) + with self._lock: + if generation.done: + return generation + generation.cancel_requested = True + if generation.status == "queued": + generation.finish("canceled") + return generation + + def delete(self, generation_id: str) -> None: + generation = self.cancel(generation_id) + with self._lock: + del self._generations[generation.id] + # A generation still running keeps writing into its directory; the worker removes it once the backend returns. + if generation.done: + self._remove_files(generation) + + def _prune(self) -> None: + now = datetime.now(timezone.utc) + with self._lock: + expired = [g for g in self._generations.values() if g.expires_at is not None and g.expires_at <= now] + for generation in expired: + del self._generations[generation.id] + for generation in expired: + self._remove_files(generation) + + def _remove_files(self, generation: Generation) -> None: + if generation.directory is not None: + shutil.rmtree(generation.directory, ignore_errors=True) + + def _work(self) -> None: + while True: + generation = self._pending.get() + with self._lock: + if generation.done: + continue + generation.status = "generating" + generation.directory = self._output_dir / generation.id + + started = time.perf_counter() + error = None + try: + generation.directory.mkdir(parents=True, exist_ok=True) + outputs = self._backend.generate(generation) + except GMSError as e: + error = e + # The helpers shared with `run` raise SystemExit, and anything escaping here would end the + # only worker and leave every later generation queued. + except BaseException as e: + logger.error(f"generation {generation.id} failed", exc_info=True) + error = GMSError("internal", f"{type(e).__name__}: {e}") + + with self._lock: + deleted = generation.id not in self._generations + if error is not None: + generation.finish("failed", error) + elif generation.cancel_requested: + generation.finish("canceled") + else: + generation.outputs = outputs + generation.usage = {"compute_seconds": round(time.perf_counter() - started, 2)} + generation.finish("complete") + if deleted or generation.status != "complete": + self._remove_files(generation) diff --git a/src/diffusers/commands/serve/manifest.py b/src/diffusers/commands/serve/manifest.py new file mode 100644 index 000000000000..054dbb8a9c93 --- /dev/null +++ b/src/diffusers/commands/serve/manifest.py @@ -0,0 +1,779 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generative Media Spec (GMS) manifests: loading, derivation from a pipeline signature, and request validation.""" + +from __future__ import annotations + +import hashlib +import inspect +import json +import math +import operator +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from huggingface_hub import hf_hub_download +from huggingface_hub.errors import EntryNotFoundError + +from ...pipelines.auto_pipeline import ( + AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, + AUTO_IMAGE2VIDEO_PIPELINES_MAPPING, + AUTO_INPAINT_PIPELINES_MAPPING, + AUTO_TEXT2AUDIO_PIPELINES_MAPPING, + AUTO_TEXT2IMAGE_PIPELINES_MAPPING, + AUTO_TEXT2VIDEO_PIPELINES_MAPPING, + AUTO_VIDEO2VIDEO_PIPELINES_MAPPING, +) +from ..run import _AUDIO_INPUT_KEYS, _IMAGE_INPUT_KEYS, _VIDEO_INPUT_KEYS + + +SPEC_VERSION = "gms/0.1" +MANIFEST_FILENAME = "gms.json" +EXTENSION_KEY = "x-diffusers" + +MAX_ADAPTERS = 4 +ADAPTER_SOURCES = ("hf",) + +# First entry of each tuple is the default when the request names no `output_format`. +OUTPUT_FORMATS = { + "image": ("image/png", "image/jpeg", "image/webp"), + "video": ("video/mp4",), + "audio": ("audio/wav",), +} + +_ERROR_STATUS = { + "invalid_request": 400, + "not_found": 404, + "rate_limited": 429, + "internal": 500, + "not_implemented": 501, +} + +_REQUEST_FIELDS = ( + "model", + "task", + "preset", + "inputs", + "parameters", + "output_format", + "response_format", + "adapters", +) + +# GMS name -> pipeline kwargs that carry the same meaning, tried in order before the name itself. +_PARAMETER_KWARGS = { + "seed": ("generator",), + "steps": ("num_inference_steps",), + "num_outputs": ("num_images_per_prompt", "num_videos_per_prompt", "num_waveforms_per_prompt"), + "fps": ("fps", "frame_rate"), + "duration": ("num_frames",), +} +_INPUT_KWARGS = {"mask": ("mask_image",)} + +# (task, classes registered for it, inputs the task requires, output modality) +_AUTO_TASKS = ( + ("text_to_image", AUTO_TEXT2IMAGE_PIPELINES_MAPPING, (), "image"), + ("image_to_image", AUTO_IMAGE2IMAGE_PIPELINES_MAPPING, ("image",), "image"), + ("inpaint", AUTO_INPAINT_PIPELINES_MAPPING, ("image", "mask"), "image"), + ("text_to_video", AUTO_TEXT2VIDEO_PIPELINES_MAPPING, (), "video"), + ("image_to_video", AUTO_IMAGE2VIDEO_PIPELINES_MAPPING, ("image",), "video"), + ("video_to_video", AUTO_VIDEO2VIDEO_PIPELINES_MAPPING, ("video",), "video"), + ("text_to_audio", AUTO_TEXT2AUDIO_PIPELINES_MAPPING, (), "audio"), +) + +_SKIPPED_KWARGS = frozenset({"output_type", "return_dict", "callback_steps", "callback_on_step_end_tensor_inputs"}) +_AUDIO_OUTPUT_KWARGS = frozenset({"audio_length_in_s", "audio_end_in_s", "num_waveforms_per_prompt"}) +_SCALAR_TYPES = {"float": "number", "int": "integer", "bool": "boolean", "str": "string"} +_PYTHON_TYPES = {"number": (int, float), "integer": int, "boolean": bool, "string": str} +_COMPARISONS = { + "eq": operator.eq, + "gt": operator.gt, + "gte": operator.ge, + "lt": operator.lt, + "lte": operator.le, +} +_FRAME_LATTICE = re.compile(r"(\d+)k\+(\d+)") +_HF_ADAPTER_PATH = re.compile(r"hf:[\w.-]+/[\w.-]+") + + +class GMSError(Exception): + def __init__(self, code: str, message: str, pointer: str | None = None): + super().__init__(message) + self.code = code + self.message = message + self.pointer = pointer + self.status = _ERROR_STATUS[code] + + def to_dict(self) -> dict[str, Any]: + error = {"code": self.code, "message": self.message} + if self.pointer is not None: + error["pointer"] = self.pointer + return {"error": error} + + +@dataclass +class Manifest: + raw: bytes + data: dict[str, Any] + source: str | None = None + + @property + def hash(self) -> str: + return f"sha256:{hashlib.sha256(self.raw).hexdigest()}" + + +@dataclass +class ResolvedRequest: + task: str + kwargs: dict[str, Any] + # pipeline kwarg -> (request input name, media value or list of media values) + media: dict[str, tuple[str, Any]] + seed: int | None + output_formats: dict[str, str] + response_format: str + adapters: list[dict[str, Any]] = field(default_factory=list) + + +def blocks_signature(blocks: Any) -> inspect.Signature: + """The call signature of a modular pipeline, which takes its keyword arguments from its blocks' inputs.""" + parameters = [] + for input_param in blocks.inputs: + if input_param.name is None: + continue + annotation = inspect.Parameter.empty if input_param.type_hint is None else input_param.type_hint + parameters.append( + inspect.Parameter( + input_param.name, inspect.Parameter.KEYWORD_ONLY, default=input_param.default, annotation=annotation + ) + ) + return inspect.Signature(parameters) + + +def load_manifest( + model: str, + pipeline_cls: type, + signature: inspect.Signature, + kwarg_names: set[str], + manifest_path: str | None = None, + revision: str | None = None, + token: str | None = None, +) -> Manifest: + if manifest_path is not None: + return _parse_manifest(Path(manifest_path).read_bytes(), origin=manifest_path) + + if Path(model).is_dir(): + published = Path(model) / MANIFEST_FILENAME + if published.is_file(): + return _parse_manifest(published.read_bytes(), origin=str(published)) + else: + try: + published = hf_hub_download(model, MANIFEST_FILENAME, revision=revision, token=token) + except EntryNotFoundError: + published = None + if published is not None: + manifest = _parse_manifest(Path(published).read_bytes(), origin=f"{model}/{MANIFEST_FILENAME}") + manifest.source = f"hf:{model}/{MANIFEST_FILENAME}" + return manifest + + model_id = Path(model).name if Path(model).is_dir() else model + data = derive_manifest(model_id, pipeline_cls, signature, kwarg_names) + return Manifest(raw=json.dumps(data).encode(), data=data) + + +def _parse_manifest(raw: bytes, origin: str) -> Manifest: + try: + data = json.loads(raw) + except json.JSONDecodeError as e: + raise SystemExit(f"GMS manifest {origin!r} is not valid JSON: {e}") from e + if not isinstance(data, dict) or not isinstance(data.get("id"), str): + raise SystemExit(f"GMS manifest {origin!r} must be a JSON object with a string `id`.") + if not isinstance(data.get("tasks"), dict) or not data["tasks"]: + raise SystemExit(f"GMS manifest {origin!r} declares no `tasks`, so there is nothing to serve.") + return Manifest(raw=raw, data=data) + + +def _annotation_schema(annotation: Any, default: Any) -> dict[str, Any] | None: + """Map a `__call__` annotation to a GMS parameter schema, or `None` when it is not a scalar or list of scalars.""" + if annotation is inspect.Parameter.empty: + text = type(default).__name__ if default is not None else "" + elif isinstance(annotation, str): + text = annotation + elif type(annotation) is type: + text = annotation.__name__ + else: + text = str(annotation) + text = text.replace("typing.", "") + + def scalar(fragment: str) -> str | None: + names = set(re.findall(r"[A-Za-z_][\w.]*", fragment)) + names -= {"Optional", "Union", "None", "NoneType", "list", "List"} + if not names or not names <= set(_SCALAR_TYPES): + return None + return next(gms for python, gms in _SCALAR_TYPES.items() if python in names) + + outside_lists = re.sub(r"\b[Ll]ist\[[^\[\]]*\]", "", text) + if outside_lists == text: + scalar_type = scalar(text) + return {"type": scalar_type} if scalar_type else None + + # `str | list[str]` is a scalar that also accepts a batch; only a bare list is an array parameter. + scalar_type = scalar(text) + if scalar_type is None: + return None + if scalar(outside_lists): + return {"type": scalar_type} + return {"type": "array", "items": {"type": scalar_type}} + + +def derive_manifest( + model_id: str, pipeline_cls: type, signature: inspect.Signature, kwarg_names: set[str] +) -> dict[str, Any]: + """Build a server-authored manifest from the pipeline call `signature`, keeping only kwargs in `kwarg_names`.""" + inputs: dict[str, dict[str, Any]] = {} + parameters: dict[str, dict[str, Any]] = {} + input_kwargs: dict[str, str] = {} + parameter_kwargs: dict[str, str] = {} + media_types = { + **dict.fromkeys(_IMAGE_INPUT_KEYS, "image"), + **dict.fromkeys(_VIDEO_INPUT_KEYS, "video"), + **dict.fromkeys(_AUDIO_INPUT_KEYS, "audio"), + "mask_image": "mask", + } + kwarg_parameters = {kwarg: name for name, kwargs in _PARAMETER_KWARGS.items() for kwarg in kwargs} + del kwarg_parameters["num_frames"] + + for kwarg, parameter in signature.parameters.items(): + if kwarg == "self" or kwarg in _SKIPPED_KWARGS or kwarg not in kwarg_names: + continue + if parameter.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): + continue + has_default = parameter.default is not inspect.Parameter.empty + default = parameter.default if has_default else None + + if kwarg in media_types: + name = "mask" if kwarg == "mask_image" else kwarg + spec: dict[str, Any] = {"type": media_types[kwarg]} + if not has_default: + spec["required"] = True + if re.search(r"\b[Ll]ist\b", str(parameter.annotation)): + spec["multiple"] = {} + inputs[name] = spec + if name != kwarg: + input_kwargs[name] = kwarg + continue + + if kwarg == "generator": + parameters["seed"] = {"type": "integer"} + parameter_kwargs["seed"] = kwarg + continue + + schema = _annotation_schema(parameter.annotation, default) + # Modular blocks often declare `prompt` without a type hint or a default to infer one from. + if schema is None and kwarg == "prompt" and parameter.annotation is inspect.Parameter.empty: + schema = {"type": "string"} + if schema is None: + continue + if schema["type"] == "string" and re.fullmatch(r"(.*_)?prompt(_\d+)?", kwarg): + inputs[kwarg] = {"type": "text", "required": True} if kwarg == "prompt" else {"type": "text"} + continue + + name = kwarg_parameters.get(kwarg, kwarg) + if name in parameters: + continue + if default is not None and isinstance(default, (bool, int, float, str, list)): + schema["default"] = default + parameters[name] = schema + if name != kwarg: + parameter_kwargs[name] = kwarg + + registered = [entry for entry in _AUTO_TASKS if pipeline_cls in entry[1].values()] + names = set(signature.parameters) + if registered: + output = registered[0][3] + elif "num_frames" in names or "Video" in pipeline_cls.__name__: + output = "video" + elif names & _AUDIO_OUTPUT_KWARGS: + output = "audio" + else: + output = "image" + + extension: dict[str, Any] = {} + if input_kwargs: + extension["inputs"] = input_kwargs + if parameter_kwargs: + extension["parameters"] = parameter_kwargs + + def task(required: tuple[str, ...], excluded: set[str]) -> dict[str, Any]: + task_inputs = { + name: {**spec, "required": True} if name in required else spec + for name, spec in inputs.items() + if name not in excluded + } + spec = {"inputs": task_inputs, "outputs": {output: {"type": output}}, "parameters": parameters} + if extension: + spec[EXTENSION_KEY] = extension + return spec + + primary = next((name for name in ("image", "video") if name in inputs), None) + tasks = {} + if registered: + # A class registered for several tasks takes each task's conditioning input optionally; an + # input another task requires is what tells the tasks apart, so it is left out of the rest. + for name, _, required, _ in registered: + if not set(required) <= set(inputs): + continue + others = {key for other in registered if other[0] != name for key in other[2]} + tasks[name] = task(required, others - set(required)) + elif primary is None: + tasks[f"text_to_{output}"] = task((), set()) + elif inputs[primary].get("required"): + tasks[f"{primary}_to_{output}"] = task((), set()) + else: + tasks[f"text_to_{output}"] = task((), {primary}) + tasks[f"{primary}_to_{output}"] = task((primary,), set()) + + modalities_in = sorted({"image" if s["type"] == "mask" else s["type"] for s in inputs.values()} - {"text"}) + manifest: dict[str, Any] = { + "id": model_id, + "modalities": {"in": modalities_in or ["text"], "out": [output]}, + "tasks": tasks, + } + if hasattr(pipeline_cls, "load_lora_weights"): + manifest["adapters"] = {"lora": {"targets": list(getattr(pipeline_cls, "_lora_loadable_modules", []))}} + return manifest + + +def _bind(name: str, aliases: dict[str, tuple[str, ...]], declared: dict[str, str], kwarg_names: set[str]): + candidates = (declared[name],) if name in declared else (*aliases.get(name, ()), name) + return next((kwarg for kwarg in candidates if kwarg in kwarg_names), None) + + +def _frame_lattice(spec: dict[str, Any]) -> tuple[int, int, float] | None: + lattice = spec.get("lattice") or {} + match = _FRAME_LATTICE.fullmatch(str(lattice.get("frames", "")).replace(" ", "")) + if match is None or not isinstance(lattice.get("fps"), (int, float)): + return None + return int(match.group(1)), int(match.group(2)), lattice["fps"] + + +def _is_active(condition: dict[str, Any], values: dict[str, Any]) -> bool: + if "all" in condition: + return all(_is_active(item, values) for item in condition["all"]) + actual = values.get(condition["parameter"]) + for name, compare in _COMPARISONS.items(): + if name not in condition: + continue + if actual is None or isinstance(actual, str) != isinstance(condition[name], str): + return False + if not compare(actual, condition[name]): + return False + return True + + +def _check_value(label: str, value: Any, spec: dict[str, Any], pointer: str) -> None: + kind = spec.get("type") + if kind == "array": + if not isinstance(value, list): + raise GMSError("invalid_request", f"{label} must be an array, got {value!r}", pointer) + if "maxItems" in spec and len(value) > spec["maxItems"]: + raise GMSError("invalid_request", f"{label} accepts at most {spec['maxItems']} items", pointer) + for index, item in enumerate(value): + _check_value(f"{label}[{index}]", item, spec["items"], f"{pointer}/{index}") + return + if kind == "object": + if not isinstance(value, dict): + raise GMSError("invalid_request", f"{label} must be an object, got {value!r}", pointer) + for key, item in value.items(): + if key not in spec["properties"]: + raise GMSError("invalid_request", f"{label} has no property {key!r}", f"{pointer}/{key}") + _check_value(f"{label}.{key}", item, spec["properties"][key], f"{pointer}/{key}") + return + + expected = _PYTHON_TYPES.get(kind) + is_bool = isinstance(value, bool) + if expected is not None and (not isinstance(value, expected) or (is_bool and kind != "boolean")): + raise GMSError("invalid_request", f"{label} must be of type {kind}, got {value!r}", pointer) + if isinstance(value, float) and not math.isfinite(value): + raise GMSError("invalid_request", f"{label} must be a finite number, got {value}", pointer) + if "enum" in spec and value not in spec["enum"]: + raise GMSError("invalid_request", f"{label} must be one of {spec['enum']}, got {value!r}", pointer) + if "min" in spec and value < spec["min"]: + raise GMSError("invalid_request", f"{label} must be at least {spec['min']}, got {value}", pointer) + if "max" in spec and value > spec["max"]: + raise GMSError("invalid_request", f"{label} must be at most {spec['max']}, got {value}", pointer) + if "multiple_of" in spec: + quotient = value / spec["multiple_of"] + if not math.isclose(quotient, round(quotient), abs_tol=1e-9): + raise GMSError( + "invalid_request", f"{label} must be a multiple of {spec['multiple_of']}, got {value}", pointer + ) + + +def _check_media(name: str, value: Any, spec: dict[str, Any], pointer: str) -> None: + if not isinstance(value, dict): + raise GMSError("invalid_request", f"{name} must be a media value object, got {value!r}", pointer) + if value.get("type") != spec["type"]: + raise GMSError( + "invalid_request", f"{name} must have type {spec['type']!r}, got {value.get('type')!r}", f"{pointer}/type" + ) + keys = set(value) + if keys == {"type", "base64", "media_type"}: + if not isinstance(value["base64"], str) or not isinstance(value["media_type"], str): + raise GMSError("invalid_request", f"{name} base64 and media_type must be strings", pointer) + return + if "url" not in keys or not keys <= {"type", "url", "media_type"}: + raise GMSError( + "invalid_request", + f"{name} must be either {{type, url, media_type?}} or {{type, base64, media_type}}, got keys {sorted(keys)}", + pointer, + ) + # A URL media value is fetched by the server, so anything but http(s) would read the server's own filesystem. + if not isinstance(value["url"], str) or not value["url"].startswith(("http://", "https://")): + raise GMSError("invalid_request", f"{name} url must be an http(s) URL", f"{pointer}/url") + + +class ServedModel: + """One manifest bound to the kwargs a backend accepts: the discovery document and request resolution.""" + + def __init__( + self, + manifest: Manifest, + kwarg_names: set[str], + output_types: set[str], + capabilities: list[str], + revision: str | None = None, + ): + self.manifest = manifest + self.id = manifest.data["id"] + self.capabilities = capabilities + self.revision = revision + self.adapters = manifest.data.get("adapters", {}) + self._kwarg_names = kwarg_names + + self.tasks: dict[str, dict[str, Any]] = {} + self._inputs: dict[str, dict[str, str | None]] = {} + self._parameters: dict[str, dict[str, str | None]] = {} + self._presets: dict[str, list[str]] = {} + for name, task in manifest.data["tasks"].items(): + declared = task.get(EXTENSION_KEY, {}) + inputs = { + key: _bind(key, _INPUT_KWARGS, declared.get("inputs", {}), kwarg_names) for key in task["inputs"] + } + parameters = { + key: _bind(key, _PARAMETER_KWARGS, declared.get("parameters", {}), kwarg_names) + for key in task.get("parameters", {}) + } + if "duration" in parameters and _frame_lattice(task["parameters"]["duration"]) is None: + parameters["duration"] = None + if any(inputs[key] is None for key, spec in task["inputs"].items() if spec.get("required")): + continue + if not any(spec["type"] in output_types for spec in task["outputs"].values()): + continue + self.tasks[name] = task + self._inputs[name] = inputs + self._parameters[name] = parameters + self._presets[name] = [ + preset + for preset, body in task.get("presets", {}).items() + if all(parameters.get(key) is not None for key in body.get("parameters", {})) + ] + + def summary(self) -> dict[str, Any]: + return { + "id": self.id, + "modalities": self.manifest.data.get("modalities", {}), + "tasks": list(self.tasks), + } + + def overlay(self) -> dict[str, Any]: + tasks = {} + limits = {} + for name, task in self.tasks.items(): + tasks[name] = {"presets": self._presets[name]} if "presets" in task else {} + disabled = {key: False for key, kwarg in self._inputs[name].items() if kwarg is None} + if disabled: + limits[name] = {"inputs": disabled} + overlay: dict[str, Any] = {"tasks": tasks, "capabilities": self.capabilities} + if limits: + overlay["limits"] = limits + if "adapters" in self.capabilities: + overlay["adapters"] = {"max": MAX_ADAPTERS, "sources": list(ADAPTER_SOURCES)} + return overlay + + def discovery_document(self) -> bytes: + # The manifest bytes are spliced in untouched: the spec defines verification as a hash of the served bytes. + head: dict[str, Any] = {"spec": SPEC_VERSION} + if self.manifest.source is not None: + head["source"] = self.manifest.source + head["hash"] = self.manifest.hash + prefix = json.dumps(head)[:-1].encode() + return ( + prefix + + b', "manifest": ' + + self.manifest.raw + + b', "server": ' + + json.dumps(self.overlay()).encode() + + b"}" + ) + + def resolve(self, body: Any) -> ResolvedRequest: + if not isinstance(body, dict): + raise GMSError("invalid_request", "request body must be a JSON object") + for key in body: + if key not in _REQUEST_FIELDS: + raise GMSError("invalid_request", f"unknown request field {key!r}", f"/{key}") + if not isinstance(body.get("model"), str): + raise GMSError("invalid_request", "model is required and must be a string", "/model") + if body["model"] != self.id: + raise GMSError("not_found", f"model {body['model']!r} is not served here; this server hosts {self.id!r}") + inputs = body.get("inputs") + if not isinstance(inputs, dict): + raise GMSError("invalid_request", "inputs is required and must be an object", "/inputs") + + name = self._select_task(body.get("task"), inputs) + task = self.tasks[name] + values, explicit = self._resolve_parameters(name, task, body) + + kwargs: dict[str, Any] = {} + media: dict[str, tuple[str, Any]] = {} + for key, value in inputs.items(): + pointer = f"/inputs/{key}" + spec = task["inputs"].get(key) + if spec is None: + raise GMSError("invalid_request", f"task {name!r} declares no input {key!r}", pointer) + kwarg = self._inputs[name][key] + if kwarg is None: + raise GMSError("invalid_request", f"input {key!r} is disabled on this server", pointer) + if "active_when" in spec and not _is_active(spec["active_when"], values): + raise GMSError( + "invalid_request", f"input {key!r} is inactive unless {json.dumps(spec['active_when'])}", pointer + ) + if spec["type"] == "text": + if not isinstance(value, str): + raise GMSError("invalid_request", f"{key} must be a string, got {value!r}", pointer) + kwargs[kwarg] = value + continue + if isinstance(value, list): + if "multiple" not in spec: + raise GMSError("invalid_request", f"{key} accepts a single media value, not a list", pointer) + limit = spec["multiple"].get("max") + if limit is not None and len(value) > limit: + raise GMSError( + "invalid_request", f"{key} accepts at most {limit} items, got {len(value)}", pointer + ) + for index, item in enumerate(value): + _check_media(key, item, spec, f"{pointer}/{index}") + else: + _check_media(key, value, spec, pointer) + media[kwarg] = (key, value) + for key, spec in task["inputs"].items(): + if spec.get("required") and key not in inputs: + raise GMSError("invalid_request", f"task {name!r} requires input {key!r}", f"/inputs/{key}") + + seed = None + for key, value in values.items(): + spec = task["parameters"][key] + kwarg = self._parameters[name][key] + pointer = f"/parameters/{key}" + if "active_when" in spec and not _is_active(spec["active_when"], values): + if key in explicit: + raise GMSError( + "invalid_request", + f"parameter {key!r} is inactive unless {json.dumps(spec['active_when'])}", + pointer, + ) + continue + if kwarg is None: + # Clients that render every control resend defaults, so a default on a parameter this + # backend cannot set is the one supplied value that changes nothing by being skipped. + if key in explicit and value != spec.get("default"): + raise GMSError( + "not_implemented", f"parameter {key!r} is not supported by this server's backend", pointer + ) + continue + if key == "seed": + if not 0 <= value < 2**63: + raise GMSError("invalid_request", f"seed must be between 0 and 2**63 - 1, got {value}", pointer) + seed = value + continue + if key == "duration": + lattice_fps = spec["lattice"]["fps"] + if values.get("fps", lattice_fps) != lattice_fps: + raise GMSError( + "invalid_request", + f"duration is declared on a {lattice_fps} fps frame lattice; fps must be {lattice_fps} " + f"when duration is set, got {values['fps']}", + "/parameters/fps", + ) + kwargs[kwarg] = self._duration_to_frames(value, spec, pointer) + fps_kwarg = _bind("fps", _PARAMETER_KWARGS, {}, self._kwarg_names) + if fps_kwarg is not None: + kwargs[fps_kwarg] = lattice_fps + continue + kwargs[kwarg] = value + + output_formats = {kind: formats[0] for kind, formats in OUTPUT_FORMATS.items()} + output_format = body.get("output_format") + if output_format is not None: + kind = output_format.split("/")[0] if isinstance(output_format, str) else None + produced = {spec["type"] for spec in task["outputs"].values()} + if kind not in produced or output_format not in OUTPUT_FORMATS.get(kind, ()): + supported = sorted(f for k in produced for f in OUTPUT_FORMATS.get(k, ())) + raise GMSError( + "invalid_request", + f"output_format must be one of {supported} for task {name!r}, got {output_format!r}", + "/output_format", + ) + output_formats[kind] = output_format + + response_format = body.get("response_format", "url") + if response_format not in ("url", "b64_json"): + raise GMSError( + "invalid_request", + f"response_format must be 'url' or 'b64_json', got {response_format!r}", + "/response_format", + ) + + return ResolvedRequest( + task=name, + kwargs=kwargs, + media=media, + seed=seed, + output_formats=output_formats, + response_format=response_format, + adapters=self._resolve_adapters(body.get("adapters")), + ) + + def seeded(self, task: str) -> bool: + return self._parameters[task].get("seed") is not None + + def _select_task(self, requested: Any, inputs: dict[str, Any]) -> str: + if requested is not None: + if not isinstance(requested, str) or requested not in self.tasks: + raise GMSError( + "invalid_request", + f"task {requested!r} is not hosted for {self.id!r}; hosted tasks: {list(self.tasks)}", + "/task", + ) + return requested + + supplied = set(inputs) + candidates = [] + for name, task in self.tasks.items(): + accepted = {key for key, kwarg in self._inputs[name].items() if kwarg is not None} + required = {key for key, spec in task["inputs"].items() if spec.get("required")} + if required <= supplied <= accepted: + candidates.append(name) + if len(candidates) == 1: + return candidates[0] + if len(self.tasks) == 1: + return next(iter(self.tasks)) + raise GMSError( + "invalid_request", + f"could not infer the task from inputs {sorted(supplied)}; set `task` to one of {list(self.tasks)}", + "/task", + ) + + def _resolve_parameters(self, name: str, task: dict[str, Any], body: dict[str, Any]): + declared = task.get("parameters", {}) + explicit = body.get("parameters", {}) + if not isinstance(explicit, dict): + raise GMSError("invalid_request", "parameters must be an object", "/parameters") + for key, value in explicit.items(): + if key not in declared: + raise GMSError("invalid_request", f"task {name!r} declares no parameter {key!r}", f"/parameters/{key}") + _check_value(key, value, declared[key], f"/parameters/{key}") + + values = {key: spec["default"] for key, spec in declared.items() if "default" in spec} + preset = body.get("preset") + if preset is not None: + if preset not in self._presets[name]: + raise GMSError( + "invalid_request", + f"preset {preset!r} is not supported for task {name!r}; supported presets: {self._presets[name]}", + "/preset", + ) + values.update(task["presets"][preset].get("parameters", {})) + values.update(explicit) + return values, set(explicit) + + def _duration_to_frames(self, seconds: float, spec: dict[str, Any], pointer: str) -> int: + step, offset, fps = _frame_lattice(spec) + position = (seconds * fps - offset) / step + snap = spec["lattice"].get("snap", "nearest") + if snap == "reject" and not math.isclose(position, round(position), abs_tol=1e-9): + raise GMSError( + "invalid_request", + f"duration {seconds} is off the {spec['lattice']['frames']} frame lattice at {fps} fps", + pointer, + ) + index = math.ceil(position - 1e-9) if snap == "up" else round(position) + return step * max(index, 0) + offset + + def _resolve_adapters(self, adapters: Any) -> list[dict[str, Any]]: + if adapters is None: + return [] + if "adapters" not in self.capabilities: + raise GMSError("not_implemented", "this server does not apply adapters", "/adapters") + if not isinstance(adapters, list): + raise GMSError("invalid_request", "adapters must be an array", "/adapters") + if len(adapters) > MAX_ADAPTERS: + raise GMSError( + "invalid_request", f"at most {MAX_ADAPTERS} adapters per request, got {len(adapters)}", "/adapters" + ) + for index, adapter in enumerate(adapters): + pointer = f"/adapters/{index}" + if not isinstance(adapter, dict): + raise GMSError("invalid_request", "an adapter must be an object", pointer) + for key in adapter: + if key not in ("type", "path", "scale", "targets"): + raise GMSError("invalid_request", f"unknown adapter field {key!r}", f"{pointer}/{key}") + if not isinstance(adapter.get("type"), str) or adapter["type"] not in self.adapters: + raise GMSError( + "invalid_request", + f"adapter type must be one of {list(self.adapters)}, got {adapter.get('type')!r}", + f"{pointer}/type", + ) + path = adapter.get("path") + # Anything but a repo id would reach `load_lora_weights` as a path on the server's own disk. + if not isinstance(path, str) or not _HF_ADAPTER_PATH.fullmatch(path): + raise GMSError( + "invalid_request", + f"adapter path must be a Hub repo id of the form 'hf:/', got {path!r}", + f"{pointer}/path", + ) + scale = adapter.get("scale", 1.0) + if isinstance(scale, bool) or not isinstance(scale, (int, float)): + raise GMSError("invalid_request", f"adapter scale must be a number, got {scale!r}", f"{pointer}/scale") + valid_targets = self.adapters[adapter["type"]].get("targets", []) + targets = adapter.get("targets", {}) + if not isinstance(targets, dict): + raise GMSError("invalid_request", "adapter targets must be an object", f"{pointer}/targets") + for target, target_scale in targets.items(): + if target not in valid_targets: + raise GMSError( + "invalid_request", + f"adapter target must be one of {valid_targets}, got {target!r}", + f"{pointer}/targets/{target}", + ) + if isinstance(target_scale, bool) or not isinstance(target_scale, (int, float)): + raise GMSError( + "invalid_request", + f"adapter target scale must be a number, got {target_scale!r}", + f"{pointer}/targets/{target}", + ) + return adapters diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index 89d411a63a1d..a0975e4a5918 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -53,4 +53,6 @@ "opencv-python": "opencv-python", "timm": "timm", "flashpack": "flashpack", + "fastapi": "fastapi", + "uvicorn": "uvicorn", } diff --git a/tests/others/test_cli_commands.py b/tests/others/test_cli_commands.py index 5b27eaf372f4..e2d2d11f3cb6 100644 --- a/tests/others/test_cli_commands.py +++ b/tests/others/test_cli_commands.py @@ -53,7 +53,7 @@ ) -AVAILABLE_COMMANDS = ("env", "fp16_safetensors", "custom_blocks", "run", "schema", "skills") +AVAILABLE_COMMANDS = ("env", "fp16_safetensors", "custom_blocks", "generate", "run", "schema", "serve", "skills") class TestRunCommand: diff --git a/tests/others/test_cli_serve.py b/tests/others/test_cli_serve.py new file mode 100644 index 000000000000..03a31909fe5f --- /dev/null +++ b/tests/others/test_cli_serve.py @@ -0,0 +1,826 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for `diffusers-cli serve`. + +No weights are loaded: the diffusers backend runs a fake pipeline and the engine backend talks to a fake +SGLang / vLLM-Omni server through an httpx mock transport. +""" + +import base64 +import inspect +import io +import json +import subprocess +import sys +import threading +from argparse import ArgumentParser, Namespace +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from huggingface_hub.utils import httpx +from PIL import Image + +from diffusers.commands.generate import GenerateCommand +from diffusers.commands.serve import ServeCommand +from diffusers.commands.serve.backends import ( + DiffusersBackend, + EngineBackend, + backend_capabilities, + backend_kwargs, + backend_output_types, +) +from diffusers.commands.serve.generations import Generation, GenerationQueue +from diffusers.commands.serve.manifest import GMSError, Manifest, ServedModel, blocks_signature, derive_manifest +from diffusers.modular_pipelines import InputParam, ModularPipelineBlocks, OutputParam + + +MODEL_ID = "test-org/fake-model" +WAIT = {"Prefer": "wait=30"} + + +class FakePipeline: + _lora_loadable_modules = ["transformer", "text_encoder"] + device = torch.device("cpu") + num_timesteps = 0 + + def __init__(self): + self.calls = [] + self.lora_calls = [] + self.release = threading.Event() + self.release.set() + self.started = threading.Event() + self.error = None + self.result = None + + def __call__( + self, + prompt: str | list[str] = None, + negative_prompt: str | list[str] | None = None, + image: Image.Image | None = None, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 4, + sigmas: list[float] | None = None, + guidance_scale: float = 3.5, + num_images_per_prompt: int | None = 1, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + return_dict: bool = True, + callback_on_step_end=None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + ): + self.calls.append( + { + "prompt": prompt, + "negative_prompt": negative_prompt, + "image": image, + "height": height, + "width": width, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "generator": generator, + } + ) + self.started.set() + self.release.wait(timeout=30) + if self.error is not None: + raise self.error + self.num_timesteps = num_inference_steps + self._interrupt = False + for step in range(num_inference_steps): + if self._interrupt: + break + callback_on_step_end(self, step, 0, {}) + if self.result is not None: + return self.result + size = (width or 32, height or 32) + return SimpleNamespace(images=[Image.new("RGB", size, "red") for _ in range(num_images_per_prompt)]) + + def load_lora_weights(self, repo, adapter_name=None): + self.lora_calls.append(("load", repo, adapter_name)) + + def set_adapters(self, names, adapter_weights=None): + self.lora_calls.append(("set", names, adapter_weights)) + + def unload_lora_weights(self): + self.lora_calls.append(("unload",)) + + +FAKE_SIGNATURE = inspect.signature(FakePipeline.__call__) + + +class FakeImageBlock(ModularPipelineBlocks): + """A one-block modular pipeline that declares its inputs the way the built-in blocks do, mostly untyped.""" + + model_name = "fake" + + @property + def inputs(self): + return [ + InputParam("prompt"), + InputParam("image"), + InputParam("width"), + InputParam("num_inference_steps", default=4), + InputParam("generator"), + ] + + @property + def intermediate_outputs(self): + return [OutputParam("images")] + + def __call__(self, components, state): + block_state = self.get_block_state(state) + size = (block_state.width or 32, 32) + block_state.images = [Image.new("RGB", size, "red")] + self.set_block_state(state, block_state) + return components, state + + +def _args(**overrides): + values = {"backend": "diffusers", "compile": None, "lora": None, "fps": 8, "sampling_rate": None} + return Namespace(**{**values, **overrides}) + + +def _served(backend="diffusers", manifest=None, pipeline_cls=FakePipeline, **arg_overrides): + args = _args(backend=backend, **arg_overrides) + signature = inspect.signature(pipeline_cls.__call__) + kwarg_names = backend_kwargs(backend, signature) + if manifest is None: + manifest = derive_manifest(MODEL_ID, pipeline_cls, signature, kwarg_names) + raw = manifest if isinstance(manifest, bytes) else json.dumps(manifest).encode() + return ServedModel( + Manifest(raw=raw, data=json.loads(raw)), + kwarg_names, + backend_output_types(backend), + backend_capabilities(args, pipeline_cls), + revision="abc123", + ) + + +@pytest.fixture +def serve(tmp_path, monkeypatch): + """Build a test client around a `DiffusersBackend` running `FakePipeline`. Returns `(client, pipeline)`.""" + pytest.importorskip("fastapi", reason="`diffusers-cli serve` needs the `diffusers[serve]` extra") + from fastapi.testclient import TestClient + + from diffusers.commands.serve.app import create_app + + def _build(manifest=None, app_options=None, **arg_overrides): + pipeline = FakePipeline() + monkeypatch.setattr("diffusers.commands.serve.backends._load_pipeline", lambda args: pipeline) + served = _served(manifest=manifest, **arg_overrides) + backend = DiffusersBackend(_args(**arg_overrides), served) + app = create_app(served, GenerationQueue(backend, str(tmp_path)), **(app_options or {})) + return TestClient(app), pipeline + + return _build + + +LATTICE_MANIFEST = { + "id": MODEL_ID, + "modalities": {"in": ["image"], "out": ["video"]}, + "tasks": { + "text_to_video": { + "inputs": { + "prompt": {"type": "text", "required": True}, + "negative_prompt": {"type": "text", "active_when": {"parameter": "guidance_scale", "gt": 1}}, + "last_image": {"type": "image"}, + }, + "outputs": {"video": {"type": "video"}}, + "parameters": { + "duration": {"type": "number", "min": 1, "max": 20, "lattice": {"frames": "8k+1", "fps": 24}}, + "width": {"type": "integer", "default": 768, "min": 512, "max": 1536, "multiple_of": 64}, + "steps": {"type": "integer", "default": 30, "min": 1, "max": 100}, + "fps": {"type": "number"}, + "guidance_scale": {"type": "number", "default": 3.0, "min": 1.0, "max": 10.0}, + "guidance_type": {"type": "string", "enum": ["cfg", "distilled"], "default": "cfg"}, + "seed": {"type": "integer"}, + "video_decoder": {"type": "string", "enum": ["conv", "diffusion"], "default": "conv"}, + }, + "presets": { + "default": {}, + "distilled": {"parameters": {"steps": 8, "guidance_scale": 1.0}}, + "quality": {"parameters": {"steps": 40, "video_decoder": "diffusion"}}, + }, + } + }, +} + + +class FakeVideoPipeline: + def __call__( + self, + prompt: str = None, + negative_prompt: str | None = None, + width: int = 768, + num_frames: int = 121, + frame_rate: float = 24.0, + num_inference_steps: int = 40, + guidance_scale: float = 4.0, + generator: torch.Generator | None = None, + ): + pass + + +class TestManifest: + def test_derive_manifest_from_signature(self): + manifest = derive_manifest(MODEL_ID, FakePipeline, FAKE_SIGNATURE, backend_kwargs("diffusers", FAKE_SIGNATURE)) + + # An optional `image` kwarg yields one task without it and one that requires it. + assert list(manifest["tasks"]) == ["text_to_image", "image_to_image"], manifest["tasks"].keys() + text_to_image = manifest["tasks"]["text_to_image"] + assert text_to_image["inputs"] == { + "prompt": {"type": "text", "required": True}, + "negative_prompt": {"type": "text"}, + } + assert manifest["tasks"]["image_to_image"]["inputs"]["image"] == {"type": "image", "required": True} + assert text_to_image["outputs"] == {"image": {"type": "image"}} + # Core vocabulary names replace the pipeline's own; tensors, callbacks and output plumbing are dropped. + assert text_to_image["parameters"] == { + "height": {"type": "integer"}, + "width": {"type": "integer"}, + "steps": {"type": "integer", "default": 4}, + "sigmas": {"type": "array", "items": {"type": "number"}}, + "guidance_scale": {"type": "number", "default": 3.5}, + "num_outputs": {"type": "integer", "default": 1}, + "seed": {"type": "integer"}, + } + assert text_to_image["x-diffusers"] == { + "parameters": {"steps": "num_inference_steps", "num_outputs": "num_images_per_prompt", "seed": "generator"} + } + assert manifest["adapters"] == {"lora": {"targets": ["transformer", "text_encoder"]}} + + @pytest.mark.parametrize( + "body, code, pointer", + [ + ({"parameters": {"cfg": 3}}, "invalid_request", "/parameters/cfg"), + ({"parameters": {"width": 700}}, "invalid_request", "/parameters/width"), + ({"parameters": {"steps": 500}}, "invalid_request", "/parameters/steps"), + ({"parameters": {"steps": "8"}}, "invalid_request", "/parameters/steps"), + ({"inputs": {}}, "invalid_request", "/inputs/prompt"), + ({"inputs": {"prompt": "a cat", "caption": "x"}}, "invalid_request", "/inputs/caption"), + ({"task": "image_to_video"}, "invalid_request", "/task"), + ({"task": []}, "invalid_request", "/task"), + ({"parameters": []}, "invalid_request", "/parameters"), + ({"parameters": {"seed": 10**30}}, "invalid_request", "/parameters/seed"), + ({"parameters": {"duration": float("nan")}}, "invalid_request", "/parameters/duration"), + # The lattice pins duration to 24 fps, so another fps would deliver a different length. + ({"parameters": {"duration": 4, "fps": 12}}, "invalid_request", "/parameters/fps"), + ({"steps": 8}, "invalid_request", "/steps"), + ({"preset": "quality"}, "invalid_request", "/preset"), + ({"model": "someone/else"}, "not_found", None), + ({"adapters": [{"type": "lora", "path": "hf:a/b"}]}, "not_implemented", "/adapters"), + ({"output_format": "video/webm"}, "invalid_request", "/output_format"), + # `negative_prompt` is only honored while guidance_scale > 1. + ( + {"inputs": {"prompt": "a cat", "negative_prompt": "blurry"}, "preset": "distilled"}, + "invalid_request", + "/inputs/negative_prompt", + ), + # `last_image` has no matching pipeline kwarg, so the server overlay disables it. + ( + {"inputs": {"prompt": "a cat", "last_image": {"type": "image", "url": "https://x/y.png"}}}, + "invalid_request", + "/inputs/last_image", + ), + # A parameter the pipeline cannot set is refused rather than dropped, unless it is the default. + ({"parameters": {"video_decoder": "diffusion"}}, "not_implemented", "/parameters/video_decoder"), + ], + ) + def test_resolve_rejects_invalid_requests(self, body, code, pointer): + served = _served(manifest=LATTICE_MANIFEST, pipeline_cls=FakeVideoPipeline) + with pytest.raises(GMSError) as error: + served.resolve({"model": MODEL_ID, "inputs": {"prompt": "a cat"}, **body}) + assert (error.value.code, error.value.pointer) == (code, pointer), error.value.message + + def test_resolve_rejects_non_http_media_urls(self): + served = _served() + body = {"model": MODEL_ID, "inputs": {"prompt": "a cat", "image": {"type": "image", "url": "/etc/passwd"}}} + with pytest.raises(GMSError) as error: + served.resolve(body) + assert error.value.pointer == "/inputs/image/url", error.value.message + + @pytest.mark.parametrize("path", ["hf:/etc/passwd", "hf:../loras/x", "/etc/passwd", "https://x/y.safetensors"]) + def test_resolve_rejects_adapter_paths_that_are_not_repo_ids(self, path): + body = {"model": MODEL_ID, "inputs": {"prompt": "a cat"}, "adapters": [{"type": "lora", "path": path}]} + with pytest.raises(GMSError) as error: + _served().resolve(body) + assert (error.value.code, error.value.pointer) == ("invalid_request", "/adapters/0/path"), error.value.message + + def test_resolve_maps_request_to_pipeline_kwargs(self): + served = _served(manifest=LATTICE_MANIFEST, pipeline_cls=FakeVideoPipeline) + resolved = served.resolve( + { + "model": MODEL_ID, + "preset": "distilled", + "inputs": {"prompt": "a cat"}, + "parameters": {"duration": 8, "seed": 7, "guidance_type": "cfg"}, + } + ) + # 8 s at 24 fps is 192 frames; the nearest point on the 8k+1 lattice is 193. The preset overrides the + # task defaults, and `guidance_type` at its default is accepted although the pipeline has no such kwarg. + assert resolved.kwargs == { + "prompt": "a cat", + "num_frames": 193, + "frame_rate": 24, + "width": 768, + "num_inference_steps": 8, + "guidance_scale": 1.0, + }, resolved.kwargs + assert resolved.seed == 7 + assert resolved.task == "text_to_video" + + def test_discovery_overlay(self): + served = _served(manifest=LATTICE_MANIFEST, pipeline_cls=FakeVideoPipeline) + assert served.overlay() == { + "tasks": {"text_to_video": {"presets": ["default", "distilled"]}}, + "capabilities": ["sse", "wait"], + "limits": {"text_to_video": {"inputs": {"last_image": False}}}, + } + + +class TestServeApp: + def test_generation_lifecycle(self, serve): + client, pipeline = serve() + body = { + "model": MODEL_ID, + "inputs": {"prompt": "a cat"}, + "parameters": {"width": 64, "height": 48, "steps": 3, "seed": 11}, + } + response = client.post("/v1/generations", json=body, headers=WAIT) + assert response.status_code == 201, response.text + generation = response.json() + assert generation["status"] == "complete", generation + assert generation["seed"] == 11 + assert generation["revision"] == "abc123" + assert generation["progress"] is None + assert generation["usage"]["compute_seconds"] >= 0 + + call = pipeline.calls[0] + assert (call["prompt"], call["width"], call["height"], call["num_inference_steps"]) == ("a cat", 64, 48, 3) + assert torch.equal(call["generator"].get_state(), torch.Generator().manual_seed(11).get_state()) + + (artifact,) = generation["outputs"]["image"] + assert {k: artifact[k] for k in ("media_type", "width", "height")} == { + "media_type": "image/png", + "width": 64, + "height": 48, + } + image = Image.open(io.BytesIO(client.get(artifact["url"]).content)) + assert (image.format, image.size) == ("PNG", (64, 48)) + + events = client.get(f"/v1/generations/{generation['id']}/events").text + assert events.count("event: progress") == 3, events + assert 'data: {"step": 3, "total_steps": 3, "fraction": 1.0}' in events + assert events.rstrip().split("\n\n")[-1].startswith("id: 4\nevent: complete\n") + # Resuming after the last progress event replays only what followed it. + resumed = client.get(f"/v1/generations/{generation['id']}/events", headers={"Last-Event-ID": "3"}).text + assert "event: progress" not in resumed and "event: complete" in resumed, resumed + + assert client.delete(f"/v1/generations/{generation['id']}").status_code == 204 + missing = client.get(f"/v1/generations/{generation['id']}") + assert (missing.status_code, missing.json()["error"]["code"]) == (404, "not_found") + assert client.get(artifact["url"]).status_code == 404 + + def test_inline_response_and_media_input(self, serve): + client, pipeline = serve() + buffer = io.BytesIO() + Image.new("RGB", (8, 8), "blue").save(buffer, format="PNG") + body = { + "model": MODEL_ID, + "inputs": { + "prompt": "a cat", + "image": { + "type": "image", + "base64": base64.b64encode(buffer.getvalue()).decode(), + "media_type": "image/png", + }, + }, + "output_format": "image/jpeg", + "response_format": "b64_json", + } + generation = client.post("/v1/generations", json=body, headers=WAIT).json() + assert generation["status"] == "complete", generation + # The task is inferred from the inputs: only `image_to_image` accepts an image. + assert pipeline.calls[0]["image"].size == (8, 8) + (artifact,) = generation["outputs"]["image"] + assert artifact["media_type"] == "image/jpeg" and "url" not in artifact + assert Image.open(io.BytesIO(base64.b64decode(artifact["base64"]))).format == "JPEG" + + def test_joint_video_and_audio_outputs(self, serve, monkeypatch): + monkeypatch.setattr( + "diffusers.commands.serve.backends.export_to_video", + lambda frames, path, fps: Path(path).write_bytes(b"mp4"), + ) + manifest = { + "id": MODEL_ID, + "tasks": { + "text_to_video": { + "inputs": {"prompt": {"type": "text", "required": True}}, + "outputs": {"video": {"type": "video"}, "audio": {"type": "audio"}}, + } + }, + } + client, pipeline = serve(manifest=manifest) + pipeline.result = SimpleNamespace( + frames=[[Image.new("RGB", (32, 24)) for _ in range(16)]], + audio=np.zeros((1, 2, 8000), dtype=np.float32), + ) + body = {"model": MODEL_ID, "inputs": {"prompt": "a cat"}} + generation = client.post("/v1/generations", json=body, headers=WAIT).json() + assert generation["status"] == "complete", generation + + # The fake pipeline takes no fps argument, so the video is encoded at the `--fps` default. + (video,) = generation["outputs"]["video"] + (audio,) = generation["outputs"]["audio"] + assert {k: v for k, v in video.items() if k != "url"} == { + "media_type": "video/mp4", + "width": 32, + "height": 24, + "duration": 2.0, + "fps": 8, + } + assert {k: v for k, v in audio.items() if k != "url"} == { + "media_type": "audio/wav", + "duration": 0.5, + "sample_rate": 16000, + } + assert client.get(audio["url"]).content[:4] == b"RIFF" + + def test_validation_error_envelope(self, serve): + client, _ = serve() + body = {"model": MODEL_ID, "inputs": {"prompt": "a cat"}, "parameters": {"cfg_scale": 3}} + response = client.post("/v1/generations", json=body) + assert response.status_code == 400 + assert response.json() == { + "error": { + "code": "invalid_request", + "message": "task 'text_to_image' declares no parameter 'cfg_scale'", + "pointer": "/parameters/cfg_scale", + } + } + + def test_failed_generation_returns_error_envelope(self, serve): + client, pipeline = serve() + body = {"model": MODEL_ID, "inputs": {"prompt": "a cat"}} + envelope = {"error": {"code": "internal", "message": "RuntimeError: CUDA out of memory"}} + + pipeline.error = RuntimeError("CUDA out of memory") + created = client.post("/v1/generations", json=body) + assert created.status_code == 201 + # The event stream closes once the generation is done, so reading it to the end waits for the failure. + events = client.get(f"/v1/generations/{created.json()['id']}/events").text + assert 'event: error\ndata: {"code": "internal", "message": "RuntimeError: CUDA out of memory"}' in events + response = client.get(f"/v1/generations/{created.json()['id']}") + assert (response.status_code, response.json()) == (500, envelope) + # There is no `failed` status: a create that waited for the failure returns the envelope as well. + waited = client.post("/v1/generations", json=body, headers=WAIT) + assert (waited.status_code, waited.json()) == (500, envelope) + + # A SystemExit from the backend must not end the worker: the next generation still runs. + pipeline.error = SystemExit("unsupported") + assert client.post("/v1/generations", json=body, headers=WAIT).status_code == 500 + pipeline.error = None + assert client.post("/v1/generations", json=body, headers=WAIT).json()["status"] == "complete" + + def test_health(self, serve): + client, _ = serve() + response = client.get("/health") + assert (response.status_code, response.json()) == (200, {"status": "ok"}) + + def test_request_id_is_generated_or_echoed(self, serve): + client, _ = serve() + assert client.get("/health").headers["x-request-id"], "a request without an id should be given one" + echoed = client.get("/v1/models/unknown", headers={"x-request-id": "req-1"}) + assert (echoed.status_code, echoed.headers["x-request-id"]) == (404, "req-1") + + def test_default_seed_applies_when_request_omits_seed(self, serve): + client, _ = serve(app_options={"default_seed": 7}) + body = {"model": MODEL_ID, "inputs": {"prompt": "a cat"}} + assert client.post("/v1/generations", json=body, headers=WAIT).json()["seed"] == 7 + body["parameters"] = {"seed": 11} + assert client.post("/v1/generations", json=body, headers=WAIT).json()["seed"] == 11 + + def test_cors_is_opt_in(self, serve): + origin = {"Origin": "http://example.com"} + client, _ = serve() + assert "access-control-allow-origin" not in client.get("/health", headers=origin).headers + client, _ = serve(app_options={"enable_cors": True}) + assert client.get("/health", headers=origin).headers["access-control-allow-origin"] == "*" + + def test_cancel_queued_and_running(self, serve): + client, pipeline = serve() + pipeline.release.clear() + body = {"model": MODEL_ID, "inputs": {"prompt": "a cat"}} + running = client.post("/v1/generations", json=body).json() + queued = client.post("/v1/generations", json=body).json() + assert queued["status"] == "queued", queued + assert pipeline.started.wait(timeout=30), "the worker never picked up the first generation" + + assert client.post(f"/v1/generations/{queued['id']}/cancel").json()["status"] == "canceled" + client.post(f"/v1/generations/{running['id']}/cancel") + pipeline.release.set() + + # The stream ends once the generation is done; a canceled generation emits no `complete` event. + events = client.get(f"/v1/generations/{running['id']}/events").text + assert "event: complete" not in events, events + canceled = client.get(f"/v1/generations/{running['id']}").json() + assert (canceled["status"], canceled["outputs"]) == ("canceled", None) + assert len(pipeline.calls) == 1, "the canceled queued generation must never reach the pipeline" + + def test_request_adapters_are_loaded_then_unloaded(self, serve): + client, pipeline = serve() + body = { + "model": MODEL_ID, + "inputs": {"prompt": "a cat"}, + "adapters": [ + {"type": "lora", "path": "hf:someone/style", "scale": 0.8}, + {"type": "lora", "path": "hf:someone/detail", "scale": 0.5, "targets": {"text_encoder": 0.1}}, + ], + } + generation = client.post("/v1/generations", json=body, headers=WAIT).json() + assert generation["status"] == "complete", generation + assert pipeline.lora_calls == [ + ("load", "someone/style", "request_0"), + ("load", "someone/detail", "request_1"), + ("set", ["request_0", "request_1"], [0.8, {"transformer": 0.5, "text_encoder": 0.1}]), + ("unload",), + ] + + def test_adapters_not_advertised_when_compiled(self, serve): + client, _ = serve(compile="{}") + discovery = client.get(f"/v1/models/{MODEL_ID}").json() + assert "adapters" not in discovery["server"]["capabilities"] + body = {"model": MODEL_ID, "inputs": {"prompt": "a cat"}, "adapters": [{"type": "lora", "path": "hf:a/b"}]} + response = client.post("/v1/generations", json=body) + assert (response.status_code, response.json()["error"]["code"]) == (501, "not_implemented") + + def test_discovery_serves_manifest_bytes_verbatim(self, serve): + manifest = derive_manifest(MODEL_ID, FakePipeline, FAKE_SIGNATURE, backend_kwargs("diffusers", FAKE_SIGNATURE)) + raw = json.dumps(manifest, indent=3).encode() + client, _ = serve(manifest=raw) + + response = client.get(f"/v1/models/{MODEL_ID}") + assert raw in response.content + discovery = response.json() + assert discovery["spec"] == "gms/0.1" + assert discovery["manifest"] == manifest + assert discovery["server"]["tasks"] == {"text_to_image": {}, "image_to_image": {}} + assert discovery["server"]["adapters"] == {"max": 4, "sources": ["hf"]} + assert client.get("/v1/models").json() == { + "data": [ + { + "id": MODEL_ID, + "modalities": {"in": ["image"], "out": ["image"]}, + "tasks": ["text_to_image", "image_to_image"], + } + ] + } + + def test_unadvertised_capabilities_return_not_implemented(self, serve): + client, _ = serve() + for response in ( + client.post("/v1/uploads", json={}), + client.get("/v1/generations/gen_x/outputs/video/stream"), + ): + assert (response.status_code, response.json()["error"]["code"]) == (501, "not_implemented") + + +class FakeEngine: + """The image and video routes SGLang and vLLM-Omni share, recording what the backend sends.""" + + def __init__(self): + self.requests = [] + self.polls = 0 + + def __call__(self, request): + content_type = request.headers.get("content-type", "") + self.requests.append((request.method, request.url.path, content_type.split(";")[0], request.content)) + path = request.url.path + if path in ("/v1/images/generations", "/v1/images/edits"): + buffer = io.BytesIO() + Image.new("RGB", (16, 8), "green").save(buffer, format="PNG") + return httpx.Response(200, json={"data": [{"b64_json": base64.b64encode(buffer.getvalue()).decode()}]}) + if path == "/v1/videos" and request.method == "POST": + return httpx.Response(200, json={"id": "video_1", "status": "queued"}) + if path == "/v1/videos/video_1" and request.method == "GET": + self.polls += 1 + if self.polls < 2: + return httpx.Response(200, json={"id": "video_1", "status": "in_progress"}) + return httpx.Response(200, json={"id": "video_1", "status": "completed", "fps": 16.0, "duration_s": 5.0}) + if path == "/v1/videos/video_1/content": + return httpx.Response(200, content=b"mp4-bytes") + if path == "/v1/videos/video_1" and request.method == "DELETE": + return httpx.Response(200, json={"id": "video_1", "deleted": True}) + return httpx.Response(404, text="no such route") + + +class FakeTextToVideoPipeline: + def __call__( + self, + prompt: str = None, + image: Image.Image = None, + height: int = 480, + width: int = 832, + num_frames: int = 81, + num_inference_steps: int = 50, + generator: torch.Generator | None = None, + ): + pass + + +class TestModularPipeline: + def test_generation_runs_a_modular_pipeline(self, tmp_path, monkeypatch): + pytest.importorskip("fastapi", reason="`diffusers-cli serve` needs the `diffusers[serve]` extra") + from fastapi.testclient import TestClient + + from diffusers.commands.serve.app import create_app + + pipeline = FakeImageBlock().init_pipeline() + monkeypatch.setattr("diffusers.commands.serve.backends._load_pipeline", lambda args: pipeline) + signature = blocks_signature(pipeline.blocks) + kwarg_names = backend_kwargs("diffusers", signature) + manifest = derive_manifest(MODEL_ID, type(pipeline), signature, kwarg_names) + assert list(manifest["tasks"]) == ["text_to_image", "image_to_image"], manifest["tasks"].keys() + assert manifest["tasks"]["text_to_image"]["inputs"] == {"prompt": {"type": "text", "required": True}} + assert manifest["tasks"]["text_to_image"]["parameters"] == { + "steps": {"type": "integer", "default": 4}, + "seed": {"type": "integer"}, + } + + raw = json.dumps(manifest).encode() + served = ServedModel( + Manifest(raw=raw, data=manifest), + kwarg_names, + backend_output_types("diffusers"), + backend_capabilities(_args(), type(pipeline)), + revision=None, + ) + backend = DiffusersBackend(_args(), served) + client = TestClient(create_app(served, GenerationQueue(backend, str(tmp_path)))) + + body = {"model": MODEL_ID, "inputs": {"prompt": "a cat"}, "parameters": {"seed": 3}} + generation = client.post("/v1/generations", json=body, headers=WAIT).json() + assert generation["status"] == "complete", generation + assert generation["seed"] == 3 + (artifact,) = generation["outputs"]["image"] + assert (artifact["media_type"], artifact["width"], artifact["height"]) == ("image/png", 32, 32) + + +class TestEngineBackend: + def _generate(self, tmp_path, monkeypatch, served, body): + monkeypatch.setattr("diffusers.commands.serve.backends._POLL_SECONDS", 0.0) + engine = FakeEngine() + client = httpx.Client(base_url="http://engine", transport=httpx.MockTransport(engine)) + backend = EngineBackend("vllm", served, client) + resolved = served.resolve({"model": MODEL_ID, **body}) + generation = Generation(served.id, None, resolved, seeded=served.seeded(resolved.task)) + generation.directory = tmp_path + return engine, backend.generate(generation) + + def test_exited_engine_is_unhealthy(self, tmp_path): + pytest.importorskip("fastapi", reason="`diffusers-cli serve` needs the `diffusers[serve]` extra") + from fastapi.testclient import TestClient + + from diffusers.commands.serve.app import create_app + + process = subprocess.Popen([sys.executable, "-c", "raise SystemExit(3)"]) + process.wait() + served = _served(backend="vllm") + backend = EngineBackend("vllm", served, httpx.Client(base_url="http://engine"), process) + response = TestClient(create_app(served, GenerationQueue(backend, str(tmp_path)))).get("/health") + expected = {"status": "unhealthy", "reason": "the vllm backend exited with code 3"} + assert (response.status_code, response.json()) == (503, expected) + + def test_derived_manifest_keeps_only_engine_kwargs(self): + manifest = derive_manifest(MODEL_ID, FakePipeline, FAKE_SIGNATURE, backend_kwargs("vllm", FAKE_SIGNATURE)) + assert "sigmas" not in manifest["tasks"]["text_to_image"]["parameters"] + assert "steps" in manifest["tasks"]["text_to_image"]["parameters"] + + def test_text_to_image_request_translation(self, tmp_path, monkeypatch): + body = { + "inputs": {"prompt": "a cat", "negative_prompt": "blurry"}, + "parameters": {"width": 1024, "height": 768, "steps": 20, "num_outputs": 2, "seed": 5}, + "output_format": "image/webp", + } + engine, outputs = self._generate(tmp_path, monkeypatch, _served(backend="vllm"), body) + + ((method, path, content_type, content),) = engine.requests + assert (method, path, content_type) == ("POST", "/v1/images/generations", "application/json") + assert json.loads(content) == { + "prompt": "a cat", + "negative_prompt": "blurry", + "num_inference_steps": 20, + "guidance_scale": 3.5, + "n": 2, + "seed": 5, + "size": "1024x768", + "response_format": "b64_json", + "output_format": "webp", + } + assert outputs == {"image": [{"file": "image-0.webp", "media_type": "image/webp", "width": 16, "height": 8}]} + assert (tmp_path / "image-0.webp").is_file() + + def test_image_input_uses_multipart_edits(self, tmp_path, monkeypatch): + image = {"type": "image", "base64": base64.b64encode(b"png-bytes").decode(), "media_type": "image/png"} + body = {"inputs": {"prompt": "make it blue", "image": image}, "parameters": {"seed": 5}} + engine, _ = self._generate(tmp_path, monkeypatch, _served(backend="sglang"), body) + + ((method, path, content_type, content),) = engine.requests + assert (method, path, content_type) == ("POST", "/v1/images/edits", "multipart/form-data") + assert b'name="image"; filename="image-0.png"' in content and b"png-bytes" in content + assert b'name="prompt"\r\n\r\nmake it blue' in content + + def test_video_job_is_polled_downloaded_and_deleted(self, tmp_path, monkeypatch): + served = _served(backend="vllm", pipeline_cls=FakeTextToVideoPipeline) + body = {"task": "text_to_video", "inputs": {"prompt": "a cat"}, "parameters": {"num_frames": 33, "seed": 5}} + engine, outputs = self._generate(tmp_path, monkeypatch, served, body) + + assert [(method, path) for method, path, _, _ in engine.requests] == [ + ("POST", "/v1/videos"), + ("GET", "/v1/videos/video_1"), + ("GET", "/v1/videos/video_1"), + ("GET", "/v1/videos/video_1/content"), + ("DELETE", "/v1/videos/video_1"), + ] + _, _, content_type, content = engine.requests[0] + assert content_type == "multipart/form-data" + assert b'name="num_frames"\r\n\r\n33' in content and b'name="size"\r\n\r\n832x480' in content + assert outputs == {"video": [{"file": "video-0.mp4", "media_type": "video/mp4", "fps": 16.0, "duration": 5.0}]} + assert (tmp_path / "video-0.mp4").read_bytes() == b"mp4-bytes" + + def test_audio_tasks_are_not_hosted(self): + manifest = { + "id": MODEL_ID, + "tasks": { + "text_to_audio": {"inputs": {"prompt": {"type": "text"}}, "outputs": {"audio": {"type": "audio"}}} + }, + } + assert _served(backend="vllm", manifest=manifest).tasks == {} + + +class TestGenerateCommand: + def _run(self, monkeypatch, client, argv): + parser = ArgumentParser() + GenerateCommand.register_subcommand(parser.add_subparsers()) + monkeypatch.setattr("diffusers.commands.generate._POLL_SECONDS", 0.01) + monkeypatch.setattr(GenerateCommand, "_client", lambda self: client) + GenerateCommand(parser.parse_args(["generate", "--url", "http://testserver", *argv])).run() + + def test_generate_saves_outputs(self, serve, tmp_path, monkeypatch): + client, pipeline = serve() + output = tmp_path / "saved" + argv = ["--inputs", '{"prompt": "a cat"}', "--parameters", '{"width": 64, "seed": 5}', "-o", str(output)] + self._run(monkeypatch, client, argv) + + assert pipeline.calls[0]["prompt"] == "a cat" + assert [path.name for path in output.iterdir()] == ["image-0.png"], list(output.iterdir()) + assert Image.open(output / "image-0.png").size == (64, 32) + + def test_local_file_is_sent_as_a_media_value(self, serve, tmp_path, monkeypatch): + client, pipeline = serve() + source = tmp_path / "cat.png" + Image.new("RGB", (8, 8), "blue").save(source) + inputs = json.dumps({"prompt": "make it grey", "image": str(source)}) + self._run(monkeypatch, client, ["--inputs", inputs, "-o", str(tmp_path / "saved")]) + + assert pipeline.calls[0]["image"].size == (8, 8) + + def test_server_error_is_reported(self, serve, tmp_path, monkeypatch): + client, _ = serve() + argv = ["--inputs", '{"prompt": "a cat"}', "--parameters", '{"cfg": 1}', "-o", str(tmp_path)] + with pytest.raises(SystemExit, match=r"HTTP 400: invalid_request: .*\(at /parameters/cfg\)"): + self._run(monkeypatch, client, argv) + + +class TestServeCommand: + def _parse(self, argv): + parser = ArgumentParser() + ServeCommand.register_subcommand(parser.add_subparsers()) + return parser.parse_args(["serve", *argv]) + + def test_engine_backend_rejects_pipeline_flags(self): + pytest.importorskip("fastapi", reason="`diffusers-cli serve` needs the `diffusers[serve]` extra") + args = self._parse(["-m", "org/model", "--backend", "vllm", "--dtype", "bf16", "--vae-tiling"]) + with pytest.raises(SystemExit, match="--dtype, --vae-tiling configure the in-process pipeline"): + ServeCommand(args).run() + + def test_context_parallel_is_rejected(self): + pytest.importorskip("fastapi", reason="`diffusers-cli serve` needs the `diffusers[serve]` extra") + with pytest.raises(SystemExit, match="--context-parallel is not supported"): + ServeCommand(self._parse(["-m", "org/model", "--context-parallel"])).run() From 3e4aa57c2842e1ff01edbde8dec80e2147c341be Mon Sep 17 00:00:00 2001 From: DN6 Date: Fri, 9 Oct 2026 13:07:51 +0530 Subject: [PATCH 2/4] add remote serve --- src/diffusers/commands/serve/remote.py | 156 +++++++++++++++++++++++++ 1 file changed, 156 insertions(+) create mode 100644 src/diffusers/commands/serve/remote.py diff --git a/src/diffusers/commands/serve/remote.py b/src/diffusers/commands/serve/remote.py new file mode 100644 index 000000000000..ebabe97ab57e --- /dev/null +++ b/src/diffusers/commands/serve/remote.py @@ -0,0 +1,156 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""`diffusers-cli serve --remote` — run the server in a Hugging Face Sandbox, reachable through the sandbox proxy.""" + +from __future__ import annotations + +import shlex +import sys +import threading +from argparse import ArgumentParser, Namespace +from pathlib import Path +from typing import Any + +from huggingface_hub import Sandbox, get_token +from huggingface_hub.cli._output import out +from huggingface_hub.utils import httpx + +from ...utils import logging +from ..run import _DEFAULT_REMOTE_DEPS, _DEFAULT_REMOTE_IMAGE, _build_task_kwargs, _kwargs_to_argv + + +logger = logging.get_logger("diffusers-cli/serve") + +# The port the server binds inside the sandbox. `generate --sandbox-id` reaches it through the proxy. +SANDBOX_PORT = 8000 +_SANDBOX_MANIFEST_PATH = "/tmp/diffusers-cli/gms.json" +_SERVE_DEPS = ("fastapi", "uvicorn") +_READY_POLL_SECONDS = 5.0 + +# Flags that say how the sandbox is set up, or that the sandbox server sets itself. +_LOCAL_KEYS = ("remote", "flavor", "dependencies", "namespace", "image", "idle_timeout", "func", "format") +_SANDBOX_SET_KEYS = ("host", "port", "public_url", "output_dir") + + +def add_remote_arguments(parser: ArgumentParser) -> None: + parser.add_argument( + "--remote", + action="store_true", + help="Run the server in a Hugging Face Sandbox instead of on the local machine.", + ) + parser.add_argument( + "--flavor", + default="a10g-small", + help="HF Sandbox hardware flavor for --remote (e.g. a10g-small, a100-large).", + ) + parser.add_argument( + "--dependencies", + action="append", + default=None, + help="Extra pip dependencies to install in the sandbox. Repeat to add multiple.", + ) + parser.add_argument( + "--namespace", + default=None, + help="HF namespace to create the sandbox under (defaults to the current user).", + ) + parser.add_argument( + "--image", + default=None, + help=f"Sandbox image for --remote (defaults to {_DEFAULT_REMOTE_IMAGE!r}). Must provide torch + CUDA.", + ) + parser.add_argument( + "--idle-timeout", + default="30m", + help="Auto-shutdown the sandbox after this much inactivity (e.g. 30m, 2h). Defaults to 30m.", + ) + + +def _announce_when_ready(sbx: Any, args: Namespace, stopped: threading.Event) -> None: + health_url = sbx.proxy_url_for(SANDBOX_PORT, "/health") + while not stopped.wait(_READY_POLL_SECONDS): + try: + ready = httpx.get(health_url, headers=sbx.proxy_headers, timeout=10.0).status_code == 200 + except httpx.HTTPError: + ready = False + if not ready: + continue + out.result( + "remote-serve", + sandbox_id=sbx.id, + url=sbx.proxy_url_for(SANDBOX_PORT), + model=args.model, + hint=f'diffusers-cli generate --sandbox-id {sbx.id} --inputs \'{{"prompt": "..."}}\'', + ) + return + + +def serve_remote(args: Namespace, task: str) -> None: + if Path(args.model).exists(): + raise SystemExit( + f"--model {args.model!r} is a local path; the sandbox can't see it. " + "Pass a Hub repo id so the sandbox can download it." + ) + if args.manifest is not None and not Path(args.manifest).is_file(): + raise SystemExit(f"--manifest {args.manifest!r} is not a file.") + + 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, + "flavor": args.flavor, + "forward_hf_token": True, + "token": hf_token, + "env": {"HF_ENABLE_PARALLEL_LOADING": "1"}, + "idle_timeout": args.idle_timeout, + # The 120s default expires while a cold node pulls the multi-GB pytorch base image or + # waits for GPU capacity, which fails the launch before it starts. + "start_timeout": 600.0, + } + if args.namespace is not None: + create_kwargs["namespace"] = args.namespace + sbx = Sandbox.create(**create_kwargs) + + def _stream(chunk: str) -> None: + sys.stderr.write(chunk) + sys.stderr.flush() + + stopped = threading.Event() + try: + dependencies = [*_DEFAULT_REMOTE_DEPS, *_SERVE_DEPS, *(args.dependencies or [])] + 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) + + task_kwargs = _build_task_kwargs(args) + for key in (*_LOCAL_KEYS, *_SANDBOX_SET_KEYS): + task_kwargs.pop(key, None) + if args.manifest is not None: + sbx.files.upload(args.manifest, _SANDBOX_MANIFEST_PATH) + task_kwargs["manifest"] = _SANDBOX_MANIFEST_PATH + 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)] + + 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.") + result = sbx.run(cli_argv, on_stdout=_stream, on_stderr=_stream, check=False, capture_output=False) + finally: + stopped.set() + sbx.kill() + + if result.exit_code != 0: + raise SystemExit(f"remote server exited with code {result.exit_code}") From 5bc39489723b18f1ce2939d31955b48fe38aa45c Mon Sep 17 00:00:00 2001 From: DN6 Date: Fri, 9 Oct 2026 13:07:55 +0530 Subject: [PATCH 3/4] update --- docs/source/en/using-diffusers/cli.md | 26 ++++++++++ src/diffusers/commands/generate.py | 34 ++++++++++--- src/diffusers/commands/serve/__init__.py | 32 +++++++++--- src/diffusers/commands/serve/app.py | 27 ++++++---- tests/others/test_cli_serve.py | 65 ++++++++++++++++++++++++ 5 files changed, 160 insertions(+), 24 deletions(-) diff --git a/docs/source/en/using-diffusers/cli.md b/docs/source/en/using-diffusers/cli.md index d7e8d6523b38..45eb53ee61d6 100644 --- a/docs/source/en/using-diffusers/cli.md +++ b/docs/source/en/using-diffusers/cli.md @@ -311,6 +311,7 @@ report no step progress, and a cancel request only takes effect once the pipelin |---|---| | `--default-seed` | Seed used for requests that do not set one. Without it, each such request gets a random seed. | | `--enable-cors` | Allow cross-origin requests from any origin, so a browser app on another host can call the server. | +| `--public-url` | Base URL clients reach the server at, used to build output file links. Set it behind a reverse proxy. | | `--log-level` | Logging level for diffusers and the HTTP server (`debug`, `info`, `warning`, `error`, `critical`). | `GET /health` returns `{"status": "ok"}`, or `503` when an SGLang or vLLM-Omni backend has exited. Every response @@ -352,6 +353,28 @@ with a JSON pointer to the offending field. {"error": {"code": "invalid_request", "message": "task 'text_to_image' declares no parameter 'cfg'", "pointer": "/parameters/cfg"}} ``` +### Remote serving (`--remote`) + +`--remote` starts the server in a [Hugging Face Sandbox](https://huggingface.co/docs/huggingface_hub/guides/sandbox) +instead of on the local machine. The command creates the sandbox, installs the dependencies, starts the server, +and streams its logs until you interrupt it; the sandbox is then deleted. + +```bash +diffusers-cli serve --model black-forest-labs/FLUX.1-dev --dtype bf16 --remote --flavor a10g-large +``` + +Once the model has loaded, the command prints the sandbox id. The server is only reachable through the sandbox +proxy, which requires your Hugging Face token, so pass the id to [`generate`](#generate) rather than calling a URL +directly: + +```bash +diffusers-cli generate --sandbox-id --inputs '{"prompt": "a cat on the moon"}' +``` + +`--flavor`, `--image`, `--dependencies`, `--namespace` and `--idle-timeout` work as they do for +[`run --remote`](#remote-execution---remote). `--model` must be a Hub repo id. A local `--manifest` file is +uploaded to the sandbox. + ### Manifests What a model accepts is described by its manifest. `serve` looks for one in this order: @@ -431,6 +454,9 @@ Outputs are saved to `~/.diffusers/cli/generate/outputs//` unless `--token` sends a bearer token in the `Authorization` header, for servers behind an authenticating proxy; no token is sent by default. Interrupting the command cancels the generation on the server. +For a server started with [`serve --remote`](#remote-serving---remote), pass `--sandbox-id ` instead of +`--url`. The request then goes through the sandbox proxy, authenticated with your Hugging Face token. + ## `custom_blocks` Package a local `ModularPipelineBlocks` subclass for upload to the Hub. Reads a Python file, AST-scans it for diff --git a/src/diffusers/commands/generate.py b/src/diffusers/commands/generate.py index f3ea85d54e1d..aeebd03c1b59 100644 --- a/src/diffusers/commands/generate.py +++ b/src/diffusers/commands/generate.py @@ -82,6 +82,7 @@ def register_subcommand(subparsers: _SubParsersAction) -> None: " $ diffusers-cli generate --url http://localhost:8000 \\\n" ' --inputs \'{"prompt": "make the fur grey", "image": "cat.png"}\' \\\n' ' --parameters \'{"steps": 8, "seed": 0}\' -o outputs/\n' + ' $ diffusers-cli generate --sandbox-id --inputs \'{"prompt": "a cat on the moon"}\'\n' "\n" "Learn more\n" " Use `diffusers-cli --help` for more information about a command.\n" @@ -91,12 +92,18 @@ def register_subcommand(subparsers: _SubParsersAction) -> None: parser: ArgumentParser = subparsers.add_parser( "generate", help="Send a generation to a Generative Media Spec server and save its outputs.", - usage="\n diffusers-cli generate --url --inputs [options]", + usage="\n diffusers-cli generate (--url | --sandbox-id ) --inputs [options]", epilog=epilog, formatter_class=RawDescriptionHelpFormatter, ) parser._optionals.title = "Options" - parser.add_argument("--url", required=True, help="Base URL of the server, e.g. http://localhost:8000.") + target = parser.add_mutually_exclusive_group(required=True) + target.add_argument("--url", default=None, help="Base URL of the server, e.g. http://localhost:8000.") + target.add_argument( + "--sandbox-id", + default=None, + help="Id of a sandbox started with `diffusers-cli serve --remote`, reached through the sandbox proxy.", + ) parser.add_argument( "--model", "-m", default=None, help="Model id to request. Defaults to the server's only model." ) @@ -130,7 +137,10 @@ def register_subcommand(subparsers: _SubParsersAction) -> None: parser.add_argument( "--token", default=None, - help="Bearer token sent to the server in the Authorization header. No token is sent by default.", + help=( + "With --url: bearer token sent in the Authorization header; none is sent by default. " + "With --sandbox-id: Hugging Face token used to reach the sandbox; defaults to the logged-in one." + ), ) parser.set_defaults(func=GenerateCommand) @@ -189,7 +199,7 @@ def run(self) -> None: out.result( self.task, - url=args.url, + url=str(client.base_url), model=model, generation=generation["id"], seed=generation.get("seed"), @@ -197,6 +207,18 @@ def run(self) -> None: ) def _client(self) -> httpx.Client: + if self.args.sandbox_id is not None: + from huggingface_hub import Sandbox + + from .serve.remote import SANDBOX_PORT + + sbx = Sandbox.connect(self.args.sandbox_id, token=self.args.token) + return httpx.Client( + base_url=sbx.proxy_url_for(SANDBOX_PORT), + headers=sbx.proxy_headers, + timeout=_TIMEOUT_SECONDS, + follow_redirects=True, + ) headers = {"Authorization": f"Bearer {self.args.token}"} if self.args.token else {} return httpx.Client(base_url=self.args.url, headers=headers, timeout=_TIMEOUT_SECONDS, follow_redirects=True) @@ -204,7 +226,7 @@ def _request(self, client: httpx.Client, method: str, path: str, **kwargs: Any) try: response = client.request(method, path, **kwargs) except httpx.HTTPError as e: - raise SystemExit(f"Could not reach {self.args.url}: {type(e).__name__}: {e}") from e + raise SystemExit(f"Could not reach {client.base_url}: {type(e).__name__}: {e}") from e if response.status_code < 400: return response.json() @@ -220,7 +242,7 @@ def _request(self, client: httpx.Client, method: str, path: str, **kwargs: Any) def _only_model(self, client: httpx.Client) -> str: models = [model["id"] for model in self._request(client, "GET", "/v1/models")["data"]] if len(models) != 1: - raise SystemExit(f"{self.args.url} serves {len(models)} models ({models}); pick one with --model.") + raise SystemExit(f"{client.base_url} serves {len(models)} models ({models}); pick one with --model.") return models[0] def _artifact_bytes(self, client: httpx.Client, artifact: dict[str, Any]) -> bytes: diff --git a/src/diffusers/commands/serve/__init__.py b/src/diffusers/commands/serve/__init__.py index e56c73a5c530..e1ddd6d8294b 100644 --- a/src/diffusers/commands/serve/__init__.py +++ b/src/diffusers/commands/serve/__init__.py @@ -30,6 +30,7 @@ from ...utils import logging from .. import BaseDiffusersCLICommand from ..run import _add_loading_arguments, _add_optimization_arguments +from .remote import add_remote_arguments, serve_remote DEFAULT_OUTPUT_DIR = str(Path.home() / ".diffusers" / "cli" / "serve" / "outputs") @@ -139,6 +140,15 @@ def register_subcommand(subparsers: _SubParsersAction) -> None: ) parser.add_argument("--host", default="127.0.0.1", help="Interface to bind. Defaults to 127.0.0.1.") parser.add_argument("--port", type=int, default=8000, help="Port to bind. Defaults to 8000.") + parser.add_argument( + "--public-url", + default=None, + metavar="URL", + help=( + "Base URL clients reach the server at, used to build output file links. Set it when the server " + "sits behind a reverse proxy; defaults to the address of each incoming request." + ), + ) parser.add_argument( "--manifest", default=None, @@ -182,6 +192,7 @@ def register_subcommand(subparsers: _SubParsersAction) -> None: default=None, help="Logging level for diffusers and the HTTP server. Defaults to the current diffusers verbosity.", ) + add_remote_arguments(parser) # `_load_pipeline` reads these `run` flags, which `serve` does not expose. parser.set_defaults(func=ServeCommand, workflow=None, offload_margin=None) @@ -190,14 +201,6 @@ def __init__(self, args: Namespace): def run(self) -> None: args = self.args - try: - import fastapi # noqa: F401 - import uvicorn - except ImportError as e: - raise SystemExit( - '`diffusers-cli serve` requires FastAPI and uvicorn. Install them with: pip install "diffusers[serve]"' - ) from e - if args.log_level is not None: logging.set_verbosity(logging.get_log_levels_dict()[args.log_level]) @@ -220,6 +223,18 @@ def run(self) -> None: f"--backend {args.backend}. Pass the engine's own flags through --backend-args instead." ) + if args.remote: + serve_remote(args, self.task) + return + + try: + import fastapi # noqa: F401 + import uvicorn + except ImportError as e: + raise SystemExit( + '`diffusers-cli serve` requires FastAPI and uvicorn. Install them with: pip install "diffusers[serve]"' + ) from e + from .app import create_app from .backends import ( DiffusersBackend, @@ -266,6 +281,7 @@ def run(self) -> None: GenerationQueue(backend, args.output_dir), default_seed=args.default_seed, enable_cors=args.enable_cors, + public_url=args.public_url, ) out.result( self.task, diff --git a/src/diffusers/commands/serve/app.py b/src/diffusers/commands/serve/app.py index 05213fc302ce..706ef4c0e4d5 100644 --- a/src/diffusers/commands/serve/app.py +++ b/src/diffusers/commands/serve/app.py @@ -22,6 +22,7 @@ import re import uuid from typing import Any +from urllib.parse import urlparse from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware @@ -46,7 +47,7 @@ def _timestamp(value: Any) -> str | None: return value.strftime("%Y-%m-%dT%H:%M:%SZ") -def _render(generation: Generation, request: Request) -> dict[str, Any]: +def _render(generation: Generation, request: Request, public_url: str | None) -> dict[str, Any]: outputs = None if generation.status == "complete": outputs = {} @@ -58,10 +59,12 @@ def _render(generation: Generation, request: Request) -> dict[str, Any]: content = (generation.directory / artifact["file"]).read_bytes() rendered["base64"] = base64.b64encode(content).decode() else: - url = request.url_for( - "get_generation_file", generation_id=generation.id, filename=artifact["file"] - ) - rendered["url"] = str(url) + route = {"generation_id": generation.id, "filename": artifact["file"]} + if public_url is None: + rendered["url"] = str(request.url_for("get_generation_file", **route)) + else: + path = request.app.url_path_for("get_generation_file", **route) + rendered["url"] = public_url.rstrip("/") + path outputs[name].append(rendered) return { @@ -83,8 +86,12 @@ def create_app( generations: GenerationQueue, default_seed: int | None = None, enable_cors: bool = False, + public_url: str | None = None, ) -> FastAPI: - app = FastAPI(title="diffusers-cli serve", docs_url=None, redoc_url=None, openapi_url=None) + # A proxy may or may not strip the path it publishes the server under before forwarding a request. + # With that path as the root path, routing accepts both forms. + root_path = urlparse(public_url).path.rstrip("/") if public_url is not None else "" + app = FastAPI(title="diffusers-cli serve", docs_url=None, redoc_url=None, openapi_url=None, root_path=root_path) @app.middleware("http") async def request_id(request: Request, call_next) -> Response: @@ -157,14 +164,14 @@ async def create_generation(request: Request) -> dict[str, Any]: await asyncio.to_thread(generation.wait, min(float(wait.group(1)), MAX_WAIT_SECONDS)) if generation.error is not None: raise generation.error - return _render(generation, request) + return _render(generation, request, public_url) @app.get("/v1/generations/{generation_id}") async def get_generation(generation_id: str, request: Request) -> dict[str, Any]: generation = generations.get(generation_id) if generation.error is not None: raise generation.error - return _render(generation, request) + return _render(generation, request, public_url) @app.delete("/v1/generations/{generation_id}", status_code=204) async def delete_generation(generation_id: str) -> Response: @@ -176,7 +183,7 @@ async def cancel_generation(generation_id: str, request: Request) -> dict[str, A generation = generations.cancel(generation_id) if generation.error is not None: raise generation.error - return _render(generation, request) + return _render(generation, request, public_url) @app.get("/v1/generations/{generation_id}/events") async def get_generation_events(generation_id: str, request: Request) -> StreamingResponse: @@ -190,7 +197,7 @@ async def stream(): events, done = await asyncio.to_thread(generation.wait_for_events, index, _KEEP_ALIVE_SECONDS) for event, data in events: index += 1 - payload = _render(generation, request) if event == "complete" else data + payload = _render(generation, request, public_url) if event == "complete" else data yield f"id: {index}\nevent: {event}\ndata: {json.dumps(payload)}\n\n" if done: return diff --git a/tests/others/test_cli_serve.py b/tests/others/test_cli_serve.py index 03a31909fe5f..380034ea42d8 100644 --- a/tests/others/test_cli_serve.py +++ b/tests/others/test_cli_serve.py @@ -516,6 +516,19 @@ def test_default_seed_applies_when_request_omits_seed(self, serve): body["parameters"] = {"seed": 11} assert client.post("/v1/generations", json=body, headers=WAIT).json()["seed"] == 11 + def test_public_url_replaces_the_request_address_in_output_links(self, serve): + client, _ = serve(app_options={"public_url": "https://proxy.example/sandbox/proxy/8000/"}) + body = {"model": MODEL_ID, "inputs": {"prompt": "a cat"}} + generation = client.post("/v1/generations", json=body, headers=WAIT).json() + expected = f"https://proxy.example/sandbox/proxy/8000/v1/generations/{generation['id']}/files/image-0.png" + assert generation["outputs"]["image"][0]["url"] == expected + + # The sandbox proxy forwards some requests with the published path still attached. + for path in ("/v1/models", "/sandbox/proxy/8000/v1/models"): + assert client.get(path).status_code == 200, path + prefixed = client.post("/sandbox/proxy/8000/v1/generations", json=body, headers=WAIT) + assert prefixed.status_code == 201, prefixed.text + def test_cors_is_opt_in(self, serve): origin = {"Origin": "http://example.com"} client, _ = serve() @@ -808,6 +821,58 @@ def test_server_error_is_reported(self, serve, tmp_path, monkeypatch): self._run(monkeypatch, client, argv) +class FakeSandbox: + id = "sbx-1" + proxy_headers = {"Authorization": "Bearer hf_x", "X-Sandbox-Token": "tok"} + + def __init__(self): + self.commands = [] + self.uploads = [] + self.killed = False + self.files = SimpleNamespace(upload=lambda local, remote: self.uploads.append((local, remote))) + + def proxy_url_for(self, port, path="/"): + return f"https://job.hf.jobs/v1/sandboxes/s/proxy/{port}{path}" + + def run(self, cmd, **kwargs): + self.commands.append(cmd) + return SimpleNamespace(exit_code=0) + + def kill(self): + self.killed = True + + +class TestServeRemote: + def test_remote_starts_the_server_in_a_sandbox(self, tmp_path, monkeypatch): + sandbox = FakeSandbox() + created = {} + monkeypatch.setattr( + "diffusers.commands.serve.remote.Sandbox", + SimpleNamespace(create=lambda **kwargs: created.update(kwargs) or sandbox), + ) + manifest = tmp_path / "gms.json" + manifest.write_text("{}") + parser = ArgumentParser() + ServeCommand.register_subcommand(parser.add_subparsers()) + argv = ["serve", "-m", "org/model", "--dtype", "bf16", "--remote", "--flavor", "a10g-large"] + argv += ["--manifest", str(manifest), "--port", "9000", "--dependencies", "extra-package"] + ServeCommand(parser.parse_args(argv)).run() + + assert created["flavor"] == "a10g-large" + install, serve_argv = sandbox.commands + assert "fastapi" in install and "extra-package" in install, install + assert serve_argv[:5] == ["diffusers-cli", "--format", "quiet", "serve", "--model"], serve_argv + options = dict(zip(serve_argv[4::2], serve_argv[5::2])) + assert options["--dtype"] == "bf16" + # The sandbox server always binds the proxied port and builds output links on the proxied address. + assert options["--port"] == "8000" + assert options["--public-url"] == "https://job.hf.jobs/v1/sandboxes/s/proxy/8000/" + assert options["--manifest"] == "/tmp/diffusers-cli/gms.json" + assert sandbox.uploads == [(str(manifest), "/tmp/diffusers-cli/gms.json")] + assert "--remote" not in serve_argv and "--flavor" not in serve_argv, serve_argv + assert sandbox.killed, "the sandbox must be killed once the server stops" + + class TestServeCommand: def _parse(self, argv): parser = ArgumentParser() From fc92295c376bf82c30c75e06608d549ffdecc5fc Mon Sep 17 00:00:00 2001 From: DN6 Date: Fri, 9 Oct 2026 23:42:55 +0530 Subject: [PATCH 4/4] update --- docs/source/en/using-diffusers/cli.md | 25 ++++++-- src/diffusers/commands/run.py | 79 +++++++++++++++++++----- src/diffusers/commands/serve/__init__.py | 11 ++-- src/diffusers/commands/serve/backends.py | 14 ++++- src/diffusers/commands/serve/remote.py | 15 ++++- tests/others/test_cli_commands.py | 37 +++++++++++ tests/others/test_cli_serve.py | 45 +++++++++++--- 7 files changed, 187 insertions(+), 39 deletions(-) diff --git a/docs/source/en/using-diffusers/cli.md b/docs/source/en/using-diffusers/cli.md index 45eb53ee61d6..36994135bdca 100644 --- a/docs/source/en/using-diffusers/cli.md +++ b/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 --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. diff --git a/src/diffusers/commands/run.py b/src/diffusers/commands/run.py index cad580d592bc..26922842927a 100644 --- a/src/diffusers/commands/run.py +++ b/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) diff --git a/src/diffusers/commands/serve/__init__.py b/src/diffusers/commands/serve/__init__.py index e1ddd6d8294b..ec35ea631bfa 100644 --- a/src/diffusers/commands/serve/__init__.py +++ b/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." ) diff --git a/src/diffusers/commands/serve/backends.py b/src/diffusers/commands/serve/backends.py index c0a779804dad..bc1eb2811eaf 100644 --- a/src/diffusers/commands/serve/backends.py +++ b/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 diff --git a/src/diffusers/commands/serve/remote.py b/src/diffusers/commands/serve/remote.py index ebabe97ab57e..7743ebf0e444 100644 --- a/src/diffusers/commands/serve/remote.py +++ b/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.") diff --git a/tests/others/test_cli_commands.py b/tests/others/test_cli_commands.py index e2d2d11f3cb6..4c0fc16cbded 100644 --- a/tests/others/test_cli_commands.py +++ b/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") diff --git a/tests/others/test_cli_serve.py b/tests/others/test_cli_serve.py index 380034ea42d8..8701d06a8942 100644 --- a/tests/others/test_cli_serve.py +++ b/tests/others/test_cli_serve.py @@ -46,6 +46,8 @@ from diffusers.commands.serve.generations import Generation, GenerationQueue from diffusers.commands.serve.manifest import GMSError, Manifest, ServedModel, blocks_signature, derive_manifest from diffusers.modular_pipelines import InputParam, ModularPipelineBlocks, OutputParam +from diffusers.utils import logging +from diffusers.utils.testing_utils import CaptureLogger MODEL_ID = "test-org/fake-model" @@ -836,7 +838,7 @@ def proxy_url_for(self, port, path="/"): def run(self, cmd, **kwargs): self.commands.append(cmd) - return SimpleNamespace(exit_code=0) + return SimpleNamespace(exit_code=0, stdout="/usr/bin/python3\n") def kill(self): self.killed = True @@ -859,10 +861,13 @@ def test_remote_starts_the_server_in_a_sandbox(self, tmp_path, monkeypatch): ServeCommand(parser.parse_args(argv)).run() assert created["flavor"] == "a10g-large" - install, serve_argv = sandbox.commands + assert created["image"] == "pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime" + install, _, serve_argv = sandbox.commands assert "fastapi" in install and "extra-package" in install, install - assert serve_argv[:5] == ["diffusers-cli", "--format", "quiet", "serve", "--model"], serve_argv - options = dict(zip(serve_argv[4::2], serve_argv[5::2])) + # The server runs under the interpreter the dependencies were installed into, not whatever is on PATH. + assert serve_argv[:3] == ["/usr/bin/python3", "-m", "diffusers.commands.diffusers_cli"], serve_argv + assert serve_argv[3:7] == ["--format", "quiet", "serve", "--model"], serve_argv + options = dict(zip(serve_argv[6::2], serve_argv[7::2])) assert options["--dtype"] == "bf16" # The sandbox server always binds the proxied port and builds output links on the proxied address. assert options["--port"] == "8000" @@ -872,6 +877,25 @@ def test_remote_starts_the_server_in_a_sandbox(self, tmp_path, monkeypatch): assert "--remote" not in serve_argv and "--flavor" not in serve_argv, serve_argv assert sandbox.killed, "the sandbox must be killed once the server stops" + @pytest.mark.parametrize( + "backend, image", [("sglang", "lmsysorg/sglang:latest"), ("vllm", "vllm/vllm-omni:latest")] + ) + def test_engine_backend_defaults_to_the_engine_image(self, monkeypatch, backend, image): + created = {} + monkeypatch.setattr( + "diffusers.commands.serve.remote.Sandbox", + SimpleNamespace(create=lambda **kwargs: created.update(kwargs) or FakeSandbox()), + ) + parser = ArgumentParser() + ServeCommand.register_subcommand(parser.add_subparsers()) + ServeCommand(parser.parse_args(["serve", "-m", "org/model", "--backend", backend, "--remote"])).run() + assert created["image"] == image + + ServeCommand( + parser.parse_args(["serve", "-m", "org/model", "--backend", backend, "--remote", "--image", "mine"]) + ).run() + assert created["image"] == "mine", "an explicit --image must win over the engine default" + class TestServeCommand: def _parse(self, argv): @@ -879,11 +903,18 @@ def _parse(self, argv): ServeCommand.register_subcommand(parser.add_subparsers()) return parser.parse_args(["serve", *argv]) - def test_engine_backend_rejects_pipeline_flags(self): + def test_engine_backend_warns_about_ignored_pipeline_flags(self, monkeypatch): pytest.importorskip("fastapi", reason="`diffusers-cli serve` needs the `diffusers[serve]` extra") + + def stop(args): + raise RuntimeError("stop after the flag check") + + monkeypatch.setattr("diffusers.commands.serve._resolve_pipeline", stop) args = self._parse(["-m", "org/model", "--backend", "vllm", "--dtype", "bf16", "--vae-tiling"]) - with pytest.raises(SystemExit, match="--dtype, --vae-tiling configure the in-process pipeline"): - ServeCommand(args).run() + with CaptureLogger(logging.get_logger("diffusers-cli/serve")) as captured: + with pytest.raises(RuntimeError, match="stop after the flag check"): + ServeCommand(args).run() + assert "--dtype, --vae-tiling configure the in-process pipeline and are ignored" in captured.out def test_context_parallel_is_rejected(self): pytest.importorskip("fastapi", reason="`diffusers-cli serve` needs the `diffusers[serve]` extra")