# 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:

1. Declares a unique `name`
2. Defines accepted parameters via `parameters()`
3. Implements `run(ctx)` to process images
4. 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`:

```python
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:

```python
import imagepipeline.modules.my_module  # noqa: F401
```

Use it in a pipeline script:

```python
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:

```python
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:

```python
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`:

```python
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()` — raise `DependencyError` if tools are missing
- `list_output_images(ctx)` — default: all images in `output_dir`

### `SubprocessModule`

Use for CLI tools (ImageMagick, gmic, darktable-cli, rembg):

```python
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_dir` if 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/`:

1. **Unit tests** — parameter validation, matching logic (no external tools)
2. **Integration tests** — run the module against a real image; skip if CLI tool missing:

```python
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)
- [ ] `@register` decorator
- [ ] `parameters()` documents all kwargs
- [ ] `run()` writes outputs preserving stems
- [ ] `check_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

```python
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 |