From 6afdfd9f4e2e2381c5d5cced0dc4262e9e4267bd Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Wed, 22 Jul 2026 17:02:44 +0530 Subject: [PATCH] Update `diffusers-cli` for agentic use (#13966) * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * pdate * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update --------- Co-authored-by: Sayak Paul --- .ai/skills/custom-blocks/SKILL.md | 157 +++ .ai/skills/diffusers-cli/SKILL.md | 68 ++ .ai/skills/diffusers-cli/run.md | 305 +++++ docs/source/en/_toctree.yml | 2 + docs/source/en/using-diffusers/cli.md | 292 +++++ src/diffusers/commands/custom_blocks.py | 42 +- src/diffusers/commands/diffusers_cli.py | 28 +- src/diffusers/commands/env.py | 7 +- src/diffusers/commands/fp16_safetensors.py | 14 +- src/diffusers/commands/run.py | 1227 ++++++++++++++++++++ src/diffusers/commands/schema.py | 287 +++++ src/diffusers/commands/skills.py | 344 ++++++ tests/others/test_cli_commands.py | 306 +++++ tests/others/test_cli_skills.py | 103 ++ 14 files changed, 3161 insertions(+), 21 deletions(-) create mode 100644 .ai/skills/custom-blocks/SKILL.md create mode 100644 .ai/skills/diffusers-cli/SKILL.md create mode 100644 .ai/skills/diffusers-cli/run.md create mode 100644 docs/source/en/using-diffusers/cli.md create mode 100644 src/diffusers/commands/run.py create mode 100644 src/diffusers/commands/schema.py create mode 100644 src/diffusers/commands/skills.py create mode 100644 tests/others/test_cli_commands.py create mode 100644 tests/others/test_cli_skills.py diff --git a/.ai/skills/custom-blocks/SKILL.md b/.ai/skills/custom-blocks/SKILL.md new file mode 100644 index 0000000000..09ac19ddb2 --- /dev/null +++ b/.ai/skills/custom-blocks/SKILL.md @@ -0,0 +1,157 @@ +--- +name: custom-blocks +description: > + Use when the user has written (or wants to write) a `ModularPipelineBlocks` + subclass in a local Python file and needs to package it into a Hub-uploadable + directory. Covers the workflow from a single `block.py` file to a published + custom-block repo that consumers can load via + `ModularPipeline.from_pretrained(, trust_remote_code=True)`. +--- + +## What this skill is for + +A `ModularPipelineBlocks` subclass is a unit of pipeline logic — input/output spec plus a `__call__` — that +slots into diffusers' modular pipeline composition. Once you have one defined locally, you almost always want to +publish it as a small Hub repo so others can `from_pretrained` it. `diffusers-cli custom_blocks` automates the +packaging step: it parses your Python file, instantiates the chosen block class, and writes a +`save_pretrained`-style directory in your cwd that's ready to push to the Hub. + +Use this skill when: + +- The user is writing a custom modular block and asks "how do I publish this?" or "package this for the Hub". +- The user has a `block.py` (or similar) file with one or more `ModularPipelineBlocks` subclasses. +- You're scaffolding a new modular pipeline repo and need the on-disk layout that `ModularPipelineBlocks.from_pretrained` + expects. + +Don't use this skill for: running an existing modular pipeline (`diffusers-cli run`), introspecting one +(`diffusers-cli schema`), or writing the block class itself — this skill packages an *already-written* block. + +## The end-to-end workflow + +``` +[you: write block.py] → diffusers-cli custom_blocks → [packaged dir in cwd] + ↓ + hf upload . + ↓ + consumers: ModularPipeline.from_pretrained(, trust_remote_code=True) + diffusers-cli schema --model --trust-remote-code + diffusers-cli run --model --trust-remote-code ... +``` + +The skill covers the middle box. The bookends (writing the block and uploading) are out of scope. + +## Command surface + +```bash +diffusers-cli custom_blocks [--block_module_name ] [--block_class_name ] +``` + +### Flags + +- `--block_module_name ` — Python file containing the block class. Defaults to `block.py` in the cwd. +- `--block_class_name ` — Which class in the file to package. Optional: if omitted, the CLI parses the + file with `ast`, finds every class that inherits from `ModularPipelineBlocks`, and uses the first one (with + an info log naming the others). Specify explicitly when the file defines more than one block and you want a + specific one. + +### What it does + +1. **AST scan**: parses `` without executing it, walks top-level `ClassDef` nodes, and collects every + class whose `bases` include `ModularPipelineBlocks`. +2. **Pick a class**: uses `--block_class_name` if given, else the first found. Errors with the list of available + classes if your name doesn't match. +3. **Load and save**: imports the file via `importlib.util.spec_from_file_location` (this does execute the + module — make sure your block.py is something you trust to run), instantiates the chosen class with no + constructor args, and calls `.save_pretrained(os.getcwd())`. + +The result is a Hub-uploadable directory laid out the way `ModularPipelineBlocks.from_pretrained` expects: +your block source, an `auto_map` in the config so consumers know to load it with `trust_remote_code=True`, +and any artifacts `save_pretrained` writes for that block class. + +## End-to-end example + +Given a `block.py` like: + +```python +from diffusers.modular_pipelines import ModularPipelineBlocks, InputParam, OutputParam + +class MyDenoiseBlock(ModularPipelineBlocks): + model_name = "my-denoise" + + @property + def inputs(self): + return [ + InputParam("latents", type_hint="torch.Tensor", required=True, description="Noisy latents."), + InputParam("guidance_scale", type_hint="float", default=7.5), + ] + + @property + def intermediate_outputs(self): + return [OutputParam("latents", type_hint="torch.Tensor")] + + def __call__(self, components, state): + # ... denoising logic ... + return components, state +``` + +Package it: + +```bash +diffusers-cli custom_blocks --block_module_name block.py +``` + +Output in cwd: + +``` +./ +├── block.py +├── modular_config.json # contains auto_map → MyDenoiseBlock +└── (any state files MyDenoiseBlock.save_pretrained writes) +``` + +Upload to the Hub: + +```bash +hf upload my-user/my-denoise-block . +``` + +Consumers can now use it: + +```python +from diffusers import ModularPipeline +pipe = ModularPipeline.from_pretrained("my-user/my-denoise-block", trust_remote_code=True) +``` + +Or via CLI: + +```bash +diffusers-cli schema --model my-user/my-denoise-block --trust-remote-code +diffusers-cli run --model my-user/my-denoise-block --trust-remote-code \ + --pipeline-kwargs '{"latents": "...", "guidance_scale": 7.5}' +``` + +## Common errors + +- **`Could not parse '': SyntaxError`** — the file isn't valid Python. Fix the syntax; the AST step runs + before any execution. +- **`block_class_name could not be retrieved. Available classes from : [ClassA, ClassB]`** — your + `--block_class_name` doesn't match any `ModularPipelineBlocks` subclass found. Pick from the list shown. +- **No classes found**: silent — the command will try to use the first entry in an empty list and raise + `IndexError`. If you hit that, double-check your class actually inherits from `ModularPipelineBlocks` + (the AST scan looks for that literal base-class name; aliased imports like `from diffusers import ... + as MPB` won't be picked up). +- **Block requires constructor args**: the command calls `()` with no args. If your block needs + `__init__` parameters, refactor to take them from `state`/`components` at `__call__` time instead, or + hardcode defaults in `__init__`. + +## Verifying the install + +If `diffusers-cli` isn't on PATH, see the install verification section of +[`../diffusers-cli/SKILL.md`](../diffusers-cli/SKILL.md#verifying-the-cli-is-installed). + +## Related + +- [`diffusers-cli` skill](../diffusers-cli/SKILL.md) — once your block is uploaded, `schema`/`run` + let you call it from the terminal without writing Python. +- diffusers' [modular pipelines docs](../../../docs/source/en/modular_diffusers) — for writing the block + class itself. diff --git a/.ai/skills/diffusers-cli/SKILL.md b/.ai/skills/diffusers-cli/SKILL.md new file mode 100644 index 0000000000..dd9ffd66da --- /dev/null +++ b/.ai/skills/diffusers-cli/SKILL.md @@ -0,0 +1,68 @@ +--- +name: diffusers-cli +description: > + Use when the user wants to run a diffusers pipeline from a terminal (one-off + generation, batch jobs, smoke-testing a new model), run on HF Sandbox + hardware via `--remote`, introspect a pipeline's input schema before + calling it, or attach a LoRA at inference time. Prefer this over writing + ad-hoc Python scripts for generation tasks. +--- + +## Overview + +`diffusers-cli` is the shipped CLI in `src/diffusers/commands/`. Subcommands relevant to agentic use: + +| Command | Purpose | +| --------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `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. | +| `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. | + +## When to read which file + +Most agentic work goes through `run`. Read the matching reference file before constructing a command: + +- **[`run.md`](run.md)** — full reference for `diffusers-cli run`. Covers `--pipeline-kwargs` + semantics and the shell-quoting gotcha, LoRA via `--lora`, optimization flags (`--dtype`, `--cpu-offload`, + `--attention-backend`, `--vae-tiling/slicing`), output handling and `--push-to` bucket uploads, the full + `--remote` HF Jobs flow (image, container command, log streaming, timing payload, artifact download), and + context parallel (`--context-parallel`) for both local-torchrun and `--remote` paths. + +The other commands are small enough that `diffusers-cli --help` is the canonical reference: + +```bash +diffusers-cli schema --help +diffusers-cli custom_blocks --help +diffusers-cli env --help +``` + +## When NOT to use this skill + +- Multi-stage workflows where you need intermediate tensor manipulation between pipelines → write Python. +- Training or fine-tuning → CLI only covers inference. +- Anything requiring `quantization_config` or other low-level loader knobs not exposed by the CLI flags → write + Python. (`device_map` is exposed as `--device-map`; see [run.md](run.md#optimization-flags).) + +## Verifying the CLI is installed + +The console entry point is registered in `pyproject.toml` (`diffusers-cli = +"diffusers.commands.diffusers_cli:main"`). If `diffusers-cli` is not on PATH after `pip install -e .`, reinstall +with `pip install -e . --force-reinstall --no-deps` and check `which diffusers-cli`. If the installed binary is +missing recent features (e.g. you see `unrecognized arguments: --lora`), reinstall. + +## Output formats + +`--format {auto, human, agent, json}` (top-level flag, must appear before the subcommand): + +- **`human`** — plain-text indented output for terminals (default when not running under an agent harness). No ANSI color. +- **`agent`** — TSV tables and `key=value` lines. Auto-selected when an agent env var is present + (`CLAUDECODE`, `CLAUDE_CODE`, `CODEX_SANDBOX`, `CURSOR_AI`, `AIDER_AI_CONTEXT`, `GH_COPILOT_AGENT`, + `AI_AGENT`). Token-cheap for LLM agents to read. +- **`json`** — compact JSON. Use for programmatic parsing (scripts, services) where type fidelity and nested + structures matter. + +`stdout` carries data; `stderr` carries hints/warnings/progress — parseable output is never polluted. + +Rule of thumb: `--format json` for scripts that will `json.loads()` the output, otherwise leave it on +auto-detect (`agent` for LLMs, `human` for terminals). diff --git a/.ai/skills/diffusers-cli/run.md b/.ai/skills/diffusers-cli/run.md new file mode 100644 index 0000000000..5eb56d8265 --- /dev/null +++ b/.ai/skills/diffusers-cli/run.md @@ -0,0 +1,305 @@ +# `diffusers-cli run` — reference + +Full surface for `diffusers-cli run`. Use this file as the source of truth when constructing a `run` +invocation. The top-level [`SKILL.md`](SKILL.md) covers when to use the CLI; this file covers how. + +## The schema → run flow + +For any model you haven't called before, run `schema` first to learn its input contract, then `run` with +the right `--pipeline-kwargs`: + +```bash +# 1. Discover what kwargs the pipeline takes (no weight download) +diffusers-cli --format json schema --model black-forest-labs/FLUX.2-klein-9B + +# 2. Run it +diffusers-cli run \ + --model black-forest-labs/FLUX.2-klein-9B \ + --pipeline-kwargs '{"prompt": "Make the cats fur grey", "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"}' \ + --dtype bf16 +``` + +`schema --format json` emits a `{task, model, pipeline_class, inputs[]}` payload where each input is +`{name, type_hint, default, required, description}`. + +## Standard vs modular detection + +`run` auto-detects which kind of pipeline it's calling: + +1. If `model_index.json` exists on the repo → `DiffusionPipeline.from_pretrained` path. +2. Otherwise → `ModularPipeline.from_pretrained` path. + +You don't need to tell it which. Modular repos must pass `--trust-remote-code` if they ship custom block code. + +## `--pipeline-kwargs` semantics + +A JSON object passed straight through to `pipeline(**kwargs)`. String values at known media-input keys are +auto-loaded before the pipeline is called: + +- **Images** (`image`, `mask_image`, `control_image`, `ip_adapter_image`, `image_2`) → `PIL.Image.Image` + via `diffusers.utils.load_image`. Accepts URLs or local paths. +- **Videos** (`video`, `control_video`) → `list[PIL.Image.Image]` via `diffusers.utils.load_video`. +- **Audio** (`initial_audio_waveforms`, `reference_audio`, `src_audio`) → `torch.Tensor` via `torchaudio.load`. + For `initial_audio_waveforms`, the file's native sample rate is auto-written to + `initial_audio_sampling_rate` if you didn't pass it explicitly. + +```bash +diffusers-cli run \ + --model black-forest-labs/FLUX.2-klein-9B \ + --pipeline-kwargs '{"image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png", "prompt": "make the fur grey", "strength": 0.6}' +``` + +Media resolution runs **before** the pipeline weights load, so a dead URL or missing file fails within +seconds instead of after a multi-minute model download. + +**Batched inputs**: any media key (and `prompt`) accepts a JSON array — each string in the list is loaded +individually, and the pipeline processes the whole list in one forward pass on the GPU. Use this when you +have several inputs that share flavor/model/dtype and want them batched: + +```bash +diffusers-cli run --model black-forest-labs/FLUX.1-Kontext-dev \ + --pipeline-kwargs '{"prompt": ["make it grey", "make it pink", "make it blue"], + "image": ["https://.../cat1.png", "https://.../cat2.png", "https://.../cat3.png"]}' +``` + +**Shell-quoting gotcha**: the JSON must be on one line (or use `\` to line-continue). A literal newline inside the +single-quoted argument lands as a raw control char inside the string and breaks `json.loads`. + +## LoRA adapters (`--lora`) + +Attach one or more LoRAs after the pipeline loads via a JSON spec. `--lora` accepts a single object or a list. + +Single adapter: + +```bash +diffusers-cli run \ + --model black-forest-labs/FLUX.2-klein-9B \ + --pipeline-kwargs '{"prompt": "a tiny grey cat"}' \ + --lora '{"lora_id": "alvdansen/littletinies", "lora_scale": 0.8}' +``` + +Multiple stacked adapters: + +```bash +diffusers-cli run \ + --model black-forest-labs/FLUX.2-klein-9B \ + --pipeline-kwargs '{"prompt": "a tiny grey cat"}' \ + --lora '[ + {"lora_id": "alvdansen/littletinies", "lora_scale": 0.5, "adapter_name": "style"}, + {"lora_id": "author/detail-boost", "lora_scale": 0.3} + ]' +``` + +Per-entry fields: `lora_id` (required), `lora_scale` (optional, default `1.0`), `adapter_name` (optional; when +omitted while stacking, auto-generated as `lora_0`, `lora_1`, …; single-entry defaults to `"default"`). + +Each entry calls `pipeline.load_lora_weights(, adapter_name=)`. After all adapters load, one +`pipeline.set_adapters(names, adapter_weights=scales)` activates them together. + +## Optimization flags + +- `--dtype {auto, bf16, fp16, fp32, …}` — pipeline weight dtype. `bf16` is the right default for modern DiTs on + A100/H100. +- `--device-map ` — component placement, forwarded to `from_pretrained(device_map=...)`. Accepts a plain + torch device (`cuda`, `cuda:0`, `cpu`, `mps`), the string `balanced` (auto-splits pipeline components across + visible GPUs), or a JSON dict `{"transformer": "cuda:0", "vae": "cuda:1"}` for explicit per-component + placement. Auto-detects if omitted (pinned to `cuda:$LOCAL_RANK` under torchrun). `balanced` and dict values + are incompatible with `--cpu-offload`. +- `--cpu-offload {model, group}` — `model` uses `enable_model_cpu_offload`, `group` uses + `enable_group_offload(offload_type="leaf_level", use_stream=True)`. Use `group` to fit a 9B+ model on a single + A100. Onload target device comes from `--device-map` (must be a plain device string in this case). +- `--attention-backend {default, flash_hub, flash_varlen_hub, flash_4_hub, sage_hub}` — hub-hosted kernels, + auto-downloaded on first use. Failures (kernel not available, CUDA arch mismatch, network) raise a clear + `SystemExit` listing the alternatives instead of silently reverting to the default. Only supported on + transformer-based pipelines; UNet pipelines get a `logger.warning` and the flag is ignored. +- `--vae-tiling` / `--vae-slicing` — lower peak VAE decode VRAM. +- `--compile [JSON]` — `torch.compile` every denoiser submodule. See [Compile](#compile) below. +- `--context-parallel` — Ulysses-style context parallelism on a DiT. See [Context parallel](#context-parallel) below. + +## Output handling + +`run` detects the pipeline return type and saves accordingly: + +- `PIL.Image` / list of them → `.png` (zero-padded, e.g. `0000.png`) +- Frame sequence (≥2 PILs or ndarrays) → `0000.mp4` (uses `--fps`, default 8) +- Numpy audio array → `0000.wav` (uses `--sampling-rate`) +- Anything else → JSON dump + +**Default output directory** is `~/.diffusers/cli/run/outputs/diffusers-run--/`. +Each run gets its own subdirectory so consecutive invocations don't overwrite each other. The same run id +is used for the local dir, the sandbox's `DIFFUSERS_CLI_RUN_ID` env var, and the sandbox input/output paths +under `--remote`, so a run is traceable end-to-end. + +Override the destination with `--output ` (file or directory). Explicit `--output` bypasses the +`diffusers-run-*` namespace — files land flat in the path you gave. + +### `--push-to` + +Upload outputs to an HF bucket. Objects land under `//`. The bucket is created +if it doesn't exist. Behavior interacts with `--output` and `--remote`: + +| flags | download to local? | bucket write? | +|---|---|---| +| (neither) | ✅ `~/.diffusers/cli/run/outputs//` | — | +| `--output /path` | ✅ `/path/` | — | +| `--push-to my/bucket` | ❌ | ✅ (bucket only) | +| `--push-to my/bucket --output /path` | ✅ | ✅ (both) | + +The rule: explicit `--push-to` means "the bucket is my destination" — skip the local download unless the +user also explicitly asked for a local target via `--output`. + +## Remote execution (`--remote`) + +Add `--remote` to run the same call inside a [Hugging Face Sandbox](https://huggingface.co/docs/huggingface_hub/en/guides/sandbox) +— an isolated cloud VM (built on HF Jobs) the CLI drives over HTTP: it uploads inputs, installs deps, runs +the pipeline, downloads outputs, then terminates the sandbox. + +```bash +diffusers-cli run \ + --model black-forest-labs/FLUX.2-klein-9B \ + --pipeline-kwargs '{"prompt": "Make the cats fur grey", "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"}' \ + --remote --flavor a100-large \ + --dtype bf16 \ + --cpu-offload group +``` + +What happens: + +1. Your HF token is picked up (from `--token` or your login) and forwarded into the sandbox as `HF_TOKEN`. +2. `--pipeline-kwargs` are parsed locally so JSON errors fail fast (no wasted sandbox time). +3. A dedicated sandbox is created on `--flavor` from a pytorch image + (`pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime` by default) that already has torch + CUDA. Any local + file paths in `--pipeline-kwargs` are uploaded into the sandbox under `/tmp/diffusers-cli/inputs//` + via native file transfer (no bucket), and the JSON paths are rewritten to point at them. +4. The small Python deps (`diffusers`, `accelerate`, `transformers`, `safetensors`, `sentencepiece`, `ftfy`) + are installed with `uv pip install --system`. Output (install + run) streams live to your terminal. +5. The sandbox CLI writes outputs to `/tmp/diffusers-cli/outputs/`; the CLI downloads every artifact back into + the local target (see [`--push-to`](#-push-to) for when the download is skipped). +6. The sandbox is terminated (unless `--keep-alive`/`--sandbox-id`), and the wallclock `run_seconds` for the + pipeline command is printed and included in the JSON payload. + +Flags: + +- `--flavor ` — sandbox hardware (e.g. `a10g-small`, `a100-large`, `4xa100-large`). +- `--timeout ` — max wallclock for the run command inside the sandbox (e.g. `30m`, `2h`). Defaults to `10m`. +- `--dependencies ` — extra pip deps (repeatable). Appends to the defaults. +- `--namespace ` — create the sandbox under a different account. +- `--push-to ` — see [`--push-to`](#-push-to) above. The upload runs inside the sandbox; an explicit + value with no `--output` makes the bucket the sole destination and skips the local download. +- `--image ` — override the sandbox image. Must ship torch + CUDA; the CLI installs the small Python + deps on top via `uv pip install --system`. Useful for pinning a specific torch or bundling extra system libs. +- `--volume [:]` — mount an HF storage bucket into the sandbox as a read-write directory. + Repeatable. Default mount is `/mnt/buckets/`. Reference mounted files from `--pipeline-kwargs` + like any other local path — no upload happens, the container reads straight from the FUSE mount. Applied + only on new sandbox creation — ignored when reconnecting via `--sandbox-id`. Enables batch workflows: + loop over bucket contents on the host, one `--sandbox-id` call per input. +- `--keep-alive` — don't terminate the sandbox after the run; its id is printed so a later run can reconnect. +- `--sandbox-id ` — reconnect to a kept-alive sandbox instead of creating a new one. Implies `--keep-alive`. +- `--idle-timeout ` — auto-shutdown after this much inactivity (default 10m). The billing backstop + for kept-alive sandboxes. Only applied on new sandbox creation — ignored when reconnecting via `--sandbox-id`. + +### Reusing a warm sandbox (`--keep-alive` / `--sandbox-id`) + +By default each `--remote` run is ephemeral: create → run → download → kill. The cold cost (dep install + +model weight download + any compile) is paid every time. For iterative work, keep one sandbox warm and reuse +it — deps, the HF weight cache, and the `torch.compile` cache all survive on its disk: + +```bash +# First run: create + keep alive. Prints sandbox_id=. +diffusers-cli run -m black-forest-labs/FLUX.2-klein-9B --dtype bf16 \ + --pipeline-kwargs '{"prompt": "a cat"}' --remote --flavor a100-large --keep-alive + +# Subsequent runs: reconnect — model already cached, only inference runs. +diffusers-cli run -m black-forest-labs/FLUX.2-klein-9B --dtype bf16 \ + --pipeline-kwargs '{"prompt": "a dog"}' --remote --flavor a100-large --sandbox-id + +# Stop it when done (or let --idle-timeout reap it). +hf sandbox kill +``` + +Notes on `--remote` argv forwarding: flags that control how the run is dispatched (`--flavor`, `--timeout`, +`--namespace`, `--dependencies`, `--image`, `--keep-alive`, `--sandbox-id`, `--idle-timeout`, `--format`) are +stripped before the argv is rebuilt for the sandbox CLI. Everything else — model, dtype, pipeline-kwargs, +optimizations, output flags — is forwarded verbatim (`--output` is repointed at the sandbox output dir). + +## Compile + +`--compile` runs `torch.compile` over every `transformer*` / `unet*` submodule on the pipeline. Prefers +regional compilation via `module.compile_repeated_blocks(**kwargs)` when the model exposes `_repeated_blocks` +— this only compiles the repeated inner blocks (the bulk of the compute) rather than the whole module, so +first-step latency is much lower. Falls back to full `torch.compile(module, **kwargs)` when no regional +metadata is declared. + +```bash +# Bare — uses fullgraph=true +diffusers-cli run --model --dtype bf16 --compile --pipeline-kwargs '...' + +# With kwargs forwarded to torch.compile +diffusers-cli run --model --dtype bf16 \ + --compile '{"mode": "max-autotune", "fullgraph": true}' \ + --pipeline-kwargs '...' +``` + +**When it's worth it**: multi-step generation (~50+ denoising steps). You pay a one-time compilation cost on +the first step, then every subsequent step is faster. + +**Under `--remote`**: on an ephemeral sandbox the compile cache doesn't survive, so the compilation cost is +paid on every run — `--compile` only breaks even on very long generations. On a `--keep-alive`/`--sandbox-id` +sandbox the compile cache persists on disk, so the cost is paid once and reused: keep a warm sandbox and +`--compile` pays off across runs. + +`--compile` is **currently not supported with `--context-parallel`** — CP shards attention across ranks while +regional compile assumes a stable single-device graph. If both are set, the CLI logs a warning and skips the +compile step; CP still runs. + +## Context parallel + +`--context-parallel` enables Ulysses CP on a DiT-based pipeline. **Locally** the user must launch via torchrun: + +```bash +torchrun --nproc-per-node=2 -m diffusers.commands.diffusers_cli run \ + --model black-forest-labs/FLUX.2-klein-9B \ + --pipeline-kwargs '{"prompt": "Make the cats fur grey"}' \ + --dtype bf16 \ + --context-parallel +``` + +**Remotely** the CLI handles the torchrun wrapping — just pass `--context-parallel` to a `--remote` invocation on +a multi-GPU flavor: + +```bash +diffusers-cli run \ + --model black-forest-labs/FLUX.2-klein-9B \ + --pipeline-kwargs '{"prompt": "Make the cats fur grey", "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"}' \ + --remote --flavor 4xa100-large \ + --dtype bf16 \ + --context-parallel +``` + +Inside the sandbox, CP swaps the entrypoint to `torchrun --nproc-per-node=gpu -m +diffusers.commands.diffusers_cli`, initializes a hybrid process group (`cpu:gloo,cuda:nccl` — NCCL for the +attention all-to-all, Gloo for `ulysses_anything`'s per-rank size coordination), pins each rank to +`cuda:{LOCAL_RANK}`, and gates output saving/printing to rank 0 only. + +**Memory note**: CP shards the sequence, **not the weights**. Every rank still holds the full transformer. Wins +are wall-clock attention speedup and headroom for very long sequences, not "fit a model that doesn't fit." For +weight sharding you'd want TP or FSDP — not exposed in the CLI yet. + +CP is DiT-only. UNet pipelines raise a clear error directing you to a DiT pipeline (FLUX, SD3, HunyuanDiT, +AuraFlow, …). + +## Output mode (`--format`) + +`--format` controls the shape of **stdout metadata** (which paths were written, timing, sandbox id, pushed +bucket URLs) — **not** the media file format. Written images are always PNG, videos MP4, audio WAV; only the +summary printed alongside them changes shape. + +The CLI auto-detects when running under an AI coding agent (Claude Code, Cursor, Aider, GH Copilot Agent — via +`CLAUDECODE`, `CLAUDE_CODE`, `CURSOR_AI`, `AIDER_AI_CONTEXT`, `GH_COPILOT_AGENT`) and switches to **agent +mode** automatically — TSV tables, `key=value` results, compact JSON dicts, no progress bars. + +Override explicitly with `--format {auto, human, agent, json, quiet}` placed **before** the subcommand: + +```bash +diffusers-cli --format json run --model --pipeline-kwargs '...' +``` diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 23826ba0e2..98773eab88 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -220,6 +220,8 @@ title: Task recipes - local: using-diffusers/write_own_pipeline title: Understanding pipelines, models and schedulers + - local: using-diffusers/cli + title: Command line interface - local: community_projects title: Projects built with Diffusers - local: conceptual/philosophy diff --git a/docs/source/en/using-diffusers/cli.md b/docs/source/en/using-diffusers/cli.md new file mode 100644 index 0000000000..4a8127888c --- /dev/null +++ b/docs/source/en/using-diffusers/cli.md @@ -0,0 +1,292 @@ + + +# Command line interface + +`diffusers-cli` is a command line client for running, inspecting, and packaging Diffusers pipelines. + +## Available commands + +| Command | Purpose | +|---|---| +| [`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. | +| [`custom_blocks`](#custom_blocks) | Package a local `ModularPipelineBlocks` subclass for the Hub. | +| [`fp16_safetensors`](#fp16_safetensors) | Convert a checkpoint to fp16 `.safetensors`. | +| [`skills`](#skills) | Install pre-authored skill bundles into your AI coding agent. | + +> [!TIP] +> This page does not provide details for all options under each subcommand. For the full, always-current list of options for any subcommand, run `diffusers-cli --help` (`diffusers-cli run --help`). + +## `env` + +Prints Python/PyTorch/Diffusers versions, CUDA info, and installed optional deps. Use it when opening an +issue so maintainers can reproduce your setup. + +```bash +diffusers-cli env +``` + +## `schema` + +Returns the pipeline's accepted inputs without downloading weights. This is useful when building `--pipeline-kwargs` + +Only the index file is fetched. Standard pipelines read `model_index.json`; modular pipelines read +`modular_model_index.json`; custom-block repos read `modular_config.json` and need `--trust-remote-code` since +loading them runs code from the Hub. + +```bash +diffusers-cli --format json schema --model black-forest-labs/FLUX.1-dev +diffusers-cli schema --model my-org/my-custom-blocks --trust-remote-code +``` + +## `run` + +Run a pipeline end-to-end. Auto-detects standard vs modular repos, auto-loads media inputs from URLs or local +paths, saves outputs by detecting the pipeline's return type, and can run remotely on a Hugging Face +Sandbox via `--remote`. + +Minimal example: + +```bash +diffusers-cli run \ + --model black-forest-labs/FLUX.1-dev \ + --dtype bf16 \ + --pipeline-kwargs '{"prompt": "an astronaut riding a horse"}' +``` + +### Passing pipeline arguments + +`--pipeline-kwargs` takes a JSON object that's forwarded to `pipeline(**kwargs)`. String values at known +media-input keys are auto-loaded: + +- Images (`image`, `mask_image`, `control_image`, `ip_adapter_image`, `image_2`) → `PIL.Image` via + `load_image`. +- Videos (`video`, `control_video`) → `list[PIL.Image]` via `load_video`. +- Audio (`initial_audio_waveforms`, `reference_audio`, `src_audio`) → `torch.Tensor` via `torchaudio.load`. + +```bash +diffusers-cli run \ + --model black-forest-labs/FLUX.2-klein-9B --dtype bf16 \ + --pipeline-kwargs '{"prompt": "make the fur grey", "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png", "strength": 0.6}' +``` + +Both media keys and text keys accept a JSON array to run a batch through a single pipeline call. Each entry +in a media list is loaded individually (URL, local path, or bucket-mount path), and diffusers processes the +whole list in one forward pass on the GPU: + +```bash +diffusers-cli run \ + --model black-forest-labs/FLUX.1-Kontext-dev --dtype bf16 \ + --pipeline-kwargs '{ + "prompt": ["make it grey", "make it pink", "make it blue"], + "image": [ + "https://.../cat1.png", + "https://.../cat2.png", + "https://.../cat3.png" + ] + }' +``` + +### Loading + +Configure how the CLI loads model weights and custom pipeline code. + +- `--dtype {auto, bfloat16, bf16, float16, fp16, float32, fp32}` — weight dtype. +- `--device-map ` — component placement. Accepts a torch device string (`cuda`, `cuda:0`, `cpu`, `mps`), + `balanced` (auto-splits components across visible GPUs), or a JSON dict for explicit per-component placement. + Auto-detected if omitted. See [device_map](../training/distributed_inference#device_map) for more details +- `--variant fp16` — pick a weight variant. +- `--revision ` — pin a specific model revision. +- `--trust-remote-code` — allow custom code from the Hub (required for repos that ship custom pipeline classes + or modular blocks). See [Community pipelines](../using-diffusers/custom_pipeline_overview) for standard custom + pipelines and [Modular Diffusers](../modular_diffusers/overview). +- `--lora ` — attach a LoRA adapter after loading. Each value is a JSON dict; repeat the flag to + stack multiple adapters. `lora_id` is required per entry; `lora_scale` defaults to `1.0`; + `adapter_name` is optional (auto-generated as `lora_` when stacking). + - Single: `--lora '{"lora_id": "alvdansen/flux-koda", "lora_scale": 0.8}'` + - Multiple: `--lora '{"lora_id": "alvdansen/flux-koda", "lora_scale": 0.6, "adapter_name": "koda"}' --lora '{"lora_id": "Shakker-Labs/FLUX.1-dev-LoRA-AntiBlur", "lora_scale": 0.4}'` + + All specs are loaded by `pipeline.load_lora_weights(...)`, then activated together with a single + `pipeline.set_adapters(names, adapter_weights=scales)` call. See [LoRA](../tutorials/using_peft_for_inference) + for a deeper walkthrough of adapter stacking, scale scheduling, and hotswapping. + +### Optimizations + +- `--cpu-offload {model, group}` — `model` calls `enable_model_cpu_offload`; `group` calls + `enable_group_offload(offload_type="leaf_level", use_stream=True)`. Onload target comes from `--device-map` + (which must be a plain device string for offload). See + [Model offloading](../optimization/memory#model-offloading) and + [Group offloading](../optimization/memory#group-offloading). +- `--attention-backend {default, flash_hub, flash_varlen_hub, flash_4_hub, sage_hub}` — Hub-hosted attention + kernels, auto-downloaded on first use. Transformer-based pipelines only; 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`. +- `--context-parallel` — Ulysses-style context parallelism on a DiT-based pipeline. Locally requires torchrun; + under `--remote` the CLI wraps `torchrun --nproc-per-node=gpu` for you. See + [Context parallelism](../training/distributed_inference#context-parallelism). + +### Outputs + +`run` detects the pipeline output type: + +- `PIL.Image`/list → `.png` (zero-padded index, e.g. `0000.png`) +- Image sequence → `0000.mp4` (`--fps` controls framerate, default 8) +- Audio array → `0000.wav` (`--sampling-rate` controls rate) +- Anything else → JSON dump + +The default output directory format is `~/.diffusers/cli/run/outputs/diffusers-run--/`. Each +run gets its own subdirectory so consecutive invocations don't overwrite. + +Override with `--output `. How the path expands depends on its shape and the batch size: + +| `--output` | 1 output | N outputs | +|---|---|---| +| _omitted_ | default dir → `0000.png` | default dir → `0000.png`, `0001.png`, `0002.png`, … | +| `./results/` (trailing `/` or an existing directory) | `./results/0000.png` | `./results/0000.png`, `./results/0001.png`, … | +| `my-cat.png` (file path) | `my-cat.png` (used verbatim) | `my-cat-0000.png`, `my-cat-0001.png`, … | + +Directory outputs always use bare padded names (`0000`, `0001`, …). Explicit file paths preserve your chosen +stem and get the padded index appended when the batch produces multiple outputs. + +Use `--push-to` to upload outputs to a +[Hugging Face storage bucket](https://huggingface.co/docs/hub/en/storage-buckets). Accepts an HF bucket +id (`/`), an `hf://buckets//[/]` +[HF URI](https://huggingface.co/docs/huggingface_hub/main/en/package_reference/hf_uris), or a browser URL +for the same — a subpath is used as a folder prefix. The bucket is created if missing; objects land under +`[/]/`. + +```bash +# HF bucket id — files at hf://buckets/alice/edit-outputs// +--push-to alice/edit-outputs + +# URI with subpath — files at hf://buckets/alice/edit-outputs/greyscale/2026-07// +--push-to hf://buckets/alice/edit-outputs/greyscale/2026-07 + +# Browser URL copy-paste from the Hub also works. +--push-to https://huggingface.co/buckets/alice/edit-outputs/tree/greyscale/2026-07 +``` + +The table below describes remote runs. For local runs, `--push-to` uploads the locally saved output. It does +not suppress local file creation. + +| `--push-to` set? | `--output` set? | Result | +|---|---|---| +| no | no | download to default local dir | +| no | yes | download to `--output` | +| yes | no | bucket only, no local download | +| yes | yes | bucket AND `--output` | + +`--format` shapes the stdout metadata (paths, timing, sandbox info) — it does not change the file format of +the media itself. Written images are always PNG, videos MP4, audio WAV. + +### Remote execution (`--remote`) + +Run the same call inside a [Hugging Face Sandbox](https://huggingface.co/docs/huggingface_hub/en/guides/sandbox) +— an isolated cloud VM the CLI drives over HTTP: it uploads inputs, installs deps, runs the pipeline, downloads +outputs, then terminates the sandbox. Requires `huggingface_hub>=1.23`. + +```bash +diffusers-cli run \ + --model black-forest-labs/FLUX.1-dev --dtype bf16 \ + --pipeline-kwargs '{"prompt": "an astronaut riding a horse"}' \ + --remote --flavor a100-large +``` + +Remote flags: + +- `--flavor ` — sandbox hardware (e.g. `a10g-small`, `h200`, `rtx-pro-6000`). +- `--timeout ` — max wallclock for the run command inside the sandbox (default `10m`). +- `--dependencies ` — extra pip deps (repeatable). Useful for pinning a diffusers branch tarball or + adding pipeline-specific extras. +- `--namespace ` — create the sandbox under a different HF org/account. +- `--image ` — override the sandbox image. Must ship torch + CUDA compatible with your `--flavor`'s + driver. +- `--volume [:]` — mount an [HF storage bucket](https://huggingface.co/docs/hub/en/storage-buckets) + into the sandbox as a read-write directory. Repeatable. Default mount path is + `/mnt/buckets/`. Reference mounted files from `--pipeline-kwargs` like any other local path. + Applied only on new sandbox creation — ignored when reconnecting via `--sandbox-id`. + +By default each `--remote` run is ephemeral (create → run → download → kill). To reuse a warm sandbox across +runs — keeping deps, the model weight cache, and the `torch.compile` cache on its disk — keep it alive and +reconnect: + +- `--keep-alive` — don't terminate the sandbox after the run; its id is printed. +- `--sandbox-id ` — reconnect to a kept-alive sandbox instead of creating a new one. +- `--idle-timeout ` — auto-shutdown after this much inactivity (default `10m`). Applied only on new sandbox creation — ignored when reconnecting via `--sandbox-id`. + +```bash +# First run keeps the sandbox alive and prints sandbox_id=. +diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \ + --pipeline-kwargs '{"prompt": "a cat"}' --remote --flavor a100-large --keep-alive + +# Reconnect for the next run — the model is already cached, so only inference runs. +diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \ + --pipeline-kwargs '{"prompt": "a dog"}' --remote --flavor a100-large --sandbox-id + +# Stop it when done (or let it timeout). +hf sandbox kill +``` + +## `custom_blocks` + +Package a local `ModularPipelineBlocks` subclass for upload to the Hub. Reads a Python file, AST-scans it for +subclasses of `ModularPipelineBlocks`, instantiates the chosen one, and calls `save_pretrained` in the current +working directory. + +```bash +# Package the first block found in ./block.py +diffusers-cli custom_blocks + +# Point at a different file / pick a specific class +diffusers-cli custom_blocks --block_module_name my_block.py --block_class_name MyDenoiseBlock +``` + +The block class must be instantiable with zero constructor args — hardcode defaults in `__init__` or read +config from the pipeline `state` at call time. + +## `fp16_safetensors` + +Convert a checkpoint on the Hub to fp16 `.safetensors` and push the result. Useful for shrinking a repo's +weight size for faster loading. See `diffusers-cli fp16_safetensors --help` for the exact args. + +## `skills` + +Install skills from the diffusers repo ([`.ai/skills/`](https://github.com/huggingface/diffusers/tree/main/.ai/skills)). + +```bash +# Install a single skill +diffusers-cli skills add "" + +# Install every skill in the registry +diffusers-cli skills add --all + +# List available skills +diffusers-cli skills list + +# Preview a skill's SKILL.md without installing +diffusers-cli skills preview diffusers-cli + +# Refetch and reinstall every managed skill +diffusers-cli skills update + +# Install to the user-level directory instead of the current project +diffusers-cli skills add diffusers-cli --global +``` + + diff --git a/src/diffusers/commands/custom_blocks.py b/src/diffusers/commands/custom_blocks.py index 25dd2cfc0b..7ebaf785ba 100644 --- a/src/diffusers/commands/custom_blocks.py +++ b/src/diffusers/commands/custom_blocks.py @@ -12,9 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" -Usage example: - TODO +"""`diffusers-cli custom_blocks` — package a local `ModularPipelineBlocks` subclass for the Hub. + +Parses `block.py` (or `--block_module_name`), instantiates the chosen block, and calls `save_pretrained` in the current +working directory. """ import ast @@ -28,7 +29,6 @@ from . import BaseDiffusersCLICommand EXPECTED_PARENT_CLASSES = ["ModularPipelineBlocks"] -CONFIG = "config.json" def conversion_command_factory(args: Namespace): @@ -38,7 +38,27 @@ def conversion_command_factory(args: Namespace): class CustomBlocksCommand(BaseDiffusersCLICommand): @staticmethod def register_subcommand(parser: ArgumentParser): - conversion_parser = parser.add_parser("custom_blocks") + from argparse import RawDescriptionHelpFormatter + + epilog = ( + "Examples\n" + " $ diffusers-cli custom_blocks\n" + " $ diffusers-cli custom_blocks --block_module_name my_block.py\n" + " $ diffusers-cli custom_blocks --block_module_name my_block.py --block_class_name MyDenoiseBlock\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" + ) + + conversion_parser = parser.add_parser( + "custom_blocks", + help="Package a local ModularPipelineBlocks subclass for the Hub.", + usage="\n diffusers-cli custom_blocks [options]", + epilog=epilog, + formatter_class=RawDescriptionHelpFormatter, + ) + conversion_parser._optionals.title = "Options" conversion_parser.add_argument( "--block_module_name", type=str, @@ -85,11 +105,6 @@ class CustomBlocksCommand(BaseDiffusersCLICommand): spec.loader.exec_module(module) getattr(module, child_class)().save_pretrained(os.getcwd()) - # or, we could create it manually. - # automap = self._create_automap(parent_class=parent_class, child_class=child_class) - # with open(CONFIG, "w") as f: - # json.dump(automap, f) - def _choose_block(self, candidates, chosen=None): for cls, base in candidates: if cls == chosen: @@ -108,10 +123,8 @@ class CustomBlocksCommand(BaseDiffusersCLICommand): if not isinstance(node, ast.ClassDef): continue - # extract all base names for this class base_names = [bname for b in node.bases if (bname := self._get_base_name(b)) is not None] - # for each allowed base that appears in the class's bases, emit a tuple for allowed in EXPECTED_PARENT_CLASSES: if allowed in base_names: results.append((node.name, allowed)) @@ -125,8 +138,3 @@ class CustomBlocksCommand(BaseDiffusersCLICommand): val = self._get_base_name(node.value) return f"{val}.{node.attr}" if val else node.attr return None - - def _create_automap(self, parent_class, child_class): - module = str(self.block_module_name).replace(".py", "").rsplit(".", 1)[-1] - auto_map = {f"{parent_class}": f"{module}.{child_class}"} - return {"auto_map": auto_map} diff --git a/src/diffusers/commands/diffusers_cli.py b/src/diffusers/commands/diffusers_cli.py index a64a0cac94..0e4f2c27fb 100644 --- a/src/diffusers/commands/diffusers_cli.py +++ b/src/diffusers/commands/diffusers_cli.py @@ -15,23 +15,47 @@ from argparse import ArgumentParser +from huggingface_hub.cli._output import OutputFormat, out + from .custom_blocks import CustomBlocksCommand from .env import EnvironmentCommand from .fp16_safetensors import FP16SafetensorsCommand +from .run import RunCommand +from .schema import SchemaCommand +from .skills import SkillsCommand def main(): - parser = ArgumentParser("Diffusers CLI tool", usage="diffusers-cli []") - commands_parser = parser.add_subparsers(help="diffusers-cli command helpers") + parser = ArgumentParser( + prog="diffusers-cli", + usage="\n diffusers-cli [--format ] [options]", + ) + parser._optionals.title = "Options" + parser.add_argument( + "--format", + choices=[m.value for m in OutputFormat], + default=OutputFormat.auto.value, + help=( + "Output format. 'auto' (default) picks 'agent' when an AI coding agent is detected " + "(via CLAUDECODE/CURSOR_AI/AIDER_AI_CONTEXT/... env vars) and 'human' otherwise. " + "Must appear before the subcommand." + ), + ) + commands_parser = parser.add_subparsers(title="Commands", metavar="") # Register commands EnvironmentCommand.register_subcommand(commands_parser) FP16SafetensorsCommand.register_subcommand(commands_parser) CustomBlocksCommand.register_subcommand(commands_parser) + RunCommand.register_subcommand(commands_parser) + SchemaCommand.register_subcommand(commands_parser) + SkillsCommand.register_subcommand(commands_parser) # Let's go args = parser.parse_args() + out.set_mode(OutputFormat(args.format)) + if not hasattr(args, "func"): parser.print_help() exit(1) diff --git a/src/diffusers/commands/env.py b/src/diffusers/commands/env.py index 29ae8bc493..cbd2d11138 100644 --- a/src/diffusers/commands/env.py +++ b/src/diffusers/commands/env.py @@ -54,7 +54,12 @@ def info_command_factory(_): class EnvironmentCommand(BaseDiffusersCLICommand): @staticmethod def register_subcommand(parser: ArgumentParser) -> None: - download_parser = parser.add_parser("env") + download_parser = parser.add_parser( + "env", + help="Print versions of diffusers and its dependencies (for bug reports).", + usage="\n diffusers-cli env", + ) + download_parser._optionals.title = "Options" download_parser.set_defaults(func=info_command_factory) def run(self) -> dict: diff --git a/src/diffusers/commands/fp16_safetensors.py b/src/diffusers/commands/fp16_safetensors.py index 6c46688dbf..ec91bba357 100644 --- a/src/diffusers/commands/fp16_safetensors.py +++ b/src/diffusers/commands/fp16_safetensors.py @@ -33,6 +33,13 @@ from . import BaseDiffusersCLICommand def conversion_command_factory(args: Namespace): + warnings.warn( + "`diffusers-cli fp16_safetensors` is deprecated and will be removed in a future version. " + "Convert weights to fp16 safetensors directly with `safetensors.torch.save_file` or via " + "`pipeline.save_pretrained(..., safe_serialization=True, variant='fp16')`.", + FutureWarning, + stacklevel=2, + ) if args.use_auth_token: warnings.warn( "The `--use_auth_token` flag is deprecated and will be removed in a future version." @@ -44,7 +51,12 @@ def conversion_command_factory(args: Namespace): class FP16SafetensorsCommand(BaseDiffusersCLICommand): @staticmethod def register_subcommand(parser: ArgumentParser): - conversion_parser = parser.add_parser("fp16_safetensors") + conversion_parser = parser.add_parser( + "fp16_safetensors", + help="[DEPRECATED] Convert a Hub checkpoint's weights to fp16 safetensors and push back as a PR.", + usage="\n diffusers-cli fp16_safetensors [options]", + ) + conversion_parser._optionals.title = "Options" conversion_parser.add_argument( "--ckpt_id", type=str, diff --git a/src/diffusers/commands/run.py b/src/diffusers/commands/run.py new file mode 100644 index 0000000000..9cd6385478 --- /dev/null +++ b/src/diffusers/commands/run.py @@ -0,0 +1,1227 @@ +# 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 run` — single agentic entry point. + +Runs any diffusers pipeline (standard or modular) by forwarding `--pipeline-kwargs` verbatim, saves the output by +detecting its runtime type, and can submit the same call to an HF Sandbox via `--remote`. +""" + +from __future__ import annotations + +import json +import os +import sys +from argparse import ArgumentParser, Namespace, _SubParsersAction +from pathlib import Path +from typing import Any + +from huggingface_hub.cli._output import out + +from diffusers.models.attention_dispatch import _HUB_KERNELS_REGISTRY +from diffusers.utils import load_image, load_video, logging + +from . import BaseDiffusersCLICommand + + +logger = logging.get_logger("diffusers-cli/run") + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +DEFAULT_OUTPUT_DIR = str(Path.home() / ".diffusers" / "cli" / "run" / "outputs") +DTYPE_CHOICES = ("auto", "float16", "fp16", "bfloat16", "bf16", "float32", "fp32") +CPU_OFFLOAD_CHOICES = ("model", "group") + + +ATTENTION_BACKEND_CHOICES = ("default", *sorted(b.value for b in _HUB_KERNELS_REGISTRY)) + +# 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 +# `diffusers.utils.load_video` → list[PIL.Image.Image]. +_IMAGE_INPUT_KEYS = ( + "image", + "mask_image", + "control_image", + "ip_adapter_image", + "image_2", +) +_VIDEO_INPUT_KEYS = ( + "video", + "control_video", +) +_AUDIO_INPUT_KEYS = ( + "initial_audio_waveforms", + "reference_audio", + "src_audio", +) + +# Pipeline attribute prefixes that identify a denoiser submodule. Matches base names +# (`transformer`, `unet`) and their numbered variants (`transformer_2`, etc.). +_DENOISER_COMPONENT_KEYS = ("transformer", "unet") + +_DEFAULT_REMOTE_DEPS = ( + "diffusers", + "accelerate", + "transformers", + "safetensors", + "sentencepiece", # required by several text-encoder tokenizers (T5, LLaMA, …) + "ftfy", # required by older CLIP text-encoder paths +) + +# Base sandbox image — provides torch + CUDA so `uv pip install --system` +# only has to add the small Python deps. cuda12.8 is the highest cuda12.x tag +# below the HF Jobs host driver's CUDA 12.9 max. +_DEFAULT_REMOTE_IMAGE = "pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime" + +# Installed console-script name invoked inside the sandbox after the deps land. +_CONTAINER_CLI_BINARY = "diffusers-cli" + +# Working directories inside the sandbox: local media from `--pipeline-kwargs` is uploaded +# under _SANDBOX_INPUTS_DIR, and the sandbox CLI is told to write its outputs under +# _SANDBOX_OUTPUTS_DIR so we can download them back afterwards. +_SANDBOX_INPUTS_DIR = "/tmp/diffusers-cli/inputs" +_SANDBOX_OUTPUTS_DIR = "/tmp/diffusers-cli/outputs" + +RUN_ID_ENV = "DIFFUSERS_CLI_RUN_ID" + +# Namespace keys that control *how* a remote run is dispatched, not what the sandbox CLI +# runs. They are stripped when forwarding argv to the sandbox. +REMOTE_KEYS = frozenset( + { + "remote", + "flavor", + "timeout", + "dependencies", + "namespace", + "image", + "keep_alive", + "sandbox_id", + "idle_timeout", + "volume", + "func", + "format", # top-level --format is a local rendering flag; never forward to the sandbox + } +) + + +# --------------------------------------------------------------------------- +# Argparse helpers +# --------------------------------------------------------------------------- + + +def _add_loading_arguments(parser: ArgumentParser) -> None: + parser.add_argument("--model", "-m", required=True, help="Model id on the Hugging Face Hub or local path.") + parser.add_argument( + "--device-map", + default=None, + help=( + "Component placement. Accepts a torch device string (`cuda`, `cuda:0`, `cpu`, `mps`), " + "`balanced` for pipeline-level auto-split across visible GPUs, or a JSON dict of " + '`{"": }` for explicit per-component placement. Auto-detected if omitted.' + ), + ) + parser.add_argument("--dtype", default="auto", choices=DTYPE_CHOICES, help="Torch dtype for pipeline weights.") + parser.add_argument("--variant", default=None, help='Optional weight variant (e.g. "fp16").') + parser.add_argument("--revision", default=None, help="Model revision (branch, tag, or commit SHA).") + parser.add_argument("--token", default=None, help="Hugging Face token for gated/private models.") + parser.add_argument("--trust-remote-code", action="store_true", help="Allow custom code from the Hub.") + parser.add_argument( + "--lora", + action="append", + default=None, + metavar="JSON", + help=( + "JSON dict describing a LoRA adapter to attach after the pipeline loads. Repeat to stack " + 'multiple adapters. Format: \'{"lora_id": "", "lora_scale": }\'. `lora_scale` ' + "defaults to 1.0; `adapter_name` is optional (auto-generated as `lora_` when stacking)." + ), + ) + + +def _add_optimization_arguments(parser: ArgumentParser) -> None: + parser.add_argument( + "--cpu-offload", + choices=CPU_OFFLOAD_CHOICES, + default=None, + help=( + "Offload pipeline components to CPU during inference. " + "'model' uses enable_model_cpu_offload, " + "'group' uses pipeline.enable_group_offload(leaf_level, use_stream=True)." + ), + ) + parser.add_argument( + "--attention-backend", + choices=ATTENTION_BACKEND_CHOICES, + default="default", + help=( + "Override the attention backend on the transformer/UNet. " + "Only Hub-hosted kernels are exposed — they auto-download on first use." + ), + ) + parser.add_argument("--vae-tiling", action="store_true", help="Enable VAE tiling (lower peak VRAM).") + parser.add_argument("--vae-slicing", action="store_true", help="Enable VAE slicing (lower peak VRAM).") + parser.add_argument( + "--context-parallel", + action="store_true", + help=( + "Enable Ulysses-style context parallelism (ulysses_anything mode). " + "Requires a DiT-based pipeline and launching the CLI under torchrun with ≥2 GPUs." + ), + ) + parser.add_argument( + "--compile", + nargs="?", + const='{"fullgraph": true}', + default=None, + metavar="JSON", + help=( + "torch.compile every denoiser submodule on the pipeline. 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)." + ), + ) + + +def _add_output_arguments(parser: ArgumentParser) -> None: + parser.add_argument( + "--output", + "-o", + default=None, + help=( + "Output file or directory. Defaults to " + "~/.diffusers/cli/run/outputs/diffusers-run--/.." + ), + ) + parser.add_argument( + "--push-to", + default=None, + help=( + "Upload the generated files to this HF bucket after saving (created if missing). Accepts " + "an HF bucket id (`/`), an `hf://buckets//[/]` " + "URI, or a browser URL for the same — a subpath is used as a folder prefix. Under --remote " + "the upload runs inside the sandbox; without an explicit --output the bucket becomes the " + "sole destination and nothing is downloaded back." + ), + ) + + +def _add_remote_arguments(parser: ArgumentParser) -> None: + parser.add_argument( + "--remote", + action="store_true", + help="Run this command 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, cpu-basic).", + ) + parser.add_argument( + "--timeout", + default="10m", + help="Max wallclock for the run command inside the sandbox (e.g. 30m, 2h). Defaults to 10m.", + ) + 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=( + "Sandbox image for --remote (defaults to " + f"{_DEFAULT_REMOTE_IMAGE!r}). Must provide torch + CUDA; the CLI installs the " + "small Python deps on top via `uv pip install --system`." + ), + ) + parser.add_argument( + "--keep-alive", + action="store_true", + help=( + "Don't terminate the sandbox after the run. Its id is printed so a later --remote run " + "can reconnect with --sandbox-id and reuse the warm deps/weights/compile cache." + ), + ) + parser.add_argument( + "--sandbox-id", + default=None, + help=( + "Reconnect to an existing sandbox (from a prior --keep-alive run) instead of creating a new " + "one, reusing its warm deps/weights/compile cache. Implies --keep-alive; stop it with " + "`hf sandbox kill `." + ), + ) + parser.add_argument( + "--idle-timeout", + default="10m", + help=( + "Auto-shutdown the sandbox after this much inactivity (e.g. 30m, 1h). Defaults to 10m. " + "Only applied on new sandbox creation — ignored when reconnecting via --sandbox-id." + ), + ) + parser.add_argument( + "--volume", + action="append", + default=None, + metavar="BUCKET_ID[:MOUNT_PATH]", + help=( + "Mount an HF bucket into the sandbox as a read-write directory. Repeatable. Format: " + "`/` (mounts at `/mnt/buckets//`) or " + "`/:/some/path` for a custom path. Reference mounted files from " + "--pipeline-kwargs like any other local path. Applied only on new sandbox creation — " + "ignored when reconnecting via --sandbox-id." + ), + ) + + +# --------------------------------------------------------------------------- +# Pipeline loading + optimization +# --------------------------------------------------------------------------- + + +def _resolve_dtype(name: str | None): + if name in (None, "auto"): + return "auto" + import torch + + mapping = { + "fp32": torch.float32, + "float32": torch.float32, + "fp16": torch.float16, + "float16": torch.float16, + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + } + if name not in mapping: + raise ValueError(f"Unknown dtype: {name}") + return mapping[name] + + +def _resolve_device_map(raw: str | None) -> str | dict: + """Parse `--device-map` into a value acceptable by `from_pretrained(device_map=...)`. + + Returns a JSON dict if the value looks like one, `"balanced"` verbatim, or a single-device string (e.g. `"cuda"`, + `"cuda:1"`, `"cpu"`, `"mps"`). Auto-detects when `raw is None`, pinning to `cuda:$LOCAL_RANK` under torchrun. + """ + if raw is None: + from diffusers.utils.torch_utils import torch_device + + if torch_device == "cuda": + local_rank = os.environ.get("LOCAL_RANK") + if local_rank is not None: + import torch + + torch.cuda.set_device(int(local_rank)) + return f"cuda:{local_rank}" + return torch_device + + if raw.strip().startswith("{"): + try: + parsed = json.loads(raw) + except json.JSONDecodeError as e: + raise SystemExit(f"--device-map must be a device string or a JSON dict: {e}") from e + if not isinstance(parsed, dict): + raise SystemExit("--device-map JSON must decode to an object.") + return parsed + + return raw + + +def _apply_cpu_offload(pipeline: Any, mode: str, device_map: str | dict) -> None: + """Apply model or group CPU offload. Requires a single-device target (not balanced or dict).""" + if not isinstance(device_map, str) or device_map == "balanced": + raise SystemExit( + "--cpu-offload requires --device-map to be a single device string (e.g. 'cuda'); " + f"got {device_map!r}. balanced/dict placement is incompatible with CPU offload." + ) + + if mode == "model": + pipeline.enable_model_cpu_offload(device=device_map) + elif mode == "group": + import torch + + pipeline.enable_group_offload( + onload_device=torch.device(device_map), + offload_type="leaf_level", + use_stream=True, + ) + + +def _set_attention_backend(pipeline: Any, backend: str) -> None: + transformer = getattr(pipeline, "transformer", None) + if transformer is None or not hasattr(transformer, "set_attention_backend"): + logger.warning( + f"--attention-backend is only supported on transformer-based pipelines; " + f"{type(pipeline).__name__} uses the legacy UNet attention path." + ) + return + try: + transformer.set_attention_backend(backend) + except (ValueError, ImportError, RuntimeError) as e: + logger.warning( + f"Attention backend {backend!r} could not be set on {type(transformer).__name__}: " + f"{type(e).__name__}: {e}. Falling back to the model's default backend." + ) + + +def _enable_context_parallel(pipeline: Any) -> None: + import torch + + if not torch.distributed.is_available(): + raise SystemExit("--context-parallel requires a torch build with distributed support.") + + if not torch.distributed.is_initialized(): + # Hybrid backend: ulysses_anything's per-rank size coordination wants Gloo on CPU + # (avoids H2D/D2H for a tiny int tensor); the main attention all-to-all stays on NCCL. + torch.distributed.init_process_group(backend="cpu:gloo,cuda:nccl") + + transformer = getattr(pipeline, "transformer", None) + if transformer is None or not hasattr(transformer, "enable_parallelism"): + raise SystemExit( + "--context-parallel requires a DiT-based pipeline. " + f"{type(pipeline).__name__} does not expose a `transformer` with `enable_parallelism`." + ) + + from diffusers import ContextParallelConfig + + transformer.enable_parallelism( + config=ContextParallelConfig( + ulysses_degree=torch.distributed.get_world_size(), + ring_degree=1, + ulysses_anything=True, + ) + ) + + +def _apply_optimizations(pipeline: Any, args: Namespace) -> None: + """Apply VAE tiling/slicing, attention backend, context-parallel, and torch.compile toggles.""" + vae = getattr(pipeline, "vae", None) + if args.vae_tiling and vae is not None and hasattr(vae, "enable_tiling"): + vae.enable_tiling() + if args.vae_slicing and vae is not None and hasattr(vae, "enable_slicing"): + vae.enable_slicing() + if args.attention_backend != "default": + _set_attention_backend(pipeline, args.attention_backend) + if args.context_parallel: + _enable_context_parallel(pipeline) + if args.compile is not None: + if args.context_parallel: + logger.warning("--compile is currently not supported with --context-parallel; skipping compile.") + else: + _compile_denoiser(pipeline, args.compile) + + +def _compile_denoiser(pipeline: Any, compile_spec: str) -> 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`. + """ + import torch + + 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.") + + for attr in dir(pipeline): + if not any(attr.startswith(key) for key in _DENOISER_COMPONENT_KEYS): + continue + module = getattr(pipeline, attr, None) + if not isinstance(module, torch.nn.Module): + continue + + if getattr(module, "_repeated_blocks", None): + # 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. + setattr(pipeline, attr, torch.compile(module, **compile_kwargs)) + + +def _load_lora(pipeline: Any, args: Namespace) -> None: + """Attach one or more LoRA adapters. Each `--lora` value is a JSON dict. + + Per-entry fields: `lora_id` (required), `lora_scale` (optional float, default 1.0), `adapter_name` (optional; + auto-generated as `lora_` when stacking). Multiple `--lora` flags stack via a single `set_adapters(...)` call at + the end. + """ + if not args.lora: + return + specs = [] + for raw in args.lora: + try: + parsed = json.loads(raw) + except json.JSONDecodeError as e: + raise SystemExit(f"--lora must be valid JSON: {e}") from e + if not isinstance(parsed, dict): + raise SystemExit(f"--lora must decode to a JSON object; got {type(parsed).__name__}.") + specs.append(parsed) + if not hasattr(pipeline, "load_lora_weights"): + raise SystemExit(f"{type(pipeline).__name__} does not support LoRA loading.") + + names: list[str] = [] + scales: list[float] = [] + for i, spec in enumerate(specs): + lora_id = spec.get("lora_id") + if not lora_id: + raise SystemExit(f"--lora entry {i} is missing 'lora_id'.") + adapter_name = spec.get("adapter_name") or (f"lora_{i}" if len(specs) > 1 else "default") + pipeline.load_lora_weights(lora_id, adapter_name=adapter_name) + names.append(adapter_name) + scales.append(float(spec.get("lora_scale", 1.0))) + + if hasattr(pipeline, "set_adapters"): + pipeline.set_adapters(names, adapter_weights=scales) + + +def _load_pipeline(args: Namespace) -> Any: + import diffusers + + # Detect modular repos by trying the standard config; `ModularPipeline` repos ship + # `modular_model_index.json` instead of `model_index.json`, so `load_config` OSErrors. + try: + diffusers.DiffusionPipeline.load_config(args.model, token=args.token, revision=args.revision) + modular = False + except OSError: + modular = True + + dtype = _resolve_dtype(args.dtype) + device_map = _resolve_device_map(args.device_map) + common_kwargs: dict[str, Any] = { + "trust_remote_code": args.trust_remote_code, + } + if dtype != "auto": + common_kwargs["torch_dtype"] = dtype + if args.variant: + common_kwargs["variant"] = args.variant + if args.token: + common_kwargs["token"] = args.token + # CPU offload sets up its own placement hooks, so leave weights on CPU at load time. + if not args.cpu_offload: + common_kwargs["device_map"] = device_map + + if modular: + # ModularPipeline.from_pretrained fetches only the pipeline config; component + # weights come in via load_components(). `revision` scopes the config fetch, + # so it stays on from_pretrained — each ComponentSpec pins its own revision, + # and forwarding a global `revision` to load_components() would override those. + pipeline = diffusers.ModularPipeline.from_pretrained( + args.model, + trust_remote_code=args.trust_remote_code, + token=args.token, + revision=args.revision, + ) + pipeline.load_components(**common_kwargs) + else: + pipeline = diffusers.DiffusionPipeline.from_pretrained(args.model, revision=args.revision, **common_kwargs) + + _load_lora(pipeline, args) + if args.cpu_offload: + _apply_cpu_offload(pipeline, args.cpu_offload, device_map) + _apply_optimizations(pipeline, args) + + return pipeline + + +# --------------------------------------------------------------------------- +# Pipeline call helpers +# --------------------------------------------------------------------------- + + +def _parse_pipeline_kwargs(raw: str | None) -> dict[str, Any]: + if not raw: + return {} + try: + parsed = json.loads(raw) + except json.JSONDecodeError as e: + raise SystemExit(f"--pipeline-kwargs must be valid JSON: {e}") from e + if not isinstance(parsed, dict): + raise SystemExit("--pipeline-kwargs must decode to a JSON object.") + return parsed + + +def _load_audio(url_or_path: str) -> tuple[Any, int]: + """Load audio from a URL or local path via torchaudio. Returns `(waveform, sampling_rate)`.""" + import torchaudio + + if url_or_path.startswith(("http://", "https://")): + import io + + import httpx + + from ..utils.constants import DIFFUSERS_REQUEST_TIMEOUT + + resp = httpx.get(url_or_path, follow_redirects=True, timeout=DIFFUSERS_REQUEST_TIMEOUT) + resp.raise_for_status() + return torchaudio.load(io.BytesIO(resp.content)) + return torchaudio.load(url_or_path) + + +def _resolve_media_inputs(call_kwargs: dict[str, Any]) -> None: + """Replace string paths/URLs at known media-input keys with loaded tensors. + + Images resolve to `PIL.Image.Image` via `load_image`; videos to `list[PIL.Image.Image]` via `load_video`; audio to + a `torch.Tensor` via `_load_audio` (also auto-sets the paired sampling-rate kwarg for `initial_audio_waveforms` if + the user didn't supply it). A `list[str]` at any key is treated as a batch: each entry is loaded and the value + becomes a list of loaded objects. Non-string, non-list values pass through untouched. + """ + + def _is_string_list(v: Any) -> bool: + return isinstance(v, list) and bool(v) and all(isinstance(x, str) for x in v) + + for key in _IMAGE_INPUT_KEYS: + value = call_kwargs.get(key) + if isinstance(value, str): + call_kwargs[key] = load_image(value) + elif _is_string_list(value): + call_kwargs[key] = [load_image(v) for v in value] + for key in _VIDEO_INPUT_KEYS: + value = call_kwargs.get(key) + if isinstance(value, str): + call_kwargs[key] = load_video(value) + elif _is_string_list(value): + call_kwargs[key] = [load_video(v) for v in value] + for key in _AUDIO_INPUT_KEYS: + value = call_kwargs.get(key) + if isinstance(value, str): + waveform, sr = _load_audio(value) + call_kwargs[key] = waveform + if key == "initial_audio_waveforms" and "initial_audio_sampling_rate" not in call_kwargs: + call_kwargs["initial_audio_sampling_rate"] = sr + elif _is_string_list(value): + pairs = [_load_audio(v) for v in value] + call_kwargs[key] = [w for w, _ in pairs] + if key == "initial_audio_waveforms" and "initial_audio_sampling_rate" not in call_kwargs: + # All batched waveforms must share a sampling rate; use the first entry's. + call_kwargs["initial_audio_sampling_rate"] = pairs[0][1] + + +def _get_generator(seed: int | None, device: str): + if seed is None: + return None + import torch + + generator_device = "cpu" if device == "mps" else device + return torch.Generator(device=generator_device).manual_seed(seed) + + +def _unwrap_pipeline_output(result: Any) -> Any: + """Unwrap a pipeline-output object into the raw payload the saver can dispatch on.""" + if hasattr(result, "images"): + return result.images + if hasattr(result, "frames"): + return result.frames[0] + if hasattr(result, "audios"): + return result.audios + return result + + +# --------------------------------------------------------------------------- +# Output saving (dispatch by type) +# --------------------------------------------------------------------------- + + +def _get_or_create_run_id() -> str: + """Return the current run's id, creating one if not yet set. + + Format: `diffusers-run--<6-char-uuid>`. Same id is reused as the local output subdirectory, the + remote bucket prefix, and the container-side `RUN_ID_ENV` so a run's artifacts are traceable end-to-end. + """ + import uuid + from datetime import datetime + + existing = os.environ.get(RUN_ID_ENV) + if existing: + return existing + run_id = f"diffusers-run-{datetime.now().strftime('%Y%m%dT%H%M%S')}-{uuid.uuid4().hex[:6]}" + os.environ[RUN_ID_ENV] = run_id + return run_id + + +def _resolve_output_paths(task: str, num: int, explicit: str | None, ext: str) -> list[Path]: + if explicit is None: + base = Path(DEFAULT_OUTPUT_DIR) / _get_or_create_run_id() + base.mkdir(parents=True, exist_ok=True) + return [base / f"{i:04d}.{ext}" for i in range(num)] + + p = Path(explicit) + if explicit.endswith(os.sep) or p.is_dir(): + p.mkdir(parents=True, exist_ok=True) + return [p / f"{i:04d}.{ext}" for i in range(num)] + + p.parent.mkdir(parents=True, exist_ok=True) + if num == 1: + return [p] + stem, suffix = p.stem, p.suffix or f".{ext}" + return [p.with_name(f"{stem}-{i:04d}{suffix}") for i in range(num)] + + +def _as_pil_list(value: Any): + try: + from PIL.Image import Image as PILImage + except ImportError: + return None + if isinstance(value, PILImage): + return [value] + if isinstance(value, (list, tuple)) and value and all(isinstance(v, PILImage) for v in value): + return list(value) + return None + + +def _as_frame_sequence(value: Any): + try: + from PIL.Image import Image as PILImage + except ImportError: + PILImage = None # type: ignore[assignment] + + if isinstance(value, (list, tuple)) and len(value) >= 2: + first = value[0] + if PILImage is not None and isinstance(first, PILImage): + return list(value) + try: + import numpy as np + + if isinstance(first, np.ndarray): + return list(value) + except ImportError: + pass + return None + + +def _as_audio_arrays(value: Any): + try: + import numpy as np + except ImportError: + return None + if isinstance(value, np.ndarray) and value.ndim <= 2: + return [value] + if isinstance(value, (list, tuple)) and value and all(isinstance(v, np.ndarray) for v in value): + return list(value) + return None + + +def _save_audio_arrays(audios, sampling_rate: int, args: Namespace, task: str) -> list[str]: + """Write each numpy audio array to a 16-bit PCM WAV at `sampling_rate` Hz. + + Uses the stdlib `wave` module so no scipy dependency is required. + """ + import wave + + import numpy as np + + paths = _resolve_output_paths(task, len(audios), args.output, ext="wav") + saved: list[str] = [] + for audio, path in zip(audios, paths): + data = np.asarray(audio) + if data.dtype.kind == "f": + data = (np.clip(data, -1.0, 1.0) * 32767).astype(np.int16) + else: + data = data.astype(np.int16) + if data.ndim == 1: + n_channels = 1 + else: + # Heuristic: shorter axis is channels (interleaved layout for `wave` is + # samples × channels, so transpose if needed). + if data.shape[0] < data.shape[-1]: + data = data.T + n_channels = data.shape[1] + with wave.open(str(path), "wb") as w: + w.setnchannels(n_channels) + w.setsampwidth(2) # 16-bit PCM + w.setframerate(sampling_rate) + w.writeframes(data.tobytes()) + saved.append(str(path)) + return saved + + +def _save_output(value: Any, args: Namespace, task: str) -> list[str]: + """Save `value` by dispatching on its runtime type.""" + pil_images = _as_pil_list(value) + if pil_images is not None: + paths = _resolve_output_paths(task, len(pil_images), args.output, ext="png") + for img, path in zip(pil_images, paths): + img.save(path) + return [str(p) for p in paths] + + frames = _as_frame_sequence(value) + if frames is not None: + from diffusers.utils import export_to_video + + path = _resolve_output_paths(task, 1, args.output, ext="mp4")[0] + export_to_video(frames, str(path), fps=args.fps) + return [str(path)] + + audios = _as_audio_arrays(value) + if audios is not None: + return _save_audio_arrays(audios, args.sampling_rate or 16000, args, task) + + path = _resolve_output_paths(task, 1, args.output, ext="json")[0] + Path(path).write_text(json.dumps(value, default=str, indent=2)) + return [str(path)] + + +# --------------------------------------------------------------------------- +# Hub bucket upload (--push-to) +# --------------------------------------------------------------------------- + + +def _parse_push_to(spec: str) -> tuple[str, str]: + """Split `--push-to` into a bucket id and an optional subpath prefix. + + Accepts an HF bucket id (`/[/]`), a canonical + `hf://buckets//[/]` URI, or a Hub web URL for the same. Non-bucket URIs (models, + datasets, spaces) are rejected — `--push-to` targets storage buckets only. + """ + from huggingface_hub import parse_hf_uri + + # Bare shorthand → canonical URI so a single parser handles every accepted form. + if not spec.startswith(("hf://", "http://", "https://")): + spec = f"hf://buckets/{spec.strip('/')}" + uri = parse_hf_uri(spec) + if not uri.is_bucket: + raise SystemExit(f"--push-to must point at a bucket; got {uri.type!r} URI {spec!r}.") + return uri.id, uri.path_in_repo + + +def _push_outputs(args: Namespace, saved_paths: list[str], task: str) -> dict[str, Any] | None: + """Upload `saved_paths` to the `--push-to` bucket. Returns a summary or None.""" + if not args.push_to: + return None + + from huggingface_hub import HfApi + + bucket_id, subpath = _parse_push_to(args.push_to) + api = HfApi(token=args.token) + api.create_bucket(bucket_id, exist_ok=True) + + run_id = _get_or_create_run_id() + prefix = f"{subpath}/{run_id}" if subpath else run_id + add = [(local, f"{prefix}/{Path(local).name}") for local in saved_paths] + api.batch_bucket_files(bucket_id, add=add) + + uploaded = [f"hf://buckets/{bucket_id}/{dest}" for _, dest in add] + return {"bucket_id": bucket_id, "uploaded": uploaded} + + +# --------------------------------------------------------------------------- +# Remote execution (HF Sandbox) +# --------------------------------------------------------------------------- + + +def _build_task_kwargs(args: Namespace) -> dict[str, Any]: + """Pick out the kwargs the sandbox CLI should invoke the task with.""" + out: dict[str, Any] = {} + for key, value in vars(args).items(): + if key in REMOTE_KEYS or value is None or value is False: + continue + out[key] = value + return out + + +def _kwargs_to_argv(task: str, task_kwargs: dict[str, Any]) -> list[str]: + """Render `task_kwargs` as the argv list the sandbox CLI's argparse will see.""" + argv: list[str] = [task] + for key, value in task_kwargs.items(): + flag = "--" + key.replace("_", "-") + if value is True: + argv.append(flag) + elif isinstance(value, list): + for item in value: + argv.extend([flag, str(item)]) + else: + argv.extend([flag, str(value)]) + return argv + + +def _duration_to_seconds(value: str) -> float: + """Parse a duration like `30s`, `10m`, `2h` (or a bare number of seconds) into seconds.""" + value = value.strip() + units = {"s": 1, "m": 60, "h": 3600} + if value and value[-1] in units: + return float(value[:-1]) * units[value[-1]] + return float(value) + + +def _upload_inputs_to_sandbox(args: Namespace, sbx: Any, run_id: str) -> None: + """Upload local media paths in `--pipeline-kwargs` into the sandbox and rewrite the JSON in place. + + Walks known image/video/audio-input keys; any string value that resolves to a local file is uploaded to + `<_SANDBOX_INPUTS_DIR>//_` and the JSON path is rewritten to that in-sandbox path. URLs, + `hf://` URIs, and non-existent paths pass through untouched. + """ + if not args.pipeline_kwargs: + return + try: + parsed = json.loads(args.pipeline_kwargs) + except json.JSONDecodeError: + return # the sandbox CLI will fail loudly with a parse error later + if not isinstance(parsed, dict): + return + + def _upload_one(key: str, index: int | None, local_str: str) -> str: + # `index` is None for scalar entries, an int for list entries (used to disambiguate names). + local = Path(local_str) + suffix = f"_{index}" if index is not None else "" + remote_path = f"{_SANDBOX_INPUTS_DIR}/{run_id}/{key}{suffix}_{local.name}" + sbx.files.upload(str(local), remote_path) + return remote_path + + uploaded = 0 + for key in (*_IMAGE_INPUT_KEYS, *_VIDEO_INPUT_KEYS, *_AUDIO_INPUT_KEYS): + value = parsed.get(key) + if isinstance(value, str) and Path(value).is_file(): + parsed[key] = _upload_one(key, None, value) + uploaded += 1 + elif isinstance(value, list): + # Batched inputs: upload each local path, leave URLs/hf:// URIs alone. + new_list = list(value) + for i, entry in enumerate(value): + if isinstance(entry, str) and Path(entry).is_file(): + new_list[i] = _upload_one(key, i, entry) + uploaded += 1 + parsed[key] = new_list + + if uploaded: + logger.info(f"uploaded {uploaded} local input file(s) to the sandbox") + args.pipeline_kwargs = json.dumps(parsed) + + +def _download_outputs_from_sandbox(sbx: Any, sandbox_dir: str, local_dir: Path) -> list[str]: + """Download every file the sandbox CLI wrote under `sandbox_dir` into `local_dir`.""" + local_dir.mkdir(parents=True, exist_ok=True) + saved: list[str] = [] + for entry in sbx.files.list(sandbox_dir): + if entry.type != "file": + continue + target = local_dir / Path(entry.path).name + sbx.files.download(entry.path, str(target)) + saved.append(str(target)) + return saved + + +def _maybe_submit_remote(args: Namespace, task: str) -> bool: + """If `--remote` was set, run this invocation inside an HF Sandbox and return True.""" + if not args.remote: + return False + + import shlex + import time + + from huggingface_hub import get_token + from huggingface_hub.utils import send_telemetry + + import diffusers + + try: + from huggingface_hub import Sandbox + except ImportError: + raise SystemExit( + "--remote requires huggingface_hub>=1.23 for HF Sandbox support. " + "Upgrade with `pip install -U huggingface_hub`." + ) + + 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." + ) + + hf_token = args.token or get_token() + run_id = _get_or_create_run_id() + + # An explicit --push-to means the bucket is the user's destination, so skip the local + # download unless they also asked for a local path via --output. + user_bucket = bool(args.push_to) + download_locally = (not user_bucket) or (args.output is not None) + local_dir = Path(args.output) if args.output else Path(DEFAULT_OUTPUT_DIR) / run_id + + use_existing_sandbox = bool(args.sandbox_id) + keep_alive = args.keep_alive or use_existing_sandbox + if use_existing_sandbox and args.volume: + logger.warning( + "--volume is ignored when reconnecting to an existing sandbox (mounts are set at creation time)." + ) + if use_existing_sandbox: + logger.info(f"reconnecting to sandbox {args.sandbox_id!r}...") + sbx = Sandbox.connect(args.sandbox_id, token=hf_token) + else: + logger.info(f"creating sandbox on flavor={args.flavor!r}...") + create_kwargs: dict[str, Any] = { + "image": args.image or _DEFAULT_REMOTE_IMAGE, + "flavor": args.flavor, + "forward_hf_token": True, + "token": hf_token, + "env": { + "HF_ENABLE_PARALLEL_LOADING": "1", + "DIFFUSERS_VERBOSITY": os.environ.get("DIFFUSERS_VERBOSITY", "info"), + }, + "idle_timeout": args.idle_timeout, + } + if args.volume: + from huggingface_hub import Volume + + volumes = [] + for spec in args.volume: + bucket_id, sep, mount_path = spec.partition(":") + if not sep: + mount_path = f"/mnt/buckets/{bucket_id}" + if bucket_id.count("/") != 1: + raise SystemExit(f"--volume: bucket id must be /, got {bucket_id!r}") + if not mount_path.startswith("/"): + raise SystemExit(f"--volume: mount path must be absolute, got {mount_path!r}") + volumes.append(Volume(type="bucket", source=bucket_id, mount_path=mount_path)) + create_kwargs["volumes"] = volumes + 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() + + exit_code = 0 + saved: list[str] = [] + run_seconds = 0.0 + try: + _upload_inputs_to_sandbox(args, sbx, run_id) + + dependencies = list(_DEFAULT_REMOTE_DEPS) + if args.dependencies: + dependencies.extend(args.dependencies) + # --break-system-packages bypasses PEP 668; harmless in a throwaway sandbox. uv is a + # near no-op when the deps are already satisfied, so this stays cheap on a reused sandbox. + 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) + + # Per-run outputs subdirectory so a reused sandbox doesn't leak files from prior runs + # into this run's download set. + sandbox_output_dir = f"{_SANDBOX_OUTPUTS_DIR}/{run_id}" + task_kwargs = _build_task_kwargs(args) + task_kwargs["output"] = sandbox_output_dir + "/" + cli_argv = _kwargs_to_argv(task, task_kwargs) + # Suppress the container CLI's own `out.result(...)` payload — the outer wrapper owns the + # final structured output for --remote runs. + format_argv = ["--format", "quiet"] + # torchrun wraps the CLI for --context-parallel so torch.distributed initializes across + # every visible GPU before the run command starts. + if args.context_parallel: + cli_argv = [ + "torchrun", + "--nproc-per-node=gpu", + "-m", + "diffusers.commands.diffusers_cli", + *format_argv, + *cli_argv, + ] + else: + cli_argv = [_CONTAINER_CLI_BINARY, *format_argv, *cli_argv] + + started = time.perf_counter() + # Per-invocation env: RUN_ID_ENV must be fresh each run. Sandbox.create-time env is + # baked in and would go stale on reused sandboxes, silently reusing the initial run's + # bucket prefix in `_push_outputs`. + result = sbx.run( + cli_argv, + env={RUN_ID_ENV: run_id}, + on_stdout=_stream, + on_stderr=_stream, + timeout=_duration_to_seconds(args.timeout), + check=False, + ) + run_seconds = time.perf_counter() - started + exit_code = result.exit_code + + if exit_code == 0 and download_locally: + saved = _download_outputs_from_sandbox(sbx, sandbox_output_dir, local_dir) + finally: + if keep_alive: + logger.info( + f"sandbox {sbx.id} kept alive — reconnect with " + f"`--remote --sandbox-id {sbx.id}`, stop with `hf sandbox kill {sbx.id}`." + ) + else: + sbx.kill() + + send_telemetry( + topic="diffusers/cli/run/remote", + library_name="diffusers", + library_version=diffusers.__version__, + ) + + payload: dict[str, Any] = { + "exit_code": exit_code, + "run_seconds": round(run_seconds, 1), + } + if keep_alive: + payload["sandbox_id"] = sbx.id + if download_locally: + payload["outputs"] = saved + if args.push_to: + bucket_id, subpath = _parse_push_to(args.push_to) + prefix = f"{subpath}/{run_id}" if subpath else run_id + payload["pushed-to"] = f"hf://buckets/{bucket_id}/{prefix}/" + out.result("remote-run", **payload) + + if exit_code != 0: + raise SystemExit(f"remote run failed with exit code {exit_code}") + return True + + +# --------------------------------------------------------------------------- +# Subcommand +# --------------------------------------------------------------------------- + + +class RunCommand(BaseDiffusersCLICommand): + task = "run" + + @staticmethod + def register_subcommand(subparsers: _SubParsersAction) -> None: + from argparse import RawDescriptionHelpFormatter + + epilog = ( + "Examples\n" + " $ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \\\n" + ' --pipeline-kwargs \'{"prompt": "a cat on the moon"}\'\n' + " $ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \\\n" + ' --pipeline-kwargs \'{"prompt": "make the fur grey", "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"}\'\n' + " $ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \\\n" + ' --pipeline-kwargs \'{"prompt": "a tiny cat"}\' \\\n' + ' --lora \'{"lora_id": "alvdansen/littletinies", "lora_scale": 0.8}\'\n' + " $ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \\\n" + ' --pipeline-kwargs \'{"prompt": "a cat"}\' --remote --flavor a100-large\n' + " $ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 --context-parallel \\\n" + ' --pipeline-kwargs \'{"prompt": "a cat"}\' --remote --flavor 4xa100-large\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( + "run", + help="Run any diffusers pipeline locally or remotely in an HF Sandbox.", + usage="\n diffusers-cli run [options]", + epilog=epilog, + formatter_class=RawDescriptionHelpFormatter, + ) + parser._optionals.title = "Options" + _add_loading_arguments(parser) + _add_optimization_arguments(parser) + parser.add_argument( + "--pipeline-kwargs", + default=None, + help=( + "JSON object of kwargs passed to the pipeline call. String values at known " + f"image-input keys ({', '.join(_IMAGE_INPUT_KEYS)}) are auto-loaded as PIL images; " + f"video-input keys ({', '.join(_VIDEO_INPUT_KEYS)}) are auto-loaded as frame lists; " + f"audio-input keys ({', '.join(_AUDIO_INPUT_KEYS)}) are auto-loaded via torchaudio." + ), + ) + parser.add_argument( + "--output-key", + default=None, + help="For modular pipelines: name of the intermediate to extract (passed as `output=` to the call).", + ) + parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility.") + parser.add_argument( + "--fps", + type=int, + default=8, + help="FPS used when the output happens to be a frame sequence.", + ) + parser.add_argument( + "--sampling-rate", + type=int, + default=None, + help="Sample rate used when the output happens to be an audio array.", + ) + _add_remote_arguments(parser) + _add_output_arguments(parser) + parser.set_defaults(func=RunCommand) + + def __init__(self, args: Namespace): + self.args = args + + def run(self) -> None: + import diffusers + + _get_or_create_run_id() # populate RUN_ID_ENV so local output dir + remote bucket prefix agree + + call_kwargs = _parse_pipeline_kwargs(self.args.pipeline_kwargs) + + if _maybe_submit_remote(self.args, self.task): + return + + # Resolve media before loading pipeline weights so dead URLs / missing files fail + # fast — cheap to fetch, expensive to load a 20GB model just to hit a 404. + _resolve_media_inputs(call_kwargs) + pipeline = _load_pipeline(self.args) + is_modular = isinstance(pipeline, diffusers.ModularPipeline) + + if self.args.output_key is not None: + call_kwargs["output"] = self.args.output_key + + device = pipeline.device.type if hasattr(pipeline, "device") else "cpu" + generator = _get_generator(self.args.seed, device) + if generator is not None: + call_kwargs["generator"] = generator + + try: + result = pipeline(**call_kwargs) + + # Under torchrun, ranks > 0 produce identical output to rank 0 (CP shards the + # transformer compute but ranks reduce to the same final tensors). Save/push/print + # from rank 0 only to avoid clobbering bucket files 4x and printing 4x. + if os.environ.get("RANK", "0") == "0": + savable = result if is_modular else _unwrap_pipeline_output(result) + saved = _save_output(savable, self.args, self.task) + pushed = _push_outputs(self.args, saved, self.task) + + out.result( + self.task, + model=self.args.model, + device=device, + pipeline_class=type(pipeline).__name__, + modular=is_modular, + outputs=saved, + pushed=pushed, + seed=self.args.seed, + output_key=self.args.output_key, + ) + finally: + import torch + + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() diff --git a/src/diffusers/commands/schema.py b/src/diffusers/commands/schema.py new file mode 100644 index 0000000000..dc5965e8ad --- /dev/null +++ b/src/diffusers/commands/schema.py @@ -0,0 +1,287 @@ +# 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 schema` — print the input schema for any pipeline repo. + +Tries `DiffusionPipeline.config_name` first (so standard repos get their `__call__` signature introspected); falls back +to `ModularPipelineBlocks.from_pretrained` for modular repos. No weights are downloaded — only the small index file +(and any custom block code if `--trust-remote-code` is set). +""" + +from __future__ import annotations + +import inspect +import re +from argparse import ArgumentParser, Namespace, _SubParsersAction +from typing import Any + +from huggingface_hub.cli._output import OutputFormat, out + +from ..utils import logging +from . import BaseDiffusersCLICommand + + +logger = logging.get_logger("diffusers-cli/schema") + + +def _schema(args: Namespace) -> None: + """Print the pipeline's input schema. + + Tries `DiffusionPipeline.config_name` (= `model_index.json`) first; if present, introspects the declared pipeline + class's `__call__` signature. Otherwise falls back to `ModularPipelineBlocks.from_pretrained` and reads the + block-declared `inputs`. No weights downloaded either way. + """ + import diffusers + + try: + index = diffusers.DiffusionPipeline.load_config(args.model, token=args.token, revision=args.revision) + except OSError: + index = None + + if index is not None: + class_name = index.get("_class_name") + if class_name is None: + raise SystemExit( + f"{diffusers.DiffusionPipeline.config_name} for {args.model!r} has no `_class_name` field." + ) + pipeline_cls = getattr(diffusers, class_name, None) + if pipeline_cls is None: + raise SystemExit( + f"Pipeline class {class_name!r} declared in {diffusers.DiffusionPipeline.config_name} " + "is not exported by the installed diffusers." + ) + + sig = inspect.signature(pipeline_cls.__call__) + descriptions = _parse_docstring_args(pipeline_cls.__call__.__doc__) if args.verbose else {} + schema: list[dict[str, Any]] = [] + for name, param in sig.parameters.items(): + if name == "self": + continue + if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): + continue + has_default = param.default is not inspect.Parameter.empty + schema.append( + { + "name": name, + "type_hint": str(param.annotation) if param.annotation is not inspect.Parameter.empty else None, + "default": param.default if has_default else None, + "required": not has_default, + "description": descriptions.get(name, ""), + } + ) + else: + kwargs: dict[str, Any] = {"trust_remote_code": args.trust_remote_code} + if args.revision: + kwargs["revision"] = args.revision + if args.token: + kwargs["token"] = args.token + + # If the repo declares custom code + external dependencies, surface them upfront so + # the user knows what to install before we hit an ImportError inside from_pretrained. + _warn_custom_block_requirements(args) + + try: + blocks = diffusers.ModularPipelineBlocks.from_pretrained(args.model, **kwargs) + except Exception as e: + hint = "\nPass --trust-remote-code if it ships custom block code." if not args.trust_remote_code else "" + raise SystemExit( + f"Could not read schema for {args.model!r}: no {diffusers.DiffusionPipeline.config_name} and " + f"loading as a modular pipeline failed with:\n {type(e).__name__}: {e}{hint}" + ) from e + + class_name = type(blocks).__name__ + schema = [ + { + "name": p.name, + "type_hint": str(p.type_hint) if p.type_hint is not None else None, + "default": p.default, + "required": p.required, + "description": p.description, + } + for p in blocks.inputs + ] + + if out.mode == OutputFormat.json: + out.dict({"task": "schema", "model": args.model, "pipeline_class": class_name, "inputs": schema}) + elif out.mode == OutputFormat.agent: + out.table(schema, headers=["name", "required", "type_hint", "default", "description"]) + else: + out.text(f"{class_name} ({args.model}) inputs:") + for entry in schema: + tag = "required" if entry["required"] else f"optional, default={entry['default']!r}" + out.text(f" {entry['name']} ({tag})") + if entry["type_hint"]: + out.text(f" type: {entry['type_hint']}") + if entry["description"]: + out.text(f" desc: {entry['description']}") + + +def _warn_custom_block_requirements(args: Namespace) -> None: + """Warn upfront when a modular block ships custom code with declared external dependencies. + + Reads `modular_config.json` if present; if it has an `auto_map` (custom code) and a non-empty `requirements` + list/dict, prints a heads-up. `from_pretrained` will otherwise fail with an `ImportError` deep in the loader stack + when a listed dep is missing. + """ + import diffusers + + try: + config = diffusers.ModularPipelineBlocks.load_config(args.model, token=args.token, revision=args.revision) + except Exception: + return # no modular_config.json or unreachable — nothing to warn about + if not isinstance(config, dict): + return + if not config.get("auto_map"): + return + requirements = config.get("requirements") + if not requirements: + return + + # `requirements` may be a dict {name: version} or (older repos) a list of [name, version] pairs. + if isinstance(requirements, dict): + pairs = list(requirements.items()) + elif isinstance(requirements, list): + pairs = [(item[0], item[1]) for item in requirements if isinstance(item, (list, tuple)) and len(item) >= 2] + else: + pairs = [] + if not pairs: + return + + formatted = ", ".join(f"{name}=={version}" for name, version in pairs) + logger.warning( + f"{args.model!r} ships custom block code with external dependencies: {formatted}. " + "You will need to install these in order to determine the pipeline schema." + ) + + +def _parse_docstring_args(docstring: str | None) -> dict[str, str]: + """Extract per-argument descriptions from a Google-style `Args:` block. + + Returns a `{name: description}` mapping. Best-effort — unrecognised formats just yield an empty dict rather than + raising. + """ + if not docstring: + return {} + + lines = docstring.expandtabs().splitlines() + start = None + section_indent = 0 + for i, line in enumerate(lines): + if line.strip() in ("Args:", "Arguments:", "Parameters:"): + start = i + 1 + section_indent = len(line) - len(line.lstrip()) + break + if start is None: + return {} + + descriptions: dict[str, str] = {} + current_name: str | None = None + current_lines: list[str] = [] + arg_indent: int | None = None + name_pattern = re.compile(r"^(\w+)\s*(?:\([^)]*\))?\s*:?\s*(.*)$") + + def _flush() -> None: + if current_name and current_lines: + descriptions[current_name] = " ".join(s.strip() for s in current_lines).strip() + + for line in lines[start:]: + if not line.strip(): + continue + indent = len(line) - len(line.lstrip()) + # A new top-level section ends the Args block. + if indent <= section_indent and line.strip().endswith(":"): + break + if arg_indent is None: + arg_indent = indent + if indent == arg_indent: + _flush() + current_lines = [] + match = name_pattern.match(line.strip()) + if match: + current_name = match.group(1) + tail = match.group(2).strip() + if tail: + current_lines.append(tail) + else: + current_name = None + elif current_name is not None and indent > arg_indent: + current_lines.append(line.strip()) + _flush() + return descriptions + + +class SchemaCommand(BaseDiffusersCLICommand): + task = "schema" + + @staticmethod + def register_subcommand(subparsers: _SubParsersAction) -> None: + from argparse import RawDescriptionHelpFormatter + + epilog = ( + "Examples\n" + " $ diffusers-cli schema -m stabilityai/stable-diffusion-xl-base-1.0\n" + " $ diffusers-cli schema -m black-forest-labs/FLUX.1-dev --verbose\n" + " $ diffusers-cli --format json schema -m stabilityai/stable-diffusion-xl-base-1.0\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( + "schema", + help="Print the input schema for a diffusers pipeline repo. No weights downloaded.", + usage="\n diffusers-cli schema [options]", + epilog=epilog, + formatter_class=RawDescriptionHelpFormatter, + ) + parser._optionals.title = "Options" + parser.add_argument( + "--model", + "-m", + required=True, + help="Model id on the Hugging Face Hub or local path.", + ) + parser.add_argument( + "--revision", + default=None, + help="Model revision (branch, tag, or commit SHA).", + ) + parser.add_argument( + "--token", + default=None, + help="Hugging Face token for gated/private models.", + ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Allow custom code from the Hub (required for modular pipelines that ship block code).", + ) + parser.add_argument( + "--verbose", + "-v", + action="store_true", + help=( + "Also include per-argument descriptions from the pipeline's __call__ docstring. " + "Modular pipelines always include block-declared descriptions; --verbose populates " + "the equivalent field for standard pipelines by parsing the Google-style Args: block." + ), + ) + parser.set_defaults(func=SchemaCommand) + + def __init__(self, args: Namespace): + self.args = args + + def run(self) -> None: + _schema(self.args) diff --git a/src/diffusers/commands/skills.py b/src/diffusers/commands/skills.py new file mode 100644 index 0000000000..60d2e40e88 --- /dev/null +++ b/src/diffusers/commands/skills.py @@ -0,0 +1,344 @@ +# 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 skills` — install Agent Skills bundles. + +Skill bundles live under `.ai/skills//` in the diffusers repo and follow the Agent Skills standard: a directory +containing `SKILL.md` (plus optional resources). Installs to `.agents/skills//` which Claude, Codex, and Cursor +all discover. +""" + +from __future__ import annotations + +import os +import shutil +from argparse import ArgumentParser, Namespace, _SubParsersAction +from pathlib import Path + +import httpx +from huggingface_hub.cli._output import out + +from ..utils import logging +from ..utils.constants import DIFFUSERS_REQUEST_TIMEOUT +from . import BaseDiffusersCLICommand + + +logger = logging.get_logger("diffusers-cli/skills") + + +_REGISTRY_BASE = "https://api.github.com/repos/huggingface/diffusers/contents/.ai/skills" +_REGISTRY_REF = "main" + +# Native skill-discovery paths per agent. Claude Code reads only `.claude/skills/`; Codex and +# Cursor read `.agents/skills/` (Cursor also honors `.claude/skills/` via compat, but installing +# to `.agents/skills/` is the portable choice for both). +_CLAUDE_SKILLS_DIR = Path(".claude") / "skills" +_AGENTS_SKILLS_DIR = Path(".agents") / "skills" + +# Env vars set by each agent when it launches the CLI. Values are the install path to use. +_AGENT_ENV_TO_DIR: dict[str, Path] = { + "CLAUDECODE": _CLAUDE_SKILLS_DIR, + "CLAUDE_CODE": _CLAUDE_SKILLS_DIR, + "CODEX_SANDBOX": _AGENTS_SKILLS_DIR, + "CURSOR_AI": _AGENTS_SKILLS_DIR, +} +# When no agent env var is set, install to every native path so whichever agent the user +# later switches to picks the skill up. +_ALL_INSTALL_DIRS: tuple[Path, ...] = (_CLAUDE_SKILLS_DIR, _AGENTS_SKILLS_DIR) + +# Empty marker dropped inside each installed skill dir so `update` can distinguish our +# installs from user-placed skills at the same paths. +_MANAGED_MARKER_FILE = ".diffusers-skill-managed" + + +# --------------------------------------------------------------------------- +# Registry fetch +# --------------------------------------------------------------------------- + + +def _registry_url(name: str = "") -> str: + """API URL for the registry root, or for a single skill bundle when `name` is given.""" + path = f"/{name}" if name else "" + return f"{_REGISTRY_BASE}{path}?ref={_REGISTRY_REF}" + + +def _fetch_json(url: str) -> list[dict]: + try: + resp = httpx.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) + resp.raise_for_status() + return resp.json() + except httpx.HTTPStatusError as e: + if e.response.status_code == 404: + raise SystemExit(f"Not found in registry: {url}") from e + raise SystemExit(f"Registry fetch failed: HTTP {e.response.status_code} {e.response.reason_phrase}") from e + except httpx.HTTPError as e: + raise SystemExit(f"Could not reach registry: {e}") from e + + +def _walk_skill_files(name: str) -> list[tuple[str, str]]: + files: list[tuple[str, str]] = [] + + def _walk(api_url: str, prefix: str) -> None: + for entry in _fetch_json(api_url): + if entry["type"] == "file": + files.append((f"{prefix}{entry['name']}", entry["download_url"])) + elif entry["type"] == "dir": + _walk(entry["url"], f"{prefix}{entry['name']}/") + + _walk(_registry_url(name), "") + return files + + +def _download_skill_bundle(name: str) -> dict[str, bytes]: + files = _walk_skill_files(name) + if not files: + raise SystemExit(f"Skill '{name}' has no files in the registry.") + bundle: dict[str, bytes] = {} + for rel_path, url in files: + resp = httpx.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) + resp.raise_for_status() + bundle[rel_path] = resp.content + return bundle + + +# --------------------------------------------------------------------------- +# Install / discovery +# --------------------------------------------------------------------------- + + +def _detect_install_dirs() -> tuple[Path, ...]: + """Pick where to install based on the launching agent. + + If we detect a specific agent from its env var, install only there. If nothing is detected, install to every native + path so any agent picks the skill up later. + """ + for env_var, skills_dir in _AGENT_ENV_TO_DIR.items(): + if os.environ.get(env_var): + return (skills_dir,) + return _ALL_INSTALL_DIRS + + +def _install_skill(name: str, bundle: dict[str, bytes], root: Path, skills_dir: Path, force: bool) -> Path: + skill_dir = root / skills_dir / name + if skill_dir.exists(): + if not force: + raise SystemExit(f"Skill already installed at {skill_dir}. Use --force to reinstall.") + shutil.rmtree(skill_dir) + skill_dir.mkdir(parents=True, exist_ok=True) + for rel_path, data in bundle.items(): + target = skill_dir / rel_path + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(data) + (skill_dir / _MANAGED_MARKER_FILE).touch() + return skill_dir + + +def _has_local_changes(skill_dir: Path, bundle: dict[str, bytes]) -> bool: + """True if the installed skill has any file that differs from `bundle` or has extra files. + + The marker file is ignored. Compares raw bytes so a whitespace-only edit still counts as dirty. + """ + on_disk: dict[str, bytes] = {} + for path in skill_dir.rglob("*"): + if not path.is_file(): + continue + rel = str(path.relative_to(skill_dir)) + if rel == _MANAGED_MARKER_FILE: + continue + on_disk[rel] = path.read_bytes() + return on_disk != bundle + + +def _discover_installed(root: Path) -> list[tuple[Path, str]]: + """Return `(skills_dir, name)` pairs for every managed install under `root`.""" + found: list[tuple[Path, str]] = [] + for skills_dir in _ALL_INSTALL_DIRS: + skills_root = root / skills_dir + if not skills_root.exists(): + continue + for d in sorted(skills_root.iterdir()): + if d.is_dir() and (d / _MANAGED_MARKER_FILE).exists(): + found.append((skills_dir, d.name)) + return found + + +class SkillsCommand(BaseDiffusersCLICommand): + @staticmethod + def register_subcommand(subparsers: _SubParsersAction) -> None: + parser: ArgumentParser = subparsers.add_parser( + "skills", + help="Manage Agent Skills for AI assistants.", + usage="\n diffusers-cli skills [options]", + ) + parser._optionals.title = "Options" + actions = parser.add_subparsers(dest="skills_action", required=True, metavar="") + + add = actions.add_parser("add", help="Download and install a skill.") + add.add_argument( + "name", + nargs="?", + default=None, + help="Skill name (e.g. diffusers-cli, custom-blocks). Omit and pass --all to install every skill.", + ) + add.add_argument( + "--all", + dest="install_all", + action="store_true", + help="Install every skill in the registry. Mutually exclusive with a positional name.", + ) + add.add_argument( + "--global", + "-g", + dest="install_global", + action="store_true", + help="Install globally (user-level) instead of in the current project directory.", + ) + add.add_argument("--force", action="store_true", help="Overwrite existing skills in the destination.") + add.set_defaults(func=SkillsCommand) + + list_action = actions.add_parser("list", help="List available skills in the registry.") + list_action.set_defaults(func=SkillsCommand) + + update = actions.add_parser("update", help="Re-download and reinstall managed skills.") + update.add_argument( + "name", + nargs="?", + default=None, + help="Optional installed skill name to update. Omit to update every managed skill.", + ) + update.add_argument( + "--global", + "-g", + dest="install_global", + action="store_true", + help="Update skills installed globally (user-level) instead of the current project.", + ) + update.add_argument( + "--force", + action="store_true", + help="Overwrite skills even if they have local modifications since install.", + ) + update.set_defaults(func=SkillsCommand) + + preview = actions.add_parser("preview", help="Print a skill's SKILL.md from the registry.") + preview.add_argument("name", help="Skill name to preview.") + preview.set_defaults(func=SkillsCommand) + + def __init__(self, args: Namespace): + self.args = args + + def run(self) -> None: + if self.args.skills_action == "add": + self._add() + elif self.args.skills_action == "list": + self._list() + elif self.args.skills_action == "update": + self._update() + elif self.args.skills_action == "preview": + self._preview() + + def _add(self) -> None: + if self.args.install_all and self.args.name: + raise SystemExit("--all and a positional skill name are mutually exclusive.") + if not self.args.install_all and not self.args.name: + raise SystemExit("Pass a skill name (e.g. diffusers-cli) or --all to install every skill.") + + root = Path.home() if self.args.install_global else Path.cwd() + install_dirs = _detect_install_dirs() + names = self._resolve_names() + + installed: list[str] = [] + failed: list[str] = [] + for name in names: + try: + bundle = _download_skill_bundle(name) + for skills_dir in install_dirs: + _install_skill(name, bundle, root, skills_dir, self.args.force) + installed.append(name) + except (SystemExit, httpx.HTTPError) as e: + # Downgrade to a warning so one broken skill doesn't abort the batch. + logger.warning(f"Skipping skill {name!r}: {e}") + failed.append(name) + + if not installed: + raise SystemExit(f"No skills installed. Failed: {failed}") + out.result( + f"Installed {len(installed)} skill(s)", + installed=", ".join(installed), + failed=", ".join(failed) if failed else None, + paths=", ".join(str(root / d) for d in install_dirs), + ) + + def _update(self) -> None: + root = Path.home() if self.args.install_global else Path.cwd() + installed = _discover_installed(root) + if self.args.name is not None: + installed = [entry for entry in installed if entry[1] == self.args.name] + if not installed: + raise SystemExit(f"No installed skill named {self.args.name!r} found under {root}.") + if not installed: + raise SystemExit(f"No managed skills found under {root}.") + + # Group by skill name so we redownload each bundle once even if it's installed to + # multiple locations (e.g. both .claude/skills/ and .agents/skills/). + by_name: dict[str, list[Path]] = {} + for skills_dir, name in installed: + by_name.setdefault(name, []).append(skills_dir) + + updated: list[str] = [] + failed: list[str] = [] + skipped: list[str] = [] + for name, dirs in sorted(by_name.items()): + try: + bundle = _download_skill_bundle(name) + for skills_dir in dirs: + skill_dir = root / skills_dir / name + if not self.args.force and _has_local_changes(skill_dir, bundle): + logger.warning( + f"Skill {name!r} at {skill_dir} has local modifications; " + "skipping. Pass --force to overwrite them." + ) + skipped.append(name) + continue + _install_skill(name, bundle, root, skills_dir, force=True) + updated.append(name) + except (SystemExit, httpx.HTTPError) as e: + logger.warning(f"Skipping skill {name!r}: {e}") + failed.append(name) + + out.result( + f"Updated {len(updated)} skill(s)", + updated=", ".join(updated), + skipped=", ".join(skipped) if skipped else None, + failed=", ".join(failed) if failed else None, + ) + + def _preview(self) -> None: + bundle = _download_skill_bundle(self.args.name) + skill_md = bundle.get("SKILL.md") + if skill_md is None: + raise SystemExit(f"Skill {self.args.name!r} has no SKILL.md in the registry.") + print(skill_md.decode()) + + def _list(self) -> None: + entries = _fetch_json(_registry_url()) + skills = [{"name": e["name"]} for e in entries if e["type"] == "dir" and not e["name"].startswith(".")] + if not skills: + raise SystemExit("No skills found in registry.") + out.table(skills, headers=["name"]) + + def _resolve_names(self) -> list[str]: + if self.args.install_all: + entries = _fetch_json(_registry_url()) + return sorted(e["name"] for e in entries if e["type"] == "dir" and not e["name"].startswith(".")) + return [self.args.name] diff --git a/tests/others/test_cli_commands.py b/tests/others/test_cli_commands.py new file mode 100644 index 0000000000..9e9ff00c41 --- /dev/null +++ b/tests/others/test_cli_commands.py @@ -0,0 +1,306 @@ +# 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. +"""Essential unit tests for the non-``skills`` CLI commands. + +One test per contract that would ship broken if regressed. Grouped by command. +""" + +import subprocess +from argparse import ArgumentParser, Namespace + +import pytest + +from diffusers.commands.custom_blocks import CustomBlocksCommand +from diffusers.commands.run import ( + _build_task_kwargs, + _kwargs_to_argv, + _load_pipeline, + _parse_pipeline_kwargs, + _resolve_dtype, + _resolve_media_inputs, + _upload_inputs_to_sandbox, +) +from diffusers.commands.schema import _parse_docstring_args +from diffusers.utils.testing_utils import require_accelerator, require_torch_gpu + + +AVAILABLE_COMMANDS = ("env", "fp16_safetensors", "custom_blocks", "run", "schema", "skills") + + +class TestRunCommand: + def test_parse_pipeline_kwargs(self): + assert _parse_pipeline_kwargs('{"prompt": "a cat", "steps": 50}') == {"prompt": "a cat", "steps": 50} + with pytest.raises(SystemExit, match="must be valid JSON"): + _parse_pipeline_kwargs('{"prompt": "unterminated') + with pytest.raises(SystemExit, match="must decode to a JSON object"): + _parse_pipeline_kwargs('["not", "an", "object"]') + + def test_resolve_dtype_unknown_raises(self): + with pytest.raises(ValueError, match="Unknown dtype"): + _resolve_dtype("dummy-dtype") + + def test_resolve_media_inputs(self, monkeypatch): + monkeypatch.setattr("diffusers.commands.run.load_image", lambda v: f"img({v})") + monkeypatch.setattr("diffusers.commands.run.load_video", lambda v: [f"frame({v})"]) + # Scalar strings load individually; string-lists load each entry (batched inputs); + # `prompt` isn't a media key so it passes through; non-string / non-string-list values + # (e.g. a pre-loaded PIL object modeled here as a dict) pass through untouched. + preloaded = {"pre": "loaded"} + kwargs = { + "image": "url1", + "control_video": "url2", + "prompt": "text", + "mask_image": ["a.png", "b.png"], + "control_image": ["c.png"], + "video": ["v1.mp4", "v2.mp4"], + "ip_adapter_image": preloaded, + } + _resolve_media_inputs(kwargs) + assert kwargs == { + "image": "img(url1)", + "control_video": ["frame(url2)"], + "prompt": "text", + "mask_image": ["img(a.png)", "img(b.png)"], + "control_image": ["img(c.png)"], + "video": [["frame(v1.mp4)"], ["frame(v2.mp4)"]], + "ip_adapter_image": preloaded, + } + + def test_resolve_media_inputs_audio_batch(self, monkeypatch): + # Batched audio: each entry loads, the paired sampling-rate kwarg defaults to the first entry's rate. + monkeypatch.setattr("diffusers.commands.run._load_audio", lambda v: (f"wave({v})", 44100)) + kwargs = {"initial_audio_waveforms": ["a.wav", "b.wav"]} + _resolve_media_inputs(kwargs) + assert kwargs == { + "initial_audio_waveforms": ["wave(a.wav)", "wave(b.wav)"], + "initial_audio_sampling_rate": 44100, + } + + def test_upload_inputs_to_sandbox_batched(self, tmp_path): + # Batched media on --remote: `_upload_inputs_to_sandbox` uploads each local path in a list + # and rewrites the JSON in place. URLs and non-existent paths pass through untouched. + import json as _json + + local_a = tmp_path / "a.png" + local_b = tmp_path / "b.png" + local_a.write_bytes(b"a") + local_b.write_bytes(b"b") + + uploads: list[tuple[str, str]] = [] + + class FakeFiles: + def upload(self, src, dst): + uploads.append((src, dst)) + + class FakeSbx: + files = FakeFiles() + + args = Namespace( + pipeline_kwargs=_json.dumps( + { + "image": str(local_a), # scalar local → uploaded + "mask_image": [str(local_a), "https://example.com/x.png", str(local_b)], # mixed batch + "prompt": "unchanged", + } + ) + ) + _upload_inputs_to_sandbox(args, FakeSbx(), run_id="rid") + + assert len(uploads) == 3 # a (scalar), a (list[0]), b (list[2]); URL skipped + rewritten = _json.loads(args.pipeline_kwargs) + assert rewritten["image"].startswith("/tmp/diffusers-cli/inputs/rid/image_") + assert rewritten["mask_image"][0].startswith("/tmp/diffusers-cli/inputs/rid/mask_image_0_") + assert rewritten["mask_image"][1] == "https://example.com/x.png" # URL untouched + assert rewritten["mask_image"][2].startswith("/tmp/diffusers-cli/inputs/rid/mask_image_2_") + assert rewritten["prompt"] == "unchanged" + + def test_kwargs_to_argv(self): + argv = _kwargs_to_argv( + "run", {"cpu_offload": "model", "vae_tiling": True, "dependencies": ["torch", "accelerate"]} + ) + assert argv[0] == "run" + assert "--cpu-offload" in argv and argv[argv.index("--cpu-offload") + 1] == "model" + assert "--vae-tiling" in argv + assert argv.count("--dependencies") == 2 and "torch" in argv and "accelerate" in argv + + def test_remote_argv_omits_hf_job_args(self): + # `_build_task_kwargs` must strip REMOTE_KEYS (flags that control how the remote run is + # dispatched, not the sandbox CLI) plus None/False values before we forward argv to the sandbox. + args = Namespace( + remote=True, + flavor="a100-large", + timeout="10m", + format="auto", + image="pytorch/pytorch:2.10.0", + func=object(), + model="x", + dtype="bf16", + device_map="cuda", # not in REMOTE_KEYS: forwarded to the sandbox + revision=None, + trust_remote_code=False, + vae_tiling=True, + ) + assert _build_task_kwargs(args) == { + "model": "x", + "dtype": "bf16", + "device_map": "cuda", + "vae_tiling": True, + } + + # ----------------------------------------------------------------------- + # Route flags through the CLI parser and call `_load_pipeline` directly to assert + # their effect on a real tiny pipeline. `hf-internal-testing/tiny-flux-pipe` is small + # enough to load without a GPU and is already used across the pipeline suite. + # ----------------------------------------------------------------------- + + pretrained_model_name_or_path = "hf-internal-testing/tiny-flux-pipe" + + def _parse_run_argv(self, extra_argv: list[str]) -> Namespace: + from diffusers.commands.run import RunCommand + + parser = ArgumentParser() + subparsers = parser.add_subparsers() + RunCommand.register_subcommand(subparsers) + return parser.parse_args( + [ + "run", + "--model", + self.pretrained_model_name_or_path, + "--pipeline-kwargs", + '{"prompt": "a cat"}', + *extra_argv, + ] + ) + + @require_torch_gpu + def test_group_offload_arg(self): + from diffusers.hooks.group_offloading import _is_group_offload_enabled + + args = self._parse_run_argv(["--cpu-offload", "group"]) + pipeline = _load_pipeline(args) + assert _is_group_offload_enabled(pipeline.transformer) + + @require_accelerator + def test_model_cpu_offload_arg(self): + import accelerate + + args = self._parse_run_argv(["--cpu-offload", "model"]) + pipeline = _load_pipeline(args) + assert isinstance(pipeline.transformer._hf_hook, accelerate.hooks.CpuOffload) + + def test_vae_tiling_arg(self): + args = self._parse_run_argv(["--vae-tiling"]) + pipeline = _load_pipeline(args) + assert pipeline.vae.use_tiling is True + + def test_vae_slicing_arg(self): + args = self._parse_run_argv(["--vae-slicing"]) + pipeline = _load_pipeline(args) + assert pipeline.vae.use_slicing is True + + @require_torch_gpu + def test_compile_arg(self): + args = self._parse_run_argv(["--compile"]) + pipeline = _load_pipeline(args) + # `_compile_denoiser` either applies `torch.compile` (which wraps the module in an + # `OptimizedModule` with an `_orig_mod` attribute), or calls `compile_repeated_blocks` + # which mutates the transformer's inner blocks in place. Either path leaves + # `_orig_mod` on either the transformer itself or on one of its repeated blocks. + has_torch_compile_wrapper = hasattr(pipeline.transformer, "_orig_mod") or any( + hasattr(m, "_orig_mod") for m in pipeline.transformer.modules() + ) + assert has_torch_compile_wrapper + + @require_torch_gpu + def test_attention_backend_arg(self): + args = self._parse_run_argv(["--attention-backend", "flash_hub"]) + pipeline = _load_pipeline(args) + # `set_attention_backend` stamps each attention processor's `_attention_backend` attr. + backends = { + m.processor._attention_backend + for m in pipeline.transformer.modules() + if hasattr(m, "processor") and hasattr(m.processor, "_attention_backend") + } + assert any(b.value == "flash_hub" for b in backends) + + +class TestSchemaCommand: + pretrained_model_name_or_path = "hf-internal-testing/tiny-flux-pipe" + + def test_parse_docstring_args(self): + docstring = """Description. + + Args: + prompt (str): The prompt text + wraps across multiple lines + for readability. + steps (int, optional): Steps to run. + """ + result = _parse_docstring_args(docstring) + assert result["steps"] == "Steps to run." + assert "wraps across multiple lines" in result["prompt"] + assert "\n" not in result["prompt"] + + def test_schema(self, monkeypatch): + # End-to-end: parse real argv → SchemaCommand.run → capture the emitted payload and + # verify it contains the pipeline class + at least a `prompt` input parsed from the + # pipeline's `__call__` signature. + from huggingface_hub.cli._output import OutputFormat, out + + from diffusers.commands.schema import SchemaCommand + + captured: dict = {} + monkeypatch.setattr(out, "dict", lambda payload: captured.update(payload)) + + parser = ArgumentParser() + subparsers = parser.add_subparsers() + SchemaCommand.register_subcommand(subparsers) + args = parser.parse_args(["schema", "-m", self.pretrained_model_name_or_path]) + + out.set_mode(OutputFormat.json) + args.func(args).run() + + assert captured["pipeline_class"] == "FluxPipeline" + assert captured["model"] == self.pretrained_model_name_or_path + input_names = [p["name"] for p in captured["inputs"]] + assert "prompt" in input_names + + +class TestCustomBlocksCommand: + def test_class_discovery(self, tmp_path): + block_py = tmp_path / "block.py" + block_py.write_text( + "class OtherBase:\n pass\n" + "class NotABlock(OtherBase):\n pass\n" + "class MyBlock(ModularPipelineBlocks):\n pass\n" + ) + cmd = CustomBlocksCommand() + assert cmd._get_class_names(block_py) == [("MyBlock", "ModularPipelineBlocks")] + + broken = tmp_path / "broken.py" + broken.write_text("class Broken(:\n pass\n") + with pytest.raises(ValueError, match="Could not parse"): + cmd._get_class_names(broken) + + +class TestCli: + def test_toplevel_help_lists_all_commands(self): + result = subprocess.run( + ["python", "-m", "diffusers.commands.diffusers_cli", "--help"], + capture_output=True, + text=True, + ) + assert result.returncode == 0 + for cmd in AVAILABLE_COMMANDS: + assert cmd in result.stdout diff --git a/tests/others/test_cli_skills.py b/tests/others/test_cli_skills.py new file mode 100644 index 0000000000..d5e87ed61d --- /dev/null +++ b/tests/others/test_cli_skills.py @@ -0,0 +1,103 @@ +# 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.commands.skills``. + +Exercises the install / discovery / agent-detection logic. The GitHub Contents API path +(``_fetch_json`` / ``_download_skill_bundle``) is left out — it needs the network. +""" + +import pytest + +from diffusers.commands.skills import ( + _AGENT_ENV_TO_DIR, + _AGENTS_SKILLS_DIR, + _ALL_INSTALL_DIRS, + _CLAUDE_SKILLS_DIR, + _MANAGED_MARKER_FILE, + _detect_install_dirs, + _discover_installed, + _install_skill, +) + + +SAMPLE_BUNDLE = { + "SKILL.md": b"---\nname: example-skill\ndescription: An example skill.\n---\n\n# Example\n", + "extra.md": b"# Extra reference file\n", + "reference/script.py": b"print('hello from a nested reference')\n", +} + + +class TestInstallSkill: + def test_writes_bundle_and_marker(self, tmp_path): + location = _install_skill("example-skill", SAMPLE_BUNDLE, tmp_path, _CLAUDE_SKILLS_DIR, force=False) + skill_dir = tmp_path / _CLAUDE_SKILLS_DIR / "example-skill" + assert location == skill_dir + assert (skill_dir / "SKILL.md").read_bytes() == SAMPLE_BUNDLE["SKILL.md"] + assert (skill_dir / "extra.md").read_bytes() == SAMPLE_BUNDLE["extra.md"] + assert (skill_dir / "reference" / "script.py").read_bytes() == SAMPLE_BUNDLE["reference/script.py"] + assert (skill_dir / _MANAGED_MARKER_FILE).exists() + + def test_errors_without_force_on_existing_dir(self, tmp_path): + _install_skill("example-skill", SAMPLE_BUNDLE, tmp_path, _CLAUDE_SKILLS_DIR, force=False) + with pytest.raises(SystemExit, match="Use --force to reinstall"): + _install_skill("example-skill", SAMPLE_BUNDLE, tmp_path, _CLAUDE_SKILLS_DIR, force=False) + + def test_force_replaces_existing(self, tmp_path): + _install_skill("example-skill", SAMPLE_BUNDLE, tmp_path, _CLAUDE_SKILLS_DIR, force=False) + new_bundle = {"SKILL.md": b"different content"} + _install_skill("example-skill", new_bundle, tmp_path, _CLAUDE_SKILLS_DIR, force=True) + skill_dir = tmp_path / _CLAUDE_SKILLS_DIR / "example-skill" + assert (skill_dir / "SKILL.md").read_bytes() == b"different content" + assert not (skill_dir / "extra.md").exists() + assert (skill_dir / _MANAGED_MARKER_FILE).exists() + + +class TestDetectInstallDirs: + def test_no_env_installs_everywhere(self, monkeypatch): + for var in _AGENT_ENV_TO_DIR: + monkeypatch.delenv(var, raising=False) + assert _detect_install_dirs() == _ALL_INSTALL_DIRS + + @pytest.mark.parametrize( + ("env_var", "expected"), + [ + ("CLAUDECODE", _CLAUDE_SKILLS_DIR), + ("CLAUDE_CODE", _CLAUDE_SKILLS_DIR), + ("CODEX_SANDBOX", _AGENTS_SKILLS_DIR), + ("CURSOR_AI", _AGENTS_SKILLS_DIR), + ], + ) + def test_env_var_picks_specific_target(self, monkeypatch, env_var, expected): + for var in _AGENT_ENV_TO_DIR: + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv(env_var, "1") + assert _detect_install_dirs() == (expected,) + + +class TestDiscoverInstalled: + def test_returns_managed_installs_from_all_paths(self, tmp_path): + _install_skill("skill-a", SAMPLE_BUNDLE, tmp_path, _CLAUDE_SKILLS_DIR, force=False) + _install_skill("skill-b", SAMPLE_BUNDLE, tmp_path, _AGENTS_SKILLS_DIR, force=False) + # User-placed skill without our marker — should be ignored. + unmanaged = tmp_path / _CLAUDE_SKILLS_DIR / "user-placed" + unmanaged.mkdir(parents=True) + (unmanaged / "SKILL.md").write_bytes(b"user content") + + assert _discover_installed(tmp_path) == [ + (_CLAUDE_SKILLS_DIR, "skill-a"), + (_AGENTS_SKILLS_DIR, "skill-b"), + ] + + def test_returns_empty_when_no_skills_dir(self, tmp_path): + assert _discover_installed(tmp_path) == []