Skip to content

Commit 1abbedd

Browse files
Copilotmudler
andauthored
feat(diffusers): implement dynamic pipeline loader to remove per-pipeline conditionals (#7365)
* Initial plan Signed-off-by: Ettore Di Giacinto <[email protected]> * Add dynamic loader for diffusers pipelines and refactor backend.py Co-authored-by: mudler <[email protected]> Signed-off-by: Ettore Di Giacinto <[email protected]> * Fix pipeline discovery error handling and test mock issue Co-authored-by: mudler <[email protected]> Signed-off-by: Ettore Di Giacinto <[email protected]> * Address code review feedback: direct imports, better error handling, improved tests Co-authored-by: mudler <[email protected]> Signed-off-by: Ettore Di Giacinto <[email protected]> * Address remaining code review feedback: specific exceptions, registry access, test imports Co-authored-by: mudler <[email protected]> Signed-off-by: Ettore Di Giacinto <[email protected]> * Add defensive fallback for DiffusionPipeline registry access Co-authored-by: mudler <[email protected]> Signed-off-by: Ettore Di Giacinto <[email protected]> * Actually use dynamic pipeline loading for all pipelines in backend Co-authored-by: mudler <[email protected]> Signed-off-by: Ettore Di Giacinto <[email protected]> * Use dynamic loader consistently for all pipelines including AutoPipelineForText2Image Co-authored-by: mudler <[email protected]> Signed-off-by: Ettore Di Giacinto <[email protected]> * Move dynamic loader tests into test.py for CI compatibility Co-authored-by: mudler <[email protected]> Signed-off-by: Ettore Di Giacinto <[email protected]> * Extend dynamic loader to discover any diffusers class type, not just DiffusionPipeline Co-authored-by: mudler <[email protected]> Signed-off-by: Ettore Di Giacinto <[email protected]> * Add AutoPipeline classes to pipeline registry for default model loading Co-authored-by: mudler <[email protected]> Signed-off-by: Ettore Di Giacinto <[email protected]> * fix(python): set pyvenv python home Signed-off-by: Ettore Di Giacinto <[email protected]> * do pyenv update during start Signed-off-by: Ettore Di Giacinto <[email protected]> * Minor changes Signed-off-by: Ettore Di Giacinto <[email protected]> --------- Signed-off-by: Ettore Di Giacinto <[email protected]> Co-authored-by: copilot-swe-agent[bot] <[email protected]> Co-authored-by: mudler <[email protected]> Co-authored-by: Ettore Di Giacinto <[email protected]> Co-authored-by: Ettore Di Giacinto <[email protected]>
1 parent 92ee8c2 commit 1abbedd

File tree

5 files changed

+1142
-156
lines changed

5 files changed

+1142
-156
lines changed

backend/python/common/libbackend.sh

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -237,7 +237,14 @@ function getBuildProfile() {
237237
# Make the venv relocatable:
238238
# - rewrite venv/bin/python{,3} to relative symlinks into $(_portable_dir)
239239
# - normalize entrypoint shebangs to /usr/bin/env python3
240+
# - optionally update pyvenv.cfg to point to the portable Python directory (only at runtime)
241+
# Usage: _makeVenvPortable [--update-pyvenv-cfg]
240242
_makeVenvPortable() {
243+
local update_pyvenv_cfg=false
244+
if [ "${1:-}" = "--update-pyvenv-cfg" ]; then
245+
update_pyvenv_cfg=true
246+
fi
247+
241248
local venv_dir="${EDIR}/venv"
242249
local vbin="${venv_dir}/bin"
243250

@@ -255,7 +262,39 @@ _makeVenvPortable() {
255262
ln -s "${rel_py}" "${vbin}/python3"
256263
ln -s "python3" "${vbin}/python"
257264

258-
# 2) Rewrite shebangs of entry points to use env, so the venv is relocatable
265+
# 2) Update pyvenv.cfg to point to the portable Python directory (only at runtime)
266+
# Use absolute path resolved at runtime so it works when the venv is copied
267+
if [ "$update_pyvenv_cfg" = "true" ]; then
268+
local pyvenv_cfg="${venv_dir}/pyvenv.cfg"
269+
if [ -f "${pyvenv_cfg}" ]; then
270+
local portable_dir="$(_portable_dir)"
271+
# Resolve to absolute path - this ensures it works when the backend is copied
272+
# Only resolve if the directory exists (it should if ensurePortablePython was called)
273+
if [ -d "${portable_dir}" ]; then
274+
portable_dir="$(cd "${portable_dir}" && pwd)"
275+
else
276+
# Fallback to relative path if directory doesn't exist yet
277+
portable_dir="../python"
278+
fi
279+
local sed_i=(sed -i)
280+
# macOS/BSD sed needs a backup suffix; GNU sed doesn't. Make it portable:
281+
if sed --version >/dev/null 2>&1; then
282+
sed_i=(sed -i)
283+
else
284+
sed_i=(sed -i '')
285+
fi
286+
# Update the home field in pyvenv.cfg
287+
# Handle both absolute paths (starting with /) and relative paths
288+
if grep -q "^home = " "${pyvenv_cfg}"; then
289+
"${sed_i[@]}" "s|^home = .*|home = ${portable_dir}|" "${pyvenv_cfg}"
290+
else
291+
# If home field doesn't exist, add it
292+
echo "home = ${portable_dir}" >> "${pyvenv_cfg}"
293+
fi
294+
fi
295+
fi
296+
297+
# 3) Rewrite shebangs of entry points to use env, so the venv is relocatable
259298
# Only touch text files that start with #! and reference the current venv.
260299
local ve_abs="${vbin}/python"
261300
local sed_i=(sed -i)
@@ -316,6 +355,7 @@ function ensureVenv() {
316355
fi
317356
fi
318357
if [ "x${PORTABLE_PYTHON}" == "xtrue" ]; then
358+
# During install, only update symlinks and shebangs, not pyvenv.cfg
319359
_makeVenvPortable
320360
fi
321361
fi
@@ -420,6 +460,11 @@ function installRequirements() {
420460
# - ${BACKEND_NAME}.py
421461
function startBackend() {
422462
ensureVenv
463+
# Update pyvenv.cfg before running to ensure paths are correct for current location
464+
# This is critical when the backend position is dynamic (e.g., copied from container)
465+
if [ "x${PORTABLE_PYTHON}" == "xtrue" ] || [ -x "$(_portable_python)" ]; then
466+
_makeVenvPortable --update-pyvenv-cfg
467+
fi
423468
if [ ! -z "${BACKEND_FILE:-}" ]; then
424469
exec "${EDIR}/venv/bin/python" "${BACKEND_FILE}" "$@"
425470
elif [ -e "${MY_DIR}/server.py" ]; then

backend/python/diffusers/README.md

Lines changed: 133 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,136 @@
1-
# Creating a separate environment for the diffusers project
1+
# LocalAI Diffusers Backend
2+
3+
This backend provides gRPC access to Hugging Face diffusers pipelines with dynamic pipeline loading.
4+
5+
## Creating a separate environment for the diffusers project
26

37
```
48
make diffusers
5-
```
9+
```
10+
11+
## Dynamic Pipeline Loader
12+
13+
The diffusers backend includes a dynamic pipeline loader (`diffusers_dynamic_loader.py`) that automatically discovers and loads diffusers pipelines at runtime. This eliminates the need for per-pipeline conditional statements - new pipelines added to diffusers become available automatically without code changes.
14+
15+
### How It Works
16+
17+
1. **Pipeline Discovery**: On first use, the loader scans the `diffusers` package to find all classes that inherit from `DiffusionPipeline`.
18+
19+
2. **Registry Caching**: Discovery results are cached for the lifetime of the process to avoid repeated scanning.
20+
21+
3. **Task Aliases**: The loader automatically derives task aliases from class names (e.g., "text-to-image", "image-to-image", "inpainting") without hardcoding.
22+
23+
4. **Multiple Resolution Methods**: Pipelines can be resolved by:
24+
- Exact class name (e.g., `StableDiffusionPipeline`)
25+
- Task alias (e.g., `text-to-image`, `img2img`)
26+
- Model ID (uses HuggingFace Hub to infer pipeline type)
27+
28+
### Usage Examples
29+
30+
```python
31+
from diffusers_dynamic_loader import (
32+
load_diffusers_pipeline,
33+
get_available_pipelines,
34+
get_available_tasks,
35+
resolve_pipeline_class,
36+
discover_diffusers_classes,
37+
get_available_classes,
38+
)
39+
40+
# List all available pipelines
41+
pipelines = get_available_pipelines()
42+
print(f"Available pipelines: {pipelines[:10]}...")
43+
44+
# List all task aliases
45+
tasks = get_available_tasks()
46+
print(f"Available tasks: {tasks}")
47+
48+
# Resolve a pipeline class by name
49+
cls = resolve_pipeline_class(class_name="StableDiffusionPipeline")
50+
51+
# Resolve by task alias
52+
cls = resolve_pipeline_class(task="stable-diffusion")
53+
54+
# Load and instantiate a pipeline
55+
pipe = load_diffusers_pipeline(
56+
class_name="StableDiffusionPipeline",
57+
model_id="runwayml/stable-diffusion-v1-5",
58+
torch_dtype=torch.float16
59+
)
60+
61+
# Load from single file
62+
pipe = load_diffusers_pipeline(
63+
class_name="StableDiffusionPipeline",
64+
model_id="/path/to/model.safetensors",
65+
from_single_file=True,
66+
torch_dtype=torch.float16
67+
)
68+
69+
# Discover other diffusers classes (schedulers, models, etc.)
70+
schedulers = discover_diffusers_classes("SchedulerMixin")
71+
print(f"Available schedulers: {list(schedulers.keys())[:5]}...")
72+
73+
# Get list of available scheduler classes
74+
scheduler_list = get_available_classes("SchedulerMixin")
75+
```
76+
77+
### Generic Class Discovery
78+
79+
The dynamic loader can discover not just pipelines but any class type from diffusers:
80+
81+
```python
82+
# Discover all scheduler classes
83+
schedulers = discover_diffusers_classes("SchedulerMixin")
84+
85+
# Discover all model classes
86+
models = discover_diffusers_classes("ModelMixin")
87+
88+
# Get a sorted list of available classes
89+
scheduler_names = get_available_classes("SchedulerMixin")
90+
```
91+
92+
### Special Pipeline Handling
93+
94+
Most pipelines are loaded dynamically through `load_diffusers_pipeline()`. Only pipelines requiring truly custom initialization logic are handled explicitly:
95+
96+
- `FluxTransformer2DModel`: Requires quantization and custom transformer loading (cannot use dynamic loader)
97+
- `WanPipeline` / `WanImageToVideoPipeline`: Uses dynamic loader with special VAE (float32 dtype)
98+
- `SanaPipeline`: Uses dynamic loader with post-load dtype conversion for VAE/text encoder
99+
- `StableVideoDiffusionPipeline`: Uses dynamic loader with CPU offload handling
100+
- `VideoDiffusionPipeline`: Alias for DiffusionPipeline with video flags
101+
102+
All other pipelines (StableDiffusionPipeline, StableDiffusionXLPipeline, FluxPipeline, etc.) are loaded purely through the dynamic loader.
103+
104+
### Error Handling
105+
106+
When a pipeline cannot be resolved, the loader provides helpful error messages listing available pipelines and tasks:
107+
108+
```
109+
ValueError: Unknown pipeline class 'NonExistentPipeline'.
110+
Available pipelines: AnimateDiffPipeline, AnimateDiffVideoToVideoPipeline, ...
111+
```
112+
113+
## Environment Variables
114+
115+
| Variable | Default | Description |
116+
|----------|---------|-------------|
117+
| `COMPEL` | `0` | Enable Compel for prompt weighting |
118+
| `XPU` | `0` | Enable Intel XPU support |
119+
| `CLIPSKIP` | `1` | Enable CLIP skip support |
120+
| `SAFETENSORS` | `1` | Use safetensors format |
121+
| `CHUNK_SIZE` | `8` | Decode chunk size for video |
122+
| `FPS` | `7` | Video frames per second |
123+
| `DISABLE_CPU_OFFLOAD` | `0` | Disable CPU offload |
124+
| `FRAMES` | `64` | Number of video frames |
125+
| `BFL_REPO` | `ChuckMcSneed/FLUX.1-dev` | Flux base repo |
126+
| `PYTHON_GRPC_MAX_WORKERS` | `1` | Max gRPC workers |
127+
128+
## Running Tests
129+
130+
```bash
131+
./test.sh
132+
```
133+
134+
The test suite includes:
135+
- Unit tests for the dynamic loader (`test_dynamic_loader.py`)
136+
- Integration tests for the gRPC backend (`test.py`)

0 commit comments

Comments
 (0)