Refactor the core diffusion pipeline into clearly separated, importable modules (e.g., data loading, conditioning, sampling loop, post‑processing) with a thin orchestration layer that wires them together.
Why
- A modular pipeline makes it easier to target specific stages for performance work.
- Clear boundaries enable experimentation with alternative implementations (e.g., faster samplers, fused ops) without touching unrelated code.
- Better structure supports upcoming features like API/worker modes and template-based runs.
What to do
- Identify the current diffusion “stages” in the codebase (data input, conditioning, sampling, decoding, post‑processing).
- Extract each stage into its own module with a small, explicit interface (input/output types, config, and lifecycle).
- Create a central orchestrator that wires stages together based on a config object, without embedding heavy logic.
- Update existing entrypoints (CLI, notebooks, API) to go through the orchestrator rather than calling stage internals directly.
- Add minimal docs and/or diagrams to explain the pipeline structure for contributors.
Acceptance criteria
- Each major stage of the diffusion process lives in a dedicated module with a clear interface.
- The orchestrator is the single public entrypoint for running a diffusion pipeline end‑to‑end.
- Existing behavior and outputs remain unchanged for current configs and templates.
- New performance experiments can be implemented by swapping modules behind the orchestrator, not by editing entrypoints.
Refactor the core diffusion pipeline into clearly separated, importable modules (e.g., data loading, conditioning, sampling loop, post‑processing) with a thin orchestration layer that wires them together.
Why
What to do
Acceptance criteria