Spaces:
Runtime error
Runtime error
Add LangSmith tracing, structured logging, checkpoint sweep; fix base repo, fp16 variant, batch seeds, set_adapters, attention slicing
80cb2a6 verified | #!/usr/bin/env python3 | |
| """Run the fine-tuned model locally. Weights come from the Hub, not from disk. | |
| pip install -r requirements-local.txt | |
| python3 -m src.infer "a steampunk owl inside a glass jar, intricate detail" | |
| python3 src/infer.py "a cyberpunk portrait" --lora-scale 0 --out base.png | |
| python3 src/infer.py "a cyberpunk portrait" --lora-scale 1 --out tuned.png | |
| # same seed, both scales -> a direct before/after against vanilla SD 1.5 | |
| Everything is pulled from the Hub on first run and cached under ~/.cache/ | |
| huggingface: SD 1.5 (~4 GB) plus the LoRA adapter (~25 MB). | |
| On an M-series Mac this uses the MPS backend. Expect roughly a minute per | |
| image at 25 steps, and keep the batch at 1 - 8 GB of unified memory does not | |
| leave much headroom. | |
| """ | |
| import argparse | |
| import sys | |
| from pathlib import Path | |
| # Keep `python3 src/infer.py ...` working alongside `python3 -m src.infer ...`: | |
| # both need the repo root importable so `src.pipeline` resolves the same way | |
| # it does from app.py. | |
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) | |
| def main() -> int: | |
| parser = argparse.ArgumentParser( | |
| description=__doc__, | |
| formatter_class=argparse.RawDescriptionHelpFormatter, | |
| ) | |
| parser.add_argument("prompt") | |
| parser.add_argument("--negative", default="", help="what to avoid") | |
| parser.add_argument("--steps", type=int, default=25) | |
| parser.add_argument("--guidance", type=float, default=7.5) | |
| parser.add_argument("--lora-scale", type=float, default=1.0, | |
| help="0 = vanilla SD 1.5, 1 = as trained") | |
| parser.add_argument("--seed", type=int, default=-1, help="-1 for random") | |
| parser.add_argument("--count", type=int, default=1) | |
| parser.add_argument("--checkpoint", default=None, | |
| help="'base' to skip the adapter entirely") | |
| parser.add_argument("--out", default="out.png") | |
| args = parser.parse_args() | |
| # Imported here so --help works without torch installed. | |
| from src.pipeline import LORA_CHECKPOINT, LORA_REPO, ModelRunner, pick_device | |
| args.checkpoint = args.checkpoint or LORA_CHECKPOINT | |
| print(f"device {pick_device()}") | |
| print(f"model {LORA_REPO} @ {args.checkpoint}") | |
| print("loading (first run downloads ~4 GB) ...", flush=True) | |
| runner = ModelRunner(checkpoint=args.checkpoint) | |
| result = runner.generate( | |
| prompt=args.prompt, | |
| negative_prompt=args.negative, | |
| steps=args.steps, | |
| guidance=args.guidance, | |
| lora_scale=args.lora_scale, | |
| seed=None if args.seed < 0 else args.seed, | |
| count=args.count, | |
| ) | |
| out = Path(args.out) | |
| if len(result.images) == 1: | |
| result.images[0].save(out) | |
| print(f"seed {result.seed} -> {out}") | |
| else: | |
| # Image i used seed+i, so each file names the seed that reproduces | |
| # it on its own with --count 1. | |
| for index, image in enumerate(result.images): | |
| path = out.with_name(f"{out.stem}-{index}{out.suffix}") | |
| image.save(path) | |
| print(f"seed {result.seed + index} -> {path}") | |
| return 0 | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) | |