File size: 1,811 Bytes
88d0acf | 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 | #!/usr/bin/env python3
"""Generate an image with a local SpaceV1.7 Image / Diffusers checkpoint."""
from __future__ import annotations
import argparse
from pathlib import Path
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", default="SpaceV1.7-Image-model", help="local folder or Hugging Face model ID")
parser.add_argument("--prompt", required=True)
parser.add_argument("--negative-prompt", default=None)
parser.add_argument("--output", default="spacev17-image.png")
parser.add_argument("--steps", type=int, default=20)
parser.add_argument("--guidance-scale", type=float, default=7.5)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--device", choices=("cuda", "cpu"), default="cuda")
args = parser.parse_args()
import torch
from diffusers import DiffusionPipeline
if args.steps < 1:
parser.error("--steps must be positive")
if args.device == "cuda" and not torch.cuda.is_available():
parser.error("CUDA was requested but no CUDA device is available; use --device cpu")
dtype = torch.float16 if args.device == "cuda" else torch.float32
pipe = DiffusionPipeline.from_pretrained(args.model, torch_dtype=dtype)
pipe = pipe.to(args.device)
generator = torch.Generator(device=args.device).manual_seed(args.seed)
result = pipe(
args.prompt,
negative_prompt=args.negative_prompt,
num_inference_steps=args.steps,
guidance_scale=args.guidance_scale,
generator=generator,
)
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
result.images[0].save(output)
print(f"Saved image: {output}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|