Spaces:
Runtime error
Runtime error
File size: 3,161 Bytes
4ff39a4 80cb2a6 4ff39a4 80cb2a6 4ff39a4 80cb2a6 4ff39a4 80cb2a6 4ff39a4 80cb2a6 4ff39a4 80cb2a6 4ff39a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | #!/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())
|