mamba_diffusion/AGENTS.md

Repository Guidelines

Project Structure & Module Organization

This repository is a flat Python training project, not a packaged library. main.py is the CLI entry point. as_mamba.py holds TrainConfig, dataset loading, the training loop, validation image generation, and SwanLab logging. mamba2_minimal.py contains the standalone Mamba-2 building blocks. train_as_mamba.sh is the multi-GPU launcher. Runtime artifacts land in outputs/ and swanlog/; treat both as generated data, not source.

Build, Test, and Development Commands

Use uv with the committed lockfile.

• uv sync installs the Python 3.12 environment from pyproject.toml and uv.lock.
• uv run python main.py --help lists all training flags and is the fastest CLI sanity check.
• uv run python main.py --device cpu --epochs 1 --steps-per-epoch 1 --batch-size 8 --num-workers 0 --val-every 0 runs a minimal local smoke test.
• bash train_as_mamba.sh launches the default distributed training job with torchrun; adjust --nproc_per_node and shell variables before using it on a new machine.
• uv run python -m compileall main.py as_mamba.py mamba2_minimal.py performs a lightweight syntax check when no test suite is available.

Coding Style & Naming Conventions

Follow the existing Python style: 4-space indentation, type hints on public functions, dataclass-based config, and small helper functions over deeply nested logic. Use snake_case for functions, variables, and CLI flags; use PascalCase for classes such as ASMamba and TrainConfig. Keep new modules top-level unless the project is restructured. There is no configured formatter or linter yet, so match surrounding code closely and keep imports grouped and readable.

Testing Guidelines

There is no dedicated tests/ directory or pytest setup yet. For changes, require at minimum:

• uv run python -m compileall main.py as_mamba.py mamba2_minimal.py
• one short training smoke run with reduced epochs and workers

If you add reusable logic, prefer extracting pure functions so a future pytest suite can cover them easily. Name any new test files test_<module>.py.

Commit & Pull Request Guidelines

Recent history favors short imperative commit subjects, usually Conventional Commit style: feat: ..., fix: ..., and scoped variants like feat(mamba): .... Keep commits focused on one training or infrastructure change. Pull requests should describe the config or behavior change, list the verification commands you ran, and include sample images or metric notes when outputs in outputs/ materially change. Do not commit .venv/, __pycache__/, or large generated artifacts.