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 8625001cef89..df25159a65e2 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. | @@ -130,10 +132,15 @@ Configure how the CLI loads model weights and custom pipeline code. - `--attention-backend` — Hub-hosted attention kernels, auto-downloaded on first use. Choices are `default` plus the Hub backends registered in Diffusers (for example `flash_hub`, `flash_varlen_hub`, `flash_4_hub`, `sage_hub`, `sage_blackwell_hub`). Run `diffusers-cli run --help` for the current list. It only supports Transformer-based pipelines only, and is ignored with a warning on legacy UNet 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, but under `--remote` the CLI wraps `torchrun --nproc-per-node=gpu` for you. See [Context parallelism](../training/distributed_inference#context-parallelism). @@ -287,6 +294,182 @@ 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. | +| `--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 +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"}} +``` + +### 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. + +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: + +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: 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. + +## `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. + +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/setup.py b/setup.py index 14eff971970c..494a7d0a27ba 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..aeebd03c1b59 --- /dev/null +++ b/src/diffusers/commands/generate.py @@ -0,0 +1,260 @@ +# 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' + ' $ 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" + " 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 | --sandbox-id ) --inputs [options]", + epilog=epilog, + formatter_class=RawDescriptionHelpFormatter, + ) + parser._optionals.title = "Options" + 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." + ) + 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=( + "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) + + 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=str(client.base_url), + model=model, + generation=generation["id"], + seed=generation.get("seed"), + outputs=saved, + ) + + 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) + + 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 {client.base_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"{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: + 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/run.py b/src/diffusers/commands/run.py index ae9a465d5bad..ebb453ead4d9 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 @@ -217,13 +225,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: @@ -468,26 +486,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): @@ -496,11 +532,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)) @@ -542,6 +584,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 @@ -807,7 +852,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) @@ -844,14 +889,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 new file mode 100644 index 000000000000..ec35ea631bfa --- /dev/null +++ b/src/diffusers/commands/serve/__init__.py @@ -0,0 +1,300 @@ +# 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 +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") + +# 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, + "compile_mode": "regional", +} + + +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( + "--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, + 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.", + ) + 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) + + def __init__(self, args: Namespace): + self.args = args + + def run(self) -> None: + args = self.args + 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": + ignored = [ + "--" + name.replace("_", "-") + for name, unset in _DIFFUSERS_ONLY_FLAGS.items() + if getattr(args, name) != unset + ] + 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." + ) + + 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, + 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, + public_url=args.public_url, + ) + 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..706ef4c0e4d5 --- /dev/null +++ b/src/diffusers/commands/serve/app.py @@ -0,0 +1,225 @@ +# 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 urllib.parse import urlparse + +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, public_url: str | None) -> 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: + 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 { + "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, + public_url: str | None = None, +) -> FastAPI: + # 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: + 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, 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, public_url) + + @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, public_url) + + @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, public_url) 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..bc1eb2811eaf --- /dev/null +++ b/src/diffusers/commands/serve/backends.py @@ -0,0 +1,482 @@ +# 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 ( + 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 + + +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") + 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 + + +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/commands/serve/remote.py b/src/diffusers/commands/serve/remote.py new file mode 100644 index 000000000000..7743ebf0e444 --- /dev/null +++ b/src/diffusers/commands/serve/remote.py @@ -0,0 +1,165 @@ +# 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") +# 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. +_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}, 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", + 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 _ENGINE_IMAGES.get(args.backend, _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) + # 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): + 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_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.") + 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}") diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index 156d2a1066ca..439e70266e45 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 ba30cfb4a92b..fbfe1f660dbe 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, @@ -54,7 +56,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: @@ -242,6 +244,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 new file mode 100644 index 000000000000..8701d06a8942 --- /dev/null +++ b/tests/others/test_cli_serve.py @@ -0,0 +1,922 @@ +# 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 +from diffusers.utils import logging +from diffusers.utils.testing_utils import CaptureLogger + + +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_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() + 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 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, stdout="/usr/bin/python3\n") + + 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" + 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 + # 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" + 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" + + @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): + parser = ArgumentParser() + ServeCommand.register_subcommand(parser.add_subparsers()) + return parser.parse_args(["serve", *argv]) + + 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 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") + with pytest.raises(SystemExit, match="--context-parallel is not supported"): + ServeCommand(self._parse(["-m", "org/model", "--context-parallel"])).run()