8.0 KiB
Module development guide
This document explains how to add new processing modules to the image pipeline framework.
Overview
Each module is a Python class that:
- Declares a unique
name - Defines accepted parameters via
parameters() - Implements
run(ctx)to process images - Registers itself with
@register
Pipeline scripts reference modules by name and pass parameters as kwargs.
Quick Start
Create a new file in imagepipeline/modules/, e.g. my_module.py:
from imagepipeline.core.context import ModuleContext
from imagepipeline.core.params import Param
from imagepipeline.modules.base import BaseModule
from imagepipeline.modules.registry import register
@register
class MyModule(BaseModule):
name = "my_module"
description = "Short description for documentation"
@classmethod
def parameters(cls) -> dict[str, Param]:
return {
"strength": Param(
"float",
default=1.0,
help="Processing strength from 0.0 to 1.0",
),
}
def run(self, ctx: ModuleContext) -> None:
ctx.output_dir.mkdir(parents=True, exist_ok=True)
for src in ctx.input_paths:
dst = ctx.output_dir / src.name
# ... write dst ...
Import the module in imagepipeline/modules/__init__.py so it registers on load:
import imagepipeline.modules.my_module # noqa: F401
Use it in a pipeline script:
with Pipeline(name="my_run", input_dir=INPUT) as p:
result = p.step("my_module", inputs="input", strength=0.8)
p.run()
ModuleContext
Every run() receives a ModuleContext:
| Field | Description |
|---|---|
input_paths |
Flat list of input images for this step |
matched_groups |
For multi-input steps: list of path groups matched by stem |
output_dir |
Directory where outputs must be written |
params |
Validated parameters |
pipeline_output_root |
Root folder of the current run |
step_id |
Internal step identifier (e.g. my_module_01) |
Single-input modules
Use ctx.input_paths — one path per image in the batch:
for src in ctx.input_paths:
dst = ctx.output_dir / src.name
process(src, dst)
Multi-input modules
When a step receives multiple input references (e.g. background + foreground), the runner matches images by filename stem across all sources.
Use ctx.matched_groups — each entry is a list of paths with the same stem:
for group in ctx.matched_groups:
background, foreground = group
dst = ctx.output_dir / foreground.name
composite(background, foreground, dst)
If a stem is missing in any source, the pipeline fails with a clear error before your module runs.
Parameters
Define parameters with Param:
Param("string", default="value", help="Description")
Param("int", required=True, help="Required integer")
Param("float", default=0.5)
Param("bool", default=False)
Param("path")
Param("list", default=[])
Param("string", choices=("a", "b", "c"))
Supported types: string, int, float, bool, path, list.
Unknown kwargs passed to p.step() raise a validation error. Missing required parameters raise as well.
Base Classes
BaseModule
Use for pure Python processing or custom logic.
Optional class attributes:
supported_input_formats— tuple of extensions (informational)description— short module description
Optional methods:
check_dependencies()— raiseDependencyErrorif tools are missinglist_output_images(ctx)— default: all images inoutput_dir
SubprocessModule
Use for CLI tools (ImageMagick, gmic, darktable-cli, rembg):
from imagepipeline.modules.base import SubprocessModule
from imagepipeline.utils.subprocess import run_command
@register
class MyCliModule(SubprocessModule):
name = "my_cli_module"
command_candidates = ("my-tool", "my-tool-fallback")
@classmethod
def check_dependencies(cls) -> None:
super().check_dependencies() # verifies command_candidates
def run(self, ctx: ModuleContext) -> None:
tool = self.resolve_command()
for src in ctx.input_paths:
dst = ctx.output_dir / src.name
run_command([tool, str(src), str(dst)])
run_command() captures stderr and raises RuntimeError with the command output on failure.
Output Conventions
- Write one output file per input, keeping the same filename (especially the stem) so downstream steps can match batches.
- Only write image files into
ctx.output_dir— the runner discovers outputs by scanning for known image extensions. - Create
ctx.output_dirif your tool does not do so automatically.
Step Folder Naming
The runner assigns output folders automatically: {module_name}_{nn}.
Using the same module twice in one pipeline produces separate folders:
imagemagick_grayscale_01/
imagemagick_grayscale_02/
darktable_style_01/
darktable_style_02/
You do not choose folder names in the module.
Format Warnings
Declare supported_input_formats on your module. When chaining modules with incompatible formats (e.g. JPEG without alpha after rembg), document expected behavior in your module's description and consider converting explicitly in run().
Testing
Add tests under tests/:
- Unit tests — parameter validation, matching logic (no external tools)
- Integration tests — run the module against a real image; skip if CLI tool missing:
import shutil
import pytest
pytestmark = pytest.mark.skipif(
not shutil.which("magick") and not shutil.which("convert"),
reason="ImageMagick not installed",
)
Checklist for New Modules
- Unique
name(snake_case) @registerdecoratorparameters()documents all kwargsrun()writes outputs preserving stemscheck_dependencies()if external tools required- Import in
imagepipeline/modules/__init__.py - Test added (unit and/or integration)
AI Modules
Optional dependencies: pip install -e ".[ai]" (numpy, Pillow, torch).
All AI modules inherit from AIModule (imagepipeline/modules/ai_base.py), which adds:
| Parameter | Default | Description |
|---|---|---|
skip_existing |
True |
Skip when output file already exists |
max_edge |
2048 |
Downscale long edge before inference (0 = full resolution) |
device |
"cpu" |
Inference device (v1: CPU only) |
Use iter_input_images(ctx, processor) in run() — it handles skip, resize, upscale, and per-image ETA logging.
Built-in AI modules
| Module | Backend | Notes |
|---|---|---|
ai_exposure |
Zero-DCE++ (PyTorch) | Weights auto-download on first run |
ai_tone_map |
HDRNet or CLAHE fallback | Set checkpoint to a creotiv-format .pth; empty uses CLAHE |
openrouter_edit |
OpenRouter API | Requires OPENROUTER_API_KEY; default model Flux Klein 4B |
comfy_flux_edit |
ComfyUI HTTP API | Experimental; export workflow to workflows/comfy/flux_klein_edit_api.json |
Example pipeline: pipelines/example_ai.py.
Adding a new AI module
from imagepipeline.modules.ai_base import AIModule
@register
class MyAIModule(AIModule):
name = "my_ai_module"
@classmethod
def parameters(cls) -> dict[str, Param]:
params = dict(super().parameters())
params["strength"] = Param("float", default=1.0)
return params
def run(self, ctx: ModuleContext) -> None:
self.configure_torch(ctx.params["device"])
def process(src: Path, dst: Path, index: int, total: int) -> None:
...
self.iter_input_images(ctx, process)
Put shared inference code under imagepipeline/ai/.
Planned Modules
These slots follow the same interface — no pipeline API changes needed:
| Module | Base class | Tool |
|---|---|---|
imagemagick_resize |
SubprocessModule | ImageMagick |
darktable_style |
SubprocessModule | darktable-cli |
rembg |
SubprocessModule | rembg |
gmic |
SubprocessModule | gmic |
composite |
BaseModule | Pillow |