# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 DeepMachine contributors.
"""Run after installing requirements.txt and downloading this release folder."""
import argparse
from pathlib import Path
from PIL import Image
from pipeline import MOSSAnyPaintPipeline

def main():
    p = argparse.ArgumentParser(description="MOSS AnyPaint arbitrary-mask editing")
    p.add_argument("--source", type=Path, required=True)
    p.add_argument("--mask", type=Path, required=True)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--prompt", required=True)
    p.add_argument("--width", type=int)
    p.add_argument("--height", type=int)
    p.add_argument("--bbox", type=int, nargs=4)
    p.add_argument("--seed", type=int, default=2026)
    p.add_argument("--steps", type=int, default=8)
    p.add_argument("--adapter", type=Path, default=Path(__file__).parent / "moss_anypaint_rank32.safetensors")
    p.add_argument("--model-dir", type=Path)
    p.add_argument("--distillation", type=Path)
    args = p.parse_args()
    source = Image.open(args.source).convert("RGB")
    size = (args.width or source.width, args.height or source.height)
    runtime = MOSSAnyPaintPipeline(args.adapter, model_dir=args.model_dir, distillation_path=args.distillation)
    image = runtime(source, Image.open(args.mask), args.prompt, canvas_size=size,
                    bbox=args.bbox, seed=args.seed, steps=args.steps)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    image.save(args.output)
    print(runtime.last_metadata)

if __name__ == "__main__":
    main()