cli.py

"""Command line tools for multidiff-modular."""

from __future__ import annotations

import argparse
from pathlib import Path

from repository import SUPPORTED_MODEL_FAMILIES, export_hf_repo


def _add_export_parser(subparsers: argparse._SubParsersAction[argparse.ArgumentParser]) -> None:
    parser = subparsers.add_parser("export-hf", help="Export a Hugging Face Modular Diffusers repo.")
    parser.add_argument("output_dir", type=Path, help="Directory to write the Hub repo files into.")
    parser.add_argument(
        "--family",
        default="flux2-klein-4b",
        choices=SUPPORTED_MODEL_FAMILIES,
        help="Flux.2 model family to target.",
    )
    parser.add_argument("--repo-id", default=None, help="Optional Hub repo id used when pushing.")
    parser.add_argument("--push-to-hub", action="store_true", help="Push the exported repo to the Hub.")
    parser.set_defaults(func=_run_export)


def _run_export(args: argparse.Namespace) -> None:
    output_dir = export_hf_repo(
        args.output_dir,
        model_family=args.family,
        repo_id=args.repo_id,
        push_to_hub=args.push_to_hub,
    )
    print(f"Exported Hugging Face repo files to {output_dir}")


def _default_panorama_inspection_output_path(input_path: Path) -> Path:
    suffix = input_path.suffix or ".png"
    return input_path.with_name(f"{input_path.stem}_panorama_seams{suffix}")


def make_panorama_inspection_image(image):
    from PIL import ImageChops

    return ImageChops.offset(image, image.width // 2, image.height // 2)


def _add_inspect_panorama_parser(subparsers: argparse._SubParsersAction[argparse.ArgumentParser]) -> None:
    parser = subparsers.add_parser(
        "inspect-panorama",
        help="Recenter panorama seams by moving the image corners to the middle.",
    )
    parser.add_argument("input", type=Path, help="PNG image to inspect.")
    parser.add_argument(
        "--output",
        "-o",
        type=Path,
        default=None,
        help="Output PNG path. Defaults to '<input>_panorama_seams.png'.",
    )
    parser.set_defaults(func=_run_inspect_panorama)


def _run_inspect_panorama(args: argparse.Namespace) -> None:
    from PIL import Image

    output_path = args.output or _default_panorama_inspection_output_path(args.input)
    with Image.open(args.input) as image:
        inspected = make_panorama_inspection_image(image)
        inspected.save(output_path)
    print(f"Wrote panorama seam inspection image to {output_path}")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="multidiff-modular")
    subparsers = parser.add_subparsers(dest="command", required=True)
    _add_export_parser(subparsers)
    _add_inspect_panorama_parser(subparsers)
    return parser


def main(argv: list[str] | None = None) -> None:
    parser = build_parser()
    args = parser.parse_args(argv)
    args.func(args)


if __name__ == "__main__":
    main()