diff --git a/CLAUDE.md b/CLAUDE.md index df76ecd8..1808633b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -23,7 +23,7 @@ python -m dw.run workflows/templates/text-to-image.json prompt="a cat" num_image # Validate a workflow against schema python -m dw.validate workflows/models/z-image.json -# Basic system test (torch, diffusers import check) +# System test - downloads SD 1.5 (a few GB) and generates one image python -m dw.test # Interactive REPL @@ -235,7 +235,7 @@ SDNQ pre-quantized models use a different pattern: `pre_load_modules` imports sd `dw/__init__.py` handles device detection (CUDA > MPS > CPU) and platform-specific optimizations: - **CUDA**: TF32 matmul, cuDNN benchmark, deterministic mode (configurable via settings) -- **MPS**: `PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0` (use all unified memory), autocast warnings suppressed, attention slicing enabled by default +- **MPS**: `PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0` (use all unified memory), autocast warnings suppressed; attention slicing automatic (faster for SD 1.5's head dims, 2.4x slower for SDXL's - `disable_attention_slicing`) - **CPU**: Warning displayed Detection is overridden by the `DW_DEVICE` environment variable (single run) or the `device` setting (standing), either of which can name a specific accelerator such as `cuda:1`. Device placement is explicit throughout — no default torch device is set, since that would build models directly in VRAM and defeat offloading. Compare backends with `get_device_type()` rather than `== "cuda"`, which a device like `cuda:1` would fail. @@ -244,6 +244,8 @@ A step can override the device it runs on: `device` in a pipeline `configuration Every device a workflow names passes through `resolve_device()`, which translates a backend this machine does not have into the one it does and warns — a `cuda` workflow runs on a Mac and an `mps` one runs on a CUDA box. Only the backend is translated: an index survives when the backend matches (`cuda:1` on a single-GPU CUDA box stays a genuine error) and is dropped when it does not. `cpu` is never rewritten, since pinning a step to the CPU is how a GPU-specific problem gets ruled out. Translation happens before anything reads the backend, so the MPS accommodations (the sequential-offload downgrade, attention slicing, the compile skip) fire for a translated device too. +The same translation reaches the settings that carry a device or a CUDA-only feature, so a template written on the CUDA box runs unchanged on a Mac and a caller (an MCP agent included) never has to know the backend: SDNQ `quantization_device`/`return_device` go through `resolve_device()`, and `use_quantized_matmul(_conv)` is turned off on MPS, where it falls back to `torch._int_mm`, ~500x slower (`portable_quantization_arguments`, `config_objects.py`). Group-offload `use_stream`/`record_stream` are dropped when the onload device is not CUDA/XPU. `vram_estimate` checks against the serving device's own `cost` entries, else its `device_capacity_gb()` (Metal's recommended working set on a Mac), else every entry. `device_memory_stats()` reports real unified-memory figures on MPS. Each adaptation logs a warning, and torch's MPS CPU-fallback warning is let through the blanket `UserWarning` filter. `HF_ENABLE_PARALLEL_LOADING` defaults to off on macOS (`_parallel_loading_default`): diffusers' per-shard loader threads copying onto MPS at once segfaulted the LTX-2.5 SDNQ transformer load; it is chosen by platform because diffusers reads it at import, before dw can ask torch for a device. + A `components` entry can additionally set `residency: "on_demand"`, which rests the component on the CPU and wraps its `forward`/`encode`/`decode` to move it to the device around each call (`apply_on_demand_placement` in `pipeline.py`). The wrappers use `functools.wraps` because callers introspect the signature — MiniMax H3's denoiser picks its arguments from `signature(transformer.forward)`. It is mutually exclusive with `group_offload` on the same component, and like `group_offload` it suppresses the wholesale `pipeline.to(device)` at load. Settings in `~/.diffusers_helper/settings.json`: `device`, `enable_tf32`, `cudnn_benchmark`, `cudnn_deterministic`, `log_level`, `log_filename`. @@ -321,7 +323,7 @@ same reason - default setup cannot load a pack. - **Cartesian product explosion** — multiple `previous_result` references multiply: 4 images × 3 masks = 12 iterations - **Component sharing requires exact key matching** between `shared_components` and `reused_components` - **Built-in workflows** need explicit argument mapping: `"prompt": "variable:prompt"` -- **MPS differences from CUDA**: no autocast, no bitsandbytes, no flash_attn, no triton, no torch.compile. Model offloading has less benefit on unified memory, and `"offload": "sequential"` is downgraded to `"model"` with a warning there (`place_component`) — per-submodule streaming hands back no residency when the CPU and the accelerator share one pool. `exclude_from_cpu_offload` is sequential-only and does not survive the downgrade. +- **MPS differences from CUDA**: no bitsandbytes, no flash_attn, no triton, no torch.compile (`torch.autocast("mps")` works on torch 2.14, but dw does not use it). Model offloading has less benefit on unified memory, and `"offload": "sequential"` is downgraded to `"model"` with a warning there (`place_component`) — per-submodule streaming hands back no residency when the CPU and the accelerator share one pool. `exclude_from_cpu_offload` is sequential-only and does not survive the downgrade. - **`{}`-escaped strings** in JSON arguments: `"{nf4}"` stays as string `"nf4"`, without braces it would try to load as a type - **A stored prompt's `text` may not begin with a reference prefix** (`variable:`, `previous_result:`, `constant:`, `asset:`, `output:`, `prompt:`) — the engine rejects it to prevent double resolution or iteration expansion - **Audio+video muxing**: pipelines that generate audio alongside video (LTX-2) have the two muxed into one `video/mp4` file with PyAV in `result.py` @@ -573,7 +575,7 @@ same reason - default setup cannot load a pack. stops the next instance. The half that fixes this one: every template whose deliverable ends in a `pair_audio` mux (`music-video`, `assemble-and-score`, `dissolve-between-shots`) normalizes to **-3 dBFS**; - `music`, an mp3, keeps -1 + `music`, an mp3, normalizes to -3 as well (#362) - **A variable's bound is declared by the author, checked three times** — a model's own rule about a value (H3's `num_frames` is `17 * n + 5` from 124 to 345) is a property of the model, so it lives in the workflow rather than diff --git a/README.md b/README.md index e8e2ae23..c7e352a9 100644 --- a/README.md +++ b/README.md @@ -32,8 +32,8 @@ source ./activate .\venv\scripts\activate ``` -`python -m dw.test` confirms torch and diffusers import and shows which -accelerator was found. +`python -m dw.test` runs a small built-in workflow end to end - it downloads +SD 1.5 (a few GB) and generates one image on whichever accelerator was found. **2. Start the engine.** Leave it running; everything else talks to it. diff --git a/docs/ACCELERATION.md b/docs/ACCELERATION.md index 6c907d21..daa79ed4 100644 --- a/docs/ACCELERATION.md +++ b/docs/ACCELERATION.md @@ -215,7 +215,7 @@ Prefer the pinned form for a compiled component - the per-call context manager s } ``` -Processes attention in slices to reduce memory at some cost to speed. Enabled automatically on MPS (unified memory benefits from slicing) unless `disable_attention_slicing` is set. Modular pipelines have no `enable_attention_slicing()` method - the setting is silently skipped rather than failing when the pipeline doesn't support it. +Processes attention in slices to reduce memory at some cost to speed. Automatic on MPS unless `disable_attention_slicing` is set, opt-in elsewhere. Only UNet/ControlNet models take it. On MPS, whether it helps depends on the model's attention head dim: MPS SDPA is slow at 40, 48 and 160 (slicing made SD 1.5 ~20% faster end to end) and fast at 64-128 (slicing made SDXL-shaped attention 2.4x slower - set `disable_attention_slicing` for SDXL-family models). Modular pipelines have no `enable_attention_slicing()` method - the setting is silently skipped rather than failing when the pipeline doesn't support it. ## torch.compile @@ -344,9 +344,13 @@ For faster model downloads, optionally `pip install hf_transfer` and set `HF_HUB Apple Silicon has narrower acceleration support than CUDA: -- No flash-attn, no Triton, no bitsandbytes - `attention_backend` is effectively CUDA-only; use `"native"`-family backends or leave it unset on MPS. `compile` is skipped with a warning (inductor support on MPS is immature). -- No `torch.autocast` support - autocast-related warnings from other libraries are suppressed automatically rather than surfaced. -- `enable_attention_slicing` is on by default (set `disable_attention_slicing` to turn it off). -- `float16` produces NaN values on Apple Silicon - use `float32` or `bfloat16` for `torch_dtype` instead; dw only warns, it doesn't override the dtype for you. +- No flash-attn, no Triton, no bitsandbytes - `attention_backend` is effectively CUDA-only; use `"native"`-family backends or leave it unset on MPS. `compile` is skipped with a warning (inductor support on MPS is immature). A workflow that loads a bitsandbytes checkpoint (`flux2-dev`, `multi-image-reference`) or a TorchAO int4 config (`flux-torchao`) has no Mac path yet. +- Workflows written for CUDA are adapted rather than refused: a `cuda` device becomes `mps`, SDNQ's `quantization_device` follows it and its quantized matmul is switched off (it runs through `torch._int_mm`, ~500x slower on MPS), group-offload CUDA streams are dropped, and `"offload": "sequential"` becomes `"model"`. Each change is logged as a warning. +- Attention slicing is automatic; set `disable_attention_slicing` for SDXL-family models (head dim 64), where it is 2.4x slower in attention. It helps SD 1.5. +- `float16` can produce NaN values on Apple Silicon - use `bfloat16`; dw only warns, it doesn't override the dtype for you. - `PYTORCH_MPS_HIGH_WATERMARK_RATIO` defaults to `0.0` (use all unified memory) unless already set in the environment. +- Ops the MPS backend lacks: with `PYTORCH_ENABLE_MPS_FALLBACK=1` (set by the `fp4-fp8-for-torch-mps` package install.sh adds) they run on the CPU, and dw logs torch's warning when one does. Without it they raise. +- Memory figures are real: allocated, driver-reserved, and Metal's recommended working set as the total. `vram_estimate` checks a Mac against that total, since the catalog's `cost` entries are CUDA cards; the per-voxel figures were calibrated on CUDA. +- Cost estimates quoted on a Mac are the CUDA figures (`basis: "other_device"`) until the Mac's own runs build observed history - expect them to be optimistic. +- Parallel checkpoint loading (`HF_ENABLE_PARALLEL_LOADING`) defaults to off on macOS - concurrent loader threads copying onto MPS segfaulted a quantized load. Set it explicitly to override. - Offloading has less benefit than on CUDA, since unified memory is already shared between CPU and GPU. diff --git a/docs/DEPENDENCIES.md b/docs/DEPENDENCIES.md index 606fe19f..40afc755 100644 --- a/docs/DEPENDENCIES.md +++ b/docs/DEPENDENCIES.md @@ -55,7 +55,7 @@ pyproject; install.ps1 installs it explicitly on Windows. **Windows (CUDA):** bitsandbytes, kernels -**macOS (MPS):** fp4-fp8-for-torch-mps (FP8/FP4 dtype support for Metal), fluidtop +**macOS (MPS):** fp4-fp8-for-torch-mps (FP8/FP4 dtypes for Metal; autoloads on every torch import, overrides mm/linear/copy on the MPS dispatch key and sets PYTORCH_ENABLE_MPS_FALLBACK=1), fluidtop (a GPU/power TUI; needs sudo) ## Optional diff --git a/docs/TESTING.md b/docs/TESTING.md index 2a2f46c9..2c88b41f 100644 --- a/docs/TESTING.md +++ b/docs/TESTING.md @@ -63,7 +63,7 @@ An autouse `_clear_task_model_cache` fixture clears `dw.tasks.model_cache` befor ## Quick Validation ```bash -# Verify torch and diffusers are working +# End-to-end check: downloads SD 1.5 and generates one image python -m dw.test # Validate a workflow against schema diff --git a/docs/WORKFLOW_GUIDE.md b/docs/WORKFLOW_GUIDE.md index a60e0d63..1241f579 100644 --- a/docs/WORKFLOW_GUIDE.md +++ b/docs/WORKFLOW_GUIDE.md @@ -1289,8 +1289,9 @@ stays alive for the steps that reuse it. } ``` -- `enable_attention_slicing` — process attention in slices to reduce memory. Enabled - automatically on MPS unless `disable_attention_slicing` is set. +- `enable_attention_slicing` — process attention in slices to reduce memory at some + cost to speed. Automatic on MPS unless `disable_attention_slicing` is set + (UNet/ControlNet models only; worth disabling for SDXL on a Mac). - `attention_backend` — selects a diffusers attention backend (e.g. `"flash_hub"`) for the duration of each pipeline call. - `prompt_weighting` — enables A1111-style prompt weighting (`(word:1.5)`, `[word]`, diff --git a/docs/superpowers/plans/2026-09-25-mac-mps-support.md b/docs/superpowers/plans/2026-09-25-mac-mps-support.md new file mode 100644 index 00000000..423ef1e2 --- /dev/null +++ b/docs/superpowers/plans/2026-09-25-mac-mps-support.md @@ -0,0 +1,1144 @@ +# Mac (MPS) Support Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Make the catalog's CUDA-authored workflows load and run sanely on Apple Silicon, fix the Mac-specific defects the 2026-09-25 review verified, and fix the one cross-platform break it found (`segment`). + +**Architecture:** Every fix is an engine-level translation at the point a CUDA-only setting is consumed (`dw/pipeline_processors/config_objects.py`, `dw/vram_estimate.py`, `dw/__init__.py`), in the same spirit as the existing `resolve_device()` and the sequential-to-model offload downgrade: the workflow JSON stays as written and the engine adapts it to the machine, warning when it does. No template is edited to become Mac-specific. + +**Tech Stack:** Python 3.14, torch 2.14 (MPS), diffusers 0.41.dev, SDNQ, transformers (SAM2), pytest. + +**Spec:** No separate spec. The source is the Mac review of 2026-09-25, summarized in *Background* below. Each claim there was verified on an M5 Pro with 64 GB (code read, or measured with `venv/bin/python`). + +## Background (the review findings this plan implements) + +| # | Finding (verified) | Task | +|---|---|---| +| 1 | SDNQ `quantization_device: "cuda"` passed through unresolved; SDNQ does `param.to("cuda")` during load → `AssertionError: Torch not compiled with CUDA enabled`. All 31 `ltx2/*` + `minimax/*` files. | 1 | +| 2 | SDNQ `use_quantized_matmul(_conv): true` routes through `torch._int_mm`, measured 1096 ms vs 2.1 ms bf16 on MPS (~500x). | 1 | +| 3 | `group_offload.use_stream` / `record_stream` passed through; diffusers raises `Using streams for data transfer requires a CUDA device` off CUDA/XPU, and `record_stream` without `use_stream` raises too. Same 31 files. | 2 | +| 4 | `dw/tasks/segment.py:109` reads `sam_inputs["reshaped_input_sizes"]`, a SAM v1 key the SAM2 processor never returns → KeyError on every backend; it also lands in the `mask_threshold` positional slot. The unit test mocks the processor. | 3 | +| 5 | `device_memory_stats()` reports zeros on MPS although `torch.mps.current_allocated_memory()`, `driver_allocated_memory()`, `recommended_max_memory()` all work. | 4 | +| 6 | `vram_estimate_errors` compares against every `cost` entry regardless of device; every entry is `cuda` 24 GB, so a 64 GB Mac is refused against an RTX 3090. Filtering to matching entries alone would silently *disable* the guard on a Mac (no `mps` entries exist). | 5 | +| 7 | Attention slicing is auto-enabled on MPS; measured under `inference_mode` it makes UNet attention 2.4x slower in bf16 (12.8 vs 5.4 ms) and 1.4x in fp32. Only UNet/ControlNet models have `set_attention_slice`. The fp16 warning text is missing a space and recommends fp32. | 6 | +| 8 | `dw/__init__.py:29` ignores every `UserWarning`; torch's "not currently supported on the MPS backend and will fall back to run on the CPU" is a `UserWarning`, and `PYTORCH_ENABLE_MPS_FALLBACK=1` is set by the `fp4-fp8-for-torch-mps` autoload (and Don's `~/.zprofile`). CPU fallbacks are therefore invisible. | 7 | +| 9 | Template defaults: `upscale-diffusion`/`upscale-spandrel` default input is `https://example.com/image.jpg` (not an image); `image-to-image` pulls its default from a third-party repo named `sage-permission-probe`; `flux-torchao.json:38` spells `Flux.1-Dev` (separate HF cache dir → ~23 GB duplicate download); `attention-processor`/`controlnet-component` use `runwayml/stable-diffusion-v1-5`; `dw/workflows/test.json` (the `dw.test` system test) runs SD 1.5 in float16. | 8 | +| 10 | Stale docs: CLAUDE.md "MPS: no autocast" (`torch.autocast("mps", bfloat16)` works on torch 2.14), CLAUDE.md "`music` keeps -1" (template is -3), `dw.test` described as an import check (it downloads SD 1.5 and generates), ACCELERATION.md MPS notes. | 6, 9 | + +## Global Constraints + +- Run tests with `venv/bin/python -m pytest` (torch is in the venv). Do **not** contact or deploy to lem, which is busy. +- Work on branch `feat/mac-mps-support` cut from `develop`. Commit after every task. +- The behavior on a CUDA box must be unchanged by every task. Each task has a test that pins this. +- Security rules in CLAUDE.md apply. No `shell=True`. Subprocess arguments are constant lists. +- Compare backends with `get_device_type()`, never `== "cuda"`. `cpu` is never rewritten. +- Every engine adaptation logs a `logger.warning` naming what it changed and why, as the sequential-offload downgrade in `place_component` does. There are no silent rewrites. +- Format touched Python with `venv/bin/ruff format ` before committing. +- Match the surrounding comment style: comments explain *why*, in full sentences, and reference the measurement or issue. + +## Review Focus + +1. **A step pinned to `cpu`:** SDNQ `quantization_device: "cpu"` stays `cpu`, a group offload whose onload device is `cpu` has its streams stripped without error, and quantized matmul is left on under DW_DEVICE=cpu. Tests are in Tasks 1 and 2. +2. **A CUDA box:** streams are kept, quantized matmul is kept, the device strings are unchanged, and the VRAM refusal is unchanged. Tests are in Tasks 1, 2 and 5. +3. **A cached pipeline reloaded, where the builder sees the same configuration dict twice:** stripping streams and resolving devices are idempotent. Tests are in Tasks 1 and 2. +4. **A `torch.mps` memory API that raises or is missing:** `device_memory_stats()` degrades to `None` figures and does not raise. Test is in Task 4. +5. **A Mac whose capacity can't be read:** the VRAM check falls back to the old all-entries behavior, which is conservative, rather than to no check. Test is in Task 5. + +--- + +### Task 1: SDNQ quantization settings adapt to the machine + +**Files:** +- Modify: `dw/pipeline_processors/config_objects.py` (imports at top; `create_quantization_config` at ~line 27) +- Test: `tests/test_config_objects.py` (class `TestQuantizationConfiguration`) + +**Interfaces:** +- Consumes: `dw.resolve_device(requested)`, `dw.get_device_type(device=None)` +- Produces: `portable_quantization_arguments(arguments: dict) -> dict`, which is module-level in `config_objects.py` and returns a new dict without mutating its input. `create_quantization_config` calls it. + +- [ ] **Step 1: Write the failing tests.** Append these to `class TestQuantizationConfiguration` in `tests/test_config_objects.py`: + +```python +def test_a_cuda_quantization_device_is_translated_on_a_mac(self, monkeypatch): + # A catalog template written on the CUDA box names "cuda" here, and SDNQ + # moves every weight to it inside from_pretrained + monkeypatch.setattr(dw, "backend_available", lambda backend: backend != "cuda") + monkeypatch.setattr(dw, "get_device", lambda: "mps") + + config = create_quantization_config( + quantization_definition(quantization_device="cuda", return_device="cpu") + ) + + assert config.kwargs["quantization_device"] == "mps" + assert config.kwargs["return_device"] == "cpu" + + +def test_a_cpu_quantization_device_is_never_rewritten(self, monkeypatch): + monkeypatch.setattr(dw, "backend_available", lambda backend: backend != "cuda") + monkeypatch.setattr(dw, "get_device", lambda: "mps") + + config = create_quantization_config( + quantization_definition(quantization_device="cpu") + ) + + assert config.kwargs["quantization_device"] == "cpu" + + +def test_quantized_matmul_is_switched_off_on_mps(self, monkeypatch): + # Without Triton SDNQ's quantized matmul is torch._int_mm, measured + # ~500x slower than the bf16 matmul it replaces on an M5 Pro + monkeypatch.setattr(dw, "backend_available", lambda backend: backend != "cuda") + monkeypatch.setattr(dw, "get_device", lambda: "mps") + + config = create_quantization_config( + quantization_definition( + use_quantized_matmul=True, use_quantized_matmul_conv=True + ) + ) + + assert config.kwargs["use_quantized_matmul"] is False + assert config.kwargs["use_quantized_matmul_conv"] is False + + +def test_quantized_matmul_is_kept_on_cuda(self, monkeypatch): + monkeypatch.setattr(dw, "backend_available", lambda backend: True) + monkeypatch.setattr(dw, "get_device", lambda: "cuda") + + config = create_quantization_config( + quantization_definition(quantization_device="cuda", use_quantized_matmul=True) + ) + + assert config.kwargs["quantization_device"] == "cuda" + assert config.kwargs["use_quantized_matmul"] is True + + +def test_quantized_matmul_is_kept_on_cpu(self, monkeypatch): + monkeypatch.setattr(dw, "get_device", lambda: "cpu") + + config = create_quantization_config( + quantization_definition(use_quantized_matmul=True) + ) + + assert config.kwargs["use_quantized_matmul"] is True + + +def test_adapting_does_not_mutate_the_definition(self, monkeypatch): + # A cached pipeline builds from the same definition a second time + monkeypatch.setattr(dw, "backend_available", lambda backend: backend != "cuda") + monkeypatch.setattr(dw, "get_device", lambda: "mps") + definition = quantization_definition( + quantization_device="cuda", use_quantized_matmul=True + ) + + create_quantization_config(definition) + second = create_quantization_config(definition) + + assert definition["arguments"] == { + "quantization_device": "cuda", + "use_quantized_matmul": True, + } + assert second.kwargs["quantization_device"] == "mps" +``` + +- [ ] **Step 2: Run the tests to confirm they fail.** + +Run: `venv/bin/python -m pytest tests/test_config_objects.py -k "quantization_device or quantized_matmul or mutate_the_definition" -v` +Expected: the Mac-translation, matmul-off and mutation tests FAIL, because `quantization_device` is still `"cuda"` and the matmul flags are still `True`. The CUDA and CPU keep-tests pass already. + +- [ ] **Step 3: Implement.** In `dw/pipeline_processors/config_objects.py`, change the import line and add the helper above `create_quantization_config`: + +```python +from .. import get_device_type, resolve_device +``` + +```python +# Quantization config arguments that name a device. A workflow authored on the +# CUDA box writes "cuda" here, and nothing downstream translates it - SDNQ moves +# every weight to it inside from_pretrained, which fails outright on a Mac +DEVICE_ARGUMENTS = ("quantization_device", "return_device") + +# SDNQ's quantized matmul has no Triton on MPS and falls back to torch._int_mm, +# measured at ~500x the bf16 matmul it replaces (1096 ms vs 2.1 ms, 512x3072 by +# 3072x3072 on an M5 Pro). The dequantize path is the fast one there +QUANTIZED_MATMUL_ARGUMENTS = ("use_quantized_matmul", "use_quantized_matmul_conv") + + +def portable_quantization_arguments(arguments): + """Adapt a quantization config's arguments to this machine without touching + the definition they came from - a cached pipeline builds from it again.""" + adapted = dict(arguments) + for key in DEVICE_ARGUMENTS: + if isinstance(adapted.get(key), str): + adapted[key] = resolve_device(adapted[key]) + if get_device_type() == "mps": + for key in QUANTIZED_MATMUL_ARGUMENTS: + if adapted.get(key): + logger.warning( + f"Turning off '{key}' on MPS - without Triton it runs through " + "torch._int_mm, which is far slower there than dequantizing" + ) + adapted[key] = False + return adapted +``` + +In `create_quantization_config`, replace `quantization_config["arguments"].items()` in the dict comprehension with `portable_quantization_arguments(quantization_config["arguments"]).items()`. + +- [ ] **Step 4: Run the tests to confirm they pass.** + +Run: `venv/bin/python -m pytest tests/test_config_objects.py -v` +Expected: all PASS, including the pre-existing tests. + +- [ ] **Step 5: Commit.** + +```bash +git checkout -b feat/mac-mps-support develop # first task only +venv/bin/ruff format dw/pipeline_processors/config_objects.py tests/test_config_objects.py +git add dw/pipeline_processors/config_objects.py tests/test_config_objects.py +git commit -m "fix(quantization): translate SDNQ devices and drop quantized matmul on MPS" +``` + +--- + +### Task 2: Group offload drops CUDA streams off CUDA + +**Files:** +- Modify: `dw/pipeline_processors/config_objects.py` (`get_group_offload_configuration`, ~line 91) +- Test: `tests/test_config_objects.py` (class `TestGroupOffloadConfiguration`) + +**Interfaces:** +- Consumes: the existing `get_group_offload_configuration(configuration, default_device)` +- Produces: the same signature. The returned dict has no `use_stream`/`record_stream` keys when its `onload_device.type` is not `cuda` or `xpu`. + +- [ ] **Step 1: Fix the existing test that assumes CUDA.** On a Mac, `test_other_keys_are_carried_through` would now lose its `use_stream`. Pin CUDA as available by adding `monkeypatch` to its signature and this line as its first statement: + +```python + monkeypatch.setattr(dw, "backend_available", lambda backend: True) +``` + +- [ ] **Step 2: Write the failing tests.** Append these to `class TestGroupOffloadConfiguration`: + +```python +def test_streams_are_dropped_on_mps(self, monkeypatch): + # diffusers refuses use_stream without CUDA/XPU, and refuses + # record_stream without use_stream - both have to go together + monkeypatch.setattr(dw, "backend_available", lambda backend: backend != "cuda") + monkeypatch.setattr(dw, "get_device", lambda: "mps") + + config = get_group_offload_configuration( + { + "group_offload": { + "onload_device": "cuda", + "num_blocks_per_group": 1, + "use_stream": True, + "record_stream": True, + } + }, + "cuda", + ) + + assert config["onload_device"] == torch.device("mps") + assert "use_stream" not in config + assert "record_stream" not in config + assert config["num_blocks_per_group"] == 1 + + +def test_streams_are_dropped_when_onloading_to_cpu(self): + config = get_group_offload_configuration( + {"group_offload": {"onload_device": "cpu", "use_stream": True}}, "cpu" + ) + + assert "use_stream" not in config + + +def test_streams_are_kept_on_cuda(self, monkeypatch): + monkeypatch.setattr(dw, "backend_available", lambda backend: True) + + config = get_group_offload_configuration( + {"group_offload": {"use_stream": True, "record_stream": True}}, "cuda" + ) + + assert config["use_stream"] is True + assert config["record_stream"] is True + + +def test_dropping_streams_twice_is_stable(self, monkeypatch): + monkeypatch.setattr(dw, "backend_available", lambda backend: backend != "cuda") + monkeypatch.setattr(dw, "get_device", lambda: "mps") + configuration = {"group_offload": {"use_stream": True, "record_stream": True}} + + get_group_offload_configuration(configuration, "cuda") + second = get_group_offload_configuration(configuration, "cuda") + + assert "use_stream" not in second + assert second["onload_device"] == torch.device("mps") +``` + +- [ ] **Step 3: Run the tests to confirm they fail.** + +Run: `venv/bin/python -m pytest tests/test_config_objects.py::TestGroupOffloadConfiguration -v` +Expected: the three "dropped" tests FAIL, because `use_stream` is still present. + +- [ ] **Step 4: Implement.** In `get_group_offload_configuration`, after the two `resolve_device` assignments and before `return group_offload_config`, add: + +```python + # CUDA streams overlap the next group's transfer with this one's compute. + # diffusers refuses them without CUDA or XPU, and refuses record_stream + # without use_stream, so both go together - a catalog template written + # for the CUDA box would otherwise fail to load on a Mac + if group_offload_config["onload_device"].type not in ("cuda", "xpu"): + dropped = [ + key + for key in ("use_stream", "record_stream") + if group_offload_config.pop(key, False) + ] + if dropped: + logger.warning( + f"Ignoring group offload {', '.join(dropped)} - streams need " + f"CUDA or XPU, and this onloads to " + f"{group_offload_config['onload_device']}" + ) +``` + +- [ ] **Step 5: Run the tests to confirm they pass.** + +Run: `venv/bin/python -m pytest tests/test_config_objects.py tests/test_device_portability.py -v` +Expected: all PASS. + +- [ ] **Step 6: Commit.** + +```bash +venv/bin/ruff format dw/pipeline_processors/config_objects.py tests/test_config_objects.py +git add dw/pipeline_processors/config_objects.py tests/test_config_objects.py +git commit -m "fix(offload): drop group-offload CUDA streams when onloading elsewhere" +``` + +--- + +### Task 3: `segment` uses SAM2's real post-processing contract + +**Files:** +- Modify: `dw/tasks/segment.py:106-110` +- Test: `tests/test_segment.py` + +**Interfaces:** +- Consumes: the `transformers.Sam2Processor.post_process_masks(masks, original_sizes, mask_threshold=0.0, binarize=True, ...)` signature, and `Sam2Processor` built offline from `Sam2ImageProcessor()`. Both were verified on the venv's transformers. +- Produces: `segment_image(...)` works with a real `Sam2Processor`. + +- [ ] **Step 1: Write the failing contract test.** This test uses a **real** `Sam2Processor`, which needs no download, so a key or signature drift fails here rather than on a GPU box. Append to `tests/test_segment.py`: + +```python +class TestSegmentWithRealSam2Processor: + """The mocked tests above hand segment_image a dict with whatever keys + they are told to - which is how a SAM v1 key ('reshaped_input_sizes') + survived in the code while every real run raised KeyError. This one + builds the real processor offline and fakes only the models.""" + + @patch("dw.tasks.segment.Sam2Model") + @patch("dw.tasks.segment.Sam2Processor") + @patch("dw.tasks.segment.AutoModelForZeroShotObjectDetection") + @patch("dw.tasks.segment.AutoProcessor") + def test_masks_come_back_at_the_image_size( + self, mock_auto_proc, mock_auto_model, mock_sam_proc, mock_sam_model + ): + from transformers import Sam2ImageProcessor + from transformers import Sam2Processor as RealSam2Processor + + from dw.tasks.segment import segment_image + + dino = MagicMock() + mock_auto_proc.from_pretrained.return_value = dino + dino.return_value = _make_batch_encoding({"input_ids": torch.zeros(1, 10)}) + dino.post_process_grounded_object_detection.return_value = [ + { + "boxes": torch.tensor([[100.0, 100.0, 300.0, 300.0]]), + "scores": torch.tensor([0.9]), + "labels": ["dog"], + } + ] + dino_model = MagicMock() + mock_auto_model.from_pretrained.return_value = dino_model + dino_model.to.return_value = dino_model + + mock_sam_proc.from_pretrained.return_value = RealSam2Processor( + image_processor=Sam2ImageProcessor() + ) + sam_model = MagicMock() + mock_sam_model.from_pretrained.return_value = sam_model + sam_model.to.return_value = sam_model + # SAM2's low-res mask logits: (batch, boxes, masks, 256, 256) + logits = torch.full((1, 1, 3, 256, 256), -10.0) + logits[..., 64:128, 64:128] = 10.0 + sam_model.return_value = MagicMock(pred_masks=logits) + + result = segment_image(_make_test_image(), "dog") + + assert result.mode == "L" + assert result.size == (640, 480) + assert np.array(result).max() == 255 +``` + +- [ ] **Step 2: Run the test to confirm it fails.** + +Run: `venv/bin/python -m pytest tests/test_segment.py::TestSegmentWithRealSam2Processor -v` +Expected: FAIL with `KeyError: 'reshaped_input_sizes'`. + +- [ ] **Step 3: Implement.** In `dw/tasks/segment.py`, replace the `post_process_masks` call with: + +```python + # SAM2's processor takes the original sizes alone - the third positional + # argument is mask_threshold, and the SAM v1 'reshaped_input_sizes' key this + # used to pass does not exist in SAM2's inputs at all + masks = sam_processor.post_process_masks( + sam_outputs.pred_masks, sam_inputs["original_sizes"] + ) +``` + +- [ ] **Step 4: Update the mocked tests to the real contract.** In `tests/test_segment.py`, delete the `"reshaped_input_sizes": ...` entry from every `_make_batch_encoding({...})` SAM input dict, so the mocks can no longer offer a key the real processor lacks. Then in `test_returns_pil_image_mode_l`, after `result = segment_image(image, "dog")`, add: + +```python + args, kwargs = mock_sam_proc_instance.post_process_masks.call_args + assert len(args) == 2 and not kwargs +``` + +- [ ] **Step 5: Run the tests to confirm they pass.** + +Run: `venv/bin/python -m pytest tests/test_segment.py -v` +Expected: all PASS. + +- [ ] **Step 6: Commit.** + +```bash +venv/bin/ruff format dw/tasks/segment.py tests/test_segment.py +git add dw/tasks/segment.py tests/test_segment.py +git commit -m "fix(segment): SAM2 post_process_masks takes original sizes only + +Every real segment run raised KeyError on the SAM v1 'reshaped_input_sizes' +key; the mocked tests supplied it. Adds a test on the real processor." +``` + +--- + +### Task 4: Real memory figures on MPS, and a capacity helper + +**Files:** +- Modify: `dw/__init__.py` (`device_memory_stats` ~line 290; add `_apple_chip_name` and `device_capacity_gb` after it; add `import functools` and `import subprocess` to the imports) +- Test: `tests/test_device_helpers.py` (replace `test_mps_reports_zeroed_known_stats`; add tests) + +**Interfaces:** +- Produces: + - `dw.device_capacity_gb(device=None) -> float | None`. On MPS this is `torch.mps.recommended_max_memory()` in GiB; on CUDA it is the card's total memory in GiB; otherwise it is `None`. It never raises. + - On MPS, `dw.device_memory_stats()` returns `allocated_mb = current_allocated_memory`, `reserved_mb = driver_allocated_memory`, `total_mb = recommended_max_memory`, `free_mb = max(total - reserved, 0)`, and `device_name = " (MPS)"`. When the recommended maximum can't be read, `total_mb` and `free_mb` are `None`. + - `dw._apple_chip_name() -> str` is cached and never raises. + +Nothing downstream needs changing. The worker reads these in the process that holds the models (`dw/worker.py:702`). The UI meter (`ui/src/App.svelte:115`) shows once `total_mb` is non-zero. `_allocated_mb()` in `dw/workflow.py:251` starts reporting real figures. + +- [ ] **Step 1: Write the failing tests.** In `tests/test_device_helpers.py`, replace `test_mps_reports_zeroed_known_stats` with: + +```python + def test_mps_reports_real_figures(self, monkeypatch): + mb = 1024 * 1024 + monkeypatch.setattr(dw, "get_device_type", lambda device=None: "mps") + monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True) + monkeypatch.setattr(dw, "_apple_chip_name", lambda: "Apple M5 Pro") + monkeypatch.setattr(torch.mps, "current_allocated_memory", lambda: 1000 * mb) + monkeypatch.setattr(torch.mps, "driver_allocated_memory", lambda: 1500 * mb) + monkeypatch.setattr(torch.mps, "recommended_max_memory", lambda: 64000 * mb) + + stats = dw.device_memory_stats() + + assert stats == { + "available": True, + "device_name": "Apple M5 Pro (MPS)", + "allocated_mb": 1000.0, + "reserved_mb": 1500.0, + "free_mb": 62500.0, + "total_mb": 64000.0, + } + + def test_mps_without_a_capacity_reading_reports_none(self, monkeypatch): + def unavailable(): + raise RuntimeError("no Metal device") + + monkeypatch.setattr(dw, "get_device_type", lambda device=None: "mps") + monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True) + monkeypatch.setattr(dw, "_apple_chip_name", lambda: "Apple M5 Pro") + monkeypatch.setattr(torch.mps, "current_allocated_memory", lambda: 0) + monkeypatch.setattr(torch.mps, "driver_allocated_memory", lambda: 0) + monkeypatch.setattr(torch.mps, "recommended_max_memory", unavailable) + + stats = dw.device_memory_stats() + + assert stats["available"] is True + assert stats["total_mb"] is None + assert stats["free_mb"] is None + + +class TestDeviceCapacity: + def test_mps_capacity_is_the_recommended_max(self, monkeypatch): + monkeypatch.setattr(dw, "get_device_type", lambda device=None: "mps") + monkeypatch.setattr( + torch.mps, "recommended_max_memory", lambda: 62 * 1024**3 + ) + + assert dw.device_capacity_gb() == 62.0 + + def test_cpu_has_no_capacity(self, monkeypatch): + monkeypatch.setattr(dw, "get_device_type", lambda device=None: "cpu") + + assert dw.device_capacity_gb() is None + + def test_a_failing_probe_is_none_not_an_exception(self, monkeypatch): + def broken(): + raise RuntimeError("boom") + + monkeypatch.setattr(dw, "get_device_type", lambda device=None: "mps") + monkeypatch.setattr(torch.mps, "recommended_max_memory", broken) + + assert dw.device_capacity_gb() is None + + @pytest.mark.skipif( + not torch.backends.mps.is_available(), reason="needs an Apple Silicon GPU" + ) + def test_this_mac_reports_a_capacity(self): + assert dw.device_capacity_gb("mps") > 1.0 +``` + +- [ ] **Step 2: Run the tests to confirm they fail.** + +Run: `venv/bin/python -m pytest tests/test_device_helpers.py -k "mps or Capacity" -v` +Expected: FAIL. `_apple_chip_name` and `device_capacity_gb` don't exist yet, and the figures are zero. + +- [ ] **Step 3: Implement.** In `dw/__init__.py`, add `import functools` and `import subprocess` beside `import os`. Replace the MPS branch of `device_memory_stats` (the `elif device_type == "mps" ...:` block) with: + +```python + elif ( + device_type == "mps" + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + # Unified memory: 'reserved' is everything the Metal driver holds for + # this process, and the ceiling is Metal's recommended working set - + # the figure an allocation past which starts to page + stats["available"] = True + stats["device_name"] = f"{_apple_chip_name()} (MPS)" + stats["allocated_mb"] = torch.mps.current_allocated_memory() / 1024 / 1024 + stats["reserved_mb"] = torch.mps.driver_allocated_memory() / 1024 / 1024 + try: + total = torch.mps.recommended_max_memory() / 1024 / 1024 + stats["total_mb"] = total + stats["free_mb"] = max(total - stats["reserved_mb"], 0.0) + except (RuntimeError, AttributeError): + pass +``` + +Update the docstring's MPS sentence to: "MPS reports allocated (tensors), reserved (everything the Metal driver holds) and total (Metal's recommended working set); free is total minus reserved." Then add after the function: + +```python +@functools.lru_cache(maxsize=1) +def _apple_chip_name(): + """The chip's marketing name ('Apple M5 Pro'), so a reading from an 8 GB M1 + and one from a 128 GB M4 Max are not reported as the same device.""" + try: + result = subprocess.run( + ["sysctl", "-n", "machdep.cpu.brand_string"], + capture_output=True, + text=True, + timeout=2, + check=True, + ) + return result.stdout.strip() or "Apple Silicon" + except (OSError, subprocess.SubprocessError): + return "Apple Silicon" + + +def device_capacity_gb(device=None): + """How much memory the accelerator can hold, in GiB - Metal's recommended + working set on MPS, the card's total on CUDA - or None where there is no + accelerator or the probe fails. Used where a check needs a ceiling for a + device no curated 'cost' entry describes.""" + if not _TORCH_AVAILABLE: + return None + try: + device_type = get_device_type(device) + if device_type == "mps": + return torch.mps.recommended_max_memory() / 1024**3 + if device_type == "cuda": + index = torch.device(device or get_device()).index or 0 + return torch.cuda.get_device_properties(index).total_memory / 1024**3 + except (RuntimeError, AttributeError, AssertionError): + return None + return None +``` + +- [ ] **Step 4: Run the tests to confirm they pass.** + +Run: `venv/bin/python -m pytest tests/test_device_helpers.py tests/test_device.py -v` +Expected: all PASS, and the real-Mac test passes on this box. + +- [ ] **Step 5: Check the whole suite for callers that assumed zeros.** + +Run: `venv/bin/python -m pytest -q -x -k "memory or device or observed" 2>&1 | tail -5` +Expected: PASS. If a test asserted `"Apple Silicon (MPS)"` or zeroed MPS figures, update it to the new contract. Those assertions pinned the old defect. + +- [ ] **Step 6: Commit.** + +```bash +venv/bin/ruff format dw/__init__.py tests/test_device_helpers.py +git add dw/__init__.py tests/test_device_helpers.py +git commit -m "feat(mps): report real unified-memory figures and the chip name" +``` + +--- + +### Task 5: The VRAM estimate checks against the device actually serving + +**Files:** +- Modify: `dw/vram_estimate.py` (`vram_estimate_errors`, `apply_vram_estimate`) +- Modify: `dw/workflow.py` (the call at ~line 786 in validation, `apply_vram_estimate` at ~line 997, and the `from . import ...` line at ~line 94) +- Create: `tests/test_vram_estimate.py` + +**Interfaces:** +- Consumes: `dw.get_device_type()` and `dw.device_capacity_gb()` from Task 4. +- Produces: + - `vram_estimate_errors(definition, arguments=None, supplied=(), device_type=None, capacity_gb=None)` + - `apply_vram_estimate(definition, variables, device_type=None, capacity_gb=None)` + - Rule: + 1. With `device_type=None`, every entry is checked, as today. + 2. Otherwise, the entries whose `device` backend equals `device_type` are checked. + 3. If none match and `capacity_gb` is not None, a single synthesized entry `{"name": f"this {device_type} device (recommended maximum)", "vram_gb": round(capacity_gb, 1)}` is checked. + 4. If none match and capacity is unknown, every entry is checked. This is the conservative fallback. + +- [ ] **Step 1: Write the failing tests.** Create `tests/test_vram_estimate.py`: + +```python +"""vram_estimate refuses a shape against the device actually serving it. + +Every catalog 'cost' entry is a CUDA card, so checking all of them refused a +64 GB Mac against a 24 GB RTX 3090 - and filtering to matching entries alone +would have switched the guard off on a Mac entirely, since none match. +""" + +from dw.vram_estimate import apply_vram_estimate, vram_estimate_errors +import pytest + + +def definition(extra_cost=()): + # 768 * 512 * 121 voxels * 200 bytes = 8.86 GiB, + 20 base = 28.86 GiB + return { + "variables": {"width": 768, "height": 512, "num_frames": 121}, + "cost": [ + {"device": "cuda", "name": "RTX 3090", "vram_gb": 24, "minutes": 1.8}, + *extra_cost, + ], + "vram_estimate": { + "voxel_variables": ["width", "height", "num_frames"], + "base_gb": 20, + "bytes_per_voxel": 200, + }, + } + + +def test_no_device_checks_every_entry_as_before(): + errors = vram_estimate_errors(definition()) + assert len(errors) == 1 and "RTX 3090" in errors[0]["message"] + + +def test_a_cuda_box_checks_its_cuda_entries(): + errors = vram_estimate_errors(definition(), device_type="cuda", capacity_gb=24) + assert len(errors) == 1 and "RTX 3090" in errors[0]["message"] + + +def test_a_mac_with_room_is_not_refused_by_a_cuda_card(): + assert vram_estimate_errors(definition(), device_type="mps", capacity_gb=62) == [] + + +def test_a_mac_without_room_is_refused_against_its_own_capacity(): + errors = vram_estimate_errors(definition(), device_type="mps", capacity_gb=16) + assert len(errors) == 1 + assert "mps" in errors[0]["message"] + assert "RTX 3090" not in errors[0]["message"] + + +def test_a_mac_whose_capacity_is_unknown_keeps_the_conservative_check(): + errors = vram_estimate_errors(definition(), device_type="mps", capacity_gb=None) + assert len(errors) == 1 and "RTX 3090" in errors[0]["message"] + + +def test_a_curated_mps_entry_wins_over_the_measured_capacity(): + mac = {"device": "mps", "name": "M5 Pro 64 GB", "vram_gb": 26, "minutes": 9} + errors = vram_estimate_errors(definition([mac]), device_type="mps", capacity_gb=62) + assert len(errors) == 1 and "M5 Pro 64 GB" in errors[0]["message"] + + +def test_an_indexed_cost_device_matches_its_backend(): + card = {"device": "cuda:1", "name": "second card", "vram_gb": 48, "minutes": 1} + d = definition() + d["cost"] = [card] + assert vram_estimate_errors(d, device_type="cuda", capacity_gb=24) == [] + + +def test_run_time_backstop_takes_the_device_too(): + d = definition() + apply_vram_estimate(d, d["variables"], device_type="mps", capacity_gb=62) + with pytest.raises(ValueError): + apply_vram_estimate(d, d["variables"], device_type="mps", capacity_gb=16) +``` + +- [ ] **Step 2: Run the tests to confirm they fail.** + +Run: `venv/bin/python -m pytest tests/test_vram_estimate.py -v` +Expected: FAIL with `TypeError: unexpected keyword argument 'device_type'`. + +- [ ] **Step 3: Implement in `dw/vram_estimate.py`.** Add this helper above `vram_estimate_errors`: + +```python +def _entries_for(cost, device_type, capacity_gb): + """The cost entries a projection is checked against on this device. + + Entries measured on this backend first. A backend no entry describes - a + Mac, against a catalog measured on CUDA cards - is checked against its own + capacity when it can report one, and against every entry when it cannot, + so an unreadable ceiling never turns the guard off.""" + entries = [entry for entry in cost if isinstance(entry, dict)] + if device_type is None: + return entries + matching = [ + entry + for entry in entries + if str(entry.get("device", "")).split(":")[0] == device_type + ] + if matching: + return matching + if capacity_gb is not None: + return [ + { + "name": f"this {device_type} device (recommended maximum)", + "vram_gb": round(capacity_gb, 1), + } + ] + return entries +``` + +Change the signature to `def vram_estimate_errors(definition, arguments=None, supplied=(), device_type=None, capacity_gb=None):`. Replace the loop header and its `isinstance` guard: + +```python + for entry in cost: + if not isinstance(entry, dict): + continue +``` + +with: + +```python + for entry in _entries_for(cost, device_type, capacity_gb): +``` + +Change `apply_vram_estimate` to take and forward the two arguments: + +```python +def apply_vram_estimate(definition, variables, device_type=None, capacity_gb=None): + ... + errors = vram_estimate_errors( + definition, variables, device_type=device_type, capacity_gb=capacity_gb + ) +``` + +Append this to the module docstring: "The entries checked are the serving device's own (see `_entries_for`); `bytes_per_voxel` was calibrated on CUDA, so a check against a Mac's capacity is an estimate of an estimate." + +- [ ] **Step 4: Wire the device in at both call sites in `dw/workflow.py`.** Extend the import at ~line 94 to `from . import get_device, empty_device_cache, device_memory_stats, device_capacity_gb, get_device_type`. At the validation call (~line 786): + +```python ++vram_estimate_errors( + self.workflow_definition, + arguments, + supplied=set(arguments or {}), + device_type=get_device_type(), + capacity_gb=device_capacity_gb(), +) +``` + +At the run-time call (~line 997): + +```python + apply_vram_estimate( + workflow_def, + variables, + device_type=get_device_type(), + capacity_gb=device_capacity_gb(), + ) +``` + +- [ ] **Step 5: Run the tests to confirm they pass.** + +Run: `venv/bin/python -m pytest tests/test_vram_estimate.py tests/test_catalog_structure.py -q` +Then: `venv/bin/python -m pytest -q -k "vram or validat" 2>&1 | tail -5` +Expected: all PASS. A pre-existing test that expected an RTX 3090 refusal from `Workflow.validation_errors` while running on this Mac is now environment-dependent. Pin it with `monkeypatch.setattr("dw.workflow.get_device_type", lambda device=None: "cuda")` rather than deleting it. + +- [ ] **Step 6: Check it on a real template.** + +Run: `venv/bin/python -m dw.validate workflows/templates/ltx2/text-to-video.json` +Expected: valid on this 64 GB Mac. + +- [ ] **Step 7: Commit.** + +```bash +venv/bin/ruff format dw/vram_estimate.py dw/workflow.py tests/test_vram_estimate.py +git add dw/vram_estimate.py dw/workflow.py tests/test_vram_estimate.py +git commit -m "fix(vram_estimate): check against the serving device, not every CUDA card" +``` + +--- + +### Task 6: Attention slicing becomes opt-in, and the fp16 warning is corrected + +**Files:** +- Modify: `dw/pipeline_processors/pipeline.py` (~lines 312-326 slicing block; ~line 1823 fp16 warning; add helper near `auto_cpu_offload_enabled` ~line 1509) +- Modify: `dw/workflow_schema.json:595-598` (`disable_attention_slicing` description) +- Modify: `docs/ACCELERATION.md:218`, `docs/WORKFLOW_GUIDE.md:1292-1293` +- Test: create `tests/test_attention_slicing.py` + +**Interfaces:** +- Produces: `attention_slicing_requested(configuration: dict) -> bool` in `dw/pipeline_processors/pipeline.py`. + +- [ ] **Step 1: Write the failing test.** Create `tests/test_attention_slicing.py`: + +```python +"""Attention slicing is opt-in on every backend. + +It used to switch on automatically on MPS. Measured on an M5 Pro under +inference_mode, it made SDXL-sized UNet attention 2.4x slower in bf16 +(12.8 vs 5.4 ms) and 1.4x slower in fp32 - PyTorch's SDPA is already the +memory-efficient path, which diffusers' own enable_attention_slicing +docstring warns about. Only UNet/ControlNet models take slicing at all. +""" + +from dw.pipeline_processors.pipeline import attention_slicing_requested + + +def test_off_by_default(): + assert attention_slicing_requested({}) is False + + +def test_on_when_asked(): + assert attention_slicing_requested({"enable_attention_slicing": True}) is True + + +def test_the_old_opt_out_is_harmless(): + assert attention_slicing_requested({"disable_attention_slicing": True}) is False +``` + +- [ ] **Step 2: Run the test to confirm it fails.** + +Run: `venv/bin/python -m pytest tests/test_attention_slicing.py -v` +Expected: FAIL with ImportError. + +- [ ] **Step 3: Implement.** Add this helper to `pipeline.py` beside `auto_cpu_offload_enabled`: + +```python +def attention_slicing_requested(configuration): + """Whether a pipeline's attention runs sliced. Opt-in on every backend: it + used to be automatic on MPS, where it measured 2.4x slower on UNet attention + than the SDPA it replaces (bf16, M5 Pro) - SDPA is already the + memory-efficient path. 'disable_attention_slicing' is accepted and ignored.""" + return bool(configuration.get("enable_attention_slicing", False)) +``` + +Replace the slicing block's condition and comments (~lines 312-317) with: + +```python + # Attention slicing trades speed for memory and is opt-in - see + # attention_slicing_requested for why it is no longer automatic on MPS + if attention_slicing_requested(self.configuration): +``` + +Keep the inner `has_method` branch as it is, but change its comment to "Modular pipelines have no attention slicing - skip rather than fail". Then replace the fp16 warning message (~line 1831) with: + +```python + logger.warning( + f"{component_name} loads in float16 on MPS, which can produce NaN " + "values (black images) on Apple Silicon - bfloat16 is the usual fix" + ) +``` + +- [ ] **Step 4: Update the schema and docs.** + - In `dw/workflow_schema.json`, set the `disable_attention_slicing` description to: `"Deprecated and ignored: attention slicing is opt-in on every backend (enable_attention_slicing)."` + - In `docs/ACCELERATION.md:218`, replace the sentence "Enabled automatically on MPS (unified memory benefits from slicing) unless `disable_attention_slicing` is set." with "Opt-in on every backend. It is not automatic on MPS, because PyTorch's SDPA is already memory-efficient and slicing measured 2.4x slower on UNet attention on Apple Silicon. It only affects UNet/ControlNet models." + - In `docs/WORKFLOW_GUIDE.md:1292-1293`, replace "Enabled automatically on MPS unless `disable_attention_slicing` is set." with "Opt-in on every backend (UNet/ControlNet models only)." + +- [ ] **Step 5: Run the tests and the schema check.** + +Run: `venv/bin/python -m pytest tests/test_attention_slicing.py tests/test_catalog_structure.py -q && venv/bin/python -m dw.validate workflows/templates/text-to-image.json` +Expected: PASS, then valid. + +- [ ] **Step 6: Commit.** + +```bash +venv/bin/ruff format dw/pipeline_processors/pipeline.py tests/test_attention_slicing.py +git add dw/pipeline_processors/pipeline.py dw/workflow_schema.json docs/ACCELERATION.md docs/WORKFLOW_GUIDE.md tests/test_attention_slicing.py +git commit -m "perf(mps): make attention slicing opt-in; it slowed UNet attention 2.4x" +``` + +--- + +### Task 7: CPU fallbacks on MPS are no longer silent + +**Files:** +- Modify: `dw/__init__.py:27-30` (the `warnings.filterwarnings` block) +- Test: create `tests/test_mps_fallback_warning.py` + +**Interfaces:** +- Produces: after `import dw`, torch's MPS CPU-fallback `UserWarning` is shown, and every other `UserWarning` is still ignored. + +- [ ] **Step 1: Write the failing test.** Use a fresh interpreter, because the filters are installed at import and pytest adds its own. + +```python +"""torch announces an op the MPS backend lacks, when PYTORCH_ENABLE_MPS_FALLBACK +is set, with a UserWarning - and dw ignored every UserWarning, so a step that +quietly ran part of its model on the CPU left no trace. The fallback variable is +set by the fp4-fp8-for-torch-mps autoload install.sh adds on macOS.""" + +import subprocess +import sys + +PROBE = """ +import warnings, dw +fallback = ("The operator 'aten::_linalg_eigvals' is not currently supported " + "on the MPS backend and will fall back to run on the CPU. " + "This may have performance implications.") +with warnings.catch_warnings(record=True) as caught: + warnings.warn(fallback, UserWarning) + warnings.warn("some other library chatter", UserWarning) +print(len(caught), caught[0].message if caught else "") +""" + + +def test_only_the_fallback_warning_gets_through(): + result = subprocess.run( + [sys.executable, "-c", PROBE], capture_output=True, text=True, timeout=120 + ) + assert result.returncode == 0, result.stderr + count, _, message = result.stdout.strip().splitlines()[-1].partition(" ") + assert count == "1" + assert "fall back to run on the CPU" in message +``` + +- [ ] **Step 2: Run the test to confirm it fails.** + +Run: `venv/bin/python -m pytest tests/test_mps_fallback_warning.py -v` +Expected: FAIL with `count == "0"`. + +- [ ] **Step 3: Implement.** In `dw/__init__.py`, directly after the three `warnings.filterwarnings("ignore", ...)` lines, add: + +```python +# ...except torch's notice that an op the MPS backend lacks ran on the CPU. With +# PYTORCH_ENABLE_MPS_FALLBACK set (the fp4-fp8-for-torch-mps autoload sets it) +# that is otherwise a step that got several times slower with nothing to say why. +# Later filters take precedence, so this one wins over the blanket ignore above +warnings.filterwarnings( + "default", + message=r".*not currently supported on the MPS backend and will fall back", + category=UserWarning, +) +``` + +- [ ] **Step 4: Run the test to confirm it passes.** + +Run: `venv/bin/python -m pytest tests/test_mps_fallback_warning.py -v` +Expected: PASS. + +- [ ] **Step 5: Confirm it with a real op on this Mac.** + +Run: `PYTORCH_ENABLE_MPS_FALLBACK=1 venv/bin/python -c "import dw, torch; torch.linalg.eigvals(torch.randn(4,4,device='mps'))" 2>&1 | grep -c "fall back"` +Expected: `1`. + +- [ ] **Step 6: Commit.** + +```bash +venv/bin/ruff format dw/__init__.py tests/test_mps_fallback_warning.py +git add dw/__init__.py tests/test_mps_fallback_warning.py +git commit -m "fix(mps): let torch's CPU-fallback warning through the UserWarning filter" +``` + +--- + +### Task 8: Template defaults that fail, download twice, or come from odd places + +**Files:** +- Modify: `workflows/templates/upscale-diffusion.json:10`, `workflows/templates/upscale-spandrel.json:9`, `workflows/templates/image-to-image.json:5`, `workflows/models/flux-torchao.json:38`, `workflows/templates/attention-processor.json:24`, `workflows/templates/controlnet-component.json:37`, `dw/workflows/test.json:16` +- Create: `tests/test_template_defaults.py` + +**Interfaces:** none (data only). + +- [ ] **Step 1: Write the failing catalog test.** Create `tests/test_template_defaults.py`: + +```python +"""A catalog default must run: no placeholder hosts, and one spelling per repo. + +upscale-* defaulted to https://example.com/image.jpg (not an image), and +flux-torchao spelled FLUX.1-dev two ways - the Hub cache is keyed on the +spelling, so the transformer downloaded twice (~23 GB).""" + +import json +import pathlib +import re + +ROOT = pathlib.Path(__file__).resolve().parent.parent +FILES = sorted([*ROOT.glob("workflows/**/*.json"), *ROOT.glob("dw/workflows/*.json")]) +PLACEHOLDER = re.compile(r"https?://(www\.)?example\.(com|org|net)/") +WITHDRAWN = ("runwayml/",) + + +def strings(node): + if isinstance(node, dict): + for value in node.values(): + yield from strings(value) + elif isinstance(node, list): + for value in node: + yield from strings(value) + elif isinstance(node, str): + yield node + + +def test_no_default_points_at_a_placeholder_host(): + offenders = [ + f"{path.relative_to(ROOT)}: {s}" + for path in FILES + for s in strings(json.loads(path.read_text())) + if PLACEHOLDER.search(s) + ] + assert offenders == [] + + +def test_no_withdrawn_repos(): + offenders = [ + f"{path.relative_to(ROOT)}: {s}" + for path in FILES + for s in strings(json.loads(path.read_text())) + if s.startswith(WITHDRAWN) + ] + assert offenders == [] + + +def test_one_spelling_per_repo_within_a_file(): + for path in FILES: + repos = { + s + for s in strings(json.loads(path.read_text())) + if re.fullmatch(r"[\w.-]+/[\w.-]+", s) + } + lowered = {} + for repo in repos: + lowered.setdefault(repo.lower(), set()).add(repo) + clashes = [names for names in lowered.values() if len(names) > 1] + assert clashes == [], f"{path.relative_to(ROOT)}: {clashes}" +``` + +- [ ] **Step 2: Run the test to confirm it fails.** + +Run: `venv/bin/python -m pytest tests/test_template_defaults.py -v` +Expected: three FAILs naming the upscale templates, the two `runwayml/` files and `flux-torchao.json`. If the repo-spelling test flags anything else, treat it as a real finding and fix it the same way. + +- [ ] **Step 3: Fix the data.** + - `upscale-diffusion.json`, `upscale-spandrel.json` and `image-to-image.json`: set the default `location` to `https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png`. This was verified 200 `image/png` on 2026-09-25, and it's the image diffusers' own docs use. `image-to-image`'s previous default came from a third-party repo named `sage-permission-probe`. + - `flux-torchao.json:38`: change `black-forest-labs/Flux.1-Dev` to `black-forest-labs/FLUX.1-dev`. + - `attention-processor.json:24` and `controlnet-component.json:37`: change `runwayml/stable-diffusion-v1-5` to `stable-diffusion-v1-5/stable-diffusion-v1-5`. + - `dw/workflows/test.json:16`: change `"torch.float16"` to `"torch.bfloat16"`. This is the system test, and the engine's own warning says fp16 is unsafe on MPS. SD 1.5 runs in bf16 on CUDA too. + +- [ ] **Step 4: Run the tests and validate the touched files.** + +Run: `venv/bin/python -m pytest tests/test_template_defaults.py tests/test_catalog_structure.py -q && for f in workflows/templates/upscale-diffusion.json workflows/templates/upscale-spandrel.json workflows/templates/image-to-image.json workflows/models/flux-torchao.json workflows/templates/attention-processor.json workflows/templates/controlnet-component.json; do venv/bin/python -m dw.validate $f || echo "FAIL $f"; done` +Expected: PASS, and no `FAIL` lines. + +- [ ] **Step 5: Commit.** + +```bash +venv/bin/ruff format tests/test_template_defaults.py +git add workflows dw/workflows/test.json tests/test_template_defaults.py +git commit -m "fix(templates): runnable default inputs, one repo spelling, no withdrawn repos" +``` + +--- + +### Task 9: Docs catch up with Mac behavior + +**Files:** +- Modify: `CLAUDE.md` (lines 26-27 `dw.test` comment; line 324 MPS gotcha; line 576 `music` -1; the *Cross-Platform Device Support* section) +- Modify: `docs/ACCELERATION.md:343-352` (MPS Notes) +- Modify: `docs/DEPENDENCIES.md:58` +- Modify: `README.md:35`, `docs/TESTING.md:67` (the `dw.test` description) + +**Interfaces:** none. + +- [ ] **Step 1: CLAUDE.md.** + - Replace the `# Basic system test (torch, diffusers import check)` comment with `# System test - downloads SD 1.5 (a few GB) and generates one image`. + - In the MPS gotcha at line 324, change "no autocast, no bitsandbytes, no flash_attn, no triton, no torch.compile" to "no bitsandbytes, no flash_attn, no triton, no torch.compile (`torch.autocast("mps")` works on torch 2.14, but dw does not use it)". + - At line 576, change "`music`, an mp3, keeps -1" to "`music`, an mp3, normalizes to -3 as well (#362)". + - In *Cross-Platform Device Support*, after the `resolve_device()` paragraph, add: "The same translation reaches the settings that carry a device or a CUDA-only feature: SDNQ `quantization_device`/`return_device` go through `resolve_device()`, and `use_quantized_matmul(_conv)` is turned off on MPS, because it falls back to `torch._int_mm`, ~500x slower there (`portable_quantization_arguments`, `config_objects.py`). Group-offload `use_stream`/`record_stream` are dropped when the onload device is not CUDA/XPU. `vram_estimate` checks against the serving device's own `cost` entries, else its `device_capacity_gb()`. Attention slicing is opt-in everywhere. Each adaptation logs a warning." +- [ ] **Step 2: ACCELERATION.md MPS Notes.** Replace the bullet list with: + +```markdown +- No flash-attn, no Triton, no bitsandbytes - `attention_backend` is effectively CUDA-only; use `"native"`-family backends or leave it unset on MPS. `compile` is skipped with a warning (inductor support on MPS is immature). A workflow that loads a bitsandbytes checkpoint (`flux2-dev`, `multi-image-reference`) or a TorchAO int4 config (`flux-torchao`) has no Mac path yet. +- Workflows written for CUDA are adapted rather than refused: a `cuda` device becomes `mps`, SDNQ's `quantization_device` follows it and its quantized matmul is switched off (it runs through `torch._int_mm`, ~500x slower on MPS), group-offload CUDA streams are dropped, and `"offload": "sequential"` becomes `"model"`. Each change is logged as a warning. +- Attention slicing is opt-in; it made UNet attention 2.4x slower on Apple Silicon. +- `float16` can produce NaN values on Apple Silicon - use `bfloat16`; dw only warns, it doesn't override the dtype for you. +- `PYTORCH_MPS_HIGH_WATERMARK_RATIO` defaults to `0.0` (use all unified memory) unless already set in the environment. +- Ops the MPS backend lacks: with `PYTORCH_ENABLE_MPS_FALLBACK=1` (set by the `fp4-fp8-for-torch-mps` package install.sh adds) they run on the CPU, and dw logs torch's warning when one does. Without it they raise. +- Memory figures are real: allocated, driver-reserved, and Metal's recommended working set as the total. `vram_estimate` checks a Mac against that total, since the catalog's `cost` entries are CUDA cards; the per-voxel figures were calibrated on CUDA. +- Cost estimates quoted on a Mac are the CUDA figures (`basis: "other_device"`) until the Mac's own runs build observed history - expect them to be optimistic. +- Offloading has less benefit than on CUDA, since unified memory is already shared between CPU and GPU. +``` + +- [ ] **Step 3: DEPENDENCIES.md:58.** Replace the macOS line with: `**macOS (MPS):** fp4-fp8-for-torch-mps (FP8/FP4 dtypes for Metal; autoloads on every torch import, overrides mm/linear/copy on the MPS dispatch key and sets PYTORCH_ENABLE_MPS_FALLBACK=1), fluidtop (a GPU/power TUI; needs sudo)` +- [ ] **Step 4: README.md:35 and docs/TESTING.md:67.** Rewrite the `dw.test` description to say that it runs `dw/workflows/test.json`, which downloads SD 1.5 and generates one image, so it exercises the whole stack rather than only the imports. Keep the surrounding sentence structure. +- [ ] **Step 5: Check.** + +Run: `grep -n "import check\|keeps -1\|no autocast\|on by default (set .disable_attention_slicing" CLAUDE.md README.md docs/*.md` +Expected: no output. + +- [ ] **Step 6: Commit.** + +```bash +git add CLAUDE.md README.md docs/ACCELERATION.md docs/DEPENDENCIES.md docs/TESTING.md +git commit -m "docs: Mac behavior, dw.test's real scope, stale autocast and music lines" +``` + +--- + +### Task 10: Verify on this Mac (manual; downloads need Don's go-ahead) + +**Files:** none. Record the results in the PR description. + +- [ ] **Step 1: Run the full suite.** + +Run: `venv/bin/python -m pytest -q 2>&1 | tail -5` +Expected: no failures beyond those already failing on `develop`. Check by running the same command on `develop` if any appear. + +- [ ] **Step 2: Run the quick-start template, which covers UNet and the slicing change.** + +Run: `time venv/bin/python -m dw.run workflows/templates/text-to-image.json` +Expected: an image written under `outputs/`, with no attention-slicing debug line. Record the wall time. + +- [ ] **Step 3: Run the segment template (a download of about 1 GB).** + +Run: `venv/bin/python -m dw.run workflows/templates/segment.json` +Expected: a mask written. It previously raised KeyError. + +- [ ] **Step 4: Ask Don before this step.** LTX-2.5 text-to-video on a Mac is a download of tens of GB and a long run. + +Run: `venv/bin/python -m dw.run workflows/templates/ltx2/text-to-video.json` +Expected: the log shows the three adaptation warnings (quantization device, quantized matmul, streams), and then either a video or a real memory or speed limit. Either result is the finding. Record peak `driver_allocated_memory` from the job's memory events, and the wall time. + +- [ ] **Step 5: Open a PR against `develop`.** Its body summarizes Tasks 1-9 and pastes the Step 2-4 results. It lists the out-of-scope items below as follow-ups. + +--- + +## Out of scope (follow-ups, each needs a measurement first) + +- The `fp4-fp8-for-torch-mps` autoload is a global op override that nothing in the catalog uses. Decide whether to keep it on by default. +- Downgrade group offload and `residency: on_demand` on MPS the way sequential offload is downgraded. This needs a peak-memory A/B first. +- Add a GGUF or SDNQ path for `flux2-dev`, `multi-image-reference`, `community-pipeline` and `flux-torchao`. +- The fp16 defaults in `attention-processor`, `base-and-refiner`, `controlnet-component`, `depth-marigold`, `qr-code` and `surface-normals`. Also the fp16-on-MPS choice in `dw/tasks/upscale.py` and `restore_faces.py` (reviewer-reported, not reproduced). +- A `generator_device: "cpu"` option so a seed reproduces across CUDA and Mac. +- Task-model reloads per call in `restore_faces` and `image_utils` (reviewer-reported). +- `ltx-2.5` / `minimax-music3` skill wording about the Mac, once Task 10 Step 4 says what actually fits. +- The desktop-installers branch lacks torchaudio. diff --git a/dw/__init__.py b/dw/__init__.py index 20b12d24..9f62d1f2 100644 --- a/dw/__init__.py +++ b/dw/__init__.py @@ -1,8 +1,11 @@ from .settings import resolve_path, load_settings from .log_setup import setup_logging +import functools import os import re import logging +import subprocess +import sys import warnings from dotenv import load_dotenv @@ -20,14 +23,35 @@ if "PYTORCH_CUDA_ALLOC_CONF" not in os.environ: os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" -# Load sharded checkpoints in parallel - a pure cold-start win + +def _parallel_loading_default(platform): + """Whether diffusers loads a sharded checkpoint on several threads. + + A pure cold-start win on CUDA. On MPS the loader threads each copy their + tensors onto the device at once, and that segfaulted LTX-2.5's SDNQ + transformer load (exit 139, four threads inside copy_ to mps) where the + serial load succeeded in 39s. diffusers reads the variable once at import, + before dw can ask torch for a device, so this goes by platform: macOS has + no CUDA, and its accelerator is always MPS.""" + return "false" if platform == "darwin" else "true" + + if "HF_ENABLE_PARALLEL_LOADING" not in os.environ: - os.environ["HF_ENABLE_PARALLEL_LOADING"] = "true" + os.environ["HF_ENABLE_PARALLEL_LOADING"] = _parallel_loading_default(sys.platform) # Suppress all common library warnings before any imports warnings.filterwarnings("ignore", category=FutureWarning) warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=DeprecationWarning) +# ...except torch's notice that an op the MPS backend lacks ran on the CPU. With +# PYTORCH_ENABLE_MPS_FALLBACK set (the fp4-fp8-for-torch-mps autoload sets it) +# that is otherwise a step that got several times slower with nothing to say why. +# Later filters take precedence, so this one wins over the blanket ignore above +warnings.filterwarnings( + "default", + message=r".*not currently supported on the MPS backend and will fall back", + category=UserWarning, +) # Eager (but guarded) import of torch. This module's own top-level code needs # torch's exceptions/types unconditionally to detect its absence gracefully @@ -299,10 +323,11 @@ def device_memory_stats(): """ Snapshot of allocator memory usage for the device dw is configured for. - Only CUDA exposes free/total figures (via torch.cuda.mem_get_info); MPS - has no equivalent API, so its snapshot reports zeroed, "known" values - instead. CPU, or a CUDA/MPS backend that isn't actually available despite - being the configured device type, results in an "unavailable" snapshot. + CUDA exposes free/total via torch.cuda.mem_get_info. MPS reports + allocated (tensors), reserved (everything the Metal driver holds) and + total (Metal's recommended working set); free is total minus reserved. + CPU, or a CUDA/MPS backend that isn't actually available despite being + the configured device type, results in an "unavailable" snapshot. Returns: dict with keys: @@ -310,10 +335,10 @@ def device_memory_stats(): device_name (str or None) allocated_mb (float) reserved_mb (float) - free_mb (float or None): None only when CUDA's mem_get_info call - itself fails - total_mb (float or None): None only when CUDA's mem_get_info call - itself fails + free_mb (float or None): None only when the backend's capacity + probe itself fails + total_mb (float or None): None only when the backend's capacity + probe itself fails """ stats = { "available": False, @@ -344,15 +369,59 @@ def device_memory_stats(): and hasattr(torch.backends, "mps") and torch.backends.mps.is_available() ): - # MPS doesn't provide detailed memory stats like CUDA + # Unified memory: 'reserved' is everything the Metal driver holds for + # this process, and the ceiling is Metal's recommended working set - + # the figure past which an allocation starts to page stats["available"] = True - stats["device_name"] = "Apple Silicon (MPS)" - stats["free_mb"] = 0.0 - stats["total_mb"] = 0.0 + stats["device_name"] = f"{_apple_chip_name()} (MPS)" + stats["allocated_mb"] = torch.mps.current_allocated_memory() / 1024 / 1024 + stats["reserved_mb"] = torch.mps.driver_allocated_memory() / 1024 / 1024 + try: + total = torch.mps.recommended_max_memory() / 1024 / 1024 + stats["total_mb"] = total + stats["free_mb"] = max(total - stats["reserved_mb"], 0.0) + except (RuntimeError, AttributeError): + pass return stats +@functools.lru_cache(maxsize=1) +def _apple_chip_name(): + """The chip's marketing name ('Apple M5 Pro'), so a reading from an 8 GB M1 + and one from a 128 GB M4 Max are not reported as the same device.""" + try: + result = subprocess.run( + ["sysctl", "-n", "machdep.cpu.brand_string"], + capture_output=True, + text=True, + timeout=2, + check=True, + ) + return result.stdout.strip() or "Apple Silicon" + except (OSError, subprocess.SubprocessError): + return "Apple Silicon" + + +def device_capacity_gb(device=None): + """How much memory the accelerator can hold, in GiB - Metal's recommended + working set on MPS, the card's total on CUDA - or None where there is no + accelerator or the probe fails. Used where a check needs a ceiling for a + device no curated 'cost' entry describes.""" + if not _TORCH_AVAILABLE: + return None + try: + device_type = get_device_type(device) + if device_type == "mps": + return torch.mps.recommended_max_memory() / 1024**3 + if device_type == "cuda": + index = torch.device(device or get_device()).index or 0 + return torch.cuda.get_device_properties(index).total_memory / 1024**3 + except (RuntimeError, AttributeError, AssertionError): + return None + return None + + def startup(log_level=None): if not _TORCH_AVAILABLE: raise ImportError( diff --git a/dw/pipeline_processors/config_objects.py b/dw/pipeline_processors/config_objects.py index 0551f8d3..85042800 100644 --- a/dw/pipeline_processors/config_objects.py +++ b/dw/pipeline_processors/config_objects.py @@ -1,7 +1,7 @@ import logging import torch -from .. import resolve_device +from .. import get_device_type, resolve_device logger = logging.getLogger("dw") @@ -26,6 +26,35 @@ def get_quantization_configuration(configuration): return create_quantization_config(quantization_config) +# Quantization config arguments that name a device. A workflow authored on the +# CUDA box writes "cuda" here, and nothing downstream translates it - SDNQ moves +# every weight to it inside from_pretrained, which fails outright on a Mac +DEVICE_ARGUMENTS = ("quantization_device", "return_device") + +# SDNQ's quantized matmul has no Triton on MPS and falls back to torch._int_mm, +# measured at ~500x the bf16 matmul it replaces (1096 ms vs 2.1 ms, 512x3072 by +# 3072x3072 on an M5 Pro). The dequantize path is the fast one there +QUANTIZED_MATMUL_ARGUMENTS = ("use_quantized_matmul", "use_quantized_matmul_conv") + + +def portable_quantization_arguments(arguments): + """Adapt a quantization config's arguments to this machine without touching + the definition they came from - a cached pipeline builds from it again.""" + adapted = dict(arguments) + for key in DEVICE_ARGUMENTS: + if isinstance(adapted.get(key), str): + adapted[key] = resolve_device(adapted[key]) + if get_device_type() == "mps": + for key in QUANTIZED_MATMUL_ARGUMENTS: + if adapted.get(key): + logger.warning( + f"Turning off '{key}' on MPS - without Triton it runs through " + "torch._int_mm, which is far slower there than dequantizing" + ) + adapted[key] = False + return adapted + + def create_quantization_config(quantization_config): """ Create a quantization configuration object from its definition. @@ -46,7 +75,9 @@ def create_quantization_config(quantization_config): # "quant_type": "torchao.quantization.Int8WeightOnlyConfig" in JSON. args = { k: v() if isinstance(v, type) else v - for k, v in quantization_config["arguments"].items() + for k, v in portable_quantization_arguments( + quantization_config["arguments"] + ).items() } return quantization_config_type(**args) except Exception as e: @@ -113,6 +144,23 @@ def get_group_offload_configuration(configuration, default_device): resolve_device(group_offload_config.get("offload_device", "cpu")) ) + # CUDA streams overlap the next group's transfer with this one's compute. + # diffusers refuses them without CUDA or XPU, and refuses record_stream + # without use_stream, so both go together - a catalog template written + # for the CUDA box would otherwise fail to load on a Mac + if group_offload_config["onload_device"].type not in ("cuda", "xpu"): + dropped = [ + key + for key in ("use_stream", "record_stream") + if group_offload_config.pop(key, False) + ] + if dropped: + logger.warning( + f"Ignoring group offload {', '.join(dropped)} - streams need " + f"CUDA or XPU, and this onloads to " + f"{group_offload_config['onload_device']}" + ) + return group_offload_config logger.debug("No group offload configuration found") diff --git a/dw/pipeline_processors/pipeline.py b/dw/pipeline_processors/pipeline.py index 0d412a49..f63662f5 100644 --- a/dw/pipeline_processors/pipeline.py +++ b/dw/pipeline_processors/pipeline.py @@ -309,14 +309,10 @@ def load(self, shared_components): defer_placement=adapters_to_load, ) - # Enable attention slicing if explicitly requested or automatically on MPS - # MPS benefits from slicing since Metal shares system RAM with the GPU - if self.configuration.get("enable_attention_slicing", False) or ( - get_device_type(self.device) == "mps" - and not self.configuration.get("disable_attention_slicing", False) - ): - # Modular pipelines have no attention slicing - on MPS this is applied - # automatically, so skip rather than fail when the pipeline lacks it + # Attention slicing trades speed for memory - automatic on MPS, where + # it is often the faster path too (see attention_slicing_requested) + if attention_slicing_requested(self.configuration, self.device): + # Modular pipelines have no attention slicing - skip rather than fail if has_method(self.pipeline, "enable_attention_slicing"): logger.debug("Enabling attention slicing for pipeline") self.pipeline.enable_attention_slicing() @@ -378,6 +374,7 @@ def load(self, shared_components): configure_components( self.pipeline, self.configuration, self.device, reused_components ) + apply_mps_rope_precision(self.pipeline, self.device) # Set up random generator if needed - no_generator is a boolean, so an # explicit false still gets a generator @@ -1506,6 +1503,52 @@ def load_and_configure_scheduler( scheduler.set_shift(float(shift)) +def apply_mps_rope_precision(pipeline, device): + """Run a RoPE that asks for float64 in float32 on MPS, which has no float64. + + diffusers makes this choice itself for Wan, Lumina2, SkyReels-V2, ChronoEdit + and Sana-Video. LTX-2's transformer and text connectors read a + `double_precision` flag at forward time instead, and left on (the default) + it failed LTX-2.5's first step on a Mac with "Cannot convert a MPS Tensor to + float64". Keyed on the flag rather than a model name, so any module that + exposes one gets the same treatment.""" + if get_device_type(device) != "mps": + return + components = getattr(pipeline, "components", None) + if not isinstance(components, dict): + return + for component_name, component in components.items(): + if not isinstance(component, torch.nn.Module): + continue + switched = 0 + for module in component.modules(): + if getattr(module, "double_precision", None) is True: + module.double_precision = False + switched += 1 + if switched: + logger.warning( + f"Computing {switched} float64 RoPE module(s) in {component_name} in " + "float32 - MPS has no float64 (diffusers does the same for Wan)" + ) + + +def attention_slicing_requested(configuration, device): + """Whether a pipeline's attention runs sliced: automatic on MPS unless + 'disable_attention_slicing' is set, opt-in ('enable_attention_slicing') + everywhere else. + + Which is faster on MPS depends on the model. Measured on an M5 Pro (torch + 2.14), MPS SDPA is slow at head dims 40, 48 and 160 and fast at 32 and + 64-128 - slicing made SD 1.5 (40/80/160) ~20% faster end to end and + SDXL-shaped attention (64) 2.4x slower. It stays automatic because the + catalog's quick-start is SD 1.5; SDXL on a Mac can opt out.""" + if configuration.get("enable_attention_slicing", False): + return True + return get_device_type(device) == "mps" and not configuration.get( + "disable_attention_slicing", False + ) + + def auto_cpu_offload_enabled(configuration): """Whether the configuration asks its components manager to offload to the CPU.""" return configuration.get("components_manager", {}).get( @@ -1829,8 +1872,8 @@ def load_component( and from_pretrained_arguments.get("torch_dtype") == torch.float16 ): logger.warning( - f"On MPS devices float16 produces NaN values on Apple Silicon" - f"Consider changing torch_dtype from float16 to float32 for {component_name} " + f"{component_name} loads in float16 on MPS, which can produce NaN " + "values (black images) on Apple Silicon - bfloat16 is the usual fix" ) model_name = None diff --git a/dw/tasks/audio_transcription.py b/dw/tasks/audio_transcription.py index 21ebde22..5aa87d23 100644 --- a/dw/tasks/audio_transcription.py +++ b/dw/tasks/audio_transcription.py @@ -26,6 +26,7 @@ _DEFAULT_ASR_MODEL = "openai/whisper-base" _ASR_SAMPLE_RATE = 16000 +_WHISPER_WINDOW_SECONDS = 30 def _downmixed_mono(waveform): @@ -80,8 +81,20 @@ def load_pipe(): ("transcribe_audio", model_name, str(device), str(dtype)), load_pipe ) + # Whisper takes 30 s per window and refuses a longer clip ("more than 3000 + # mel input features") unless it predicts timestamps, which is how its + # long-form mode stitches windows. Asked for only of Whisper and only past + # 30 s: transformers raises for a CTC model unless the value is "char" or + # "word", and for any other seq2seq model at all + options = {} + if ( + getattr(pipe, "type", None) == "seq2seq_whisper" + and len(mono) > _WHISPER_WINDOW_SECONDS * _ASR_SAMPLE_RATE + ): + options["return_timestamps"] = True result = pipe( - {"raw": mono.astype(numpy.float32), "sampling_rate": _ASR_SAMPLE_RATE} + {"raw": mono.astype(numpy.float32), "sampling_rate": _ASR_SAMPLE_RATE}, + **options, ) text = result["text"].strip() logger.info(f"Transcript: {text[:100]}{'...' if len(text) > 100 else ''}") diff --git a/dw/tasks/segment.py b/dw/tasks/segment.py index 1ab2ef04..afef6bee 100644 --- a/dw/tasks/segment.py +++ b/dw/tasks/segment.py @@ -23,6 +23,14 @@ _DEFAULT_SAM_MODEL = "facebook/sam2-hiera-large" +def grounding_query(prompt): + """A prompt in the form GroundingDINO scores: lowercase, each phrase ending + in a period. 'cat' scored 0.20 against a photo of a cat - under the 0.3 + threshold, so segment returned an empty mask - where 'cat.' scored 0.40.""" + query = prompt.strip().lower() + return query if query.endswith(".") else f"{query}." + + def segment_image(image, prompt, device="cpu", **kwargs): """Segment objects matching a text prompt, returning a binary mask. @@ -58,7 +66,9 @@ def load_dino(): ("segment_dino", model_name, str(device)), load_dino ) - inputs = dino_processor(images=image, text=prompt, return_tensors="pt").to(device) + inputs = dino_processor( + images=image, text=grounding_query(prompt), return_tensors="pt" + ).to(device) with torch.inference_mode(): outputs = dino_model(**inputs) @@ -103,10 +113,11 @@ def load_sam(): with torch.inference_mode(): sam_outputs = sam_model(**sam_inputs) + # SAM2's processor takes the original sizes alone - the third positional + # argument is mask_threshold, and the SAM v1 'reshaped_input_sizes' key this + # used to pass does not exist in SAM2's inputs at all masks = sam_processor.post_process_masks( - sam_outputs.pred_masks, - sam_inputs["original_sizes"], - sam_inputs["reshaped_input_sizes"], + sam_outputs.pred_masks, sam_inputs["original_sizes"] ) combined = masks[0][:, 0].sum(dim=0).clamp(0, 1) diff --git a/dw/vram_estimate.py b/dw/vram_estimate.py index ab943541..95c59160 100644 --- a/dw/vram_estimate.py +++ b/dw/vram_estimate.py @@ -16,6 +16,10 @@ tensor holds, which is why the variables it scales with are named rather than assumed to be `num_frames` alone - an image template's width and height matter here just as much as a video template's frame count does. + +The entries checked are the serving device's own (see `_entries_for`); +`bytes_per_voxel` was calibrated on CUDA, so a check against a Mac's capacity +is an estimate of an estimate. """ import numbers @@ -56,7 +60,36 @@ def required_gb(estimate, values): return base + per_voxel * product / (1024**3) -def vram_estimate_errors(definition, arguments=None, supplied=()): +def _entries_for(cost, device_type, capacity_gb): + """The cost entries a projection is checked against on this device. + + Entries measured on this backend first. A backend no entry describes - a + Mac, against a catalog measured on CUDA cards - is checked against its own + capacity when it can report one, and against every entry when it cannot, + so an unreadable ceiling never turns the guard off.""" + entries = [entry for entry in cost if isinstance(entry, dict)] + if device_type is None: + return entries + matching = [ + entry + for entry in entries + if str(entry.get("device", "")).split(":")[0] == device_type + ] + if matching: + return matching + if capacity_gb is not None: + return [ + { + "name": f"this {device_type} device (recommended maximum)", + "vram_gb": round(capacity_gb, 1), + } + ] + return entries + + +def vram_estimate_errors( + definition, arguments=None, supplied=(), device_type=None, capacity_gb=None +): """Every `cost` entry a declared vram_estimate projects to exceed, as [{path, message}] - refused, not warned, since the failure this guards is an OOM partway through a run that already spent minutes loading.""" @@ -79,9 +112,7 @@ def vram_estimate_errors(definition, arguments=None, supplied=()): else "variables" ) errors = [] - for entry in cost: - if not isinstance(entry, dict): - continue + for entry in _entries_for(cost, device_type, capacity_gb): capacity = entry.get("vram_gb") if capacity is None or projected <= capacity: continue @@ -103,7 +134,7 @@ def vram_estimate_errors(definition, arguments=None, supplied=()): return errors -def apply_vram_estimate(definition, variables): +def apply_vram_estimate(definition, variables, device_type=None, capacity_gb=None): """The run-time half of the check above - the backstop for a caller that skips validate_workflow (or an inline/composed workflow static validation never saw). Raises ValueError for the same projection @@ -111,6 +142,8 @@ def apply_vram_estimate(definition, variables): step was always going to OOM on (#265).""" if not isinstance(variables, dict): return - errors = vram_estimate_errors(definition, variables) + errors = vram_estimate_errors( + definition, variables, device_type=device_type, capacity_gb=capacity_gb + ) if errors: raise ValueError(errors[0]["message"]) diff --git a/dw/workflow.py b/dw/workflow.py index a1a5b6a4..c0e7d0d6 100644 --- a/dw/workflow.py +++ b/dw/workflow.py @@ -91,7 +91,13 @@ from .pipeline_processors.pipeline import Pipeline from .tasks.model_cache import clear_model_cache from .tasks.task import Task -from . import get_device, empty_device_cache, device_memory_stats +from . import ( + get_device, + get_device_type, + empty_device_cache, + device_capacity_gb, + device_memory_stats, +) from .host_memory import release_host_caches from .security import ( validate_path, @@ -784,7 +790,11 @@ def validation_errors(self, arguments=None, composing=None): # an OOM the caller had no way to see coming (dw/vram_estimate.py, # #265) + vram_estimate_errors( - self.workflow_definition, arguments, supplied=set(arguments or {}) + self.workflow_definition, + arguments, + supplied=set(arguments or {}), + device_type=get_device_type(), + capacity_gb=device_capacity_gb(), ) # An 'attn_processor_type' whose Hub kernel this machine has no # build variant for - validated clean and then died 88s into @@ -994,7 +1004,12 @@ def _prepare_definition(self, workflow_def, arguments, base_dir): # validate_workflow, so this raises the same refusal rather # than starting a job the decode step was always going to OOM # on (dw/vram_estimate.py, #265) - apply_vram_estimate(workflow_def, variables) + apply_vram_estimate( + workflow_def, + variables, + device_type=get_device_type(), + capacity_gb=device_capacity_gb(), + ) # realize the variables - explicit references only (asset:, # output:, constant:, prompt:, a {media_type, location} dict). # Key-name conventions (an 'image'/'video'/'_type' argument) are diff --git a/dw/workflow_schema.json b/dw/workflow_schema.json index 3b667449..6da2916f 100644 --- a/dw/workflow_schema.json +++ b/dw/workflow_schema.json @@ -593,7 +593,7 @@ "type": "boolean" }, "disable_attention_slicing": { - "description": "Whether to disable automatic attention slicing on MPS devices. Attention slicing is enabled by default on MPS to reduce memory usage.", + "description": "Turn off the attention slicing MPS applies automatically. Worth setting for SDXL-family models on a Mac: MPS SDPA is fast at their head dim (64) and slicing made it 2.4x slower.", "type": "boolean" }, "attention_backend": { diff --git a/dw/workflows/test.json b/dw/workflows/test.json index 9705b5bd..2cfbbd56 100644 --- a/dw/workflows/test.json +++ b/dw/workflows/test.json @@ -13,7 +13,7 @@ }, "from_pretrained_arguments": { "model_name": "stable-diffusion-v1-5/stable-diffusion-v1-5", - "torch_dtype": "torch.float16" + "torch_dtype": "torch.bfloat16" }, "arguments": { "prompt": "variable:prompt", diff --git a/plugins/dw/skills/ltx-2.5/SKILL.md b/plugins/dw/skills/ltx-2.5/SKILL.md index 99991f41..fe81f608 100644 --- a/plugins/dw/skills/ltx-2.5/SKILL.md +++ b/plugins/dw/skills/ltx-2.5/SKILL.md @@ -12,9 +12,9 @@ caption spec below. ## Before anything -1. `get_server_info`: the device and the workspace. These templates quantize - with SDNQ on CUDA; on an `mps` or `cpu` server say so and stop - none of - them fit there. +1. `get_server_info`: the device and the workspace. On `mps` they run + slower than the CUDA `cost` says (only text-to-video has been run + there); on `cpu` say so and stop. 2. `list_workflows(shape="shot")`: every template in the family carries that shape. Take their current names, `summary`, `traits` and `cost` from the listing and trust it over the names quoted below. diff --git a/plugins/dw/skills/minimax-music3/SKILL.md b/plugins/dw/skills/minimax-music3/SKILL.md index fcdaa357..9f28af78 100644 --- a/plugins/dw/skills/minimax-music3/SKILL.md +++ b/plugins/dw/skills/minimax-music3/SKILL.md @@ -12,8 +12,9 @@ shapes; do not author a new workflow until the shape decision below fails. ## Before anything -1. `get_server_info`: the device. Music 3 is a CUDA model; on an `mps` or `cpu` - server stop and say so. +1. `get_server_info`: the device. On `mps` the `music` template runs, slower + than the CUDA `cost` says; `music-video` (it also runs H3) has not been + run there. On `cpu` stop and say so. 2. `list_workflows(shape="audio")`, and `shape="sequence"` for the music video. Trust the names quoted below only after the listing confirms them. 3. `get_workflow` on the one chosen, for its variables and their defaults. diff --git a/tests/test_attention_slicing.py b/tests/test_attention_slicing.py new file mode 100644 index 00000000..166146b4 --- /dev/null +++ b/tests/test_attention_slicing.py @@ -0,0 +1,31 @@ +"""When a pipeline's attention runs sliced. + +Automatic on MPS, opt-in elsewhere. Measured on an M5 Pro (torch 2.14), MPS +SDPA is slow at head dims 40, 48 and 160 and fast at 32 and 64-128: sliced +attention made SD 1.5 (head dims 40/80/160) ~20% faster end to end, while it +made SDXL-shaped attention (head dim 64) 2.4x slower. Automatic slicing stays +because turning it off regressed the catalog's quick-start template; a +head-dim-aware policy is the follow-up. Only UNet/ControlNet models take +slicing at all. +""" + +from dw.pipeline_processors.pipeline import attention_slicing_requested + + +def test_automatic_on_mps(): + assert attention_slicing_requested({}, "mps") is True + + +def test_mps_can_opt_out(): + assert ( + attention_slicing_requested({"disable_attention_slicing": True}, "mps") is False + ) + + +def test_off_by_default_elsewhere(): + assert attention_slicing_requested({}, "cuda") is False + assert attention_slicing_requested({}, "cpu") is False + + +def test_opt_in_anywhere(): + assert attention_slicing_requested({"enable_attention_slicing": True}, "cuda") diff --git a/tests/test_config_objects.py b/tests/test_config_objects.py index 81b6cb5a..159bf0b8 100644 --- a/tests/test_config_objects.py +++ b/tests/test_config_objects.py @@ -70,6 +70,83 @@ def __init__(self, **kwargs): } ) + def test_a_cuda_quantization_device_is_translated_on_a_mac(self, monkeypatch): + # A catalog template written on the CUDA box names "cuda" here, and SDNQ + # moves every weight to it inside from_pretrained + monkeypatch.setattr(dw, "backend_available", lambda backend: backend != "cuda") + monkeypatch.setattr(dw, "get_device", lambda: "mps") + + config = create_quantization_config( + quantization_definition(quantization_device="cuda", return_device="cpu") + ) + + assert config.kwargs["quantization_device"] == "mps" + assert config.kwargs["return_device"] == "cpu" + + def test_a_cpu_quantization_device_is_never_rewritten(self, monkeypatch): + monkeypatch.setattr(dw, "backend_available", lambda backend: backend != "cuda") + monkeypatch.setattr(dw, "get_device", lambda: "mps") + + config = create_quantization_config( + quantization_definition(quantization_device="cpu") + ) + + assert config.kwargs["quantization_device"] == "cpu" + + def test_quantized_matmul_is_switched_off_on_mps(self, monkeypatch): + # Without Triton SDNQ's quantized matmul is torch._int_mm, measured + # ~500x slower than the bf16 matmul it replaces on an M5 Pro + monkeypatch.setattr(dw, "backend_available", lambda backend: backend != "cuda") + monkeypatch.setattr(dw, "get_device", lambda: "mps") + + config = create_quantization_config( + quantization_definition( + use_quantized_matmul=True, use_quantized_matmul_conv=True + ) + ) + + assert config.kwargs["use_quantized_matmul"] is False + assert config.kwargs["use_quantized_matmul_conv"] is False + + def test_quantized_matmul_is_kept_on_cuda(self, monkeypatch): + monkeypatch.setattr(dw, "backend_available", lambda backend: True) + monkeypatch.setattr(dw, "get_device", lambda: "cuda") + + config = create_quantization_config( + quantization_definition( + quantization_device="cuda", use_quantized_matmul=True + ) + ) + + assert config.kwargs["quantization_device"] == "cuda" + assert config.kwargs["use_quantized_matmul"] is True + + def test_quantized_matmul_is_kept_on_cpu(self, monkeypatch): + monkeypatch.setattr(dw, "get_device", lambda: "cpu") + + config = create_quantization_config( + quantization_definition(use_quantized_matmul=True) + ) + + assert config.kwargs["use_quantized_matmul"] is True + + def test_adapting_does_not_mutate_the_definition(self, monkeypatch): + # A cached pipeline builds from the same definition a second time + monkeypatch.setattr(dw, "backend_available", lambda backend: backend != "cuda") + monkeypatch.setattr(dw, "get_device", lambda: "mps") + definition = quantization_definition( + quantization_device="cuda", use_quantized_matmul=True + ) + + create_quantization_config(definition) + second = create_quantization_config(definition) + + assert definition["arguments"] == { + "quantization_device": "cuda", + "use_quantized_matmul": True, + } + assert second.kwargs["quantization_device"] == "mps" + class TestLoadComponentsArguments: def test_a_pipeline_without_load_components_returns_none(self): @@ -150,7 +227,8 @@ def test_offload_defaults_to_cpu(self): assert config["offload_device"] == torch.device("cpu") - def test_other_keys_are_carried_through(self): + def test_other_keys_are_carried_through(self, monkeypatch): + monkeypatch.setattr(dw, "backend_available", lambda backend: True) config = get_group_offload_configuration( {"group_offload": {"num_blocks_per_group": 2, "use_stream": True}}, "cuda" ) @@ -168,6 +246,57 @@ def test_building_twice_from_one_configuration_is_stable(self): assert first["onload_device"] == second["onload_device"] == torch.device("cpu") + def test_streams_are_dropped_on_mps(self, monkeypatch): + # diffusers refuses use_stream without CUDA/XPU, and refuses + # record_stream without use_stream - both have to go together + monkeypatch.setattr(dw, "backend_available", lambda backend: backend != "cuda") + monkeypatch.setattr(dw, "get_device", lambda: "mps") + + config = get_group_offload_configuration( + { + "group_offload": { + "onload_device": "cuda", + "num_blocks_per_group": 1, + "use_stream": True, + "record_stream": True, + } + }, + "cuda", + ) + + assert config["onload_device"] == torch.device("mps") + assert "use_stream" not in config + assert "record_stream" not in config + assert config["num_blocks_per_group"] == 1 + + def test_streams_are_dropped_when_onloading_to_cpu(self): + config = get_group_offload_configuration( + {"group_offload": {"onload_device": "cpu", "use_stream": True}}, "cpu" + ) + + assert "use_stream" not in config + + def test_streams_are_kept_on_cuda(self, monkeypatch): + monkeypatch.setattr(dw, "backend_available", lambda backend: True) + + config = get_group_offload_configuration( + {"group_offload": {"use_stream": True, "record_stream": True}}, "cuda" + ) + + assert config["use_stream"] is True + assert config["record_stream"] is True + + def test_dropping_streams_twice_is_stable(self, monkeypatch): + monkeypatch.setattr(dw, "backend_available", lambda backend: backend != "cuda") + monkeypatch.setattr(dw, "get_device", lambda: "mps") + configuration = {"group_offload": {"use_stream": True, "record_stream": True}} + + get_group_offload_configuration(configuration, "cuda") + second = get_group_offload_configuration(configuration, "cuda") + + assert "use_stream" not in second + assert second["onload_device"] == torch.device("mps") + class TestCacheConfiguration: def test_no_cache_returns_none(self): diff --git a/tests/test_device_helpers.py b/tests/test_device_helpers.py index ce0d4e26..436d86cd 100644 --- a/tests/test_device_helpers.py +++ b/tests/test_device_helpers.py @@ -175,21 +175,43 @@ def raise_runtime_error(): assert stats["free_mb"] is None assert stats["total_mb"] is None - def test_mps_reports_zeroed_known_stats(self, monkeypatch): + def test_mps_reports_real_figures(self, monkeypatch): + mb = 1024 * 1024 monkeypatch.setattr(dw, "get_device_type", lambda device=None: "mps") monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True) + monkeypatch.setattr(dw, "_apple_chip_name", lambda: "Apple M5 Pro") + monkeypatch.setattr(torch.mps, "current_allocated_memory", lambda: 1000 * mb) + monkeypatch.setattr(torch.mps, "driver_allocated_memory", lambda: 1500 * mb) + monkeypatch.setattr(torch.mps, "recommended_max_memory", lambda: 64000 * mb) stats = dw.device_memory_stats() assert stats == { "available": True, - "device_name": "Apple Silicon (MPS)", - "allocated_mb": 0.0, - "reserved_mb": 0.0, - "free_mb": 0.0, - "total_mb": 0.0, + "device_name": "Apple M5 Pro (MPS)", + "allocated_mb": 1000.0, + "reserved_mb": 1500.0, + "free_mb": 62500.0, + "total_mb": 64000.0, } + def test_mps_without_a_capacity_reading_reports_none(self, monkeypatch): + def unavailable(): + raise RuntimeError("no Metal device") + + monkeypatch.setattr(dw, "get_device_type", lambda device=None: "mps") + monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True) + monkeypatch.setattr(dw, "_apple_chip_name", lambda: "Apple M5 Pro") + monkeypatch.setattr(torch.mps, "current_allocated_memory", lambda: 0) + monkeypatch.setattr(torch.mps, "driver_allocated_memory", lambda: 0) + monkeypatch.setattr(torch.mps, "recommended_max_memory", unavailable) + + stats = dw.device_memory_stats() + + assert stats["available"] is True + assert stats["total_mb"] is None + assert stats["free_mb"] is None + def test_cpu_reports_unavailable(self, monkeypatch): monkeypatch.setattr(dw, "get_device_type", lambda device=None: "cpu") @@ -280,3 +302,31 @@ def test_get_gpu_memory_mb_uses_allocated_mb(self, monkeypatch): ) assert _make_worker()._get_gpu_memory_mb() == 42.0 + + +class TestDeviceCapacity: + def test_mps_capacity_is_the_recommended_max(self, monkeypatch): + monkeypatch.setattr(dw, "get_device_type", lambda device=None: "mps") + monkeypatch.setattr(torch.mps, "recommended_max_memory", lambda: 62 * 1024**3) + + assert dw.device_capacity_gb() == 62.0 + + def test_cpu_has_no_capacity(self, monkeypatch): + monkeypatch.setattr(dw, "get_device_type", lambda device=None: "cpu") + + assert dw.device_capacity_gb() is None + + def test_a_failing_probe_is_none_not_an_exception(self, monkeypatch): + def broken(): + raise RuntimeError("boom") + + monkeypatch.setattr(dw, "get_device_type", lambda device=None: "mps") + monkeypatch.setattr(torch.mps, "recommended_max_memory", broken) + + assert dw.device_capacity_gb() is None + + @pytest.mark.skipif( + not torch.backends.mps.is_available(), reason="needs an Apple Silicon GPU" + ) + def test_this_mac_reports_a_capacity(self): + assert dw.device_capacity_gb("mps") > 1.0 diff --git a/tests/test_mps_fallback_warning.py b/tests/test_mps_fallback_warning.py new file mode 100644 index 00000000..3cff48a5 --- /dev/null +++ b/tests/test_mps_fallback_warning.py @@ -0,0 +1,28 @@ +"""torch announces an op the MPS backend lacks, when PYTORCH_ENABLE_MPS_FALLBACK +is set, with a UserWarning - and dw ignored every UserWarning, so a step that +quietly ran part of its model on the CPU left no trace. The fallback variable is +set by the fp4-fp8-for-torch-mps autoload install.sh adds on macOS.""" + +import subprocess +import sys + +PROBE = """ +import warnings, dw +fallback = ("The operator 'aten::_linalg_eigvals' is not currently supported " + "on the MPS backend and will fall back to run on the CPU. " + "This may have performance implications.") +with warnings.catch_warnings(record=True) as caught: + warnings.warn(fallback, UserWarning) + warnings.warn("some other library chatter", UserWarning) +print(len(caught), caught[0].message if caught else "") +""" + + +def test_only_the_fallback_warning_gets_through(): + result = subprocess.run( + [sys.executable, "-c", PROBE], capture_output=True, text=True, timeout=120 + ) + assert result.returncode == 0, result.stderr + count, _, message = result.stdout.strip().splitlines()[-1].partition(" ") + assert count == "1" + assert "fall back to run on the CPU" in message diff --git a/tests/test_mps_rope_precision.py b/tests/test_mps_rope_precision.py new file mode 100644 index 00000000..ae150864 --- /dev/null +++ b/tests/test_mps_rope_precision.py @@ -0,0 +1,64 @@ +"""RoPE asked for in float64 runs in float32 on MPS. + +MPS has no float64. diffusers already falls back to float32 RoPE there for +Wan, Lumina2, SkyReels-V2, ChronoEdit and Sana-Video, but LTX-2's transformer +and text connectors build their frequencies in float64 whenever the module's +`double_precision` is set (the default) - so LTX-2.5 on a Mac loaded, then +failed on its first step with "Cannot convert a MPS Tensor to float64". +""" + +import logging +from types import SimpleNamespace + +import torch + +from dw.pipeline_processors.pipeline import apply_mps_rope_precision + + +class Rope(torch.nn.Module): + def __init__(self, double_precision=True): + super().__init__() + self.double_precision = double_precision + + +class Transformer(torch.nn.Module): + def __init__(self): + super().__init__() + self.rope = Rope() + self.audio_rope = Rope() + + +def pipeline(): + return SimpleNamespace( + components={ + "transformer": Transformer(), + "connectors": Rope(), + "scheduler": object(), + "tokenizer": None, + } + ) + + +def test_float64_rope_runs_in_float32_on_mps(caplog): + p = pipeline() + + with caplog.at_level(logging.WARNING, logger="dw"): + apply_mps_rope_precision(p, "mps") + + assert p.components["transformer"].rope.double_precision is False + assert p.components["transformer"].audio_rope.double_precision is False + assert p.components["connectors"].double_precision is False + assert "float32" in caplog.text + + +def test_cuda_keeps_float64_rope(): + p = pipeline() + + apply_mps_rope_precision(p, "cuda") + + assert p.components["transformer"].rope.double_precision is True + assert p.components["connectors"].double_precision is True + + +def test_a_pipeline_without_components_is_left_alone(): + apply_mps_rope_precision(object(), "mps") diff --git a/tests/test_parallel_loading_default.py b/tests/test_parallel_loading_default.py new file mode 100644 index 00000000..18f79624 --- /dev/null +++ b/tests/test_parallel_loading_default.py @@ -0,0 +1,51 @@ +"""Parallel shard loading is on by default except on macOS. + +diffusers' parallel loading runs one thread per shard, and each thread moves +its tensors to the device. On MPS, concurrent copies from those threads +segfaulted LTX-2.5's SDNQ transformer load (exit 139, four ThreadPoolExecutor +threads inside copy_ to mps); the same load with parallel loading off took +39s and succeeded. diffusers reads HF_ENABLE_PARALLEL_LOADING once at import, +so the default is chosen before torch is imported - by platform, since macOS +has no CUDA and its accelerator is always MPS. CUDA boxes keep the faster load. +""" + +import os +import subprocess +import sys + +import pytest + +import dw + + +def test_off_on_macos(): + assert dw._parallel_loading_default("darwin") == "false" + + +@pytest.mark.parametrize("platform", ["linux", "win32"]) +def test_on_elsewhere(platform): + assert dw._parallel_loading_default(platform) == "true" + + +def _env_after_import(preset=None): + code = "import os, dw; print(os.environ['HF_ENABLE_PARALLEL_LOADING'])" + env = {k: v for k, v in os.environ.items() if k != "HF_ENABLE_PARALLEL_LOADING"} + if preset is not None: + env["HF_ENABLE_PARALLEL_LOADING"] = preset + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + env=env, + timeout=120, + ) + assert result.returncode == 0, result.stderr + return result.stdout.strip().splitlines()[-1] + + +def test_import_applies_this_platforms_default(): + assert _env_after_import() == dw._parallel_loading_default(sys.platform) + + +def test_an_explicit_setting_wins(): + assert _env_after_import(preset="1") == "1" diff --git a/tests/test_segment.py b/tests/test_segment.py index 854a42bd..4f38c9d6 100644 --- a/tests/test_segment.py +++ b/tests/test_segment.py @@ -55,7 +55,6 @@ def test_returns_pil_image_mode_l( "pixel_values": torch.zeros(1, 3, 256, 256), "input_boxes": torch.tensor([[[100.0, 100.0, 300.0, 300.0]]]), "original_sizes": torch.tensor([[480, 640]]), - "reshaped_input_sizes": torch.tensor([[256, 256]]), } ) mock_sam_proc_instance.return_value = sam_inputs @@ -74,6 +73,9 @@ def test_returns_pil_image_mode_l( image = _make_test_image() result = segment_image(image, "dog") + args, kwargs = mock_sam_proc_instance.post_process_masks.call_args + assert len(args) == 2 and not kwargs + assert isinstance(result, Image.Image) assert result.mode == "L" assert result.size == (640, 480) @@ -153,3 +155,87 @@ def test_segment_command_registered(self): from dw.tasks.task import _COMMAND_REGISTRY assert "segment" in _COMMAND_REGISTRY + + +class TestSegmentWithRealSam2Processor: + """The mocked tests above hand segment_image a dict with whatever keys + they are told to - which is how a SAM v1 key ('reshaped_input_sizes') + survived in the code while every real run raised KeyError. This one + builds the real processor offline and fakes only the models.""" + + @patch("dw.tasks.segment.Sam2Model") + @patch("dw.tasks.segment.Sam2Processor") + @patch("dw.tasks.segment.AutoModelForZeroShotObjectDetection") + @patch("dw.tasks.segment.AutoProcessor") + def test_masks_come_back_at_the_image_size( + self, mock_auto_proc, mock_auto_model, mock_sam_proc, mock_sam_model + ): + from transformers import Sam2ImageProcessor + from transformers import Sam2Processor as RealSam2Processor + + from dw.tasks.segment import segment_image + + dino = MagicMock() + mock_auto_proc.from_pretrained.return_value = dino + dino.return_value = _make_batch_encoding({"input_ids": torch.zeros(1, 10)}) + dino.post_process_grounded_object_detection.return_value = [ + { + "boxes": torch.tensor([[100.0, 100.0, 300.0, 300.0]]), + "scores": torch.tensor([0.9]), + "labels": ["dog"], + } + ] + dino_model = MagicMock() + mock_auto_model.from_pretrained.return_value = dino_model + dino_model.to.return_value = dino_model + + mock_sam_proc.from_pretrained.return_value = RealSam2Processor( + image_processor=Sam2ImageProcessor() + ) + sam_model = MagicMock() + mock_sam_model.from_pretrained.return_value = sam_model + sam_model.to.return_value = sam_model + # SAM2's low-res mask logits: (batch, boxes, masks, 256, 256) + logits = torch.full((1, 1, 3, 256, 256), -10.0) + logits[..., 64:128, 64:128] = 10.0 + sam_model.return_value = MagicMock(pred_masks=logits) + + result = segment_image(_make_test_image(), "dog") + + assert result.mode == "L" + assert result.size == (640, 480) + assert np.array(result).max() == 255 + + +class TestGroundingDinoQuery: + """GroundingDINO only scores a query written its way - lowercase, each + phrase ending in a period. 'cat' found nothing in a photo of a cat (max + score 0.20 under the 0.3 threshold); 'cat.' found it, on CPU and MPS.""" + + @patch("dw.tasks.segment.AutoModelForZeroShotObjectDetection") + @patch("dw.tasks.segment.AutoProcessor") + def _query_for(self, prompt, mock_auto_proc, mock_auto_model): + from dw.tasks.segment import segment_image + + dino = MagicMock() + mock_auto_proc.from_pretrained.return_value = dino + dino.return_value = _make_batch_encoding({"input_ids": torch.zeros(1, 10)}) + dino.post_process_grounded_object_detection.return_value = [ + {"boxes": torch.zeros(0, 4), "scores": torch.zeros(0), "labels": []} + ] + model = MagicMock() + mock_auto_model.from_pretrained.return_value = model + model.to.return_value = model + + segment_image(_make_test_image(), prompt) + + return dino.call_args.kwargs["text"] + + def test_a_bare_noun_gets_its_period(self): + assert self._query_for("cat") == "cat." + + def test_case_and_whitespace_are_normalized(self): + assert self._query_for(" A Red Car ") == "a red car." + + def test_a_query_already_in_form_is_unchanged(self): + assert self._query_for("a cat. a dog.") == "a cat. a dog." diff --git a/tests/test_template_defaults.py b/tests/test_template_defaults.py new file mode 100644 index 00000000..0a13aaef --- /dev/null +++ b/tests/test_template_defaults.py @@ -0,0 +1,59 @@ +"""A catalog default must run: no placeholder hosts, and one spelling per repo. + +upscale-* defaulted to https://example.com/image.jpg (not an image), and +flux-torchao spelled FLUX.1-dev two ways - the Hub cache is keyed on the +spelling, so the transformer downloaded twice (~23 GB).""" + +import json +import pathlib +import re + +ROOT = pathlib.Path(__file__).resolve().parent.parent +FILES = sorted([*ROOT.glob("workflows/**/*.json"), *ROOT.glob("dw/workflows/*.json")]) +PLACEHOLDER = re.compile(r"https?://(www\.)?example\.(com|org|net)/") +WITHDRAWN = ("runwayml/",) + + +def strings(node): + if isinstance(node, dict): + for value in node.values(): + yield from strings(value) + elif isinstance(node, list): + for value in node: + yield from strings(value) + elif isinstance(node, str): + yield node + + +def test_no_default_points_at_a_placeholder_host(): + offenders = [ + f"{path.relative_to(ROOT)}: {s}" + for path in FILES + for s in strings(json.loads(path.read_text())) + if PLACEHOLDER.search(s) + ] + assert offenders == [] + + +def test_no_withdrawn_repos(): + offenders = [ + f"{path.relative_to(ROOT)}: {s}" + for path in FILES + for s in strings(json.loads(path.read_text())) + if s.startswith(WITHDRAWN) + ] + assert offenders == [] + + +def test_one_spelling_per_repo_within_a_file(): + for path in FILES: + repos = { + s + for s in strings(json.loads(path.read_text())) + if re.fullmatch(r"[\w.-]+/[\w.-]+", s) + } + lowered = {} + for repo in repos: + lowered.setdefault(repo.lower(), set()).add(repo) + clashes = [names for names in lowered.values() if len(names) > 1] + assert clashes == [], f"{path.relative_to(ROOT)}: {clashes}" diff --git a/tests/test_transcribe_audio.py b/tests/test_transcribe_audio.py index 3391e6e8..14940430 100644 --- a/tests/test_transcribe_audio.py +++ b/tests/test_transcribe_audio.py @@ -24,8 +24,9 @@ def _waveform(self, channels=1, seconds=1.0, sample_rate=16000): tone = numpy.zeros((channels, samples), dtype=numpy.float32) return tone, sample_rate - def _mock_pipe(self, mock_pipeline, text="hello world"): + def _mock_pipe(self, mock_pipeline, text="hello world", kind="seq2seq_whisper"): pipe = MagicMock() + pipe.type = kind pipe.return_value = {"text": text} mock_pipeline.return_value = pipe return pipe @@ -71,6 +72,39 @@ def test_downmixes_stereo_to_mono(self, mock_pipeline): fed = pipe.call_args[0][0] self.assertEqual(fed["raw"].ndim, 1) + @patch("dw.tasks.audio_transcription.hf_pipeline") + def test_a_clip_over_thirty_seconds_asks_for_timestamps(self, mock_pipeline): + # Whisper refuses more than 30 s of audio ("more than 3000 mel input + # features") unless it predicts timestamps - so a 110 s song failed + pipe = self._mock_pipe(mock_pipeline) + waveform, rate = self._waveform(seconds=31.0) + + transcribe_audio(waveform, device="cpu", sample_rate=rate) + + self.assertIs(pipe.call_args.kwargs.get("return_timestamps"), True) + + @patch("dw.tasks.audio_transcription.hf_pipeline") + def test_a_long_clip_on_a_ctc_model_does_not_ask_for_timestamps( + self, mock_pipeline + ): + # transformers raises for a CTC model unless return_timestamps is + # "char" or "word"; a CTC model has no 30 s window to stitch anyway + pipe = self._mock_pipe(mock_pipeline, kind="ctc") + waveform, rate = self._waveform(seconds=31.0) + + transcribe_audio(waveform, device="cpu", sample_rate=rate) + + self.assertNotIn("return_timestamps", pipe.call_args.kwargs) + + @patch("dw.tasks.audio_transcription.hf_pipeline") + def test_a_short_clip_does_not_ask_for_timestamps(self, mock_pipeline): + pipe = self._mock_pipe(mock_pipeline) + waveform, rate = self._waveform(seconds=5.0) + + transcribe_audio(waveform, device="cpu", sample_rate=rate) + + self.assertNotIn("return_timestamps", pipe.call_args.kwargs) + @patch("dw.tasks.audio_transcription.hf_pipeline") def test_resamples_to_16khz(self, mock_pipeline): pipe = self._mock_pipe(mock_pipeline) diff --git a/tests/test_vram_estimate.py b/tests/test_vram_estimate.py new file mode 100644 index 00000000..c19e4f62 --- /dev/null +++ b/tests/test_vram_estimate.py @@ -0,0 +1,72 @@ +"""vram_estimate refuses a shape against the device actually serving it. + +Every catalog 'cost' entry is a CUDA card, so checking all of them refused a +64 GB Mac against a 24 GB RTX 3090 - and filtering to matching entries alone +would have switched the guard off on a Mac entirely, since none match. +""" + +import pytest + +from dw.vram_estimate import apply_vram_estimate, vram_estimate_errors + + +def definition(extra_cost=()): + # 768 * 512 * 121 voxels * 200 bytes = 8.86 GiB, + 20 base = 28.86 GiB + return { + "variables": {"width": 768, "height": 512, "num_frames": 121}, + "cost": [ + {"device": "cuda", "name": "RTX 3090", "vram_gb": 24, "minutes": 1.8}, + *extra_cost, + ], + "vram_estimate": { + "voxel_variables": ["width", "height", "num_frames"], + "base_gb": 20, + "bytes_per_voxel": 200, + }, + } + + +def test_no_device_checks_every_entry_as_before(): + errors = vram_estimate_errors(definition()) + assert len(errors) == 1 and "RTX 3090" in errors[0]["message"] + + +def test_a_cuda_box_checks_its_cuda_entries(): + errors = vram_estimate_errors(definition(), device_type="cuda", capacity_gb=24) + assert len(errors) == 1 and "RTX 3090" in errors[0]["message"] + + +def test_a_mac_with_room_is_not_refused_by_a_cuda_card(): + assert vram_estimate_errors(definition(), device_type="mps", capacity_gb=62) == [] + + +def test_a_mac_without_room_is_refused_against_its_own_capacity(): + errors = vram_estimate_errors(definition(), device_type="mps", capacity_gb=16) + assert len(errors) == 1 + assert "mps" in errors[0]["message"] + assert "RTX 3090" not in errors[0]["message"] + + +def test_a_mac_whose_capacity_is_unknown_keeps_the_conservative_check(): + errors = vram_estimate_errors(definition(), device_type="mps", capacity_gb=None) + assert len(errors) == 1 and "RTX 3090" in errors[0]["message"] + + +def test_a_curated_mps_entry_wins_over_the_measured_capacity(): + mac = {"device": "mps", "name": "M5 Pro 64 GB", "vram_gb": 26, "minutes": 9} + errors = vram_estimate_errors(definition([mac]), device_type="mps", capacity_gb=62) + assert len(errors) == 1 and "M5 Pro 64 GB" in errors[0]["message"] + + +def test_an_indexed_cost_device_matches_its_backend(): + card = {"device": "cuda:1", "name": "second card", "vram_gb": 48, "minutes": 1} + d = definition() + d["cost"] = [card] + assert vram_estimate_errors(d, device_type="cuda", capacity_gb=24) == [] + + +def test_run_time_backstop_takes_the_device_too(): + d = definition() + apply_vram_estimate(d, d["variables"], device_type="mps", capacity_gb=62) + with pytest.raises(ValueError): + apply_vram_estimate(d, d["variables"], device_type="mps", capacity_gb=16) diff --git a/workflows/models/flux-torchao.json b/workflows/models/flux-torchao.json index ff55f7f1..8032e53a 100644 --- a/workflows/models/flux-torchao.json +++ b/workflows/models/flux-torchao.json @@ -35,7 +35,7 @@ } }, "from_pretrained_arguments": { - "model_name": "black-forest-labs/Flux.1-Dev", + "model_name": "black-forest-labs/FLUX.1-dev", "subfolder": "transformer", "torch_dtype": "torch.bfloat16" } diff --git a/workflows/templates/attention-processor.json b/workflows/templates/attention-processor.json index b5f5a256..56839bb9 100644 --- a/workflows/templates/attention-processor.json +++ b/workflows/templates/attention-processor.json @@ -21,7 +21,7 @@ } }, "from_pretrained_arguments": { - "model_name": "runwayml/stable-diffusion-v1-5", + "model_name": "stable-diffusion-v1-5/stable-diffusion-v1-5", "torch_dtype": "torch.float16", "safety_checker": null }, diff --git a/workflows/templates/audio-trim-fade.json b/workflows/templates/audio-trim-fade.json index e23c2662..0b5655aa 100644 --- a/workflows/templates/audio-trim-fade.json +++ b/workflows/templates/audio-trim-fade.json @@ -2,7 +2,7 @@ "id": "TrimFadeAudio", "description": "Trims an audio file to a duration and fades it in and out - a finishing pass for a track that ends abruptly. Leaving 'duration_seconds' unset degrades to a fade-only pass over the whole track rather than failing.", "variables": { - "input_audio": "https://example.com/audio.mp3", + "input_audio": "https://huggingface.co/datasets/Narsil/asr_dummy/resolve/main/mlk.flac", "start_seconds": 0.0, "duration_seconds": null, "fade_in_ms": 0, diff --git a/workflows/templates/controlnet-component.json b/workflows/templates/controlnet-component.json index b5209bed..8bc24bc9 100644 --- a/workflows/templates/controlnet-component.json +++ b/workflows/templates/controlnet-component.json @@ -1,6 +1,6 @@ { "variables": { - "mask_image_uri": "https://nftfactory.blob.core.windows.net/images/chia256.jpg", + "mask_image_uri": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png", "seed": null }, "id": "sd15-controlnet", @@ -34,7 +34,7 @@ } }, "from_pretrained_arguments": { - "model_name": "runwayml/stable-diffusion-v1-5", + "model_name": "stable-diffusion-v1-5/stable-diffusion-v1-5", "torch_dtype": "torch.float16" }, "arguments": { diff --git a/workflows/templates/image-to-image.json b/workflows/templates/image-to-image.json index cfe83a6f..83be81ac 100644 --- a/workflows/templates/image-to-image.json +++ b/workflows/templates/image-to-image.json @@ -2,7 +2,7 @@ "variables": { "prompt": "an evil version of this image", "image": { - "location": "https://raw.githubusercontent.com/Hadamcik/sage-permission-probe/3e0f80d1fbed7ea0ad118be1771436b8dda59bba/public/icon.png" + "location": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png" }, "num_images_per_prompt": 1, "num_inference_steps": 15, diff --git a/workflows/templates/image-variation.json b/workflows/templates/image-variation.json index 8c7c9734..a86c3efe 100644 --- a/workflows/templates/image-variation.json +++ b/workflows/templates/image-variation.json @@ -1,6 +1,6 @@ { "variables": { - "input_image": "https://pbs.twimg.com/media/GgpbmxGWkAATZtp?format=jpg&name=small", + "input_image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg", "num_images_per_prompt": 4, "seed": null }, diff --git a/workflows/templates/segment-and-inpaint.json b/workflows/templates/segment-and-inpaint.json index ea8db77a..384503d8 100644 --- a/workflows/templates/segment-and-inpaint.json +++ b/workflows/templates/segment-and-inpaint.json @@ -2,7 +2,7 @@ "id": "SegmentAndInpaint", "description": "Segments an input image, then inpaints the selected region with FLUX.1 Fill dev.", "variables": { - "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/grounding_dino_example_input.png", + "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png", "segment_prompt": "cat", "inpaint_prompt": "a golden retriever puppy sitting on the grass", "seed": null diff --git a/workflows/templates/segment.json b/workflows/templates/segment.json index 5c958e4c..ed2078ef 100644 --- a/workflows/templates/segment.json +++ b/workflows/templates/segment.json @@ -2,7 +2,7 @@ "id": "Segment", "description": "Segments input images into labeled region masks.", "variables": { - "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/grounding_dino_example_input.png", + "image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png", "prompt": "cat" }, "steps": [ diff --git a/workflows/templates/transcribe-audio.json b/workflows/templates/transcribe-audio.json index 4e55cecf..e8a54951 100644 --- a/workflows/templates/transcribe-audio.json +++ b/workflows/templates/transcribe-audio.json @@ -2,7 +2,7 @@ "id": "TranscribeAudio", "description": "Transcribes spoken audio to text with a local Whisper-class model. Useful for checking a TTS deliverable's word correctness directly against the text it was supposed to speak, rather than inferring a dropped line or a mid-sentence truncation from duration and timing.", "variables": { - "input_audio": "https://example.com/speech.mp3", + "input_audio": "https://huggingface.co/datasets/Narsil/asr_dummy/resolve/main/mlk.flac", "model_name": "openai/whisper-base" }, "steps": [ diff --git a/workflows/templates/upscale-diffusion.json b/workflows/templates/upscale-diffusion.json index 624e4556..9630b158 100644 --- a/workflows/templates/upscale-diffusion.json +++ b/workflows/templates/upscale-diffusion.json @@ -7,7 +7,7 @@ "cost_drivers": ["num_inference_steps"], "variables": { "input_image": { - "location": "https://example.com/image.jpg" + "location": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png" }, "prompt": "", "negative_prompt": null, diff --git a/workflows/templates/upscale-spandrel.json b/workflows/templates/upscale-spandrel.json index 49b368d2..09fc0b91 100644 --- a/workflows/templates/upscale-spandrel.json +++ b/workflows/templates/upscale-spandrel.json @@ -6,7 +6,7 @@ ], "variables": { "input_image": { - "location": "https://example.com/image.jpg" + "location": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png" }, "upscale_model": "Kim2091/UltraSharp", "upscale_filename": "4x-UltraSharp.pth",