diff --git a/docs/ACCELERATION.md b/docs/ACCELERATION.md index 8791f4d7..ebfb3f24 100644 --- a/docs/ACCELERATION.md +++ b/docs/ACCELERATION.md @@ -136,6 +136,18 @@ Select the attention implementation diffusers uses for the duration of each pipe Common values: `"flash"`, `"flash_hub"`, `"sage"`, `"sage_hub"`, `"native"`, `"flex"`. The full set is diffusers' `AttentionBackendName` enum - availability depends on what's installed (`flash-attn`, `sageattention`, etc.) and the platform. `_hub`-suffixed backends are fetched from the Hugging Face Hub kernel registry on first use, which needs the `kernels` package installed (`pip install kernels`) - it is not a dw dependency, and no bundled workflow sets a backend, so each runs on a plain install. diffusers also refuses a Hub kernel whose publisher it does not list as trusted - `sage_hub` among them, since SageAttention 2 comes from `SageAttention/sage-attention` - until `DIFFUSERS_TRUST_REMOTE_KERNELS=true`, which allows remote code from any such publisher. On an RTX 3090 `sage_hub` measured about 30% faster per MiniMax-H3 step (13.8 s against 19.2 s at 960x544x124, 54.7 s against 81.2 s at 345 frames); the local `sage` backend needs `sageattention>=2.1.1`, which PyPI does not carry. +The MiniMax-H3 baselines expose this as a variable: `attention_backend` on +[video-with-audio.json](../workflows/templates/minimax/video-with-audio.json), +[video-with-audio-768p.json](../workflows/templates/minimax/video-with-audio-768p.json), +[image-to-video.json](../workflows/templates/minimax/image-to-video.json) and +[reference-to-video.json](../workflows/templates/minimax/reference-to-video.json), pinned +on the transformer and `null` by default. Pass `attention_backend=sage_hub` to a box with +the `kernels` package and `DIFFUSERS_TRUST_REMOTE_KERNELS=true` in its `.env`; on MPS +leave it null. The other H3 templates take the same key under `components.transformer` +(`transformer_ref` for a ref2va workflow). LTX-2.5 cannot use it: its transformer +passes an attention mask, and sage attention refuses one ("`attn_mask` is not supported +for sage attention", measured 2026-10-10), so the LTX templates expose no backend. + A component can also pin its backend persistently instead, via `set_attention_backend`: ```json @@ -232,7 +244,11 @@ Three older per-component knobs, set in the pipeline `configuration` beside ## Memory Offloading -`offload` (`"model"` or `"sequential"`) and `group_offload` trade speed for VRAM by streaming weights between system memory and the accelerator instead of keeping everything resident. `"model"` moves whole submodules and costs the least speed; `"sequential"` moves individual layers and is the slowest but uses the least memory; block/leaf-level `group_offload` sits between the two and is what a modular pipeline's self-loaded components use, since they aren't reachable in time for `offload`. Full configuration syntax is in [WORKFLOW_GUIDE.md](WORKFLOW_GUIDE.md#memory-offloading). Omit both for the fastest run, when VRAM allows it. +`offload` (`"model"` or `"sequential"`) and `group_offload` trade speed for VRAM by streaming weights between system memory and the accelerator instead of keeping everything resident. `"model"` moves whole submodules and costs the least speed; `"sequential"` moves individual layers and is the slowest but uses the least memory; block/leaf-level `group_offload` sits between the two and is what a modular pipeline's self-loaded components use, since they aren't reachable in time for `offload`. Full configuration syntax is in [WORKFLOW_GUIDE.md](WORKFLOW_GUIDE.md#memory-offloading). Omit both for the fastest run, when VRAM allows it. Measured on an RTX 3090 (RECIPES_24GB.md): sequential +offload costs about 6x per step on FLUX dev and 3.4x on Z-Image against model offload; +a block-level, streamed `group_offload` of the transformer from pinned host memory +matches the resident per-step cost while holding one block on the card, at ~22 GB of +pinned host RAM per FLUX pipeline - the placement the unquantized FLUX templates use. `"residency": "on_demand"` on a component is the cheap case of the same trade: the model rests in system memory and is moved to the device whole around each of its own calls. That is a bad deal for anything called once per step, and a good one for a VAE called twice a run - it frees the VAE's VRAM for the denoise loop at the cost of two transfers, where group offloading the same VAE would restream it once per decode tile. See [On-demand components](WORKFLOW_GUIDE.md#on-demand-components). diff --git a/docs/LORAS.md b/docs/LORAS.md index c08a8c04..788cf6cb 100644 --- a/docs/LORAS.md +++ b/docs/LORAS.md @@ -48,6 +48,10 @@ LoRA (Low-Rank Adaptation) models apply lightweight style or subject modificatio Any other property (e.g. `revision`) is forwarded as-is to the underlying `load_lora_weights()` call. +`scale` and `alpha` are set in place on the loaded adapter, so a server run that changes +only them reuses the warm pipeline instead of reloading it; changing `model_name` or +`weight_name` reloads. + ## Multiple LoRAs Stack multiple LoRAs. They are blended via weighted adapter composition: diff --git a/docs/RECIPES_24GB.md b/docs/RECIPES_24GB.md index 6385aae5..5d1346f1 100644 --- a/docs/RECIPES_24GB.md +++ b/docs/RECIPES_24GB.md @@ -12,23 +12,79 @@ The general recipe, in order of impact: ## Flux dev (12B) -The bf16 transformer is ~24GB - it does not fit alongside the T5 encoder, and it does -not fit alongside the VAE either: a resident bf16 transformer runs the denoise and then -fails at decode. `offload: "model"` is what makes bf16 work on 24GB, and on an RTX 3090 -it is also the fastest configuration measured (1024x1024, 28 steps, pipeline loaded): - -| Approach | Config | Measured on RTX 3090 | -| -------- | ------ | -------------------- | -| **bf16 + model offload** | [flux-dev.json](../workflows/models/flux-dev.json) | 55s per image (72s cold) | -| **bf16 + model offload + compile** | [flux-dev-compile.json](../workflows/models/flux-dev-compile.json) - `compile: {repeated_blocks: true}` on the transformer | 52s per image (64s cold) | -| **int8 TorchAO** | transformer `quant_type: "torchao.quantization.Int8WeightOnlyConfig"` + `compile`, with or without `cache: first_block` | over 60s per *denoising step* - do not use on Ampere | -| **float8 TorchAO** (RTX 40-series+) | `Float8DynamicActivationFloat8WeightConfig` + `compile` | needs compute capability 8.9+ (Ada); unmeasured | -| **GGUF Q8** | [flux-gguf.json](../workflows/models/flux-gguf.json) - `from_single_file` Q8_0 transformer + `offload: "model"` | unmeasured; smallest VRAM, best quality retention | - -Measured 2026-09-07. A number in this table came from a run on the named card; a row -without one is a configuration that loads, not a speed claim. - -**Examples:** [flux-dev-compile.json](../workflows/models/flux-dev-compile.json), [flux-gguf.json](../workflows/models/flux-gguf.json), [step-caching.json](../workflows/templates/step-caching.json) +The bf16 transformer is ~22 GB, so on 24 GB it cannot sit beside its own activations +with any margin, and it cannot sit beside the T5 encoder or the VAE at all. Three +placements, measured on an RTX 3090 (2026-10-10, cold load included): + +| Placement | 1024x1024, 25 steps, 1 image, dev + LoRA | 768x768, 25 steps, 4 images, Krea | 1232x1632, 50 steps, Fill | +| --------- | ----------------------------------------- | --------------------------------- | ------------------------- | +| `offload: "sequential"` | 273 s (10.8 s/step) | not measured - same per-step cost | not measured | +| `offload: "model"` | 61 s (1.7 s/step) at **23.3 GB reserved, 491 MB free** | **OOM** at 23.45 GiB | not attempted | +| **transformer `group_offload`, block-level, streamed, pinned host copies** | 78 s (1.5 s/step), 17.7 GB peak reserved, 6 GB free | 131 s (3.2 s/step) | 152 s denoise (3.0 s/step), 1 s decode | + +Model offload is the fastest placement that loads, and the one #580 (9c708a12) retired: +at 1024x1024 it runs with a few hundred MB to spare and fails on anything larger - +a batch, a second encoder resident for prompt weighting, a cache's state. Sequential +offload is safe everywhere and 6x slower per step. The streamed group offload is both: +the transformer's blocks are copied in one at a time from pinned host memory under a +CUDA stream, so each step costs the same as the resident model's, and the card holds +one block plus the activations. It is what every unquantized FLUX template now uses: + +```json +"configuration": { + "component_type": "FluxPipeline", + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } +} +``` + +Two things the shape depends on. The non-offloaded components are placed explicitly: +with no pipeline-level `offload`, nothing else moves them, and a T5 or VAE left on the +CPU costs a minute per prompt encode and per decode (measured: 70 s before the first +step, 185 s decoding). And the host copies are pinned - diffusers' default, +`low_cpu_mem_usage: false` - which holds ~22 GB of unswappable host memory per loaded +FLUX pipeline and takes ~35 s once per load; with `low_cpu_mem_usage: true` the copies +come from pageable memory at 8 s/step, barely ahead of sequential. A host with less +than 32 GB free should take `offload: "sequential"` instead. + +Compile and caching stack on top as before: `compile: {repeated_blocks: true}` measured +52 s vs 55 s per image on the old model-offload placement; `cache: first_block` is +[step-caching.json](../workflows/templates/step-caching.json). TorchAO int8 is over 60 s +per *denoising step* on Ampere - do not use it there; float8 needs compute capability +8.9+ (Ada) and is unmeasured; GGUF Q8 ([flux-gguf.json](../workflows/models/flux-gguf.json)) +loads with the least VRAM and is unmeasured for speed. + +**Examples:** [lora.json](../workflows/templates/lora.json), [flux-dev-compile.json](../workflows/models/flux-dev-compile.json), [flux-gguf.json](../workflows/models/flux-gguf.json), [step-caching.json](../workflows/templates/step-caching.json) + +## Z-Image Turbo (6B) + +Fits a 24 GB card in bf16 with model offload, and that is the configuration to use. +Measured on an RTX 3090 (960x544, 9 steps, 4 images per prompt, cold load included, +2026-10-10): + +| Offload | Load | Denoise | Decode | Total | Peak VRAM reserved | +| ------- | ---- | ------- | ------ | ----- | ------------------ | +| `sequential` | 18 s | 84 s | 2 s | 107 s | 1.1 GB | +| **`model`** | 4 s | 25 s | 9 s | 41 s | 13.0 GB | +| none (resident) | 11 s | 17 s | 2 s | 33 s | 23.7 GB, 95 MB free at decode | + +Sequential offload streams the transformer's layers every step and costs 3.4x on the +denoise for a model that does not need the room. Resident is faster still but decodes +four images with nothing to spare, so [z-image.json](../workflows/models/z-image.json) +and the Z-Image portrait steps of the MiniMax templates use `"offload": "model"`. The +decode's 9 s under model offload is the VAE swapping in; `vae.enable_slicing` would +trade it for per-image decodes. ## Qwen-Image (20B) diff --git a/docs/RELEASING.md b/docs/RELEASING.md index a6d47862..f1fa844a 100644 --- a/docs/RELEASING.md +++ b/docs/RELEASING.md @@ -7,6 +7,47 @@ notes from commits at tag time (see below). This section is a scratch pad for items a branch's author wants the next release note to name; clear it when a release ships. +### 0.11.1 + + + +A speed release. The unquantized FLUX templates run about 3.5x faster on +a 24 GB card, Z-Image Turbo about 2.6x, a LoRA scale or scheduler shift +change no longer reloads the model, and switching workflows keeps the +models the next one shares. Every number below was measured on an RTX 3090; +the placement tables are in docs/RECIPES_24GB.md. + +**Workflow speed** + +- The 14 unquantized FLUX templates and `models/flux-krea` stream the + transformer block by block from pinned host memory (`group_offload`) + instead of `offload: "sequential"`: a 25-step 1024x1024 image takes 78 s + instead of 273 s (1.5 s/step against 10.8), with 6 GB of VRAM to spare - + enough for the 4-image batches `offload: "model"` runs out of memory on. + It holds ~22 GB of pinned host RAM per loaded FLUX pipeline. +- Z-Image Turbo (`models/z-image` and the MiniMax portrait steps) takes + `offload: "model"`: four 960x544 images in 41 s instead of 107 s. +- The four MiniMax-H3 family baselines expose `attention_backend` (null by + default; `sage_hub` measured about 30% faster per step on a 3090; it needs + the `kernels` package and `DIFFUSERS_TRUST_REMOTE_KERNELS=true`, see + docs/ACCELERATION.md). LTX-2.5 cannot use sage attention (its attention + mask is refused). +- A LoRA `scale`/`alpha` change or a scheduler `shift` change re-applies on + the warm pipeline instead of reloading it. Removing an alpha or shift still + reloads, and a `pipeline_reference` runs at the referenced step's own values. +- Switching workflows in the worker keeps the warm pipelines the next + workflow loads anyway, instead of releasing everything first. +- `restore_faces` builds its face helper once per run instead of once per + frame, and picture frames decode with libav threading, so video inputs + process faster. +- Catalog defaults corrected: compose-workflows video steps run 5 steps + (was 20) and the prompt-weighting schnell workflow runs 4 steps at guidance + 0 (was 20). +- The six 9-step minimax ref2va templates drop a dead `first_block` cache, + and minimax/shots-batch drops its `release_pipeline`: a no-op inside the + run, but it unloaded H3 at job end, so H3 now stays resident after a + shots-batch job, which is the win. + ### 0.11.0 diff --git a/docs/SERVER.md b/docs/SERVER.md index b6e52603..d4a95dbb 100644 --- a/docs/SERVER.md +++ b/docs/SERVER.md @@ -222,7 +222,9 @@ Among the free cards a job fits, dispatch prefers the one it has an affinity for: a rerun goes to the card that ran the original, and any job to the card whose worker last ran the same workflow (the same file, or an inline definition's `id`) - that worker's pipelines and step cache are warm, and a -worker frees both when the workflow changes. A busy preferred card is not +worker's step cache and any model the next workflow does not load are freed +when the workflow changes, while a model both workflows load (same weights +identity) stays warm. `clear_memory` still empties the card by hand. A busy preferred card is not waited for. A validation's cache probe and its `plan.estimate` ask about the card the job would be dispatched to now, and the estimate names it in `priced_for`. Memory is per card: `GET /api/memory` and `POST diff --git a/docs/TASKS.md b/docs/TASKS.md index c88c724c..026f9298 100644 --- a/docs/TASKS.md +++ b/docs/TASKS.md @@ -2632,7 +2632,7 @@ Restore and enhance faces in images using spandrel-compatible face restoration m | `eye_dist_threshold` | No | Skip faces with eye distance below this (default: 5) | | `upsample_img` | No | Pre-upscaled background image (e.g., from a prior upscale step) | -Models are loaded via spandrel, so any `.pth`/`.safetensors` face restoration weights work. CodeFormer requires `pip install spandrel-extra-arches` (non-commercial license). +Models are loaded via spandrel, so any `.pth`/`.safetensors` face restoration weights work. CodeFormer requires `pip install spandrel-extra-arches` (non-commercial license). The detector and parser are loaded once per settings/device and reused across frames; `release_models` drops them with the restorer. **Example:** [restore-faces.json](../workflows/templates/restore-faces.json) — Generate a portrait, then restore faces with GFPGAN v1.4. diff --git a/docs/WORKER_GUIDE.md b/docs/WORKER_GUIDE.md index c05822cc..a90c0594 100644 --- a/docs/WORKER_GUIDE.md +++ b/docs/WORKER_GUIDE.md @@ -28,6 +28,13 @@ The worker cleans up automatically between runs (garbage collection + GPU cache clearing). If memory grows unexpectedly, `GET /api/memory` reports the current reading and `POST /api/memory/clear` resets it. +Switching to a different workflow releases only the models the next workflow +does not load: the worker prepares the incoming definition, and a pipeline +both workflows load (same weights identity) stays warm. The step cache and +the task model cache are still emptied, and if the incoming definition cannot +be prepared everything is released. `clear_memory` (`POST /api/memory/clear`) +remains the way to empty the card by hand. + ## Troubleshooting **Worker crashes**: `JobManager` detects it and starts a fresh worker on the diff --git a/docs/WORKFLOW_GUIDE.md b/docs/WORKFLOW_GUIDE.md index 286f5bb9..94a43ab4 100644 --- a/docs/WORKFLOW_GUIDE.md +++ b/docs/WORKFLOW_GUIDE.md @@ -112,6 +112,10 @@ inverted latents: ``` `reference_name` must name a step earlier in the same workflow that has a `pipeline`. +Two steps whose pipelines differ only in a LoRA's `scale` or `alpha` or a scheduler's +`shift` share one warm model, and a reference always runs at the referenced step's own +values; `release_pipeline` on either step frees the shared model, as it does for two +identical steps. See [workflows/templates/community-pipeline.json](../workflows/templates/community-pipeline.json) for a full example. ### Task Steps @@ -1173,7 +1177,7 @@ Control how models use memory: ``` - `"model"` — Moves entire models between CPU and GPU. Good balance of speed and memory. -- `"sequential"` — Moves individual layers. Slowest but uses least GPU memory. On MPS it is downgraded to `"model"` with a warning: with unified memory there is no separate pool to keep small, so the per-layer copies cost speed and save nothing. +- `"sequential"` — Moves individual layers. Slowest but uses least GPU memory (measured at 6x the per-step cost of model offload on FLUX dev and 3.4x on Z-Image; a streamed block-level `group_offload` of the transformer is the usual better answer - see [RECIPES_24GB.md](RECIPES_24GB.md#flux-dev-12b)). On MPS it is downgraded to `"model"` with a warning: with unified memory there is no separate pool to keep small, so the per-layer copies cost speed and save nothing. `exclude_from_cpu_offload` names components the sweep should leave alone. - Omit for no offloading (fastest, requires enough VRAM). @@ -1325,6 +1329,13 @@ The step-level `release_pipeline` flag unloads the step's pipeline after its res saved. A later `pipeline_reference` to a released step is an error, and the process-wide step cache will not retain it. +What keeps a pipeline warm: the server's worker keeps a pipeline loaded across runs for as +long as what it *loads* is unchanged. A LoRA's `scale` and `alpha` and a scheduler's (or +`audio_scheduler`'s) `shift` are applied in place on the warm pipeline, so iterating on +them costs no reload; removing (nulling) an alpha or a shift reloads, so the checkpoint +default comes back. Changing a LoRA's `model_name` or `weight_name`, the quantization, +the placement or the scheduler type reloads it. + #### Releasing task models mid-workflow Task models - the checkpoints behind `text_generation`, `segment`, `depth_estimator` and diff --git a/docs/superpowers/plans/2026-10-10-workflow-speed.md b/docs/superpowers/plans/2026-10-10-workflow-speed.md new file mode 100644 index 00000000..77651954 --- /dev/null +++ b/docs/superpowers/plans/2026-10-10-workflow-speed.md @@ -0,0 +1,1076 @@ +# Workflow Speed 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:** Cut wall time per job without touching model quality: stop reloading warm models when nothing about their weights changed, stop re-building per-frame helpers, decode video with all cores, and fix catalog defaults that multiply step counts for nothing. + +**Architecture:** Three engine changes and one catalog change. (1) The pipeline cache key becomes a *weights identity*: LoRA scale/alpha and scheduler shift leave the hash, and a cache hit re-applies them with the in-place setters the engine already has. (2) A workflow switch in the worker evicts only pipelines the next workflow will not load, instead of everything. (3) Two task-level sinks (face-restore helper rebuilt per frame, single-threaded PyAV decode) are fixed in place. (4) Stale catalog defaults are corrected and pinned by a test. + +**Tech Stack:** Python 3, pytest (`venv/bin/python -m pytest`, xdist `-n auto` in pytest.ini; use `-n0` when naming files), ruff, PyAV, diffusers, peft, facexlib. + +**Spec:** The findings write-up in this session (speed audit of 2026-10-10). The ranked list: (11) LoRA scale/shift reloads, (7) workflow switch unloads all, (12) `release_pipeline` on a last step, (1) stale step counts, (6) stale first_block cache on turbo ref2va, (9) per-frame task sinks, (8/9) PyAV decode threads. Items needing a GPU measurement are listed under *Out of scope* at the end and are not tasks. + +## Global Constraints + +- Branch: `feat/workflow-speed` off `develop`. Never commit to `develop` or `master`. +- Gate for every task: `venv/bin/python -m pytest -n0 -q` green, then `venv/bin/ruff format --check dw dw_mcp tests scripts && venv/bin/ruff check dw dw_mcp tests scripts` clean. Run the full suite (`venv/bin/python -m pytest tests/ -q`) once at the end of the branch. +- No `eval()`, `exec()`, `shell=True`. No new `av.open` outside `dw/media.py` (`tests/test_media_layering.py` enforces it). +- Model knowledge lives in the catalog and `plugins/dw/`, never in engine code. +- Every `av.open` that decodes picture frames sets `stream.thread_type = "AUTO"`; none that only reads headers does. +- Docs that describe a changed behaviour change in the same task (`docs/WORKFLOW_GUIDE.md`, `docs/ACCELERATION.md`, `docs/WORKER_GUIDE.md`). +- `docs/ARCHITECTURE.md` has a module count ratchet (`tests/test_arch_metrics.py`). No task adds a module, so the count must not change. + +## Review Focus + +1. A LoRA whose `model_name` changes between runs must still reload (only scale and alpha are runtime-mutable). Test in Task 4. +2. A LoRA nulled by a variable (`model_name: null`) keeps its *index* for adapter naming, so re-applying scales on a hit must skip nulled entries while still naming the survivors by their original index. Test in Task 4. +3. A cache hit for a pipeline with no `loras` and no `scheduler.shift` must apply nothing and call nothing on the model. Test in Task 4. +4. A workflow switch whose new definition fails to prepare (bad variable, missing asset) must fall back to the full cleanup, never leave the old models resident silently. Test in Task 5. +5. Face restore on a video must produce the same per-frame result with a cached helper as with a fresh one, so the helper's per-image state must be cleared before each frame. Test in Task 2. + +--- + +### Task 1: Catalog defaults that multiply cost + +**Files:** +- Modify: `workflows/templates/compose-workflows.json:10` +- Modify: `workflows/templates/prompt-weighting.json:24-25` +- Modify: `workflows/templates/minimax/shots-batch.json:100` +- Modify: `workflows/templates/minimax/reference-to-video.json`, `chain-matched-to-audio.json`, `chain-video-continuity.json`, `composable-references.json`, `voice-timbre-reference.json`, `generated-subject-reference.json` (remove the `cache` block from the ModularPipeline step's `configuration`) +- Create: `tests/test_catalog_speed_defaults.py` + +**Interfaces:** +- Consumes: nothing. +- Produces: a test that pins the three rules below for every template under `workflows/templates/`. + +Rules the test pins: +1. A minimax template whose `lora_weight_name` default contains `turbo` and whose `num_inference_steps` default is at most 9 carries no `cache` block (RECIPES_24GB.md: first_block never skips on the turbo schedules). +2. `compose-workflows.json`'s `video_num_inference_steps` default equals the `num_inference_steps` default of `minimax/image-to-video.json` (the turbo LoRA's passes plus one). +3. Any step whose `model_name` contains `FLUX.1-schnell` runs at most 4 steps and guidance 0 (schnell is 4-step distilled; its guidance input is unused). + +- [ ] **Step 1: Write the failing test** + +```python +"""Catalog defaults that silently multiply a run's cost (speed audit 2026-10-10).""" + +import glob +import json +import os + +import pytest + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +TEMPLATES = sorted( + glob.glob(os.path.join(REPO_ROOT, "workflows", "templates", "**", "*.json"), recursive=True) +) + + +def _load(path): + with open(path, encoding="utf-8") as handle: + return json.load(handle) + + +def _default(definition, name): + variable = definition.get("variables", {}).get(name) + return variable.get("default") if isinstance(variable, dict) else variable + + +def _pipeline_steps(definition): + for step in definition.get("steps", []): + if "pipeline" in step: + yield step + + +MINIMAX = [p for p in TEMPLATES if os.sep + "minimax" + os.sep in p] + + +@pytest.mark.parametrize("path", MINIMAX, ids=os.path.basename) +def test_turbo_schedules_carry_no_block_cache(path): + definition = _load(path) + weight = _default(definition, "lora_weight_name") or "" + steps = _default(definition, "num_inference_steps") + if "turbo" not in weight or steps is None or steps > 9: + pytest.skip("not a turbo schedule") + for step in _pipeline_steps(definition): + configuration = step["pipeline"].get("configuration", {}) + assert "cache" not in configuration, ( + f"{os.path.basename(path)} step '{step['name']}' carries a cache block on a " + f"{steps}-step turbo schedule; RECIPES_24GB.md says it never skips there" + ) + + +def test_compose_workflows_video_steps_match_the_turbo_default(): + compose = _load(os.path.join(REPO_ROOT, "workflows", "templates", "compose-workflows.json")) + i2v = _load( + os.path.join(REPO_ROOT, "workflows", "templates", "minimax", "image-to-video.json") + ) + assert _default(compose, "video_num_inference_steps") == _default( + i2v, "num_inference_steps" + ) + + +@pytest.mark.parametrize("path", TEMPLATES, ids=lambda p: os.path.relpath(p, REPO_ROOT)) +def test_schnell_runs_its_distilled_schedule(path): + definition = _load(path) + for step in _pipeline_steps(definition): + pipeline = step["pipeline"] + model = pipeline.get("from_pretrained_arguments", {}).get("model_name", "") + if "FLUX.1-schnell" not in str(model): + continue + arguments = pipeline.get("arguments", {}) + assert arguments.get("num_inference_steps", 4) <= 4, path + assert arguments.get("guidance_scale", 0.0) == 0.0, path +``` + +- [ ] **Step 2: Run it to verify it fails** + +Run: `venv/bin/python -m pytest tests/test_catalog_speed_defaults.py -n0 -q` +Expected: failures for the six ref2va templates (cache block present), compose-workflows (20 != 5), prompt-weighting (20 steps, guidance 4.0). + +- [ ] **Step 3: Fix the templates** + +- `compose-workflows.json:10`: `"video_num_inference_steps": 20` -> `5`. +- `prompt-weighting.json:24-25`: `"num_inference_steps": 20` -> `4`, `"guidance_scale": 4.0` -> `0.0`. +- In each of the six ref2va templates, delete the `"cache": { "type": "first_block", "threshold": 0.1 }` block from the ModularPipeline step's `configuration` (watch the trailing comma on the previous key). Do not touch `image-to-video.json`; its description already explains the absence. +- `shots-batch.json:100`: delete `"release_pipeline": true,`. It is the only step, so the release frees nothing within the run and only forces the next job to reload H3; Task 5 makes the worker drop it when a different workflow needs the card. +- If a template's `description` string mentions the cache or the release you removed, edit that sentence too (grep `first_block` and `release` in the changed files). + +- [ ] **Step 4: Run the new test and the catalog tests** + +Run: `venv/bin/python -m pytest tests/test_catalog_speed_defaults.py tests/test_catalog_structure.py tests/test_catalog_shape.py tests/test_h3_schedule.py tests/test_plugin_skills.py -n0 -q` +Expected: PASS. If `test_catalog_structure` complains about a description that no longer matches the derived shape or a `cost` entry, fix the description; do not change the `cost` minutes (they are measured values, now upper bounds). + +- [ ] **Step 5: Validate the edited JSON against the schema** + +Run: `for f in workflows/templates/compose-workflows.json workflows/templates/prompt-weighting.json workflows/templates/minimax/shots-batch.json workflows/templates/minimax/reference-to-video.json workflows/templates/minimax/chain-matched-to-audio.json workflows/templates/minimax/chain-video-continuity.json workflows/templates/minimax/composable-references.json workflows/templates/minimax/voice-timbre-reference.json workflows/templates/minimax/generated-subject-reference.json; do venv/bin/python -m dw.validate "$f" || echo "FAILED $f"; done` +Expected: no `FAILED` line. + +- [ ] **Step 6: Ruff and commit** + +```bash +venv/bin/ruff format --check dw dw_mcp tests scripts && venv/bin/ruff check dw dw_mcp tests scripts +git add workflows/templates tests/test_catalog_speed_defaults.py +git commit -m "fix(catalog): drop stale step counts, dead turbo caches and a no-op release" +``` + +--- + +### Task 2: Face restore keeps its detector between frames + +**Files:** +- Modify: `dw/tasks/restore_faces.py:70-80` +- Test: `tests/test_restore_faces_helper_cache.py` (create) + +**Interfaces:** +- Consumes: `dw.tasks.model_cache.cached_model(key, factory)`. +- Produces: nothing new; `restore_faces(image, model_name, device, **kwargs)` keeps its signature. + +Background: `FaceRestoreHelper(...)` loads RetinaFace (and ParseNet when `use_parse`) in its constructor. `per_frame` calls `restore_faces` once per frame, so a video reloads detector weights per frame. The helper keeps per-image state (`all_landmarks_5`, `cropped_faces`, `restored_faces`, affine matrices, `det_faces`, `pad_input_imgs`) that `clean_all()` resets. + +- [ ] **Step 1: Write the failing test** + +```python +"""restore_faces builds its facexlib helper once per (settings, device), not per frame.""" + +import sys +import types +from unittest.mock import MagicMock + +import numpy as np +import pytest +from PIL import Image + +from dw.tasks import model_cache + + +@pytest.fixture +def fake_facexlib(monkeypatch): + """A facexlib whose helper records constructions and detects no faces.""" + constructed = [] + + class FakeHelper: + def __init__(self, **kwargs): + constructed.append(kwargs) + self.cleaned = 0 + + def clean_all(self): + self.cleaned += 1 + + def read_image(self, image): + self.image = image + + def get_face_landmarks_5(self, **kwargs): + return 0 + + module = types.ModuleType("facexlib.utils.face_restoration_helper") + module.FaceRestoreHelper = FakeHelper + package = types.ModuleType("facexlib") + utils = types.ModuleType("facexlib.utils") + monkeypatch.setitem(sys.modules, "facexlib", package) + monkeypatch.setitem(sys.modules, "facexlib.utils", utils) + monkeypatch.setitem(sys.modules, "facexlib.utils.face_restoration_helper", module) + model_cache.clear_model_cache() + yield constructed + model_cache.clear_model_cache() + + +def test_two_frames_share_one_helper(fake_facexlib, monkeypatch): + from dw.tasks import restore_faces as module + + descriptor = MagicMock(supports_half=False) + monkeypatch.setattr(module, "_load_face_model", lambda path, device: descriptor) + monkeypatch.setattr( + "dw.tasks.upscale.resolve_model_path", lambda model_name, filename: "weights.pth" + ) + frame = Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)) + + module.restore_faces(frame, "m", device="cpu") + module.restore_faces(frame, "m", device="cpu") + + assert len(fake_facexlib) == 1 + helper_kwargs = fake_facexlib[0] + assert helper_kwargs["det_model"] == "retinaface_resnet50" + + +def test_helper_state_is_cleared_before_each_frame(fake_facexlib, monkeypatch): + from dw.tasks import restore_faces as module + + descriptor = MagicMock(supports_half=False) + monkeypatch.setattr(module, "_load_face_model", lambda path, device: descriptor) + monkeypatch.setattr( + "dw.tasks.upscale.resolve_model_path", lambda model_name, filename: "weights.pth" + ) + frame = Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)) + module.restore_faces(frame, "m", device="cpu") + module.restore_faces(frame, "m", device="cpu") + + helper = model_cache._cache[ + ("restore_faces_helper", 1, 512, True, "cpu") + ] + assert helper.cleaned == 2 + + +def test_different_settings_build_different_helpers(fake_facexlib, monkeypatch): + from dw.tasks import restore_faces as module + + descriptor = MagicMock(supports_half=False) + monkeypatch.setattr(module, "_load_face_model", lambda path, device: descriptor) + monkeypatch.setattr( + "dw.tasks.upscale.resolve_model_path", lambda model_name, filename: "weights.pth" + ) + frame = Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)) + module.restore_faces(frame, "m", device="cpu", face_size=512) + module.restore_faces(frame, "m", device="cpu", face_size=256) + assert len(fake_facexlib) == 2 +``` + +Read `dw/tasks/restore_faces.py` first: if `_load_face_model` or `resolve_model_path` is imported or named differently from what the test patches, adjust the test's patch targets to the real names, not the implementation. + +- [ ] **Step 2: Run it to verify it fails** + +Run: `venv/bin/python -m pytest tests/test_restore_faces_helper_cache.py -n0 -q` +Expected: `test_two_frames_share_one_helper` fails with 2 constructions; the cache-key test fails with KeyError. + +- [ ] **Step 3: Cache the helper** + +Replace the helper construction at `dw/tasks/restore_faces.py:70-77` with: + +```python + # facexlib's helper loads RetinaFace (and ParseNet with use_parse) in + # its constructor. per_frame calls this once per frame, so a video was + # reloading detector weights for every frame. The helper keeps per-image + # state between calls; clean_all() resets it before each read + def build_helper(): + return FaceRestoreHelper( + upscale_factor=upscale_factor, + face_size=face_size, + crop_ratio=(1, 1), + det_model="retinaface_resnet50", + use_parse=use_parse, + device=torch.device(device), + ) + + face_helper = cached_model( + ("restore_faces_helper", upscale_factor, face_size, use_parse, str(device)), + build_helper, + ) + face_helper.clean_all() +``` + +Keep the existing `read_image` call that follows. + +- [ ] **Step 4: Run the tests** + +Run: `venv/bin/python -m pytest tests/test_restore_faces_helper_cache.py tests/test_restore_faces_video.py tests/test_face_track.py -n0 -q` +Expected: PASS. + +- [ ] **Step 5: Note it in docs/TASKS.md** + +Find the `restore_faces` entry in `docs/TASKS.md` and add one sentence: "The detector and parser are loaded once per settings/device and reused across frames; `release_models` drops them with the restorer." + +- [ ] **Step 6: Ruff and commit** + +```bash +venv/bin/ruff format --check dw dw_mcp tests scripts && venv/bin/ruff check dw dw_mcp tests scripts +git add dw/tasks/restore_faces.py tests/test_restore_faces_helper_cache.py docs/TASKS.md +git commit -m "perf(restore_faces): build the facexlib helper once, not per frame" +``` + +--- + +### Task 3: Frame decode uses every core + +**Files:** +- Modify: `dw/media.py` — the picture-decoding opens: `video_shape` (~341, only its counting branch decodes), `read_frames_at` (~385), `read_frame_range` (~481), the frame counter (~515), `decode_audio_video` (~533), `decode_rgb_frames` (~561), and `read_thumbnails_and_track` (~577) if it decodes pictures. +- Test: `tests/test_media_decode_threads.py` (create) + +**Interfaces:** +- Produces: a module-level helper in `dw/media.py`: + +```python +def _threaded(stream): + """A video stream set to decode on every core. libav defaults a + software h264 decode to one thread; AUTO lets it pick frame and + slice threading for the codec. Header-only reads do not call this.""" + stream.thread_type = "AUTO" + return stream +``` + +- [ ] **Step 1: Write the failing test** + +```python +"""Every picture-decoding open in dw/media.py decodes on every core.""" + +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from dw import media + + +class FakeFrame: + def __init__(self): + self.pts = 0 + + def to_ndarray(self, format="rgb24"): + return np.zeros((2, 2, 3), dtype=np.uint8) + + +def _container(frames=1): + container = MagicMock() + stream = MagicMock() + stream.average_rate = 8 + stream.frames = 0 + stream.start_time = 0 + stream.time_base = 1 + del stream.thread_type # set by the code under test, asserted below + container.streams.video = [stream] + container.streams.audio = [] + container.decode.return_value = [FakeFrame() for _ in range(frames)] + container.__enter__.return_value = container + container.__exit__.return_value = False + return container, stream + + +@pytest.mark.parametrize( + "call", + [ + lambda: media.decode_rgb_frames("x.mp4"), + lambda: media.decode_audio_video("x.mp4"), + lambda: media.video_shape("x.mp4"), + ], +) +def test_decoders_set_auto_threading(call): + container, stream = _container() + with patch.object(media.av, "open", return_value=container): + call() + assert stream.thread_type == "AUTO" +``` + +Read each target function before running: `decode_audio_video` builds an `AudioResampler` only with an audio stream, and `video_shape` only decodes when `stream.frames` is falsy (the fake sets 0). If a function needs more fake attributes (`frame.time_base`, `stream.codec_context`), add them to the fake rather than changing the function. + +- [ ] **Step 2: Run it to verify it fails** + +Run: `venv/bin/python -m pytest tests/test_media_decode_threads.py -n0 -q` +Expected: FAIL, `thread_type` not set (AttributeError or MagicMock inequality). + +- [ ] **Step 3: Add `_threaded` and apply it** + +Add `_threaded` near the top of `dw/media.py` (after the imports). In each picture-decoding function, wrap the stream once where it is taken: `stream = _threaded(container.streams.video[0])`, and in `decode_rgb_frames` change `container.decode(video=0)` to `container.decode(_threaded(container.streams.video[0]))`. In `video_shape`, call `_threaded(stream)` only inside the `if count is None:` branch. Leave `probe_media`, `_stream_info`, the duration helpers and the audio-only opens alone. + +- [ ] **Step 4: Run the media tests** + +Run: `venv/bin/python -m pytest tests/test_media_decode_threads.py tests/test_media_layering.py tests/test_video_utils.py tests/test_concat_videos.py tests/test_dissolve_videos.py -n0 -q` +Expected: PASS. (If a named file does not exist, drop it from the list; `ls tests | grep media` for the real names.) + +- [ ] **Step 5: Ruff and commit** + +```bash +venv/bin/ruff format --check dw dw_mcp tests scripts && venv/bin/ruff check dw dw_mcp tests scripts +git add dw/media.py tests/test_media_decode_threads.py +git commit -m "perf(media): decode picture frames with libav AUTO threading" +``` + +--- + +### Task 4: LoRA scale, alpha and scheduler shift no longer reload the pipeline + +**Files:** +- Modify: `dw/step_cache.py:124-144` (`pipeline_cache_key`) +- Modify: `dw/pipeline_processors/adapters.py:104-155` (`load_loras`) +- Modify: `dw/pipeline_processors/components.py:~420-470` (`load_and_configure_scheduler`) +- Modify: `dw/pipeline_processors/pipeline.py` (`Pipeline`: new method `apply_runtime_settings`) +- Modify: `dw/pipeline_ownership.py:252-293` (`wrap_resident`) +- Modify: `docs/WORKFLOW_GUIDE.md` (the pipeline reuse / `release_pipeline` area near line 1319) and `docs/LORAS.md` +- Test: `tests/test_runtime_settings.py` (create); extend `tests/test_step_cache.py` + +**Interfaces:** +- Produces, in `dw/step_cache.py`: + +```python +RUNTIME_LORA_KEYS = ("scale", "alpha") + +def weights_identity(pipeline_definition): + """The pipeline definition with every runtime-mutable setting removed: + a LoRA's scale and alpha (set_adapters / set_adapter_alpha) and a + scheduler's shift (set_shift) change nothing about what is loaded.""" +``` + +`pipeline_cache_key` hashes `weights_identity(...)` instead of the raw filtered dict. + +- Produces, in `dw/pipeline_processors/adapters.py`: + +```python +def adapter_settings(loras): + """(adapter_names, adapter_weights, alphas) for the active entries of + `loras`, naming an adapter by its own adapter_name or its ORIGINAL index + (a nulled entry keeps its slot). Reads, never pops.""" + +def apply_adapter_settings(loras, pipeline): + """Set each active LoRA's alpha then every scale on a pipeline that + already holds them. No-op with no active entry.""" +``` + +`load_loras` keeps its signature and its pop-then-load behaviour but computes names/weights/alphas through `adapter_settings` *before* popping, so the two paths cannot disagree. + +- Produces, in `dw/pipeline_processors/components.py`: + +```python +def apply_scheduler_shift(scheduler_definition, pipeline, component_name="scheduler"): + """Set the shift a scheduler definition names on an already-loaded + scheduler; nothing when the definition is None or names no shift.""" +``` + +`load_and_configure_scheduler` calls it after loading the scheduler type. + +- Produces, in `dw/pipeline_processors/pipeline.py`: + +```python + def apply_runtime_settings(self): + """Re-apply what a cache hit may have changed without changing the + weights: scheduler and audio_scheduler shift, LoRA scales and + alphas. Called by wrap_resident; a fresh load applies the same + values on its way through load().""" +``` + +- [ ] **Step 1: Write the failing key tests** + +Append to `tests/test_step_cache.py`: + +```python +from dw.step_cache import pipeline_cache_key, weights_identity + + +def _definition(**overrides): + base = { + "configuration": {"component_type": "FluxPipeline"}, + "from_pretrained_arguments": {"model_name": "black-forest-labs/FLUX.1-dev"}, + "loras": [ + {"model_name": "a/b", "weight_name": "x.safetensors", "scale": 1.0, "alpha": 16}, + ], + "scheduler": {"shift": 6.0}, + "audio_scheduler": {"shift": 3.0}, + "arguments": {"prompt": "p"}, + } + base.update(overrides) + return base + + +def test_lora_scale_and_alpha_do_not_change_the_pipeline_key(): + a = _definition() + b = _definition(loras=[{"model_name": "a/b", "weight_name": "x.safetensors", "scale": 0.5, "alpha": 8}]) + assert pipeline_cache_key(a) == pipeline_cache_key(b) + + +def test_scheduler_shift_does_not_change_the_pipeline_key(): + a = _definition() + b = _definition(scheduler={"shift": 12.0}, audio_scheduler={"shift": 1.0}) + assert pipeline_cache_key(a) == pipeline_cache_key(b) + + +def test_lora_model_name_still_changes_the_pipeline_key(): + a = _definition() + b = _definition(loras=[{"model_name": "a/c", "weight_name": "x.safetensors", "scale": 1.0}]) + assert pipeline_cache_key(a) != pipeline_cache_key(b) + + +def test_scheduler_type_still_changes_the_pipeline_key(): + a = _definition(scheduler={"scheduler_type": "X", "shift": 6.0}) + b = _definition(scheduler={"scheduler_type": "Y", "shift": 6.0}) + assert pipeline_cache_key(a) != pipeline_cache_key(b) + + +def test_weights_identity_leaves_the_definition_untouched(): + definition = _definition() + weights_identity(definition) + assert definition["loras"][0]["scale"] == 1.0 + assert definition["scheduler"]["shift"] == 6.0 +``` + +- [ ] **Step 2: Run them to verify they fail** + +Run: `venv/bin/python -m pytest tests/test_step_cache.py -n0 -q` +Expected: ImportError on `weights_identity`. + +- [ ] **Step 3: Implement `weights_identity` and use it in the key** + +In `dw/step_cache.py`, above `pipeline_cache_key`: + +```python +RUNTIME_LORA_KEYS = ("scale", "alpha") + + +def weights_identity(pipeline_definition): + """The pipeline definition with every runtime-mutable setting removed. + + A LoRA's scale and alpha (set_adapters / set_adapter_alpha) and a + scheduler's shift (set_shift) change nothing about what is loaded, so a + run that only moves one of them must hit the warm pipeline and re-apply + the value (Pipeline.apply_runtime_settings), not reload the stack. + Everything else - the model, the LoRA files, quantization, placement, + the scheduler type - still changes the key. Returns a new dict; the + definition is the workflow's and is not edited. + """ + identity = { + k: v + for k, v in pipeline_definition.items() + if k not in ("arguments", "seed", "chain") + } + loras = identity.get("loras") + if isinstance(loras, list): + identity["loras"] = [ + {k: v for k, v in lora.items() if k not in RUNTIME_LORA_KEYS} + if isinstance(lora, dict) + else lora + for lora in loras + ] + for name in ("scheduler", "audio_scheduler"): + scheduler = identity.get(name) + if isinstance(scheduler, dict): + identity[name] = {k: v for k, v in scheduler.items() if k != "shift"} + return identity +``` + +and make `pipeline_cache_key` serialize `weights_identity(pipeline_definition)`; update its docstring to say scale, alpha and shift are excluded and why. + +- [ ] **Step 4: Run the key tests** + +Run: `venv/bin/python -m pytest tests/test_step_cache.py -n0 -q` +Expected: PASS. + +- [ ] **Step 5: Write the failing setter tests** + +Create `tests/test_runtime_settings.py`: + +```python +"""A cache hit re-applies LoRA scale/alpha and scheduler shift in place.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from dw.pipeline_processors.adapters import adapter_settings, apply_adapter_settings +from dw.pipeline_processors.components import apply_scheduler_shift + + +def test_adapter_settings_name_by_original_index_and_skip_nulled(): + loras = [ + {"model_name": None, "scale": 0.3}, + {"model_name": "a/b", "scale": 0.5, "alpha": 8}, + {"model_name": "c/d", "adapter_name": "turbo"}, + ] + names, weights, alphas = adapter_settings(loras) + assert names == ["1", "turbo"] + assert weights == [0.5, 1.0] + assert alphas == {"1": 8} + assert loras[1]["scale"] == 0.5 # read, not popped + + +def test_apply_adapter_settings_sets_alpha_then_scales(): + pipeline = MagicMock() + calls = [] + with patch( + "dw.pipeline_processors.adapters.set_adapter_alpha", + side_effect=lambda p, n, a: calls.append(("alpha", n, a)), + ): + pipeline.set_adapters.side_effect = lambda n, w: calls.append(("scales", n, w)) + apply_adapter_settings( + [{"model_name": "a/b", "scale": 0.5, "alpha": 8}], pipeline + ) + assert calls == [("alpha", "0", 8), ("scales", ["0"], [0.5])] + + +def test_apply_adapter_settings_is_a_no_op_without_active_loras(): + pipeline = MagicMock() + apply_adapter_settings([{"model_name": None}], pipeline) + apply_adapter_settings(None, pipeline) + pipeline.set_adapters.assert_not_called() + + +def test_apply_scheduler_shift_sets_the_shift(): + pipeline = MagicMock() + apply_scheduler_shift({"shift": 12}, pipeline, "scheduler") + pipeline.scheduler.set_shift.assert_called_once_with(12.0) + + +def test_apply_scheduler_shift_without_a_shift_touches_nothing(): + pipeline = MagicMock() + apply_scheduler_shift(None, pipeline, "scheduler") + apply_scheduler_shift({"scheduler_type": "X"}, pipeline, "scheduler") + pipeline.scheduler.set_shift.assert_not_called() + + +def test_wrap_resident_reapplies_runtime_settings(): + from dw.pipeline_ownership import wrap_resident + + model = MagicMock() + cached = MagicMock() + cached.pipeline = model + step_definition = { + "name": "s", + "pipeline": { + "configuration": {"component_type": "X", "no_generator": True}, + "loras": [{"model_name": "a/b", "scale": 0.25}], + "scheduler": {"shift": 9}, + }, + } + with patch("dw.pipeline_ownership.emit_phase"): + wrapper = wrap_resident(cached, step_definition, 1, "cpu", "/out", "p") + assert wrapper.pipeline is model + model.scheduler.set_shift.assert_called_once_with(9.0) + model.set_adapters.assert_called_once_with(["0"], [0.25]) + + +def test_wrap_resident_with_nothing_mutable_calls_nothing(): + from dw.pipeline_ownership import wrap_resident + + model = MagicMock() + cached = MagicMock() + cached.pipeline = model + step_definition = { + "name": "s", + "pipeline": {"configuration": {"component_type": "X", "no_generator": True}}, + } + with patch("dw.pipeline_ownership.emit_phase"): + wrap_resident(cached, step_definition, 1, "cpu", "/out", "p") + model.set_adapters.assert_not_called() + model.scheduler.set_shift.assert_not_called() +``` + +Check how `get_component(pipeline, "scheduler")` resolves on a MagicMock (it may use `getattr` or a components dict); if it does not return `pipeline.scheduler`, give the mock what it reads rather than changing the helper. `wrap_resident` imports `emit_phase`; if it is imported under another name, patch that. + +- [ ] **Step 6: Run them to verify they fail** + +Run: `venv/bin/python -m pytest tests/test_runtime_settings.py -n0 -q` +Expected: ImportError on `adapter_settings`. + +- [ ] **Step 7: Implement the setters** + +In `dw/pipeline_processors/adapters.py`: + +```python +def adapter_settings(loras): + """(adapter_names, adapter_weights, alphas) for the active entries of + `loras`. An adapter is named by its own adapter_name or its ORIGINAL + index - a nulled entry keeps its slot, so the name a hit re-applies to + is the name the load gave. Reads the entries, never pops them.""" + names, weights, alphas = [], [], {} + for i, lora in enumerate(loras or []): + if not isinstance(lora, dict) or lora.get("model_name") is None: + continue + name = lora.get("adapter_name") or str(i) + names.append(name) + scale = lora.get("scale") + weights.append(1.0 if scale is None else float(scale)) + if lora.get("alpha") is not None: + alphas[name] = lora["alpha"] + return names, weights, alphas + + +def apply_adapter_settings(loras, pipeline): + """Set each active LoRA's alpha, then every scale, on a pipeline that + already holds the adapters. Alpha first: set_adapters is what recomputes + each layer's scaling from it. No-op with no active entry.""" + names, weights, alphas = adapter_settings(loras) + if not names: + return + for name, alpha in alphas.items(): + set_adapter_alpha(pipeline, name, alpha) + logger.info(f"Setting adapter weights: {list(zip(names, weights))}") + pipeline.set_adapters(names, weights) +``` + +Rewrite `load_loras` so the loop takes `names, weights, alphas = adapter_settings(loras)` once, pops `model_name`, `adapter_name`, `scale`, `alpha` from each active entry as it does now (they must not reach `load_lora_weights`), loads with the name from `adapter_settings`, then finishes with the same alpha-then-`set_adapters` sequence (call `apply_adapter_settings` on a snapshot taken before the pops, or keep the names/weights/alphas locals - either is fine, but the names must be the ones `adapter_settings` computed). + +In `dw/pipeline_processors/components.py`, move the tail of `load_and_configure_scheduler` (from `shift = scheduler_definition.get("shift", None)` to the `set_shift` call) into `apply_scheduler_shift(scheduler_definition, pipeline, component_name="scheduler")`, which returns at once when `scheduler_definition` is None or has no `shift`, and call it from `load_and_configure_scheduler` after the type is loaded. Keep both `ValueError`s. + +In `dw/pipeline_processors/pipeline.py`, on `Pipeline`: + +```python + def apply_runtime_settings(self): + """Re-apply what a run may have changed without changing the weights: + scheduler and audio_scheduler shift, LoRA scales and alphas. The + pipeline cache keys on weights_identity, so a hit can carry any of + these at a new value; load() applies the same values on its way + through, so a cold load and a hit end in the same state.""" + apply_scheduler_shift(self.pipeline_definition.get("scheduler"), self.pipeline) + apply_scheduler_shift( + self.pipeline_definition.get("audio_scheduler"), + self.pipeline, + "audio_scheduler", + ) + apply_adapter_settings(self.pipeline_definition.get("loras"), self.pipeline) +``` + +Import `apply_scheduler_shift` from `.components` and `apply_adapter_settings` from `.adapters` beside the existing imports. + +In `dw/pipeline_ownership.py` `wrap_resident`, after the generator block and before `emit_phase("cached", ...)`: + +```python + # The key the hit matched on is the weights identity: a scale, alpha or + # shift may differ from the values the resident model last ran with + new_pipeline_wrapper.apply_runtime_settings() +``` + +- [ ] **Step 8: Run the setter tests and the neighbours** + +Run: `venv/bin/python -m pytest tests/test_runtime_settings.py tests/test_step_cache.py tests/test_lora_alpha.py tests/test_lora_disable.py tests/test_h3_adapters.py tests/test_h3_schedule.py tests/test_pipeline_ownership.py tests/test_workflow_run.py -n0 -q` +Expected: PASS. (Drop a name that does not exist; `ls tests | grep -E "ownership|workflow_run|lora"`.) If a test asserted that a scale change produces a *different* key or a reload, read it: it was pinning the behaviour this task replaces, and its assertion flips. + +- [ ] **Step 9: Docs** + +`docs/WORKFLOW_GUIDE.md`, next to the `release_pipeline` description (~line 1324), add a short paragraph "What keeps a pipeline warm": the server's worker keeps a pipeline loaded across runs as long as what it *loads* is unchanged; a LoRA's `scale` and `alpha` and a scheduler's `shift` are applied in place on the warm pipeline, so iterating on them costs no reload; changing a LoRA's `model_name` or `weight_name`, the quantization, placement or scheduler type reloads. `docs/LORAS.md`: one sentence at the scale/alpha description saying the same. `docs/ARCHITECTURE.md`: if its map has a row for the pipeline cache key (grep `pipeline_cache_key`), extend the rule text with "keys on the weights identity (scale, alpha, shift excluded)". + +- [ ] **Step 10: Ruff and commit** + +```bash +venv/bin/ruff format --check dw dw_mcp tests scripts && venv/bin/ruff check dw dw_mcp tests scripts +git add dw/step_cache.py dw/pipeline_processors/adapters.py dw/pipeline_processors/components.py dw/pipeline_processors/pipeline.py dw/pipeline_ownership.py tests/test_runtime_settings.py tests/test_step_cache.py docs/WORKFLOW_GUIDE.md docs/LORAS.md docs/ARCHITECTURE.md +git commit -m "perf(cache): LoRA scale/alpha and scheduler shift re-apply on a warm pipeline instead of reloading" +``` + +--- + +### Task 5: A workflow switch keeps the pipelines the next workflow loads + +**Files:** +- Modify: `dw/workflow_run.py` (new `pipeline_keys(workflow, arguments)` beside `cache_hits`, ~line 445) +- Modify: `dw/worker.py:349-367` (`_prepare_workflow`) and `:638-709` (`_cleanup_all` gains a `keep` parameter) +- Modify: `docs/WORKER_GUIDE.md` and `docs/SERVER.md:~224` (the "Workflow changed - releasing cached models" behaviour) +- Test: `tests/test_worker_execute.py` (extend) and `tests/test_workflow_run_keys.py` (create) + +**Interfaces:** +- Produces, in `dw/workflow_run.py`: + +```python +def pipeline_keys(workflow, arguments): + """The set of pipeline cache keys a run of `workflow` with `arguments` + would load its steps under - the same table run() and cache_hits() + take (step_pipeline_keys over the prepared definition). Prepares the + definition exactly as a run does and executes nothing. Sub-workflow + steps load their own pipelines later and are not in the set.""" +``` + +- Produces, in `dw/worker.py`: `_cleanup_all(self, keep=frozenset())` - every pipeline whose key is in `keep` survives; the rest, the task model cache and the step cache go as today. With `keep` empty the behaviour is exactly today's. Returns the same summary string, now naming how many pipelines were kept. + +- `_prepare_workflow` on an identity change: compute `keep = pipeline_keys(new_workflow, command.get("arguments") or {}) & set(self.loaded_pipelines)`; on any exception, log at debug and use an empty set. If `keep` is non-empty, reply `Output(message="Workflow changed - keeping N warm model(s) the new workflow loads, releasing the rest...")`, else the existing message; then `self._reply(Output(message=self._cleanup_all(keep=keep)))`. + +- [ ] **Step 1: Write the failing `pipeline_keys` test** + +Create `tests/test_workflow_run_keys.py`, reusing `build_test_workflow_and_call_count_spy` from `tests/test_workflow_step_cache.py` (see its `TestCacheHits` class at ~line 1003 for the try/finally that stops its patchers): + +```python +"""pipeline_keys() answers the worker's question on a workflow switch: which +resident pipelines does the next workflow load anyway.""" + +import copy + +from dw import workflow_run as workflow_run_module +from dw.step_cache import step_cache, step_pipeline_keys +from tests.test_workflow_step_cache import build_test_workflow_and_call_count_spy + + +def test_pipeline_keys_match_what_a_run_keys_on(tmp_path): + step_cache.clear() + workflow, _ = build_test_workflow_and_call_count_spy(str(tmp_path)) + try: + expected_def, _, _ = workflow_run_module.prepare_definition( + workflow, + copy.deepcopy(workflow.workflow_definition), + {}, + workflow_run_module.run_base_dir(workflow), + ) + expected = set(step_pipeline_keys(expected_def["steps"]).values()) + assert expected + assert workflow_run_module.pipeline_keys(workflow, {}) == expected + finally: + for p in workflow._test_patcher: + p.stop() + + +def test_an_unseeded_workflow_still_answers(tmp_path): + step_cache.clear() + workflow, _ = build_test_workflow_and_call_count_spy(str(tmp_path)) + del workflow.workflow_definition["seed"] + try: + assert workflow_run_module.pipeline_keys(workflow, {}) + finally: + for p in workflow._test_patcher: + p.stop() + + +def test_the_probe_writes_nothing(tmp_path): + step_cache.clear() + workflow, call_count = build_test_workflow_and_call_count_spy(str(tmp_path)) + try: + workflow_run_module.pipeline_keys(workflow, {}) + assert list(tmp_path.iterdir()) == [] + assert call_count() == 0 + finally: + for p in workflow._test_patcher: + p.stop() +``` + +If `tests/test_workflow_step_cache.py` cannot be imported as `tests.test_workflow_step_cache` (check `tests/__init__.py` exists - it does), import it the way the other tests do. + +Unlike `cache_hits`, `pipeline_keys` must not return `[]` for an unseeded workflow: a warm model is worth keeping whether or not step results are cacheable. + +- [ ] **Step 2: Run it to verify it fails** + +Run: `venv/bin/python -m pytest tests/test_workflow_run_keys.py -n0 -q` +Expected: ImportError on `pipeline_keys`. + +- [ ] **Step 3: Implement `pipeline_keys`** + +In `dw/workflow_run.py`, after `cache_hits`: + +```python +def pipeline_keys(workflow, arguments): + """The set of pipeline cache keys a run of `workflow` with `arguments` + loads its steps under - the table run() and cache_hits() take + (step_pipeline_keys over the prepared definition). Prepares the + definition exactly as a run does and executes nothing: the worker asks + this before releasing the previous workflow's models, so a model the + next workflow loads anyway stays warm. A sub-workflow step loads its + own pipelines later and is not in the set; an unseeded workflow still + answers, since warmth does not depend on the step cache. + """ + output_root_token = activate_output_root(workflow.output_dir) + try: + workflow_def = copy.deepcopy(workflow.workflow_definition) + base_dir = run_base_dir(workflow) + workflow_def, _, _ = prepare_definition( + workflow, workflow_def, arguments or {}, base_dir + ) + return set(step_pipeline_keys(workflow_def.get("steps", [])).values()) + finally: + deactivate_output_root(output_root_token) +``` + +`step_pipeline_keys` is already imported for `cache_hits`; if not, import it from `.step_cache`. + +- [ ] **Step 4: Run it** + +Run: `venv/bin/python -m pytest tests/test_workflow_run_keys.py -n0 -q` +Expected: PASS. + +- [ ] **Step 5: Write the failing worker tests** + +In `tests/test_worker_execute.py`, after `test_workflow_switch_evicts_cache_and_untouched_keys_dropped`: + +```python +def test_workflow_switch_keeps_pipelines_the_new_workflow_loads(): + worker = _make_worker() + _execute(worker, StubWorkflow()) + worker.loaded_pipelines["shared-key"] = object() + worker.loaded_pipelines["old-only-key"] = object() + + with patch("dw.worker.workflow_run.pipeline_keys", return_value={"shared-key"}): + messages = _execute( + worker, StubWorkflow(), command=snapshot_command("/w/other.json") + ) + + assert "shared-key" in worker.loaded_pipelines + assert "old-only-key" not in worker.loaded_pipelines + outputs = [m["message"] for m in messages if m["type"] == "output"] + assert any("keeping 1 warm" in m for m in outputs) + + +def test_workflow_switch_falls_back_to_a_full_release_when_keys_cannot_be_taken(): + worker = _make_worker() + _execute(worker, StubWorkflow()) + worker.loaded_pipelines["shared-key"] = object() + + with patch( + "dw.worker.workflow_run.pipeline_keys", side_effect=ValueError("bad variable") + ): + _execute(worker, StubWorkflow(), command=snapshot_command("/w/other.json")) + + assert worker.loaded_pipelines == {} + + +def test_full_cleanup_keep_spares_named_keys_only(): + worker = _make_worker() + worker.loaded_pipelines = {"a": object(), "b": object()} + worker._cleanup_all(keep={"a"}) + assert set(worker.loaded_pipelines) == {"a"} +``` + +Read `_execute`, `_make_worker` and `snapshot_command` at the top of the file first; `StubWorkflow` may need the `workflow_definition` attribute `pipeline_keys` is patched around (it is patched, so it does not). The existing `test_workflow_switch_evicts_cache_and_untouched_keys_dropped` patches `_cleanup_all` and asserts one call; it still holds (the call now carries `keep=`), so leave it or tighten it to `cleanup.assert_called_once_with(keep=set())`. + +- [ ] **Step 6: Run them to verify they fail** + +Run: `venv/bin/python -m pytest tests/test_worker_execute.py -n0 -q` +Expected: the three new tests fail (`keep` unexpected keyword; `shared-key` evicted). + +- [ ] **Step 7: Implement the worker change** + +`_cleanup_all(self, keep=frozenset())`: at `dw/worker.py:654-658` replace + +```python + self.loaded_pipelines.clear() + self.shared_components.clear() + self.prior_step_keys.clear() +``` + +with + +```python + kept = {k: v for k, v in self.loaded_pipelines.items() if k in keep} + dropped = len(self.loaded_pipelines) - len(kept) + # Rebuilt rather than cleared in place: the dict object is shared with + # nothing, and a kept entry must keep its identity + self.loaded_pipelines.clear() + self.loaded_pipelines.update(kept) + # Every run republishes shared components from the pipelines it hits + # (workflow.py: publish_shared_components on a cache hit) + self.shared_components.clear() + # A step's prior key addresses an entry that is now gone - unless it + # was kept, in which case a redefinition next run still releases it + self.prior_step_keys = { + name: key for name, key in self.prior_step_keys.items() if key in kept + } +``` + +Leave `clear_model_cache()` and `step_cache.clear()` as they are. Append `f", kept {len(kept)} warm"` to the summary when `kept` is non-empty, and use `dropped` in the "Full cleanup complete" log line. + +`_prepare_workflow`: + +```python + if identity != self.workflow_identity: + if self.workflow_identity is not None: + keep = self._keys_worth_keeping(command, job.workflow) + if keep: + self._reply( + Output( + message=f"Workflow changed - keeping {len(keep)} warm " + "model(s) the new workflow loads, releasing the rest..." + ) + ) + else: + self._reply( + Output(message="Workflow changed - releasing cached models...") + ) + self._reply(Output(message=self._cleanup_all(keep=keep))) + self.workflow_identity = identity +``` + +with + +```python + def _keys_worth_keeping(self, command, workflow): + """The resident pipeline keys the next workflow loads anyway. Any + failure to prepare the definition answers the empty set - the run + itself will report what is wrong, and a full release is the safe + default on a card that cannot hold two stacks.""" + try: + wanted = workflow_run.pipeline_keys( + workflow, command.get("arguments") or {} + ) + except Exception as e: + logger.debug(f"Could not take the next workflow's pipeline keys: {e}") + return set() + return wanted & set(self.loaded_pipelines) +``` + +Check whether `_prepare_workflow` runs with the job's asset dir already active (`_activate_job` is phase 1, so yes); `pipeline_keys` may realize `asset:` variables and needs it. + +- [ ] **Step 8: Run the worker tests** + +Run: `venv/bin/python -m pytest tests/test_worker_execute.py tests/test_worker_pool.py tests/test_workflow_run_keys.py -n0 -q` +Expected: PASS. + +- [ ] **Step 9: Docs** + +`docs/WORKER_GUIDE.md` and `docs/SERVER.md` (~line 224, "that worker's pipelines and step cache are warm"): state that switching workflows releases only the models the next workflow does not load; a model both load stays warm (same weights identity, see Task 4's paragraph in WORKFLOW_GUIDE). Mention that `clear_memory` is still the way to empty the card by hand. + +- [ ] **Step 10: Ruff and commit** + +```bash +venv/bin/ruff format --check dw dw_mcp tests scripts && venv/bin/ruff check dw dw_mcp tests scripts +git add dw/worker.py dw/workflow_run.py tests/test_worker_execute.py tests/test_workflow_run_keys.py docs/WORKER_GUIDE.md docs/SERVER.md +git commit -m "perf(worker): keep warm pipelines a switched workflow loads anyway" +``` + +--- + +### Task 6: Full-suite gate and release note + +**Files:** +- Modify: `docs/RELEASING.md` ("Next" section) + +- [ ] **Step 1: Full suite** + +Run: `venv/bin/python -m pytest tests/ -q` +Expected: all pass. The memory notes that `test_audio_utils`, `test_worker` lifecycle and `test_download_watch` can flake under xdist; re-run a failing one alone with `-n0` before treating it as a regression. + +- [ ] **Step 2: Ruff on the whole tree** + +Run: `venv/bin/ruff format --check dw dw_mcp tests scripts && venv/bin/ruff check dw dw_mcp tests scripts` + +- [ ] **Step 3: Release note** + +Under "Next" in `docs/RELEASING.md`, one bullet per task in user terms: warm pipelines survive a LoRA scale/alpha or shift change and a workflow switch; face restore and video decode faster on video inputs; compose-workflows and prompt-weighting defaults corrected; dead caches and a no-op release removed from the minimax templates. + +- [ ] **Step 4: Commit** + +```bash +git add docs/RELEASING.md +git commit -m "docs(release): note the workflow speed changes" +``` + +--- + +## Out of scope (needs a GPU measurement on lem before changing) + +These came out of the same audit and are deliberately not tasks here. Each changes output or depends on a number this plan cannot produce on a Mac: + +- Sequential -> model offload across the 19 Flux / Z-Image / Krea / Qwen-edit workflows (RECIPES_24GB measured model offload fastest for Flux; needs a per-family VRAM check at each template's canvas). +- `attention_backend: sage_hub` variable on the H3 and LTX templates (28-33% per step measured; needs `kernels` and `DIFFUSERS_TRUST_REMOTE_KERNELS=true`). +- Qwen-Image-2.1 Turbo8 LoRA (catalog `trial`) and a Flux dev 8-step LoRA; quality judgement needed. +- Kandinsky 6 text encoder int8 / on_demand residency; compile on its base transformer. +- qr-code at 150 steps; SD 1.5 templates in float32. +- Background video encode and skipping the post-write LUFS probe for in-memory intermediates; batching upscale tiles and RIFE pairs; SAM2 video propagation. +- Prompt and reference embedding cache per resident pipeline (chained-segments, restore-long, extend-clip). +- #815 on-disk cache of SDNQ-quantized weights. diff --git a/dw/media.py b/dw/media.py index 83df7700..2dbe87e8 100644 --- a/dw/media.py +++ b/dw/media.py @@ -33,6 +33,15 @@ logger = logging.getLogger("dw") + +def _threaded(stream): + """A video stream set to decode on every core. libav defaults a + software h264 decode to one thread; AUTO lets it pick frame and + slice threading for the codec. Header-only reads do not call this.""" + stream.thread_type = "AUTO" + return stream + + # The most a soundtrack may be as base64 before the gallery route refuses to # extract it whole - the twin of dw_mcp/media.py's MAX_RETURNED_BYTES (the # MCP package's cap on any inline payload). Two constants because dw_mcp @@ -345,7 +354,7 @@ def video_shape(path): fps = float(stream.average_rate) if stream.average_rate else None count = int(stream.frames) if stream.frames else None if count is None: - count = sum(1 for _ in container.decode(stream)) + count = sum(1 for _ in container.decode(_threaded(stream))) return { "frame_count": count, "fps": fps, @@ -383,6 +392,8 @@ def read_frames(path, indexes, fit=None): wanted = sorted(set(int(i) for i in indexes)) found = {} with av.open(path) as container: + # Not _threaded: frame threads hold frames in flight, and a backward + # seek after a decoded frame then yields stale ones (tests/test_media_decode_threads.py) stream = container.streams.video[0] fps = float(stream.average_rate) if stream.average_rate else None start_pts = stream.start_time if stream.start_time is not None else 0 @@ -479,6 +490,8 @@ def read_frame_range(path, start, stop): ) out = None with av.open(path) as container: + # Not _threaded: frame threads hold frames in flight, and a backward + # seek after a decoded frame then yields stale ones (tests/test_media_decode_threads.py) stream = container.streams.video[0] fps = float(stream.average_rate) if stream.average_rate else None start_pts = stream.start_time if stream.start_time is not None else 0 @@ -513,7 +526,7 @@ def count_video_frames(path): """The picture frame count of `path`, by decoding every frame rather than reading the header. Nothing is held but the frame in hand.""" with av.open(path) as container: - return sum(1 for _ in container.decode(container.streams.video[0])) + return sum(1 for _ in container.decode(_threaded(container.streams.video[0]))) def decode_audio_video(handle): @@ -531,7 +544,7 @@ def decode_audio_video(handle): sample_rate = None with av.open(handle) as container: - video_stream = container.streams.video[0] + video_stream = _threaded(container.streams.video[0]) frame_rate = ( float(video_stream.average_rate) if video_stream.average_rate else None ) @@ -559,7 +572,10 @@ def decode_rgb_frames(path): """Every picture frame of `path` as an (height, width, 3) uint8 array, in order.""" with av.open(path) as container: - return [frame.to_ndarray(format="rgb24") for frame in container.decode(video=0)] + return [ + frame.to_ndarray(format="rgb24") + for frame in container.decode(_threaded(container.streams.video[0])) + ] def read_thumbnails_and_track(path, thumb_width, thumb_height): @@ -575,7 +591,9 @@ def read_thumbnails_and_track(path, thumb_width, thumb_height): thumbs = [] chunks = [] with av.open(path) as container: - video = container.streams.video[0] if container.streams.video else None + video = ( + _threaded(container.streams.video[0]) if container.streams.video else None + ) audio = container.streams.audio[0] if container.streams.audio else None if video is None and audio is None: raise ValueError(f"{path} has neither a video nor an audio stream") diff --git a/dw/pipeline_ownership.py b/dw/pipeline_ownership.py index d26b480c..8c2354de 100644 --- a/dw/pipeline_ownership.py +++ b/dw/pipeline_ownership.py @@ -62,6 +62,11 @@ def __init__(self, prior_step_keys=None): # pipeline_reference still address pipelines by the step that made # them self.keys_by_step = {} + # Step name -> the step's pipeline definition as the workflow holds + # it (never a loading copy), so a pipeline_reference can put the + # referenced step's own scale, alpha and shift back on a model a + # step differing only in those shares with it + self.definitions_by_step = {} # Step name -> _Deferred for each cache hit whose pipeline was not # resident and so was not loaded self.deferred = {} @@ -92,9 +97,16 @@ def load_key(self, step_definition): ) return self.running[step_name] - def record(self, step_name, key): - """Record which cache key a step's pipeline lives under this run.""" + def record(self, step_name, key, pipeline_definition=None): + """Record which cache key a step's pipeline lives under this run, + and the definition it was loaded or wrapped from.""" self.keys_by_step[step_name] = key + if pipeline_definition is not None: + self.definitions_by_step[step_name] = pipeline_definition + + def definition_for(self, step_name): + """The pipeline definition `step_name` recorded this run, or None.""" + return self.definitions_by_step.get(step_name) def key_for(self, step_name): """The cache key `step_name` recorded this run, or None.""" @@ -286,6 +298,10 @@ def wrap_resident( new_pipeline_wrapper.pipeline_definition.get("seed", default_seed) ) + # The key the hit matched on is the weights identity: a scale, alpha or + # shift may differ from the values the resident model last ran with + new_pipeline_wrapper.apply_runtime_settings() + # A cache hit and a cold load look identical from the outside - # same step, same dot - and they differ by minutes emit_phase("cached", detail=new_pipeline_wrapper.name) diff --git a/dw/pipeline_processors/adapters.py b/dw/pipeline_processors/adapters.py index 54d85161..3e11a224 100644 --- a/dw/pipeline_processors/adapters.py +++ b/dw/pipeline_processors/adapters.py @@ -101,11 +101,103 @@ def active_loras(loras): ] +def adapter_settings(loras): + """(adapter_names, adapter_weights, alphas) for the active entries of + `loras`. An adapter is named by its own adapter_name or its ORIGINAL + index - a nulled entry keeps its slot, so the name a hit re-applies to + is the name the load gave. Reads the entries, never pops them. + + `or` on adapter_name, because a variable nulled by the caller arrives as + a present None; float() on scale, because the schema takes a 'variable:' + reference there, and a variable declared as a string default substitutes + as one.""" + names, weights, alphas = [], [], {} + for i, lora in enumerate(loras or []): + if not isinstance(lora, dict) or lora.get("model_name") is None: + continue + name = lora.get("adapter_name") or str(i) + names.append(name) + scale = lora.get("scale") + weights.append(1.0 if scale is None else float(scale)) + if lora.get("alpha") is not None: + alphas[name] = lora["alpha"] + return names, weights, alphas + + +def set_adapter_scales(pipeline, names, weights): + """Set each named adapter's scale on every loaded layer that carries it, + without re-activating the adapters. + + The warm-pipeline path. diffusers' set_adapters() activates as well as + scales, and activation goes through peft's set_adapter(), which sets + requires_grad on the adapter weights - after a run under model offload + those are inference tensors, and torch refuses ("Setting requires_grad=True + on inference tensor outside InferenceMode" - seen on a 3090 the first time + a scale changed on a warm Flux pipeline, 2026-10-10). + A pipeline that already ran holds the adapters active, so a hit only + needs each layer's set_scale(), which recomputes scale * alpha / rank and + touches no tensor. + + Raises: + ValueError: if an adapter is carried by no loaded layer - a scale + that applied to nothing would silently run the old value + """ + touched = {name: 0 for name in names} + for module in _lora_layers(pipeline): + scaling = getattr(module, "scaling", None) + if not isinstance(scaling, dict): + continue + for name, weight in zip(names, weights): + if name in scaling: + module.set_scale(name, weight) + touched[name] += 1 + missing = [name for name, count in touched.items() if not count] + if missing: + raise ValueError( + f"Cannot set the scale of adapter(s) {missing} - no loaded layer " + "carries them. Nothing would have been scaled." + ) + + +def apply_adapter_settings(loras, pipeline, activate=False): + """Set each active LoRA's alpha, then every scale, on a pipeline that + already holds the adapters. Alpha first: the scale pass is what recomputes + each layer's scaling from it. No-op with no active entry. + + `activate` is the cold path, right after load_lora_weights: diffusers' + set_adapters() activates the adapters and sets their scales in one call. + A warm pipeline (a cache hit) has them active already and takes the + scale-only path - see set_adapter_scales for why it must. + """ + names, weights, alphas = adapter_settings(loras) + if not names: + return + for name, alpha in alphas.items(): + set_adapter_alpha(pipeline, name, alpha) + logger.info(f"Setting adapter weights: {list(zip(names, weights))}") + if activate: + # Positionally - diffusers' mixin calls the second parameter + # 'adapter_weights' while custom pipelines that delegate to the model + # (ostris/Krea2OstrisEdit) call it 'weights' + pipeline.set_adapters(names, weights) + else: + set_adapter_scales(pipeline, names, weights) + + def load_loras(loras, pipeline): - """Load and configure LoRA models.""" - adapter_names = [] - adapter_weights = [] - alphas = {} + """Load and configure LoRA models. + + Names, scales and alphas come from adapter_settings, read before the + entries are popped, so a load and a later cache hit's + apply_adapter_settings set the same values on the same names.""" + # A snapshot of each entry's runtime keys, taken before the pops below + settings = [ + {k: lora.get(k) for k in ("model_name", "adapter_name", "scale", "alpha")} + if isinstance(lora, dict) + else lora + for lora in loras or [] + ] + names = iter(adapter_settings(settings)[0]) for i, lora in enumerate(loras or []): if isinstance(lora, dict) and lora.get("model_name") is None: @@ -117,40 +209,15 @@ def load_loras(loras, pipeline): logger.info(f"Loading LoRA: {model_name}") emit_phase("loading", detail=f"LoRA: {model_name}") - # Use provided adapter_name or generate from index - `or`, because a - # variable nulled by the caller arrives as a present None - adapter_name = lora.pop("adapter_name", None) or str(i) - adapter_names.append(adapter_name) + # Popped before the load: everything left in the dict is a keyword + # argument to load_lora_weights; the scale and alpha are applied to + # the layers afterwards rather than passed to it + for key in ("adapter_name", "scale", "alpha"): + lora.pop(key, None) - # Extract scale for adapter weights - float() because the schema takes a - # 'variable:' reference here, and a variable declared as a string default - # substitutes as one - scale = lora.pop("scale", None) - scale = 1.0 if scale is None else float(scale) - adapter_weights.append(scale) + pipeline.load_lora_weights(model_name, adapter_name=next(names), **lora) - # Popped before the load: everything left in the dict is a keyword - # argument to load_lora_weights, and the alpha is applied to the layers - # afterwards rather than passed to it - alpha = lora.pop("alpha", None) - if alpha is not None: - alphas[adapter_name] = alpha - - # Load the LoRA with the adapter name - pipeline.load_lora_weights(model_name, adapter_name=adapter_name, **lora) - - for adapter_name, alpha in alphas.items(): - set_adapter_alpha(pipeline, adapter_name, alpha) - - # Set adapter weights for all loaded LoRAs - if adapter_names: - logger.info( - f"Setting adapter weights: {list(zip(adapter_names, adapter_weights))}" - ) - # Positionally - diffusers' mixin calls the second parameter 'adapter_weights' - # while custom pipelines that delegate to the model (ostris/Krea2OstrisEdit) - # call it 'weights' - pipeline.set_adapters(adapter_names, adapter_weights) + apply_adapter_settings(settings, pipeline, activate=True) def load_ip_adapter(ip_adapter_definition, pipeline): diff --git a/dw/pipeline_processors/components.py b/dw/pipeline_processors/components.py index f8d24a71..e5757543 100644 --- a/dw/pipeline_processors/components.py +++ b/dw/pipeline_processors/components.py @@ -447,6 +447,23 @@ def load_and_configure_scheduler( ), ) + apply_scheduler_shift(scheduler_definition, pipeline, component_name) + + +def apply_scheduler_shift(scheduler_definition, pipeline, component_name="scheduler"): + """Set the shift a scheduler definition names on an already-loaded + scheduler; nothing when the definition is None or names no shift. + + A load sets it after any scheduler replacement, and a pipeline cache hit + sets it again (Pipeline.apply_runtime_settings), because the cache key + leaves the shift out - a run that only moves it reuses the warm pipeline. + + Raises: + ValueError: if the scheduler is registered but not loaded, or has no + set_shift() + """ + if scheduler_definition is None: + return shift = scheduler_definition.get("shift", None) if shift is None: return diff --git a/dw/pipeline_processors/pipeline.py b/dw/pipeline_processors/pipeline.py index f7f31829..223083d8 100644 --- a/dw/pipeline_processors/pipeline.py +++ b/dw/pipeline_processors/pipeline.py @@ -9,8 +9,14 @@ get_quantization_configuration, get_cache_configuration, ) -from .adapters import active_loras, load_ip_adapter, load_loras +from .adapters import ( + active_loras, + apply_adapter_settings, + load_ip_adapter, + load_loras, +) from .components import ( + apply_scheduler_shift, apply_sdnq_optimizations, configure_components, enable_cache_on_transformer, @@ -135,6 +141,22 @@ def _loading_copy(pipeline_definition): return copied +def apply_runtime_settings(pipeline_definition, model): + """Set what a pipeline definition states that changes no weights - + scheduler and audio_scheduler shift, LoRA scales and alphas - on a model + that already holds those weights. + + Takes the definition as the workflow holds it, never a Pipeline's + loading copy: load_loras pops each entry's model_name from that copy, + and an entry without one reads as switched off and would be skipped. + """ + apply_scheduler_shift(pipeline_definition.get("scheduler"), model) + apply_scheduler_shift( + pipeline_definition.get("audio_scheduler"), model, "audio_scheduler" + ) + apply_adapter_settings(pipeline_definition.get("loras"), model) + + class Pipeline: """ Manages pipeline initialization, configuration, and execution. @@ -303,6 +325,15 @@ def _check_trusted_block(block, what): for entry in value: Pipeline._check_trusted_block(entry, key) + def apply_runtime_settings(self): + """Re-apply what a run may have changed without changing the weights: + scheduler and audio_scheduler shift, LoRA scales and alphas. The + pipeline cache keys on weights_identity, so a hit can carry any of + these at a new value; load() applies the same values on its way + through, so a cold load and a hit end in the same state. Only for a + wrapper that has not loaded (a hit): load() pops from this copy.""" + apply_runtime_settings(self.pipeline_definition, self.pipeline) + def load(self, shared_components): """ Load and configure the pipeline with all components. diff --git a/dw/step_cache.py b/dw/step_cache.py index f784eedf..ca726338 100644 --- a/dw/step_cache.py +++ b/dw/step_cache.py @@ -121,26 +121,101 @@ def component_names(pipeline_definition, key): return list(pipeline_definition.get(key, [])) + list(configuration.get(key, [])) +RUNTIME_LORA_KEYS = ("scale", "alpha") + + +def weights_identity(pipeline_definition): + """The pipeline definition with every runtime-mutable setting removed. + + A LoRA's scale and alpha (set_adapters / set_adapter_alpha) and a + scheduler's shift (set_shift) change nothing about what is loaded, so a + run that only moves one of them must hit the warm pipeline and re-apply + the value (Pipeline.apply_runtime_settings), not reload the stack. + Everything else - the model, the LoRA files, quantization, placement, + the scheduler type - still changes the key. Returns a new dict; the + definition is the workflow's and is not edited. + + An alpha's or a shift's *presence* stays in the identity (as True, never + the value). A hit re-applies only what the definition states, and nothing + sets an alpha or shift back to the checkpoint's own: a run that drops one + (or nulls the variable behind it) must reload to get the default back, + where 16 -> 8 just re-applies. A scale needs no marker - an absent scale + re-applies as 1.0. + """ + identity = { + k: v + for k, v in pipeline_definition.items() + if k not in ("arguments", "seed", "chain") + } + loras = identity.get("loras") + if isinstance(loras, list): + identity["loras"] = [ + _without_runtime(lora, RUNTIME_LORA_KEYS, "alpha") + if isinstance(lora, dict) + else lora + for lora in loras + ] + for name in ("scheduler", "audio_scheduler"): + scheduler = identity.get(name) + if isinstance(scheduler, dict): + identity[name] = _without_runtime(scheduler, ("shift",), "shift") + return identity + + +def _without_runtime(block, runtime_keys, marked): + """`block` without its runtime keys, but with `marked` as True when the + block states it - see weights_identity for why its presence counts.""" + stripped = {k: v for k, v in block.items() if k not in runtime_keys} + if block.get(marked) is not None: + stripped[marked] = True + return stripped + + +def pipeline_definition_key(pipeline_definition): + """A hash of everything in a pipeline definition but what varies per call + (arguments, seed, chain) - runtime settings included, unlike + pipeline_cache_key. + + The step cache keys a borrowing step on this, not on the pipeline cache + key: a step that runs on another step's pipeline (pipeline_reference, + reused_components) produces a different result when only its source's + LoRA scale or shift moves, though the loaded weights - and so the warm + pipeline - are the same. + """ + serialized = json.dumps( + { + k: v + for k, v in pipeline_definition.items() + if k not in ("arguments", "seed", "chain") + }, + sort_keys=True, + default=str, + ) + return hashlib.sha256(serialized.encode()).hexdigest() + + def pipeline_cache_key(pipeline_definition): """Stable identity for a loaded pipeline. - Hashes everything that shapes loading - configuration, components, - quantization, loras - and excludes what varies per call (arguments, seed, - chain), so a cache hit means "this exact model stack is already loaded". - Keying the cache by identity instead of step name means two workflows - whose steps happen to share a name can no longer collide, and a rerun of - an edited workflow keeps every pipeline whose definition did not change. + Hashes the weights identity - everything that shapes loading: + configuration, components, quantization, the LoRA files, the scheduler + type - and excludes what varies per call (arguments, seed, chain), so a + cache hit means "this exact model stack is already loaded". A LoRA's + scale and alpha and a scheduler's shift are excluded too: each is set in + place on the loaded model, and a hit re-applies them + (Pipeline.apply_runtime_settings), so nudging one between runs costs no + reload. Keying the cache by identity instead of step name means two + workflows whose steps happen to share a name can no longer collide, and + a rerun of an edited workflow keeps every pipeline whose definition did + not change. Computed after variable substitution but the excluded keys keep realized per-run values (images, generators) out of the hash; realized types and dtypes stringify stably via default=str. """ - load_definition = { - k: v - for k, v in pipeline_definition.items() - if k not in ("arguments", "seed", "chain") - } - serialized = json.dumps(load_definition, sort_keys=True, default=str) + serialized = json.dumps( + weights_identity(pipeline_definition), sort_keys=True, default=str + ) return hashlib.sha256(serialized.encode()).hexdigest() @@ -170,14 +245,34 @@ def step_pipeline_keys(steps): loaded - a deferred cache hit, or the cache_hits probe. A key that has to agree across those cases is read from this table, not re-hashed. """ + return _closure_keys(steps, pipeline_cache_key) + + +def step_definition_keys(steps): + """Step name -> effective definition key for every pipeline step, taken + before any step runs: step_pipeline_keys over pipeline_definition_key. + + What the step cache folds in for a borrowed pipeline + (borrowed_pipeline_keys), so a source step whose only change is a LoRA + scale, alpha or shift - the same warm pipeline, a different result - is + a miss for the steps that borrow it. Pipeline ownership keeps to + step_pipeline_keys. + """ + return _closure_keys(steps, pipeline_definition_key) + + +def _closure_keys(steps, own_key_of): + """Step name -> the step's own key (`own_key_of` its pipeline + definition), or, for a step that reuses components, the hash of the + sorted own keys over its reuse closure - see step_pipeline_keys.""" keys = {} - # Step name -> the own pipeline_cache_keys of its reuse closure + # Step name -> the own keys of its reuse closure closures = {} for index, step_data in enumerate(steps): if "pipeline" not in step_data: continue pipeline_definition = step_data["pipeline"] - own_key = pipeline_cache_key(pipeline_definition) + own_key = own_key_of(pipeline_definition) closure = {own_key} reused = component_names(pipeline_definition, "reused_components") for name in reused: @@ -215,10 +310,15 @@ def borrowed_pipeline_keys(steps, index, pipeline_keys): step's model changes - its own definition never names the model. Folding the source steps' keys into the lookup snapshot makes a source model change a miss here too, rather than a stale hit that republishes a - now-wrong pipeline. The keys are effective keys (step_pipeline_keys), so - a change further up a borrow chain than the direct source is a miss too. + now-wrong pipeline. The keys are effective keys, so a change further up + a borrow chain than the direct source is a miss too. + + The step cache passes step_definition_keys, which also moves with a + source's LoRA scale, alpha or shift - a change that keeps the warm + pipeline but changes what the borrower produces. Callers that only need + which steps are borrowed may pass step_pipeline_keys. - The keys come from `pipeline_keys` (step_pipeline_keys), computed from + The keys come from `pipeline_keys` (one of those tables), computed from the definition before any step ran, never from `steps[:index]` as they stand now: a load edits its own definition in place, and a source step that loaded before this lookup (a cold run) and one that did not (a diff --git a/dw/tasks/restore_faces.py b/dw/tasks/restore_faces.py index 3267c90e..7adb161c 100644 --- a/dw/tasks/restore_faces.py +++ b/dw/tasks/restore_faces.py @@ -66,15 +66,25 @@ def load_descriptor(): ("restore_faces", model_name, filename, str(device)), load_descriptor ) - # Set up facexlib helper - face_helper = FaceRestoreHelper( - upscale_factor=upscale_factor, - face_size=face_size, - crop_ratio=(1, 1), - det_model="retinaface_resnet50", - use_parse=use_parse, - device=torch.device(device), + # facexlib's helper loads RetinaFace (and ParseNet with use_parse) in + # its constructor. per_frame calls this once per frame, so a video was + # reloading detector weights for every frame. The helper keeps per-image + # state between calls; clean_all() resets it before each read + def build_helper(): + return FaceRestoreHelper( + upscale_factor=upscale_factor, + face_size=face_size, + crop_ratio=(1, 1), + det_model="retinaface_resnet50", + use_parse=use_parse, + device=torch.device(device), + ) + + face_helper = cached_model( + ("restore_faces_helper", upscale_factor, face_size, use_parse, str(device)), + build_helper, ) + face_helper.clean_all() # Convert PIL to BGR numpy (facexlib format) input_bgr = np.array(image.convert("RGB"))[:, :, ::-1].copy() diff --git a/dw/worker.py b/dw/worker.py index 7ce52ec1..0759cf5c 100644 --- a/dw/worker.py +++ b/dw/worker.py @@ -357,15 +357,38 @@ def _prepare_workflow(self, command: Dict[str, Any], job: "_Job"): # what runs out of memory if identity != self.workflow_identity: if self.workflow_identity is not None: - self._reply( - Output(message="Workflow changed - releasing cached models...") - ) - self._reply(Output(message=self._cleanup_all())) + keep = self._keys_worth_keeping(command, job.workflow) + if keep: + self._reply( + Output( + message=f"Workflow changed - keeping {len(keep)} warm " + "model(s) the new workflow loads, releasing the rest..." + ) + ) + else: + self._reply( + Output(message="Workflow changed - releasing cached models...") + ) + self._reply(Output(message=self._cleanup_all(keep=keep))) self.workflow_identity = identity self._reply(WorkflowLoaded(workflow_name=job.workflow.name)) self._reply(Output(message=f"Executing workflow: {job.workflow.name}")) + def _keys_worth_keeping(self, command, workflow): + """The resident pipeline keys the next workflow loads anyway. Any + failure to prepare the definition answers the empty set - the run + itself will report what is wrong, and a full release is the safe + default on a card that cannot hold two stacks.""" + try: + wanted = workflow_run.pipeline_keys( + workflow, command.get("arguments") or {} + ) + except Exception as e: + logger.debug(f"Could not take the next workflow's pipeline keys: {e}") + return set() + return wanted & set(self.loaded_pipelines) + def _run_job(self, command: Dict[str, Any], job: "_Job"): """Phase 3: run the workflow, streaming its events, then reclaim.""" workflow = job.workflow @@ -635,10 +658,11 @@ def _cleanup_between_runs(self): logger.debug("Inter-run cleanup complete") - def _cleanup_all(self): + def _cleanup_all(self, keep=frozenset()): """ Complete cleanup - clear all cached models and components. - Called when workflow changes or on shutdown. + Called when workflow changes or on shutdown. Pipelines whose key is + in `keep` survive; with `keep` empty everything goes. Returns: One line saying what host memory looks like on the other side of @@ -652,11 +676,20 @@ def _cleanup_all(self): logger.info("Performing full cleanup") # Clear pipeline cache and any models task handlers cached + kept = {k: v for k, v in self.loaded_pipelines.items() if k in keep} + dropped = len(self.loaded_pipelines) - len(kept) + # Cleared and refilled in place, so the kept entries keep their + # identity self.loaded_pipelines.clear() + self.loaded_pipelines.update(kept) + # Every run republishes shared components from the pipelines it hits + # (workflow.py: publish_shared_components on a cache hit) self.shared_components.clear() - # The keys addressed entries that are now gone; keeping them would - # have a later step chase a release that has already happened - self.prior_step_keys.clear() + # A step's prior key addresses an entry that is now gone - unless it + # was kept, in which case a redefinition next run still releases it + self.prior_step_keys = { + name: key for name, key in self.prior_step_keys.items() if key in kept + } clear_model_cache() # Drop cached step results too - stale results would otherwise # survive a memory clear and keep getting served for steps whose @@ -703,9 +736,10 @@ def _cleanup_all(self): f"Released cached models: host RSS {_mb(stats.get('rss_mb'))}, " f"{_mb(stats.get('available_mb'))} available" + (f" ({released:.0f} MB returned to the OS)" if released else "") + + (f", kept {len(kept)} warm" if kept else "") ) - logger.info(f"Full cleanup complete. {summary}") + logger.info(f"Full cleanup complete ({dropped} pipeline(s) dropped). {summary}") return summary def _parent_is_dead(self) -> bool: diff --git a/dw/workflow.py b/dw/workflow.py index 7e77649b..2e7e6cfe 100644 --- a/dw/workflow.py +++ b/dw/workflow.py @@ -27,7 +27,7 @@ resolve_variable_values, set_variables, ) -from .pipeline_processors.pipeline import Pipeline +from .pipeline_processors.pipeline import Pipeline, apply_runtime_settings from .tasks.task import Task from . import get_device from .pipeline_ownership import ( @@ -730,7 +730,7 @@ def _pipeline_action( # Pipelines are cached by what they load, not what step loads them ownership = self.pipeline_ownership cache_key = ownership.load_key(step_definition) - ownership.record(step_name, cache_key) + ownership.record(step_name, cache_key, step_definition["pipeline"]) get_context().touch_pipeline(cache_key) if cache_hit and cache_key not in previous_pipelines: @@ -798,6 +798,12 @@ def _reference_action( "earlier pipeline step in this run (or it was released)" ) previous_pipeline = previous_pipelines[referenced_key] + # Steps whose pipelines differ only in scale, alpha or shift share + # one model, and the last of them to run left its values on it: put + # the referenced step's own back + referenced_definition = self.pipeline_ownership.definition_for(reference_name) + if referenced_definition is not None: + apply_runtime_settings(referenced_definition, previous_pipeline.pipeline) return Pipeline( pipeline_reference, default_seed, diff --git a/dw/workflow_run.py b/dw/workflow_run.py index 058cd44b..ca913a05 100644 --- a/dw/workflow_run.py +++ b/dw/workflow_run.py @@ -62,6 +62,7 @@ reference_resolves_to, referenced_result_names, step_cache, + step_definition_keys, step_pipeline_keys, ) from .subfolders import step_subfolder @@ -138,6 +139,9 @@ class StepLoop: hits: set = field(default_factory=set) # Final result is the workflow return value last_result: object = None + # step_definition_keys of `steps`, taken before any step loaded - what + # the step cache folds in for a borrowed pipeline (cache_lookup) + definition_keys: dict = field(default_factory=dict) def run_base_dir(workflow): @@ -367,14 +371,16 @@ def cache_lookup( step_seed, hits_this_run, cache_enabled, - pipeline_keys, + definition_keys, ): """Whether the step cache serves step `index`, as (cached_result or None, the step_data snapshot the entry is keyed on or None, whether a later step still reads this one's result, the names later steps still reference). Shared by run() and cache_hits() - see - prepare_definition for why. `pipeline_keys` is step_pipeline_keys - of `steps`, taken before the first step ran. + prepare_definition for why. `definition_keys` is step_definition_keys + of `steps`, taken before the first step ran - not the pipeline + ownership table, whose keys leave out a source's LoRA scale, alpha and + shift, each of which changes what a borrowing step produces. """ # What later steps still read, which decides both whether this # step's result has to be kept alive after the step (release_unreferenced_results @@ -411,7 +417,7 @@ def cache_lookup( # model - only the source step's does - so without this a # source model change would leave the borrowing step's # snapshot unchanged and serve a stale hit - borrowed = borrowed_pipeline_keys(steps, index, pipeline_keys) + borrowed = borrowed_pipeline_keys(steps, index, definition_keys) if borrowed: step_data_snapshot["__borrowed_pipelines__"] = borrowed except Exception as ex: @@ -465,7 +471,7 @@ def cache_hits(workflow, arguments): realize_args(steps, base_dir) # The same table run() takes at the same point, so a borrowed # key here is the key the run stored its entry under - pipeline_keys = step_pipeline_keys(steps) + definition_keys = step_definition_keys(steps) hits_this_run = set() hits = [] for index, step_data in enumerate(steps): @@ -479,7 +485,7 @@ def cache_hits(workflow, arguments): step_seed, hits_this_run, True, - pipeline_keys, + definition_keys, ) if cached_result is not None: hits_this_run.add(step_data["name"]) @@ -489,6 +495,33 @@ def cache_hits(workflow, arguments): deactivate_output_root(output_root_token) +def pipeline_keys(workflow, arguments): + """The set of pipeline cache keys a run of `workflow` with `arguments` + loads its steps under - the table run() and cache_hits() take + (step_pipeline_keys over the prepared definition). Prepares the + definition exactly as a run does and executes nothing: the worker asks + this before releasing the previous workflow's models, so a model the + next workflow loads anyway stays warm. A sub-workflow step loads its + own pipelines later and is not in the set; an unseeded workflow still + answers, since warmth does not depend on the step cache. The probe + realizes the steps because the run keys on the realized definition, at + the price of loading step media once per switch (the price cache_hits + already pays). + """ + output_root_token = activate_output_root(workflow.output_dir) + try: + workflow_def = copy.deepcopy(workflow.workflow_definition) + base_dir = run_base_dir(workflow) + workflow_def, _, _ = prepare_definition( + workflow, workflow_def, arguments or {}, base_dir + ) + steps = workflow_def.get("steps", []) + realize_args(steps, base_dir) + return set(step_pipeline_keys(steps).values()) + finally: + deactivate_output_root(output_root_token) + + def prepare_run(workflow, record): """Phase B: the definition as this run works from it, and its seed. @@ -678,6 +711,7 @@ def begin_steps(workflow, prepared, previous_pipelines, run_context): prepared.cache_enabled, run_context, pipelines, + definition_keys=step_definition_keys(steps), ) @@ -875,7 +909,7 @@ def run_step(workflow, loop, index, step_data, record): step_seed, loop.hits, loop.cache_enabled, - workflow.pipeline_ownership.running, + loop.definition_keys, ) # The last step of a composed child whose parent does the # saving (#92) - its files are written once, by the parent diff --git a/plugins/dw/.claude-plugin/plugin.json b/plugins/dw/.claude-plugin/plugin.json index 149d97c9..7a6c095c 100644 --- a/plugins/dw/.claude-plugin/plugin.json +++ b/plugins/dw/.claude-plugin/plugin.json @@ -1,7 +1,7 @@ { "name": "dw", "description": "Compose MiniMax H3 video, MiniMax Music 3, LTX-2.5 and Kandinsky 6.0 workflows over a dw MCP server, and cut a multi-episode series from them: which template fits which shape, the hard rules, cost, and how to judge the output. Prompt format comes from the vendors' own guides.", - "version": "0.11.0", + "version": "0.12.0-alpha.1", "author": { "name": "Don Kackman" }, diff --git a/pyproject.toml b/pyproject.toml index 1aa8fc60..66011ddf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ name = "diffusers-workflow" # runtime (static in TOML because importing dw at build time would drag # torch into the build environment). Release tags must match it - see # docs/RELEASING.md -version = "0.11.0" +version = "0.12.0-alpha.1" description = "Declarative workflow engine and web UI for Hugging Face Diffusers" readme = "README.md" requires-python = ">=3.10" diff --git a/tests/test_catalog_speed_defaults.py b/tests/test_catalog_speed_defaults.py new file mode 100644 index 00000000..fd142326 --- /dev/null +++ b/tests/test_catalog_speed_defaults.py @@ -0,0 +1,76 @@ +"""Catalog defaults that silently multiply a run's cost (speed audit 2026-10-10).""" + +import glob +import json +import os + +import pytest + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +TEMPLATES = sorted( + glob.glob( + os.path.join(REPO_ROOT, "workflows", "templates", "**", "*.json"), + recursive=True, + ) +) + + +def _load(path): + with open(path, encoding="utf-8") as handle: + return json.load(handle) + + +def _default(definition, name): + variable = definition.get("variables", {}).get(name) + return variable.get("default") if isinstance(variable, dict) else variable + + +def _pipeline_steps(definition): + for step in definition.get("steps", []): + if "pipeline" in step: + yield step + + +MINIMAX = [p for p in TEMPLATES if os.sep + "minimax" + os.sep in p] + + +@pytest.mark.parametrize("path", MINIMAX, ids=os.path.basename) +def test_turbo_schedules_carry_no_block_cache(path): + definition = _load(path) + weight = _default(definition, "lora_weight_name") or "" + steps = _default(definition, "num_inference_steps") + if "turbo" not in weight or steps is None or steps > 9: + pytest.skip("not a turbo schedule") + for step in _pipeline_steps(definition): + configuration = step["pipeline"].get("configuration", {}) + assert "cache" not in configuration, ( + f"{os.path.basename(path)} step '{step['name']}' carries a cache block on a " + f"{steps}-step turbo schedule; RECIPES_24GB.md says it never skips there" + ) + + +def test_compose_workflows_video_steps_match_the_turbo_default(): + compose = _load( + os.path.join(REPO_ROOT, "workflows", "templates", "compose-workflows.json") + ) + i2v = _load( + os.path.join( + REPO_ROOT, "workflows", "templates", "minimax", "image-to-video.json" + ) + ) + assert _default(compose, "video_num_inference_steps") == _default( + i2v, "num_inference_steps" + ) + + +@pytest.mark.parametrize("path", TEMPLATES, ids=lambda p: os.path.relpath(p, REPO_ROOT)) +def test_schnell_runs_its_distilled_schedule(path): + definition = _load(path) + for step in _pipeline_steps(definition): + pipeline = step["pipeline"] + model = pipeline.get("from_pretrained_arguments", {}).get("model_name", "") + if "FLUX.1-schnell" not in str(model): + continue + arguments = pipeline.get("arguments", {}) + assert arguments.get("num_inference_steps", 4) <= 4, path + assert arguments.get("guidance_scale", 0.0) == 0.0, path diff --git a/tests/test_catalog_structure.py b/tests/test_catalog_structure.py index f4cde760..c9ff3ef9 100644 --- a/tests/test_catalog_structure.py +++ b/tests/test_catalog_structure.py @@ -471,7 +471,10 @@ def test_no_stale_entry_in_the_allowlist(): # Then to 10_650 for the `relevel-clip` utility template (#794), measured at # 10_632: a new catalog entry, the CPU-only fix for a rendered clip that # clips, where the only other route was a full GPU re-render. -COMPACT_BUDGET = 10_650 +# Then to 10_700 for the `attention_backend` variable on the four MiniMax-H3 +# family baselines (2026-10-10), measured at 10_676 with six: the knob behind the +# measured 28-33% per-step gain from sage_hub, null by default. +COMPACT_BUDGET = 10_700 FILTERED_BUDGET = 1_500 diff --git a/tests/test_media_decode_threads.py b/tests/test_media_decode_threads.py new file mode 100644 index 00000000..6cb91e9e --- /dev/null +++ b/tests/test_media_decode_threads.py @@ -0,0 +1,81 @@ +"""Every picture-decoding open in dw/media.py decodes on every core.""" + +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from dw import media + + +class FakeFrame: + def __init__(self): + self.pts = 0 + + def to_ndarray(self, format="rgb24"): + return np.zeros((2, 2, 3), dtype=np.uint8) + + +def _container(frames=1): + container = MagicMock() + stream = MagicMock() + stream.average_rate = 8 + stream.frames = 0 + stream.start_time = 0 + stream.time_base = 1 + stream.width = 2 + stream.height = 2 + del stream.thread_type # set by the code under test, asserted below + container.streams.video = [stream] + container.streams.audio = [] + container.decode.return_value = [FakeFrame() for _ in range(frames)] + container.__enter__.return_value = container + container.__exit__.return_value = False + return container, stream + + +@pytest.mark.parametrize( + "call", + [ + lambda: media.decode_rgb_frames("x.mp4"), + lambda: media.decode_audio_video("x.mp4"), + lambda: media.video_shape("x.mp4"), + lambda: media.count_video_frames("x.mp4"), + lambda: media.read_thumbnails_and_track("x.mp4", 2, 2), + ], +) +def test_decoders_set_auto_threading(call): + container, stream = _container() + with patch.object(media.av, "open", return_value=container): + try: + call() + except Exception: + pass # the fake need not satisfy the rest of the function + assert stream.thread_type == "AUTO" + + +def test_video_shape_header_read_does_not_thread(): + container, stream = _container() + stream.frames = 10 + with patch.object(media.av, "open", return_value=container): + media.video_shape("x.mp4") + assert not hasattr(stream, "thread_type") + + +@pytest.mark.parametrize( + "call", + [ + lambda: media.read_frames("x.mp4", [0]), + lambda: media.read_frame_range("x.mp4", 0, 1), + ], +) +def test_seeking_readers_stay_single_threaded(call): + """Frame threads hold frames in flight; a backward seek after a decoded + frame then yields stale ones (the recovery seek of read_frames).""" + container, stream = _container() + with patch.object(media.av, "open", return_value=container): + try: + call() + except Exception: + pass + assert not hasattr(stream, "thread_type") diff --git a/tests/test_restore_faces_helper_cache.py b/tests/test_restore_faces_helper_cache.py new file mode 100644 index 00000000..e58f8042 --- /dev/null +++ b/tests/test_restore_faces_helper_cache.py @@ -0,0 +1,93 @@ +"""restore_faces builds its facexlib helper once per (settings, device), not per frame.""" + +import sys +import types +from unittest.mock import MagicMock + +import numpy as np +import pytest +from PIL import Image + +from dw.tasks import model_cache + + +@pytest.fixture +def fake_facexlib(monkeypatch): + """A facexlib whose helper records constructions and detects no faces.""" + constructed = [] + + class FakeHelper: + def __init__(self, **kwargs): + constructed.append(kwargs) + self.cleaned = 0 + + def clean_all(self): + self.cleaned += 1 + + def read_image(self, image): + self.image = image + + def get_face_landmarks_5(self, **kwargs): + return 0 + + module = types.ModuleType("facexlib.utils.face_restoration_helper") + module.FaceRestoreHelper = FakeHelper + package = types.ModuleType("facexlib") + utils = types.ModuleType("facexlib.utils") + monkeypatch.setitem(sys.modules, "facexlib", package) + monkeypatch.setitem(sys.modules, "facexlib.utils", utils) + monkeypatch.setitem(sys.modules, "facexlib.utils.face_restoration_helper", module) + model_cache.clear_model_cache() + yield constructed + model_cache.clear_model_cache() + + +def test_two_frames_share_one_helper(fake_facexlib, monkeypatch): + from dw.tasks import restore_faces as module + + descriptor = MagicMock(supports_half=False) + monkeypatch.setattr(module, "_load_face_model", lambda path, device: descriptor) + monkeypatch.setattr( + "dw.tasks.upscale.resolve_model_path", + lambda model_name, filename: "weights.pth", + ) + frame = Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)) + + module.restore_faces(frame, "m", device="cpu") + module.restore_faces(frame, "m", device="cpu") + + assert len(fake_facexlib) == 1 + helper_kwargs = fake_facexlib[0] + assert helper_kwargs["det_model"] == "retinaface_resnet50" + + +def test_helper_state_is_cleared_before_each_frame(fake_facexlib, monkeypatch): + from dw.tasks import restore_faces as module + + descriptor = MagicMock(supports_half=False) + monkeypatch.setattr(module, "_load_face_model", lambda path, device: descriptor) + monkeypatch.setattr( + "dw.tasks.upscale.resolve_model_path", + lambda model_name, filename: "weights.pth", + ) + frame = Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)) + module.restore_faces(frame, "m", device="cpu") + module.restore_faces(frame, "m", device="cpu") + + helper = model_cache._cache[("restore_faces_helper", 1, 512, True, "cpu")] + assert helper.cleaned == 2 + + +def test_different_settings_build_different_helpers(fake_facexlib, monkeypatch): + from dw.tasks import restore_faces as module + + descriptor = MagicMock(supports_half=False) + monkeypatch.setattr(module, "_load_face_model", lambda path, device: descriptor) + monkeypatch.setattr( + "dw.tasks.upscale.resolve_model_path", + lambda model_name, filename: "weights.pth", + ) + frame = Image.fromarray(np.zeros((8, 8, 3), dtype=np.uint8)) + module.restore_faces(frame, "m", device="cpu", face_size=512) + module.restore_faces(frame, "m", device="cpu", face_size=256) + assert len(fake_facexlib) == 2 diff --git a/tests/test_runtime_settings.py b/tests/test_runtime_settings.py new file mode 100644 index 00000000..ee0a229b --- /dev/null +++ b/tests/test_runtime_settings.py @@ -0,0 +1,151 @@ +"""A cache hit re-applies LoRA scale/alpha and scheduler shift in place.""" + +from unittest.mock import MagicMock + +import pytest + +from dw.pipeline_processors.adapters import adapter_settings, apply_adapter_settings +from dw.pipeline_processors.components import apply_scheduler_shift + + +def test_adapter_settings_name_by_original_index_and_skip_nulled(): + loras = [ + {"model_name": None, "scale": 0.3}, + {"model_name": "a/b", "scale": 0.5, "alpha": 8}, + {"model_name": "c/d", "adapter_name": "turbo"}, + ] + names, weights, alphas = adapter_settings(loras) + assert names == ["1", "turbo"] + assert weights == [0.5, 1.0] + assert alphas == {"1": 8} + assert loras[1]["scale"] == 0.5 # read, not popped + + +def test_the_cold_path_activates_through_set_adapters(): + pipeline = MagicMock() + apply_adapter_settings( + [{"model_name": "a/b", "scale": 0.5}], pipeline, activate=True + ) + pipeline.set_adapters.assert_called_once_with(["0"], [0.5]) + + +class _FakeLoraLayer: + """A peft-shaped layer: the scaling dict set_scale() recomputes.""" + + def __init__(self, adapters, alpha=8.0, rank=16): + self.lora_alpha = {name: alpha for name in adapters} + self.r = {name: rank for name in adapters} + self.scaling = {name: alpha / rank for name in adapters} + self.activated = 0 + + def set_scale(self, adapter, scale): + if adapter in self.scaling: + self.scaling[adapter] = scale * self.lora_alpha[adapter] / self.r[adapter] + + def set_adapter(self, names): + self.activated += 1 + + +def test_the_warm_path_sets_scales_without_touching_activation(monkeypatch): + from dw.pipeline_processors import adapters as module + + layers = [_FakeLoraLayer(["turbo"]), _FakeLoraLayer(["turbo", "other"])] + monkeypatch.setattr(module, "_lora_layers", lambda pipeline: iter(layers)) + pipeline = MagicMock() + + apply_adapter_settings( + [{"model_name": "a/b", "adapter_name": "turbo", "scale": 0.25}], pipeline + ) + + assert [layer.scaling["turbo"] for layer in layers] == [0.125, 0.125] + assert layers[1].scaling["other"] == 0.5 + assert all(layer.activated == 0 for layer in layers) + pipeline.set_adapters.assert_not_called() + + +def test_the_warm_path_sets_alpha_before_the_scale(monkeypatch): + from dw.pipeline_processors import adapters as module + + layer = _FakeLoraLayer(["0"], alpha=8.0, rank=16) + monkeypatch.setattr(module, "_lora_layers", lambda pipeline: iter([layer])) + pipeline = MagicMock() + + apply_adapter_settings([{"model_name": "a/b", "scale": 0.5, "alpha": 4}], pipeline) + + # The scale pass recomputes from the new alpha: 0.5 * 4 / 16, not 0.5 * 8 / 16 + assert layer.lora_alpha["0"] == 4.0 + assert layer.scaling["0"] == 0.125 + pipeline.set_adapters.assert_not_called() + + +def test_the_warm_path_refuses_an_adapter_no_layer_carries(monkeypatch): + from dw.pipeline_processors import adapters as module + + monkeypatch.setattr( + module, "_lora_layers", lambda pipeline: iter([_FakeLoraLayer(["other"])]) + ) + with pytest.raises(ValueError, match="turbo"): + apply_adapter_settings( + [{"model_name": "a/b", "adapter_name": "turbo"}], MagicMock() + ) + + +def test_apply_adapter_settings_is_a_no_op_without_active_loras(): + pipeline = MagicMock() + apply_adapter_settings([{"model_name": None}], pipeline) + apply_adapter_settings(None, pipeline) + pipeline.set_adapters.assert_not_called() + + +def test_apply_scheduler_shift_sets_the_shift(): + pipeline = MagicMock() + apply_scheduler_shift({"shift": 12}, pipeline, "scheduler") + pipeline.scheduler.set_shift.assert_called_once_with(12.0) + + +def test_apply_scheduler_shift_without_a_shift_touches_nothing(): + pipeline = MagicMock() + apply_scheduler_shift(None, pipeline, "scheduler") + apply_scheduler_shift({"scheduler_type": "X"}, pipeline, "scheduler") + pipeline.scheduler.set_shift.assert_not_called() + + +def test_wrap_resident_reapplies_runtime_settings(monkeypatch): + from dw.pipeline_ownership import wrap_resident + from dw.pipeline_processors import adapters as module + + layer = _FakeLoraLayer(["0"]) + monkeypatch.setattr(module, "_lora_layers", lambda pipeline: iter([layer])) + + model = MagicMock() + cached = MagicMock() + cached.pipeline = model + step_definition = { + "name": "s", + "pipeline": { + "configuration": {"component_type": "X", "no_generator": True}, + "loras": [{"model_name": "a/b", "scale": 0.25}], + "scheduler": {"shift": 9}, + }, + } + wrapper = wrap_resident(cached, step_definition, 1, "cpu", "/out", "p") + assert wrapper.pipeline is model + model.scheduler.set_shift.assert_called_once_with(9.0) + assert layer.scaling["0"] == 0.25 * 8.0 / 16 + assert layer.activated == 0 + model.set_adapters.assert_not_called() + + +def test_wrap_resident_with_nothing_mutable_calls_nothing(): + from dw.pipeline_ownership import wrap_resident + + model = MagicMock() + cached = MagicMock() + cached.pipeline = model + step_definition = { + "name": "s", + "pipeline": {"configuration": {"component_type": "X", "no_generator": True}}, + } + wrap_resident(cached, step_definition, 1, "cpu", "/out", "p") + model.set_adapters.assert_not_called() + model.scheduler.set_shift.assert_not_called() diff --git a/tests/test_step_cache.py b/tests/test_step_cache.py index d9744c4b..9cdd5593 100644 --- a/tests/test_step_cache.py +++ b/tests/test_step_cache.py @@ -12,6 +12,8 @@ audio_replaced_downstream, normalized_downstream, _result_bytes, + pipeline_cache_key, + weights_identity, ) @@ -600,9 +602,22 @@ def test_normalized_downstream_false_when_nothing_references_it(): # lora, prompt-weighting and step-caching from model to sequential offload: # offload is part of how a pipeline is loaded, so its key moves with it. # +# Every unquantized FLUX row (fourteen templates) and the Z-Image portrait +# steps of four minimax templates moved on 2026-10-10 when sequential offload +# left the catalog: FLUX streams its transformer block by block from pinned +# host memory, Z-Image takes model offload (docs/RECIPES_24GB.md). +# # Nine minimax rows were regenerated when the T2VA/FL2VA templates moved from # the 8-step turbo LoRAs to the 4-step v1.2 file (2026-10-10): the adapter and -# the scheduler shift are part of how the pipeline is loaded. +# the scheduler shift were then part of how the pipeline is loaded. +# +# Thirty-eight rows across the minimax and ltx2 templates were regenerated +# when the key moved to the weights identity (2026-10-10): a LoRA's scale +# and alpha and a scheduler's shift are set in place on a warm pipeline, so +# they left the hash, and every step that states one moved. Thirty of them +# moved again the same day when an alpha's or a shift's presence (True, +# never its value) went back in: a hit cannot restore a checkpoint default, +# so dropping one has to reload. TEMPLATE_PIPELINE_KEYS = { "assemble-and-score.json": {}, "attention-processor.json": { @@ -626,16 +641,16 @@ def test_normalized_downstream_false_when_nothing_references_it(): }, "compose-workflows.json": {}, "consistent-set.json": { - "base": "ee282369a48c066f07fd13d0e7fc77b54f0342b518d8d5927b7900d9f759b6c1", - "first": "ee282369a48c066f07fd13d0e7fc77b54f0342b518d8d5927b7900d9f759b6c1", - "second": "ee282369a48c066f07fd13d0e7fc77b54f0342b518d8d5927b7900d9f759b6c1", - "third": "ee282369a48c066f07fd13d0e7fc77b54f0342b518d8d5927b7900d9f759b6c1", + "base": "241e4cc77f00d837873d01e3ea9ec7105513af8362908d8229c740cbca997737", + "first": "241e4cc77f00d837873d01e3ea9ec7105513af8362908d8229c740cbca997737", + "second": "241e4cc77f00d837873d01e3ea9ec7105513af8362908d8229c740cbca997737", + "third": "241e4cc77f00d837873d01e3ea9ec7105513af8362908d8229c740cbca997737", }, "controlnet-component.json": { "main": "23e479d111c0c2883755e0d32ef48b686d89f8967b535524ed8c6cb6f863bb9d" }, "controlnet.json": { - "FluxCanny": "cfe38a8a886222292666951457910698f54fb42d631de87009f4319e267bb6f4" + "FluxCanny": "d5e1abfe8058ed9b2235c2b93aef6e155faaaedfc6cb3de849911fd577ea0f5b" }, "depth-marigold.json": { "depth": "9338475d3fb48eab6e3dc8eb53cb07a916470906e338bb02a71e3ed0416a447b" @@ -643,30 +658,30 @@ def test_normalized_downstream_false_when_nothing_references_it(): "describe-and-regenerate.json": {}, "dissolve-between-shots.json": {}, "embed-metadata.json": { - "generate": "5326e88c70c970b932201c47ab65899f5fd17a307d583f1add96a24792e17585" + "generate": "bfe2e09d70b6bf6e7799266cf735abd779e5a9e3c283b8287d962eab17d95425" }, "expand-prompt.json": {}, "generate-speech.json": {}, "image-edit.json": { - "edit": "ee282369a48c066f07fd13d0e7fc77b54f0342b518d8d5927b7900d9f759b6c1" + "edit": "241e4cc77f00d837873d01e3ea9ec7105513af8362908d8229c740cbca997737" }, "image-processors.json": {}, "image-to-image.json": { - "main": "bf8dcfe42d46a7031793e8aea4460cef13b511ca93b1a40cb78b50a98c50f4f8" + "main": "d94c9c21dc83d4781a92dfda09ef657d142f1a62e6cef1a2f26edd87d14d6480" }, "image-to-text.json": {}, "image-variation.json": { - "main": "eb340f983349eb499e6fe01e29c65baf3fefd807e1f6a6a74f561768225fc496", + "main": "1e30480fbe56f5c5c73f6f3f60876efceac5fbc2ecca4a3883e9e3a6b01dc876", "prior": "3253c892afe8ff8d02c464e2b175c519cfac4585fef0e9d2231683719db61a81", }, "inpaint.json": { - "fill": "a52f958d8c596182fb83182bad33df3da491b63a3ad758ec5807103b401f92eb" + "fill": "8ece64de576884be947f5af07c5713b92d3a1e641035434ef805629f12ae8c77" }, "interpolate-frames.json": { "generate_video": "e35d0272e90c76b6730987d5c60a8dfda18290ef71276eb72246e1468ad16dfb" }, "ip-adapter.json": { - "main": "d75dffe2b60f5b3b65ad094896aa982029a830d14d149c0eb099caa8acf78291" + "main": "a67c769c0ee881e36aa15b62f8226a233c1e4b93542a8caaf5eab5715cb9f556" }, "kandinsky6/generate-and-upscale.json": { "text_to_video": "3095873d2e63833c2e8c8048379035840694ca108309e544ff17a0722b4ee51a", @@ -679,19 +694,19 @@ def test_normalized_downstream_false_when_nothing_references_it(): "text_to_video": "3095873d2e63833c2e8c8048379035840694ca108309e544ff17a0722b4ee51a" }, "lora-styles.json": { - "couple": "e0d4d24cbe4caec22ee9e2afc43e9cfd772971430ed633cd289ff362e8213816", - "font-design": "8f28fcccd557b862030508fa1db006f365afbf9b394f981488673df4c46eb548", - "home-decoration": "8f0b0e4608b143d62d89641cddb7a4ff8b97c0d3cd4ba476f973fe6701bc36cc", - "portrait-illustration": "23e72f3a1c737781910ef385cfc156730bfbe9ecc36a67c53249c70e5dbfa404", - "portrait-photography": "6fd461304bdd784ecd9e94aea9a2f06f66ef61426ff9dc6cb552bf08e69d0520", - "ppt-templates": "b03b08bd86d82d633a07ce1e62326223609340ba17db59b8aebaa9c43a9db8b5", - "sandstorm-visual-effect": "0ec17d1beaaa8edafcbc5df532db4bef528694619f398fdaaf19f772c0d44871", - "sparklers-visual-effect": "d7dec11561b3645b6c316159e06acd20c6ff84883712c5ba1b252c1de88424ae", - "storyboard": "25c647a6d3e2013861ffafc783d8e15232f344f0c8579030e9d46b94bc5383d4", - "visual-identity-design": "004a829d2a7c1d55b36f279ffc67057a165ecbf8a109da55f93f6acb09545afa", + "couple": "4ccc3e1cf80d2117f2be5a9b29c345a8d105cc901ef7f0e87330bd918d12c03b", + "font-design": "430a83558b5a8e37726837e558c55aeca8b0d97c977a33ab3c022072fccece86", + "home-decoration": "aa084137494a66c5ebbcb582451a439c46fb13a46ff18126ab3cf6718984aa71", + "portrait-illustration": "af0647ee216ae5223ee690d49d38279d12394721bb03acf43f3e56e350544e02", + "portrait-photography": "2abc410efe284258a9de41aa39bf5757e67dad6c120da66e3bcb77b4705c5acf", + "ppt-templates": "5e599c25436980f20285584ab53a625342ea5a23cfd7ffa0916140cd419608ec", + "sandstorm-visual-effect": "cfe11545ca180142a36773ecb86a8a0d9a879899ff55f7de5e23ff6d7e7a034f", + "sparklers-visual-effect": "1ae7ca40cbd46a32ca6d0d396ed2882fb578f3b6b50ec21f429350176734ec65", + "storyboard": "43728e3912d7dddffe8009e803c56c9b0885a9bc9ca524ad50df429dc11fcec7", + "visual-identity-design": "ed6bb9742190c26d75bf5444ddff876a3198c66a8ee114f484c3b0f9298bb2ff", }, "lora.json": { - "txt2img": "45cddf5b7b38cd2ba06e7906368e45eec2db70ae11ad12edbee1f5604878094e" + "txt2img": "04031c5e01597b46f9ceb50f717e036013d5f6e6a497ae3aaef585f7121bcf4e" }, "ltx2/chained-segments.json": { "chained_image_to_video": "3a36def4e09fe10832bc1821c7cb6fd82d9e050ba149e5b6e7afa6ec66f60f68" @@ -712,7 +727,7 @@ def test_normalized_downstream_false_when_nothing_references_it(): }, "ltx2/generative-upscale.json": { "low_resolution": "fab6e3ec8feb0909dc44660990636418bdccbbb045154e1a56b26f235435acf0", - "upscaled": "eb12a8878241511cecd56d4c41a168e9b59a11803656bef3948c704d18fba112", + "upscaled": "0e3eeb97dc88f6a3ab74ea3cc6439fd266ade29d55459af400b905751777d2c3", }, "ltx2/image-to-video.json": { "image_to_video": "3a36def4e09fe10832bc1821c7cb6fd82d9e050ba149e5b6e7afa6ec66f60f68" @@ -721,7 +736,7 @@ def test_normalized_downstream_false_when_nothing_references_it(): "keyframes_to_video": "ea27aeb8d07376a1295ea2d102ddc61a8d64fa70aeae26e02f6f48e825a6d9ee" }, "ltx2/reference-sheet.json": { - "shot": "f561e8c08eb096c02bcd252244e6dc94c6eb82f133fea559221c4acce364b3f8" + "shot": "c3a7874569313ba50b184de0f22a52dc3f42f34fa5e53831c3586d948bf94372" }, "ltx2/refine-clip.json": { "refine": "cdc51f4b605ee60ec9359390b4a9fb7f60762b419637699527f5a7b1457c1f22", @@ -731,15 +746,15 @@ def test_normalized_downstream_false_when_nothing_references_it(): "refine": "4b54a53f5c05de50d0b3094addf68b2d843135fdc02518b0ab02fb7e225b2552" }, "ltx2/restore-deblur.json": { - "restored": "46a91b9693a5d42b7e43229989cfff47d7f58c26cec910ed3d18afcb06257396" + "restored": "676c6d3535647dca49d39838de6cf5962c1c817eefb4c750f73be169fd60aa78" }, "ltx2/restore-decompression.json": { - "restored": "afcb2e67a7c03447b5860bfe19f4785480e214eb54a39a61a53ccbf0e51717f6" + "restored": "ad261336e57932166606b7e3897163eb085ca8d85f6f855965b5def21950f581" }, "ltx2/restore-long.json": { - "restore@w0": "46a91b9693a5d42b7e43229989cfff47d7f58c26cec910ed3d18afcb06257396", - "restore@w1": "46a91b9693a5d42b7e43229989cfff47d7f58c26cec910ed3d18afcb06257396", - "restore@w2": "46a91b9693a5d42b7e43229989cfff47d7f58c26cec910ed3d18afcb06257396", + "restore@w0": "676c6d3535647dca49d39838de6cf5962c1c817eefb4c750f73be169fd60aa78", + "restore@w1": "676c6d3535647dca49d39838de6cf5962c1c817eefb4c750f73be169fd60aa78", + "restore@w2": "676c6d3535647dca49d39838de6cf5962c1c817eefb4c750f73be169fd60aa78", }, "ltx2/text-to-video.json": { "text_to_video": "c4da9234234392f4df1a76b9d3ffb2cbff735540e278709da11e43f9b4bc7313" @@ -750,96 +765,96 @@ def test_normalized_downstream_false_when_nothing_references_it(): "upscale": "caa78ea250fb765b6aff640c03219cfbc3a9c02366fd31997776520a20c8ee75", }, "ltx2/upscale-clip.json": { - "upscaled": "0f0195201e3be833527a0bc1ef04d29a1dd048a8877678ab1769f663fd27575f" + "upscaled": "94618b4210efdaaf17afda812fbcfd9566bdaba1220a5b2518b178a154a9677c" }, "minimax/chain-matched-and-aligned.json": { - "audio_aligned_chained_reference_to_video_audio": "29913e317c4b9d427e617b120cec9213fc99a8b5c859f2d173372d11788a88e8" + "audio_aligned_chained_reference_to_video_audio": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6" }, "minimax/chain-matched-to-audio.json": { - "chained_reference_to_video_audio": "1bce70dfbd9418b3dbe7ca8abd6f109dd8cdaffbdc802edce327454f36c3f600" + "chained_reference_to_video_audio": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6" }, "minimax/chain-video-continuity.json": { - "chained_reference_to_video_audio": "1bce70dfbd9418b3dbe7ca8abd6f109dd8cdaffbdc802edce327454f36c3f600" + "chained_reference_to_video_audio": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6" }, "minimax/chained-segments.json": { - "chained_keyframe_to_video_audio": "2e770f1c5be668ab4af2174502107bf82cb2bcea5d500f31e5f231fe149cffe9" + "chained_keyframe_to_video_audio": "5fb73e7ccf6a8526a42eb04c4ac2fe777ed528e3e3305c787d4d48abacb158c8" }, "minimax/composable-references.json": { - "video_reference_to_video_audio": "1bce70dfbd9418b3dbe7ca8abd6f109dd8cdaffbdc802edce327454f36c3f600" + "video_reference_to_video_audio": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6" }, "minimax/dialogue-short.json": { - "draw_character_a": "65e14c94aa1bfa3a41f44e1a278ac84b09804577b2cf78c1a5a066366db5e20f", - "draw_character_b": "65e14c94aa1bfa3a41f44e1a278ac84b09804577b2cf78c1a5a066366db5e20f", - "shot@button": "29913e317c4b9d427e617b120cec9213fc99a8b5c859f2d173372d11788a88e8", - "shot@cold_open": "29913e317c4b9d427e617b120cec9213fc99a8b5c859f2d173372d11788a88e8", - "shot@deflect": "29913e317c4b9d427e617b120cec9213fc99a8b5c859f2d173372d11788a88e8", - "shot@react": "29913e317c4b9d427e617b120cec9213fc99a8b5c859f2d173372d11788a88e8", - "shot@tag": "29913e317c4b9d427e617b120cec9213fc99a8b5c859f2d173372d11788a88e8", + "draw_character_a": "a3861857f0d412d15cd569c065aac00fcf7c0af9c2e68a7944877c1eb0f796bf", + "draw_character_b": "a3861857f0d412d15cd569c065aac00fcf7c0af9c2e68a7944877c1eb0f796bf", + "shot@button": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6", + "shot@cold_open": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6", + "shot@deflect": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6", + "shot@react": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6", + "shot@tag": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6", }, "minimax/enhance-prompt-with-image.json": { - "keyframe_to_video_audio": "2e770f1c5be668ab4af2174502107bf82cb2bcea5d500f31e5f231fe149cffe9" + "keyframe_to_video_audio": "5fb73e7ccf6a8526a42eb04c4ac2fe777ed528e3e3305c787d4d48abacb158c8" }, "minimax/enhance-prompt.json": { - "text_to_video_audio": "26a7a9b6677e0dc8c5a8ed0a01c4e6017be2923f79b5ab016a620b7fa7a60d7c" + "text_to_video_audio": "68977225a70231ba3dd93248624501e4d8335c39eed9d4b2f0c970037af89bac" }, "minimax/first-and-last-frame.json": { - "first_and_last_frame_to_video_audio": "2e770f1c5be668ab4af2174502107bf82cb2bcea5d500f31e5f231fe149cffe9" + "first_and_last_frame_to_video_audio": "5fb73e7ccf6a8526a42eb04c4ac2fe777ed528e3e3305c787d4d48abacb158c8" }, "minimax/generated-subject-reference.json": { - "draw_subject": "65e14c94aa1bfa3a41f44e1a278ac84b09804577b2cf78c1a5a066366db5e20f", - "reference_to_video_audio": "1bce70dfbd9418b3dbe7ca8abd6f109dd8cdaffbdc802edce327454f36c3f600", + "draw_subject": "a3861857f0d412d15cd569c065aac00fcf7c0af9c2e68a7944877c1eb0f796bf", + "reference_to_video_audio": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6", }, "minimax/image-to-video.json": { - "keyframe_to_video_audio": "2e770f1c5be668ab4af2174502107bf82cb2bcea5d500f31e5f231fe149cffe9" + "keyframe_to_video_audio": "5fb73e7ccf6a8526a42eb04c4ac2fe777ed528e3e3305c787d4d48abacb158c8" }, "minimax/last-frame-only.json": { - "last_frame_to_video_audio": "2e770f1c5be668ab4af2174502107bf82cb2bcea5d500f31e5f231fe149cffe9" + "last_frame_to_video_audio": "5fb73e7ccf6a8526a42eb04c4ac2fe777ed528e3e3305c787d4d48abacb158c8" }, "minimax/music-video-cuts.json": {}, "minimax/music-video.json": { - "draw_singer": "65e14c94aa1bfa3a41f44e1a278ac84b09804577b2cf78c1a5a066366db5e20f", - "shot@closeup": "29913e317c4b9d427e617b120cec9213fc99a8b5c859f2d173372d11788a88e8", - "shot@finale": "29913e317c4b9d427e617b120cec9213fc99a8b5c859f2d173372d11788a88e8", - "shot@room": "29913e317c4b9d427e617b120cec9213fc99a8b5c859f2d173372d11788a88e8", - "shot@wide_open": "29913e317c4b9d427e617b120cec9213fc99a8b5c859f2d173372d11788a88e8", + "draw_singer": "a3861857f0d412d15cd569c065aac00fcf7c0af9c2e68a7944877c1eb0f796bf", + "shot@closeup": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6", + "shot@finale": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6", + "shot@room": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6", + "shot@wide_open": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6", "write_song": "1abd10e3c899fb5bbc490baa64467ed46d5005f78fa8f94bc777f3ddcaf81fb5", }, "minimax/music.json": { "generate_music": "1abd10e3c899fb5bbc490baa64467ed46d5005f78fa8f94bc777f3ddcaf81fb5" }, "minimax/reference-to-video.json": { - "reference_to_video_audio": "1bce70dfbd9418b3dbe7ca8abd6f109dd8cdaffbdc802edce327454f36c3f600" + "reference_to_video_audio": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6" }, "minimax/shots-batch.json": { - "shot@shot_1": "26a7a9b6677e0dc8c5a8ed0a01c4e6017be2923f79b5ab016a620b7fa7a60d7c", - "shot@shot_2": "26a7a9b6677e0dc8c5a8ed0a01c4e6017be2923f79b5ab016a620b7fa7a60d7c", - "shot@shot_3": "26a7a9b6677e0dc8c5a8ed0a01c4e6017be2923f79b5ab016a620b7fa7a60d7c", - "shot@shot_4": "26a7a9b6677e0dc8c5a8ed0a01c4e6017be2923f79b5ab016a620b7fa7a60d7c", - "shot@shot_5": "26a7a9b6677e0dc8c5a8ed0a01c4e6017be2923f79b5ab016a620b7fa7a60d7c", + "shot@shot_1": "68977225a70231ba3dd93248624501e4d8335c39eed9d4b2f0c970037af89bac", + "shot@shot_2": "68977225a70231ba3dd93248624501e4d8335c39eed9d4b2f0c970037af89bac", + "shot@shot_3": "68977225a70231ba3dd93248624501e4d8335c39eed9d4b2f0c970037af89bac", + "shot@shot_4": "68977225a70231ba3dd93248624501e4d8335c39eed9d4b2f0c970037af89bac", + "shot@shot_5": "68977225a70231ba3dd93248624501e4d8335c39eed9d4b2f0c970037af89bac", }, "minimax/storyboard.json": { - "board_1_launch": "65e14c94aa1bfa3a41f44e1a278ac84b09804577b2cf78c1a5a066366db5e20f", - "board_2_gutter": "65e14c94aa1bfa3a41f44e1a278ac84b09804577b2cf78c1a5a066366db5e20f", - "board_3_shore": "65e14c94aa1bfa3a41f44e1a278ac84b09804577b2cf78c1a5a066366db5e20f", - "voyage": "ee8e8884e8f049bf7de6aabe053a8c389ab6cd0c42108a04cc6fd1049e671300", + "board_1_launch": "a3861857f0d412d15cd569c065aac00fcf7c0af9c2e68a7944877c1eb0f796bf", + "board_2_gutter": "a3861857f0d412d15cd569c065aac00fcf7c0af9c2e68a7944877c1eb0f796bf", + "board_3_shore": "a3861857f0d412d15cd569c065aac00fcf7c0af9c2e68a7944877c1eb0f796bf", + "voyage": "0f30a4f21253ffdec029b3e2523c077f68629579bb3ae57cc69c09b638665987", }, "minimax/video-with-audio-768p.json": { - "text_to_video_audio": "26a7a9b6677e0dc8c5a8ed0a01c4e6017be2923f79b5ab016a620b7fa7a60d7c" + "text_to_video_audio": "68977225a70231ba3dd93248624501e4d8335c39eed9d4b2f0c970037af89bac" }, "minimax/video-with-audio.json": { - "text_to_video_audio": "26a7a9b6677e0dc8c5a8ed0a01c4e6017be2923f79b5ab016a620b7fa7a60d7c" + "text_to_video_audio": "68977225a70231ba3dd93248624501e4d8335c39eed9d4b2f0c970037af89bac" }, "minimax/voice-timbre-reference.json": { - "reference_to_video_audio": "1bce70dfbd9418b3dbe7ca8abd6f109dd8cdaffbdc802edce327454f36c3f600" + "reference_to_video_audio": "95c8ba90efeba9603384619f1db3f275251f078370add6b2a3e7dab8481aacc6" }, "multi-image-reference.json": { "txt2img": "6e563b308f1508e0f24a660dcf25636774e2b70c4ccc723e8975421a7224b48d" }, "outpaint.json": { - "outpaint": "a52f958d8c596182fb83182bad33df3da491b63a3ad758ec5807103b401f92eb" + "outpaint": "8ece64de576884be947f5af07c5713b92d3a1e641035434ef805629f12ae8c77" }, "prompt-weighting.json": { - "txt2img": "74fd88b30ffbf3cd95b1b971e90ce8da33e7370a160e838478fc5d5a6d3bc9df" + "txt2img": "745df144eb83ac12caa01de34dc7238b8595b28b70db7c960f4d1c0366b2f2f4" }, "qr-code.json": { "main": "f52e468931a9e9d546696cb9fccc365e7347b47e07730470b7a3ac419722fce5" @@ -850,11 +865,11 @@ def test_normalized_downstream_false_when_nothing_references_it(): "generate": "8655c1f1635a7f9a1d3ce02a9cbf6f0f9e8b27130052d6e2758fcb7838270c14" }, "segment-and-inpaint.json": { - "inpaint": "a52f958d8c596182fb83182bad33df3da491b63a3ad758ec5807103b401f92eb" + "inpaint": "8ece64de576884be947f5af07c5713b92d3a1e641035434ef805629f12ae8c77" }, "segment.json": {}, "step-caching.json": { - "txt2img": "69c2b1c0bf23b291404b655919a94e7ed2b766d962dc209ca1e5df86f619a9dd" + "txt2img": "3e3ea6a2cb1d2d1b1abaec9de159758b137ac0bcd3197a1bac0f7bbcead7ac24" }, "sub-workflow.json": {}, "surface-normals.json": { @@ -1011,3 +1026,104 @@ def test_level_spread_warning_suppressed_when_audio_replaced(monkeypatch): seen.append(events) assert any("level_spread" in str(e) for e in seen[0]) assert not any("level_spread" in str(e) for e in seen[1]) + + +def _definition(**overrides): + base = { + "configuration": {"component_type": "FluxPipeline"}, + "from_pretrained_arguments": {"model_name": "black-forest-labs/FLUX.1-dev"}, + "loras": [ + { + "model_name": "a/b", + "weight_name": "x.safetensors", + "scale": 1.0, + "alpha": 16, + }, + ], + "scheduler": {"shift": 6.0}, + "audio_scheduler": {"shift": 3.0}, + "arguments": {"prompt": "p"}, + } + base.update(overrides) + return base + + +def test_lora_scale_and_alpha_do_not_change_the_pipeline_key(): + a = _definition() + b = _definition( + loras=[ + { + "model_name": "a/b", + "weight_name": "x.safetensors", + "scale": 0.5, + "alpha": 8, + } + ] + ) + assert pipeline_cache_key(a) == pipeline_cache_key(b) + + +def test_scheduler_shift_does_not_change_the_pipeline_key(): + a = _definition() + b = _definition(scheduler={"shift": 12.0}, audio_scheduler={"shift": 1.0}) + assert pipeline_cache_key(a) == pipeline_cache_key(b) + + +def test_lora_model_name_still_changes_the_pipeline_key(): + a = _definition() + b = _definition( + loras=[{"model_name": "a/c", "weight_name": "x.safetensors", "scale": 1.0}] + ) + assert pipeline_cache_key(a) != pipeline_cache_key(b) + + +def test_scheduler_type_still_changes_the_pipeline_key(): + a = _definition(scheduler={"scheduler_type": "X", "shift": 6.0}) + b = _definition(scheduler={"scheduler_type": "Y", "shift": 6.0}) + assert pipeline_cache_key(a) != pipeline_cache_key(b) + + +def test_weights_identity_leaves_the_definition_untouched(): + definition = _definition() + weights_identity(definition) + assert definition["loras"][0]["scale"] == 1.0 + assert definition["scheduler"]["shift"] == 6.0 + + +def test_lora_alpha_value_change_keeps_the_pipeline_key(): + a = _definition() + b = _definition( + loras=[{"model_name": "a/b", "weight_name": "x.safetensors", "alpha": 8}] + ) + assert pipeline_cache_key(a) == pipeline_cache_key(b) + + +def test_lora_alpha_removed_changes_the_pipeline_key(): + """Nothing re-applies the checkpoint's own alpha on a hit, so dropping + the override has to reload.""" + a = _definition() + b = _definition( + loras=[{"model_name": "a/b", "weight_name": "x.safetensors", "alpha": None}] + ) + c = _definition(loras=[{"model_name": "a/b", "weight_name": "x.safetensors"}]) + assert pipeline_cache_key(a) != pipeline_cache_key(b) + assert pipeline_cache_key(b) == pipeline_cache_key(c) + + +@pytest.mark.parametrize("scheduler", ["scheduler", "audio_scheduler"]) +def test_scheduler_shift_removed_changes_the_pipeline_key(scheduler): + a = _definition() + b = _definition(**{scheduler: {"shift": None}}) + c = _definition(**{scheduler: {}}) + assert pipeline_cache_key(a) != pipeline_cache_key(b) + assert pipeline_cache_key(b) == pipeline_cache_key(c) + + +def test_pipeline_definition_key_moves_with_runtime_settings(): + from dw.step_cache import pipeline_definition_key + + a = _definition() + b = _definition(scheduler={"shift": 12.0}) + c = _definition(arguments={"prompt": "other"}, seed=7) + assert pipeline_definition_key(a) != pipeline_definition_key(b) + assert pipeline_definition_key(a) == pipeline_definition_key(c) diff --git a/tests/test_worker_execute.py b/tests/test_worker_execute.py index 81e604b7..e3966634 100644 --- a/tests/test_worker_execute.py +++ b/tests/test_worker_execute.py @@ -109,7 +109,52 @@ def test_workflow_switch_evicts_cache_and_untouched_keys_dropped(): StubWorkflow(), command=snapshot_command("/w/other.json"), ) - cleanup.assert_called_once() + cleanup.assert_called_once_with(keep=set()) + + +def test_workflow_switch_keeps_pipelines_the_new_workflow_loads(): + class TouchesShared(StubWorkflow): + def run(self, arguments, previous_pipelines=None, context=None, **kwargs): + # The new workflow uses the warm pipeline, so the post-run + # eviction of untouched keys leaves it alone + assert "shared-key" in previous_pipelines + context.touch_pipeline("shared-key") + return [] + + worker = _make_worker() + _execute(worker, StubWorkflow()) + worker.loaded_pipelines["shared-key"] = object() + worker.loaded_pipelines["old-only-key"] = object() + + with patch.object(workflow_run, "pipeline_keys", return_value={"shared-key"}): + messages = _execute( + worker, TouchesShared(), command=snapshot_command("/w/other.json") + ) + + assert "shared-key" in worker.loaded_pipelines + assert "old-only-key" not in worker.loaded_pipelines + outputs = [m["message"] for m in messages if m["type"] == "output"] + assert any("keeping 1 warm" in m for m in outputs) + + +def test_workflow_switch_falls_back_to_a_full_release_when_keys_cannot_be_taken(): + worker = _make_worker() + _execute(worker, StubWorkflow()) + worker.loaded_pipelines["shared-key"] = object() + + with patch.object( + workflow_run, "pipeline_keys", side_effect=ValueError("bad variable") + ): + _execute(worker, StubWorkflow(), command=snapshot_command("/w/other.json")) + + assert worker.loaded_pipelines == {} + + +def test_full_cleanup_keep_spares_named_keys_only(): + worker = _make_worker() + worker.loaded_pipelines = {"a": object(), "b": object()} + worker._cleanup_all(keep={"a"}) + assert set(worker.loaded_pipelines) == {"a"} def test_inline_workflow_definition_executes(tmp_path): diff --git a/tests/test_workflow_run_keys.py b/tests/test_workflow_run_keys.py new file mode 100644 index 00000000..bd612a5f --- /dev/null +++ b/tests/test_workflow_run_keys.py @@ -0,0 +1,42 @@ +"""pipeline_keys() answers the worker's question on a workflow switch: which +resident pipelines does the next workflow load anyway.""" + +from dw import workflow_run as workflow_run_module +from dw.step_cache import step_cache +from tests.test_workflow_step_cache import build_test_workflow_and_call_count_spy + + +def test_pipeline_keys_match_what_a_run_keys_on(tmp_path): + step_cache.clear() + workflow, _ = build_test_workflow_and_call_count_spy(str(tmp_path)) + try: + pipelines = {} + workflow.run({}, pipelines) + assert pipelines + assert workflow_run_module.pipeline_keys(workflow, {}) == set(pipelines) + finally: + for p in workflow._test_patcher: + p.stop() + + +def test_an_unseeded_workflow_still_answers(tmp_path): + step_cache.clear() + workflow, _ = build_test_workflow_and_call_count_spy(str(tmp_path)) + del workflow.workflow_definition["seed"] + try: + assert workflow_run_module.pipeline_keys(workflow, {}) + finally: + for p in workflow._test_patcher: + p.stop() + + +def test_the_probe_writes_nothing(tmp_path): + step_cache.clear() + workflow, call_count = build_test_workflow_and_call_count_spy(str(tmp_path)) + try: + workflow_run_module.pipeline_keys(workflow, {}) + assert list(tmp_path.iterdir()) == [] + assert call_count() == 0 + finally: + for p in workflow._test_patcher: + p.stop() diff --git a/tests/test_workflow_step_cache.py b/tests/test_workflow_step_cache.py index b38cf499..fb493437 100644 --- a/tests/test_workflow_step_cache.py +++ b/tests/test_workflow_step_cache.py @@ -13,12 +13,26 @@ from dw.events import RunContext from dw.step_cache import step_cache from dw.step import Step +from dw.pipeline_processors import adapters from dw.pipeline_processors.config_objects import get_group_offload_configuration from dw.pipeline_processors.pipeline import Pipeline from dw import workflow_run as workflow_run_module from dw.workflow import Workflow +@pytest.fixture(autouse=True) +def _scales_through_the_mock(): + """The warm path sets LoRA scales layer by layer (set_adapter_scales); + the mocked models here have no peft layers, so route it to the mock's + set_adapters, which is what these tests read back.""" + with patch.object( + adapters, + "set_adapter_scales", + lambda model, names, weights: model.set_adapters(names, weights), + ): + yield + + def _mock_pipeline_load(self, shared_components): self.pipeline = MagicMock() @@ -1152,3 +1166,109 @@ def test_a_change_to_an_intermediate_borrower_misses_every_step_below_it( [{"model_b": "model-b"}, {"model_b": "model-b2"}], ) assert executed[1] == ["B", "D"] + + +def _pipeline_reference_with_runtime_settings(): + definition = _pipeline_reference_workflow_def() + definition["variables"].update({"lora_scale": 1.0, "shift": 6.0}) + source = definition["steps"][0]["pipeline"] + source["loras"] = [{"model_name": "a/b", "scale": "variable:lora_scale"}] + source["scheduler"] = {"shift": "variable:shift"} + return definition + + +@pytest.mark.parametrize( + "changed", [{"lora_scale": 0.5}, {"shift": 9.0}], ids=["lora_scale", "shift"] +) +def test_a_step_borrowing_a_pipeline_misses_when_the_source_runtime_setting_changes( + tmp_path, changed +): + """The source keeps its warm pipeline - its scale or shift is applied in + place - but what the borrower produces changes, so its entry must miss.""" + first = {"lora_scale": 1.0, "shift": 6.0, "prompt_b": "x"} + order = _run_twice_recording_order( + tmp_path, + _pipeline_reference_with_runtime_settings(), + [first, {**first, **changed}], + ) + assert order == ["A", "B", "A", "B"] + + +def test_a_step_borrowing_a_pipeline_hits_when_the_source_runtime_settings_hold( + tmp_path, +): + arguments = {"lora_scale": 0.5, "shift": 9.0, "prompt_b": "x"} + order = _run_twice_recording_order( + tmp_path, _pipeline_reference_with_runtime_settings(), [arguments, arguments] + ) + assert order == ["A", "B"] + + +def test_a_pipeline_reference_runs_with_the_referenced_steps_runtime_settings( + tmp_path, +): + """A and B differ only in scale and shift, so they share one warm model; + B's hit applies its own values to it. C references A and must run at + A's values, not at whatever B left on the shared model.""" + step_cache.clear() + + def pipeline(scale, shift): + return { + "configuration": {"component_type": "{MockPipeline}"}, + "from_pretrained_arguments": {"model_name": "model-a"}, + "loras": [{"model_name": "a/b", "scale": scale}], + "scheduler": {"shift": shift}, + "arguments": {"prompt": "p"}, + } + + definition = { + "id": "test_pipeline_reference_runtime_settings", + "seed": 42, + "steps": [ + {"name": "A", "pipeline": pipeline(1.0, 6.0)}, + # Saves, so elision keeps it though nothing reads its result + { + "name": "B", + "pipeline": pipeline(0.5, 9.0), + "result": {"content_type": "image/jpeg"}, + }, + { + "name": "C", + "pipeline_reference": { + "reference_name": "A", + "arguments": {"prompt": "q"}, + }, + }, + ], + } + workflow = Workflow(definition, str(tmp_path), "test.json") + seen = {} + + def fake_step_run(self, previous_results, previous_pipelines, step_action): + model = step_action.pipeline + seen[self.name] = ( + model.set_adapters.call_args, + model.scheduler.set_shift.call_args, + ) + return FakeResult() + + patchers = [ + patch.object(Step, "run", fake_step_run), + patch.object(Pipeline, "load", _mock_pipeline_load), + ] + for p in patchers: + p.start() + try: + workflow.run({}) + finally: + for p in patchers: + p.stop() + + assert workflow.pipeline_ownership.key_for( + "A" + ) == workflow.pipeline_ownership.key_for("B") + adapters, shift = seen["B"] + assert adapters.args == (["0"], [0.5]) and shift.args == (9.0,) + adapters, shift = seen["C"] + assert adapters.args == (["0"], [1.0]) + assert shift.args == (6.0,) diff --git a/ui/e2e/serve_fixture.py b/ui/e2e/serve_fixture.py index 38ede9a6..b6886da2 100644 --- a/ui/e2e/serve_fixture.py +++ b/ui/e2e/serve_fixture.py @@ -8,7 +8,11 @@ The prompt directory is passed explicitly: the server would otherwise discover ./prompts in the working directory - the real library. The scratch directory is also the workspace, so the asset library the server creates and -uploads into lands there rather than in the checkout. +uploads into lands there rather than in the checkout. The settings root +(DIFFUSERS_HELPER_ROOT) is scratch too: it holds the job history the plan +prices from, and the developer's real one gave a box that had run a catalog +model an inherited-cost warning CI never sees, so a spec expecting a clean +validation passed in CI and failed locally. Extra arguments (--port ...) pass through to dw.serve. The process replaces itself with the server so Playwright's shutdown signal reaches it directly. @@ -48,6 +52,8 @@ with open(os.path.join(run_dir, "still.png"), "wb") as handle: handle.write(ONE_PIXEL_PNG) +os.environ["DIFFUSERS_HELPER_ROOT"] = os.path.join(scratch, "helper") + os.chdir(root) os.execv( sys.executable, diff --git a/workflows/models/flux-krea.json b/workflows/models/flux-krea.json index 9ab6780c..59578615 100644 --- a/workflows/models/flux-krea.json +++ b/workflows/models/flux-krea.json @@ -18,7 +18,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-Krea-dev", diff --git a/workflows/models/z-image.json b/workflows/models/z-image.json index 782cb562..9c968927 100644 --- a/workflows/models/z-image.json +++ b/workflows/models/z-image.json @@ -23,7 +23,7 @@ "pipeline": { "configuration": { "component_type": "ZImagePipeline", - "offload": "sequential" + "offload": "model" }, "from_pretrained_arguments": { "model_name": "Tongyi-MAI/Z-Image-Turbo", diff --git a/workflows/templates/compose-workflows.json b/workflows/templates/compose-workflows.json index 8a633ba2..135d7c75 100644 --- a/workflows/templates/compose-workflows.json +++ b/workflows/templates/compose-workflows.json @@ -7,7 +7,7 @@ "image_prompt": "portrait | wide angle shot of eyes off to one side of frame, lucid dream-like 3d model of a marmot, game asset, blender, looking off in distance ::8 style | glowing ::8 background | forest, vivid neon wonderland, particles, blue, green, orange ::7 parameters | rule of thirds, golden ratio, asymmetric composition, hyper- maximalist, octane render, photorealism, cinematic realism, unreal engine, 8k ::7 --ar 16:9 --s 1000", "video_generation_workflow": "./minimax/image-to-video.json", "video_prompt": "the marmot blinks and looks around", - "video_num_inference_steps": 20, + "video_num_inference_steps": 5, "seed": null }, "id": "txt2img2vid", diff --git a/workflows/templates/consistent-set.json b/workflows/templates/consistent-set.json index 7968b3e0..3d6e9b61 100644 --- a/workflows/templates/consistent-set.json +++ b/workflows/templates/consistent-set.json @@ -24,7 +24,19 @@ "pipeline": { "configuration": { "component_type": "FluxKontextPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-Kontext-dev", @@ -49,7 +61,19 @@ "pipeline": { "configuration": { "component_type": "FluxKontextPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-Kontext-dev", @@ -73,7 +97,19 @@ "pipeline": { "configuration": { "component_type": "FluxKontextPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-Kontext-dev", @@ -97,7 +133,19 @@ "pipeline": { "configuration": { "component_type": "FluxKontextPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-Kontext-dev", diff --git a/workflows/templates/controlnet.json b/workflows/templates/controlnet.json index 82567537..8c68b2a7 100644 --- a/workflows/templates/controlnet.json +++ b/workflows/templates/controlnet.json @@ -31,7 +31,19 @@ "pipeline": { "configuration": { "component_type": "FluxControlPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-Canny-dev", diff --git a/workflows/templates/embed-metadata.json b/workflows/templates/embed-metadata.json index 7f1419a5..bbd4a7b7 100644 --- a/workflows/templates/embed-metadata.json +++ b/workflows/templates/embed-metadata.json @@ -12,7 +12,19 @@ "name": "generate", "pipeline": { "configuration": { - "offload": "sequential", + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + }, "component_type": "FluxPipeline" }, "from_pretrained_arguments": { diff --git a/workflows/templates/image-edit.json b/workflows/templates/image-edit.json index a6caca8a..5e19393b 100644 --- a/workflows/templates/image-edit.json +++ b/workflows/templates/image-edit.json @@ -15,7 +15,19 @@ "pipeline": { "configuration": { "component_type": "FluxKontextPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-Kontext-dev", diff --git a/workflows/templates/image-to-image.json b/workflows/templates/image-to-image.json index 83be81ac..2eda6778 100644 --- a/workflows/templates/image-to-image.json +++ b/workflows/templates/image-to-image.json @@ -21,7 +21,19 @@ "pipeline": { "configuration": { "component_type": "FluxImg2ImgPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", diff --git a/workflows/templates/image-variation.json b/workflows/templates/image-variation.json index a86c3efe..9b3deb71 100644 --- a/workflows/templates/image-variation.json +++ b/workflows/templates/image-variation.json @@ -32,7 +32,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", diff --git a/workflows/templates/inpaint.json b/workflows/templates/inpaint.json index 691681eb..4eea1907 100644 --- a/workflows/templates/inpaint.json +++ b/workflows/templates/inpaint.json @@ -13,7 +13,19 @@ "pipeline": { "configuration": { "component_type": "FluxFillPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-Fill-dev", diff --git a/workflows/templates/ip-adapter.json b/workflows/templates/ip-adapter.json index 27b9eae9..9f728a01 100644 --- a/workflows/templates/ip-adapter.json +++ b/workflows/templates/ip-adapter.json @@ -12,7 +12,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", diff --git a/workflows/templates/lora-styles.json b/workflows/templates/lora-styles.json index d928a390..225409f7 100644 --- a/workflows/templates/lora-styles.json +++ b/workflows/templates/lora-styles.json @@ -18,7 +18,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", @@ -51,7 +63,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", @@ -84,7 +108,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", @@ -117,7 +153,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", @@ -150,7 +198,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", @@ -183,7 +243,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", @@ -216,7 +288,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", @@ -249,7 +333,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", @@ -282,7 +378,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", @@ -315,7 +423,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-dev", diff --git a/workflows/templates/lora.json b/workflows/templates/lora.json index 3929fcfc..93af491f 100644 --- a/workflows/templates/lora.json +++ b/workflows/templates/lora.json @@ -16,7 +16,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential", + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + }, "vae": { "enable_slicing": true, "enable_tiling": true diff --git a/workflows/templates/minimax/chain-matched-to-audio.json b/workflows/templates/minimax/chain-matched-to-audio.json index 31b9769a..a38fb395 100644 --- a/workflows/templates/minimax/chain-matched-to-audio.json +++ b/workflows/templates/minimax/chain-matched-to-audio.json @@ -68,10 +68,6 @@ "pre_load_modules": [ "sdnq" ], - "cache": { - "type": "first_block", - "threshold": 0.1 - }, "load_components": { "dtype": "torch.bfloat16", "quantization_config": { diff --git a/workflows/templates/minimax/chain-video-continuity.json b/workflows/templates/minimax/chain-video-continuity.json index 1104a5cd..05f8cfd5 100644 --- a/workflows/templates/minimax/chain-video-continuity.json +++ b/workflows/templates/minimax/chain-video-continuity.json @@ -71,10 +71,6 @@ "pre_load_modules": [ "sdnq" ], - "cache": { - "type": "first_block", - "threshold": 0.1 - }, "load_components": { "dtype": "torch.bfloat16", "quantization_config": { diff --git a/workflows/templates/minimax/composable-references.json b/workflows/templates/minimax/composable-references.json index b6202d5b..0a8c9b9b 100644 --- a/workflows/templates/minimax/composable-references.json +++ b/workflows/templates/minimax/composable-references.json @@ -67,10 +67,6 @@ "pre_load_modules": [ "sdnq" ], - "cache": { - "type": "first_block", - "threshold": 0.1 - }, "load_components": { "dtype": "torch.bfloat16", "quantization_config": { diff --git a/workflows/templates/minimax/dialogue-short.json b/workflows/templates/minimax/dialogue-short.json index e6ab36c3..6fc3fd58 100644 --- a/workflows/templates/minimax/dialogue-short.json +++ b/workflows/templates/minimax/dialogue-short.json @@ -171,7 +171,7 @@ "pipeline": { "configuration": { "component_type": "ZImagePipeline", - "offload": "sequential" + "offload": "model" }, "from_pretrained_arguments": { "model_name": "Tongyi-MAI/Z-Image-Turbo", @@ -199,7 +199,7 @@ "pipeline": { "configuration": { "component_type": "ZImagePipeline", - "offload": "sequential" + "offload": "model" }, "from_pretrained_arguments": { "model_name": "Tongyi-MAI/Z-Image-Turbo", diff --git a/workflows/templates/minimax/generated-subject-reference.json b/workflows/templates/minimax/generated-subject-reference.json index 7ee9c821..97e0226c 100644 --- a/workflows/templates/minimax/generated-subject-reference.json +++ b/workflows/templates/minimax/generated-subject-reference.json @@ -64,7 +64,7 @@ "pipeline": { "configuration": { "component_type": "ZImagePipeline", - "offload": "sequential" + "offload": "model" }, "from_pretrained_arguments": { "model_name": "Tongyi-MAI/Z-Image-Turbo", @@ -93,10 +93,6 @@ "pre_load_modules": [ "sdnq" ], - "cache": { - "type": "first_block", - "threshold": 0.1 - }, "load_components": { "dtype": "torch.bfloat16", "quantization_config": { diff --git a/workflows/templates/minimax/image-to-video.json b/workflows/templates/minimax/image-to-video.json index ca90fa79..6f219b23 100644 --- a/workflows/templates/minimax/image-to-video.json +++ b/workflows/templates/minimax/image-to-video.json @@ -1,6 +1,6 @@ { "id": "MiniMaxH3I2V", - "description": "Conditioning generation on a supplied still. The same model is loaded under its 'fl2va' workflow and given only an 'image', so the picture becomes the video's first frame - the frame-conditioned form of the baseline. Also shows the acceleration that applies to any of these pipelines: a step-distillation LoRA, which composes with per-component quantization and offloading to cut the step count. The 'first_block' cache is deliberately absent - measured on the 9-step turbo schedule (and fewer steps on today's 4-step one only widens the gap), consecutive steps differ too much for it to ever skip, so it only held cache state for nothing; the 20-step ref2va examples are where it can engage.", + "description": "Conditioning generation on a supplied still. The same model is loaded under its 'fl2va' workflow and given only an 'image', so the picture becomes the video's first frame - the frame-conditioned form of the baseline. Also shows the acceleration that applies to any of these pipelines: a step-distillation LoRA, which composes with per-component quantization and offloading to cut the step count. The 'first_block' cache is deliberately absent - measured on the 9-step turbo schedule (and fewer steps on today's 4-step one only widens the gap), consecutive steps differ too much for it to ever skip, so it only held cache state for nothing; it only pays off on a 20-step undistilled ref2va schedule, which no shipped template runs.", "summary": "H3 video with audio from a supplied still as first frame, accelerated by a step-distillation LoRA.", "cost_drivers": [ "num_frames", @@ -25,7 +25,8 @@ "lora_model_name": "lightx2v/Minimax-h3-Turbo", "lora_weight_name": "minimax_h3_fl2v_turbo_4step_v1.2_768p_bf16.safetensors", "lora_adapter_name": "turbo", - "seed": null + "seed": null, + "attention_backend": null }, "variable_constraints": { "num_frames": { @@ -115,6 +116,7 @@ }, "components": { "transformer": { + "attention_backend": "variable:attention_backend", "group_offload": { "offload_type": "block_level", "num_blocks_per_group": 2, diff --git a/workflows/templates/minimax/music-video.json b/workflows/templates/minimax/music-video.json index 30da871b..008add00 100644 --- a/workflows/templates/minimax/music-video.json +++ b/workflows/templates/minimax/music-video.json @@ -120,7 +120,7 @@ "pipeline": { "configuration": { "component_type": "ZImagePipeline", - "offload": "sequential" + "offload": "model" }, "from_pretrained_arguments": { "model_name": "Tongyi-MAI/Z-Image-Turbo", diff --git a/workflows/templates/minimax/reference-to-video.json b/workflows/templates/minimax/reference-to-video.json index 939f4272..9a4dbf84 100644 --- a/workflows/templates/minimax/reference-to-video.json +++ b/workflows/templates/minimax/reference-to-video.json @@ -43,7 +43,8 @@ "lora_weight_name": "minimax_h3_ref2v_turbo_8step_v1.0_768p_bf16.safetensors", "lora_adapter_name": "turbo", "weights_dtype": "{int4}", - "seed": null + "seed": null, + "attention_backend": null }, "variable_constraints": { "num_frames": { @@ -75,10 +76,6 @@ "pre_load_modules": [ "sdnq" ], - "cache": { - "type": "first_block", - "threshold": 0.1 - }, "load_components": { "dtype": "torch.bfloat16", "quantization_config": { @@ -137,6 +134,7 @@ }, "components": { "transformer_ref": { + "attention_backend": "variable:attention_backend", "group_offload": { "offload_type": "block_level", "num_blocks_per_group": 1, diff --git a/workflows/templates/minimax/shots-batch.json b/workflows/templates/minimax/shots-batch.json index 85f09382..80b986c9 100644 --- a/workflows/templates/minimax/shots-batch.json +++ b/workflows/templates/minimax/shots-batch.json @@ -97,7 +97,6 @@ { "name": "shot", "for_each": "variable:shots", - "release_pipeline": true, "pipeline": { "configuration": { "component_type": "ModularPipeline", diff --git a/workflows/templates/minimax/storyboard.json b/workflows/templates/minimax/storyboard.json index 1dc6e88d..7888e7bd 100644 --- a/workflows/templates/minimax/storyboard.json +++ b/workflows/templates/minimax/storyboard.json @@ -74,7 +74,7 @@ "pipeline": { "configuration": { "component_type": "ZImagePipeline", - "offload": "sequential" + "offload": "model" }, "from_pretrained_arguments": { "model_name": "Tongyi-MAI/Z-Image-Turbo", @@ -100,7 +100,7 @@ "pipeline": { "configuration": { "component_type": "ZImagePipeline", - "offload": "sequential" + "offload": "model" }, "from_pretrained_arguments": { "model_name": "Tongyi-MAI/Z-Image-Turbo", @@ -127,7 +127,7 @@ "pipeline": { "configuration": { "component_type": "ZImagePipeline", - "offload": "sequential" + "offload": "model" }, "from_pretrained_arguments": { "model_name": "Tongyi-MAI/Z-Image-Turbo", diff --git a/workflows/templates/minimax/video-with-audio-768p.json b/workflows/templates/minimax/video-with-audio-768p.json index 4eb20dae..77f83b3c 100644 --- a/workflows/templates/minimax/video-with-audio-768p.json +++ b/workflows/templates/minimax/video-with-audio-768p.json @@ -41,7 +41,8 @@ "lora_model_name": "lightx2v/Minimax-h3-Turbo", "lora_weight_name": "minimax_h3_fl2v_turbo_4step_v1.2_768p_bf16.safetensors", "lora_adapter_name": "turbo", - "seed": 42 + "seed": 42, + "attention_backend": null }, "variable_constraints": { "num_frames": { @@ -131,6 +132,7 @@ }, "components": { "transformer": { + "attention_backend": "variable:attention_backend", "group_offload": { "offload_type": "block_level", "num_blocks_per_group": 2, diff --git a/workflows/templates/minimax/video-with-audio.json b/workflows/templates/minimax/video-with-audio.json index de1ddd5a..a0dd1d9a 100644 --- a/workflows/templates/minimax/video-with-audio.json +++ b/workflows/templates/minimax/video-with-audio.json @@ -41,7 +41,8 @@ "lora_model_name": "lightx2v/Minimax-h3-Turbo", "lora_weight_name": "minimax_h3_fl2v_turbo_4step_v1.2_768p_bf16.safetensors", "lora_adapter_name": "turbo", - "seed": 42 + "seed": 42, + "attention_backend": null }, "variable_constraints": { "num_frames": { @@ -131,6 +132,7 @@ }, "components": { "transformer": { + "attention_backend": "variable:attention_backend", "group_offload": { "offload_type": "block_level", "num_blocks_per_group": 2, diff --git a/workflows/templates/minimax/voice-timbre-reference.json b/workflows/templates/minimax/voice-timbre-reference.json index 04182fc1..2e81b91b 100644 --- a/workflows/templates/minimax/voice-timbre-reference.json +++ b/workflows/templates/minimax/voice-timbre-reference.json @@ -83,10 +83,6 @@ "pre_load_modules": [ "sdnq" ], - "cache": { - "type": "first_block", - "threshold": 0.1 - }, "load_components": { "dtype": "torch.bfloat16", "quantization_config": { diff --git a/workflows/templates/outpaint.json b/workflows/templates/outpaint.json index 7d688129..797603fd 100644 --- a/workflows/templates/outpaint.json +++ b/workflows/templates/outpaint.json @@ -25,7 +25,19 @@ "pipeline": { "configuration": { "component_type": "FluxFillPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-Fill-dev", diff --git a/workflows/templates/prompt-weighting.json b/workflows/templates/prompt-weighting.json index 0531edc8..89e70ddd 100644 --- a/workflows/templates/prompt-weighting.json +++ b/workflows/templates/prompt-weighting.json @@ -13,7 +13,19 @@ "configuration": { "component_type": "FluxPipeline", "prompt_weighting": true, - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-schnell", @@ -21,8 +33,8 @@ }, "arguments": { "prompt": "variable:prompt", - "num_inference_steps": 20, - "guidance_scale": 4.0, + "num_inference_steps": 4, + "guidance_scale": 0.0, "height": 1280, "width": 896 } diff --git a/workflows/templates/segment-and-inpaint.json b/workflows/templates/segment-and-inpaint.json index 384503d8..6689610f 100644 --- a/workflows/templates/segment-and-inpaint.json +++ b/workflows/templates/segment-and-inpaint.json @@ -44,7 +44,19 @@ "pipeline": { "configuration": { "component_type": "FluxFillPipeline", - "offload": "sequential" + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + } }, "from_pretrained_arguments": { "model_name": "black-forest-labs/FLUX.1-Fill-dev", diff --git a/workflows/templates/step-caching.json b/workflows/templates/step-caching.json index e8f8505c..a91f6d47 100644 --- a/workflows/templates/step-caching.json +++ b/workflows/templates/step-caching.json @@ -14,7 +14,19 @@ "pipeline": { "configuration": { "component_type": "FluxPipeline", - "offload": "sequential", + "components": { + "transformer": { + "group_offload": { + "offload_type": "block_level", + "num_blocks_per_group": 1, + "use_stream": true, + "record_stream": true + } + }, + "text_encoder": { "device": "cuda" }, + "text_encoder_2": { "device": "cuda" }, + "vae": { "device": "cuda" } + }, "cache": { "type": "first_block", "threshold": 0.05