inline-core 1.1.1__tar.gz → 1.1.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {inline_core-1.1.1 → inline_core-1.1.2}/CLAUDE.md +20 -3
- {inline_core-1.1.1 → inline_core-1.1.2}/PKG-INFO +5 -1
- {inline_core-1.1.1 → inline_core-1.1.2}/README.md +31 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/pyproject.toml +10 -1
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/device/detect.py +16 -0
- inline_core-1.1.2/src/inline_core/device/memory.py +361 -0
- inline_core-1.1.2/src/inline_core/device/policy.py +167 -0
- inline_core-1.1.2/src/inline_core/graph/loader_runners.py +100 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/graph/primitives.py +10 -2
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/graph/registry.py +8 -4
- inline_core-1.1.2/src/inline_core/models/loaders.py +471 -0
- inline_core-1.1.2/src/inline_core/models/sampling.py +174 -0
- inline_core-1.1.2/src/inline_core/models/zimage/primitives.py +578 -0
- inline_core-1.1.2/src/inline_core/models/zimage/requirements.py +260 -0
- inline_core-1.1.2/src/inline_core/models/zimage/runner.py +945 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/server/__main__.py +8 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/server/app.py +39 -13
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/server/bootstrap.py +7 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/generation.py +19 -2
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/graph_build.py +23 -2
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/models.py +108 -38
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/moodboard.py +17 -2
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_hidden_nodes.py +13 -6
- inline_core-1.1.2/tests/test_loader_runners.py +83 -0
- inline_core-1.1.2/tests/test_loaders.py +84 -0
- inline_core-1.1.2/tests/test_memory_policy.py +238 -0
- inline_core-1.1.2/tests/test_sampling.py +121 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_studio_generation.py +23 -0
- inline_core-1.1.2/tests/test_studio_models.py +84 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_studio_rpc.py +14 -3
- inline_core-1.1.2/tests/test_zimage_primitives.py +321 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_zimage_resolve.py +38 -6
- inline_core-1.1.2/tests/test_zimage_runner.py +481 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/uv.lock +146 -1
- {inline_core-1.1.1 → inline_core-1.1.2}/webui.sh +82 -1
- inline_core-1.1.1/src/inline_core/device/memory.py +0 -168
- inline_core-1.1.1/src/inline_core/device/policy.py +0 -82
- inline_core-1.1.1/src/inline_core/models/zimage/requirements.py +0 -180
- inline_core-1.1.1/src/inline_core/models/zimage/runner.py +0 -386
- inline_core-1.1.1/tests/test_memory_policy.py +0 -87
- inline_core-1.1.1/tests/test_studio_models.py +0 -94
- inline_core-1.1.1/tests/test_zimage_runner.py +0 -161
- {inline_core-1.1.1 → inline_core-1.1.2}/.gitignore +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/.python-version +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/main.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/scripts/reference.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/components/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/components/conditioning.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/components/interfaces.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/config.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/device/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/device/auto.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/device/types.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/errors.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/graph/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/graph/cache.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/graph/descriptor.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/graph/executor.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/graph/runners.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/graph/schema.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/graph/topo.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/graph/validate.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/importer/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/importer/comfy.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/media.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/models/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/models/catalog.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/models/zimage/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/parallel/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/parallel/config.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/parallel/group.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/parallel/launch.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/parallel/protocol.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/parallel/registry.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/parallel/worker.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/runtime/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/runtime/context.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/runtime/file_store.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/runtime/progress.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/runtime/run.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/runtime/store.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/sampling/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/sampling/batch.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/server/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/server/assets.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/server/frontend.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/server/manager.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/server/rpc.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/server/run_store.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/server/serialize.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/assets.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/config.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/fal.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/frames.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/handlers.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/schema.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/store.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/timeline/__init__.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/timeline/compose.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/timeline/ffmpeg.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/timeline/render.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/studio/timeline/resolve.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/src/inline_core/takes.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/helpers.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_cache.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_catalog.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_comfy_import.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_config.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_executor.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_file_store.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_frontend_serving.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_parallel_group.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_primitives.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_rpc_bridge.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_run_store.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_schema.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_server.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_studio_assets.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_studio_fal.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_studio_frames.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_studio_moodboard.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_studio_schema.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_studio_store.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_studio_timeline.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_take_bytes.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_topo.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_validate.py +0 -0
- {inline_core-1.1.1 → inline_core-1.1.2}/tests/test_xfuser_sampler.py +0 -0
|
@@ -120,9 +120,26 @@ between nodes and are never takes.
|
|
|
120
120
|
- **Server bind** — `INLINE_HOST` (default `127.0.0.1`), `INLINE_PORT` (default `8848`).
|
|
121
121
|
- **Model overrides** — e.g. `INLINE_ZIMAGE_MODEL` (a single `.safetensors` file path, a local
|
|
122
122
|
diffusers dir, or a HF repo id for Z-Image). Auto-resolved from `diffusion_models/` when unset.
|
|
123
|
-
- **Memory** —
|
|
124
|
-
|
|
125
|
-
|
|
123
|
+
- **Memory** — by default prefer the GPU, and **auto-fit the model to it**. When a runner hands the
|
|
124
|
+
policy the model's on-disk sizes (`DevicePolicy.set_footprint`, a `ModelFootprint`), the policy sizes
|
|
125
|
+
the weights against total VRAM (minus an activation headroom) and picks the lightest plan that fits:
|
|
126
|
+
full-precision **resident** → else int8 **resident** (torchao halves the transformer + text encoder,
|
|
127
|
+
no CPU offload — int8 + accelerate's `enable_model_cpu_offload` deadlock together) → else unquantized
|
|
128
|
+
`SEQUENTIAL` submodule streaming. So a card that can't hold Z-Image full-precision (a T4)
|
|
129
|
+
**auto-int8s with no flag**; `set_footprint(None)` / an unsizable whole-pipeline folder falls back to
|
|
130
|
+
the coarse total-VRAM buckets. Capacity is TOTAL VRAM (a fixed device property), not live-free, so the
|
|
131
|
+
plan — and the pipeline cache key it feeds — is stable across runs. The runner does a **pre-flight
|
|
132
|
+
check** (`DevicePolicy.fit_estimate`): a model too big for VRAM+RAM fails with a clean node error
|
|
133
|
+
before any load, instead of a host-RAM OOM-kill that would take the shared server down. Weights stream
|
|
134
|
+
straight to the GPU on load (`device=`/`device_map` in `models/loaders.py`) so peak host RAM ≈ one
|
|
135
|
+
tensor, and switching checkpoints **evicts** the previous model (`loaders.unload_components`,
|
|
136
|
+
`_evict_stale`) rather than stacking. `webui.sh` always sets `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`
|
|
137
|
+
to cut fragmentation OOMs. **Smart memory** (`INLINE_SMART_MEMORY=1`, `--smart-memory`) and
|
|
138
|
+
`INLINE_ALLOW_CPU_OFFLOAD=1` remain as explicit overrides when no footprint is set.
|
|
139
|
+
- **Compute dtype** — bf16 on the GPU by default, but **fp16 on cards without bf16 acceleration**
|
|
140
|
+
(Turing/Volta — compute capability < 8.0, e.g. a T4): same footprint, but it uses the fp16 tensor
|
|
141
|
+
cores instead of bf16's slow path. The **VAE stays upcast** (bf16, or fp32 when the denoiser is
|
|
142
|
+
fp16) because fp16 VAE decode can overflow to black/NaN images (`_compute_dtype` in `device/`).
|
|
126
143
|
- **Multi-GPU** — `INLINE_PARALLEL` (e.g. `pipefusion=2`, `pipefusion=2,ulysses=2`); degrees multiply
|
|
127
144
|
to the world size, which must equal the GPU count.
|
|
128
145
|
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: inline-core
|
|
3
|
-
Version: 1.1.
|
|
3
|
+
Version: 1.1.2
|
|
4
4
|
Summary: The generation engine behind Inline.
|
|
5
5
|
Requires-Python: >=3.11
|
|
6
6
|
Requires-Dist: numpy>=1.26
|
|
@@ -19,7 +19,9 @@ Requires-Dist: diffusers>=0.36; extra == 'parallel'
|
|
|
19
19
|
Requires-Dist: huggingface-hub>=0.23; extra == 'parallel'
|
|
20
20
|
Requires-Dist: nvidia-ml-py>=12; extra == 'parallel'
|
|
21
21
|
Requires-Dist: safetensors>=0.4; extra == 'parallel'
|
|
22
|
+
Requires-Dist: scipy>=1.11; extra == 'parallel'
|
|
22
23
|
Requires-Dist: torch>=2.2; extra == 'parallel'
|
|
24
|
+
Requires-Dist: torchao>=0.14; extra == 'parallel'
|
|
23
25
|
Requires-Dist: transformers>=4.44; extra == 'parallel'
|
|
24
26
|
Requires-Dist: xfuser>=0.4; extra == 'parallel'
|
|
25
27
|
Provides-Extra: server
|
|
@@ -31,5 +33,7 @@ Requires-Dist: accelerate>=0.30; extra == 'zimage'
|
|
|
31
33
|
Requires-Dist: diffusers>=0.36; extra == 'zimage'
|
|
32
34
|
Requires-Dist: huggingface-hub>=0.23; extra == 'zimage'
|
|
33
35
|
Requires-Dist: safetensors>=0.4; extra == 'zimage'
|
|
36
|
+
Requires-Dist: scipy>=1.11; extra == 'zimage'
|
|
34
37
|
Requires-Dist: torch>=2.2; extra == 'zimage'
|
|
38
|
+
Requires-Dist: torchao>=0.14; extra == 'zimage'
|
|
35
39
|
Requires-Dist: transformers>=4.44; extra == 'zimage'
|
|
@@ -136,6 +136,37 @@ Enabling it:
|
|
|
136
136
|
The device policy and the worker-group IPC are in place and tested with a stub worker; the real
|
|
137
137
|
xfuser denoise lands with the GPU-side Z-Image runner.
|
|
138
138
|
|
|
139
|
+
## Memory: fit a big model onto a small GPU
|
|
140
|
+
|
|
141
|
+
When a model is too large to hold full-precision on your GPU — or larger than your system RAM — Inline
|
|
142
|
+
Core fits it to the machine automatically, **with no flags**. This is the low-VRAM path, and it aims to
|
|
143
|
+
"just run" instead of making you tune profiles.
|
|
144
|
+
|
|
145
|
+
- **Weights stream straight to the GPU.** Each component loads directly onto the GPU from its
|
|
146
|
+
memory-mapped `.safetensors` — the engine never materializes a full copy in system RAM first. So you
|
|
147
|
+
do **not** need roughly 2× the model size in RAM to load it, and a load can no longer exhaust host RAM
|
|
148
|
+
and take the server down.
|
|
149
|
+
- **Auto-fit picks the lightest plan that fits.** The policy sizes the model against your VRAM and
|
|
150
|
+
chooses, in order: full-precision **resident** → **int8 resident** (weight-only quantization halves the
|
|
151
|
+
big weights so they fit on the GPU) → **CPU-offload streaming** (submodules move on/off the GPU). A card
|
|
152
|
+
that can't hold the model full-precision quantizes to int8 on its own — no need to know a flag.
|
|
153
|
+
- **Load-time fit check.** A model too large for VRAM + RAM fails fast with a clear message on the node,
|
|
154
|
+
instead of crashing mid-load. The model popup also shows the estimate up front (fits / will run int8 /
|
|
155
|
+
won't fit).
|
|
156
|
+
- **Switching models frees the old one.** Loading a different checkpoint evicts the previous weights
|
|
157
|
+
rather than stacking them, so you can move between models without accumulating VRAM.
|
|
158
|
+
- **Generation fits too, not just loading.** Prompts encode under `no_grad` so the text-encoder forward
|
|
159
|
+
never retains its activation graph, and the encoder is parked on the CPU during denoise to free its
|
|
160
|
+
VRAM; large images (1024² and up) decode the VAE in tiles instead of one full-frame pass. So the
|
|
161
|
+
memory-heavy steps — encode and decode — stay bounded rather than spiking host RAM or VRAM and taking
|
|
162
|
+
the server down at high resolution.
|
|
163
|
+
- **Fragmentation-resistant allocator.** `webui.sh` runs with PyTorch expandable CUDA segments, which
|
|
164
|
+
avoids the "allocation failed with VRAM still free" fragmentation OOMs common on low-VRAM cards.
|
|
165
|
+
|
|
166
|
+
You can still steer it by hand: `--lowvram` for the tight-VRAM profile, `--smart-memory` to force the
|
|
167
|
+
offload machinery, `--vram-budget GB` / `INLINE_VRAM_BUDGET_GB` to cap the assumed VRAM, or
|
|
168
|
+
`INLINE_PROFILE` to pin `gpu-max | lowvram | cpu`. `GET /v1/health` reports the live VRAM/RAM budget.
|
|
169
|
+
|
|
139
170
|
## Run
|
|
140
171
|
|
|
141
172
|
The easy path is `webui.sh`, which maps friendly flags onto the engine's `INLINE_*` env knobs:
|
|
@@ -2,7 +2,7 @@
|
|
|
2
2
|
# PyPI distribution name. The import package is `inline_core` (src/inline_core); PyPI treats the
|
|
3
3
|
# hyphen and underscore as equivalent, so `pip install inline-core` installs the `inline_core` package.
|
|
4
4
|
name = "inline-core"
|
|
5
|
-
version = "1.1.
|
|
5
|
+
version = "1.1.2"
|
|
6
6
|
description = "The generation engine behind Inline."
|
|
7
7
|
requires-python = ">=3.11"
|
|
8
8
|
dependencies = [
|
|
@@ -21,6 +21,15 @@ zimage = [
|
|
|
21
21
|
"transformers>=4.44",
|
|
22
22
|
"accelerate>=0.30",
|
|
23
23
|
"safetensors>=0.4",
|
|
24
|
+
# Smart memory (`webui.sh --smart-memory`) quantizes the big weights to int8 with torchao's
|
|
25
|
+
# weight-only quant — torch-native and movable, so it coexists with CPU offload. >=0.14 for the
|
|
26
|
+
# config-object API (Int8WeightOnlyConfig; the "int8_weight_only" string form was dropped).
|
|
27
|
+
# Best-effort at load time: an absent/incompatible torchao logs and falls back, never crashes.
|
|
28
|
+
"torchao>=0.14",
|
|
29
|
+
# scipy backs the Beta sigma schedule (FlowMatchEulerDiscreteScheduler.use_beta_sigmas). diffusers
|
|
30
|
+
# gates that option on scipy at scheduler construction; the sampling registry falls back to the
|
|
31
|
+
# model default spacing if it is ever absent, but declare it so Beta works out of the box.
|
|
32
|
+
"scipy>=1.11",
|
|
24
33
|
# Explicit model downloads (the node's model popup) go through huggingface_hub. It ships
|
|
25
34
|
# transitively with diffusers/transformers, but pin it here since we call snapshot_download
|
|
26
35
|
# directly and rely on local_dir= to land files in models/ (never the hidden HF cache).
|
|
@@ -23,6 +23,22 @@ def available_devices() -> tuple[Device, ...]:
|
|
|
23
23
|
return (Device(DeviceKind.CPU),)
|
|
24
24
|
|
|
25
25
|
|
|
26
|
+
def cuda_supports_bf16(device: Device) -> bool:
|
|
27
|
+
"""True when the CUDA GPU has native bf16 acceleration — Ampere or newer (compute capability
|
|
28
|
+
>= 8.0). Turing/Volta (T4, V100) can run bf16 but only through a slow unaccelerated path, so the
|
|
29
|
+
policy prefers fp16 there (same memory footprint, but it uses the fp16 tensor cores). Unknown /
|
|
30
|
+
non-CUDA falls back to True so bf16 stays the default on anything we can't measure."""
|
|
31
|
+
if device.kind is not DeviceKind.CUDA:
|
|
32
|
+
return True
|
|
33
|
+
try:
|
|
34
|
+
import torch
|
|
35
|
+
|
|
36
|
+
major, _ = torch.cuda.get_device_capability(device.index)
|
|
37
|
+
return major >= 8
|
|
38
|
+
except Exception:
|
|
39
|
+
return True
|
|
40
|
+
|
|
41
|
+
|
|
26
42
|
def has_nvlink() -> bool:
|
|
27
43
|
"""Best-effort: True only when NVLink is confirmed between two GPUs. Conservative (False) on any
|
|
28
44
|
uncertainty, so the policy defaults to the PCIe-friendly PipeFusion split."""
|
|
@@ -0,0 +1,361 @@
|
|
|
1
|
+
"""The memory-aware device policy: measure RAM/VRAM, pick a profile, own dtype/offload/quant/tiling.
|
|
2
|
+
|
|
3
|
+
Profiles: gpu-max (ample VRAM), lowvram (tight VRAM), cpu (fp32, tiling, int8 to fit RAM). By
|
|
4
|
+
default we always prefer the GPU: even under lowvram, weights stay resident on the GPU (tiling +
|
|
5
|
+
attention
|
|
6
|
+
slicing + int8 do the memory saving) and we do NOT auto-offload to CPU — offloading is slow and
|
|
7
|
+
defeats "use the GPU we have".
|
|
8
|
+
|
|
9
|
+
Smart memory (INLINE_SMART_MEMORY=1, `webui.sh --smart-memory`) is the opt-in escape hatch for a
|
|
10
|
+
model that simply does not fit resident: it spreads the model across VRAM + RAM + CPU by streaming
|
|
11
|
+
components on/off the GPU (a graduated OffloadMode — MODEL, or SEQUENTIAL on a very small GPU) and
|
|
12
|
+
quantizes the big weights to int8 so the offloaded half also fits in RAM. Slower per image, but it
|
|
13
|
+
runs where full-resident OOMs. Set INLINE_ALLOW_CPU_OFFLOAD=1 for the older bare model-offload knob
|
|
14
|
+
without quantization. Override the profile/budget with INLINE_PROFILE and INLINE_VRAM_BUDGET_GB.
|
|
15
|
+
Detection is lazy so the core imports without torch or psutil; an unavailable measurement keeps the
|
|
16
|
+
policy conservative.
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
from __future__ import annotations
|
|
20
|
+
|
|
21
|
+
import os
|
|
22
|
+
|
|
23
|
+
from .detect import available_devices, cuda_supports_bf16, has_nvlink
|
|
24
|
+
from .policy import (
|
|
25
|
+
AttentionBackend,
|
|
26
|
+
DevicePolicy,
|
|
27
|
+
FitEstimate,
|
|
28
|
+
ModelFootprint,
|
|
29
|
+
OffloadMode,
|
|
30
|
+
Parallel,
|
|
31
|
+
Placement,
|
|
32
|
+
Profile,
|
|
33
|
+
Quantization,
|
|
34
|
+
)
|
|
35
|
+
from .types import Device, DeviceKind, DType
|
|
36
|
+
|
|
37
|
+
_GPU_MAX_MIN_VRAM_GB = 16.0 # at or above -> gpu-max, else lowvram
|
|
38
|
+
_QUANT_VRAM_GB = 10.0 # lowvram below this -> int8
|
|
39
|
+
_QUANT_RAM_GB = 48.0 # cpu below this -> int8
|
|
40
|
+
# Smart memory: at/above this VRAM the int8 model fits RESIDENT on the GPU (no offload); below it,
|
|
41
|
+
# even int8 won't fit, so fall to SEQUENTIAL submodule streaming (unquantized — torchao int8 and CPU
|
|
42
|
+
# offload deadlock together) so a very small GPU can still run.
|
|
43
|
+
_SMART_RESIDENT_MIN_VRAM_GB = 6.0
|
|
44
|
+
|
|
45
|
+
# Fit estimate (size-aware placement): reserve this much VRAM beyond the weights for denoise
|
|
46
|
+
# activations + the CUDA context + allocator fragmentation, and treat int8 weight-only quant as
|
|
47
|
+
# ~half the fp16 weight bytes. Deliberately generous so the estimate errs toward a lighter plan.
|
|
48
|
+
_ACTIVATION_HEADROOM_GB = 2.5
|
|
49
|
+
_INT8_FACTOR = 0.5
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def _system_ram_gb() -> float | None:
|
|
53
|
+
try:
|
|
54
|
+
return os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") / 1e9
|
|
55
|
+
except (ValueError, OSError, AttributeError):
|
|
56
|
+
pass
|
|
57
|
+
try:
|
|
58
|
+
import psutil
|
|
59
|
+
|
|
60
|
+
return psutil.virtual_memory().total / 1e9
|
|
61
|
+
except ModuleNotFoundError:
|
|
62
|
+
return None
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def _vram_gb(device: Device) -> float | None:
|
|
66
|
+
if device.kind is not DeviceKind.CUDA:
|
|
67
|
+
return None
|
|
68
|
+
try:
|
|
69
|
+
import torch
|
|
70
|
+
|
|
71
|
+
return torch.cuda.mem_get_info(device.index)[1] / 1e9
|
|
72
|
+
except Exception:
|
|
73
|
+
return None
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _free_vram_gb(device: Device) -> float | None:
|
|
77
|
+
"""Live *free* VRAM (mem_get_info[0]) — for diagnostics/UI only. NOT used to choose the plan:
|
|
78
|
+
residency-dependent free readings would make the quant/offload decision (and thus the pipeline
|
|
79
|
+
cache key) oscillate. Capacity decisions use total VRAM (``_vram_gb``)."""
|
|
80
|
+
if device.kind is not DeviceKind.CUDA:
|
|
81
|
+
return None
|
|
82
|
+
try:
|
|
83
|
+
import torch
|
|
84
|
+
|
|
85
|
+
return torch.cuda.mem_get_info(device.index)[0] / 1e9
|
|
86
|
+
except Exception:
|
|
87
|
+
return None
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
def _free_ram_gb() -> float | None:
|
|
91
|
+
"""Live *available* system RAM (psutil), for the UI + a load guard."""
|
|
92
|
+
try:
|
|
93
|
+
import psutil
|
|
94
|
+
|
|
95
|
+
return psutil.virtual_memory().available / 1e9
|
|
96
|
+
except Exception:
|
|
97
|
+
return None
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def _env_profile() -> Profile | None:
|
|
101
|
+
value = os.environ.get("INLINE_PROFILE", "").strip().lower()
|
|
102
|
+
return next((p for p in Profile if p.value == value), None)
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
def _env_budget() -> float | None:
|
|
106
|
+
value = os.environ.get("INLINE_VRAM_BUDGET_GB", "").strip()
|
|
107
|
+
try:
|
|
108
|
+
return float(value) if value else None
|
|
109
|
+
except ValueError:
|
|
110
|
+
return None
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
def _env_flag(name: str) -> bool:
|
|
114
|
+
return os.environ.get(name, "").strip().lower() in ("1", "true", "yes", "on")
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
def _env_allow_offload() -> bool:
|
|
118
|
+
"""Opt back into bare CPU offload under lowvram. Off by default: we keep weights on the GPU."""
|
|
119
|
+
return _env_flag("INLINE_ALLOW_CPU_OFFLOAD")
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
def _env_smart_memory() -> bool:
|
|
123
|
+
"""Opt into smart memory: graduated CPU offload + int8 quant so a too-big model fits across
|
|
124
|
+
VRAM + RAM. Off by default. Set by `webui.sh --smart-memory`."""
|
|
125
|
+
return _env_flag("INLINE_SMART_MEMORY")
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
def _env_parallel() -> Parallel | None:
|
|
129
|
+
"""Parse INLINE_PARALLEL like `pipefusion=2,ulysses=2` into a degree spec."""
|
|
130
|
+
value = os.environ.get("INLINE_PARALLEL", "").strip()
|
|
131
|
+
if not value:
|
|
132
|
+
return None
|
|
133
|
+
valid = {"pipefusion", "ulysses", "ring", "cfg", "tensor"}
|
|
134
|
+
degrees: dict[str, int] = {}
|
|
135
|
+
for part in value.split(","):
|
|
136
|
+
key, _, raw = part.partition("=")
|
|
137
|
+
if key.strip() in valid:
|
|
138
|
+
try:
|
|
139
|
+
degrees[key.strip()] = int(raw)
|
|
140
|
+
except ValueError:
|
|
141
|
+
continue
|
|
142
|
+
return Parallel(**degrees) if degrees else None
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
class MemoryPolicy(DevicePolicy):
|
|
146
|
+
def __init__(
|
|
147
|
+
self,
|
|
148
|
+
device: Device | None = None,
|
|
149
|
+
*,
|
|
150
|
+
ram_gb: float | None = None,
|
|
151
|
+
vram_gb: float | None = None,
|
|
152
|
+
profile: Profile | None = None,
|
|
153
|
+
devices: tuple[Device, ...] | None = None,
|
|
154
|
+
nvlink: bool | None = None,
|
|
155
|
+
parallel: Parallel | None = None,
|
|
156
|
+
allow_offload: bool | None = None,
|
|
157
|
+
smart_memory: bool | None = None,
|
|
158
|
+
supports_bf16: bool | None = None,
|
|
159
|
+
) -> None:
|
|
160
|
+
self._devices = devices if devices is not None else available_devices()
|
|
161
|
+
self._device = device or self._devices[0]
|
|
162
|
+
self._ram_gb = ram_gb if ram_gb is not None else _system_ram_gb()
|
|
163
|
+
self._vram_gb = vram_gb if vram_gb is not None else _vram_gb(self._device)
|
|
164
|
+
self._profile = profile or _env_profile() or self._choose_profile()
|
|
165
|
+
self._nvlink = nvlink if nvlink is not None else has_nvlink()
|
|
166
|
+
self._parallel = parallel if parallel is not None else _env_parallel()
|
|
167
|
+
self._smart = smart_memory if smart_memory is not None else _env_smart_memory()
|
|
168
|
+
self._allow_offload = allow_offload if allow_offload is not None else _env_allow_offload()
|
|
169
|
+
self._bf16 = (
|
|
170
|
+
supports_bf16 if supports_bf16 is not None else cuda_supports_bf16(self._device)
|
|
171
|
+
)
|
|
172
|
+
self._footprint: ModelFootprint | None = None
|
|
173
|
+
self._fit: FitEstimate | None = None
|
|
174
|
+
|
|
175
|
+
# --- size-aware fit (auto-fit the model to the device without a --smart-memory flag) ---------
|
|
176
|
+
|
|
177
|
+
def set_footprint(self, footprint: ModelFootprint | None) -> None:
|
|
178
|
+
"""Record the model's on-disk sizes so ``profile``/``quantization``/offload fit the device.
|
|
179
|
+
Called by the runner per run (files can change). Falls back to the coarse buckets when the
|
|
180
|
+
footprint is None or unmeasurable."""
|
|
181
|
+
self._footprint = footprint
|
|
182
|
+
self._fit = self._compute_fit(footprint) if footprint is not None else None
|
|
183
|
+
|
|
184
|
+
def estimate_fit(self, footprint: ModelFootprint) -> FitEstimate | None:
|
|
185
|
+
"""A pure fit verdict — does NOT mutate the policy, so the UI thread can call it while a run
|
|
186
|
+
holds the policy on the worker thread."""
|
|
187
|
+
return self._compute_fit(footprint)
|
|
188
|
+
|
|
189
|
+
def fit_estimate(self) -> FitEstimate | None:
|
|
190
|
+
return self._fit
|
|
191
|
+
|
|
192
|
+
def _compute_fit(self, fp: ModelFootprint) -> FitEstimate | None:
|
|
193
|
+
"""Pick the lightest plan whose weights fit the GPU: full-precision resident, else int8
|
|
194
|
+
resident, else CPU-offload streaming (unquantized — int8 + offload deadlock). Capacity is
|
|
195
|
+
TOTAL VRAM (a fixed device property) minus activation headroom, so the decision is stable
|
|
196
|
+
across runs and doesn't bust the pipeline cache. Returns None (→ coarse buckets) off-CUDA or
|
|
197
|
+
when sizes are unavailable (e.g. a whole-pipeline folder)."""
|
|
198
|
+
if self._device.kind is not DeviceKind.CUDA or fp.total_bytes <= 0:
|
|
199
|
+
return None
|
|
200
|
+
budget = _env_budget() or self._vram_gb
|
|
201
|
+
if budget is None:
|
|
202
|
+
return None
|
|
203
|
+
cap = max(0.0, budget - _ACTIVATION_HEADROOM_GB)
|
|
204
|
+
big = (fp.diffusion_bytes + fp.text_encoder_bytes) / 1e9
|
|
205
|
+
vae = fp.vae_bytes / 1e9
|
|
206
|
+
full = big + vae
|
|
207
|
+
int8 = big * _INT8_FACTOR + vae
|
|
208
|
+
forced = _env_profile() is not None # explicit --profile pins the profile; fit picks quant
|
|
209
|
+
|
|
210
|
+
def prof(auto: Profile) -> Profile:
|
|
211
|
+
return self._profile if forced else auto
|
|
212
|
+
|
|
213
|
+
if full <= cap:
|
|
214
|
+
return FitEstimate(
|
|
215
|
+
"resident", Quantization.NONE, OffloadMode.NONE, prof(Profile.GPU_MAX),
|
|
216
|
+
full, budget, True, "Full-precision weights fit in VRAM.",
|
|
217
|
+
)
|
|
218
|
+
if int8 <= cap:
|
|
219
|
+
return FitEstimate(
|
|
220
|
+
"int8", Quantization.INT8, OffloadMode.NONE, prof(Profile.LOWVRAM),
|
|
221
|
+
int8, budget, True,
|
|
222
|
+
"Weights are int8-quantized to fit this GPU's VRAM.",
|
|
223
|
+
)
|
|
224
|
+
# int8 still won't fit resident -> CPU-offload streaming. Only viable if the (unquantized)
|
|
225
|
+
# model fits in system RAM, since sequential offload holds the off-GPU weights there.
|
|
226
|
+
ram = self._ram_gb
|
|
227
|
+
if ram is not None and full > ram:
|
|
228
|
+
return FitEstimate(
|
|
229
|
+
"wont-fit", Quantization.NONE, OffloadMode.SEQUENTIAL, prof(Profile.LOWVRAM),
|
|
230
|
+
full, budget, False,
|
|
231
|
+
"Model is too large for this GPU's VRAM and this machine's RAM.",
|
|
232
|
+
)
|
|
233
|
+
return FitEstimate(
|
|
234
|
+
"offload", Quantization.NONE, OffloadMode.SEQUENTIAL, prof(Profile.LOWVRAM),
|
|
235
|
+
int8, budget, True, "Weights stream between GPU and RAM (slower).",
|
|
236
|
+
)
|
|
237
|
+
|
|
238
|
+
def vram_budget_mb(self) -> int | None:
|
|
239
|
+
budget = _env_budget() or self._vram_gb
|
|
240
|
+
return int(budget * 1024) if budget is not None else None
|
|
241
|
+
|
|
242
|
+
def free_vram_mb(self) -> int | None:
|
|
243
|
+
gb = _free_vram_gb(self._device)
|
|
244
|
+
return int(gb * 1024) if gb is not None else None
|
|
245
|
+
|
|
246
|
+
def free_ram_mb(self) -> int | None:
|
|
247
|
+
gb = _free_ram_gb()
|
|
248
|
+
return int(gb * 1024) if gb is not None else None
|
|
249
|
+
|
|
250
|
+
def _choose_profile(self) -> Profile:
|
|
251
|
+
if self._device.kind is DeviceKind.CPU:
|
|
252
|
+
return Profile.CPU
|
|
253
|
+
if self._device.kind is DeviceKind.MPS:
|
|
254
|
+
return Profile.GPU_MAX # unified memory; offload semantics differ
|
|
255
|
+
budget = _env_budget() or self._vram_gb
|
|
256
|
+
if budget is not None and budget < _GPU_MAX_MIN_VRAM_GB:
|
|
257
|
+
return Profile.LOWVRAM
|
|
258
|
+
return Profile.GPU_MAX
|
|
259
|
+
|
|
260
|
+
@property
|
|
261
|
+
def profile(self) -> Profile:
|
|
262
|
+
"""The effective profile: the size-aware fit's profile when a footprint is set, else the
|
|
263
|
+
profile chosen at init from the coarse VRAM buckets / env override."""
|
|
264
|
+
return self._fit.profile if self._fit is not None else self._profile
|
|
265
|
+
|
|
266
|
+
def placement(self, role: str) -> Placement:
|
|
267
|
+
if self.profile is Profile.CPU:
|
|
268
|
+
return Placement(self._device, DType.FP32)
|
|
269
|
+
offload_mode = self._offload_mode()
|
|
270
|
+
dtype = self._compute_dtype(role)
|
|
271
|
+
if role == "denoiser":
|
|
272
|
+
parallel = self._denoiser_parallel()
|
|
273
|
+
if parallel is not None:
|
|
274
|
+
cuda = tuple(d for d in self._devices if d.kind is DeviceKind.CUDA)
|
|
275
|
+
return Placement(
|
|
276
|
+
self._device,
|
|
277
|
+
dtype,
|
|
278
|
+
offload_mode=offload_mode,
|
|
279
|
+
devices=cuda[: parallel.world_size],
|
|
280
|
+
parallel=parallel,
|
|
281
|
+
)
|
|
282
|
+
return Placement(self._device, dtype, offload_mode=offload_mode)
|
|
283
|
+
|
|
284
|
+
def _compute_dtype(self, role: str) -> DType:
|
|
285
|
+
"""The weight/compute dtype for a GPU role. bf16 by default, but fp16 on GPUs without bf16
|
|
286
|
+
acceleration (Turing/Volta) — same footprint, yet it uses their fp16 tensor cores instead of
|
|
287
|
+
bf16's slow path. The VAE is the exception: fp16 decode can overflow to black/NaN images, so
|
|
288
|
+
it stays upcast (bf16, or fp32 on a card that also lacks bf16) — it is tiny, so the cost is
|
|
289
|
+
nil.
|
|
290
|
+
|
|
291
|
+
**int8 overrides the fp16 preference to bf16.** torchao weight-only int8 only supports a
|
|
292
|
+
bf16 compute dtype — with fp16 the quantization silently no-ops, so the "int8" weights load
|
|
293
|
+
at *full* fp16 size and blow the VRAM budget (a T4 then OOMs mid-load). The int8 matmul still
|
|
294
|
+
runs on the card's int8 tensor cores; only the residual bf16 activations pay the slow path.
|
|
295
|
+
bf16 also has fp32's exponent range, so the VAE no longer needs the fp32 anti-overflow
|
|
296
|
+
upcast — it rides along at bf16."""
|
|
297
|
+
if self.quantization() is Quantization.INT8:
|
|
298
|
+
return DType.BF16
|
|
299
|
+
if self._bf16:
|
|
300
|
+
return DType.BF16
|
|
301
|
+
return DType.FP32 if role == "vae" else DType.FP16
|
|
302
|
+
|
|
303
|
+
def _offload_mode(self) -> OffloadMode:
|
|
304
|
+
"""Whether (and how) to stream weights to CPU. Default: never — prefer the GPU we have, and
|
|
305
|
+
let tiling/slicing/int8 do the saving with weights resident.
|
|
306
|
+
|
|
307
|
+
Smart memory keeps weights RESIDENT and quantizes to int8 instead of offloading: int8 halves
|
|
308
|
+
the model so a lowvram GPU holds it, and that avoids the slow — and, with torchao int8,
|
|
309
|
+
hang-prone — accelerate CPU-offload path. Only a GPU too small for even int8-resident
|
|
310
|
+
streams submodules (SEQUENTIAL, unquantized). The older INLINE_ALLOW_CPU_OFFLOAD flag still
|
|
311
|
+
does bare (unquantized) MODEL offload for anyone who wants it."""
|
|
312
|
+
if self._fit is not None:
|
|
313
|
+
return self._fit.offload_mode
|
|
314
|
+
if self._profile is not Profile.LOWVRAM:
|
|
315
|
+
return OffloadMode.NONE
|
|
316
|
+
if self._smart:
|
|
317
|
+
if _below(self._vram_gb, _SMART_RESIDENT_MIN_VRAM_GB):
|
|
318
|
+
return OffloadMode.SEQUENTIAL
|
|
319
|
+
return OffloadMode.NONE
|
|
320
|
+
return OffloadMode.MODEL if self._allow_offload else OffloadMode.NONE
|
|
321
|
+
|
|
322
|
+
def _denoiser_parallel(self) -> Parallel | None:
|
|
323
|
+
"""Split the denoiser across GPUs when there are 2+. An explicit override wins; else auto:
|
|
324
|
+
PipeFusion on PCIe (no NVLink), sequence-parallel (Ulysses) on NVLink."""
|
|
325
|
+
if self._parallel is not None:
|
|
326
|
+
return self._parallel if self._parallel.world_size > 1 else None
|
|
327
|
+
cuda = [d for d in self._devices if d.kind is DeviceKind.CUDA]
|
|
328
|
+
if len(cuda) < 2:
|
|
329
|
+
return None
|
|
330
|
+
return Parallel(ulysses=len(cuda)) if self._nvlink else Parallel(pipefusion=len(cuda))
|
|
331
|
+
|
|
332
|
+
def attention_backend(self) -> AttentionBackend:
|
|
333
|
+
return AttentionBackend.SDPA
|
|
334
|
+
|
|
335
|
+
def vae_tiling(self) -> bool:
|
|
336
|
+
return self.profile in (Profile.LOWVRAM, Profile.CPU)
|
|
337
|
+
|
|
338
|
+
def attention_slicing(self) -> bool:
|
|
339
|
+
return self.profile in (Profile.LOWVRAM, Profile.CPU)
|
|
340
|
+
|
|
341
|
+
def quantization(self) -> Quantization:
|
|
342
|
+
# A size-aware fit (when set) owns the quant choice — int8 auto-engages when full precision
|
|
343
|
+
# won't fit, no --smart-memory flag needed.
|
|
344
|
+
if self._fit is not None:
|
|
345
|
+
return self._fit.quant
|
|
346
|
+
# Smart memory quantizes to int8 so the halved model fits RESIDENT on the GPU (see
|
|
347
|
+
# _offload_mode). The one exception: a GPU so small it must SEQUENTIAL-offload instead runs
|
|
348
|
+
# unquantized — torchao int8 + CPU offload deadlock together.
|
|
349
|
+
if self._smart and self._profile is Profile.LOWVRAM:
|
|
350
|
+
if self._offload_mode() is OffloadMode.SEQUENTIAL:
|
|
351
|
+
return Quantization.NONE
|
|
352
|
+
return Quantization.INT8
|
|
353
|
+
if self._profile is Profile.LOWVRAM and _below(self._vram_gb, _QUANT_VRAM_GB):
|
|
354
|
+
return Quantization.INT8
|
|
355
|
+
if self._profile is Profile.CPU and _below(self._ram_gb, _QUANT_RAM_GB):
|
|
356
|
+
return Quantization.INT8
|
|
357
|
+
return Quantization.NONE
|
|
358
|
+
|
|
359
|
+
|
|
360
|
+
def _below(measured: float | None, threshold: float) -> bool:
|
|
361
|
+
return measured is not None and measured < threshold
|
|
@@ -0,0 +1,167 @@
|
|
|
1
|
+
"""The device and memory policy interface. Components never self-assign a device; they ask here."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from abc import ABC, abstractmethod
|
|
6
|
+
from dataclasses import dataclass
|
|
7
|
+
from enum import Enum
|
|
8
|
+
|
|
9
|
+
from .types import Device, DType
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class Profile(str, Enum):
|
|
13
|
+
GPU_MAX = "gpu-max"
|
|
14
|
+
LOWVRAM = "lowvram"
|
|
15
|
+
CPU = "cpu"
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class AttentionBackend(str, Enum):
|
|
19
|
+
FLASH = "flash"
|
|
20
|
+
XFORMERS = "xformers"
|
|
21
|
+
SDPA = "sdpa"
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class OffloadMode(str, Enum):
|
|
25
|
+
"""How a component's weights are split between GPU VRAM and CPU RAM.
|
|
26
|
+
|
|
27
|
+
NONE weights stay resident on the GPU (the default — prefer the GPU we have).
|
|
28
|
+
MODEL diffusers ``enable_model_cpu_offload``: only the *active* component (text encoder,
|
|
29
|
+
then transformer, then VAE) sits on the GPU; the rest waits in CPU RAM. Peak VRAM ≈
|
|
30
|
+
the largest single component. Fast — the standard low-VRAM fit.
|
|
31
|
+
SEQUENTIAL diffusers ``enable_sequential_cpu_offload``: submodules stream on/off the GPU
|
|
32
|
+
layer-by-layer. Lowest peak VRAM, slowest — for GPUs too small for MODEL offload.
|
|
33
|
+
"""
|
|
34
|
+
|
|
35
|
+
NONE = "none"
|
|
36
|
+
MODEL = "model"
|
|
37
|
+
SEQUENTIAL = "sequential"
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class Quantization(str, Enum):
|
|
41
|
+
NONE = "none"
|
|
42
|
+
INT8 = "int8" # torch-native weight-only, portable
|
|
43
|
+
NF4 = "nf4" # bitsandbytes, cuda-only
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
@dataclass(frozen=True)
|
|
47
|
+
class Parallel:
|
|
48
|
+
"""How the denoiser is split across GPUs (xDiT degrees). Product of degrees = world size (GPUs).
|
|
49
|
+
PipeFusion is PCIe-friendly; Ulysses/Ring want NVLink; CFG applies only to guided models."""
|
|
50
|
+
|
|
51
|
+
pipefusion: int = 1
|
|
52
|
+
ulysses: int = 1
|
|
53
|
+
ring: int = 1
|
|
54
|
+
cfg: int = 1
|
|
55
|
+
tensor: int = 1
|
|
56
|
+
|
|
57
|
+
@property
|
|
58
|
+
def world_size(self) -> int:
|
|
59
|
+
return self.pipefusion * self.ulysses * self.ring * self.cfg * self.tensor
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
@dataclass(frozen=True)
|
|
63
|
+
class Placement:
|
|
64
|
+
"""Where and how a component runs. Chosen by the policy, never by the component."""
|
|
65
|
+
|
|
66
|
+
device: Device
|
|
67
|
+
dtype: DType
|
|
68
|
+
offload_mode: OffloadMode = OffloadMode.NONE
|
|
69
|
+
# Multi-GPU (denoiser only): the device group + how it is split. Empty/None = single device.
|
|
70
|
+
devices: tuple[Device, ...] = ()
|
|
71
|
+
parallel: Parallel | None = None
|
|
72
|
+
|
|
73
|
+
@property
|
|
74
|
+
def offload(self) -> bool:
|
|
75
|
+
"""Whether any CPU offload is in effect (either MODEL or SEQUENTIAL). Kept as a bool so the
|
|
76
|
+
many call sites that only care "is this streaming to CPU?" stay simple."""
|
|
77
|
+
return self.offload_mode is not OffloadMode.NONE
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
@dataclass(frozen=True)
|
|
81
|
+
class ModelFootprint:
|
|
82
|
+
"""On-disk byte sizes of a model's big components — the input to a memory fit estimate. The
|
|
83
|
+
files ship already 16-bit, so a size is a good proxy for the fp16-resident weight footprint;
|
|
84
|
+
int8 halves the two big weights. The policy uses this to pick dtype/quant/offload to the GPU
|
|
85
|
+
instead of a coarse VRAM bucket."""
|
|
86
|
+
|
|
87
|
+
diffusion_bytes: int = 0
|
|
88
|
+
text_encoder_bytes: int = 0
|
|
89
|
+
vae_bytes: int = 0
|
|
90
|
+
|
|
91
|
+
@property
|
|
92
|
+
def total_bytes(self) -> int:
|
|
93
|
+
return self.diffusion_bytes + self.text_encoder_bytes + self.vae_bytes
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
@dataclass(frozen=True)
|
|
97
|
+
class FitEstimate:
|
|
98
|
+
"""The policy's verdict for a given ModelFootprint on this device: what plan fits, how much VRAM
|
|
99
|
+
it needs, and whether it fits. Drives the runner's pre-flight check and the UI warning."""
|
|
100
|
+
|
|
101
|
+
plan: str # "resident" | "int8" | "offload" | "wont-fit"
|
|
102
|
+
quant: Quantization
|
|
103
|
+
offload_mode: OffloadMode
|
|
104
|
+
profile: Profile
|
|
105
|
+
required_vram_gb: float
|
|
106
|
+
total_vram_gb: float | None
|
|
107
|
+
fits: bool
|
|
108
|
+
note: str
|
|
109
|
+
|
|
110
|
+
|
|
111
|
+
class DevicePolicy(ABC):
|
|
112
|
+
"""Owns dtype, device, offload, attention backend, and tiling for a worker."""
|
|
113
|
+
|
|
114
|
+
@property
|
|
115
|
+
@abstractmethod
|
|
116
|
+
def profile(self) -> Profile: ...
|
|
117
|
+
|
|
118
|
+
@abstractmethod
|
|
119
|
+
def placement(self, role: str) -> Placement:
|
|
120
|
+
"""Placement for a component role: text_encoder, denoiser, vae, and so on."""
|
|
121
|
+
|
|
122
|
+
@abstractmethod
|
|
123
|
+
def attention_backend(self) -> AttentionBackend: ...
|
|
124
|
+
|
|
125
|
+
@abstractmethod
|
|
126
|
+
def vae_tiling(self) -> bool:
|
|
127
|
+
"""Whether to tile VAE decode to cap peak memory."""
|
|
128
|
+
|
|
129
|
+
def attention_slicing(self) -> bool:
|
|
130
|
+
"""Whether to slice attention to cap peak memory. Default off."""
|
|
131
|
+
return False
|
|
132
|
+
|
|
133
|
+
def quantization(self) -> Quantization:
|
|
134
|
+
"""Weight quantization to fit low memory. Default none."""
|
|
135
|
+
return Quantization.NONE
|
|
136
|
+
|
|
137
|
+
# --- optional memory-fit surface (a size-aware policy overrides these) -----------------------
|
|
138
|
+
# Kept off the abstract contract with no-op/None defaults so simple policies (AutoDevicePolicy,
|
|
139
|
+
# test stubs) need not implement them; a memory-aware policy uses the model's on-disk sizes to
|
|
140
|
+
# fit dtype/quant/offload to the actual device.
|
|
141
|
+
|
|
142
|
+
def set_footprint(self, footprint: ModelFootprint | None) -> None:
|
|
143
|
+
"""Receive the model's on-disk component sizes so later placement/quantization calls fit the
|
|
144
|
+
device. Default: ignore (placement decided without a size estimate)."""
|
|
145
|
+
return None
|
|
146
|
+
|
|
147
|
+
def estimate_fit(self, footprint: ModelFootprint) -> FitEstimate | None:
|
|
148
|
+
"""A pure fit verdict for a footprint, WITHOUT mutating the policy (safe to call from the UI
|
|
149
|
+
thread while a run holds the policy). Default None (no estimate)."""
|
|
150
|
+
return None
|
|
151
|
+
|
|
152
|
+
def fit_estimate(self) -> FitEstimate | None:
|
|
153
|
+
"""The verdict for the footprint last given to ``set_footprint`` (for the pre-flight check).
|
|
154
|
+
Default None."""
|
|
155
|
+
return None
|
|
156
|
+
|
|
157
|
+
def vram_budget_mb(self) -> int | None:
|
|
158
|
+
"""Usable VRAM budget in MB (total or the env override), or None off-GPU. Default None."""
|
|
159
|
+
return None
|
|
160
|
+
|
|
161
|
+
def free_vram_mb(self) -> int | None:
|
|
162
|
+
"""Live free VRAM in MB, or None off-GPU/unmeasurable. Default None."""
|
|
163
|
+
return None
|
|
164
|
+
|
|
165
|
+
def free_ram_mb(self) -> int | None:
|
|
166
|
+
"""Live free system RAM in MB, or None when unmeasurable. Default None."""
|
|
167
|
+
return None
|